diff --git a/datafusion/datasource-arrow/src/source.rs b/datafusion/datasource-arrow/src/source.rs index bf92faa9a110..db0958ba1753 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,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 99b4ef6272b2..d3d81e6e6958 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1290,8 +1290,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 61b1ea3ff104..96a2db9a2aa8 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(&self, serializer: S) -> std::result::Result + 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(deserializer: D) -> std::result::Result + 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(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_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 @@ -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(deserializer: D) -> std::result::Result @@ -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::()? 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 233b5fee1b29..84cabae5daba 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -1987,6 +1987,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 { @@ -2785,6 +2787,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 { + 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 04708dec6439..5cdd42a0acc3 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) -> 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 = @@ -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]