Skip to content

proto: layered DataSource / FileSource / DataSink nodes (precursor to #26555, #26807) #26824

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
adriangb wants to merge 4 commits into
base: main
Choose a base branch
from
+4,000 −940
Open

proto: layered DataSource / FileSource / DataSink nodes (precursor to #26555, #26807) #26824

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
2b3a28e
proto: add layered DataSource/FileSource/DataSink messages
adriangb Oct 27, 2026
e9f9146
datasource, proto: layered proto hooks; deprecate whole-plan hooks
adriangb Oct 27, 2026
19d2ff3
proto: tests for layered and legacy data source/sink decoding
adriangb Oct 27, 2026
009fa28
docs: upgrade guide for layered proto data source/sink nodes
adriangb Oct 27, 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
62 changes: 35 additions & 27 deletions datafusion/datasource-arrow/src/source.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -408,16 +408,15 @@ impl FileSource for ArrowSource {
datafusion_physical_plan::apply_expression_roots(self.projection.source.iter(), f)
}

/// Emit an `ArrowScan` node wrapping the shared base config and recording
/// which Arrow IPC format (file or stream) this source reads.
/// Emit an `ArrowSourceNode` recording which Arrow IPC format (file or
/// stream) this source reads.
#[cfg(feature = "proto")]
fn try_to_proto(
fn try_to_proto_file_source(
&self,
base: &FileScanConfig,
ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>,
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
_ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>,
) -> Result<Option<datafusion_proto_models::protobuf::FileSourceNode>> {
use datafusion_proto_models::protobuf;
use protobuf::physical_plan_node::PhysicalPlanType;
use protobuf::file_source_node::FileSourceType;

// Exhaustive destructure: adding a field to `ArrowSource` without
// deciding how it is serialized is a compile error, not a silent
Expand All @@ -437,23 +436,18 @@ impl FileSource for ArrowSource {
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,
},
)),
Ok(Some(protobuf::FileSourceNode {
file_source_type: Some(FileSourceType::Arrow(protobuf::ArrowSourceNode {
format: format as i32,
})),
}))
}
}

#[cfg(feature = "proto")]
impl ArrowSource {
/// Reconstructs a `DataSourceExec` from a protobuf `ArrowScan`.
///
/// Payloads encoded before the `format` field existed leave it unset;
/// those decode as the IPC file format, matching the historical behavior.
/// Reconstructs a `DataSourceExec` from a legacy protobuf `ArrowScan`.
#[deprecated(since = "56.0.0", note = "use `try_from_proto_file_source` instead")]
pub fn try_from_proto(
node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>,
Expand All @@ -469,31 +463,45 @@ impl ArrowSource {
"PhysicalPlanNode is not an ArrowScan"
);
};

let protobuf::ArrowScanExecNode { base_conf, format } = scan;

let base_conf = base_conf.as_ref().ok_or_else(|| {
let base_conf = scan.base_conf.as_ref().ok_or_else(|| {
datafusion_common::internal_datafusion_err!(
"ArrowScanExecNode is missing required field 'base_conf'"
)
})?;
let source_node = protobuf::ArrowSourceNode::from(scan);
let source = Self::try_from_proto_file_source(&source_node, base_conf, ctx)?;
let scan_conf = FileScanConfig::try_from_proto(base_conf, ctx, source)?;
Ok(DataSourceExec::from_data_source(scan_conf))
}

/// Reconstructs an `ArrowSource` from its protobuf representation. `base`
/// is the enclosing scan's shared config, used for the table schema.
///
/// 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_file_source(
node: &datafusion_proto_models::protobuf::ArrowSourceNode,
base: &datafusion_proto_models::protobuf::FileScanExecConf,
_ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>,
) -> Result<Arc<dyn FileSource>> {
use datafusion_datasource::file_scan_config::FileScanConfig;
use datafusion_proto_models::protobuf;

let protobuf::ArrowSourceNode { format } = node;
let format = protobuf::ArrowIpcFormat::try_from(*format).map_err(|_| {
datafusion_common::internal_datafusion_err!(
"Unknown ArrowIpcFormat: {}",
format
)
})?;

let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?;
let source = match format {
let table_schema = FileScanConfig::parse_table_schema_from_proto(base)?;
Ok(Arc::new(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
47 changes: 27 additions & 20 deletions datafusion/datasource-avro/src/source.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -179,15 +179,14 @@ impl FileSource for AvroSource {
datafusion_physical_plan::apply_expression_roots(self.projection.source.iter(), f)
}

/// Emit an `AvroScan` node wrapping the shared base config.
/// Emit an (empty) `AvroSourceNode`.
#[cfg(feature = "proto")]
fn try_to_proto(
fn try_to_proto_file_source(
&self,
base: &FileScanConfig,
ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>,
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
_ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>,
) -> Result<Option<datafusion_proto_models::protobuf::FileSourceNode>> {
use datafusion_proto_models::protobuf;
use protobuf::physical_plan_node::PhysicalPlanType;
use protobuf::file_source_node::FileSourceType;

// Exhaustive destructure: adding a field to `AvroSource` without
// deciding how it is serialized is a compile error, not a silent
Expand All @@ -203,18 +202,16 @@ impl FileSource for AvroSource {
metrics: _,
} = self;

let node = protobuf::AvroScanExecNode {
base_conf: Some(base.try_to_proto(ctx)?),
};
Ok(Some(protobuf::PhysicalPlanNode {
physical_plan_type: Some(PhysicalPlanType::AvroScan(node)),
Ok(Some(protobuf::FileSourceNode {
file_source_type: Some(FileSourceType::Avro(protobuf::AvroSourceNode {})),
}))
}
}

#[cfg(feature = "proto")]
impl AvroSource {
/// Reconstructs a `DataSourceExec` from a protobuf `AvroScan`.
/// Reconstructs a `DataSourceExec` from a legacy protobuf `AvroScan`.
#[deprecated(since = "56.0.0", note = "use `try_from_proto_file_source` instead")]
pub fn try_from_proto(
node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>,
Expand All @@ -229,21 +226,31 @@ impl AvroSource {
"PhysicalPlanNode is not an AvroScan"
);
};

let protobuf::AvroScanExecNode { base_conf } = scan;

let base_conf = base_conf.as_ref().ok_or_else(|| {
let base_conf = scan.base_conf.as_ref().ok_or_else(|| {
datafusion_common::internal_datafusion_err!(
"AvroScanExecNode is missing required field 'base_conf'"
)
})?;

let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?;
let source = Arc::new(AvroSource::new(table_schema));

let source = Self::try_from_proto_file_source(
&protobuf::AvroSourceNode::from(scan),
base_conf,
ctx,
)?;
let conf = FileScanConfig::try_from_proto(base_conf, ctx, source)?;
Ok(DataSourceExec::from_data_source(conf))
}

/// Reconstructs an `AvroSource` from its protobuf representation. `base`
/// is the enclosing scan's shared config, used for the table schema.
pub fn try_from_proto_file_source(
node: &datafusion_proto_models::protobuf::AvroSourceNode,
base: &datafusion_proto_models::protobuf::FileScanExecConf,
_ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>,
) -> Result<Arc<dyn FileSource>> {
let datafusion_proto_models::protobuf::AvroSourceNode {} = node;
let table_schema = FileScanConfig::parse_table_schema_from_proto(base)?;
Ok(Arc::new(AvroSource::new(table_schema)))
}
}

mod private {
Expand Down
30 changes: 13 additions & 17 deletions datafusion/datasource-csv/src/file_format.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -853,13 +853,11 @@ impl DataSink for CsvSink {
}

#[cfg(feature = "proto")]
fn try_to_proto(
fn try_to_proto_sink(
&self,
exec: &DataSinkExec,
ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>,
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
_ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>,
) -> Result<Option<datafusion_proto_models::protobuf::DataSinkNode>> {
use datafusion_proto_models::protobuf;
use protobuf::physical_plan_node::PhysicalPlanType;

// Keep an exhaustive guard in the active hook while centralizing the
// field mapping in the exhaustive `TryFrom<&CsvSink>` below.
Expand All @@ -868,17 +866,10 @@ impl DataSink for CsvSink {
writer_options: _,
} = self;

let input = ctx.encode_child(exec.input())?;
let sort_order = exec.encode_sort_order(ctx)?;
let sink = protobuf::CsvSink::try_from(self)?;
let node = protobuf::CsvSinkExecNode {
input: Some(Box::new(input)),
sink: Some(sink),
sink_schema: Some(exec.schema().as_ref().try_into()?),
sort_order,
};
Ok(Some(protobuf::PhysicalPlanNode {
physical_plan_type: Some(PhysicalPlanType::CsvSink(Box::new(node))),
Ok(Some(protobuf::DataSinkNode {
sink_type: Some(protobuf::data_sink_node::SinkType::Csv(
protobuf::CsvSink::try_from(self)?,
)),
}))
}
}
Expand Down Expand Up @@ -928,7 +919,12 @@ impl TryFrom<&datafusion_proto_models::protobuf::CsvSink> for CsvSink {

#[cfg(feature = "proto")]
impl CsvSink {
/// Reconstructs a [`DataSinkExec`] containing a `CsvSink` from protobuf.
/// Reconstructs a [`DataSinkExec`] containing a `CsvSink` from a legacy
/// protobuf plan node.
#[deprecated(
since = "56.0.0",
note = "decode `DataSinkNode` with `TryFrom<&protobuf::CsvSink>` instead"
)]
pub fn try_from_proto(
node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>,
Expand Down
76 changes: 43 additions & 33 deletions datafusion/datasource-csv/src/source.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -335,26 +335,24 @@ impl FileSource for CsvSource {
datafusion_physical_plan::apply_expression_roots(self.projection.source.iter(), f)
}

/// Emit a `CsvScan` node wrapping the shared base config and CSV options.
/// Emit a `CsvSourceNode` carrying the CSV options.
#[cfg(feature = "proto")]
fn try_to_proto(
fn try_to_proto_file_source(
&self,
base: &FileScanConfig,
ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>,
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
_ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>,
) -> Result<Option<datafusion_proto_models::protobuf::FileSourceNode>> {
use datafusion_proto_models::protobuf;
use protobuf::physical_plan_node::PhysicalPlanType;
use protobuf::file_source_node::FileSourceType;

let node = protobuf::CsvScanExecNode {
base_conf: Some(base.try_to_proto(ctx)?),
let node = protobuf::CsvSourceNode {
has_header: self.has_header(),
delimiter: proto_byte_to_string(self.delimiter(), "delimiter")?,
quote: proto_byte_to_string(self.quote(), "quote")?,
optional_escape: self
.escape()
.map(|escape| {
Ok::<_, DataFusionError>(
protobuf::csv_scan_exec_node::OptionalEscape::Escape(
protobuf::csv_source_node::OptionalEscape::Escape(
proto_byte_to_string(escape, "escape")?,
),
)
Expand All @@ -364,7 +362,7 @@ impl FileSource for CsvSource {
.comment()
.map(|comment| {
Ok::<_, DataFusionError>(
protobuf::csv_scan_exec_node::OptionalComment::Comment(
protobuf::csv_source_node::OptionalComment::Comment(
proto_byte_to_string(comment, "comment")?,
),
)
Expand All @@ -374,8 +372,8 @@ impl FileSource for CsvSource {
truncate_rows: self.truncate_rows(),
terminator: self.terminator().map(|terminator| vec![terminator]),
};
Ok(Some(protobuf::PhysicalPlanNode {
physical_plan_type: Some(PhysicalPlanType::CsvScan(node)),
Ok(Some(protobuf::FileSourceNode {
file_source_type: Some(FileSourceType::Csv(node)),
}))
}
}
Expand Down Expand Up @@ -606,14 +604,12 @@ fn proto_bytes_to_byte(bytes: &[u8], description: &str) -> Result<u8> {

#[cfg(feature = "proto")]
impl CsvSource {
/// Reconstructs a `DataSourceExec` from a protobuf `CsvScan`.
///
/// Payloads without a terminator use the default newline terminator.
/// Reconstructs a `DataSourceExec` from a legacy protobuf `CsvScan`.
#[deprecated(since = "56.0.0", note = "use `try_from_proto_file_source` instead")]
pub fn try_from_proto(
node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>,
) -> Result<Arc<dyn ExecutionPlan>> {
use datafusion_common::config::CsvOptions;
use datafusion_datasource::file_scan_config::FileScanConfig;
use datafusion_datasource::source::DataSourceExec;
use datafusion_proto_models::protobuf;
Expand All @@ -623,51 +619,65 @@ impl CsvSource {
else {
return datafusion_common::internal_err!("PhysicalPlanNode is not a CsvScan");
};

let base_conf = scan.base_conf.as_ref().ok_or_else(|| {
datafusion_common::internal_datafusion_err!(
"CsvScanExecNode is missing required field 'base_conf'"
)
})?;
let source_node = protobuf::CsvSourceNode::from(scan);
let source = Self::try_from_proto_file_source(&source_node, base_conf, ctx)?;
let conf = FileScanConfig::try_from_proto(base_conf, ctx, source)?;
Ok(DataSourceExec::from_data_source(conf))
}

let escape = match &scan.optional_escape {
Some(protobuf::csv_scan_exec_node::OptionalEscape::Escape(escape)) => {
/// Reconstructs a `CsvSource` from its protobuf representation. `base`
/// is the enclosing scan's shared config, used for the table schema.
///
/// Payloads without a terminator use the default newline terminator.
pub fn try_from_proto_file_source(
node: &datafusion_proto_models::protobuf::CsvSourceNode,
base: &datafusion_proto_models::protobuf::FileScanExecConf,
_ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>,
) -> Result<Arc<dyn FileSource>> {
use datafusion_common::config::CsvOptions;
use datafusion_datasource::file_scan_config::FileScanConfig;
use datafusion_proto_models::protobuf;

let escape = match &node.optional_escape {
Some(protobuf::csv_source_node::OptionalEscape::Escape(escape)) => {
Some(proto_str_to_byte(escape, "escape")?)
}
None => None,
};
let comment = match &scan.optional_comment {
Some(protobuf::csv_scan_exec_node::OptionalComment::Comment(comment)) => {
let comment = match &node.optional_comment {
Some(protobuf::csv_source_node::OptionalComment::Comment(comment)) => {
Some(proto_str_to_byte(comment, "comment")?)
}
None => None,
};
let terminator = scan
let terminator = node
.terminator
.as_deref()
.map(|terminator| proto_bytes_to_byte(terminator, "terminator"))
.transpose()?;

let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?;
let table_schema = FileScanConfig::parse_table_schema_from_proto(base)?;

let csv_options = CsvOptions {
has_header: Some(scan.has_header),
delimiter: proto_str_to_byte(&scan.delimiter, "delimiter")?,
quote: proto_str_to_byte(&scan.quote, "quote")?,
newlines_in_values: Some(scan.newlines_in_values),
truncated_rows: Some(scan.truncate_rows),
has_header: Some(node.has_header),
delimiter: proto_str_to_byte(&node.delimiter, "delimiter")?,
quote: proto_str_to_byte(&node.quote, "quote")?,
newlines_in_values: Some(node.newlines_in_values),
truncated_rows: Some(node.truncate_rows),
..Default::default()
};
let source = Arc::new(
Ok(Arc::new(
CsvSource::new(table_schema)
.with_csv_options(csv_options)
.with_escape(escape)
.with_comment(comment)
.with_terminator(terminator),
);

let conf = FileScanConfig::try_from_proto(base_conf, ctx, source)?;
Ok(DataSourceExec::from_data_source(conf))
))
}
}

Expand Down
Loading
Loading