Skip to content

Commit 0bbad2d

Browse files
committed
fix(substrait): Preserve pushed-down table scan offsets and limits
1 parent 6047791 commit 0bbad2d

5 files changed

Lines changed: 116 additions & 24 deletions

File tree

‎Cargo.lock‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎datafusion/substrait/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ datafusion-functions-aggregate = { workspace = true }
5454
serde_json = { workspace = true }
5555
tokio = { workspace = true }
5656
insta = { workspace = true }
57+
tempfile = { workspace = true }
5758

5859
[features]
5960
default = ["physical"]

‎datafusion/substrait/src/logical_plan/producer/rel/fetch_rel.rs‎

Lines changed: 5 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
use crate::logical_plan::producer::SubstraitProducer;
1919
use datafusion::common::DFSchema;
20-
use datafusion::logical_expr::{Limit, LogicalPlan, lit};
20+
use datafusion::logical_expr::Limit;
2121
use std::sync::Arc;
2222
use substrait::proto::rel::RelType;
2323
use substrait::proto::{FetchRel, Rel, fetch_rel};
@@ -29,26 +29,10 @@ pub fn from_limit(
2929
let input = producer.handle_plan(limit.input.as_ref())?;
3030
let empty_schema = Arc::new(DFSchema::empty());
3131

32-
// A provider-level skip pushdown (see `push_down_limit`) may have moved
33-
// some or all of this `Limit`'s skip into the child `TableScan::skip`,
34-
// reducing `limit.skip` accordingly (down to `None`/0 when the scan
35-
// handles the whole skip itself). Substrait's `ReadRel` has no field of
36-
// its own for it, so it must be folded back into this `FetchRel`'s
37-
// offset — otherwise it is silently lost when the plan is serialized.
38-
let scan_offset = match limit.input.as_ref() {
39-
LogicalPlan::TableScan(scan) => scan.skip,
40-
_ => None,
41-
};
42-
let skip_expr = match (limit.skip.as_deref(), scan_offset) {
43-
(Some(limit_skip), Some(scan_skip)) => {
44-
Some(limit_skip.clone() + lit(scan_skip as i64))
45-
}
46-
(Some(skip), None) => Some(skip.clone()),
47-
(None, Some(skip)) => Some(lit(skip as i64)),
48-
(None, None) => None,
49-
};
50-
let offset_mode = skip_expr
51-
.map(|expr| producer.handle_expr(&expr, &empty_schema))
32+
let offset_mode = limit
33+
.skip
34+
.as_ref()
35+
.map(|expr| producer.handle_expr(expr.as_ref(), &empty_schema))
5236
.transpose()?
5337
.map(Box::new)
5438
.map(fetch_rel::OffsetMode::OffsetExpr);

‎datafusion/substrait/src/logical_plan/producer/rel/read_rel.rs‎

Lines changed: 31 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,15 +18,15 @@
1818
use crate::logical_plan::producer::{SubstraitProducer, to_substrait_named_struct};
1919
use datafusion::common::{DFSchema, ToDFSchema, substrait_datafusion_err};
2020
use datafusion::logical_expr::utils::conjunction;
21-
use datafusion::logical_expr::{EmptyRelation, Expr, TableScan, Values};
21+
use datafusion::logical_expr::{EmptyRelation, Expr, TableScan, Values, lit};
2222
use datafusion::scalar::ScalarValue;
2323
use std::sync::Arc;
2424
use substrait::proto::expression::MaskExpression;
2525
use substrait::proto::expression::mask_expression::{StructItem, StructSelect};
2626
use substrait::proto::expression::nested::Struct as NestedStruct;
2727
use substrait::proto::read_rel::{NamedTable, ReadType, VirtualTable};
2828
use substrait::proto::rel::RelType;
29-
use substrait::proto::{ReadRel, Rel};
29+
use substrait::proto::{FetchRel, ReadRel, Rel, fetch_rel};
3030

3131
/// Converts rows of arbitrary expressions into Substrait nested structs.
3232
///
@@ -95,7 +95,7 @@ pub fn from_table_scan(
9595
Some(Box::new(filter_expr))
9696
};
9797

98-
Ok(Box::new(Rel {
98+
let read = Box::new(Rel {
9999
rel_type: Some(RelType::Read(Box::new(ReadRel {
100100
common: None,
101101
base_schema: Some(base_schema),
@@ -108,6 +108,34 @@ pub fn from_table_scan(
108108
advanced_extension: None,
109109
})),
110110
}))),
111+
});
112+
113+
// ReadRel cannot represent pushed-down offsets or limits. Keep them with
114+
// the scan, even when there is no parent Limit or a Projection separates it.
115+
if scan.skip.is_none() && scan.fetch.is_none() {
116+
return Ok(read);
117+
}
118+
let empty_schema = Arc::new(DFSchema::empty());
119+
let offset_mode = scan
120+
.skip
121+
.map(|skip| producer.handle_expr(&lit(skip as i64), &empty_schema))
122+
.transpose()?
123+
.map(Box::new)
124+
.map(fetch_rel::OffsetMode::OffsetExpr);
125+
let count_mode = scan
126+
.fetch
127+
.map(|fetch| producer.handle_expr(&lit(fetch as i64), &empty_schema))
128+
.transpose()?
129+
.map(Box::new)
130+
.map(fetch_rel::CountMode::CountExpr);
131+
Ok(Box::new(Rel {
132+
rel_type: Some(RelType::Fetch(Box::new(FetchRel {
133+
common: None,
134+
input: Some(read),
135+
offset_mode,
136+
count_mode,
137+
advanced_extension: None,
138+
}))),
111139
}))
112140
}
113141

‎datafusion/substrait/tests/cases/roundtrip_logical_plan.rs‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -320,6 +320,84 @@ async fn select_without_limit() -> Result<()> {
320320
roundtrip_fill_na("SELECT * FROM data OFFSET 5+7-2").await
321321
}
322322

323+
#[tokio::test]
324+
async fn roundtrip_table_scan_offset_and_fetch() -> Result<()> {
325+
use datafusion::arrow::array::Int64Array;
326+
use datafusion::dataframe::DataFrame;
327+
use datafusion::datasource::provider_as_source;
328+
use datafusion::logical_expr::TableScanBuilder;
329+
330+
// A single CSV partition preserves input order and supports exact skip pushdown.
331+
let temp = tempfile::tempdir()?;
332+
let path = temp.path().join("data.csv");
333+
std::fs::write(&path, "x\n0\n1\n2\n3\n4\n5\n")?;
334+
let ctx = SessionContext::new();
335+
ctx.register_csv("t", path.to_str().unwrap(), CsvReadOptions::new())
336+
.await?;
337+
let source = provider_as_source(ctx.table_provider("t").await?);
338+
339+
for skip in [Some(2), None, Some(0), Some(8)] {
340+
for fetch in [None, Some(0), Some(2)] {
341+
for project in [false, true] {
342+
for parent in [
343+
None,
344+
Some((0, Some(2))),
345+
Some((1, Some(3))),
346+
Some((1, None)),
347+
] {
348+
let scan = TableScanBuilder::new("t", source.clone())
349+
.with_skip(skip)
350+
.with_fetch(fetch)
351+
.build()?;
352+
let mut builder =
353+
LogicalPlanBuilder::from(LogicalPlan::TableScan(scan));
354+
if project {
355+
builder = builder.project(vec![col("x")])?;
356+
}
357+
if let Some((skip, fetch)) = parent {
358+
builder = builder.limit(skip, fetch)?;
359+
}
360+
let plan = builder.build()?;
361+
let proto = to_substrait_plan(&plan, &ctx.state())?;
362+
let restored = from_substrait_plan(&ctx.state(), &proto).await?;
363+
let mut expected = (0_i64..6)
364+
.skip(skip.unwrap_or(0))
365+
.take(fetch.unwrap_or(usize::MAX))
366+
.collect::<Vec<_>>();
367+
if let Some((skip, fetch)) = parent {
368+
expected = expected
369+
.into_iter()
370+
.skip(skip)
371+
.take(fetch.unwrap_or(usize::MAX))
372+
.collect();
373+
}
374+
for (stage, plan) in [("original", plan), ("roundtrip", restored)] {
375+
let batches = DataFrame::new(ctx.state(), plan).collect().await?;
376+
let actual = batches
377+
.iter()
378+
.flat_map(|batch| {
379+
batch
380+
.column(0)
381+
.as_any()
382+
.downcast_ref::<Int64Array>()
383+
.unwrap()
384+
.values()
385+
.iter()
386+
.copied()
387+
})
388+
.collect::<Vec<_>>();
389+
assert_eq!(
390+
actual, expected,
391+
"{stage}: skip={skip:?}, fetch={fetch:?}, project={project}, parent={parent:?}"
392+
);
393+
}
394+
}
395+
}
396+
}
397+
}
398+
Ok(())
399+
}
400+
323401
#[tokio::test]
324402
async fn select_with_limit_offset() -> Result<()> {
325403
roundtrip("SELECT * FROM data LIMIT 200 OFFSET 10").await?;

0 commit comments

Comments
 (0)