[Story] Tracking Polars API Gaps for Remote GPU Query Execution
- Dominant language
- C++
- Stars
- 9.8k
- Forks
- 1.1k
- Avg merge
- 3d 6m
- Merged PRs (30d)
- 278
Description
## Description
This issue tracks the upstream Polars changes needed to support GPU-resident query execution on remote nodes via Polars' [`LazyFrame.execute` and `QueryResult` APIs](https://github.com/pola-rs/polars/pull/27525).
To feed a remote, GPU-resident result back into a query, cudf-polars currently relies on Polars internals that are not part of a documented public contract. This issue collects those gaps in one place so we can evaluate the full set of requirements before proposing any upstream changes.
Nothing here is currently blocking. Each item has a working, albeit brittle, workaround today.
### Gaps
- [ ] A supported way to pass distributed context (`rank` and `nranks`) to an IO source (`register_io_source` / `PythonScan`). Currently, users subclass [`RankAwareSource`](https://github.com/rapidsai/cudf/pull/22867), and we reach into the closure Polars wraps the source in and rebind it with a `functools.partial` that injects `rank`/`nranks` for the call. This depends on undocumented closure layout and mutates a user object.
- [ ] A stable contract for `PythonScan` options and pushed predicates. Currently, we parse undocumented internals directly during translation, including the options tuple, `source_type` strings, and predicate encoding (`("polars", node_id)`).
- [ ] A supported way to forward the pushed predicate to an IO source for read-time predicate pushdown. The `register_io_source` contract hands the source a `pl.Expr` (Polars deserializes it before calling), but on the GPU path cudf-polars only has the pushed predicate as the internal `("polars", node_id)` encoding / its translated cudf-polars form, with no public API to materialize a `pl.Expr` from it. So we currently apply the predicate on the GPU *after* the source runs, instead of letting a source (e.g. a remote or columnar reader) filter at read time.
- [ ] Expose `validate_schema` in the `PythonScan` IR. The flag is currently not propagated into the IR, so cudf-polars always validates the source output against the declared schema, diverging from `validate_schema=False` sources and from Polars' coercion path (for example, datetime `us` vs `ns`).
- [ ] Byte-based IO source batch-size hint. `register_io_source` currently passes `batch_size` as a row count, while cudf-polars sizes streaming partitions in bytes (`target_partition_size`). Converting bytes to rows only works for fixed-width schemas; variable-width types (strings, lists, structs) require guessing. A byte-based hint would map directly to our partition target and work for all schemas.
- [ ] Extension support for `pl.LazyFrame.execute(engine=...)`. Currently, `execute()` is effectively hardcoded to:
```python
df = self.collect(...)
return SingleNodeQueryResult(df)
```
As a result, even `engine="gpu"` gathers the result into a host `DataFrame` before wrapping it. We need:
* a pluggable execution path, similar to `.collect(engine="gpu", ...)`, and
* a way for an engine to return its own `QueryResult` implementation, instead of `SingleNodeQueryResult`.
Contributor guide
Assessment
This issue has not been assessed yet.