Skip to content

fix(proto): preserve Arrow IPC stream format #24224

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
buraksenn wants to merge 5 commits into
base: main
Choose a base branch
from
Open

fix(proto): preserve Arrow IPC stream format #24224

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
2c7b08b
fix: preserve Arrow IPC format
buraksenn Aug 10, 2026
a896b6a
apply review to delete unspecified format enum
buraksenn Aug 10, 2026
abe90a6
Merge remote-tracking branch 'upstream/main' into 24196-preserve-arro…
buraksenn Aug 10, 2026
30811b6
Merge remote-tracking branch 'upstream/main' into 24196-preserve-arro…
buraksenn Aug 10, 2026
9ea9b9d
Merge remote-tracking branch 'upstream/main' into 24196-preserve-arro…
buraksenn Aug 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 24 additions & 8 deletions datafusion/datasource-arrow/src/source.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
},
)),
}))
Expand All @@ -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<'_>,
Expand All @@ -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))
}
}
Expand Down
10 changes: 10 additions & 0 deletions datafusion/proto-models/proto/datafusion.proto
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Comment on lines +1294 to +1300

@Rich-T-kid Rich-T-kid Aug 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

if ARROW_IPC_FORMAT_UNSPECIFIED exist for the sake of backwards compatiability cant the default value be set to ARROW_IPC_FORMAT_FILE. this way we dont need a third variant especially since it will be interpreted as ARROW_IPC_FORMAT_FILE anyway?

@buraksenn buraksenn Aug 10, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks @Rich-T-kid, yes that was the case. I wanted to put it to not break anything but as you've said its interpreted as file anyways. Applied the change now and it simplified the code as well.


message ArrowScanExecNode {
FileScanExecConf base_conf = 1;
ArrowIpcFormat format = 2;
}

message MemoryScanExecNode {
Expand Down
90 changes: 90 additions & 0 deletions datafusion/proto-models/src/generated/pbjson.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand All @@ -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()
}
}
Expand All @@ -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>
Expand All @@ -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)),
}
}
Expand All @@ -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 => {
Expand All @@ -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(),
})
}
}
Expand Down
32 changes: 32 additions & 0 deletions datafusion/proto-models/src/generated/prost.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1987,6 +1987,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 {
Expand Down Expand Up @@ -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<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 {
Expand Down
Loading
Loading