Skip to content

[FEA] Support grouped LazyFrame.rolling in cuDF-Polars #23590

Description

@rjzamora

LazyFrame.rolling(index_column=..., group_by=...).agg(...) is supported by the GPU engine for single-partition/in-memory execution, but streaming multi-partition execution is not yet supported.

def test_lazyframe_rolling_grouped(engine: GPUEngine) -> None:
    """``LazyFrame.rolling`` (grouped).

    Status
    ------
    - In-memory OK
    - Streaming NOT supported

    """
    lf = pl.LazyFrame(
        {
            "g": ["A", "A", "A", "B", "B"],
            "idx": [1, 2, 3, 1, 2],
            "val": [10, 20, 30, 40, 50],
        }
    ).sort("g", "idx")
    q = lf.rolling(index_column="idx", period="2i", group_by="g", closed="right").agg(
        s=pl.col("val").sum(),
        n=pl.len(),
    )
    assert_gpu_result_equal(q, engine=engine)

Note: The first implementation may be able to shuffle groups together and evaluate locally, but very large groups may need a more distributed strategy later. We can open a new/distinct issue to track the latter case if needed.

Metadata

Metadata

Assignees

No one assigned

    Labels

    cudf-polarsIssues specific to cudf-polarsfeature requestNew feature or request

    Type

    No type

    Projects

    Status
    Todo

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions