diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index d94253a84aa5f..a8162f137ed0a 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs @@ -34,7 +34,7 @@ use arrow::datatypes::{DataType, SchemaRef}; use datafusion_common::config::{ConfigOptions, CsvOptions}; use datafusion_common::tree_node::{TreeNode, TransformedResult}; use datafusion_common::{create_array, DataFusionError, NullEquality, Result, TableReference}; -use datafusion_datasource::file_scan_config::FileScanConfigBuilder; +use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; use datafusion_datasource::source::DataSourceExec; use datafusion_expr_common::operator::Operator; use datafusion_expr::{JoinType, SortExpr}; @@ -57,6 +57,8 @@ use datafusion_physical_optimizer::sanity_checker::SanityCheckPlan; use datafusion_physical_optimizer::enforce_sorting::replace_with_order_preserving_variants::{replace_with_order_preserving_variants, OrderPreservationContext}; use datafusion_physical_optimizer::enforce_sorting::sort_pushdown::{SortPushDown, assign_initial_requirements, pushdown_sorts}; use datafusion_physical_optimizer::ensure_requirements::EnsureRequirements; +use datafusion_physical_optimizer::limit_pushdown::LimitPushdown; +use datafusion_physical_optimizer::projection_pushdown::ProjectionPushdown; use datafusion_physical_optimizer::output_requirements::OutputRequirementExec; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion::prelude::*; @@ -2297,6 +2299,55 @@ async fn test_remove_unnecessary_spm2() -> Result<()> { Ok(()) } +#[test] +fn test_spm_fetch_preserves_ordering_through_child_rewrite() -> Result<()> { + let schema = create_test_schema()?; + let ordering: LexOrdering = [sort_expr("non_nullable_col", &schema)].into(); + let source = parquet_exec_with_sort(Arc::clone(&schema), vec![ordering.clone()]); + let projection = projection_exec( + vec![ + (col("nullable_col", &schema)?, "nullable_col".to_string()), + ( + col("non_nullable_col", &schema)?, + "non_nullable_col".to_string(), + ), + ], + source, + )?; + let plan = sort_preserving_merge_exec_with_fetch(ordering.clone(), projection, 100); + + let optimized = PlanWithCorrespondingSort::new_default(plan) + .transform_up(ensure_sorting)? + .data; + let optimized = check_integrity(optimized)?.plan; + let limit = optimized + .downcast_ref::() + .expect("SPM fetch should become a local limit"); + assert_eq!(limit.fetch(), 100); + assert_eq!(limit.required_ordering().as_ref(), Some(&ordering)); + + let config = ConfigOptions::new(); + let optimized = ProjectionPushdown::new().optimize(optimized, &config)?; + let limit = optimized + .downcast_ref::() + .expect("projection rewrite should retain the local limit"); + assert_eq!(limit.required_ordering().as_ref(), Some(&ordering)); + assert!(limit.input().is::()); + + let optimized = LimitPushdown::new().optimize(optimized, &config)?; + let source = optimized + .downcast_ref::() + .expect("limit should be pushed into the parquet scan"); + let config = source + .data_source() + .downcast_ref::() + .expect("parquet scan should use FileScanConfig"); + assert_eq!(config.limit, Some(100)); + assert!(config.preserve_order); + + Ok(()) +} + #[tokio::test] async fn test_change_wrong_sorting() -> Result<()> { let schema = create_test_schema()?; diff --git a/datafusion/physical-expr-common/src/sort_expr.rs b/datafusion/physical-expr-common/src/sort_expr.rs index 72e877234752f..6e8dbccdb7c0e 100644 --- a/datafusion/physical-expr-common/src/sort_expr.rs +++ b/datafusion/physical-expr-common/src/sort_expr.rs @@ -265,7 +265,8 @@ pub fn sort_exprs_try_to_proto>( /// [`LexRequirement`], because callers differ in what an empty list means: /// `LexOrdering::new` / `LexRequirement::new` return `None` for it, which is /// "no ordering declared" for a scan and an error for an operator that requires -/// one. +/// one. Callers with the former convention can use +/// [`optional_ordering_try_from_proto`] instead. /// /// [`PhysicalSortExprNode`]: datafusion_proto_models::protobuf::PhysicalSortExprNode #[cfg(feature = "proto")] @@ -279,6 +280,25 @@ pub fn sort_exprs_try_from_proto( .collect() } +/// Serialize an optional [`LexOrdering`], encoding `None` as an empty list. +#[cfg(feature = "proto")] +pub fn optional_ordering_try_to_proto( + ordering: Option<&LexOrdering>, + ctx: &crate::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, +) -> Result> { + sort_exprs_try_to_proto(ordering.into_iter().flatten(), ctx) +} + +/// Counterpart of [`optional_ordering_try_to_proto`]: an empty list decodes +/// as `None`. +#[cfg(feature = "proto")] +pub fn optional_ordering_try_from_proto( + nodes: &[datafusion_proto_models::protobuf::PhysicalSortExprNode], + ctx: &crate::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, +) -> Result> { + Ok(LexOrdering::new(sort_exprs_try_from_proto(nodes, ctx)?)) +} + impl PartialEq for PhysicalSortExpr { fn eq(&self, other: &Self) -> bool { self.options == other.options && self.expr.eq(&other.expr) diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs index 6efaf76457919..c66d5310a1c44 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs @@ -461,13 +461,17 @@ pub fn ensure_sorting( } else if is_sort_preserving_merge(&requirements.plan) && child_node.plan.output_partitioning().partition_count() <= 1 { - // This `SortPreservingMergeExec` is unnecessary, input already has a - // single partition and no fetch is required. - let mut child_node = requirements.children.swap_remove(0); + // This `SortPreservingMergeExec` is unnecessary because its input has a + // single partition. + let child_node = requirements.children.swap_remove(0); if let Some(fetch) = requirements.plan.fetch() { - // Add the limit exec if the original SPM had a fetch: - child_node.plan = - Arc::new(LocalLimitExec::new(Arc::clone(&child_node.plan), fetch)); + let mut limit = LocalLimitExec::new(Arc::clone(&child_node.plan), fetch); + limit.set_required_ordering(requirements.plan.output_ordering().cloned()); + return Ok(Transformed::yes(PlanContext::new( + Arc::new(limit), + false, + vec![child_node], + ))); } return Ok(Transformed::yes(child_node)); } diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index eb358b10b4bfd..b959a2823a723 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -699,20 +699,16 @@ impl ExecutionPlan for SymmetricHashJoinExec { }) .transpose()?; let expr_ctx = ctx.expr_ctx(); - let encode_sort_exprs = - |exprs: Option<&LexOrdering>| -> Result> { - exprs.map_or_else( - || Ok(vec![]), - |exprs| { - datafusion_physical_expr_common::sort_expr::sort_exprs_try_to_proto( - exprs.iter(), - &expr_ctx, - ) - }, - ) - }; - let left_sort_exprs = encode_sort_exprs(self.left_sort_exprs())?; - let right_sort_exprs = encode_sort_exprs(self.right_sort_exprs())?; + let left_sort_exprs = + datafusion_physical_expr_common::sort_expr::optional_ordering_try_to_proto( + self.left_sort_exprs(), + &expr_ctx, + )?; + let right_sort_exprs = + datafusion_physical_expr_common::sort_expr::optional_ordering_try_to_proto( + self.right_sort_exprs(), + &expr_ctx, + )?; Ok(Some(protobuf::PhysicalPlanNode { physical_plan_type: Some( @@ -876,20 +872,16 @@ impl SymmetricHashJoinExec { )) }) .transpose()?; - let decode_sort_exprs = |sort_exprs: &[protobuf::PhysicalSortExprNode], - schema: &Schema| - -> Result> { - let sort_exprs = - datafusion_physical_expr_common::sort_expr::sort_exprs_try_from_proto( - sort_exprs, - &ctx.expr_ctx(schema), - )?; - Ok(LexOrdering::new(sort_exprs)) - }; let left_sort_exprs = - decode_sort_exprs(&sym_join.left_sort_exprs, left_schema.as_ref())?; + datafusion_physical_expr_common::sort_expr::optional_ordering_try_from_proto( + &sym_join.left_sort_exprs, + &ctx.expr_ctx(left_schema.as_ref()), + )?; let right_sort_exprs = - decode_sort_exprs(&sym_join.right_sort_exprs, right_schema.as_ref())?; + datafusion_physical_expr_common::sort_expr::optional_ordering_try_from_proto( + &sym_join.right_sort_exprs, + &ctx.expr_ctx(right_schema.as_ref()), + )?; Self::try_new( left, diff --git a/datafusion/physical-plan/src/limit.rs b/datafusion/physical-plan/src/limit.rs index ddce680fc18ad..68f3b77d89def 100644 --- a/datafusion/physical-plan/src/limit.rs +++ b/datafusion/physical-plan/src/limit.rs @@ -35,7 +35,7 @@ use crate::{ use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::{Result, assert_eq_or_internal_err, internal_err}; +use datafusion_common::{Result, assert_eq_or_internal_err}; use datafusion_execution::TaskContext; use datafusion_physical_expr::LexOrdering; @@ -54,8 +54,8 @@ pub struct GlobalLimitExec { fetch: Option, /// Execution metrics metrics: ExecutionPlanMetricsSet, - /// Does the limit have to preserve the order of its input, and if so what is it? - /// Some optimizations may reorder the input if no particular sort is required + /// Input ordering that must be preserved so limit pushdown does not change + /// which rows are returned. required_ordering: Option, cache: Arc, } @@ -172,11 +172,10 @@ impl ExecutionPlan for GlobalLimitExec { mut children: Vec>, ) -> Result> { check_if_same_properties!(self, children); - Ok(Arc::new(GlobalLimitExec::new( - children.swap_remove(0), - self.skip, - self.fetch, - ))) + let mut new_limit = + GlobalLimitExec::new(children.swap_remove(0), self.skip, self.fetch); + new_limit.set_required_ordering(self.required_ordering.clone()); + Ok(Arc::new(new_limit)) } fn with_new_children_and_same_properties( @@ -250,8 +249,13 @@ impl ExecutionPlan for GlobalLimitExec { &self, ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, ) -> Result> { + use datafusion_physical_expr_common::sort_expr::optional_ordering_try_to_proto; use datafusion_proto_models::protobuf; let input = ctx.encode_child(self.input())?; + let required_ordering = optional_ordering_try_to_proto( + self.required_ordering.as_ref(), + &ctx.expr_ctx(), + )?; Ok(Some(protobuf::PhysicalPlanNode { physical_plan_type: Some( protobuf::physical_plan_node::PhysicalPlanType::GlobalLimit(Box::new( @@ -262,6 +266,7 @@ impl ExecutionPlan for GlobalLimitExec { Some(n) => n as i64, _ => -1, // no limit }, + required_ordering, }, )), ), @@ -275,6 +280,7 @@ impl GlobalLimitExec { node: &datafusion_proto_models::protobuf::PhysicalPlanNode, ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, ) -> Result> { + use datafusion_physical_expr_common::sort_expr::optional_ordering_try_from_proto; use datafusion_proto_models::protobuf; let limit = crate::expect_plan_variant!( node, @@ -291,11 +297,13 @@ impl GlobalLimitExec { } else { None }; - Ok(Arc::new(GlobalLimitExec::new( - input, - limit.skip as usize, - fetch, - ))) + let required_ordering = optional_ordering_try_from_proto( + &limit.required_ordering, + &ctx.expr_ctx(input.schema().as_ref()), + )?; + let mut exec = GlobalLimitExec::new(input, limit.skip as usize, fetch); + exec.set_required_ordering(required_ordering); + Ok(Arc::new(exec)) } } @@ -308,8 +316,8 @@ pub struct LocalLimitExec { fetch: usize, /// Execution metrics metrics: ExecutionPlanMetricsSet, - /// If the child plan is a sort node, after the sort node is removed during - /// physical optimization, we should add the required ordering to the limit node + /// Input ordering that must be preserved so limit pushdown does not change + /// which rows are returned. required_ordering: Option, cache: Arc, } @@ -400,16 +408,12 @@ impl ExecutionPlan for LocalLimitExec { fn with_new_children( self: Arc, - children: Vec>, + mut children: Vec>, ) -> Result> { check_if_same_properties!(self, children); - match children.len() { - 1 => Ok(Arc::new(LocalLimitExec::new( - Arc::clone(&children[0]), - self.fetch, - ))), - _ => internal_err!("LocalLimitExec wrong number of children"), - } + let mut new_limit = LocalLimitExec::new(children.swap_remove(0), self.fetch); + new_limit.set_required_ordering(self.required_ordering.clone()); + Ok(Arc::new(new_limit)) } fn with_new_children_and_same_properties( @@ -478,14 +482,20 @@ impl ExecutionPlan for LocalLimitExec { &self, ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, ) -> Result> { + use datafusion_physical_expr_common::sort_expr::optional_ordering_try_to_proto; use datafusion_proto_models::protobuf; let input = ctx.encode_child(self.input())?; + let required_ordering = optional_ordering_try_to_proto( + self.required_ordering.as_ref(), + &ctx.expr_ctx(), + )?; Ok(Some(protobuf::PhysicalPlanNode { physical_plan_type: Some( protobuf::physical_plan_node::PhysicalPlanType::LocalLimit(Box::new( protobuf::LocalLimitExecNode { input: Some(Box::new(input)), fetch: self.fetch() as u32, + required_ordering, }, )), ), @@ -499,6 +509,7 @@ impl LocalLimitExec { node: &datafusion_proto_models::protobuf::PhysicalPlanNode, ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, ) -> Result> { + use datafusion_physical_expr_common::sort_expr::optional_ordering_try_from_proto; use datafusion_proto_models::protobuf; let limit = crate::expect_plan_variant!( node, @@ -507,7 +518,13 @@ impl LocalLimitExec { ); let input = ctx.decode_required_child(limit.input.as_deref(), "LocalLimitExec", "input")?; - Ok(Arc::new(LocalLimitExec::new(input, limit.fetch as usize))) + let required_ordering = optional_ordering_try_from_proto( + &limit.required_ordering, + &ctx.expr_ctx(input.schema().as_ref()), + )?; + let mut exec = LocalLimitExec::new(input, limit.fetch as usize); + exec.set_required_ordering(required_ordering); + Ok(Arc::new(exec)) } } @@ -646,10 +663,11 @@ mod tests { use crate::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; use arrow::array::RecordBatchOptions; + use arrow::compute::SortOptions; use arrow::datatypes::Schema; use datafusion_common::stats::Precision; - use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::expressions::col; + use datafusion_physical_expr::{PhysicalExpr, PhysicalSortExpr}; #[tokio::test] async fn limit() -> Result<()> { @@ -839,6 +857,35 @@ mod tests { Ok(()) } + #[test] + fn with_new_children_preserves_required_ordering() -> Result<()> { + let source = test::scan_partitioned(1); + let schema = source.schema(); + let ordering = LexOrdering::new(vec![PhysicalSortExpr { + expr: col("i", &schema)?, + options: SortOptions { + descending: true, + nulls_first: false, + }, + }]); + + let mut global = GlobalLimitExec::new(Arc::clone(&source), 0, Some(10)); + global.set_required_ordering(ordering.clone()); + let rebuilt = + Arc::new(global).with_new_children(vec![test::scan_partitioned(1)])?; + let rebuilt = rebuilt.downcast_ref::().unwrap(); + assert_eq!(rebuilt.required_ordering(), &ordering); + + let mut local = LocalLimitExec::new(source, 10); + local.set_required_ordering(ordering.clone()); + let rebuilt = + Arc::new(local).with_new_children(vec![test::scan_partitioned(1)])?; + let rebuilt = rebuilt.downcast_ref::().unwrap(); + assert_eq!(rebuilt.required_ordering(), &ordering); + + Ok(()) + } + #[test] fn test_row_number_statistics_for_global_limit() -> Result<()> { let row_count = row_number_statistics_for_global_limit(0, Some(10))?; diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index cbc41a7c5713e..06cd74990f4fa 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1472,11 +1472,15 @@ message GlobalLimitExecNode { uint32 skip = 2; // Maximum number of rows to fetch; negative means no limit int64 fetch = 3; + // Ordering the limit must preserve; empty means none + repeated PhysicalSortExprNode required_ordering = 4; } message LocalLimitExecNode { PhysicalPlanNode input = 1; uint32 fetch = 2; + // Ordering the limit must preserve; empty means none + repeated PhysicalSortExprNode required_ordering = 3; } message SortExecNode { diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 7f9b9eddc5ff5..d52e9aa227fbf 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -8783,6 +8783,9 @@ impl serde::Serialize for GlobalLimitExecNode { if self.fetch != 0 { len += 1; } + if !self.required_ordering.is_empty() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.GlobalLimitExecNode", len)?; if let Some(v) = self.input.as_ref() { struct_ser.serialize_field("input", v)?; @@ -8795,6 +8798,9 @@ impl serde::Serialize for GlobalLimitExecNode { #[allow(clippy::needless_borrows_for_generic_args)] struct_ser.serialize_field("fetch", ToString::to_string(&self.fetch).as_str())?; } + if !self.required_ordering.is_empty() { + struct_ser.serialize_field("requiredOrdering", &self.required_ordering)?; + } struct_ser.end() } } @@ -8808,6 +8814,8 @@ impl<'de> serde::Deserialize<'de> for GlobalLimitExecNode { "input", "skip", "fetch", + "required_ordering", + "requiredOrdering", ]; #[allow(clippy::enum_variant_names)] @@ -8815,6 +8823,7 @@ impl<'de> serde::Deserialize<'de> for GlobalLimitExecNode { Input, Skip, Fetch, + RequiredOrdering, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -8839,6 +8848,7 @@ impl<'de> serde::Deserialize<'de> for GlobalLimitExecNode { "input" => Ok(GeneratedField::Input), "skip" => Ok(GeneratedField::Skip), "fetch" => Ok(GeneratedField::Fetch), + "requiredOrdering" | "required_ordering" => Ok(GeneratedField::RequiredOrdering), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -8861,6 +8871,7 @@ impl<'de> serde::Deserialize<'de> for GlobalLimitExecNode { let mut input__ = None; let mut skip__ = None; let mut fetch__ = None; + let mut required_ordering__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Input => { @@ -8885,12 +8896,19 @@ impl<'de> serde::Deserialize<'de> for GlobalLimitExecNode { Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) ; } + GeneratedField::RequiredOrdering => { + if required_ordering__.is_some() { + return Err(serde::de::Error::duplicate_field("requiredOrdering")); + } + required_ordering__ = Some(map_.next_value()?); + } } } Ok(GlobalLimitExecNode { input: input__, skip: skip__.unwrap_or_default(), fetch: fetch__.unwrap_or_default(), + required_ordering: required_ordering__.unwrap_or_default(), }) } } @@ -12593,6 +12611,9 @@ impl serde::Serialize for LocalLimitExecNode { if self.fetch != 0 { len += 1; } + if !self.required_ordering.is_empty() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.LocalLimitExecNode", len)?; if let Some(v) = self.input.as_ref() { struct_ser.serialize_field("input", v)?; @@ -12600,6 +12621,9 @@ impl serde::Serialize for LocalLimitExecNode { if self.fetch != 0 { struct_ser.serialize_field("fetch", &self.fetch)?; } + if !self.required_ordering.is_empty() { + struct_ser.serialize_field("requiredOrdering", &self.required_ordering)?; + } struct_ser.end() } } @@ -12612,12 +12636,15 @@ impl<'de> serde::Deserialize<'de> for LocalLimitExecNode { const FIELDS: &[&str] = &[ "input", "fetch", + "required_ordering", + "requiredOrdering", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { Input, Fetch, + RequiredOrdering, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -12641,6 +12668,7 @@ impl<'de> serde::Deserialize<'de> for LocalLimitExecNode { match value { "input" => Ok(GeneratedField::Input), "fetch" => Ok(GeneratedField::Fetch), + "requiredOrdering" | "required_ordering" => Ok(GeneratedField::RequiredOrdering), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -12662,6 +12690,7 @@ impl<'de> serde::Deserialize<'de> for LocalLimitExecNode { { let mut input__ = None; let mut fetch__ = None; + let mut required_ordering__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Input => { @@ -12678,11 +12707,18 @@ impl<'de> serde::Deserialize<'de> for LocalLimitExecNode { Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) ; } + GeneratedField::RequiredOrdering => { + if required_ordering__.is_some() { + return Err(serde::de::Error::duplicate_field("requiredOrdering")); + } + required_ordering__ = Some(map_.next_value()?); + } } } Ok(LocalLimitExecNode { input: input__, fetch: fetch__.unwrap_or_default(), + required_ordering: required_ordering__.unwrap_or_default(), }) } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index f7633483080f1..3c3c973efe229 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -2232,6 +2232,9 @@ pub struct GlobalLimitExecNode { /// Maximum number of rows to fetch; negative means no limit #[prost(int64, tag = "3")] pub fetch: i64, + /// Ordering the limit must preserve; empty means none + #[prost(message, repeated, tag = "4")] + pub required_ordering: ::prost::alloc::vec::Vec, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct LocalLimitExecNode { @@ -2239,6 +2242,9 @@ pub struct LocalLimitExecNode { pub input: ::core::option::Option<::prost::alloc::boxed::Box>, #[prost(uint32, tag = "2")] pub fetch: u32, + /// Ordering the limit must preserve; empty means none + #[prost(message, repeated, tag = "3")] + pub required_ordering: ::prost::alloc::vec::Vec, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct SortExecNode { diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 19a5ca337d7f6..1830a0a4168ba 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -56,6 +56,7 @@ use datafusion::physical_expr::{ }; use datafusion::physical_optimizer::PhysicalOptimizerRule; use datafusion::physical_optimizer::filter_pushdown::FilterPushdown; +use datafusion::physical_optimizer::limit_pushdown::LimitPushdown; use datafusion::physical_plan::aggregates::{ AggregateExec, AggregateMode, LimitOptions, PhysicalGroupBy, }; @@ -439,6 +440,108 @@ fn roundtrip_global_skip_no_limit() -> Result<()> { ))) } +/// Sort key at index 1, so a decoder that misbinds column name vs index +/// cannot pass. +fn limit_test_schema() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Int64, false), + ])) +} + +/// Non-default sort options, so a decode that falls back to defaults cannot +/// pass. +fn limit_required_ordering(schema: &Schema) -> Result> { + Ok(LexOrdering::new(vec![PhysicalSortExpr { + expr: col("b", schema)?, + options: SortOptions { + descending: true, + nulls_first: false, + }, + }])) +} + +#[test] +fn roundtrip_limit_with_required_ordering() -> Result<()> { + let schema = limit_test_schema(); + let required_ordering = limit_required_ordering(&schema)?; + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + + let mut global = + GlobalLimitExec::new(Arc::new(EmptyExec::new(Arc::clone(&schema))), 3, Some(25)); + global.set_required_ordering(required_ordering.clone()); + let decoded = + roundtrip_test_and_return(Arc::new(global), &ctx, &codec, &proto_converter)?; + let decoded = decoded + .downcast_ref::() + .expect("expected GlobalLimitExec"); + assert_eq!(decoded.required_ordering(), &required_ordering); + + let mut local = LocalLimitExec::new(Arc::new(EmptyExec::new(schema)), 25); + local.set_required_ordering(required_ordering.clone()); + let decoded = + roundtrip_test_and_return(Arc::new(local), &ctx, &codec, &proto_converter)?; + let decoded = decoded + .downcast_ref::() + .expect("expected LocalLimitExec"); + assert_eq!(decoded.required_ordering(), &required_ordering); + Ok(()) +} + +/// A limit's `required_ordering` is the only record that an `ORDER BY ... LIMIT` +/// whose sort node was optimized away is order-sensitive, so it must survive +/// serde all the way into the scan's `preserve_order` flag. +#[test] +fn roundtrip_limit_required_ordering_reaches_data_source() -> Result<()> { + let file_schema = limit_test_schema(); + let make_scan = || { + let file_source = Arc::new(ParquetSource::new(Arc::clone(&file_schema))); + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( + "/path/to/file.parquet".to_string(), + 1024, + )])]) + .build(); + DataSourceExec::from_data_source(scan_config) + }; + let scan_after_limit_pushdown = |limit: GlobalLimitExec| -> Result { + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let decoded = + roundtrip_test_and_return(Arc::new(limit), &ctx, &codec, &proto_converter)?; + + // Child replacement must not erase the decoded ordering before pushdown. + let rebuilt = decoded.with_new_children(vec![make_scan()])?; + + let optimized = + LimitPushdown::new().optimize(rebuilt, &ConfigOptions::default())?; + let scan = optimized + .downcast_ref::() + .expect("limit should be absorbed into the scan"); + Ok(scan + .data_source() + .downcast_ref::() + .expect("expected FileScanConfig") + .clone()) + }; + + let mut limit = GlobalLimitExec::new(make_scan(), 0, Some(10)); + limit.set_required_ordering(limit_required_ordering(&file_schema)?); + let scan_config = scan_after_limit_pushdown(limit)?; + assert_eq!(scan_config.limit, Some(10)); + assert!(scan_config.preserve_order); + + let scan_config = + scan_after_limit_pushdown(GlobalLimitExec::new(make_scan(), 0, Some(10)))?; + assert_eq!(scan_config.limit, Some(10)); + assert!(!scan_config.preserve_order); + Ok(()) +} + #[test] fn roundtrip_hash_join() -> Result<()> { let field_a = Field::new("col", DataType::Int64, false);