kosiew commented on code in PR #24094:
URL: https://github.com/apache/datafusion/pull/24094#discussion_r3719961769
##########
datafusion/physical-plan/src/union.rs:
##########
@@ -141,6 +135,179 @@ fn conform_stream_schema(
}
}
+/// Coerces a single child's declared output schema to `schema`, re-stamping
+/// every batch it produces to match. [`UnionExec::try_new`] and
+/// [`InterleaveExec::try_new`] insert this above any child whose own schema
+/// disagrees with the computed union schema, so that the coercion is visible
+/// in the plan tree (e.g. in `EXPLAIN`) instead of happening invisibly inside
+/// the union operator's own `execute()`.
+///
+/// A genuine data type mismatch (as opposed to a nullability-only one) is
+/// rejected eagerly, at construction time, by `EquivalenceProperties::
+/// with_new_schema` below -- unlike the old purely-runtime approach, a
+/// hand-built union/interleave with mismatched child types now fails in
+/// `try_new` rather than in `execute()`.
+///
+/// This node is a strict 1:1, order-preserving passthrough of `input` (it
+/// only ever changes a batch's declared schema, never its rows), so every
+/// `ExecutionPlan` method below that isn't about the schema itself just
+/// delegates straight to `input`.
+///
+/// See <https://github.com/apache/datafusion/issues/15394>.
+#[derive(Debug)]
+struct CoerceSchemaExec {
+ input: Arc<dyn ExecutionPlan>,
+ cache: Arc<PlanProperties>,
+ metrics: ExecutionPlanMetricsSet,
+}
+
+impl CoerceSchemaExec {
+ /// Wraps `input` in a [`CoerceSchemaExec`] targeting `schema` if its own
+ /// schema disagrees with `schema`, otherwise returns it unchanged.
+ fn wrap_if_needed(
+ input: Arc<dyn ExecutionPlan>,
+ schema: &SchemaRef,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ if &input.schema() == schema {
+ Ok(input)
+ } else {
+ Ok(Arc::new(Self::new(input, schema)?))
+ }
+ }
+
+ fn new(input: Arc<dyn ExecutionPlan>, schema: &SchemaRef) -> Result<Self> {
+ let eq_properties = input
+ .equivalence_properties()
+ .clone()
+ .with_new_schema(Arc::clone(schema))?;
+ let output_partitioning = input.output_partitioning().clone();
+ let cache = PlanProperties::new(
+ eq_properties,
+ output_partitioning,
+ emission_type_from_children(std::iter::once(&input)),
+ boundedness_from_children(std::iter::once(&input)),
+ );
+ Ok(Self {
+ input,
+ cache: Arc::new(cache),
+ metrics: ExecutionPlanMetricsSet::new(),
+ })
+ }
+}
+
+impl DisplayAs for CoerceSchemaExec {
+ fn fmt_as(
+ &self,
+ t: DisplayFormatType,
+ f: &mut std::fmt::Formatter,
+ ) -> std::fmt::Result {
+ match t {
+ DisplayFormatType::Default | DisplayFormatType::Verbose => {
+ write!(f, "CoerceSchemaExec")
+ }
+ DisplayFormatType::TreeRender => Ok(()),
+ }
+ }
+}
+
+impl ExecutionPlan for CoerceSchemaExec {
+ fn name(&self) -> &'static str {
+ "CoerceSchemaExec"
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.cache
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![&self.input]
+ }
+
+ fn maintains_input_order(&self) -> Vec<bool> {
+ vec![true]
+ }
+
+ // A 1:1 passthrough never combines partitions, so re-deriving the cache
+ // (rather than collapsing back to the raw child via `wrap_if_needed`) is
+ // always safe here -- it just keeps this node from vanishing and
+ // changing arity out from under a caller mid-rewrite.
+ fn with_new_children(
+ self: Arc<Self>,
+ mut children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ assert_or_internal_err!(
+ children.len() == 1,
+ "CoerceSchemaExec expects exactly one child"
+ );
+ Ok(Arc::new(Self::new(children.remove(0), &self.schema())?))
+ }
+
+ fn execute(
+ &self,
+ partition: usize,
+ context: Arc<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ let baseline_metrics = BaselineMetrics::new(&self.metrics, partition);
+ let stream = self.input.execute(partition, context)?;
+ let stream = conform_stream_schema(self.schema(), stream);
+ Ok(Box::pin(ObservedStream::new(
+ stream,
+ baseline_metrics,
+ None,
+ )))
+ }
+
+ fn metrics(&self) -> Option<MetricsSet> {
+ Some(self.metrics.clone_inner())
+ }
+
+ fn benefits_from_input_partitioning(&self) -> Vec<bool> {
+ vec![false]
+ }
+
+ fn supports_limit_pushdown(&self) -> bool {
+ true
+ }
+
+ fn cardinality_effect(&self) -> CardinalityEffect {
+ CardinalityEffect::Equal
+ }
+
+ fn child_stats_requests(&self, partition: Option<usize>) ->
Vec<ChildStats> {
+ vec![ChildStats::At(partition)]
+ }
+
+ fn statistics_from_inputs(
+ &self,
+ input_stats: &[Arc<Statistics>],
+ _args: &StatisticsArgs,
+ ) -> Result<Arc<Statistics>> {
+ Ok(Arc::clone(&input_stats[0]))
+ }
+
+ fn gather_filters_for_pushdown(
+ &self,
+ _phase: FilterPushdownPhase,
+ parent_filters: Vec<Arc<dyn PhysicalExpr>>,
+ _config: &ConfigOptions,
+ ) -> Result<FilterDescription> {
+ FilterDescription::from_children(parent_filters, &self.children())
+ }
+
+ #[cfg(feature = "proto")]
Review Comment:
Could we add a protobuf round-trip regression test for nullable and
non-nullable `UNION` and `INTERLEAVE` inputs?
`CoerceSchemaExec::try_to_proto` intentionally leaves the wrapper out of the
serialized plan, while `UnionExec::try_from_proto` and
`InterleaveExec::try_from_proto` rebuild it through `try_new`. It would be
helpful to verify that the decoded plans contain the expected coercion and that
their emitted batches expose the final nullable schema.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
