kosiew commented on code in PR #24222:
URL: https://github.com/apache/datafusion/pull/24222#discussion_r3811317045
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -1695,6 +1695,289 @@ 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() {
Review Comment:
Could we add a regression test for `usize::MAX` here, and the equivalent
test for `SortPreservingMergeExec`? The current encoder uses `usize as i64`, so
on a 64-bit target `usize::MAX` becomes `-1`, which then decodes as an absent
fetch / unlimited plan. This looks like pre-existing behavior, so I don't think
it needs to block this PR, but this new fetch-focused test tier seems like a
good place to expose it and make the intended overflow behavior explicit.
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -1695,6 +1695,289 @@ 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);
+ }
+
+ #[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() {
Review Comment:
Could we also consider a decode-side test for rejecting an invalid dynamic
filter? For example, populate `dynamic_filter` with `column_node(...)` and
assert the downcast error. The current test covers the encode-side presence of
the dynamic filter, and this would exercise the corresponding plan-owned
invalid-wire path without needing `datafusion-proto`. This is non-blocking.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
