Repository navigation
Conversation
github-actions
Bot
added
documentation
adriangb
left a comment
adriangb
left a comment
There was a problem hiding this comment.
This PR is stacked on #26555. The design points from #26555 (review) also apply here.
With the precursor #26824:
PhysicalPlanNode
├─ ... (other ExecutionPlan variants)
└─ data_sink: DataSinkExecNode {input, sort_order}
└─ DataSinkNode (oneof)
├─ json / csv / parquet
└─ extension {name, payload}
DataSinkExecencodesinputandsort_order. The sink encodes only its own fields.- This PR then needs one
ExtensionDataSinkFromPrototrait with&[u8]input, a providedregister_proto, and one decode arm.
| fn try_from_proto( | ||
| node: &PhysicalExtensionDataSinkNode, | ||
| ctx: &ExecutionPlanDecodeCtx<'_>, | ||
| ) -> Result<Arc<dyn DataSink>>; | ||
| } |
There was a problem hiding this comment.
Same point as in #26555: pass payload: &[u8]. sort_order belongs to DataSinkExec, not to the sink.
| fn try_from_proto( | |
| node: &PhysicalExtensionDataSinkNode, | |
| ctx: &ExecutionPlanDecodeCtx<'_>, | |
| ) -> Result<Arc<dyn DataSink>>; | |
| } | |
| fn try_from_proto( | |
| payload: &[u8], | |
| ctx: &ExecutionPlanDecodeCtx<'_>, | |
| ) -> Result<Arc<dyn DataSink>>; |
| pub fn extension_data_sink_node<T: ExtensionDataSinkFromProto>( | ||
| ctx: &ExecutionPlanEncodeCtx<'_>, | ||
| exec: &DataSinkExec, | ||
| payload: Vec<u8>, | ||
| ) -> Result<PhysicalPlanNode> { | ||
| let sink = PhysicalExtensionDataSinkNode { | ||
| node: payload, | ||
| sink_name: Some(T::NAME.to_string()), | ||
| sort_order: exec.encode_sort_order(ctx)?, | ||
| }; | ||
| ctx.extension_node::<DataSinkExec>(sink.encode_to_vec(), vec![exec.input()]) | ||
| } |
There was a problem hiding this comment.
The sink receives &DataSinkExec only to encode the wrapper's input and sort order. That gives the sink access to its parent.
With DataSink::try_to_proto_sink -> DataSinkNode, DataSinkExec encodes input and sort_order, and this function can be removed.
| // The shared wrapper may remain registered if sink registration fails. | ||
| // This is safe: it still requires a registered sink decoder and rejects | ||
| // unknown sinks rather than decoding their payloads through a plan codec. | ||
| DataSinkExec::register_proto(registry)?; | ||
| registry | ||
| .register_decoder::<dyn DataSink, T, DataSinkDecoder>(T::NAME, T::try_from_proto) | ||
| } |
There was a problem hiding this comment.
Same as #26555: validate first, then register the wrapper.
| // The shared wrapper may remain registered if sink registration fails. | |
| // This is safe: it still requires a registered sink decoder and rejects | |
| // unknown sinks rather than decoding their payloads through a plan codec. | |
| DataSinkExec::register_proto(registry)?; | |
| registry | |
| .register_decoder::<dyn DataSink, T, DataSinkDecoder>(T::NAME, T::try_from_proto) | |
| } | |
| if T::NAME.is_empty() { | |
| return plan_err!("extension DataSink name must not be empty"); | |
| } | |
| registry | |
| .register_decoder::<dyn DataSink, T, DataSinkDecoder>(T::NAME, T::try_from_proto)?; | |
| DataSinkExec::register_proto(registry) | |
| } |
|
|
||
| /// Every registered extension sink name, in arbitrary order. | ||
| /// Names registered for other extension kinds are not included. | ||
| pub fn data_sink_names(registry: &ProtoDecoderRegistry) -> impl Iterator<Item = &str> { |
There was a problem hiding this comment.
Same as data_source_names: used outside tests only by the error message at L169.
| pub fn data_sink_names(registry: &ProtoDecoderRegistry) -> impl Iterator<Item = &str> { | |
| pub(crate) fn data_sink_names(registry: &ProtoDecoderRegistry) -> impl Iterator<Item = &str> { |
| let mut names: Vec<&str> = registry | ||
| .as_ref() | ||
| .map(|registry| data_sink_names(registry).collect()) | ||
| .unwrap_or_default(); | ||
| names.sort_unstable(); | ||
| let names = if names.is_empty() { | ||
| "none".to_string() | ||
| } else { | ||
| names.join(", ") | ||
| }; | ||
| plan_err!( | ||
| "No decoder is registered for the extension DataSink '{name}'. Register \ | ||
| the sink in the ProtoDecoderRegistry attached to the decoding session's \ | ||
| SessionConfig. Registered extension data sinks: {names}" | ||
| ) | ||
| } |
There was a problem hiding this comment.
This duplicates registry.rs L180–197 except for the noun. Put it in one shared helper, for example fn unknown_extension_err(kind: &str, name: &str, names: impl Iterator<Item = &str>) -> DataFusionError.
| pub mod registry; | ||
| pub mod sink_registry; |
There was a problem hiding this comment.
Same as #26555: keep these modules private. All items are also re-exported below, so this is a second public path to the same items.
| bytes node = 1; | ||
| // Namespaced sink type name. Once the DataSinkExec wrapper decoder is | ||
| // registered, sink lookup and decoding errors do not fall back to a plan codec. | ||
| optional string sink_name = 2; | ||
| PhysicalSortExprNodeCollection sort_order = 3; |
There was a problem hiding this comment.
sort_order belongs to DataSinkExec, not to the sink payload. In the precursor (#26824) it moves to DataSinkExecNode, and the sink message is ExtensionDataSinkNode { string name = 1; bytes payload = 2; }.
Which issue does this PR close?
Stacked on #26555; merge that PR first. No separate issue.
The branch includes the rebased DataSource change. Sink-only diff. The upstream PR targets
mainbecause GitHub cannot use a fork branch as the base of an upstream PR.Rationale for this change
Custom
DataSinkimplementations can encode themselves throughtry_to_proto, but decoding still requires an application-ownedPhysicalExtensionCodec. Sinks from independent libraries therefore need codec composition.Give sinks the same session-scoped, per-type decoder registration as sources in #26555.
What changes are included in this PR?
ExtensionDataSinkFromProto,register_data_sink,data_sink_names, andextension_data_sink_nodealongside the source APIs.DataSinkExecas the shared extension-plan decoder. Its typed payload contains the sink name, sink bytes, and optional sort requirements; the existing extension-plan node carries exactly one input plan. DataFusion reconstructs the input, ordering, and execution-plan wrapper around the decoded sink.What is the testing strategy for this PR?
Five registration unit tests and twelve integration tests cover receiving-session state and actual writes after a binary round trip, input and ordering preservation, independent sink types, name isolation, invalid payloads/inputs/orderings, codec dispatch boundaries, legacy codecs, and JSON serialization (with the
jsonfeature).Validation: 570 affected-crate tests and 14 doctests passed. Formatting, workspace Clippy with warnings denied, rustdoc with warnings denied, protobuf regeneration, and documentation formatting completed. The broader
dev/rust_lint.shsuite could not start because/usr/bin/python3lacks PyYAML.Are there any user-facing changes?
New opt-in sink registration APIs behind the existing
protofeature, documented in the 56.0.0 upgrade guide. Register sinks on every reader before writers emit the new wrapper payload. Existing built-in sink formats are unchanged.AI-assisted implementation with Codex.