diff --git a/crates/core/src/expr/limit.rs b/crates/core/src/expr/limit.rs index 37347acd5..f53737b5e 100644 --- a/crates/core/src/expr/limit.rs +++ b/crates/core/src/expr/limit.rs @@ -22,6 +22,7 @@ use pyo3::IntoPyObjectExt; use pyo3::prelude::*; use crate::common::df_schema::PyDFSchema; +use crate::expr::PyExpr; use crate::expr::logical_node::LogicalNode; use crate::sql::logical::PyLogicalPlan; @@ -64,19 +65,23 @@ impl Display for PyLimit { #[pymethods] impl PyLimit { - // NOTE: Upstream now has expressions for skip and fetch - // TODO: Do we still want to expose these? - // REF: https://github.com/apache/datafusion/pull/12836 - - // /// Retrieves the skip value for this `Limit` - // fn skip(&self) -> usize { - // self.limit.skip - // } + // Retrieves the skip expression for this `Limit`, if any. + // + // `LIMIT`/`OFFSET` were changed upstream to support arbitrary + // expressions (not just constants), see + // https://github.com/apache/datafusion/pull/13028. Callers that expect + // a simple literal (the common case, e.g. `OFFSET 5`) should evaluate + // the returned `PyExpr` via `Expr.python_value()`. + fn skip(&self) -> PyResult> { + Ok(self.limit.skip.as_deref().cloned().map(PyExpr::from)) + } - // /// Retrieves the fetch value for this `Limit` - // fn fetch(&self) -> Option { - // self.limit.fetch - // } + // Retrieves the fetch expression for this `Limit`, if any. + // + // See the note on `skip` above regarding expression-based limits. + fn fetch(&self) -> PyResult> { + Ok(self.limit.fetch.as_deref().cloned().map(PyExpr::from)) + } /// Retrieves the input `LogicalPlan` to this `Limit` node fn input(&self) -> PyResult> { diff --git a/python/tests/test_expr.py b/python/tests/test_expr.py index 9246dd694..ef006bd91 100644 --- a/python/tests/test_expr.py +++ b/python/tests/test_expr.py @@ -115,6 +115,8 @@ def test_limit(test_ctx): 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() @@ -122,6 +124,8 @@ def test_limit(test_ctx): 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 def test_aggregate_query(test_ctx):