From f1be5716b70f3da19b615e039c2be65095319d2c Mon Sep 17 00:00:00 2001 From: buraksenn Date: Sat, 8 Aug 2026 14:32:05 +0300 Subject: [PATCH 1/4] fix: serialize Global/LocalLimitExec required_ordering --- datafusion/physical-plan/src/limit.rs | 109 ++++++++++++++--- .../proto-models/proto/datafusion.proto | 4 + .../proto-models/src/generated/pbjson.rs | 36 ++++++ .../proto-models/src/generated/prost.rs | 6 + .../tests/cases/roundtrip_physical_plan.rs | 110 ++++++++++++++++++ 5 files changed, 249 insertions(+), 16 deletions(-) diff --git a/datafusion/physical-plan/src/limit.rs b/datafusion/physical-plan/src/limit.rs index ddce680fc18ad..356ce16ca11ee 100644 --- a/datafusion/physical-plan/src/limit.rs +++ b/datafusion/physical-plan/src/limit.rs @@ -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( @@ -252,6 +251,7 @@ impl ExecutionPlan for GlobalLimitExec { ) -> Result> { use datafusion_proto_models::protobuf; let input = ctx.encode_child(self.input())?; + let required_ordering = encode_required_ordering(self.required_ordering(), ctx)?; Ok(Some(protobuf::PhysicalPlanNode { physical_plan_type: Some( protobuf::physical_plan_node::PhysicalPlanType::GlobalLimit(Box::new( @@ -262,6 +262,7 @@ impl ExecutionPlan for GlobalLimitExec { Some(n) => n as i64, _ => -1, // no limit }, + required_ordering, }, )), ), @@ -291,14 +292,49 @@ impl GlobalLimitExec { } else { None }; - Ok(Arc::new(GlobalLimitExec::new( - input, - limit.skip as usize, - fetch, - ))) + let required_ordering = + decode_required_ordering(&limit.required_ordering, &input, ctx)?; + let mut exec = GlobalLimitExec::new(input, limit.skip as usize, fetch); + exec.set_required_ordering(required_ordering); + Ok(Arc::new(exec)) } } +/// Serialize a limit's `required_ordering` as a flat sort-expression list; +/// `None` maps to an empty list. +#[cfg(feature = "proto")] +fn encode_required_ordering( + required_ordering: &Option, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, +) -> Result> { + required_ordering.as_ref().map_or_else( + || Ok(vec![]), + |ordering| { + datafusion_physical_expr_common::sort_expr::sort_exprs_try_to_proto( + ordering.iter(), + &ctx.expr_ctx(), + ) + }, + ) +} + +/// Reconstruct a limit's `required_ordering` against the input's schema; an +/// empty list maps to `None`. +#[cfg(feature = "proto")] +fn decode_required_ordering( + nodes: &[datafusion_proto_models::protobuf::PhysicalSortExprNode], + input: &Arc, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, +) -> Result> { + let input_schema = input.schema(); + let sort_exprs = + datafusion_physical_expr_common::sort_expr::sort_exprs_try_from_proto( + nodes, + &ctx.expr_ctx(input_schema.as_ref()), + )?; + Ok(LexOrdering::new(sort_exprs)) +} + /// LocalLimitExec applies a limit to a single partition #[derive(Debug, Clone)] pub struct LocalLimitExec { @@ -404,10 +440,12 @@ impl ExecutionPlan for LocalLimitExec { ) -> Result> { check_if_same_properties!(self, children); match children.len() { - 1 => Ok(Arc::new(LocalLimitExec::new( - Arc::clone(&children[0]), - self.fetch, - ))), + 1 => { + let mut new_limit = + LocalLimitExec::new(Arc::clone(&children[0]), self.fetch); + new_limit.set_required_ordering(self.required_ordering.clone()); + Ok(Arc::new(new_limit)) + } _ => internal_err!("LocalLimitExec wrong number of children"), } } @@ -480,12 +518,14 @@ impl ExecutionPlan for LocalLimitExec { ) -> Result> { use datafusion_proto_models::protobuf; let input = ctx.encode_child(self.input())?; + let required_ordering = encode_required_ordering(self.required_ordering(), 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, }, )), ), @@ -507,7 +547,11 @@ 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 = + decode_required_ordering(&limit.required_ordering, &input, ctx)?; + let mut exec = LocalLimitExec::new(input, limit.fetch as usize); + exec.set_required_ordering(required_ordering); + Ok(Arc::new(exec)) } } @@ -646,10 +690,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 +884,38 @@ 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, + }, + }]); + + // A fresh child never shares the old child's properties `Arc`, so + // these rebuilds take the reconstruction path rather than the + // same-properties fast path. + 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..2656df17594a6 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,115 @@ fn roundtrip_global_skip_no_limit() -> Result<()> { ))) } +/// A single-column ordering with non-default sort options, so a decode that +/// falls back to defaults cannot pass the roundtrip assertions. +fn limit_required_ordering(schema: &Schema) -> Result> { + Ok(LexOrdering::new(vec![PhysicalSortExpr { + expr: col("a", schema)?, + options: SortOptions { + descending: true, + nulls_first: false, + }, + }])) +} + +#[test] +fn roundtrip_local_limit_with_required_ordering() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let required_ordering = limit_required_ordering(&schema)?; + let mut limit = LocalLimitExec::new(Arc::new(EmptyExec::new(schema)), 25); + limit.set_required_ordering(required_ordering.clone()); + + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let result = + roundtrip_test_and_return(Arc::new(limit), &ctx, &codec, &proto_converter)?; + let result = result.downcast_ref::().unwrap(); + assert_eq!(result.required_ordering(), &required_ordering); + Ok(()) +} + +#[test] +fn roundtrip_global_limit_with_required_ordering() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let required_ordering = limit_required_ordering(&schema)?; + let mut limit = GlobalLimitExec::new(Arc::new(EmptyExec::new(schema)), 3, Some(25)); + limit.set_required_ordering(required_ordering.clone()); + + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let result = + roundtrip_test_and_return(Arc::new(limit), &ctx, &codec, &proto_converter)?; + let result = result.downcast_ref::().unwrap(); + assert_eq!(result.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. Exercise the full +/// consumer path: decode a plan that still carries the limit, rebuild the +/// limit's child as optimizer passes do, run `LimitPushdown`, and check that +/// the scan comes out order-preserving. +#[test] +fn roundtrip_limit_required_ordering_reaches_data_source() -> Result<()> { + let file_schema = + Arc::new(Schema::new(vec![Field::new("col", DataType::Int64, false)])); + 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)?; + + // An optimizer pass that changes the child subtree rebuilds the + // limit around the new child. + 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(LexOrdering::new(vec![PhysicalSortExpr { + expr: Arc::new(Column::new("col", 0)), + options: SortOptions { + descending: true, + nulls_first: false, + }, + }])); + let scan_config = scan_after_limit_pushdown(limit)?; + assert_eq!(scan_config.limit, Some(10)); + assert!(scan_config.preserve_order); + + // Control: with no required ordering the scan must stay order-free. + 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); From a9604496590f2df5c3a4882ef2d0569499fb6973 Mon Sep 17 00:00:00 2001 From: buraksenn Date: Sat, 8 Aug 2026 16:17:56 +0300 Subject: [PATCH 2/4] second pass for ordering serialization --- .../physical-expr-common/src/sort_expr.rs | 22 ++++- .../src/joins/symmetric_hash_join.rs | 44 +++++----- datafusion/physical-plan/src/limit.rs | 81 +++++++------------ .../tests/cases/roundtrip_physical_plan.rs | 73 +++++++---------- 4 files changed, 93 insertions(+), 127 deletions(-) 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-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 356ce16ca11ee..9000083290e9b 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; @@ -249,9 +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 = encode_required_ordering(self.required_ordering(), ctx)?; + 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( @@ -276,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, @@ -292,49 +297,16 @@ impl GlobalLimitExec { } else { None }; - let required_ordering = - decode_required_ordering(&limit.required_ordering, &input, ctx)?; + 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)) } } -/// Serialize a limit's `required_ordering` as a flat sort-expression list; -/// `None` maps to an empty list. -#[cfg(feature = "proto")] -fn encode_required_ordering( - required_ordering: &Option, - ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, -) -> Result> { - required_ordering.as_ref().map_or_else( - || Ok(vec![]), - |ordering| { - datafusion_physical_expr_common::sort_expr::sort_exprs_try_to_proto( - ordering.iter(), - &ctx.expr_ctx(), - ) - }, - ) -} - -/// Reconstruct a limit's `required_ordering` against the input's schema; an -/// empty list maps to `None`. -#[cfg(feature = "proto")] -fn decode_required_ordering( - nodes: &[datafusion_proto_models::protobuf::PhysicalSortExprNode], - input: &Arc, - ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, -) -> Result> { - let input_schema = input.schema(); - let sort_exprs = - datafusion_physical_expr_common::sort_expr::sort_exprs_try_from_proto( - nodes, - &ctx.expr_ctx(input_schema.as_ref()), - )?; - Ok(LexOrdering::new(sort_exprs)) -} - /// LocalLimitExec applies a limit to a single partition #[derive(Debug, Clone)] pub struct LocalLimitExec { @@ -436,18 +408,13 @@ impl ExecutionPlan for LocalLimitExec { fn with_new_children( self: Arc, - children: Vec>, + mut children: Vec>, ) -> Result> { + // `check_if_same_properties!` has already verified the child count. check_if_same_properties!(self, children); - match children.len() { - 1 => { - let mut new_limit = - LocalLimitExec::new(Arc::clone(&children[0]), self.fetch); - new_limit.set_required_ordering(self.required_ordering.clone()); - Ok(Arc::new(new_limit)) - } - _ => 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( @@ -516,9 +483,13 @@ 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 = encode_required_ordering(self.required_ordering(), ctx)?; + 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( @@ -539,6 +510,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, @@ -547,8 +519,10 @@ impl LocalLimitExec { ); let input = ctx.decode_required_child(limit.input.as_deref(), "LocalLimitExec", "input")?; - let required_ordering = - decode_required_ordering(&limit.required_ordering, &input, ctx)?; + 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)) @@ -896,8 +870,7 @@ mod tests { }, }]); - // A fresh child never shares the old child's properties `Arc`, so - // these rebuilds take the reconstruction path rather than the + // Fresh children force the reconstruction path rather than the // same-properties fast path. let mut global = GlobalLimitExec::new(Arc::clone(&source), 0, Some(10)); global.set_required_ordering(ordering.clone()); diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 2656df17594a6..ec3d4109725c6 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -440,11 +440,20 @@ fn roundtrip_global_skip_no_limit() -> Result<()> { ))) } -/// A single-column ordering with non-default sort options, so a decode that -/// falls back to defaults cannot pass the roundtrip assertions. +/// 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("a", schema)?, + expr: col("b", schema)?, options: SortOptions { descending: true, nulls_first: false, @@ -453,48 +462,27 @@ fn limit_required_ordering(schema: &Schema) -> Result> { } #[test] -fn roundtrip_local_limit_with_required_ordering() -> Result<()> { - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); +fn roundtrip_limit_with_required_ordering() -> Result<()> { + let schema = limit_test_schema(); let required_ordering = limit_required_ordering(&schema)?; - let mut limit = LocalLimitExec::new(Arc::new(EmptyExec::new(schema)), 25); - limit.set_required_ordering(required_ordering.clone()); - let ctx = SessionContext::new(); - let codec = DefaultPhysicalExtensionCodec {}; - let proto_converter = DefaultPhysicalProtoConverter {}; - let result = - roundtrip_test_and_return(Arc::new(limit), &ctx, &codec, &proto_converter)?; - let result = result.downcast_ref::().unwrap(); - assert_eq!(result.required_ordering(), &required_ordering); - Ok(()) -} + // `roundtrip_test`'s `Debug` equality covers `required_ordering`. + let mut global = + GlobalLimitExec::new(Arc::new(EmptyExec::new(Arc::clone(&schema))), 3, Some(25)); + global.set_required_ordering(required_ordering.clone()); + roundtrip_test(Arc::new(global))?; -#[test] -fn roundtrip_global_limit_with_required_ordering() -> Result<()> { - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - let required_ordering = limit_required_ordering(&schema)?; - let mut limit = GlobalLimitExec::new(Arc::new(EmptyExec::new(schema)), 3, Some(25)); - limit.set_required_ordering(required_ordering.clone()); - - let ctx = SessionContext::new(); - let codec = DefaultPhysicalExtensionCodec {}; - let proto_converter = DefaultPhysicalProtoConverter {}; - let result = - roundtrip_test_and_return(Arc::new(limit), &ctx, &codec, &proto_converter)?; - let result = result.downcast_ref::().unwrap(); - assert_eq!(result.required_ordering(), &required_ordering); - Ok(()) + let mut local = LocalLimitExec::new(Arc::new(EmptyExec::new(schema)), 25); + local.set_required_ordering(required_ordering); + roundtrip_test(Arc::new(local)) } /// A limit's `required_ordering` is the only record that an `ORDER BY ... LIMIT` -/// whose sort node was optimized away is order-sensitive. Exercise the full -/// consumer path: decode a plan that still carries the limit, rebuild the -/// limit's child as optimizer passes do, run `LimitPushdown`, and check that -/// the scan comes out order-preserving. +/// 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 = - Arc::new(Schema::new(vec![Field::new("col", DataType::Int64, false)])); + let file_schema = limit_test_schema(); let make_scan = || { let file_source = Arc::new(ParquetSource::new(Arc::clone(&file_schema))); let scan_config = @@ -513,8 +501,7 @@ fn roundtrip_limit_required_ordering_reaches_data_source() -> Result<()> { let decoded = roundtrip_test_and_return(Arc::new(limit), &ctx, &codec, &proto_converter)?; - // An optimizer pass that changes the child subtree rebuilds the - // limit around the new child. + // Rebuild the limit as an optimizer pass replacing its child would. let rebuilt = decoded.with_new_children(vec![make_scan()])?; let optimized = @@ -530,13 +517,7 @@ fn roundtrip_limit_required_ordering_reaches_data_source() -> Result<()> { }; let mut limit = GlobalLimitExec::new(make_scan(), 0, Some(10)); - limit.set_required_ordering(LexOrdering::new(vec![PhysicalSortExpr { - expr: Arc::new(Column::new("col", 0)), - options: SortOptions { - descending: true, - nulls_first: false, - }, - }])); + 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); From 3f351f269be6648ee36bef46441403be0ad918ac Mon Sep 17 00:00:00 2001 From: buraksenn Date: Sat, 8 Aug 2026 22:13:24 +0300 Subject: [PATCH 3/4] enf-fix spm --- .../physical_optimizer/enforce_sorting.rs | 73 ++++++++++++++++++- .../enforce_sorting/mod.rs | 17 +++-- 2 files changed, 83 insertions(+), 7 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index d94253a84aa5f..9fb3bbfda06ff 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,75 @@ async fn test_remove_unnecessary_spm2() -> Result<()> { Ok(()) } +#[test] +fn test_spm_fetch_replacement_context_integrity() -> 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 plan = sort_preserving_merge_exec_with_fetch(ordering.clone(), source, 100); + + let optimized = PlanWithCorrespondingSort::new_default(plan) + .transform_up(ensure_sorting)? + .data; + let optimized = check_integrity(optimized)?; + let limit = optimized + .plan + .downcast_ref::() + .expect("SPM fetch should become a local limit"); + assert_eq!(limit.fetch(), 100); + assert_eq!(limit.required_ordering().as_ref(), Some(&ordering)); + + 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.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-optimizer/src/ensure_requirements/enforce_sorting/mod.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs index 6efaf76457919..aa40ed7fa4106 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,18 @@ 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)); + // Preserve both the fetch and its ordering requirement. + 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)); } From e6a8172b486787d1f8d4c014c16864f67f7d2def Mon Sep 17 00:00:00 2001 From: buraksenn Date: Sat, 8 Aug 2026 22:27:59 +0300 Subject: [PATCH 4/4] adjust comments --- .../physical_optimizer/enforce_sorting.rs | 22 +---------------- .../enforce_sorting/mod.rs | 1 - datafusion/physical-plan/src/limit.rs | 11 ++++----- .../tests/cases/roundtrip_physical_plan.rs | 24 ++++++++++++++----- 4 files changed, 23 insertions(+), 35 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index 9fb3bbfda06ff..a8162f137ed0a 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs @@ -2299,27 +2299,6 @@ async fn test_remove_unnecessary_spm2() -> Result<()> { Ok(()) } -#[test] -fn test_spm_fetch_replacement_context_integrity() -> 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 plan = sort_preserving_merge_exec_with_fetch(ordering.clone(), source, 100); - - let optimized = PlanWithCorrespondingSort::new_default(plan) - .transform_up(ensure_sorting)? - .data; - let optimized = check_integrity(optimized)?; - let limit = optimized - .plan - .downcast_ref::() - .expect("SPM fetch should become a local limit"); - assert_eq!(limit.fetch(), 100); - assert_eq!(limit.required_ordering().as_ref(), Some(&ordering)); - - Ok(()) -} - #[test] fn test_spm_fetch_preserves_ordering_through_child_rewrite() -> Result<()> { let schema = create_test_schema()?; @@ -2344,6 +2323,7 @@ fn test_spm_fetch_preserves_ordering_through_child_rewrite() -> Result<()> { 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(); 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 aa40ed7fa4106..c66d5310a1c44 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs @@ -465,7 +465,6 @@ pub fn ensure_sorting( // single partition. let child_node = requirements.children.swap_remove(0); if let Some(fetch) = requirements.plan.fetch() { - // Preserve both the fetch and its ordering requirement. 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( diff --git a/datafusion/physical-plan/src/limit.rs b/datafusion/physical-plan/src/limit.rs index 9000083290e9b..68f3b77d89def 100644 --- a/datafusion/physical-plan/src/limit.rs +++ b/datafusion/physical-plan/src/limit.rs @@ -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, } @@ -316,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, } @@ -410,7 +410,6 @@ impl ExecutionPlan for LocalLimitExec { self: Arc, mut children: Vec>, ) -> Result> { - // `check_if_same_properties!` has already verified the child count. check_if_same_properties!(self, children); let mut new_limit = LocalLimitExec::new(children.swap_remove(0), self.fetch); new_limit.set_required_ordering(self.required_ordering.clone()); @@ -870,8 +869,6 @@ mod tests { }, }]); - // Fresh children force the reconstruction path rather than the - // same-properties fast path. let mut global = GlobalLimitExec::new(Arc::clone(&source), 0, Some(10)); global.set_required_ordering(ordering.clone()); let rebuilt = diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index ec3d4109725c6..1830a0a4168ba 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -465,16 +465,29 @@ fn limit_required_ordering(schema: &Schema) -> Result> { 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 {}; - // `roundtrip_test`'s `Debug` equality covers `required_ordering`. let mut global = GlobalLimitExec::new(Arc::new(EmptyExec::new(Arc::clone(&schema))), 3, Some(25)); global.set_required_ordering(required_ordering.clone()); - roundtrip_test(Arc::new(global))?; + 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); - roundtrip_test(Arc::new(local)) + 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` @@ -501,7 +514,7 @@ fn roundtrip_limit_required_ordering_reaches_data_source() -> Result<()> { let decoded = roundtrip_test_and_return(Arc::new(limit), &ctx, &codec, &proto_converter)?; - // Rebuild the limit as an optimizer pass replacing its child would. + // Child replacement must not erase the decoded ordering before pushdown. let rebuilt = decoded.with_new_children(vec![make_scan()])?; let optimized = @@ -522,7 +535,6 @@ fn roundtrip_limit_required_ordering_reaches_data_source() -> Result<()> { assert_eq!(scan_config.limit, Some(10)); assert!(scan_config.preserve_order); - // Control: with no required ordering the scan must stay order-free. let scan_config = scan_after_limit_pushdown(GlobalLimitExec::new(make_scan(), 0, Some(10)))?; assert_eq!(scan_config.limit, Some(10));