Repository navigation
Conversation
github-actions
Bot
added
documentation
86c375a to
8c43f67
Compare
|
cc @adriangb I want to tackle this one myself in order to get some hands-on exp with the approach here. Still iterating on this, so not looking for a review. |
8c43f67 to
bceb61d
Compare
| expect_plan_variant!(node, PhysicalPlanType::Extension, "Extension"); | ||
| if !extension.inputs.is_empty() { | ||
| return plan_err!("DataSourceExec expects no input plans"); | ||
| } |
There was a problem hiding this comment.
The fact that DataSources are entities nested under DataSourceExec (an ExecutionPlan impl) complicates things a bit.
Having that kind of nesting makes things a bit weird, as ExecutionPlans are also serializable entities themselves.
One weird thing that happens because of this:
DataSourceExec nodes baked by Parquet or CSV have their own proto representation, which are handled outside of the registry. It's just when the DataSource is baked by something other than Parquet or CSV when this impl ExtensionPlanFromProto for DataSourceExec mechanisms kicks in. This makes the approach of [de]serializing custom DataSources and built-in DataSource significantly different.
Something clean that comes to mind is: rather than separating [de]serialization logic for built-in entities VS custom entities, use this registry mechanism for everything, and dogfood it with every built-in implementation that needs to go over the wire. This is a very large refactor though...
There was a problem hiding this comment.
yeah, this is a wart that exists from when we had ParquetExecutionPlan, pre DataSource abstraction. Changing it is indeed a big refactor and a very breaking change for the proto wire bytes, hence why we’ve avoided it.
2ac0ba0 to
1472814
Compare
|
@adriangb this is ready for a review. I've tried to make it as consistent as possible with your previous PRs, so I think there should not be a lot of surprises. |
1472814 to
774dfd5
Compare
|
I'm working on a follow up for |
There was a problem hiding this comment.
Thanks @gabotechs. The registry logic is solid (details below). My concern is the shape of the API it locks in for 56.0.0.
Problem
Encode and decode work at different levels:
| Hook | Added in | Returns | Level |
|---|---|---|---|
DataSource::try_to_proto |
55.0.0 | PhysicalPlanNode |
plan |
FileSource::try_to_proto |
55.0.0 | PhysicalPlanNode |
plan |
ExtensionDataSourceFromProto::try_from_proto |
this PR | Arc<dyn DataSource> |
source |
Today / this PR:
PhysicalPlanNode
├─ ... (other ExecutionPlan variants)
├─ parquet_scan (built-in: one variant per format)
├─ csv_scan
├─ memory_scan
└─ extension {plan_name = "datafusion.DataSourceExec"}
└─ PhysicalExtensionDataSourceNode {source_name, bytes}
Proposed (#26824):
PhysicalPlanNode
├─ ... (other ExecutionPlan variants)
├─ parquet_scan, ... (deprecated: decode only)
└─ data_source: DataSourceExecNode
└─ DataSourceNode (oneof)
├─ file: FileScanNode
│ ├─ base: FileScanExecConf
│ └─ FileSourceNode (oneof)
│ ├─ parquet / csv / ...
│ └─ extension {name, payload}
├─ memory
└─ extension {name, payload}
Results of the current shape:
- A source must build its own wrapper. It calls
extension_data_source_node, which must be public. The source can also return anyPhysicalPlanNode. - Two ways to serialize a source. Built-ins and custom sources do not share a path (the thread here).
- Custom
FileSourcehas no path. The workaround is a markerDataSourcetype. datafusion-distributed already does this inwork_unit_file_scan.rs. - The error is not clear. See the MRE below.
Proposal: precursor PR #26824 (draft)
- Add a proto layer that mirrors the traits:
DataSourceNode,FileSourceNodeandDataSinkNode. Each is aoneofwith anextension {name, payload}arm. - Add new hooks that return their own message. Keep the hooks added in 55.0.0 and mark them
#[deprecated]. No breaking change. - Always write the new form. In
datafusion-proto, translate legacy variants into the new form on decode. (Plans are not guaranteed compatible across versions.) - After that lands, this PR is smaller:
- one
ExtensionDataSourceFromPrototrait, which takes&[u8]and has a providedregister_proto; - one decode arm for
DataSourceNode::Extension; - no wrapper registration, no
extension_data_source_node, no public modules.
- one
Timing: I suggest we land the precursor and this PR before 56.0.0. Then the hooks added in 55.0.0 change only once, through deprecation.
Stress test results (head 774dfd5)
| Check | Result |
|---|---|
| cargo-mutants (registry.rs, sink_registry.rs) | 13/13 caught, 2 unviable |
| Hand mutants | 9/10 caught. Survivor: remove names.sort_unstable() (no test asserts the order) |
| Garbage or truncated payloads | Error, no panic ✅ |
| Source nested in Projection / Repartition / Union | Byte-identical round trip ✅ |
| Name collisions, empty names, re-registration | Correct ✅ |
MRE: the reader has not registered the wrapper.
// writer: register_data_source::<VSource>(&mut reg) → encodes VSource{v: 42}
// reader: no registry, codec = a lenient user PhysicalExtensionCodec
let plan = physical_plan_from_bytes_with_ctx(&bytes, &ctx, &lenient_codec)?;
collect(plan, task_ctx).await?; // → [Some(-1)] (expected [Some(42)])
// reader with no registry and no codec:
// "No decoder is registered for the extension ExecutionPlan 'datafusion.DataSourceExec'"
// — the name 'stress.VSource' is never in the messageWith a DataSourceNode::Extension arm, the reader always knows it has a source called stress.VSource.
| fn try_from_proto( | ||
| node: &PhysicalExtensionDataSourceNode, | ||
| ctx: &ExecutionPlanDecodeCtx<'_>, | ||
| ) -> Result<Arc<dyn DataSource>>; |
There was a problem hiding this comment.
The input is &PhysicalExtensionDataSourceNode, but implementors only need the bytes. source_name is internal plumbing. Pass the payload bytes instead, so the message can change (for example, to add a version field) without breaking implementors:
| fn try_from_proto( | |
| node: &PhysicalExtensionDataSourceNode, | |
| ctx: &ExecutionPlanDecodeCtx<'_>, | |
| ) -> Result<Arc<dyn DataSource>>; | |
| fn try_from_proto( | |
| payload: &[u8], | |
| ctx: &ExecutionPlanDecodeCtx<'_>, | |
| ) -> Result<Arc<dyn DataSource>>; |
Also add a provided register_proto here, as ExtensionPlanFromProto (physical-plan/src/proto/registry.rs:136) and ExtensionPhysicalExprFromProto (physical-expr/src/proto.rs:141) do. Then register_data_source does not need to be public.
| pub fn extension_data_source_node<T: ExtensionDataSourceFromProto>( | ||
| ctx: &ExecutionPlanEncodeCtx<'_>, | ||
| payload: Vec<u8>, | ||
| ) -> Result<PhysicalPlanNode> { | ||
| let source = PhysicalExtensionDataSourceNode { | ||
| node: payload, | ||
| source_name: Some(T::NAME.to_string()), | ||
| }; | ||
| ctx.extension_node::<DataSourceExec>(source.encode_to_vec(), vec![]) |
There was a problem hiding this comment.
This helper is public only because the source must build the DataSourceExec wrapper itself (DataSource::try_to_proto returns PhysicalPlanNode).
With the precursor (#26824) (DataSource::try_to_proto_source -> DataSourceNode), DataSourceExec builds the wrapper. Then this function can be removed or made private.
| // The shared wrapper may remain registered if source registration fails. | ||
| // This is safe: it still requires a registered source decoder and rejects | ||
| // unknown sources rather than decoding their payloads through a plan codec. | ||
| DataSourceExec::register_proto(registry)?; | ||
| registry.register_decoder::<dyn DataSource, T, DataSourceDecoder>( | ||
| T::NAME, | ||
| T::try_from_proto, | ||
| ) | ||
| } |
There was a problem hiding this comment.
Registration is not atomic. If T::NAME is empty or taken, datafusion.DataSourceExec stays registered (register_rejects_an_empty_source_name tests this).
The simplest fix is to validate first:
| // The shared wrapper may remain registered if source registration fails. | |
| // This is safe: it still requires a registered source decoder and rejects | |
| // unknown sources rather than decoding their payloads through a plan codec. | |
| DataSourceExec::register_proto(registry)?; | |
| registry.register_decoder::<dyn DataSource, T, DataSourceDecoder>( | |
| T::NAME, | |
| T::try_from_proto, | |
| ) | |
| } | |
| if T::NAME.is_empty() { | |
| return plan_err!("extension DataSource name must not be empty"); | |
| } | |
| registry.register_decoder::<dyn DataSource, T, DataSourceDecoder>( | |
| T::NAME, | |
| T::try_from_proto, | |
| )?; | |
| DataSourceExec::register_proto(registry) | |
| } |
A collision can still leave the wrapper unregistered, but that is harmless. The precursor removes this problem: no wrapper registration is needed.
|
|
||
| /// Every registered extension source name, in arbitrary order. | ||
| /// Names registered for other extension kinds are not included. | ||
| pub fn data_source_names(registry: &ProtoDecoderRegistry) -> impl Iterator<Item = &str> { |
There was a problem hiding this comment.
Is this needed in the public API? Outside tests, it is used only in the error message at L183. ProtoDecoderRegistry::names::<dyn DataSource>() is already public.
| pub fn data_source_names(registry: &ProtoDecoderRegistry) -> impl Iterator<Item = &str> { | |
| pub(crate) fn data_source_names(registry: &ProtoDecoderRegistry) -> impl Iterator<Item = &str> { |
| .as_ref() | ||
| .map(|registry| data_source_names(registry).collect()) | ||
| .unwrap_or_default(); | ||
| names.sort_unstable(); |
There was a problem hiding this comment.
Mutation testing: deleting this line survives. No test asserts the order of names in the message. Add an assert_eq! on the full message in the test that registers two names.
| /// The same imports compile without attempting to register a built-in: | ||
| /// | ||
| /// ``` | ||
| /// use datafusion_datasource::memory::MemorySourceConfig; | ||
| /// #[allow(unused_imports)] | ||
| /// use datafusion_datasource::proto::{ProtoDecoderRegistry, register_data_source}; | ||
| /// let _ = (MemorySourceConfig::try_from_proto, ProtoDecoderRegistry::new()); | ||
| /// ``` |
There was a problem hiding this comment.
This doctest exists only to use the imports from the compile_fail doctest above. I suggest removing it. The compile_fail doctest alone shows the point.
| /// [`FileGroup`], gated on the `proto` feature. | ||
| #[cfg(feature = "proto")] | ||
| mod proto; | ||
| pub mod proto; |
There was a problem hiding this comment.
This makes FileRange / PartitionedFile / FileGroup conversions public as a side effect. They become semver surface.
Keep mod proto; private and re-export only the trait (for example, from crate::source)? Also, pub mod registry in proto.rs gives a second path to the same items.
| message PhysicalExtensionDataSourceNode { | ||
| bytes node = 1; | ||
| // Namespaced source type name. The decoding session must register a source | ||
| // decoder for this name; source payloads are never passed to a plan codec. | ||
| optional string source_name = 2; |
There was a problem hiding this comment.
optional, but the decoder rejects a missing or empty value. Writerequired in practicein the comment, or use a plainstring.bytes nodeholds the source payload, not a node. Call itpayload.
The precursor's ExtensionDataSourceNode { string name = 1; bytes payload = 2; } makes both changes.
Which issue does this PR close?
Follow-up to #24631 and #24628. No separate issue.
Rationale for this change
Custom
DataSourceimplementations can encode themselves throughtry_to_proto, but still need aPhysicalExtensionCodecto decode. Applications combining sources from multiple libraries therefore have to compose their codecs.Extend the session-scoped decoder registry to sources so each source can own its decoder and register it by name, following the extension execution-plan and expression APIs.
What changes are included in this PR?
ExtensionDataSourceFromProto,register_data_source, andextension_data_source_node. Constructors receive a typedPhysicalExtensionDataSourceNodeand returnArc<dyn DataSource>.DataSourceExecas a shared extension-plan decoder. Its payload contains the source name and bytes; the wrapper looks up the source decoder and rebuildsDataSourceExecthrough the existing plan registry dispatch. No new top-level plan variant is needed. Source names are independent of plan and expression names.What is the testing strategy for this PR?
Six registry unit tests and fourteen integration tests cover registration, execution after decoding, reading-session state, malformed wrapper and source payloads, codec dispatch boundaries, legacy codec compatibility, shared names across extension kinds, and binary/JSON round trips.
Validation: 541 affected-crate tests and 14 doctests passed; formatting, workspace-wide Clippy (
cargo clippy --all-targets --all-features -- -D warnings), rustdoc with warnings denied, protobuf regeneration, and documentation formatting passed.Are there any user-facing changes?
New source-registry APIs behind the existing
protofeature. Built-in source serialization is unchanged. Readers must register source decoders before writers emit the new wrapper payload. Documented in the 56.0.0 upgrade guide.AI-assisted implementation with Codex.