kumarUjjawal commented on code in PR #23752: URL: https://github.com/apache/datafusion/pull/23752#discussion_r3649561289
########## datafusion/datasource/src/file_sink_config/proto.rs: ########## @@ -0,0 +1,232 @@ +// 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 conversion for the format-independent [`FileSinkConfig`]. + +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_expr::dml::InsertOp; +use datafusion_physical_expr::PhysicalSortExpr; +use datafusion_physical_expr_common::sort_expr::LexRequirement; +use datafusion_physical_plan::proto::ExecutionPlanDecodeCtx; +use datafusion_proto_models::protobuf; +use object_store::ObjectMeta; +use object_store::path::Path; + +use crate::file_groups::FileGroup; +use crate::file_sink_config::{FileOutputMode, FileSinkConfig}; +use crate::{ListingTableUrl, PartitionedFile}; + +impl FileSinkConfig { + /// Serialize this shared file-sink configuration without format-specific + /// writer options. + pub fn to_proto(&self) -> Result<protobuf::FileSinkConfig> { + let file_groups = self + .file_group + .iter() + .map(partitioned_file_to_proto) + .collect::<Result<Vec<_>>>()?; + let table_paths = self + .table_paths + .iter() + .map(ToString::to_string) + .collect::<Vec<_>>(); + let table_partition_cols = self + .table_partition_cols + .iter() + .map(|(name, data_type)| { + Ok(protobuf::PartitionColumn { + name: name.to_owned(), + arrow_type: Some(data_type.try_into()?), + }) + }) + .collect::<Result<Vec<_>>>()?; + let file_output_mode = match self.file_output_mode { + FileOutputMode::Automatic => protobuf::FileOutputMode::Automatic, + FileOutputMode::SingleFile => protobuf::FileOutputMode::SingleFile, + FileOutputMode::Directory => protobuf::FileOutputMode::Directory, + }; + + Ok(protobuf::FileSinkConfig { + object_store_url: self.object_store_url.to_string(), + file_groups, + table_paths, + output_schema: Some(self.output_schema.as_ref().try_into()?), + table_partition_cols, + keep_partition_by_columns: self.keep_partition_by_columns, + insert_op: self.insert_op as i32, + file_extension: self.file_extension.clone(), + file_output_mode: file_output_mode.into(), + }) + } + + /// Reconstruct a shared file-sink configuration from protobuf. + pub fn from_proto(conf: &protobuf::FileSinkConfig) -> Result<Self> { + let file_group = FileGroup::new( + conf.file_groups + .iter() + .map(partitioned_file_from_proto) + .collect::<Result<Vec<_>>>()?, + ); + let table_paths = conf + .table_paths + .iter() + .map(ListingTableUrl::parse) + .collect::<Result<Vec<_>>>()?; + let table_partition_cols = conf + .table_partition_cols + .iter() + .map(|protobuf::PartitionColumn { name, arrow_type }| { + let data_type = arrow_type + .as_ref() + .ok_or_else(|| { + internal_datafusion_err!( + "PartitionColumn is missing required field 'arrow_type'" + ) + })? + .try_into()?; + Ok((name.clone(), data_type)) + }) + .collect::<Result<Vec<_>>>()?; + let insert_op = match conf.insert_op() { Review Comment: Unknown enum values are silently treated as defaults. This can turn an unknown write operation into append, or an unknown output mode into automatic. Please return an error for invalid values and add regression tests ########## datafusion/proto/tests/cases/roundtrip_physical_plan.rs: ########## @@ -1956,6 +1958,122 @@ fn roundtrip_analyze() -> Result<()> { )) } +#[derive(Debug)] +struct ProtoHookSink { + schema: SchemaRef, +} + +impl DisplayAs for ProtoHookSink { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "ProtoHookSink") + } +} + +#[async_trait] +impl DataSink for ProtoHookSink { + fn schema(&self) -> &SchemaRef { + &self.schema + } + + async fn write_all( + &self, + _data: SendableRecordBatchStream, + _context: &Arc<TaskContext>, + ) -> Result<u64> { + unreachable!("serialization test does not execute the sink") + } + + fn try_to_proto( + &self, + input: PhysicalPlanNode, + sort_order: Option<protobuf::PhysicalSortExprNodeCollection>, + sink_schema: &Schema, + ) -> Result<Option<PhysicalPlanNode>> { + assert!(matches!( + input.physical_plan_type, + Some(protobuf::physical_plan_node::PhysicalPlanType::PlaceholderRow(_)) + )); + assert_eq!( + sort_order + .as_ref() + .map(|ordering| ordering.physical_sort_expr_nodes.len()), + Some(1) + ); + assert_eq!(sink_schema.fields().len(), 1); + + Ok(Some(PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Empty( + protobuf::EmptyExecNode { + schema: Some(sink_schema.try_into()?), + partitions: 1, + }, + ), + ), + })) + } +} + +#[test] +fn data_sink_exec_delegates_to_sink_proto_hook() -> Result<()> { + let input_schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int64, + false, + )])); + let input = Arc::new(PlaceholderRowExec::new(Arc::clone(&input_schema))); + let sink = Arc::new(ProtoHookSink { + schema: Arc::clone(&input_schema), + }); + let sort_order = [PhysicalSortRequirement::new( + Arc::new(Column::new("value", 0)), + Some(SortOptions::default()), + )] + .into(); + let plan = Arc::new(DataSinkExec::new(input, sink, Some(sort_order))); + + let node = PhysicalPlanNode::try_from_physical_plan( + plan, + &DefaultPhysicalExtensionCodec {}, + )?; + + assert!(matches!( + node.physical_plan_type, + Some(protobuf::physical_plan_node::PhysicalPlanType::Empty(_)) + )); + Ok(()) +} + +#[test] +fn file_sink_config_conversion_preserves_compatibility_api() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "partition", + DataType::Utf8, + false, + )])); + let config = FileSinkConfig { + original_url: "file:///tmp/output".to_string(), + object_store_url: ObjectStoreUrl::local_filesystem(), + file_group: FileGroup::new(vec![PartitionedFile::new("/tmp/output", 1)]), + table_paths: vec![ListingTableUrl::parse("file:///tmp/output")?], + output_schema: schema, + table_partition_cols: vec![("partition".to_string(), DataType::Utf8)], + insert_op: InsertOp::Overwrite, + keep_partition_by_columns: true, + file_extension: "parquet".to_string(), + file_output_mode: FileOutputMode::Directory, + }; + + let direct = config.to_proto()?; + let compatibility = protobuf::FileSinkConfig::try_from_proto(&config)?; + assert_eq!(direct, compatibility); + + let direct = FileSinkConfig::from_proto(&direct)?; Review Comment: Both decode paths use the same implementation, so this test can still pass when decoding drops or changes a field. Please compare the decoded configuration with the original and add malformed-input tests ########## datafusion/datasource/src/file_sink_config/proto.rs: ########## @@ -0,0 +1,232 @@ +// 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 conversion for the format-independent [`FileSinkConfig`]. + +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_expr::dml::InsertOp; +use datafusion_physical_expr::PhysicalSortExpr; +use datafusion_physical_expr_common::sort_expr::LexRequirement; +use datafusion_physical_plan::proto::ExecutionPlanDecodeCtx; +use datafusion_proto_models::protobuf; +use object_store::ObjectMeta; +use object_store::path::Path; + +use crate::file_groups::FileGroup; +use crate::file_sink_config::{FileOutputMode, FileSinkConfig}; +use crate::{ListingTableUrl, PartitionedFile}; + +impl FileSinkConfig { + /// Serialize this shared file-sink configuration without format-specific + /// writer options. + pub fn to_proto(&self) -> Result<protobuf::FileSinkConfig> { + let file_groups = self + .file_group + .iter() + .map(partitioned_file_to_proto) + .collect::<Result<Vec<_>>>()?; + let table_paths = self + .table_paths + .iter() + .map(ToString::to_string) + .collect::<Vec<_>>(); + let table_partition_cols = self + .table_partition_cols + .iter() + .map(|(name, data_type)| { + Ok(protobuf::PartitionColumn { + name: name.to_owned(), + arrow_type: Some(data_type.try_into()?), + }) + }) + .collect::<Result<Vec<_>>>()?; + let file_output_mode = match self.file_output_mode { + FileOutputMode::Automatic => protobuf::FileOutputMode::Automatic, + FileOutputMode::SingleFile => protobuf::FileOutputMode::SingleFile, + FileOutputMode::Directory => protobuf::FileOutputMode::Directory, + }; + + Ok(protobuf::FileSinkConfig { + object_store_url: self.object_store_url.to_string(), + file_groups, + table_paths, + output_schema: Some(self.output_schema.as_ref().try_into()?), + table_partition_cols, + keep_partition_by_columns: self.keep_partition_by_columns, + insert_op: self.insert_op as i32, + file_extension: self.file_extension.clone(), + file_output_mode: file_output_mode.into(), + }) + } + + /// Reconstruct a shared file-sink configuration from protobuf. + pub fn from_proto(conf: &protobuf::FileSinkConfig) -> Result<Self> { + let file_group = FileGroup::new( + conf.file_groups + .iter() + .map(partitioned_file_from_proto) + .collect::<Result<Vec<_>>>()?, + ); + let table_paths = conf + .table_paths + .iter() + .map(ListingTableUrl::parse) + .collect::<Result<Vec<_>>>()?; + let table_partition_cols = conf + .table_partition_cols + .iter() + .map(|protobuf::PartitionColumn { name, arrow_type }| { + let data_type = arrow_type + .as_ref() + .ok_or_else(|| { + internal_datafusion_err!( + "PartitionColumn is missing required field 'arrow_type'" + ) + })? + .try_into()?; + Ok((name.clone(), data_type)) + }) + .collect::<Result<Vec<_>>>()?; + let insert_op = match conf.insert_op() { + protobuf::InsertOp::Append => InsertOp::Append, + protobuf::InsertOp::Overwrite => InsertOp::Overwrite, + protobuf::InsertOp::Replace => InsertOp::Replace, + }; + let file_output_mode = match conf.file_output_mode() { + protobuf::FileOutputMode::Automatic => FileOutputMode::Automatic, + protobuf::FileOutputMode::SingleFile => FileOutputMode::SingleFile, + protobuf::FileOutputMode::Directory => FileOutputMode::Directory, + }; + let output_schema = conf.output_schema.as_ref().ok_or_else(|| { + internal_datafusion_err!( + "FileSinkConfig is missing required field 'output_schema'" + ) + })?; + + Ok(Self { + original_url: String::default(), + object_store_url: ObjectStoreUrl::parse(&conf.object_store_url)?, + file_group, + table_paths, + output_schema: Arc::new(output_schema.try_into()?), + table_partition_cols, + insert_op, + keep_partition_by_columns: conf.keep_partition_by_columns, + file_extension: conf.file_extension.clone(), + file_output_mode, + }) + } +} + +/// Decode a sink's optional required output ordering against its input schema. +pub fn parse_sink_sort_order( + collection: Option<&protobuf::PhysicalSortExprNodeCollection>, + ctx: &ExecutionPlanDecodeCtx<'_>, + schema: &Schema, +) -> Result<Option<LexRequirement>> { + let Some(collection) = collection else { + return Ok(None); + }; + let sort_exprs = collection + .physical_sort_expr_nodes + .iter() + .map(|node| { + let expr = node.expr.as_ref().ok_or_else(|| { + internal_datafusion_err!("Unexpected empty physical expression") + })?; + Ok(PhysicalSortExpr { + expr: ctx.decode_expr(expr, schema)?, + options: SortOptions { + descending: !node.asc, + nulls_first: node.nulls_first, + }, + }) + }) + .collect::<Result<Vec<_>>>()?; + Ok(LexRequirement::new(sort_exprs.into_iter().map(Into::into))) +} + +fn partitioned_file_to_proto( Review Comment: #23683 now provides these shared PartitionedFile helpers. Once that gets merged in the main we can merge main to this pr so we can remove this duplicate code, and reuse those helpers. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
