datafusion_physical_plan/execution_plan.rs
1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements. See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership. The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License. You may obtain a copy of the License at
8//
9// http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied. See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18pub use crate::display::{DefaultDisplay, DisplayAs, DisplayFormatType, VerboseDisplay};
19use crate::distribution_requirements::InputDistributionRequirements;
20use crate::filter_pushdown::{
21 ChildPushdownResult, FilterDescription, FilterPushdownPhase,
22 FilterPushdownPropagation,
23};
24pub use crate::metrics::Metric;
25pub use crate::ordering::InputOrderMode;
26use crate::sort_pushdown::SortOrderPushdownResult;
27pub use crate::stream::EmptyRecordBatchStream;
28
29use arrow_schema::Schema;
30pub use datafusion_common::hash_utils;
31use datafusion_common::tree_node::{
32 Transformed, TransformedResult, TreeNode, TreeNodeRecursion,
33};
34pub use datafusion_common::utils::project_schema;
35pub use datafusion_common::{ColumnStatistics, Statistics, internal_err};
36pub use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream};
37pub use datafusion_expr::{Accumulator, ColumnarValue};
38use datafusion_physical_expr::projection::ProjectionExpr;
39pub use datafusion_physical_expr::window::WindowExpr;
40pub use datafusion_physical_expr::{
41 Distribution, Partitioning, PhysicalExpr, expressions,
42};
43
44use std::any::Any;
45use std::collections::HashSet;
46use std::fmt::Debug;
47use std::sync::{Arc, LazyLock};
48
49use crate::coalesce_partitions::CoalescePartitionsExec;
50use crate::display::DisplayableExecutionPlan;
51use crate::metrics::MetricsSet;
52use crate::projection::ProjectionExec;
53use crate::repartition::RepartitionExec;
54use crate::sorts::sort_preserving_merge::SortPreservingMergeExec;
55use crate::statistics::{ChildStats, StatisticsArgs};
56use crate::stream::RecordBatchStreamAdapter;
57
58use arrow::array::{Array, RecordBatch};
59use arrow::datatypes::SchemaRef;
60use datafusion_common::config::ConfigOptions;
61use datafusion_common::{
62 Constraints, DataFusionError, Result, assert_eq_or_internal_err,
63 assert_or_internal_err, exec_err,
64};
65use datafusion_common_runtime::JoinSet;
66use datafusion_execution::TaskContext;
67use datafusion_physical_expr::EquivalenceProperties;
68use datafusion_physical_expr_common::sort_expr::{
69 LexOrdering, OrderingRequirements, PhysicalSortExpr,
70};
71
72use futures::stream::{StreamExt, TryStreamExt};
73
74/// Represent nodes in the DataFusion Physical Plan.
75///
76/// Calling [`execute`] produces an `async` [`SendableRecordBatchStream`] of
77/// [`RecordBatch`] that incrementally computes a partition of the
78/// `ExecutionPlan`'s output from its input. See [`Partitioning`] for more
79/// details on partitioning.
80///
81/// Methods such as [`Self::schema`] and [`Self::properties`] communicate
82/// properties of the output to the DataFusion optimizer, and methods such as
83/// [`required_input_distribution`] and [`required_input_ordering`] express
84/// requirements of the `ExecutionPlan` from its input.
85///
86/// [`ExecutionPlan`] can be displayed in a simplified form using the
87/// return value from [`displayable`] in addition to the (normally
88/// quite verbose) `Debug` output.
89///
90/// [`execute`]: ExecutionPlan::execute
91/// [`required_input_distribution`]: ExecutionPlan::required_input_distribution
92/// [`required_input_ordering`]: ExecutionPlan::required_input_ordering
93///
94/// # Examples
95///
96/// See [`datafusion-examples`] for examples, including
97/// [`memory_pool_execution_plan.rs`] which shows how to implement a custom
98/// `ExecutionPlan` with memory tracking and spilling support.
99///
100/// [`datafusion-examples`]: https://github.com/apache/datafusion/tree/main/datafusion-examples
101/// [`memory_pool_execution_plan.rs`]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs
102pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync {
103 /// Short name for the ExecutionPlan, such as 'DataSourceExec'.
104 ///
105 /// Implementation note: this method can just proxy to
106 /// [`static_name`](ExecutionPlan::static_name) if no special action is
107 /// needed. It doesn't provide a default implementation like that because
108 /// this method doesn't require the `Sized` constrain to allow a wilder
109 /// range of use cases.
110 fn name(&self) -> &str;
111
112 /// Short name for the ExecutionPlan, such as 'DataSourceExec'.
113 /// Like [`name`](ExecutionPlan::name) but can be called without an instance.
114 fn static_name() -> &'static str
115 where
116 Self: Sized,
117 {
118 let full_name = std::any::type_name::<Self>();
119 let maybe_start_idx = full_name.rfind(':');
120 match maybe_start_idx {
121 Some(start_idx) => &full_name[start_idx + 1..],
122 None => "UNKNOWN",
123 }
124 }
125
126 /// Returns the plan that provides this plan's public
127 /// [`ExecutionPlan`] downcast identity.
128 ///
129 /// This hook is for wrapper nodes that delegate their public downcast
130 /// identity to another plan while adding cross-cutting behavior such as
131 /// instrumentation. The default implementation returns `None`, meaning this
132 /// plan's concrete type is used for type introspection.
133 ///
134 /// Most `ExecutionPlan` implementations should use the default `None`;
135 /// override this only for wrapper plans that intentionally delegate their
136 /// public downcast identity to another plan.
137 ///
138 /// The `is` and `downcast_ref` helpers follow the returned delegate instead
139 /// of checking the current concrete type, making intermediate delegating
140 /// wrappers invisible to normal downcast-based inspection.
141 ///
142 /// Implementations that opt in should return the delegate plan, not `self`.
143 ///
144 /// This is independent from [`Self::children`] and should not be used for
145 /// plan traversal or optimizer rewrites.
146 fn downcast_delegate(&self) -> Option<&dyn ExecutionPlan> {
147 None
148 }
149
150 /// Get the schema for this execution plan
151 fn schema(&self) -> SchemaRef {
152 Arc::clone(self.properties().schema())
153 }
154
155 /// Return properties of the output of the `ExecutionPlan`, such as output
156 /// ordering(s), partitioning information etc.
157 ///
158 /// This information is available via methods on [`ExecutionPlanProperties`]
159 /// trait, which is implemented for all `ExecutionPlan`s.
160 fn properties(&self) -> &Arc<PlanProperties>;
161
162 /// Returns an error if this individual node does not conform to its invariants.
163 /// These invariants are typically only checked in debug mode.
164 ///
165 /// A default set of invariants is provided in the [check_default_invariants] function.
166 /// The default implementation of `check_invariants` calls this function.
167 /// Extension nodes can provide their own invariants.
168 fn check_invariants(&self, check: InvariantLevel) -> Result<()> {
169 check_default_invariants(self, check)
170 }
171
172 /// Returns the dynamic expressions produced by this plan node.
173 ///
174 /// A dynamic expression is produced when this node updates or completes its
175 /// runtime state during execution. Expressions that this node only consumes
176 /// must not be returned. This method is shallow and does not include dynamic
177 /// expressions produced by child plans.
178 ///
179 /// Each returned expression must have a [`PhysicalExpr::expression_id`]
180 /// since all dynamic expressions such as [`DynamicFilterPhysicalExpr`]
181 /// have an expression id.
182 ///
183 /// [`DynamicFilterPhysicalExpr`]: datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr
184 fn dynamic_expressions_produced(&self) -> Vec<Arc<dyn PhysicalExpr>> {
185 Vec::new()
186 }
187
188 /// Specifies simple per-child input distribution requirements.
189 ///
190 /// Deprecated: override [`Self::input_distribution_requirements`] instead.
191 ///
192 /// By default, each child has [`Distribution::UnspecifiedDistribution`].
193 #[deprecated(since = "55.0.0", note = "Use input_distribution_requirements")]
194 fn required_input_distribution(&self) -> Vec<Distribution> {
195 vec![Distribution::UnspecifiedDistribution; self.children().len()]
196 }
197
198 /// Specifies the input distribution requirements for this plan.
199 ///
200 /// The default implementation wraps [`Self::required_input_distribution`].
201 /// Override this method for richer requirements, such as allowing alternate
202 /// satisfaction policies or requiring multiple children to be co-partitioned.
203 /// See [`InputDistributionRequirements`] for details.
204 fn input_distribution_requirements(&self) -> InputDistributionRequirements {
205 #[expect(
206 deprecated,
207 reason = "compatibility shim for external ExecutionPlan implementations"
208 )]
209 InputDistributionRequirements::new(self.required_input_distribution())
210 }
211
212 /// Specifies the ordering required for all of the children of this
213 /// `ExecutionPlan`.
214 ///
215 /// For each child, it's the local ordering requirement within
216 /// each partition rather than the global ordering
217 ///
218 /// NOTE that checking `!is_empty()` does **not** check for a
219 /// required input ordering. Instead, the correct check is that at
220 /// least one entry must be `Some`
221 fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
222 vec![None; self.children().len()]
223 }
224
225 /// Returns `false` if this `ExecutionPlan`'s implementation may reorder
226 /// rows within or between partitions.
227 ///
228 /// For example, Projection, Filter, and Limit maintain the order
229 /// of inputs -- they may transform values (Projection) or not
230 /// produce the same number of rows that went in (Filter and
231 /// Limit), but the rows that are produced go in the same way.
232 ///
233 /// DataFusion uses this metadata to apply certain optimizations
234 /// such as automatically repartitioning correctly.
235 ///
236 /// The default implementation returns `false`
237 ///
238 /// WARNING: if you override this default, you *MUST* ensure that
239 /// the `ExecutionPlan`'s maintains the ordering invariant or else
240 /// DataFusion may produce incorrect results.
241 fn maintains_input_order(&self) -> Vec<bool> {
242 vec![false; self.children().len()]
243 }
244
245 /// Specifies whether the `ExecutionPlan` benefits from increased
246 /// parallelization at its input for each child.
247 ///
248 /// If returns `true`, the `ExecutionPlan` would benefit from partitioning
249 /// its corresponding child (and thus from more parallelism). For
250 /// `ExecutionPlan` that do very little work the overhead of extra
251 /// parallelism may outweigh any benefits
252 ///
253 /// The default implementation returns `true` unless this `ExecutionPlan`
254 /// has signalled it requires a single child input partition.
255 fn benefits_from_input_partitioning(&self) -> Vec<bool> {
256 // By default try to maximize parallelism with more CPUs if
257 // possible
258 self.input_distribution_requirements()
259 .per_child_distributions()
260 .map(|dist| !matches!(dist, Distribution::SinglePartition))
261 .collect()
262 }
263
264 /// Get a list of children `ExecutionPlan`s that act as inputs to this plan.
265 /// The returned list will be empty for leaf nodes such as scans, will contain
266 /// a single value for unary nodes, or two values for binary nodes (such as
267 /// joins).
268 fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>>;
269
270 /// Returns a clone of the existing plan with the children replaced,
271 /// skipping recomputation of plan properties when the options indicate
272 /// the new children's properties are unchanged.
273 ///
274 /// Callers should typically call [`replace_children_if_necessary`] and
275 /// not invoke this method directly.
276 fn replace_children(
277 self: Arc<Self>,
278 children: Vec<Arc<dyn ExecutionPlan>>,
279 options: ReplaceChildrenOptions,
280 ) -> Result<Arc<dyn ExecutionPlan>> {
281 #[expect(deprecated)]
282 match options.children_properties {
283 ChildrenPropertiesMode::Keep => {
284 self.with_new_children_and_same_properties(children)
285 }
286 ChildrenPropertiesMode::Recompute => self.with_new_children(children),
287 }
288 }
289
290 /// Apply a closure `f` to each root expression that this node owns and uses
291 /// during execution, either by evaluating it or updating it dynamically.
292 ///
293 /// An expression must not be visited solely because it describes an input or
294 /// output property, such as cached ordering, partitioning, or equivalence
295 /// metadata. However, these may be traversed indirectly. For example,
296 /// `RepartitionExec` visits the partitioning expressions it evaluates and
297 /// `SortExec` visits the sort expressions it evaluates to order rows.
298 ///
299 /// This method is shallow: it must not visit expression children or expressions
300 /// owned by child execution plans.
301 ///
302 /// Similarly to other [`TreeNode`] APIs, the closure can return
303 /// [`TreeNodeRecursion::Stop`] to stop iteration, otherwise iteration
304 /// should continue. Note that [`TreeNodeRecursion::Continue`] and
305 /// [`TreeNodeRecursion::Jump`] are equivalent because this method is not
306 /// recursive.
307 ///
308 ///
309 /// # Example Usage
310 /// ```
311 /// # use std::sync::Arc;
312 /// # use datafusion_physical_plan::ExecutionPlan;
313 /// # use datafusion_common::tree_node::TreeNodeRecursion;
314 /// # fn example(plan: Arc<dyn ExecutionPlan>) -> datafusion_common::Result<()> {
315 /// // Count the number of expressions
316 /// let mut count = 0;
317 /// plan.apply_expressions(&mut |_expr| {
318 /// count += 1;
319 /// Ok(TreeNodeRecursion::Continue)
320 /// })?;
321 /// # Ok(())
322 /// # }
323 /// ```
324 ///
325 /// # Implementation Examples
326 ///
327 /// ## Node with expressions (e.g., FilterExec, ProjectionExec)
328 ///
329 /// Use [`apply_expression_roots`] to implement this method. It abstracts away the
330 /// [`TreeNodeRecursion`] iteration from implementors.
331 /// ```ignore
332 /// fn apply_expressions(
333 /// &self,
334 /// f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
335 /// ) -> Result<TreeNodeRecursion> {
336 /// apply_expression_roots([&self.predicate], f)
337 /// }
338 /// ```
339 ///
340 /// ## Node with no expressions (e.g., EmptyExec, MemoryExec)
341 /// ```ignore
342 /// fn apply_expressions(
343 /// &self,
344 /// _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
345 /// ) -> Result<TreeNodeRecursion> {
346 /// Ok(TreeNodeRecursion::Continue)
347 /// }
348 /// ```
349 fn apply_expressions(
350 &self,
351 f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
352 ) -> Result<TreeNodeRecursion>;
353
354 /// Deprecated.
355 ///
356 /// DataFusion will remove this method in the future in favor of
357 /// [`ExecutionPlan::replace_children`].
358 ///
359 /// Note that this method is still required by the trait; implementations
360 /// should delegate to [`ExecutionPlan::replace_children`] with
361 /// [`ChildrenPropertiesMode::Recompute`].
362 ///
363 /// # Example Implementation
364 /// ```
365 /// # #![allow(deprecated)]
366 /// # use std::fmt;
367 /// # use std::sync::Arc;
368 /// # use datafusion_common::Result;
369 /// # use datafusion_common::tree_node::TreeNodeRecursion;
370 /// # use datafusion_execution::{SendableRecordBatchStream, TaskContext};
371 /// # use datafusion_physical_expr::PhysicalExpr;
372 /// # use datafusion_physical_plan::{
373 /// # ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan,
374 /// # PlanProperties, ReplaceChildrenOptions,
375 /// # };
376 /// # #[derive(Debug)]
377 /// # struct MyExec {
378 /// # input: Arc<dyn ExecutionPlan>,
379 /// # }
380 /// # impl DisplayAs for MyExec {
381 /// # fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
382 /// # write!(f, "MyExec")
383 /// # }
384 /// # }
385 /// impl ExecutionPlan for MyExec {
386 /// // ...
387 /// # fn name(&self) -> &'static str {
388 /// # "MyExec"
389 /// # }
390 /// # fn properties(&self) -> &Arc<PlanProperties> {
391 /// # self.input.properties()
392 /// # }
393 /// # fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
394 /// # vec![&self.input]
395 /// # }
396 /// # fn apply_expressions(
397 /// # &self,
398 /// # _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
399 /// # ) -> Result<TreeNodeRecursion> {
400 /// # Ok(TreeNodeRecursion::Continue)
401 /// # }
402 /// # fn execute(
403 /// # &self,
404 /// # _partition: usize,
405 /// # _context: Arc<TaskContext>,
406 /// # ) -> Result<SendableRecordBatchStream> {
407 /// # unimplemented!()
408 /// # }
409 /// fn replace_children(
410 /// self: Arc<Self>,
411 /// mut children: Vec<Arc<dyn ExecutionPlan>>,
412 /// _options: ReplaceChildrenOptions,
413 /// ) -> Result<Arc<dyn ExecutionPlan>> {
414 /// Ok(Arc::new(MyExec {
415 /// input: children.swap_remove(0),
416 /// }))
417 /// }
418 ///
419 /// fn with_new_children(
420 /// self: Arc<Self>,
421 /// children: Vec<Arc<dyn ExecutionPlan>>,
422 /// ) -> Result<Arc<dyn ExecutionPlan>> {
423 /// // call into `replace_children` with `ReplaceChildrenOptions`
424 /// self.replace_children(
425 /// children,
426 /// ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
427 /// )
428 /// }
429 /// }
430 /// ```
431 #[deprecated(
432 since = "55.0.0",
433 note = "Use `ExecutionPlan::replace_children` with `ReplaceChildrenOptions`"
434 )]
435 fn with_new_children(
436 self: Arc<Self>,
437 children: Vec<Arc<dyn ExecutionPlan>>,
438 ) -> Result<Arc<dyn ExecutionPlan>>;
439
440 /// Deprecated. Implement [`ExecutionPlan::replace_children`] instead.
441 #[deprecated(
442 since = "55.0.0",
443 note = "Use `ExecutionPlan::replace_children` with `ReplaceChildrenOptions`"
444 )]
445 #[expect(deprecated)]
446 fn with_new_children_and_same_properties(
447 self: Arc<Self>,
448 children: Vec<Arc<dyn ExecutionPlan>>,
449 ) -> Result<Arc<dyn ExecutionPlan>> {
450 self.with_new_children(children)
451 }
452
453 /// Reset any internal state within this [`ExecutionPlan`].
454 ///
455 /// This method is called when an [`ExecutionPlan`] needs to be re-executed,
456 /// such as in recursive queries. Unlike [`ExecutionPlan::replace_children`], this method
457 /// ensures that any stateful components (e.g., [`DynamicFilterPhysicalExpr`])
458 /// are reset to their initial state.
459 ///
460 /// The default implementation simply calls [`ExecutionPlan::replace_children`] with the existing children,
461 /// effectively creating a new instance of the [`ExecutionPlan`] with the same children but without
462 /// necessarily resetting any internal state. Implementations that require resetting of some
463 /// internal state should override this method to provide the necessary logic.
464 ///
465 /// This method should *not* reset state recursively for children, as it is expected that
466 /// it will be called from within a walk of the execution plan tree so that it will be called on each child later
467 /// or was already called on each child.
468 ///
469 /// Note to implementers: unlike [`ExecutionPlan::replace_children`] this method does not accept new children as an argument,
470 /// thus it is expected that any cached plan properties will remain valid after the reset.
471 ///
472 /// [`DynamicFilterPhysicalExpr`]: datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr
473 fn reset_state(self: Arc<Self>) -> Result<Arc<dyn ExecutionPlan>> {
474 let children = self.children().into_iter().cloned().collect();
475 self.replace_children(
476 children,
477 ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep),
478 )
479 }
480
481 /// If supported, attempt to increase the partitioning of this `ExecutionPlan` to
482 /// produce `target_partitions` partitions.
483 ///
484 /// If the `ExecutionPlan` does not support changing its partitioning,
485 /// returns `Ok(None)` (the default).
486 ///
487 /// If the `ExecutionPlan` can increase its partitioning, but not to
488 /// `target_partitions`, it may return an ExecutionPlan with fewer
489 /// partitions. This might happen, for example, if each new partition would
490 /// be too small to be efficiently processed individually.
491 ///
492 /// The DataFusion optimizer attempts to use as many threads as possible by
493 /// repartitioning its inputs to match the target number of threads
494 /// available (`target_partitions`). Some data sources, such as the built in
495 /// CSV and Parquet readers, implement this method as they are able to read
496 /// from their input files in parallel, regardless of how the source data is
497 /// split amongst files.
498 fn repartitioned(
499 &self,
500 _target_partitions: usize,
501 _config: &ConfigOptions,
502 ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
503 Ok(None)
504 }
505
506 /// Begin execution of `partition`, returning a [`Stream`] of
507 /// [`RecordBatch`]es.
508 ///
509 /// # Notes
510 ///
511 /// The `execute` method itself is not `async` but it returns an `async`
512 /// [`futures::stream::Stream`]. This `Stream` should incrementally compute
513 /// the output, `RecordBatch` by `RecordBatch` (in a streaming fashion).
514 /// Most `ExecutionPlan`s should not do any work before the first
515 /// `RecordBatch` is requested from the stream.
516 ///
517 /// [`RecordBatchStreamAdapter`] can be used to convert an `async`
518 /// [`Stream`] into a [`SendableRecordBatchStream`].
519 ///
520 /// Using `async` `Streams` allows for network I/O during execution and
521 /// takes advantage of Rust's built in support for `async` continuations and
522 /// crate ecosystem.
523 ///
524 /// [`Stream`]: futures::stream::Stream
525 /// [`StreamExt`]: futures::stream::StreamExt
526 /// [`TryStreamExt`]: futures::stream::TryStreamExt
527 /// [`RecordBatchStreamAdapter`]: crate::stream::RecordBatchStreamAdapter
528 ///
529 /// # Error handling
530 ///
531 /// Any error that occurs during execution is sent as an `Err` in the output
532 /// stream.
533 ///
534 /// `ExecutionPlan` implementations in DataFusion cancel additional work
535 /// immediately once an error occurs. The rationale is that if the overall
536 /// query will return an error, any additional work such as continued
537 /// polling of inputs will be wasted as it will be thrown away.
538 ///
539 /// # Cancellation / Aborting Execution
540 ///
541 /// The [`Stream`] that is returned must ensure that any allocated resources
542 /// are freed when the stream itself is dropped. This is particularly
543 /// important for [`spawn`]ed tasks or threads. Unless care is taken to
544 /// "abort" such tasks, they may continue to consume resources even after
545 /// the plan is dropped, generating intermediate results that are never
546 /// used.
547 /// Thus, [`spawn`] is disallowed, and instead use [`SpawnedTask`].
548 ///
549 /// To enable timely cancellation, the [`Stream`] that is returned must not
550 /// block the CPU indefinitely and must yield back to the tokio runtime regularly.
551 /// In a typical [`ExecutionPlan`], this automatically happens unless there are
552 /// special circumstances; e.g. when the computational complexity of processing a
553 /// batch is superlinear. See this [general guideline][async-guideline] for more context
554 /// on this point, which explains why one should avoid spending a long time without
555 /// reaching an `await`/yield point in asynchronous runtimes.
556 /// This can be achieved by using the utilities from the [`coop`](crate::coop) module, by
557 /// manually returning [`Poll::Pending`] and setting up wakers appropriately, or by calling
558 /// [`tokio::task::yield_now()`] when appropriate.
559 /// In special cases that warrant manual yielding, determination for "regularly" may be
560 /// made using the [Tokio task budget](https://docs.rs/tokio/latest/tokio/task/coop/index.html),
561 /// a timer (being careful with the overhead-heavy system call needed to take the time), or by
562 /// counting rows or batches.
563 ///
564 /// The [cancellation benchmark] tracks some cases of how quickly queries can
565 /// be cancelled.
566 ///
567 /// For more details see [`SpawnedTask`], [`JoinSet`] and [`RecordBatchReceiverStreamBuilder`]
568 /// for structures to help ensure all background tasks are cancelled.
569 ///
570 /// [`spawn`]: tokio::task::spawn
571 /// [cancellation benchmark]: https://github.com/apache/datafusion/blob/main/benchmarks/README.md#cancellation
572 /// [`JoinSet`]: datafusion_common_runtime::JoinSet
573 /// [`SpawnedTask`]: datafusion_common_runtime::SpawnedTask
574 /// [`RecordBatchReceiverStreamBuilder`]: crate::stream::RecordBatchReceiverStreamBuilder
575 /// [`Poll::Pending`]: std::task::Poll::Pending
576 /// [async-guideline]: https://ryhl.io/blog/async-what-is-blocking/
577 ///
578 /// # Implementation Examples
579 ///
580 /// While `async` `Stream`s have a non trivial learning curve, the
581 /// [`futures`] crate provides [`StreamExt`] and [`TryStreamExt`]
582 /// which help simplify many common operations.
583 ///
584 /// Here are some common patterns:
585 ///
586 /// ## Return Precomputed `RecordBatch`
587 ///
588 /// We can return a precomputed `RecordBatch` as a `Stream`:
589 ///
590 /// ```
591 /// # use std::sync::Arc;
592 /// # use arrow::array::RecordBatch;
593 /// # use arrow::datatypes::SchemaRef;
594 /// # use datafusion_common::Result;
595 /// # use datafusion_execution::{SendableRecordBatchStream, TaskContext};
596 /// # use datafusion_physical_plan::memory::MemoryStream;
597 /// # use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
598 /// struct MyPlan {
599 /// batch: RecordBatch,
600 /// }
601 ///
602 /// impl MyPlan {
603 /// fn execute(
604 /// &self,
605 /// partition: usize,
606 /// context: Arc<TaskContext>,
607 /// ) -> Result<SendableRecordBatchStream> {
608 /// // use functions from futures crate to convert the batch into a stream
609 /// let fut = futures::future::ready(Ok(self.batch.clone()));
610 /// let stream = futures::stream::once(fut);
611 /// Ok(Box::pin(RecordBatchStreamAdapter::new(
612 /// self.batch.schema(),
613 /// stream,
614 /// )))
615 /// }
616 /// }
617 /// ```
618 ///
619 /// ## Lazily (async) Compute `RecordBatch`
620 ///
621 /// We can also lazily compute a `RecordBatch` when the returned `Stream` is polled
622 ///
623 /// ```
624 /// # use std::sync::Arc;
625 /// # use arrow::array::RecordBatch;
626 /// # use arrow::datatypes::SchemaRef;
627 /// # use datafusion_common::Result;
628 /// # use datafusion_execution::{SendableRecordBatchStream, TaskContext};
629 /// # use datafusion_physical_plan::memory::MemoryStream;
630 /// # use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
631 /// struct MyPlan {
632 /// schema: SchemaRef,
633 /// }
634 ///
635 /// /// Returns a single batch when the returned stream is polled
636 /// async fn get_batch() -> Result<RecordBatch> {
637 /// todo!()
638 /// }
639 ///
640 /// impl MyPlan {
641 /// fn execute(
642 /// &self,
643 /// partition: usize,
644 /// context: Arc<TaskContext>,
645 /// ) -> Result<SendableRecordBatchStream> {
646 /// let fut = get_batch();
647 /// let stream = futures::stream::once(fut);
648 /// Ok(Box::pin(RecordBatchStreamAdapter::new(
649 /// self.schema.clone(),
650 /// stream,
651 /// )))
652 /// }
653 /// }
654 /// ```
655 ///
656 /// ## Lazily (async) create a Stream
657 ///
658 /// If you need to create the return `Stream` using an `async` function,
659 /// you can do so by flattening the result:
660 ///
661 /// ```
662 /// # use std::sync::Arc;
663 /// # use arrow::array::RecordBatch;
664 /// # use arrow::datatypes::SchemaRef;
665 /// # use futures::TryStreamExt;
666 /// # use datafusion_common::Result;
667 /// # use datafusion_execution::{SendableRecordBatchStream, TaskContext};
668 /// # use datafusion_physical_plan::memory::MemoryStream;
669 /// # use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
670 /// struct MyPlan {
671 /// schema: SchemaRef,
672 /// }
673 ///
674 /// /// async function that returns a stream
675 /// async fn get_batch_stream() -> Result<SendableRecordBatchStream> {
676 /// todo!()
677 /// }
678 ///
679 /// impl MyPlan {
680 /// fn execute(
681 /// &self,
682 /// partition: usize,
683 /// context: Arc<TaskContext>,
684 /// ) -> Result<SendableRecordBatchStream> {
685 /// // A future that yields a stream
686 /// let fut = get_batch_stream();
687 /// // Use TryStreamExt::try_flatten to flatten the stream of streams
688 /// let stream = futures::stream::once(fut).try_flatten();
689 /// Ok(Box::pin(RecordBatchStreamAdapter::new(
690 /// self.schema.clone(),
691 /// stream,
692 /// )))
693 /// }
694 /// }
695 /// ```
696 fn execute(
697 &self,
698 partition: usize,
699 context: Arc<TaskContext>,
700 ) -> Result<SendableRecordBatchStream>;
701
702 /// Return a snapshot of the set of [`Metric`]s for this
703 /// [`ExecutionPlan`]. If no `Metric`s are available, return None.
704 ///
705 /// While the values of the metrics in the returned
706 /// [`MetricsSet`]s may change as execution progresses, the
707 /// specific metrics will not.
708 ///
709 /// Once `self.execute()` has returned (technically the future is
710 /// resolved) for all available partitions, the set of metrics
711 /// should be complete. If this function is called prior to
712 /// `execute()` new metrics may appear in subsequent calls.
713 fn metrics(&self) -> Option<MetricsSet> {
714 None
715 }
716
717 /// Returns statistics for a specific partition of this `ExecutionPlan` node.
718 ///
719 /// Deprecated: use [`StatisticsContext::compute`] instead.
720 ///
721 /// [`StatisticsContext::compute`]: crate::statistics::StatisticsContext::compute
722 #[deprecated(since = "55.0.0", note = "Use StatisticsContext::compute instead")]
723 fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
724 if let Some(idx) = partition {
725 // Validate partition index
726 let partition_count = self.properties().partitioning.partition_count();
727 assert_or_internal_err!(
728 idx < partition_count,
729 "Invalid partition index: {}, the partition count is {}",
730 idx,
731 partition_count
732 );
733 }
734 Ok(Arc::new(Statistics::new_unknown(&self.schema())))
735 }
736
737 /// Returns statistics for a specific partition of this `ExecutionPlan` node,
738 /// given pre-computed child statistics.
739 ///
740 /// If statistics are not available, should return [`Statistics::new_unknown`]
741 /// (the default), not an error.
742 /// If `args.partition()` is `None`, it returns statistics for all partitions.
743 ///
744 /// Implementations should not call [`StatisticsContext::compute`] from within
745 /// this method; child statistics are provided via `input_stats`.
746 ///
747 /// Use [`StatisticsContext::compute`] to initiate a full plan-tree walk.
748 ///
749 /// [`StatisticsContext::compute`]: crate::statistics::StatisticsContext::compute
750 fn statistics_from_inputs(
751 &self,
752 _input_stats: &[Arc<Statistics>],
753 args: &StatisticsArgs,
754 ) -> Result<Arc<Statistics>> {
755 #[expect(deprecated)]
756 self.partition_statistics(args.partition())
757 }
758
759 /// Returns, per child, which statistics the [`StatisticsContext`] should resolve
760 /// before calling [`Self::statistics_from_inputs`].
761 ///
762 /// One entry per child (same order as [`Self::children`]): [`ChildStats::At`]
763 /// requests the child's statistics at a partition (`None` = overall);
764 /// [`ChildStats::Skip`] omits a child whose statistics this node does not need
765 /// (a `Statistics::new_unknown` placeholder fills its `input_stats` slot).
766 ///
767 /// The default skips every child, so a node that derives nothing from its
768 /// children (for example one that only overrides the deprecated
769 /// [`Self::partition_statistics`]) triggers no child traversal. A node that reads
770 /// `input_stats` in [`Self::statistics_from_inputs`] must override this to declare
771 /// the children it uses.
772 ///
773 /// [`StatisticsContext`]: crate::statistics::StatisticsContext
774 fn child_stats_requests(&self, _partition: Option<usize>) -> Vec<ChildStats> {
775 self.children().iter().map(|_| ChildStats::Skip).collect()
776 }
777
778 /// Returns `true` if a limit can be safely pushed down through this
779 /// `ExecutionPlan` node.
780 ///
781 /// If this method returns `true`, and the query plan contains a limit at
782 /// the output of this node, DataFusion will push the limit to the input
783 /// of this node.
784 fn supports_limit_pushdown(&self) -> bool {
785 false
786 }
787
788 /// Returns a fetching variant of this `ExecutionPlan` node, if it supports
789 /// fetch limits. Returns `None` otherwise.
790 ///
791 /// See physical optimizer rule [`limit_pushdown`] for details.
792 ///
793 /// [`limit_pushdown`]: https://docs.rs/datafusion/latest/datafusion/physical_optimizer/limit_pushdown/index.html
794 fn with_fetch(&self, _limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> {
795 None
796 }
797
798 /// Gets the fetch count for the operator, `None` means there is no fetch.
799 fn fetch(&self) -> Option<usize> {
800 None
801 }
802
803 /// Gets the effect on cardinality, if known
804 fn cardinality_effect(&self) -> CardinalityEffect {
805 CardinalityEffect::Unknown
806 }
807
808 /// Attempts to push down the given projection into the input of this `ExecutionPlan`.
809 ///
810 /// If the operator supports this optimization, the resulting plan will be:
811 /// `self_new <- projection <- source`, starting from `projection <- self <- source`.
812 /// Otherwise, it returns the current `ExecutionPlan` as-is.
813 ///
814 /// Returns `Ok(Some(...))` if pushdown is applied, `Ok(None)` if it is not supported
815 /// or not possible, or `Err` on failure.
816 fn try_swapping_with_projection(
817 &self,
818 _projection: &ProjectionExec,
819 ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
820 Ok(None)
821 }
822
823 /// Collect filters that this node can push down to its children.
824 /// Filters that are being pushed down from parents are passed in,
825 /// and the node may generate additional filters to push down.
826 /// For example, given the plan FilterExec -> HashJoinExec -> DataSourceExec,
827 /// what will happen is that we recurse down the plan calling `ExecutionPlan::gather_filters_for_pushdown`:
828 /// 1. `FilterExec::gather_filters_for_pushdown` is called with no parent
829 /// filters so it only returns that `FilterExec` wants to push down its own predicate.
830 /// 2. `HashJoinExec::gather_filters_for_pushdown` is called with the filter from
831 /// `FilterExec`, which it only allows to push down to one side of the join (unless it's on the join key)
832 /// but it also adds its own filters (e.g. pushing down a bloom filter of the hash table to the scan side of the join).
833 /// 3. `DataSourceExec::gather_filters_for_pushdown` is called with both filters from `HashJoinExec`
834 /// and `FilterExec`, however `DataSourceExec::gather_filters_for_pushdown` doesn't actually do anything
835 /// since it has no children and no additional filters to push down.
836 /// It's only once [`ExecutionPlan::handle_child_pushdown_result`] is called on `DataSourceExec` as we recurse
837 /// up the plan that `DataSourceExec` can actually bind the filters.
838 ///
839 /// The default implementation bars all parent filters from being pushed down and adds no new filters.
840 /// This is the safest option, making filter pushdown opt-in on a per-node basis.
841 ///
842 /// There are two different phases in filter pushdown, which some operators may handle the same and some differently.
843 /// Depending on the phase the operator may or may not be allowed to modify the plan.
844 /// See [`FilterPushdownPhase`] for more details.
845 ///
846 /// Implementations must preserve the order of `parent_filters` in the
847 /// returned child [`FilterDescription`]: each child parent-filter result is
848 /// matched back to the corresponding input parent filter by position.
849 /// Unsupported filters should therefore be marked unsupported in place,
850 /// rather than removed or appended after supported filters.
851 fn gather_filters_for_pushdown(
852 &self,
853 _phase: FilterPushdownPhase,
854 parent_filters: Vec<Arc<dyn PhysicalExpr>>,
855 _config: &ConfigOptions,
856 ) -> Result<FilterDescription> {
857 Ok(FilterDescription::all_unsupported(
858 &parent_filters,
859 &self.children(),
860 ))
861 }
862
863 /// Handle the result of a child pushdown.
864 ///
865 /// This method is called as we recurse back up the plan tree after pushing
866 /// filters down to child nodes via [`ExecutionPlan::gather_filters_for_pushdown`].
867 /// It allows the current node to process the results of filter pushdown from
868 /// its children, deciding whether to absorb filters, modify the plan, or pass
869 /// filters back up to its parent.
870 ///
871 /// **Purpose and Context:**
872 /// Filter pushdown is a critical optimization in DataFusion that aims to
873 /// reduce the amount of data processed by applying filters as early as
874 /// possible in the query plan. This method is part of the second phase of
875 /// filter pushdown, where results are propagated back up the tree after
876 /// being pushed down. Each node can inspect the pushdown results from its
877 /// children and decide how to handle any unapplied filters, potentially
878 /// optimizing the plan structure or filter application.
879 ///
880 /// **Behavior in Different Nodes:**
881 /// - For a `DataSourceExec`, this often means absorbing the filters to apply
882 /// them during the scan phase (late materialization), reducing the data
883 /// read from the source.
884 /// - A `FilterExec` may absorb any filters its children could not handle,
885 /// combining them with its own predicate. If no filters remain (i.e., the
886 /// predicate becomes trivially true), it may remove itself from the plan
887 /// altogether. It typically marks parent filters as supported, indicating
888 /// they have been handled.
889 /// - A `HashJoinExec` might ignore the pushdown result if filters need to
890 /// be applied during the join operation. It passes the parent filters back
891 /// up wrapped in [`FilterPushdownPropagation::if_any`], discarding
892 /// any self-filters from children.
893 ///
894 /// **Example Walkthrough:**
895 /// Consider a query plan: `FilterExec (f1) -> HashJoinExec -> DataSourceExec`.
896 /// 1. **Downward Phase (`gather_filters_for_pushdown`):** Starting at
897 /// `FilterExec`, the filter `f1` is gathered and pushed down to
898 /// `HashJoinExec`. `HashJoinExec` may allow `f1` to pass to one side of
899 /// the join or add its own filters (e.g., a min-max filter from the build side),
900 /// then pushes filters to `DataSourceExec`. `DataSourceExec`, being a leaf node,
901 /// has no children to push to, so it prepares to handle filters in the
902 /// upward phase.
903 /// 2. **Upward Phase (`handle_child_pushdown_result`):** Starting at
904 /// `DataSourceExec`, it absorbs applicable filters from `HashJoinExec`
905 /// for late materialization during scanning, marking them as supported.
906 /// `HashJoinExec` receives the result, decides whether to apply any
907 /// remaining filters during the join, and passes unhandled filters back
908 /// up to `FilterExec`. `FilterExec` absorbs any unhandled filters,
909 /// updates its predicate if necessary, or removes itself if the predicate
910 /// becomes trivial (e.g., `lit(true)`), and marks filters as supported
911 /// for its parent.
912 ///
913 /// The default implementation is a no-op that passes the result of pushdown
914 /// from the children to its parent transparently, ensuring no filters are
915 /// lost if a node does not override this behavior.
916 ///
917 /// **Notes for Implementation:**
918 /// When returning filters via [`FilterPushdownPropagation`], the order of
919 /// filters need not match the order they were passed in via
920 /// `child_pushdown_result`. However, preserving the order is recommended for
921 /// debugging and ease of reasoning about the resulting plans.
922 ///
923 /// **Helper Methods for Customization:**
924 /// There are various helper methods to simplify implementing this method:
925 /// - [`FilterPushdownPropagation::if_any`]: Marks all parent filters as
926 /// supported as long as at least one child supports them.
927 /// - [`FilterPushdownPropagation::if_all`]: Marks all parent filters as
928 /// supported as long as all children support them.
929 /// - [`FilterPushdownPropagation::with_parent_pushdown_result`]: Allows adding filters
930 /// to the propagation result, indicating which filters are supported by
931 /// the current node.
932 /// - [`FilterPushdownPropagation::with_updated_node`]: Allows updating the
933 /// current node in the propagation result, used if the node
934 /// has modified its plan based on the pushdown results.
935 ///
936 /// **Filter Pushdown Phases:**
937 /// There are two different phases in filter pushdown (`Pre` and others),
938 /// which some operators may handle differently. Depending on the phase, the
939 /// operator may or may not be allowed to modify the plan. See
940 /// [`FilterPushdownPhase`] for more details on phase-specific behavior.
941 ///
942 /// [`PushedDownPredicate::supported`]: crate::filter_pushdown::PushedDownPredicate::supported
943 fn handle_child_pushdown_result(
944 &self,
945 _phase: FilterPushdownPhase,
946 child_pushdown_result: ChildPushdownResult,
947 _config: &ConfigOptions,
948 ) -> Result<FilterPushdownPropagation<Arc<dyn ExecutionPlan>>> {
949 Ok(FilterPushdownPropagation::if_all(child_pushdown_result))
950 }
951
952 /// Injects arbitrary run-time state into this execution plan, returning a new plan
953 /// instance that incorporates that state *if* it is relevant to the concrete
954 /// node implementation.
955 ///
956 /// This is a generic entry point: the `state` can be any type wrapped in
957 /// `Arc<dyn Any + Send + Sync>`. A node that cares about the state should
958 /// down-cast it to the concrete type it expects and, if successful, return a
959 /// modified copy of itself that captures the provided value. If the state is
960 /// not applicable, the default behaviour is to return `None` so that parent
961 /// nodes can continue propagating the attempt further down the plan tree.
962 ///
963 /// For example, [`WorkTableExec`](crate::work_table::WorkTableExec)
964 /// down-casts the supplied state to an `Arc<WorkTable>`
965 /// in order to wire up the working table used during recursive-CTE execution.
966 /// Similar patterns can be followed by custom nodes that need late-bound
967 /// dependencies or shared state.
968 fn with_new_state(
969 &self,
970 _state: Arc<dyn Any + Send + Sync>,
971 ) -> Option<Arc<dyn ExecutionPlan>> {
972 None
973 }
974
975 /// Try to push down sort ordering requirements to this node.
976 ///
977 /// This method is called during sort pushdown optimization to determine if this
978 /// node can optimize for a requested sort ordering. Implementations should:
979 ///
980 /// - Return [`SortOrderPushdownResult::Exact`] if the node can guarantee the exact
981 /// ordering (allowing the Sort operator to be removed)
982 /// - Return [`SortOrderPushdownResult::Inexact`] if the node can optimize for the
983 /// ordering but cannot guarantee perfect sorting (Sort operator is kept)
984 /// - Return [`SortOrderPushdownResult::Unsupported`] if the node cannot optimize
985 /// for the ordering
986 ///
987 /// For transparent nodes (that preserve ordering), implement this to delegate to
988 /// children and wrap the result with a new instance of this node.
989 ///
990 /// Default implementation returns `Unsupported`.
991 fn try_pushdown_sort(
992 &self,
993 _order: &[PhysicalSortExpr],
994 ) -> Result<SortOrderPushdownResult<Arc<dyn ExecutionPlan>>> {
995 Ok(SortOrderPushdownResult::Unsupported)
996 }
997
998 /// Returns a variant of this `ExecutionPlan` that is aware of order-sensitivity.
999 ///
1000 /// This is used to signal to data sources that the output ordering must be
1001 /// preserved, even if it might be more efficient to ignore it (e.g. by
1002 /// skipping some row groups in Parquet).
1003 ///
1004 fn with_preserve_order(
1005 &self,
1006 _preserve_order: bool,
1007 ) -> Option<Arc<dyn ExecutionPlan>> {
1008 None
1009 }
1010
1011 /// Serialize this plan to its protobuf representation, if it knows how.
1012 ///
1013 /// This is the `ExecutionPlan` analog of
1014 /// [`PhysicalExpr::try_to_proto`].
1015 ///
1016 /// * `Ok(None)` (the default) — "I don't serialize myself"; the caller
1017 /// (`datafusion-proto`) falls back to the central downcast chain. Every
1018 /// un-migrated plan keeps its existing behavior.
1019 /// * `Ok(Some(node))` — fully serialized; the caller must not fall back.
1020 /// * `Err(_)` — a real failure (e.g. a child failed to serialize).
1021 ///
1022 /// Only *self-contained* plans should override this — see [`crate::proto`]
1023 /// for the session-dependency boundary.
1024 #[cfg(feature = "proto")]
1025 fn try_to_proto(
1026 &self,
1027 _ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>,
1028 ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
1029 Ok(None)
1030 }
1031}
1032
1033/// Options for [`ExecutionPlan::replace_children`]
1034#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1035pub struct ReplaceChildrenOptions {
1036 /// Describes how plan properties should be handled for the replacement
1037 /// children.
1038 pub children_properties: ChildrenPropertiesMode,
1039}
1040
1041impl ReplaceChildrenOptions {
1042 /// Create new options for [`ExecutionPlan::replace_children`].
1043 pub const fn new(children_properties: ChildrenPropertiesMode) -> Self {
1044 Self {
1045 children_properties,
1046 }
1047 }
1048}
1049
1050/// Indicates whether the plan properties of the new children must be recomputed.
1051///
1052/// Part of [`ReplaceChildrenOptions`].
1053#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1054pub enum ChildrenPropertiesMode {
1055 /// The plan properties of the new children are identical to the properties
1056 /// of the existing children, so we can skip recomputation.
1057 Keep,
1058 /// The plan properties of the new children are different from the properties
1059 /// of the existing children, so we must recompute the properties from scratch.
1060 Recompute,
1061}
1062
1063/// Allows a type to be treated as a reference to an
1064/// [`Arc<dyn PhysicalExpr>`].
1065///
1066/// Used by [`apply_expression_roots`].
1067pub trait AsPhysicalExprRef {
1068 /// Returns the referenced physical expression.
1069 fn as_physical_expr_ref(&self) -> &Arc<dyn PhysicalExpr>;
1070}
1071
1072/// Allows an [`Arc<dyn PhysicalExpr>`] to be treated as a reference to itself.
1073///
1074/// This is needed because `Arc<dyn PhysicalExpr>` does not implement
1075/// `AsRef<Arc<dyn PhysicalExpr>>`.
1076impl AsPhysicalExprRef for Arc<dyn PhysicalExpr> {
1077 fn as_physical_expr_ref(&self) -> &Arc<dyn PhysicalExpr> {
1078 self
1079 }
1080}
1081
1082/// Allows a [`ProjectionExpr`] to be treated as a reference to its
1083/// [`Arc<dyn PhysicalExpr>`].
1084impl AsPhysicalExprRef for ProjectionExpr {
1085 fn as_physical_expr_ref(&self) -> &Arc<dyn PhysicalExpr> {
1086 self.as_ref()
1087 }
1088}
1089
1090impl<T> AsPhysicalExprRef for &T
1091where
1092 T: AsPhysicalExprRef + ?Sized,
1093{
1094 fn as_physical_expr_ref(&self) -> &Arc<dyn PhysicalExpr> {
1095 (*self).as_physical_expr_ref()
1096 }
1097}
1098
1099/// Applies `f` to a shallow sequence of physical expression roots.
1100///
1101/// [`TreeNodeRecursion::Stop`] stops iteration and is returned immediately.
1102/// [`TreeNodeRecursion::Jump`] is normalized to [`TreeNodeRecursion::Continue`]
1103/// because this function does not visit expression children.
1104pub fn apply_expression_roots<I>(
1105 roots: I,
1106 f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
1107) -> Result<TreeNodeRecursion>
1108where
1109 I: IntoIterator,
1110 I::Item: AsPhysicalExprRef,
1111{
1112 for root in roots {
1113 match f(root.as_physical_expr_ref())? {
1114 TreeNodeRecursion::Stop => return Ok(TreeNodeRecursion::Stop),
1115 TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => {}
1116 }
1117 }
1118 Ok(TreeNodeRecursion::Continue)
1119}
1120
1121/// Returns whether `plan` contains a physical expression with `expression_id`.
1122///
1123/// This traverses both the execution plan and the children of each expression root
1124/// reported by [`ExecutionPlan::apply_expressions`].
1125pub(crate) fn plan_contains_expression_id(
1126 plan: &Arc<dyn ExecutionPlan>,
1127 expression_id: u64,
1128) -> Result<bool> {
1129 let mut found = false;
1130 plan.apply(|node| {
1131 node.apply_expressions(&mut |root| {
1132 root.apply(|expr| {
1133 if expr.expression_id() == Some(expression_id) {
1134 found = true;
1135 Ok(TreeNodeRecursion::Stop)
1136 } else {
1137 Ok(TreeNodeRecursion::Continue)
1138 }
1139 })
1140 })?;
1141
1142 Ok(if found {
1143 TreeNodeRecursion::Stop
1144 } else {
1145 TreeNodeRecursion::Continue
1146 })
1147 })?;
1148 Ok(found)
1149}
1150
1151impl dyn ExecutionPlan {
1152 /// Returns `true` if the plan is of type `T`.
1153 ///
1154 /// If this plan provides a [`ExecutionPlan::downcast_delegate`], delegates
1155 /// to it.
1156 ///
1157 /// Prefer this over `downcast_ref::<T>().is_some()`. Works correctly when
1158 /// called on `Arc<dyn ExecutionPlan>` via auto-deref.
1159 pub fn is<T: ExecutionPlan>(&self) -> bool {
1160 match self.downcast_delegate() {
1161 Some(delegate) => delegate.is::<T>(),
1162 None => (self as &dyn Any).is::<T>(),
1163 }
1164 }
1165
1166 /// Attempts to downcast this plan to a concrete type `T`, returning `None`
1167 /// if the plan is not of that type.
1168 ///
1169 /// If this plan provides a [`ExecutionPlan::downcast_delegate`], delegates
1170 /// to it.
1171 ///
1172 /// Works correctly when called on `Arc<dyn ExecutionPlan>` via auto-deref,
1173 /// unlike `(&arc as &dyn Any).downcast_ref::<T>()` which would attempt to
1174 /// downcast the `Arc` itself.
1175 pub fn downcast_ref<T: ExecutionPlan>(&self) -> Option<&T> {
1176 match self.downcast_delegate() {
1177 Some(delegate) => delegate.downcast_ref::<T>(),
1178 None => (self as &dyn Any).downcast_ref(),
1179 }
1180 }
1181}
1182
1183/// [`ExecutionPlan`] Invariant Level
1184///
1185/// What set of assertions ([Invariant]s) holds for a particular `ExecutionPlan`
1186///
1187/// [Invariant]: https://en.wikipedia.org/wiki/Invariant_(mathematics)#Invariants_in_computer_science
1188#[derive(Clone, Copy)]
1189pub enum InvariantLevel {
1190 /// Invariants that are always true for the [`ExecutionPlan`] node
1191 /// such as the number of expected children.
1192 Always,
1193 /// Invariants that must hold true for the [`ExecutionPlan`] node
1194 /// to be "executable", such as ordering and/or distribution requirements
1195 /// being fulfilled.
1196 Executable,
1197}
1198
1199/// Extension trait provides an easy API to fetch various properties of
1200/// [`ExecutionPlan`] objects based on [`ExecutionPlan::properties`].
1201pub trait ExecutionPlanProperties {
1202 /// Specifies how the output of this `ExecutionPlan` is split into
1203 /// partitions.
1204 fn output_partitioning(&self) -> &Partitioning;
1205
1206 /// If the output of this `ExecutionPlan` within each partition is sorted,
1207 /// returns `Some(keys)` describing the ordering. A `None` return value
1208 /// indicates no assumptions should be made on the output ordering.
1209 ///
1210 /// For example, `SortExec` (obviously) produces sorted output as does
1211 /// `SortPreservingMergeStream`. Less obviously, `Projection` produces sorted
1212 /// output if its input is sorted as it does not reorder the input rows.
1213 fn output_ordering(&self) -> Option<&LexOrdering>;
1214
1215 /// Boundedness information of the stream corresponding to this `ExecutionPlan`.
1216 /// For more details, see [`Boundedness`].
1217 fn boundedness(&self) -> Boundedness;
1218
1219 /// Indicates how the stream of this `ExecutionPlan` emits its results.
1220 /// For more details, see [`EmissionType`].
1221 fn pipeline_behavior(&self) -> EmissionType;
1222
1223 /// Get the [`EquivalenceProperties`] within the plan.
1224 ///
1225 /// Equivalence properties tell DataFusion what columns are known to be
1226 /// equal, during various optimization passes. By default, this returns "no
1227 /// known equivalences" which is always correct, but may cause DataFusion to
1228 /// unnecessarily resort data.
1229 ///
1230 /// If this ExecutionPlan makes no changes to the schema of the rows flowing
1231 /// through it or how columns within each row relate to each other, it
1232 /// should return the equivalence properties of its input. For
1233 /// example, since [`FilterExec`] may remove rows from its input, but does not
1234 /// otherwise modify them, it preserves its input equivalence properties.
1235 /// However, since `ProjectionExec` may calculate derived expressions, it
1236 /// needs special handling.
1237 ///
1238 /// See also [`ExecutionPlan::maintains_input_order`] and [`Self::output_ordering`]
1239 /// for related concepts.
1240 ///
1241 /// [`FilterExec`]: crate::filter::FilterExec
1242 fn equivalence_properties(&self) -> &EquivalenceProperties;
1243}
1244
1245impl ExecutionPlanProperties for Arc<dyn ExecutionPlan> {
1246 fn output_partitioning(&self) -> &Partitioning {
1247 self.properties().output_partitioning()
1248 }
1249
1250 fn output_ordering(&self) -> Option<&LexOrdering> {
1251 self.properties().output_ordering()
1252 }
1253
1254 fn boundedness(&self) -> Boundedness {
1255 self.properties().boundedness
1256 }
1257
1258 fn pipeline_behavior(&self) -> EmissionType {
1259 self.properties().emission_type
1260 }
1261
1262 fn equivalence_properties(&self) -> &EquivalenceProperties {
1263 self.properties().equivalence_properties()
1264 }
1265}
1266
1267impl ExecutionPlanProperties for &dyn ExecutionPlan {
1268 fn output_partitioning(&self) -> &Partitioning {
1269 self.properties().output_partitioning()
1270 }
1271
1272 fn output_ordering(&self) -> Option<&LexOrdering> {
1273 self.properties().output_ordering()
1274 }
1275
1276 fn boundedness(&self) -> Boundedness {
1277 self.properties().boundedness
1278 }
1279
1280 fn pipeline_behavior(&self) -> EmissionType {
1281 self.properties().emission_type
1282 }
1283
1284 fn equivalence_properties(&self) -> &EquivalenceProperties {
1285 self.properties().equivalence_properties()
1286 }
1287}
1288
1289/// Represents whether a stream of data **generated** by an operator is bounded (finite)
1290/// or unbounded (infinite).
1291///
1292/// This is used to determine whether an execution plan will eventually complete
1293/// processing all its data (bounded) or could potentially run forever (unbounded).
1294///
1295/// For unbounded streams, it also tracks whether the operator requires finite memory
1296/// to process the stream or if memory usage could grow unbounded.
1297///
1298/// Boundedness of the output stream is based on the boundedness of the input stream and the nature of
1299/// the operator. For example, limit or topk with fetch operator can convert an unbounded stream to a bounded stream.
1300#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1301pub enum Boundedness {
1302 /// The data stream is bounded (finite) and will eventually complete
1303 Bounded,
1304 /// The data stream is unbounded (infinite) and could run forever
1305 Unbounded {
1306 /// Whether this operator requires infinite memory to process the unbounded stream.
1307 /// If false, the operator can process an infinite stream with bounded memory.
1308 /// If true, memory usage may grow unbounded while processing the stream.
1309 ///
1310 /// For example, `Median` requires infinite memory to compute the median of an unbounded stream.
1311 /// `Min/Max` requires infinite memory if the stream is unordered, but can be computed with bounded memory if the stream is ordered.
1312 requires_infinite_memory: bool,
1313 },
1314}
1315
1316impl Boundedness {
1317 pub fn is_unbounded(&self) -> bool {
1318 matches!(self, Boundedness::Unbounded { .. })
1319 }
1320}
1321
1322/// Represents how an operator emits its output records.
1323///
1324/// This is used to determine whether an operator emits records incrementally as they arrive,
1325/// only emits a final result at the end, or can do both. Note that it generates the output -- record batch with `batch_size` rows
1326/// but it may still buffer data internally until it has enough data to emit a record batch or the source is exhausted.
1327///
1328/// For example, in the following plan:
1329/// ```text
1330/// SortExec [EmissionType::Final]
1331/// |_ on: [col1 ASC]
1332/// FilterExec [EmissionType::Incremental]
1333/// |_ pred: col2 > 100
1334/// DataSourceExec [EmissionType::Incremental]
1335/// |_ file: "data.csv"
1336/// ```
1337/// - DataSourceExec emits records incrementally as it reads from the file
1338/// - FilterExec processes and emits filtered records incrementally as they arrive
1339/// - SortExec must wait for all input records before it can emit the sorted result,
1340/// since it needs to see all values to determine their final order
1341///
1342/// Left joins can emit both incrementally and finally:
1343/// - Incrementally emit matches as they are found
1344/// - Finally emit non-matches after all input is processed
1345#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1346pub enum EmissionType {
1347 /// Records are emitted incrementally as they arrive and are processed
1348 Incremental,
1349 /// Records are only emitted once all input has been processed
1350 Final,
1351 /// Records can be emitted both incrementally and as a final result
1352 Both,
1353}
1354
1355/// Represents whether an operator's `Stream` has been implemented to actively cooperate with the
1356/// Tokio scheduler or not. Please refer to the [`coop`](crate::coop) module for more details.
1357#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1358pub enum SchedulingType {
1359 /// The stream generated by [`execute`](ExecutionPlan::execute) does not actively participate in
1360 /// cooperative scheduling. This means the implementation of the `Stream` returned by
1361 /// [`ExecutionPlan::execute`] does not contain explicit task budget consumption such as
1362 /// [`tokio::task::coop::consume_budget`].
1363 ///
1364 /// `NonCooperative` is the default value and is acceptable for most operators. Please refer to
1365 /// the [`coop`](crate::coop) module for details on when it may be useful to use
1366 /// `Cooperative` instead.
1367 NonCooperative,
1368 /// The stream generated by [`execute`](ExecutionPlan::execute) actively participates in
1369 /// cooperative scheduling by consuming task budget when it was able to produce a
1370 /// [`RecordBatch`].
1371 Cooperative,
1372}
1373
1374/// Represents how an operator's stream drives [`RecordBatch`] production
1375/// relative to downstream demand.
1376///
1377/// This is execution-topology metadata for optimizers. It distinguishes streams
1378/// whose batch production is driven directly by downstream calls to
1379/// `Stream::poll_next` from streams that may also drive input or output
1380/// production independently, such as by spawning tasks or buffering batches
1381/// ahead of demand.
1382#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1383pub enum EvaluationType {
1384 /// The stream generated by [`execute`](ExecutionPlan::execute) is
1385 /// demand-driven: it produces [`RecordBatch`]es in response to downstream
1386 /// calls to `Stream::poll_next`.
1387 ///
1388 /// Filter, projection, and join operators are examples of lazy operators.
1389 ///
1390 /// Lazy operators are also known as demand-driven operators.
1391 Lazy,
1392 /// The stream generated by [`execute`](ExecutionPlan::execute) may drive
1393 /// input or output [`RecordBatch`] production ahead of, or independently
1394 /// from, downstream calls to `Stream::poll_next`.
1395 ///
1396 /// Eager operators commonly poll input streams from spawned Tokio tasks,
1397 /// buffer batches ahead of demand, or otherwise create an independent
1398 /// child-polling pipeline. Eager work may start when `execute` creates the
1399 /// stream or when the returned stream is first polled; that timing is an
1400 /// implementation detail.
1401 ///
1402 /// Repartition, coalesce partitions, sort-preserving merge, buffer, and
1403 /// analyze operators are examples of eager operators.
1404 ///
1405 /// Eager operators are also known as a data-driven operators.
1406 Eager,
1407}
1408
1409/// Utility to determine an operator's boundedness based on its children's boundedness.
1410///
1411/// Assumes boundedness can be inferred from child operators:
1412/// - Unbounded (requires_infinite_memory: true) takes precedence.
1413/// - Unbounded (requires_infinite_memory: false) is considered next.
1414/// - Otherwise, the operator is bounded.
1415///
1416/// **Note:** This is a general-purpose utility and may not apply to
1417/// all multi-child operators. Ensure your operator's behavior aligns
1418/// with these assumptions before using.
1419pub(crate) fn boundedness_from_children<'a>(
1420 children: impl IntoIterator<Item = &'a Arc<dyn ExecutionPlan>>,
1421) -> Boundedness {
1422 let mut unbounded_with_finite_mem = false;
1423
1424 for child in children {
1425 match child.boundedness() {
1426 Boundedness::Unbounded {
1427 requires_infinite_memory: true,
1428 } => {
1429 return Boundedness::Unbounded {
1430 requires_infinite_memory: true,
1431 };
1432 }
1433 Boundedness::Unbounded {
1434 requires_infinite_memory: false,
1435 } => {
1436 unbounded_with_finite_mem = true;
1437 }
1438 Boundedness::Bounded => {}
1439 }
1440 }
1441
1442 if unbounded_with_finite_mem {
1443 Boundedness::Unbounded {
1444 requires_infinite_memory: false,
1445 }
1446 } else {
1447 Boundedness::Bounded
1448 }
1449}
1450
1451/// Determines the emission type of an operator based on its children's pipeline behavior.
1452///
1453/// The precedence of emission types is:
1454/// - `Final` has the highest precedence.
1455/// - `Both` is next: if any child emits both incremental and final results, the parent inherits this behavior unless a `Final` is present.
1456/// - `Incremental` is the default if all children emit incremental results.
1457///
1458/// **Note:** This is a general-purpose utility and may not apply to
1459/// all multi-child operators. Verify your operator's behavior aligns
1460/// with these assumptions.
1461pub(crate) fn emission_type_from_children<'a>(
1462 children: impl IntoIterator<Item = &'a Arc<dyn ExecutionPlan>>,
1463) -> EmissionType {
1464 let mut inc_and_final = false;
1465
1466 for child in children {
1467 match child.pipeline_behavior() {
1468 EmissionType::Final => return EmissionType::Final,
1469 EmissionType::Both => inc_and_final = true,
1470 EmissionType::Incremental => continue,
1471 }
1472 }
1473
1474 if inc_and_final {
1475 EmissionType::Both
1476 } else {
1477 EmissionType::Incremental
1478 }
1479}
1480
1481/// Stores plan properties used in query optimization.
1482///
1483/// Serves as a cache for these properties, which are often
1484/// expensive to compute.
1485#[derive(Debug, Clone)]
1486pub struct PlanProperties {
1487 /// See [ExecutionPlanProperties::equivalence_properties]
1488 pub eq_properties: EquivalenceProperties,
1489 /// See [ExecutionPlanProperties::output_partitioning]
1490 pub partitioning: Partitioning,
1491 /// See [ExecutionPlanProperties::pipeline_behavior]
1492 pub emission_type: EmissionType,
1493 /// See [ExecutionPlanProperties::boundedness]
1494 pub boundedness: Boundedness,
1495 pub evaluation_type: EvaluationType,
1496 pub scheduling_type: SchedulingType,
1497 /// See [ExecutionPlanProperties::output_ordering]
1498 output_ordering: Option<LexOrdering>,
1499}
1500
1501impl PlanProperties {
1502 /// Construct a new `PlanPropertiesCache` from the
1503 pub fn new(
1504 eq_properties: EquivalenceProperties,
1505 partitioning: Partitioning,
1506 emission_type: EmissionType,
1507 boundedness: Boundedness,
1508 ) -> Self {
1509 // Output ordering can be derived from `eq_properties`.
1510 let output_ordering = eq_properties.output_ordering();
1511 Self {
1512 eq_properties,
1513 partitioning,
1514 emission_type,
1515 boundedness,
1516 evaluation_type: EvaluationType::Lazy,
1517 scheduling_type: SchedulingType::NonCooperative,
1518 output_ordering,
1519 }
1520 }
1521
1522 /// Overwrite output partitioning with its new value.
1523 pub fn with_partitioning(mut self, partitioning: Partitioning) -> Self {
1524 self.partitioning = partitioning;
1525 self
1526 }
1527
1528 /// Set equivalence properties having mut reference.
1529 pub fn set_eq_properties(&mut self, eq_properties: EquivalenceProperties) {
1530 // Changing equivalence properties also changes output ordering, so
1531 // make sure to overwrite it:
1532 self.output_ordering = eq_properties.output_ordering();
1533 self.eq_properties = eq_properties;
1534 }
1535
1536 /// Overwrite equivalence properties with its new value.
1537 pub fn with_eq_properties(mut self, eq_properties: EquivalenceProperties) -> Self {
1538 self.set_eq_properties(eq_properties);
1539 self
1540 }
1541
1542 /// Overwrite boundedness with its new value.
1543 pub fn with_boundedness(mut self, boundedness: Boundedness) -> Self {
1544 self.boundedness = boundedness;
1545 self
1546 }
1547
1548 /// Overwrite emission type with its new value.
1549 pub fn with_emission_type(mut self, emission_type: EmissionType) -> Self {
1550 self.emission_type = emission_type;
1551 self
1552 }
1553
1554 /// Set the [`SchedulingType`].
1555 ///
1556 /// Defaults to [`SchedulingType::NonCooperative`]
1557 pub fn with_scheduling_type(mut self, scheduling_type: SchedulingType) -> Self {
1558 self.scheduling_type = scheduling_type;
1559 self
1560 }
1561
1562 /// Set the [`EvaluationType`].
1563 ///
1564 /// Defaults to [`EvaluationType::Lazy`]
1565 pub fn with_evaluation_type(mut self, drive_type: EvaluationType) -> Self {
1566 self.evaluation_type = drive_type;
1567 self
1568 }
1569
1570 /// Set constraints having mut reference.
1571 pub fn set_constraints(&mut self, constraints: Constraints) {
1572 self.eq_properties.set_constraints(constraints);
1573 }
1574
1575 /// Overwrite constraints with its new value.
1576 pub fn with_constraints(mut self, constraints: Constraints) -> Self {
1577 self.set_constraints(constraints);
1578 self
1579 }
1580
1581 pub fn equivalence_properties(&self) -> &EquivalenceProperties {
1582 &self.eq_properties
1583 }
1584
1585 pub fn output_partitioning(&self) -> &Partitioning {
1586 &self.partitioning
1587 }
1588
1589 pub fn output_ordering(&self) -> Option<&LexOrdering> {
1590 self.output_ordering.as_ref()
1591 }
1592
1593 /// Get schema of the node.
1594 pub(crate) fn schema(&self) -> &SchemaRef {
1595 self.eq_properties.schema()
1596 }
1597}
1598
1599macro_rules! check_len {
1600 ($target:expr, $func_name:ident, $expected_len:expr) => {
1601 let actual_len = $target.$func_name().len();
1602 assert_eq_or_internal_err!(
1603 actual_len,
1604 $expected_len,
1605 "{}::{} returned Vec with incorrect size: {} != {}",
1606 $target.name(),
1607 stringify!($func_name),
1608 actual_len,
1609 $expected_len
1610 );
1611 };
1612}
1613
1614/// All dynamic expressions must have an expression id.
1615fn check_dynamic_expression_invariants<P: ExecutionPlan + ?Sized>(
1616 plan: &P,
1617) -> Result<()> {
1618 let mut produced_ids = HashSet::new();
1619 for expr in plan.dynamic_expressions_produced() {
1620 let Some(expression_id) = expr.expression_id() else {
1621 return internal_err!(
1622 "{}::dynamic_expressions_produced returned an expression without an expression ID",
1623 plan.name()
1624 );
1625 };
1626 assert_or_internal_err!(
1627 produced_ids.insert(expression_id),
1628 "{}::dynamic_expressions_produced returned duplicate expression ID {expression_id}",
1629 plan.name()
1630 );
1631 }
1632 Ok(())
1633}
1634
1635/// Checks a set of invariants that apply to all ExecutionPlan implementations.
1636/// Returns an error if the given node does not conform.
1637pub fn check_default_invariants<P: ExecutionPlan + ?Sized>(
1638 plan: &P,
1639 check: InvariantLevel,
1640) -> Result<(), DataFusionError> {
1641 let children_len = plan.children().len();
1642
1643 check_len!(plan, maintains_input_order, children_len);
1644 check_len!(plan, required_input_ordering, children_len);
1645 check_len!(plan, benefits_from_input_partitioning, children_len);
1646 plan.input_distribution_requirements()
1647 .check_invariants(plan, check)?;
1648 check_dynamic_expression_invariants(plan)?;
1649
1650 Ok(())
1651}
1652
1653/// Indicate whether a data exchange is needed for the input of `plan`.
1654///
1655/// This identifies physical operators that redistribute child partitions or
1656/// gather multiple child partitions into one output partition:
1657///
1658/// 1. RepartitionExec for non-round-robin repartitioning
1659/// 2. CoalescePartitionsExec for collapsing multiple partitions into one without ordering guarantee
1660/// 3. SortPreservingMergeExec for collapsing multiple sorted partitions into one with ordering guarantee
1661#[expect(clippy::needless_pass_by_value)]
1662pub fn need_data_exchange(plan: Arc<dyn ExecutionPlan>) -> bool {
1663 if let Some(repartition) = plan.downcast_ref::<RepartitionExec>() {
1664 !matches!(repartition.partitioning(), Partitioning::RoundRobinBatch(_))
1665 } else if let Some(coalesce) = plan.downcast_ref::<CoalescePartitionsExec>() {
1666 coalesce.input().output_partitioning().partition_count() > 1
1667 } else if let Some(sort_preserving_merge) =
1668 plan.downcast_ref::<SortPreservingMergeExec>()
1669 {
1670 sort_preserving_merge
1671 .input()
1672 .output_partitioning()
1673 .partition_count()
1674 > 1
1675 } else {
1676 false
1677 }
1678}
1679
1680/// Returns a plan with the given children, skipping as much work as possible.
1681///
1682/// This helper is the single entry point for "rebuild a plan from new
1683/// children" and applies three layers of short-circuits, from cheapest to
1684/// most expensive:
1685///
1686/// 1. **Same child pointers** — if every `children[i]` is `Arc::ptr_eq` to the
1687/// corresponding existing child, the original `plan` is returned
1688/// unchanged (no allocation, no [`ExecutionPlan::replace_children`]
1689/// call).
1690/// 2. **Same child properties** — if the children's `PlanProperties` Arcs
1691/// match (via [`has_same_children_properties`]), the plan's own
1692/// `PlanProperties` cache can be reused. This calls
1693/// [`ExecutionPlan::replace_children`] with [`ChildrenPropertiesMode::Keep`],
1694/// which swaps the child pointers without recomputing `PlanProperties`.
1695/// 3. **Full recompute** — otherwise, delegate to
1696/// [`ExecutionPlan::replace_children`] with [`ChildrenPropertiesMode::Recompute`],
1697/// which recomputes `PlanProperties` from scratch.
1698///
1699/// The size of `children` must be equal to the size of `ExecutionPlan::children()`.
1700pub fn replace_children_if_necessary(
1701 plan: Arc<dyn ExecutionPlan>,
1702 children: Vec<Arc<dyn ExecutionPlan>>,
1703) -> Result<Arc<dyn ExecutionPlan>> {
1704 let old_children = plan.children();
1705 assert_eq_or_internal_err!(
1706 children.len(),
1707 old_children.len(),
1708 "Wrong number of children"
1709 );
1710 if !children.is_empty() {
1711 // Layer 1: same child pointers → return the plan unchanged.
1712 if children
1713 .iter()
1714 .zip(old_children.iter())
1715 .all(|(c1, c2)| Arc::ptr_eq(c1, c2))
1716 {
1717 return Ok(plan);
1718 }
1719 // Layer 2: same child properties → reuse `PlanProperties` cache.
1720 if has_same_children_properties(plan.as_ref(), &children)? {
1721 return plan.replace_children(
1722 children,
1723 ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep),
1724 );
1725 }
1726 }
1727 // Layer 3: full recompute.
1728 plan.replace_children(
1729 children,
1730 ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
1731 )
1732}
1733
1734#[deprecated(since = "55.0.0", note = "Use `replace_children_if_necessary`")]
1735pub fn with_new_children_if_necessary(
1736 plan: Arc<dyn ExecutionPlan>,
1737 children: Vec<Arc<dyn ExecutionPlan>>,
1738) -> Result<Arc<dyn ExecutionPlan>> {
1739 replace_children_if_necessary(plan, children)
1740}
1741
1742/// Return a [`DisplayableExecutionPlan`] wrapper around an
1743/// [`ExecutionPlan`] which can be displayed in various easier to
1744/// understand ways.
1745///
1746/// See examples on [`DisplayableExecutionPlan`]
1747pub fn displayable(plan: &dyn ExecutionPlan) -> DisplayableExecutionPlan<'_> {
1748 DisplayableExecutionPlan::new(plan)
1749}
1750
1751/// Execute the [ExecutionPlan] and collect the results in memory
1752pub async fn collect(
1753 plan: Arc<dyn ExecutionPlan>,
1754 context: Arc<TaskContext>,
1755) -> Result<Vec<RecordBatch>> {
1756 let stream = execute_stream(plan, context)?;
1757 crate::common::collect(stream).await
1758}
1759
1760/// Execute the [ExecutionPlan] and return a single stream of `RecordBatch`es.
1761///
1762/// See [collect] to buffer the `RecordBatch`es in memory.
1763///
1764/// # Aborting Execution
1765///
1766/// Dropping the stream will abort the execution of the query, and free up
1767/// any allocated resources
1768#[expect(
1769 clippy::needless_pass_by_value,
1770 reason = "Public API that historically takes owned Arcs"
1771)]
1772pub fn execute_stream(
1773 plan: Arc<dyn ExecutionPlan>,
1774 context: Arc<TaskContext>,
1775) -> Result<SendableRecordBatchStream> {
1776 match plan.output_partitioning().partition_count() {
1777 0 => Ok(Box::pin(EmptyRecordBatchStream::new(plan.schema()))),
1778 1 => plan.execute(0, context),
1779 2.. => {
1780 // merge into a single partition
1781 let plan = CoalescePartitionsExec::new(Arc::clone(&plan));
1782 // CoalescePartitionsExec must produce a single partition
1783 assert_eq!(1, plan.properties().output_partitioning().partition_count());
1784 plan.execute(0, context)
1785 }
1786 }
1787}
1788
1789/// Execute the [ExecutionPlan] and collect the results in memory
1790pub async fn collect_partitioned(
1791 plan: Arc<dyn ExecutionPlan>,
1792 context: Arc<TaskContext>,
1793) -> Result<Vec<Vec<RecordBatch>>> {
1794 // Avoid `JoinSet::spawn` for single partition
1795 if plan.output_partitioning().partition_count() == 1 {
1796 let stream = plan.execute(0, context)?;
1797 let batches: Vec<RecordBatch> = stream.try_collect().await?;
1798 return Ok(vec![batches]);
1799 }
1800
1801 let streams = execute_stream_partitioned(plan, context)?;
1802
1803 let mut join_set = JoinSet::new();
1804 // Execute the plan and collect the results into batches.
1805 streams.into_iter().enumerate().for_each(|(idx, stream)| {
1806 join_set.spawn(async move {
1807 let result: Result<Vec<RecordBatch>> = stream.try_collect().await;
1808 (idx, result)
1809 });
1810 });
1811
1812 let mut batches = vec![];
1813 // Note that currently this doesn't identify the thread that panicked
1814 //
1815 // TODO: Replace with [join_next_with_id](https://docs.rs/tokio/latest/tokio/task/struct.JoinSet.html#method.join_next_with_id
1816 // once it is stable
1817 while let Some(result) = join_set.join_next().await {
1818 match result {
1819 Ok((idx, res)) => batches.push((idx, res?)),
1820 Err(e) => {
1821 if e.is_panic() {
1822 std::panic::resume_unwind(e.into_panic());
1823 } else {
1824 unreachable!();
1825 }
1826 }
1827 }
1828 }
1829
1830 batches.sort_by_key(|(idx, _)| *idx);
1831 let batches = batches.into_iter().map(|(_, batch)| batch).collect();
1832
1833 Ok(batches)
1834}
1835
1836/// Execute the [ExecutionPlan] and return a vec with one stream per output
1837/// partition
1838///
1839/// # Aborting Execution
1840///
1841/// Dropping the stream will abort the execution of the query, and free up
1842/// any allocated resources
1843#[expect(
1844 clippy::needless_pass_by_value,
1845 reason = "Public API that historically takes owned Arcs"
1846)]
1847pub fn execute_stream_partitioned(
1848 plan: Arc<dyn ExecutionPlan>,
1849 context: Arc<TaskContext>,
1850) -> Result<Vec<SendableRecordBatchStream>> {
1851 let num_partitions = plan.output_partitioning().partition_count();
1852 let mut streams = Vec::with_capacity(num_partitions);
1853 for i in 0..num_partitions {
1854 streams.push(plan.execute(i, Arc::clone(&context))?);
1855 }
1856 Ok(streams)
1857}
1858
1859/// Executes an input stream and ensures that the resulting stream adheres to
1860/// the `not null` constraints specified in the `sink_schema`.
1861///
1862/// # Arguments
1863///
1864/// * `input` - An execution plan
1865/// * `sink_schema` - The schema to be applied to the output stream
1866/// * `partition` - The partition index to be executed
1867/// * `context` - The task context
1868///
1869/// # Returns
1870///
1871/// * `Result<SendableRecordBatchStream>` - A stream of `RecordBatch`es if successful
1872///
1873/// This function first executes the given input plan for the specified partition
1874/// and context. It then checks if there are any columns in the input that might
1875/// violate the `not null` constraints specified in the `sink_schema`. If there are
1876/// such columns, it wraps the resulting stream to enforce the `not null` constraints
1877/// by invoking the [`check_not_null_constraints`] function on each batch of the stream.
1878#[expect(
1879 clippy::needless_pass_by_value,
1880 reason = "Public API that historically takes owned Arcs"
1881)]
1882pub fn execute_input_stream(
1883 input: Arc<dyn ExecutionPlan>,
1884 sink_schema: SchemaRef,
1885 partition: usize,
1886 context: Arc<TaskContext>,
1887) -> Result<SendableRecordBatchStream> {
1888 let input_stream = input.execute(partition, context)?;
1889
1890 debug_assert_eq!(sink_schema.fields().len(), input.schema().fields().len());
1891
1892 // Find input columns that may violate the not null constraint.
1893 let risky_columns: Vec<_> = sink_schema
1894 .fields()
1895 .iter()
1896 .zip(input.schema().fields().iter())
1897 .enumerate()
1898 .filter_map(|(idx, (sink_field, input_field))| {
1899 (!sink_field.is_nullable() && input_field.is_nullable()).then_some(idx)
1900 })
1901 .collect();
1902
1903 if risky_columns.is_empty() {
1904 Ok(input_stream)
1905 } else {
1906 // Check not null constraint on the input stream
1907 Ok(Box::pin(RecordBatchStreamAdapter::new(
1908 sink_schema,
1909 input_stream
1910 .map(move |batch| check_not_null_constraints(batch?, &risky_columns)),
1911 )))
1912 }
1913}
1914
1915/// Checks a `RecordBatch` for `not null` constraints on specified columns.
1916///
1917/// # Arguments
1918///
1919/// * `batch` - The `RecordBatch` to be checked
1920/// * `column_indices` - A vector of column indices that should be checked for
1921/// `not null` constraints.
1922///
1923/// # Returns
1924///
1925/// * `Result<RecordBatch>` - The original `RecordBatch` if all constraints are met
1926///
1927/// This function iterates over the specified column indices and ensures that none
1928/// of the columns contain null values. If any column contains null values, an error
1929/// is returned.
1930pub fn check_not_null_constraints(
1931 batch: RecordBatch,
1932 column_indices: &Vec<usize>,
1933) -> Result<RecordBatch> {
1934 for &index in column_indices {
1935 if batch.num_columns() <= index {
1936 return exec_err!(
1937 "Invalid batch column count {} expected > {}",
1938 batch.num_columns(),
1939 index
1940 );
1941 }
1942
1943 if batch
1944 .column(index)
1945 .logical_nulls()
1946 .map(|nulls| nulls.null_count())
1947 .unwrap_or_default()
1948 > 0
1949 {
1950 return exec_err!(
1951 "Invalid batch column at '{}' has null but schema specifies non-nullable",
1952 index
1953 );
1954 }
1955 }
1956
1957 Ok(batch)
1958}
1959
1960/// Make plan ready to be re-executed returning its clone with state reset for all nodes.
1961///
1962/// Some plans will change their internal states after execution, making them unable to be executed again.
1963/// This function uses [`ExecutionPlan::reset_state`] to reset any internal state within the plan.
1964///
1965/// An example is `CrossJoinExec`, which loads the left table into memory and stores it in the plan.
1966/// However, if the data of the left table is derived from the work table, it will become outdated
1967/// as the work table changes. When the next iteration executes this plan again, we must clear the left table.
1968///
1969/// # Limitations
1970///
1971/// While this function enables plan reuse, it does not allow the same plan to be executed if it (OR):
1972///
1973/// * uses dynamic filters,
1974/// * represents a recursive query.
1975///
1976pub fn reset_plan_states(plan: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn ExecutionPlan>> {
1977 plan.transform_up(|plan| {
1978 let new_plan = Arc::clone(&plan).reset_state()?;
1979 Ok(Transformed::yes(new_plan))
1980 })
1981 .data()
1982}
1983
1984/// Check if the `plan` children has the same properties as passed `children`.
1985/// In this case plan can avoid self properties re-computation when its children
1986/// replace is requested.
1987/// The size of `children` must be equal to the size of `ExecutionPlan::children()`.
1988pub fn has_same_children_properties(
1989 plan: &dyn ExecutionPlan,
1990 children: &[Arc<dyn ExecutionPlan>],
1991) -> Result<bool> {
1992 let old_children = plan.children();
1993 assert_eq_or_internal_err!(
1994 children.len(),
1995 old_children.len(),
1996 "Wrong number of children"
1997 );
1998 for (lhs, rhs) in old_children.iter().zip(children.iter()) {
1999 if !Arc::ptr_eq(lhs.properties(), rhs.properties()) {
2000 return Ok(false);
2001 }
2002 }
2003 Ok(true)
2004}
2005
2006/// Helper macro to avoid properties re-computation if passed children properties
2007/// the same as plan already has. Could be used to implement fast-path for method
2008/// [`ExecutionPlan::with_new_children`].
2009///
2010/// New call sites should route through [`replace_children_if_necessary`],
2011/// which applies this check together with the child-pointer short-circuit
2012/// (see [`replace_children_if_necessary`] for the layered policy). This
2013/// macro remains for direct-caller sites that have not been migrated yet.
2014#[macro_export]
2015macro_rules! check_if_same_properties {
2016 ($plan: expr, $children: expr) => {
2017 if $crate::execution_plan::has_same_children_properties(
2018 $plan.as_ref(),
2019 &$children,
2020 )? {
2021 return ::std::sync::Arc::clone(&$plan)
2022 .with_new_children_and_same_properties($children);
2023 }
2024 };
2025}
2026
2027/// Helper macro to validate that replacement children match a plan's existing
2028/// child count.
2029///
2030/// This is useful for [`ExecutionPlan::replace_children`] implementations that
2031/// need to preserve the same child-count validation behavior.
2032#[macro_export]
2033macro_rules! validate_child_count {
2034 ($plan: expr, $children: expr) => {
2035 datafusion_common::assert_eq_or_internal_err!(
2036 $children.len(),
2037 $plan.children().len(),
2038 "Wrong number of children"
2039 );
2040 };
2041}
2042
2043/// Utility function yielding a string representation of the given [`ExecutionPlan`].
2044pub fn get_plan_string(plan: &Arc<dyn ExecutionPlan>) -> Vec<String> {
2045 let formatted = displayable(plan.as_ref()).indent(true).to_string();
2046 let actual: Vec<&str> = formatted.trim().lines().collect();
2047 actual.iter().map(|elem| (*elem).to_string()).collect()
2048}
2049
2050/// Indicates the effect an execution plan operator will have on the cardinality
2051/// of its input stream
2052pub enum CardinalityEffect {
2053 /// Unknown effect. This is the default
2054 Unknown,
2055 /// The operator is guaranteed to produce exactly one row for
2056 /// each input row
2057 Equal,
2058 /// The operator may produce fewer output rows than it receives input rows
2059 LowerEqual,
2060 /// The operator may produce more output rows than it receives input rows
2061 GreaterEqual,
2062}
2063
2064/// Can be used in contexts where properties have not yet been initialized properly.
2065pub(crate) fn stub_properties() -> Arc<PlanProperties> {
2066 static STUB_PROPERTIES: LazyLock<Arc<PlanProperties>> = LazyLock::new(|| {
2067 Arc::new(PlanProperties::new(
2068 EquivalenceProperties::new(Arc::new(Schema::empty())),
2069 Partitioning::UnknownPartitioning(1),
2070 EmissionType::Final,
2071 Boundedness::Bounded,
2072 ))
2073 });
2074
2075 Arc::clone(&STUB_PROPERTIES)
2076}
2077
2078#[cfg(test)]
2079mod tests {
2080
2081 use super::*;
2082 use crate::buffer::BufferExec;
2083 use crate::test::exec::MockExec;
2084 use crate::{DisplayAs, DisplayFormatType, ExecutionPlan};
2085
2086 use arrow::array::{DictionaryArray, Int32Array, NullArray, RunArray};
2087 use arrow::datatypes::{DataType, Field, Schema};
2088 use datafusion_physical_expr::expressions::{DynamicFilterPhysicalExpr, lit};
2089
2090 #[derive(Debug)]
2091 pub struct EmptyExec {
2092 dynamic_expressions: Vec<Arc<dyn PhysicalExpr>>,
2093 }
2094
2095 impl EmptyExec {
2096 pub fn new(_schema: SchemaRef) -> Self {
2097 Self {
2098 dynamic_expressions: vec![],
2099 }
2100 }
2101
2102 fn with_dynamic_expressions(
2103 mut self,
2104 dynamic_expressions: Vec<Arc<dyn PhysicalExpr>>,
2105 ) -> Self {
2106 self.dynamic_expressions = dynamic_expressions;
2107 self
2108 }
2109 }
2110
2111 impl DisplayAs for EmptyExec {
2112 fn fmt_as(
2113 &self,
2114 _t: DisplayFormatType,
2115 _f: &mut std::fmt::Formatter,
2116 ) -> std::fmt::Result {
2117 unimplemented!()
2118 }
2119 }
2120
2121 impl ExecutionPlan for EmptyExec {
2122 fn name(&self) -> &'static str {
2123 Self::static_name()
2124 }
2125
2126 fn properties(&self) -> &Arc<PlanProperties> {
2127 unimplemented!()
2128 }
2129
2130 fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
2131 vec![]
2132 }
2133
2134 fn replace_children(
2135 self: Arc<Self>,
2136 _: Vec<Arc<dyn ExecutionPlan>>,
2137 _: ReplaceChildrenOptions,
2138 ) -> Result<Arc<dyn ExecutionPlan>> {
2139 unimplemented!()
2140 }
2141
2142 fn apply_expressions(
2143 &self,
2144 _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
2145 ) -> Result<TreeNodeRecursion> {
2146 Ok(TreeNodeRecursion::Continue)
2147 }
2148
2149 fn with_new_children(
2150 self: Arc<Self>,
2151 children: Vec<Arc<dyn ExecutionPlan>>,
2152 ) -> Result<Arc<dyn ExecutionPlan>> {
2153 self.replace_children(
2154 children,
2155 ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
2156 )
2157 }
2158
2159 fn dynamic_expressions_produced(&self) -> Vec<Arc<dyn PhysicalExpr>> {
2160 self.dynamic_expressions.iter().map(Arc::clone).collect()
2161 }
2162
2163 fn execute(
2164 &self,
2165 _partition: usize,
2166 _context: Arc<TaskContext>,
2167 ) -> Result<SendableRecordBatchStream> {
2168 unimplemented!()
2169 }
2170
2171 fn statistics_from_inputs(
2172 &self,
2173 _input_stats: &[Arc<Statistics>],
2174 _args: &StatisticsArgs,
2175 ) -> Result<Arc<Statistics>> {
2176 unimplemented!()
2177 }
2178 }
2179
2180 #[test]
2181 fn test_dynamic_expression_invariants() -> Result<()> {
2182 let schema = Arc::new(Schema::empty());
2183 let dynamic: Arc<dyn PhysicalExpr> =
2184 Arc::new(DynamicFilterPhysicalExpr::new(vec![], lit(true)));
2185 let valid = EmptyExec::new(Arc::clone(&schema))
2186 .with_dynamic_expressions(vec![Arc::clone(&dynamic)]);
2187 check_default_invariants(&valid, InvariantLevel::Always)?;
2188
2189 let missing_id =
2190 EmptyExec::new(Arc::clone(&schema)).with_dynamic_expressions(vec![lit(true)]);
2191 let error = check_default_invariants(&missing_id, InvariantLevel::Always)
2192 .unwrap_err()
2193 .strip_backtrace();
2194 assert!(error.contains("without an expression ID"), "{error}");
2195
2196 let duplicate = EmptyExec::new(schema)
2197 .with_dynamic_expressions(vec![Arc::clone(&dynamic), dynamic]);
2198 let error = check_default_invariants(&duplicate, InvariantLevel::Always)
2199 .unwrap_err()
2200 .strip_backtrace();
2201 assert!(error.contains("duplicate expression ID"), "{error}");
2202
2203 Ok(())
2204 }
2205
2206 #[derive(Debug)]
2207 pub struct RenamedEmptyExec;
2208
2209 impl RenamedEmptyExec {
2210 pub fn new(_schema: SchemaRef) -> Self {
2211 Self
2212 }
2213 }
2214
2215 impl DisplayAs for RenamedEmptyExec {
2216 fn fmt_as(
2217 &self,
2218 _t: DisplayFormatType,
2219 _f: &mut std::fmt::Formatter,
2220 ) -> std::fmt::Result {
2221 unimplemented!()
2222 }
2223 }
2224
2225 impl ExecutionPlan for RenamedEmptyExec {
2226 fn name(&self) -> &'static str {
2227 Self::static_name()
2228 }
2229
2230 fn static_name() -> &'static str
2231 where
2232 Self: Sized,
2233 {
2234 "MyRenamedEmptyExec"
2235 }
2236
2237 fn properties(&self) -> &Arc<PlanProperties> {
2238 unimplemented!()
2239 }
2240
2241 fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
2242 vec![]
2243 }
2244
2245 fn apply_expressions(
2246 &self,
2247 _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
2248 ) -> Result<TreeNodeRecursion> {
2249 Ok(TreeNodeRecursion::Continue)
2250 }
2251
2252 fn replace_children(
2253 self: Arc<Self>,
2254 _: Vec<Arc<dyn ExecutionPlan>>,
2255 _: ReplaceChildrenOptions,
2256 ) -> Result<Arc<dyn ExecutionPlan>> {
2257 unimplemented!()
2258 }
2259
2260 fn with_new_children(
2261 self: Arc<Self>,
2262 children: Vec<Arc<dyn ExecutionPlan>>,
2263 ) -> Result<Arc<dyn ExecutionPlan>> {
2264 self.replace_children(
2265 children,
2266 ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
2267 )
2268 }
2269
2270 fn execute(
2271 &self,
2272 _partition: usize,
2273 _context: Arc<TaskContext>,
2274 ) -> Result<SendableRecordBatchStream> {
2275 unimplemented!()
2276 }
2277
2278 fn statistics_from_inputs(
2279 &self,
2280 _input_stats: &[Arc<Statistics>],
2281 _args: &StatisticsArgs,
2282 ) -> Result<Arc<Statistics>> {
2283 unimplemented!()
2284 }
2285 }
2286
2287 #[derive(Debug)]
2288 struct DowncastDelegatingExec(Arc<dyn ExecutionPlan>);
2289
2290 impl DisplayAs for DowncastDelegatingExec {
2291 fn fmt_as(
2292 &self,
2293 _t: DisplayFormatType,
2294 _f: &mut std::fmt::Formatter,
2295 ) -> std::fmt::Result {
2296 unimplemented!()
2297 }
2298 }
2299
2300 impl ExecutionPlan for DowncastDelegatingExec {
2301 fn name(&self) -> &'static str {
2302 Self::static_name()
2303 }
2304
2305 fn properties(&self) -> &Arc<PlanProperties> {
2306 unimplemented!()
2307 }
2308
2309 fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
2310 vec![]
2311 }
2312
2313 fn apply_expressions(
2314 &self,
2315 f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
2316 ) -> Result<TreeNodeRecursion> {
2317 self.0.apply_expressions(f)
2318 }
2319
2320 fn replace_children(
2321 self: Arc<Self>,
2322 _: Vec<Arc<dyn ExecutionPlan>>,
2323 _: ReplaceChildrenOptions,
2324 ) -> Result<Arc<dyn ExecutionPlan>> {
2325 unimplemented!()
2326 }
2327
2328 fn with_new_children(
2329 self: Arc<Self>,
2330 children: Vec<Arc<dyn ExecutionPlan>>,
2331 ) -> Result<Arc<dyn ExecutionPlan>> {
2332 self.replace_children(
2333 children,
2334 ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
2335 )
2336 }
2337
2338 fn downcast_delegate(&self) -> Option<&dyn ExecutionPlan> {
2339 Some(self.0.as_ref())
2340 }
2341
2342 fn execute(
2343 &self,
2344 _partition: usize,
2345 _context: Arc<TaskContext>,
2346 ) -> Result<SendableRecordBatchStream> {
2347 unimplemented!()
2348 }
2349
2350 fn partition_statistics(
2351 &self,
2352 _partition: Option<usize>,
2353 ) -> Result<Arc<Statistics>> {
2354 unimplemented!()
2355 }
2356 }
2357 /// Test leaf plan with a real [`PlanProperties`] cache. Different instances
2358 /// can share the same cache Arc by cloning `cache`.
2359 #[derive(Debug, Clone)]
2360 struct WithChildrenTestLeaf {
2361 cache: Arc<PlanProperties>,
2362 }
2363
2364 impl WithChildrenTestLeaf {
2365 fn new(cache: Arc<PlanProperties>) -> Self {
2366 Self { cache }
2367 }
2368 }
2369
2370 impl DisplayAs for WithChildrenTestLeaf {
2371 fn fmt_as(
2372 &self,
2373 _t: DisplayFormatType,
2374 _f: &mut std::fmt::Formatter,
2375 ) -> std::fmt::Result {
2376 unimplemented!()
2377 }
2378 }
2379
2380 impl ExecutionPlan for WithChildrenTestLeaf {
2381 fn name(&self) -> &'static str {
2382 "WithChildrenTestLeaf"
2383 }
2384 fn properties(&self) -> &Arc<PlanProperties> {
2385 &self.cache
2386 }
2387 fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
2388 vec![]
2389 }
2390 fn apply_expressions(
2391 &self,
2392 _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
2393 ) -> Result<TreeNodeRecursion> {
2394 Ok(TreeNodeRecursion::Continue)
2395 }
2396
2397 fn replace_children(
2398 self: Arc<Self>,
2399 _: Vec<Arc<dyn ExecutionPlan>>,
2400 _: ReplaceChildrenOptions,
2401 ) -> Result<Arc<dyn ExecutionPlan>> {
2402 Ok(self)
2403 }
2404 fn with_new_children(
2405 self: Arc<Self>,
2406 children: Vec<Arc<dyn ExecutionPlan>>,
2407 ) -> Result<Arc<dyn ExecutionPlan>> {
2408 self.replace_children(
2409 children,
2410 ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
2411 )
2412 }
2413 fn execute(
2414 &self,
2415 _partition: usize,
2416 _context: Arc<TaskContext>,
2417 ) -> Result<SendableRecordBatchStream> {
2418 unimplemented!()
2419 }
2420 }
2421
2422 /// Test unary plan that counts which of `with_new_children` (full
2423 /// recompute) vs `with_new_children_and_same_properties` (fast path) is
2424 /// taken.
2425 #[derive(Debug, Clone)]
2426 struct WithChildrenTestParent {
2427 input: Arc<dyn ExecutionPlan>,
2428 cache: Arc<PlanProperties>,
2429 recompute_calls: Arc<std::sync::atomic::AtomicUsize>,
2430 fast_path_calls: Arc<std::sync::atomic::AtomicUsize>,
2431 }
2432
2433 impl WithChildrenTestParent {
2434 fn new(input: Arc<dyn ExecutionPlan>) -> Self {
2435 let cache = Arc::new(PlanProperties::new(
2436 EquivalenceProperties::new(Arc::new(Schema::empty())),
2437 Partitioning::UnknownPartitioning(1),
2438 EmissionType::Final,
2439 Boundedness::Bounded,
2440 ));
2441 Self {
2442 input,
2443 cache,
2444 recompute_calls: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
2445 fast_path_calls: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
2446 }
2447 }
2448 }
2449
2450 impl DisplayAs for WithChildrenTestParent {
2451 fn fmt_as(
2452 &self,
2453 _t: DisplayFormatType,
2454 _f: &mut std::fmt::Formatter,
2455 ) -> std::fmt::Result {
2456 unimplemented!()
2457 }
2458 }
2459
2460 impl ExecutionPlan for WithChildrenTestParent {
2461 fn name(&self) -> &'static str {
2462 "WithChildrenTestParent"
2463 }
2464 fn properties(&self) -> &Arc<PlanProperties> {
2465 &self.cache
2466 }
2467 fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
2468 vec![&self.input]
2469 }
2470 fn apply_expressions(
2471 &self,
2472 _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
2473 ) -> Result<TreeNodeRecursion> {
2474 Ok(TreeNodeRecursion::Continue)
2475 }
2476
2477 fn replace_children(
2478 self: Arc<Self>,
2479 mut children: Vec<Arc<dyn ExecutionPlan>>,
2480 options: ReplaceChildrenOptions,
2481 ) -> Result<Arc<dyn ExecutionPlan>> {
2482 match options.children_properties {
2483 ChildrenPropertiesMode::Keep => {
2484 self.fast_path_calls
2485 .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
2486 Ok(Arc::new(Self {
2487 input: children.swap_remove(0),
2488 ..Self::clone(&*self)
2489 }))
2490 }
2491 ChildrenPropertiesMode::Recompute => {
2492 self.recompute_calls
2493 .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
2494 // Full recompute: allocate a fresh `PlanProperties` Arc so this
2495 // path is observable via `Arc::ptr_eq` on properties.
2496 let new_input = children.swap_remove(0);
2497 let cache = Arc::new(PlanProperties::new(
2498 EquivalenceProperties::new(Arc::new(Schema::empty())),
2499 Partitioning::UnknownPartitioning(1),
2500 EmissionType::Final,
2501 Boundedness::Bounded,
2502 ));
2503 Ok(Arc::new(Self {
2504 input: new_input,
2505 cache,
2506 recompute_calls: Arc::clone(&self.recompute_calls),
2507 fast_path_calls: Arc::clone(&self.fast_path_calls),
2508 }))
2509 }
2510 }
2511 }
2512 fn with_new_children(
2513 self: Arc<Self>,
2514 children: Vec<Arc<dyn ExecutionPlan>>,
2515 ) -> Result<Arc<dyn ExecutionPlan>> {
2516 self.replace_children(
2517 children,
2518 ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
2519 )
2520 }
2521 fn with_new_children_and_same_properties(
2522 self: Arc<Self>,
2523 children: Vec<Arc<dyn ExecutionPlan>>,
2524 ) -> Result<Arc<dyn ExecutionPlan>> {
2525 self.replace_children(
2526 children,
2527 ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep),
2528 )
2529 }
2530 fn execute(
2531 &self,
2532 _partition: usize,
2533 _context: Arc<TaskContext>,
2534 ) -> Result<SendableRecordBatchStream> {
2535 unimplemented!()
2536 }
2537 }
2538
2539 /// Test unary plan that does **not** override
2540 /// `with_new_children_and_same_properties`. Used to verify the default
2541 /// trait fallback still routes through `with_new_children` (which is
2542 /// the semantics-preserving path for downstream / external
2543 /// `ExecutionPlan` implementations that haven't opted into the
2544 /// fast path yet).
2545 #[derive(Debug, Clone)]
2546 struct WithChildrenTestParentDefault {
2547 input: Arc<dyn ExecutionPlan>,
2548 cache: Arc<PlanProperties>,
2549 recompute_calls: Arc<std::sync::atomic::AtomicUsize>,
2550 }
2551
2552 impl WithChildrenTestParentDefault {
2553 fn new(input: Arc<dyn ExecutionPlan>) -> Self {
2554 let cache = Arc::new(PlanProperties::new(
2555 EquivalenceProperties::new(Arc::new(Schema::empty())),
2556 Partitioning::UnknownPartitioning(1),
2557 EmissionType::Final,
2558 Boundedness::Bounded,
2559 ));
2560 Self {
2561 input,
2562 cache,
2563 recompute_calls: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
2564 }
2565 }
2566 }
2567
2568 impl DisplayAs for WithChildrenTestParentDefault {
2569 fn fmt_as(
2570 &self,
2571 _t: DisplayFormatType,
2572 _f: &mut std::fmt::Formatter,
2573 ) -> std::fmt::Result {
2574 unimplemented!()
2575 }
2576 }
2577
2578 impl ExecutionPlan for WithChildrenTestParentDefault {
2579 fn name(&self) -> &'static str {
2580 "WithChildrenTestParentDefault"
2581 }
2582 fn properties(&self) -> &Arc<PlanProperties> {
2583 &self.cache
2584 }
2585 fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
2586 vec![&self.input]
2587 }
2588 fn apply_expressions(
2589 &self,
2590 _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
2591 ) -> Result<TreeNodeRecursion> {
2592 Ok(TreeNodeRecursion::Continue)
2593 }
2594 fn with_new_children(
2595 self: Arc<Self>,
2596 mut children: Vec<Arc<dyn ExecutionPlan>>,
2597 ) -> Result<Arc<dyn ExecutionPlan>> {
2598 self.recompute_calls
2599 .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
2600 let new_input = children.swap_remove(0);
2601 let cache = Arc::new(PlanProperties::new(
2602 EquivalenceProperties::new(Arc::new(Schema::empty())),
2603 Partitioning::UnknownPartitioning(1),
2604 EmissionType::Final,
2605 Boundedness::Bounded,
2606 ));
2607 Ok(Arc::new(Self {
2608 input: new_input,
2609 cache,
2610 recompute_calls: Arc::clone(&self.recompute_calls),
2611 }))
2612 }
2613 // Intentionally does **not** override
2614 // `with_new_children_and_same_properties` — relies on the trait
2615 // default that falls back to `with_new_children`.
2616 fn execute(
2617 &self,
2618 _partition: usize,
2619 _context: Arc<TaskContext>,
2620 ) -> Result<SendableRecordBatchStream> {
2621 unimplemented!()
2622 }
2623 }
2624
2625 /// Cover the three short-circuit layers of
2626 /// [`replace_children_if_necessary`].
2627 #[test]
2628 fn test_replace_children_if_necessary_layers() -> Result<()> {
2629 use std::sync::atomic::Ordering;
2630
2631 // Two leaves that share the same `PlanProperties` Arc but sit behind
2632 // distinct `Arc<dyn ExecutionPlan>` pointers.
2633 let leaf_props = Arc::new(PlanProperties::new(
2634 EquivalenceProperties::new(Arc::new(Schema::empty())),
2635 Partitioning::UnknownPartitioning(1),
2636 EmissionType::Final,
2637 Boundedness::Bounded,
2638 ));
2639 let leaf_a: Arc<dyn ExecutionPlan> =
2640 Arc::new(WithChildrenTestLeaf::new(Arc::clone(&leaf_props)));
2641 let leaf_b: Arc<dyn ExecutionPlan> =
2642 Arc::new(WithChildrenTestLeaf::new(Arc::clone(&leaf_props)));
2643 // A third leaf with a *different* `PlanProperties` Arc — for layer 3.
2644 let leaf_c_props = Arc::new(PlanProperties::new(
2645 EquivalenceProperties::new(Arc::new(Schema::empty())),
2646 Partitioning::UnknownPartitioning(1),
2647 EmissionType::Final,
2648 Boundedness::Bounded,
2649 ));
2650 let leaf_c: Arc<dyn ExecutionPlan> =
2651 Arc::new(WithChildrenTestLeaf::new(leaf_c_props));
2652
2653 let parent = Arc::new(WithChildrenTestParent::new(Arc::clone(&leaf_a)));
2654 let parent_dyn: Arc<dyn ExecutionPlan> = Arc::clone(&parent) as _;
2655 let orig_props = Arc::clone(parent.properties());
2656
2657 // Layer 1: same child pointer → returns the original plan Arc verbatim.
2658 let out = replace_children_if_necessary(
2659 Arc::clone(&parent_dyn),
2660 vec![Arc::clone(&leaf_a)],
2661 )?;
2662 assert!(Arc::ptr_eq(&out, &parent_dyn));
2663 assert_eq!(parent.recompute_calls.load(Ordering::SeqCst), 0);
2664 assert_eq!(parent.fast_path_calls.load(Ordering::SeqCst), 0);
2665
2666 // Layer 2: distinct child Arc, but children share the same
2667 // `PlanProperties` Arc → fast path, parent's `PlanProperties` cache
2668 // Arc is reused (not reallocated).
2669 assert!(!Arc::ptr_eq(&leaf_a, &leaf_b));
2670 assert!(Arc::ptr_eq(leaf_a.properties(), leaf_b.properties()));
2671 let out = replace_children_if_necessary(
2672 Arc::clone(&parent_dyn),
2673 vec![Arc::clone(&leaf_b)],
2674 )?;
2675 assert!(Arc::ptr_eq(out.properties(), &orig_props));
2676 assert_eq!(parent.recompute_calls.load(Ordering::SeqCst), 0);
2677 assert_eq!(parent.fast_path_calls.load(Ordering::SeqCst), 1);
2678
2679 // Layer 3: child's `PlanProperties` Arc differs → full recompute.
2680 assert!(!Arc::ptr_eq(leaf_a.properties(), leaf_c.properties()));
2681 let out = replace_children_if_necessary(
2682 Arc::clone(&parent_dyn),
2683 vec![Arc::clone(&leaf_c)],
2684 )?;
2685 assert!(!Arc::ptr_eq(out.properties(), &orig_props));
2686 assert_eq!(parent.recompute_calls.load(Ordering::SeqCst), 1);
2687 assert_eq!(parent.fast_path_calls.load(Ordering::SeqCst), 1);
2688
2689 Ok(())
2690 }
2691
2692 /// A plan that does not override `with_new_children_and_same_properties`
2693 /// (per @kosiew's review on #23332) must still be routed through
2694 /// `with_new_children` when the helper hits the "same properties"
2695 /// branch. The default trait implementation forwards to
2696 /// `with_new_children`, so downstream / external `ExecutionPlan`
2697 /// implementations keep the semantics-preserving path.
2698 #[test]
2699 fn test_replace_children_if_necessary_default_fallback() -> Result<()> {
2700 use std::sync::atomic::Ordering;
2701
2702 let leaf_props = Arc::new(PlanProperties::new(
2703 EquivalenceProperties::new(Arc::new(Schema::empty())),
2704 Partitioning::UnknownPartitioning(1),
2705 EmissionType::Final,
2706 Boundedness::Bounded,
2707 ));
2708 let leaf_a: Arc<dyn ExecutionPlan> =
2709 Arc::new(WithChildrenTestLeaf::new(Arc::clone(&leaf_props)));
2710 let leaf_b: Arc<dyn ExecutionPlan> =
2711 Arc::new(WithChildrenTestLeaf::new(Arc::clone(&leaf_props)));
2712 assert!(!Arc::ptr_eq(&leaf_a, &leaf_b));
2713 assert!(Arc::ptr_eq(leaf_a.properties(), leaf_b.properties()));
2714
2715 let parent = Arc::new(WithChildrenTestParentDefault::new(Arc::clone(&leaf_a)));
2716 let parent_dyn: Arc<dyn ExecutionPlan> = Arc::clone(&parent) as _;
2717
2718 // Using the same child means we return the original plan Arc verbatim, so even when
2719 // the `replace_children` `ChildrenPropertiesMode::Keep` path is not defined,
2720 // we do not recompute.
2721 let out = replace_children_if_necessary(
2722 Arc::clone(&parent_dyn),
2723 vec![Arc::clone(&leaf_a)],
2724 )?;
2725 assert!(Arc::ptr_eq(&out, &parent_dyn));
2726 assert_eq!(parent.recompute_calls.load(Ordering::SeqCst), 0);
2727
2728 // Using a distinct child but the same `PlanProperties` Arc means the helper
2729 // attempts to enter the Keep branch. If it does not exist, we fall back
2730 // to recomputation.
2731 let out = replace_children_if_necessary(
2732 Arc::clone(&parent_dyn),
2733 vec![Arc::clone(&leaf_b)],
2734 )?;
2735 // `with_new_children` was invoked exactly once via the default.
2736 assert_eq!(parent.recompute_calls.load(Ordering::SeqCst), 1);
2737 // The returned plan has a freshly-recomputed `PlanProperties` Arc,
2738 // so it differs from the parent's original cache. This confirms
2739 // the fallback ran and did not short-circuit.
2740 assert!(!Arc::ptr_eq(out.properties(), parent.properties()));
2741
2742 Ok(())
2743 }
2744
2745 /// A test node that holds a fixed list of expressions, used to test
2746 /// `apply_expressions` behavior.
2747 #[derive(Debug)]
2748 struct MultiExprExec {
2749 exprs: Vec<Arc<dyn PhysicalExpr>>,
2750 children: Vec<Arc<dyn ExecutionPlan>>,
2751 }
2752
2753 impl DisplayAs for MultiExprExec {
2754 fn fmt_as(
2755 &self,
2756 _t: DisplayFormatType,
2757 _f: &mut std::fmt::Formatter,
2758 ) -> std::fmt::Result {
2759 unimplemented!()
2760 }
2761 }
2762
2763 impl ExecutionPlan for MultiExprExec {
2764 fn name(&self) -> &'static str {
2765 "MultiExprExec"
2766 }
2767
2768 fn properties(&self) -> &Arc<PlanProperties> {
2769 unimplemented!()
2770 }
2771
2772 fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
2773 self.children.iter().collect()
2774 }
2775
2776 fn with_new_children(
2777 self: Arc<Self>,
2778 _: Vec<Arc<dyn ExecutionPlan>>,
2779 ) -> Result<Arc<dyn ExecutionPlan>> {
2780 unimplemented!()
2781 }
2782
2783 fn apply_expressions(
2784 &self,
2785 f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
2786 ) -> Result<TreeNodeRecursion> {
2787 apply_expression_roots(&self.exprs, f)
2788 }
2789
2790 fn execute(
2791 &self,
2792 _partition: usize,
2793 _context: Arc<TaskContext>,
2794 ) -> Result<SendableRecordBatchStream> {
2795 unimplemented!()
2796 }
2797
2798 fn partition_statistics(
2799 &self,
2800 _partition: Option<usize>,
2801 ) -> Result<Arc<Statistics>> {
2802 unimplemented!()
2803 }
2804 }
2805
2806 /// Returns a simple literal `Arc<dyn PhysicalExpr>` for use in tests.
2807 fn lit_expr(val: i64) -> Arc<dyn PhysicalExpr> {
2808 use datafusion_physical_expr::expressions::Literal;
2809 Arc::new(Literal::new(datafusion_common::ScalarValue::Int64(Some(
2810 val,
2811 ))))
2812 }
2813
2814 /// `apply_expressions` visits all expressions when `f` always returns `Continue`.
2815 #[test]
2816 fn test_apply_expressions_continue_visits_all() -> Result<()> {
2817 let plan = MultiExprExec {
2818 exprs: vec![lit_expr(1), lit_expr(2), lit_expr(3)],
2819 children: vec![],
2820 };
2821 let mut visited = 0usize;
2822 plan.apply_expressions(&mut |_expr| {
2823 visited += 1;
2824 Ok(TreeNodeRecursion::Continue)
2825 })?;
2826 assert_eq!(visited, 3);
2827 Ok(())
2828 }
2829
2830 #[test]
2831 fn test_apply_expressions_stop_halts_early() -> Result<()> {
2832 let plan = MultiExprExec {
2833 exprs: vec![lit_expr(1), lit_expr(2), lit_expr(3)],
2834 children: vec![],
2835 };
2836 let mut visited = 0usize;
2837 let tnr = plan.apply_expressions(&mut |_expr| {
2838 visited += 1;
2839 Ok(TreeNodeRecursion::Stop)
2840 })?;
2841 // Only the first expression is visited; the rest are skipped.
2842 assert_eq!(visited, 1);
2843 assert_eq!(tnr, TreeNodeRecursion::Stop);
2844 Ok(())
2845 }
2846
2847 #[test]
2848 fn test_apply_expressions_jump_visits_next_root() -> Result<()> {
2849 let plan = MultiExprExec {
2850 exprs: vec![lit_expr(1), lit_expr(2), lit_expr(3)],
2851 children: vec![],
2852 };
2853 let mut visited = 0usize;
2854 let tnr = plan.apply_expressions(&mut |_expr| {
2855 visited += 1;
2856 Ok(TreeNodeRecursion::Jump)
2857 })?;
2858 assert_eq!(visited, 3);
2859 assert_eq!(tnr, TreeNodeRecursion::Continue);
2860 Ok(())
2861 }
2862
2863 #[test]
2864 fn test_apply_expressions_does_not_recurse() -> Result<()> {
2865 use datafusion_physical_expr::expressions::NegativeExpr;
2866
2867 let child: Arc<dyn ExecutionPlan> = Arc::new(MultiExprExec {
2868 exprs: vec![lit_expr(2)],
2869 children: vec![],
2870 });
2871 let nested: Arc<dyn PhysicalExpr> = Arc::new(NegativeExpr::new(lit_expr(1)));
2872 let plan = MultiExprExec {
2873 exprs: vec![nested],
2874 children: vec![child],
2875 };
2876
2877 let mut visited = 0;
2878 plan.apply_expressions(&mut |expr| {
2879 visited += 1;
2880 assert!(expr.is::<NegativeExpr>());
2881 Ok(TreeNodeRecursion::Continue)
2882 })?;
2883 assert_eq!(visited, 1);
2884 Ok(())
2885 }
2886
2887 #[test]
2888 fn test_apply_expressions_callback_can_retain_arc() -> Result<()> {
2889 let expected = lit_expr(1);
2890 let plan = MultiExprExec {
2891 exprs: vec![Arc::clone(&expected)],
2892 children: vec![],
2893 };
2894 let mut retained = None;
2895 plan.apply_expressions(&mut |expr| {
2896 retained = Some(Arc::clone(expr));
2897 Ok(TreeNodeRecursion::Continue)
2898 })?;
2899 drop(plan);
2900
2901 assert!(Arc::ptr_eq(
2902 &expected,
2903 retained
2904 .as_ref()
2905 .expect("callback should retain expression")
2906 ));
2907 Ok(())
2908 }
2909
2910 #[test]
2911 fn test_execution_plan_name() {
2912 let schema1 = Arc::new(Schema::empty());
2913 let default_name_exec = EmptyExec::new(schema1);
2914 assert_eq!(default_name_exec.name(), "EmptyExec");
2915
2916 let schema2 = Arc::new(Schema::empty());
2917 let renamed_exec = RenamedEmptyExec::new(schema2);
2918 assert_eq!(renamed_exec.name(), "MyRenamedEmptyExec");
2919 assert_eq!(RenamedEmptyExec::static_name(), "MyRenamedEmptyExec");
2920 }
2921
2922 #[test]
2923 fn test_execution_plan_downcast_delegates_to_downcast_delegate() {
2924 let schema = Arc::new(Schema::empty());
2925 let inner: Arc<dyn ExecutionPlan> = Arc::new(EmptyExec::new(schema));
2926 let wrapped: Arc<dyn ExecutionPlan> = Arc::new(DowncastDelegatingExec(inner));
2927 let nested: Arc<dyn ExecutionPlan> =
2928 Arc::new(DowncastDelegatingExec(Arc::clone(&wrapped)));
2929
2930 for plan in [wrapped.as_ref(), nested.as_ref()] {
2931 assert!(!plan.is::<DowncastDelegatingExec>());
2932 assert!(plan.downcast_ref::<DowncastDelegatingExec>().is_none());
2933 assert!(plan.is::<EmptyExec>());
2934 assert!(plan.downcast_ref::<EmptyExec>().is_some());
2935 assert!(!plan.is::<RenamedEmptyExec>());
2936 assert!(plan.downcast_ref::<RenamedEmptyExec>().is_none());
2937 }
2938 }
2939
2940 /// A compilation test to ensure that the `ExecutionPlan::name()` method can
2941 /// be called from a trait object.
2942 /// Related ticket: https://github.com/apache/datafusion/pull/11047
2943 #[expect(unused)]
2944 fn use_execution_plan_as_trait_object(plan: &dyn ExecutionPlan) {
2945 let _ = plan.name();
2946 }
2947
2948 #[test]
2949 fn buffer_exec_does_not_need_data_exchange() {
2950 let schema = Arc::new(Schema::empty());
2951 let input: Arc<dyn ExecutionPlan> = Arc::new(MockExec::new(vec![], schema));
2952 let buffer: Arc<dyn ExecutionPlan> = Arc::new(BufferExec::new(input, 1024));
2953
2954 assert!(!need_data_exchange(buffer));
2955 }
2956
2957 #[test]
2958 fn test_check_not_null_constraints_accept_non_null() -> Result<()> {
2959 check_not_null_constraints(
2960 RecordBatch::try_new(
2961 Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])),
2962 vec![Arc::new(Int32Array::from(vec![Some(1), Some(2), Some(3)]))],
2963 )?,
2964 &vec![0],
2965 )?;
2966 Ok(())
2967 }
2968
2969 #[test]
2970 fn test_check_not_null_constraints_reject_null() -> Result<()> {
2971 let result = check_not_null_constraints(
2972 RecordBatch::try_new(
2973 Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])),
2974 vec![Arc::new(Int32Array::from(vec![Some(1), None, Some(3)]))],
2975 )?,
2976 &vec![0],
2977 );
2978 assert!(result.is_err());
2979 assert_eq!(
2980 result.err().unwrap().strip_backtrace(),
2981 "Execution error: Invalid batch column at '0' has null but schema specifies non-nullable",
2982 );
2983 Ok(())
2984 }
2985
2986 #[test]
2987 fn test_check_not_null_constraints_with_run_end_array() -> Result<()> {
2988 // some null value inside REE array
2989 let run_ends = Int32Array::from(vec![1, 2, 3, 4]);
2990 let values = Int32Array::from(vec![Some(0), None, Some(1), None]);
2991 let run_end_array = RunArray::try_new(&run_ends, &values)?;
2992 let result = check_not_null_constraints(
2993 RecordBatch::try_new(
2994 Arc::new(Schema::new(vec![Field::new(
2995 "a",
2996 run_end_array.data_type().to_owned(),
2997 true,
2998 )])),
2999 vec![Arc::new(run_end_array)],
3000 )?,
3001 &vec![0],
3002 );
3003 assert!(result.is_err());
3004 assert_eq!(
3005 result.err().unwrap().strip_backtrace(),
3006 "Execution error: Invalid batch column at '0' has null but schema specifies non-nullable",
3007 );
3008 Ok(())
3009 }
3010
3011 #[test]
3012 fn test_check_not_null_constraints_with_dictionary_array_with_null() -> Result<()> {
3013 let values = Arc::new(Int32Array::from(vec![Some(1), None, Some(3), Some(4)]));
3014 let keys = Int32Array::from(vec![0, 1, 2, 3]);
3015 let dictionary = DictionaryArray::new(keys, values);
3016 let result = check_not_null_constraints(
3017 RecordBatch::try_new(
3018 Arc::new(Schema::new(vec![Field::new(
3019 "a",
3020 dictionary.data_type().to_owned(),
3021 true,
3022 )])),
3023 vec![Arc::new(dictionary)],
3024 )?,
3025 &vec![0],
3026 );
3027 assert!(result.is_err());
3028 assert_eq!(
3029 result.err().unwrap().strip_backtrace(),
3030 "Execution error: Invalid batch column at '0' has null but schema specifies non-nullable",
3031 );
3032 Ok(())
3033 }
3034
3035 #[test]
3036 fn test_check_not_null_constraints_with_dictionary_masking_null() -> Result<()> {
3037 // some null value marked out by dictionary array
3038 let values = Arc::new(Int32Array::from(vec![
3039 Some(1),
3040 None, // this null value is masked by dictionary keys
3041 Some(3),
3042 Some(4),
3043 ]));
3044 let keys = Int32Array::from(vec![0, /*1,*/ 2, 3]);
3045 let dictionary = DictionaryArray::new(keys, values);
3046 check_not_null_constraints(
3047 RecordBatch::try_new(
3048 Arc::new(Schema::new(vec![Field::new(
3049 "a",
3050 dictionary.data_type().to_owned(),
3051 true,
3052 )])),
3053 vec![Arc::new(dictionary)],
3054 )?,
3055 &vec![0],
3056 )?;
3057 Ok(())
3058 }
3059
3060 #[test]
3061 fn test_check_not_null_constraints_on_null_type() -> Result<()> {
3062 // null value of Null type
3063 let result = check_not_null_constraints(
3064 RecordBatch::try_new(
3065 Arc::new(Schema::new(vec![Field::new("a", DataType::Null, true)])),
3066 vec![Arc::new(NullArray::new(3))],
3067 )?,
3068 &vec![0],
3069 );
3070 assert!(result.is_err());
3071 assert_eq!(
3072 result.err().unwrap().strip_backtrace(),
3073 "Execution error: Invalid batch column at '0' has null but schema specifies non-nullable",
3074 );
3075 Ok(())
3076 }
3077}