NVIDIA / NVIDIA/cudf

[FEA] Allow engines to return a custom QueryResult from `execute()`

Open
#23,041 0 comments 0 reactions 1 assignee Claimed by @madsbk View on GitHub
cudf-polars feature request
Dominant language
C++
Stars
9.8k
Forks
1.1k
Avg merge
3d 6m
Merged PRs (30d)
278

Description

**Is your feature request related to a problem? Please describe.**

`LazyFrame.execute()` calls `self.collect(engine=engine)` and wraps the result in `SingleNodeQueryResult(df)`. The GPU engine runs the query, but the result is always copied to a host `DataFrame` before being returned. There is no way for the engine to return its own `QueryResult` subclass that keeps data device-resident.

**Describe the solution you'd like**

`execute()` should give the engine the opportunity to return its own `QueryResult`, similar to how `collect(engine=...)` dispatches to the engine today. The GPU engine would return a `GPUQueryResult` or similar that keeps the result in device memory and implements `.lazy()` to return a `LazyFrame` whose root is a `PythonScan`. See #22867 which added support `PythonScan` sources.

This would enable query chaining where the

```python
q1_result = q1.execute(engine=pl.GPUEngine(...))
q2 = result.lazy().filter(pl.col("a") > 0).group_by("b").agg(pl.sum("c"))
q2_result = q2.execute(engine=pl.GPUEngine(...))
```

**Describe alternatives you've considered**

No workarounds AFAICT.

**Additional context**

See [cudf#23039](https://github.com/rapidsai/cudf/issues/23039) for a concrete user workload (iterative graph traversal) that needs this.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.