diff --git a/datafusion/datasource/src/file.rs b/datafusion/datasource/src/file.rs index 44fdf8dc16d74..691bb314b7c03 100644 --- a/datafusion/datasource/src/file.rs +++ b/datafusion/datasource/src/file.rs @@ -358,7 +358,7 @@ pub trait FileSource: Any + Send + Sync { /// `base` is the shared [`FileScanConfig`] this source is wrapped in; the /// format-agnostic parts (file groups, schema, statistics, ordering, /// projection, …) are encoded via - /// [`FileScanConfig::to_proto_conf`](crate::file_scan_config::FileScanConfig::to_proto_conf), + /// [`FileScanConfig::try_to_proto`](crate::file_scan_config::FileScanConfig::try_to_proto), /// and the concrete source appends its format-specific fields (e.g. CSV /// delimiter/quote) around it. /// diff --git a/datafusion/datasource/src/file_scan_config/mod.rs b/datafusion/datasource/src/file_scan_config/mod.rs index f9535d81417dd..d02ecff914396 100644 --- a/datafusion/datasource/src/file_scan_config/mod.rs +++ b/datafusion/datasource/src/file_scan_config/mod.rs @@ -21,7 +21,7 @@ pub(crate) mod sort_pushdown; /// Shared `FileScanConfig` <-> proto conversion, gated on the `proto` feature. -/// Attaches inherent `to_proto_conf` / `from_proto_conf` / `parse_table_schema_from_proto` +/// Attaches inherent `try_to_proto` / `try_from_proto` / `parse_table_schema_from_proto` /// helpers to [`FileScanConfig`] used by every file source's `try_to_proto` hook. #[cfg(feature = "proto")] pub(crate) mod proto; diff --git a/datafusion/datasource/src/file_scan_config/proto.rs b/datafusion/datasource/src/file_scan_config/proto.rs index 89a764ad143aa..aac2364b5f030 100644 --- a/datafusion/datasource/src/file_scan_config/proto.rs +++ b/datafusion/datasource/src/file_scan_config/proto.rs @@ -25,33 +25,29 @@ //! instead of the raw `PhysicalExtensionCodec` + //! `PhysicalProtoConverterExtension`. Every //! `FileSource::try_to_proto` hook (CSV, JSON, Arrow, Parquet, Avro) builds its -//! `*ScanExecNode` around [`FileScanConfig::to_proto_conf`] and decodes with -//! [`FileScanConfig::from_proto_conf`], keeping a single copy of the shared +//! `*ScanExecNode` around [`FileScanConfig::try_to_proto`] and decodes with +//! [`FileScanConfig::try_from_proto`], keeping a single copy of the shared //! wire logic. The wire format is byte-for-byte identical to the old central //! serializer. //! -//! Child physical expressions (sort orderings, hash/range partitioning, and -//! projection expressions) are (de)serialized through `ctx.encode_expr` / -//! `ctx.decode_expr`; `Schema`, `Statistics`, `Constraints`, and `ScalarValue` -//! go through `datafusion-proto-common`. Nothing here needs the raw codec. +//! Sort orderings and output partitioning ride the shared +//! `ctx.encode_sort_exprs` / `ctx.encode_partitioning` helpers (and their decode +//! siblings); projection expressions go through `ctx.encode_expr` / +//! `ctx.decode_expr`; file groups through +//! [`FileGroup::try_to_proto`](crate::file_groups::FileGroup::try_to_proto); +//! and `Schema`, `Statistics`, `Constraints` and `ScalarValue` through +//! `datafusion-proto-common`. Nothing here needs the raw codec. use std::sync::Arc; -use arrow::compute::SortOptions; use arrow::datatypes::Schema; -use chrono::{TimeZone, Utc}; use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_physical_expr::projection::{ProjectionExpr, ProjectionExprs}; -use datafusion_physical_expr::{ - LexOrdering, Partitioning, PhysicalSortExpr, RangePartitioning, SplitPoint, -}; +use datafusion_physical_expr::{LexOrdering, Partitioning, PhysicalSortExpr}; use datafusion_physical_plan::proto::{ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx}; use datafusion_proto_models::protobuf; -use object_store::ObjectMeta; -use object_store::path::Path; -use crate::PartitionedFile; use crate::file::FileSource; use crate::file_groups::FileGroup; use crate::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; @@ -64,44 +60,42 @@ impl FileScanConfig { /// Each concrete [`FileSource::try_to_proto`] /// wraps the returned value in its own `*ScanExecNode`. Byte-compatible with /// the former `serialize_file_scan_config` in `datafusion-proto`. - pub fn to_proto_conf( + /// + /// Note this inherent method shadows + /// [`DataSource::try_to_proto`](crate::source::DataSource::try_to_proto) for + /// callers holding a concrete `FileScanConfig`: that one serializes the whole + /// `DataSourceExec` node (by delegating to the file source), this one only the + /// shared `FileScanExecConf` spine. Spell the trait method + /// `DataSource::try_to_proto(&config, ctx)` when you want the plan node. + pub fn try_to_proto( &self, ctx: &ExecutionPlanEncodeCtx<'_>, ) -> Result { let file_groups = self .file_groups .iter() - .map(file_group_to_proto) + .map(FileGroup::try_to_proto) .collect::>>()?; - // Sort orderings: only the child expressions need the ctx; the - // asc/nulls_first wrapping is plain data inlined into a - // `PhysicalSortExprNode` (same shape as `sorts/sort.rs`). + let expr_ctx = ctx.expr_ctx(); let mut output_ordering = vec![]; for order in &self.output_ordering { - let nodes = order - .iter() - .map(|sort_expr| { - Ok(protobuf::PhysicalSortExprNode { - expr: Some(Box::new(ctx.encode_expr(&sort_expr.expr)?)), - asc: !sort_expr.options.descending, - nulls_first: sort_expr.options.nulls_first, - }) - }) - .collect::>>()?; output_ordering.push(protobuf::PhysicalSortExprNodeCollection { - physical_sort_expr_nodes: nodes, + physical_sort_expr_nodes: order + .iter() + .map(|sort_expr| sort_expr.try_to_proto(&expr_ctx)) + .collect::>>()?, }); } let output_partitioning = self .output_partitioning .as_ref() - .map(|p| partitioning_to_proto(p, ctx)) + .map(|partitioning| partitioning.try_to_proto(&expr_ctx)) .transpose()?; // Fields must be added to the schema so that they can persist in the - // protobuf, and then removed from the schema in `from_proto_conf`. + // protobuf, and then removed from the schema in `try_from_proto`. let mut fields = self .file_schema() .fields() @@ -156,7 +150,7 @@ impl FileScanConfig { /// table schema via [`FileScanConfig::parse_table_schema_from_proto`]). /// /// Byte-compatible with the former `parse_protobuf_file_scan_config`. - pub fn from_proto_conf( + pub fn try_from_proto( conf: &protobuf::FileScanExecConf, ctx: &ExecutionPlanDecodeCtx<'_>, file_source: Arc, @@ -185,7 +179,7 @@ impl FileScanConfig { let file_groups = conf .file_groups .iter() - .map(file_group_from_proto) + .map(FileGroup::try_from_proto) .collect::>>()?; let object_store_url = match conf.object_store_url.is_empty() { @@ -193,18 +187,23 @@ impl FileScanConfig { true => ObjectStoreUrl::local_filesystem(), }; + let expr_ctx = ctx.expr_ctx(&schema); let mut output_ordering = vec![]; for node_collection in &conf.output_ordering { - let sort_exprs = parse_sort_exprs( - &node_collection.physical_sort_expr_nodes, - ctx, - &schema, - )?; + let sort_exprs = node_collection + .physical_sort_expr_nodes + .iter() + .map(|sort_expr| PhysicalSortExpr::try_from_proto(sort_expr, &expr_ctx)) + .collect::>>()?; output_ordering.extend(LexOrdering::new(sort_exprs)); } - let output_partitioning = - partitioning_from_proto(conf.output_partitioning.as_ref(), ctx, &schema)?; + let output_partitioning = conf + .output_partitioning + .as_ref() + .map(|partitioning| Partitioning::try_from_proto(partitioning, &expr_ctx)) + .transpose()? + .flatten(); // Parse projection expressions if present and apply to the file source. let file_source = if let Some(proto_projection_exprs) = &conf.projection_exprs { @@ -244,7 +243,7 @@ impl FileScanConfig { /// Parse a [`TableSchema`] (file schema + partition columns) from a /// [`protobuf::FileScanExecConf`]. File sources use this to rebuild their - /// concrete source before calling [`FileScanConfig::from_proto_conf`]. + /// concrete source before calling [`FileScanConfig::try_from_proto`]. /// /// Byte-compatible with the former `parse_table_schema_from_proto`. pub fn parse_table_schema_from_proto( @@ -295,225 +294,3 @@ fn parse_file_scan_schema(conf: &protobuf::FileScanExecConf) -> Result, - schema: &Schema, -) -> Result> { - nodes - .iter() - .map(|sort_expr| { - let expr = sort_expr.expr.as_ref().ok_or_else(|| { - internal_datafusion_err!("Unexpected empty physical expression") - })?; - Ok(PhysicalSortExpr { - expr: ctx.decode_expr(expr, schema)?, - options: SortOptions { - descending: !sort_expr.asc, - nulls_first: sort_expr.nulls_first, - }, - }) - }) - .collect() -} - -/// Inlined equivalent of `datafusion-proto`'s `serialize_partitioning`. Only -/// child physical expressions and `ScalarValue`s need the ctx; the -/// `protobuf::Partitioning` wrapping is built directly here. -fn partitioning_to_proto( - partitioning: &Partitioning, - ctx: &ExecutionPlanEncodeCtx<'_>, -) -> Result { - let partition_method = match partitioning { - Partitioning::RoundRobinBatch(n) => { - protobuf::partitioning::PartitionMethod::RoundRobin(*n as u64) - } - Partitioning::Hash(exprs, n) => { - let hash_expr = ctx.encode_expressions(exprs)?; - protobuf::partitioning::PartitionMethod::Hash( - protobuf::PhysicalHashRepartition { - hash_expr, - partition_count: *n as u64, - }, - ) - } - Partitioning::Range(range) => { - let sort_expr = range - .ordering() - .iter() - .map(|sort_expr| { - Ok(protobuf::PhysicalSortExprNode { - expr: Some(Box::new(ctx.encode_expr(&sort_expr.expr)?)), - asc: !sort_expr.options.descending, - nulls_first: sort_expr.options.nulls_first, - }) - }) - .collect::>>()?; - 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::>>()?; - Ok(protobuf::PhysicalRangeSplitPoint { value }) - }) - .collect::>>()?; - protobuf::partitioning::PartitionMethod::Range( - protobuf::PhysicalRangePartitioning { - sort_expr, - split_point, - }, - ) - } - Partitioning::UnknownPartitioning(n) => { - protobuf::partitioning::PartitionMethod::Unknown(*n as u64) - } - }; - Ok(protobuf::Partitioning { - partition_method: Some(partition_method), - }) -} - -/// Inlined equivalent of `datafusion-proto`'s `parse_protobuf_partitioning`. -fn partitioning_from_proto( - partitioning: Option<&protobuf::Partitioning>, - ctx: &ExecutionPlanDecodeCtx<'_>, - schema: &Schema, -) -> Result> { - let Some(partitioning) = partitioning else { - return Ok(None); - }; - let Some(partition_method) = partitioning.partition_method.as_ref() else { - return Ok(None); - }; - let partitioning = match partition_method { - protobuf::partitioning::PartitionMethod::RoundRobin(n) => { - Partitioning::RoundRobinBatch(*n as usize) - } - protobuf::partitioning::PartitionMethod::Hash(hash) => { - let exprs = hash - .hash_expr - .iter() - .map(|expr| ctx.decode_expr(expr, schema)) - .collect::>>()?; - Partitioning::Hash(exprs, hash.partition_count as usize) - } - protobuf::partitioning::PartitionMethod::Unknown(n) => { - Partitioning::UnknownPartitioning(*n as usize) - } - protobuf::partitioning::PartitionMethod::Range(range) => { - let sort_exprs = parse_sort_exprs(&range.sort_expr, ctx, schema)?; - 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 Err(internal_datafusion_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| { - datafusion_common::ScalarValue::try_from(value) - .map_err(Into::into) - }) - .collect::>>()?; - Ok(SplitPoint::new(values)) - }) - .collect::>>()?; - Partitioning::Range(RangePartitioning::try_new(ordering, split_points)?) - } - }; - Ok(Some(partitioning)) -} - -fn file_group_to_proto(group: &FileGroup) -> Result { - Ok(protobuf::FileGroup { - files: group - .files() - .iter() - .map(partitioned_file_to_proto) - .collect::>>()?, - }) -} - -fn file_group_from_proto(group: &protobuf::FileGroup) -> Result { - let files = group - .files - .iter() - .map(partitioned_file_from_proto) - .collect::>>()?; - Ok(FileGroup::new(files)) -} - -pub(crate) fn partitioned_file_to_proto( - pf: &PartitionedFile, -) -> Result { - let last_modified = pf.object_meta.last_modified; - let last_modified_ns = last_modified.timestamp_nanos_opt().ok_or_else(|| { - DataFusionError::Plan(format!( - "Invalid timestamp on PartitionedFile::ObjectMeta: {last_modified}" - )) - })? as u64; - Ok(protobuf::PartitionedFile { - arrow_schema: pf - .arrow_schema - .as_ref() - .map(|s| s.as_ref().try_into()) - .transpose()?, - path: pf.object_meta.location.as_ref().to_owned(), - size: pf.object_meta.size, - last_modified_ns, - partition_values: pf - .partition_values - .iter() - .map(|v| v.try_into()) - .collect::, _>>()?, - range: pf.range.as_ref().map(|range| protobuf::FileRange { - start: range.start, - end: range.end, - }), - statistics: pf.statistics.as_ref().map(|s| s.as_ref().into()), - }) -} - -pub(crate) fn partitioned_file_from_proto( - val: &protobuf::PartitionedFile, -) -> Result { - let mut pf = PartitionedFile::new_from_meta(ObjectMeta { - location: Path::parse(val.path.as_str()) - .map_err(|e| internal_datafusion_err!("Invalid object_store path: {e}"))?, - last_modified: Utc.timestamp_nanos(val.last_modified_ns as i64), - size: val.size, - e_tag: None, - version: None, - }) - .with_partition_values( - val.partition_values - .iter() - .map(|v| v.try_into()) - .collect::, _>>()?, - ); - if let Some(proto_schema) = val.arrow_schema.as_ref() { - pf = pf.with_arrow_schema(Arc::new( - proto_schema.try_into().map_err(DataFusionError::from)?, - )); - } - if let Some(range) = val.range.as_ref() { - pf = pf.with_range(range.start, range.end); - } - if let Some(proto_stats) = val.statistics.as_ref() { - pf = pf.with_statistics(Arc::new(proto_stats.try_into()?)); - } - Ok(pf) -} diff --git a/datafusion/datasource/src/mod.rs b/datafusion/datasource/src/mod.rs index 7c8cae337f1eb..e718064c819da 100644 --- a/datafusion/datasource/src/mod.rs +++ b/datafusion/datasource/src/mod.rs @@ -41,6 +41,10 @@ pub mod file_stream; pub mod memory; pub mod morsel; pub mod projection; +/// Protobuf conversions for [`FileRange`], [`PartitionedFile`] and +/// [`FileGroup`](crate::file_groups::FileGroup), gated on the `proto` feature. +#[cfg(feature = "proto")] +pub mod proto; pub mod schema_adapter; pub mod sink; pub mod source; diff --git a/datafusion/datasource/src/proto.rs b/datafusion/datasource/src/proto.rs new file mode 100644 index 0000000000000..2796aa224c5de --- /dev/null +++ b/datafusion/datasource/src/proto.rs @@ -0,0 +1,218 @@ +// 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. + +//! Protobuf conversions for the file-scan leaf types owned by this crate: +//! [`FileRange`], [`PartitionedFile`] and [`FileGroup`]. +//! +//! These are the single copy of that wire logic. `datafusion-proto`'s +//! `TryFromProto` implementations for the same types are thin shims that +//! delegate here, so the format cannot drift between the central serializer and +//! the per-source `try_to_proto` hooks. +//! +//! None of these conversions need a codec or an encode/decode context: every +//! field is plain data or goes through `datafusion-proto-common`. + +use std::sync::Arc; + +use chrono::{TimeZone, Utc}; +use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; +use datafusion_proto_models::protobuf; +use object_store::ObjectMeta; +use object_store::path::Path; + +use crate::file_groups::FileGroup; +use crate::{FileRange, PartitionedFile}; + +impl FileRange { + /// Serialize this range into its protobuf representation. + pub fn try_to_proto(&self) -> Result { + Ok(protobuf::FileRange { + start: self.start, + end: self.end, + }) + } + + /// Reconstruct a [`FileRange`] from its protobuf representation. + pub fn try_from_proto(range: &protobuf::FileRange) -> Result { + Ok(FileRange { + start: range.start, + end: range.end, + }) + } +} + +impl PartitionedFile { + /// Serialize this file into its protobuf representation. + pub fn try_to_proto(&self) -> Result { + let last_modified = self.object_meta.last_modified; + let last_modified_ns = last_modified.timestamp_nanos_opt().ok_or_else(|| { + DataFusionError::Plan(format!( + "Invalid timestamp on PartitionedFile::ObjectMeta: {last_modified}" + )) + })? as u64; + Ok(protobuf::PartitionedFile { + arrow_schema: self + .arrow_schema + .as_ref() + .map(|s| s.as_ref().try_into()) + .transpose()?, + path: self.object_meta.location.as_ref().to_owned(), + size: self.object_meta.size, + last_modified_ns, + partition_values: self + .partition_values + .iter() + .map(|v| v.try_into()) + .collect::, _>>()?, + range: self + .range + .as_ref() + .map(FileRange::try_to_proto) + .transpose()?, + statistics: self.statistics.as_ref().map(|s| s.as_ref().into()), + }) + } + + /// Reconstruct a [`PartitionedFile`] from its protobuf representation. + pub fn try_from_proto(file: &protobuf::PartitionedFile) -> Result { + let mut pf = PartitionedFile::new_from_meta(ObjectMeta { + location: Path::parse(file.path.as_str()).map_err(|e| { + internal_datafusion_err!("Invalid object_store path: {e}") + })?, + last_modified: Utc.timestamp_nanos(file.last_modified_ns as i64), + size: file.size, + e_tag: None, + version: None, + }) + .with_partition_values( + file.partition_values + .iter() + .map(|v| v.try_into()) + .collect::, _>>()?, + ); + if let Some(proto_schema) = file.arrow_schema.as_ref() { + pf = pf.with_arrow_schema(Arc::new( + proto_schema.try_into().map_err(DataFusionError::from)?, + )); + } + if let Some(range) = file.range.as_ref() { + let range = FileRange::try_from_proto(range)?; + pf = pf.with_range(range.start, range.end); + } + if let Some(proto_stats) = file.statistics.as_ref() { + pf = pf.with_statistics(Arc::new(proto_stats.try_into()?)); + } + Ok(pf) + } +} + +impl FileGroup { + /// Serialize this group into its protobuf representation. + pub fn try_to_proto(&self) -> Result { + Ok(protobuf::FileGroup { + files: self + .files() + .iter() + .map(PartitionedFile::try_to_proto) + .collect::>>()?, + }) + } + + /// Reconstruct a [`FileGroup`] from its protobuf representation. + pub fn try_from_proto(group: &protobuf::FileGroup) -> Result { + Ok(FileGroup::new( + group + .files + .iter() + .map(PartitionedFile::try_from_proto) + .collect::>>()?, + )) + } +} + +#[cfg(test)] +mod tests { + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::{ScalarValue, Statistics}; + + use super::*; + + #[test] + fn partitioned_file_roundtrip_preserves_all_fields() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let pf = PartitionedFile::new_from_meta(ObjectMeta { + location: Path::parse("foo/bar.parquet")?, + last_modified: Utc.timestamp_nanos(1_000_000_000), + size: 1234, + e_tag: None, + version: None, + }) + .with_partition_values(vec![ScalarValue::from("2024-01-01")]) + .with_range(10, 20) + .with_arrow_schema(Arc::clone(&schema)) + .with_statistics(Arc::new(Statistics::new_unknown(&schema))); + + let decoded = PartitionedFile::try_from_proto(&pf.try_to_proto()?)?; + + assert_eq!(decoded.object_meta.location, pf.object_meta.location); + assert_eq!(decoded.object_meta.size, pf.object_meta.size); + assert_eq!( + decoded.object_meta.last_modified, + pf.object_meta.last_modified + ); + assert_eq!(decoded.partition_values, pf.partition_values); + assert_eq!(decoded.range, pf.range); + assert_eq!(decoded.arrow_schema.as_deref(), Some(schema.as_ref())); + // Statistics on `PartitionedFile` span the full table schema, and + // `PartitionedFile::with_statistics` re-derives the partition column + // entries, so the decoded statistics are not compared field by field + // here; this is pre-existing behavior of the wire format. + assert!(decoded.statistics.is_some()); + Ok(()) + } + + #[test] + fn partitioned_file_from_proto_rejects_invalid_path() { + let proto = protobuf::PartitionedFile { + path: "foo//bar.parquet".to_string(), + ..Default::default() + }; + + let err = PartitionedFile::try_from_proto(&proto).unwrap_err(); + assert!( + err.to_string().contains("Invalid object_store path"), + "unexpected error: {err}" + ); + } + + #[test] + fn file_group_roundtrip() -> Result<()> { + let group = FileGroup::new(vec![ + PartitionedFile::new("a.parquet", 1), + PartitionedFile::new("b.parquet", 2), + ]); + + let decoded = FileGroup::try_from_proto(&group.try_to_proto()?)?; + + assert_eq!(decoded.len(), 2); + assert_eq!( + decoded.files()[1].object_meta.location, + group.files()[1].object_meta.location + ); + Ok(()) + } +} diff --git a/datafusion/physical-expr-common/src/sort_expr.rs b/datafusion/physical-expr-common/src/sort_expr.rs index 84ffb92eaa600..82f7a6e5fc33f 100644 --- a/datafusion/physical-expr-common/src/sort_expr.rs +++ b/datafusion/physical-expr-common/src/sort_expr.rs @@ -183,6 +183,49 @@ 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 { + 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 { + 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, + }, + }) + } +} + impl PartialEq for PhysicalSortExpr { fn eq(&self, other: &Self) -> bool { self.options == other.options && self.expr.eq(&other.expr) diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs index 59d36c4efc1bb..9a24658411300 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -515,6 +515,146 @@ 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 { + 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 = range + .ordering() + .iter() + .map(|sort_expr| sort_expr.try_to_proto(ctx)) + .collect::>>()?; + 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::>>()?; + Ok(protobuf::PhysicalRangeSplitPoint { value }) + }) + .collect::>>()?; + protobuf::partitioning::PartitionMethod::Range( + protobuf::PhysicalRangePartitioning { + sort_expr, + split_point, + }, + ) + } + Partitioning::UnknownPartitioning(n) => { + protobuf::partitioning::PartitionMethod::Unknown(*n as u64) + } + }; + 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> { + 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::>>()?; + 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 = range + .sort_expr + .iter() + .map(|sort_expr| PhysicalSortExpr::try_from_proto(sort_expr, ctx)) + .collect::>>()?; + 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::>>()?; + Ok(SplitPoint::new(values)) + }) + .collect::>>()?; + 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::try_from(count).map_err(|_| { + datafusion_common::internal_datafusion_err!( + "Partition count {count} exceeds usize::MAX" + ) + }) +} + impl PartialEq for Partitioning { fn eq(&self, other: &Partitioning) -> bool { match (self, other) { diff --git a/datafusion/physical-plan/src/proto.rs b/datafusion/physical-plan/src/proto.rs index 1731203f6c767..5883f1d1fec81 100644 --- a/datafusion/physical-plan/src/proto.rs +++ b/datafusion/physical-plan/src/proto.rs @@ -64,6 +64,12 @@ use datafusion_common::{Result, internal_datafusion_err}; use datafusion_execution::TaskContext; use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF}; use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr_common::physical_expr::proto_decode::{ + PhysicalExprDecode, PhysicalExprDecodeCtx, +}; +use datafusion_physical_expr_common::physical_expr::proto_encode::{ + PhysicalExprEncode, PhysicalExprEncodeCtx, +}; use datafusion_proto_models::protobuf::{PhysicalExprNode, PhysicalPlanNode}; use crate::ExecutionPlan; @@ -190,6 +196,25 @@ impl<'a> ExecutionPlanEncodeCtx<'a> { pub fn encode_udwf(&self, udwf: &WindowUDF) -> Result>> { 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. +impl PhysicalExprEncode for ExecutionPlanEncodeCtx<'_> { + fn encode(&self, expr: &Arc) -> Result { + self.encode_expr(expr) + } } /// Context handed to a plan's `try_from_proto` associated function. @@ -286,6 +311,28 @@ impl<'a> ExecutionPlanDecodeCtx<'a> { ) -> Result> { self.decoder.decode_udwf(name, payload) } + + /// An expression-level decode context backed by this plan context, bound to + /// `input_schema`. + /// + /// The decode counterpart of + /// [`ExecutionPlanEncodeCtx::expr_ctx`], for calling conversions such as + /// [`Partitioning::try_from_proto`](datafusion_physical_expr::Partitioning::try_from_proto). + pub fn expr_ctx<'s>(&'s self, input_schema: &'s Schema) -> PhysicalExprDecodeCtx<'s> { + PhysicalExprDecodeCtx::new(input_schema, self) + } +} + +/// Lets [`ExecutionPlanDecodeCtx`] back a [`PhysicalExprDecodeCtx`], so +/// expression-level conversions can be reused from plan hooks. +impl PhysicalExprDecode for ExecutionPlanDecodeCtx<'_> { + fn decode( + &self, + node: &PhysicalExprNode, + schema: &Schema, + ) -> Result> { + self.decode_expr(node, schema) + } } /// Assert that a [`PhysicalPlanNode`] carries the expected `PhysicalPlanType` diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 3473aad9b3fc0..873f35fd6aed9 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -1713,64 +1713,14 @@ impl ExecutionPlan for RepartitionExec { let input = ctx.encode_child(self.input())?; - // Keep the existing protobuf wire representation unchanged. - let partition_method = match self.partitioning() { - Partitioning::RoundRobinBatch(n) => { - protobuf::partitioning::PartitionMethod::RoundRobin(*n as u64) - } - Partitioning::Hash(exprs, n) => { - let hash_expr = ctx.encode_expressions(exprs)?; - protobuf::partitioning::PartitionMethod::Hash( - protobuf::PhysicalHashRepartition { - hash_expr, - partition_count: *n as u64, - }, - ) - } - Partitioning::Range(range) => { - let sort_expr = range - .ordering() - .iter() - .map(|sort_expr| { - Ok(protobuf::PhysicalSortExprNode { - expr: Some(Box::new(ctx.encode_expr(&sort_expr.expr)?)), - asc: !sort_expr.options.descending, - nulls_first: sort_expr.options.nulls_first, - }) - }) - .collect::>>()?; - 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::>>()?; - Ok(protobuf::PhysicalRangeSplitPoint { value }) - }) - .collect::>>()?; - protobuf::partitioning::PartitionMethod::Range( - protobuf::PhysicalRangePartitioning { - sort_expr, - split_point, - }, - ) - } - Partitioning::UnknownPartitioning(n) => { - protobuf::partitioning::PartitionMethod::Unknown(*n as u64) - } - }; + let partitioning = self.partitioning().try_to_proto(&ctx.expr_ctx())?; Ok(Some(protobuf::PhysicalPlanNode { physical_plan_type: Some( protobuf::physical_plan_node::PhysicalPlanType::Repartition(Box::new( protobuf::RepartitionExecNode { input: Some(Box::new(input)), - partitioning: Some(protobuf::Partitioning { - partition_method: Some(partition_method), - }), + partitioning: Some(partitioning), preserve_order: self.preserve_order(), }, )), @@ -1800,84 +1750,23 @@ impl RepartitionExec { )?; let input_schema = input.schema(); - let partition_method = repart + let partitioning = repart .partitioning .as_ref() - .and_then(|p| p.partition_method.as_ref()) + .map(|partitioning| { + Partitioning::try_from_proto( + partitioning, + &ctx.expr_ctx(input_schema.as_ref()), + ) + }) + .transpose()? + .flatten() .ok_or_else(|| { datafusion_common::internal_datafusion_err!( "RepartitionExec is missing required field 'partitioning'" ) })?; - let partitioning = match partition_method { - protobuf::partitioning::PartitionMethod::RoundRobin(n) => { - Partitioning::RoundRobinBatch(*n as usize) - } - protobuf::partitioning::PartitionMethod::Hash(hash) => { - let exprs = hash - .hash_expr - .iter() - .map(|expr| ctx.decode_expr(expr, input_schema.as_ref())) - .collect::>>()?; - let partition_count = - usize::try_from(hash.partition_count).map_err(|_| { - datafusion_common::internal_datafusion_err!( - "Hash partition count {} exceeds usize::MAX", - hash.partition_count - ) - })?; - Partitioning::Hash(exprs, partition_count) - } - protobuf::partitioning::PartitionMethod::Unknown(n) => { - Partitioning::UnknownPartitioning(*n as usize) - } - protobuf::partitioning::PartitionMethod::Range(range) => { - let sort_exprs = range - .sort_expr - .iter() - .map(|sort_expr| { - let expr = sort_expr.expr.as_ref().ok_or_else(|| { - datafusion_common::internal_datafusion_err!( - "Unexpected empty physical expression" - ) - })?; - Ok(PhysicalSortExpr { - expr: ctx.decode_expr(expr, input_schema.as_ref())?, - options: SortOptions { - descending: !sort_expr.asc, - nulls_first: sort_expr.nulls_first, - }, - }) - }) - .collect::>>()?; - let sort_expr_count = sort_exprs.len(); - let ordering = LexOrdering::new(sort_exprs).ok_or_else(|| { - datafusion_common::internal_datafusion_err!( - "Range partitioning requires non-empty ordering" - ) - })?; - if ordering.len() != sort_expr_count { - return datafusion_common::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::>>()?; - Ok(SplitPoint::new(values)) - }) - .collect::>>()?; - Partitioning::Range(RangePartitioning::try_new(ordering, split_points)?) - } - }; - let mut repart_exec = RepartitionExec::try_new(input, partitioning)?; if repart.preserve_order { repart_exec = repart_exec.with_preserve_order(); diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index dd88646054d95..fcefe402b175a 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -23,10 +23,7 @@ use arrow::array::RecordBatch; use arrow::compute::SortOptions; use arrow::datatypes::{Field, Schema}; use arrow::ipc::reader::StreamReader; -use chrono::{TimeZone, Utc}; -use datafusion_common::{ - DataFusionError, Result, ScalarValue, internal_datafusion_err, not_impl_err, -}; +use datafusion_common::{DataFusionError, Result, internal_datafusion_err, not_impl_err}; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_groups::FileGroup; use datafusion_datasource::file_scan_config::FileScanConfig; @@ -43,7 +40,7 @@ use datafusion_expr::dml::InsertOp; use datafusion_physical_expr::expressions::{LambdaExpr, LambdaVariable}; use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr; use datafusion_physical_expr::{ - HigherOrderFunctionExpr, LexOrdering, PhysicalSortExpr, ScalarFunctionExpr, + HigherOrderFunctionExpr, PhysicalSortExpr, ScalarFunctionExpr, }; use datafusion_physical_plan::expressions::{ BinaryExpr, CaseExpr, CastExpr, Column, InListExpr, IsNotNullExpr, IsNullExpr, @@ -52,12 +49,8 @@ use datafusion_physical_plan::expressions::{ use datafusion_physical_plan::joins::HashExpr; use datafusion_physical_plan::proto::ExecutionPlanDecodeCtx; use datafusion_physical_plan::windows::{create_window_expr, schema_add_window_field}; -use datafusion_physical_plan::{ - Partitioning, PhysicalExpr, RangePartitioning, SplitPoint, WindowExpr, -}; +use datafusion_physical_plan::{Partitioning, PhysicalExpr, WindowExpr}; use datafusion_proto_common::common::proto_error; -use object_store::ObjectMeta; -use object_store::path::Path; use super::{ ConverterPlanDecoder, DefaultPhysicalProtoConverter, PhysicalExtensionCodec, @@ -426,83 +419,20 @@ pub fn parse_protobuf_partitioning( input_schema: &Schema, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - match partitioning { - Some(protobuf::Partitioning { partition_method }) => match partition_method { - Some(protobuf::partitioning::PartitionMethod::RoundRobin( - partition_count, - )) => Ok(Some(Partitioning::RoundRobinBatch( - *partition_count as usize, - ))), - Some(protobuf::partitioning::PartitionMethod::Hash(hash_repartition)) => { - parse_protobuf_hash_partitioning( - Some(hash_repartition), - ctx, - input_schema, - proto_converter, - ) - } - Some(protobuf::partitioning::PartitionMethod::Range(range_partitioning)) => { - Ok(Some(parse_protobuf_range_partitioning( - range_partitioning, - ctx, - input_schema, - proto_converter, - )?)) - } - Some(protobuf::partitioning::PartitionMethod::Unknown(partition_count)) => { - Ok(Some(Partitioning::UnknownPartitioning( - *partition_count as usize, - ))) - } - None => Ok(None), - }, - None => Ok(None), - } -} - -fn parse_protobuf_range_partitioning( - range_partitioning: &protobuf::PhysicalRangePartitioning, - ctx: &PhysicalPlanDecodeContext<'_>, - input_schema: &Schema, - proto_converter: &dyn PhysicalProtoConverterExtension, -) -> Result { - let sort_exprs = parse_physical_sort_exprs( - &range_partitioning.sort_expr, + let decoder = ConverterDecoder { ctx, - input_schema, proto_converter, - )?; - 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 Err(internal_datafusion_err!( - "Range partitioning ordering must not contain duplicate expressions" - )); - } - let split_points = range_partitioning - .split_point - .iter() - .map(parse_protobuf_range_split_point) - .collect::>()?; - Ok(Partitioning::Range(RangePartitioning::try_new( - ordering, - split_points, - )?)) -} - -fn parse_protobuf_range_split_point( - split_point: &protobuf::PhysicalRangeSplitPoint, -) -> Result { - let values = split_point - .value - .iter() - .map(|value| ScalarValue::try_from(value).map_err(Into::into)) - .collect::>()?; - Ok(SplitPoint::new(values)) + }; + let decode_ctx = + datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx::new( + input_schema, + &decoder, + ); + partitioning + .map(|partitioning| Partitioning::try_from_proto(partitioning, &decode_ctx)) + .transpose() + .map(Option::flatten) } - pub fn parse_protobuf_file_scan_schema( proto: &protobuf::FileScanExecConf, ) -> Result> { @@ -510,6 +440,10 @@ pub fn parse_protobuf_file_scan_schema( } /// Parses a TableSchema from protobuf, extracting the file schema and partition columns +#[deprecated( + since = "55.0.0", + note = "use `FileScanConfig::parse_table_schema_from_proto` instead" +)] pub fn parse_table_schema_from_proto( proto: &protobuf::FileScanExecConf, ) -> Result { @@ -526,7 +460,7 @@ pub fn parse_protobuf_file_scan_config( ctx, proto_converter, }; - FileScanConfig::from_proto_conf( + FileScanConfig::try_from_proto( proto, &ExecutionPlanDecodeCtx::new(&decoder), file_source, @@ -545,61 +479,30 @@ pub fn parse_record_batches(buf: &[u8]) -> Result> { Ok(batches) } +/// Thin shim over [`PartitionedFile::try_from_proto`], which owns the wire logic. impl TryFromProto<&protobuf::PartitionedFile> for PartitionedFile { type Error = DataFusionError; fn try_from_proto(val: &protobuf::PartitionedFile) -> Result { - let mut pf = PartitionedFile::new_from_meta(ObjectMeta { - location: Path::parse(val.path.as_str()) - .map_err(|e| proto_error(format!("Invalid object_store path: {e}")))?, - last_modified: Utc.timestamp_nanos(val.last_modified_ns as i64), - size: val.size, - e_tag: None, - version: None, - }) - .with_partition_values( - val.partition_values - .iter() - .map(|v| v.try_into()) - .collect::, _>>()?, - ); - if let Some(proto_schema) = val.arrow_schema.as_ref() { - pf = pf.with_arrow_schema(Arc::new( - proto_schema.try_into().map_err(DataFusionError::from)?, - )); - } - if let Some(range) = val.range.as_ref() { - let file_range = FileRange::try_from_proto(range)?; - pf = pf.with_range(file_range.start, file_range.end); - } - if let Some(proto_stats) = val.statistics.as_ref() { - pf = pf.with_statistics(Arc::new(proto_stats.try_into()?)); - } - Ok(pf) + PartitionedFile::try_from_proto(val) } } +/// Thin shim over [`FileRange::try_from_proto`], which owns the wire logic. impl TryFromProto<&protobuf::FileRange> for FileRange { type Error = DataFusionError; fn try_from_proto(value: &protobuf::FileRange) -> Result { - Ok(FileRange { - start: value.start, - end: value.end, - }) + FileRange::try_from_proto(value) } } +/// Thin shim over [`FileGroup::try_from_proto`], which owns the wire logic. impl TryFromProto<&protobuf::FileGroup> for FileGroup { type Error = DataFusionError; fn try_from_proto(val: &protobuf::FileGroup) -> Result { - let files = val - .files - .iter() - .map(PartitionedFile::try_from_proto) - .collect::, _>>()?; - Ok(FileGroup::new(files)) + FileGroup::try_from_proto(val) } } @@ -721,6 +624,10 @@ impl datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprD #[cfg(test)] mod tests { + use chrono::{TimeZone, Utc}; + use object_store::ObjectMeta; + use object_store::path::Path; + use super::*; #[test] diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index b70275e84fc2d..77e05cdc40109 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -100,7 +100,7 @@ use crate::convert::TryFromProto; use crate::convert_required; use crate::physical_plan::from_proto::{ parse_physical_expr_with_converter, parse_physical_sort_exprs, - parse_protobuf_file_scan_config, parse_record_batches, parse_table_schema_from_proto, + parse_protobuf_file_scan_config, parse_record_batches, }; use crate::physical_plan::to_proto::{ serialize_file_scan_config, serialize_physical_expr_with_converter, @@ -309,7 +309,7 @@ mod file_scan_config_serde { codec: &self.codec, proto_converter: &self.converter, }; - config.to_proto_conf(&ExecutionPlanEncodeCtx::new(&encoder)) + config.try_to_proto(&ExecutionPlanEncodeCtx::new(&encoder)) } fn decode(&self, conf: &protobuf::FileScanExecConf) -> Result { @@ -327,7 +327,7 @@ mod file_scan_config_serde { ctx: &physical_decode_ctx, proto_converter: &self.converter, }; - FileScanConfig::from_proto_conf( + FileScanConfig::try_from_proto( conf, &ExecutionPlanDecodeCtx::new(&decoder), file_source, @@ -1381,8 +1381,9 @@ pub trait PhysicalPlanNodeExt: Sized { }; // Parse table schema with partition columns - let table_schema = - parse_table_schema_from_proto(scan.base_conf.as_ref().unwrap())?; + let table_schema = FileScanConfig::parse_table_schema_from_proto( + scan.base_conf.as_ref().unwrap(), + )?; let csv_options = CsvOptions { has_header: Some(scan.has_header), @@ -1416,7 +1417,7 @@ pub trait PhysicalPlanNodeExt: Sized { proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { let base_conf = scan.base_conf.as_ref().unwrap(); - let table_schema = parse_table_schema_from_proto(base_conf)?; + let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?; let scan_conf = parse_protobuf_file_scan_config( base_conf, ctx, @@ -1435,7 +1436,7 @@ pub trait PhysicalPlanNodeExt: Sized { let base_conf = scan.base_conf.as_ref().ok_or_else(|| { internal_datafusion_err!("base_conf in ArrowScanExecNode is missing.") })?; - let table_schema = parse_table_schema_from_proto(base_conf)?; + let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?; let scan_conf = parse_protobuf_file_scan_config( base_conf, ctx, @@ -1490,7 +1491,7 @@ pub trait PhysicalPlanNodeExt: Sized { } // Parse table schema with partition columns - let table_schema = parse_table_schema_from_proto(base_conf)?; + let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?; let object_store_url = match base_conf.object_store_url.is_empty() { false => ObjectStoreUrl::parse(&base_conf.object_store_url)?, true => ObjectStoreUrl::local_filesystem(), @@ -1537,8 +1538,9 @@ pub trait PhysicalPlanNodeExt: Sized { ) -> Result> { #[cfg(feature = "avro")] { - let table_schema = - parse_table_schema_from_proto(scan.base_conf.as_ref().unwrap())?; + let table_schema = FileScanConfig::parse_table_schema_from_proto( + scan.base_conf.as_ref().unwrap(), + )?; let conf = parse_protobuf_file_scan_config( scan.base_conf.as_ref().unwrap(), ctx, diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index 56cacabccc859..04252176d2845 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -37,9 +37,7 @@ use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use datafusion_physical_plan::proto::ExecutionPlanEncodeCtx; use datafusion_physical_plan::udaf::AggregateFunctionExpr; use datafusion_physical_plan::windows::{PlainAggregateWindowExpr, WindowUDFExpr}; -use datafusion_physical_plan::{ - Partitioning, PhysicalExpr, RangePartitioning, SplitPoint, WindowExpr, -}; +use datafusion_physical_plan::{Partitioning, PhysicalExpr, WindowExpr}; use super::{ ConverterPlanEncoder, DefaultPhysicalProtoConverter, PhysicalExtensionCodec, @@ -358,117 +356,33 @@ pub fn serialize_partitioning( codec: &dyn PhysicalExtensionCodec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result { - let serialized_partitioning = match partitioning { - Partitioning::RoundRobinBatch(partition_count) => protobuf::Partitioning { - partition_method: Some(protobuf::partitioning::PartitionMethod::RoundRobin( - *partition_count as u64, - )), - }, - Partitioning::Hash(exprs, partition_count) => { - let serialized_exprs = - serialize_physical_exprs(exprs, codec, proto_converter)?; - protobuf::Partitioning { - partition_method: Some(protobuf::partitioning::PartitionMethod::Hash( - protobuf::PhysicalHashRepartition { - hash_expr: serialized_exprs, - partition_count: *partition_count as u64, - }, - )), - } - } - Partitioning::Range(range) => protobuf::Partitioning { - partition_method: Some(protobuf::partitioning::PartitionMethod::Range( - serialize_range_partitioning(range, codec, proto_converter)?, - )), - }, - Partitioning::UnknownPartitioning(partition_count) => protobuf::Partitioning { - partition_method: Some(protobuf::partitioning::PartitionMethod::Unknown( - *partition_count as u64, - )), - }, + let encoder = ConverterEncoder { + codec, + proto_converter, }; - Ok(serialized_partitioning) -} - -fn serialize_range_partitioning( - range: &RangePartitioning, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, -) -> Result { - Ok(protobuf::PhysicalRangePartitioning { - sort_expr: serialize_physical_sort_exprs( - range.ordering().iter().cloned(), - codec, - proto_converter, - )?, - split_point: range - .split_points() - .iter() - .map(serialize_range_split_point) - .collect::>()?, - }) -} - -fn serialize_range_split_point( - split_point: &SplitPoint, -) -> Result { - Ok(protobuf::PhysicalRangeSplitPoint { - value: split_point - .values() - .iter() - .map(|value| { - TryInto::::try_into(value) - .map_err(Into::into) - }) - .collect::>()?, - }) + partitioning.try_to_proto( + &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx::new(&encoder), + ) } - +/// Thin shim over [`PartitionedFile::try_to_proto`], which owns the wire logic. impl TryFromProto<&PartitionedFile> for protobuf::PartitionedFile { type Error = DataFusionError; fn try_from_proto(pf: &PartitionedFile) -> Result { - let last_modified = pf.object_meta.last_modified; - let last_modified_ns = last_modified.timestamp_nanos_opt().ok_or_else(|| { - DataFusionError::Plan(format!( - "Invalid timestamp on PartitionedFile::ObjectMeta: {last_modified}" - )) - })? as u64; - Ok(protobuf::PartitionedFile { - arrow_schema: pf - .arrow_schema - .as_ref() - .map(|s| s.as_ref().try_into()) - .transpose()?, - path: pf.object_meta.location.as_ref().to_owned(), - size: pf.object_meta.size, - last_modified_ns, - partition_values: pf - .partition_values - .iter() - .map(|v| v.try_into()) - .collect::, _>>()?, - range: pf - .range - .as_ref() - .map(protobuf::FileRange::try_from_proto) - .transpose()?, - statistics: pf.statistics.as_ref().map(|s| s.as_ref().into()), - }) + pf.try_to_proto() } } +/// Thin shim over [`FileRange::try_to_proto`], which owns the wire logic. impl TryFromProto<&FileRange> for protobuf::FileRange { type Error = DataFusionError; fn try_from_proto(value: &FileRange) -> Result { - Ok(protobuf::FileRange { - start: value.start, - end: value.end, - }) + value.try_to_proto() } } +/// Thin shim over [`PartitionedFile::try_to_proto`], which owns the wire logic. impl TryFromProto<&[PartitionedFile]> for protobuf::FileGroup { type Error = DataFusionError; @@ -476,8 +390,8 @@ impl TryFromProto<&[PartitionedFile]> for protobuf::FileGroup { Ok(protobuf::FileGroup { files: gr .iter() - .map(protobuf::PartitionedFile::try_from_proto) - .collect::, _>>()?, + .map(PartitionedFile::try_to_proto) + .collect::>>()?, }) } } @@ -491,7 +405,7 @@ pub fn serialize_file_scan_config( codec, proto_converter, }; - conf.to_proto_conf(&ExecutionPlanEncodeCtx::new(&encoder)) + conf.try_to_proto(&ExecutionPlanEncodeCtx::new(&encoder)) } pub fn serialize_maybe_filter( diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md b/docs/source/library-user-guide/upgrading/55.0.0.md index 6097c8dc717df..a4218b592e4c2 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -894,6 +894,27 @@ let plan = deserialize_bytes(&proto_bytes)?; See [PR #23827](https://github.com/apache/datafusion/pull/23827) for details. +### `parse_table_schema_from_proto` is deprecated + +`datafusion_proto::physical_plan::from_proto::parse_table_schema_from_proto` is +deprecated in favor of the equivalent method on `FileScanConfig`, which lives in +`datafusion-datasource` alongside the rest of the shared file-scan protobuf +conversions: + +```rust,ignore +// Before +let table_schema = parse_table_schema_from_proto(base_conf)?; + +// After +let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?; +``` + +`PartitionedFile`, `FileGroup` and `FileRange` also gained inherent +`try_to_proto` / `try_from_proto` methods (in `datafusion_datasource::proto`, +behind the `proto` feature) which now own that wire logic; +`datafusion-proto`'s `TryFromProto` implementations for those types delegate to +them and keep working unchanged. + ### `MSRV` updated to 1.94.0 The Minimum Supported Rust Version (MSRV) has been updated to [`1.94.0`].