diff --git a/Cargo.lock b/Cargo.lock index db8c6e1252dec..76b8bdb92a6a4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1920,6 +1920,8 @@ dependencies = [ "datafusion-physical-expr-adapter", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-common", + "datafusion-proto-models", "datafusion-session", "flate2", "futures", @@ -1954,6 +1956,7 @@ dependencies = [ "datafusion-expr", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "futures", "itertools 0.15.0", @@ -1973,6 +1976,7 @@ dependencies = [ "datafusion-datasource", "datafusion-physical-expr-adapter", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "futures", "object_store", @@ -1992,6 +1996,7 @@ dependencies = [ "datafusion-expr", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "futures", "object_store", @@ -2013,6 +2018,7 @@ dependencies = [ "datafusion-expr", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "futures", "object_store", @@ -2043,6 +2049,8 @@ dependencies = [ "datafusion-physical-expr-adapter", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-common", + "datafusion-proto-models", "datafusion-pruning", "datafusion-session", "futures", diff --git a/datafusion/datasource-arrow/Cargo.toml b/datafusion/datasource-arrow/Cargo.toml index 2718e424c6386..6f50135403d69 100644 --- a/datafusion/datasource-arrow/Cargo.toml +++ b/datafusion/datasource-arrow/Cargo.toml @@ -42,6 +42,7 @@ datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } datafusion-physical-expr-common = { workspace = true } datafusion-physical-plan = { workspace = true } +datafusion-proto-models = { workspace = true, optional = true } datafusion-session = { workspace = true } futures = { workspace = true } itertools = { workspace = true } @@ -65,3 +66,10 @@ path = "src/mod.rs" # This feature is deprecated, as core functionality in the SpillManager requires all features # it enabled, and will be removed in a future version. compression = [] +# Enables `FileSource::try_to_proto` on `ArrowSource` and the `ArrowScan` decode +# entry point. Mirrors the `proto` feature on `datafusion-datasource`. +proto = [ + "dep:datafusion-proto-models", + "datafusion-datasource/proto", + "datafusion-physical-plan/proto", +] diff --git a/datafusion/datasource-arrow/src/source.rs b/datafusion/datasource-arrow/src/source.rs index 59c020c779ca2..45cfc95bf9c88 100644 --- a/datafusion/datasource-arrow/src/source.rs +++ b/datafusion/datasource-arrow/src/source.rs @@ -392,6 +392,67 @@ impl FileSource for ArrowSource { fn projection(&self) -> Option<&ProjectionExprs> { Some(&self.projection.source) } + + /// Emit an `ArrowScan` node wrapping the shared base config. Arrow is + /// base_conf-only (no format-specific fields). Kept byte-identical to the + /// former `try_from_data_source_exec` ArrowScan branch in `datafusion-proto`. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + base: &FileScanConfig, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + use protobuf::physical_plan_node::PhysicalPlanType; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::ArrowScan( + protobuf::ArrowScanExecNode { + base_conf: Some(base.to_proto_conf(ctx)?), + }, + )), + })) + } +} + +#[cfg(feature = "proto")] +impl ArrowSource { + /// Reconstruct a `DataSourceExec` (wrapping a `FileScanConfig` over an + /// `ArrowSource`) from an `ArrowScan` [`PhysicalPlanNode`]. + /// + /// The inverse of [`FileSource::try_to_proto`] on `ArrowSource`; kept + /// byte-compatible with the former `try_into_arrow_scan_physical_plan` in + /// `datafusion-proto`. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_datasource::file_scan_config::FileScanConfig; + use datafusion_datasource::source::DataSourceExec; + use datafusion_proto_models::protobuf; + + let scan = match &node.physical_plan_type { + Some(protobuf::physical_plan_node::PhysicalPlanType::ArrowScan(scan)) => scan, + _ => { + return datafusion_common::internal_err!( + "PhysicalPlanNode is not an ArrowScan" + ); + } + }; + + let base_conf = scan.base_conf.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "base_conf in ArrowScanExecNode is missing." + ) + })?; + + let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?; + let source = Arc::new(ArrowSource::new_file_source(table_schema)); + let scan_conf = FileScanConfig::from_proto_conf(base_conf, ctx, source)?; + Ok(DataSourceExec::from_data_source(scan_conf)) + } } /// `FileOpener` wrapper for both Arrow IPC file and stream formats diff --git a/datafusion/datasource-avro/Cargo.toml b/datafusion/datasource-avro/Cargo.toml index adc2be1cb8f24..70b675d63f427 100644 --- a/datafusion/datasource-avro/Cargo.toml +++ b/datafusion/datasource-avro/Cargo.toml @@ -30,6 +30,15 @@ version.workspace = true [package.metadata.docs.rs] all-features = true +[features] +# Enables `FileSource::try_to_proto` on `AvroSource` and the `AvroScan` decode +# entry point. Mirrors the `proto` feature on `datafusion-datasource`. +proto = [ + "dep:datafusion-proto-models", + "datafusion-datasource/proto", + "datafusion-physical-plan/proto", +] + [dependencies] arrow = { workspace = true } arrow-avro = { workspace = true } @@ -39,6 +48,7 @@ datafusion-common = { workspace = true, features = ["object_store"] } datafusion-datasource = { workspace = true } datafusion-physical-expr-adapter = { workspace = true } datafusion-physical-plan = { workspace = true } +datafusion-proto-models = { workspace = true, optional = true } datafusion-session = { workspace = true } futures = { workspace = true } object_store = { workspace = true } diff --git a/datafusion/datasource-avro/src/source.rs b/datafusion/datasource-avro/src/source.rs index e3be9d8a401d0..9c2b70a8d811a 100644 --- a/datafusion/datasource-avro/src/source.rs +++ b/datafusion/datasource-avro/src/source.rs @@ -168,6 +168,67 @@ impl FileSource for AvroSource { // Avro OCF does not support safe byte-range splitting in this reader path. false } + + /// Emit an `AvroScan` node wrapping the shared base config. Avro carries no + /// format-specific fields, so this is just the base conf. Kept + /// byte-identical to the former `try_from_data_source_exec` AvroScan branch + /// in `datafusion-proto`. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + base: &FileScanConfig, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + use protobuf::physical_plan_node::PhysicalPlanType; + + let node = protobuf::AvroScanExecNode { + base_conf: Some(base.to_proto_conf(ctx)?), + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::AvroScan(node)), + })) + } +} + +#[cfg(feature = "proto")] +impl AvroSource { + /// Reconstruct a `DataSourceExec` (wrapping a `FileScanConfig` over an + /// `AvroSource`) from an `AvroScan` + /// [`PhysicalPlanNode`](datafusion_proto_models::protobuf::PhysicalPlanNode). + /// + /// The inverse of [`FileSource::try_to_proto`] on `AvroSource`; kept + /// byte-compatible with the former `try_into_avro_scan_physical_plan` in + /// `datafusion-proto`. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_datasource::source::DataSourceExec; + use datafusion_proto_models::protobuf; + + let scan = match &node.physical_plan_type { + Some(protobuf::physical_plan_node::PhysicalPlanType::AvroScan(scan)) => scan, + _ => { + return datafusion_common::internal_err!( + "PhysicalPlanNode is not an AvroScan" + ); + } + }; + + let base_conf = scan.base_conf.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "AvroScanExecNode is missing required field 'base_conf'" + ) + })?; + + // Parse table schema with partition columns, then rebuild the source. + let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?; + let source = Arc::new(AvroSource::new(table_schema)); + + let conf = FileScanConfig::from_proto_conf(base_conf, ctx, source)?; + Ok(DataSourceExec::from_data_source(conf)) + } } mod private { diff --git a/datafusion/datasource-csv/Cargo.toml b/datafusion/datasource-csv/Cargo.toml index 295092512742b..4e8f81b464390 100644 --- a/datafusion/datasource-csv/Cargo.toml +++ b/datafusion/datasource-csv/Cargo.toml @@ -30,6 +30,15 @@ version.workspace = true [package.metadata.docs.rs] all-features = true +[features] +# Enables `FileSource::try_to_proto` on `CsvSource` and the `CsvScan` decode +# entry point. Mirrors the `proto` feature on `datafusion-datasource`. +proto = [ + "dep:datafusion-proto-models", + "datafusion-datasource/proto", + "datafusion-physical-plan/proto", +] + [dependencies] arrow = { workspace = true } async-trait = { workspace = true } @@ -41,6 +50,7 @@ datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } datafusion-physical-expr-common = { workspace = true } datafusion-physical-plan = { workspace = true } +datafusion-proto-models = { workspace = true, optional = true } datafusion-session = { workspace = true } futures = { workspace = true } object_store = { workspace = true } diff --git a/datafusion/datasource-csv/src/file_format.rs b/datafusion/datasource-csv/src/file_format.rs index 9fdd688037682..a6c06b4171999 100644 --- a/datafusion/datasource-csv/src/file_format.rs +++ b/datafusion/datasource-csv/src/file_format.rs @@ -826,6 +826,104 @@ impl DataSink for CsvSink { ) -> Result { FileSink::write_all(self, data, context).await } + + /// Emit a `CsvSink` node wrapping the shared sink config. Kept + /// byte-identical to the former `try_from_data_sink_exec` CsvSink branch in + /// `datafusion-proto`. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + input: datafusion_proto_models::protobuf::PhysicalPlanNode, + sort_order: Option< + datafusion_proto_models::protobuf::PhysicalSortExprNodeCollection, + >, + sink_schema: &Schema, + ) -> Result> { + use datafusion_proto_models::protobuf; + use protobuf::physical_plan_node::PhysicalPlanType; + + let sink = protobuf::CsvSink { + config: Some(self.config.to_proto()?), + writer_options: Some(self.writer_options().try_into()?), + }; + let node = protobuf::CsvSinkExecNode { + input: Some(Box::new(input)), + sink: Some(sink), + sink_schema: Some(sink_schema.try_into()?), + sort_order, + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::CsvSink(Box::new(node))), + })) + } +} + +#[cfg(feature = "proto")] +impl CsvSink { + /// Reconstruct a `DataSinkExec` (wrapping a `CsvSink`) from a `CsvSink` + /// [`PhysicalPlanNode`]. + /// + /// The inverse of [`DataSink::try_to_proto`] on `CsvSink`; kept + /// byte-compatible with the former `try_into_csv_sink_physical_plan` in + /// `datafusion-proto`. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_datasource::file_sink_config::parse_sink_sort_order; + use datafusion_proto_models::protobuf; + + let sink_node = match &node.physical_plan_type { + Some(protobuf::physical_plan_node::PhysicalPlanType::CsvSink(sink)) => { + sink.as_ref() + } + _ => { + return datafusion_common::internal_err!( + "PhysicalPlanNode is not a CsvSink" + ); + } + }; + + let input = ctx.decode_required_child( + sink_node.input.as_deref(), + "CsvSinkExecNode", + "input", + )?; + + let proto_sink = sink_node.sink.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "CsvSinkExecNode is missing required field 'sink'" + ) + })?; + let config = + FileSinkConfig::from_proto(proto_sink.config.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "CsvSink is missing required field 'config'" + ) + })?)?; + let writer_options = proto_sink + .writer_options + .as_ref() + .ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "CsvSink is missing required field 'writer_options'" + ) + })? + .try_into()?; + let data_sink = CsvSink::new(config, writer_options); + + let sink_schema = input.schema(); + let sort_order = + parse_sink_sort_order(sink_node.sort_order.as_ref(), ctx, &sink_schema)?; + + Ok(Arc::new(DataSinkExec::new( + input, + Arc::new(data_sink), + sort_order, + ))) + } } #[cfg(test)] diff --git a/datafusion/datasource-csv/src/source.rs b/datafusion/datasource-csv/src/source.rs index 25ec311880405..632d452f59cbc 100644 --- a/datafusion/datasource-csv/src/source.rs +++ b/datafusion/datasource-csv/src/source.rs @@ -308,6 +308,51 @@ impl FileSource for CsvSource { DisplayFormatType::TreeRender => Ok(()), } } + + /// Emit a `CsvScan` node wrapping the shared base config. Kept byte-identical + /// to the former `try_from_data_source_exec` CsvScan branch in + /// `datafusion-proto`. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + base: &FileScanConfig, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + use protobuf::physical_plan_node::PhysicalPlanType; + + let node = protobuf::CsvScanExecNode { + base_conf: Some(base.to_proto_conf(ctx)?), + has_header: self.has_header(), + delimiter: proto_byte_to_string(self.delimiter(), "delimiter")?, + quote: proto_byte_to_string(self.quote(), "quote")?, + optional_escape: self + .escape() + .map(|escape| { + Ok::<_, DataFusionError>( + protobuf::csv_scan_exec_node::OptionalEscape::Escape( + proto_byte_to_string(escape, "escape")?, + ), + ) + }) + .transpose()?, + optional_comment: self + .comment() + .map(|comment| { + Ok::<_, DataFusionError>( + protobuf::csv_scan_exec_node::OptionalComment::Comment( + proto_byte_to_string(comment, "comment")?, + ), + ) + }) + .transpose()?, + newlines_in_values: self.newlines_in_values(), + truncate_rows: self.truncate_rows(), + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::CsvScan(node)), + })) + } } impl FileOpener for CsvOpener { @@ -501,3 +546,104 @@ pub async fn plan_to_csv( Ok(()) } + +/// Convert a single byte to its one-character UTF-8 string form for the wire. +/// Local copy of `datafusion-proto`'s `byte_to_string`; kept byte-identical. +#[cfg(feature = "proto")] +fn proto_byte_to_string(b: u8, description: &str) -> Result { + let bytes = &[b]; + let s = std::str::from_utf8(bytes).map_err(|_| { + datafusion_common::internal_datafusion_err!( + "Invalid CSV {description}: can not represent {bytes:0x?} as utf8" + ) + })?; + Ok(s.to_owned()) +} + +/// Convert a one-character string from the wire back to a single byte. +/// Local copy of `datafusion-proto`'s `str_to_byte`; kept byte-identical. +#[cfg(feature = "proto")] +fn proto_str_to_byte(s: &str, description: &str) -> Result { + datafusion_common::assert_eq_or_internal_err!( + s.len(), + 1, + "Invalid CSV {description}: expected single character, got {s}" + ); + Ok(s.as_bytes()[0]) +} + +#[cfg(feature = "proto")] +impl CsvSource { + /// Reconstruct a `DataSourceExec` (wrapping a `FileScanConfig` over a + /// `CsvSource`) from a `CsvScan` [`PhysicalPlanNode`]. + /// + /// The inverse of [`FileSource::try_to_proto`] on `CsvSource`; kept + /// byte-compatible with the former `try_into_csv_scan_physical_plan` in + /// `datafusion-proto`. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_common::config::CsvOptions; + use datafusion_datasource::file_compression_type::FileCompressionType; + use datafusion_datasource::file_scan_config::{ + FileScanConfig, FileScanConfigBuilder, + }; + use datafusion_datasource::source::DataSourceExec; + use datafusion_proto_models::protobuf; + + let scan = match &node.physical_plan_type { + Some(protobuf::physical_plan_node::PhysicalPlanType::CsvScan(scan)) => scan, + _ => { + return datafusion_common::internal_err!( + "PhysicalPlanNode is not a CsvScan" + ); + } + }; + + let base_conf = scan.base_conf.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "CsvScanExecNode is missing required field 'base_conf'" + ) + })?; + + let escape = match &scan.optional_escape { + Some(protobuf::csv_scan_exec_node::OptionalEscape::Escape(escape)) => { + Some(proto_str_to_byte(escape, "escape")?) + } + None => None, + }; + let comment = match &scan.optional_comment { + Some(protobuf::csv_scan_exec_node::OptionalComment::Comment(comment)) => { + Some(proto_str_to_byte(comment, "comment")?) + } + None => None, + }; + + // Parse table schema with partition columns. + let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?; + + let csv_options = CsvOptions { + has_header: Some(scan.has_header), + delimiter: proto_str_to_byte(&scan.delimiter, "delimiter")?, + quote: proto_str_to_byte(&scan.quote, "quote")?, + newlines_in_values: Some(scan.newlines_in_values), + ..Default::default() + }; + let source = Arc::new( + CsvSource::new(table_schema) + .with_csv_options(csv_options) + .with_escape(escape) + .with_comment(comment), + ); + + let conf = FileScanConfigBuilder::from(FileScanConfig::from_proto_conf( + base_conf, ctx, source, + )?) + .with_file_compression_type(FileCompressionType::UNCOMPRESSED) + .build(); + Ok(DataSourceExec::from_data_source(conf)) + } +} diff --git a/datafusion/datasource-json/Cargo.toml b/datafusion/datasource-json/Cargo.toml index b5947ea5c4c67..78d3016c9a549 100644 --- a/datafusion/datasource-json/Cargo.toml +++ b/datafusion/datasource-json/Cargo.toml @@ -30,6 +30,15 @@ version.workspace = true [package.metadata.docs.rs] all-features = true +[features] +# Enables `FileSource::try_to_proto` on `JsonSource` and the `JsonScan` decode +# entry point. Mirrors the `proto` feature on `datafusion-datasource`. +proto = [ + "dep:datafusion-proto-models", + "datafusion-datasource/proto", + "datafusion-physical-plan/proto", +] + [dependencies] arrow = { workspace = true } async-trait = { workspace = true } @@ -41,6 +50,7 @@ datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } datafusion-physical-expr-common = { workspace = true } datafusion-physical-plan = { workspace = true } +datafusion-proto-models = { workspace = true, optional = true } datafusion-session = { workspace = true } futures = { workspace = true } object_store = { workspace = true } diff --git a/datafusion/datasource-json/src/file_format.rs b/datafusion/datasource-json/src/file_format.rs index 1854fddfb84b3..4dbf02a065385 100644 --- a/datafusion/datasource-json/src/file_format.rs +++ b/datafusion/datasource-json/src/file_format.rs @@ -490,6 +490,105 @@ impl DataSink for JsonSink { ) -> Result { FileSink::write_all(self, data, context).await } + + /// Emit a `JsonSink` node wrapping the shared sink config. Kept + /// byte-identical to the former `try_from_data_sink_exec` JsonSink branch in + /// `datafusion-proto`. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + input: datafusion_proto_models::protobuf::PhysicalPlanNode, + sort_order: Option< + datafusion_proto_models::protobuf::PhysicalSortExprNodeCollection, + >, + sink_schema: &Schema, + ) -> Result> { + use datafusion_proto_models::protobuf; + use protobuf::physical_plan_node::PhysicalPlanType; + + let sink = protobuf::JsonSink { + config: Some(self.config.to_proto()?), + writer_options: Some(self.writer_options().try_into()?), + }; + let node = protobuf::JsonSinkExecNode { + input: Some(Box::new(input)), + sink: Some(sink), + sink_schema: Some(sink_schema.try_into()?), + sort_order, + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::JsonSink(Box::new(node))), + })) + } +} + +#[cfg(feature = "proto")] +impl JsonSink { + /// Reconstruct a `DataSinkExec` (wrapping a `JsonSink`) from a `JsonSink` + /// [`PhysicalPlanNode`]. + /// + /// The inverse of [`DataSink::try_to_proto`] on `JsonSink`; kept + /// byte-compatible with the former `try_into_json_sink_physical_plan` in + /// `datafusion-proto`. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_datasource::file_sink_config::parse_sink_sort_order; + use datafusion_datasource::sink::DataSinkExec; + use datafusion_proto_models::protobuf; + + let sink_node = match &node.physical_plan_type { + Some(protobuf::physical_plan_node::PhysicalPlanType::JsonSink(sink)) => { + sink.as_ref() + } + _ => { + return datafusion_common::internal_err!( + "PhysicalPlanNode is not a JsonSink" + ); + } + }; + + let input = ctx.decode_required_child( + sink_node.input.as_deref(), + "JsonSinkExecNode", + "input", + )?; + + let proto_sink = sink_node.sink.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "JsonSinkExecNode is missing required field 'sink'" + ) + })?; + let config = + FileSinkConfig::from_proto(proto_sink.config.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "JsonSink is missing required field 'config'" + ) + })?)?; + let writer_options = proto_sink + .writer_options + .as_ref() + .ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "JsonSink is missing required field 'writer_options'" + ) + })? + .try_into()?; + let data_sink = JsonSink::new(config, writer_options); + + let sink_schema = input.schema(); + let sort_order = + parse_sink_sort_order(sink_node.sort_order.as_ref(), ctx, &sink_schema)?; + + Ok(Arc::new(DataSinkExec::new( + input, + Arc::new(data_sink), + sort_order, + ))) + } } #[derive(Debug)] diff --git a/datafusion/datasource-json/src/source.rs b/datafusion/datasource-json/src/source.rs index 8632d6b942bc1..df30e04f62638 100644 --- a/datafusion/datasource-json/src/source.rs +++ b/datafusion/datasource-json/src/source.rs @@ -231,6 +231,75 @@ impl FileSource for JsonSource { fn file_type(&self) -> &str { "json" } + + /// Emit a `JsonScan` node wrapping the shared base config. JSON has no + /// format-specific fields, so the node is `base_conf` only. Kept + /// byte-identical to the former `try_from_data_source_exec` JsonScan branch + /// in `datafusion-proto`. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + base: &FileScanConfig, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + use protobuf::physical_plan_node::PhysicalPlanType; + + let node = protobuf::JsonScanExecNode { + base_conf: Some(base.to_proto_conf(ctx)?), + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::JsonScan(node)), + })) + } +} + +#[cfg(feature = "proto")] +impl JsonSource { + /// Reconstruct a `DataSourceExec` (wrapping a `FileScanConfig` over a + /// `JsonSource`) from a `JsonScan` [`PhysicalPlanNode`]. + /// + /// The inverse of [`FileSource::try_to_proto`] on `JsonSource`; kept + /// byte-compatible with the former `try_into_json_scan_physical_plan` in + /// `datafusion-proto`. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_datasource::file_scan_config::{ + FileScanConfig, FileScanConfigBuilder, + }; + use datafusion_datasource::source::DataSourceExec; + use datafusion_proto_models::protobuf; + + let scan = match &node.physical_plan_type { + Some(protobuf::physical_plan_node::PhysicalPlanType::JsonScan(scan)) => scan, + _ => { + return datafusion_common::internal_err!( + "PhysicalPlanNode is not a JsonScan" + ); + } + }; + + let base_conf = scan.base_conf.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "JsonScanExecNode is missing required field 'base_conf'" + ) + })?; + + // Parse table schema with partition columns. + let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?; + + let source = Arc::new(JsonSource::new(table_schema)); + + let conf = FileScanConfigBuilder::from(FileScanConfig::from_proto_conf( + base_conf, ctx, source, + )?) + .build(); + Ok(DataSourceExec::from_data_source(conf)) + } } impl FileOpener for JsonOpener { diff --git a/datafusion/datasource-parquet/Cargo.toml b/datafusion/datasource-parquet/Cargo.toml index 32424069c17a0..f3f66a1995b43 100644 --- a/datafusion/datasource-parquet/Cargo.toml +++ b/datafusion/datasource-parquet/Cargo.toml @@ -46,6 +46,8 @@ datafusion-physical-expr = { workspace = true } datafusion-physical-expr-adapter = { workspace = true } datafusion-physical-expr-common = { workspace = true } datafusion-physical-plan = { workspace = true } +datafusion-proto-common = { workspace = true, optional = true } +datafusion-proto-models = { workspace = true, optional = true } datafusion-pruning = { workspace = true } datafusion-session = { workspace = true } futures = { workspace = true } @@ -74,6 +76,16 @@ name = "datafusion_datasource_parquet" path = "src/mod.rs" [features] +# Enables `FileSource::try_to_proto` on `ParquetSource` and the `ParquetScan` +# decode entry point. Mirrors the `proto` feature on `datafusion-datasource-csv`, +# but additionally pulls in `datafusion-proto-common` for the +# `TableParquetOptions` <-> proto conversion. +proto = [ + "dep:datafusion-proto-models", + "dep:datafusion-proto-common", + "datafusion-datasource/proto", + "datafusion-physical-plan/proto", +] parquet_encryption = [ "parquet/encryption", "datafusion-common/parquet_encryption", diff --git a/datafusion/datasource-parquet/src/sink.rs b/datafusion/datasource-parquet/src/sink.rs index f15f67aab0a87..a1470203d35a6 100644 --- a/datafusion/datasource-parquet/src/sink.rs +++ b/datafusion/datasource-parquet/src/sink.rs @@ -40,6 +40,8 @@ use datafusion_datasource::write::{ use datafusion_execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation}; use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; +#[cfg(feature = "proto")] +use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::metrics::{ ElapsedComputeFutureExt, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, MetricsSet, Time, @@ -411,6 +413,105 @@ impl DataSink for ParquetSink { ) -> Result { FileSink::write_all(self, data, context).await } + + /// Emit a `ParquetSink` node wrapping the shared sink config. Kept + /// byte-identical to the former `try_from_data_sink_exec` ParquetSink branch + /// in `datafusion-proto`. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + input: datafusion_proto_models::protobuf::PhysicalPlanNode, + sort_order: Option< + datafusion_proto_models::protobuf::PhysicalSortExprNodeCollection, + >, + sink_schema: &Schema, + ) -> Result> { + use datafusion_proto_models::protobuf; + use protobuf::physical_plan_node::PhysicalPlanType; + + let sink = protobuf::ParquetSink { + config: Some(self.config.to_proto()?), + parquet_options: Some(self.parquet_options().try_into()?), + }; + let node = protobuf::ParquetSinkExecNode { + input: Some(Box::new(input)), + sink: Some(sink), + sink_schema: Some(sink_schema.try_into()?), + sort_order, + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::ParquetSink(Box::new(node))), + })) + } +} + +#[cfg(feature = "proto")] +impl ParquetSink { + /// Reconstruct a `DataSinkExec` (wrapping a `ParquetSink`) from a + /// `ParquetSink` [`PhysicalPlanNode`]. + /// + /// The inverse of [`DataSink::try_to_proto`] on `ParquetSink`; kept + /// byte-compatible with the former `try_into_parquet_sink_physical_plan` in + /// `datafusion-proto`. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_datasource::file_sink_config::parse_sink_sort_order; + use datafusion_datasource::sink::DataSinkExec; + use datafusion_proto_models::protobuf; + + let sink_node = match &node.physical_plan_type { + Some(protobuf::physical_plan_node::PhysicalPlanType::ParquetSink(sink)) => { + sink.as_ref() + } + _ => { + return datafusion_common::internal_err!( + "PhysicalPlanNode is not a ParquetSink" + ); + } + }; + + let input = ctx.decode_required_child( + sink_node.input.as_deref(), + "ParquetSinkExecNode", + "input", + )?; + + let proto_sink = sink_node.sink.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ParquetSinkExecNode is missing required field 'sink'" + ) + })?; + let config = + FileSinkConfig::from_proto(proto_sink.config.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ParquetSink is missing required field 'config'" + ) + })?)?; + let parquet_options = proto_sink + .parquet_options + .as_ref() + .ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ParquetSink is missing required field 'parquet_options'" + ) + })? + .try_into()?; + let data_sink = ParquetSink::new(config, parquet_options); + + let sink_schema = input.schema(); + let sort_order = + parse_sink_sort_order(sink_node.sort_order.as_ref(), ctx, &sink_schema)?; + + Ok(Arc::new(DataSinkExec::new( + input, + Arc::new(data_sink), + sort_order, + ))) + } } /// Consumes a stream of [ArrowLeafColumn] via a channel and serializes them using an [ArrowColumnWriter] diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index 3443b08475e0d..67f3e6c70627f 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -1047,6 +1047,154 @@ impl FileSource for ParquetSource { inner: Arc::new(new_source) as Arc, }) } + + /// Emit a `ParquetScan` node wrapping the shared base config plus the + /// Parquet-specific predicate and `TableParquetOptions`. Kept + /// byte-identical to the former `try_from_data_source_exec` ParquetScan + /// branch in `datafusion-proto`. + /// + /// The predicate rides `ctx.encode_expr` (no raw codec); the options go + /// through `datafusion-proto-common`'s `TableParquetOptions` conversion. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + base: &FileScanConfig, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> datafusion_common::Result< + Option, + > { + use datafusion_proto_models::protobuf; + use protobuf::physical_plan_node::PhysicalPlanType; + + let predicate = self + .filter() + .map(|pred| ctx.encode_expr(&pred)) + .transpose()?; + + let node = protobuf::ParquetScanExecNode { + base_conf: Some(base.to_proto_conf(ctx)?), + predicate, + parquet_options: Some(self.table_parquet_options().try_into()?), + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::ParquetScan(node)), + })) + } +} + +#[cfg(feature = "proto")] +impl ParquetSource { + /// Reconstruct a `DataSourceExec` (wrapping a `FileScanConfig` over a + /// `ParquetSource`) from a `ParquetScan` + /// [`PhysicalPlanNode`]. + /// + /// The inverse of [`FileSource::try_to_proto`] on `ParquetSource`; kept + /// byte-compatible with the former `try_into_parquet_scan_physical_plan` + /// in `datafusion-proto`. + /// + /// The predicate is decoded via `ctx.decode_expr` against the (possibly + /// projected) predicate schema; the `TableParquetOptions` are decoded + /// through `datafusion-proto-common` with no raw codec. The Parquet reader + /// factory is rebuilt as a `CachedParquetFileReaderFactory` from the decode + /// context's runtime environment, exactly as the old central decoder did. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + ) -> datafusion_common::Result> { + use crate::CachedParquetFileReaderFactory; + use arrow::datatypes::Schema; + use datafusion_common::config::TableParquetOptions; + use datafusion_datasource::file_scan_config::FileScanConfig; + use datafusion_datasource::source::DataSourceExec; + use datafusion_execution::object_store::ObjectStoreUrl; + use datafusion_proto_models::protobuf; + + let scan = match &node.physical_plan_type { + Some(protobuf::physical_plan_node::PhysicalPlanType::ParquetScan(scan)) => { + scan + } + _ => { + return datafusion_common::internal_err!( + "PhysicalPlanNode is not a ParquetScan" + ); + } + }; + + let base_conf = scan.base_conf.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ParquetScanExecNode is missing required field 'base_conf'" + ) + })?; + + // Full (file + partition columns) schema off the base conf, matching + // the former `parse_protobuf_file_scan_schema`. + let schema: Arc = Arc::new( + base_conf + .schema + .as_ref() + .ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "FileScanExecConf is missing required field 'schema'" + ) + })? + .try_into()?, + ); + + // Check if there's a projection and use projected schema for predicate parsing + let predicate_schema = if !base_conf.projection.is_empty() { + // Create projected schema for parsing the predicate + let projected_fields: Vec<_> = base_conf + .projection + .iter() + .map(|&i| schema.field(i as usize).clone()) + .collect(); + Arc::new(Schema::new(projected_fields)) + } else { + schema + }; + + let predicate = scan + .predicate + .as_ref() + .map(|expr| ctx.decode_expr(expr, predicate_schema.as_ref())) + .transpose()?; + + let mut options = TableParquetOptions::default(); + if let Some(table_options) = scan.parquet_options.as_ref() { + options = table_options.try_into()?; + } + + // Parse table schema with partition columns + 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(), + }; + let store = ctx + .task_ctx() + .runtime_env() + .object_store(object_store_url)?; + let metadata_cache = ctx + .task_ctx() + .runtime_env() + .cache_manager + .get_file_metadata_cache(); + let reader_factory = + Arc::new(CachedParquetFileReaderFactory::new(store, metadata_cache)); + + let mut source = ParquetSource::new(table_schema) + .with_parquet_file_reader_factory(reader_factory) + .with_table_parquet_options(options); + + if let Some(predicate) = predicate { + source = source.with_predicate(predicate); + } + let base_config = + FileScanConfig::from_proto_conf(base_conf, ctx, Arc::new(source))?; + Ok(DataSourceExec::from_data_source(base_config)) + } } /// Returns the a [`TableSchema`] containing a [`RowNumber`] virtual column and a [`Column`] expression referencing its row index column. diff --git a/datafusion/datasource/Cargo.toml b/datafusion/datasource/Cargo.toml index 2ac42ed900095..714d414047d89 100644 --- a/datafusion/datasource/Cargo.toml +++ b/datafusion/datasource/Cargo.toml @@ -34,6 +34,15 @@ 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", + "dep:datafusion-proto-common", + "datafusion-physical-plan/proto", +] [dependencies] arrow = { workspace = true } @@ -56,6 +65,8 @@ datafusion-physical-expr = { workspace = true } datafusion-physical-expr-adapter = { workspace = true } datafusion-physical-expr-common = { workspace = true } datafusion-physical-plan = { workspace = true } +datafusion-proto-common = { workspace = true, optional = 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..44fdf8dc16d74 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::to_proto_conf`](crate::file_scan_config::FileScanConfig::to_proto_conf), + /// 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 b73d100e056f4..29507e0ebd76e 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 `to_proto_conf` / `from_proto_conf` / `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, @@ -1148,6 +1154,18 @@ impl DataSource for FileScanConfig { Some(Arc::new(SharedWorkSource::from_config(self)) as Arc) } + + /// Serialize this file scan by delegating to the concrete + /// [`FileSource`](crate::file::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 { 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..549f13c7c95ec --- /dev/null +++ b/datafusion/datasource/src/file_scan_config/proto.rs @@ -0,0 +1,521 @@ +// 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`] / [`ExecutionPlanDecodeCtx`] 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 +//! 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. + +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_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}; +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`](crate::file::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( + &self, + ctx: &ExecutionPlanEncodeCtx<'_>, + ) -> Result { + let file_groups = self + .file_groups + .iter() + .map(file_group_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 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, + }); + } + + let output_partitioning = self + .output_partitioning + .as_ref() + .map(|p| partitioning_to_proto(p, 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`. + 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, + partitioned_by_file_group: Some(self.partitioned_by_file_group), + 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 from_proto_conf( + 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(file_group_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 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, + )?; + output_ordering.extend(LexOrdering::new(sort_exprs)); + } + + let output_partitioning = + partitioning_from_proto(conf.output_partitioning.as_ref(), ctx, &schema)?; + + // 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 mut 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)); + if conf.partitioned_by_file_group.unwrap_or(false) { + config_builder = config_builder.with_partitioned_by_file_group(true); + } + 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::from_proto_conf`]. + /// + /// 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)) +} + +fn parse_sort_exprs( + nodes: &[protobuf::PhysicalSortExprNode], + ctx: &ExecutionPlanDecodeCtx<'_>, + 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/file_sink_config.rs b/datafusion/datasource/src/file_sink_config/mod.rs similarity index 92% rename from datafusion/datasource/src/file_sink_config.rs rename to datafusion/datasource/src/file_sink_config/mod.rs index 1abce86a3565f..f20578079e4fc 100644 --- a/datafusion/datasource/src/file_sink_config.rs +++ b/datafusion/datasource/src/file_sink_config/mod.rs @@ -15,6 +15,18 @@ // specific language governing permissions and limitations // under the License. +/// Shared `FileSinkConfig` <-> proto conversion, gated on the `proto` feature. +/// Attaches inherent `to_proto` / `from_proto` helpers to [`FileSinkConfig`] +/// (plus a `parse_sink_sort_order` helper) used by every file sink's +/// `try_to_proto` / `try_from_proto` hook. +#[cfg(feature = "proto")] +pub(crate) mod proto; + +/// Re-exported so each file sink's `try_from_proto` can decode the shared, +/// optional required output ordering without naming the internal `proto` module. +#[cfg(feature = "proto")] +pub use proto::parse_sink_sort_order; + use std::sync::Arc; use crate::ListingTableUrl; diff --git a/datafusion/datasource/src/file_sink_config/proto.rs b/datafusion/datasource/src/file_sink_config/proto.rs new file mode 100644 index 0000000000000..a22f2cc1d9694 --- /dev/null +++ b/datafusion/datasource/src/file_sink_config/proto.rs @@ -0,0 +1,198 @@ +// 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 [`FileSinkConfig`] spine. +//! +//! This is the relocated body of `datafusion-proto`'s `FileSinkConfig` <-> proto +//! conversion (the `TryFromProto<&FileSinkConfig>` / `TryFromProto<&protobuf::FileSinkConfig>` +//! impls). Every concrete sink's `DataSink::try_to_proto` hook (CSV, JSON, +//! Parquet) builds its `*SinkExecNode` around [`FileSinkConfig::to_proto`] and +//! decodes with [`FileSinkConfig::from_proto`], keeping a single copy of the +//! shared wire logic. The wire format is byte-for-byte identical to the old +//! central serializer. +//! +//! Unlike the scan spine, [`FileSinkConfig`] carries no child physical +//! expressions, so this conversion needs neither the `PhysicalExtensionCodec` +//! nor the [`ExecutionPlanEncodeCtx`]/[`ExecutionPlanDecodeCtx`]: `Schema`, +//! `DataType`, and `PartitionedFile` all go through `datafusion-proto-common` +//! (and the shared scan helpers). The sink's optional required-ordering *does* +//! carry exprs, so [`parse_sink_sort_order`] rides the decode ctx (the encode +//! side lives in `DataSinkExec::try_to_proto`). + +use std::sync::Arc; + +use arrow::compute::SortOptions; +use arrow::datatypes::Schema; +use datafusion_common::{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 crate::ListingTableUrl; +use crate::file_groups::FileGroup; +use crate::file_scan_config::proto::{ + partitioned_file_from_proto, partitioned_file_to_proto, +}; +use crate::file_sink_config::{FileOutputMode, FileSinkConfig}; + +impl FileSinkConfig { + /// Serialize the shared, format-agnostic part of a file sink into a + /// [`protobuf::FileSinkConfig`]. + /// + /// Each concrete sink's `DataSink::try_to_proto` wraps the returned value in + /// its own `*Sink` node. Byte-compatible with the former + /// `TryFromProto<&FileSinkConfig> for protobuf::FileSinkConfig` in + /// `datafusion-proto`. + pub fn to_proto(&self) -> Result { + let file_groups = self + .file_group + .iter() + .map(partitioned_file_to_proto) + .collect::>>()?; + let table_paths = self + .table_paths + .iter() + .map(ToString::to_string) + .collect::>(); + 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::>>()?; + 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.to_string(), + file_output_mode: file_output_mode.into(), + }) + } + + /// Reconstruct a [`FileSinkConfig`] from a [`protobuf::FileSinkConfig`]. + /// + /// Byte-compatible with the former + /// `TryFromProto<&protobuf::FileSinkConfig> for FileSinkConfig` in + /// `datafusion-proto`. + pub fn from_proto(conf: &protobuf::FileSinkConfig) -> Result { + let file_group = FileGroup::new( + conf.file_groups + .iter() + .map(partitioned_file_from_proto) + .collect::>>()?, + ); + let table_paths = conf + .table_paths + .iter() + .map(ListingTableUrl::parse) + .collect::>>()?; + 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::>>()?; + 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(FileSinkConfig { + 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 from its serialized +/// `PhysicalSortExprNodeCollection`, resolving the child expressions against +/// `schema` (the decoded input's schema). Byte-compatible with the former +/// `parse_physical_sort_exprs` + `LexRequirement::new` dance in the central +/// `try_into_*_sink_physical_plan` decoders. +/// +/// The encode counterpart is inlined in +/// [`DataSinkExec::try_to_proto`](crate::sink::DataSinkExec) (it needs the +/// [`ExecutionPlanEncodeCtx`](datafusion_physical_plan::proto::ExecutionPlanEncodeCtx)). +pub fn parse_sink_sort_order( + collection: Option<&protobuf::PhysicalSortExprNodeCollection>, + ctx: &ExecutionPlanDecodeCtx<'_>, + schema: &Schema, +) -> Result> { + 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::>>()?; + Ok(LexRequirement::new(sort_exprs.into_iter().map(Into::into))) +} diff --git a/datafusion/datasource/src/memory.rs b/datafusion/datasource/src/memory.rs index a4e30d7f0bd82..31d76e48c3634 100644 --- a/datafusion/datasource/src/memory.rs +++ b/datafusion/datasource/src/memory.rs @@ -256,6 +256,94 @@ impl DataSource for MemorySourceConfig { }) .transpose() } + + /// Emit a `MemoryScan` node describing these in-memory partitions. Kept + /// byte-identical to the former `try_from_data_source_exec` + /// `MemorySourceConfig` branch in `datafusion-proto` (#22419). + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_common::DataFusionError; + use datafusion_proto_models::protobuf; + use protobuf::physical_plan_node::PhysicalPlanType; + + let partitions = self + .partitions + .iter() + .map(|batches| proto_serialize_record_batches(batches)) + .collect::>>()?; + + let projection = self.projection.as_ref().map_or_else(Vec::new, |v| { + v.iter().map(|x| *x as u32).collect::>() + }); + + let sort_information = self + .sort_information + .iter() + .map(|ordering| { + let physical_sort_expr_nodes = ordering + .iter() + .map(|sort_expr| { + Ok::<_, DataFusionError>(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::>>()?; + Ok::<_, DataFusionError>(protobuf::PhysicalSortExprNodeCollection { + physical_sort_expr_nodes, + }) + }) + .collect::>>()?; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::MemoryScan( + protobuf::MemoryScanExecNode { + partitions, + schema: Some((&*self.schema).try_into()?), + projection, + sort_information, + show_sizes: self.show_sizes, + fetch: self.fetch.map(|f| f as u32), + }, + )), + })) + } +} + +/// Serialize a partition's record batches to Arrow IPC stream bytes. Local copy +/// of `datafusion-proto`'s `serialize_record_batches`; kept byte-identical so the +/// `MemoryScan` wire format is unchanged. The batch IPC bytes are fully portable +/// (pure Arrow), so no `datafusion-proto` type is required here. +#[cfg(feature = "proto")] +fn proto_serialize_record_batches(batches: &[RecordBatch]) -> Result> { + if batches.is_empty() { + return Ok(vec![]); + } + let schema = batches[0].schema(); + let mut buf = Vec::new(); + let mut writer = arrow::ipc::writer::StreamWriter::try_new(&mut buf, &schema)?; + for batch in batches { + writer.write(batch)?; + } + writer.finish()?; + Ok(buf) +} + +/// Parse a partition's Arrow IPC stream bytes back into record batches. Local +/// copy of `datafusion-proto`'s `parse_record_batches`; kept byte-identical. +#[cfg(feature = "proto")] +fn proto_parse_record_batches(buf: &[u8]) -> Result> { + if buf.is_empty() { + return Ok(vec![]); + } + let reader = arrow::ipc::reader::StreamReader::try_new(buf, None)?; + reader + .map(|batch| batch.map_err(Into::into)) + .collect::>>() } impl MemorySourceConfig { @@ -607,6 +695,88 @@ impl MemorySourceConfig { } } +#[cfg(feature = "proto")] +impl MemorySourceConfig { + /// Reconstruct a `DataSourceExec` wrapping a `MemorySourceConfig` from a + /// `MemoryScan` [`PhysicalPlanNode`]. The inverse of + /// [`DataSource::try_to_proto`] on `MemorySourceConfig`; kept + /// byte-compatible with the former `try_into_memory_scan_physical_plan` in + /// `datafusion-proto` (#22419). + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_common::{DataFusionError, internal_datafusion_err}; + use datafusion_physical_expr::PhysicalSortExpr; + use datafusion_proto_models::protobuf; + + let scan = match &node.physical_plan_type { + Some(protobuf::physical_plan_node::PhysicalPlanType::MemoryScan(scan)) => { + scan + } + _ => { + return datafusion_common::internal_err!( + "PhysicalPlanNode is not a MemoryScan" + ); + } + }; + + let partitions = scan + .partitions + .iter() + .map(|p| proto_parse_record_batches(p)) + .collect::>>()?; + + let proto_schema = scan.schema.as_ref().ok_or_else(|| { + internal_datafusion_err!("schema in MemoryScanExecNode is missing.") + })?; + let schema: SchemaRef = SchemaRef::new(proto_schema.try_into()?); + + let projection = if !scan.projection.is_empty() { + Some( + scan.projection + .iter() + .map(|i| *i as usize) + .collect::>(), + ) + } else { + None + }; + + let mut sort_information = vec![]; + for ordering in &scan.sort_information { + let sort_exprs = ordering + .physical_sort_expr_nodes + .iter() + .map(|sort_expr| { + let expr_node = sort_expr.expr.as_deref().ok_or_else(|| { + internal_datafusion_err!( + "PhysicalSortExprNode is missing required field 'expr'" + ) + })?; + Ok::<_, DataFusionError>(PhysicalSortExpr { + expr: ctx.decode_expr(expr_node, schema.as_ref())?, + options: arrow::compute::SortOptions { + descending: !sort_expr.asc, + nulls_first: sort_expr.nulls_first, + }, + }) + }) + .collect::>>()?; + sort_information.extend(LexOrdering::new(sort_exprs)); + } + + let source = Self::try_new(&partitions, schema, projection)? + .with_limit(scan.fetch.map(|f| f as usize)) + .with_show_sizes(scan.show_sizes); + let source = source.try_with_sort_information(sort_information)?; + + Ok(DataSourceExec::from_data_source(source)) + } +} + /// For use in repartitioning, track the total size and original partition index. /// /// Do not implement clone, in order to avoid unnecessary copying during repartitioning. diff --git a/datafusion/datasource/src/sink.rs b/datafusion/datasource/src/sink.rs index e3df1ad6381f4..9e8e78dd66836 100644 --- a/datafusion/datasource/src/sink.rs +++ b/datafusion/datasource/src/sink.rs @@ -70,6 +70,43 @@ pub trait DataSink: Any + DisplayAs + Debug + Send + Sync { data: SendableRecordBatchStream, context: &Arc, ) -> Result; + + /// Serialize this sink into a full [`PhysicalPlanNode`] (a `DataSinkExec` + /// wrapping this sink), if it knows how. + /// + /// This is the `DataSink` analog of + /// [`FileSource::try_to_proto`](crate::file::FileSource::try_to_proto): + /// [`DataSinkExec::try_to_proto`] does the shared work of encoding the child + /// plan and the required output ordering, then hands the concrete sink the + /// finished pieces so it only appends its own format-specific fields around + /// the shared [`FileSinkConfig`](crate::file_sink_config::FileSinkConfig) + /// spine via [`FileSinkConfig::to_proto`](crate::file_sink_config::FileSinkConfig::to_proto). + /// + /// * `input` — the already-encoded child plan node. + /// * `sort_order` — the already-encoded required output ordering, if any. + /// * `sink_schema` — the exec's (count) output schema; serialized into the + /// node's `sink_schema` field for byte compatibility (the decoder + /// recomputes it from the input, so it is otherwise informational). + /// + /// The sink carries no child physical expressions of its own, so no encode + /// ctx is needed here (contrast with `FileSource::try_to_proto`). + /// + /// * `Ok(None)` (the default) — this sink has no proto hook; 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, + _input: datafusion_proto_models::protobuf::PhysicalPlanNode, + _sort_order: Option< + datafusion_proto_models::protobuf::PhysicalSortExprNodeCollection, + >, + _sink_schema: &Schema, + ) -> Result> { + Ok(None) + } } impl dyn DataSink { @@ -260,6 +297,47 @@ impl ExecutionPlan for DataSinkExec { fn metrics(&self) -> Option { self.sink.metrics() } + + /// Encodes the shared parts of a sink node — the child plan and the optional + /// required output ordering — then delegates the format-specific wrapping to + /// the concrete [`DataSink::try_to_proto`]. Keeps the format-specific wire + /// logic in the format crate, mirroring `DataSourceExec::try_to_proto`. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr::PhysicalSortExpr; + use datafusion_proto_models::protobuf; + + let input = ctx.encode_child(self.input())?; + + // Required output ordering: 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 sort_order = match self.sort_order() { + Some(requirements) => { + let nodes = requirements + .iter() + .map(|requirement| { + let expr: PhysicalSortExpr = requirement.to_owned().into(); + Ok(protobuf::PhysicalSortExprNode { + expr: Some(Box::new(ctx.encode_expr(&expr.expr)?)), + asc: !expr.options.descending, + nulls_first: expr.options.nulls_first, + }) + }) + .collect::>>()?; + Some(protobuf::PhysicalSortExprNodeCollection { + physical_sort_expr_nodes: nodes, + }) + } + None => None, + }; + + self.sink() + .try_to_proto(input, sort_order, self.schema().as_ref()) + } } /// Create a output record batch with a count diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index 9eb92e7e3525d..b55c426100654 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`](datafusion_physical_plan::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`](crate::file::FileSource::try_to_proto) + /// through the shared [`FileScanConfig`](crate::file_scan_config::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`] @@ -549,6 +573,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`](crate::file::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-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index d7c72253ecc0c..011c80bf31f45 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -2006,6 +2006,427 @@ impl ExecutionPlan for AggregateExec { Ok(result) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let input = ctx.encode_child(self.input())?; + + // Group-by carries (expr, name) pairs plus a flattened groups bitmap and a + // null_expr list. The exprs ride the ctx; the names and bitmap are plain + // data serialized exactly as on the wire (bitmap flattened row-major). + let group_by = self.group_expr(); + let group_expr = + ctx.encode_expressions(group_by.expr().iter().map(|(e, _)| e))?; + let group_expr_name = group_by + .expr() + .iter() + .map(|(_, name)| name.to_owned()) + .collect::>(); + let null_expr = + ctx.encode_expressions(group_by.null_expr().iter().map(|(e, _)| e))?; + let groups: Vec = group_by.groups().iter().flatten().copied().collect(); + + // Each aggregate is a UDAF + arg exprs + ordering + flags + human_display. + // The UDAF rides the ctx (bytes-only), never the codec. + let aggr_expr = self + .aggr_expr() + .iter() + .map(|aggr| encode_aggregate_expr(aggr, ctx)) + .collect::>>()?; + let aggr_expr_name = self + .aggr_expr() + .iter() + .map(|aggr| aggr.name().to_string()) + .collect::>(); + + // Optional per-aggregate filter exprs (MaybeFilter wraps Option). + let filter_expr = self + .filter_expr() + .iter() + .map(|maybe| { + Ok(protobuf::MaybeFilter { + expr: maybe.as_ref().map(|e| ctx.encode_expr(e)).transpose()?, + }) + }) + .collect::>>()?; + + // AggregateMode: inline exhaustive by-NAME match (proto vs common differ in + // numbering, so never cast). + let mode = match self.mode() { + AggregateMode::Partial => protobuf::AggregateMode::Partial, + AggregateMode::Final => protobuf::AggregateMode::Final, + AggregateMode::FinalPartitioned => protobuf::AggregateMode::FinalPartitioned, + AggregateMode::Single => protobuf::AggregateMode::Single, + AggregateMode::SinglePartitioned => { + protobuf::AggregateMode::SinglePartitioned + } + AggregateMode::PartialReduce => protobuf::AggregateMode::PartialReduce, + }; + + let limit = self.limit_options().map(|config| protobuf::AggLimit { + limit: config.limit() as u64, + descending: config.descending(), + }); + + let dynamic_filter = match self.dynamic_filter_expr() { + Some(df) => { + let df_expr: Arc = + Arc::clone(df) as Arc; + Some(ctx.encode_expr(&df_expr)?) + } + None => None, + }; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Aggregate(Box::new( + protobuf::AggregateExecNode { + group_expr, + group_expr_name, + aggr_expr, + filter_expr, + aggr_expr_name, + mode: mode as i32, + input: Some(Box::new(input)), + input_schema: Some(self.input_schema().as_ref().try_into()?), + null_expr, + groups, + limit, + has_grouping_set: group_by.has_grouping_set(), + dynamic_filter, + }, + )), + ), + })) + } +} + +/// Wire-format marker prefixing an encoded `human_display`/`alias` pair. Kept +/// byte-identical to the encoding in `datafusion-proto` so the two formats stay +/// interchangeable. +#[cfg(feature = "proto")] +const HUMAN_DISPLAY_ALIAS_PREFIX: &str = "\u{1f}datafusion_human_display_alias_v1:"; + +/// Encode a `human_display` + visible `alias` into the single `human_display` +/// wire string. Inverse of [`split_human_display_alias`]. +#[cfg(feature = "proto")] +fn encode_human_display_alias(human_display: &str, alias: &str) -> String { + format!( + "{HUMAN_DISPLAY_ALIAS_PREFIX}{}:{alias}{human_display}", + alias.len() + ) +} + +/// Split an encoded `human_display` wire string back into `(human_display, +/// Option)`. Falls back to `(input, None)` for a plain (non-aliased) +/// string or any malformed prefix. Inverse of [`encode_human_display_alias`]. +#[cfg(feature = "proto")] +fn split_human_display_alias<'a>( + human_display: &'a str, + name: &'a str, +) -> (&'a str, Option<&'a str>) { + if let Some(encoded) = human_display.strip_prefix(HUMAN_DISPLAY_ALIAS_PREFIX) + && let Some((alias_len, encoded)) = encoded.split_once(':') + && let Ok(alias_len) = alias_len.parse::() + && let Some(alias) = encoded.get(..alias_len) + && let Some(human_display) = encoded.get(alias_len..) + && alias == name + && !human_display.is_empty() + { + return (human_display, Some(alias)); + } + + (human_display, None) +} + +/// Serialize a single [`AggregateFunctionExpr`] into a `PhysicalExprNode`. +/// +/// Arg exprs and the ordering requirement ride the encode ctx; the ordering is +/// inlined as `PhysicalSortExprNode` (expr through the ctx, asc/nulls_first are +/// plain data) exactly like `SortExec`. The UDAF is serialized through +/// [`ExecutionPlanEncodeCtx::encode_udaf`](crate::proto::ExecutionPlanEncodeCtx::encode_udaf) +/// (bytes-only) and never touches the extension codec. +#[cfg(feature = "proto")] +fn encode_aggregate_expr( + aggr_expr: &Arc, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, +) -> Result { + use datafusion_proto_models::protobuf; + + let arg_exprs = aggr_expr.expressions(); + let expr = ctx.encode_expressions(arg_exprs.iter())?; + let ordering_req = aggr_expr + .order_bys() + .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 name = aggr_expr.fun().name().to_string(); + // UDAF payload goes straight into `fun_definition`; the + // `(!buf.is_empty()).then_some(buf)` semantics are already applied by the ctx. + let fun_definition = ctx.encode_udaf(aggr_expr.fun())?; + + let human_display = match (aggr_expr.human_display(), aggr_expr.human_display_alias()) + { + (Some(display), Some(alias)) => encode_human_display_alias(display, alias), + (Some(display), None) => display.to_string(), + (None, _) => String::new(), + }; + + Ok(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::AggregateExpr( + protobuf::PhysicalAggregateExprNode { + aggregate_function: Some( + protobuf::physical_aggregate_expr_node::AggregateFunction::UserDefinedAggrFunction(name), + ), + expr, + ordering_req, + distinct: aggr_expr.is_distinct(), + ignore_nulls: aggr_expr.ignore_nulls(), + fun_definition, + human_display, + }, + )), + }) +} + +#[cfg(feature = "proto")] +impl AggregateExec { + /// Reconstruct an [`AggregateExec`] from its protobuf representation. + /// + /// The inverse of [`ExecutionPlan::try_to_proto`]. Group-by exprs decode + /// against the child's schema; aggregate arg/ordering/filter exprs and the + /// carried-on-the-wire aggregate input schema decode against that physical + /// schema. Each aggregate's UDAF is rebuilt via + /// [`ExecutionPlanDecodeCtx::decode_udaf`](crate::proto::ExecutionPlanDecodeCtx::decode_udaf) + /// (which owns the payload → codec / registry → codec lookup-order policy), + /// then reassembled with [`AggregateExprBuilder`]. + /// + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr::PhysicalSortExpr; + use datafusion_physical_expr::aggregate::AggregateExprBuilder; + use datafusion_proto_models::protobuf; + use protobuf::physical_aggregate_expr_node::AggregateFunction; + use protobuf::physical_expr_node::ExprType; + + let hash_agg = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Aggregate, + "AggregateExec", + ); + + let input = ctx.decode_required_child( + hash_agg.input.as_deref(), + "AggregateExec", + "input", + )?; + + // AggregateMode: inline exhaustive by-NAME match (never a numeric cast). + let agg_mode = + match protobuf::AggregateMode::try_from(hash_agg.mode).map_err(|_| { + datafusion_common::internal_datafusion_err!( + "Received an AggregateNode message with unknown AggregateMode {}", + hash_agg.mode + ) + })? { + protobuf::AggregateMode::Partial => AggregateMode::Partial, + protobuf::AggregateMode::Final => AggregateMode::Final, + protobuf::AggregateMode::FinalPartitioned => { + AggregateMode::FinalPartitioned + } + protobuf::AggregateMode::Single => AggregateMode::Single, + protobuf::AggregateMode::SinglePartitioned => { + AggregateMode::SinglePartitioned + } + protobuf::AggregateMode::PartialReduce => AggregateMode::PartialReduce, + }; + + let num_expr = hash_agg.group_expr.len(); + + // Group-by / null exprs decode against the CHILD input schema. + let input_schema = input.schema(); + let group_expr = hash_agg + .group_expr + .iter() + .zip(hash_agg.group_expr_name.iter()) + .map(|(expr, name)| { + Ok(( + ctx.decode_expr(expr, input_schema.as_ref())?, + name.to_string(), + )) + }) + .collect::>>()?; + + let null_expr = hash_agg + .null_expr + .iter() + .zip(hash_agg.group_expr_name.iter()) + .map(|(expr, name)| { + Ok(( + ctx.decode_expr(expr, input_schema.as_ref())?, + name.to_string(), + )) + }) + .collect::>>()?; + + let groups: Vec> = if !hash_agg.groups.is_empty() { + hash_agg + .groups + .chunks(num_expr) + .map(|g| g.to_vec()) + .collect() + } else { + vec![] + }; + + let has_grouping_set = hash_agg.has_grouping_set; + + // The aggregate's own physical schema is carried on the wire; aggregate + // arg / ordering / filter exprs decode against it. + let input_schema_proto = hash_agg.input_schema.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "input_schema in AggregateNode is missing." + ) + })?; + let physical_schema: SchemaRef = SchemaRef::new(input_schema_proto.try_into()?); + + let physical_filter_expr = hash_agg + .filter_expr + .iter() + .map(|expr| { + expr.expr + .as_ref() + .map(|e| ctx.decode_expr(e, physical_schema.as_ref())) + .transpose() + }) + .collect::>>()?; + + let physical_aggr_expr: Vec> = hash_agg + .aggr_expr + .iter() + .zip(hash_agg.aggr_expr_name.iter()) + .map(|(expr, name)| { + let expr_type = expr.expr_type.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "Unexpected empty aggregate physical expression" + ) + })?; + let ExprType::AggregateExpr(agg_node) = expr_type else { + return internal_err!( + "Invalid aggregate expression for AggregateExec" + ); + }; + + let input_phy_expr = agg_node + .expr + .iter() + .map(|e| ctx.decode_expr(e, physical_schema.as_ref())) + .collect::>>()?; + + // ordering_req: inline PhysicalSortExprNode (mirrors SortExec). + let order_bys = agg_node + .ordering_req + .iter() + .map(|e| { + let expr_node = e.expr.as_deref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "AggregateExec ordering expression is missing its inner expr" + ) + })?; + Ok(PhysicalSortExpr { + expr: ctx.decode_expr(expr_node, physical_schema.as_ref())?, + options: arrow::compute::SortOptions { + descending: !e.asc, + nulls_first: e.nulls_first, + }, + }) + }) + .collect::>>()?; + + let Some(AggregateFunction::UserDefinedAggrFunction(udaf_name)) = + agg_node.aggregate_function.as_ref() + else { + return internal_err!( + "Invalid AggregateExpr, missing aggregate_function" + ); + }; + + // UDAF rides the ctx (bytes-only); the payload → codec / registry → + // codec lookup-order policy lives inside `decode_udaf`. + let agg_udf = + ctx.decode_udaf(udaf_name, agg_node.fun_definition.as_deref())?; + + let (human_display, human_display_alias) = + split_human_display_alias(&agg_node.human_display, name); + let builder = AggregateExprBuilder::new(agg_udf, input_phy_expr) + .schema(Arc::clone(&physical_schema)) + .alias(name) + .with_ignore_nulls(agg_node.ignore_nulls) + .with_distinct(agg_node.distinct) + .order_by(order_bys) + .human_display(human_display); + let builder = if let Some(alias) = human_display_alias { + builder.human_display_alias(alias) + } else { + builder + }; + builder.build().map(Arc::new) + }) + .collect::>>()?; + + let agg = AggregateExec::try_new( + agg_mode, + PhysicalGroupBy::new(group_expr, null_expr, groups, has_grouping_set), + physical_aggr_expr, + physical_filter_expr, + input, + Arc::clone(&physical_schema), + )?; + + let agg = if let Some(limit_proto) = &hash_agg.limit { + let limit = limit_proto.limit as usize; + let limit_options = match limit_proto.descending { + Some(descending) => LimitOptions::new_with_order(limit, descending), + None => LimitOptions::new(limit), + }; + agg.with_limit_options(Some(limit_options)) + } else { + agg + }; + + let agg = if let Some(dynamic_filter_proto) = &hash_agg.dynamic_filter { + let dynamic_filter_expr = + ctx.decode_expr(dynamic_filter_proto, physical_schema.as_ref())?; + let df = (dynamic_filter_expr as Arc) + .downcast::() + .map_err(|_| { + datafusion_common::internal_datafusion_err!( + "AggregateExec dynamic_filter did not decode to a DynamicFilterPhysicalExpr" + ) + })?; + agg.with_dynamic_filter_expr(df)? + } else { + agg + }; + + Ok(Arc::new(agg)) + } } /// Creates the output schema for an [`AggregateExec`] containing the group by columns followed diff --git a/datafusion/physical-plan/src/analyze.rs b/datafusion/physical-plan/src/analyze.rs index 27e0f5e923d85..532174abbd3a3 100644 --- a/datafusion/physical-plan/src/analyze.rs +++ b/datafusion/physical-plan/src/analyze.rs @@ -297,6 +297,106 @@ impl ExecutionPlan for AnalyzeExec { futures::stream::once(output), ))) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + let (has_metric_categories, metric_categories) = match self.metric_categories() { + Some(cats) => (true, cats.iter().map(|c| c.to_string()).collect()), + None => (false, vec![]), + }; + // Inline by-NAME match: the wire field is an i32; map the common enum to + // the proto enum by variant (never a numeric cast) so a byte-identical + // discriminant is written and a future variant is a compile error. + let format = match self.format() { + ExplainFormat::Indent => protobuf::ExplainFormat::Indent, + ExplainFormat::Tree => protobuf::ExplainFormat::Tree, + ExplainFormat::PostgresJSON => protobuf::ExplainFormat::Pgjson, + ExplainFormat::Graphviz => protobuf::ExplainFormat::Graphviz, + } as i32; + let schema = self.schema().as_ref().try_into()?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Analyze(Box::new( + protobuf::AnalyzeExecNode { + verbose: self.verbose(), + show_statistics: self.show_statistics(), + input: Some(Box::new(input)), + schema: Some(schema), + has_metric_categories, + metric_categories, + format, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl AnalyzeExec { + /// Reconstruct an [`AnalyzeExec`] from its protobuf representation. + /// + /// The inverse of [`ExecutionPlan::try_to_proto`]: it takes the whole + /// [`PhysicalPlanNode`], decodes the single child recursively via the + /// [`ExecutionPlanDecodeCtx`], and reads back the flags, metric categories + /// and output format. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + /// [`ExecutionPlanDecodeCtx`]: crate::proto::ExecutionPlanDecodeCtx + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let analyze = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Analyze, + "AnalyzeExec", + ); + let input = + ctx.decode_required_child(analyze.input.as_deref(), "AnalyzeExec", "input")?; + let metric_categories = if analyze.has_metric_categories { + let cats: Result> = analyze + .metric_categories + .iter() + .map(|s| s.parse::()) + .collect(); + Some(cats?) + } else { + None + }; + let pb_format = + protobuf::ExplainFormat::try_from(analyze.format).map_err(|_| { + DataFusionError::Internal(format!( + "Received an AnalyzeExecNode message with unknown ExplainFormat {}", + analyze.format + )) + })?; + let format = match pb_format { + protobuf::ExplainFormat::Indent => ExplainFormat::Indent, + protobuf::ExplainFormat::Tree => ExplainFormat::Tree, + protobuf::ExplainFormat::Pgjson => ExplainFormat::PostgresJSON, + protobuf::ExplainFormat::Graphviz => ExplainFormat::Graphviz, + }; + let schema = analyze.schema.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "AnalyzeExec is missing required field 'schema'" + ) + })?; + let schema = Arc::new(arrow::datatypes::Schema::try_from(schema)?); + Ok(Arc::new( + AnalyzeExec::builder(analyze.verbose, analyze.show_statistics, input, schema) + .with_metric_categories(metric_categories) + .with_format(format) + .build(), + )) + } } /// Creates the output of AnalyzeExec as a RecordBatch diff --git a/datafusion/physical-plan/src/async_func.rs b/datafusion/physical-plan/src/async_func.rs index efac83bbbe5ba..4f7f41a97b3bd 100644 --- a/datafusion/physical-plan/src/async_func.rs +++ b/datafusion/physical-plan/src/async_func.rs @@ -246,6 +246,80 @@ impl ExecutionPlan for AsyncFuncExec { fn metrics(&self) -> Option { Some(self.metrics.clone_inner()) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + let async_exprs = + ctx.encode_expressions(self.async_exprs.iter().map(|e| &e.func))?; + let async_expr_names = self.async_exprs.iter().map(|e| e.name.clone()).collect(); + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::AsyncFunc(Box::new( + protobuf::AsyncFuncExecNode { + input: Some(Box::new(input)), + async_exprs, + async_expr_names, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl AsyncFuncExec { + /// Reconstruct an [`AsyncFuncExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`], taking the whole + /// [`PhysicalPlanNode`] so every plan's `try_from_proto` shares one + /// signature. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let async_func = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::AsyncFunc, + "AsyncFuncExec", + ); + let input = ctx.decode_required_child( + async_func.input.as_deref(), + "AsyncFuncExec", + "input", + )?; + + if async_func.async_exprs.len() != async_func.async_expr_names.len() { + return datafusion_common::internal_err!( + "AsyncFuncExecNode async_exprs length does not match async_expr_names" + ); + } + + let input_schema = input.schema(); + let async_exprs = async_func + .async_exprs + .iter() + .zip(async_func.async_expr_names.iter()) + .map(|(expr, name)| { + let physical_expr = ctx.decode_expr(expr, input_schema.as_ref())?; + Ok(Arc::new(AsyncFuncExpr::try_new( + name.clone(), + physical_expr, + input_schema.as_ref(), + )?)) + }) + .collect::>>()?; + + Ok(Arc::new(AsyncFuncExec::try_new(async_exprs, input)?)) + } } struct CoalesceInputStream { diff --git a/datafusion/physical-plan/src/buffer.rs b/datafusion/physical-plan/src/buffer.rs index 5e220b7e48544..603759f8a1bf4 100644 --- a/datafusion/physical-plan/src/buffer.rs +++ b/datafusion/physical-plan/src/buffer.rs @@ -290,6 +290,48 @@ impl ExecutionPlan for BufferExec { Ok(Arc::new(Self::new(new_input, self.capacity)) as Arc) }) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Buffer(Box::new( + protobuf::BufferExecNode { + input: Some(Box::new(input)), + capacity: self.capacity() as u64, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl BufferExec { + /// Reconstruct a [`BufferExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]. + /// + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let buffer = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Buffer, + "BufferExec", + ); + let input = + ctx.decode_required_child(buffer.input.as_deref(), "BufferExec", "input")?; + Ok(Arc::new(BufferExec::new(input, buffer.capacity as usize))) + } } /// Represents anything that occupies a capacity in a [MemoryBufferedStream]. diff --git a/datafusion/physical-plan/src/coalesce_batches.rs b/datafusion/physical-plan/src/coalesce_batches.rs index 59b3138b55430..b4929d890fe3e 100644 --- a/datafusion/physical-plan/src/coalesce_batches.rs +++ b/datafusion/physical-plan/src/coalesce_batches.rs @@ -284,6 +284,61 @@ impl ExecutionPlan for CoalesceBatchesExec { ) as Arc) }) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::CoalesceBatches( + Box::new(protobuf::CoalesceBatchesExecNode { + input: Some(Box::new(input)), + target_batch_size: self.target_batch_size() as u32, + fetch: self.fetch().map(|n| n as u32), + }), + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +#[expect(deprecated)] +impl CoalesceBatchesExec { + /// Reconstruct a [`CoalesceBatchesExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]: it takes the whole + /// [`PhysicalPlanNode`] so every plan's `try_from_proto` shares one + /// signature. The child plan is decoded recursively via the + /// [`ExecutionPlanDecodeCtx`]. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + /// [`ExecutionPlanDecodeCtx`]: crate::proto::ExecutionPlanDecodeCtx + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let coalesce_batches = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::CoalesceBatches, + "CoalesceBatchesExec", + ); + let input = ctx.decode_required_child( + coalesce_batches.input.as_deref(), + "CoalesceBatchesExec", + "input", + )?; + Ok(Arc::new( + CoalesceBatchesExec::new(input, coalesce_batches.target_batch_size as usize) + .with_fetch(coalesce_batches.fetch.map(|f| f as usize)), + )) + } } /// Stream for [`CoalesceBatchesExec`]. See [`CoalesceBatchesExec`] for more details. diff --git a/datafusion/physical-plan/src/coalesce_partitions.rs b/datafusion/physical-plan/src/coalesce_partitions.rs index 0a8c5f78882c5..82dba551ddfba 100644 --- a/datafusion/physical-plan/src/coalesce_partitions.rs +++ b/datafusion/physical-plan/src/coalesce_partitions.rs @@ -339,6 +339,56 @@ impl ExecutionPlan for CoalescePartitionsExec { } }) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Merge(Box::new( + protobuf::CoalescePartitionsExecNode { + input: Some(Box::new(input)), + fetch: self.fetch().map(|f| f as u32), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl CoalescePartitionsExec { + /// Reconstruct a [`CoalescePartitionsExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]. Note the proto + /// variant is named `Merge` (node [`CoalescePartitionsExecNode`]). + /// + /// [`CoalescePartitionsExecNode`]: datafusion_proto_models::protobuf::CoalescePartitionsExecNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let merge = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Merge, + "CoalescePartitionsExec", + ); + let input = ctx.decode_required_child( + merge.input.as_deref(), + "CoalescePartitionsExec", + "input", + )?; + Ok(Arc::new( + CoalescePartitionsExec::new(input) + .with_fetch(merge.fetch.map(|f| f as usize)), + )) + } } #[cfg(test)] diff --git a/datafusion/physical-plan/src/coop.rs b/datafusion/physical-plan/src/coop.rs index 7bd84a3a6b392..914055b681dce 100644 --- a/datafusion/physical-plan/src/coop.rs +++ b/datafusion/physical-plan/src/coop.rs @@ -361,6 +361,50 @@ impl ExecutionPlan for CooperativeExec { } } } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Cooperative(Box::new( + protobuf::CooperativeExecNode { + input: Some(Box::new(input)), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl CooperativeExec { + /// Reconstruct a [`CooperativeExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]. + /// + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let cooperative = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Cooperative, + "CooperativeExec", + ); + let input = ctx.decode_required_child( + cooperative.input.as_deref(), + "CooperativeExec", + "input", + )?; + Ok(Arc::new(CooperativeExec::new(input))) + } } /// Creates a [`CooperativeStream`] wrapper around the given [`RecordBatchStream`]. diff --git a/datafusion/physical-plan/src/empty.rs b/datafusion/physical-plan/src/empty.rs index a8f4af5b3d34d..e4fc3cb37977c 100644 --- a/datafusion/physical-plan/src/empty.rs +++ b/datafusion/physical-plan/src/empty.rs @@ -181,6 +181,55 @@ impl ExecutionPlan for EmptyExec { Ok(Arc::new(stats)) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + // EmptyExec is a leaf plan carrying only its schema on the wire. + let schema = self.schema().as_ref().try_into()?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Empty( + protobuf::EmptyExecNode { + schema: Some(schema), + }, + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl EmptyExec { + /// Reconstruct an [`EmptyExec`] from its protobuf representation. + /// + /// The inverse of [`ExecutionPlan::try_to_proto`]: it takes the whole + /// [`PhysicalPlanNode`] and reads back the schema (this leaf plan has no + /// children or expressions to decode). + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + _ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let empty = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Empty, + "EmptyExec", + ); + let schema = empty.schema.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "EmptyExec is missing required field 'schema'" + ) + })?; + let schema = Arc::new(arrow::datatypes::Schema::try_from(schema)?); + Ok(Arc::new(EmptyExec::new(schema))) + } } #[cfg(test)] diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index 76abf73e0ebbe..13c413d174c19 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -757,6 +757,27 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { ) -> Option> { None } + + /// Serialize this plan to its protobuf representation, if it knows how. + /// + /// This is the `ExecutionPlan` analog of + /// [`PhysicalExpr::try_to_proto`](datafusion_physical_expr_common::physical_expr::PhysicalExpr::try_to_proto). + /// + /// * `Ok(None)` (the default) — "I don't serialize myself"; the caller + /// (`datafusion-proto`) falls back to the central downcast chain. Every + /// un-migrated plan keeps its existing behavior. + /// * `Ok(Some(node))` — fully serialized; the caller must not fall back. + /// * `Err(_)` — a real failure (e.g. a child failed to serialize). + /// + /// Only *self-contained* plans should override this — see [`crate::proto`] + /// for the session-dependency boundary. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + Ok(None) + } } impl dyn ExecutionPlan { diff --git a/datafusion/physical-plan/src/explain.rs b/datafusion/physical-plan/src/explain.rs index 98eac3d28b5df..8ec379aa13829 100644 --- a/datafusion/physical-plan/src/explain.rs +++ b/datafusion/physical-plan/src/explain.rs @@ -185,6 +185,209 @@ impl ExecutionPlan for ExplainExec { futures::stream::iter(vec![Ok(record_batch)]), ))) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + // ExplainExec is a leaf plan: it carries its schema, the stringified + // plans, and the verbose flag on the wire (no exec children). + let schema = self.schema().as_ref().try_into()?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Explain( + protobuf::ExplainExecNode { + schema: Some(schema), + stringified_plans: self + .stringified_plans + .iter() + .map(stringified_plan_to_proto) + .collect(), + verbose: self.verbose, + }, + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl ExplainExec { + /// Reconstruct an [`ExplainExec`] from its protobuf representation. + /// + /// The inverse of [`ExecutionPlan::try_to_proto`]: it takes the whole + /// [`PhysicalPlanNode`] and reads back the schema, the stringified plans and + /// the verbose flag (this leaf plan has no children to decode). + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + _ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let explain = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Explain, + "ExplainExec", + ); + let schema = explain.schema.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ExplainExec is missing required field 'schema'" + ) + })?; + let schema = Arc::new(arrow::datatypes::Schema::try_from(schema)?); + Ok(Arc::new(ExplainExec::new( + schema, + explain + .stringified_plans + .iter() + .map(stringified_plan_from_proto) + .collect(), + explain.verbose, + ))) + } +} + +/// Convert a [`StringifiedPlan`] into its protobuf representation. +/// +/// Inlined here (rather than reused from `datafusion-proto`) because the +/// `StringifiedPlan` <-> proto conversion lives in `datafusion-proto`, which +/// sits *above* `datafusion-physical-plan` in the crate graph. The conversion is +/// pure (a match over [`PlanType`](datafusion_common::display::PlanType) and the +/// pure prost model types), so it can be expressed against +/// `datafusion-proto-models` directly. The exhaustive match keeps the wire +/// format byte-identical and makes any added `PlanType` variant a compile error. +#[cfg(feature = "proto")] +fn stringified_plan_to_proto( + stringified_plan: &StringifiedPlan, +) -> datafusion_proto_models::protobuf::StringifiedPlan { + use datafusion_common::display::PlanType; + use datafusion_proto_models::datafusion_common::EmptyMessage; + use datafusion_proto_models::protobuf; + use protobuf::plan_type::PlanTypeEnum::{ + AnalyzedLogicalPlan, FinalAnalyzedLogicalPlan, FinalLogicalPlan, + FinalPhysicalPlan, FinalPhysicalPlanWithSchema, FinalPhysicalPlanWithStats, + InitialLogicalPlan, InitialPhysicalPlan, InitialPhysicalPlanWithSchema, + InitialPhysicalPlanWithStats, OptimizedLogicalPlan, OptimizedPhysicalPlan, + PhysicalPlanError, + }; + + protobuf::StringifiedPlan { + plan_type: match stringified_plan.clone().plan_type { + PlanType::InitialLogicalPlan => Some(protobuf::PlanType { + plan_type_enum: Some(InitialLogicalPlan(EmptyMessage {})), + }), + PlanType::AnalyzedLogicalPlan { analyzer_name } => Some(protobuf::PlanType { + plan_type_enum: Some(AnalyzedLogicalPlan( + protobuf::AnalyzedLogicalPlanType { analyzer_name }, + )), + }), + PlanType::FinalAnalyzedLogicalPlan => Some(protobuf::PlanType { + plan_type_enum: Some(FinalAnalyzedLogicalPlan(EmptyMessage {})), + }), + PlanType::OptimizedLogicalPlan { optimizer_name } => { + Some(protobuf::PlanType { + plan_type_enum: Some(OptimizedLogicalPlan( + protobuf::OptimizedLogicalPlanType { optimizer_name }, + )), + }) + } + PlanType::FinalLogicalPlan => Some(protobuf::PlanType { + plan_type_enum: Some(FinalLogicalPlan(EmptyMessage {})), + }), + PlanType::InitialPhysicalPlan => Some(protobuf::PlanType { + plan_type_enum: Some(InitialPhysicalPlan(EmptyMessage {})), + }), + PlanType::OptimizedPhysicalPlan { optimizer_name } => { + Some(protobuf::PlanType { + plan_type_enum: Some(OptimizedPhysicalPlan( + protobuf::OptimizedPhysicalPlanType { optimizer_name }, + )), + }) + } + PlanType::FinalPhysicalPlan => Some(protobuf::PlanType { + plan_type_enum: Some(FinalPhysicalPlan(EmptyMessage {})), + }), + PlanType::InitialPhysicalPlanWithStats => Some(protobuf::PlanType { + plan_type_enum: Some(InitialPhysicalPlanWithStats(EmptyMessage {})), + }), + PlanType::InitialPhysicalPlanWithSchema => Some(protobuf::PlanType { + plan_type_enum: Some(InitialPhysicalPlanWithSchema(EmptyMessage {})), + }), + PlanType::FinalPhysicalPlanWithStats => Some(protobuf::PlanType { + plan_type_enum: Some(FinalPhysicalPlanWithStats(EmptyMessage {})), + }), + PlanType::FinalPhysicalPlanWithSchema => Some(protobuf::PlanType { + plan_type_enum: Some(FinalPhysicalPlanWithSchema(EmptyMessage {})), + }), + PlanType::PhysicalPlanError => Some(protobuf::PlanType { + plan_type_enum: Some(PhysicalPlanError(EmptyMessage {})), + }), + }, + plan: stringified_plan.plan.to_string(), + } +} + +/// Reconstruct a [`StringifiedPlan`] from its protobuf representation. The +/// inverse of [`stringified_plan_to_proto`]; see that function for why this is +/// inlined here rather than reused from `datafusion-proto`. +#[cfg(feature = "proto")] +fn stringified_plan_from_proto( + stringified_plan: &datafusion_proto_models::protobuf::StringifiedPlan, +) -> StringifiedPlan { + use datafusion_common::display::PlanType; + use datafusion_proto_models::protobuf::plan_type::PlanTypeEnum::{ + AnalyzedLogicalPlan, FinalAnalyzedLogicalPlan, FinalLogicalPlan, + FinalPhysicalPlan, FinalPhysicalPlanWithSchema, FinalPhysicalPlanWithStats, + InitialLogicalPlan, InitialPhysicalPlan, InitialPhysicalPlanWithSchema, + InitialPhysicalPlanWithStats, OptimizedLogicalPlan, OptimizedPhysicalPlan, + PhysicalPlanError, + }; + use datafusion_proto_models::protobuf::{ + AnalyzedLogicalPlanType, OptimizedLogicalPlanType, OptimizedPhysicalPlanType, + }; + + StringifiedPlan { + plan_type: match stringified_plan + .plan_type + .as_ref() + .and_then(|pt| pt.plan_type_enum.as_ref()) + .unwrap_or_else(|| { + panic!( + "Cannot create protobuf::StringifiedPlan from {stringified_plan:?}" + ) + }) { + InitialLogicalPlan(_) => PlanType::InitialLogicalPlan, + AnalyzedLogicalPlan(AnalyzedLogicalPlanType { analyzer_name }) => { + PlanType::AnalyzedLogicalPlan { + analyzer_name: analyzer_name.clone(), + } + } + FinalAnalyzedLogicalPlan(_) => PlanType::FinalAnalyzedLogicalPlan, + OptimizedLogicalPlan(OptimizedLogicalPlanType { optimizer_name }) => { + PlanType::OptimizedLogicalPlan { + optimizer_name: optimizer_name.clone(), + } + } + FinalLogicalPlan(_) => PlanType::FinalLogicalPlan, + InitialPhysicalPlan(_) => PlanType::InitialPhysicalPlan, + InitialPhysicalPlanWithStats(_) => PlanType::InitialPhysicalPlanWithStats, + InitialPhysicalPlanWithSchema(_) => PlanType::InitialPhysicalPlanWithSchema, + OptimizedPhysicalPlan(OptimizedPhysicalPlanType { optimizer_name }) => { + PlanType::OptimizedPhysicalPlan { + optimizer_name: optimizer_name.clone(), + } + } + FinalPhysicalPlan(_) => PlanType::FinalPhysicalPlan, + FinalPhysicalPlanWithStats(_) => PlanType::FinalPhysicalPlanWithStats, + FinalPhysicalPlanWithSchema(_) => PlanType::FinalPhysicalPlanWithSchema, + PhysicalPlanError(_) => PlanType::PhysicalPlanError, + }, + plan: Arc::new(stringified_plan.plan.clone()), + } } /// If this plan should be shown, given the previous plan that was diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs index d23dd380423d1..cf8cc7686c42d 100644 --- a/datafusion/physical-plan/src/filter.rs +++ b/datafusion/physical-plan/src/filter.rs @@ -803,6 +803,102 @@ impl ExecutionPlan for FilterExec { .ok() }) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + let expr = ctx.encode_expr(self.predicate())?; + // Preserve the exact wire format: `None` (full projection) is serialized + // as the identity projection `[0, 1, .., num_fields - 1]` so that it is + // distinguishable from an explicit projection on decode. + let projection = match self.projection() { + None => (0..self.input().schema().fields().len()) + .map(|i| i as u32) + .collect(), + Some(v) => v.iter().map(|x| *x as u32).collect(), + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Filter(Box::new( + protobuf::FilterExecNode { + input: Some(Box::new(input)), + expr: Some(expr), + default_filter_selectivity: self.default_selectivity() as u32, + projection, + batch_size: self.batch_size() as u32, + fetch: self.fetch().map(|f| f as u32), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl FilterExec { + /// Reconstruct a [`FilterExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]: it takes the whole + /// [`PhysicalPlanNode`] so every plan's `try_from_proto` shares one + /// signature. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let filter = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Filter, + "FilterExec", + ); + let input = + ctx.decode_required_child(filter.input.as_deref(), "FilterExec", "input")?; + let predicate = ctx.decode_required_expr( + filter.expr.as_ref(), + input.schema().as_ref(), + "FilterExec", + "expr", + )?; + + let filter_selectivity = filter.default_filter_selectivity.try_into(); + // Preserve the `None` state across proto boundaries. Proto cannot distinguish + // between `None` (full projection) and `Some(vec![])` (empty projection) since + // both serialize as an empty list. If all columns are included, we reconstruct + // `None` to avoid losing this semantic distinction on deserialization. + let num_fields = input.schema().fields().len(); + let mut is_full_projection = filter.projection.len() == num_fields; + let mut projection_vec: Vec = Vec::with_capacity(filter.projection.len()); + for (i, idx) in filter.projection.iter().enumerate() { + let idx = *idx as usize; + is_full_projection &= idx == i; + projection_vec.push(idx); + } + let projection = if is_full_projection { + None + } else { + Some(projection_vec) + }; + let filter = FilterExecBuilder::new(predicate, input) + .apply_projection(projection)? + .with_batch_size(filter.batch_size as usize) + .with_fetch(filter.fetch.map(|f| f as usize)) + .build()?; + match filter_selectivity { + Ok(filter_selectivity) => Ok(Arc::new( + filter.with_default_selectivity(filter_selectivity)?, + )), + Err(_) => Err(datafusion_common::internal_datafusion_err!( + "filter_selectivity in PhysicalPlanNode is invalid " + )), + } + } } impl EmbeddedProjection for FilterExec { diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index 79295ba2fb556..f7bc0d9b2a6a3 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -424,6 +424,62 @@ impl ExecutionPlan for CrossJoinExec { Arc::new(new_right), )))) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let left = ctx.encode_child(self.left())?; + let right = ctx.encode_child(self.right())?; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::CrossJoin(Box::new( + protobuf::CrossJoinExecNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl CrossJoinExec { + /// Reconstruct a [`CrossJoinExec`] from its protobuf representation. + /// + /// The inverse of [`ExecutionPlan::try_to_proto`]. + /// + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let crossjoin = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::CrossJoin, + "CrossJoinExec", + ); + + let left = ctx.decode_required_child( + crossjoin.left.as_deref(), + "CrossJoinExec", + "left", + )?; + let right = ctx.decode_required_child( + crossjoin.right.as_deref(), + "CrossJoinExec", + "right", + )?; + + Ok(Arc::new(CrossJoinExec::new(left, right))) + } } /// [left/right]_col_count are required in case the column statistics are None diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 50a90b0f54633..ad0bb1e5a3a75 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -1761,6 +1761,316 @@ impl ExecutionPlan for HashJoinExec { .ok() .map(|exec| Arc::new(exec) as _) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let left = ctx.encode_child(self.left())?; + let right = ctx.encode_child(self.right())?; + + let on = self + .on() + .iter() + .map(|(l, r)| -> Result { + Ok(protobuf::JoinOn { + left: Some(ctx.encode_expr(l)?), + right: Some(ctx.encode_expr(r)?), + }) + }) + .collect::>>()?; + + // The `*::from_proto` conversions live in `datafusion-proto`, which sits + // above this crate in the dependency graph and is therefore unreachable. + // Inline an exhaustive match from the `datafusion_common` enum to the + // proto enum for each of the four enums this node carries. + let join_type = match self.join_type() { + JoinType::Inner => protobuf::JoinType::Inner, + JoinType::Left => protobuf::JoinType::Left, + JoinType::Right => protobuf::JoinType::Right, + JoinType::Full => protobuf::JoinType::Full, + JoinType::LeftSemi => protobuf::JoinType::Leftsemi, + JoinType::RightSemi => protobuf::JoinType::Rightsemi, + JoinType::LeftAnti => protobuf::JoinType::Leftanti, + JoinType::RightAnti => protobuf::JoinType::Rightanti, + JoinType::LeftMark => protobuf::JoinType::Leftmark, + JoinType::RightMark => protobuf::JoinType::Rightmark, + }; + let null_equality = match self.null_equality() { + NullEquality::NullEqualsNothing => protobuf::NullEquality::NullEqualsNothing, + NullEquality::NullEqualsNull => protobuf::NullEquality::NullEqualsNull, + }; + let partition_mode = match self.partition_mode() { + PartitionMode::CollectLeft => protobuf::PartitionMode::CollectLeft, + PartitionMode::Partitioned => protobuf::PartitionMode::Partitioned, + PartitionMode::Auto => protobuf::PartitionMode::Auto, + }; + + let filter = self + .filter() + .map(|f| -> Result { + let expression = ctx.encode_expr(f.expression())?; + let column_indices = f + .column_indices() + .iter() + .map(|i| { + let side = match i.side { + JoinSide::Left => protobuf::JoinSide::LeftSide, + JoinSide::Right => protobuf::JoinSide::RightSide, + JoinSide::None => protobuf::JoinSide::None, + }; + protobuf::ColumnIndex { + index: i.index as u32, + side: side.into(), + } + }) + .collect(); + // `JoinFilter::schema` is `arrow::Schema`; the proto field is + // `datafusion_proto_common::protobuf_common::Schema`, so this + // resolves to the `TryFrom<&Schema>` impl in datafusion-proto-common + // (a dependency of this crate under the `proto` feature). + let schema = f.schema().as_ref().try_into()?; + Ok(protobuf::JoinFilter { + expression: Some(expression), + column_indices, + schema: Some(schema), + }) + }) + .transpose()?; + + let dynamic_filter = self + .dynamic_filter_expr() + .map(|df| { + let df_expr: Arc = + Arc::clone(df) as Arc; + ctx.encode_expr(&df_expr) + }) + .transpose()?; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::HashJoin(Box::new( + protobuf::HashJoinExecNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + on, + join_type: join_type.into(), + partition_mode: partition_mode.into(), + null_equality: null_equality.into(), + filter, + // Proto3 `repeated` cannot distinguish `None` from + // `Some(vec![])`. `Some(vec![])` (reachable via + // `try_embed_projection` for e.g. `SELECT count(1) … JOIN …`) + // changes the output schema, so it is encoded with the + // single-element sentinel `[u32::MAX]` (never a valid column + // index); every other state is sent as-is. See + // `try_from_proto` for the matching decoder. + projection: match self.projection.as_ref() { + None => Vec::new(), + Some(v) if v.is_empty() => vec![u32::MAX], + Some(v) => v.iter().map(|x| *x as u32).collect(), + }, + null_aware: self.null_aware, + dynamic_filter, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl HashJoinExec { + /// Reconstruct a [`HashJoinExec`] from its protobuf representation. + /// + /// The inverse of [`ExecutionPlan::try_to_proto`]. It takes the whole + /// [`PhysicalPlanNode`] (so every plan's `try_from_proto` shares one + /// signature) and drives child-plan / child-expression recursion through the + /// [`ExecutionPlanDecodeCtx`]. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + /// [`ExecutionPlanDecodeCtx`]: crate::proto::ExecutionPlanDecodeCtx + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_common::internal_datafusion_err; + use datafusion_proto_models::protobuf; + use std::any::Any; + + let hashjoin = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::HashJoin, + "HashJoinExec", + ); + + let left = + ctx.decode_required_child(hashjoin.left.as_deref(), "HashJoinExec", "left")?; + let right = ctx.decode_required_child( + hashjoin.right.as_deref(), + "HashJoinExec", + "right", + )?; + let left_schema = left.schema(); + let right_schema = right.schema(); + + // The `on` pairs decode against their respective child schemas. + let on: Vec<(PhysicalExprRef, PhysicalExprRef)> = hashjoin + .on + .iter() + .map(|col| { + let l = ctx.decode_required_expr( + col.left.as_ref(), + left_schema.as_ref(), + "HashJoinExec", + "on.left", + )?; + let r = ctx.decode_required_expr( + col.right.as_ref(), + right_schema.as_ref(), + "HashJoinExec", + "on.right", + )?; + Ok((l, r)) + }) + .collect::>()?; + + // Inline exhaustive proto -> `datafusion_common` enum conversions (the + // `from_proto` helpers live in the unreachable datafusion-proto crate). + let join_type = + match protobuf::JoinType::try_from(hashjoin.join_type).map_err(|_| { + internal_datafusion_err!( + "HashJoinExec: unknown JoinType {}", + hashjoin.join_type + ) + })? { + protobuf::JoinType::Inner => JoinType::Inner, + protobuf::JoinType::Left => JoinType::Left, + protobuf::JoinType::Right => JoinType::Right, + protobuf::JoinType::Full => JoinType::Full, + protobuf::JoinType::Leftsemi => JoinType::LeftSemi, + protobuf::JoinType::Rightsemi => JoinType::RightSemi, + protobuf::JoinType::Leftanti => JoinType::LeftAnti, + protobuf::JoinType::Rightanti => JoinType::RightAnti, + protobuf::JoinType::Leftmark => JoinType::LeftMark, + protobuf::JoinType::Rightmark => JoinType::RightMark, + }; + let null_equality = match protobuf::NullEquality::try_from(hashjoin.null_equality) + .map_err(|_| { + internal_datafusion_err!( + "HashJoinExec: unknown NullEquality {}", + hashjoin.null_equality + ) + })? { + protobuf::NullEquality::NullEqualsNothing => NullEquality::NullEqualsNothing, + protobuf::NullEquality::NullEqualsNull => NullEquality::NullEqualsNull, + }; + let partition_mode = match protobuf::PartitionMode::try_from( + hashjoin.partition_mode, + ) + .map_err(|_| { + internal_datafusion_err!( + "HashJoinExec: unknown PartitionMode {}", + hashjoin.partition_mode + ) + })? { + protobuf::PartitionMode::CollectLeft => PartitionMode::CollectLeft, + protobuf::PartitionMode::Partitioned => PartitionMode::Partitioned, + protobuf::PartitionMode::Auto => PartitionMode::Auto, + }; + + let filter = hashjoin + .filter + .as_ref() + .map(|f| -> Result { + // The `JoinFilter` carries its own intermediate schema; the filter + // expression and column indices are relative to it. The proto + // schema type is datafusion-proto-common's, so `try_into` resolves + // to that crate's `TryFrom<&Schema>` impl. + let schema: Schema = f + .schema + .as_ref() + .ok_or_else(|| { + internal_datafusion_err!( + "HashJoinExec: JoinFilter missing schema" + ) + })? + .try_into()?; + let expression = ctx.decode_required_expr( + f.expression.as_ref(), + &schema, + "HashJoinExec", + "filter.expression", + )?; + let column_indices = f + .column_indices + .iter() + .map(|i| { + let side = + match protobuf::JoinSide::try_from(i.side).map_err(|_| { + internal_datafusion_err!( + "HashJoinExec: unknown JoinSide {}", + i.side + ) + })? { + protobuf::JoinSide::LeftSide => JoinSide::Left, + protobuf::JoinSide::RightSide => JoinSide::Right, + protobuf::JoinSide::None => JoinSide::None, + }; + Ok(ColumnIndex { + index: i.index as usize, + side, + }) + }) + .collect::>>()?; + Ok(JoinFilter::new( + expression, + column_indices, + Arc::new(schema), + )) + }) + .transpose()?; + + // Preserve the empty-projection sentinel written by `try_to_proto`. + let projection = match hashjoin.projection.as_slice() { + [] => None, + [u32::MAX] => Some(Vec::new()), + indices => Some(indices.iter().map(|i| *i as usize).collect()), + }; + + let mut hash_join = HashJoinExec::try_new( + left, + right, + on, + filter, + &join_type, + projection, + partition_mode, + null_equality, + hashjoin.null_aware, + )?; + + if let Some(dynamic_filter_proto) = &hashjoin.dynamic_filter { + // The dynamic filter is a `DynamicFilterPhysicalExpr` over the probe + // (right) side; decode against the right schema then downcast. + let dynamic_filter_expr = + ctx.decode_expr(dynamic_filter_proto, right_schema.as_ref())?; + let df = (dynamic_filter_expr as Arc) + .downcast::() + .map_err(|_| { + internal_datafusion_err!( + "HashJoinExec dynamic_filter did not decode to a DynamicFilterPhysicalExpr" + ) + })?; + hash_join = hash_join.with_dynamic_filter_expr(df)?; + } + + Ok(Arc::new(hash_join)) + } } /// Determines which sides of a join are "preserved" for filter pushdown. diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index 3cae05a3a815a..3bbe2d5f6c20d 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -761,6 +761,208 @@ impl ExecutionPlan for NestedLoopJoinExec { try_embed_projection(projection, self) } } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let left = ctx.encode_child(self.left())?; + let right = ctx.encode_child(self.right())?; + + // The `*::from_proto` conversions live in `datafusion-proto`, which sits + // above this crate in the dependency graph and is therefore unreachable. + // Inline an exhaustive match from the `datafusion_common` enum to the + // proto enum (a numeric cast would corrupt the value, as the two enums + // number their variants differently). + let join_type = match self.join_type() { + JoinType::Inner => protobuf::JoinType::Inner, + JoinType::Left => protobuf::JoinType::Left, + JoinType::Right => protobuf::JoinType::Right, + JoinType::Full => protobuf::JoinType::Full, + JoinType::LeftSemi => protobuf::JoinType::Leftsemi, + JoinType::RightSemi => protobuf::JoinType::Rightsemi, + JoinType::LeftAnti => protobuf::JoinType::Leftanti, + JoinType::RightAnti => protobuf::JoinType::Rightanti, + JoinType::LeftMark => protobuf::JoinType::Leftmark, + JoinType::RightMark => protobuf::JoinType::Rightmark, + }; + + let filter = self + .filter() + .map(|f| -> Result { + let expression = ctx.encode_expr(f.expression())?; + let column_indices = f + .column_indices() + .iter() + .map(|i| { + let side = match i.side { + JoinSide::Left => protobuf::JoinSide::LeftSide, + JoinSide::Right => protobuf::JoinSide::RightSide, + JoinSide::None => protobuf::JoinSide::None, + }; + protobuf::ColumnIndex { + index: i.index as u32, + side: side.into(), + } + }) + .collect(); + // `JoinFilter::schema` is `arrow::Schema`; the proto field is + // `datafusion_proto_common::protobuf_common::Schema`, so this + // resolves to the `TryFrom<&Schema>` impl in datafusion-proto-common. + let schema = f.schema().as_ref().try_into()?; + Ok(protobuf::JoinFilter { + expression: Some(expression), + column_indices, + schema: Some(schema), + }) + }) + .transpose()?; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::NestedLoopJoin(Box::new( + protobuf::NestedLoopJoinExecNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + join_type: join_type.into(), + filter, + // Proto3 `repeated` cannot distinguish `None` from + // `Some(vec![])`; the single-element sentinel `[u32::MAX]` + // (never a valid column index) encodes the empty-but-present + // projection. See `try_from_proto` for the decoder. + projection: match self.projection.as_ref() { + None => Vec::new(), + Some(v) if v.is_empty() => vec![u32::MAX], + Some(v) => v.iter().map(|x| *x as u32).collect(), + }, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl NestedLoopJoinExec { + /// Reconstruct a [`NestedLoopJoinExec`] from its protobuf representation. + /// + /// The inverse of [`ExecutionPlan::try_to_proto`]; it takes the whole + /// [`PhysicalPlanNode`] and drives child-plan / child-expression recursion + /// through the [`ExecutionPlanDecodeCtx`]. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + /// [`ExecutionPlanDecodeCtx`]: crate::proto::ExecutionPlanDecodeCtx + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let join = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::NestedLoopJoin, + "NestedLoopJoinExec", + ); + + let left = ctx.decode_required_child( + join.left.as_deref(), + "NestedLoopJoinExec", + "left", + )?; + let right = ctx.decode_required_child( + join.right.as_deref(), + "NestedLoopJoinExec", + "right", + )?; + + // Inline exhaustive proto -> `datafusion_common` enum conversion (the + // `from_proto` helper lives in the unreachable datafusion-proto crate). + let join_type = + match protobuf::JoinType::try_from(join.join_type).map_err(|_| { + internal_datafusion_err!( + "NestedLoopJoinExec: unknown JoinType {}", + join.join_type + ) + })? { + protobuf::JoinType::Inner => JoinType::Inner, + protobuf::JoinType::Left => JoinType::Left, + protobuf::JoinType::Right => JoinType::Right, + protobuf::JoinType::Full => JoinType::Full, + protobuf::JoinType::Leftsemi => JoinType::LeftSemi, + protobuf::JoinType::Rightsemi => JoinType::RightSemi, + protobuf::JoinType::Leftanti => JoinType::LeftAnti, + protobuf::JoinType::Rightanti => JoinType::RightAnti, + protobuf::JoinType::Leftmark => JoinType::LeftMark, + protobuf::JoinType::Rightmark => JoinType::RightMark, + }; + + let filter = join + .filter + .as_ref() + .map(|f| -> Result { + // The `JoinFilter` carries its own intermediate schema; the filter + // expression and column indices are relative to it. The proto + // schema type is datafusion-proto-common's, so `try_into` resolves + // to that crate's `TryFrom<&Schema>` impl. + let schema: Schema = f + .schema + .as_ref() + .ok_or_else(|| { + internal_datafusion_err!( + "NestedLoopJoinExec: JoinFilter missing schema" + ) + })? + .try_into()?; + let expression = ctx.decode_required_expr( + f.expression.as_ref(), + &schema, + "NestedLoopJoinExec", + "filter.expression", + )?; + let column_indices = f + .column_indices + .iter() + .map(|i| { + let side = + match protobuf::JoinSide::try_from(i.side).map_err(|_| { + internal_datafusion_err!( + "NestedLoopJoinExec: unknown JoinSide {}", + i.side + ) + })? { + protobuf::JoinSide::LeftSide => JoinSide::Left, + protobuf::JoinSide::RightSide => JoinSide::Right, + protobuf::JoinSide::None => JoinSide::None, + }; + Ok(ColumnIndex { + index: i.index as usize, + side, + }) + }) + .collect::>>()?; + Ok(JoinFilter::new( + expression, + column_indices, + Arc::new(schema), + )) + }) + .transpose()?; + + // Preserve the empty-projection sentinel written by `try_to_proto`. + let projection = match join.projection.as_slice() { + [] => None, + [u32::MAX] => Some(Vec::new()), + indices => Some(indices.iter().map(|i| *i as usize).collect()), + }; + + Ok(Arc::new(NestedLoopJoinExec::try_new( + left, right, filter, &join_type, projection, + )?)) + } } impl EmbeddedProjection for NestedLoopJoinExec { diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index 2dc7065eee04e..82b155f6b7139 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -642,4 +642,277 @@ impl ExecutionPlan for SortMergeJoinExec { self.null_equality, )?))) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let left = ctx.encode_child(self.left())?; + let right = ctx.encode_child(self.right())?; + + let on = self + .on() + .iter() + .map(|(l, r)| -> Result { + Ok(protobuf::JoinOn { + left: Some(ctx.encode_expr(l)?), + right: Some(ctx.encode_expr(r)?), + }) + }) + .collect::>>()?; + + // The `*::from_proto` conversions live in `datafusion-proto`, which sits + // above this crate in the dependency graph and is therefore unreachable. + // Inline an exhaustive match from the `datafusion_common` enum to the + // proto enum, by name (the two enums are numbered differently, so a + // numeric cast would silently corrupt the value). + let join_type = match self.join_type() { + JoinType::Inner => protobuf::JoinType::Inner, + JoinType::Left => protobuf::JoinType::Left, + JoinType::Right => protobuf::JoinType::Right, + JoinType::Full => protobuf::JoinType::Full, + JoinType::LeftSemi => protobuf::JoinType::Leftsemi, + JoinType::RightSemi => protobuf::JoinType::Rightsemi, + JoinType::LeftAnti => protobuf::JoinType::Leftanti, + JoinType::RightAnti => protobuf::JoinType::Rightanti, + JoinType::LeftMark => protobuf::JoinType::Leftmark, + JoinType::RightMark => protobuf::JoinType::Rightmark, + }; + let null_equality = match self.null_equality() { + NullEquality::NullEqualsNothing => protobuf::NullEquality::NullEqualsNothing, + NullEquality::NullEqualsNull => protobuf::NullEquality::NullEqualsNull, + }; + + let filter = self + .filter() + .as_ref() + .map(|f| -> Result { + let expression = ctx.encode_expr(f.expression())?; + let column_indices = f + .column_indices() + .iter() + .map(|i| { + let side = match i.side { + JoinSide::Left => protobuf::JoinSide::LeftSide, + JoinSide::Right => protobuf::JoinSide::RightSide, + JoinSide::None => protobuf::JoinSide::None, + }; + protobuf::ColumnIndex { + index: i.index as u32, + side: side.into(), + } + }) + .collect(); + // `JoinFilter::schema` is `arrow::Schema`; the proto field is + // datafusion-proto-common's `Schema`, so `try_into` resolves to + // that crate's `TryFrom<&Schema>` impl. + let schema = f.schema().as_ref().try_into()?; + Ok(protobuf::JoinFilter { + expression: Some(expression), + column_indices, + schema: Some(schema), + }) + }) + .transpose()?; + + // Each key's `SortOptions` is plain data. The wire type carries an + // (unused for this node) expr slot, kept `None` to preserve the byte + // layout; only asc/nulls_first are meaningful. + let sort_options = self + .sort_options() + .iter() + .map(|o| protobuf::SortExprNode { + expr: None, + asc: !o.descending, + nulls_first: o.nulls_first, + }) + .collect(); + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::SortMergeJoin(Box::new( + protobuf::SortMergeJoinExecNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + on, + join_type: join_type.into(), + null_equality: null_equality.into(), + filter, + sort_options, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl SortMergeJoinExec { + /// Reconstruct a [`SortMergeJoinExec`] from its protobuf representation. + /// + /// The inverse of [`ExecutionPlan::try_to_proto`]: it takes the whole + /// [`PhysicalPlanNode`] and drives child-plan / child-expression recursion + /// through the [`ExecutionPlanDecodeCtx`]. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + /// [`ExecutionPlanDecodeCtx`]: crate::proto::ExecutionPlanDecodeCtx + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use crate::joins::utils::ColumnIndex; + use arrow::datatypes::Schema; + use datafusion_common::internal_datafusion_err; + use datafusion_proto_models::protobuf; + + let sort_join = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::SortMergeJoin, + "SortMergeJoinExec", + ); + + let left = ctx.decode_required_child( + sort_join.left.as_deref(), + "SortMergeJoinExec", + "left", + )?; + let right = ctx.decode_required_child( + sort_join.right.as_deref(), + "SortMergeJoinExec", + "right", + )?; + let left_schema = left.schema(); + let right_schema = right.schema(); + + // The `on` pairs decode against their respective child schemas. + let on: JoinOn = sort_join + .on + .iter() + .map(|col| { + let l = ctx.decode_required_expr( + col.left.as_ref(), + left_schema.as_ref(), + "SortMergeJoinExec", + "on.left", + )?; + let r = ctx.decode_required_expr( + col.right.as_ref(), + right_schema.as_ref(), + "SortMergeJoinExec", + "on.right", + )?; + Ok((l, r)) + }) + .collect::>()?; + + // Inline exhaustive proto -> `datafusion_common` enum conversions (the + // `from_proto` helpers live in the unreachable datafusion-proto crate); + // matched by name because the proto/common enums are numbered + // differently. + let join_type = + match protobuf::JoinType::try_from(sort_join.join_type).map_err(|_| { + internal_datafusion_err!( + "SortMergeJoinExec: unknown JoinType {}", + sort_join.join_type + ) + })? { + protobuf::JoinType::Inner => JoinType::Inner, + protobuf::JoinType::Left => JoinType::Left, + protobuf::JoinType::Right => JoinType::Right, + protobuf::JoinType::Full => JoinType::Full, + protobuf::JoinType::Leftsemi => JoinType::LeftSemi, + protobuf::JoinType::Rightsemi => JoinType::RightSemi, + protobuf::JoinType::Leftanti => JoinType::LeftAnti, + protobuf::JoinType::Rightanti => JoinType::RightAnti, + protobuf::JoinType::Leftmark => JoinType::LeftMark, + protobuf::JoinType::Rightmark => JoinType::RightMark, + }; + let null_equality = match protobuf::NullEquality::try_from( + sort_join.null_equality, + ) + .map_err(|_| { + internal_datafusion_err!( + "SortMergeJoinExec: unknown NullEquality {}", + sort_join.null_equality + ) + })? { + protobuf::NullEquality::NullEqualsNothing => NullEquality::NullEqualsNothing, + protobuf::NullEquality::NullEqualsNull => NullEquality::NullEqualsNull, + }; + + let filter = sort_join + .filter + .as_ref() + .map(|f| -> Result { + // The `JoinFilter` carries its own intermediate schema; the + // expression and column indices are relative to it. The proto + // schema type is datafusion-proto-common's, so `try_into` + // resolves to that crate's `TryFrom<&Schema>` impl. + let schema: Schema = f + .schema + .as_ref() + .ok_or_else(|| { + internal_datafusion_err!( + "SortMergeJoinExec: JoinFilter missing schema" + ) + })? + .try_into()?; + let expression = ctx.decode_required_expr( + f.expression.as_ref(), + &schema, + "SortMergeJoinExec", + "filter.expression", + )?; + let column_indices = f + .column_indices + .iter() + .map(|i| { + let side = + match protobuf::JoinSide::try_from(i.side).map_err(|_| { + internal_datafusion_err!( + "SortMergeJoinExec: unknown JoinSide {}", + i.side + ) + })? { + protobuf::JoinSide::LeftSide => JoinSide::Left, + protobuf::JoinSide::RightSide => JoinSide::Right, + protobuf::JoinSide::None => JoinSide::None, + }; + Ok(ColumnIndex { + index: i.index as usize, + side, + }) + }) + .collect::>>()?; + Ok(JoinFilter::new( + expression, + column_indices, + Arc::new(schema), + )) + }) + .transpose()?; + + let sort_options = sort_join + .sort_options + .iter() + .map(|o| SortOptions { + descending: !o.asc, + nulls_first: o.nulls_first, + }) + .collect(); + + Ok(Arc::new(SortMergeJoinExec::try_new( + left, + right, + on, + filter, + join_type, + sort_options, + null_equality, + )?)) + } } diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 52a1aa056d244..beb68cc1ec37d 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -653,6 +653,342 @@ impl ExecutionPlan for SymmetricHashJoinExec { ) .map(|e| Some(Arc::new(e) as _)) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let left = ctx.encode_child(self.left())?; + let right = ctx.encode_child(self.right())?; + + let on = self + .on() + .iter() + .map(|(l, r)| -> Result { + Ok(protobuf::JoinOn { + left: Some(ctx.encode_expr(l)?), + right: Some(ctx.encode_expr(r)?), + }) + }) + .collect::>>()?; + + // The `*::from_proto` conversions live in `datafusion-proto`, which sits + // above this crate in the dependency graph and is therefore unreachable. + // Inline an exhaustive by-name match from the `datafusion_common` enum to + // the proto enum for each enum this node carries. + let join_type = match self.join_type() { + JoinType::Inner => protobuf::JoinType::Inner, + JoinType::Left => protobuf::JoinType::Left, + JoinType::Right => protobuf::JoinType::Right, + JoinType::Full => protobuf::JoinType::Full, + JoinType::LeftSemi => protobuf::JoinType::Leftsemi, + JoinType::RightSemi => protobuf::JoinType::Rightsemi, + JoinType::LeftAnti => protobuf::JoinType::Leftanti, + JoinType::RightAnti => protobuf::JoinType::Rightanti, + JoinType::LeftMark => protobuf::JoinType::Leftmark, + JoinType::RightMark => protobuf::JoinType::Rightmark, + }; + let null_equality = match self.null_equality() { + NullEquality::NullEqualsNothing => protobuf::NullEquality::NullEqualsNothing, + NullEquality::NullEqualsNull => protobuf::NullEquality::NullEqualsNull, + }; + let partition_mode = match self.partition_mode() { + StreamJoinPartitionMode::SinglePartition => { + protobuf::StreamPartitionMode::SinglePartition + } + StreamJoinPartitionMode::Partitioned => { + protobuf::StreamPartitionMode::PartitionedExec + } + }; + + let filter = self + .filter() + .map(|f| -> Result { + let expression = ctx.encode_expr(f.expression())?; + let column_indices = f + .column_indices() + .iter() + .map(|i| { + let side = match i.side { + JoinSide::Left => protobuf::JoinSide::LeftSide, + JoinSide::Right => protobuf::JoinSide::RightSide, + JoinSide::None => protobuf::JoinSide::None, + }; + protobuf::ColumnIndex { + index: i.index as u32, + side: side.into(), + } + }) + .collect(); + // `JoinFilter::schema` is `arrow::Schema`; the proto field is + // datafusion-proto-common's `Schema`, so this resolves to the + // `TryFrom<&Schema>` impl in datafusion-proto-common. + let schema = f.schema().as_ref().try_into()?; + Ok(protobuf::JoinFilter { + expression: Some(expression), + column_indices, + schema: Some(schema), + }) + }) + .transpose()?; + + // A `PhysicalSortExpr` is a `PhysicalExpr` + `SortOptions`. The expr goes + // through the ctx; the asc/nulls_first wrapping into a + // `PhysicalSortExprNode` is plain data inlined here. `None` and an empty + // ordering both encode to an empty `repeated`, matching the original + // `.unwrap_or(vec![])` wire format. + let encode_sort_exprs = + |exprs: Option<&LexOrdering>| -> Result> { + exprs + .map(|exprs| { + exprs + .iter() + .map(|expr| { + Ok(protobuf::PhysicalSortExprNode { + expr: Some(Box::new(ctx.encode_expr(&expr.expr)?)), + asc: !expr.options.descending, + nulls_first: expr.options.nulls_first, + }) + }) + .collect::>>() + }) + .transpose() + .map(|v| v.unwrap_or_default()) + }; + let left_sort_exprs = encode_sort_exprs(self.left_sort_exprs())?; + let right_sort_exprs = encode_sort_exprs(self.right_sort_exprs())?; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::SymmetricHashJoin( + Box::new(protobuf::SymmetricHashJoinExecNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + on, + join_type: join_type.into(), + partition_mode: partition_mode.into(), + null_equality: null_equality.into(), + filter, + left_sort_exprs, + right_sort_exprs, + }), + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl SymmetricHashJoinExec { + /// Reconstruct a [`SymmetricHashJoinExec`] from its protobuf representation. + /// + /// The inverse of [`ExecutionPlan::try_to_proto`]. It takes the whole + /// [`PhysicalPlanNode`] (so every plan's `try_from_proto` shares one + /// signature) and drives child-plan / child-expression recursion through the + /// [`ExecutionPlanDecodeCtx`]. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + /// [`ExecutionPlanDecodeCtx`]: crate::proto::ExecutionPlanDecodeCtx + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_common::internal_datafusion_err; + use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; + use datafusion_proto_models::protobuf; + + let sym_join = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::SymmetricHashJoin, + "SymmetricHashJoinExec", + ); + + let left = ctx.decode_required_child( + sym_join.left.as_deref(), + "SymmetricHashJoinExec", + "left", + )?; + let right = ctx.decode_required_child( + sym_join.right.as_deref(), + "SymmetricHashJoinExec", + "right", + )?; + let left_schema = left.schema(); + let right_schema = right.schema(); + + // The `on` pairs decode against their respective child schemas. + let on: Vec<(PhysicalExprRef, PhysicalExprRef)> = sym_join + .on + .iter() + .map(|col| { + let l = ctx.decode_required_expr( + col.left.as_ref(), + left_schema.as_ref(), + "SymmetricHashJoinExec", + "on.left", + )?; + let r = ctx.decode_required_expr( + col.right.as_ref(), + right_schema.as_ref(), + "SymmetricHashJoinExec", + "on.right", + )?; + Ok((l, r)) + }) + .collect::>()?; + + // Inline exhaustive proto -> `datafusion_common` enum conversions (the + // `from_proto` helpers live in the unreachable datafusion-proto crate). + let join_type = + match protobuf::JoinType::try_from(sym_join.join_type).map_err(|_| { + internal_datafusion_err!( + "SymmetricHashJoinExec: unknown JoinType {}", + sym_join.join_type + ) + })? { + protobuf::JoinType::Inner => JoinType::Inner, + protobuf::JoinType::Left => JoinType::Left, + protobuf::JoinType::Right => JoinType::Right, + protobuf::JoinType::Full => JoinType::Full, + protobuf::JoinType::Leftsemi => JoinType::LeftSemi, + protobuf::JoinType::Rightsemi => JoinType::RightSemi, + protobuf::JoinType::Leftanti => JoinType::LeftAnti, + protobuf::JoinType::Rightanti => JoinType::RightAnti, + protobuf::JoinType::Leftmark => JoinType::LeftMark, + protobuf::JoinType::Rightmark => JoinType::RightMark, + }; + let null_equality = match protobuf::NullEquality::try_from(sym_join.null_equality) + .map_err(|_| { + internal_datafusion_err!( + "SymmetricHashJoinExec: unknown NullEquality {}", + sym_join.null_equality + ) + })? { + protobuf::NullEquality::NullEqualsNothing => NullEquality::NullEqualsNothing, + protobuf::NullEquality::NullEqualsNull => NullEquality::NullEqualsNull, + }; + let partition_mode = + match protobuf::StreamPartitionMode::try_from(sym_join.partition_mode) + .map_err(|_| { + internal_datafusion_err!( + "SymmetricHashJoinExec: unknown StreamPartitionMode {}", + sym_join.partition_mode + ) + })? { + protobuf::StreamPartitionMode::SinglePartition => { + StreamJoinPartitionMode::SinglePartition + } + protobuf::StreamPartitionMode::PartitionedExec => { + StreamJoinPartitionMode::Partitioned + } + }; + + let filter = sym_join + .filter + .as_ref() + .map(|f| -> Result { + // The `JoinFilter` carries its own intermediate schema; the filter + // expression and column indices are relative to it. The proto + // schema type is datafusion-proto-common's, so `try_into` resolves + // to that crate's `TryFrom<&Schema>` impl. + let schema: Schema = f + .schema + .as_ref() + .ok_or_else(|| { + internal_datafusion_err!( + "SymmetricHashJoinExec: JoinFilter missing schema" + ) + })? + .try_into()?; + let expression = ctx.decode_required_expr( + f.expression.as_ref(), + &schema, + "SymmetricHashJoinExec", + "filter.expression", + )?; + let column_indices = f + .column_indices + .iter() + .map(|i| { + let side = + match protobuf::JoinSide::try_from(i.side).map_err(|_| { + internal_datafusion_err!( + "SymmetricHashJoinExec: unknown JoinSide {}", + i.side + ) + })? { + protobuf::JoinSide::LeftSide => JoinSide::Left, + protobuf::JoinSide::RightSide => JoinSide::Right, + protobuf::JoinSide::None => JoinSide::None, + }; + Ok(ColumnIndex { + index: i.index as usize, + side, + }) + }) + .collect::>>()?; + Ok(JoinFilter::new( + expression, + column_indices, + Arc::new(schema), + )) + }) + .transpose()?; + + // Each side's ordering decodes against the matching child schema. An + // empty `repeated` maps back to `None` (via `LexOrdering::new`), the + // inverse of the encoder's empty-vec-for-`None`. + let decode_sort_exprs = |protos: &[protobuf::PhysicalSortExprNode], + schema: &Schema, + field: &str| + -> Result> { + let exprs = protos + .iter() + .map(|sort_expr| { + let expr_node = sort_expr.expr.as_deref().ok_or_else(|| { + internal_datafusion_err!( + "SymmetricHashJoinExec: {field} missing its inner expr" + ) + })?; + Ok(PhysicalSortExpr { + expr: ctx.decode_expr(expr_node, schema)?, + options: arrow::compute::SortOptions { + descending: !sort_expr.asc, + nulls_first: sort_expr.nulls_first, + }, + }) + }) + .collect::>>()?; + Ok(LexOrdering::new(exprs)) + }; + let left_sort_exprs = decode_sort_exprs( + &sym_join.left_sort_exprs, + left_schema.as_ref(), + "left_sort_exprs", + )?; + let right_sort_exprs = decode_sort_exprs( + &sym_join.right_sort_exprs, + right_schema.as_ref(), + "right_sort_exprs", + )?; + + SymmetricHashJoinExec::try_new( + left, + right, + on, + filter, + &join_type, + null_equality, + left_sort_exprs, + right_sort_exprs, + partition_mode, + ) + .map(|e| Arc::new(e) as _) + } } /// A stream that issues [RecordBatch]es as they arrive from the right of the join. diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index 6cc6e44c32cc3..0658b4df9fcb7 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -84,6 +84,8 @@ pub mod metrics; pub mod operator_statistics; pub mod placeholder_row; pub mod projection; +#[cfg(feature = "proto")] +pub mod proto; pub mod recursive_query; pub mod repartition; pub mod scalar_subquery; diff --git a/datafusion/physical-plan/src/limit.rs b/datafusion/physical-plan/src/limit.rs index 1e4b5e5bb6426..a3aea479d5032 100644 --- a/datafusion/physical-plan/src/limit.rs +++ b/datafusion/physical-plan/src/limit.rs @@ -234,6 +234,67 @@ impl ExecutionPlan for GlobalLimitExec { fn supports_limit_pushdown(&self) -> bool { true } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::GlobalLimit(Box::new( + protobuf::GlobalLimitExecNode { + input: Some(Box::new(input)), + skip: self.skip() as u32, + fetch: match self.fetch() { + Some(n) => n as i64, + _ => -1, // no limit + }, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl GlobalLimitExec { + /// Reconstruct a [`GlobalLimitExec`] from its protobuf representation. + /// + /// The inverse of [`ExecutionPlan::try_to_proto`]: it takes the whole + /// [`PhysicalPlanNode`], decodes the single child recursively, and restores + /// the skip/fetch (a negative wire `fetch` decodes back to `None`). + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let limit = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::GlobalLimit, + "GlobalLimitExec", + ); + let input = ctx.decode_required_child( + limit.input.as_deref(), + "GlobalLimitExec", + "input", + )?; + let fetch = if limit.fetch >= 0 { + Some(limit.fetch as usize) + } else { + None + }; + Ok(Arc::new(GlobalLimitExec::new( + input, + limit.skip as usize, + fetch, + ))) + } } /// LocalLimitExec applies a limit to a single partition @@ -403,6 +464,50 @@ impl ExecutionPlan for LocalLimitExec { fn cardinality_effect(&self) -> CardinalityEffect { CardinalityEffect::LowerEqual } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::LocalLimit(Box::new( + protobuf::LocalLimitExecNode { + input: Some(Box::new(input)), + fetch: self.fetch() as u32, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl LocalLimitExec { + /// Reconstruct a [`LocalLimitExec`] from its protobuf representation. + /// + /// The inverse of [`ExecutionPlan::try_to_proto`]: it takes the whole + /// [`PhysicalPlanNode`] and decodes the single child recursively. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let limit = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::LocalLimit, + "LocalLimitExec", + ); + let input = + ctx.decode_required_child(limit.input.as_deref(), "LocalLimitExec", "input")?; + Ok(Arc::new(LocalLimitExec::new(input, limit.fetch as usize))) + } } /// A Limit stream skips `skip` rows, and then fetch up to `fetch` rows. diff --git a/datafusion/physical-plan/src/placeholder_row.rs b/datafusion/physical-plan/src/placeholder_row.rs index 64b192d58d238..37be2feb40b5e 100644 --- a/datafusion/physical-plan/src/placeholder_row.rs +++ b/datafusion/physical-plan/src/placeholder_row.rs @@ -182,6 +182,56 @@ impl ExecutionPlan for PlaceholderRowExec { None, ))) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + // PlaceholderRowExec is a leaf plan carrying only its schema on the wire. + let schema = self.schema().as_ref().try_into()?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::PlaceholderRow( + protobuf::PlaceholderRowExecNode { + schema: Some(schema), + }, + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl PlaceholderRowExec { + /// Reconstruct a [`PlaceholderRowExec`] from its protobuf representation. + /// + /// The inverse of [`ExecutionPlan::try_to_proto`]: it takes the whole + /// [`PhysicalPlanNode`] and reads back the schema (this leaf plan has no + /// children or expressions to decode; partition count is not serialized and + /// defaults to 1, matching the prior behavior). + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + _ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let placeholder = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::PlaceholderRow, + "PlaceholderRowExec", + ); + let schema = placeholder.schema.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "PlaceholderRowExec is missing required field 'schema'" + ) + })?; + let schema = Arc::new(Schema::try_from(schema)?); + Ok(Arc::new(PlaceholderRowExec::new(schema))) + } } #[cfg(test)] diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index 16b0a5ad7e4b5..39444132565ec 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -494,6 +494,71 @@ impl ExecutionPlan for ProjectionExec { .ok() }) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + let expr = ctx.encode_expressions(self.expr().iter().map(|p| &p.expr))?; + let expr_name = self.expr().iter().map(|p| p.alias.clone()).collect(); + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Projection(Box::new( + protobuf::ProjectionExecNode { + input: Some(Box::new(input)), + expr, + expr_name, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl ProjectionExec { + /// Reconstruct a [`ProjectionExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]: it takes the whole + /// [`PhysicalPlanNode`] so every plan's `try_from_proto` shares one + /// signature. Child plans and expressions are decoded recursively via the + /// [`ExecutionPlanDecodeCtx`]. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + /// [`ExecutionPlanDecodeCtx`]: crate::proto::ExecutionPlanDecodeCtx + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let projection = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Projection, + "ProjectionExec", + ); + let input = ctx.decode_required_child( + projection.input.as_deref(), + "ProjectionExec", + "input", + )?; + let input_schema = input.schema(); + let exprs = projection + .expr + .iter() + .zip(projection.expr_name.iter()) + .map(|(expr, name)| { + Ok(ProjectionExpr { + expr: ctx.decode_expr(expr, input_schema.as_ref())?, + alias: name.to_string(), + }) + }) + .collect::>>()?; + Ok(Arc::new(ProjectionExec::try_new(exprs, input)?)) + } } impl ProjectionStream { diff --git a/datafusion/physical-plan/src/proto.rs b/datafusion/physical-plan/src/proto.rs new file mode 100644 index 0000000000000..f3b3c14f184a8 --- /dev/null +++ b/datafusion/physical-plan/src/proto.rs @@ -0,0 +1,339 @@ +// 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. + +//! Serialization hooks for [`ExecutionPlan`], mirroring the +//! `try_to_proto`/`try_from_proto` pattern used for `PhysicalExpr`. +//! +//! # Why the indirection +//! +//! An `ExecutionPlan` must be able to (de)serialize its child plans and its +//! child physical expressions recursively. The concrete recursion lives in +//! `datafusion-proto` (it owns the extension codec, the session context and the +//! central converter), but `datafusion-proto` sits *above* `datafusion-physical-plan` +//! in the crate graph. To let a plan drive that recursion without a dependency +//! cycle, this module defines: +//! +//! * [`ExecutionPlanEncodeCtx`] / [`ExecutionPlanDecodeCtx`] — the stable, +//! concrete context types a plan author interacts with. New capabilities can +//! be added here without changing every plan's hook signature. +//! * [`ExecutionPlanEncode`] / [`ExecutionPlanDecode`] — internal dispatch +//! traits, *defined* here but *implemented* in `datafusion-proto`, that the +//! context types delegate to. This is the dependency inversion that keeps the +//! proto types flowing in one direction only. +//! +//! `datafusion-physical-plan` depends on the pure prost types in +//! `datafusion-proto-models` (feature `proto`), never on `datafusion-proto`. +//! +//! # Function-carrying plans +//! +//! Plans that reference UD(A/W)Fs (`AggregateExec`, the window execs, …) also +//! ride the hook: the context exposes typed, *bytes-only* function serde — +//! [`encode_udaf`](ExecutionPlanEncodeCtx::encode_udaf) / +//! [`decode_udaf`](ExecutionPlanDecodeCtx::decode_udaf) and the udf/udwf +//! siblings. These take/return `datafusion-expr` types plus `Vec` and never +//! name a proto type, so the `PhysicalExtensionCodec` (which only +//! `datafusion-proto` can name) stays fully encapsulated behind the adapter that +//! backs these traits. The lookup-order policy (payload → codec; else registry → +//! codec fallback) lives once, in that adapter, rather than in every plan. +//! +//! This is possible because `datafusion-physical-plan` sits *above* +//! `datafusion-expr` in the crate graph; the expression-side ctx (in +//! `physical-expr-common`, *below* `datafusion-expr`) cannot do this, which is +//! why `ScalarFunctionExpr` remains special-cased there. +//! +//! [`ExecutionPlan`]: crate::ExecutionPlan + +use std::sync::Arc; + +use arrow::datatypes::Schema; +use datafusion_common::{Result, internal_datafusion_err}; +use datafusion_execution::TaskContext; +use datafusion_expr::execution_props::ScalarSubqueryResults; +use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF}; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_proto_models::protobuf::{PhysicalExprNode, PhysicalPlanNode}; + +use crate::ExecutionPlan; + +/// Internal dispatch trait backing [`ExecutionPlanEncodeCtx`]. +/// +/// Implemented by `datafusion-proto`. Plan authors never name this trait; they +/// call methods on [`ExecutionPlanEncodeCtx`] instead. +pub trait ExecutionPlanEncode { + /// Serialize a child execution plan (recursing through the central + /// serializer, so the child's own `try_to_proto` hook is honored). + fn encode_plan(&self, plan: &Arc) -> Result; + + /// Serialize a physical expression owned by the plan. + fn encode_expr(&self, expr: &Arc) -> Result; + + /// Serialize a scalar UDF to an opaque payload. `None` means "decodable by + /// name alone" (built-ins). Bytes-only: no proto types cross this boundary. + fn encode_udf(&self, udf: &ScalarUDF) -> Result>>; + + /// Serialize an aggregate UDF to an opaque payload. `None` means "decodable + /// by name alone". + fn encode_udaf(&self, udaf: &AggregateUDF) -> Result>>; + + /// Serialize a window UDF to an opaque payload. `None` means "decodable by + /// name alone". + fn encode_udwf(&self, udwf: &WindowUDF) -> Result>>; +} + +/// Internal dispatch trait backing [`ExecutionPlanDecodeCtx`]. +/// +/// Implemented by `datafusion-proto`. Plan authors never name this trait; they +/// call methods on [`ExecutionPlanDecodeCtx`] instead. +pub trait ExecutionPlanDecode { + /// Deserialize a child execution plan (recursing through the central + /// deserializer, so the child's own `try_from_proto` is honored). + fn decode_plan(&self, node: &PhysicalPlanNode) -> Result>; + + /// Deserialize a physical expression against `input_schema`. + fn decode_expr( + &self, + node: &PhysicalExprNode, + input_schema: &Schema, + ) -> Result>; + + /// Deserialize a child plan within a scope where `results` is the active + /// scalar-subquery-results container, so `ScalarSubqueryExpr` nodes anywhere + /// in the child subtree resolve against it. Used only by `ScalarSubqueryExec`. + fn decode_plan_with_scalar_subquery_results( + &self, + node: &PhysicalPlanNode, + results: ScalarSubqueryResults, + ) -> Result>; + + /// The session task context, used by plans that need the function registry + /// or session configuration. Never exposes the proto extension codec. + fn task_ctx(&self) -> &TaskContext; + + /// Reconstruct a scalar UDF from its name and optional payload. Encapsulates + /// the lookup-order policy (payload → codec; else registry → codec fallback) + /// so no plan re-derives it. Bytes-only: no proto types cross this boundary. + fn decode_udf(&self, name: &str, payload: Option<&[u8]>) -> Result>; + + /// Reconstruct an aggregate UDF from its name and optional payload. + fn decode_udaf( + &self, + name: &str, + payload: Option<&[u8]>, + ) -> Result>; + + /// Reconstruct a window UDF from its name and optional payload. + fn decode_udwf(&self, name: &str, payload: Option<&[u8]>) -> Result>; +} + +/// Context handed to [`ExecutionPlan::try_to_proto`](crate::ExecutionPlan::try_to_proto). +/// +/// Provides the primitives a plan needs to serialize its children and +/// expressions without naming `datafusion-proto`. +pub struct ExecutionPlanEncodeCtx<'a> { + encoder: &'a dyn ExecutionPlanEncode, +} + +impl<'a> ExecutionPlanEncodeCtx<'a> { + /// Create a new encode context wrapping an [`ExecutionPlanEncode`] + /// implementation (supplied by `datafusion-proto`). + pub fn new(encoder: &'a dyn ExecutionPlanEncode) -> Self { + Self { encoder } + } + + /// Serialize a single child plan. + pub fn encode_child( + &self, + plan: &Arc, + ) -> Result { + self.encoder.encode_plan(plan) + } + + /// Serialize an iterator of child plans. + pub fn encode_children<'b, I>(&self, plans: I) -> Result> + where + I: IntoIterator>, + { + plans.into_iter().map(|p| self.encode_child(p)).collect() + } + + /// Serialize a single physical expression. + pub fn encode_expr(&self, expr: &Arc) -> Result { + self.encoder.encode_expr(expr) + } + + /// Serialize an iterator of physical expressions. + pub fn encode_expressions<'b, I>(&self, exprs: I) -> Result> + where + I: IntoIterator>, + { + exprs.into_iter().map(|e| self.encode_expr(e)).collect() + } + + /// Serialize a scalar UDF to an opaque payload (`None` = built-in, decodable + /// by name). No proto types cross this boundary. + pub fn encode_udf(&self, udf: &ScalarUDF) -> Result>> { + self.encoder.encode_udf(udf) + } + + /// Serialize an aggregate UDF to an opaque payload (`None` = decodable by + /// name). + pub fn encode_udaf(&self, udaf: &AggregateUDF) -> Result>> { + self.encoder.encode_udaf(udaf) + } + + /// Serialize a window UDF to an opaque payload (`None` = decodable by name). + pub fn encode_udwf(&self, udwf: &WindowUDF) -> Result>> { + self.encoder.encode_udwf(udwf) + } +} + +/// Context handed to a plan's `try_from_proto` associated function. +/// +/// Provides the primitives a plan needs to deserialize its children and +/// expressions without naming `datafusion-proto`. +pub struct ExecutionPlanDecodeCtx<'a> { + decoder: &'a dyn ExecutionPlanDecode, +} + +impl<'a> ExecutionPlanDecodeCtx<'a> { + /// Create a new decode context wrapping an [`ExecutionPlanDecode`] + /// implementation (supplied by `datafusion-proto`). + pub fn new(decoder: &'a dyn ExecutionPlanDecode) -> Self { + Self { decoder } + } + + /// Deserialize a single child plan. + pub fn decode_child( + &self, + node: &PhysicalPlanNode, + ) -> Result> { + self.decoder.decode_plan(node) + } + + /// Deserialize a slice of child plans (the repeated-`inputs` shape used by + /// `UnionExec`, `InterleaveExec`, …). Symmetric with + /// [`ExecutionPlanEncodeCtx::encode_children`]. + pub fn decode_children<'b, I>(&self, nodes: I) -> Result>> + where + I: IntoIterator, + { + nodes.into_iter().map(|n| self.decode_child(n)).collect() + } + + /// Deserialize a child plan within a scope carrying `results` as the active + /// scalar-subquery-results container (used only by `ScalarSubqueryExec`, so + /// that `ScalarSubqueryExpr` nodes in the child subtree resolve against it). + pub fn decode_child_with_scalar_subquery_results( + &self, + node: &PhysicalPlanNode, + results: ScalarSubqueryResults, + ) -> Result> { + self.decoder + .decode_plan_with_scalar_subquery_results(node, results) + } + + /// Deserialize a required child plan, producing a uniform "missing required + /// field" error when the optional wire field is absent. + pub fn decode_required_child( + &self, + node: Option<&PhysicalPlanNode>, + plan_name: &str, + field: &str, + ) -> Result> { + let node = node.ok_or_else(|| { + internal_datafusion_err!("{plan_name} is missing required field '{field}'") + })?; + self.decode_child(node) + } + + /// Deserialize a physical expression against `input_schema`. + pub fn decode_expr( + &self, + node: &PhysicalExprNode, + input_schema: &Schema, + ) -> Result> { + self.decoder.decode_expr(node, input_schema) + } + + /// Deserialize a required physical expression against `input_schema`. + pub fn decode_required_expr( + &self, + node: Option<&PhysicalExprNode>, + input_schema: &Schema, + plan_name: &str, + field: &str, + ) -> Result> { + let node = node.ok_or_else(|| { + internal_datafusion_err!("{plan_name} is missing required field '{field}'") + })?; + self.decode_expr(node, input_schema) + } + + /// The session task context (function registry + session config). Never + /// exposes the proto extension codec. + pub fn task_ctx(&self) -> &TaskContext { + self.decoder.task_ctx() + } + + /// Reconstruct a scalar UDF from its name and optional payload. The + /// lookup-order policy is owned by `datafusion-proto`; no proto types cross + /// this boundary. + pub fn decode_udf( + &self, + name: &str, + payload: Option<&[u8]>, + ) -> Result> { + self.decoder.decode_udf(name, payload) + } + + /// Reconstruct an aggregate UDF from its name and optional payload. + pub fn decode_udaf( + &self, + name: &str, + payload: Option<&[u8]>, + ) -> Result> { + self.decoder.decode_udaf(name, payload) + } + + /// Reconstruct a window UDF from its name and optional payload. + pub fn decode_udwf( + &self, + name: &str, + payload: Option<&[u8]>, + ) -> Result> { + self.decoder.decode_udwf(name, payload) + } +} + +/// Assert that a [`PhysicalPlanNode`] carries the expected `PhysicalPlanType` +/// variant, returning a reference to the inner payload, else an `internal_err!`. +/// Mirrors `expect_expr_variant!` on the expression side. Field access on the +/// result auto-derefs through the `Box` that boxed variants use. +#[macro_export] +macro_rules! expect_plan_variant { + ($node:expr, $variant:path, $plan_name:literal $(,)?) => {{ + match &$node.physical_plan_type { + Some($variant(inner)) => inner, + _ => { + return ::datafusion_common::internal_err!(concat!( + "PhysicalPlanNode is not a ", + $plan_name + )); + } + } + }}; +} diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 1617af3a68baa..72ed82731af22 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -1539,6 +1539,195 @@ impl ExecutionPlan for RepartitionExec { cache: new_properties.into(), }))) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let input = ctx.encode_child(self.input())?; + + // Inlined equivalent of datafusion-proto's `serialize_partitioning`. + // Only child physical expressions and `ScalarValue`s need serializing, + // both of which are reachable from `datafusion-physical-plan`, so the + // `protobuf::Partitioning` wrapping is built directly here. The proto + // wire format is 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) + } + }; + + 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), + }), + preserve_order: self.preserve_order(), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl RepartitionExec { + /// Reconstruct a [`RepartitionExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]. Inlines the + /// equivalent of datafusion-proto's `parse_protobuf_partitioning`; the + /// `protobuf::Partitioning` wrapping is read directly here since only child + /// expressions and `ScalarValue`s need decoding. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let repart = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Repartition, + "RepartitionExec", + ); + let input = ctx.decode_required_child( + repart.input.as_deref(), + "RepartitionExec", + "input", + )?; + let input_schema = input.schema(); + + let partition_method = repart + .partitioning + .as_ref() + .and_then(|p| p.partition_method.as_ref()) + .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::>>()?; + Partitioning::Hash(exprs, hash.partition_count as usize) + } + protobuf::partitioning::PartitionMethod::Unknown(n) => { + Partitioning::UnknownPartitioning(*n as usize) + } + protobuf::partitioning::PartitionMethod::Range(range) => { + use arrow::compute::SortOptions; + use datafusion_physical_expr::{RangePartitioning, SplitPoint}; + + 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| { + datafusion_common::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(); + } + Ok(Arc::new(repart_exec)) + } } impl RepartitionExec { diff --git a/datafusion/physical-plan/src/scalar_subquery.rs b/datafusion/physical-plan/src/scalar_subquery.rs index dd44d09c386c5..62fde8e78ca0b 100644 --- a/datafusion/physical-plan/src/scalar_subquery.rs +++ b/datafusion/physical-plan/src/scalar_subquery.rs @@ -243,6 +243,68 @@ impl ExecutionPlan for ScalarSubqueryExec { fn cardinality_effect(&self) -> CardinalityEffect { CardinalityEffect::Equal } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + // Index is positional (recovered on decode via enumerate), not on the wire. + let subqueries = + ctx.encode_children(self.subqueries().iter().map(|sq| &sq.plan))?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::ScalarSubquery(Box::new( + protobuf::ScalarSubqueryExecNode { + input: Some(Box::new(input)), + subqueries, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl ScalarSubqueryExec { + /// Reconstruct a [`ScalarSubqueryExec`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let sq = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::ScalarSubquery, + "ScalarSubqueryExec", + ); + // The subquery-results container must be active while decoding the input, + // so ScalarSubqueryExpr nodes in the input subtree resolve against it. + let results = ScalarSubqueryResults::new(sq.subqueries.len()); + let input_node = sq.input.as_deref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ScalarSubqueryExec is missing required field 'input'" + ) + })?; + let input = + ctx.decode_child_with_scalar_subquery_results(input_node, results.clone())?; + let subqueries = sq + .subqueries + .iter() + .enumerate() + .map(|(index, plan_node)| { + Ok(ScalarSubqueryLink { + plan: ctx.decode_child(plan_node)?, + index: SubqueryIndex::new(index), + }) + }) + .collect::>>()?; + Ok(Arc::new(ScalarSubqueryExec::new( + input, subqueries, results, + ))) + } } /// Wait for the subquery execution future to complete. diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 792c432155a8b..87d47ef41071e 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -1542,6 +1542,134 @@ impl ExecutionPlan for SortExec { updated_node: Some(new_sort), }) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + // A PhysicalSortExpr is a PhysicalExpr + SortOptions. The expr is encoded + // through the ctx; the SortOptions wrapping into a PhysicalSortExprNode + // (asc/nulls_first) is plain data and inlined here. + let expr = self + .expr() + .iter() + .map(|sort_expr| { + let sort_node = Box::new(protobuf::PhysicalSortExprNode { + expr: Some(Box::new(ctx.encode_expr(&sort_expr.expr)?)), + asc: !sort_expr.options.descending, + nulls_first: sort_expr.options.nulls_first, + }); + Ok(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::Sort( + sort_node, + )), + }) + }) + .collect::>>()?; + let dynamic_filter = match self.dynamic_filter_expr() { + Some(df) => { + let df_expr: Arc = df; + Some(ctx.encode_expr(&df_expr)?) + } + None => None, + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Sort(Box::new( + protobuf::SortExecNode { + input: Some(Box::new(input)), + expr, + fetch: match self.fetch() { + Some(n) => n as i64, + None => -1, + }, + preserve_partitioning: self.preserve_partitioning(), + dynamic_filter, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl SortExec { + /// Reconstruct a [`SortExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]: it takes the whole + /// [`PhysicalPlanNode`] and decodes the child plan and sort expressions + /// recursively through the [`ExecutionPlanDecodeCtx`]. The `SortOptions` + /// (asc/nulls_first) are read directly off the plain `PhysicalSortExprNode`. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + /// [`ExecutionPlanDecodeCtx`]: crate::proto::ExecutionPlanDecodeCtx + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + use protobuf::physical_expr_node::ExprType; + let sort = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Sort, + "SortExec", + ); + let input = + ctx.decode_required_child(sort.input.as_deref(), "SortExec", "input")?; + let input_schema = input.schema(); + let exprs = sort + .expr + .iter() + .map(|expr| { + let Some(ExprType::Sort(sort_expr)) = expr.expr_type.as_ref() else { + return datafusion_common::internal_err!( + "SortExec expr must be a sort expression" + ); + }; + let expr_node = sort_expr.expr.as_deref().ok_or_else(|| { + internal_datafusion_err!( + "SortExec sort expression is missing its inner expr" + ) + })?; + Ok(PhysicalSortExpr { + expr: ctx.decode_expr(expr_node, input_schema.as_ref())?, + options: arrow::compute::SortOptions { + descending: !sort_expr.asc, + nulls_first: sort_expr.nulls_first, + }, + }) + }) + .collect::>>()?; + let Some(ordering) = LexOrdering::new(exprs) else { + return datafusion_common::internal_err!("SortExec requires an ordering"); + }; + let fetch = (sort.fetch >= 0).then_some(sort.fetch as usize); + let new_sort = SortExec::new(ordering, input) + .with_fetch(fetch) + .with_preserve_partitioning(sort.preserve_partitioning); + + let new_sort = if let Some(df_proto) = &sort.dynamic_filter { + let df_expr = + ctx.decode_expr(df_proto, new_sort.input().schema().as_ref())?; + let df = (df_expr as Arc) + .downcast::() + .map_err(|_| { + internal_datafusion_err!( + "SortExec dynamic_filter did not decode to a DynamicFilterPhysicalExpr" + ) + })?; + new_sort.with_dynamic_filter_expr(df)? + } else { + new_sort + }; + + Ok(Arc::new(new_sort)) + } } #[cfg(test)] diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index 77a7d8f8f2e11..0d6aeaeac3ec6 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -414,6 +414,100 @@ impl ExecutionPlan for SortPreservingMergeExec { .with_fetch(self.fetch()), ))) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + let expr = self + .expr() + .iter() + .map(|e| { + Ok(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::Sort( + Box::new(protobuf::PhysicalSortExprNode { + expr: Some(Box::new(ctx.encode_expr(&e.expr)?)), + asc: !e.options.descending, + nulls_first: e.options.nulls_first, + }), + )), + }) + }) + .collect::>>()?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::SortPreservingMerge( + Box::new(protobuf::SortPreservingMergeExecNode { + input: Some(Box::new(input)), + expr, + fetch: self.fetch().map(|f| f as i64).unwrap_or(-1), + }), + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl SortPreservingMergeExec { + /// Reconstruct a [`SortPreservingMergeExec`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use arrow::compute::SortOptions; + use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; + use datafusion_proto_models::protobuf; + let spm = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::SortPreservingMerge, + "SortPreservingMergeExec", + ); + let input = ctx.decode_required_child( + spm.input.as_deref(), + "SortPreservingMergeExec", + "input", + )?; + let input_schema = input.schema(); + let exprs = spm + .expr + .iter() + .map(|e| { + let sort = match &e.expr_type { + Some(protobuf::physical_expr_node::ExprType::Sort(s)) => s, + _ => { + return internal_err!( + "SortPreservingMergeExec expression is not a sort expression" + ); + } + }; + let expr = ctx.decode_required_expr( + sort.expr.as_deref(), + input_schema.as_ref(), + "SortPreservingMergeExec", + "sort expression", + )?; + Ok(PhysicalSortExpr { + expr, + options: SortOptions { + descending: !sort.asc, + nulls_first: sort.nulls_first, + }, + }) + }) + .collect::>>()?; + let Some(ordering) = LexOrdering::new(exprs) else { + return internal_err!("SortPreservingMergeExec requires an ordering"); + }; + let fetch = (spm.fetch >= 0).then_some(spm.fetch as usize); + Ok(Arc::new( + SortPreservingMergeExec::new(ordering, input).with_fetch(fetch), + )) + } } #[cfg(test)] diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index d6f664c0059bc..f1170d17c4b8d 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -450,6 +450,50 @@ impl ExecutionPlan for UnionExec { // on all children (either pushed down or via FilterExec) Ok(propagation) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let inputs = ctx.encode_children(self.inputs())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Union( + protobuf::UnionExecNode { inputs }, + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl UnionExec { + /// Reconstruct a [`UnionExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]: each child in + /// `node.inputs` is decoded recursively via the [`ExecutionPlanDecodeCtx`]. + /// + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + /// [`ExecutionPlanDecodeCtx`]: crate::proto::ExecutionPlanDecodeCtx + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let union = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Union, + "UnionExec", + ); + let inputs = union + .inputs + .iter() + .map(|input| ctx.decode_child(input)) + .collect::>>()?; + UnionExec::try_new(inputs) + } } /// Combines multiple input streams by interleaving them. @@ -653,6 +697,50 @@ impl ExecutionPlan for InterleaveExec { fn benefits_from_input_partitioning(&self) -> Vec { vec![false; self.children().len()] } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let inputs = ctx.encode_children(self.inputs())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Interleave( + protobuf::InterleaveExecNode { inputs }, + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl InterleaveExec { + /// Reconstruct an [`InterleaveExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]: each child in + /// `node.inputs` is decoded recursively via the [`ExecutionPlanDecodeCtx`]. + /// + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + /// [`ExecutionPlanDecodeCtx`]: crate::proto::ExecutionPlanDecodeCtx + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let interleave = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Interleave, + "InterleaveExec", + ); + let inputs = interleave + .inputs + .iter() + .map(|input| ctx.decode_child(input)) + .collect::>>()?; + Ok(Arc::new(InterleaveExec::try_new(inputs)?)) + } } /// If all the input partitions have the same Hash partition spec with the first_input_partition diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index c31d0dd23fa68..24008b2407058 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -277,6 +277,128 @@ impl ExecutionPlan for UnnestExec { fn metrics(&self) -> Option { Some(self.metrics.clone_inner()) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + // `self.schema()` is `arrow::Schema`; the proto field is + // datafusion-proto-common's `Schema`, so `try_into` resolves to that + // crate's `TryFrom<&Schema>` impl. + let schema = self.schema().as_ref().try_into()?; + let list_type_columns = self + .list_column_indices() + .iter() + .map(|c| protobuf::ListUnnest { + index_in_input_schema: c.index_in_input_schema as _, + depth: c.depth as _, + }) + .collect(); + let struct_type_columns = self + .struct_column_indices() + .iter() + .map(|c| *c as _) + .collect(); + // `UnnestOptions::from_proto` lives in datafusion-proto (unreachable + // here), so inline the same field-by-field conversion. The nested + // `Column`s ride datafusion-proto-common's `From<&Column>` impl. + let options = self.options(); + let options = protobuf::UnnestOptions { + preserve_nulls: options.preserve_nulls, + recursions: options + .recursions + .iter() + .map(|r| protobuf::RecursionUnnestOption { + input_column: Some((&r.input_column).into()), + output_column: Some((&r.output_column).into()), + depth: r.depth as u32, + }) + .collect(), + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Unnest(Box::new( + protobuf::UnnestExecNode { + input: Some(Box::new(input)), + schema: Some(schema), + list_type_columns, + struct_type_columns, + options: Some(options), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl UnnestExec { + /// Reconstruct an [`UnnestExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]. + /// + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let unnest = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Unnest, + "UnnestExec", + ); + let input = + ctx.decode_required_child(unnest.input.as_deref(), "UnnestExec", "input")?; + // The proto field type is datafusion-proto-common's `Schema`, so + // `try_into` resolves to that crate's `TryFrom<&Schema>` impl. + let schema: Schema = unnest + .schema + .as_ref() + .ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "UnnestExec is missing required field 'schema'" + ) + })? + .try_into()?; + let list_column_indices = unnest + .list_type_columns + .iter() + .map(|c| ListUnnest { + index_in_input_schema: c.index_in_input_schema as _, + depth: c.depth as _, + }) + .collect(); + let struct_column_indices = + unnest.struct_type_columns.iter().map(|c| *c as _).collect(); + let options = unnest.options.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "UnnestExec is missing required field 'options'" + ) + })?; + let options = UnnestOptions { + preserve_nulls: options.preserve_nulls, + recursions: options + .recursions + .iter() + .map(|r| datafusion_common::RecursionUnnestOption { + input_column: r.input_column.as_ref().unwrap().into(), + output_column: r.output_column.as_ref().unwrap().into(), + depth: r.depth as usize, + }) + .collect(), + }; + Ok(Arc::new(UnnestExec::new( + input, + list_column_indices, + struct_column_indices, + Arc::new(schema), + options, + )?)) + } } #[derive(Clone, Debug)] diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs index a9d580f4c687d..2322f5932293a 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -388,6 +388,54 @@ impl ExecutionPlan for BoundedWindowAggExec { fn cardinality_effect(&self) -> CardinalityEffect { CardinalityEffect::Equal } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use super::window_agg_exec::encode_physical_window_expr; + use datafusion_proto_common::protobuf_common::EmptyMessage; + use datafusion_proto_models::protobuf; + use protobuf::window_agg_exec_node::InputOrderMode as ProtoInputOrderMode; + + let input = ctx.encode_child(self.input())?; + let window_expr = self + .window_expr() + .iter() + .map(|e| encode_physical_window_expr(e, ctx)) + .collect::>>()?; + let partition_keys = self + .partition_keys() + .iter() + .map(|e| ctx.encode_expr(e)) + .collect::>>()?; + // A `Some(input_order_mode)` is what tells the shared `Window` decode + // arm to rebuild a `BoundedWindowAggExec` rather than a `WindowAggExec`. + let input_order_mode = match &self.input_order_mode { + InputOrderMode::Linear => ProtoInputOrderMode::Linear(EmptyMessage {}), + InputOrderMode::PartiallySorted(columns) => { + ProtoInputOrderMode::PartiallySorted( + protobuf::PartiallySortedInputOrderMode { + columns: columns.iter().map(|c| *c as u64).collect(), + }, + ) + } + InputOrderMode::Sorted => ProtoInputOrderMode::Sorted(EmptyMessage {}), + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Window(Box::new( + protobuf::WindowAggExecNode { + input: Some(Box::new(input)), + window_expr, + partition_keys, + input_order_mode: Some(input_order_mode), + }, + )), + ), + })) + } } /// Trait that specifies how we search for (or calculate) partitions. It has two diff --git a/datafusion/physical-plan/src/windows/window_agg_exec.rs b/datafusion/physical-plan/src/windows/window_agg_exec.rs index 72474c6a55483..42d9e70a8d6c1 100644 --- a/datafusion/physical-plan/src/windows/window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs @@ -304,6 +304,415 @@ impl ExecutionPlan for WindowAggExec { fn cardinality_effect(&self) -> CardinalityEffect { CardinalityEffect::Equal } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + let window_expr = self + .window_expr() + .iter() + .map(|e| encode_physical_window_expr(e, ctx)) + .collect::>>()?; + let partition_keys = self + .partition_keys() + .iter() + .map(|e| ctx.encode_expr(e)) + .collect::>>()?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Window(Box::new( + protobuf::WindowAggExecNode { + input: Some(Box::new(input)), + window_expr, + partition_keys, + // `None` distinguishes a `WindowAggExec` from a + // `BoundedWindowAggExec` on the shared `Window` variant. + input_order_mode: None, + }, + )), + ), + })) + } +} + +/// Reconstruct a window plan from its protobuf representation. +/// +/// Both [`WindowAggExec`] and [`BoundedWindowAggExec`] serialize to the single +/// `PhysicalPlanType::Window(WindowAggExecNode)` proto variant; they are told +/// apart by the optional `input_order_mode` field (`None` => `WindowAggExec`, +/// `Some` => `BoundedWindowAggExec`). This one associated function therefore +/// owns the decode for *both* plans and the central dispatch routes the `Window` +/// arm here. +#[cfg(feature = "proto")] +impl WindowAggExec { + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use super::BoundedWindowAggExec; + use crate::InputOrderMode; + use datafusion_proto_models::protobuf; + use protobuf::window_agg_exec_node::InputOrderMode as ProtoInputOrderMode; + + let window_agg = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Window, + "WindowAggExec", + ); + let input = ctx.decode_required_child( + window_agg.input.as_deref(), + "WindowAggExec", + "input", + )?; + let input_schema = input.schema(); + + let window_expr = window_agg + .window_expr + .iter() + .map(|w| decode_physical_window_expr(w, ctx, input_schema.as_ref())) + .collect::>>()?; + + let partition_keys = window_agg + .partition_keys + .iter() + .map(|e| ctx.decode_expr(e, input_schema.as_ref())) + .collect::>>()?; + + if let Some(input_order_mode) = window_agg.input_order_mode.as_ref() { + let input_order_mode = match input_order_mode { + ProtoInputOrderMode::Linear(_) => InputOrderMode::Linear, + ProtoInputOrderMode::PartiallySorted( + protobuf::PartiallySortedInputOrderMode { columns }, + ) => InputOrderMode::PartiallySorted( + columns.iter().map(|c| *c as usize).collect(), + ), + ProtoInputOrderMode::Sorted(_) => InputOrderMode::Sorted, + }; + Ok(Arc::new(BoundedWindowAggExec::try_new( + window_expr, + input, + input_order_mode, + !partition_keys.is_empty(), + )?)) + } else { + Ok(Arc::new(WindowAggExec::try_new( + window_expr, + input, + !partition_keys.is_empty(), + )?)) + } + } +} + +/// Serialize a single [`WindowExpr`] to its protobuf representation. +/// +/// Shared by [`WindowAggExec`] and [`BoundedWindowAggExec`]. UD(A)Fs used as +/// window functions are encoded to opaque bytes through the ctx +/// (`encode_udaf`/`encode_udwf`); the extension codec never leaks here. +#[cfg(feature = "proto")] +pub(crate) fn encode_physical_window_expr( + window_expr: &Arc, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, +) -> Result { + use super::{PlainAggregateWindowExpr, StandardWindowExpr, WindowUDFExpr}; + use datafusion_common::not_impl_err; + use datafusion_physical_expr::window::SlidingAggregateWindowExpr; + use datafusion_proto_models::protobuf::{self, physical_window_expr_node}; + + let expr = window_expr.as_any(); + let mut args = window_expr.expressions().to_vec(); + let window_frame = window_expr.get_window_frame(); + + let (window_function, fun_definition, ignore_nulls, distinct) = + if let Some(plain) = expr.downcast_ref::() { + let aggr_expr = plain.get_aggregate_expr(); + ( + physical_window_expr_node::WindowFunction::UserDefinedAggrFunction( + aggr_expr.fun().name().to_string(), + ), + ctx.encode_udaf(aggr_expr.fun())?, + aggr_expr.ignore_nulls(), + aggr_expr.is_distinct(), + ) + } else if let Some(sliding) = expr.downcast_ref::() { + let aggr_expr = sliding.get_aggregate_expr(); + ( + physical_window_expr_node::WindowFunction::UserDefinedAggrFunction( + aggr_expr.fun().name().to_string(), + ), + ctx.encode_udaf(aggr_expr.fun())?, + aggr_expr.ignore_nulls(), + aggr_expr.is_distinct(), + ) + } else if let Some(udf_window_expr) = expr.downcast_ref::() { + if let Some(e) = udf_window_expr + .get_standard_func_expr() + .as_any() + .downcast_ref::() + { + // `WindowUDFExpr::args` returns the full, unfiltered argument list so + // every argument survives the round-trip. + args = e.args().to_vec(); + ( + physical_window_expr_node::WindowFunction::UserDefinedWindowFunction( + e.fun().name().to_string(), + ), + ctx.encode_udwf(e.fun().as_ref())?, + // `WindowUDFExpr` has no ignore_nulls/distinct. + false, + false, + ) + } else { + return not_impl_err!( + "User-defined window function not supported: {window_expr:?}" + ); + } + } else { + return not_impl_err!("WindowExpr not supported: {window_expr:?}"); + }; + + let args = ctx.encode_expressions(&args)?; + let partition_by = ctx.encode_expressions(window_expr.partition_by())?; + let order_by = window_expr + .order_by() + .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 window_frame = encode_window_frame(window_frame.as_ref())?; + + Ok(protobuf::PhysicalWindowExprNode { + args, + partition_by, + order_by, + window_frame: Some(window_frame), + window_function: Some(window_function), + name: window_expr.name().to_string(), + fun_definition, + ignore_nulls, + distinct, + }) +} + +/// Reconstruct a single [`WindowExpr`] from its protobuf representation. +#[cfg(feature = "proto")] +fn decode_physical_window_expr( + proto: &datafusion_proto_models::protobuf::PhysicalWindowExprNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + input_schema: &arrow::datatypes::Schema, +) -> Result> { + use super::{create_window_expr, schema_add_window_field}; + use arrow::compute::SortOptions; + use datafusion_common::{internal_datafusion_err, internal_err}; + use datafusion_expr::WindowFunctionDefinition; + use datafusion_proto_models::protobuf::physical_window_expr_node; + + let args = proto + .args + .iter() + .map(|e| ctx.decode_expr(e, input_schema)) + .collect::>>()?; + let partition_by = proto + .partition_by + .iter() + .map(|e| ctx.decode_expr(e, input_schema)) + .collect::>>()?; + let order_by = proto + .order_by + .iter() + .map(|sort_expr| { + let expr = sort_expr.expr.as_ref().ok_or_else(|| { + internal_datafusion_err!( + "Missing expr in window order_by sort expression" + ) + })?; + Ok(PhysicalSortExpr { + expr: ctx.decode_expr(expr, input_schema)?, + options: SortOptions { + descending: !sort_expr.asc, + nulls_first: sort_expr.nulls_first, + }, + }) + }) + .collect::>>()?; + let window_frame = proto + .window_frame + .as_ref() + .map(decode_window_frame) + .transpose()? + .ok_or_else(|| { + internal_datafusion_err!("Missing required field 'window_frame' in protobuf") + })?; + + let fun = match proto.window_function.as_ref() { + Some(physical_window_expr_node::WindowFunction::UserDefinedAggrFunction( + name, + )) => WindowFunctionDefinition::AggregateUDF( + ctx.decode_udaf(name, proto.fun_definition.as_deref())?, + ), + Some(physical_window_expr_node::WindowFunction::UserDefinedWindowFunction( + name, + )) => WindowFunctionDefinition::WindowUDF( + ctx.decode_udwf(name, proto.fun_definition.as_deref())?, + ), + None => { + return internal_err!("Missing required field 'window_function' in protobuf"); + } + }; + + let name = proto.name.clone(); + // TODO: Remove extended_schema if functions are all UDAF + let extended_schema = schema_add_window_field(&args, input_schema, &fun, &name)?; + create_window_expr( + &fun, + name, + &args, + &partition_by, + &order_by, + Arc::new(window_frame), + extended_schema, + proto.ignore_nulls, + proto.distinct, + None, + ) +} + +/// Serialize a [`WindowFrame`](datafusion_expr::WindowFrame) inline, by-name. +/// +/// The `datafusion-proto` `TryFromProto` conversions live above this crate, so +/// the frame/units/bound enums are matched exhaustively by name here (never a +/// numeric cast) to stay wire-identical and to make a new variant a compile +/// error. +#[cfg(feature = "proto")] +fn encode_window_frame( + window_frame: &datafusion_expr::WindowFrame, +) -> Result { + use datafusion_expr::WindowFrameUnits; + use datafusion_proto_models::protobuf; + + let units = match window_frame.units { + WindowFrameUnits::Rows => protobuf::WindowFrameUnits::Rows, + WindowFrameUnits::Range => protobuf::WindowFrameUnits::Range, + WindowFrameUnits::Groups => protobuf::WindowFrameUnits::Groups, + }; + Ok(protobuf::WindowFrame { + window_frame_units: units.into(), + start_bound: Some(encode_window_frame_bound(&window_frame.start_bound)?), + end_bound: Some(protobuf::window_frame::EndBound::Bound( + encode_window_frame_bound(&window_frame.end_bound)?, + )), + }) +} + +#[cfg(feature = "proto")] +fn encode_window_frame_bound( + bound: &datafusion_expr::WindowFrameBound, +) -> Result { + use datafusion_expr::WindowFrameBound; + use datafusion_proto_common::protobuf_common; + use datafusion_proto_models::protobuf; + + // Bind the conversion target explicitly so the `?` error type is unambiguous + // and stays wire-identical to the `TryFrom<&ScalarValue>` proto encoding. + let encode_value = + |v: &datafusion_common::ScalarValue| -> Result { + Ok(v.try_into()?) + }; + Ok(match bound { + WindowFrameBound::CurrentRow => protobuf::WindowFrameBound { + window_frame_bound_type: protobuf::WindowFrameBoundType::CurrentRow.into(), + bound_value: None, + }, + WindowFrameBound::Preceding(v) => protobuf::WindowFrameBound { + window_frame_bound_type: protobuf::WindowFrameBoundType::Preceding.into(), + bound_value: Some(encode_value(v)?), + }, + WindowFrameBound::Following(v) => protobuf::WindowFrameBound { + window_frame_bound_type: protobuf::WindowFrameBoundType::Following.into(), + bound_value: Some(encode_value(v)?), + }, + }) +} + +#[cfg(feature = "proto")] +fn decode_window_frame( + window_frame: &datafusion_proto_models::protobuf::WindowFrame, +) -> Result { + use datafusion_common::internal_datafusion_err; + use datafusion_expr::{WindowFrame, WindowFrameBound, WindowFrameUnits}; + use datafusion_proto_models::protobuf; + + let units = + match protobuf::WindowFrameUnits::try_from(window_frame.window_frame_units) + .map_err(|_| { + internal_datafusion_err!( + "Received a WindowFrame message with unknown WindowFrameUnits {}", + window_frame.window_frame_units + ) + })? { + protobuf::WindowFrameUnits::Rows => WindowFrameUnits::Rows, + protobuf::WindowFrameUnits::Range => WindowFrameUnits::Range, + protobuf::WindowFrameUnits::Groups => WindowFrameUnits::Groups, + }; + let start_bound = + decode_window_frame_bound(window_frame.start_bound.as_ref().ok_or_else( + || internal_datafusion_err!("Missing start_bound in WindowFrame"), + )?)?; + let end_bound = window_frame + .end_bound + .as_ref() + .map(|end_bound| match end_bound { + protobuf::window_frame::EndBound::Bound(end_bound) => { + decode_window_frame_bound(end_bound) + } + }) + .transpose()? + .unwrap_or(WindowFrameBound::CurrentRow); + Ok(WindowFrame::new_bounds(units, start_bound, end_bound)) +} + +#[cfg(feature = "proto")] +fn decode_window_frame_bound( + bound: &datafusion_proto_models::protobuf::WindowFrameBound, +) -> Result { + use datafusion_common::{ScalarValue, internal_datafusion_err}; + use datafusion_expr::WindowFrameBound; + use datafusion_proto_common::protobuf_common; + use datafusion_proto_models::protobuf; + + // Bind the conversion source explicitly so the `?` error type is unambiguous. + let decode_value = |x: &protobuf_common::ScalarValue| -> Result { + Ok(ScalarValue::try_from(x)?) + }; + let bound_type = protobuf::WindowFrameBoundType::try_from( + bound.window_frame_bound_type, + ) + .map_err(|_| { + internal_datafusion_err!( + "Received a WindowFrameBound message with unknown WindowFrameBoundType {}", + bound.window_frame_bound_type + ) + })?; + match bound_type { + protobuf::WindowFrameBoundType::CurrentRow => Ok(WindowFrameBound::CurrentRow), + protobuf::WindowFrameBoundType::Preceding => match &bound.bound_value { + Some(x) => Ok(WindowFrameBound::Preceding(decode_value(x)?)), + None => Ok(WindowFrameBound::Preceding(ScalarValue::UInt64(None))), + }, + protobuf::WindowFrameBoundType::Following => match &bound.bound_value { + Some(x) => Ok(WindowFrameBound::Following(decode_value(x)?)), + None => Ok(WindowFrameBound::Following(ScalarValue::UInt64(None))), + }, + } } /// Compute the window aggregate columns diff --git a/datafusion/proto/Cargo.toml b/datafusion/proto/Cargo.toml index cfff8a949418a..3b4f47314accc 100644 --- a/datafusion/proto/Cargo.toml +++ b/datafusion/proto/Cargo.toml @@ -54,12 +54,12 @@ chrono = { workspace = true } datafusion-catalog = { workspace = true } datafusion-catalog-listing = { workspace = true } datafusion-common = { workspace = true } -datafusion-datasource = { workspace = true } -datafusion-datasource-arrow = { workspace = true } -datafusion-datasource-avro = { workspace = true, optional = true } -datafusion-datasource-csv = { workspace = true } -datafusion-datasource-json = { workspace = true } -datafusion-datasource-parquet = { workspace = true, optional = true } +datafusion-datasource = { workspace = true, features = ["proto"] } +datafusion-datasource-arrow = { workspace = true, features = ["proto"] } +datafusion-datasource-avro = { workspace = true, optional = true, features = ["proto"] } +datafusion-datasource-csv = { workspace = true, features = ["proto"] } +datafusion-datasource-json = { workspace = true, features = ["proto"] } +datafusion-datasource-parquet = { workspace = true, optional = true, features = ["proto"] } datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } datafusion-functions-table = { workspace = true } diff --git a/datafusion/proto/src/common.rs b/datafusion/proto/src/common.rs index bff017edbc998..16d138c9aafff 100644 --- a/datafusion/proto/src/common.rs +++ b/datafusion/proto/src/common.rs @@ -17,6 +17,9 @@ use datafusion_common::{Result, assert_eq_or_internal_err, internal_datafusion_err}; +// Dead once CsvSource moved to the try_to_proto hook (#22419); a follow-up will +// remove these CSV byte helpers along with the old CSV scan code. +#[allow(dead_code)] pub(crate) fn str_to_byte(s: &String, description: &str) -> Result { assert_eq_or_internal_err!( s.len(), @@ -26,6 +29,9 @@ pub(crate) fn str_to_byte(s: &String, description: &str) -> Result { Ok(s.as_bytes()[0]) } +// Dead once CsvSource moved to the try_to_proto hook (#22419); the old CSV +// decode arm that still uses it will be deleted in the same cleanup pass. +#[allow(dead_code)] pub(crate) fn byte_to_string(b: u8, description: &str) -> Result { let b = &[b]; let b = std::str::from_utf8(b).map_err(|_| { diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index 53ff4a41d466e..f1144a35eac67 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -32,10 +32,6 @@ use datafusion_datasource::file_groups::FileGroup; use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; use datafusion_datasource::file_sink_config::FileSinkConfig; use datafusion_datasource::{FileRange, ListingTableUrl, PartitionedFile, TableSchema}; -use datafusion_datasource_csv::file_format::CsvSink; -use datafusion_datasource_json::file_format::JsonSink; -#[cfg(feature = "parquet")] -use datafusion_datasource_parquet::file_format::ParquetSink; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{FunctionRegistry, TaskContext}; use datafusion_expr::WindowFunctionDefinition; @@ -63,7 +59,7 @@ use super::{ }; use crate::convert::TryFromProto; use crate::protobuf::physical_expr_node::ExprType; -use crate::{convert_required, convert_required_proto, protobuf}; +use crate::{convert_required, protobuf}; use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr; /// Parses a physical sort expression from a protobuf. @@ -670,41 +666,7 @@ impl TryFromProto<&protobuf::FileGroup> for FileGroup { Ok(FileGroup::new(files)) } } - -impl TryFromProto<&protobuf::JsonSink> for JsonSink { - type Error = DataFusionError; - - fn try_from_proto(value: &protobuf::JsonSink) -> Result { - Ok(Self::new( - convert_required_proto!(FileSinkConfig, value.config)?, - convert_required!(value.writer_options)?, - )) - } -} - #[cfg(feature = "parquet")] -impl TryFromProto<&protobuf::ParquetSink> for ParquetSink { - type Error = DataFusionError; - - fn try_from_proto(value: &protobuf::ParquetSink) -> Result { - Ok(Self::new( - convert_required_proto!(FileSinkConfig, value.config)?, - convert_required!(value.parquet_options)?, - )) - } -} - -impl TryFromProto<&protobuf::CsvSink> for CsvSink { - type Error = DataFusionError; - - fn try_from_proto(value: &protobuf::CsvSink) -> Result { - Ok(Self::new( - convert_required_proto!(FileSinkConfig, value.config)?, - convert_required!(value.writer_options)?, - )) - } -} - impl TryFromProto<&protobuf::FileSinkConfig> for FileSinkConfig { type Error = DataFusionError; diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 72f6e5af5bff2..fed740c333d1c 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -24,19 +24,10 @@ use std::sync::Arc; use arrow::compute::SortOptions; use arrow::datatypes::{IntervalMonthDayNanoType, Schema, SchemaRef}; use datafusion_catalog::memory::MemorySourceConfig; -use datafusion_common::config::CsvOptions; -use datafusion_common::display::StringifiedPlan; -use datafusion_common::format::ExplainFormat; use datafusion_common::{ DataFusionError, JoinType, NullEquality, Result, internal_datafusion_err, internal_err, not_impl_err, }; -#[cfg(feature = "parquet")] -use datafusion_datasource::file::FileSource; -use datafusion_datasource::file_compression_type::FileCompressionType; -use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; -use datafusion_datasource::sink::DataSinkExec; -use datafusion_datasource::source::{DataSource, DataSourceExec}; use datafusion_datasource_arrow::source::ArrowSource; #[cfg(feature = "avro")] use datafusion_datasource_avro::source::AvroSource; @@ -45,26 +36,16 @@ use datafusion_datasource_csv::source::CsvSource; use datafusion_datasource_json::file_format::JsonSink; use datafusion_datasource_json::source::JsonSource; #[cfg(feature = "parquet")] -use datafusion_datasource_parquet::CachedParquetFileReaderFactory; -#[cfg(feature = "parquet")] use datafusion_datasource_parquet::file_format::ParquetSink; #[cfg(feature = "parquet")] use datafusion_datasource_parquet::source::ParquetSource; -#[cfg(feature = "parquet")] -use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{FunctionRegistry, TaskContext}; -use datafusion_expr::execution_props::{ScalarSubqueryResults, SubqueryIndex}; +use datafusion_expr::execution_props::ScalarSubqueryResults; use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF}; use datafusion_functions_table::generate_series::{ Empty, GenSeriesArgs, GenerateSeriesTable, GenericSeriesState, TimestampValue, }; -use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr}; -use datafusion_physical_expr::async_scalar_function::AsyncFuncExpr; -use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr; -use datafusion_physical_expr::{LexOrdering, LexRequirement, PhysicalExprRef}; -use datafusion_physical_plan::aggregates::{ - AggregateExec, AggregateMode, LimitOptions, PhysicalGroupBy, -}; +use datafusion_physical_plan::aggregates::AggregateExec; use datafusion_physical_plan::analyze::AnalyzeExec; use datafusion_physical_plan::async_func::AsyncFuncExec; use datafusion_physical_plan::buffer::BufferExec; @@ -74,51 +55,37 @@ use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::coop::CooperativeExec; use datafusion_physical_plan::empty::EmptyExec; use datafusion_physical_plan::explain::ExplainExec; -use datafusion_physical_plan::expressions::PhysicalSortExpr; -use datafusion_physical_plan::filter::{FilterExec, FilterExecBuilder}; +use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::utils::{ColumnIndex, JoinFilter}; use datafusion_physical_plan::joins::{ - CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, - StreamJoinPartitionMode, SymmetricHashJoinExec, + CrossJoinExec, HashJoinExec, NestedLoopJoinExec, SortMergeJoinExec, + SymmetricHashJoinExec, }; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::memory::LazyMemoryExec; -use datafusion_physical_plan::metrics::MetricCategory; use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; -use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; +use datafusion_physical_plan::projection::ProjectionExec; +use datafusion_physical_plan::proto::{ + ExecutionPlanDecode, ExecutionPlanDecodeCtx, ExecutionPlanEncode, + ExecutionPlanEncodeCtx, +}; use datafusion_physical_plan::repartition::RepartitionExec; -use datafusion_physical_plan::scalar_subquery::{ScalarSubqueryExec, ScalarSubqueryLink}; +use datafusion_physical_plan::scalar_subquery::ScalarSubqueryExec; use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::union::{InterleaveExec, UnionExec}; -use datafusion_physical_plan::unnest::{ListUnnest, UnnestExec}; -use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowAggExec}; -use datafusion_physical_plan::{ExecutionPlan, InputOrderMode, PhysicalExpr, WindowExpr}; +use datafusion_physical_plan::unnest::UnnestExec; +use datafusion_physical_plan::windows::WindowAggExec; +use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr}; use prost::Message; use prost::bytes::BufMut; -use self::from_proto::parse_protobuf_partitioning; -use self::to_proto::serialize_partitioning; -use crate::common::{byte_to_string, str_to_byte}; -use crate::convert::{FromProto, TryFromProto}; +use crate::convert::FromProto; use crate::convert_required; -use crate::physical_plan::from_proto::{ - parse_physical_expr_with_converter, parse_physical_sort_expr, - parse_physical_sort_exprs, parse_physical_window_expr, - parse_protobuf_file_scan_config, parse_record_batches, parse_table_schema_from_proto, -}; -use crate::physical_plan::to_proto::{ - serialize_file_scan_config, serialize_maybe_filter, serialize_physical_aggr_expr, - serialize_physical_expr_with_converter, serialize_physical_sort_exprs, - serialize_physical_window_expr, serialize_record_batches, -}; -use crate::protobuf::physical_aggregate_expr_node::AggregateFunction; -use crate::protobuf::physical_expr_node::ExprType; +use crate::physical_plan::from_proto::parse_physical_expr_with_converter; +use crate::physical_plan::to_proto::serialize_physical_expr_with_converter; use crate::protobuf::physical_plan_node::PhysicalPlanType; -use crate::protobuf::{ - self, ListUnnest as ProtoListUnnest, SortExprNode, SortMergeJoinExecNode, - proto_error, window_agg_exec_node, -}; +use crate::protobuf::{self, SortMergeJoinExecNode, proto_error}; pub mod from_proto; pub mod to_proto; @@ -132,6 +99,10 @@ fn encode_human_display_alias(human_display: &str, alias: &str) -> String { ) } +// The production decoder now lives alongside `AggregateExec` in +// `datafusion-physical-plan` (`aggregates/mod.rs`); this copy is retained for the +// encode/decode round-trip unit test below. +#[allow(dead_code)] fn split_human_display_alias<'a>( human_display: &'a str, name: &'a str, @@ -312,126 +283,150 @@ pub trait PhysicalPlanNodeExt: Sized { self.node(), )) })?; + // Decode context for plans migrated to the `try_from_proto` pattern + // (#22419). Arms for migrated plans are one-liners delegating to the + // plan's own crate; un-migrated arms keep their inline bodies. + let plan_decoder = ConverterPlanDecoder { + ctx, + proto_converter, + }; + let decode_ctx = ExecutionPlanDecodeCtx::new(&plan_decoder); match plan { - PhysicalPlanType::Explain(explain) => { - self.try_into_explain_physical_plan(explain, ctx, proto_converter) + PhysicalPlanType::Explain(_) => { + ExplainExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Projection(projection) => { - self.try_into_projection_physical_plan(projection, ctx, proto_converter) + PhysicalPlanType::Projection(_) => { + ProjectionExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Filter(filter) => { - self.try_into_filter_physical_plan(filter, ctx, proto_converter) + PhysicalPlanType::Filter(_) => { + FilterExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::CsvScan(scan) => { - self.try_into_csv_scan_physical_plan(scan, ctx, proto_converter) + PhysicalPlanType::CsvScan(_) => { + CsvSource::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::JsonScan(scan) => { - self.try_into_json_scan_physical_plan(scan, ctx, proto_converter) + PhysicalPlanType::JsonScan(_) => { + JsonSource::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::ParquetScan(scan) => { - self.try_into_parquet_scan_physical_plan(scan, ctx, proto_converter) + PhysicalPlanType::ParquetScan(_) => { + #[cfg(feature = "parquet")] + { + ParquetSource::try_from_proto(self.node(), &decode_ctx) + } + #[cfg(not(feature = "parquet"))] + panic!( + "Unable to process a Parquet PhysicalPlan when `parquet` feature is not enabled" + ) } - PhysicalPlanType::AvroScan(scan) => { - self.try_into_avro_scan_physical_plan(scan, ctx, proto_converter) + PhysicalPlanType::AvroScan(_) => { + #[cfg(feature = "avro")] + { + AvroSource::try_from_proto(self.node(), &decode_ctx) + } + #[cfg(not(feature = "avro"))] + panic!( + "Unable to process a Avro PhysicalPlan when `avro` feature is not enabled" + ) } - PhysicalPlanType::MemoryScan(scan) => { - self.try_into_memory_scan_physical_plan(scan, ctx, proto_converter) + PhysicalPlanType::MemoryScan(_) => { + MemorySourceConfig::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::ArrowScan(scan) => { - self.try_into_arrow_scan_physical_plan(scan, ctx, proto_converter) + PhysicalPlanType::ArrowScan(_) => { + ArrowSource::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::CoalesceBatches(coalesce_batches) => self - .try_into_coalesce_batches_physical_plan( - coalesce_batches, - ctx, - proto_converter, - ), - PhysicalPlanType::Merge(merge) => { - self.try_into_merge_physical_plan(merge, ctx, proto_converter) + #[expect(deprecated)] + PhysicalPlanType::CoalesceBatches(_) => { + CoalesceBatchesExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Repartition(repart) => { - self.try_into_repartition_physical_plan(repart, ctx, proto_converter) + PhysicalPlanType::Merge(_) => { + CoalescePartitionsExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::GlobalLimit(limit) => { - self.try_into_global_limit_physical_plan(limit, ctx, proto_converter) + PhysicalPlanType::Repartition(_) => { + RepartitionExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::LocalLimit(limit) => { - self.try_into_local_limit_physical_plan(limit, ctx, proto_converter) + PhysicalPlanType::GlobalLimit(_) => { + GlobalLimitExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Window(window_agg) => { - self.try_into_window_physical_plan(window_agg, ctx, proto_converter) + PhysicalPlanType::LocalLimit(_) => { + LocalLimitExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Aggregate(hash_agg) => { - self.try_into_aggregate_physical_plan(hash_agg, ctx, proto_converter) + PhysicalPlanType::Window(_) => { + WindowAggExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::HashJoin(hashjoin) => { - self.try_into_hash_join_physical_plan(hashjoin, ctx, proto_converter) + PhysicalPlanType::Aggregate(_) => { + AggregateExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::SymmetricHashJoin(sym_join) => self - .try_into_symmetric_hash_join_physical_plan( - sym_join, - ctx, - proto_converter, - ), - PhysicalPlanType::Union(union) => { - self.try_into_union_physical_plan(union, ctx, proto_converter) + PhysicalPlanType::HashJoin(_) => { + HashJoinExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Interleave(interleave) => { - self.try_into_interleave_physical_plan(interleave, ctx, proto_converter) + PhysicalPlanType::SymmetricHashJoin(_) => { + SymmetricHashJoinExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::CrossJoin(crossjoin) => { - self.try_into_cross_join_physical_plan(crossjoin, ctx, proto_converter) + PhysicalPlanType::Union(_) => { + UnionExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Empty(empty) => { - self.try_into_empty_physical_plan(empty, ctx, proto_converter) + PhysicalPlanType::Interleave(_) => { + InterleaveExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::PlaceholderRow(placeholder) => { - self.try_into_placeholder_row_physical_plan(placeholder, ctx) + PhysicalPlanType::CrossJoin(_) => { + CrossJoinExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Sort(sort) => { - self.try_into_sort_physical_plan(sort, ctx, proto_converter) + PhysicalPlanType::Empty(_) => { + EmptyExec::try_from_proto(self.node(), &decode_ctx) + } + PhysicalPlanType::PlaceholderRow(_) => { + PlaceholderRowExec::try_from_proto(self.node(), &decode_ctx) + } + PhysicalPlanType::Sort(_) => { + SortExec::try_from_proto(self.node(), &decode_ctx) + } + PhysicalPlanType::SortPreservingMerge(_) => { + SortPreservingMergeExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::SortPreservingMerge(sort) => self - .try_into_sort_preserving_merge_physical_plan(sort, ctx, proto_converter), PhysicalPlanType::Extension(extension) => { self.try_into_extension_physical_plan(extension, ctx, proto_converter) } - PhysicalPlanType::NestedLoopJoin(join) => { - self.try_into_nested_loop_join_physical_plan(join, ctx, proto_converter) + PhysicalPlanType::NestedLoopJoin(_) => { + NestedLoopJoinExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Analyze(analyze) => { - self.try_into_analyze_physical_plan(analyze, ctx, proto_converter) + PhysicalPlanType::Analyze(_) => { + AnalyzeExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::JsonSink(sink) => { - self.try_into_json_sink_physical_plan(sink, ctx, proto_converter) + PhysicalPlanType::JsonSink(_) => { + JsonSink::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::CsvSink(sink) => { - self.try_into_csv_sink_physical_plan(sink, ctx, proto_converter) + PhysicalPlanType::CsvSink(_) => { + CsvSink::try_from_proto(self.node(), &decode_ctx) } - #[cfg_attr(not(feature = "parquet"), allow(unused_variables))] - PhysicalPlanType::ParquetSink(sink) => { - self.try_into_parquet_sink_physical_plan(sink, ctx, proto_converter) + PhysicalPlanType::ParquetSink(_) => { + #[cfg(feature = "parquet")] + { + ParquetSink::try_from_proto(self.node(), &decode_ctx) + } + #[cfg(not(feature = "parquet"))] + panic!( + "Unable to process a Parquet PhysicalPlan when `parquet` feature is not enabled" + ) } - PhysicalPlanType::Unnest(unnest) => { - self.try_into_unnest_physical_plan(unnest, ctx, proto_converter) + PhysicalPlanType::Unnest(_) => { + UnnestExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Cooperative(cooperative) => { - self.try_into_cooperative_physical_plan(cooperative, ctx, proto_converter) + PhysicalPlanType::Cooperative(_) => { + CooperativeExec::try_from_proto(self.node(), &decode_ctx) } PhysicalPlanType::GenerateSeries(generate_series) => { self.try_into_generate_series_physical_plan(generate_series) } - PhysicalPlanType::SortMergeJoin(sort_join) => { - self.try_into_sort_join(sort_join, ctx, proto_converter) + PhysicalPlanType::SortMergeJoin(_) => { + SortMergeJoinExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::AsyncFunc(async_func) => { - self.try_into_async_func_physical_plan(async_func, ctx, proto_converter) + PhysicalPlanType::AsyncFunc(_) => { + AsyncFuncExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Buffer(buffer) => { - self.try_into_buffer_physical_plan(buffer, ctx, proto_converter) + PhysicalPlanType::Buffer(_) => { + BufferExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::ScalarSubquery(sq) => { - self.try_into_scalar_subquery_physical_plan(sq, ctx, proto_converter) + PhysicalPlanType::ScalarSubquery(_) => { + ScalarSubqueryExec::try_from_proto(self.node(), &decode_ctx) } } } @@ -444,217 +439,26 @@ pub trait PhysicalPlanNodeExt: Sized { let plan_clone = Arc::clone(&plan); let plan = plan.as_ref(); - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_explain_exec(exec, codec); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_projection_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_analyze_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_filter_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(limit) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_global_limit_exec( - limit, - codec, - proto_converter, - ); - } - - if let Some(limit) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_local_limit_exec( - limit, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_hash_join_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_symmetric_hash_join_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_sort_merge_join_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_cross_join_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_aggregate_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(empty) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_empty_exec(empty, codec); - } - - if let Some(empty) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_placeholder_row_exec( - empty, codec, - ); - } - - #[expect(deprecated)] - if let Some(coalesce_batches) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_coalesce_batches_exec( - coalesce_batches, - codec, - proto_converter, - ); - } - - if let Some(data_source_exec) = plan.downcast_ref::() - && let Some(node) = protobuf::PhysicalPlanNode::try_from_data_source_exec( - data_source_exec, - codec, - proto_converter, - )? - { - return Ok(node); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_coalesce_partitions_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_repartition_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_sort_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(union) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_union_exec( - union, - codec, - proto_converter, - ); - } - - if let Some(interleave) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_interleave_exec( - interleave, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_sort_preserving_merge_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_nested_loop_join_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_window_agg_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_bounded_window_agg_exec( - exec, - codec, - proto_converter, - ); + // Self-serializing plans handle themselves via the `try_to_proto` hook + // (#22419). `Ok(None)` means "not migrated" and falls through to the + // central downcast chain below. + // + // The hook is a trait method dispatched on the concrete type, so — like + // the `downcast_ref::()` arms below — it must first follow + // `downcast_delegate()` so a wrapper plan serializes as its delegate. + let encoder = ConverterPlanEncoder { + codec, + proto_converter, + }; + let encode_ctx = ExecutionPlanEncodeCtx::new(&encoder); + let mut hook_target = plan; + while let Some(delegate) = hook_target.downcast_delegate() { + hook_target = delegate; } - - if let Some(exec) = plan.downcast_ref::() - && let Some(node) = protobuf::PhysicalPlanNode::try_from_data_sink_exec( - exec, - codec, - proto_converter, - )? - { + if let Some(node) = hook_target.try_to_proto(&encode_ctx)? { return Ok(node); } - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_unnest_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_cooperative_exec( - exec, - codec, - proto_converter, - ); - } - if let Some(exec) = plan.downcast_ref::() && let Some(node) = protobuf::PhysicalPlanNode::try_from_lazy_memory_exec(exec)? @@ -662,30 +466,6 @@ pub trait PhysicalPlanNodeExt: Sized { return Ok(node); } - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_async_func_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_buffer_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_scalar_subquery_exec( - exec, - codec, - proto_converter, - ); - } - let mut buf: Vec = vec![]; match codec.try_encode(Arc::clone(&plan_clone), &mut buf, proto_converter) { Ok(_) => { @@ -713,2976 +493,213 @@ pub trait PhysicalPlanNodeExt: Sized { ), } } - - fn try_into_explain_physical_plan( - &self, - explain: &protobuf::ExplainExecNode, - _ctx: &PhysicalPlanDecodeContext<'_>, - _proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - Ok(Arc::new(ExplainExec::new( - Arc::new(explain.schema.as_ref().unwrap().try_into()?), - explain - .stringified_plans - .iter() - .map(StringifiedPlan::from_proto) - .collect(), - explain.verbose, - ))) - } - - fn try_into_projection_physical_plan( + fn try_into_extension_physical_plan( &self, - projection: &protobuf::ProjectionExecNode, + extension: &protobuf::PhysicalExtensionNode, ctx: &PhysicalPlanDecodeContext<'_>, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - let input: Arc = - into_physical_plan(&projection.input, ctx, proto_converter)?; - let exprs = projection - .expr + let inputs: Vec> = extension + .inputs .iter() - .zip(projection.expr_name.iter()) - .map(|(expr, name)| { - Ok(( - proto_converter.proto_to_physical_expr( - expr, - input.schema().as_ref(), - ctx, - )?, - name.to_string(), - )) - }) - .collect::, String)>>>()?; - let proj_exprs: Vec = exprs - .into_iter() - .map(|(expr, alias)| ProjectionExpr { expr, alias }) - .collect(); - Ok(Arc::new(ProjectionExec::try_new(proj_exprs, input)?)) - } - - fn try_into_filter_physical_plan( - &self, - filter: &protobuf::FilterExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&filter.input, ctx, proto_converter)?; - - let predicate = filter - .expr - .as_ref() - .map(|expr| { - proto_converter.proto_to_physical_expr(expr, input.schema().as_ref(), ctx) - }) - .transpose()? - .ok_or_else(|| { - internal_datafusion_err!( - "filter (FilterExecNode) in PhysicalPlanNode is missing." - ) - })?; - - let filter_selectivity = filter.default_filter_selectivity.try_into(); - // Preserve the `None` state across proto boundaries. Proto cannot distinguish - // between `None` (full projection) and `Some(vec![])` (empty projection) since - // both serialize as an empty list. If all columns are included, we reconstruct - // `None` to avoid losing this semantic distinction on deserialization. - let num_fields = input.schema().fields().len(); - let mut is_full_projection = filter.projection.len() == num_fields; - let mut projection_vec: Vec = Vec::with_capacity(filter.projection.len()); - for (i, idx) in filter.projection.iter().enumerate() { - let idx = *idx as usize; - is_full_projection &= idx == i; - projection_vec.push(idx); - } - let projection = if is_full_projection { - None - } else { - Some(projection_vec) - }; - let filter = FilterExecBuilder::new(predicate, input) - .apply_projection(projection)? - .with_batch_size(filter.batch_size as usize) - .with_fetch(filter.fetch.map(|f| f as usize)) - .build()?; - match filter_selectivity { - Ok(filter_selectivity) => Ok(Arc::new( - filter.with_default_selectivity(filter_selectivity)?, - )), - Err(_) => Err(internal_datafusion_err!( - "filter_selectivity in PhysicalPlanNode is invalid " - )), - } - } - - fn try_into_csv_scan_physical_plan( - &self, - scan: &protobuf::CsvScanExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let escape = - if let Some(protobuf::csv_scan_exec_node::OptionalEscape::Escape(escape)) = - &scan.optional_escape - { - Some(str_to_byte(escape, "escape")?) - } else { - None - }; - - let comment = if let Some( - protobuf::csv_scan_exec_node::OptionalComment::Comment(comment), - ) = &scan.optional_comment - { - Some(str_to_byte(comment, "comment")?) - } else { - None - }; - - // Parse table schema with partition columns - let table_schema = - parse_table_schema_from_proto(scan.base_conf.as_ref().unwrap())?; - - let csv_options = CsvOptions { - has_header: Some(scan.has_header), - delimiter: str_to_byte(&scan.delimiter, "delimiter")?, - quote: str_to_byte(&scan.quote, "quote")?, - newlines_in_values: Some(scan.newlines_in_values), - ..Default::default() - }; - let source = Arc::new( - CsvSource::new(table_schema) - .with_csv_options(csv_options) - .with_escape(escape) - .with_comment(comment), - ); - - let conf = FileScanConfigBuilder::from(parse_protobuf_file_scan_config( - scan.base_conf.as_ref().unwrap(), - ctx, - proto_converter, - source, - )?) - .with_file_compression_type(FileCompressionType::UNCOMPRESSED) - .build(); - Ok(DataSourceExec::from_data_source(conf)) - } - - fn try_into_json_scan_physical_plan( - &self, - scan: &protobuf::JsonScanExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - 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 scan_conf = parse_protobuf_file_scan_config( - base_conf, - ctx, - proto_converter, - Arc::new(JsonSource::new(table_schema)), - )?; - Ok(DataSourceExec::from_data_source(scan_conf)) - } + .map(|i| proto_converter.proto_to_execution_plan(i, ctx)) + .collect::>()?; - fn try_into_arrow_scan_physical_plan( - &self, - scan: &protobuf::ArrowScanExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - 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 scan_conf = parse_protobuf_file_scan_config( - base_conf, - ctx, + let extension_node = ctx.codec().try_decode( + extension.node.as_slice(), + &inputs, + ctx.task_ctx(), proto_converter, - Arc::new(ArrowSource::new_file_source(table_schema)), )?; - Ok(DataSourceExec::from_data_source(scan_conf)) - } - - #[cfg_attr(not(feature = "parquet"), expect(unused_variables))] - fn try_into_parquet_scan_physical_plan( - &self, - scan: &protobuf::ParquetScanExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - #[cfg(feature = "parquet")] - { - let schema = from_proto::parse_protobuf_file_scan_schema( - scan.base_conf.as_ref().unwrap(), - )?; - - // Check if there's a projection and use projected schema for predicate parsing - let base_conf = scan.base_conf.as_ref().unwrap(); - let predicate_schema = if !base_conf.projection.is_empty() { - // Create projected schema for parsing the predicate - let projected_fields: Vec<_> = base_conf - .projection - .iter() - .map(|&i| schema.field(i as usize).clone()) - .collect(); - Arc::new(Schema::new(projected_fields)) - } else { - schema - }; - - let predicate = scan - .predicate - .as_ref() - .map(|expr| { - proto_converter.proto_to_physical_expr( - expr, - predicate_schema.as_ref(), - ctx, - ) - }) - .transpose()?; - let mut options = datafusion_common::config::TableParquetOptions::default(); - - if let Some(table_options) = scan.parquet_options.as_ref() { - options = table_options.try_into()?; - } - // Parse table schema with partition columns - let table_schema = 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(), - }; - let store = ctx - .task_ctx() - .runtime_env() - .object_store(object_store_url)?; - let metadata_cache = ctx - .task_ctx() - .runtime_env() - .cache_manager - .get_file_metadata_cache(); - let reader_factory = - Arc::new(CachedParquetFileReaderFactory::new(store, metadata_cache)); - - let mut source = ParquetSource::new(table_schema) - .with_parquet_file_reader_factory(reader_factory) - .with_table_parquet_options(options); - - if let Some(predicate) = predicate { - source = source.with_predicate(predicate); - } - let base_config = parse_protobuf_file_scan_config( - base_conf, - ctx, - proto_converter, - Arc::new(source), - )?; - Ok(DataSourceExec::from_data_source(base_config)) - } - #[cfg(not(feature = "parquet"))] - panic!( - "Unable to process a Parquet PhysicalPlan when `parquet` feature is not enabled" - ) + Ok(extension_node) } - - #[cfg_attr(not(feature = "avro"), expect(unused_variables))] - fn try_into_avro_scan_physical_plan( - &self, - scan: &protobuf::AvroScanExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - #[cfg(feature = "avro")] - { - let table_schema = - 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, - proto_converter, - Arc::new(AvroSource::new(table_schema)), - )?; - Ok(DataSourceExec::from_data_source(conf)) + fn generate_series_name_to_str(name: protobuf::GenerateSeriesName) -> &'static str { + match name { + protobuf::GenerateSeriesName::GsGenerateSeries => "generate_series", + protobuf::GenerateSeriesName::GsRange => "range", } - - #[cfg(not(feature = "avro"))] - panic!("Unable to process a Avro PhysicalPlan when `avro` feature is not enabled") } - - fn try_into_memory_scan_physical_plan( + fn try_into_sort_join( &self, - scan: &protobuf::MemoryScanExecNode, + sort_join: &SortMergeJoinExecNode, ctx: &PhysicalPlanDecodeContext<'_>, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - let partitions = scan - .partitions - .iter() - .map(|p| parse_record_batches(p)) - .collect::>>()?; + let left = into_physical_plan(&sort_join.left, ctx, proto_converter)?; + let left_schema = left.schema(); + let right = into_physical_plan(&sort_join.right, ctx, proto_converter)?; + let right_schema = right.schema(); - let proto_schema = scan.schema.as_ref().ok_or_else(|| { - internal_datafusion_err!("schema in MemoryScanExecNode is missing.") - })?; - let schema: SchemaRef = SchemaRef::new(proto_schema.try_into()?); + let filter = sort_join + .filter + .as_ref() + .map(|f| { + let schema = f + .schema + .as_ref() + .ok_or_else(|| proto_error("Missing JoinFilter schema"))? + .try_into()?; - let projection = if !scan.projection.is_empty() { - Some( - scan.projection + let expression = proto_converter.proto_to_physical_expr( + f.expression.as_ref().ok_or_else(|| { + proto_error("Unexpected empty filter expression") + })?, + &schema, + ctx, + )?; + let column_indices = f + .column_indices .iter() - .map(|i| *i as usize) - .collect::>(), - ) - } else { - None - }; - - let mut sort_information = vec![]; - for ordering in &scan.sort_information { - let sort_exprs = parse_physical_sort_exprs( - &ordering.physical_sort_expr_nodes, - ctx, - &schema, - proto_converter, - )?; - sort_information.extend(LexOrdering::new(sort_exprs)); - } - - let source = MemorySourceConfig::try_new(&partitions, schema, projection)? - .with_limit(scan.fetch.map(|f| f as usize)) - .with_show_sizes(scan.show_sizes); - - let source = source.try_with_sort_information(sort_information)?; - - Ok(DataSourceExec::from_data_source(source)) - } + .map(|i| { + let side = + protobuf::JoinSide::try_from(i.side).map_err(|_| { + proto_error(format!( + "Received a SortMergeJoinExecNode message with JoinSide in Filter {}", + i.side + )) + })?; - fn try_into_coalesce_batches_physical_plan( - &self, - coalesce_batches: &protobuf::CoalesceBatchesExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&coalesce_batches.input, ctx, proto_converter)?; - Ok(Arc::new( - #[expect(deprecated)] - CoalesceBatchesExec::new(input, coalesce_batches.target_batch_size as usize) - .with_fetch(coalesce_batches.fetch.map(|f| f as usize)), - )) - } + Ok(ColumnIndex { + index: i.index as usize, + side: side.into(), + }) + }) + .collect::>>()?; - fn try_into_merge_physical_plan( - &self, - merge: &protobuf::CoalescePartitionsExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&merge.input, ctx, proto_converter)?; - Ok(Arc::new( - CoalescePartitionsExec::new(input) - .with_fetch(merge.fetch.map(|f| f as usize)), - )) - } + Ok(JoinFilter::new( + expression, + column_indices, + Arc::new(schema), + )) + }) + .map_or(Ok(None), |v: Result| v.map(Some))?; - fn try_into_repartition_physical_plan( - &self, - repart: &protobuf::RepartitionExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&repart.input, ctx, proto_converter)?; - let partitioning = parse_protobuf_partitioning( - repart.partitioning.as_ref(), - ctx, - input.schema().as_ref(), - proto_converter, - )?; - let mut repart_exec = RepartitionExec::try_new(input, partitioning.unwrap())?; - if repart.preserve_order { - repart_exec = repart_exec.with_preserve_order(); - } - Ok(Arc::new(repart_exec)) - } - - fn try_into_global_limit_physical_plan( - &self, - limit: &protobuf::GlobalLimitExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&limit.input, ctx, proto_converter)?; - let fetch = if limit.fetch >= 0 { - Some(limit.fetch as usize) - } else { - None - }; - Ok(Arc::new(GlobalLimitExec::new( - input, - limit.skip as usize, - fetch, - ))) - } - - fn try_into_local_limit_physical_plan( - &self, - limit: &protobuf::LocalLimitExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&limit.input, ctx, proto_converter)?; - Ok(Arc::new(LocalLimitExec::new(input, limit.fetch as usize))) - } - - fn try_into_window_physical_plan( - &self, - window_agg: &protobuf::WindowAggExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&window_agg.input, ctx, proto_converter)?; - let input_schema = input.schema(); - - let physical_window_expr: Vec> = window_agg - .window_expr - .iter() - .map(|window_expr| { - parse_physical_window_expr( - window_expr, - ctx, - input_schema.as_ref(), - proto_converter, - ) - }) - .collect::, _>>()?; - - let partition_keys = window_agg - .partition_keys - .iter() - .map(|expr| { - proto_converter.proto_to_physical_expr(expr, input.schema().as_ref(), ctx) - }) - .collect::>>>()?; - - if let Some(input_order_mode) = window_agg.input_order_mode.as_ref() { - let input_order_mode = match input_order_mode { - window_agg_exec_node::InputOrderMode::Linear(_) => InputOrderMode::Linear, - window_agg_exec_node::InputOrderMode::PartiallySorted( - protobuf::PartiallySortedInputOrderMode { columns }, - ) => InputOrderMode::PartiallySorted( - columns.iter().map(|c| *c as usize).collect(), - ), - window_agg_exec_node::InputOrderMode::Sorted(_) => InputOrderMode::Sorted, - }; - - Ok(Arc::new(BoundedWindowAggExec::try_new( - physical_window_expr, - input, - input_order_mode, - !partition_keys.is_empty(), - )?)) - } else { - Ok(Arc::new(WindowAggExec::try_new( - physical_window_expr, - input, - !partition_keys.is_empty(), - )?)) - } - } - - fn try_into_aggregate_physical_plan( - &self, - hash_agg: &protobuf::AggregateExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&hash_agg.input, ctx, proto_converter)?; - let mode = protobuf::AggregateMode::try_from(hash_agg.mode).map_err(|_| { - proto_error(format!( - "Received a AggregateNode message with unknown AggregateMode {}", - hash_agg.mode - )) - })?; - let agg_mode: AggregateMode = match mode { - protobuf::AggregateMode::Partial => AggregateMode::Partial, - protobuf::AggregateMode::Final => AggregateMode::Final, - protobuf::AggregateMode::FinalPartitioned => AggregateMode::FinalPartitioned, - protobuf::AggregateMode::Single => AggregateMode::Single, - protobuf::AggregateMode::SinglePartitioned => { - AggregateMode::SinglePartitioned - } - protobuf::AggregateMode::PartialReduce => AggregateMode::PartialReduce, - }; - - let num_expr = hash_agg.group_expr.len(); - - let group_expr = hash_agg - .group_expr - .iter() - .zip(hash_agg.group_expr_name.iter()) - .map(|(expr, name)| { - proto_converter - .proto_to_physical_expr(expr, input.schema().as_ref(), ctx) - .map(|expr| (expr, name.to_string())) - }) - .collect::, _>>()?; - - let null_expr = hash_agg - .null_expr - .iter() - .zip(hash_agg.group_expr_name.iter()) - .map(|(expr, name)| { - proto_converter - .proto_to_physical_expr(expr, input.schema().as_ref(), ctx) - .map(|expr| (expr, name.to_string())) - }) - .collect::, _>>()?; - - let groups: Vec> = if !hash_agg.groups.is_empty() { - hash_agg - .groups - .chunks(num_expr) - .map(|g| g.to_vec()) - .collect::>>() - } else { - vec![] - }; - - let has_grouping_set = hash_agg.has_grouping_set; - - let input_schema = hash_agg.input_schema.as_ref().ok_or_else(|| { - internal_datafusion_err!("input_schema in AggregateNode is missing.") - })?; - let physical_schema: SchemaRef = SchemaRef::new(input_schema.try_into()?); - - let physical_filter_expr = hash_agg - .filter_expr - .iter() - .map(|expr| { - expr.expr - .as_ref() - .map(|e| { - proto_converter.proto_to_physical_expr(e, &physical_schema, ctx) - }) - .transpose() - }) - .collect::, _>>()?; - - let physical_aggr_expr: Vec> = hash_agg - .aggr_expr - .iter() - .zip(hash_agg.aggr_expr_name.iter()) - .map(|(expr, name)| { - let expr_type = expr.expr_type.as_ref().ok_or_else(|| { - proto_error("Unexpected empty aggregate physical expression") - })?; - - match expr_type { - ExprType::AggregateExpr(agg_node) => { - let input_phy_expr: Vec> = agg_node - .expr - .iter() - .map(|e| { - proto_converter.proto_to_physical_expr( - e, - &physical_schema, - ctx, - ) - }) - .collect::>>()?; - let order_bys = agg_node - .ordering_req - .iter() - .map(|e| { - parse_physical_sort_expr( - e, - ctx, - &physical_schema, - proto_converter, - ) - }) - .collect::>()?; - agg_node - .aggregate_function - .as_ref() - .map(|func| match func { - AggregateFunction::UserDefinedAggrFunction(udaf_name) => { - let agg_udf = match &agg_node.fun_definition { - Some(buf) => { - ctx.codec().try_decode_udaf(udaf_name, buf)? - } - None => ctx.task_ctx().udaf(udaf_name).or_else( - |_| { - ctx.codec() - .try_decode_udaf(udaf_name, &[]) - }, - )?, - }; - - let (human_display, human_display_alias) = - split_human_display_alias( - &agg_node.human_display, - name, - ); - let builder = AggregateExprBuilder::new( - agg_udf, - input_phy_expr, - ) - .schema(Arc::clone(&physical_schema)) - .alias(name) - .with_ignore_nulls(agg_node.ignore_nulls) - .with_distinct(agg_node.distinct) - .order_by(order_bys) - .human_display(human_display); - let builder = if let Some(alias) = human_display_alias - { - builder.human_display_alias(alias) - } else { - builder - }; - builder.build().map(Arc::new) - } - }) - .transpose()? - .ok_or_else(|| { - proto_error( - "Invalid AggregateExpr, missing aggregate_function", - ) - }) - } - _ => internal_err!("Invalid aggregate expression for AggregateExec"), - } - }) - .collect::, _>>()?; - - let physical_schema_ref = Arc::clone(&physical_schema); - let agg = AggregateExec::try_new( - agg_mode, - PhysicalGroupBy::new(group_expr, null_expr, groups, has_grouping_set), - physical_aggr_expr, - physical_filter_expr, - input, - physical_schema, - )?; - - let agg = if let Some(limit_proto) = &hash_agg.limit { - let limit = limit_proto.limit as usize; - let limit_options = match limit_proto.descending { - Some(descending) => LimitOptions::new_with_order(limit, descending), - None => LimitOptions::new(limit), - }; - agg.with_limit_options(Some(limit_options)) - } else { - agg - }; - - let agg = if let Some(dynamic_filter_proto) = &hash_agg.dynamic_filter { - let dynamic_filter_expr = proto_converter.proto_to_physical_expr( - dynamic_filter_proto, - physical_schema_ref.as_ref(), - ctx, - )?; - let df = (dynamic_filter_expr as Arc) - .downcast::() - .map_err(|_| { - internal_datafusion_err!( - "AggregateExec dynamic_filter did not decode to a DynamicFilterPhysicalExpr" - ) - })?; - agg.with_dynamic_filter_expr(df)? - } else { - agg - }; - - Ok(Arc::new(agg)) - } - - fn try_into_hash_join_physical_plan( - &self, - hashjoin: &protobuf::HashJoinExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let left: Arc = - into_physical_plan(&hashjoin.left, ctx, proto_converter)?; - let right: Arc = - into_physical_plan(&hashjoin.right, ctx, proto_converter)?; - let left_schema = left.schema(); - let right_schema = right.schema(); - let on: Vec<(PhysicalExprRef, PhysicalExprRef)> = hashjoin - .on - .iter() - .map(|col| { - let left = proto_converter.proto_to_physical_expr( - &col.left.clone().unwrap(), - left_schema.as_ref(), - ctx, - )?; - let right = proto_converter.proto_to_physical_expr( - &col.right.clone().unwrap(), - right_schema.as_ref(), - ctx, - )?; - Ok((left, right)) - }) - .collect::>()?; - let join_type = - protobuf::JoinType::try_from(hashjoin.join_type).map_err(|_| { - proto_error(format!( - "Received a HashJoinNode message with unknown JoinType {}", - hashjoin.join_type - )) - })?; - let null_equality = protobuf::NullEquality::try_from(hashjoin.null_equality) - .map_err(|_| { - proto_error(format!( - "Received a HashJoinNode message with unknown NullEquality {}", - hashjoin.null_equality - )) - })?; - let filter = hashjoin - .filter - .as_ref() - .map(|f| { - let schema = f - .schema - .as_ref() - .ok_or_else(|| proto_error("Missing JoinFilter schema"))? - .try_into()?; - - let expression = proto_converter.proto_to_physical_expr( - f.expression.as_ref().ok_or_else(|| { - proto_error("Unexpected empty filter expression") - })?, - &schema, - ctx, - )?; - let column_indices = f.column_indices - .iter() - .map(|i| { - let side = protobuf::JoinSide::try_from(i.side) - .map_err(|_| proto_error(format!( - "Received a HashJoinNode message with JoinSide in Filter {}", - i.side)) - )?; - - Ok(ColumnIndex { - index: i.index as usize, - side: side.into(), - }) - }) - .collect::>>()?; - - Ok(JoinFilter::new(expression, column_indices, Arc::new(schema))) - }) - .map_or(Ok(None), |v: Result| v.map(Some))?; - - let partition_mode = protobuf::PartitionMode::try_from(hashjoin.partition_mode) - .map_err(|_| { - proto_error(format!( - "Received a HashJoinNode message with unknown PartitionMode {}", - hashjoin.partition_mode - )) - })?; - let partition_mode = match partition_mode { - protobuf::PartitionMode::CollectLeft => PartitionMode::CollectLeft, - protobuf::PartitionMode::Partitioned => PartitionMode::Partitioned, - protobuf::PartitionMode::Auto => PartitionMode::Auto, - }; - // Proto3 `repeated` cannot distinguish `None` from `Some(vec![])`. The latter - // is reachable via `try_embed_projection` for `SELECT count(1) … JOIN …` and - // changes the join's output schema, so the encoder reserves the single-element - // sentinel `[u32::MAX]` (never a valid column index) to mean "explicitly empty"; - // every other state is sent as-is. See `try_from_hash_join_exec`. - let projection = match hashjoin.projection.as_slice() { - [] => None, - [u32::MAX] => Some(Vec::new()), - indices => Some(indices.iter().map(|i| *i as usize).collect()), - }; - let mut hash_join = HashJoinExec::try_new( - left, - right, - on, - filter, - &JoinType::from_proto(join_type), - projection, - partition_mode, - NullEquality::from_proto(null_equality), - hashjoin.null_aware, - )?; - - if let Some(dynamic_filter_proto) = &hashjoin.dynamic_filter { - let dynamic_filter_expr = proto_converter.proto_to_physical_expr( - dynamic_filter_proto, - right_schema.as_ref(), - ctx, - )?; - let df = (dynamic_filter_expr as Arc) - .downcast::() - .map_err(|_| { - internal_datafusion_err!( - "HashJoinExec dynamic_filter did not decode to a DynamicFilterPhysicalExpr" - ) - })?; - hash_join = hash_join.with_dynamic_filter_expr(df)?; - } - - Ok(Arc::new(hash_join)) - } - - fn try_into_symmetric_hash_join_physical_plan( - &self, - sym_join: &protobuf::SymmetricHashJoinExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let left = into_physical_plan(&sym_join.left, ctx, proto_converter)?; - let right = into_physical_plan(&sym_join.right, ctx, proto_converter)?; - let left_schema = left.schema(); - let right_schema = right.schema(); - let on = sym_join - .on - .iter() - .map(|col| { - let left = proto_converter.proto_to_physical_expr( - &col.left.clone().unwrap(), - left_schema.as_ref(), - ctx, - )?; - let right = proto_converter.proto_to_physical_expr( - &col.right.clone().unwrap(), - right_schema.as_ref(), - ctx, - )?; - Ok((left, right)) - }) - .collect::>()?; let join_type = - protobuf::JoinType::try_from(sym_join.join_type).map_err(|_| { - proto_error(format!( - "Received a SymmetricHashJoin message with unknown JoinType {}", - sym_join.join_type - )) - })?; - let null_equality = protobuf::NullEquality::try_from(sym_join.null_equality) - .map_err(|_| { + protobuf::JoinType::try_from(sort_join.join_type).map_err(|_| { proto_error(format!( - "Received a SymmetricHashJoin message with unknown NullEquality {}", - sym_join.null_equality + "Received a SortMergeJoinExecNode message with unknown JoinType {}", + sort_join.join_type )) })?; - let filter = sym_join - .filter - .as_ref() - .map(|f| { - let schema = f - .schema - .as_ref() - .ok_or_else(|| proto_error("Missing JoinFilter schema"))? - .try_into()?; - - let expression = proto_converter.proto_to_physical_expr( - f.expression.as_ref().ok_or_else(|| { - proto_error("Unexpected empty filter expression") - })?, - &schema, - ctx, - )?; - let column_indices = f.column_indices - .iter() - .map(|i| { - let side = protobuf::JoinSide::try_from(i.side) - .map_err(|_| proto_error(format!( - "Received a HashJoinNode message with JoinSide in Filter {}", - i.side)) - )?; - - Ok(ColumnIndex { - index: i.index as usize, - side: side.into(), - }) - }) - .collect::>()?; - - Ok(JoinFilter::new(expression, column_indices, Arc::new(schema))) - }) - .map_or(Ok(None), |v: Result| v.map(Some))?; - - let left_sort_exprs = parse_physical_sort_exprs( - &sym_join.left_sort_exprs, - ctx, - &left_schema, - proto_converter, - )?; - let left_sort_exprs = LexOrdering::new(left_sort_exprs); - - let right_sort_exprs = parse_physical_sort_exprs( - &sym_join.right_sort_exprs, - ctx, - &right_schema, - proto_converter, - )?; - let right_sort_exprs = LexOrdering::new(right_sort_exprs); - - let partition_mode = protobuf::StreamPartitionMode::try_from( - sym_join.partition_mode, - ) - .map_err(|_| { - proto_error(format!( - "Received a SymmetricHashJoin message with unknown PartitionMode {}", - sym_join.partition_mode - )) - })?; - let partition_mode = match partition_mode { - protobuf::StreamPartitionMode::SinglePartition => { - StreamJoinPartitionMode::SinglePartition - } - protobuf::StreamPartitionMode::PartitionedExec => { - StreamJoinPartitionMode::Partitioned - } - }; - SymmetricHashJoinExec::try_new( - left, - right, - on, - filter, - &JoinType::from_proto(join_type), - NullEquality::from_proto(null_equality), - left_sort_exprs, - right_sort_exprs, - partition_mode, - ) - .map(|e| Arc::new(e) as _) - } - - fn try_into_union_physical_plan( - &self, - union: &protobuf::UnionExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let mut inputs: Vec> = vec![]; - for input in &union.inputs { - inputs.push(proto_converter.proto_to_execution_plan(input, ctx)?); - } - UnionExec::try_new(inputs) - } - - fn try_into_interleave_physical_plan( - &self, - interleave: &protobuf::InterleaveExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let mut inputs: Vec> = vec![]; - for input in &interleave.inputs { - inputs.push(proto_converter.proto_to_execution_plan(input, ctx)?); - } - Ok(Arc::new(InterleaveExec::try_new(inputs)?)) - } - - fn try_into_cross_join_physical_plan( - &self, - crossjoin: &protobuf::CrossJoinExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let left: Arc = - into_physical_plan(&crossjoin.left, ctx, proto_converter)?; - let right: Arc = - into_physical_plan(&crossjoin.right, ctx, proto_converter)?; - Ok(Arc::new(CrossJoinExec::new(left, right))) - } - - fn try_into_empty_physical_plan( - &self, - empty: &protobuf::EmptyExecNode, - _ctx: &PhysicalPlanDecodeContext<'_>, - _proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let schema = Arc::new(convert_required!(empty.schema)?); - Ok(Arc::new(EmptyExec::new(schema))) - } - - fn try_into_placeholder_row_physical_plan( - &self, - placeholder: &protobuf::PlaceholderRowExecNode, - _ctx: &PhysicalPlanDecodeContext<'_>, - ) -> Result> { - let schema = Arc::new(convert_required!(placeholder.schema)?); - Ok(Arc::new(PlaceholderRowExec::new(schema))) - } - - fn try_into_sort_physical_plan( - &self, - sort: &protobuf::SortExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let node = self.node(); - let input = into_physical_plan(&sort.input, ctx, proto_converter)?; - let exprs = sort - .expr - .iter() - .map(|expr| { - let expr = expr.expr_type.as_ref().ok_or_else(|| { - proto_error(format!( - "physical_plan::from_proto() Unexpected expr {node:?}" - )) - })?; - if let ExprType::Sort(sort_expr) = expr { - let expr = sort_expr - .expr - .as_ref() - .ok_or_else(|| { - proto_error(format!( - "physical_plan::from_proto() Unexpected sort expr {node:?}" - )) - })? - .as_ref(); - Ok(PhysicalSortExpr { - expr: proto_converter.proto_to_physical_expr( - expr, - input.schema().as_ref(), - ctx, - )?, - options: SortOptions { - descending: !sort_expr.asc, - nulls_first: sort_expr.nulls_first, - }, - }) - } else { - internal_err!( - "physical_plan::from_proto() {node:?}" - ) - } - }) - .collect::>>()?; - let Some(ordering) = LexOrdering::new(exprs) else { - return internal_err!("SortExec requires an ordering"); - }; - let fetch = (sort.fetch >= 0).then_some(sort.fetch as _); - let new_sort = SortExec::new(ordering, input) - .with_fetch(fetch) - .with_preserve_partitioning(sort.preserve_partitioning); - - let new_sort = if let Some(dynamic_filter_proto) = &sort.dynamic_filter { - let dynamic_filter_expr = proto_converter.proto_to_physical_expr( - dynamic_filter_proto, - new_sort.input().schema().as_ref(), - ctx, - )?; - let df = (dynamic_filter_expr as Arc) - .downcast::() - .map_err(|_| { - internal_datafusion_err!( - "SortExec dynamic_filter did not decode to a DynamicFilterPhysicalExpr" - ) - })?; - new_sort.with_dynamic_filter_expr(df)? - } else { - new_sort - }; - - Ok(Arc::new(new_sort)) - } - - fn try_into_sort_preserving_merge_physical_plan( - &self, - sort: &protobuf::SortPreservingMergeExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let node = self.node(); - let input = into_physical_plan(&sort.input, ctx, proto_converter)?; - let exprs = sort - .expr - .iter() - .map(|expr| { - let expr = expr.expr_type.as_ref().ok_or_else(|| { - proto_error(format!( - "physical_plan::from_proto() Unexpected expr {node:?}" - )) - })?; - if let ExprType::Sort(sort_expr) = expr { - let expr = sort_expr - .expr - .as_ref() - .ok_or_else(|| { - proto_error(format!( - "physical_plan::from_proto() Unexpected sort expr {node:?}" - )) - })? - .as_ref(); - Ok(PhysicalSortExpr { - expr: proto_converter.proto_to_physical_expr( - expr, - input.schema().as_ref(), - ctx, - )?, - options: SortOptions { - descending: !sort_expr.asc, - nulls_first: sort_expr.nulls_first, - }, - }) - } else { - internal_err!("physical_plan::from_proto() {node:?}") - } - }) - .collect::>>()?; - let Some(ordering) = LexOrdering::new(exprs) else { - return internal_err!("SortExec requires an ordering"); - }; - let fetch = (sort.fetch >= 0).then_some(sort.fetch as _); - Ok(Arc::new( - SortPreservingMergeExec::new(ordering, input).with_fetch(fetch), - )) - } - - fn try_into_extension_physical_plan( - &self, - extension: &protobuf::PhysicalExtensionNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let inputs: Vec> = extension - .inputs - .iter() - .map(|i| proto_converter.proto_to_execution_plan(i, ctx)) - .collect::>()?; - - let extension_node = ctx.codec().try_decode( - extension.node.as_slice(), - &inputs, - ctx.task_ctx(), - proto_converter, - )?; - - Ok(extension_node) - } - - fn try_into_nested_loop_join_physical_plan( - &self, - join: &protobuf::NestedLoopJoinExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let left: Arc = - into_physical_plan(&join.left, ctx, proto_converter)?; - let right: Arc = - into_physical_plan(&join.right, ctx, proto_converter)?; - let join_type = protobuf::JoinType::try_from(join.join_type).map_err(|_| { - proto_error(format!( - "Received a NestedLoopJoinExecNode message with unknown JoinType {}", - join.join_type - )) - })?; - let filter = join - .filter - .as_ref() - .map(|f| { - let schema = f - .schema - .as_ref() - .ok_or_else(|| proto_error("Missing JoinFilter schema"))? - .try_into()?; - - let expression = proto_converter - .proto_to_physical_expr( - f.expression.as_ref().ok_or_else(|| { - proto_error("Unexpected empty filter expression") - })?, - &schema, - ctx, - )?; - let column_indices = f.column_indices - .iter() - .map(|i| { - let side = protobuf::JoinSide::try_from(i.side) - .map_err(|_| proto_error(format!( - "Received a NestedLoopJoinExecNode message with JoinSide in Filter {}", - i.side)) - )?; - - Ok(ColumnIndex { - index: i.index as usize, - side: side.into(), - }) - }) - .collect::>>()?; - - Ok(JoinFilter::new(expression, column_indices, Arc::new(schema))) - }) - .map_or(Ok(None), |v: Result| v.map(Some))?; - - // See `try_into_hash_join_physical_plan` for the rationale behind the - // `[u32::MAX]` sentinel; `NestedLoopJoinExec` has the same `Option>` - // projection field and shares the proto3 `repeated` ambiguity. - let projection = match join.projection.as_slice() { - [] => None, - [u32::MAX] => Some(Vec::new()), - indices => Some(indices.iter().map(|i| *i as usize).collect()), - }; - - Ok(Arc::new(NestedLoopJoinExec::try_new( - left, - right, - filter, - &JoinType::from_proto(join_type), - projection, - )?)) - } - - fn try_into_analyze_physical_plan( - &self, - analyze: &protobuf::AnalyzeExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&analyze.input, ctx, proto_converter)?; - let metric_categories = if analyze.has_metric_categories { - let cats: Result> = analyze - .metric_categories - .iter() - .map(|s| s.parse::()) - .collect(); - Some(cats?) - } else { - None - }; - let pb_format = - protobuf::ExplainFormat::try_from(analyze.format).map_err(|_| { - DataFusionError::Internal(format!( - "Received an AnalyzeExecNode message with unknown ExplainFormat {}", - analyze.format - )) - })?; - let format = match pb_format { - protobuf::ExplainFormat::Indent => ExplainFormat::Indent, - protobuf::ExplainFormat::Tree => ExplainFormat::Tree, - protobuf::ExplainFormat::Pgjson => ExplainFormat::PostgresJSON, - protobuf::ExplainFormat::Graphviz => ExplainFormat::Graphviz, - }; - Ok(Arc::new( - AnalyzeExec::builder( - analyze.verbose, - analyze.show_statistics, - input, - Arc::new(convert_required!(analyze.schema)?), - ) - .with_metric_categories(metric_categories) - .with_format(format) - .build(), - )) - } - - fn try_into_json_sink_physical_plan( - &self, - sink: &protobuf::JsonSinkExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input = into_physical_plan(&sink.input, ctx, proto_converter)?; - - let data_sink = JsonSink::try_from_proto( - sink.sink - .as_ref() - .ok_or_else(|| proto_error("Missing required field in protobuf"))?, - )?; - let sink_schema = input.schema(); - let sort_order = sink - .sort_order - .as_ref() - .map(|collection| { - parse_physical_sort_exprs( - &collection.physical_sort_expr_nodes, - ctx, - &sink_schema, - proto_converter, - ) - .map(|sort_exprs| { - LexRequirement::new(sort_exprs.into_iter().map(Into::into)) - }) - }) - .transpose()? - .flatten(); - Ok(Arc::new(DataSinkExec::new( - input, - Arc::new(data_sink), - sort_order, - ))) - } - - fn try_into_csv_sink_physical_plan( - &self, - sink: &protobuf::CsvSinkExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input = into_physical_plan(&sink.input, ctx, proto_converter)?; - - let data_sink = CsvSink::try_from_proto( - sink.sink - .as_ref() - .ok_or_else(|| proto_error("Missing required field in protobuf"))?, - )?; - let sink_schema = input.schema(); - let sort_order = sink - .sort_order - .as_ref() - .map(|collection| { - parse_physical_sort_exprs( - &collection.physical_sort_expr_nodes, - ctx, - &sink_schema, - proto_converter, - ) - .map(|sort_exprs| { - LexRequirement::new(sort_exprs.into_iter().map(Into::into)) - }) - }) - .transpose()? - .flatten(); - Ok(Arc::new(DataSinkExec::new( - input, - Arc::new(data_sink), - sort_order, - ))) - } - - #[cfg_attr(not(feature = "parquet"), expect(unused_variables))] - fn try_into_parquet_sink_physical_plan( - &self, - sink: &protobuf::ParquetSinkExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - #[cfg(feature = "parquet")] - { - let input = into_physical_plan(&sink.input, ctx, proto_converter)?; - - let data_sink = ParquetSink::try_from_proto( - sink.sink - .as_ref() - .ok_or_else(|| proto_error("Missing required field in protobuf"))?, - )?; - let sink_schema = input.schema(); - let sort_order = sink - .sort_order - .as_ref() - .map(|collection| { - parse_physical_sort_exprs( - &collection.physical_sort_expr_nodes, - ctx, - &sink_schema, - proto_converter, - ) - .map(|sort_exprs| { - LexRequirement::new(sort_exprs.into_iter().map(Into::into)) - }) - }) - .transpose()? - .flatten(); - Ok(Arc::new(DataSinkExec::new( - input, - Arc::new(data_sink), - sort_order, - ))) - } - #[cfg(not(feature = "parquet"))] - panic!("Trying to use ParquetSink without `parquet` feature enabled"); - } - - fn try_into_unnest_physical_plan( - &self, - unnest: &protobuf::UnnestExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input = into_physical_plan(&unnest.input, ctx, proto_converter)?; - - Ok(Arc::new(UnnestExec::new( - input, - unnest - .list_type_columns - .iter() - .map(|c| ListUnnest { - index_in_input_schema: c.index_in_input_schema as _, - depth: c.depth as _, - }) - .collect(), - unnest.struct_type_columns.iter().map(|c| *c as _).collect(), - Arc::new(convert_required!(unnest.schema)?), - unnest - .options - .as_ref() - .map(datafusion_common::UnnestOptions::from_proto) - .ok_or_else(|| proto_error("Missing required field in protobuf"))?, - )?)) - } - - fn generate_series_name_to_str(name: protobuf::GenerateSeriesName) -> &'static str { - match name { - protobuf::GenerateSeriesName::GsGenerateSeries => "generate_series", - protobuf::GenerateSeriesName::GsRange => "range", - } - } - fn try_into_sort_join( - &self, - sort_join: &SortMergeJoinExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let left = into_physical_plan(&sort_join.left, ctx, proto_converter)?; - let left_schema = left.schema(); - let right = into_physical_plan(&sort_join.right, ctx, proto_converter)?; - let right_schema = right.schema(); - - let filter = sort_join - .filter - .as_ref() - .map(|f| { - let schema = f - .schema - .as_ref() - .ok_or_else(|| proto_error("Missing JoinFilter schema"))? - .try_into()?; - - let expression = proto_converter.proto_to_physical_expr( - f.expression.as_ref().ok_or_else(|| { - proto_error("Unexpected empty filter expression") - })?, - &schema, - ctx, - )?; - let column_indices = f - .column_indices - .iter() - .map(|i| { - let side = - protobuf::JoinSide::try_from(i.side).map_err(|_| { - proto_error(format!( - "Received a SortMergeJoinExecNode message with JoinSide in Filter {}", - i.side - )) - })?; - - Ok(ColumnIndex { - index: i.index as usize, - side: side.into(), - }) - }) - .collect::>>()?; - - Ok(JoinFilter::new( - expression, - column_indices, - Arc::new(schema), - )) - }) - .map_or(Ok(None), |v: Result| v.map(Some))?; - - let join_type = - protobuf::JoinType::try_from(sort_join.join_type).map_err(|_| { - proto_error(format!( - "Received a SortMergeJoinExecNode message with unknown JoinType {}", - sort_join.join_type - )) - })?; - - let null_equality = protobuf::NullEquality::try_from(sort_join.null_equality) - .map_err(|_| { - proto_error(format!( - "Received a SortMergeJoinExecNode message with unknown NullEquality {}", - sort_join.null_equality - )) - })?; - - let sort_options = sort_join - .sort_options - .iter() - .map(|e| SortOptions { - descending: !e.asc, - nulls_first: e.nulls_first, - }) - .collect(); - let on = sort_join - .on - .iter() - .map(|col| { - let left = proto_converter.proto_to_physical_expr( - &col.left.clone().unwrap(), - left_schema.as_ref(), - ctx, - )?; - let right = proto_converter.proto_to_physical_expr( - &col.right.clone().unwrap(), - right_schema.as_ref(), - ctx, - )?; - Ok((left, right)) - }) - .collect::>()?; - - Ok(Arc::new(SortMergeJoinExec::try_new( - left, - right, - on, - filter, - JoinType::from_proto(join_type), - sort_options, - NullEquality::from_proto(null_equality), - )?)) - } - - fn try_into_generate_series_physical_plan( - &self, - generate_series: &protobuf::GenerateSeriesNode, - ) -> Result> { - let schema: SchemaRef = Arc::new(convert_required!(generate_series.schema)?); - - let args = match &generate_series.args { - Some(protobuf::generate_series_node::Args::ContainsNull(args)) => { - GenSeriesArgs::ContainsNull { - name: protobuf::PhysicalPlanNode::generate_series_name_to_str( - args.name(), - ), - } - } - Some(protobuf::generate_series_node::Args::Int64Args(args)) => { - GenSeriesArgs::Int64Args { - start: args.start, - end: args.end, - step: args.step, - include_end: args.include_end, - name: protobuf::PhysicalPlanNode::generate_series_name_to_str( - args.name(), - ), - } - } - Some(protobuf::generate_series_node::Args::TimestampArgs(args)) => { - let step_proto = args.step.as_ref().ok_or_else(|| { - internal_datafusion_err!("Missing step in TimestampArgs") - })?; - let step = IntervalMonthDayNanoType::make_value( - step_proto.months, - step_proto.days, - step_proto.nanos, - ); - GenSeriesArgs::TimestampArgs { - start: args.start, - end: args.end, - step, - tz: args.tz.as_ref().map(|s| Arc::from(s.as_str())), - include_end: args.include_end, - name: protobuf::PhysicalPlanNode::generate_series_name_to_str( - args.name(), - ), - } - } - Some(protobuf::generate_series_node::Args::DateArgs(args)) => { - let step_proto = args.step.as_ref().ok_or_else(|| { - internal_datafusion_err!("Missing step in DateArgs") - })?; - let step = IntervalMonthDayNanoType::make_value( - step_proto.months, - step_proto.days, - step_proto.nanos, - ); - GenSeriesArgs::DateArgs { - start: args.start, - end: args.end, - step, - include_end: args.include_end, - name: protobuf::PhysicalPlanNode::generate_series_name_to_str( - args.name(), - ), - } - } - None => return internal_err!("Missing args in GenerateSeriesNode"), - }; - - let table = GenerateSeriesTable::new(Arc::clone(&schema), args); - let generator = table.as_generator(generate_series.target_batch_size as usize)?; - - Ok(Arc::new(LazyMemoryExec::try_new(schema, vec![generator])?)) - } - - fn try_into_cooperative_physical_plan( - &self, - field_stream: &protobuf::CooperativeExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input = into_physical_plan(&field_stream.input, ctx, proto_converter)?; - Ok(Arc::new(CooperativeExec::new(input))) - } - - fn try_into_async_func_physical_plan( - &self, - async_func: &protobuf::AsyncFuncExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&async_func.input, ctx, proto_converter)?; - - if async_func.async_exprs.len() != async_func.async_expr_names.len() { - return internal_err!( - "AsyncFuncExecNode async_exprs length does not match async_expr_names" - ); - } - - let async_exprs = async_func - .async_exprs - .iter() - .zip(async_func.async_expr_names.iter()) - .map(|(expr, name)| { - let physical_expr = proto_converter.proto_to_physical_expr( - expr, - input.schema().as_ref(), - ctx, - )?; - - Ok(Arc::new(AsyncFuncExpr::try_new( - name.clone(), - physical_expr, - input.schema().as_ref(), - )?)) - }) - .collect::>>()?; - - Ok(Arc::new(AsyncFuncExec::try_new(async_exprs, input)?)) - } - - fn try_into_buffer_physical_plan( - &self, - buffer: &protobuf::BufferExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&buffer.input, ctx, proto_converter)?; - - Ok(Arc::new(BufferExec::new(input, buffer.capacity as usize))) - } - - fn try_into_scalar_subquery_physical_plan( - &self, - sq: &protobuf::ScalarSubqueryExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - // First, deserialize the main input plan. We set up the subquery results - // container first, so that ScalarSubqueryExpr nodes can reference it. - let subquery_results = ScalarSubqueryResults::new(sq.subqueries.len()); - let input_ctx = ctx.with_scalar_subquery_results(subquery_results.clone()); - let input = into_physical_plan(&sq.input, &input_ctx, proto_converter)?; - - // Now deserialize the subquery children. - let subqueries: Vec = sq - .subqueries - .iter() - .enumerate() - .map(|(index, sq_plan)| { - let plan = - sq_plan.try_into_physical_plan_with_context(ctx, proto_converter)?; - Ok(ScalarSubqueryLink { - plan, - index: SubqueryIndex::new(index), - }) - }) - .collect::>>()?; - - Ok(Arc::new(ScalarSubqueryExec::new( - input, - subqueries, - subquery_results, - ))) - } - - fn try_from_explain_exec( - exec: &ExplainExec, - _codec: &dyn PhysicalExtensionCodec, - ) -> Result { - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Explain( - protobuf::ExplainExecNode { - schema: Some(exec.schema().as_ref().try_into()?), - stringified_plans: exec - .stringified_plans() - .iter() - .map(protobuf::StringifiedPlan::from_proto) - .collect(), - verbose: exec.verbose(), - }, - )), - }) - } - - fn try_from_projection_exec( - exec: &ProjectionExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; - let expr = exec - .expr() - .iter() - .map(|proj_expr| { - proto_converter.physical_expr_to_proto(&proj_expr.expr, codec) - }) - .collect::>>()?; - let expr_name = exec - .expr() - .iter() - .map(|proj_expr| proj_expr.alias.clone()) - .collect(); - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Projection(Box::new( - protobuf::ProjectionExecNode { - input: Some(Box::new(input)), - expr, - expr_name, - }, - ))), - }) - } - - fn try_from_analyze_exec( - exec: &AnalyzeExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; - let (has_metric_categories, metric_categories) = match exec.metric_categories() { - Some(cats) => (true, cats.iter().map(|c| c.to_string()).collect()), - None => (false, vec![]), - }; - let format = match exec.format() { - ExplainFormat::Indent => protobuf::ExplainFormat::Indent, - ExplainFormat::Tree => protobuf::ExplainFormat::Tree, - ExplainFormat::PostgresJSON => protobuf::ExplainFormat::Pgjson, - ExplainFormat::Graphviz => protobuf::ExplainFormat::Graphviz, - } as i32; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Analyze(Box::new( - protobuf::AnalyzeExecNode { - verbose: exec.verbose(), - show_statistics: exec.show_statistics(), - input: Some(Box::new(input)), - schema: Some(exec.schema().as_ref().try_into()?), - has_metric_categories, - metric_categories, - format, - }, - ))), - }) - } - - fn try_from_filter_exec( - exec: &FilterExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Filter(Box::new( - protobuf::FilterExecNode { - input: Some(Box::new(input)), - expr: Some( - proto_converter - .physical_expr_to_proto(exec.predicate(), codec)?, - ), - default_filter_selectivity: exec.default_selectivity() as u32, - projection: match exec.projection() { - None => (0..exec.input().schema().fields().len()) - .map(|i| i as u32) - .collect(), - Some(v) => v.iter().map(|x| *x as u32).collect(), - }, - batch_size: exec.batch_size() as u32, - fetch: exec.fetch().map(|f| f as u32), - }, - ))), - }) - } - - fn try_from_global_limit_exec( - limit: &GlobalLimitExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - limit.input().to_owned(), - codec, - proto_converter, - )?; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::GlobalLimit(Box::new( - protobuf::GlobalLimitExecNode { - input: Some(Box::new(input)), - skip: limit.skip() as u32, - fetch: match limit.fetch() { - Some(n) => n as i64, - _ => -1, // no limit - }, - }, - ))), - }) - } - - fn try_from_local_limit_exec( - limit: &LocalLimitExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - limit.input().to_owned(), - codec, - proto_converter, - )?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::LocalLimit(Box::new( - protobuf::LocalLimitExecNode { - input: Some(Box::new(input)), - fetch: limit.fetch() as u32, - }, - ))), - }) - } - - fn try_from_hash_join_exec( - exec: &HashJoinExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let left = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.left().to_owned(), - codec, - proto_converter, - )?; - let right = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.right().to_owned(), - codec, - proto_converter, - )?; - let on: Vec = exec - .on() - .iter() - .map(|tuple| { - let l = proto_converter.physical_expr_to_proto(&tuple.0, codec)?; - let r = proto_converter.physical_expr_to_proto(&tuple.1, codec)?; - Ok::<_, DataFusionError>(protobuf::JoinOn { - left: Some(l), - right: Some(r), - }) - }) - .collect::>()?; - let join_type = protobuf::JoinType::from_proto(exec.join_type().to_owned()); - let null_equality = protobuf::NullEquality::from_proto(exec.null_equality()); - let filter = exec - .filter() - .as_ref() - .map(|f| { - let expression = - proto_converter.physical_expr_to_proto(f.expression(), codec)?; - let column_indices = f - .column_indices() - .iter() - .map(|i| { - let side: protobuf::JoinSide = i.side.to_owned().into(); - protobuf::ColumnIndex { - index: i.index as u32, - side: side.into(), - } - }) - .collect(); - let schema = f.schema().as_ref().try_into()?; - Ok(protobuf::JoinFilter { - expression: Some(expression), - column_indices, - schema: Some(schema), - }) - }) - .map_or(Ok(None), |v: Result| v.map(Some))?; - - let partition_mode = match exec.partition_mode() { - PartitionMode::CollectLeft => protobuf::PartitionMode::CollectLeft, - PartitionMode::Partitioned => protobuf::PartitionMode::Partitioned, - PartitionMode::Auto => protobuf::PartitionMode::Auto, - }; - - let dynamic_filter = exec - .dynamic_filter_expr() - .map(|df| { - let df_expr: Arc = - Arc::clone(df) as Arc; - proto_converter.physical_expr_to_proto(&df_expr, codec) - }) - .transpose()?; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::HashJoin(Box::new( - protobuf::HashJoinExecNode { - left: Some(Box::new(left)), - right: Some(Box::new(right)), - on, - join_type: join_type.into(), - partition_mode: partition_mode.into(), - null_equality: null_equality.into(), - filter, - // Send `Some(vec![])` as `[u32::MAX]` (never a valid index) so the - // wire format can distinguish it from `None` (which stays empty). - // See `try_into_hash_join_physical_plan` for the matching decoder. - projection: match exec.projection.as_ref() { - None => Vec::new(), - Some(v) if v.is_empty() => vec![u32::MAX], - Some(v) => v.iter().map(|x| *x as u32).collect(), - }, - null_aware: exec.null_aware, - dynamic_filter, - }, - ))), - }) - } - - fn try_from_symmetric_hash_join_exec( - exec: &SymmetricHashJoinExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let left = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.left().to_owned(), - codec, - proto_converter, - )?; - let right = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.right().to_owned(), - codec, - proto_converter, - )?; - let on = exec - .on() - .iter() - .map(|tuple| { - let l = proto_converter.physical_expr_to_proto(&tuple.0, codec)?; - let r = proto_converter.physical_expr_to_proto(&tuple.1, codec)?; - Ok::<_, DataFusionError>(protobuf::JoinOn { - left: Some(l), - right: Some(r), - }) - }) - .collect::>()?; - let join_type = protobuf::JoinType::from_proto(exec.join_type().to_owned()); - let null_equality = protobuf::NullEquality::from_proto(exec.null_equality()); - let filter = exec - .filter() - .as_ref() - .map(|f| { - let expression = - proto_converter.physical_expr_to_proto(f.expression(), codec)?; - let column_indices = f - .column_indices() - .iter() - .map(|i| { - let side: protobuf::JoinSide = i.side.to_owned().into(); - protobuf::ColumnIndex { - index: i.index as u32, - side: side.into(), - } - }) - .collect(); - let schema = f.schema().as_ref().try_into()?; - Ok(protobuf::JoinFilter { - expression: Some(expression), - column_indices, - schema: Some(schema), - }) - }) - .map_or(Ok(None), |v: Result| v.map(Some))?; - - let partition_mode = match exec.partition_mode() { - StreamJoinPartitionMode::SinglePartition => { - protobuf::StreamPartitionMode::SinglePartition - } - StreamJoinPartitionMode::Partitioned => { - protobuf::StreamPartitionMode::PartitionedExec - } - }; - - let left_sort_exprs = exec - .left_sort_exprs() - .map(|exprs| { - exprs - .iter() - .map(|expr| { - Ok(protobuf::PhysicalSortExprNode { - expr: Some(Box::new( - proto_converter - .physical_expr_to_proto(&expr.expr, codec)?, - )), - asc: !expr.options.descending, - nulls_first: expr.options.nulls_first, - }) - }) - .collect::>>() - }) - .transpose()? - .unwrap_or(vec![]); - - let right_sort_exprs = exec - .right_sort_exprs() - .map(|exprs| { - exprs - .iter() - .map(|expr| { - Ok(protobuf::PhysicalSortExprNode { - expr: Some(Box::new( - proto_converter - .physical_expr_to_proto(&expr.expr, codec)?, - )), - asc: !expr.options.descending, - nulls_first: expr.options.nulls_first, - }) - }) - .collect::>>() - }) - .transpose()? - .unwrap_or(vec![]); - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::SymmetricHashJoin(Box::new( - protobuf::SymmetricHashJoinExecNode { - left: Some(Box::new(left)), - right: Some(Box::new(right)), - on, - join_type: join_type.into(), - partition_mode: partition_mode.into(), - null_equality: null_equality.into(), - left_sort_exprs, - right_sort_exprs, - filter, - }, - ))), - }) - } - - fn try_from_sort_merge_join_exec( - exec: &SortMergeJoinExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let left = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.left().to_owned(), - codec, - proto_converter, - )?; - let right = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.right().to_owned(), - codec, - proto_converter, - )?; - let on = exec - .on() - .iter() - .map(|tuple| { - let l = proto_converter.physical_expr_to_proto(&tuple.0, codec)?; - let r = proto_converter.physical_expr_to_proto(&tuple.1, codec)?; - Ok::<_, DataFusionError>(protobuf::JoinOn { - left: Some(l), - right: Some(r), - }) - }) - .collect::>()?; - let join_type = protobuf::JoinType::from_proto(exec.join_type().to_owned()); - let null_equality = protobuf::NullEquality::from_proto(exec.null_equality()); - let filter = exec - .filter() - .as_ref() - .map(|f| { - let expression = - proto_converter.physical_expr_to_proto(f.expression(), codec)?; - let column_indices = f - .column_indices() - .iter() - .map(|i| { - let side: protobuf::JoinSide = i.side.to_owned().into(); - protobuf::ColumnIndex { - index: i.index as u32, - side: side.into(), - } - }) - .collect(); - let schema = f.schema().as_ref().try_into()?; - Ok(protobuf::JoinFilter { - expression: Some(expression), - column_indices, - schema: Some(schema), - }) - }) - .map_or(Ok(None), |v: Result| v.map(Some))?; - - let sort_options = exec - .sort_options() - .iter() - .map( - |SortOptions { - descending, - nulls_first, - }| { - SortExprNode { - expr: None, - asc: !*descending, - nulls_first: *nulls_first, - } - }, - ) - .collect(); - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::SortMergeJoin(Box::new( - SortMergeJoinExecNode { - left: Some(Box::new(left)), - right: Some(Box::new(right)), - on, - join_type: join_type.into(), - null_equality: null_equality.into(), - filter, - sort_options, - }, - ))), - }) - } - - fn try_from_cross_join_exec( - exec: &CrossJoinExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let left = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.left().to_owned(), - codec, - proto_converter, - )?; - let right = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.right().to_owned(), - codec, - proto_converter, - )?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::CrossJoin(Box::new( - protobuf::CrossJoinExecNode { - left: Some(Box::new(left)), - right: Some(Box::new(right)), - }, - ))), - }) - } - - fn try_from_aggregate_exec( - exec: &AggregateExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let groups: Vec = exec - .group_expr() - .groups() - .iter() - .flatten() - .copied() - .collect(); - - let group_names = exec - .group_expr() - .expr() - .iter() - .map(|expr| expr.1.to_owned()) - .collect(); - - let filter = exec - .filter_expr() - .iter() - .map(|expr| serialize_maybe_filter(expr.to_owned(), codec, proto_converter)) - .collect::>>()?; - - let agg = exec - .aggr_expr() - .iter() - .map(|expr| { - serialize_physical_aggr_expr(expr.to_owned(), codec, proto_converter) - }) - .collect::>>()?; - - let agg_names = exec - .aggr_expr() - .iter() - .map(|expr| expr.name().to_string()) - .collect::>(); - - let agg_mode = match exec.mode() { - AggregateMode::Partial => protobuf::AggregateMode::Partial, - AggregateMode::Final => protobuf::AggregateMode::Final, - AggregateMode::FinalPartitioned => protobuf::AggregateMode::FinalPartitioned, - AggregateMode::Single => protobuf::AggregateMode::Single, - AggregateMode::SinglePartitioned => { - protobuf::AggregateMode::SinglePartitioned - } - AggregateMode::PartialReduce => protobuf::AggregateMode::PartialReduce, - }; - let input_schema = exec.input_schema(); - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; - - let null_expr = exec - .group_expr() - .null_expr() - .iter() - .map(|expr| proto_converter.physical_expr_to_proto(&expr.0, codec)) - .collect::>>()?; - - let group_expr = exec - .group_expr() - .expr() - .iter() - .map(|expr| proto_converter.physical_expr_to_proto(&expr.0, codec)) - .collect::>>()?; - - let limit = exec.limit_options().map(|config| protobuf::AggLimit { - limit: config.limit() as u64, - descending: config.descending(), - }); - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Aggregate(Box::new( - protobuf::AggregateExecNode { - group_expr, - group_expr_name: group_names, - aggr_expr: agg, - filter_expr: filter, - aggr_expr_name: agg_names, - mode: agg_mode as i32, - input: Some(Box::new(input)), - input_schema: Some(input_schema.as_ref().try_into()?), - null_expr, - groups, - limit, - has_grouping_set: exec.group_expr().has_grouping_set(), - dynamic_filter: exec - .dynamic_filter_expr() - .map(|df| { - let df_expr: Arc = - Arc::clone(df) as Arc; - proto_converter.physical_expr_to_proto(&df_expr, codec) - }) - .transpose()?, - }, - ))), - }) - } - - fn try_from_empty_exec( - empty: &EmptyExec, - _codec: &dyn PhysicalExtensionCodec, - ) -> Result { - let schema = empty.schema().as_ref().try_into()?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Empty(protobuf::EmptyExecNode { - schema: Some(schema), - })), - }) - } - - fn try_from_placeholder_row_exec( - empty: &PlaceholderRowExec, - _codec: &dyn PhysicalExtensionCodec, - ) -> Result { - let schema = empty.schema().as_ref().try_into()?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::PlaceholderRow( - protobuf::PlaceholderRowExecNode { - schema: Some(schema), - }, - )), - }) - } - - #[expect(deprecated)] - fn try_from_coalesce_batches_exec( - coalesce_batches: &CoalesceBatchesExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - coalesce_batches.input().to_owned(), - codec, - proto_converter, - )?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::CoalesceBatches(Box::new( - protobuf::CoalesceBatchesExecNode { - input: Some(Box::new(input)), - target_batch_size: coalesce_batches.target_batch_size() as u32, - fetch: coalesce_batches.fetch().map(|n| n as u32), - }, - ))), - }) - } - - fn try_from_data_source_exec( - data_source_exec: &DataSourceExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let data_source = data_source_exec.data_source(); - if let Some(maybe_csv) = data_source.downcast_ref::() { - let source = maybe_csv.file_source(); - if let Some(csv_config) = source.downcast_ref::() { - return Ok(Some(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::CsvScan( - protobuf::CsvScanExecNode { - base_conf: Some(serialize_file_scan_config( - maybe_csv, - codec, - proto_converter, - )?), - has_header: csv_config.has_header(), - delimiter: byte_to_string( - csv_config.delimiter(), - "delimiter", - )?, - quote: byte_to_string(csv_config.quote(), "quote")?, - optional_escape: if let Some(escape) = csv_config.escape() { - Some( - protobuf::csv_scan_exec_node::OptionalEscape::Escape( - byte_to_string(escape, "escape")?, - ), - ) - } else { - None - }, - optional_comment: if let Some(comment) = csv_config.comment() - { - Some(protobuf::csv_scan_exec_node::OptionalComment::Comment( - byte_to_string(comment, "comment")?, - )) - } else { - None - }, - newlines_in_values: csv_config.newlines_in_values(), - truncate_rows: csv_config.truncate_rows(), - }, - )), - })); - } - } - - if let Some(scan_conf) = data_source.downcast_ref::() { - let source = scan_conf.file_source(); - if let Some(_json_source) = source.downcast_ref::() { - return Ok(Some(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::JsonScan( - protobuf::JsonScanExecNode { - base_conf: Some(serialize_file_scan_config( - scan_conf, - codec, - proto_converter, - )?), - }, - )), - })); - } - } - - if let Some(scan_conf) = data_source.downcast_ref::() { - let source = scan_conf.file_source(); - if let Some(_arrow_source) = source.downcast_ref::() { - return Ok(Some(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::ArrowScan( - protobuf::ArrowScanExecNode { - base_conf: Some(serialize_file_scan_config( - scan_conf, - codec, - proto_converter, - )?), - }, - )), - })); - } - } - - #[cfg(feature = "parquet")] - if let Some((maybe_parquet, conf)) = - data_source_exec.downcast_to_file_source::() - { - let predicate = conf - .filter() - .map(|pred| proto_converter.physical_expr_to_proto(&pred, codec)) - .transpose()?; - return Ok(Some(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::ParquetScan( - protobuf::ParquetScanExecNode { - base_conf: Some(serialize_file_scan_config( - maybe_parquet, - codec, - proto_converter, - )?), - predicate, - parquet_options: Some(conf.table_parquet_options().try_into()?), - }, - )), - })); - } - - #[cfg(feature = "avro")] - if let Some(maybe_avro) = data_source.downcast_ref::() { - let source = maybe_avro.file_source(); - if source.downcast_ref::().is_some() { - return Ok(Some(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::AvroScan( - protobuf::AvroScanExecNode { - base_conf: Some(serialize_file_scan_config( - maybe_avro, - codec, - proto_converter, - )?), - }, - )), - })); - } - } - - if let Some(source_conf) = data_source.downcast_ref::() { - let proto_partitions = source_conf - .partitions() - .iter() - .map(|p| serialize_record_batches(p)) - .collect::>>()?; - - let proto_schema: protobuf::Schema = - source_conf.original_schema().as_ref().try_into()?; - - let proto_projection = source_conf - .projection() - .as_ref() - .map_or_else(Vec::new, |v| { - v.iter().map(|x| *x as u32).collect::>() - }); - - let proto_sort_information = source_conf - .sort_information() - .iter() - .map(|ordering| { - let sort_exprs = serialize_physical_sort_exprs( - ordering.to_owned(), - codec, - proto_converter, - )?; - Ok::<_, DataFusionError>(protobuf::PhysicalSortExprNodeCollection { - physical_sort_expr_nodes: sort_exprs, - }) - }) - .collect::, _>>()?; - - return Ok(Some(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::MemoryScan( - protobuf::MemoryScanExecNode { - partitions: proto_partitions, - schema: Some(proto_schema), - projection: proto_projection, - sort_information: proto_sort_information, - show_sizes: source_conf.show_sizes(), - fetch: source_conf.fetch().map(|f| f as u32), - }, - )), - })); - } - - Ok(None) - } - - fn try_from_coalesce_partitions_exec( - exec: &CoalescePartitionsExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Merge(Box::new( - protobuf::CoalescePartitionsExecNode { - input: Some(Box::new(input)), - fetch: exec.fetch().map(|f| f as u32), - }, - ))), - }) - } - - fn try_from_repartition_exec( - exec: &RepartitionExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; - - let pb_partitioning = - serialize_partitioning(exec.partitioning(), codec, proto_converter)?; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Repartition(Box::new( - protobuf::RepartitionExecNode { - input: Some(Box::new(input)), - partitioning: Some(pb_partitioning), - preserve_order: exec.preserve_order(), - }, - ))), - }) - } - - fn try_from_sort_exec( - exec: &SortExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = proto_converter.execution_plan_to_proto(exec.input(), codec)?; - let expr = exec - .expr() - .iter() - .map(|expr| { - let sort_expr = Box::new(protobuf::PhysicalSortExprNode { - expr: Some(Box::new( - proto_converter.physical_expr_to_proto(&expr.expr, codec)?, - )), - asc: !expr.options.descending, - nulls_first: expr.options.nulls_first, - }); - Ok(protobuf::PhysicalExprNode { - expr_id: None, - expr_type: Some(ExprType::Sort(sort_expr)), - }) - }) - .collect::>>()?; - let dynamic_filter = exec - .dynamic_filter_expr() - .map(|df| { - let df_expr: Arc = df as Arc; - proto_converter.physical_expr_to_proto(&df_expr, codec) - }) - .transpose()?; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Sort(Box::new( - protobuf::SortExecNode { - input: Some(Box::new(input)), - expr, - fetch: match exec.fetch() { - Some(n) => n as i64, - _ => -1, - }, - preserve_partitioning: exec.preserve_partitioning(), - dynamic_filter, - }, - ))), - }) - } - - fn try_from_union_exec( - union: &UnionExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let mut inputs: Vec = vec![]; - for input in union.inputs() { - inputs.push( - protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - input.to_owned(), - codec, - proto_converter, - )?, - ); - } - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Union(protobuf::UnionExecNode { - inputs, - })), - }) - } - - fn try_from_interleave_exec( - interleave: &InterleaveExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let mut inputs: Vec = vec![]; - for input in interleave.inputs() { - inputs.push( - protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - input.to_owned(), - codec, - proto_converter, - )?, - ); - } - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Interleave( - protobuf::InterleaveExecNode { inputs }, - )), - }) - } - - fn try_from_sort_preserving_merge_exec( - exec: &SortPreservingMergeExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; - let expr = exec - .expr() - .iter() - .map(|expr| { - let sort_expr = Box::new(protobuf::PhysicalSortExprNode { - expr: Some(Box::new( - proto_converter.physical_expr_to_proto(&expr.expr, codec)?, - )), - asc: !expr.options.descending, - nulls_first: expr.options.nulls_first, - }); - Ok(protobuf::PhysicalExprNode { - expr_id: None, - expr_type: Some(ExprType::Sort(sort_expr)), - }) - }) - .collect::>>()?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::SortPreservingMerge(Box::new( - protobuf::SortPreservingMergeExecNode { - input: Some(Box::new(input)), - expr, - fetch: exec.fetch().map(|f| f as i64).unwrap_or(-1), - }, - ))), - }) - } - - fn try_from_nested_loop_join_exec( - exec: &NestedLoopJoinExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let left = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.left().to_owned(), - codec, - proto_converter, - )?; - let right = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.right().to_owned(), - codec, - proto_converter, - )?; - - let join_type = protobuf::JoinType::from_proto(exec.join_type().to_owned()); - let filter = exec - .filter() - .as_ref() - .map(|f| { - let expression = - proto_converter.physical_expr_to_proto(f.expression(), codec)?; - let column_indices = f - .column_indices() - .iter() - .map(|i| { - let side: protobuf::JoinSide = i.side.to_owned().into(); - protobuf::ColumnIndex { - index: i.index as u32, - side: side.into(), - } - }) - .collect(); - let schema = f.schema().as_ref().try_into()?; - Ok(protobuf::JoinFilter { - expression: Some(expression), - column_indices, - schema: Some(schema), - }) - }) - .map_or(Ok(None), |v: Result| v.map(Some))?; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::NestedLoopJoin(Box::new( - protobuf::NestedLoopJoinExecNode { - left: Some(Box::new(left)), - right: Some(Box::new(right)), - join_type: join_type.into(), - filter, - // `[u32::MAX]` sentinel distinguishes `Some(vec![])` from `None`; - // see `try_from_hash_join_exec`. - projection: match exec.projection().as_ref() { - None => Vec::new(), - Some(v) if v.is_empty() => vec![u32::MAX], - Some(v) => v.iter().map(|x| *x as u32).collect(), - }, - }, - ))), - }) - } - fn try_from_window_agg_exec( - exec: &WindowAggExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; + let null_equality = protobuf::NullEquality::try_from(sort_join.null_equality) + .map_err(|_| { + proto_error(format!( + "Received a SortMergeJoinExecNode message with unknown NullEquality {}", + sort_join.null_equality + )) + })?; - let window_expr = exec - .window_expr() + let sort_options = sort_join + .sort_options .iter() - .map(|e| serialize_physical_window_expr(e, codec, proto_converter)) - .collect::>>()?; - - let partition_keys = exec - .partition_keys() + .map(|e| SortOptions { + descending: !e.asc, + nulls_first: e.nulls_first, + }) + .collect(); + let on = sort_join + .on .iter() - .map(|e| proto_converter.physical_expr_to_proto(e, codec)) - .collect::>>()?; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Window(Box::new( - protobuf::WindowAggExecNode { - input: Some(Box::new(input)), - window_expr, - partition_keys, - input_order_mode: None, - }, - ))), - }) - } - - fn try_from_bounded_window_agg_exec( - exec: &BoundedWindowAggExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; + .map(|col| { + let left = proto_converter.proto_to_physical_expr( + &col.left.clone().unwrap(), + left_schema.as_ref(), + ctx, + )?; + let right = proto_converter.proto_to_physical_expr( + &col.right.clone().unwrap(), + right_schema.as_ref(), + ctx, + )?; + Ok((left, right)) + }) + .collect::>()?; - let window_expr = exec - .window_expr() - .iter() - .map(|e| serialize_physical_window_expr(e, codec, proto_converter)) - .collect::>>()?; + Ok(Arc::new(SortMergeJoinExec::try_new( + left, + right, + on, + filter, + JoinType::from_proto(join_type), + sort_options, + NullEquality::from_proto(null_equality), + )?)) + } - let partition_keys = exec - .partition_keys() - .iter() - .map(|e| proto_converter.physical_expr_to_proto(e, codec)) - .collect::>>()?; + fn try_into_generate_series_physical_plan( + &self, + generate_series: &protobuf::GenerateSeriesNode, + ) -> Result> { + let schema: SchemaRef = Arc::new(convert_required!(generate_series.schema)?); - let input_order_mode = match &exec.input_order_mode { - InputOrderMode::Linear => { - window_agg_exec_node::InputOrderMode::Linear(protobuf::EmptyMessage {}) + let args = match &generate_series.args { + Some(protobuf::generate_series_node::Args::ContainsNull(args)) => { + GenSeriesArgs::ContainsNull { + name: protobuf::PhysicalPlanNode::generate_series_name_to_str( + args.name(), + ), + } } - InputOrderMode::PartiallySorted(columns) => { - window_agg_exec_node::InputOrderMode::PartiallySorted( - protobuf::PartiallySortedInputOrderMode { - columns: columns.iter().map(|c| *c as u64).collect(), - }, - ) + Some(protobuf::generate_series_node::Args::Int64Args(args)) => { + GenSeriesArgs::Int64Args { + start: args.start, + end: args.end, + step: args.step, + include_end: args.include_end, + name: protobuf::PhysicalPlanNode::generate_series_name_to_str( + args.name(), + ), + } } - InputOrderMode::Sorted => { - window_agg_exec_node::InputOrderMode::Sorted(protobuf::EmptyMessage {}) + Some(protobuf::generate_series_node::Args::TimestampArgs(args)) => { + let step_proto = args.step.as_ref().ok_or_else(|| { + internal_datafusion_err!("Missing step in TimestampArgs") + })?; + let step = IntervalMonthDayNanoType::make_value( + step_proto.months, + step_proto.days, + step_proto.nanos, + ); + GenSeriesArgs::TimestampArgs { + start: args.start, + end: args.end, + step, + tz: args.tz.as_ref().map(|s| Arc::from(s.as_str())), + include_end: args.include_end, + name: protobuf::PhysicalPlanNode::generate_series_name_to_str( + args.name(), + ), + } } - }; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Window(Box::new( - protobuf::WindowAggExecNode { - input: Some(Box::new(input)), - window_expr, - partition_keys, - input_order_mode: Some(input_order_mode), - }, - ))), - }) - } - - fn try_from_data_sink_exec( - exec: &DataSinkExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: protobuf::PhysicalPlanNode = - protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; - let sort_order = match exec.sort_order() { - Some(requirements) => { - let expr = requirements - .iter() - .map(|requirement| { - let expr: PhysicalSortExpr = requirement.to_owned().into(); - let sort_expr = protobuf::PhysicalSortExprNode { - expr: Some(Box::new( - proto_converter - .physical_expr_to_proto(&expr.expr, codec)?, - )), - asc: !expr.options.descending, - nulls_first: expr.options.nulls_first, - }; - Ok(sort_expr) - }) - .collect::>>()?; - Some(protobuf::PhysicalSortExprNodeCollection { - physical_sort_expr_nodes: expr, - }) + Some(protobuf::generate_series_node::Args::DateArgs(args)) => { + let step_proto = args.step.as_ref().ok_or_else(|| { + internal_datafusion_err!("Missing step in DateArgs") + })?; + let step = IntervalMonthDayNanoType::make_value( + step_proto.months, + step_proto.days, + step_proto.nanos, + ); + GenSeriesArgs::DateArgs { + start: args.start, + end: args.end, + step, + include_end: args.include_end, + name: protobuf::PhysicalPlanNode::generate_series_name_to_str( + args.name(), + ), + } } - None => None, + None => return internal_err!("Missing args in GenerateSeriesNode"), }; - if let Some(sink) = exec.sink().downcast_ref::() { - return Ok(Some(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::JsonSink(Box::new( - protobuf::JsonSinkExecNode { - input: Some(Box::new(input)), - sink: Some(protobuf::JsonSink::try_from_proto(sink)?), - sink_schema: Some(exec.schema().as_ref().try_into()?), - sort_order, - }, - ))), - })); - } - - if let Some(sink) = exec.sink().downcast_ref::() { - return Ok(Some(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::CsvSink(Box::new( - protobuf::CsvSinkExecNode { - input: Some(Box::new(input)), - sink: Some(protobuf::CsvSink::try_from_proto(sink)?), - sink_schema: Some(exec.schema().as_ref().try_into()?), - sort_order, - }, - ))), - })); - } - - #[cfg(feature = "parquet")] - if let Some(sink) = exec.sink().downcast_ref::() { - return Ok(Some(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::ParquetSink(Box::new( - protobuf::ParquetSinkExecNode { - input: Some(Box::new(input)), - sink: Some(protobuf::ParquetSink::try_from_proto(sink)?), - sink_schema: Some(exec.schema().as_ref().try_into()?), - sort_order, - }, - ))), - })); - } - - // If unknown DataSink then let extension handle it - Ok(None) - } - - fn try_from_unnest_exec( - exec: &UnnestExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Unnest(Box::new( - protobuf::UnnestExecNode { - input: Some(Box::new(input)), - schema: Some(exec.schema().try_into()?), - list_type_columns: exec - .list_column_indices() - .iter() - .map(|c| ProtoListUnnest { - index_in_input_schema: c.index_in_input_schema as _, - depth: c.depth as _, - }) - .collect(), - struct_type_columns: exec - .struct_column_indices() - .iter() - .map(|c| *c as _) - .collect(), - options: Some(protobuf::UnnestOptions::from_proto(exec.options())), - }, - ))), - }) - } - - fn try_from_cooperative_exec( - exec: &CooperativeExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; + let table = GenerateSeriesTable::new(Arc::clone(&schema), args); + let generator = table.as_generator(generate_series.target_batch_size as usize)?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Cooperative(Box::new( - protobuf::CooperativeExecNode { - input: Some(Box::new(input)), - }, - ))), - }) + Ok(Arc::new(LazyMemoryExec::try_new(schema, vec![generator])?)) } - fn str_to_generate_series_name(name: &str) -> Result { match name { "generate_series" => Ok(protobuf::GenerateSeriesName::GsGenerateSeries), @@ -3805,90 +822,6 @@ pub trait PhysicalPlanNodeExt: Sized { Ok(None) } - - fn try_from_async_func_exec( - exec: &AsyncFuncExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - Arc::clone(exec.input()), - codec, - proto_converter, - )?; - - let mut async_exprs = vec![]; - let mut async_expr_names = vec![]; - - for async_expr in exec.async_exprs() { - async_exprs - .push(proto_converter.physical_expr_to_proto(&async_expr.func, codec)?); - async_expr_names.push(async_expr.name.clone()) - } - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::AsyncFunc(Box::new( - protobuf::AsyncFuncExecNode { - input: Some(Box::new(input)), - async_exprs, - async_expr_names, - }, - ))), - }) - } - - fn try_from_buffer_exec( - exec: &BufferExec, - extension_codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - Arc::clone(exec.input()), - extension_codec, - proto_converter, - )?; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Buffer(Box::new( - protobuf::BufferExecNode { - input: Some(Box::new(input)), - capacity: exec.capacity() as u64, - }, - ))), - }) - } - - fn try_from_scalar_subquery_exec( - exec: &ScalarSubqueryExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - Arc::clone(exec.input()), - codec, - proto_converter, - )?; - let subqueries = exec - .subqueries() - .iter() - .map(|sq| { - protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - Arc::clone(&sq.plan), - codec, - proto_converter, - ) - }) - .collect::>>()?; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::ScalarSubquery(Box::new( - protobuf::ScalarSubqueryExecNode { - input: Some(Box::new(input)), - subqueries, - }, - ))), - }) - } } impl PhysicalPlanNodeExt for protobuf::PhysicalPlanNode { @@ -4369,3 +1302,129 @@ fn into_physical_plan( Err(proto_error("Missing required field in protobuf")) } } + +/// Adapter backing [`ExecutionPlanEncodeCtx`] for plans migrated to the +/// `try_to_proto` hook (#22419). Routes child-plan and child-expr encoding back +/// through the central converter so nested plans honor their own hooks. +struct ConverterPlanEncoder<'a> { + codec: &'a dyn PhysicalExtensionCodec, + proto_converter: &'a dyn PhysicalProtoConverterExtension, +} + +impl ExecutionPlanEncode for ConverterPlanEncoder<'_> { + fn encode_plan( + &self, + plan: &Arc, + ) -> Result { + self.proto_converter + .execution_plan_to_proto(plan, self.codec) + } + + fn encode_expr( + &self, + expr: &Arc, + ) -> Result { + self.proto_converter + .physical_expr_to_proto(expr, self.codec) + } + + // Bytes-only function serde. `(!buf.is_empty()).then_some(buf)` preserves the + // existing `fun_definition` wire semantics (empty payload == encode-by-name). + fn encode_udf(&self, udf: &ScalarUDF) -> Result>> { + let mut buf = vec![]; + self.codec.try_encode_udf(udf, &mut buf)?; + Ok((!buf.is_empty()).then_some(buf)) + } + + fn encode_udaf(&self, udaf: &AggregateUDF) -> Result>> { + let mut buf = vec![]; + self.codec.try_encode_udaf(udaf, &mut buf)?; + Ok((!buf.is_empty()).then_some(buf)) + } + + fn encode_udwf(&self, udwf: &WindowUDF) -> Result>> { + let mut buf = vec![]; + self.codec.try_encode_udwf(udwf, &mut buf)?; + Ok((!buf.is_empty()).then_some(buf)) + } +} + +/// Adapter backing [`ExecutionPlanDecodeCtx`] for plans migrated to the +/// `try_from_proto` pattern (#22419). Routes child-plan and child-expr decoding +/// back through the central converter, and exposes the session task context +/// (never the extension codec). +struct ConverterPlanDecoder<'a, 'ctx> { + ctx: &'a PhysicalPlanDecodeContext<'ctx>, + proto_converter: &'a dyn PhysicalProtoConverterExtension, +} + +impl ExecutionPlanDecode for ConverterPlanDecoder<'_, '_> { + fn decode_plan( + &self, + node: &protobuf::PhysicalPlanNode, + ) -> Result> { + self.proto_converter.proto_to_execution_plan(node, self.ctx) + } + + fn decode_expr( + &self, + node: &protobuf::PhysicalExprNode, + input_schema: &Schema, + ) -> Result> { + self.proto_converter + .proto_to_physical_expr(node, input_schema, self.ctx) + } + + fn decode_plan_with_scalar_subquery_results( + &self, + node: &protobuf::PhysicalPlanNode, + results: ScalarSubqueryResults, + ) -> Result> { + let scoped = self.ctx.with_scalar_subquery_results(results); + self.proto_converter.proto_to_execution_plan(node, &scoped) + } + + fn task_ctx(&self) -> &TaskContext { + self.ctx.task_ctx() + } + + // Lookup-order policy, owned here so no plan re-derives it: an explicit + // payload is decoded by the codec; otherwise resolve by name from the + // registry, falling back to the codec with an empty buffer. + fn decode_udf(&self, name: &str, payload: Option<&[u8]>) -> Result> { + match payload { + Some(buf) => self.ctx.codec().try_decode_udf(name, buf), + None => self + .ctx + .task_ctx() + .udf(name) + .or_else(|_| self.ctx.codec().try_decode_udf(name, &[])), + } + } + + fn decode_udaf( + &self, + name: &str, + payload: Option<&[u8]>, + ) -> Result> { + match payload { + Some(buf) => self.ctx.codec().try_decode_udaf(name, buf), + None => self + .ctx + .task_ctx() + .udaf(name) + .or_else(|_| self.ctx.codec().try_decode_udaf(name, &[])), + } + } + + fn decode_udwf(&self, name: &str, payload: Option<&[u8]>) -> Result> { + match payload { + Some(buf) => self.ctx.codec().try_decode_udwf(name, buf), + None => self + .ctx + .task_ctx() + .udwf(name) + .or_else(|_| self.ctx.codec().try_decode_udwf(name, &[])), + } + } +} diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 8acb891d6a94d..ea2e758726d6a 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -35,9 +35,9 @@ use datafusion::datasource::listing::{ }; use datafusion::datasource::object_store::ObjectStoreUrl; use datafusion::datasource::physical_plan::{ - ArrowSource, FileGroup, FileOutputMode, FileScanConfig, FileScanConfigBuilder, - FileSinkConfig, ParquetSource, wrap_partition_type_in_dict, - wrap_partition_value_in_dict, + ArrowSource, CsvSource, FileGroup, FileOutputMode, FileScanConfig, + FileScanConfigBuilder, FileSinkConfig, JsonSource, ParquetSource, + wrap_partition_type_in_dict, wrap_partition_value_in_dict, }; use datafusion::datasource::sink::DataSinkExec; use datafusion::datasource::source::DataSourceExec; @@ -1148,6 +1148,70 @@ fn roundtrip_arrow_scan() -> Result<()> { roundtrip_test(DataSourceExec::from_data_source(scan_config)) } +// Exercises the JsonSource hook (#22419) — see roundtrip_csv_scan. +#[test] +fn roundtrip_json_scan() -> Result<()> { + let file_schema = + Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); + let file_source = Arc::new(JsonSource::new(TableSchema::from(&file_schema))); + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( + "/path/to/file.json".to_string(), + 1024, + )])]) + .build(); + roundtrip_test(DataSourceExec::from_data_source(scan_config)) +} + +// Exercises the AvroSource hook (#22419) — see roundtrip_csv_scan. +#[cfg(feature = "avro")] +#[test] +fn roundtrip_avro_scan() -> Result<()> { + use datafusion_datasource_avro::source::AvroSource; + + let file_schema = + Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); + let file_source = Arc::new(AvroSource::new(TableSchema::from(&file_schema))); + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( + "/path/to/file.avro".to_string(), + 1024, + )])]) + .build(); + roundtrip_test(DataSourceExec::from_data_source(scan_config)) +} + +// Exercises the CsvSource `try_to_proto` / `try_from_proto` hook (#22419): the +// central CSV encode/decode arms were removed, so a green roundtrip here proves +// the DataSource -> FileSource hook chain is reached in both directions. +#[test] +fn roundtrip_csv_scan() -> Result<()> { + use datafusion::common::config::CsvOptions; + + let file_schema = + Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); + let table_schema = TableSchema::from(&file_schema); + let file_source = + Arc::new(CsvSource::new(table_schema).with_csv_options(CsvOptions { + has_header: Some(true), + delimiter: b',', + quote: b'"', + ..Default::default() + })); + + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( + "/path/to/file.csv".to_string(), + 1024, + )])]) + .build(); + + roundtrip_test(DataSourceExec::from_data_source(scan_config)) +} + #[tokio::test] async fn roundtrip_parquet_exec_with_table_partition_cols() -> Result<()> { let mut file_group =