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/15 11:04:44 UTC
(datafusion) 01/01: fix(proto): preserve Arrow IPC stream format (#24224)
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-24224-b9399dc228697db8ef27b24b67b5a1ddf46ff073
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit f956fd1128d1ba419747dbd3ff10aac695ecf158
Author: Burak Şen <bu...@gmail.com>
AuthorDate: Sat Aug 15 11:04:19 2026 +0000
fix(proto): preserve Arrow IPC stream format (#24224)
## Which issue does this PR close?
- Closes #24196.
## Rationale for this change
`ArrowSource` distinguishes between Arrow IPC file and stream formats,
but this format was not serialized. As a result, a stream scan
round-tripped through protobuf as a file scan, selecting the wrong
opener and potentially allowing unsupported range-based repartitioning.
## What changes are included in this PR?
- Add an `ArrowIpcFormat` discriminator to `ArrowScanExecNode`.
- Serialize and restore the `ArrowSource` IPC format.
- Decode payloads without the new field as file format, preserving the
previous behavior.
- Regenerate the prost and pbjson models.
- Add regression coverage for file, stream, and older payloads without
the format field.
## Are these changes tested?
Yes:
- `cargo test -p datafusion-proto --test proto_integration
roundtrip_arrow`
- `cargo test -p datafusion-proto --test proto_integration
arrow_scan_without_format_field_decodes_as_file_format`
- `cargo fmt --all`
- Targeted all-feature clippy for `datafusion-datasource-arrow` and
`datafusion-proto`
## Are there any user-facing changes?
Arrow IPC stream scans now preserve their format across protobuf round
trips.
The binary protobuf change is additive. Older payloads continue to
decode as file format. Stream preservation requires both producer and
consumer to include this change because older versions do not carry or
read the discriminator. The generated Rust `ArrowScanExecNode` struct
gains a `format` field.
---
datafusion/datasource-arrow/src/source.rs | 32 ++++++--
datafusion/proto-models/proto/datafusion.proto | 10 +++
datafusion/proto-models/src/generated/pbjson.rs | 90 +++++++++++++++++++++
datafusion/proto-models/src/generated/prost.rs | 32 ++++++++
datafusion/proto/tests/cases/plans/sources.rs | 101 +++++++++++++++++++++++-
5 files changed, 256 insertions(+), 9 deletions(-)
diff --git a/datafusion/datasource-arrow/src/source.rs b/datafusion/datasource-arrow/src/source.rs
index dba99a9758..dd973c03ae 100644
--- a/datafusion/datasource-arrow/src/source.rs
+++ b/datafusion/datasource-arrow/src/source.rs
@@ -403,10 +403,8 @@ impl FileSource for ArrowSource {
datafusion_physical_plan::apply_expression_roots(self.projection.source.iter(), f)
}
- /// Emit an `ArrowScan` node wrapping the shared base config.
- ///
- /// Decoding defaults to the IPC file format because protobuf does not
- /// distinguish it from the IPC stream format.
+ /// Emit an `ArrowScan` node wrapping the shared base config and recording
+ /// which Arrow IPC format (file or stream) this source reads.
#[cfg(feature = "proto")]
fn try_to_proto(
&self,
@@ -416,10 +414,16 @@ impl FileSource for ArrowSource {
use datafusion_proto_models::protobuf;
use protobuf::physical_plan_node::PhysicalPlanType;
+ let format = match self.format {
+ ArrowFormat::File => protobuf::ArrowIpcFormat::File,
+ ArrowFormat::Stream => protobuf::ArrowIpcFormat::Stream,
+ };
+
Ok(Some(protobuf::PhysicalPlanNode {
physical_plan_type: Some(PhysicalPlanType::ArrowScan(
protobuf::ArrowScanExecNode {
base_conf: Some(base.try_to_proto(ctx)?),
+ format: format as i32,
},
)),
}))
@@ -430,8 +434,8 @@ impl FileSource for ArrowSource {
impl ArrowSource {
/// Reconstructs a `DataSourceExec` from a protobuf `ArrowScan`.
///
- /// Defaults to the IPC file format because protobuf does not distinguish it
- /// from the IPC stream format.
+ /// Payloads encoded before the `format` field existed leave it unset;
+ /// those decode as the IPC file format, matching the historical behavior.
pub fn try_from_proto(
node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>,
@@ -455,9 +459,21 @@ impl ArrowSource {
)
})?;
+ let format = protobuf::ArrowIpcFormat::try_from(scan.format).map_err(|_| {
+ datafusion_common::internal_datafusion_err!(
+ "Unknown ArrowIpcFormat: {}",
+ scan.format
+ )
+ })?;
+
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::try_from_proto(base_conf, ctx, source)?;
+ let source = match format {
+ protobuf::ArrowIpcFormat::Stream => {
+ ArrowSource::new_stream_file_source(table_schema)
+ }
+ protobuf::ArrowIpcFormat::File => ArrowSource::new_file_source(table_schema),
+ };
+ let scan_conf = FileScanConfig::try_from_proto(base_conf, ctx, Arc::new(source))?;
Ok(DataSourceExec::from_data_source(scan_conf))
}
}
diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto
index 43a90264c2..f29d63c476 100644
--- a/datafusion/proto-models/proto/datafusion.proto
+++ b/datafusion/proto-models/proto/datafusion.proto
@@ -1296,8 +1296,18 @@ message AvroScanExecNode {
FileScanExecConf base_conf = 1;
}
+// Identifies which Arrow IPC format an ArrowScanExecNode reads.
+enum ArrowIpcFormat {
+ // Arrow IPC file format (with footer, supports range-based parallel reading).
+ // This is the default for payloads encoded before the format field existed.
+ ARROW_IPC_FORMAT_FILE = 0;
+ // Arrow IPC stream format (without footer, sequential reading only)
+ ARROW_IPC_FORMAT_STREAM = 1;
+}
+
message ArrowScanExecNode {
FileScanExecConf base_conf = 1;
+ ArrowIpcFormat format = 2;
}
message MemoryScanExecNode {
diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs
index 908f9752b7..19121b2dbe 100644
--- a/datafusion/proto-models/src/generated/pbjson.rs
+++ b/datafusion/proto-models/src/generated/pbjson.rs
@@ -1445,6 +1445,77 @@ impl<'de> serde::Deserialize<'de> for AnalyzedLogicalPlanType {
deserializer.deserialize_struct("datafusion.AnalyzedLogicalPlanType", FIELDS, GeneratedVisitor)
}
}
+impl serde::Serialize for ArrowIpcFormat {
+ #[allow(deprecated)]
+ fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
+ where
+ S: serde::Serializer,
+ {
+ let variant = match self {
+ Self::File => "ARROW_IPC_FORMAT_FILE",
+ Self::Stream => "ARROW_IPC_FORMAT_STREAM",
+ };
+ serializer.serialize_str(variant)
+ }
+}
+impl<'de> serde::Deserialize<'de> for ArrowIpcFormat {
+ #[allow(deprecated)]
+ fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
+ where
+ D: serde::Deserializer<'de>,
+ {
+ const FIELDS: &[&str] = &[
+ "ARROW_IPC_FORMAT_FILE",
+ "ARROW_IPC_FORMAT_STREAM",
+ ];
+
+ struct GeneratedVisitor;
+
+ impl serde::de::Visitor<'_> for GeneratedVisitor {
+ type Value = ArrowIpcFormat;
+
+ fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ write!(formatter, "expected one of: {:?}", &FIELDS)
+ }
+
+ fn visit_i64<E>(self, v: i64) -> std::result::Result<Self::Value, E>
+ where
+ E: serde::de::Error,
+ {
+ i32::try_from(v)
+ .ok()
+ .and_then(|x| x.try_into().ok())
+ .ok_or_else(|| {
+ serde::de::Error::invalid_value(serde::de::Unexpected::Signed(v), &self)
+ })
+ }
+
+ fn visit_u64<E>(self, v: u64) -> std::result::Result<Self::Value, E>
+ where
+ E: serde::de::Error,
+ {
+ i32::try_from(v)
+ .ok()
+ .and_then(|x| x.try_into().ok())
+ .ok_or_else(|| {
+ serde::de::Error::invalid_value(serde::de::Unexpected::Unsigned(v), &self)
+ })
+ }
+
+ fn visit_str<E>(self, value: &str) -> std::result::Result<Self::Value, E>
+ where
+ E: serde::de::Error,
+ {
+ match value {
+ "ARROW_IPC_FORMAT_FILE" => Ok(ArrowIpcFormat::File),
+ "ARROW_IPC_FORMAT_STREAM" => Ok(ArrowIpcFormat::Stream),
+ _ => Err(serde::de::Error::unknown_variant(value, FIELDS)),
+ }
+ }
+ }
+ deserializer.deserialize_any(GeneratedVisitor)
+ }
+}
impl serde::Serialize for ArrowScanExecNode {
#[allow(deprecated)]
fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
@@ -1456,10 +1527,18 @@ impl serde::Serialize for ArrowScanExecNode {
if self.base_conf.is_some() {
len += 1;
}
+ if self.format != 0 {
+ len += 1;
+ }
let mut struct_ser = serializer.serialize_struct("datafusion.ArrowScanExecNode", len)?;
if let Some(v) = self.base_conf.as_ref() {
struct_ser.serialize_field("baseConf", v)?;
}
+ if self.format != 0 {
+ let v = ArrowIpcFormat::try_from(self.format)
+ .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.format)))?;
+ struct_ser.serialize_field("format", &v)?;
+ }
struct_ser.end()
}
}
@@ -1472,11 +1551,13 @@ impl<'de> serde::Deserialize<'de> for ArrowScanExecNode {
const FIELDS: &[&str] = &[
"base_conf",
"baseConf",
+ "format",
];
#[allow(clippy::enum_variant_names)]
enum GeneratedField {
BaseConf,
+ Format,
}
impl<'de> serde::Deserialize<'de> for GeneratedField {
fn deserialize<D>(deserializer: D) -> std::result::Result<GeneratedField, D::Error>
@@ -1499,6 +1580,7 @@ impl<'de> serde::Deserialize<'de> for ArrowScanExecNode {
{
match value {
"baseConf" | "base_conf" => Ok(GeneratedField::BaseConf),
+ "format" => Ok(GeneratedField::Format),
_ => Err(serde::de::Error::unknown_field(value, FIELDS)),
}
}
@@ -1519,6 +1601,7 @@ impl<'de> serde::Deserialize<'de> for ArrowScanExecNode {
V: serde::de::MapAccess<'de>,
{
let mut base_conf__ = None;
+ let mut format__ = None;
while let Some(k) = map_.next_key()? {
match k {
GeneratedField::BaseConf => {
@@ -1527,10 +1610,17 @@ impl<'de> serde::Deserialize<'de> for ArrowScanExecNode {
}
base_conf__ = map_.next_value()?;
}
+ GeneratedField::Format => {
+ if format__.is_some() {
+ return Err(serde::de::Error::duplicate_field("format"));
+ }
+ format__ = Some(map_.next_value::<ArrowIpcFormat>()? as i32);
+ }
}
}
Ok(ArrowScanExecNode {
base_conf: base_conf__,
+ format: format__.unwrap_or_default(),
})
}
}
diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs
index ba00577ab9..a6fe825ff6 100644
--- a/datafusion/proto-models/src/generated/prost.rs
+++ b/datafusion/proto-models/src/generated/prost.rs
@@ -1996,6 +1996,8 @@ pub struct AvroScanExecNode {
pub struct ArrowScanExecNode {
#[prost(message, optional, tag = "1")]
pub base_conf: ::core::option::Option<FileScanExecConf>,
+ #[prost(enumeration = "ArrowIpcFormat", tag = "2")]
+ pub format: i32,
}
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct MemoryScanExecNode {
@@ -2794,6 +2796,36 @@ impl InsertOp {
}
}
}
+/// Identifies which Arrow IPC format an ArrowScanExecNode reads.
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)]
+#[repr(i32)]
+pub enum ArrowIpcFormat {
+ /// Arrow IPC file format (with footer, supports range-based parallel reading).
+ /// This is the default for payloads encoded before the format field existed.
+ File = 0,
+ /// Arrow IPC stream format (without footer, sequential reading only)
+ Stream = 1,
+}
+impl ArrowIpcFormat {
+ /// String value of the enum field names used in the ProtoBuf definition.
+ ///
+ /// The values are not transformed in any way and thus are considered stable
+ /// (if the ProtoBuf definition does not change) and safe for programmatic use.
+ pub fn as_str_name(&self) -> &'static str {
+ match self {
+ Self::File => "ARROW_IPC_FORMAT_FILE",
+ Self::Stream => "ARROW_IPC_FORMAT_STREAM",
+ }
+ }
+ /// Creates an enum from field names used in the ProtoBuf definition.
+ pub fn from_str_name(value: &str) -> ::core::option::Option<Self> {
+ match value {
+ "ARROW_IPC_FORMAT_FILE" => Some(Self::File),
+ "ARROW_IPC_FORMAT_STREAM" => Some(Self::Stream),
+ _ => None,
+ }
+ }
+}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)]
#[repr(i32)]
pub enum PartitionMode {
diff --git a/datafusion/proto/tests/cases/plans/sources.rs b/datafusion/proto/tests/cases/plans/sources.rs
index 04708dec64..5cdd42a0ac 100644
--- a/datafusion/proto/tests/cases/plans/sources.rs
+++ b/datafusion/proto/tests/cases/plans/sources.rs
@@ -57,10 +57,15 @@ use datafusion_datasource::{TableSchema, TableSchemaBuilder};
use datafusion_expr::ColumnarValue;
use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx;
use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx;
+use datafusion_proto::bytes::{
+ physical_plan_from_bytes_with_proto_converter,
+ physical_plan_to_bytes_with_proto_converter,
+};
use datafusion_proto::physical_plan::{
AsExecutionPlan, DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter,
PhysicalExtensionCodec, PhysicalProtoConverterExtension,
};
+use datafusion_proto::protobuf;
use datafusion_proto::protobuf::PhysicalPlanNode;
use prost::Message;
use std::collections::HashMap;
@@ -156,6 +161,22 @@ fn roundtrip_parquet_exec_attaches_cached_reader_factory_after_roundtrip() -> Re
Ok(())
}
+/// Returns `FileSource::file_type` of a `DataSourceExec` file scan, e.g.
+/// "arrow" vs "arrow_stream". The two Arrow IPC formats print identically in
+/// plan debug output, so roundtrip tests must inspect the source directly.
+fn scan_file_type(plan: &Arc<dyn ExecutionPlan>) -> Result<String> {
+ let data_source = plan.downcast_ref::<DataSourceExec>().ok_or_else(|| {
+ internal_datafusion_err!("Expected DataSourceExec after roundtrip")
+ })?;
+ let file_scan = data_source
+ .data_source()
+ .downcast_ref::<FileScanConfig>()
+ .ok_or_else(|| {
+ internal_datafusion_err!("Expected FileScanConfig after roundtrip")
+ })?;
+ Ok(file_scan.file_source().file_type().to_string())
+}
+
#[test]
fn roundtrip_arrow_scan() -> Result<()> {
let file_schema =
@@ -177,7 +198,85 @@ fn roundtrip_arrow_scan() -> Result<()> {
})
.build();
- roundtrip_test(DataSourceExec::from_data_source(scan_config))
+ let roundtripped = roundtrip_test_and_return(
+ DataSourceExec::from_data_source(scan_config),
+ &SessionContext::new(),
+ &DefaultPhysicalExtensionCodec {},
+ &DefaultPhysicalProtoConverter {},
+ )?;
+ assert_eq!(scan_file_type(&roundtripped)?, "arrow");
+ Ok(())
+}
+
+#[test]
+fn roundtrip_arrow_stream_scan() -> Result<()> {
+ let file_schema =
+ Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)]));
+ let file_source = Arc::new(ArrowSource::new_stream_file_source(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.arrows".to_string(),
+ 1024,
+ )])])
+ .build();
+
+ let roundtripped = roundtrip_test_and_return(
+ DataSourceExec::from_data_source(scan_config),
+ &SessionContext::new(),
+ &DefaultPhysicalExtensionCodec {},
+ &DefaultPhysicalProtoConverter {},
+ )?;
+ assert_eq!(scan_file_type(&roundtripped)?, "arrow_stream");
+ Ok(())
+}
+
+#[test]
+fn arrow_scan_without_format_field_decodes_as_file_format() -> Result<()> {
+ // Payloads encoded before `ArrowScanExecNode.format` existed carry no
+ // format discriminator; they must keep decoding as the IPC file format.
+ let file_schema =
+ Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)]));
+ let file_source = Arc::new(ArrowSource::new_file_source(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.arrow".to_string(),
+ 1024,
+ )])])
+ .build();
+
+ let codec = DefaultPhysicalExtensionCodec {};
+ let proto_converter = DefaultPhysicalProtoConverter {};
+ let bytes = physical_plan_to_bytes_with_proto_converter(
+ DataSourceExec::from_data_source(scan_config),
+ &codec,
+ &proto_converter,
+ )?;
+
+ let mut node = PhysicalPlanNode::decode(bytes.as_ref()).map_err(|e| {
+ internal_datafusion_err!("Failed to decode PhysicalPlanNode: {e}")
+ })?;
+ match node.physical_plan_type.as_mut() {
+ Some(protobuf::physical_plan_node::PhysicalPlanType::ArrowScan(scan)) => {
+ scan.format = protobuf::ArrowIpcFormat::File as i32;
+ }
+ other => return internal_err!("Expected ArrowScan node, got {other:?}"),
+ }
+
+ let ctx = SessionContext::new();
+ let decoded = physical_plan_from_bytes_with_proto_converter(
+ &node.encode_to_vec(),
+ ctx.task_ctx().as_ref(),
+ &codec,
+ &proto_converter,
+ )?;
+ assert_eq!(scan_file_type(&decoded)?, "arrow");
+ Ok(())
}
#[test]
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org