Skip to content
Merged
Show file tree
Hide file tree
Changes from 5 commits
Commits
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
6 changes: 4 additions & 2 deletions datafusion/datasource/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -34,8 +34,10 @@ all-features = true
backtrace = ["datafusion-common/backtrace"]
compression = ["async-compression", "liblzma", "bzip2", "flate2", "zstd", "tokio-util"]
default = ["compression"]
# Enables protobuf conversions for datasource types and serialization hooks.
# Off by default so consumers that never serialize plans pay nothing.
# Enables protobuf conversions for datasource types, source serialization hooks,
# and the shared `FileScanConfig` <-> proto conversion. Off by default so
# consumers that never serialize plans pay nothing. Mirrors the `proto` feature
# on `datafusion-physical-plan`.
proto = [
"dep:datafusion-proto-models",
"datafusion-physical-plan/proto",
Expand Down
25 changes: 25 additions & 0 deletions datafusion/datasource/src/file.rs
Original file line number Diff line number Diff line change
Expand Up @@ -351,6 +351,31 @@ pub trait FileSource: Any + Send + Sync {
fn schema_adapter_factory(&self) -> Option<Arc<dyn SchemaAdapterFactory>> {
None
}

/// Serialize this file source into a full [`PhysicalPlanNode`] (a
/// `DataSourceExec` wrapping the `FileScanConfig`), if it knows how.
///
/// `base` is the shared [`FileScanConfig`] this source is wrapped in; the
/// format-agnostic parts (file groups, schema, statistics, ordering,
/// projection, …) are encoded via
/// [`FileScanConfig::try_to_proto`](crate::file_scan_config::FileScanConfig::try_to_proto),
/// and the concrete source appends its format-specific fields (e.g. CSV
/// delimiter/quote) around it.
///
/// * `Ok(None)` (the default) — this source has no proto hook yet; the
/// caller falls back to the central downcast chain in `datafusion-proto`.
/// * `Ok(Some(node))` — fully serialized; the caller must not fall back.
///
/// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode
/// [`FileScanConfig`]: crate::file_scan_config::FileScanConfig
#[cfg(feature = "proto")]
fn try_to_proto(
&self,
_base: &FileScanConfig,
_ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>,
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
Ok(None)
}
}

impl dyn FileSource {
Expand Down
127 changes: 127 additions & 0 deletions datafusion/datasource/src/file_scan_config/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,13 @@

pub(crate) mod sort_pushdown;

/// Shared `FileScanConfig` <-> proto conversion, gated on the `proto` feature.
/// Attaches inherent `try_to_proto` / `try_from_proto` /
/// `parse_table_schema_from_proto` helpers to [`FileScanConfig`] used by every
/// file source's `try_to_proto` hook.
#[cfg(feature = "proto")]
mod proto;

use crate::file_groups::FileGroup;
use crate::{
PartitionedFile, display::FileGroupsDisplay, file::FileSource,
Expand Down Expand Up @@ -1176,6 +1183,18 @@ impl DataSource for FileScanConfig {

Some(Arc::new(SharedWorkSource::from_config(self)) as Arc<dyn Any + Send + Sync>)
}

/// Serialize this file scan by delegating to the concrete
/// [`FileSource`]'s
/// [`try_to_proto`](crate::file::FileSource::try_to_proto) hook, passing
/// `self` as the shared spine it needs to emit the base config.
#[cfg(feature = "proto")]
fn try_to_proto(
&self,
ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>,
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
self.file_source().try_to_proto(self, ctx)
}
}

impl FileScanConfig {
Expand Down Expand Up @@ -1569,12 +1588,18 @@ mod tests {
use datafusion_common::{Result, assert_batches_eq, internal_err};
use datafusion_execution::TaskContext;
use datafusion_expr::SortExpr;
#[cfg(feature = "proto")]
use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF};
use datafusion_physical_expr::create_physical_sort_expr;
use datafusion_physical_expr::expressions::Literal;
use datafusion_physical_expr::projection::ProjectionExpr;
use datafusion_physical_expr::projection::ProjectionExprs;
use datafusion_physical_plan::ExecutionPlan;
use datafusion_physical_plan::execution_plan::collect;
#[cfg(feature = "proto")]
use datafusion_physical_plan::proto::{ExecutionPlanEncode, ExecutionPlanEncodeCtx};
#[cfg(feature = "proto")]
use datafusion_proto_models::protobuf::{PhysicalExprNode, PhysicalPlanNode};
use futures::FutureExt as _;
use futures::StreamExt as _;
use futures::stream;
Expand Down Expand Up @@ -1633,6 +1658,108 @@ mod tests {
}
}

#[cfg(feature = "proto")]
#[derive(Clone)]
struct ProtoHookSource {
metrics: ExecutionPlanMetricsSet,
table_schema: TableSchema,
}

#[cfg(feature = "proto")]
impl ProtoHookSource {
fn new(table_schema: TableSchema) -> Self {
Self {
metrics: ExecutionPlanMetricsSet::new(),
table_schema,
}
}
}

#[cfg(feature = "proto")]
impl FileSource for ProtoHookSource {
fn create_file_opener(
&self,
_object_store: Arc<dyn ObjectStore>,
_base_config: &FileScanConfig,
_partition: usize,
) -> Result<Arc<dyn crate::file_stream::FileOpener>> {
internal_err!("not needed for proto delegation test")
}

fn table_schema(&self) -> &TableSchema {
&self.table_schema
}

fn with_batch_size(&self, _batch_size: usize) -> Arc<dyn FileSource> {
Arc::new(self.clone())
}

fn metrics(&self) -> &ExecutionPlanMetricsSet {
&self.metrics
}

fn file_type(&self) -> &str {
"proto-hook-test"
}

fn try_to_proto(
&self,
_base: &FileScanConfig,
_ctx: &ExecutionPlanEncodeCtx<'_>,
) -> Result<Option<PhysicalPlanNode>> {
Ok(Some(PhysicalPlanNode::default()))
}
}

#[cfg(feature = "proto")]
struct UnusedPlanEncoder;

#[cfg(feature = "proto")]
impl ExecutionPlanEncode for UnusedPlanEncoder {
fn encode_plan(
&self,
_plan: &Arc<dyn ExecutionPlan>,
) -> Result<PhysicalPlanNode> {
internal_err!("not needed for proto delegation test")
}

fn encode_expr(&self, _expr: &Arc<dyn PhysicalExpr>) -> Result<PhysicalExprNode> {
internal_err!("not needed for proto delegation test")
}

fn encode_udf(&self, _udf: &ScalarUDF) -> Result<Option<Vec<u8>>> {
internal_err!("not needed for proto delegation test")
}

fn encode_udaf(&self, _udaf: &AggregateUDF) -> Result<Option<Vec<u8>>> {
internal_err!("not needed for proto delegation test")
}

fn encode_udwf(&self, _udwf: &WindowUDF) -> Result<Option<Vec<u8>>> {
internal_err!("not needed for proto delegation test")
}
}

#[cfg(feature = "proto")]
#[test]
fn data_source_exec_delegates_proto_to_file_source() -> Result<()> {
let schema = Arc::new(Schema::new(vec![Field::new(
"value",
DataType::Int32,
false,
)]));
let source = Arc::new(ProtoHookSource::new(TableSchema::from(&schema)));
let config =
FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source)
.build();
let exec = DataSourceExec::from_data_source(config);
let encoder = UnusedPlanEncoder;
let ctx = ExecutionPlanEncodeCtx::new(&encoder);

assert_eq!(exec.try_to_proto(&ctx)?, Some(PhysicalPlanNode::default()));
Ok(())
}

#[test]
fn physical_plan_config_no_projection_tab_cols_as_field() {
let file_schema = aggr_test_schema();
Expand Down
Loading
Loading