refactor(proto): put Partitioning / sort-expression serde on the types

datafusion/physical-expr-common/src/sort_expr.rs

@@ -183,6 +183,102 @@ impl PhysicalSortExpr {
     }
 }
 
+/// Protobuf conversions for [`PhysicalSortExpr`].
+///
+/// This is the flat [`PhysicalSortExprNode`] representation used wherever the
+/// wire format stores an ordering (scan output orderings, range partitioning,
+/// window frames, …). It is *not* the `PhysicalExprNode::Sort` wrapping that
+/// `SortExec` uses for its own `expr` field.
+///
+/// [`PhysicalSortExprNode`]: datafusion_proto_models::protobuf::PhysicalSortExprNode
+#[cfg(feature = "proto")]
+impl PhysicalSortExpr {
+    /// Serialize this sort expression, encoding its child expression through
+    /// `ctx`.
+    pub fn try_to_proto(
+        &self,
+        ctx: &crate::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>,
+    ) -> Result<datafusion_proto_models::protobuf::PhysicalSortExprNode> {
+        Ok(datafusion_proto_models::protobuf::PhysicalSortExprNode {
+            expr: Some(Box::new(ctx.encode_child(&self.expr)?)),
+            asc: !self.options.descending,
+            nulls_first: self.options.nulls_first,
+        })
+    }
+
+    /// Reconstruct a [`PhysicalSortExpr`] from its protobuf representation.
+    pub fn try_from_proto(
+        node: &datafusion_proto_models::protobuf::PhysicalSortExprNode,
+        ctx: &crate::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>,
+    ) -> Result<Self> {
+        let expr = ctx.decode_required_expression(
+            node.expr.as_deref(),
+            "PhysicalSortExpr",
+            "expr",
+        )?;
+        Ok(PhysicalSortExpr {
+            expr,
+            options: SortOptions {
+                descending: !node.asc,
+                nulls_first: node.nulls_first,
+            },
+        })
+    }
+}
+
+/// Serialize a sequence of sort expressions into the flat
+/// [`PhysicalSortExprNode`] list the wire format uses for an ordering.
+///
+/// Accepts anything that yields [`PhysicalSortExpr`]s by value or by reference,
+/// so a [`LexOrdering`], a `&[PhysicalSortExpr]`, or a [`LexRequirement`]
+/// mapped through [`PhysicalSortExpr::from`] all work:
+///
+/// ```ignore
+/// let nodes = sort_exprs_try_to_proto(ordering.iter(), ctx)?;
+/// let nodes = sort_exprs_try_to_proto(
+///     requirement.iter().map(|req| PhysicalSortExpr::from(req.clone())),
+///     ctx,
+/// )?;
+/// ```
+///
+/// The `PhysicalSortExprNodeCollection` message some plans use is just this
+/// list in a wrapper, so those callers wrap the result themselves rather than
+/// this function guessing which shape they mean.
+///
+/// [`PhysicalSortExprNode`]: datafusion_proto_models::protobuf::PhysicalSortExprNode
+#[cfg(feature = "proto")]
+pub fn sort_exprs_try_to_proto<E: std::borrow::Borrow<PhysicalSortExpr>>(
+    exprs: impl IntoIterator<Item = E>,
+    ctx: &crate::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>,
+) -> Result<Vec<datafusion_proto_models::protobuf::PhysicalSortExprNode>> {
+    exprs
+        .into_iter()
+        .map(|expr| expr.borrow().try_to_proto(ctx))
+        .collect()
+}
+
+/// Reconstruct a sequence of sort expressions from the flat
+/// [`PhysicalSortExprNode`] list, the counterpart of
+/// [`sort_exprs_try_to_proto`].
+///
+/// Returns the expressions rather than a [`LexOrdering`] or a
+/// [`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.
+///
+/// [`PhysicalSortExprNode`]: datafusion_proto_models::protobuf::PhysicalSortExprNode
+#[cfg(feature = "proto")]
+pub fn sort_exprs_try_from_proto(
+    nodes: &[datafusion_proto_models::protobuf::PhysicalSortExprNode],
+    ctx: &crate::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>,
+) -> Result<Vec<PhysicalSortExpr>> {
+    nodes
+        .iter()
+        .map(|node| PhysicalSortExpr::try_from_proto(node, ctx))
+        .collect()
+}
+
 impl PartialEq for PhysicalSortExpr {
     fn eq(&self, other: &Self) -> bool {
         self.options == other.options && self.expr.eq(&other.expr)

datafusion/physical-expr/src/partitioning.rs

@@ -25,6 +25,10 @@ pub use datafusion_common::SplitPoint;
 use datafusion_common::{Result, validate_range_split_points};
 use datafusion_physical_expr_common::physical_expr::format_physical_expr_list;
 use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
+#[cfg(feature = "proto")]
+use datafusion_physical_expr_common::sort_expr::{
+    sort_exprs_try_from_proto, sort_exprs_try_to_proto,
+};
 use std::fmt;
 use std::fmt::Display;
 use std::sync::Arc;
@@ -515,6 +519,156 @@ impl Partitioning {
     }
 }
 
+/// Protobuf conversions for [`Partitioning`].
+///
+/// Child expressions (hash keys, range orderings) and `ScalarValue` split
+/// points are (de)serialized through the expression-level context, so this is
+/// the single copy of the partitioning wire format: `RepartitionExec` and
+/// `datafusion-proto`'s central serializer route through it, and the remaining
+/// per-plan migrations (`FileScanConfig` and friends) are meant to do the same
+/// rather than grow another copy.
+///
+/// [`protobuf::Partitioning`]: datafusion_proto_models::protobuf::Partitioning
+#[cfg(feature = "proto")]
+impl Partitioning {
+    /// Serialize this partitioning into its protobuf representation.
+    pub fn try_to_proto(
+        &self,
+        ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>,
+    ) -> Result<datafusion_proto_models::protobuf::Partitioning> {
+        use datafusion_proto_models::protobuf;
+
+        let partition_method = match self {
+            Partitioning::RoundRobinBatch(n) => {
+                protobuf::partitioning::PartitionMethod::RoundRobin(wire_partition_count(
+                    *n,
+                )?)
+            }
+            Partitioning::Hash(exprs, n) => {
+                protobuf::partitioning::PartitionMethod::Hash(
+                    protobuf::PhysicalHashRepartition {
+                        hash_expr: ctx.encode_children_expressions(exprs)?,
+                        partition_count: wire_partition_count(*n)?,
+                    },
+                )
+            }
+            Partitioning::Range(range) => {
+                let sort_expr = sort_exprs_try_to_proto(range.ordering().iter(), ctx)?;
+                let split_point = range
+                    .split_points()
+                    .iter()
+                    .map(|split_point| {
+                        let value = split_point
+                            .values()
+                            .iter()
+                            .map(|value| value.try_into().map_err(Into::into))
+                            .collect::<Result<Vec<_>>>()?;
+                        Ok(protobuf::PhysicalRangeSplitPoint { value })
+                    })
+                    .collect::<Result<Vec<_>>>()?;
+                protobuf::partitioning::PartitionMethod::Range(
+                    protobuf::PhysicalRangePartitioning {
+                        sort_expr,
+                        split_point,
+                    },
+                )
+            }
+            Partitioning::UnknownPartitioning(n) => {
+                protobuf::partitioning::PartitionMethod::Unknown(wire_partition_count(
+                    *n,
+                )?)
+            }
+        };
+        Ok(protobuf::Partitioning {
+            partition_method: Some(partition_method),
+        })
+    }
+
+    /// Reconstruct a [`Partitioning`] from its protobuf representation.
+    ///
+    /// Returns `Ok(None)` when the message carries no `partition_method`, which
+    /// the wire format uses to mean "no output partitioning declared"; callers
+    /// for which it is required should turn that into their own error.
+    pub fn try_from_proto(
+        node: &datafusion_proto_models::protobuf::Partitioning,
+        ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>,
+    ) -> Result<Option<Self>> {
+        use datafusion_common::{ScalarValue, internal_datafusion_err, internal_err};
+        use datafusion_proto_models::protobuf;
+
+        let Some(partition_method) = node.partition_method.as_ref() else {
+            return Ok(None);
+        };
+        let partitioning = match partition_method {
+            protobuf::partitioning::PartitionMethod::RoundRobin(n) => {
+                Partitioning::RoundRobinBatch(partition_count(*n)?)
+            }
+            protobuf::partitioning::PartitionMethod::Hash(hash) => {
+                let exprs = hash
+                    .hash_expr
+                    .iter()
+                    .map(|expr| ctx.decode(expr))
+                    .collect::<Result<Vec<_>>>()?;
+                Partitioning::Hash(exprs, partition_count(hash.partition_count)?)
+            }
+            protobuf::partitioning::PartitionMethod::Unknown(n) => {
+                Partitioning::UnknownPartitioning(partition_count(*n)?)
+            }
+            protobuf::partitioning::PartitionMethod::Range(range) => {
+                let sort_exprs = sort_exprs_try_from_proto(&range.sort_expr, ctx)?;
+                let sort_expr_count = sort_exprs.len();
+                let ordering = LexOrdering::new(sort_exprs).ok_or_else(|| {
+                    internal_datafusion_err!(
+                        "Range partitioning requires non-empty ordering"
+                    )
+                })?;
+                if ordering.len() != sort_expr_count {
+                    return internal_err!(
+                        "Range partitioning ordering must not contain duplicate expressions"
+                    );
+                }
+                let split_points = range
+                    .split_point
+                    .iter()
+                    .map(|split_point| {
+                        let values = split_point
+                            .value
+                            .iter()
+                            .map(|value| ScalarValue::try_from(value).map_err(Into::into))
+                            .collect::<Result<Vec<_>>>()?;
+                        Ok(SplitPoint::new(values))
+                    })
+                    .collect::<Result<Vec<_>>>()?;
+                Partitioning::Range(RangePartitioning::try_new(ordering, split_points)?)
+            }
+        };
+        Ok(Some(partitioning))
+    }
+}
+
+/// Narrow a wire partition count to `usize`.
+#[cfg(feature = "proto")]
+fn partition_count(count: u64) -> Result<usize> {
+    usize::try_from(count).map_err(|_| {
+        datafusion_common::internal_datafusion_err!(
+            "Partition count {count} exceeds usize::MAX"
+        )
+    })
+}
+
+/// Widen a partition count to its `u64` wire representation.
+///
+/// The mirror of [`partition_count`]: an out-of-range count is an error on both
+/// sides rather than a silent truncation on the way out.
+#[cfg(feature = "proto")]
+fn wire_partition_count(count: usize) -> Result<u64> {
+    u64::try_from(count).map_err(|_| {
+        datafusion_common::internal_datafusion_err!(
+            "Partition count {count} exceeds u64::MAX"
+        )
+    })
+}
+
 impl PartialEq for Partitioning {
     fn eq(&self, other: &Partitioning) -> bool {
         match (self, other) {
@@ -1138,3 +1292,113 @@ mod tests {
         Ok(())
     }
 }
+
+#[cfg(all(test, feature = "proto"))]
+mod ordering_proto_tests {
+    use std::sync::Arc;
+
+    use arrow::compute::SortOptions;
+    use arrow::datatypes::{DataType, Field, Schema};
+    use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx;
+    use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx;
+    use datafusion_physical_expr_common::sort_expr::{
+        LexRequirement, PhysicalSortExpr, PhysicalSortRequirement,
+        sort_exprs_try_from_proto, sort_exprs_try_to_proto,
+    };
+
+    use crate::expressions::Column;
+    use crate::proto_test_util::{StubDecoder, StubEncoder};
+
+    fn schema() -> Schema {
+        Schema::new(vec![Field::new("a", DataType::Int32, false)])
+    }
+
+    fn sort_expr(descending: bool, nulls_first: bool) -> PhysicalSortExpr {
+        PhysicalSortExpr::new(
+            Arc::new(Column::new("a", 0)),
+            SortOptions {
+                descending,
+                nulls_first,
+            },
+        )
+    }
+
+    #[test]
+    fn sort_exprs_round_trip_preserves_options_and_order() {
+        let encoder = StubEncoder::ok();
+        let encode_ctx = PhysicalExprEncodeCtx::new(&encoder);
+        let exprs = vec![sort_expr(true, false), sort_expr(false, true)];
+
+        let nodes = sort_exprs_try_to_proto(&exprs, &encode_ctx).unwrap();
+        // `asc` is the inverse of `descending` on the wire.
+        assert_eq!(
+            nodes
+                .iter()
+                .map(|node| (node.asc, node.nulls_first))
+                .collect::<Vec<_>>(),
+            vec![(false, false), (true, true)]
+        );
+
+        let schema = schema();
+        let decoder = StubDecoder::ok();
+        let decode_ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);
+        let decoded = sort_exprs_try_from_proto(&nodes, &decode_ctx).unwrap();
+        assert_eq!(
+            decoded.iter().map(|expr| expr.options).collect::<Vec<_>>(),
+            exprs.iter().map(|expr| expr.options).collect::<Vec<_>>()
+        );
+    }
+
+    #[test]
+    fn sort_exprs_accepts_owned_requirements() {
+        let encoder = StubEncoder::ok();
+        let encode_ctx = PhysicalExprEncodeCtx::new(&encoder);
+        let requirement = LexRequirement::from([PhysicalSortRequirement::new(
+            Arc::new(Column::new("a", 0)),
+            Some(SortOptions {
+                descending: true,
+                nulls_first: true,
+            }),
+        )]);
+
+        let nodes = sort_exprs_try_to_proto(
+            requirement
+                .iter()
+                .map(|req| PhysicalSortExpr::from(req.clone())),
+            &encode_ctx,
+        )
+        .unwrap();
+
+        assert_eq!(nodes.len(), 1);
+        assert!(!nodes[0].asc);
+        assert!(nodes[0].nulls_first);
+    }
+
+    #[test]
+    fn sort_exprs_propagate_encode_errors() {
+        let encoder = StubEncoder::failing_on(2);
+        let encode_ctx = PhysicalExprEncodeCtx::new(&encoder);
+        let exprs = vec![sort_expr(false, false), sort_expr(true, true)];
+
+        let err = sort_exprs_try_to_proto(&exprs, &encode_ctx).unwrap_err();
+        assert!(err.to_string().contains("stub encode failure on call 2"));
+    }
+
+    #[test]
+    fn sort_exprs_reject_missing_inner_expr() {
+        let encoder = StubEncoder::ok();
+        let encode_ctx = PhysicalExprEncodeCtx::new(&encoder);
+        let mut nodes =
+            sort_exprs_try_to_proto(&[sort_expr(false, false)], &encode_ctx).unwrap();
+        nodes[0].expr = None;
+
+        let schema = schema();
+        let decoder = StubDecoder::ok();
+        let decode_ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);
+        let err = sort_exprs_try_from_proto(&nodes, &decode_ctx).unwrap_err();
+        assert!(
+            err.to_string()
+                .contains("PhysicalSortExpr is missing required field 'expr'")
+        );
+    }
+}