Proto: migrate MemorySourceConfig to per-source try_to_proto / try_from_proto hooks

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<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
+        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::<Result<Vec<_>>>()?;
+
+        // 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<Arc<dyn datafusion_physical_plan::ExecutionPlan>> {
+        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::<Result<Vec<_>>>()?;
+
+        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<Vec<u8>> {
+    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<Vec<RecordBatch>> {
+    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.

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<Vec<RecordBatch>> {
     if buf.is_empty() {
         return Ok(vec![]);

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<Arc<dyn ExecutionPlan>> {
-        let partitions = scan
-            .partitions
-            .iter()
-            .map(|p| parse_record_batches(p))
-            .collect::<Result<Vec<_>>>()?;
-
-        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::<MemorySourceConfig>() {
-            let proto_partitions = source_conf
-                .partitions()
-                .iter()
-                .map(|p| serialize_record_batches(p))
-                .collect::<Result<Vec<_>>>()?;
-
-            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::<Result<Vec<_>, _>>()?;
-
-            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)
     }
 

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<Vec<u8>> {
     if batches.is_empty() {
         return Ok(vec![]);

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::<DataSourceExec>()
+        .expect("expected DataSourceExec");
+    let decoded_source = decoded
+        .data_source()
+        .downcast_ref::<MemorySourceConfig>()
+        .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();