Skip to content

[PyLimit] adding skip and fetch with datafusion expr support - #1674

Open
hakunamatata-sb wants to merge 1 commit into
apache:mainfrom
hakunamatata-sb:main
Open

[PyLimit] adding skip and fetch with datafusion expr support#1674
hakunamatata-sb wants to merge 1 commit into
apache:mainfrom
hakunamatata-sb:main

Conversation

@hakunamatata-sb

Copy link
Copy Markdown

Which issue does this PR close?

Closes #1673

Rationale for this change

PyLimit (the Python wrapper for a LogicalPlan::Limit node) currently exposes no way to read the actual LIMIT/OFFSET value from Python. skip() and fetch() were removed in #905 ("Upgrade to Datafusion 43") because upstream changed Limit.skip/Limit.fetch from usize/Option<usize> to Option<Box<Expr>> (apache/datafusion#13028 — note the existing TODO comment in limit.rs cites #12836, which is unrelated; #13028 is the actual causal PR), and the old methods no longer compiled against the new field types. Rather than update them, they were deleted, leaving a TODO and no replacement.

This means any consumer walking a logical plan from Python — e.g. a custom SQL compiler/backend built on datafusion-python — cannot determine what LIMIT/OFFSET a query specified, even for a plain literal like LIMIT 10. The value is still present and correctly parsed internally (Display for PyLimit prints it fine via self.limit.skip/self.limit.fetch), it's just inaccessible as structured data. The only current workaround is regex-parsing repr()/str() output.

What changes are included in this PR?

  • crates/core/src/expr/limit.rs: restore skip()/fetch() on PyLimit, now returning Option<PyExpr> instead of the old Option<usize>, matching the new upstream field types (Option<Box<Expr>>). Conversion follows the same pattern already used elsewhere in this crate for Expr/Vec<Expr> fields (e.g. PyProjection::projections(), PyTableScan::py_filters()):

    fn skip(&self) -> PyResult<Option<PyExpr>> {
        Ok(self.limit.skip.as_deref().cloned().map(PyExpr::from))
    }
    
    fn fetch(&self) -> PyResult<Option<PyExpr>> {
        Ok(self.limit.fetch.as_deref().cloned().map(PyExpr::from))
    }

    This does not attempt to simplify/resolve non-literal expressions (e.g. LIMIT $1, computed expressions) — it exposes whatever Expr variant the planner produced, unchanged. Callers should check the variant (e.g. Expr.variant_name() == "Literal") before assuming a value is resolvable, mirroring how DataFusion's own physical planner relies on SimplifyExpressions and errors if it can't fold to a constant.

  • python/tests/test_expr.py: updated test_limit to assert on the new accessors directly instead of only string-matching repr():

    def test_limit(test_ctx):
        df = test_ctx.sql("select c1 from test LIMIT 10")
        plan = df.logical_plan()
    
        plan = plan.to_variant()
        assert isinstance(plan, Limit)
        assert "Skip: None" in str(plan)
        assert plan.skip() is None
        assert plan.fetch().python_value().as_py() == 10
    
        df = test_ctx.sql("select c1 from test LIMIT 10 OFFSET 5")
        plan = df.logical_plan()
    
        plan = plan.to_variant()
        assert isinstance(plan, Limit)
        assert "Skip: Some(Literal(Int64(5), None))" in str(plan)
        assert plan.skip().python_value().as_py() == 5
        assert plan.fetch().python_value().as_py() == 10

    Note: Expr.python_value() returns a PyArrow scalar, whose __eq__ only compares against other PyArrow scalars (pa.scalar(10) == 10 is False) — .as_py() is used to get a plain Python value for comparison.

Testing performed:

  • cargo check -p datafusion-python — compiles clean.
  • maturin develop — builds the real extension (not just type-checked), installs editable into a venv with this repo's pinned dev deps.
  • git submodule update --init testing was required to populate test fixture data before running tests.
  • pytest python/tests/test_expr.py::test_limit -v — passes against the live build.
  • pytest python/tests/test_expr.py -v — full suite, 176 tests, all pass, no regressions.

Are there any user-facing changes?

Yes. This adds two new public methods to datafusion.expr.Limit:

  • skip() -> Optional[Expr]
  • fetch() -> Optional[Expr]

There are no removals or signature changes to existing methods, so this is additive only — no breaking changes to public APIs. (No api change label needed.)

@hakunamatata-sb

Copy link
Copy Markdown
Author

@onestn @timsaucer requesting review and approval for workflows.

@timsaucer

Copy link
Copy Markdown
Member

This is a good start, but if you want to really use these functions then I propose we add a full class wrapper for Limit instead of just the Rust exposed function. That's a bigger change because we currently don't have an expression module but I could see the repo going that way.

If that's too large a lift, then maybe we open an issue to track wrapper coverage for all expressions. I'm a bit torn, because as a user I haven't needed access to these parts of the Expressions. @kosiew do you have thoughts on the idea?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

PyLimit (LogicalPlan Limit node) has no way to read skip/fetch — regressed in #905

2 participants