From 69b0a9a77226052dd0168447a6716aa44e0fb6d0 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Tue, 17 Feb 2026 07:25:11 -0800 Subject: [PATCH 1/6] Add config option for JIT parquet filtering --- python/cudf_polars/cudf_polars/dsl/ir.py | 7 +++++-- .../cudf_polars/experimental/benchmarks/utils.py | 6 +++--- .../cudf_polars/experimental/rapidsmpf/io.py | 14 ++++++++++++-- python/cudf_polars/cudf_polars/utils/config.py | 14 ++++++++++++++ python/cudf_polars/tests/test_config.py | 11 ++++++++++- python/cudf_polars/tests/test_parquet_filters.py | 12 +++++++++++- 6 files changed, 55 insertions(+), 9 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 441131b4e2fd..8899ee5f8aef 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -815,9 +815,12 @@ def read_csv_header( ), stream=stream, ) - parquet_reader_options = plc.io.parquet.ParquetReaderOptions.builder( + builder = plc.io.parquet.ParquetReaderOptions.builder( plc.io.SourceInfo(paths) - ).build() + ) + if filters is not None and parquet_options.use_jit_filter: + builder.use_jit_filter(use_jit_filter=True) + parquet_reader_options = builder.build() if with_columns is not None: parquet_reader_options.set_column_names(with_columns) if filters is not None: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py index 11a1a5690ba8..56714d9f8eed 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py @@ -36,7 +36,7 @@ duckdb_err = None except ImportError as e: - duckdb = None + duckdb = None # type: ignore[assignment] duckdb_err = e try: @@ -1373,7 +1373,7 @@ def print_duckdb_plan( else: tbl_names = PDSH_TABLE_NAMES - with duckdb.connect() as conn: + with duckdb.connect() as conn: # type: ignore[attr-defined] for name in tbl_names: pattern = (Path(dataset_path) / name).as_posix() + suffix conn.execute( @@ -1409,7 +1409,7 @@ def execute_duckdb_query( tbl_names = PDSDS_TABLE_NAMES else: tbl_names = PDSH_TABLE_NAMES - with duckdb.connect() as conn: + with duckdb.connect() as conn: # type: ignore[attr-defined] for name in tbl_names: pattern = (Path(dataset_path) / name).as_posix() + suffix conn.execute( diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py index 7fc9e88ef968..2109042d2b54 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py @@ -544,6 +544,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. @@ -562,6 +563,8 @@ def make_rapidsmpf_read_parquet_node( The statistics collector. partition_info The partition information. + parquet_options + The Parquet options. Returns ------- @@ -573,9 +576,15 @@ def make_rapidsmpf_read_parquet_node( # Build ParquetReaderOptions try: stream = context.get_stream_from_pool() - parquet_reader_options = plc.io.parquet.ParquetReaderOptions.builder( + builder = plc.io.parquet.ParquetReaderOptions.builder( plc.io.SourceInfo(ir.paths) - ).build() + ) + 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.build() if ir.with_columns is not None: parquet_reader_options.set_column_names(ir.with_columns) @@ -692,6 +701,7 @@ def _( ch_in, rec.state["stats"], partition_info, + parquet_options, ) if native_node is not None and ch_in is not None: diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 7048f80cb822..91ad09310e87 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -286,6 +286,11 @@ class ParquetOptions: Whether to use the native rapidsmpf node for parquet reading. This option is only used when the rapidsmpf runtime is enabled. Default is True. + 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" @@ -327,6 +332,13 @@ class ParquetOptions: default=True, ) ) + 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): @@ -343,6 +355,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_blocksize(cluster: str) -> int: diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 6d2150becee2..415bd012339a 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -213,15 +213,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: @@ -479,6 +485,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() @@ -490,6 +497,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") @@ -589,6 +597,7 @@ def test_cardinality_factor_compat() -> 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 6aac64544536..2bfa5767c369 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-2025, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations @@ -58,3 +58,13 @@ def test_scan_by_hand(expr, selection, pq_file, chunked): assert_gpu_result_equal( q, engine=pl.GPUEngine(raise_on_fail=True, parquet_options={"chunked": chunked}) ) + + +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( + raise_on_fail=True, parquet_options={"use_jit_filter": True} + ), + ) From 9c12fecbc11b87951857ed8e7187fb3d01d961fb Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Tue, 17 Feb 2026 19:27:38 -0800 Subject: [PATCH 2/6] check style --- .../cudf_polars/experimental/benchmarks/utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py index 56714d9f8eed..11a1a5690ba8 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py @@ -36,7 +36,7 @@ duckdb_err = None except ImportError as e: - duckdb = None # type: ignore[assignment] + duckdb = None duckdb_err = e try: @@ -1373,7 +1373,7 @@ def print_duckdb_plan( else: tbl_names = PDSH_TABLE_NAMES - with duckdb.connect() as conn: # type: ignore[attr-defined] + with duckdb.connect() as conn: for name in tbl_names: pattern = (Path(dataset_path) / name).as_posix() + suffix conn.execute( @@ -1409,7 +1409,7 @@ def execute_duckdb_query( tbl_names = PDSDS_TABLE_NAMES else: tbl_names = PDSH_TABLE_NAMES - with duckdb.connect() as conn: # type: ignore[attr-defined] + with duckdb.connect() as conn: for name in tbl_names: pattern = (Path(dataset_path) / name).as_posix() + suffix conn.execute( From 6a3a2ab419d45e5d8a5f70dacf3c675ee53af3c2 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Tue, 16 Jun 2026 17:05:41 -0700 Subject: [PATCH 3/6] Fix test_jit_filter leaving DefaultSingletonEngine alive Use executor='in-memory' to avoid triggering DefaultSingletonEngine, which wasn't being shut down before subsequent spmd-small tests could run on the same pytest-xdist worker. --- python/cudf_polars/tests/test_parquet_filters.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/tests/test_parquet_filters.py b/python/cudf_polars/tests/test_parquet_filters.py index a080a81af02d..1b60f2a32b29 100644 --- a/python/cudf_polars/tests/test_parquet_filters.py +++ b/python/cudf_polars/tests/test_parquet_filters.py @@ -88,6 +88,8 @@ def test_jit_filter(pq_file): assert_gpu_result_equal( q, engine=pl.GPUEngine( - raise_on_fail=True, parquet_options={"use_jit_filter": True} + executor="in-memory", + raise_on_fail=True, + parquet_options={"use_jit_filter": True}, ), ) From 0827b5e692ecfdbb7b2725f33db18d298cf2498f Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Tue, 16 Jun 2026 21:32:56 -0700 Subject: [PATCH 4/6] Fix CSV scan UnicodeDecodeError on ASCII-locale systems Explicitly pass encoding='utf-8' to path.open() in the CSV blank-line skip loop. Without this, the default encoding on systems with LANG=C (e.g. CI containers) is ASCII, which fails when the CSV file contains non-ASCII data (e.g. UTF-8 encoded strings). --- python/cudf_polars/cudf_polars/dsl/ir.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index c688a3094d62..ab41cab4820c 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -728,7 +728,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 = ( From 49cc035aadb6bceeb5de2c0dd1db1367f90b87d5 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 25 Jun 2026 20:10:15 +0000 Subject: [PATCH 5/6] wrong API --- python/cudf_polars/cudf_polars/streaming/actor_graph/io.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 c999695ede2d..318ab9209405 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py @@ -485,7 +485,7 @@ def make_rapidsmpf_read_parquet_node( # Build ParquetReaderOptions try: - stream = context.get_stream_from_pool() + stream = context.br().stream_pool.get_stream() builder = plc.io.parquet.ParquetReaderOptions.builder( plc.io.SourceInfo(ir.paths) ) From ed849b1da66753171ea800fc1aa1dabee18e1441 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Mon, 6 Jul 2026 11:02:39 +0000 Subject: [PATCH 6/6] utf8-encode --- python/cudf_polars/cudf_polars/dsl/ir.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index ab41cab4820c..7d310da247a9 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -684,7 +684,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: