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/19 18:24:50 UTC
(datafusion) 01/01: test(physical-plan): add colocated field-level plan serde tests (#24222)
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-24222-c429919c7c8bd81e4e687531a609e2380ed60649
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit e2e161d534e577be2fa98ba498cccd2b24048ca9
Author: Adrian Garcia Badaracco <17...@users.noreply.github.com>
AuthorDate: Wed Aug 19 18:24:27 2026 +0000
test(physical-plan): add colocated field-level plan serde tests (#24222)
Adds the second, colocated test tier proposed in #24171 (d), next to the
serde hooks the plans already own.
`datafusion-physical-plan` gains `proto_test_util`, the plan-level
sibling of `datafusion_physical_expr::proto_test_util`: stub
`ExecutionPlanEncode` / `ExecutionPlanDecode` implementations a test
supplies itself, so a plan's `try_to_proto` / `try_from_proto` can be
driven without depending on `datafusion-proto` (which sits above this
crate). The stubs count their calls and can fail on the Nth one, so the
`?` arms are covered too.
On top of that, colocated `proto_tests` modules for the fields where
being wrong is expensive and `Debug`-comparing round-trip tests are
blind:
* `SortExec` / `SortPreservingMergeExec` — `fetch` presence semantics
(absent -> `None`, never `Some(0)`), `preserve_partitioning`, the
`asc`/`descending` inversion, the TopK dynamic filter, and the reject
paths. `SortPreservingMergeExec` also pins that
`enable_round_robin_repartition` is deliberately not on the wire.
* `HashJoinExec` — the three `projection` states proto3 cannot express
directly (the `[u32::MAX]` sentinel), `fetch` presence (the field
dropped in #24165), the checked `u64` -> `usize` conversion, and
`PartitionMode`.
* `NestedLoopJoinExec` — the same projection sentinel, which is written
out a second time in that file and so is tested a second time here.
* `joins::proto` — exhaustive by-name round trips for `JoinType`,
`JoinSide` and `NullEquality`, plus a guard that the two numberings
really do differ, so nobody "simplifies" the matches into a cast.
Per #24171 (c) the new tests were verified against a deliberately broken
encode side rather than assumed to bite.
This supplements the central round-trip tests, it does not replace them:
those still prove the real `PhysicalExtensionCodec` works, that dispatch
reaches the hook, and that bytes survive bytes.
---------
Co-authored-by: Claude <no...@anthropic.com>
---
.../physical-plan/src/joins/hash_join/exec.rs | 291 ++++++++++++++++
.../physical-plan/src/joins/nested_loop_join.rs | 96 ++++++
datafusion/physical-plan/src/joins/proto.rs | 93 +++++
datafusion/physical-plan/src/lib.rs | 3 +
datafusion/physical-plan/src/proto_test_util.rs | 382 +++++++++++++++++++++
datafusion/physical-plan/src/sorts/sort.rs | 324 +++++++++++++++++
.../src/sorts/sort_preserving_merge.rs | 266 ++++++++++++++
7 files changed, 1455 insertions(+)
diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs
index 977a36578d..5506199cb3 100644
--- a/datafusion/physical-plan/src/joins/hash_join/exec.rs
+++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs
@@ -2062,6 +2062,297 @@ impl HashJoinExec {
}
}
+/// Field-level tests for the `try_to_proto` / `try_from_proto` hooks.
+///
+/// These cover the three states of `projection` that proto3 cannot express
+/// directly (the `[u32::MAX]` sentinel), `fetch` presence semantics — the field
+/// dropped in #24165, which the central `Debug`-comparing round-trip tests
+/// could not see — and the by-name `PartitionMode` mapping, whose discriminants
+/// deliberately differ between the plan and the wire.
+#[cfg(all(test, feature = "proto"))]
+mod proto_tests {
+ use super::*;
+ use crate::proto::{ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx};
+ use crate::proto_test_util::{
+ StubPlanDecoder, StubPlanEncoder, UnreachablePlanDecoder, column_node,
+ encoded_child_node, stub_child,
+ };
+ use datafusion_physical_expr::expressions::Column;
+ use datafusion_proto_models::protobuf;
+
+ /// An inner hash join on `a = a` between two stub children.
+ fn join_builder() -> HashJoinExecBuilder {
+ let on: Vec<(PhysicalExprRef, PhysicalExprRef)> =
+ vec![(Arc::new(Column::new("a", 0)), Arc::new(Column::new("a", 0)))];
+ HashJoinExecBuilder::new(stub_child(), stub_child(), on, JoinType::Inner)
+ }
+
+ /// Encode `plan` with a stub encoder, returning the `HashJoinExecNode`.
+ fn encode(
+ plan: &HashJoinExec,
+ encoder: &StubPlanEncoder,
+ ) -> protobuf::HashJoinExecNode {
+ let ctx = ExecutionPlanEncodeCtx::new(encoder);
+ let node = plan
+ .try_to_proto(&ctx)
+ .unwrap()
+ .expect("HashJoinExec should encode to Some(node)");
+ match node.physical_plan_type {
+ Some(protobuf::physical_plan_node::PhysicalPlanType::HashJoin(join)) => *join,
+ other => panic!("expected a HashJoin node, got {other:?}"),
+ }
+ }
+
+ /// Encode a join whose only non-default state is its projection.
+ fn encode_projection(projection: Option<Vec<usize>>) -> Vec<u32> {
+ let plan = join_builder().with_projection(projection).build().unwrap();
+ encode(&plan, &StubPlanEncoder::ok()).projection
+ }
+
+ /// A hand-built `HashJoinExecNode` wrapped in its `PhysicalPlanNode`.
+ fn join_node(node: protobuf::HashJoinExecNode) -> protobuf::PhysicalPlanNode {
+ protobuf::PhysicalPlanNode {
+ physical_plan_type: Some(
+ protobuf::physical_plan_node::PhysicalPlanType::HashJoin(Box::new(node)),
+ ),
+ }
+ }
+
+ /// A decodable `HashJoinExecNode`: two children, one join key, no filter.
+ fn decodable_node() -> protobuf::HashJoinExecNode {
+ protobuf::HashJoinExecNode {
+ left: Some(Box::new(encoded_child_node())),
+ right: Some(Box::new(encoded_child_node())),
+ on: vec![protobuf::JoinOn {
+ left: Some(column_node("a", 0)),
+ right: Some(column_node("a", 0)),
+ }],
+ join_type: protobuf::JoinType::Inner.into(),
+ partition_mode: protobuf::PartitionMode::Partitioned.into(),
+ null_equality: protobuf::NullEquality::NullEqualsNothing.into(),
+ filter: None,
+ projection: vec![],
+ null_aware: false,
+ dynamic_filter: None,
+ fetch: None,
+ }
+ }
+
+ /// Decode `node` with a stub decoder.
+ fn decode(node: protobuf::HashJoinExecNode) -> Arc<dyn ExecutionPlan> {
+ let decoder = StubPlanDecoder::ok();
+ let ctx = ExecutionPlanDecodeCtx::new(&decoder);
+ HashJoinExec::try_from_proto(&join_node(node), &ctx).unwrap()
+ }
+
+ /// View a decoded plan as the `HashJoinExec` it must be.
+ fn as_join(plan: &Arc<dyn ExecutionPlan>) -> &HashJoinExec {
+ plan.downcast_ref::<HashJoinExec>()
+ .expect("decoded plan should be a HashJoinExec")
+ }
+
+ /// "No projection" is the empty repeated field.
+ #[test]
+ fn try_to_proto_encodes_an_absent_projection_as_empty() {
+ assert_eq!(encode_projection(None), Vec::<u32>::new());
+ }
+
+ /// An *empty* projection changes the output schema, so it cannot share the
+ /// "absent" encoding: it goes out as the `[u32::MAX]` sentinel.
+ #[test]
+ fn try_to_proto_encodes_an_empty_projection_as_the_sentinel() {
+ assert_eq!(encode_projection(Some(vec![])), vec![u32::MAX]);
+ }
+
+ #[test]
+ fn try_to_proto_encodes_a_non_empty_projection_as_is() {
+ assert_eq!(encode_projection(Some(vec![0, 2])), vec![0, 2]);
+ }
+
+ #[test]
+ fn try_from_proto_decodes_an_empty_projection_field_as_absent() {
+ let plan = decode(decodable_node());
+
+ assert!(as_join(&plan).projection.is_none());
+ }
+
+ #[test]
+ fn try_from_proto_decodes_the_sentinel_as_an_empty_projection() {
+ let mut node = decodable_node();
+ node.projection = vec![u32::MAX];
+
+ let plan = decode(node);
+ let projection = as_join(&plan)
+ .projection
+ .as_ref()
+ .expect("the sentinel decodes to Some(empty)");
+ assert!(projection.is_empty());
+ }
+
+ #[test]
+ fn try_from_proto_decodes_a_non_empty_projection_as_is() {
+ let mut node = decodable_node();
+ node.projection = vec![0, 2];
+
+ let plan = decode(node);
+ assert_eq!(
+ as_join(&plan).projection.as_deref(),
+ Some([0, 2].as_slice())
+ );
+ }
+
+ /// The regression guard for #24165: an unlimited join must not encode as
+ /// `Some(0)`, and a limited one must keep its limit.
+ #[test]
+ fn try_to_proto_encodes_fetch_by_presence() {
+ let unlimited = join_builder().build().unwrap();
+ assert_eq!(encode(&unlimited, &StubPlanEncoder::ok()).fetch, None);
+
+ let limited = join_builder().with_fetch(Some(10)).build().unwrap();
+ assert_eq!(encode(&limited, &StubPlanEncoder::ok()).fetch, Some(10));
+
+ let empty = join_builder().with_fetch(Some(0)).build().unwrap();
+ assert_eq!(encode(&empty, &StubPlanEncoder::ok()).fetch, Some(0));
+ }
+
+ /// A message written before `fetch` existed has no value on the wire, and
+ /// must decode to "no limit" rather than to `Some(0)` ("no rows").
+ #[test]
+ fn try_from_proto_decodes_an_absent_fetch_as_no_limit() {
+ assert_eq!(as_join(&decode(decodable_node())).fetch, None);
+ }
+
+ #[test]
+ fn try_from_proto_decodes_a_present_fetch() {
+ let mut node = decodable_node();
+ node.fetch = Some(10);
+ assert_eq!(as_join(&decode(node)).fetch, Some(10));
+
+ let mut node = decodable_node();
+ node.fetch = Some(0);
+ assert_eq!(as_join(&decode(node)).fetch, Some(0));
+ }
+
+ /// `fetch` is a `u64` on the wire and a `usize` in the plan. The conversion
+ /// is checked rather than `as usize`, so a value too large for the target
+ /// is reported instead of truncated — on a 64-bit target every `u64` fits,
+ /// and the largest one must come back intact rather than wrapping.
+ #[test]
+ #[cfg(target_pointer_width = "64")]
+ fn try_from_proto_decodes_the_largest_fetch_without_truncating() {
+ let plan = decode({
+ let mut node = decodable_node();
+ node.fetch = Some(u64::MAX);
+ node
+ });
+
+ assert_eq!(as_join(&plan).fetch, Some(usize::MAX));
+ }
+
+ /// The plan-side and wire-side `PartitionMode` discriminants differ, so the
+ /// mapping has to be by name in both directions.
+ #[test]
+ fn partition_mode_round_trips_by_name() {
+ for (mode, wire) in [
+ (
+ PartitionMode::CollectLeft,
+ protobuf::PartitionMode::CollectLeft,
+ ),
+ (
+ PartitionMode::Partitioned,
+ protobuf::PartitionMode::Partitioned,
+ ),
+ (PartitionMode::Auto, protobuf::PartitionMode::Auto),
+ ] {
+ let plan = join_builder().with_partition_mode(mode).build().unwrap();
+ let node = encode(&plan, &StubPlanEncoder::ok());
+ assert_eq!(node.partition_mode, i32::from(wire), "encoding {mode:?}");
+
+ let mut decodable = decodable_node();
+ decodable.partition_mode = node.partition_mode;
+ assert_eq!(*as_join(&decode(decodable)).partition_mode(), mode);
+ }
+ }
+
+ #[test]
+ fn try_from_proto_rejects_an_unknown_partition_mode() {
+ let decoder = StubPlanDecoder::ok();
+ let ctx = ExecutionPlanDecodeCtx::new(&decoder);
+ let mut node = decodable_node();
+ node.partition_mode = 42;
+
+ let err = HashJoinExec::try_from_proto(&join_node(node), &ctx).unwrap_err();
+ assert!(
+ err.to_string()
+ .contains("HashJoinExec: unknown PartitionMode 42")
+ );
+ }
+
+ /// `null_aware` is only legal on a `LeftAnti` join, so the round trip has
+ /// to carry the join type with it.
+ #[test]
+ fn null_aware_round_trips() {
+ let plan = join_builder()
+ .with_type(JoinType::LeftAnti)
+ .with_null_aware(true)
+ .build()
+ .unwrap();
+ assert!(encode(&plan, &StubPlanEncoder::ok()).null_aware);
+
+ let mut node = decodable_node();
+ node.join_type = protobuf::JoinType::Leftanti.into();
+ node.null_aware = true;
+ assert!(as_join(&decode(node)).null_aware);
+ }
+
+ #[test]
+ fn try_to_proto_encodes_both_children_and_both_key_sides() {
+ let encoder = StubPlanEncoder::ok();
+ let node = encode(&join_builder().build().unwrap(), &encoder);
+
+ assert_eq!(encoder.plan_calls(), 2);
+ assert_eq!(encoder.expr_calls(), 2);
+ assert_eq!(node.left, Some(Box::new(encoded_child_node())));
+ assert_eq!(node.right, Some(Box::new(encoded_child_node())));
+ }
+
+ #[test]
+ fn try_to_proto_propagates_child_encode_error() {
+ let plan = join_builder().build().unwrap();
+ let encoder = StubPlanEncoder::failing_on_plan(2);
+ let ctx = ExecutionPlanEncodeCtx::new(&encoder);
+
+ let err = plan.try_to_proto(&ctx).unwrap_err();
+ assert!(
+ err.to_string()
+ .contains("stub plan encode failure on call 2")
+ );
+ }
+
+ #[test]
+ fn try_from_proto_rejects_a_different_plan_variant() {
+ let decoder = UnreachablePlanDecoder::new();
+ let ctx = ExecutionPlanDecodeCtx::new(&decoder);
+
+ let err = HashJoinExec::try_from_proto(&encoded_child_node(), &ctx).unwrap_err();
+ assert!(err.to_string().contains("not a HashJoinExec"));
+ }
+
+ #[test]
+ fn try_from_proto_rejects_a_missing_child() {
+ let decoder = UnreachablePlanDecoder::new();
+ let ctx = ExecutionPlanDecodeCtx::new(&decoder);
+ let mut node = decodable_node();
+ node.left = None;
+
+ let err = HashJoinExec::try_from_proto(&join_node(node), &ctx).unwrap_err();
+ assert!(
+ err.to_string()
+ .contains("HashJoinExec is missing required field 'left'")
+ );
+ }
+}
+
/// Determines which sides of a join are "preserved" for filter pushdown.
///
/// A preserved side means filters on that side's columns can be safely pushed
diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs
index eb1df638c7..0de1dfdaa2 100644
--- a/datafusion/physical-plan/src/joins/nested_loop_join.rs
+++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs
@@ -885,6 +885,102 @@ impl NestedLoopJoinExec {
}
}
+/// Field-level tests for the `try_to_proto` / `try_from_proto` hooks.
+///
+/// `projection` carries the same three states as on `HashJoinExec`, encoded the
+/// same way; the mapping is written out again here, so it is tested again here.
+#[cfg(all(test, feature = "proto"))]
+mod proto_tests {
+ use super::*;
+ use crate::proto::{ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx};
+ use crate::proto_test_util::{
+ StubPlanDecoder, StubPlanEncoder, UnreachablePlanDecoder, encoded_child_node,
+ stub_child,
+ };
+ use datafusion_proto_models::protobuf;
+
+ /// Encode an inner nested loop join with the given projection.
+ fn encode_projection(projection: Option<Vec<usize>>) -> Vec<u32> {
+ let plan = NestedLoopJoinExec::try_new(
+ stub_child(),
+ stub_child(),
+ None,
+ &JoinType::Inner,
+ projection,
+ )
+ .unwrap();
+ let encoder = StubPlanEncoder::ok();
+ let ctx = ExecutionPlanEncodeCtx::new(&encoder);
+ let node = plan
+ .try_to_proto(&ctx)
+ .unwrap()
+ .expect("NestedLoopJoinExec should encode to Some(node)");
+ match node.physical_plan_type {
+ Some(protobuf::physical_plan_node::PhysicalPlanType::NestedLoopJoin(
+ join,
+ )) => join.projection,
+ other => panic!("expected a NestedLoopJoin node, got {other:?}"),
+ }
+ }
+
+ /// A hand-built `NestedLoopJoinExecNode` wrapped in its `PhysicalPlanNode`.
+ fn join_node(projection: Vec<u32>) -> protobuf::PhysicalPlanNode {
+ protobuf::PhysicalPlanNode {
+ physical_plan_type: Some(
+ protobuf::physical_plan_node::PhysicalPlanType::NestedLoopJoin(Box::new(
+ protobuf::NestedLoopJoinExecNode {
+ left: Some(Box::new(encoded_child_node())),
+ right: Some(Box::new(encoded_child_node())),
+ join_type: protobuf::JoinType::Inner.into(),
+ filter: None,
+ projection,
+ },
+ )),
+ ),
+ }
+ }
+
+ /// Decode a node with the given projection field, returning the plan's
+ /// reconstructed projection.
+ fn decode_projection(projection: Vec<u32>) -> Option<Vec<usize>> {
+ let decoder = StubPlanDecoder::ok();
+ let ctx = ExecutionPlanDecodeCtx::new(&decoder);
+ let plan =
+ NestedLoopJoinExec::try_from_proto(&join_node(projection), &ctx).unwrap();
+ plan.downcast_ref::<NestedLoopJoinExec>()
+ .expect("decoded plan should be a NestedLoopJoinExec")
+ .projection
+ .as_ref()
+ .map(|p| p.to_vec())
+ }
+
+ #[test]
+ fn projection_states_survive_the_encode_side() {
+ assert_eq!(encode_projection(None), Vec::<u32>::new());
+ // An empty projection changes the output schema, so it must not share
+ // the "absent" encoding.
+ assert_eq!(encode_projection(Some(vec![])), vec![u32::MAX]);
+ assert_eq!(encode_projection(Some(vec![0, 1])), vec![0, 1]);
+ }
+
+ #[test]
+ fn projection_states_survive_the_decode_side() {
+ assert_eq!(decode_projection(vec![]), None);
+ assert_eq!(decode_projection(vec![u32::MAX]), Some(vec![]));
+ assert_eq!(decode_projection(vec![0, 1]), Some(vec![0, 1]));
+ }
+
+ #[test]
+ fn try_from_proto_rejects_a_different_plan_variant() {
+ let decoder = UnreachablePlanDecoder::new();
+ let ctx = ExecutionPlanDecodeCtx::new(&decoder);
+
+ let err =
+ NestedLoopJoinExec::try_from_proto(&encoded_child_node(), &ctx).unwrap_err();
+ assert!(err.to_string().contains("not a NestedLoopJoinExec"));
+ }
+}
+
impl EmbeddedProjection for NestedLoopJoinExec {
fn with_projection(&self, projection: Option<Vec<usize>>) -> Result<Self> {
self.with_projection(projection)
diff --git a/datafusion/physical-plan/src/joins/proto.rs b/datafusion/physical-plan/src/joins/proto.rs
index 2272828b69..e91e00f897 100644
--- a/datafusion/physical-plan/src/joins/proto.rs
+++ b/datafusion/physical-plan/src/joins/proto.rs
@@ -159,3 +159,96 @@ pub(crate) fn join_filter_from_proto(
Arc::new(schema),
))
}
+
+/// Field-level tests for the shared join enum conversions.
+///
+/// The proto enums and the `datafusion_common` ones are numbered differently,
+/// so these round trips are what stands between a by-name match and a numeric
+/// cast that silently turns a `LEFT SEMI` join into a `FULL` one. Every variant
+/// is listed explicitly: adding one to either side without teaching the
+/// conversion about it fails to compile here.
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ const JOIN_TYPES: [JoinType; 10] = [
+ JoinType::Inner,
+ JoinType::Left,
+ JoinType::Right,
+ JoinType::Full,
+ JoinType::LeftSemi,
+ JoinType::RightSemi,
+ JoinType::LeftAnti,
+ JoinType::RightAnti,
+ JoinType::LeftMark,
+ JoinType::RightMark,
+ ];
+
+ #[test]
+ fn join_type_round_trips_by_name() {
+ for join_type in JOIN_TYPES {
+ let wire = i32::from(join_type_to_proto(join_type));
+ assert_eq!(
+ join_type_from_proto(wire, "test").unwrap(),
+ join_type,
+ "{join_type:?} did not survive the round trip"
+ );
+ }
+ }
+
+ /// The two numberings are not interchangeable, which is why the conversion
+ /// cannot be a cast. If this ever stops holding the by-name matches are
+ /// still correct — but the cast that replaces them would not be.
+ #[test]
+ fn join_type_numbering_differs_between_the_two_enums() {
+ assert!(
+ JOIN_TYPES
+ .iter()
+ .any(|&jt| i32::from(join_type_to_proto(jt)) != jt as i32),
+ "expected at least one JoinType to be numbered differently on the wire"
+ );
+ }
+
+ #[test]
+ fn join_type_from_proto_rejects_an_unknown_value() {
+ let err = join_type_from_proto(99, "TestExec").unwrap_err();
+ assert!(err.to_string().contains("TestExec: unknown JoinType 99"));
+ }
+
+ #[test]
+ fn join_side_round_trips_by_name() {
+ for side in [JoinSide::Left, JoinSide::Right, JoinSide::None] {
+ let wire = i32::from(join_side_to_proto(side));
+ assert_eq!(join_side_from_proto(wire, "test").unwrap(), side);
+ }
+ }
+
+ #[test]
+ fn join_side_from_proto_rejects_an_unknown_value() {
+ let err = join_side_from_proto(99, "TestExec").unwrap_err();
+ assert!(err.to_string().contains("TestExec: unknown JoinSide 99"));
+ }
+
+ #[test]
+ fn null_equality_round_trips_by_name() {
+ for null_equality in [
+ NullEquality::NullEqualsNothing,
+ NullEquality::NullEqualsNull,
+ ] {
+ let wire = i32::from(null_equality_to_proto(null_equality));
+ assert_eq!(
+ null_equality_from_proto(wire, "test").unwrap(),
+ null_equality
+ );
+ }
+ }
+
+ #[test]
+ fn null_equality_from_proto_rejects_an_unknown_value() {
+ let err = null_equality_from_proto(99, "TestExec").unwrap_err();
+ assert!(
+ err.to_string()
+ .contains("TestExec: unknown NullEquality 99")
+ );
+ }
+}
diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs
index 9e50a93b21..673d43f9c1 100644
--- a/datafusion/physical-plan/src/lib.rs
+++ b/datafusion/physical-plan/src/lib.rs
@@ -93,6 +93,9 @@ pub mod placeholder_row;
pub mod projection;
#[cfg(feature = "proto")]
pub mod proto;
+/// Shared test helpers for the colocated `try_to_proto` / `try_from_proto` unit tests
+#[cfg(all(test, feature = "proto"))]
+pub(crate) mod proto_test_util;
pub mod recursive_query;
pub mod repartition;
pub mod scalar_subquery;
diff --git a/datafusion/physical-plan/src/proto_test_util.rs b/datafusion/physical-plan/src/proto_test_util.rs
new file mode 100644
index 0000000000..c1f27ad7fe
--- /dev/null
+++ b/datafusion/physical-plan/src/proto_test_util.rs
@@ -0,0 +1,382 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Shared test helpers for the colocated `try_to_proto` / `try_from_proto`
+//! plan unit tests.
+//!
+//! These let a test drive a plan's serde hooks without depending on
+//! `datafusion-proto` (which sits above this crate and would create a
+//! dependency cycle): the dispatch inversion in [`crate::proto`] means a test
+//! can supply its own [`ExecutionPlanEncode`] / [`ExecutionPlanDecode`].
+//!
+//! This is the plan-level sibling of `datafusion_physical_expr::proto_test_util`.
+//!
+//! # What this tier is for
+//!
+//! Colocated tests prove a plan handles *its own fields*: enum conversions that
+//! must be by-name, the `[u32::MAX]` empty-projection sentinel, `fetch`
+//! presence semantics (absent → `None`, not `Some(0)`). They live next to the
+//! field so they rot when someone adds one, and they can assert on wire state
+//! that a plan's `Debug` output never shows.
+//!
+//! They do *not* replace the central round-trip tests in `datafusion-proto`,
+//! which prove things this tier structurally cannot: that the real
+//! `PhysicalExtensionCodec` works, that dispatch actually reaches the hook, and
+//! that bytes survive bytes.
+
+use std::cell::Cell;
+use std::sync::Arc;
+
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use datafusion_common::{DataFusionError, Result, internal_err};
+use datafusion_execution::TaskContext;
+use datafusion_expr::physical_planning_context::ScalarSubqueryResults;
+use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF};
+use datafusion_physical_expr::PhysicalExpr;
+use datafusion_physical_expr::expressions::Column;
+use datafusion_proto_models::protobuf::{
+ self, PhysicalExprNode, PhysicalPlanNode, physical_expr_node,
+};
+
+use crate::ExecutionPlan;
+use crate::empty::EmptyExec;
+use crate::proto::{ExecutionPlanDecode, ExecutionPlanEncode};
+
+/// The schema shared by the stub child plans: `a: Int32, b: Int32`.
+///
+/// Two columns so a test can build an ordering (or a join key pair) whose
+/// members stay distinguishable through the hooks.
+pub(crate) fn stub_schema() -> SchemaRef {
+ Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int32, true),
+ Field::new("b", DataType::Int32, true),
+ ]))
+}
+
+/// A child plan to hang the plan under test off of.
+pub(crate) fn stub_child() -> Arc<dyn ExecutionPlan> {
+ Arc::new(EmptyExec::new(stub_schema()))
+}
+
+/// A proto node for a `Column`, as a stand-in child node when building a plan's
+/// proto representation by hand.
+pub(crate) fn column_node(name: &str, index: u32) -> PhysicalExprNode {
+ PhysicalExprNode {
+ expr_id: None,
+ expr_type: Some(physical_expr_node::ExprType::Column(
+ protobuf::PhysicalColumn {
+ name: name.to_string(),
+ index,
+ },
+ )),
+ }
+}
+
+/// A proto node for a sort expression over `name`, as written by the sort
+/// plans' `try_to_proto`.
+pub(crate) fn sort_expr_node(
+ name: &str,
+ index: u32,
+ asc: bool,
+ nulls_first: bool,
+) -> PhysicalExprNode {
+ PhysicalExprNode {
+ expr_id: None,
+ expr_type: Some(physical_expr_node::ExprType::Sort(Box::new(
+ protobuf::PhysicalSortExprNode {
+ expr: Some(Box::new(column_node(name, index))),
+ asc,
+ nulls_first,
+ },
+ ))),
+ }
+}
+
+/// The placeholder node [`StubPlanEncoder`] emits for every child plan.
+///
+/// Distinct from `PhysicalPlanNode::default()` so a test can tell "the hook
+/// encoded the child" apart from "the field was left at its default".
+pub(crate) fn encoded_child_node() -> PhysicalPlanNode {
+ PhysicalPlanNode {
+ physical_plan_type: Some(protobuf::physical_plan_node::PhysicalPlanType::Empty(
+ protobuf::EmptyExecNode {
+ schema: None,
+ partitions: 0,
+ },
+ )),
+ }
+}
+
+/// Encoder stub for driving `try_to_proto`.
+///
+/// Emits a recognizable placeholder for each child plan and expression, counts
+/// the calls, and can fail on the Nth call so the `ctx.encode_child(..)?` /
+/// `ctx.encode_expr(..)?` error arms are exercised too.
+pub(crate) struct StubPlanEncoder {
+ plan_calls: Cell<usize>,
+ expr_calls: Cell<usize>,
+ fail_plan_on: Option<usize>,
+ fail_expr_on: Option<usize>,
+}
+
+impl StubPlanEncoder {
+ /// Always succeeds.
+ pub(crate) fn ok() -> Self {
+ Self {
+ plan_calls: Cell::new(0),
+ expr_calls: Cell::new(0),
+ fail_plan_on: None,
+ fail_expr_on: None,
+ }
+ }
+
+ /// Fails on the `call`-th child-plan encode (1-based).
+ pub(crate) fn failing_on_plan(call: usize) -> Self {
+ Self {
+ fail_plan_on: Some(call),
+ ..Self::ok()
+ }
+ }
+
+ /// Fails on the `call`-th expression encode (1-based).
+ pub(crate) fn failing_on_expr(call: usize) -> Self {
+ Self {
+ fail_expr_on: Some(call),
+ ..Self::ok()
+ }
+ }
+
+ /// How many child plans were encoded.
+ pub(crate) fn plan_calls(&self) -> usize {
+ self.plan_calls.get()
+ }
+
+ /// How many expressions were encoded.
+ pub(crate) fn expr_calls(&self) -> usize {
+ self.expr_calls.get()
+ }
+}
+
+impl ExecutionPlanEncode for StubPlanEncoder {
+ fn encode_plan(&self, _plan: &Arc<dyn ExecutionPlan>) -> Result<PhysicalPlanNode> {
+ let call = self.plan_calls.get() + 1;
+ self.plan_calls.set(call);
+ if Some(call) == self.fail_plan_on {
+ return Err(DataFusionError::Internal(format!(
+ "stub plan encode failure on call {call}"
+ )));
+ }
+ Ok(encoded_child_node())
+ }
+
+ fn encode_expr(&self, _expr: &Arc<dyn PhysicalExpr>) -> Result<PhysicalExprNode> {
+ let call = self.expr_calls.get() + 1;
+ self.expr_calls.set(call);
+ if Some(call) == self.fail_expr_on {
+ return Err(DataFusionError::Internal(format!(
+ "stub expr encode failure on call {call}"
+ )));
+ }
+ Ok(column_node("child", 0))
+ }
+
+ fn encode_udf(&self, _udf: &ScalarUDF) -> Result<Option<Vec<u8>>> {
+ Ok(None)
+ }
+
+ fn encode_udaf(&self, _udaf: &AggregateUDF) -> Result<Option<Vec<u8>>> {
+ Ok(None)
+ }
+
+ fn encode_udwf(&self, _udwf: &WindowUDF) -> Result<Option<Vec<u8>>> {
+ Ok(None)
+ }
+}
+
+/// Decoder stub for driving `try_from_proto`.
+///
+/// Returns a fixed [`stub_child`] plan for every child node and decodes column
+/// nodes for real (so a test can still tell two sort keys apart), counting
+/// calls and optionally failing on the Nth one.
+pub(crate) struct StubPlanDecoder {
+ task_ctx: Arc<TaskContext>,
+ plan_calls: Cell<usize>,
+ expr_calls: Cell<usize>,
+ fail_plan_on: Option<usize>,
+ fail_expr_on: Option<usize>,
+}
+
+impl StubPlanDecoder {
+ /// Always succeeds.
+ pub(crate) fn ok() -> Self {
+ Self {
+ task_ctx: Arc::new(TaskContext::default()),
+ plan_calls: Cell::new(0),
+ expr_calls: Cell::new(0),
+ fail_plan_on: None,
+ fail_expr_on: None,
+ }
+ }
+
+ /// Fails on the `call`-th child-plan decode (1-based).
+ pub(crate) fn failing_on_plan(call: usize) -> Self {
+ Self {
+ fail_plan_on: Some(call),
+ ..Self::ok()
+ }
+ }
+
+ /// Fails on the `call`-th expression decode (1-based).
+ pub(crate) fn failing_on_expr(call: usize) -> Self {
+ Self {
+ fail_expr_on: Some(call),
+ ..Self::ok()
+ }
+ }
+
+ /// How many child plans were decoded.
+ pub(crate) fn plan_calls(&self) -> usize {
+ self.plan_calls.get()
+ }
+
+ /// How many expressions were decoded.
+ pub(crate) fn expr_calls(&self) -> usize {
+ self.expr_calls.get()
+ }
+}
+
+impl ExecutionPlanDecode for StubPlanDecoder {
+ fn decode_plan(&self, _node: &PhysicalPlanNode) -> Result<Arc<dyn ExecutionPlan>> {
+ let call = self.plan_calls.get() + 1;
+ self.plan_calls.set(call);
+ if Some(call) == self.fail_plan_on {
+ return Err(DataFusionError::Internal(format!(
+ "stub plan decode failure on call {call}"
+ )));
+ }
+ Ok(stub_child())
+ }
+
+ fn decode_plan_with_scalar_subquery_results(
+ &self,
+ node: &PhysicalPlanNode,
+ _results: ScalarSubqueryResults,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.decode_plan(node)
+ }
+
+ fn decode_expr(
+ &self,
+ node: &PhysicalExprNode,
+ _input_schema: &Schema,
+ ) -> Result<Arc<dyn PhysicalExpr>> {
+ let call = self.expr_calls.get() + 1;
+ self.expr_calls.set(call);
+ if Some(call) == self.fail_expr_on {
+ return Err(DataFusionError::Internal(format!(
+ "stub expr decode failure on call {call}"
+ )));
+ }
+ match &node.expr_type {
+ Some(physical_expr_node::ExprType::Column(c)) => {
+ Ok(Arc::new(Column::new(&c.name, c.index as usize)))
+ }
+ _ => Ok(Arc::new(Column::new("a", 0))),
+ }
+ }
+
+ fn task_ctx(&self) -> &TaskContext {
+ &self.task_ctx
+ }
+
+ fn decode_udf(&self, name: &str, _payload: Option<&[u8]>) -> Result<Arc<ScalarUDF>> {
+ internal_err!("stub decoder cannot decode the scalar UDF {name}")
+ }
+
+ fn decode_udaf(
+ &self,
+ name: &str,
+ _payload: Option<&[u8]>,
+ ) -> Result<Arc<AggregateUDF>> {
+ internal_err!("stub decoder cannot decode the aggregate UDF {name}")
+ }
+
+ fn decode_udwf(&self, name: &str, _payload: Option<&[u8]>) -> Result<Arc<WindowUDF>> {
+ internal_err!("stub decoder cannot decode the window UDF {name}")
+ }
+}
+
+/// Decoder that must never run: asserts that the reject paths of a
+/// `try_from_proto` (wrong node variant, missing required child) bail out
+/// before any decoding happens.
+pub(crate) struct UnreachablePlanDecoder {
+ task_ctx: Arc<TaskContext>,
+}
+
+impl UnreachablePlanDecoder {
+ pub(crate) fn new() -> Self {
+ Self {
+ task_ctx: Arc::new(TaskContext::default()),
+ }
+ }
+}
+
+impl ExecutionPlanDecode for UnreachablePlanDecoder {
+ fn decode_plan(&self, _node: &PhysicalPlanNode) -> Result<Arc<dyn ExecutionPlan>> {
+ unreachable!("decode_plan must not be reached when the node is rejected")
+ }
+
+ fn decode_plan_with_scalar_subquery_results(
+ &self,
+ _node: &PhysicalPlanNode,
+ _results: ScalarSubqueryResults,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ unreachable!("decode_plan must not be reached when the node is rejected")
+ }
+
+ fn decode_expr(
+ &self,
+ _node: &PhysicalExprNode,
+ _input_schema: &Schema,
+ ) -> Result<Arc<dyn PhysicalExpr>> {
+ unreachable!("decode_expr must not be reached when the node is rejected")
+ }
+
+ fn task_ctx(&self) -> &TaskContext {
+ &self.task_ctx
+ }
+
+ fn decode_udf(&self, _name: &str, _payload: Option<&[u8]>) -> Result<Arc<ScalarUDF>> {
+ unreachable!("decode_udf must not be reached when the node is rejected")
+ }
+
+ fn decode_udaf(
+ &self,
+ _name: &str,
+ _payload: Option<&[u8]>,
+ ) -> Result<Arc<AggregateUDF>> {
+ unreachable!("decode_udaf must not be reached when the node is rejected")
+ }
+
+ fn decode_udwf(
+ &self,
+ _name: &str,
+ _payload: Option<&[u8]>,
+ ) -> Result<Arc<WindowUDF>> {
+ unreachable!("decode_udwf must not be reached when the node is rejected")
+ }
+}
diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs
index f4f1a5e5b4..bd47668cf3 100644
--- a/datafusion/physical-plan/src/sorts/sort.rs
+++ b/datafusion/physical-plan/src/sorts/sort.rs
@@ -1712,6 +1712,330 @@ impl SortExec {
}
}
+/// Field-level tests for the `try_to_proto` / `try_from_proto` hooks.
+///
+/// These sit next to the fields they cover so they rot when a field is added,
+/// and they assert on the wire representation directly — `fetch`, in
+/// particular, is not printed by `SortExec`'s `Debug` impl, which is how a
+/// dropped `fetch` survived the central round-trip tests in #24165.
+#[cfg(all(test, feature = "proto"))]
+mod proto_tests {
+ use super::*;
+ use crate::proto::{ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx};
+ use crate::proto_test_util::{
+ StubPlanDecoder, StubPlanEncoder, UnreachablePlanDecoder, column_node,
+ encoded_child_node, sort_expr_node, stub_child,
+ };
+ use arrow::compute::SortOptions;
+ use datafusion_physical_expr::expressions::Column;
+ use datafusion_proto_models::protobuf;
+
+ /// A `SortExec` over `a ASC NULLS LAST` with the given fetch.
+ fn sort_fixture(fetch: Option<usize>) -> SortExec {
+ let input = stub_child();
+ let ordering = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: false,
+ nulls_first: false,
+ },
+ )])
+ .unwrap();
+ SortExec::new(ordering, input).with_fetch(fetch)
+ }
+
+ /// Encode `plan` with a stub encoder, returning the `SortExecNode`.
+ fn encode(plan: &SortExec, encoder: &StubPlanEncoder) -> protobuf::SortExecNode {
+ let ctx = ExecutionPlanEncodeCtx::new(encoder);
+ let node = plan
+ .try_to_proto(&ctx)
+ .unwrap()
+ .expect("SortExec should encode to Some(node)");
+ match node.physical_plan_type {
+ Some(protobuf::physical_plan_node::PhysicalPlanType::Sort(sort)) => *sort,
+ other => panic!("expected a Sort node, got {other:?}"),
+ }
+ }
+
+ /// A hand-built `SortExecNode` wrapped in its `PhysicalPlanNode`.
+ fn sort_node(node: protobuf::SortExecNode) -> protobuf::PhysicalPlanNode {
+ protobuf::PhysicalPlanNode {
+ physical_plan_type: Some(
+ protobuf::physical_plan_node::PhysicalPlanType::Sort(Box::new(node)),
+ ),
+ }
+ }
+
+ /// A decodable `SortExecNode`: one child, one sort key, no dynamic filter.
+ fn decodable_node(fetch: i64, preserve_partitioning: bool) -> protobuf::SortExecNode {
+ protobuf::SortExecNode {
+ input: Some(Box::new(encoded_child_node())),
+ expr: vec![sort_expr_node("a", 0, true, false)],
+ fetch,
+ preserve_partitioning,
+ dynamic_filter: None,
+ }
+ }
+
+ /// Decode `node`, returning the `SortExec`.
+ fn decode(node: protobuf::SortExecNode, decoder: &StubPlanDecoder) -> Arc<SortExec> {
+ let ctx = ExecutionPlanDecodeCtx::new(decoder);
+ SortExec::try_from_proto(&sort_node(node), &ctx)
+ .unwrap()
+ .downcast_ref::<SortExec>()
+ .expect("decoded plan should be a SortExec")
+ .clone()
+ .into()
+ }
+
+ #[test]
+ fn try_to_proto_encodes_absent_fetch_as_negative_one() {
+ let encoder = StubPlanEncoder::ok();
+ let node = encode(&sort_fixture(None), &encoder);
+
+ assert_eq!(node.fetch, -1);
+ // No fetch means no TopK dynamic filter to carry either.
+ assert_eq!(node.dynamic_filter, None);
+ assert_eq!(encoder.plan_calls(), 1);
+ }
+
+ #[test]
+ fn try_to_proto_encodes_present_fetch() {
+ let encoder = StubPlanEncoder::ok();
+ let node = encode(&sort_fixture(Some(10)), &encoder);
+
+ assert_eq!(node.fetch, 10);
+ }
+
+ /// `Some(0)` must not collapse into the "absent" encoding: a `LIMIT 0`
+ /// sort returns no rows, an unlimited one returns all of them.
+ #[test]
+ fn try_to_proto_distinguishes_zero_fetch_from_absent_fetch() {
+ let encoder = StubPlanEncoder::ok();
+ let node = encode(&sort_fixture(Some(0)), &encoder);
+
+ assert_eq!(node.fetch, 0);
+ }
+
+ /// The other end of the range does *not* survive: `fetch` goes onto the
+ /// wire as `usize as i64`, so on a 64-bit target `usize::MAX` wraps to
+ /// `-1` — the very value that means "absent" — and reads back as an
+ /// unlimited sort. That is pre-existing behavior of the `i64` wire field
+ /// rather than something this tier introduces; pinning it keeps any change
+ /// to the encoding a deliberate one instead of a silent fix.
+ #[test]
+ #[cfg(target_pointer_width = "64")]
+ fn try_to_proto_wraps_usize_max_fetch_into_the_absent_encoding() {
+ let encoder = StubPlanEncoder::ok();
+ let node = encode(&sort_fixture(Some(usize::MAX)), &encoder);
+
+ assert_eq!(node.fetch, -1);
+
+ // ... and so the limit is gone by the time the node is read back.
+ let decoder = StubPlanDecoder::ok();
+ assert_eq!(
+ decode(decodable_node(node.fetch, false), &decoder).fetch(),
+ None
+ );
+ }
+
+ #[test]
+ fn try_to_proto_encodes_preserve_partitioning() {
+ let encoder = StubPlanEncoder::ok();
+ let plan = sort_fixture(None).with_preserve_partitioning(true);
+
+ assert!(encode(&plan, &encoder).preserve_partitioning);
+ assert!(
+ !encode(&sort_fixture(None), &StubPlanEncoder::ok()).preserve_partitioning
+ );
+ }
+
+ /// The wire format stores `asc`, the plan stores `descending`; the
+ /// inversion has to survive both directions.
+ #[test]
+ fn try_to_proto_inverts_descending_into_asc() {
+ let input = stub_child();
+ let ordering = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+ let encoder = StubPlanEncoder::ok();
+ let node = encode(&SortExec::new(ordering, input), &encoder);
+
+ let sort_expr = match node.expr[0].expr_type.as_ref().unwrap() {
+ protobuf::physical_expr_node::ExprType::Sort(sort) => sort,
+ other => panic!("expected a Sort expr node, got {other:?}"),
+ };
+ assert!(!sort_expr.asc);
+ assert!(sort_expr.nulls_first);
+ }
+
+ /// A fetch turns the sort into a TopK, which produces a dynamic filter that
+ /// has to ride along on the wire.
+ #[test]
+ fn try_to_proto_encodes_the_topk_dynamic_filter() {
+ let encoder = StubPlanEncoder::ok();
+ let node = encode(&sort_fixture(Some(3)), &encoder);
+
+ assert!(node.dynamic_filter.is_some());
+ // One call for the sort key, one for the dynamic filter.
+ assert_eq!(encoder.expr_calls(), 2);
+ }
+
+ #[test]
+ fn try_to_proto_propagates_child_encode_error() {
+ let encoder = StubPlanEncoder::failing_on_plan(1);
+ let ctx = ExecutionPlanEncodeCtx::new(&encoder);
+
+ let err = sort_fixture(None).try_to_proto(&ctx).unwrap_err();
+ assert!(err.to_string().contains("stub plan encode failure"));
+ }
+
+ #[test]
+ fn try_to_proto_propagates_expr_encode_error() {
+ let encoder = StubPlanEncoder::failing_on_expr(1);
+ let ctx = ExecutionPlanEncodeCtx::new(&encoder);
+
+ let err = sort_fixture(None).try_to_proto(&ctx).unwrap_err();
+ assert!(err.to_string().contains("stub expr encode failure"));
+ }
+
+ #[test]
+ fn try_from_proto_decodes_negative_fetch_as_absent() {
+ let decoder = StubPlanDecoder::ok();
+ let plan = decode(decodable_node(-1, false), &decoder);
+
+ assert_eq!(plan.fetch(), None);
+ assert_eq!(decoder.plan_calls(), 1);
+ assert_eq!(decoder.expr_calls(), 1);
+ }
+
+ #[test]
+ fn try_from_proto_decodes_zero_fetch_as_some_zero() {
+ let decoder = StubPlanDecoder::ok();
+
+ assert_eq!(decode(decodable_node(0, false), &decoder).fetch(), Some(0));
+ }
+
+ #[test]
+ fn try_from_proto_decodes_present_fetch() {
+ let decoder = StubPlanDecoder::ok();
+
+ assert_eq!(decode(decodable_node(7, false), &decoder).fetch(), Some(7));
+ }
+
+ #[test]
+ fn try_from_proto_restores_preserve_partitioning() {
+ let decoder = StubPlanDecoder::ok();
+
+ assert!(decode(decodable_node(-1, true), &decoder).preserve_partitioning());
+ assert!(!decode(decodable_node(-1, false), &decoder).preserve_partitioning());
+ }
+
+ #[test]
+ fn try_from_proto_restores_sort_options() {
+ let decoder = StubPlanDecoder::ok();
+ let mut node = decodable_node(-1, false);
+ node.expr = vec![sort_expr_node("a", 0, false, true)];
+
+ let plan = decode(node, &decoder);
+ let sort_expr = plan.expr().first();
+ assert!(sort_expr.options.descending);
+ assert!(sort_expr.options.nulls_first);
+ }
+
+ #[test]
+ fn try_from_proto_rejects_a_different_plan_variant() {
+ let decoder = UnreachablePlanDecoder::new();
+ let ctx = ExecutionPlanDecodeCtx::new(&decoder);
+
+ let err = SortExec::try_from_proto(&encoded_child_node(), &ctx).unwrap_err();
+ assert!(err.to_string().contains("not a SortExec"));
+ }
+
+ #[test]
+ fn try_from_proto_rejects_a_missing_input() {
+ let decoder = UnreachablePlanDecoder::new();
+ let ctx = ExecutionPlanDecodeCtx::new(&decoder);
+ let mut node = decodable_node(-1, false);
+ node.input = None;
+
+ let err = SortExec::try_from_proto(&sort_node(node), &ctx).unwrap_err();
+ assert!(
+ err.to_string()
+ .contains("SortExec is missing required field 'input'")
+ );
+ }
+
+ #[test]
+ fn try_from_proto_rejects_a_non_sort_expression() {
+ let decoder = StubPlanDecoder::ok();
+ let ctx = ExecutionPlanDecodeCtx::new(&decoder);
+ let mut node = decodable_node(-1, false);
+ node.expr = vec![column_node("a", 0)];
+
+ let err = SortExec::try_from_proto(&sort_node(node), &ctx).unwrap_err();
+ assert!(
+ err.to_string()
+ .contains("SortExec expr must be a sort expression")
+ );
+ }
+
+ #[test]
+ fn try_from_proto_rejects_an_empty_ordering() {
+ let decoder = StubPlanDecoder::ok();
+ let ctx = ExecutionPlanDecodeCtx::new(&decoder);
+ let mut node = decodable_node(-1, false);
+ node.expr = vec![];
+
+ let err = SortExec::try_from_proto(&sort_node(node), &ctx).unwrap_err();
+ assert!(err.to_string().contains("SortExec requires an ordering"));
+ }
+
+ /// `dynamic_filter` is plan-owned state, not just another expression: the
+ /// node has to come back as a `DynamicFilterPhysicalExpr`, and one holding
+ /// anything else is rejected rather than quietly dropped. This is the
+ /// decode-side counterpart of the encode test above.
+ #[test]
+ fn try_from_proto_rejects_a_dynamic_filter_of_the_wrong_type() {
+ let decoder = StubPlanDecoder::ok();
+ let ctx = ExecutionPlanDecodeCtx::new(&decoder);
+ let mut node = decodable_node(3, false);
+ node.dynamic_filter = Some(column_node("a", 0));
+
+ let err = SortExec::try_from_proto(&sort_node(node), &ctx).unwrap_err();
+ assert!(err.to_string().contains(
+ "SortExec dynamic_filter did not decode to a DynamicFilterPhysicalExpr"
+ ));
+ // The sort key and the dynamic filter both reached the codec.
+ assert_eq!(decoder.expr_calls(), 2);
+ }
+
+ #[test]
+ fn try_from_proto_propagates_child_decode_error() {
+ let decoder = StubPlanDecoder::failing_on_plan(1);
+ let ctx = ExecutionPlanDecodeCtx::new(&decoder);
+
+ let err = SortExec::try_from_proto(&sort_node(decodable_node(-1, false)), &ctx)
+ .unwrap_err();
+ assert!(err.to_string().contains("stub plan decode failure"));
+ }
+
+ #[test]
+ fn try_from_proto_propagates_expr_decode_error() {
+ let decoder = StubPlanDecoder::failing_on_expr(1);
+ let ctx = ExecutionPlanDecodeCtx::new(&decoder);
+
+ let err = SortExec::try_from_proto(&sort_node(decodable_node(-1, false)), &ctx)
+ .unwrap_err();
+ assert!(err.to_string().contains("stub expr decode failure"));
+ }
+}
+
#[cfg(test)]
mod tests {
use std::collections::HashMap;
diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs
index ad17f2c213..09c1489002 100644
--- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs
+++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs
@@ -561,6 +561,272 @@ impl SortPreservingMergeExec {
}
}
+/// Field-level tests for the `try_to_proto` / `try_from_proto` hooks.
+///
+/// `SortPreservingMergeExec` had no proto coverage at all until #24172 (see
+/// #24171); these cover the fields that its `Debug` output does not show, so a
+/// dropped `fetch` fails here rather than passing silently.
+#[cfg(all(test, feature = "proto"))]
+mod proto_tests {
+ use super::*;
+ use crate::proto::{ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx};
+ use crate::proto_test_util::{
+ StubPlanDecoder, StubPlanEncoder, UnreachablePlanDecoder, column_node,
+ encoded_child_node, sort_expr_node, stub_child,
+ };
+ use arrow::compute::SortOptions;
+ use datafusion_physical_expr::expressions::Column;
+ use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
+ use datafusion_proto_models::protobuf;
+
+ /// A `SortPreservingMergeExec` over `a ASC NULLS LAST` with the given fetch.
+ fn spm_fixture(fetch: Option<usize>) -> SortPreservingMergeExec {
+ let ordering = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: false,
+ nulls_first: false,
+ },
+ )])
+ .unwrap();
+ SortPreservingMergeExec::new(ordering, stub_child()).with_fetch(fetch)
+ }
+
+ /// Encode `plan` with a stub encoder, returning the `SortPreservingMergeExecNode`.
+ fn encode(
+ plan: &SortPreservingMergeExec,
+ encoder: &StubPlanEncoder,
+ ) -> protobuf::SortPreservingMergeExecNode {
+ let ctx = ExecutionPlanEncodeCtx::new(encoder);
+ let node = plan
+ .try_to_proto(&ctx)
+ .unwrap()
+ .expect("SortPreservingMergeExec should encode to Some(node)");
+ match node.physical_plan_type {
+ Some(
+ protobuf::physical_plan_node::PhysicalPlanType::SortPreservingMerge(spm),
+ ) => *spm,
+ other => panic!("expected a SortPreservingMerge node, got {other:?}"),
+ }
+ }
+
+ /// A hand-built `SortPreservingMergeExecNode` wrapped in its `PhysicalPlanNode`.
+ fn spm_node(
+ node: protobuf::SortPreservingMergeExecNode,
+ ) -> protobuf::PhysicalPlanNode {
+ protobuf::PhysicalPlanNode {
+ physical_plan_type: Some(
+ protobuf::physical_plan_node::PhysicalPlanType::SortPreservingMerge(
+ Box::new(node),
+ ),
+ ),
+ }
+ }
+
+ /// A decodable `SortPreservingMergeExecNode`: one child, one sort key.
+ fn decodable_node(fetch: i64) -> protobuf::SortPreservingMergeExecNode {
+ protobuf::SortPreservingMergeExecNode {
+ input: Some(Box::new(encoded_child_node())),
+ expr: vec![sort_expr_node("a", 0, true, false)],
+ fetch,
+ }
+ }
+
+ /// Decode `node`, returning the `SortPreservingMergeExec`.
+ fn decode(
+ node: protobuf::SortPreservingMergeExecNode,
+ decoder: &StubPlanDecoder,
+ ) -> Arc<SortPreservingMergeExec> {
+ let ctx = ExecutionPlanDecodeCtx::new(decoder);
+ SortPreservingMergeExec::try_from_proto(&spm_node(node), &ctx)
+ .unwrap()
+ .downcast_ref::<SortPreservingMergeExec>()
+ .expect("decoded plan should be a SortPreservingMergeExec")
+ .clone()
+ .into()
+ }
+
+ #[test]
+ fn try_to_proto_encodes_absent_fetch_as_negative_one() {
+ let encoder = StubPlanEncoder::ok();
+ let node = encode(&spm_fixture(None), &encoder);
+
+ assert_eq!(node.fetch, -1);
+ assert_eq!(encoder.plan_calls(), 1);
+ assert_eq!(encoder.expr_calls(), 1);
+ }
+
+ #[test]
+ fn try_to_proto_encodes_present_fetch() {
+ let encoder = StubPlanEncoder::ok();
+
+ assert_eq!(encode(&spm_fixture(Some(11)), &encoder).fetch, 11);
+ }
+
+ /// `Some(0)` must not collapse into the "absent" encoding.
+ #[test]
+ fn try_to_proto_distinguishes_zero_fetch_from_absent_fetch() {
+ let encoder = StubPlanEncoder::ok();
+
+ assert_eq!(encode(&spm_fixture(Some(0)), &encoder).fetch, 0);
+ }
+
+ /// ... and `usize::MAX` does not survive it at all: `fetch` goes onto the
+ /// wire as `usize as i64`, so on a 64-bit target it wraps to `-1`, the
+ /// "absent" value, and reads back as an unlimited merge. Same pre-existing
+ /// `i64` wire behavior as `SortExec`, pinned here for the same reason.
+ #[test]
+ #[cfg(target_pointer_width = "64")]
+ fn try_to_proto_wraps_usize_max_fetch_into_the_absent_encoding() {
+ let encoder = StubPlanEncoder::ok();
+ let node = encode(&spm_fixture(Some(usize::MAX)), &encoder);
+
+ assert_eq!(node.fetch, -1);
+
+ let decoder = StubPlanDecoder::ok();
+ assert_eq!(decode(decodable_node(node.fetch), &decoder).fetch(), None);
+ }
+
+ #[test]
+ fn try_to_proto_inverts_descending_into_asc() {
+ let ordering = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+ let plan = SortPreservingMergeExec::new(ordering, stub_child());
+ let encoder = StubPlanEncoder::ok();
+
+ let node = encode(&plan, &encoder);
+ let sort_expr = match node.expr[0].expr_type.as_ref().unwrap() {
+ protobuf::physical_expr_node::ExprType::Sort(sort) => sort,
+ other => panic!("expected a Sort expr node, got {other:?}"),
+ };
+ assert!(!sort_expr.asc);
+ assert!(sort_expr.nulls_first);
+ }
+
+ #[test]
+ fn try_to_proto_propagates_child_encode_error() {
+ let encoder = StubPlanEncoder::failing_on_plan(1);
+ let ctx = ExecutionPlanEncodeCtx::new(&encoder);
+
+ let err = spm_fixture(None).try_to_proto(&ctx).unwrap_err();
+ assert!(err.to_string().contains("stub plan encode failure"));
+ }
+
+ #[test]
+ fn try_to_proto_propagates_expr_encode_error() {
+ let encoder = StubPlanEncoder::failing_on_expr(1);
+ let ctx = ExecutionPlanEncodeCtx::new(&encoder);
+
+ let err = spm_fixture(None).try_to_proto(&ctx).unwrap_err();
+ assert!(err.to_string().contains("stub expr encode failure"));
+ }
+
+ #[test]
+ fn try_from_proto_decodes_negative_fetch_as_absent() {
+ let decoder = StubPlanDecoder::ok();
+
+ assert_eq!(decode(decodable_node(-1), &decoder).fetch(), None);
+ }
+
+ #[test]
+ fn try_from_proto_decodes_zero_fetch_as_some_zero() {
+ let decoder = StubPlanDecoder::ok();
+
+ assert_eq!(decode(decodable_node(0), &decoder).fetch(), Some(0));
+ }
+
+ #[test]
+ fn try_from_proto_decodes_present_fetch() {
+ let decoder = StubPlanDecoder::ok();
+
+ assert_eq!(decode(decodable_node(4), &decoder).fetch(), Some(4));
+ }
+
+ #[test]
+ fn try_from_proto_restores_sort_options() {
+ let decoder = StubPlanDecoder::ok();
+ let mut node = decodable_node(-1);
+ node.expr = vec![sort_expr_node("a", 0, false, true)];
+
+ let plan = decode(node, &decoder);
+ let sort_expr = plan.expr().first();
+ assert!(sort_expr.options.descending);
+ assert!(sort_expr.options.nulls_first);
+ }
+
+ /// `enable_round_robin_repartition` is deliberately not on the wire: a
+ /// decoded plan always comes back with the default. Pinning that here keeps
+ /// the omission a decision rather than an accident.
+ #[test]
+ fn try_from_proto_leaves_round_robin_repartition_at_its_default() {
+ let decoder = StubPlanDecoder::ok();
+ let plan = decode(decodable_node(-1), &decoder);
+
+ assert!(plan.enable_round_robin_repartition);
+ }
+
+ #[test]
+ fn try_from_proto_rejects_a_different_plan_variant() {
+ let decoder = UnreachablePlanDecoder::new();
+ let ctx = ExecutionPlanDecodeCtx::new(&decoder);
+
+ let err = SortPreservingMergeExec::try_from_proto(&encoded_child_node(), &ctx)
+ .unwrap_err();
+ assert!(err.to_string().contains("not a SortPreservingMergeExec"));
+ }
+
+ #[test]
+ fn try_from_proto_rejects_a_missing_input() {
+ let decoder = UnreachablePlanDecoder::new();
+ let ctx = ExecutionPlanDecodeCtx::new(&decoder);
+ let mut node = decodable_node(-1);
+ node.input = None;
+
+ let err =
+ SortPreservingMergeExec::try_from_proto(&spm_node(node), &ctx).unwrap_err();
+ assert!(
+ err.to_string()
+ .contains("SortPreservingMergeExec is missing required field 'input'")
+ );
+ }
+
+ #[test]
+ fn try_from_proto_rejects_a_non_sort_expression() {
+ let decoder = StubPlanDecoder::ok();
+ let ctx = ExecutionPlanDecodeCtx::new(&decoder);
+ let mut node = decodable_node(-1);
+ node.expr = vec![column_node("a", 0)];
+
+ let err =
+ SortPreservingMergeExec::try_from_proto(&spm_node(node), &ctx).unwrap_err();
+ assert!(
+ err.to_string()
+ .contains("expression is not a sort expression")
+ );
+ }
+
+ #[test]
+ fn try_from_proto_rejects_an_empty_ordering() {
+ let decoder = StubPlanDecoder::ok();
+ let ctx = ExecutionPlanDecodeCtx::new(&decoder);
+ let mut node = decodable_node(-1);
+ node.expr = vec![];
+
+ let err =
+ SortPreservingMergeExec::try_from_proto(&spm_node(node), &ctx).unwrap_err();
+ assert!(
+ err.to_string()
+ .contains("SortPreservingMergeExec requires an ordering")
+ );
+ }
+}
+
#[cfg(test)]
mod tests {
use std::collections::HashSet;
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org