[Story] JIT Compilation in cudf-polars
- Dominant language
- C++
- Stars
- 9.8k
- Forks
- 1.1k
- Avg merge
- 3d 6m
- Merged PRs (30d)
- 278
Description
## Motivation
This story tracks adding JIT support across cudf_polars. Instead of evaluating expressions through a chain of kernel launches with intermediate results, it can compile the entire expression tree into a single optimized CUDA kernel. cudf_polars already already supports converting expressions to libcudf AST nodes (see [`cudf_polars/dsl/to_ast.py`](https://github.com/rapidsai/cudf/blob/main/python/cudf_polars/cudf_polars/dsl/to_ast.py)), so we should explore where it's beneficial in cudf-polars to JIT-compile these expressions/filters.
## Background
libcudf exposes JIT compilation through three APIs:
1. **Parquet post-read filtering** (`use_jit_filter` in `ParquetReaderOptions`) - Compiles filter predicates applied during Parquet reads
2. **Stream compaction** (`plc.stream_compaction.filter` with AST predicates) - Compiles filter expressions for general filtering operations
3. **Expression evaluation** (`plc.transform.compute_column_jit`) - Compiles expression trees for column transformations
The tradeoff is compilation time versus fused kernel execution. JIT takes tens to hundreds of milliseconds to compile on first execution, but the compiled kernel launches as a single fused operation instead of multiple separate kernels. This reduces kernel launch overhead and avoids materializing intermediate results. The compiled kernels get cached in `.cudf/$VERSION/$COMPUTE_CAPABILITY/` and reused. The cache is thread-safe, so this works with RapidsMPF.
## Tracking PRs
- #20697
- #20790
- #21300
## Parquet Post-Read Filtering
When reading Parquet files with filters, libcudf does two-stage filtering. First, predicate pushdown at the row group level. Second, for rows that pass the row group filter, libcudf evaluates the predicate on individual rows. This second stage is where `use_jit_filter` comes into play.
To control whether a query uses JIT filtering, there's a `use_jit_filter` configuration option is available in `ParquetOptions`. It defaults to `False` and is controlled via `CUDF_POLARS__PARQUET_OPTIONS__USE_JIT_FILTER=True` or `parquet_options={"use_jit_filter": True}` in `pl.GPUEngine(...)`
### Benchmarks
The benchmark ([benchmark_jit_parquet_filter.py](https://github.com/user-attachments/files/25355715/benchmark_jit_parquet_filter.py)) replicates libcudf's `parquet_reader_filter.cpp` benchmark to test performance across different filter complexities. It generates Parquet files with 10 million rows and 32 columns - one filter column with a counting sequence `[0, 1, 2, ..., n-1]` and 31 other columns cycling through bool, float32, float64, and string types. The filter expression is multiple range checks combined with OR: `(col >= min_0 AND col <= max) OR (col >= min_1 AND col <= max) OR ...`. The number of predicates varies (4, 8, 12, 16) to test how JIT scales with complexity. All runs use 50% selectivity and 30% nulls.
The JIT cache at `.cudf///` is cleared before each configuration and the file system page cache is dropped before each iteration with `sync && sudo sysctl vm.drop_caches=3`. Each configuration runs 5 iterations in both orders (JIT first vs non-JIT first) to account for initialization effects.
**JIT-first then non-JIT:**
```bash
python benchmark_jit_parquet_filter.py \
--predicate-intensity 4 8 12 16 \
--num-rows 10000000 \
--selectivity 0.5 \
--iterations 5 \
--drop-caches \
--jit-first \
--output jit_first_clean.json
```
[jit_first_clean.json](https://github.com/user-attachments/files/25355778/jit_first_clean.json)
| Predicates | non-JIT | JIT | Speedup |
|------------|---------|-----|---------|
| 4 | 1.85s | 2.03s | -10.1% |
| 8 | 1.89s | 1.91s | -1.0% |
| 12 | 1.87s | 1.88s | -0.6% |
| 16 | 1.86s | 1.88s | -0.9% |
**non-JIT first then JIT:**
```bash
python benchmark_jit_parquet_filter.py \
--predicate-intensity 4 8 12 16 \
--num-rows 10000000 \
--selectivity 0.5 \
--iterations 5 \
--drop-caches \
--output nojit_first_clean.json
```
[nojit_first_clean.json](https://github.com/user-attachments/files/25355779/nojit_first_clean.json)
| Predicates | Non-JIT | JIT | Speedup |
|------------|---------|-----|---------|
| 4 | 2.06s | 1.87s | +9.2% |
| 8 | 1.89s | 1.87s | +1.0% |
| 12 | 2.03s | 2.02s | +0.8% |
| 16 | 1.97s | 1.97s | +0.2% |
Whichever configuration runs first takes a around +1 sec on the first iteration. After which, JIT and non-JIT runs perform about the same. The first JIT iteration includes compilation time, but that overhead is minimal (under 100ms).
PDS-H queries on Scale Factor 100 data using RapidsMPF show similar results. Neither Q6 nor Q19 (both of which have more complex filter conditions) show any meaningful difference. We probably won't be able to see any benefits of JIT filtering for parquet reads because in most workflows because evaluating the filter wont be a significant slice of total time (SCAN + FILTER).
Example running Q6 at SF100 with JIT filtering
```bash
CUDF_POLARS__PARQUET_OPTIONS__USE_JIT_FILTER=True \
python python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py \
--path data/tables/scale-100.0 \
--executor streaming \s
--iterations 5 \
--runtime rapidsmpf \
--native-parquet 6
```
## Stream Compaction Filtering [TODO]
The Filter IR node in cudf_polars currently materializes a boolean mask and then applies it with a separate filter operation. We could use `plc.stream_compaction.filter(predicate_table, ast, filter_table)` to evaluate the predicate during filtering without materializing the boolean mask.
## Expression Evaluation [TODO]
For evaluating expressions, cudf_polars uses a visitor pattern that calls individual (py)libcudf operations for each expression node, materializing intermediate columns at each `do_evaluate`. We could use `plc.transform.compute_column_jit` to compile the entire expression tree into a single fused kernel.
Contributor guide
Assessment
This issue has not been assessed yet.