You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/08/09 02:37:52 UTC
(datafusion) branch gh-readonly-queue/main/pr-24187-cc61491adde99acc59f334c48398ae4ea0f8d3fd created (now abc5ce7e0e)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-24187-cc61491adde99acc59f334c48398ae4ea0f8d3fd
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at abc5ce7e0e Proto: migrate MemorySourceConfig to per-source try_to_proto / try_from_proto hooks (#24187)
This branch includes the following new commits:
new abc5ce7e0e Proto: migrate MemorySourceConfig to per-source try_to_proto / try_from_proto hooks (#24187)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: Proto: migrate MemorySourceConfig to per-source try_to_proto / try_from_proto hooks (#24187)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24187-cc61491adde99acc59f334c48398ae4ea0f8d3fd
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit abc5ce7e0ec74c113081d1a23fbe130c1ba95cd0
Author: Adrian Garcia Badaracco <17...@users.noreply.github.com>
AuthorDate: Sat Aug 8 21:37:29 2026 -0500
Proto: migrate MemorySourceConfig to per-source try_to_proto / try_from_proto hooks (#24187)
## Which issue does this PR close?
- Closes #23518.
- Part of #23494.
## Rationale for this change
Continues the per-plan proto hook migration (#23494, pattern established
in #23495 and #23683): each plan/source owns its wire logic via
`try_to_proto` / `try_from_proto` hooks instead of the central
downcast/match chains in `datafusion-proto`.
## What changes are included in this PR?
- `DataSource::try_to_proto` implemented for `MemorySourceConfig` and
inherent `MemorySourceConfig::try_from_proto` added in
`datafusion/datasource/src/memory.rs` (behind the `proto` feature). The
record-batch IPC serde is pure Arrow and is inlined locally.
- The `MemoryScan` decode arm in `datafusion-proto` now delegates to
`MemorySourceConfig::try_from_proto`, and the central `MemoryScan`
encode arm in `try_from_data_source_exec` is **deleted** — proof that
the hook is the only path.
- `try_into_memory_scan_physical_plan` becomes a deprecated shim
delegating to the new hook; the now-unused `serialize_record_batches` /
`parse_record_batches` helpers are deprecated.
The wire format is byte-for-byte identical: I verified locally that a
memory scan plan (with sort information, fetch, and `show_sizes=false`)
encodes to identical bytes before and after this change, and that bytes
encoded by each version decode correctly with the other.
## Are these changes tested?
Yes. Existing roundtrip tests (`roundtrip_memory_source`,
`roundtrip_memory_source_empty_projection`) now exercise the hooks since
the central arms are gone. A new test
`roundtrip_memory_source_sort_information_and_fetch` covers
`sort_information`, `fetch`, and `show_sizes`, asserting on the decoded
`MemorySourceConfig` directly rather than only the display string.
## Are there any user-facing changes?
No behavior changes.
`PhysicalPlanNodeExt::try_into_memory_scan_physical_plan`,
`parse_record_batches`, and `serialize_record_batches` are deprecated
(still functional).
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---------
Co-authored-by: Claude Fable 5 <no...@anthropic.com>
---
datafusion/datasource/src/memory.rs | 142 +++++++++++++++++++++
datafusion/proto/src/physical_plan/from_proto.rs | 4 +
datafusion/proto/src/physical_plan/mod.rs | 106 +++------------
datafusion/proto/src/physical_plan/to_proto.rs | 4 +
.../proto/tests/cases/roundtrip_physical_plan.rs | 53 ++++++++
5 files changed, 220 insertions(+), 89 deletions(-)
diff --git a/datafusion/datasource/src/memory.rs b/datafusion/datasource/src/memory.rs
index 255dd76cbd..15ea2600f1 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<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.
diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs
index 9a3a845d7c..b88c3cf28f 100644
--- a/datafusion/proto/src/physical_plan/from_proto.rs
+++ b/datafusion/proto/src/physical_plan/from_proto.rs
@@ -459,6 +459,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![]);
diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs
index ee2d431a41..3cfd7fc318 100644
--- a/datafusion/proto/src/physical_plan/mod.rs
+++ b/datafusion/proto/src/physical_plan/mod.rs
@@ -28,7 +28,7 @@ use datafusion_common::{
};
use datafusion_datasource::file_scan_config::FileScanConfig;
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;
@@ -46,7 +46,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;
@@ -88,12 +87,11 @@ use prost::bytes::BufMut;
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};
@@ -1095,8 +1093,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)
@@ -1454,48 +1452,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(
@@ -2515,53 +2490,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)
}
diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs
index c8a7ea383a..aa10e1c5aa 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<Vec<u8>> {
if batches.is_empty() {
return Ok(vec![]);
diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs
index 9efbd90f15..c3dcaa8814 100644
--- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs
+++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs
@@ -3526,6 +3526,59 @@ 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.clone());
+
+ 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.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());
+ Ok(())
+}
+
#[tokio::test]
async fn roundtrip_listing_table_with_schema_metadata() -> Result<()> {
let ctx = SessionContext::new();
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org