feat: expose offset in python API (#437)

datafusion/tests/test_dataframe.py

@@ -115,6 +115,17 @@ def test_limit(df):
     assert len(result.column(1)) == 1
 
 
+def test_limit_with_offset(df):
+    # only 3 rows, but limit past the end to ensure that offset is working
+    df = df.limit(5, offset=2)
+
+    # execute and collect the first (and only) batch
+    result = df.collect()[0]
+
+    assert len(result.column(0)) == 1
+    assert len(result.column(1)) == 1
+
+
 def test_with_column(df):
     df = df.with_column("c", column("a") + column("b"))
 

datafusion/tests/test_expr.py

@@ -81,6 +81,14 @@ def test_limit(test_ctx):
 
     plan = plan.to_variant()
     assert isinstance(plan, Limit)
+    assert plan.skip() == 0
+
+    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 plan.skip() == 5
 
 
 def test_aggregate_query(test_ctx):

src/dataframe.rs

@@ -136,8 +136,9 @@ impl PyDataFrame {
         Ok(Self::new(df))
     }
 
-    fn limit(&self, count: usize) -> PyResult<Self> {
-        let df = self.df.as_ref().clone().limit(0, Some(count))?;
+    #[pyo3(signature = (count, offset=0))]
+    fn limit(&self, count: usize, offset: usize) -> PyResult<Self> {
+        let df = self.df.as_ref().clone().limit(offset, Some(count))?;
         Ok(Self::new(df))
     }