From 0bbad2dfe0bb0908763ea9b1c0e4479eb26cd19d Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sun, 4 Oct 2026 10:05:57 -0700 Subject: [PATCH] fix(substrait): Preserve pushed-down table scan offsets and limits --- Cargo.lock | 1 + datafusion/substrait/Cargo.toml | 1 + .../logical_plan/producer/rel/fetch_rel.rs | 26 ++----- .../src/logical_plan/producer/rel/read_rel.rs | 34 +++++++- .../tests/cases/roundtrip_logical_plan.rs | 78 +++++++++++++++++++ 5 files changed, 116 insertions(+), 24 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index da096ef3a34ae..64e266f95d7b3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2768,6 +2768,7 @@ dependencies = [ "prost", "serde_json", "substrait", + "tempfile", "tokio", "url", ] diff --git a/datafusion/substrait/Cargo.toml b/datafusion/substrait/Cargo.toml index a0f203cec8db6..d6403da052266 100644 --- a/datafusion/substrait/Cargo.toml +++ b/datafusion/substrait/Cargo.toml @@ -54,6 +54,7 @@ datafusion-functions-aggregate = { workspace = true } serde_json = { workspace = true } tokio = { workspace = true } insta = { workspace = true } +tempfile = { workspace = true } [features] default = ["physical"] diff --git a/datafusion/substrait/src/logical_plan/producer/rel/fetch_rel.rs b/datafusion/substrait/src/logical_plan/producer/rel/fetch_rel.rs index 389cea3be8b25..c5ef57a8ad0ac 100644 --- a/datafusion/substrait/src/logical_plan/producer/rel/fetch_rel.rs +++ b/datafusion/substrait/src/logical_plan/producer/rel/fetch_rel.rs @@ -17,7 +17,7 @@ use crate::logical_plan::producer::SubstraitProducer; use datafusion::common::DFSchema; -use datafusion::logical_expr::{Limit, LogicalPlan, lit}; +use datafusion::logical_expr::Limit; use std::sync::Arc; use substrait::proto::rel::RelType; use substrait::proto::{FetchRel, Rel, fetch_rel}; @@ -29,26 +29,10 @@ pub fn from_limit( let input = producer.handle_plan(limit.input.as_ref())?; let empty_schema = Arc::new(DFSchema::empty()); - // A provider-level skip pushdown (see `push_down_limit`) may have moved - // some or all of this `Limit`'s skip into the child `TableScan::skip`, - // reducing `limit.skip` accordingly (down to `None`/0 when the scan - // handles the whole skip itself). Substrait's `ReadRel` has no field of - // its own for it, so it must be folded back into this `FetchRel`'s - // offset — otherwise it is silently lost when the plan is serialized. - let scan_offset = match limit.input.as_ref() { - LogicalPlan::TableScan(scan) => scan.skip, - _ => None, - }; - let skip_expr = match (limit.skip.as_deref(), scan_offset) { - (Some(limit_skip), Some(scan_skip)) => { - Some(limit_skip.clone() + lit(scan_skip as i64)) - } - (Some(skip), None) => Some(skip.clone()), - (None, Some(skip)) => Some(lit(skip as i64)), - (None, None) => None, - }; - let offset_mode = skip_expr - .map(|expr| producer.handle_expr(&expr, &empty_schema)) + let offset_mode = limit + .skip + .as_ref() + .map(|expr| producer.handle_expr(expr.as_ref(), &empty_schema)) .transpose()? .map(Box::new) .map(fetch_rel::OffsetMode::OffsetExpr); diff --git a/datafusion/substrait/src/logical_plan/producer/rel/read_rel.rs b/datafusion/substrait/src/logical_plan/producer/rel/read_rel.rs index 900273bf8e6d7..37a2ba5b152df 100644 --- a/datafusion/substrait/src/logical_plan/producer/rel/read_rel.rs +++ b/datafusion/substrait/src/logical_plan/producer/rel/read_rel.rs @@ -18,7 +18,7 @@ use crate::logical_plan::producer::{SubstraitProducer, to_substrait_named_struct}; use datafusion::common::{DFSchema, ToDFSchema, substrait_datafusion_err}; use datafusion::logical_expr::utils::conjunction; -use datafusion::logical_expr::{EmptyRelation, Expr, TableScan, Values}; +use datafusion::logical_expr::{EmptyRelation, Expr, TableScan, Values, lit}; use datafusion::scalar::ScalarValue; use std::sync::Arc; use substrait::proto::expression::MaskExpression; @@ -26,7 +26,7 @@ use substrait::proto::expression::mask_expression::{StructItem, StructSelect}; use substrait::proto::expression::nested::Struct as NestedStruct; use substrait::proto::read_rel::{NamedTable, ReadType, VirtualTable}; use substrait::proto::rel::RelType; -use substrait::proto::{ReadRel, Rel}; +use substrait::proto::{FetchRel, ReadRel, Rel, fetch_rel}; /// Converts rows of arbitrary expressions into Substrait nested structs. /// @@ -95,7 +95,7 @@ pub fn from_table_scan( Some(Box::new(filter_expr)) }; - Ok(Box::new(Rel { + let read = Box::new(Rel { rel_type: Some(RelType::Read(Box::new(ReadRel { common: None, base_schema: Some(base_schema), @@ -108,6 +108,34 @@ pub fn from_table_scan( advanced_extension: None, })), }))), + }); + + // ReadRel cannot represent pushed-down offsets or limits. Keep them with + // the scan, even when there is no parent Limit or a Projection separates it. + if scan.skip.is_none() && scan.fetch.is_none() { + return Ok(read); + } + let empty_schema = Arc::new(DFSchema::empty()); + let offset_mode = scan + .skip + .map(|skip| producer.handle_expr(&lit(skip as i64), &empty_schema)) + .transpose()? + .map(Box::new) + .map(fetch_rel::OffsetMode::OffsetExpr); + let count_mode = scan + .fetch + .map(|fetch| producer.handle_expr(&lit(fetch as i64), &empty_schema)) + .transpose()? + .map(Box::new) + .map(fetch_rel::CountMode::CountExpr); + Ok(Box::new(Rel { + rel_type: Some(RelType::Fetch(Box::new(FetchRel { + common: None, + input: Some(read), + offset_mode, + count_mode, + advanced_extension: None, + }))), })) } diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index 6d783af16e461..9bb6a9cccb2c7 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -320,6 +320,84 @@ async fn select_without_limit() -> Result<()> { roundtrip_fill_na("SELECT * FROM data OFFSET 5+7-2").await } +#[tokio::test] +async fn roundtrip_table_scan_offset_and_fetch() -> Result<()> { + use datafusion::arrow::array::Int64Array; + use datafusion::dataframe::DataFrame; + use datafusion::datasource::provider_as_source; + use datafusion::logical_expr::TableScanBuilder; + + // A single CSV partition preserves input order and supports exact skip pushdown. + let temp = tempfile::tempdir()?; + let path = temp.path().join("data.csv"); + std::fs::write(&path, "x\n0\n1\n2\n3\n4\n5\n")?; + let ctx = SessionContext::new(); + ctx.register_csv("t", path.to_str().unwrap(), CsvReadOptions::new()) + .await?; + let source = provider_as_source(ctx.table_provider("t").await?); + + for skip in [Some(2), None, Some(0), Some(8)] { + for fetch in [None, Some(0), Some(2)] { + for project in [false, true] { + for parent in [ + None, + Some((0, Some(2))), + Some((1, Some(3))), + Some((1, None)), + ] { + let scan = TableScanBuilder::new("t", source.clone()) + .with_skip(skip) + .with_fetch(fetch) + .build()?; + let mut builder = + LogicalPlanBuilder::from(LogicalPlan::TableScan(scan)); + if project { + builder = builder.project(vec![col("x")])?; + } + if let Some((skip, fetch)) = parent { + builder = builder.limit(skip, fetch)?; + } + let plan = builder.build()?; + let proto = to_substrait_plan(&plan, &ctx.state())?; + let restored = from_substrait_plan(&ctx.state(), &proto).await?; + let mut expected = (0_i64..6) + .skip(skip.unwrap_or(0)) + .take(fetch.unwrap_or(usize::MAX)) + .collect::>(); + if let Some((skip, fetch)) = parent { + expected = expected + .into_iter() + .skip(skip) + .take(fetch.unwrap_or(usize::MAX)) + .collect(); + } + for (stage, plan) in [("original", plan), ("roundtrip", restored)] { + let batches = DataFrame::new(ctx.state(), plan).collect().await?; + let actual = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + assert_eq!( + actual, expected, + "{stage}: skip={skip:?}, fetch={fetch:?}, project={project}, parent={parent:?}" + ); + } + } + } + } + } + Ok(()) +} + #[tokio::test] async fn select_with_limit_offset() -> Result<()> { roundtrip("SELECT * FROM data LIMIT 200 OFFSET 10").await?;