Skip to content

Commit

Permalink
feat: expose offset in python API
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud committed Jul 24, 2023
1 parent 5793db3 commit 42a00ea
Show file tree
Hide file tree
Showing 3 changed files with 22 additions and 2 deletions.
11 changes: 11 additions & 0 deletions datafusion/tests/test_dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"))

Expand Down
8 changes: 8 additions & 0 deletions datafusion/tests/test_expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
5 changes: 3 additions & 2 deletions src/dataframe.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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))
}

Expand Down

0 comments on commit 42a00ea

Please sign in to comment.