From 2c7b08b90148e6a1b80aba4711f67f7224430cf2 Mon Sep 17 00:00:00 2001 From: Burak Sen Date: Mon, 10 Aug 2026 16:56:43 +0300 Subject: [PATCH 1/2] fix: preserve Arrow IPC format --- datafusion/datasource-arrow/src/source.rs | 34 +++++-- .../proto-models/proto/datafusion.proto | 12 +++ .../proto-models/src/generated/pbjson.rs | 93 ++++++++++++++++++ .../proto-models/src/generated/prost.rs | 36 +++++++ .../tests/cases/roundtrip_physical_plan.rs | 96 ++++++++++++++++++- 5 files changed, 262 insertions(+), 9 deletions(-) diff --git a/datafusion/datasource-arrow/src/source.rs b/datafusion/datasource-arrow/src/source.rs index bf92faa9a110..2ad3fb1b0191 100644 --- a/datafusion/datasource-arrow/src/source.rs +++ b/datafusion/datasource-arrow/src/source.rs @@ -409,10 +409,8 @@ impl FileSource for ArrowSource { ) } - /// 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, @@ -422,10 +420,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, }, )), })) @@ -436,8 +440,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<'_>, @@ -461,9 +465,23 @@ 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::Unspecified | 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 5a8bede19582..825acf9569bb 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1289,8 +1289,20 @@ message AvroScanExecNode { FileScanExecConf base_conf = 1; } +// Identifies which Arrow IPC format an ArrowScanExecNode reads. +enum ArrowIpcFormat { + // Unset in payloads encoded before this field existed; decoded as the + // IPC file format for backward compatibility. + ARROW_IPC_FORMAT_UNSPECIFIED = 0; + // Arrow IPC file format (with footer, supports range-based parallel reading) + ARROW_IPC_FORMAT_FILE = 1; + // Arrow IPC stream format (without footer, sequential reading only) + ARROW_IPC_FORMAT_STREAM = 2; +} + 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 9c82ceb4a2de..85bfa7273bdf 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -1428,6 +1428,80 @@ impl<'de> serde::Deserialize<'de> for AnalyzedLogicalPlanType { deserializer.deserialize_struct("datafusion.AnalyzedLogicalPlanType", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for ArrowIpcFormat { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + let variant = match self { + Self::Unspecified => "ARROW_IPC_FORMAT_UNSPECIFIED", + 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(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "ARROW_IPC_FORMAT_UNSPECIFIED", + "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(self, v: i64) -> std::result::Result + 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(self, v: u64) -> std::result::Result + 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(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "ARROW_IPC_FORMAT_UNSPECIFIED" => Ok(ArrowIpcFormat::Unspecified), + "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(&self, serializer: S) -> std::result::Result @@ -1439,10 +1513,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() } } @@ -1455,11 +1537,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(deserializer: D) -> std::result::Result @@ -1482,6 +1566,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)), } } @@ -1502,6 +1587,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 => { @@ -1510,10 +1596,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::()? 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 df8f4677c9a6..dbd746829a6d 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -1985,6 +1985,8 @@ pub struct AvroScanExecNode { pub struct ArrowScanExecNode { #[prost(message, optional, tag = "1")] pub base_conf: ::core::option::Option, + #[prost(enumeration = "ArrowIpcFormat", tag = "2")] + pub format: i32, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct MemoryScanExecNode { @@ -2774,6 +2776,40 @@ 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 { + /// Unset in payloads encoded before this field existed; decoded as the + /// IPC file format for backward compatibility. + Unspecified = 0, + /// Arrow IPC file format (with footer, supports range-based parallel reading) + File = 1, + /// Arrow IPC stream format (without footer, sequential reading only) + Stream = 2, +} +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::Unspecified => "ARROW_IPC_FORMAT_UNSPECIFIED", + 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 { + match value { + "ARROW_IPC_FORMAT_UNSPECIFIED" => Some(Self::Unspecified), + "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/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 19249cf97bb4..9a8c89080042 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -1347,6 +1347,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) -> Result { + let data_source = plan.downcast_ref::().ok_or_else(|| { + internal_datafusion_err!("Expected DataSourceExec after roundtrip") + })?; + let file_scan = data_source + .data_source() + .downcast_ref::() + .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 = @@ -1368,7 +1384,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::Unspecified 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] From a896b6a4211522f84838bbf8544431d540810c16 Mon Sep 17 00:00:00 2001 From: Burak Sen Date: Mon, 10 Aug 2026 20:18:04 +0300 Subject: [PATCH 2/2] apply review to delete unspecified format enum --- datafusion/datasource-arrow/src/source.rs | 4 +--- datafusion/proto-models/proto/datafusion.proto | 10 ++++------ datafusion/proto-models/src/generated/pbjson.rs | 3 --- datafusion/proto-models/src/generated/prost.rs | 12 ++++-------- .../proto/tests/cases/roundtrip_physical_plan.rs | 2 +- 5 files changed, 10 insertions(+), 21 deletions(-) diff --git a/datafusion/datasource-arrow/src/source.rs b/datafusion/datasource-arrow/src/source.rs index 2ad3fb1b0191..db0958ba1753 100644 --- a/datafusion/datasource-arrow/src/source.rs +++ b/datafusion/datasource-arrow/src/source.rs @@ -477,9 +477,7 @@ impl ArrowSource { protobuf::ArrowIpcFormat::Stream => { ArrowSource::new_stream_file_source(table_schema) } - protobuf::ArrowIpcFormat::Unspecified | protobuf::ArrowIpcFormat::File => { - ArrowSource::new_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 825acf9569bb..787a87f94501 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1291,13 +1291,11 @@ message AvroScanExecNode { // Identifies which Arrow IPC format an ArrowScanExecNode reads. enum ArrowIpcFormat { - // Unset in payloads encoded before this field existed; decoded as the - // IPC file format for backward compatibility. - ARROW_IPC_FORMAT_UNSPECIFIED = 0; - // Arrow IPC file format (with footer, supports range-based parallel reading) - ARROW_IPC_FORMAT_FILE = 1; + // 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 = 2; + ARROW_IPC_FORMAT_STREAM = 1; } message ArrowScanExecNode { diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 85bfa7273bdf..3b0f5743503a 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -1435,7 +1435,6 @@ impl serde::Serialize for ArrowIpcFormat { S: serde::Serializer, { let variant = match self { - Self::Unspecified => "ARROW_IPC_FORMAT_UNSPECIFIED", Self::File => "ARROW_IPC_FORMAT_FILE", Self::Stream => "ARROW_IPC_FORMAT_STREAM", }; @@ -1449,7 +1448,6 @@ impl<'de> serde::Deserialize<'de> for ArrowIpcFormat { D: serde::Deserializer<'de>, { const FIELDS: &[&str] = &[ - "ARROW_IPC_FORMAT_UNSPECIFIED", "ARROW_IPC_FORMAT_FILE", "ARROW_IPC_FORMAT_STREAM", ]; @@ -1492,7 +1490,6 @@ impl<'de> serde::Deserialize<'de> for ArrowIpcFormat { E: serde::de::Error, { match value { - "ARROW_IPC_FORMAT_UNSPECIFIED" => Ok(ArrowIpcFormat::Unspecified), "ARROW_IPC_FORMAT_FILE" => Ok(ArrowIpcFormat::File), "ARROW_IPC_FORMAT_STREAM" => Ok(ArrowIpcFormat::Stream), _ => Err(serde::de::Error::unknown_variant(value, FIELDS)), diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index dbd746829a6d..e541938ae425 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -2780,13 +2780,11 @@ impl InsertOp { #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] #[repr(i32)] pub enum ArrowIpcFormat { - /// Unset in payloads encoded before this field existed; decoded as the - /// IPC file format for backward compatibility. - Unspecified = 0, - /// Arrow IPC file format (with footer, supports range-based parallel reading) - File = 1, + /// 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 = 2, + Stream = 1, } impl ArrowIpcFormat { /// String value of the enum field names used in the ProtoBuf definition. @@ -2795,7 +2793,6 @@ impl ArrowIpcFormat { /// (if the ProtoBuf definition does not change) and safe for programmatic use. pub fn as_str_name(&self) -> &'static str { match self { - Self::Unspecified => "ARROW_IPC_FORMAT_UNSPECIFIED", Self::File => "ARROW_IPC_FORMAT_FILE", Self::Stream => "ARROW_IPC_FORMAT_STREAM", } @@ -2803,7 +2800,6 @@ impl ArrowIpcFormat { /// Creates an enum from field names used in the ProtoBuf definition. pub fn from_str_name(value: &str) -> ::core::option::Option { match value { - "ARROW_IPC_FORMAT_UNSPECIFIED" => Some(Self::Unspecified), "ARROW_IPC_FORMAT_FILE" => Some(Self::File), "ARROW_IPC_FORMAT_STREAM" => Some(Self::Stream), _ => None, diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 9a8c89080042..1d949498e0a1 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -1449,7 +1449,7 @@ fn arrow_scan_without_format_field_decodes_as_file_format() -> Result<()> { })?; match node.physical_plan_type.as_mut() { Some(protobuf::physical_plan_node::PhysicalPlanType::ArrowScan(scan)) => { - scan.format = protobuf::ArrowIpcFormat::Unspecified as i32; + scan.format = protobuf::ArrowIpcFormat::File as i32; } other => return internal_err!("Expected ArrowScan node, got {other:?}"), }