You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/08/10 12:09:26 UTC
(datafusion) 01/01: fix(proto): serialize Global/LocalLimitExec required_ordering (#24183)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24183-0f68e23bde46c19d2bc944132cc37bf9055c6f13
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit f4c8ba1e44e64e25ae69c477d73b62be350d17cb
Author: Burak Şen <bu...@gmail.com>
AuthorDate: Mon Aug 10 15:09:02 2026 +0300
fix(proto): serialize Global/LocalLimitExec required_ordering (#24183)
## Which issue does this PR close?
- Closes #24173
## Rationale for this change
`GlobalLimitExec`/`LocalLimitExec` `required_ordering` is not
serialized, so a roundtrip loses the only record that a pushed-down
`LIMIT` is order-sensitive. Re-running `LimitPushdown` on the decoded
plan then sets `preserve_order: false` on the scan, which is free to
read files out of order — `ORDER BY ... LIMIT` can return the wrong
rows.
## What changes are included in this PR?
- Add `required_ordering` to `GlobalLimitExecNode` (field 4) and
`LocalLimitExecNode` (field 3); empty list means `None`
- New
`optional_ordering_try_to_proto`/`optional_ordering_try_from_proto`
helpers in `physical-expr-common`, also reused by
`SymmetricHashJoinExec` serde which hand-rolled the same pattern
- Preserve `required_ordering` in both execs' `with_new_children`
## Are these changes tested?
Yes new tests that would fail on main if not fixed
## Are there any user-facing changes?
No
---
.../tests/physical_optimizer/enforce_sorting.rs | 53 ++++++++++-
datafusion/physical-expr-common/src/sort_expr.rs | 22 ++++-
.../src/ensure_requirements/enforce_sorting/mod.rs | 16 ++--
.../physical-plan/src/joins/symmetric_hash_join.rs | 44 ++++-----
datafusion/physical-plan/src/limit.rs | 96 ++++++++++++++-----
datafusion/proto-models/proto/datafusion.proto | 4 +
datafusion/proto-models/src/generated/pbjson.rs | 36 +++++++
datafusion/proto-models/src/generated/prost.rs | 6 ++
.../proto/tests/cases/roundtrip_physical_plan.rs | 103 +++++++++++++++++++++
9 files changed, 322 insertions(+), 58 deletions(-)
diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs
index d94253a84a..a8162f137e 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::<LocalLimitExec>()
+ .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::<LocalLimitExec>()
+ .expect("projection rewrite should retain the local limit");
+ assert_eq!(limit.required_ordering().as_ref(), Some(&ordering));
+ assert!(limit.input().is::<DataSourceExec>());
+
+ let optimized = LimitPushdown::new().optimize(optimized, &config)?;
+ let source = optimized
+ .downcast_ref::<DataSourceExec>()
+ .expect("limit should be pushed into the parquet scan");
+ let config = source
+ .data_source()
+ .downcast_ref::<FileScanConfig>()
+ .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 72e8772347..6e8dbccdb7 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<E: std::borrow::Borrow<PhysicalSortExpr>>(
/// [`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<Vec<datafusion_proto_models::protobuf::PhysicalSortExprNode>> {
+ 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<Option<LexOrdering>> {
+ 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 6efaf76457..c66d5310a1 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 86f721b711..589042a828 100644
--- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs
+++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs
@@ -709,20 +709,16 @@ impl ExecutionPlan for SymmetricHashJoinExec {
})
.transpose()?;
let expr_ctx = ctx.expr_ctx();
- let encode_sort_exprs =
- |exprs: Option<&LexOrdering>| -> Result<Vec<protobuf::PhysicalSortExprNode>> {
- 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(
@@ -886,20 +882,16 @@ impl SymmetricHashJoinExec {
))
})
.transpose()?;
- let decode_sort_exprs = |sort_exprs: &[protobuf::PhysicalSortExprNode],
- schema: &Schema|
- -> Result<Option<LexOrdering>> {
- 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 5bc5bc48f1..9dbdf17dbc 100644
--- a/datafusion/physical-plan/src/limit.rs
+++ b/datafusion/physical-plan/src/limit.rs
@@ -36,7 +36,7 @@ use crate::{
use arrow::datatypes::SchemaRef;
use arrow::record_batch::RecordBatch;
use datafusion_common::tree_node::TreeNodeRecursion;
-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, PhysicalExpr};
@@ -55,8 +55,8 @@ pub struct GlobalLimitExec {
fetch: Option<usize>,
/// 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<LexOrdering>,
cache: Arc<PlanProperties>,
}
@@ -180,11 +180,10 @@ impl ExecutionPlan for GlobalLimitExec {
mut children: Vec<Arc<dyn ExecutionPlan>>,
) -> Result<Arc<dyn ExecutionPlan>> {
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(
@@ -258,8 +257,13 @@ impl ExecutionPlan for GlobalLimitExec {
&self,
ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>,
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
+ 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(
@@ -270,6 +274,7 @@ impl ExecutionPlan for GlobalLimitExec {
Some(n) => n as i64,
_ => -1, // no limit
},
+ required_ordering,
},
)),
),
@@ -283,6 +288,7 @@ impl GlobalLimitExec {
node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>,
) -> Result<Arc<dyn ExecutionPlan>> {
+ use datafusion_physical_expr_common::sort_expr::optional_ordering_try_from_proto;
use datafusion_proto_models::protobuf;
let limit = crate::expect_plan_variant!(
node,
@@ -299,11 +305,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))
}
}
@@ -316,8 +324,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<LexOrdering>,
cache: Arc<PlanProperties>,
}
@@ -415,16 +423,12 @@ impl ExecutionPlan for LocalLimitExec {
fn with_new_children(
self: Arc<Self>,
- children: Vec<Arc<dyn ExecutionPlan>>,
+ mut children: Vec<Arc<dyn ExecutionPlan>>,
) -> Result<Arc<dyn ExecutionPlan>> {
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(
@@ -493,14 +497,20 @@ impl ExecutionPlan for LocalLimitExec {
&self,
ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>,
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
+ 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,
},
)),
),
@@ -514,6 +524,7 @@ impl LocalLimitExec {
node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>,
) -> Result<Arc<dyn ExecutionPlan>> {
+ use datafusion_physical_expr_common::sort_expr::optional_ordering_try_from_proto;
use datafusion_proto_models::protobuf;
let limit = crate::expect_plan_variant!(
node,
@@ -522,7 +533,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))
}
}
@@ -661,9 +678,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::expressions::col;
+ use datafusion_physical_expr::{PhysicalExpr, PhysicalSortExpr};
#[tokio::test]
async fn limit() -> Result<()> {
@@ -853,6 +872,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::<GlobalLimitExec>().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::<LocalLimitExec>().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 5a8bede195..4bb2797651 100644
--- a/datafusion/proto-models/proto/datafusion.proto
+++ b/datafusion/proto-models/proto/datafusion.proto
@@ -1480,11 +1480,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 9c82ceb4a2..e04124a61c 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<D>(deserializer: D) -> std::result::Result<GeneratedField, D::Error>
@@ -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(),
})
}
}
@@ -12614,6 +12632,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)?;
@@ -12621,6 +12642,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()
}
}
@@ -12633,12 +12657,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<D>(deserializer: D) -> std::result::Result<GeneratedField, D::Error>
@@ -12662,6 +12689,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)),
}
}
@@ -12683,6 +12711,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 => {
@@ -12699,11 +12728,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 df8f4677c9..51e1a6fa92 100644
--- a/datafusion/proto-models/src/generated/prost.rs
+++ b/datafusion/proto-models/src/generated/prost.rs
@@ -2241,6 +2241,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<PhysicalSortExprNode>,
}
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct LocalLimitExecNode {
@@ -2248,6 +2251,9 @@ pub struct LocalLimitExecNode {
pub input: ::core::option::Option<::prost::alloc::boxed::Box<PhysicalPlanNode>>,
#[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<PhysicalSortExprNode>,
}
#[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 3f03e58656..4e33934c6b 100644
--- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs
+++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs
@@ -57,6 +57,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,
};
@@ -450,6 +451,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<Schema> {
+ 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<Option<LexOrdering>> {
+ 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::<GlobalLimitExec>()
+ .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::<LocalLimitExec>()
+ .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<FileScanConfig> {
+ 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::<DataSourceExec>()
+ .expect("limit should be absorbed into the scan");
+ Ok(scan
+ .data_source()
+ .downcast_ref::<FileScanConfig>()
+ .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);
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org