[FEA] cudf-polars: GPU accelerated UDFs via `Expr.map_batches`
- Dominant language
- C++
- Stars
- 9.8k
- Forks
- 1.1k
- Avg merge
- 3d 6m
- Merged PRs (30d)
- 278
Description
FR: Implement GPU accelerated UDFs via `Expr.map_batches` and accept/return CuPy (zero-copy via `__cuda_array_interface__`).
- UDFs currently raise `ComputeError` with `raise_on_fail=True` or CPU silent fallback. Let's fix this.
- Will likely need a `pl.GPUEngine` parameter(s), e.g. `pl.GPUEngine(udf_passthrough=True)`. `GPUEngine` passes extra kwargs to cuDF, so no Polars API change required.
- PRO : Lazy Eval stays on GPU, end-to-end. GPU operations keep their speedup. The host-device round-trip removal.
- CON 1: Requires a small internal polars change, not API. The Rust IR visitor won't expose the UDF node to Python ([`expr_nodes.rs#L816`](https://github.com/pola-rs/polars/blob/2ce88b4df93f9f7c97e4ef7ba666bc1da674e250/crates/polars-python/src/lazyframe/visitor/expr_nodes.rs#L816)).
- CON 2 (cuDF-wide): Null handling, same as all other cuDF libraries.
- CON 3 (GPU & CPU): UDFs are not pushed through the optimizer. Same as the CPU engine, so no divergence. Possibly fixable with a new IR node, but that's a separate, more complex FR.
- Covers the `map_batches` / `map_elements` rows in tracker [#23151](https://github.com/rapidsai/cudf/issues/23151).
Use cases: iterative solvers/operations, pipelines with custom CuPy/Numba kernels, anything with GPU-winning ops in the middle of a lazy eval.
Contributor guide
Research direction
Start with the Expr.map_batches path and the Rust IR visitor in crates/polars-python/src/lazyframe/visitor/expr_nodes.rs around line 816, then review the proposed pl.GPUEngine(udf_passthrough=True) behavior. Done means CuPy inputs and outputs use __cuda_array_interface__ without host round-trips, lazy GPU evaluation stays on GPU, and UDFs no longer error or silently fall back to CPU.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, rust
- Domain
- data-engineering
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100