diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 91c4eaf55473..b645a6b8b330 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -902,7 +902,7 @@ def do_evaluate( def read_csv_header( path: Path | str, sep: str ) -> list[str]: # pragma: no cover - with Path(path).open() as f: + with Path(path).open(encoding="utf-8") as f: for line in f: stripped = line.strip() if stripped: @@ -946,7 +946,7 @@ def read_csv_header( for p in paths: skiprows = reader_options["skip_rows"] path = Path(p) - with path.open() as f: + with path.open(encoding="utf-8") as f: while f.readline() == "\n": skiprows += 1 options = ( @@ -1016,11 +1016,14 @@ def read_csv_header( ), stream=stream, ) - parquet_reader_options = ( - plc.io.parquet.ParquetReaderOptions.builder(plc.io.SourceInfo(paths)) - .decimal_width(plc.TypeId.DECIMAL128) - .build() + builder = plc.io.parquet.ParquetReaderOptions.builder( + plc.io.SourceInfo(paths) ) + if filters is not None and parquet_options.use_jit_filter: + builder.use_jit_filter(use_jit_filter=True) + parquet_reader_options = builder.decimal_width( + plc.TypeId.DECIMAL128 + ).build() if with_columns is not None: parquet_reader_options.set_column_names(with_columns) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py index 9f1cd87c59b6..4c5773874aa5 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py @@ -73,6 +73,7 @@ StatsCollector, ) from cudf_polars.streaming.io import FusedScan, SplitScan + from cudf_polars.utils.config import ParquetOptions class Lineariser: @@ -656,6 +657,7 @@ def make_rapidsmpf_read_parquet_node( ch_out: Channel[TableChunk], stats: StatsCollector, partition_info: PartitionInfo, + parquet_options: ParquetOptions, ) -> Any | None: """ Make a RapidsMPF read parquet node. @@ -676,6 +678,8 @@ def make_rapidsmpf_read_parquet_node( The statistics collector. partition_info The partition information. + parquet_options + The Parquet options. Returns ------- @@ -687,11 +691,14 @@ def make_rapidsmpf_read_parquet_node( # Build ParquetReaderOptions try: stream = context.br().stream_pool.get_stream() - parquet_reader_options = ( - plc.io.parquet.ParquetReaderOptions.builder(plc.io.SourceInfo(ir.paths)) - .decimal_width(plc.TypeId.DECIMAL128) - .build() + builder = plc.io.parquet.ParquetReaderOptions.builder( + plc.io.SourceInfo(ir.paths) ) + if ( + ir.predicate is not None and parquet_options.use_jit_filter + ): # pragma: no cover; no test yet + builder.use_jit_filter(use_jit_filter=True) + parquet_reader_options = builder.decimal_width(plc.TypeId.DECIMAL128).build() if ir.with_columns is not None: parquet_reader_options.set_column_names(ir.with_columns) @@ -789,6 +796,7 @@ def _( ch_in, rec.state["stats"], partition_info, + parquet_options, ) # Need metadata node, because the native read_parquet diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 35100c5c38f2..8e3bcdcc3d9f 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -215,6 +215,11 @@ class ParquetOptions: Whether to use the native rapidsmpf node for parquet reading. This option is only used by the streaming executor. Default is False. + use_jit_filter + Whether to use JIT compilation for post-read filtering in Parquet scans. + When enabled, filter predicates are JIT-compiled to CUDA kernels for + improved performance on large datasets with complex filters. + Default is False. """ _env_prefix = "CUDF_POLARS__PARQUET_OPTIONS" @@ -256,6 +261,13 @@ class ParquetOptions: default=False, ) ) + use_jit_filter: bool = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__USE_JIT_FILTER", + _bool_converter, + default=False, + ) + ) def __post_init__(self) -> None: # noqa: D105 if not isinstance(self.chunked, bool): @@ -272,6 +284,8 @@ def __post_init__(self) -> None: # noqa: D105 raise TypeError("max_row_group_samples must be an int") if not isinstance(self.use_rapidsmpf_native, bool): raise TypeError("use_rapidsmpf_native must be a bool") + if not isinstance(self.use_jit_filter, bool): + raise TypeError("use_jit_filter must be a bool") def default_target_partition_size(min_device_size: int | None) -> int: diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index ada8c727dd34..7a9811c7d146 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -212,15 +212,21 @@ def test_parquet_options(executor: str) -> None: ) assert config.parquet_options.chunked is True assert config.parquet_options.n_output_chunks == 1 + assert config.parquet_options.use_jit_filter is False config = ConfigOptions.from_polars_engine( pl.GPUEngine( executor=executor, - parquet_options={"chunked": False, "n_output_chunks": 16}, + parquet_options={ + "chunked": False, + "n_output_chunks": 16, + "use_jit_filter": True, + }, ) ) assert config.parquet_options.chunked is False assert config.parquet_options.n_output_chunks == 16 + assert config.parquet_options.use_jit_filter is True def test_parquet_options_from_none() -> None: @@ -325,6 +331,7 @@ def test_parquet_options_from_env(monkeypatch: pytest.MonkeyPatch) -> None: m.setenv("CUDF_POLARS__PARQUET_OPTIONS__MAX_FOOTER_SAMPLES", "0") m.setenv("CUDF_POLARS__PARQUET_OPTIONS__MAX_ROW_GROUP_SAMPLES", "0") m.setenv("CUDF_POLARS__PARQUET_OPTIONS__USE_RAPIDSMPF_NATIVE", "0") + m.setenv("CUDF_POLARS__PARQUET_OPTIONS__USE_JIT_FILTER", "1") # Test default engine = pl.GPUEngine() @@ -336,6 +343,7 @@ def test_parquet_options_from_env(monkeypatch: pytest.MonkeyPatch) -> None: assert config.parquet_options.max_footer_samples == 0 assert config.parquet_options.max_row_group_samples == 0 assert config.parquet_options.use_rapidsmpf_native is False + assert config.parquet_options.use_jit_filter is True with monkeypatch.context() as m: m.setenv("CUDF_POLARS__PARQUET_OPTIONS__CHUNKED", "foo") @@ -412,6 +420,7 @@ def test_fallback_mode_default(monkeypatch: pytest.MonkeyPatch) -> None: "max_footer_samples", "max_row_group_samples", "use_rapidsmpf_native", + "use_jit_filter", ], ) def test_validate_parquet_options(option: str) -> None: diff --git a/python/cudf_polars/tests/test_parquet_filters.py b/python/cudf_polars/tests/test_parquet_filters.py index d77a627fdb79..1b60f2a32b29 100644 --- a/python/cudf_polars/tests/test_parquet_filters.py +++ b/python/cudf_polars/tests/test_parquet_filters.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations @@ -81,3 +81,15 @@ def test_parquet_filter_boolean_column(engine: pl.GPUEngine, tmp_path): df.write_parquet(tmp_path / "df.parquet") q = pl.scan_parquet(tmp_path / "df.parquet").filter(pl.col("y")) assert_gpu_result_equal(q, engine=engine) + + +def test_jit_filter(pq_file): + q = pq_file.filter((pl.col("a") >= 2) & (pl.col("a") <= 4)).select("a", "c") + assert_gpu_result_equal( + q, + engine=pl.GPUEngine( + executor="in-memory", + raise_on_fail=True, + parquet_options={"use_jit_filter": True}, + ), + )