diff --git a/datafusion-examples/examples/custom_data_source/custom_datasource.rs b/datafusion-examples/examples/custom_data_source/custom_datasource.rs index 91fb219f34736..e612ac5505532 100644 --- a/datafusion-examples/examples/custom_data_source/custom_datasource.rs +++ b/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"])?; diff --git a/datafusion/catalog-listing/src/table.rs b/datafusion/catalog-listing/src/table.rs index 3ff86662f6b0b..a72d2c2bc7d40 100644 --- a/datafusion/catalog-listing/src/table.rs +++ b/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 = match skip { + Some(skip) => Arc::new(GlobalLimitExec::new(plan, skip, limit)), + None => plan, + }; + Ok(ScanResult::new(plan)) } diff --git a/datafusion/catalog/src/default_table_source.rs b/datafusion/catalog/src/default_table_source.rs index 3342db54de92f..f3a19d66091a6 100644 --- a/datafusion/catalog/src/default_table_source.rs +++ b/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> { self.table_provider.get_logical_plan() } diff --git a/datafusion/core/src/datasource/listing/table.rs b/datafusion/core/src/datasource/listing/table.rs index 14be8a18cebe2..c10ae42477f4d 100644 --- a/datafusion/core/src/datasource/listing/table.rs +++ b/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() { diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index d394b3e302e5f..f11b39157ce0f 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/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()) diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 75d20b506da41..964c6de396259 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/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, projection: Option>, ) -> Result { - 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, table_source: Arc, projection: Option>, filters: Vec, ) -> Result { - 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, table_source: Arc, projection: Option>, filters: Vec, fetch: Option, - ) -> Result { - Self::scan_with_filters_inner( - table_name, - table_source, - projection, - filters, - fetch, - ) - } - - fn scan_with_filters_inner( - table_name: impl Into, - table_source: Arc, - projection: Option>, - filters: Vec, - fetch: Option, ) -> Result { 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 { // 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>, filters: Vec, ) -> Result { - 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, fetch: Option, ) -> Result { - 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>, + table_schema: &Schema, + projection: Option>, + filters: Vec, + fetch: Option, + skip: Option, +) -> Result { + 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 { // TODO should we take SchemaRef and avoid cloning? let table_schema = Arc::new(table_schema.clone()); diff --git a/datafusion/expr/src/logical_plan/display.rs b/datafusion/expr/src/logical_plan/display.rs index 53c027d0f765e..2dcaefe9a2af5 100644 --- a/datafusion/expr/src/logical_plan/display.rs +++ b/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, .. }) => { diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 1807c12ba1f89..92c1c03c14b9a 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -2109,6 +2109,7 @@ impl LogicalPlan { projection, filters, fetch, + skip, .. }) => { let projected_fields = match projection { @@ -2176,6 +2177,10 @@ impl LogicalPlan { write!(f, ", fetch={n}")?; } + if let Some(n) = skip { + write!(f, ", skip={n}")?; + } + Ok(()) } LogicalPlan::Projection(Projection { expr, .. }) => { @@ -3129,12 +3134,18 @@ pub struct TableScan { pub filters: Vec, /// Optional number of rows to read pub fetch: Option, + /// Optional number of rows to skip + pub skip: Option, /// Statistics the planner would like the provider to answer for this /// scan, typically attached by a custom optimizer rule from the /// surrounding plan (e.g. Min/Max for sort keys). /// /// A [`BTreeSet`], not a `Vec` to keep the resulting plan deterministic. - pub statistics_requests: BTreeSet, + /// + // Boxed to keep this rarely-populated field from growing every + // `TableScan` (and thus `LogicalPlan`) by its own size; + // see `test_size_of_logical_plan`. + pub statistics_requests: Box>, } impl Debug for TableScan { @@ -3146,6 +3157,7 @@ impl Debug for TableScan { .field("projected_schema", &self.projected_schema) .field("filters", &self.filters) .field("fetch", &self.fetch) + .field("skip", &self.skip) .finish_non_exhaustive() } } @@ -3238,7 +3250,9 @@ pub struct TableScanBuilder { projection: Option>, filters: Vec, fetch: Option, - statistics_requests: BTreeSet, + skip: Option, + #[expect(clippy::box_collection)] // additional indirection for smaller size_of() + statistics_requests: Box>, } impl TableScanBuilder { @@ -3253,7 +3267,8 @@ impl TableScanBuilder { projection: None, filters: vec![], fetch: None, - statistics_requests: BTreeSet::new(), + skip: None, + statistics_requests: Box::default(), } } @@ -3275,13 +3290,19 @@ impl TableScanBuilder { self } + /// Set the number of rows to skip. + pub fn with_skip(mut self, skip: Option) -> Self { + self.skip = skip; + self + } + /// Set the statistics requests for the scan. See /// [`TableScan::statistics_requests`]. pub fn with_statistics_requests( mut self, statistics_requests: BTreeSet, ) -> Self { - self.statistics_requests = statistics_requests; + self.statistics_requests = Box::new(statistics_requests); self } @@ -3294,6 +3315,7 @@ impl TableScanBuilder { projection, filters, fetch, + skip, statistics_requests, } = self; @@ -3335,6 +3357,7 @@ impl TableScanBuilder { projected_schema, filters, fetch, + skip, statistics_requests, }) } @@ -3348,6 +3371,7 @@ impl From for TableScanBuilder { projection: scan.projection, filters: scan.filters, fetch: scan.fetch, + skip: scan.skip, statistics_requests: scan.statistics_requests, } } @@ -6414,7 +6438,8 @@ mod tests { projected_schema: Arc::clone(&schema), filters: vec![], fetch: None, - statistics_requests: BTreeSet::new(), + skip: None, + statistics_requests: Box::default(), })); let col = schema.field_names()[0].clone(); @@ -6445,7 +6470,8 @@ mod tests { projected_schema: Arc::clone(&unique_schema), filters: vec![], fetch: None, - statistics_requests: BTreeSet::new(), + skip: None, + statistics_requests: Box::default(), })); let col = schema.field_names()[0].clone(); diff --git a/datafusion/expr/src/logical_plan/tree_node.rs b/datafusion/expr/src/logical_plan/tree_node.rs index 4431f2b719b80..ab796e88d7ed5 100644 --- a/datafusion/expr/src/logical_plan/tree_node.rs +++ b/datafusion/expr/src/logical_plan/tree_node.rs @@ -699,6 +699,7 @@ impl LogicalPlan { projected_schema, filters, fetch, + skip: offset, statistics_requests, }) => filters.map_elements(f)?.update_data(|filters| { LogicalPlan::TableScan(TableScan { @@ -708,6 +709,7 @@ impl LogicalPlan { projected_schema, filters, fetch, + skip: offset, statistics_requests, }) }), diff --git a/datafusion/expr/src/table_source.rs b/datafusion/expr/src/table_source.rs index 65dce8f3c8b0b..f870b5e7ea949 100644 --- a/datafusion/expr/src/table_source.rs +++ b/datafusion/expr/src/table_source.rs @@ -116,6 +116,13 @@ pub trait TableSource: Any + Sync + Send { .collect()) } + /// Tests whether the underlying table provider can guarantee that a scan + /// omits *exactly* the first `skip` rows it would otherwise have + /// produced. + fn supports_skip_pushdown(&self) -> bool { + false + } + /// Get the Logical plan of this table provider, if available. /// /// For example, a view may have a logical plan, but a CSV file does not. diff --git a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs index e2f4aa40b519e..4c6bbdc83a0ef 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -1018,7 +1018,8 @@ mod tests { use arrow::datatypes::{DataType, Field, Schema}; use datafusion_expr::builder::table_source; use datafusion_expr::{ - and, binary_expr, col, cube, grouping_set, out_ref_col, rollup, table_scan, + TableScanBuilder, and, binary_expr, col, cube, grouping_set, out_ref_col, rollup, + table_scan, }; macro_rules! assert_optimized_plan_equal { @@ -2982,14 +2983,11 @@ mod tests { DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))), true, )])); - let subquery = LogicalPlanBuilder::scan_with_filters( - "sq", - subquery_table_source, - None, - vec![], - )? - .unnest_column("arr")? - .build()?; + let subquery_table_scan = + TableScanBuilder::new("sq", subquery_table_source).build()?; + let subquery = LogicalPlanBuilder::table_scan(subquery_table_scan)? + .unnest_column("arr")? + .build()?; let table_scan = test_table_scan()?; let plan = LogicalPlanBuilder::from(table_scan) .filter(exists(Arc::new(subquery)))? @@ -3017,15 +3015,12 @@ mod tests { DataType::List(Arc::new(Field::new_list_field(DataType::UInt32, true))), true, )])); - let subquery = LogicalPlanBuilder::scan_with_filters( - "sq", - subquery_table_source, - None, - vec![], - )? - .unnest_column("a")? - .filter(col("a").eq(out_ref_col(DataType::UInt32, "test.b")))? - .build()?; + let subquery_table_scan = + TableScanBuilder::new("sq", subquery_table_source).build()?; + let subquery = LogicalPlanBuilder::table_scan(subquery_table_scan)? + .unnest_column("a")? + .filter(col("a").eq(out_ref_col(DataType::UInt32, "test.b")))? + .build()?; let plan = LogicalPlanBuilder::from(table_scan) .filter(exists(Arc::new(subquery)))? .project(vec![col("test.b")])? diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs index 1972174a45f8c..64effaea82fa2 100644 --- a/datafusion/optimizer/src/push_down_filter.rs +++ b/datafusion/optimizer/src/push_down_filter.rs @@ -3235,7 +3235,8 @@ mod tests { projection, source: Arc::new(test_provider), fetch: None, - statistics_requests: std::collections::BTreeSet::new(), + skip: None, + statistics_requests: Box::default(), }); Ok(LogicalPlanBuilder::from(table_scan)) diff --git a/datafusion/optimizer/src/push_down_limit.rs b/datafusion/optimizer/src/push_down_limit.rs index 8e07d9dc93406..499a2c993c0e9 100644 --- a/datafusion/optimizer/src/push_down_limit.rs +++ b/datafusion/optimizer/src/push_down_limit.rs @@ -125,6 +125,23 @@ fn rewrite_limit(mut limit: Limit) -> Result> { }; match Arc::unwrap_or_clone(limit.input) { + LogicalPlan::TableScan(mut scan) + if skip > 0 && scan.source.supports_skip_pushdown() => + { + // The source guarantees it will omit exactly the first `skip` + // rows itself, so the remaining `Limit` only needs to trim to + // `fetch` — its skip becomes 0. + scan.skip = Some(scan.skip.unwrap_or(0).saturating_add(skip)); + let new_fetch = if fetch != 0 { + scan.fetch + .map(|existing_fetch| min(existing_fetch.saturating_sub(skip), fetch)) + .or(Some(fetch)) + } else { + Some(0) + }; + scan.fetch = new_fetch; + transformed_limit(0, fetch, LogicalPlan::TableScan(scan)) + } LogicalPlan::TableScan(mut scan) => { let rows_needed = if fetch != 0 { fetch + skip } else { 0 }; let new_fetch = scan @@ -313,10 +330,11 @@ mod test { use crate::test::*; use crate::OptimizerContext; + use arrow::datatypes::{Schema, SchemaRef}; use datafusion_common::DFSchemaRef; use datafusion_expr::{ - Expr, Extension, UserDefinedLogicalNodeCore, col, exists, - logical_plan::builder::LogicalPlanBuilder, + Expr, Extension, TableScanBuilder, TableSource, UserDefinedLogicalNodeCore, col, + exists, logical_plan::builder::LogicalPlanBuilder, }; use datafusion_functions_aggregate::expr_fn::max; @@ -497,6 +515,136 @@ mod test { ) } + /// A `TableSource` that declares it will honor `skip` exactly, so + /// `push_down_limit` is allowed to push `skip` into `TableScan::skip` + /// and elide the outer `Limit`'s skip. + #[derive(Debug)] + struct SkipPushdownTableSource { + schema: SchemaRef, + } + + impl TableSource for SkipPushdownTableSource { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn supports_skip_pushdown(&self) -> bool { + true + } + } + + fn skip_pushdown_table_scan() -> Result { + let schema = Arc::new(Schema::new(test_table_scan_fields())); + let source = Arc::new(SkipPushdownTableSource { schema }); + Ok(LogicalPlan::TableScan( + TableScanBuilder::new("test", source).build()?, + )) + } + + #[test] + fn limit_pushdown_skip_supported() -> Result<()> { + let table_scan = skip_pushdown_table_scan()?; + + let plan = LogicalPlanBuilder::from(table_scan) + .limit(10, Some(1000))? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Limit: skip=0, fetch=1000 + TableScan: test, fetch=1000, skip=10 + " + ) + } + + fn skip_pushdown_table_scan_with( + fetch: Option, + skip: Option, + ) -> Result { + let LogicalPlan::TableScan(scan) = skip_pushdown_table_scan()? else { + unreachable!() + }; + Ok(LogicalPlan::TableScan( + TableScanBuilder::from(scan) + .with_fetch(fetch) + .with_skip(skip) + .build()?, + )) + } + + /// A scan that already skips/fetches: the new skip is added on top of + /// the existing one and the existing fetch shrinks by the new skip. + #[test] + fn limit_pushdown_skip_into_scan_with_skip_and_fetch() -> Result<()> { + let table_scan = skip_pushdown_table_scan_with(Some(15), Some(3))?; + + let plan = LogicalPlanBuilder::from(table_scan) + .limit(10, Some(1000))? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Limit: skip=0, fetch=1000 + TableScan: test, fetch=5, skip=13 + " + ) + } + + /// Skipping past the scan's existing fetch leaves nothing to read + #[test] + fn limit_pushdown_skip_beyond_scan_fetch() -> Result<()> { + let table_scan = skip_pushdown_table_scan_with(Some(5), None)?; + + let plan = LogicalPlanBuilder::from(table_scan) + .limit(10, Some(3))? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Limit: skip=0, fetch=3 + TableScan: test, fetch=0, skip=10 + " + ) + } + + #[test] + fn limit_pushdown_skip_with_fetch_zero() -> Result<()> { + let table_scan = skip_pushdown_table_scan()?; + + let plan = LogicalPlanBuilder::from(table_scan) + .limit(10, Some(0))? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Limit: skip=0, fetch=0 + TableScan: test, fetch=0, skip=10 + " + ) + } + + /// Adding to an existing skip must saturate instead of overflowing + #[test] + fn limit_pushdown_skip_saturates() -> Result<()> { + let table_scan = skip_pushdown_table_scan_with(None, Some(usize::MAX - 1))?; + + let plan = LogicalPlanBuilder::from(table_scan) + .limit(10, Some(3))? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Limit: skip=0, fetch=3 + TableScan: test, fetch=3, skip=18446744073709551615 + " + ) + } + #[test] fn limit_pushdown_multiple_limits() -> Result<()> { let table_scan = test_table_scan()?; diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index 4cb4ddebe4be6..04b9464c9be66 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -103,6 +103,10 @@ message ListingTableScanNode { datafusion_common.ArrowFormat arrow = 16; } repeated SortExprNodeCollection file_sort_order = 13; + // Optional number of rows to read. + optional uint64 fetch = 17; + // Optional number of rows to skip. + optional uint64 skip = 18; } message ViewTableScanNode { @@ -112,6 +116,10 @@ message ViewTableScanNode { datafusion_common.Schema schema = 3; ProjectionColumns projection = 4; string definition = 5; + // Optional number of rows to read. + optional uint64 fetch = 7; + // Optional number of rows to skip. + optional uint64 skip = 8; } // Logical Plan to Scan a CustomTableProvider registered at runtime @@ -122,6 +130,10 @@ message CustomTableScanNode { datafusion_common.Schema schema = 3; repeated LogicalExprNode filters = 4; bytes custom_table_data = 5; + // Optional number of rows to read. + optional uint64 fetch = 7; + // Optional number of rows to skip. + optional uint64 skip = 8; } message ProjectionNode { diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 0f42dea134975..c3b5d1d3740cb 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -5704,6 +5704,12 @@ impl serde::Serialize for CustomTableScanNode { if !self.custom_table_data.is_empty() { len += 1; } + if self.fetch.is_some() { + len += 1; + } + if self.skip.is_some() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.CustomTableScanNode", len)?; if let Some(v) = self.table_name.as_ref() { struct_ser.serialize_field("tableName", v)?; @@ -5722,6 +5728,16 @@ impl serde::Serialize for CustomTableScanNode { #[allow(clippy::needless_borrows_for_generic_args)] struct_ser.serialize_field("customTableData", pbjson::private::base64::encode(&self.custom_table_data).as_str())?; } + if let Some(v) = self.fetch.as_ref() { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("fetch", ToString::to_string(&v).as_str())?; + } + if let Some(v) = self.skip.as_ref() { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("skip", ToString::to_string(&v).as_str())?; + } struct_ser.end() } } @@ -5739,6 +5755,8 @@ impl<'de> serde::Deserialize<'de> for CustomTableScanNode { "filters", "custom_table_data", "customTableData", + "fetch", + "skip", ]; #[allow(clippy::enum_variant_names)] @@ -5748,6 +5766,8 @@ impl<'de> serde::Deserialize<'de> for CustomTableScanNode { Schema, Filters, CustomTableData, + Fetch, + Skip, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -5774,6 +5794,8 @@ impl<'de> serde::Deserialize<'de> for CustomTableScanNode { "schema" => Ok(GeneratedField::Schema), "filters" => Ok(GeneratedField::Filters), "customTableData" | "custom_table_data" => Ok(GeneratedField::CustomTableData), + "fetch" => Ok(GeneratedField::Fetch), + "skip" => Ok(GeneratedField::Skip), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -5798,6 +5820,8 @@ impl<'de> serde::Deserialize<'de> for CustomTableScanNode { let mut schema__ = None; let mut filters__ = None; let mut custom_table_data__ = None; + let mut fetch__ = None; + let mut skip__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::TableName => { @@ -5832,6 +5856,22 @@ impl<'de> serde::Deserialize<'de> for CustomTableScanNode { Some(map_.next_value::<::pbjson::private::BytesDeserialize<_>>()?.0) ; } + GeneratedField::Fetch => { + if fetch__.is_some() { + return Err(serde::de::Error::duplicate_field("fetch")); + } + fetch__ = + map_.next_value::<::std::option::Option<::pbjson::private::NumberDeserialize<_>>>()?.map(|x| x.0) + ; + } + GeneratedField::Skip => { + if skip__.is_some() { + return Err(serde::de::Error::duplicate_field("skip")); + } + skip__ = + map_.next_value::<::std::option::Option<::pbjson::private::NumberDeserialize<_>>>()?.map(|x| x.0) + ; + } } } Ok(CustomTableScanNode { @@ -5840,6 +5880,8 @@ impl<'de> serde::Deserialize<'de> for CustomTableScanNode { schema: schema__, filters: filters__.unwrap_or_default(), custom_table_data: custom_table_data__.unwrap_or_default(), + fetch: fetch__, + skip: skip__, }) } } @@ -13087,6 +13129,12 @@ impl serde::Serialize for ListingTableScanNode { if !self.file_sort_order.is_empty() { len += 1; } + if self.fetch.is_some() { + len += 1; + } + if self.skip.is_some() { + len += 1; + } if self.file_format_type.is_some() { len += 1; } @@ -13115,6 +13163,16 @@ impl serde::Serialize for ListingTableScanNode { if !self.file_sort_order.is_empty() { struct_ser.serialize_field("fileSortOrder", &self.file_sort_order)?; } + if let Some(v) = self.fetch.as_ref() { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("fetch", ToString::to_string(&v).as_str())?; + } + if let Some(v) = self.skip.as_ref() { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("skip", ToString::to_string(&v).as_str())?; + } if let Some(v) = self.file_format_type.as_ref() { match v { listing_table_scan_node::FileFormatType::Csv(v) => { @@ -13156,6 +13214,8 @@ impl<'de> serde::Deserialize<'de> for ListingTableScanNode { "tablePartitionCols", "file_sort_order", "fileSortOrder", + "fetch", + "skip", "csv", "parquet", "avro", @@ -13173,6 +13233,8 @@ impl<'de> serde::Deserialize<'de> for ListingTableScanNode { Filters, TablePartitionCols, FileSortOrder, + Fetch, + Skip, Csv, Parquet, Avro, @@ -13207,6 +13269,8 @@ impl<'de> serde::Deserialize<'de> for ListingTableScanNode { "filters" => Ok(GeneratedField::Filters), "tablePartitionCols" | "table_partition_cols" => Ok(GeneratedField::TablePartitionCols), "fileSortOrder" | "file_sort_order" => Ok(GeneratedField::FileSortOrder), + "fetch" => Ok(GeneratedField::Fetch), + "skip" => Ok(GeneratedField::Skip), "csv" => Ok(GeneratedField::Csv), "parquet" => Ok(GeneratedField::Parquet), "avro" => Ok(GeneratedField::Avro), @@ -13239,6 +13303,8 @@ impl<'de> serde::Deserialize<'de> for ListingTableScanNode { let mut filters__ = None; let mut table_partition_cols__ = None; let mut file_sort_order__ = None; + let mut fetch__ = None; + let mut skip__ = None; let mut file_format_type__ = None; while let Some(k) = map_.next_key()? { match k { @@ -13290,6 +13356,22 @@ impl<'de> serde::Deserialize<'de> for ListingTableScanNode { } file_sort_order__ = Some(map_.next_value()?); } + GeneratedField::Fetch => { + if fetch__.is_some() { + return Err(serde::de::Error::duplicate_field("fetch")); + } + fetch__ = + map_.next_value::<::std::option::Option<::pbjson::private::NumberDeserialize<_>>>()?.map(|x| x.0) + ; + } + GeneratedField::Skip => { + if skip__.is_some() { + return Err(serde::de::Error::duplicate_field("skip")); + } + skip__ = + map_.next_value::<::std::option::Option<::pbjson::private::NumberDeserialize<_>>>()?.map(|x| x.0) + ; + } GeneratedField::Csv => { if file_format_type__.is_some() { return Err(serde::de::Error::duplicate_field("csv")); @@ -13336,6 +13418,8 @@ impl<'de> serde::Deserialize<'de> for ListingTableScanNode { filters: filters__.unwrap_or_default(), table_partition_cols: table_partition_cols__.unwrap_or_default(), file_sort_order: file_sort_order__.unwrap_or_default(), + fetch: fetch__, + skip: skip__, file_format_type: file_format_type__, }) } @@ -29278,6 +29362,12 @@ impl serde::Serialize for ViewTableScanNode { if !self.definition.is_empty() { len += 1; } + if self.fetch.is_some() { + len += 1; + } + if self.skip.is_some() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.ViewTableScanNode", len)?; if let Some(v) = self.table_name.as_ref() { struct_ser.serialize_field("tableName", v)?; @@ -29294,6 +29384,16 @@ impl serde::Serialize for ViewTableScanNode { if !self.definition.is_empty() { struct_ser.serialize_field("definition", &self.definition)?; } + if let Some(v) = self.fetch.as_ref() { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("fetch", ToString::to_string(&v).as_str())?; + } + if let Some(v) = self.skip.as_ref() { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("skip", ToString::to_string(&v).as_str())?; + } struct_ser.end() } } @@ -29310,6 +29410,8 @@ impl<'de> serde::Deserialize<'de> for ViewTableScanNode { "schema", "projection", "definition", + "fetch", + "skip", ]; #[allow(clippy::enum_variant_names)] @@ -29319,6 +29421,8 @@ impl<'de> serde::Deserialize<'de> for ViewTableScanNode { Schema, Projection, Definition, + Fetch, + Skip, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -29345,6 +29449,8 @@ impl<'de> serde::Deserialize<'de> for ViewTableScanNode { "schema" => Ok(GeneratedField::Schema), "projection" => Ok(GeneratedField::Projection), "definition" => Ok(GeneratedField::Definition), + "fetch" => Ok(GeneratedField::Fetch), + "skip" => Ok(GeneratedField::Skip), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -29369,6 +29475,8 @@ impl<'de> serde::Deserialize<'de> for ViewTableScanNode { let mut schema__ = None; let mut projection__ = None; let mut definition__ = None; + let mut fetch__ = None; + let mut skip__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::TableName => { @@ -29401,6 +29509,22 @@ impl<'de> serde::Deserialize<'de> for ViewTableScanNode { } definition__ = Some(map_.next_value()?); } + GeneratedField::Fetch => { + if fetch__.is_some() { + return Err(serde::de::Error::duplicate_field("fetch")); + } + fetch__ = + map_.next_value::<::std::option::Option<::pbjson::private::NumberDeserialize<_>>>()?.map(|x| x.0) + ; + } + GeneratedField::Skip => { + if skip__.is_some() { + return Err(serde::de::Error::duplicate_field("skip")); + } + skip__ = + map_.next_value::<::std::option::Option<::pbjson::private::NumberDeserialize<_>>>()?.map(|x| x.0) + ; + } } } Ok(ViewTableScanNode { @@ -29409,6 +29533,8 @@ impl<'de> serde::Deserialize<'de> for ViewTableScanNode { schema: schema__, projection: projection__, definition: definition__.unwrap_or_default(), + fetch: fetch__, + skip: skip__, }) } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index 31e060707a399..ec01eb1e67998 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -123,6 +123,12 @@ pub struct ListingTableScanNode { pub table_partition_cols: ::prost::alloc::vec::Vec, #[prost(message, repeated, tag = "13")] pub file_sort_order: ::prost::alloc::vec::Vec, + /// Optional number of rows to read. + #[prost(uint64, optional, tag = "17")] + pub fetch: ::core::option::Option, + /// Optional number of rows to skip. + #[prost(uint64, optional, tag = "18")] + pub skip: ::core::option::Option, #[prost( oneof = "listing_table_scan_node::FileFormatType", tags = "10, 11, 12, 15, 16" @@ -159,6 +165,12 @@ pub struct ViewTableScanNode { pub projection: ::core::option::Option, #[prost(string, tag = "5")] pub definition: ::prost::alloc::string::String, + /// Optional number of rows to read. + #[prost(uint64, optional, tag = "7")] + pub fetch: ::core::option::Option, + /// Optional number of rows to skip. + #[prost(uint64, optional, tag = "8")] + pub skip: ::core::option::Option, } /// Logical Plan to Scan a CustomTableProvider registered at runtime #[derive(Clone, PartialEq, ::prost::Message)] @@ -173,6 +185,12 @@ pub struct CustomTableScanNode { pub filters: ::prost::alloc::vec::Vec, #[prost(bytes = "vec", tag = "5")] pub custom_table_data: ::prost::alloc::vec::Vec, + /// Optional number of rows to read. + #[prost(uint64, optional, tag = "7")] + pub fetch: ::core::option::Option, + /// Optional number of rows to skip. + #[prost(uint64, optional, tag = "8")] + pub skip: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct ProjectionNode { diff --git a/datafusion/proto/src/lib.rs b/datafusion/proto/src/lib.rs index 71feae506dc6f..7a965fcebeb4d 100644 --- a/datafusion/proto/src/lib.rs +++ b/datafusion/proto/src/lib.rs @@ -65,7 +65,7 @@ //! # use datafusion_expr::{col, lit, Expr}; //! # use datafusion_proto::bytes::Serializeable; //! # fn main() -> Result<()>{ -//! // Create a new `Expr` a < 32 +//! // Create a new `Expr` a < 5 //! let expr = col("a").lt(lit(5i32)); //! //! // Convert it to bytes (for sending over the network, etc.) diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index 171043c233f3b..3de9237f062ab 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -696,13 +696,16 @@ impl AsLogicalPlan for LogicalPlanNode { projection = Some(column_indices); } - LogicalPlanBuilder::scan_with_filters( + let table_scan = TableScanBuilder::new( table_name, provider_as_source(Arc::new(provider)), - projection, - filters, - )? - .build() + ) + .with_projection(projection) + .with_filters(filters) + .with_fetch(scan.fetch.map(|f| f as usize)) + .with_skip(scan.skip.map(|o| o as usize)) + .build()?; + LogicalPlanBuilder::table_scan(table_scan)?.build() } LogicalPlanType::CustomScan(scan) => { let schema: Schema = convert_required!(scan.schema)?; @@ -730,13 +733,14 @@ impl AsLogicalPlan for LogicalPlanNode { ctx, )?; - LogicalPlanBuilder::scan_with_filters( - table_name, - provider_as_source(provider), - projection, - filters, - )? - .build() + let table_scan = + TableScanBuilder::new(table_name, provider_as_source(provider)) + .with_projection(projection) + .with_filters(filters) + .with_fetch(scan.fetch.map(|f| f as usize)) + .with_skip(scan.skip.map(|o| o as usize)) + .build()?; + LogicalPlanBuilder::table_scan(table_scan)?.build() } LogicalPlanType::Sort(sort) => { let input: LogicalPlan = @@ -1210,12 +1214,15 @@ impl AsLogicalPlan for LogicalPlanNode { let table_name = from_table_reference(scan.table_name.as_ref(), "ViewScan")?; - LogicalPlanBuilder::scan( + let table_scan = TableScanBuilder::new( table_name, provider_as_source(Arc::new(provider)), - projection, - )? - .build() + ) + .with_projection(projection) + .with_fetch(scan.fetch.map(|f| f as usize)) + .with_skip(scan.skip.map(|o| o as usize)) + .build()?; + LogicalPlanBuilder::table_scan(table_scan)?.build() } LogicalPlanType::Prepare(prepare) => { let input: LogicalPlan = @@ -1352,13 +1359,12 @@ impl AsLogicalPlan for LogicalPlanNode { let provider = Arc::new(EmptyTable::new(Arc::clone(&schema))); - LogicalPlanBuilder::scan_with_filters( - table_name, - provider_as_source(provider), - projection, - filters, - )? - .build() + let table_scan = + TableScanBuilder::new(table_name, provider_as_source(provider)) + .with_projection(projection) + .with_filters(filters) + .build()?; + LogicalPlanBuilder::table_scan(table_scan)?.build() } LogicalPlanType::Dml(dml_node) => { let table_name = @@ -1408,6 +1414,8 @@ impl AsLogicalPlan for LogicalPlanNode { source, filters, projection, + fetch, + skip, .. }) => { let provider = source_as_provider(source)?; @@ -1540,6 +1548,8 @@ impl AsLogicalPlan for LogicalPlanNode { projection, filters, file_sort_order: exprs_vec, + fetch: fetch.map(|f| f as u64), + skip: skip.map(|o| o as u64), }, )), }) @@ -1563,6 +1573,8 @@ impl AsLogicalPlan for LogicalPlanNode { .definition() .map(|s| s.to_string()) .unwrap_or_default(), + fetch: fetch.map(|f| f as u64), + skip: skip.map(|o| o as u64), }, ))), }) @@ -1610,6 +1622,8 @@ impl AsLogicalPlan for LogicalPlanNode { schema: Some(schema), filters, custom_table_data: bytes, + fetch: fetch.map(|f| f as u64), + skip: skip.map(|o| o as u64), }); let node = LogicalPlanNode { logical_plan_type: Some(scan), diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs index eb6207cf42c50..a479512c3fe96 100644 --- a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs @@ -617,6 +617,43 @@ async fn roundtrip_logical_plan_sort() -> Result<()> { Ok(()) } +#[tokio::test] +async fn roundtrip_logical_plan_limit_skip() -> Result<()> { + let ctx = SessionContext::new(); + + let schema = Schema::new(vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Decimal128(15, 2), true), + ]); + + ctx.register_csv( + "t1", + "tests/testdata/test.csv", + CsvReadOptions::default().schema(&schema), + ) + .await?; + + let query = "SELECT a, b FROM t1 LIMIT 5 OFFSET 3"; + let plan = ctx.sql(query).await?.into_optimized_plan()?; + + // Sanity check that the `skip` was actually pushed into the `TableScan` + // (ListingTable opts into `supports_skip_pushdown`), so this test + // exercises the new `ListingTableScanNode.skip` wire field rather than + // trivially passing because nothing needed to round-trip. + let plan_str = plan.to_string(); + assert_eq!( + plan_str, + "Limit: skip=0, fetch=5\n TableScan: t1 projection=[a, b], fetch=5, skip=3", + "expected 'fetch=5' and 'skip=3' to be pushed into the scan, got: {plan_str}" + ); + + let bytes = logical_plan_to_bytes(&plan)?; + let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?; + assert_eq!(plan_str, logical_round_trip.to_string()); + + Ok(()) +} + #[tokio::test] async fn roundtrip_logical_plan_dml() -> Result<()> { let ctx = SessionContext::new(); diff --git a/datafusion/session/src/table.rs b/datafusion/session/src/table.rs index 999ec16b8932f..19f9f23340549 100644 --- a/datafusion/session/src/table.rs +++ b/datafusion/session/src/table.rs @@ -46,7 +46,8 @@ use datafusion_physical_plan::ExecutionPlan; /// /// 1. [`Self::schema`]: The schema (columns and their types) of the table /// 2. [`Self::supports_filters_pushdown`]: Should filters be pushed into this scan -/// 2. [`Self::scan`]: An [`ExecutionPlan`] that can read data +/// 3. [`Self::supports_skip_pushdown`]: Should `skip` (offset) be pushed into this scan +/// 4. [`Self::scan`]: An [`ExecutionPlan`] that can read data /// /// [`RecordBatch`]: https://docs.rs/arrow/latest/arrow/record_batch/struct.RecordBatch.html /// [`CatalogProvider`]: super::CatalogProvider @@ -184,6 +185,11 @@ pub trait TableProvider: Any + Debug + Sync + Send { /// /// As noted above, columns referenced only by pushed-down filters may be /// absent from `projection`. + /// + /// # Note + /// + /// Overriding [`TableProvider::scan_with_args`] will give you access to more arguments, + /// e.g. [`ScanArgs::skip`] async fn scan( &self, state: &dyn Session, @@ -209,6 +215,30 @@ pub trait TableProvider: Any + Debug + Sync + Send { /// A [`ScanResult`] containing the [`ExecutionPlan`] for scanning the table /// /// See [`Self::scan`] for detailed documentation about projection, filters, and limits. + /// + /// # Evaluation Order + /// + /// The logical evaluation order is [`ScanArgs::filters`], then [`ScanArgs::skip`] + /// and [`ScanArgs::limit`], and finally `projection`. + /// + /// Note that `limit` and `skip` apply to the filtered result, not to the unfiltered + /// input, and `projection` affects only which columns are returned, not which rows qualify. + /// + /// For example, if a scan receives: + /// + /// - `projection = [a]` + /// - `filters = [b > 5]` + /// - `limit = Some(3)` + /// - `skip = Some(5)` + /// + /// It must logically produce results equivalent to: + /// + /// ```text + /// PROJECTION a (OFFSET 5 LIMIT 3 (SCAN WHERE b > 5)) + /// ``` + /// + /// As noted above, columns referenced only by pushed-down filters may be + /// absent from `projection`. // Hand-written `#[async_trait]` expansion to reduce compile time. See // . fn scan_with_args<'a, 'life0, 'life1, 'async_trait>( @@ -231,6 +261,17 @@ pub trait TableProvider: Any + Debug + Sync + Send { Box::pin(async move { Ok(plan.await?.into()) }) } + /// Specify if DataFusion should provide the offset/skip to the + /// TableProvider to apply *during* the scan. + /// + /// # Note + /// + /// Make sure [`TableProvider::scan_with_args`] is overridden too + /// and [`ScanArgs::skip`] is used! + fn supports_skip_pushdown(&self) -> bool { + false + } + /// Specify if DataFusion should provide filter expressions to the /// TableProvider to apply *during* the scan. /// @@ -538,6 +579,7 @@ pub struct ScanArgs<'a> { filters: Option<&'a [Expr]>, projection: Option<&'a [usize]>, limit: Option, + skip: Option, statistics_requests: &'a [StatisticsRequest], } @@ -601,6 +643,25 @@ impl<'a> ScanArgs<'a> { self.limit } + /// Set the number of rows to skip from the scan. + /// + /// If specified, the scan should skip this many rows. This is typically + /// used to optimize queries with `OFFSET` clauses. + /// + /// # Arguments + /// * `skip` - Optional number of rows to skip + pub fn with_skip(mut self, skip: Option) -> Self { + self.skip = skip; + self + } + + /// Get the number of rows to skip from the scan. + /// + /// Returns the row offsets/skip, or `None` if no such was specified. + pub fn skip(&self) -> Option { + self.skip + } + /// Specifies the statistics the caller may use when optimizing the query. /// /// This is intended to allow the `TableProvider` to cheaply provide diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 5b2f1193c5044..9b1b979e3fde3 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -2347,7 +2347,10 @@ impl Unparser<'_> { } fn is_scan_with_pushdown(scan: &TableScan) -> bool { - scan.projection.is_some() || !scan.filters.is_empty() || scan.fetch.is_some() + scan.projection.is_some() + || !scan.filters.is_empty() + || scan.fetch.is_some() + || scan.skip.is_some() } /// Returns true if a plan, when used as the direct child of a SubqueryAlias, @@ -2621,8 +2624,13 @@ impl Unparser<'_> { builder = builder.filter(filter)?; } - if let Some(fetch) = table_scan.fetch { - builder = builder.limit(0, Some(fetch))?; + match (table_scan.skip, table_scan.fetch) { + (Some(offset), Some(fetch)) => { + builder = builder.limit(offset, Some(fetch))? + } + (Some(offset), None) => builder = builder.limit(offset, None)?, + (None, Some(fetch)) => builder = builder.limit(0, Some(fetch))?, + (None, None) => (), } // If the table scan has an alias but no projection or filters, it means no column references are rebased. diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs index 75a536d3ec869..c30e8d3e070c3 100644 --- a/datafusion/sql/tests/cases/plan_to_sql.rs +++ b/datafusion/sql/tests/cases/plan_to_sql.rs @@ -28,7 +28,7 @@ use datafusion_expr::test::function_stub::{ }; use datafusion_expr::{ ColumnarValue, EmptyRelation, Expr, Extension, LogicalPlan, LogicalPlanBuilder, - ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Union, + ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, TableScanBuilder, Union, UserDefinedLogicalNode, UserDefinedLogicalNodeCore, Volatility, WindowFrame, WindowFunctionDefinition, cast, col, exists, in_subquery, lit, scalar_subquery, table_scan, wildcard, @@ -56,6 +56,7 @@ use std::{fmt, vec}; use crate::common::{MockContextProvider, MockSessionState}; use datafusion_expr::builder::{ project, subquery_alias, table_scan_with_filter_and_fetch, table_scan_with_filters, + table_source, }; use datafusion_functions::core::planner::CoreFunctionPlanner; use datafusion_functions::unicode::planner::UnicodeFunctionPlanner; @@ -4703,3 +4704,24 @@ fn snowflake_flatten_multiple_unnest_cross_join() -> Result<(), DataFusionError> ); Ok(()) } + +#[test] +fn test_table_scan_with_skip() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("age", DataType::Utf8, false), + ]); + let scan = |skip, fetch| -> Result { + let scan = TableScanBuilder::new("t1", table_source(&schema)) + .with_skip(skip) + .with_fetch(fetch) + .build()?; + let plan = LogicalPlanBuilder::table_scan(scan)?.build()?; + Ok(plan_to_sql(&plan)?.to_string()) + }; + + assert_snapshot!(scan(Some(5), Some(10))?, @"SELECT * FROM t1 LIMIT 10 OFFSET 5"); + assert_snapshot!(scan(Some(5), None)?, @"SELECT * FROM t1 OFFSET 5"); + assert_snapshot!(scan(None, Some(10))?, @"SELECT * FROM t1 LIMIT 10"); + Ok(()) +} diff --git a/datafusion/sqllogictest/test_files/skip_pushdown.slt b/datafusion/sqllogictest/test_files/skip_pushdown.slt new file mode 100644 index 0000000000000..d0c0963f21c74 --- /dev/null +++ b/datafusion/sqllogictest/test_files/skip_pushdown.slt @@ -0,0 +1,839 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +########## +## skip pushdown tests +## +## `push_down_limit` pushes a `LIMIT ... skip ...` skip into +## `TableScan::skip` whenever the underlying `TableSource` reports +## `supports_skip_pushdown() == true` (today: `ListingTable`, i.e. +## CSV/JSON/Parquet/Arrow files). Providers that do not opt in (e.g. the +## in-memory `MemTable` created by `CREATE TABLE ... AS VALUES`) keep the +## skip in a `Limit` node above the scan instead. +########## + +# Source data: 20 rows, split into 5 files of 4 rows each so that skip +# values can be chosen to land mid-file and to span a file boundary. +statement ok +CREATE TABLE skip_src AS +SELECT i AS id, i * 10 AS val FROM generate_series(0, 19) t(i); + +query II +SELECT id, val FROM skip_src ORDER BY id +---- +0 0 +1 10 +2 20 +3 30 +4 40 +5 50 +6 60 +7 70 +8 80 +9 90 +10 100 +11 110 +12 120 +13 130 +14 140 +15 150 +16 160 +17 170 +18 180 +19 190 + +# File 0: ids 0-3 +query I +COPY (SELECT * FROM skip_src WHERE id BETWEEN 0 AND 3) +TO 'test_files/scratch/skip_pushdown/csv/part-0.csv' +STORED AS CSV OPTIONS ('format.has_header' 'true'); +---- +4 + +query I +COPY (SELECT * FROM skip_src WHERE id BETWEEN 0 AND 3) +TO 'test_files/scratch/skip_pushdown/json/part-0.json' +STORED AS JSON; +---- +4 + +query I +COPY (SELECT * FROM skip_src WHERE id BETWEEN 0 AND 3) +TO 'test_files/scratch/skip_pushdown/parquet/part-0.parquet' +STORED AS PARQUET; +---- +4 + +query I +COPY (SELECT * FROM skip_src WHERE id BETWEEN 0 AND 3) +TO 'test_files/scratch/skip_pushdown/arrow/part-0.arrow' +STORED AS ARROW; +---- +4 + +# File 1: ids 4-7 +query I +COPY (SELECT * FROM skip_src WHERE id BETWEEN 4 AND 7) +TO 'test_files/scratch/skip_pushdown/csv/part-1.csv' +STORED AS CSV OPTIONS ('format.has_header' 'true'); +---- +4 + +query I +COPY (SELECT * FROM skip_src WHERE id BETWEEN 4 AND 7) +TO 'test_files/scratch/skip_pushdown/json/part-1.json' +STORED AS JSON; +---- +4 + +query I +COPY (SELECT * FROM skip_src WHERE id BETWEEN 4 AND 7) +TO 'test_files/scratch/skip_pushdown/parquet/part-1.parquet' +STORED AS PARQUET; +---- +4 + +query I +COPY (SELECT * FROM skip_src WHERE id BETWEEN 4 AND 7) +TO 'test_files/scratch/skip_pushdown/arrow/part-1.arrow' +STORED AS ARROW; +---- +4 + +# File 2: ids 8-11 +query I +COPY (SELECT * FROM skip_src WHERE id BETWEEN 8 AND 11) +TO 'test_files/scratch/skip_pushdown/csv/part-2.csv' +STORED AS CSV OPTIONS ('format.has_header' 'true'); +---- +4 + +query I +COPY (SELECT * FROM skip_src WHERE id BETWEEN 8 AND 11) +TO 'test_files/scratch/skip_pushdown/json/part-2.json' +STORED AS JSON; +---- +4 + +query I +COPY (SELECT * FROM skip_src WHERE id BETWEEN 8 AND 11) +TO 'test_files/scratch/skip_pushdown/parquet/part-2.parquet' +STORED AS PARQUET; +---- +4 + +query I +COPY (SELECT * FROM skip_src WHERE id BETWEEN 8 AND 11) +TO 'test_files/scratch/skip_pushdown/arrow/part-2.arrow' +STORED AS ARROW; +---- +4 + +# File 3: ids 12-15 +query I +COPY (SELECT * FROM skip_src WHERE id BETWEEN 12 AND 15) +TO 'test_files/scratch/skip_pushdown/csv/part-3.csv' +STORED AS CSV OPTIONS ('format.has_header' 'true'); +---- +4 + +query I +COPY (SELECT * FROM skip_src WHERE id BETWEEN 12 AND 15) +TO 'test_files/scratch/skip_pushdown/json/part-3.json' +STORED AS JSON; +---- +4 + +query I +COPY (SELECT * FROM skip_src WHERE id BETWEEN 12 AND 15) +TO 'test_files/scratch/skip_pushdown/parquet/part-3.parquet' +STORED AS PARQUET; +---- +4 + +query I +COPY (SELECT * FROM skip_src WHERE id BETWEEN 12 AND 15) +TO 'test_files/scratch/skip_pushdown/arrow/part-3.arrow' +STORED AS ARROW; +---- +4 + +# File 4: ids 16-19 +query I +COPY (SELECT * FROM skip_src WHERE id BETWEEN 16 AND 19) +TO 'test_files/scratch/skip_pushdown/csv/part-4.csv' +STORED AS CSV OPTIONS ('format.has_header' 'true'); +---- +4 + +query I +COPY (SELECT * FROM skip_src WHERE id BETWEEN 16 AND 19) +TO 'test_files/scratch/skip_pushdown/json/part-4.json' +STORED AS JSON; +---- +4 + +query I +COPY (SELECT * FROM skip_src WHERE id BETWEEN 16 AND 19) +TO 'test_files/scratch/skip_pushdown/parquet/part-4.parquet' +STORED AS PARQUET; +---- +4 + +query I +COPY (SELECT * FROM skip_src WHERE id BETWEEN 16 AND 19) +TO 'test_files/scratch/skip_pushdown/arrow/part-4.arrow' +STORED AS ARROW; +---- +4 + +statement ok +CREATE EXTERNAL TABLE skip_csv (id BIGINT, val BIGINT) +STORED AS CSV +LOCATION 'test_files/scratch/skip_pushdown/csv/' +OPTIONS ('format.has_header' 'true'); + +statement ok +CREATE EXTERNAL TABLE skip_json (id BIGINT, val BIGINT) +STORED AS JSON +LOCATION 'test_files/scratch/skip_pushdown/json/'; + +statement ok +CREATE EXTERNAL TABLE skip_parquet (id BIGINT, val BIGINT) +STORED AS PARQUET +LOCATION 'test_files/scratch/skip_pushdown/parquet/'; + +statement ok +CREATE EXTERNAL TABLE skip_arrow (id BIGINT, val BIGINT) +STORED AS ARROW +LOCATION 'test_files/scratch/skip_pushdown/arrow/'; + +#################### +# supports_skip_pushdown(): MemTable does not opt in, so the skip stays +# on the `Limit` node and `TableScan` gets no `skip` attribute. +#################### + +query TT +EXPLAIN SELECT id FROM skip_src OFFSET 5 LIMIT 3 +---- +logical_plan +01)Limit: skip=5, fetch=3 +02)--TableScan: skip_src projection=[id], fetch=8 +physical_plan +01)GlobalLimitExec: skip=5, fetch=3 +02)--CoalescePartitionsExec: fetch=8 +03)----DataSourceExec: partitions=4, partition_sizes=[1, 0, 0, 0], fetch=8 + +#################### +# supports_skip_pushdown(): ListingTable (CSV/JSON/Parquet/Arrow) opts +# in, so the skip is folded into `TableScan::skip` and the outer `Limit` +# no longer needs to skip anything. +#################### + +statement ok +set datafusion.explain.logical_plan_only = true; + +query TT +EXPLAIN SELECT id FROM skip_csv OFFSET 5 LIMIT 3 +---- +logical_plan +01)Limit: skip=0, fetch=3 +02)--TableScan: skip_csv projection=[id], fetch=3, skip=5 + +query TT +EXPLAIN SELECT id FROM skip_json OFFSET 5 LIMIT 3 +---- +logical_plan +01)Limit: skip=0, fetch=3 +02)--TableScan: skip_json projection=[id], fetch=3, skip=5 + +query TT +EXPLAIN SELECT id FROM skip_parquet OFFSET 5 LIMIT 3 +---- +logical_plan +01)Limit: skip=0, fetch=3 +02)--TableScan: skip_parquet projection=[id], fetch=3, skip=5 + +query TT +EXPLAIN SELECT id FROM skip_arrow OFFSET 5 LIMIT 3 +---- +logical_plan +01)Limit: skip=0, fetch=3 +02)--TableScan: skip_arrow projection=[id], fetch=3, skip=5 + +statement ok +reset datafusion.explain.logical_plan_only; + +# Physical plan with a single partition: the skip is applied by a +# `GlobalLimitExec` wrapped directly around the file scan (built inside +# `ListingTable::scan_with_args`). +# +# NB: `OFFSET 15 LIMIT 10` is chosen so that `skip + limit` (25) exceeds +# the table's 20 rows. `ListingTable` stops *listing* files once file-level +# statistics show enough rows to satisfy a smaller limit (see +# `get_files_with_limit`), and which files that leaves out is not +# deterministic (files are discovered concurrently). Requiring all rows +# keeps the file set -- and thus this plan -- deterministic. +statement ok +set datafusion.execution.target_partitions = 1; + +query TT +EXPLAIN SELECT id FROM skip_parquet OFFSET 15 LIMIT 10 +---- +logical_plan +01)Limit: skip=0, fetch=10 +02)--TableScan: skip_parquet projection=[id], fetch=10, skip=15 +physical_plan +01)GlobalLimitExec: skip=15, fetch=10 +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/skip_pushdown/parquet/part-0.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/skip_pushdown/parquet/part-1.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/skip_pushdown/parquet/part-2.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/skip_pushdown/parquet/part-3.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/skip_pushdown/parquet/part-4.parquet]]}, projection=[id], limit=25, file_type=parquet + +# Physical plan with multiple partitions: `GlobalLimitExec` requires a +# single input partition, so a `CoalescePartitionsExec` is inserted below +# it to merge the per-file partitions before the skip is applied. +statement ok +set datafusion.execution.target_partitions = 4; + +query TT +EXPLAIN SELECT id FROM skip_parquet OFFSET 15 LIMIT 10 +---- +logical_plan +01)Limit: skip=0, fetch=10 +02)--TableScan: skip_parquet projection=[id], fetch=10, skip=15 +physical_plan +01)GlobalLimitExec: skip=15, fetch=10 +02)--CoalescePartitionsExec: fetch=25 +03)----DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/skip_pushdown/parquet/part-0.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/skip_pushdown/parquet/part-1.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/skip_pushdown/parquet/part-2.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/skip_pushdown/parquet/part-3.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/skip_pushdown/parquet/part-4.parquet]]}, projection=[id], limit=25, file_type=parquet + +statement ok +reset datafusion.execution.target_partitions; + +#################### +# Correctness: skip must skip exactly the right number of rows, for +# every source format and across a range of `target_partitions` settings +# (1 = no parallelism, 8 = more partitions than files so some partitions +# read no file at all). +# +# The `COUNT`-based checks below intentionally avoid asserting exact row +# *values* without an `ORDER BY`: which of the 5 underlying files +# `ListingTable` decides to read for a given `LIMIT` is not deterministic +# (see the note above `get_files_with_limit`), so only counts, distinct +# counts and set-membership are checked. The final `ORDER BY` query per +# table cross-checks the actual values deterministically -- note that an +# `ORDER BY` blocks skip pushdown into `TableScan` (a `Sort` sits +# between the `Limit` and the scan), so it exercises the plain +# `GlobalLimitExec` path rather than the pushdown path exercised above. +#################### + +# skip_csv + +statement ok +set datafusion.execution.target_partitions = 1; + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_csv OFFSET 0 LIMIT 3) +---- +3 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_csv OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(DISTINCT id) FROM (SELECT id FROM skip_csv OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(*) FROM ((SELECT id FROM skip_csv OFFSET 7 LIMIT 6) EXCEPT (SELECT id FROM skip_src)) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_csv OFFSET 19 LIMIT 5) +---- +1 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_csv OFFSET 20) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_csv OFFSET 25 LIMIT 3) +---- +0 + +query II +SELECT id, val FROM skip_csv ORDER BY id OFFSET 7 LIMIT 6 +---- +7 70 +8 80 +9 90 +10 100 +11 110 +12 120 + +statement ok +set datafusion.execution.target_partitions = 8; + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_csv OFFSET 0 LIMIT 3) +---- +3 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_csv OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(DISTINCT id) FROM (SELECT id FROM skip_csv OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(*) FROM ((SELECT id FROM skip_csv OFFSET 7 LIMIT 6) EXCEPT (SELECT id FROM skip_src)) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_csv OFFSET 19 LIMIT 5) +---- +1 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_csv OFFSET 20) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_csv OFFSET 25 LIMIT 3) +---- +0 + +query II +SELECT id, val FROM skip_csv ORDER BY id OFFSET 7 LIMIT 6 +---- +7 70 +8 80 +9 90 +10 100 +11 110 +12 120 + +# skip_json + +statement ok +set datafusion.execution.target_partitions = 1; + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_json OFFSET 0 LIMIT 3) +---- +3 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_json OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(DISTINCT id) FROM (SELECT id FROM skip_json OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(*) FROM ((SELECT id FROM skip_json OFFSET 7 LIMIT 6) EXCEPT (SELECT id FROM skip_src)) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_json OFFSET 19 LIMIT 5) +---- +1 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_json OFFSET 20) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_json OFFSET 25 LIMIT 3) +---- +0 + +query II +SELECT id, val FROM skip_json ORDER BY id OFFSET 7 LIMIT 6 +---- +7 70 +8 80 +9 90 +10 100 +11 110 +12 120 + +statement ok +set datafusion.execution.target_partitions = 8; + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_json OFFSET 0 LIMIT 3) +---- +3 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_json OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(DISTINCT id) FROM (SELECT id FROM skip_json OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(*) FROM ((SELECT id FROM skip_json OFFSET 7 LIMIT 6) EXCEPT (SELECT id FROM skip_src)) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_json OFFSET 19 LIMIT 5) +---- +1 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_json OFFSET 20) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_json OFFSET 25 LIMIT 3) +---- +0 + +query II +SELECT id, val FROM skip_json ORDER BY id OFFSET 7 LIMIT 6 +---- +7 70 +8 80 +9 90 +10 100 +11 110 +12 120 + +# skip_parquet + +statement ok +set datafusion.execution.target_partitions = 1; + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_parquet OFFSET 0 LIMIT 3) +---- +3 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_parquet OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(DISTINCT id) FROM (SELECT id FROM skip_parquet OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(*) FROM ((SELECT id FROM skip_parquet OFFSET 7 LIMIT 6) EXCEPT (SELECT id FROM skip_src)) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_parquet OFFSET 19 LIMIT 5) +---- +1 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_parquet OFFSET 20) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_parquet OFFSET 25 LIMIT 3) +---- +0 + +query II +SELECT id, val FROM skip_parquet ORDER BY id OFFSET 7 LIMIT 6 +---- +7 70 +8 80 +9 90 +10 100 +11 110 +12 120 + +statement ok +set datafusion.execution.target_partitions = 8; + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_parquet OFFSET 0 LIMIT 3) +---- +3 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_parquet OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(DISTINCT id) FROM (SELECT id FROM skip_parquet OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(*) FROM ((SELECT id FROM skip_parquet OFFSET 7 LIMIT 6) EXCEPT (SELECT id FROM skip_src)) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_parquet OFFSET 19 LIMIT 5) +---- +1 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_parquet OFFSET 20) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_parquet OFFSET 25 LIMIT 3) +---- +0 + +query II +SELECT id, val FROM skip_parquet ORDER BY id OFFSET 7 LIMIT 6 +---- +7 70 +8 80 +9 90 +10 100 +11 110 +12 120 + +# skip_arrow + +statement ok +set datafusion.execution.target_partitions = 1; + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_arrow OFFSET 0 LIMIT 3) +---- +3 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_arrow OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(DISTINCT id) FROM (SELECT id FROM skip_arrow OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(*) FROM ((SELECT id FROM skip_arrow OFFSET 7 LIMIT 6) EXCEPT (SELECT id FROM skip_src)) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_arrow OFFSET 19 LIMIT 5) +---- +1 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_arrow OFFSET 20) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_arrow OFFSET 25 LIMIT 3) +---- +0 + +query II +SELECT id, val FROM skip_arrow ORDER BY id OFFSET 7 LIMIT 6 +---- +7 70 +8 80 +9 90 +10 100 +11 110 +12 120 + +statement ok +set datafusion.execution.target_partitions = 8; + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_arrow OFFSET 0 LIMIT 3) +---- +3 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_arrow OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(DISTINCT id) FROM (SELECT id FROM skip_arrow OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(*) FROM ((SELECT id FROM skip_arrow OFFSET 7 LIMIT 6) EXCEPT (SELECT id FROM skip_src)) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_arrow OFFSET 19 LIMIT 5) +---- +1 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_arrow OFFSET 20) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_arrow OFFSET 25 LIMIT 3) +---- +0 + +query II +SELECT id, val FROM skip_arrow ORDER BY id OFFSET 7 LIMIT 6 +---- +7 70 +8 80 +9 90 +10 100 +11 110 +12 120 + +#################### +# Edge cases +#################### + +statement ok +set datafusion.execution.target_partitions = 1; + +statement ok +set datafusion.explain.logical_plan_only = true; + +# Nested OFFSETs are merged by `push_down_limit` and pushed as one skip +query TT +EXPLAIN SELECT id FROM (SELECT id FROM skip_parquet OFFSET 3 LIMIT 10) OFFSET 4 LIMIT 3 +---- +logical_plan +01)Limit: skip=0, fetch=3 +02)--TableScan: skip_parquet projection=[id], fetch=3, skip=7 + +query I +SELECT COUNT(*) FROM (SELECT id FROM (SELECT id FROM skip_parquet OFFSET 3 LIMIT 10) OFFSET 4 LIMIT 3) +---- +3 + +# OFFSET / LIMIT near i64::MAX: `skip + limit` must saturate, not overflow +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_parquet OFFSET 9223372036854775807 LIMIT 3) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM skip_parquet OFFSET 18 LIMIT 9223372036854775807) +---- +2 + +# Parameters only become literals after binding +statement ok +PREPARE skip_q(BIGINT, BIGINT) AS +SELECT COUNT(*) FROM (SELECT id FROM skip_parquet LIMIT $1 OFFSET $2); + +query I +EXECUTE skip_q(3, 18); +---- +2 + +query I +EXECUTE skip_q(3, 20); +---- +0 + +statement ok +DEALLOCATE skip_q; + +# Hive partitioned table: filters on the partition column are `Exact`, so the +# skip is pushed below them and must apply to the pruned file set only +query I +COPY (SELECT id, val, id / 4 AS part FROM skip_src) +TO 'test_files/scratch/skip_pushdown/hive/' +STORED AS PARQUET +PARTITIONED BY (part); +---- +20 + +statement ok +CREATE EXTERNAL TABLE skip_hive (id BIGINT, val BIGINT, part BIGINT) +STORED AS PARQUET +PARTITIONED BY (part) +LOCATION 'test_files/scratch/skip_pushdown/hive/'; + +query TT +EXPLAIN SELECT id FROM skip_hive WHERE part = 2 OFFSET 1 LIMIT 2 +---- +logical_plan +01)Limit: skip=0, fetch=2 +02)--TableScan: skip_hive projection=[id], full_filters=[skip_hive.part = Int64(2)], fetch=2, skip=1 + +statement ok +reset datafusion.explain.logical_plan_only; + +query I +SELECT id FROM skip_hive WHERE part = 2 OFFSET 1 LIMIT 2 +---- +9 +10 + +statement ok +DROP TABLE skip_hive; + +# Config reset +# The SLT runner sets `target_partitions` to 4 instead of using the +# default, so reset it explicitly. +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +DROP TABLE skip_src; + +statement ok +DROP TABLE skip_csv; + +statement ok +DROP TABLE skip_json; + +statement ok +DROP TABLE skip_parquet; + +statement ok +DROP TABLE skip_arrow; diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs index 865707a11d541..f871fd4f4a2e4 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs @@ -26,7 +26,7 @@ use datafusion::common::{ use datafusion::datasource::provider_as_source; use datafusion::logical_expr::utils::split_conjunction_owned; use datafusion::logical_expr::{ - EmptyRelation, Expr, LogicalPlan, LogicalPlanBuilder, Values, + EmptyRelation, Expr, LogicalPlan, LogicalPlanBuilder, TableScanBuilder, Values, }; use std::sync::Arc; use substrait::proto::expression::MaskExpression; @@ -61,13 +61,13 @@ pub async fn from_read_rel( _ => return plan_err!("No table named '{table_ref}'"), }; - LogicalPlanBuilder::scan_with_filters( + let table_scan = TableScanBuilder::new( table_ref, provider_as_source(Arc::clone(&provider)), - None, - filters, - )? - .build()? + ) + .with_filters(filters) + .build()?; + LogicalPlanBuilder::table_scan(table_scan)?.build()? }; ensure_schema_compatibility(plan.schema(), schema.clone())?; 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 e878b3816ff42..389cea3be8b25 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; +use datafusion::logical_expr::{Limit, LogicalPlan, lit}; use std::sync::Arc; use substrait::proto::rel::RelType; use substrait::proto::{FetchRel, Rel, fetch_rel}; @@ -28,10 +28,27 @@ pub fn from_limit( ) -> datafusion::common::Result> { let input = producer.handle_plan(limit.input.as_ref())?; let empty_schema = Arc::new(DFSchema::empty()); - let offset_mode = limit - .skip - .as_ref() - .map(|expr| producer.handle_expr(expr.as_ref(), &empty_schema)) + + // 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)) .transpose()? .map(Box::new) .map(fetch_rel::OffsetMode::OffsetExpr); diff --git a/docs/source/library-user-guide/building-logical-plans.md b/docs/source/library-user-guide/building-logical-plans.md index 6efd97879ac4d..e058ebe1dac75 100644 --- a/docs/source/library-user-guide/building-logical-plans.md +++ b/docs/source/library-user-guide/building-logical-plans.md @@ -91,7 +91,7 @@ There are several functions that can be used to create a new builder, such as - `empty` - create an empty plan with no fields - `values` - create a plan from a set of literal values - `scan` - create a plan representing a table scan -- `scan_with_filters` - create a plan representing a table scan with filters +- `table_scan` - create a plan representing a table scan with optional projection/filters/fetch/skip Once the builder is created, transformation methods can be called to declare that further operations should be performed on the plan. Note that all we are doing at this stage is building up the logical plan structure. No query diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index be71743f1afa9..1f1855f76f149 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -432,7 +432,7 @@ part of `StatisticsRegistry::default_with_builtin_providers()`. Use the walk instead: -```rust +```rust,ignore StatisticsContext::new_with_registry(registry) .compute_extended(plan, &StatisticsArgs::new())?; // or .compute(...) for core Statistics ``` @@ -718,3 +718,73 @@ impl FileSource for MySource { // ... } ``` + +### `TableProvider::scan_with_args` supports skip pushdown + +The `ScanArgs` struct now carries a new `skip` (number of rows to skip) field +alongside `limit`. `TableProvider::scan_with_args()` by default forwards to `scan()`, so +existing `TableProvider` implementations keep compiling and working +unchanged; implementations that want to take advantage of `skip` (or future +`ScanArgs` fields) can override `scan_with_args` in addition to `scan`. + +A new `TableProvider::supports_skip_pushdown() -> bool` (default `false`) +lets a provider guarantee it honors `ScanArgs::skip` exactly — that a scan +built with `skip` omits _exactly_ the first `skip` rows it would +otherwise have produced, not just approximately. `ListingTable` overrides this +to `true` and honors `skip`. + +DataFusion's SQL/DataFrame optimizer (`push_down_limit`) now populates +`LogicalPlan::TableScan::skip` from a query's `LIMIT ... OFFSET ...` clause, +but only when the scan's source reports `supports_skip_pushdown() == true`. +In that case the skip is moved entirely into the scan and the remaining +`Limit` above it no longer re-applies it; every other provider keeps today's +behavior unchanged, with the skip always enforced above the scan. `skip` is +also reachable directly, independent of the optimizer, by constructing a +`TableScan` (for example via `LogicalPlanBuilder::scan_with_filters_fetch_skip`) +or by calling `scan_with_args` with `ScanArgs::with_skip`. + +This pushdown is also preserved end-to-end through plan serialization: +`datafusion-substrait`'s producer folds a scan's `skip` into the `FetchRel` +it emits for the wrapping `Limit` (Substrait's `ReadRel` has no field of its +own for it), and `datafusion-proto`'s `ListingTableScanNode`, `ViewTableScanNode`, +and `CustomTableScanNode` messages gained optional `fetch`/`skip` fields so a +`TableScan`'s own `fetch`/`skip` — previously silently dropped — now +round-trips exactly. + +#### Who is affected + +- Custom `TableProvider` implementations that want skip pushdown: override + `scan_with_args` to read `ScanArgs::skip()` and honor it exactly, and + override `supports_skip_pushdown` to return `true`. Once you do, + `LIMIT ... OFFSET ...` queries against that provider will have their skip pushed + into the scan automatically. +- Callers that already build a scan through `ScanArgs`/`scan_with_args` (for + example custom query planners) can now also set `skip`. + +#### API breaks + +The _public_ `TableScan` struct has one new _public_ field - `skip` and +one field whose type is changed - `statistics_requests`. The addition of `skip` +field increases the size of `TableScan` and thus the size of `LogicalPlan` enum. To keep +the size of `LogicalPlan` reasonable the type of the `TableScan::statistics_requests` +field is changed from `BTreeSet` (24 bytes) to +`Box>` (8 bytes). +The users' applications are recommended to use `TableScanBuilder` to construct new +instances of `TableScan` instead of instantiating directly `TableScan { ... }` to avoid +the API breaks. + +#### Example + +```rust,ignore +let plan = provider + .scan_with_args( + &state, + ScanArgs::default() + .with_projection(projection) + .with_filters(Some(&filters)) + .with_limit(limit) + .with_skip(skip), + ) + .await? + .into_inner(); +```