feat: Push down offset/skip to TableScan

datafusion-examples/examples/custom_data_source/custom_datasource.rs

@@ -31,7 +31,7 @@ use datafusion::common::tree_node::TreeNodeRecursion;
 use datafusion::datasource::{TableProvider, TableType, provider_as_source};
 use datafusion::error::Result;
 use datafusion::execution::context::TaskContext;
-use datafusion::logical_expr::LogicalPlanBuilder;
+use datafusion::logical_expr::{LogicalPlanBuilder, TableScanBuilder};
 use datafusion::physical_expr::EquivalenceProperties;
 use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
 use datafusion::physical_plan::memory::MemoryStream;
@@ -92,14 +92,10 @@ async fn search_accounts(
     // create local execution context
     let ctx = SessionContext::new();
 
+    let table_scan =
+        TableScanBuilder::new("accounts", provider_as_source(Arc::new(db))).build()?;
     // create logical plan composed of a single TableScan
-    let logical_plan = LogicalPlanBuilder::scan_with_filters(
-        "accounts",
-        provider_as_source(Arc::new(db)),
-        None,
-        vec![],
-    )?
-    .build()?;
+    let logical_plan = LogicalPlanBuilder::table_scan(table_scan)?.build()?;
 
     let mut dataframe = DataFrame::new(ctx.state(), logical_plan)
         .select_columns(&["id", "bank_account"])?;

datafusion/catalog-listing/src/table.rs

@@ -53,6 +53,7 @@ use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory;
 use datafusion_physical_expr_common::sort_expr::LexOrdering;
 use datafusion_physical_plan::ExecutionPlan;
 use datafusion_physical_plan::empty::EmptyExec;
+use datafusion_physical_plan::limit::GlobalLimitExec;
 use futures::future::BoxFuture;
 use futures::{Stream, StreamExt, TryStreamExt, future, stream};
 use object_store::ObjectStore;
@@ -79,7 +80,7 @@ pub struct ListFilesResult {
 /// * Reading multiple files as a single table
 /// * Hive style partitioning (e.g., directories named `date=2024-06-01`)
 /// * Merges schemas from files with compatible but not identical schemas (see [`ListingTableConfig::file_schema`])
-/// * `limit`, `filter` and `projection` pushdown for formats that support it (e.g.,
+/// * `limit`, `skip`, `filter` and `projection` pushdown for formats that support it (e.g.,
 ///   Parquet)
 /// * Statistics collection and pruning based on file metadata
 /// * Pre-existing sort order (see [`ListingOptions::file_sort_order`])
@@ -578,6 +579,10 @@ impl TableProvider for ListingTable {
             .collect()
     }
 
+    fn supports_skip_pushdown(&self) -> bool {
+        true
+    }
+
     fn get_table_definition(&self) -> Option<&str> {
         self.definition.as_deref()
     }
@@ -620,6 +625,10 @@ impl ListingTable {
         let projection = args.projection().map(|p| p.to_vec());
         let filters = args.filters().map(|f| f.to_vec()).unwrap_or_default();
         let limit = args.limit();
+        let skip = args.skip();
+        // The scan must read enough rows to satisfy `skip + limit`, not
+        // just `limit`, before any rows are skipped below.
+        let inflated_limit = limit.map(|l| l.saturating_add(skip.unwrap_or(0)));
 
         // extract types of partition columns
         let table_partition_cols = self
@@ -647,7 +656,7 @@ impl ListingTable {
         // or before applying non-partition filters.
         let statistic_file_limit =
             if filters.is_empty() && declared_output_partitioning.is_none() {
-                limit
+                inflated_limit
             } else {
                 None
             };
@@ -763,7 +772,7 @@ impl ListingTable {
             .with_constraints(self.constraints.clone())
             .with_statistics(statistics)
             .with_projection_indices(projection)?
-            .with_limit(limit)
+            .with_limit(inflated_limit)
             .with_output_ordering(output_ordering)
             .with_output_partitioning(output_partitioning)
             .with_expr_adapter(self.expr_adapter_factory.clone())
@@ -776,6 +785,11 @@ impl ListingTable {
             .create_physical_plan(state, scan_config)
             .await?;
 
+        let plan: Arc<dyn ExecutionPlan> = match skip {
+            Some(skip) => Arc::new(GlobalLimitExec::new(plan, skip, limit)),
+            None => plan,
+        };
+
         Ok(ScanResult::new(plan))
     }
 

datafusion/catalog/src/default_table_source.rs

@@ -70,6 +70,10 @@ impl TableSource for DefaultTableSource {
         self.table_provider.supports_filters_pushdown(filter)
     }
 
+    fn supports_skip_pushdown(&self) -> bool {
+        self.table_provider.supports_skip_pushdown()
+    }
+
     fn get_logical_plan(&'_ self) -> Option<Cow<'_, datafusion_expr::LogicalPlan>> {
         self.table_provider.get_logical_plan()
     }

datafusion/core/src/datasource/listing/table.rs

@@ -158,7 +158,7 @@ mod tests {
     };
     use arrow::{compute::SortOptions, record_batch::RecordBatch};
     use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit};
-    use datafusion_catalog::TableProvider;
+    use datafusion_catalog::{ScanArgs, TableProvider};
     use datafusion_catalog_listing::{
         ListingOptions, ListingTable, ListingTableConfig, SchemaSource,
     };
@@ -315,6 +315,46 @@ mod tests {
         Ok(())
     }
 
+    #[cfg(feature = "parquet")]
+    #[tokio::test]
+    async fn scan_with_args_skip_skips_correct_rows() -> Result<()> {
+        for target_partition in [1_usize, 8, 16] {
+            let ctx = SessionContext::new_with_config(
+                SessionConfig::new().with_target_partitions(target_partition),
+            );
+
+            let table = load_table(&ctx, "alltypes_plain.parquet").await?;
+
+            // Full scan (no skip) establishes the expected row order.
+            let full_exec = table
+                .scan_with_args(&ctx.state(), ScanArgs::default())
+                .await?
+                .into_inner();
+            let full_batches = collect(full_exec, ctx.task_ctx()).await?;
+            let full =
+                arrow::compute::concat_batches(&full_batches[0].schema(), &full_batches)?;
+            assert_eq!(full.num_rows(), 8);
+            let expected = full.slice(2, 3);
+
+            // `LIMIT 3 OFFSET 2` should return exactly rows [2, 5), in order —
+            // not just the first 3 rows.
+            let skip_exec = table
+                .scan_with_args(
+                    &ctx.state(),
+                    ScanArgs::default().with_limit(Some(3)).with_skip(Some(2)),
+                )
+                .await?
+                .into_inner();
+            let skip_batches = collect(skip_exec, ctx.task_ctx()).await?;
+            let actual =
+                arrow::compute::concat_batches(&skip_batches[0].schema(), &skip_batches)?;
+
+            assert_eq!(actual.num_rows(), 3);
+            assert_eq!(batches_to_string(&[expected]), batches_to_string(&[actual]));
+        }
+        Ok(())
+    }
+
     #[cfg(feature = "parquet")]
     #[tokio::test]
     async fn test_try_create_output_ordering() {

datafusion/core/src/physical_planner.rs

@@ -587,6 +587,7 @@ impl DefaultPhysicalPlanner {
                     projection,
                     filters,
                     fetch,
+                    skip: offset,
                     projected_schema,
                     statistics_requests,
                     ..
@@ -604,6 +605,7 @@ impl DefaultPhysicalPlanner {
                         .with_projection(projection.as_deref())
                         .with_filters(Some(&filters_vec))
                         .with_limit(*fetch)
+                        .with_skip(*offset)
                         .with_statistics_requests(&stats_requests);
                     let res = source.scan_with_args(session_state, opts).await?;
                     Arc::clone(res.plan())

datafusion/expr/src/logical_plan/builder.rs

@@ -45,8 +45,8 @@ use crate::utils::{
 };
 use crate::{
     BinaryExpr, DmlStatement, ExplainOption, Expr, ExprSchemable, Operator,
-    RecursiveQuery, Statement, TableProviderFilterPushDown, TableSource, WriteOp, and,
-    binary_expr, lit,
+    RecursiveQuery, Statement, TableProviderFilterPushDown, TableScan, TableSource,
+    WriteOp, and, binary_expr, lit,
 };
 
 use super::dml::InsertOp;
@@ -417,7 +417,10 @@ impl LogicalPlanBuilder {
         table_source: Arc<dyn TableSource>,
         projection: Option<Vec<usize>>,
     ) -> Result<Self> {
-        Self::scan_with_filters(table_name, table_source, projection, vec![])
+        let table_scan = TableScanBuilder::new(table_name, table_source)
+            .with_projection(projection)
+            .build()?;
+        Self::table_scan(table_scan)
     }
 
     /// Create a [CopyTo] for copying the contents of this builder to the specified file(s)
@@ -485,45 +488,39 @@ impl LogicalPlanBuilder {
     }
 
     /// Convert a table provider into a builder with a TableScan
+    #[deprecated(since = "56.0.0", note = "Use table_scan([`TableScan`]) instead")]
     pub fn scan_with_filters(
         table_name: impl Into<TableReference>,
         table_source: Arc<dyn TableSource>,
         projection: Option<Vec<usize>>,
         filters: Vec<Expr>,
     ) -> Result<Self> {
-        Self::scan_with_filters_inner(table_name, table_source, projection, filters, None)
+        let table_scan = TableScanBuilder::new(table_name, table_source)
+            .with_projection(projection)
+            .with_filters(filters)
+            .build()?;
+        Self::table_scan(table_scan)
     }
 
     /// Convert a table provider into a builder with a TableScan with filter and fetch
+    #[deprecated(since = "56.0.0", note = "Use table_scan([`TableScan`]) instead")]
     pub fn scan_with_filters_fetch(
         table_name: impl Into<TableReference>,
         table_source: Arc<dyn TableSource>,
         projection: Option<Vec<usize>>,
         filters: Vec<Expr>,
         fetch: Option<usize>,
-    ) -> Result<Self> {
-        Self::scan_with_filters_inner(
-            table_name,
-            table_source,
-            projection,
-            filters,
-            fetch,
-        )
-    }
-
-    fn scan_with_filters_inner(
-        table_name: impl Into<TableReference>,
-        table_source: Arc<dyn TableSource>,
-        projection: Option<Vec<usize>>,
-        filters: Vec<Expr>,
-        fetch: Option<usize>,
     ) -> Result<Self> {
         let table_scan = TableScanBuilder::new(table_name, table_source)
             .with_projection(projection)
             .with_filters(filters)
             .with_fetch(fetch)
             .build()?;
+        Self::table_scan(table_scan)
+    }
 
+    /// Convert a [`TableScan`] into a builder
+    pub fn table_scan(table_scan: TableScan) -> Result<Self> {
         // Inline TableScan
         if table_scan.filters.is_empty()
             && let Some(p) = table_scan.source.get_logical_plan()
@@ -2212,11 +2209,7 @@ pub fn table_scan_with_filters(
     projection: Option<Vec<usize>>,
     filters: Vec<Expr>,
 ) -> Result<LogicalPlanBuilder> {
-    let table_source = table_source(table_schema);
-    let name = name
-        .map(|n| n.into())
-        .unwrap_or_else(|| TableReference::bare(UNNAMED_TABLE));
-    LogicalPlanBuilder::scan_with_filters(name, table_source, projection, filters)
+    table_scan_with_filter_and_fetch(name, table_schema, projection, filters, None)
 }
 
 /// Create a LogicalPlanBuilder representing a scan of a table with the provided name and schema,
@@ -2229,19 +2222,40 @@ pub fn table_scan_with_filter_and_fetch(
     filters: Vec<Expr>,
     fetch: Option<usize>,
 ) -> Result<LogicalPlanBuilder> {
-    let table_source = table_source(table_schema);
-    let name = name
-        .map(|n| n.into())
-        .unwrap_or_else(|| TableReference::bare(UNNAMED_TABLE));
-    LogicalPlanBuilder::scan_with_filters_fetch(
+    table_scan_with_filter_and_fetch_and_skip(
         name,
-        table_source,
+        table_schema,
         projection,
         filters,
         fetch,
+        None,
     )
 }
 
+/// Create a LogicalPlanBuilder representing a scan of a table with the provided name and schema,
+/// filters, inlined fetch and offset.
+/// This is mostly used for testing and documentation.
+pub fn table_scan_with_filter_and_fetch_and_skip(
+    name: Option<impl Into<TableReference>>,
+    table_schema: &Schema,
+    projection: Option<Vec<usize>>,
+    filters: Vec<Expr>,
+    fetch: Option<usize>,
+    skip: Option<usize>,
+) -> Result<LogicalPlanBuilder> {
+    let table_source = table_source(table_schema);
+    let name = name
+        .map(|n| n.into())
+        .unwrap_or_else(|| TableReference::bare(UNNAMED_TABLE));
+    let table_scan = TableScanBuilder::new(name, table_source)
+        .with_projection(projection)
+        .with_filters(filters)
+        .with_fetch(fetch)
+        .with_skip(skip)
+        .build()?;
+    LogicalPlanBuilder::table_scan(table_scan)
+}
+
 pub fn table_source(table_schema: &Schema) -> Arc<dyn TableSource> {
     // TODO should we take SchemaRef and avoid cloning?
     let table_schema = Arc::new(table_schema.clone());

datafusion/expr/src/logical_plan/display.rs

@@ -390,6 +390,7 @@ impl<'a, 'b> PgJsonVisitor<'a, 'b> {
                 table_name,
                 filters,
                 fetch,
+                skip,
                 ..
             }) => {
                 let mut object = pg_fields!(
@@ -449,6 +450,10 @@ impl<'a, 'b> PgJsonVisitor<'a, 'b> {
                     object.push("Fetch", serde_json::Value::Number((*f).into()));
                 }
 
+                if let Some(s) = skip {
+                    object.push("Skip", serde_json::Value::Number((*s).into()));
+                }
+
                 object
             }
             LogicalPlan::Projection(Projection { expr, .. }) => {