diff --git a/datafusion/catalog-listing/src/table.rs b/datafusion/catalog-listing/src/table.rs index 6c294fe077db..7aa68801901b 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`, `offset`, `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`]) @@ -552,6 +553,10 @@ impl TableProvider for ListingTable { .collect() } + fn supports_offset_pushdown(&self) -> bool { + true + } + fn get_table_definition(&self) -> Option<&str> { self.definition.as_deref() } @@ -594,6 +599,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 offset = args.offset(); + // The scan must read enough rows to satisfy `offset + limit`, not + // just `limit`, before any rows are skipped below. + let inflated_limit = limit.map(|l| l.saturating_add(offset.unwrap_or(0))); // extract types of partition columns let table_partition_cols = self @@ -621,7 +630,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 }; @@ -737,7 +746,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()) @@ -750,6 +759,11 @@ impl ListingTable { .create_physical_plan(state, scan_config) .await?; + let plan: Arc = match offset { + 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 3342db54de92..0343242ddf0b 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_offset_pushdown(&self) -> bool { + self.table_provider.supports_offset_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 982766dc8851..194de0a89aaa 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,48 @@ mod tests { Ok(()) } + #[cfg(feature = "parquet")] + #[tokio::test] + async fn scan_with_args_offset_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 offset) 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 offset_exec = table + .scan_with_args( + &ctx.state(), + ScanArgs::default().with_limit(Some(3)).with_offset(Some(2)), + ) + .await? + .into_inner(); + let offset_batches = collect(offset_exec, ctx.task_ctx()).await?; + let actual = arrow::compute::concat_batches( + &offset_batches[0].schema(), + &offset_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 757d540c77b8..681b97f61c90 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -589,6 +589,7 @@ impl DefaultPhysicalPlanner { projection, filters, fetch, + offset, projected_schema, statistics_requests, .. @@ -606,6 +607,7 @@ impl DefaultPhysicalPlanner { .with_projection(projection.as_deref()) .with_filters(Some(&filters_vec)) .with_limit(*fetch) + .with_offset(*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 36aa67bbe7e3..935a3a378515 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -491,7 +491,14 @@ impl LogicalPlanBuilder { projection: Option>, filters: Vec, ) -> Result { - Self::scan_with_filters_inner(table_name, table_source, projection, filters, None) + Self::scan_with_filters_inner( + table_name, + table_source, + projection, + filters, + None, + None, + ) } /// Convert a table provider into a builder with a TableScan with filter and fetch @@ -508,6 +515,26 @@ impl LogicalPlanBuilder { projection, filters, fetch, + None, + ) + } + + /// Convert a table provider into a builder with a TableScan with filter, fetch and offset + pub fn scan_with_filters_fetch_offset( + table_name: impl Into, + table_source: Arc, + projection: Option>, + filters: Vec, + fetch: Option, + offset: Option, + ) -> Result { + Self::scan_with_filters_inner( + table_name, + table_source, + projection, + filters, + fetch, + offset, ) } @@ -517,11 +544,13 @@ impl LogicalPlanBuilder { projection: Option>, filters: Vec, fetch: Option, + offset: Option, ) -> Result { let table_scan = TableScanBuilder::new(table_name, table_source) .with_projection(projection) .with_filters(filters) .with_fetch(fetch) + .with_offset(offset) .build()?; // Inline TableScan @@ -2229,17 +2258,39 @@ pub fn table_scan_with_filter_and_fetch( projection: Option>, filters: Vec, fetch: Option, +) -> Result { + table_scan_with_filter_and_fetch_and_offset( + name, + 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_offset( + name: Option>, + table_schema: &Schema, + projection: Option>, + filters: Vec, + fetch: Option, + offset: 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( + LogicalPlanBuilder::scan_with_filters_fetch_offset( name, table_source, projection, filters, fetch, + offset, ) } diff --git a/datafusion/expr/src/logical_plan/display.rs b/datafusion/expr/src/logical_plan/display.rs index 1d95bf894bfa..67e9454972f4 100644 --- a/datafusion/expr/src/logical_plan/display.rs +++ b/datafusion/expr/src/logical_plan/display.rs @@ -342,6 +342,7 @@ impl<'a, 'b> PgJsonVisitor<'a, 'b> { table_name, filters, fetch, + offset, .. }) => { let mut object = json!({ @@ -395,6 +396,10 @@ impl<'a, 'b> PgJsonVisitor<'a, 'b> { object["Fetch"] = serde_json::Value::Number((*f).into()); } + if let Some(o) = offset { + object["Offset"] = serde_json::Value::Number((*o).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 de6d1667c75b..48ea12f567d8 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, + offset, .. }) => { let projected_fields = match projection { @@ -2176,6 +2177,10 @@ impl LogicalPlan { write!(f, ", fetch={n}")?; } + if let Some(n) = offset { + write!(f, ", offset={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 offset: 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("offset", &self.offset) .finish_non_exhaustive() } } @@ -3238,7 +3250,9 @@ pub struct TableScanBuilder { projection: Option>, filters: Vec, fetch: Option, - statistics_requests: BTreeSet, + offset: 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(), + offset: None, + statistics_requests: Box::default(), } } @@ -3275,13 +3290,19 @@ impl TableScanBuilder { self } + /// Set the number of rows to skip. + pub fn with_offset(mut self, offset: Option) -> Self { + self.offset = offset; + 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, + offset, statistics_requests, } = self; @@ -3335,6 +3357,7 @@ impl TableScanBuilder { projected_schema, filters, fetch, + offset, statistics_requests, }) } @@ -3348,6 +3371,7 @@ impl From for TableScanBuilder { projection: scan.projection, filters: scan.filters, fetch: scan.fetch, + offset: scan.offset, statistics_requests: scan.statistics_requests, } } @@ -6410,7 +6434,8 @@ mod tests { projected_schema: Arc::clone(&schema), filters: vec![], fetch: None, - statistics_requests: BTreeSet::new(), + offset: None, + statistics_requests: Box::default(), })); let col = schema.field_names()[0].clone(); @@ -6441,7 +6466,8 @@ mod tests { projected_schema: Arc::clone(&unique_schema), filters: vec![], fetch: None, - statistics_requests: BTreeSet::new(), + offset: 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 ee43666736fe..8a539e19b3c0 100644 --- a/datafusion/expr/src/logical_plan/tree_node.rs +++ b/datafusion/expr/src/logical_plan/tree_node.rs @@ -697,6 +697,7 @@ impl LogicalPlan { projected_schema, filters, fetch, + offset, statistics_requests, }) => filters.map_elements(f)?.update_data(|filters| { LogicalPlan::TableScan(TableScan { @@ -706,6 +707,7 @@ impl LogicalPlan { projected_schema, filters, fetch, + offset, statistics_requests, }) }), diff --git a/datafusion/expr/src/table_source.rs b/datafusion/expr/src/table_source.rs index 65dce8f3c8b0..3dea20c7f619 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 `offset` rows it would otherwise have + /// produced. + fn supports_offset_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/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs index b34827a87b0b..151b64ee7a6c 100644 --- a/datafusion/optimizer/src/push_down_filter.rs +++ b/datafusion/optimizer/src/push_down_filter.rs @@ -3164,7 +3164,8 @@ mod tests { projection, source: Arc::new(test_provider), fetch: None, - statistics_requests: std::collections::BTreeSet::new(), + offset: 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 79c18fbdeb4e..be1ae5edf886 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_offset_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.offset = Some(scan.offset.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 @@ -305,10 +322,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; @@ -489,6 +507,49 @@ mod test { ) } + /// A `TableSource` that declares it will honor `offset` exactly, so + /// `push_down_limit` is allowed to push `skip` into `TableScan::offset` + /// and elide the outer `Limit`'s skip. + #[derive(Debug)] + struct OffsetPushdownTableSource { + schema: SchemaRef, + } + + impl TableSource for OffsetPushdownTableSource { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn supports_offset_pushdown(&self) -> bool { + true + } + } + + fn offset_pushdown_table_scan() -> Result { + let schema = Arc::new(Schema::new(test_table_scan_fields())); + let source = Arc::new(OffsetPushdownTableSource { schema }); + Ok(LogicalPlan::TableScan( + TableScanBuilder::new("test", source).build()?, + )) + } + + #[test] + fn limit_pushdown_offset_supported() -> Result<()> { + let table_scan = offset_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, offset=10 + " + ) + } + #[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 15dc272fabca..4b2e93a02bcc 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 offset = 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 offset = 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 offset = 8; } message ProjectionNode { diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index cada03c9a6d0..540b2d6861f9 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.offset.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.offset.as_ref() { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("offset", 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", + "offset", ]; #[allow(clippy::enum_variant_names)] @@ -5748,6 +5766,8 @@ impl<'de> serde::Deserialize<'de> for CustomTableScanNode { Schema, Filters, CustomTableData, + Fetch, + Offset, } 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), + "offset" => Ok(GeneratedField::Offset), _ => 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 offset__ = 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::Offset => { + if offset__.is_some() { + return Err(serde::de::Error::duplicate_field("offset")); + } + offset__ = + 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__, + offset: offset__, }) } } @@ -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.offset.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.offset.as_ref() { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("offset", 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", + "offset", "csv", "parquet", "avro", @@ -13173,6 +13233,8 @@ impl<'de> serde::Deserialize<'de> for ListingTableScanNode { Filters, TablePartitionCols, FileSortOrder, + Fetch, + Offset, 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), + "offset" => Ok(GeneratedField::Offset), "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 offset__ = 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::Offset => { + if offset__.is_some() { + return Err(serde::de::Error::duplicate_field("offset")); + } + offset__ = + 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__, + offset: offset__, file_format_type: file_format_type__, }) } @@ -29220,6 +29304,12 @@ impl serde::Serialize for ViewTableScanNode { if !self.definition.is_empty() { len += 1; } + if self.fetch.is_some() { + len += 1; + } + if self.offset.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)?; @@ -29236,6 +29326,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.offset.as_ref() { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("offset", ToString::to_string(&v).as_str())?; + } struct_ser.end() } } @@ -29252,6 +29352,8 @@ impl<'de> serde::Deserialize<'de> for ViewTableScanNode { "schema", "projection", "definition", + "fetch", + "offset", ]; #[allow(clippy::enum_variant_names)] @@ -29261,6 +29363,8 @@ impl<'de> serde::Deserialize<'de> for ViewTableScanNode { Schema, Projection, Definition, + Fetch, + Offset, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -29287,6 +29391,8 @@ impl<'de> serde::Deserialize<'de> for ViewTableScanNode { "schema" => Ok(GeneratedField::Schema), "projection" => Ok(GeneratedField::Projection), "definition" => Ok(GeneratedField::Definition), + "fetch" => Ok(GeneratedField::Fetch), + "offset" => Ok(GeneratedField::Offset), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -29311,6 +29417,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 offset__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::TableName => { @@ -29343,6 +29451,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::Offset => { + if offset__.is_some() { + return Err(serde::de::Error::duplicate_field("offset")); + } + offset__ = + map_.next_value::<::std::option::Option<::pbjson::private::NumberDeserialize<_>>>()?.map(|x| x.0) + ; + } } } Ok(ViewTableScanNode { @@ -29351,6 +29475,8 @@ impl<'de> serde::Deserialize<'de> for ViewTableScanNode { schema: schema__, projection: projection__, definition: definition__.unwrap_or_default(), + fetch: fetch__, + offset: offset__, }) } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index 4bb4af1e8532..dd41bb6c2abf 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 offset: ::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 offset: ::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 offset: ::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 71feae506dc6..7a965fcebeb4 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 37b42fc9a3ba..9db9e8c9525a 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -696,11 +696,13 @@ impl AsLogicalPlan for LogicalPlanNode { projection = Some(column_indices); } - LogicalPlanBuilder::scan_with_filters( + LogicalPlanBuilder::scan_with_filters_fetch_offset( table_name, provider_as_source(Arc::new(provider)), projection, filters, + scan.fetch.map(|f| f as usize), + scan.offset.map(|o| o as usize), )? .build() } @@ -730,11 +732,13 @@ impl AsLogicalPlan for LogicalPlanNode { ctx, )?; - LogicalPlanBuilder::scan_with_filters( + LogicalPlanBuilder::scan_with_filters_fetch_offset( table_name, provider_as_source(provider), projection, filters, + scan.fetch.map(|f| f as usize), + scan.offset.map(|o| o as usize), )? .build() } @@ -1210,10 +1214,13 @@ impl AsLogicalPlan for LogicalPlanNode { let table_name = from_table_reference(scan.table_name.as_ref(), "ViewScan")?; - LogicalPlanBuilder::scan( + LogicalPlanBuilder::scan_with_filters_fetch_offset( table_name, provider_as_source(Arc::new(provider)), projection, + vec![], + scan.fetch.map(|f| f as usize), + scan.offset.map(|o| o as usize), )? .build() } @@ -1408,6 +1415,8 @@ impl AsLogicalPlan for LogicalPlanNode { source, filters, projection, + fetch, + offset, .. }) => { let provider = source_as_provider(source)?; @@ -1540,6 +1549,8 @@ impl AsLogicalPlan for LogicalPlanNode { projection, filters, file_sort_order: exprs_vec, + fetch: fetch.map(|f| f as u64), + offset: offset.map(|o| o as u64), }, )), }) @@ -1563,6 +1574,8 @@ impl AsLogicalPlan for LogicalPlanNode { .definition() .map(|s| s.to_string()) .unwrap_or_default(), + fetch: fetch.map(|f| f as u64), + offset: offset.map(|o| o as u64), }, ))), }) @@ -1610,6 +1623,8 @@ impl AsLogicalPlan for LogicalPlanNode { schema: Some(schema), filters, custom_table_data: bytes, + fetch: fetch.map(|f| f as u64), + offset: offset.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 932c4188115f..f4f95d9b2d09 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_offset() -> 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 offset was actually pushed into the `TableScan` + // (ListingTable opts into `supports_offset_pushdown`), so this test + // exercises the new `ListingTableScanNode.offset` 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, offset=3", + "expected 'fetch=5' and 'offset=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 cd824938c003..88a0d2751cc3 100644 --- a/datafusion/session/src/table.rs +++ b/datafusion/session/src/table.rs @@ -184,6 +184,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::offset`] async fn scan( &self, state: &dyn Session, @@ -231,6 +236,17 @@ pub trait TableProvider: Any + Debug + Sync + Send { Box::pin(async move { Ok(plan.await?.into()) }) } + /// Specify if DataFusion should provide the offset to the + /// TableProvider to apply *during* the scan. + /// + /// # Note + /// + /// Make sure [`TableProvider::scan_with_args`] is overridden too + /// and [`ScanArgs::offset`] is used! + fn supports_offset_pushdown(&self) -> bool { + false + } + /// Specify if DataFusion should provide filter expressions to the /// TableProvider to apply *during* the scan. /// @@ -473,6 +489,7 @@ pub struct ScanArgs<'a> { filters: Option<&'a [Expr]>, projection: Option<&'a [usize]>, limit: Option, + offset: Option, statistics_requests: &'a [StatisticsRequest], } @@ -536,6 +553,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 + /// * `offset` - Optional number of rows to skip + pub fn with_offset(mut self, offset: Option) -> Self { + self.offset = offset; + self + } + + /// Get the number of rows to skip from the scan. + /// + /// Returns the row offset, or `None` if no offset was specified. + pub fn offset(&self) -> Option { + self.offset + } + /// 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 18af08fc1836..a961bfa46c1c 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -2585,8 +2585,13 @@ impl Unparser<'_> { builder = builder.filter(filter)?; } - if let Some(fetch) = table_scan.fetch { - builder = builder.limit(0, Some(fetch))?; + match (table_scan.offset, 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/sqllogictest/test_files/offset_pushdown.slt b/datafusion/sqllogictest/test_files/offset_pushdown.slt new file mode 100644 index 000000000000..67820eed9744 --- /dev/null +++ b/datafusion/sqllogictest/test_files/offset_pushdown.slt @@ -0,0 +1,752 @@ +# 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. + +########## +## OFFSET pushdown tests +## +## `push_down_limit` pushes a `LIMIT ... OFFSET ...` skip into +## `TableScan::offset` whenever the underlying `TableSource` reports +## `supports_offset_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 OFFSET +# values can be chosen to land mid-file and to span a file boundary. +statement ok +CREATE TABLE offset_src AS +SELECT i AS id, i * 10 AS val FROM generate_series(0, 19) t(i); + +query II +SELECT id, val FROM offset_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 offset_src WHERE id BETWEEN 0 AND 3) +TO 'test_files/scratch/offset_pushdown/csv/part-0.csv' +STORED AS CSV OPTIONS ('format.has_header' 'true'); +---- +4 + +query I +COPY (SELECT * FROM offset_src WHERE id BETWEEN 0 AND 3) +TO 'test_files/scratch/offset_pushdown/json/part-0.json' +STORED AS JSON; +---- +4 + +query I +COPY (SELECT * FROM offset_src WHERE id BETWEEN 0 AND 3) +TO 'test_files/scratch/offset_pushdown/parquet/part-0.parquet' +STORED AS PARQUET; +---- +4 + +query I +COPY (SELECT * FROM offset_src WHERE id BETWEEN 0 AND 3) +TO 'test_files/scratch/offset_pushdown/arrow/part-0.arrow' +STORED AS ARROW; +---- +4 + +# File 1: ids 4-7 +query I +COPY (SELECT * FROM offset_src WHERE id BETWEEN 4 AND 7) +TO 'test_files/scratch/offset_pushdown/csv/part-1.csv' +STORED AS CSV OPTIONS ('format.has_header' 'true'); +---- +4 + +query I +COPY (SELECT * FROM offset_src WHERE id BETWEEN 4 AND 7) +TO 'test_files/scratch/offset_pushdown/json/part-1.json' +STORED AS JSON; +---- +4 + +query I +COPY (SELECT * FROM offset_src WHERE id BETWEEN 4 AND 7) +TO 'test_files/scratch/offset_pushdown/parquet/part-1.parquet' +STORED AS PARQUET; +---- +4 + +query I +COPY (SELECT * FROM offset_src WHERE id BETWEEN 4 AND 7) +TO 'test_files/scratch/offset_pushdown/arrow/part-1.arrow' +STORED AS ARROW; +---- +4 + +# File 2: ids 8-11 +query I +COPY (SELECT * FROM offset_src WHERE id BETWEEN 8 AND 11) +TO 'test_files/scratch/offset_pushdown/csv/part-2.csv' +STORED AS CSV OPTIONS ('format.has_header' 'true'); +---- +4 + +query I +COPY (SELECT * FROM offset_src WHERE id BETWEEN 8 AND 11) +TO 'test_files/scratch/offset_pushdown/json/part-2.json' +STORED AS JSON; +---- +4 + +query I +COPY (SELECT * FROM offset_src WHERE id BETWEEN 8 AND 11) +TO 'test_files/scratch/offset_pushdown/parquet/part-2.parquet' +STORED AS PARQUET; +---- +4 + +query I +COPY (SELECT * FROM offset_src WHERE id BETWEEN 8 AND 11) +TO 'test_files/scratch/offset_pushdown/arrow/part-2.arrow' +STORED AS ARROW; +---- +4 + +# File 3: ids 12-15 +query I +COPY (SELECT * FROM offset_src WHERE id BETWEEN 12 AND 15) +TO 'test_files/scratch/offset_pushdown/csv/part-3.csv' +STORED AS CSV OPTIONS ('format.has_header' 'true'); +---- +4 + +query I +COPY (SELECT * FROM offset_src WHERE id BETWEEN 12 AND 15) +TO 'test_files/scratch/offset_pushdown/json/part-3.json' +STORED AS JSON; +---- +4 + +query I +COPY (SELECT * FROM offset_src WHERE id BETWEEN 12 AND 15) +TO 'test_files/scratch/offset_pushdown/parquet/part-3.parquet' +STORED AS PARQUET; +---- +4 + +query I +COPY (SELECT * FROM offset_src WHERE id BETWEEN 12 AND 15) +TO 'test_files/scratch/offset_pushdown/arrow/part-3.arrow' +STORED AS ARROW; +---- +4 + +# File 4: ids 16-19 +query I +COPY (SELECT * FROM offset_src WHERE id BETWEEN 16 AND 19) +TO 'test_files/scratch/offset_pushdown/csv/part-4.csv' +STORED AS CSV OPTIONS ('format.has_header' 'true'); +---- +4 + +query I +COPY (SELECT * FROM offset_src WHERE id BETWEEN 16 AND 19) +TO 'test_files/scratch/offset_pushdown/json/part-4.json' +STORED AS JSON; +---- +4 + +query I +COPY (SELECT * FROM offset_src WHERE id BETWEEN 16 AND 19) +TO 'test_files/scratch/offset_pushdown/parquet/part-4.parquet' +STORED AS PARQUET; +---- +4 + +query I +COPY (SELECT * FROM offset_src WHERE id BETWEEN 16 AND 19) +TO 'test_files/scratch/offset_pushdown/arrow/part-4.arrow' +STORED AS ARROW; +---- +4 + +statement ok +CREATE EXTERNAL TABLE offset_csv (id BIGINT, val BIGINT) +STORED AS CSV +LOCATION 'test_files/scratch/offset_pushdown/csv/' +OPTIONS ('format.has_header' 'true'); + +statement ok +CREATE EXTERNAL TABLE offset_json (id BIGINT, val BIGINT) +STORED AS JSON +LOCATION 'test_files/scratch/offset_pushdown/json/'; + +statement ok +CREATE EXTERNAL TABLE offset_parquet (id BIGINT, val BIGINT) +STORED AS PARQUET +LOCATION 'test_files/scratch/offset_pushdown/parquet/'; + +statement ok +CREATE EXTERNAL TABLE offset_arrow (id BIGINT, val BIGINT) +STORED AS ARROW +LOCATION 'test_files/scratch/offset_pushdown/arrow/'; + +#################### +# supports_offset_pushdown(): MemTable does not opt in, so the skip stays +# on the `Limit` node and `TableScan` gets no `offset` attribute. +#################### + +query TT +EXPLAIN SELECT id FROM offset_src OFFSET 5 LIMIT 3 +---- +logical_plan +01)Limit: skip=5, fetch=3 +02)--TableScan: offset_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_offset_pushdown(): ListingTable (CSV/JSON/Parquet/Arrow) opts +# in, so the skip is folded into `TableScan::offset` 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 offset_csv OFFSET 5 LIMIT 3 +---- +logical_plan +01)Limit: skip=0, fetch=3 +02)--TableScan: offset_csv projection=[id], fetch=3, offset=5 + +query TT +EXPLAIN SELECT id FROM offset_json OFFSET 5 LIMIT 3 +---- +logical_plan +01)Limit: skip=0, fetch=3 +02)--TableScan: offset_json projection=[id], fetch=3, offset=5 + +query TT +EXPLAIN SELECT id FROM offset_parquet OFFSET 5 LIMIT 3 +---- +logical_plan +01)Limit: skip=0, fetch=3 +02)--TableScan: offset_parquet projection=[id], fetch=3, offset=5 + +query TT +EXPLAIN SELECT id FROM offset_arrow OFFSET 5 LIMIT 3 +---- +logical_plan +01)Limit: skip=0, fetch=3 +02)--TableScan: offset_arrow projection=[id], fetch=3, offset=5 + +statement ok +reset datafusion.explain.logical_plan_only; + +# Physical plan with a single partition: the offset 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 `offset + 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 offset_parquet OFFSET 15 LIMIT 10 +---- +logical_plan +01)Limit: skip=0, fetch=10 +02)--TableScan: offset_parquet projection=[id], fetch=10, offset=15 +physical_plan +01)GlobalLimitExec: skip=15, fetch=10 +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/offset_pushdown/parquet/part-0.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/offset_pushdown/parquet/part-1.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/offset_pushdown/parquet/part-2.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/offset_pushdown/parquet/part-3.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/offset_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 offset_parquet OFFSET 15 LIMIT 10 +---- +logical_plan +01)Limit: skip=0, fetch=10 +02)--TableScan: offset_parquet projection=[id], fetch=10, offset=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/offset_pushdown/parquet/part-0.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/offset_pushdown/parquet/part-1.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/offset_pushdown/parquet/part-2.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/offset_pushdown/parquet/part-3.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/offset_pushdown/parquet/part-4.parquet]]}, projection=[id], limit=25, file_type=parquet + +statement ok +reset datafusion.execution.target_partitions; + +#################### +# Correctness: OFFSET 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 offset 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. +#################### + +# offset_csv + +statement ok +set datafusion.execution.target_partitions = 1; + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_csv OFFSET 0 LIMIT 3) +---- +3 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_csv OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(DISTINCT id) FROM (SELECT id FROM offset_csv OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(*) FROM ((SELECT id FROM offset_csv OFFSET 7 LIMIT 6) EXCEPT (SELECT id FROM offset_src)) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_csv OFFSET 19 LIMIT 5) +---- +1 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_csv OFFSET 20) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_csv OFFSET 25 LIMIT 3) +---- +0 + +query II +SELECT id, val FROM offset_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 offset_csv OFFSET 0 LIMIT 3) +---- +3 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_csv OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(DISTINCT id) FROM (SELECT id FROM offset_csv OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(*) FROM ((SELECT id FROM offset_csv OFFSET 7 LIMIT 6) EXCEPT (SELECT id FROM offset_src)) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_csv OFFSET 19 LIMIT 5) +---- +1 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_csv OFFSET 20) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_csv OFFSET 25 LIMIT 3) +---- +0 + +query II +SELECT id, val FROM offset_csv ORDER BY id OFFSET 7 LIMIT 6 +---- +7 70 +8 80 +9 90 +10 100 +11 110 +12 120 + +# offset_json + +statement ok +set datafusion.execution.target_partitions = 1; + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_json OFFSET 0 LIMIT 3) +---- +3 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_json OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(DISTINCT id) FROM (SELECT id FROM offset_json OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(*) FROM ((SELECT id FROM offset_json OFFSET 7 LIMIT 6) EXCEPT (SELECT id FROM offset_src)) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_json OFFSET 19 LIMIT 5) +---- +1 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_json OFFSET 20) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_json OFFSET 25 LIMIT 3) +---- +0 + +query II +SELECT id, val FROM offset_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 offset_json OFFSET 0 LIMIT 3) +---- +3 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_json OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(DISTINCT id) FROM (SELECT id FROM offset_json OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(*) FROM ((SELECT id FROM offset_json OFFSET 7 LIMIT 6) EXCEPT (SELECT id FROM offset_src)) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_json OFFSET 19 LIMIT 5) +---- +1 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_json OFFSET 20) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_json OFFSET 25 LIMIT 3) +---- +0 + +query II +SELECT id, val FROM offset_json ORDER BY id OFFSET 7 LIMIT 6 +---- +7 70 +8 80 +9 90 +10 100 +11 110 +12 120 + +# offset_parquet + +statement ok +set datafusion.execution.target_partitions = 1; + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_parquet OFFSET 0 LIMIT 3) +---- +3 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_parquet OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(DISTINCT id) FROM (SELECT id FROM offset_parquet OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(*) FROM ((SELECT id FROM offset_parquet OFFSET 7 LIMIT 6) EXCEPT (SELECT id FROM offset_src)) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_parquet OFFSET 19 LIMIT 5) +---- +1 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_parquet OFFSET 20) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_parquet OFFSET 25 LIMIT 3) +---- +0 + +query II +SELECT id, val FROM offset_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 offset_parquet OFFSET 0 LIMIT 3) +---- +3 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_parquet OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(DISTINCT id) FROM (SELECT id FROM offset_parquet OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(*) FROM ((SELECT id FROM offset_parquet OFFSET 7 LIMIT 6) EXCEPT (SELECT id FROM offset_src)) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_parquet OFFSET 19 LIMIT 5) +---- +1 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_parquet OFFSET 20) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_parquet OFFSET 25 LIMIT 3) +---- +0 + +query II +SELECT id, val FROM offset_parquet ORDER BY id OFFSET 7 LIMIT 6 +---- +7 70 +8 80 +9 90 +10 100 +11 110 +12 120 + +# offset_arrow + +statement ok +set datafusion.execution.target_partitions = 1; + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_arrow OFFSET 0 LIMIT 3) +---- +3 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_arrow OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(DISTINCT id) FROM (SELECT id FROM offset_arrow OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(*) FROM ((SELECT id FROM offset_arrow OFFSET 7 LIMIT 6) EXCEPT (SELECT id FROM offset_src)) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_arrow OFFSET 19 LIMIT 5) +---- +1 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_arrow OFFSET 20) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_arrow OFFSET 25 LIMIT 3) +---- +0 + +query II +SELECT id, val FROM offset_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 offset_arrow OFFSET 0 LIMIT 3) +---- +3 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_arrow OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(DISTINCT id) FROM (SELECT id FROM offset_arrow OFFSET 7 LIMIT 6) +---- +6 + +query I +SELECT COUNT(*) FROM ((SELECT id FROM offset_arrow OFFSET 7 LIMIT 6) EXCEPT (SELECT id FROM offset_src)) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_arrow OFFSET 19 LIMIT 5) +---- +1 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_arrow OFFSET 20) +---- +0 + +query I +SELECT COUNT(*) FROM (SELECT id FROM offset_arrow OFFSET 25 LIMIT 3) +---- +0 + +query II +SELECT id, val FROM offset_arrow ORDER BY id OFFSET 7 LIMIT 6 +---- +7 70 +8 80 +9 90 +10 100 +11 110 +12 120 + +# 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 offset_src; + +statement ok +DROP TABLE offset_csv; + +statement ok +DROP TABLE offset_json; + +statement ok +DROP TABLE offset_parquet; + +statement ok +DROP TABLE offset_arrow; 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 e878b3816ff4..5d7b12855dd8 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,25 @@ 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 offset pushdown (see `push_down_limit`) may have moved + // some or all of this `Limit`'s skip into the child `TableScan::offset`, + // 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.offset, + _ => None, + }; + let skip_expr = match (limit.skip.as_deref(), scan_offset) { + (Some(skip), Some(offset)) => Some(skip.clone() + lit(offset as i64)), + (Some(skip), None) => Some(skip.clone()), + (None, Some(offset)) => Some(lit(offset 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/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index 7c3e3fff67b6..d63ffc7537f5 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -302,7 +302,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 ``` @@ -512,3 +512,73 @@ Wire compatibility is directional: - A 55.0 reader must not consume an alias-preserving 56.0 MERGE payload. It ignores the unknown field but cannot preserve the qualifier required by the expressions, which can cause resolution failure or incorrect rebinding. + +### `TableProvider::scan_with_args` supports offset pushdown + +The `ScanArgs` struct now carries a new `offset` (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 `offset` (or future +`ScanArgs` fields) can override `scan_with_args` in addition to `scan`. + +A new `TableProvider::supports_offset_pushdown() -> bool` (default `false`) +lets a provider guarantee it honors `ScanArgs::offset` exactly — that a scan +built with `offset` omits _exactly_ the first `offset` rows it would +otherwise have produced, not just approximately. `ListingTable` overrides this +to `true` and honors `offset`. + +DataFusion's SQL/DataFrame optimizer (`push_down_limit`) now populates +`LogicalPlan::TableScan::offset` from a query's `LIMIT ... OFFSET ...` clause, +but only when the scan's source reports `supports_offset_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. `offset` is +also reachable directly, independent of the optimizer, by constructing a +`TableScan` (for example via `LogicalPlanBuilder::scan_with_filters_fetch_offset`) +or by calling `scan_with_args` with `ScanArgs::with_offset`. + +This pushdown is also preserved end-to-end through plan serialization: +`datafusion-substrait`'s producer folds a scan's `offset` 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`/`offset` fields so a +`TableScan`'s own `fetch`/`offset` — previously silently dropped — now +round-trips exactly. + +#### Who is affected + +- Custom `TableProvider` implementations that want offset pushdown: override + `scan_with_args` to read `ScanArgs::offset()` and honor it exactly, and + override `supports_offset_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 `with_offset`. + +#### API breaks + +The _public_ `TableScan` struct has one new _public_ field - `offset` and +one field whose type is changed - `statistics_requests`. The addition of `offset` +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_offset(offset), + ) + .await? + .into_inner(); +```