fix(substrait): Preserve pushed-down table scan offsets and limits

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"]

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);

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

@@ -18,15 +18,15 @@
 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;
 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,
+        }))),
     }))
 }
 

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::<Vec<_>>();
+                    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::<Int64Array>()
+                                    .unwrap()
+                                    .values()
+                                    .iter()
+                                    .copied()
+                            })
+                            .collect::<Vec<_>>();
+                        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?;