Add engine.execute() returning a PersistedQueryResult

[madsbk]

Description

This PR adds an execute() API to the streaming cudf-polars engines (SPMD, Dask, and Ray) that returns a PersistedQueryResult. Instead of gathering the output to the client as a single DataFrame, each rank's partition stays GPU-resident in the process that produced it.

PersistedQueryResult.lazy() converts the persisted result back into a LazyFrame, allowing additional queries to be chained without an intermediate host round-trip:

import polars as pl
from cudf_polars.engine.ray import RayEngine

with RayEngine() as engine:
    # Runs on the GPU, the result stays there (no host copy).
    result = engine.execute(
        pl.scan_parquet("/data/dataset/*.parquet")
        .filter(pl.col("amount") > 100)
        .group_by("customer_id")
        .agg(pl.col("amount").sum())
    )

    # Chain more work off the persisted result, still on the GPU.
    df = (
        result.lazy()
        .sort("amount")
        .head(10)
        .collect(engine=engine)
    )

This provides the foundation for supporting Polars' LazyFrame.execute() / QueryResult APIs on GPU engines while keeping intermediate results GPU-resident and rank-local.

The implementation is built entirely on the existing RankAwareSource I/O-plugin infrastructure. Each rank's output partition is exposed through a registered scan source, and workers read only their local partition when the result is collected. Reads are move-on-read, so a PersistedQueryResult can only be collected once (see the design note).

Limitation

We expose this as engine.execute() because there is currently no way to hook into Polars' own pl.LazyFrame.execute(engine=...) call. This is one of the Polars gaps tracked in #22917.

The Approach

The lifecycle of a persisted result spans three phases. Each rank keeps its own output partition GPU-resident in a process-local store keyed by (query_id, rank) (see cudf_polars.engine.rank_local_store), so nothing crosses the process boundary until the caller explicitly collects.

Execute (engine.execute(lf) -> PersistedQueryResult):

• The engine translates lf to IR, creates a query_id, and asks the backend to execute the query on each rank.
• Each rank evaluates its partition and stores the surviving output locally.
• The engine returns a PersistedQueryResult referencing the producing ranks.

Collect (PersistedQueryResult.lazy() -> collect(engine=...)):

• PersistedQueryResult.lazy() exposes the stored partitions as a Polars IO-plugin LazyFrame so downstream query nodes can chain onto it normally.
• Collecting with the producing engine scans each owning rank and removes its partition from the rank-local store.

Release (GC finalizer / PersistedQueryResult.release() / context manager):

• A live result holds a weakref.finalize that calls PersistedBackend.drop_persisted, broadcasting a store drop for the query_id to every rank. It is idempotent and never raises, so a collected, released, or reset result all clean up safely.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: d6cdb034-2685-4243-a6a2-c72d9c2de2b3

📥 Commits

Reviewing files that changed from the base of the PR and between 8db4297 and 4b99b43.

📒 Files selected for processing (2)

• docs/cudf/source/cudf_polars/execute.md
• python/cudf_polars/cudf_polars/engine/rank_local_store.py

🚧 Files skipped from review as they are similar to previous changes (2)

• docs/cudf/source/cudf_polars/execute.md
• python/cudf_polars/cudf_polars/engine/rank_local_store.py

---

📝 Walkthrough

Summary by CodeRabbit

• New Features
  • Added experimental query execution that returns GPU-resident results for SPMD, Dask, and Ray engines.
  • Persisted results can be lazily collected, chained into additional queries, projected, and filtered.
  • Added explicit result release and context-manager support for deterministic GPU memory cleanup.
• Bug Fixes
  • Improved handling and presentation of unsupported operations.
  • Corrected duplicated partition handling across distributed ranks.
• Documentation
  • Added comprehensive execution and persisted-results guidance, including usage limitations and supported engines.
• Tests
  • Added extensive coverage for execution, chaining, cleanup, failures, and distributed behavior.

Walkthrough

This PR adds experimental engine.execute() support to cudf-polars, returning GPU-resident persisted results across SPMD, Dask, and Ray. It adds rank-local storage, lazy persisted sources, replicated-output handling, cleanup, serialization helpers, documentation, and tests.

Changes

Persisted execute() pipeline

Layer / File(s)	Summary
Translation and GPU output handling python/cudf_polars/cudf_polars/dsl/translate.py, .../callback.py, .../engine/core.py	Centralizes unsupported-operation errors, preserves GPU DataFrames from rank execution, and adds replicated-output helpers.
Persisted result storage and loading .../engine/rank_local_store.py, .../engine/persisted_result.py	Adds rank-local partition storage, lazy loading, one-shot consumption, cleanup, and PersistedQueryResult.
Replication metadata .../streaming/rank_aware_source.py, .../streaming/actor_graph/*	Adds duplication signaling to rank-aware sources and channel metadata, with explicit SPMD conversion back to Polars.
SPMD, Dask, and Ray backends .../engine/spmd.py, .../engine/dask.py, .../engine/ray.py	Adds persisted execution, engine methods, rank/worker/actor cleanup, and replicated-output handling.
Serializable executor configuration .../utils/config.py, .../streaming/benchmarks/utils.py	Adds shared removal of process-local executor state before serialization.
Documentation docs/cudf/source/cudf_polars/*, docs/cudf/source/conf.py	Documents the execute workflow, API, navigation, and Sphinx reference exclusions.
Validation and fixtures python/cudf_polars/tests/*	Adds persisted-result, lifecycle, duplicated-output, distributed cleanup, configuration, and environment-specific tests.

Estimated code review effort: 4 (Complex) | ~75 minutes

Possibly related issues

• [Story] Tracking Polars API Gaps for Remote GPU Query Execution #22917: Covers GPU-resident remote query execution and persisted-result integration implemented here.
• [FEA] engine.execute(): support scanning a persisted result more than once #23115: Relates to the single-consumption PersistedQueryResult lifecycle implemented here.

Possibly related PRs

• rapidsai/cudf#23179: Overlaps in engine/core.py execution entry points while adding Quent tracing context passthrough.

Suggested reviewers: wence-, tomaugspurger, nirandaperera

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: adding engine.execute() that returns a PersistedQueryResult.
Description check	✅ Passed	The description is directly related to the change and accurately describes the new persisted-result execute API.
Docstring Coverage	✅ Passed	Docstring coverage is 80.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[TomAugspurger]

Should we have some kind of base .execute() method on StreamingEngine? Same for drop_persisted?

[wence-]

Small things

[wence-]

Note, it's not really an ordinary LazyFrame because you can't collect it on the CPU, no?

[madsbk]

good point, updated to:

`result.lazy()` supports the usual `LazyFrame` API, so you can add further
operations and pass the resulting query directly to `engine.execute()` without
collecting it first.

Only materialization is special. Like any persisted result, it must run on the
engine that produced it, as described above, and cannot be collected directly
on the host. The query executes on the GPU, and its output remains there as a
new persisted result. This makes it useful for building multi-step pipelines
without transferring intermediate results to host memory:

[wence-]

This is bad (the self-join thing). I guess the problem is that Polar's common subplan elimination doesn't see that these are the same thing twice?

[madsbk]

Yes, CSE will sometimes collapse the duplicate scan, so a self-join can happen to work, but not in the general case. Support for repeated reads is tracked in #23115.

[wence-]

Question: can we somehow have this thing tied to an engine, rather than being process-global?

[madsbk]

I don't think so. I initially explored several ways to keep the store on the engine, but persisted partitions are loaded by PersistedSource, which runs during query execution in the worker deep inside the Polars IR.

That callback receives only the serialized scan arguments, the engine uid and query_id. It has no reference to the engine object, which lives in the client process.

[TomAugspurger]

I think all my questions have been addressed.

[madsbk]

/merge