diff --git a/Cargo.lock b/Cargo.lock index a41734e064d4c..fa1b8adcbd9f6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1942,6 +1942,7 @@ dependencies = [ "datafusion-physical-expr-adapter", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "flate2", "futures", diff --git a/datafusion/datasource/Cargo.toml b/datafusion/datasource/Cargo.toml index 2ac42ed900095..77b10a171f3fc 100644 --- a/datafusion/datasource/Cargo.toml +++ b/datafusion/datasource/Cargo.toml @@ -34,6 +34,14 @@ all-features = true backtrace = ["datafusion-common/backtrace"] compression = ["async-compression", "liblzma", "bzip2", "flate2", "zstd", "tokio-util"] default = ["compression"] +# Enables `DataSource::try_to_proto` / `FileSource::try_to_proto` serialization +# hooks and the shared `FileScanConfig` <-> proto conversion. Off by default so +# consumers that never serialize plans pay nothing. Mirrors the `proto` feature +# on `datafusion-physical-plan`. +proto = [ + "dep:datafusion-proto-models", + "datafusion-physical-plan/proto", +] [dependencies] arrow = { workspace = true } @@ -56,6 +64,7 @@ datafusion-physical-expr = { workspace = true } datafusion-physical-expr-adapter = { workspace = true } datafusion-physical-expr-common = { workspace = true } datafusion-physical-plan = { workspace = true } +datafusion-proto-models = { workspace = true, optional = true } datafusion-session = { workspace = true } flate2 = { workspace = true, optional = true } futures = { workspace = true } diff --git a/datafusion/datasource/src/file.rs b/datafusion/datasource/src/file.rs index 07460b23694b7..691bb314b7c03 100644 --- a/datafusion/datasource/src/file.rs +++ b/datafusion/datasource/src/file.rs @@ -351,6 +351,31 @@ pub trait FileSource: Any + Send + Sync { fn schema_adapter_factory(&self) -> Option> { None } + + /// Serialize this file source into a full [`PhysicalPlanNode`] (a + /// `DataSourceExec` wrapping the `FileScanConfig`), if it knows how. + /// + /// `base` is the shared [`FileScanConfig`] this source is wrapped in; the + /// format-agnostic parts (file groups, schema, statistics, ordering, + /// projection, …) are encoded via + /// [`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. + /// + /// * `Ok(None)` (the default) — this source has no proto hook yet; the + /// caller falls back to the central downcast chain in `datafusion-proto`. + /// * `Ok(Some(node))` — fully serialized; the caller must not fall back. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`FileScanConfig`]: crate::file_scan_config::FileScanConfig + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _base: &FileScanConfig, + _ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + Ok(None) + } } impl dyn FileSource { diff --git a/datafusion/datasource/src/file_scan_config/mod.rs b/datafusion/datasource/src/file_scan_config/mod.rs index 962df06302386..d02ecff914396 100644 --- a/datafusion/datasource/src/file_scan_config/mod.rs +++ b/datafusion/datasource/src/file_scan_config/mod.rs @@ -20,6 +20,12 @@ pub(crate) mod sort_pushdown; +/// Shared `FileScanConfig` <-> proto conversion, gated on the `proto` feature. +/// 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; + use crate::file_groups::FileGroup; use crate::{ PartitionedFile, display::FileGroupsDisplay, file::FileSource, @@ -1175,6 +1181,18 @@ impl DataSource for FileScanConfig { Some(Arc::new(SharedWorkSource::from_config(self)) as Arc) } + + /// Serialize this file scan by delegating to the concrete + /// [`FileSource`]'s + /// [`try_to_proto`](crate::file::FileSource::try_to_proto) hook, passing + /// `self` as the shared spine it needs to emit the base config. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + self.file_source().try_to_proto(self, ctx) + } } impl FileScanConfig { @@ -1566,12 +1584,18 @@ mod tests { use datafusion_common::{Result, assert_batches_eq, internal_err}; use datafusion_execution::TaskContext; use datafusion_expr::SortExpr; + #[cfg(feature = "proto")] + use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF}; use datafusion_physical_expr::create_physical_sort_expr; use datafusion_physical_expr::expressions::Literal; use datafusion_physical_expr::projection::ProjectionExpr; use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::execution_plan::collect; + #[cfg(feature = "proto")] + use datafusion_physical_plan::proto::{ExecutionPlanEncode, ExecutionPlanEncodeCtx}; + #[cfg(feature = "proto")] + use datafusion_proto_models::protobuf::{PhysicalExprNode, PhysicalPlanNode}; use futures::FutureExt as _; use futures::StreamExt as _; use futures::stream; @@ -1630,6 +1654,108 @@ mod tests { } } + #[cfg(feature = "proto")] + #[derive(Clone)] + struct ProtoHookSource { + metrics: ExecutionPlanMetricsSet, + table_schema: TableSchema, + } + + #[cfg(feature = "proto")] + impl ProtoHookSource { + fn new(table_schema: TableSchema) -> Self { + Self { + metrics: ExecutionPlanMetricsSet::new(), + table_schema, + } + } + } + + #[cfg(feature = "proto")] + impl FileSource for ProtoHookSource { + fn create_file_opener( + &self, + _object_store: Arc, + _base_config: &FileScanConfig, + _partition: usize, + ) -> Result> { + internal_err!("not needed for proto delegation test") + } + + fn table_schema(&self) -> &TableSchema { + &self.table_schema + } + + fn with_batch_size(&self, _batch_size: usize) -> Arc { + Arc::new(self.clone()) + } + + fn metrics(&self) -> &ExecutionPlanMetricsSet { + &self.metrics + } + + fn file_type(&self) -> &str { + "proto-hook-test" + } + + fn try_to_proto( + &self, + _base: &FileScanConfig, + _ctx: &ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + Ok(Some(PhysicalPlanNode::default())) + } + } + + #[cfg(feature = "proto")] + struct UnusedPlanEncoder; + + #[cfg(feature = "proto")] + impl ExecutionPlanEncode for UnusedPlanEncoder { + fn encode_plan( + &self, + _plan: &Arc, + ) -> Result { + internal_err!("not needed for proto delegation test") + } + + fn encode_expr(&self, _expr: &Arc) -> Result { + internal_err!("not needed for proto delegation test") + } + + fn encode_udf(&self, _udf: &ScalarUDF) -> Result>> { + internal_err!("not needed for proto delegation test") + } + + fn encode_udaf(&self, _udaf: &AggregateUDF) -> Result>> { + internal_err!("not needed for proto delegation test") + } + + fn encode_udwf(&self, _udwf: &WindowUDF) -> Result>> { + internal_err!("not needed for proto delegation test") + } + } + + #[cfg(feature = "proto")] + #[test] + fn data_source_exec_delegates_proto_to_file_source() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + false, + )])); + let source = Arc::new(ProtoHookSource::new(TableSchema::from(&schema))); + let config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + .build(); + let exec = DataSourceExec::from_data_source(config); + let encoder = UnusedPlanEncoder; + let ctx = ExecutionPlanEncodeCtx::new(&encoder); + + assert_eq!(exec.try_to_proto(&ctx)?, Some(PhysicalPlanNode::default())); + Ok(()) + } + #[test] fn physical_plan_config_no_projection_tab_cols_as_field() { let file_schema = aggr_test_schema(); diff --git a/datafusion/datasource/src/file_scan_config/proto.rs b/datafusion/datasource/src/file_scan_config/proto.rs new file mode 100644 index 0000000000000..aac2364b5f030 --- /dev/null +++ b/datafusion/datasource/src/file_scan_config/proto.rs @@ -0,0 +1,296 @@ +// 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 serialization of the format-agnostic [`FileScanConfig`] spine. +//! +//! This is the relocated body of `datafusion-proto`'s +//! `serialize_file_scan_config` / `parse_protobuf_file_scan_config`, ported to +//! ride the +//! [`ExecutionPlanEncodeCtx`](datafusion_physical_plan::proto::ExecutionPlanEncodeCtx) / +//! [`ExecutionPlanDecodeCtx`](datafusion_physical_plan::proto::ExecutionPlanDecodeCtx) +//! instead of the raw `PhysicalExtensionCodec` + +//! `PhysicalProtoConverterExtension`. Every +//! `FileSource::try_to_proto` hook (CSV, JSON, Arrow, Parquet, Avro) builds its +//! `*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. +//! +//! 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::datatypes::Schema; +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}; +use datafusion_physical_plan::proto::{ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx}; +use datafusion_proto_models::protobuf; + +use crate::file::FileSource; +use crate::file_groups::FileGroup; +use crate::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; +use crate::table_schema::TableSchema; + +impl FileScanConfig { + /// Serialize the shared, format-agnostic part of a file scan into a + /// [`protobuf::FileScanExecConf`]. + /// + /// 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`. + /// + /// 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(FileGroup::try_to_proto) + .collect::>>()?; + + let expr_ctx = ctx.expr_ctx(); + let mut output_ordering = vec![]; + for order in &self.output_ordering { + output_ordering.push(protobuf::PhysicalSortExprNodeCollection { + 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(|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 `try_from_proto`. + let mut fields = self + .file_schema() + .fields() + .iter() + .cloned() + .collect::>(); + fields.extend(self.table_partition_cols().iter().cloned()); + let schema = + Schema::new(fields).with_metadata(self.file_schema().metadata.clone()); + + let projection_exprs = self + .file_source() + .projection() + .as_ref() + .map(|projection_exprs| { + Ok::<_, DataFusionError>(protobuf::ProjectionExprs { + projections: projection_exprs + .iter() + .map(|expr| { + Ok(protobuf::ProjectionExpr { + alias: expr.alias.to_string(), + expr: Some(ctx.encode_expr(&expr.expr)?), + }) + }) + .collect::>>()?, + }) + }) + .transpose()?; + + Ok(protobuf::FileScanExecConf { + file_groups, + statistics: Some((&self.statistics()).into()), + limit: self.limit.map(|l| protobuf::ScanLimit { limit: l as u32 }), + projection: vec![], + schema: Some((&schema).try_into()?), + table_partition_cols: self + .table_partition_cols() + .iter() + .map(|x| x.name().clone()) + .collect::>(), + object_store_url: self.object_store_url.to_string(), + output_ordering, + constraints: Some(self.constraints.clone().into()), + batch_size: self.batch_size.map(|s| s as u64), + projection_exprs, + output_partitioning, + }) + } + + /// Reconstruct a [`FileScanConfig`] from a [`protobuf::FileScanExecConf`] + /// and a `file_source` the caller has already rebuilt (typically from the + /// table schema via [`FileScanConfig::parse_table_schema_from_proto`]). + /// + /// Byte-compatible with the former `parse_protobuf_file_scan_config`. + pub fn try_from_proto( + conf: &protobuf::FileScanExecConf, + ctx: &ExecutionPlanDecodeCtx<'_>, + file_source: Arc, + ) -> Result { + let schema = parse_file_scan_schema(conf)?; + + let constraints = conf + .constraints + .as_ref() + .ok_or_else(|| { + internal_datafusion_err!( + "FileScanExecConf is missing required field 'constraints'" + ) + })? + .try_into()?; + let statistics = conf + .statistics + .as_ref() + .ok_or_else(|| { + internal_datafusion_err!( + "FileScanExecConf is missing required field 'statistics'" + ) + })? + .try_into()?; + + let file_groups = conf + .file_groups + .iter() + .map(FileGroup::try_from_proto) + .collect::>>()?; + + let object_store_url = match conf.object_store_url.is_empty() { + false => ObjectStoreUrl::parse(&conf.object_store_url)?, + 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 = 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 = 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 { + let projection_exprs: Vec = proto_projection_exprs + .projections + .iter() + .map(|proto_expr| { + let expr = ctx.decode_expr( + proto_expr.expr.as_ref().ok_or_else(|| { + internal_datafusion_err!("ProjectionExpr missing expr field") + })?, + &schema, + )?; + Ok(ProjectionExpr::new(expr, proto_expr.alias.clone())) + }) + .collect::>>()?; + + let projection_exprs = ProjectionExprs::new(projection_exprs); + + file_source + .try_pushdown_projection(&projection_exprs)? + .unwrap_or(file_source) + } else { + file_source + }; + + let config_builder = FileScanConfigBuilder::new(object_store_url, file_source) + .with_file_groups(file_groups) + .with_constraints(constraints) + .with_statistics(statistics) + .with_limit(conf.limit.as_ref().map(|sl| sl.limit as usize)) + .with_output_ordering(output_ordering) + .with_output_partitioning(output_partitioning) + .with_batch_size(conf.batch_size.map(|s| s as usize)); + Ok(config_builder.build()) + } + + /// Parse a [`TableSchema`] (file schema + partition columns) from a + /// [`protobuf::FileScanExecConf`]. File sources use this to rebuild their + /// 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( + conf: &protobuf::FileScanExecConf, + ) -> Result { + let schema = parse_file_scan_schema(conf)?; + + // Reacquire the partition column types from the schema before removing + // them below. + let table_partition_cols = conf + .table_partition_cols + .iter() + .map(|col| Ok(Arc::new(schema.field_with_name(col)?.clone()))) + .collect::>>()?; + + // Remove partition columns from the schema after recreating + // table_partition_cols because the partition columns are not in the + // file. They are present to allow the partition column types to be + // reconstructed after serde. + let file_schema = Arc::new( + Schema::new( + schema + .fields() + .iter() + .filter(|field| !table_partition_cols.contains(field)) + .cloned() + .collect::>(), + ) + .with_metadata(schema.metadata.clone()), + ); + + Ok(TableSchema::builder(file_schema) + .with_table_partition_cols(table_partition_cols) + .build()) + } +} + +/// Parse the full (file + partition columns) schema off the base conf. +fn parse_file_scan_schema(conf: &protobuf::FileScanExecConf) -> Result> { + let schema: Schema = conf + .schema + .as_ref() + .ok_or_else(|| { + internal_datafusion_err!( + "FileScanExecConf is missing required field 'schema'" + ) + })? + .try_into()?; + Ok(Arc::new(schema)) +} 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/datasource/src/source.rs b/datafusion/datasource/src/source.rs index c280470bb0d0b..1fd5f865c45ab 100644 --- a/datafusion/datasource/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -264,6 +264,30 @@ pub trait DataSource: Any + Send + Sync + Debug { fn open_with_args(&self, args: OpenArgs) -> Result { self.open(args.partition, args.context) } + + /// Serialize this data source to a full [`PhysicalPlanNode`] (a + /// `DataSourceExec` wrapping this source), if it knows how. + /// + /// This is the `DataSource` analog of + /// [`ExecutionPlan::try_to_proto`]. + /// [`DataSourceExec::try_to_proto`](crate::source::DataSourceExec) delegates + /// to this hook, which for file scans forwards to + /// [`FileSource::try_to_proto`] + /// through the shared [`FileScanConfig`] + /// spine. + /// + /// * `Ok(None)` (the default) — "I don't serialize myself"; the caller falls + /// back to the central downcast chain in `datafusion-proto`. + /// * `Ok(Some(node))` — fully serialized; the caller must not fall back. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + Ok(None) + } } /// Arguments for [`DataSource::open_with_args`] @@ -553,6 +577,18 @@ impl ExecutionPlan for DataSourceExec { new_exec.execution_state = Arc::new(OnceLock::new()); Ok(Arc::new(new_exec)) } + + /// Delegates serialization to the wrapped [`DataSource`]. For file scans the + /// concrete [`FileSource`] emits the node via its + /// own `try_to_proto` hook, keeping the format-specific wire logic in the + /// format crate. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + self.data_source().try_to_proto(ctx) + } } impl DataSourceExec { 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/Cargo.toml b/datafusion/proto/Cargo.toml index cfff8a949418a..037be27769f4d 100644 --- a/datafusion/proto/Cargo.toml +++ b/datafusion/proto/Cargo.toml @@ -54,7 +54,7 @@ chrono = { workspace = true } datafusion-catalog = { workspace = true } datafusion-catalog-listing = { workspace = true } datafusion-common = { workspace = true } -datafusion-datasource = { workspace = true } +datafusion-datasource = { workspace = true, features = ["proto"] } datafusion-datasource-arrow = { workspace = true } datafusion-datasource-avro = { workspace = true, optional = true } datafusion-datasource-csv = { workspace = true } diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index d5a1e0efac6b6..fcefe402b175a 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -23,13 +23,10 @@ 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, FileScanConfigBuilder}; +use datafusion_datasource::file_scan_config::FileScanConfig; use datafusion_datasource::file_sink_config::FileSinkConfig; use datafusion_datasource::{FileRange, ListingTableUrl, PartitionedFile, TableSchema}; use datafusion_datasource_csv::file_format::CsvSink; @@ -41,27 +38,23 @@ use datafusion_execution::{FunctionRegistry, TaskContext}; use datafusion_expr::WindowFunctionDefinition; use datafusion_expr::dml::InsertOp; use datafusion_physical_expr::expressions::{LambdaExpr, LambdaVariable}; -use datafusion_physical_expr::projection::{ProjectionExpr, ProjectionExprs}; 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, LikeExpr, Literal, NegativeExpr, NotExpr, TryCastExpr, UnKnownColumn, }; 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::{ - DefaultPhysicalProtoConverter, PhysicalExtensionCodec, PhysicalPlanDecodeContext, - PhysicalProtoConverterExtension, + ConverterPlanDecoder, DefaultPhysicalProtoConverter, PhysicalExtensionCodec, + PhysicalPlanDecodeContext, PhysicalProtoConverterExtension, }; use crate::convert::TryFromProto; use crate::protobuf::physical_expr_node::ExprType; @@ -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,36 +440,14 @@ 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 { - let schema: Arc = parse_protobuf_file_scan_schema(proto)?; - - // Reacquire the partition column types from the schema before removing them below. - let table_partition_cols = proto - .table_partition_cols - .iter() - .map(|col| Ok(Arc::new(schema.field_with_name(col)?.clone()))) - .collect::>>()?; - - // Remove partition columns from the schema after recreating table_partition_cols - // because the partition columns are not in the file. They are present to allow - // the partition column types to be reconstructed after serde. - let file_schema = Arc::new( - Schema::new( - schema - .fields() - .iter() - .filter(|field| !table_partition_cols.contains(field)) - .cloned() - .collect::>(), - ) - .with_metadata(schema.metadata.clone()), - ); - - Ok(TableSchema::builder(file_schema) - .with_table_partition_cols(table_partition_cols) - .build()) + FileScanConfig::parse_table_schema_from_proto(proto) } pub fn parse_protobuf_file_scan_config( @@ -548,76 +456,15 @@ pub fn parse_protobuf_file_scan_config( proto_converter: &dyn PhysicalProtoConverterExtension, file_source: Arc, ) -> Result { - let schema: Arc = parse_protobuf_file_scan_schema(proto)?; - - let constraints = convert_required!(proto.constraints)?; - let statistics = convert_required!(proto.statistics)?; - - let file_groups = proto - .file_groups - .iter() - .map(FileGroup::try_from_proto) - .collect::, _>>()?; - - let object_store_url = match proto.object_store_url.is_empty() { - false => ObjectStoreUrl::parse(&proto.object_store_url)?, - true => ObjectStoreUrl::local_filesystem(), - }; - - let mut output_ordering = vec![]; - for node_collection in &proto.output_ordering { - let sort_exprs = parse_physical_sort_exprs( - &node_collection.physical_sort_expr_nodes, - ctx, - &schema, - proto_converter, - )?; - output_ordering.extend(LexOrdering::new(sort_exprs)); - } - let output_partitioning = parse_protobuf_partitioning( - proto.output_partitioning.as_ref(), + let decoder = ConverterPlanDecoder { ctx, - &schema, proto_converter, - )?; - - // Parse projection expressions if present and apply to file source - let file_source = if let Some(proto_projection_exprs) = &proto.projection_exprs { - let projection_exprs: Vec = proto_projection_exprs - .projections - .iter() - .map(|proto_expr| { - let expr = proto_converter.proto_to_physical_expr( - proto_expr.expr.as_ref().ok_or_else(|| { - internal_datafusion_err!("ProjectionExpr missing expr field") - })?, - &schema, - ctx, - )?; - Ok(ProjectionExpr::new(expr, proto_expr.alias.clone())) - }) - .collect::>>()?; - - let projection_exprs = ProjectionExprs::new(projection_exprs); - - // Apply projection to file source - file_source - .try_pushdown_projection(&projection_exprs)? - .unwrap_or(file_source) - } else { - file_source }; - - let config = FileScanConfigBuilder::new(object_store_url, file_source) - .with_file_groups(file_groups) - .with_constraints(constraints) - .with_statistics(statistics) - .with_limit(proto.limit.as_ref().map(|sl| sl.limit as usize)) - .with_output_ordering(output_ordering) - .with_output_partitioning(output_partitioning) - .with_batch_size(proto.batch_size.map(|s| s as usize)) - .build(); - Ok(config) + FileScanConfig::try_from_proto( + proto, + &ExecutionPlanDecodeCtx::new(&decoder), + file_source, + ) } pub fn parse_record_batches(buf: &[u8]) -> Result> { @@ -632,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) } } @@ -808,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 7ee173cb36868..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, @@ -121,6 +121,388 @@ fn encode_human_display_alias(human_display: &str, alias: &str) -> String { ) } +#[cfg(test)] +mod file_scan_config_serde { + use super::*; + use arrow::datatypes::{DataType, Field}; + use datafusion_common::{Constraint, Constraints, ScalarValue, Statistics}; + use datafusion_datasource::file::FileSource; + use datafusion_datasource::file_groups::FileGroup; + use datafusion_datasource::file_stream::FileOpener; + use datafusion_datasource::{PartitionedFile, TableSchema}; + use datafusion_execution::object_store::ObjectStoreUrl; + use datafusion_physical_expr::expressions::Column; + use datafusion_physical_expr::projection::{ + ProjectionExpr as FileProjectionExpr, ProjectionExprs as FileProjectionExprs, + }; + use datafusion_physical_expr::{ + LexOrdering, Partitioning, RangePartitioning, SplitPoint, + }; + use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; + use object_store::ObjectStore; + + #[derive(Clone)] + struct SerdeTestSource { + metrics: ExecutionPlanMetricsSet, + table_schema: TableSchema, + projection: Option, + } + + impl SerdeTestSource { + fn new( + table_schema: TableSchema, + projection: Option, + ) -> Self { + Self { + metrics: ExecutionPlanMetricsSet::new(), + table_schema, + projection, + } + } + } + + impl FileSource for SerdeTestSource { + fn create_file_opener( + &self, + _object_store: Arc, + _base_config: &FileScanConfig, + _partition: usize, + ) -> Result> { + internal_err!("not needed for FileScanConfig serde tests") + } + + fn table_schema(&self) -> &TableSchema { + &self.table_schema + } + + fn with_batch_size(&self, _batch_size: usize) -> Arc { + Arc::new(self.clone()) + } + + fn metrics(&self) -> &ExecutionPlanMetricsSet { + &self.metrics + } + + fn file_type(&self) -> &str { + "serde-test" + } + + fn try_pushdown_projection( + &self, + projection: &FileProjectionExprs, + ) -> Result>> { + Ok(Some(Arc::new(Self { + projection: Some(projection.clone()), + ..self.clone() + }))) + } + + fn projection(&self) -> Option<&FileProjectionExprs> { + self.projection.as_ref() + } + } + + fn populated_projection() -> FileProjectionExprs { + FileProjectionExprs::new(vec![FileProjectionExpr::new( + Arc::new(Column::new("value", 0)), + "projected_value", + )]) + } + + fn test_config(output_partitioning: Option) -> FileScanConfig { + test_config_with_projection(output_partitioning, Some(populated_projection())) + } + + fn test_config_with_projection( + output_partitioning: Option, + projection: Option, + ) -> FileScanConfig { + let file_schema = Arc::new( + Schema::new(vec![ + Field::new("value", DataType::Int32, false), + Field::new("label", DataType::Utf8, true), + ]) + .with_metadata(HashMap::from([( + "serde_test_key".to_string(), + "serde_test_value".to_string(), + )])), + ); + let table_schema = TableSchema::builder(Arc::clone(&file_schema)) + .with_table_partition_cols(vec![Arc::new(Field::new( + "part", + DataType::Utf8, + false, + ))]) + .build(); + let table_statistics = Statistics::new_unknown(table_schema.table_schema()); + let source = Arc::new(SerdeTestSource::new(table_schema, projection)); + let first_file = PartitionedFile::new("data/part=a/file.arrow", 1024) + .with_partition_values(vec![ScalarValue::Utf8(Some("a".to_string()))]) + .with_range(10, 900) + .with_arrow_schema(Arc::clone(&file_schema)) + .with_statistics(Arc::new(table_statistics.clone())); + let second_file = PartitionedFile::new("data/part=b/file.arrow", 2048) + .with_partition_values(vec![ScalarValue::Utf8(Some("b".to_string()))]); + let third_file = PartitionedFile::new("data/part=c/file.arrow", 4096) + .with_partition_values(vec![ScalarValue::Utf8(Some("c".to_string()))]) + .with_arrow_schema(Arc::clone(&file_schema)); + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(Arc::new( + Column::new("value", 0), + ))]) + .expect("single expression ordering"); + + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + .with_file_groups(vec![ + FileGroup::new(vec![first_file, second_file]), + FileGroup::new(vec![third_file]), + ]) + .with_constraints(Constraints::new_unverified(vec![Constraint::PrimaryKey( + vec![0], + )])) + .with_statistics(table_statistics) + .with_limit(Some(17)) + .with_batch_size(Some(256)) + .with_output_ordering(vec![ordering]) + .with_output_partitioning(output_partitioning) + .build() + } + + fn hash_partitioning() -> Partitioning { + Partitioning::Hash(vec![Arc::new(Column::new("value", 0))], 3) + } + + fn range_partitioning() -> Partitioning { + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(Arc::new( + Column::new("value", 0), + ))]) + .expect("single expression ordering"); + Partitioning::Range(RangePartitioning::new( + ordering, + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(10))])], + )) + } + + fn decode_source(conf: &protobuf::FileScanExecConf) -> Result> { + Ok(Arc::new(SerdeTestSource::new( + FileScanConfig::parse_table_schema_from_proto(conf)?, + None, + ))) + } + + struct FileScanSerdeHarness { + codec: DefaultPhysicalExtensionCodec, + converter: DefaultPhysicalProtoConverter, + task_ctx: TaskContext, + } + + impl FileScanSerdeHarness { + fn new() -> Self { + Self { + codec: DefaultPhysicalExtensionCodec {}, + converter: DefaultPhysicalProtoConverter {}, + task_ctx: TaskContext::default(), + } + } + + fn encode(&self, config: &FileScanConfig) -> Result { + let encoder = ConverterPlanEncoder { + codec: &self.codec, + proto_converter: &self.converter, + }; + config.try_to_proto(&ExecutionPlanEncodeCtx::new(&encoder)) + } + + fn decode(&self, conf: &protobuf::FileScanExecConf) -> Result { + self.decode_with_source(conf, decode_source(conf)?) + } + + fn decode_with_source( + &self, + conf: &protobuf::FileScanExecConf, + file_source: Arc, + ) -> Result { + let physical_decode_ctx = + PhysicalPlanDecodeContext::new(&self.task_ctx, &self.codec); + let decoder = ConverterPlanDecoder { + ctx: &physical_decode_ctx, + proto_converter: &self.converter, + }; + FileScanConfig::try_from_proto( + conf, + &ExecutionPlanDecodeCtx::new(&decoder), + file_source, + ) + } + } + + #[test] + fn new_file_scan_config_serde_roundtrips_all_partitioning_variants() -> Result<()> { + let serde = FileScanSerdeHarness::new(); + + for config in [ + test_config(None), + test_config(Some(Partitioning::RoundRobinBatch(2))), + test_config(Some(hash_partitioning())), + test_config(Some(range_partitioning())), + test_config(Some(Partitioning::UnknownPartitioning(4))), + ] { + let encoded = serde.encode(&config)?; + let reencoded = serde.encode(&serde.decode(&encoded)?)?; + assert_eq!(reencoded.output_partitioning, encoded.output_partitioning); + } + + Ok(()) + } + + #[test] + fn new_file_scan_config_serde_preserves_complete_fixture() -> Result<()> { + let serde = FileScanSerdeHarness::new(); + let config = test_config(None); + let decoded = serde.decode(&serde.encode(&config)?)?; + + assert_eq!(decoded.constraints, config.constraints); + assert_eq!( + decoded.file_schema().metadata, + config.file_schema().metadata + ); + assert_eq!(decoded.file_groups.len(), 2); + assert_eq!(decoded.file_groups[0].len(), 2); + assert_eq!(decoded.file_groups[1].len(), 1); + assert!(decoded.file_groups[0].files()[0].arrow_schema.is_some()); + assert!(decoded.file_groups[0].files()[1].arrow_schema.is_none()); + + Ok(()) + } + + #[test] + fn new_file_scan_config_serde_preserves_projection_presence() -> Result<()> { + let serde = FileScanSerdeHarness::new(); + + let absent = serde.encode(&test_config_with_projection(None, None))?; + assert!(absent.projection_exprs.is_none()); + assert!(serde.decode(&absent)?.file_source().projection().is_none()); + + let empty = serde.encode(&test_config_with_projection( + None, + Some(FileProjectionExprs::new(vec![])), + ))?; + assert!( + empty + .projection_exprs + .as_ref() + .is_some_and(|projection| projection.projections.is_empty()) + ); + assert!( + serde + .decode(&empty)? + .file_source() + .projection() + .is_some_and(|projection| projection.as_ref().is_empty()) + ); + + Ok(()) + } + + #[test] + fn new_file_scan_config_decode_rejects_malformed_required_fields() -> Result<()> { + let serde = FileScanSerdeHarness::new(); + let valid = serde.encode(&test_config(None))?; + let file_source = decode_source(&valid)?; + + for (field, malformed) in [ + ( + "schema", + protobuf::FileScanExecConf { + schema: None, + ..valid.clone() + }, + ), + ( + "constraints", + protobuf::FileScanExecConf { + constraints: None, + ..valid.clone() + }, + ), + ( + "statistics", + protobuf::FileScanExecConf { + statistics: None, + ..valid.clone() + }, + ), + ] { + let err = serde + .decode_with_source(&malformed, Arc::clone(&file_source)) + .expect_err("missing required field must fail"); + assert!(err.to_string().contains(field), "unexpected error: {err}"); + } + + let mut missing_projection_expr = valid.clone(); + missing_projection_expr + .projection_exprs + .as_mut() + .expect("test config has projection expressions") + .projections[0] + .expr = None; + let err = serde + .decode_with_source(&missing_projection_expr, file_source) + .expect_err("missing projection expression must fail"); + assert!( + err.to_string() + .contains("ProjectionExpr missing expr field"), + "unexpected error: {err}" + ); + + Ok(()) + } + + #[test] + fn new_file_scan_config_decode_rejects_invalid_range_ordering() -> Result<()> { + let serde = FileScanSerdeHarness::new(); + let mut proto = serde.encode(&test_config(Some(range_partitioning())))?; + + let mut duplicate_ordering = proto.clone(); + let range = match duplicate_ordering + .output_partitioning + .as_mut() + .and_then(|p| p.partition_method.as_mut()) + { + Some(protobuf::partitioning::PartitionMethod::Range(range)) => range, + other => panic!("expected range partitioning, got {other:?}"), + }; + range.sort_expr.push(range.sort_expr[0].clone()); + + let err = serde + .decode(&duplicate_ordering) + .expect_err("duplicate range ordering must fail"); + assert!( + err.to_string().contains("duplicate expressions"), + "unexpected error: {err}" + ); + + let range = match proto + .output_partitioning + .as_mut() + .and_then(|p| p.partition_method.as_mut()) + { + Some(protobuf::partitioning::PartitionMethod::Range(range)) => range, + other => panic!("expected range partitioning, got {other:?}"), + }; + range.sort_expr.clear(); + + let err = serde + .decode(&proto) + .expect_err("empty range ordering must fail"); + assert!( + err.to_string().contains("requires non-empty ordering"), + "unexpected error: {err}" + ); + + Ok(()) + } +} + #[cfg(test)] mod tests { use super::*; @@ -999,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), @@ -1034,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, @@ -1053,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, @@ -1108,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(), @@ -1155,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 e13923dbb9519..04252176d2845 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -34,20 +34,18 @@ use datafusion_expr::WindowFrame; use datafusion_physical_expr::window::{SlidingAggregateWindowExpr, StandardWindowExpr}; use datafusion_physical_expr::{HigherOrderFunctionExpr, ScalarFunctionExpr}; 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::{ - DefaultPhysicalProtoConverter, PhysicalExtensionCodec, + ConverterPlanEncoder, DefaultPhysicalProtoConverter, PhysicalExtensionCodec, PhysicalProtoConverterExtension, encode_human_display_alias, }; use crate::convert::TryFromProto; use crate::protobuf::{ - self, PhysicalSortExprNode, PhysicalSortExprNodeCollection, - physical_aggregate_expr_node, physical_window_expr_node, + self, PhysicalSortExprNode, physical_aggregate_expr_node, physical_window_expr_node, }; #[expect(clippy::needless_pass_by_value)] @@ -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::>>()?, }) } } @@ -487,84 +401,11 @@ pub fn serialize_file_scan_config( codec: &dyn PhysicalExtensionCodec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result { - let file_groups = conf - .file_groups - .iter() - .map(|p| protobuf::FileGroup::try_from_proto(p.files())) - .collect::, _>>()?; - - let mut output_orderings = vec![]; - for order in &conf.output_ordering { - let ordering = - serialize_physical_sort_exprs(order.to_vec(), codec, proto_converter)?; - output_orderings.push(ordering) - } - let output_partitioning = conf - .output_partitioning - .as_ref() - .map(|partitioning| serialize_partitioning(partitioning, codec, proto_converter)) - .transpose()?; - - // Fields must be added to the schema so that they can persist in the protobuf, - // and then they are to be removed from the schema in `parse_protobuf_file_scan_config` - let mut fields = conf - .file_schema() - .fields() - .iter() - .cloned() - .collect::>(); - fields.extend(conf.table_partition_cols().iter().cloned()); - - let schema = Arc::new( - Schema::new(fields.clone()).with_metadata(conf.file_schema().metadata.clone()), - ); - - let projection_exprs = conf - .file_source - .projection() - .as_ref() - .map(|projection_exprs| { - let projections = projection_exprs.iter().cloned().collect::>(); - Ok::<_, DataFusionError>(protobuf::ProjectionExprs { - projections: projections - .into_iter() - .map(|expr| { - Ok(protobuf::ProjectionExpr { - alias: expr.alias.to_string(), - expr: Some( - proto_converter - .physical_expr_to_proto(&expr.expr, codec)?, - ), - }) - }) - .collect::>>()?, - }) - }) - .transpose()?; - - Ok(protobuf::FileScanExecConf { - file_groups, - statistics: Some((&conf.statistics()).into()), - limit: conf.limit.map(|l| protobuf::ScanLimit { limit: l as u32 }), - projection: vec![], - schema: Some(schema.as_ref().try_into()?), - table_partition_cols: conf - .table_partition_cols() - .iter() - .map(|x| x.name().clone()) - .collect::>(), - object_store_url: conf.object_store_url.to_string(), - output_ordering: output_orderings - .into_iter() - .map(|e| PhysicalSortExprNodeCollection { - physical_sort_expr_nodes: e, - }) - .collect::>(), - constraints: Some(conf.constraints.clone().into()), - batch_size: conf.batch_size.map(|s| s as u64), - projection_exprs, - output_partitioning, - }) + let encoder = ConverterPlanEncoder { + codec, + proto_converter, + }; + 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`].