adriangb commented on code in PR #24003:
URL: https://github.com/apache/datafusion/pull/24003#discussion_r3685576918
##########
datafusion/physical-expr/src/partitioning.rs:
##########
@@ -515,6 +519,138 @@ 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`,
+/// `FileScanConfig` and `datafusion-proto`'s central serializer all route
+/// through it.
+///
+/// [`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(*n as u64)
+ }
+ Partitioning::Hash(exprs, n) => {
+ protobuf::partitioning::PartitionMethod::Hash(
+ protobuf::PhysicalHashRepartition {
+ hash_expr: ctx.encode_children_expressions(exprs)?,
+ partition_count: *n as u64,
+ },
+ )
+ }
Review Comment:
Fixed. Encode now goes through a `wire_partition_count` helper
(`u64::try_from`, internal error on overflow), mirroring the `partition_count`
narrowing on the decode side, so an out-of-range count is an error in both
directions rather than a silent truncation. Applied to all three sites
(round-robin, hash, unknown).
##########
datafusion/physical-expr/src/partitioning.rs:
##########
@@ -515,6 +519,138 @@ 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`,
+/// `FileScanConfig` and `datafusion-proto`'s central serializer all route
+/// through it.
Review Comment:
Reworded — `FileScanConfig` is now described as a next-step consumer
(#23683), not a current one: the doc says `RepartitionExec` and
`datafusion-proto`'s central serializer route through this, and the remaining
per-plan migrations are meant to do the same rather than grow another copy.
##########
datafusion/physical-expr/src/partitioning.rs:
##########
@@ -515,6 +519,138 @@ 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`,
+/// `FileScanConfig` and `datafusion-proto`'s central serializer all route
+/// through it.
+///
+/// [`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(*n as u64)
+ }
+ Partitioning::Hash(exprs, n) => {
+ protobuf::partitioning::PartitionMethod::Hash(
+ protobuf::PhysicalHashRepartition {
+ hash_expr: ctx.encode_children_expressions(exprs)?,
+ partition_count: *n as u64,
+ },
+ )
+ }
+ 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(*n as u64)
+ }
+ };
Review Comment:
Same as the sibling comment above — both encode sites now use the checked
`wire_partition_count` helper.
##########
datafusion/physical-plan/src/proto.rs:
##########
@@ -190,6 +196,25 @@ impl<'a> ExecutionPlanEncodeCtx<'a> {
pub fn encode_udwf(&self, udwf: &WindowUDF) -> Result<Option<Vec<u8>>> {
self.encoder.encode_udwf(udwf)
}
+
+ /// An expression-level encode context backed by this plan context.
+ ///
+ /// Lets a plan hand `ctx` to expression-level conversions that own their
own
+ /// wire logic — e.g.
+ ///
[`Partitioning::try_to_proto`](datafusion_physical_expr::Partitioning::try_to_proto)
+ /// and
+ ///
[`PhysicalSortExpr::try_to_proto`](datafusion_physical_expr::PhysicalSortExpr::try_to_proto).
+ pub fn expr_ctx(&self) -> PhysicalExprEncodeCtx<'_> {
+ PhysicalExprEncodeCtx::new(self)
+ }
+}
+
+/// Lets [`ExecutionPlanEncodeCtx`] back an [`PhysicalExprEncodeCtx`], so
+/// expression-level conversions can be reused from plan hooks.
Review Comment:
Fixed — now reads "back a `PhysicalExprEncodeCtx`".
--
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]
