NVIDIA / NVIDIA/cudf

[FEA] cudf-polars: GPU accelerated UDFs via `Expr.map_batches`

Open
#23,190 4 comments 0 reactions 0 assignees View on GitHub
feature request
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.