From dfac1fb77aa7968cbf939cd294b7dc068376ef45 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Sat, 8 Aug 2026 10:59:17 -0500 Subject: [PATCH 1/4] Add MemorySourceConfig proto hooks in datafusion-datasource Implement DataSource::try_to_proto for MemorySourceConfig and the inherent MemorySourceConfig::try_from_proto, moving the MemoryScan wire logic into the crate that owns the type. The record-batch IPC serde is pure Arrow and is inlined locally. Co-Authored-By: Claude Fable 5 --- datafusion/datasource/src/memory.rs | 142 ++++++++++++++++++++++++++++ 1 file changed, 142 insertions(+) diff --git a/datafusion/datasource/src/memory.rs b/datafusion/datasource/src/memory.rs index 255dd76cbd6b4..15ea2600f1a36 100644 --- a/datafusion/datasource/src/memory.rs +++ b/datafusion/datasource/src/memory.rs @@ -256,6 +256,58 @@ impl DataSource for MemorySourceConfig { }) .transpose() } + + /// Serialize this `MemorySourceConfig` as a `MemoryScanExecNode` wrapped + /// in a [`PhysicalPlanNode`]. Byte-compatible with the former central + /// `MemoryScan` arm in `datafusion-proto`. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr_common::sort_expr::sort_exprs_try_to_proto; + use datafusion_proto_models::protobuf; + + let partitions = self + .partitions + .iter() + .map(|batches| record_batches_to_ipc_bytes(batches)) + .collect::>>()?; + + // Proto3 can't tell `None` from `Some(vec![])`; encode the latter + // as the `[u32::MAX]` sentinel, matching the join/filter nodes. + let 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(), + }; + + let mut sort_information = Vec::with_capacity(self.sort_information.len()); + for ordering in &self.sort_information { + let physical_sort_expr_nodes = + sort_exprs_try_to_proto(ordering.iter(), &ctx.expr_ctx())?; + sort_information.push(protobuf::PhysicalSortExprNodeCollection { + physical_sort_expr_nodes, + }); + } + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::MemoryScan( + protobuf::MemoryScanExecNode { + partitions, + schema: Some(self.schema.as_ref().try_into()?), + projection, + sort_information, + show_sizes: self.show_sizes, + fetch: self.fetch.map(|f| f as u32), + }, + ), + ), + })) + } } impl MemorySourceConfig { @@ -607,6 +659,96 @@ impl MemorySourceConfig { } } +#[cfg(feature = "proto")] +impl MemorySourceConfig { + /// Reconstruct a [`DataSourceExec`] wrapping a `MemorySourceConfig` from + /// its protobuf representation. Byte-compatible with the former central + /// `MemoryScan` arm in `datafusion-proto`. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_common::internal_datafusion_err; + use datafusion_physical_expr_common::sort_expr::sort_exprs_try_from_proto; + use datafusion_proto_models::protobuf; + + let scan = datafusion_physical_plan::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::MemoryScan, + "MemorySourceConfig", + ); + + let partitions = scan + .partitions + .iter() + .map(|buf| record_batches_from_ipc_bytes(buf)) + .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()?); + + // Preserve the empty-projection sentinel written by `try_to_proto`. + let projection = match scan.projection.as_slice() { + [] => None, + [u32::MAX] => Some(Vec::new()), + indices => Some(indices.iter().map(|i| *i as usize).collect()), + }; + + let mut sort_information = vec![]; + for ordering in &scan.sort_information { + let sort_exprs = sort_exprs_try_from_proto( + &ordering.physical_sort_expr_nodes, + &ctx.expr_ctx(&schema), + )?; + 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) + .try_with_sort_information(sort_information)?; + + Ok(DataSourceExec::from_data_source(source)) + } +} + +/// Encode record batches as Arrow IPC stream bytes; an empty slice encodes to +/// an empty buffer. +#[cfg(feature = "proto")] +fn record_batches_to_ipc_bytes(batches: &[RecordBatch]) -> Result> { + use arrow::ipc::writer::StreamWriter; + + if batches.is_empty() { + return Ok(vec![]); + } + let schema = batches[0].schema(); + let mut buf = Vec::new(); + let mut writer = StreamWriter::try_new(&mut buf, &schema)?; + for batch in batches { + writer.write(batch)?; + } + writer.finish()?; + Ok(buf) +} + +/// Inverse of [`record_batches_to_ipc_bytes`]. +#[cfg(feature = "proto")] +fn record_batches_from_ipc_bytes(buf: &[u8]) -> Result> { + use arrow::ipc::reader::StreamReader; + + if buf.is_empty() { + return Ok(vec![]); + } + let reader = StreamReader::try_new(buf, None)?; + let mut batches = Vec::new(); + for batch in reader { + batches.push(batch?); + } + Ok(batches) +} + /// For use in repartitioning, track the total size and original partition index. /// /// Do not implement clone, in order to avoid unnecessary copying during repartitioning. From 7f600868158c03933dd77a62b959b91a0f658cdc Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Sat, 8 Aug 2026 10:59:23 -0500 Subject: [PATCH 2/4] Repoint MemoryScan serde to the MemorySourceConfig hooks Decode dispatches to MemorySourceConfig::try_from_proto; the central MemoryScan encode arm in try_from_data_source_exec is deleted so the DataSource::try_to_proto hook is the only path. The old decode helper becomes a deprecated shim delegating to the new hook, and the now-unused record-batch IPC helpers are deprecated. Co-Authored-By: Claude Fable 5 --- .../proto/src/physical_plan/from_proto.rs | 4 + datafusion/proto/src/physical_plan/mod.rs | 106 +++--------------- .../proto/src/physical_plan/to_proto.rs | 4 + 3 files changed, 25 insertions(+), 89 deletions(-) diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index 645854295bc00..1b0124d4f1b95 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -455,6 +455,10 @@ pub fn parse_protobuf_file_scan_config( ) } +#[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `MemorySourceConfig` deserializes its record batches itself via `MemorySourceConfig::try_from_proto`" +)] pub fn parse_record_batches(buf: &[u8]) -> Result> { if buf.is_empty() { return Ok(vec![]); diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 7e162bf95454a..2007b82c0fca2 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -32,7 +32,7 @@ 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::source::DataSourceExec; use datafusion_datasource_arrow::source::ArrowSource; #[cfg(feature = "avro")] use datafusion_datasource_avro::source::AvroSource; @@ -54,7 +54,6 @@ use datafusion_expr::{AggregateUDF, HigherOrderUDF, ScalarUDF, WindowUDF}; use datafusion_functions_table::generate_series::{ Empty, GenSeriesArgs, GenerateSeriesTable, GenericSeriesState, TimestampValue, }; -use datafusion_physical_expr::LexOrdering; use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; use datafusion_physical_plan::aggregates::AggregateExec; @@ -97,12 +96,11 @@ use prost::bytes::BufMut; use crate::common::{byte_to_string, str_to_byte}; use crate::convert_required; use crate::physical_plan::from_proto::{ - parse_physical_expr_with_converter, parse_physical_sort_exprs, - parse_protobuf_file_scan_config, parse_record_batches, parse_table_schema_from_proto, + parse_physical_expr_with_converter, parse_protobuf_file_scan_config, + parse_table_schema_from_proto, }; use crate::physical_plan::to_proto::{ serialize_file_scan_config, serialize_physical_expr_with_converter, - serialize_physical_sort_exprs, serialize_record_batches, }; use crate::protobuf::physical_plan_node::PhysicalPlanType; use crate::protobuf::{self, SortMergeJoinExecNode, proto_error}; @@ -1096,8 +1094,8 @@ pub trait PhysicalPlanNodeExt: Sized { PhysicalPlanType::AvroScan(scan) => { self.try_into_avro_scan_physical_plan(scan, ctx, proto_converter) } - 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) @@ -1536,48 +1534,25 @@ pub trait PhysicalPlanNodeExt: Sized { panic!("Unable to process a Avro PhysicalPlan when `avro` feature is not enabled") } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `MemorySourceConfig` deserializes itself via `MemorySourceConfig::try_from_proto`" + )] fn try_into_memory_scan_physical_plan( &self, scan: &protobuf::MemoryScanExecNode, ctx: &PhysicalPlanDecodeContext<'_>, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - let partitions = scan - .partitions - .iter() - .map(|p| 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()?); - - // Preserve the empty-projection sentinel written by `try_from_data_source_exec`. - let projection = match scan.projection.as_slice() { - [] => None, - [u32::MAX] => Some(Vec::new()), - indices => Some(indices.iter().map(|i| *i as usize).collect()), + let node = protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::MemoryScan(scan.clone())), }; - - 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)) + let decoder = ConverterPlanDecoder { + ctx, + proto_converter, + }; + let decode_ctx = ExecutionPlanDecodeCtx::new(&decoder); + MemorySourceConfig::try_from_proto(&node, &decode_ctx) } #[deprecated( @@ -2678,53 +2653,6 @@ pub trait PhysicalPlanNodeExt: Sized { } } - 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()?; - - // Proto3 can't tell `None` from `Some(vec![])`; encode the latter - // as the `[u32::MAX]` sentinel, matching the join/filter nodes. - let proto_projection = match source_conf.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(), - }; - - 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) } diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index c8a7ea383a69f..aa10e1c5aa91f 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -427,6 +427,10 @@ pub fn serialize_maybe_filter( } } +#[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `MemorySourceConfig` serializes its record batches itself via `DataSource::try_to_proto`" +)] pub fn serialize_record_batches(batches: &[RecordBatch]) -> Result> { if batches.is_empty() { return Ok(vec![]); From d37e23deb1de6f9389470013a88f88f4eb6d1ec9 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Sat, 8 Aug 2026 10:59:24 -0500 Subject: [PATCH 3/4] Add memory scan roundtrip test for sort information and fetch The display-string roundtrip does not cover every field, so the decoded source is downcast and its fetch, show_sizes and sort_information are asserted directly. Co-Authored-By: Claude Fable 5 --- .../tests/cases/roundtrip_physical_plan.rs | 50 +++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index afd4057d0457f..2531b085bfc7d 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -3456,6 +3456,56 @@ async fn roundtrip_memory_source() -> Result<()> { roundtrip_test(plan) } +#[tokio::test] +async fn roundtrip_memory_source_sort_information_and_fetch() -> Result<()> { + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::datasource::source::DataSource as _; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Utf8, false), + Field::new("b", DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(arrow::array::StringArray::from(vec!["Tom", "Bob"])), + Arc::new(arrow::array::Int64Array::from(vec![18i64, 21i64])), + ], + )?; + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new( + col("b", &schema)?, + SortOptions { + descending: true, + nulls_first: false, + }, + )]) + .unwrap(); + let source = MemorySourceConfig::try_new(&[vec![batch]], Arc::clone(&schema), None)? + .with_limit(Some(1)) + .with_show_sizes(false) + .try_with_sort_information(vec![ordering])?; + let exec_plan = DataSourceExec::from_data_source(source); + + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let decoded = roundtrip_test_and_return(exec_plan, &ctx, &codec, &proto_converter)?; + + // The string representation does not include every field; check the + // decoded source directly. + let decoded = decoded + .downcast_ref::() + .expect("expected DataSourceExec"); + let decoded_source = decoded + .data_source() + .downcast_ref::() + .expect("expected MemorySourceConfig"); + assert_eq!(decoded_source.fetch(), Some(1)); + assert!(!decoded_source.show_sizes()); + assert_eq!(decoded_source.sort_information().len(), 1); + Ok(()) +} + #[tokio::test] async fn roundtrip_listing_table_with_schema_metadata() -> Result<()> { let ctx = SessionContext::new(); From 2dabb54231ee22226ea4c175cc72d2af688b2bd4 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Sat, 8 Aug 2026 21:56:54 -0400 Subject: [PATCH 4/4] Compare all decoded MemorySourceConfig fields in roundtrip test Review feedback: assert equality of partitions, schema, projection and the complete sort information against the original source, not just counts. Co-Authored-By: Claude Fable 5 --- datafusion/proto/tests/cases/roundtrip_physical_plan.rs | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 2531b085bfc7d..66c40a1f9bbe5 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -3484,7 +3484,7 @@ async fn roundtrip_memory_source_sort_information_and_fetch() -> Result<()> { .with_limit(Some(1)) .with_show_sizes(false) .try_with_sort_information(vec![ordering])?; - let exec_plan = DataSourceExec::from_data_source(source); + let exec_plan = DataSourceExec::from_data_source(source.clone()); let ctx = SessionContext::new(); let codec = DefaultPhysicalExtensionCodec {}; @@ -3500,9 +3500,12 @@ async fn roundtrip_memory_source_sort_information_and_fetch() -> Result<()> { .data_source() .downcast_ref::() .expect("expected MemorySourceConfig"); + assert_eq!(decoded_source.partitions(), source.partitions()); + assert_eq!(decoded_source.original_schema(), source.original_schema()); + assert_eq!(decoded_source.projection(), source.projection()); + assert_eq!(decoded_source.sort_information(), source.sort_information()); assert_eq!(decoded_source.fetch(), Some(1)); assert!(!decoded_source.show_sizes()); - assert_eq!(decoded_source.sort_information().len(), 1); Ok(()) }