datafusion_physical_expr_common/physical_expr.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
18use std::any::Any;
19use std::fmt;
20use std::fmt::{Debug, Display, Formatter};
21use std::hash::{Hash, Hasher};
22use std::sync::Arc;
23
24use crate::utils::scatter;
25
26use arrow::array::{Array, ArrayRef, BooleanArray, new_empty_array};
27use arrow::compute::filter_record_batch;
28use arrow::datatypes::{DataType, Field, FieldRef, Schema};
29use arrow::record_batch::RecordBatch;
30use datafusion_common::tree_node::{
31 Transformed, TransformedResult, TreeNode, TreeNodeRecursion,
32};
33use datafusion_common::{
34 Result, ScalarValue, assert_eq_or_internal_err, exec_err, not_impl_err,
35};
36use datafusion_expr_common::columnar_value::ColumnarValue;
37use datafusion_expr_common::interval_arithmetic::Interval;
38use datafusion_expr_common::placement::ExpressionPlacement;
39use datafusion_expr_common::sort_properties::ExprProperties;
40#[expect(deprecated)]
41use datafusion_expr_common::statistics::Distribution;
42
43use itertools::izip;
44
45/// Shared [`PhysicalExpr`].
46pub type PhysicalExprRef = Arc<dyn PhysicalExpr>;
47
48/// [`PhysicalExpr`]s represent expressions such as `A + 1` or `CAST(c1 AS int)`.
49///
50/// `PhysicalExpr` knows its type, nullability and can be evaluated directly on
51/// a [`RecordBatch`] (see [`Self::evaluate`]).
52///
53/// `PhysicalExpr` are the physical counterpart to [`Expr`] used in logical
54/// planning. They are typically created from [`Expr`] by a [`PhysicalPlanner`]
55/// invoked from a higher level API
56///
57/// Some important examples of `PhysicalExpr` are:
58/// * [`Column`]: Represents a column at a given index in a RecordBatch
59///
60/// To create `PhysicalExpr` from `Expr`, see
61/// * [`SessionContext::create_physical_expr`]: A high level API
62/// * [`create_physical_expr`]: A low level API
63///
64/// # Formatting `PhysicalExpr` as strings
65/// There are three ways to format `PhysicalExpr` as a string:
66/// * [`Debug`]: Standard Rust debugging format (e.g. `Constant { value: ... }`)
67/// * [`Display`]: Detailed SQL-like format that shows expression structure (e.g. (`Utf8 ("foobar")`). This is often used for debugging and tests
68/// * [`Self::fmt_sql`]: SQL-like human readable format (e.g. ('foobar')`), See also [`sql_fmt`]
69///
70/// [`SessionContext::create_physical_expr`]: https://docs.rs/datafusion/latest/datafusion/execution/context/struct.SessionContext.html#method.create_physical_expr
71/// [`PhysicalPlanner`]: https://docs.rs/datafusion/latest/datafusion/physical_planner/trait.PhysicalPlanner.html
72/// [`Expr`]: https://docs.rs/datafusion/latest/datafusion/logical_expr/enum.Expr.html
73/// [`create_physical_expr`]: https://docs.rs/datafusion/latest/datafusion/physical_expr/fn.create_physical_expr.html
74/// [`Column`]: https://docs.rs/datafusion/latest/datafusion/physical_expr/expressions/struct.Column.html
75pub trait PhysicalExpr: Any + Send + Sync + Display + Debug + DynEq + DynHash {
76 /// Get the data type of this expression, given the schema of the input.
77 /// Returns an error if the data type cannot be determined, ex. if the
78 /// schema is missing a required field.
79 fn data_type(&self, input_schema: &Schema) -> Result<DataType> {
80 Ok(self.return_field(input_schema)?.data_type().to_owned())
81 }
82 /// Determine whether this expression is nullable, given the schema of the input
83 fn nullable(&self, input_schema: &Schema) -> Result<bool> {
84 Ok(self.return_field(input_schema)?.is_nullable())
85 }
86 /// Evaluate an expression against a RecordBatch
87 fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue>;
88 /// The output field associated with this expression
89 fn return_field(&self, input_schema: &Schema) -> Result<FieldRef> {
90 Ok(Arc::new(Field::new(
91 format!("{self}"),
92 self.data_type(input_schema)?,
93 self.nullable(input_schema)?,
94 )))
95 }
96 /// Evaluate an expression against a RecordBatch after first applying a validity array
97 ///
98 /// # Errors
99 ///
100 /// Returns an `Err` if the expression could not be evaluated or if the length of the
101 /// `selection` validity array and the number of row in `batch` is not equal.
102 fn evaluate_selection(
103 &self,
104 batch: &RecordBatch,
105 selection: &BooleanArray,
106 ) -> Result<ColumnarValue> {
107 let row_count = batch.num_rows();
108 if row_count != selection.len() {
109 return exec_err!(
110 "Selection array length does not match batch row count: {} != {row_count}",
111 selection.len()
112 );
113 }
114
115 // First, check if we can avoid filtering altogether.
116 if selection.null_count() == 0 && !selection.has_false() {
117 // All values from the `selection` filter are true and match the input batch.
118 // No need to perform any filtering.
119 return self.evaluate(batch);
120 }
121
122 // Next, prepare the result array for each 'true' row in the selection vector.
123 let filtered_result = if !selection.has_true() {
124 // Do not call `evaluate` when the selection is empty.
125 // `evaluate_selection` is used to conditionally evaluate expressions.
126 // When the expression in question is fallible, evaluating it with an empty
127 // record batch may trigger a runtime error (e.g. division by zero).
128 //
129 // Instead, create an empty array matching the expected return type.
130 let datatype = self.data_type(batch.schema_ref().as_ref())?;
131 ColumnarValue::Array(new_empty_array(&datatype))
132 } else {
133 // If we reach this point, there's no other option than to filter the batch.
134 // This is a fairly costly operation since it requires creating partial copies
135 // (worst case of length `row_count - 1`) of all the arrays in the record batch.
136 // The resulting `filtered_batch` will contain one row per true in `selection`.
137 let filtered_batch = filter_record_batch(batch, selection)?;
138 self.evaluate(&filtered_batch)?
139 };
140
141 // Finally, scatter the filtered result array so that the indices match the input rows again.
142 match &filtered_result {
143 ColumnarValue::Array(a) => {
144 scatter(selection, a.as_ref()).map(ColumnarValue::Array)
145 }
146 ColumnarValue::Scalar(ScalarValue::Boolean(value)) => {
147 // When the scalar is true or false, skip the scatter process
148 if let Some(v) = value {
149 if *v {
150 Ok(ColumnarValue::from(Arc::new(selection.clone()) as ArrayRef))
151 } else {
152 Ok(filtered_result)
153 }
154 } else {
155 let array = BooleanArray::from(vec![None; row_count]);
156 scatter(selection, &array).map(ColumnarValue::Array)
157 }
158 }
159 ColumnarValue::Scalar(_) => Ok(filtered_result),
160 }
161 }
162
163 /// Get a list of child PhysicalExpr that provide the input for this expr.
164 fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>>;
165
166 /// Returns a new PhysicalExpr where all children were replaced by new exprs.
167 ///
168 /// If the implementation returns a [`PhysicalExpr::expression_id`], then
169 /// the identifier should be preserved by the new expression.
170 fn with_new_children(
171 self: Arc<Self>,
172 children: Vec<Arc<dyn PhysicalExpr>>,
173 ) -> Result<Arc<dyn PhysicalExpr>>;
174
175 /// Computes the output interval for the expression, given the input
176 /// intervals.
177 ///
178 /// # Parameters
179 ///
180 /// * `children` are the intervals for the children (inputs) of this
181 /// expression.
182 ///
183 /// # Returns
184 ///
185 /// A `Result` containing the output interval for the expression in
186 /// case of success, or an error object in case of failure.
187 ///
188 /// Note that the output bounds must form an **envelope** that contains all
189 /// possible outputs of the expression given the input bounds. While
190 /// expressions should output the tightest possible bounds, they do not need
191 /// to be exact and can be conservative.
192 ///
193 /// # Example
194 ///
195 /// If the expression is `a + b`, and the input intervals are `a: [1, 2]`
196 /// and `b: [3, 4]`, then the output interval would be `[4, 6]`.
197 ///
198 /// If the expression is `sin(a)`, it is correct (though not precise) to
199 /// produce the interval `[-1, 1]` for any input interval for `a`.
200 fn evaluate_bounds(&self, _children: &[&Interval]) -> Result<Interval> {
201 not_impl_err!("Not implemented for {self}")
202 }
203
204 /// Updates bounds for child expressions, given a known interval for this
205 /// expression.
206 ///
207 /// This is used to propagate constraints down through an expression tree.
208 ///
209 /// # Parameters
210 ///
211 /// * `interval` is the currently known interval for this expression.
212 /// * `children` are the current intervals for the children of this expression.
213 ///
214 /// # Returns
215 ///
216 /// A `Result` containing a `Vec` of new intervals for the children (in order)
217 /// in case of success, or an error object in case of failure.
218 ///
219 /// If constraint propagation reveals an infeasibility for any child, returns
220 /// [`None`]. If none of the children intervals change as a result of
221 /// propagation, may return an empty vector instead of cloning `children`.
222 /// This is the default (and conservative) return value.
223 ///
224 /// # Example
225 ///
226 /// If the expression is `a + b`, the current `interval` is `[4, 5]` and the
227 /// inputs `a` and `b` are respectively given as `[0, 2]` and `[-∞, 4]`, then
228 /// propagation would return `[0, 2]` and `[2, 4]` as `b` must be at least
229 /// `2` to make the output at least `4`.
230 fn propagate_constraints(
231 &self,
232 _interval: &Interval,
233 _children: &[&Interval],
234 ) -> Result<Option<Vec<Interval>>> {
235 Ok(Some(vec![]))
236 }
237
238 /// Computes the output statistics for the expression, given the input
239 /// statistics.
240 ///
241 /// # Parameters
242 ///
243 /// * `children` are the statistics for the children (inputs) of this
244 /// expression.
245 ///
246 /// # Returns
247 ///
248 /// A `Result` containing the output statistics for the expression in
249 /// case of success, or an error object in case of failure.
250 ///
251 /// Expressions (should) implement this function and utilize the independence
252 /// assumption, match on children distribution types and compute the output
253 /// statistics accordingly. The default implementation simply creates an
254 /// unknown output distribution by combining input ranges. This logic loses
255 /// distribution information, but is a safe default.
256 #[deprecated(
257 since = "54.0.0",
258 note = "Part of the unused Statistics V2 framework; see https://github.com/apache/datafusion/pull/22071"
259 )]
260 #[expect(deprecated)]
261 fn evaluate_statistics(&self, children: &[&Distribution]) -> Result<Distribution> {
262 let children_ranges = children
263 .iter()
264 .map(|c| c.range())
265 .collect::<Result<Vec<_>>>()?;
266 let children_ranges_refs = children_ranges.iter().collect::<Vec<_>>();
267 let output_interval = self.evaluate_bounds(children_ranges_refs.as_slice())?;
268 let dt = output_interval.data_type();
269 if dt.eq(&DataType::Boolean) {
270 let p = if output_interval.eq(&Interval::TRUE) {
271 ScalarValue::new_one(&dt)
272 } else if output_interval.eq(&Interval::FALSE) {
273 ScalarValue::new_zero(&dt)
274 } else {
275 ScalarValue::try_from(&dt)
276 }?;
277 Distribution::new_bernoulli(p)
278 } else {
279 Distribution::new_from_interval(output_interval)
280 }
281 }
282
283 /// Updates children statistics using the given parent statistic for this
284 /// expression.
285 ///
286 /// This is used to propagate statistics down through an expression tree.
287 ///
288 /// # Parameters
289 ///
290 /// * `parent` is the currently known statistics for this expression.
291 /// * `children` are the current statistics for the children of this expression.
292 ///
293 /// # Returns
294 ///
295 /// A `Result` containing a `Vec` of new statistics for the children (in order)
296 /// in case of success, or an error object in case of failure.
297 ///
298 /// If statistics propagation reveals an infeasibility for any child, returns
299 /// [`None`]. If none of the children statistics change as a result of
300 /// propagation, may return an empty vector instead of cloning `children`.
301 /// This is the default (and conservative) return value.
302 ///
303 /// Expressions (should) implement this function and apply Bayes rule to
304 /// reconcile and update parent/children statistics. This involves utilizing
305 /// the independence assumption, and matching on distribution types. The
306 /// default implementation simply creates an unknown distribution if it can
307 /// narrow the range by propagating ranges. This logic loses distribution
308 /// information, but is a safe default.
309 #[deprecated(
310 since = "54.0.0",
311 note = "Part of the unused Statistics V2 framework; see https://github.com/apache/datafusion/pull/22071"
312 )]
313 #[expect(deprecated)]
314 fn propagate_statistics(
315 &self,
316 parent: &Distribution,
317 children: &[&Distribution],
318 ) -> Result<Option<Vec<Distribution>>> {
319 let children_ranges = children
320 .iter()
321 .map(|c| c.range())
322 .collect::<Result<Vec<_>>>()?;
323 let children_ranges_refs = children_ranges.iter().collect::<Vec<_>>();
324 let parent_range = parent.range()?;
325 let Some(propagated_children) =
326 self.propagate_constraints(&parent_range, children_ranges_refs.as_slice())?
327 else {
328 return Ok(None);
329 };
330 izip!(propagated_children.into_iter(), children_ranges, children)
331 .map(|(new_interval, old_interval, child)| {
332 if new_interval == old_interval {
333 // We weren't able to narrow the range, preserve the old statistics.
334 Ok((*child).clone())
335 } else if new_interval.data_type().eq(&DataType::Boolean) {
336 let dt = old_interval.data_type();
337 let p = if new_interval.eq(&Interval::TRUE) {
338 ScalarValue::new_one(&dt)
339 } else if new_interval.eq(&Interval::FALSE) {
340 ScalarValue::new_zero(&dt)
341 } else {
342 unreachable!("Given that we have a range reduction for a boolean interval, we should have certainty")
343 }?;
344 Distribution::new_bernoulli(p)
345 } else {
346 Distribution::new_from_interval(new_interval)
347 }
348 })
349 .collect::<Result<_>>()
350 .map(Some)
351 }
352
353 /// Calculates the properties of this [`PhysicalExpr`] based on its
354 /// children's properties (i.e. order and range), recursively aggregating
355 /// the information from its children. In cases where the [`PhysicalExpr`]
356 /// has no children (e.g., `Literal` or `Column`), these properties should
357 /// be specified externally, as the function defaults to unknown properties.
358 fn get_properties(&self, _children: &[ExprProperties]) -> Result<ExprProperties> {
359 Ok(ExprProperties::new_unknown())
360 }
361
362 /// Format this `PhysicalExpr` in nice human readable "SQL" format
363 ///
364 /// Specifically, this format is designed to be readable by humans, at the
365 /// expense of details. Use `Display` or `Debug` for more detailed
366 /// representation.
367 ///
368 /// See the [`fmt_sql`] function for an example of printing `PhysicalExpr`s as SQL.
369 fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result;
370
371 /// Take a snapshot of this `PhysicalExpr`, if it is dynamic.
372 ///
373 /// "Dynamic" in this case means containing references to structures that may change
374 /// during plan execution, such as hash tables.
375 ///
376 /// This method is used to capture the current state of `PhysicalExpr`s that may contain
377 /// dynamic references to other operators in order to serialize it over the wire
378 /// or treat it via downcast matching.
379 ///
380 /// You should not call this method directly as it does not handle recursion.
381 /// Instead use [`snapshot_physical_expr`] to handle recursion and capture the
382 /// full state of the `PhysicalExpr`.
383 ///
384 /// This is expected to return "simple" expressions that do not have mutable state
385 /// and are composed of DataFusion's built-in `PhysicalExpr` implementations.
386 /// Callers however should *not* assume anything about the returned expressions
387 /// since callers and implementers may not agree on what "simple" or "built-in"
388 /// means.
389 /// In other words, if you need to serialize a `PhysicalExpr` across the wire
390 /// you should call this method and then try to serialize the result,
391 /// but you should handle unknown or unexpected `PhysicalExpr` implementations gracefully
392 /// just as if you had not called this method at all.
393 ///
394 /// In particular, consider:
395 /// * A `PhysicalExpr` that references the current state of a `datafusion::physical_plan::TopK`
396 /// that is involved in a query with `SELECT * FROM t1 ORDER BY a LIMIT 10`.
397 /// This function may return something like `a >= 12`.
398 /// * A `PhysicalExpr` that references the current state of a `datafusion::physical_plan::joins::HashJoinExec`
399 /// from a query such as `SELECT * FROM t1 JOIN t2 ON t1.a = t2.b`.
400 /// This function may return something like `t2.b IN (1, 5, 7)`.
401 ///
402 /// A system or function that can only deal with a hardcoded set of `PhysicalExpr` implementations
403 /// or needs to serialize this state to bytes may not be able to handle these dynamic references.
404 /// In such cases, we should return a simplified version of the `PhysicalExpr` that does not
405 /// contain these dynamic references.
406 ///
407 /// Systems that implement remote execution of plans, e.g. serialize a portion of the query plan
408 /// and send it across the wire to a remote executor may want to call this method after
409 /// every batch on the source side and broadcast / update the current snapshot to the remote executor.
410 ///
411 /// Note for implementers: this method should *not* handle recursion.
412 /// Recursion is handled in [`snapshot_physical_expr`].
413 fn snapshot(&self) -> Result<Option<Arc<dyn PhysicalExpr>>> {
414 // By default, we return None to indicate that this PhysicalExpr does not
415 // have any dynamic references or state.
416 // This is a safe default behavior.
417 Ok(None)
418 }
419
420 /// Returns the generation of this `PhysicalExpr` for snapshotting purposes.
421 /// The generation is an arbitrary u64 that can be used to track changes
422 /// in the state of the `PhysicalExpr` over time without having to do an exhaustive comparison.
423 /// This is useful to avoid unnecessary computation or serialization if there are no changes to the expression.
424 /// In particular, dynamic expressions that may change over time; this allows cheap checks for changes.
425 /// Static expressions that do not change over time should return 0, as does the default implementation.
426 /// You should not call this method directly as it does not handle recursion.
427 /// Instead use [`snapshot_generation`] to handle recursion and capture the
428 /// full state of the `PhysicalExpr`.
429 fn snapshot_generation(&self) -> u64 {
430 // By default, we return 0 to indicate that this PhysicalExpr does not
431 // have any dynamic references or state.
432 // Since the recursive algorithm XORs the generations of all children the overall
433 // generation will be 0 if no children have a non-zero generation, meaning that
434 // static expressions will always return 0.
435 0
436 }
437
438 /// Returns true if the expression node is volatile, i.e. whether it can return
439 /// different results when evaluated multiple times with the same input.
440 ///
441 /// Note: unlike [`is_volatile`], this function does not consider inputs:
442 /// - `random()` returns `true`,
443 /// - `a + random()` returns `false` (because the operation `+` itself is not volatile.)
444 ///
445 /// The default to this function was set to `false` when it was created
446 /// to avoid imposing API churn on implementers, but this is not a safe default in general.
447 /// It is highly recommended that volatile expressions implement this method and return `true`.
448 /// This default may be removed in the future if it causes problems or we decide to
449 /// eat the cost of the breaking change and require all implementers to make a choice.
450 fn is_volatile_node(&self) -> bool {
451 false
452 }
453
454 /// Returns placement information for this expression.
455 ///
456 /// This is used by optimizers to make decisions about expression placement,
457 /// such as whether to push expressions down through projections.
458 ///
459 /// The default implementation returns [`ExpressionPlacement::KeepInPlace`].
460 fn placement(&self) -> ExpressionPlacement {
461 ExpressionPlacement::KeepInPlace
462 }
463
464 /// Return a stable, globally-unique identifier for this [`PhysicalExpr`], if it
465 /// has one.
466 ///
467 /// This identifier tracks which expressions which are connected (e.g. `DynamicFilterPhysicalExpr`
468 /// where two expressions may be different but store the same mutable inner state). Tracking
469 /// connected expressions helps preserve referential integrity within plan nodes
470 /// during serialization and deserialization.
471 ///
472 /// This id must be preserved across [`PhysicalExpr::with_new_children`] or any other
473 /// methods which may want to preserve identity.
474 ///
475 /// Default is `None`: the expression has no identity worth preserving across a
476 /// serialization boundary.
477 fn expression_id(&self) -> Option<u64> {
478 None
479 }
480
481 /// Serialize this expression to a [`PhysicalExprNode`] proto message.
482 ///
483 /// Returning `Ok(None)` means "this expression does not know how to
484 /// serialize itself"; the caller (typically `datafusion-proto`) will fall
485 /// back to its existing codec / extension paths. This matches today's
486 /// behavior for expressions that aren't built into `datafusion-proto`.
487 ///
488 /// Returning `Ok(Some(node))` means the expression has serialized itself
489 /// fully; the caller should not try any further fallback path.
490 ///
491 /// Returning `Err(_)` means a real serialization failure (e.g. the
492 /// expression knows it should serialize but a child failed).
493 ///
494 /// The motivating use case is letting expressions with private state
495 /// (e.g. `DynamicFilterPhysicalExpr`'s `RwLock`-protected inner fields)
496 /// reach into their own internals for `try_to_proto`/`try_from_proto`
497 /// without having to expose `pub` accessors to `datafusion-proto`. See
498 /// <https://github.com/apache/datafusion/issues/21835>.
499 ///
500 /// The `try_` prefix matches the fallible `try_from_proto` decode
501 /// constructors; both sides of the round-trip are fallible and named
502 /// consistently.
503 ///
504 /// [`PhysicalExprNode`]: datafusion_proto_models::protobuf::PhysicalExprNode
505 #[cfg(feature = "proto")]
506 fn try_to_proto(
507 &self,
508 _ctx: &proto_encode::PhysicalExprEncodeCtx<'_>,
509 ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalExprNode>> {
510 Ok(None)
511 }
512}
513
514/// Encode-side context for [`PhysicalExpr::try_to_proto`].
515///
516/// Expression authors only ever see [`proto_encode::PhysicalExprEncodeCtx`]:
517/// a concrete struct with stable methods. Internally it dispatches to a
518/// [`proto_encode::PhysicalExprEncode`] implementor that lives in
519/// `datafusion-proto`, which is what lets `physical-expr-common` stay free
520/// of `datafusion-proto` as a dep.
521///
522/// More specialized helpers (e.g. encoding UDFs/UDAFs/UDWFs through the
523/// extension codec) can be added to the context as expressions migrate;
524/// today they're not required because the encoder forwards to the existing
525/// codec via the proto converter.
526#[cfg(feature = "proto")]
527pub mod proto_encode {
528 use std::sync::Arc;
529
530 use datafusion_common::Result;
531 use datafusion_proto_models::protobuf::PhysicalExprNode;
532
533 use super::PhysicalExpr;
534
535 /// Encoder context handed to [`super::PhysicalExpr::try_to_proto`].
536 ///
537 /// Wraps an internal [`PhysicalExprEncode`] trait object so callers see a
538 /// stable concrete type while implementations can evolve in
539 /// `datafusion-proto`.
540 pub struct PhysicalExprEncodeCtx<'a> {
541 encoder: &'a dyn PhysicalExprEncode,
542 }
543
544 impl<'a> PhysicalExprEncodeCtx<'a> {
545 /// Construct a new encode context. Typically called by
546 /// `datafusion-proto`; expression authors receive `&PhysicalExprEncodeCtx`.
547 pub fn new(encoder: &'a dyn PhysicalExprEncode) -> Self {
548 Self { encoder }
549 }
550
551 /// Encode a child expression. Routes through the configured encoder
552 /// so dedup-aware encoding is preserved.
553 pub fn encode_child(
554 &self,
555 expr: &Arc<dyn PhysicalExpr>,
556 ) -> Result<PhysicalExprNode> {
557 self.encoder.encode(expr)
558 }
559
560 /// Encode a sequence of child expressions, preserving order.
561 ///
562 /// Convenience wrapper over [`Self::encode_child`] for expressions
563 /// holding a `repeated` proto field (e.g. the `list` of an `InList`).
564 /// The first encode error short-circuits.
565 pub fn encode_children_expressions<'b, I>(
566 &self,
567 exprs: I,
568 ) -> Result<Vec<PhysicalExprNode>>
569 where
570 I: IntoIterator<Item = &'b Arc<dyn PhysicalExpr>>,
571 {
572 exprs
573 .into_iter()
574 .map(|expr| self.encode_child(expr))
575 .collect()
576 }
577 }
578
579 /// Internal dispatch trait. Implementors live in `datafusion-proto` and
580 /// wrap the existing `PhysicalExtensionCodec` +
581 /// `PhysicalProtoConverterExtension` plumbing. Expression authors should
582 /// use [`PhysicalExprEncodeCtx`] instead of calling this directly.
583 pub trait PhysicalExprEncode {
584 /// Encode an expression to a protobuf node.
585 fn encode(&self, expr: &Arc<dyn PhysicalExpr>) -> Result<PhysicalExprNode>;
586 }
587}
588
589/// Decode-side counterpart to [`proto_encode`].
590///
591/// Expression authors implement an associated `try_from_proto` on their
592/// concrete type, with the signature
593///
594/// ```ignore
595/// fn try_from_proto(
596/// node: &PhysicalExprNode,
597/// ctx: &PhysicalExprDecodeCtx<'_>,
598/// ) -> Result<Arc<dyn PhysicalExpr>>
599/// ```
600///
601/// It takes the whole [`PhysicalExprNode`] — the exact inverse of what
602/// [`PhysicalExpr::try_to_proto`] returns — so the constructor can also see
603/// outer-node fields such as `expr_id`. The central match in
604/// `datafusion-proto` dispatches `ExprType` variants to these constructors.
605///
606/// As with the encode side, the public surface is a struct (not a `&dyn`
607/// trait) so future fields/helpers (registries for third-party expressions,
608/// schema-resolution caches, etc.) can be added without changing the
609/// signature every expression depends on.
610///
611/// [`PhysicalExprNode`]: datafusion_proto_models::protobuf::PhysicalExprNode
612#[cfg(feature = "proto")]
613pub mod proto_decode {
614 use std::sync::Arc;
615
616 use arrow::datatypes::Schema;
617 use datafusion_common::Result;
618 use datafusion_proto_models::protobuf::PhysicalExprNode;
619
620 use super::PhysicalExpr;
621
622 /// Open the outer [`PhysicalExprNode`] and assert it carries the expected
623 /// `ExprType` variant, returning the inner payload (auto-derefs through
624 /// `Box`) or bailing with an `Internal` error.
625 ///
626 /// Every `try_from_proto` starts with the same six-line `match`:
627 ///
628 /// ```ignore
629 /// let try_cast = match &node.expr_type {
630 /// Some(protobuf::physical_expr_node::ExprType::TryCast(x)) => x.as_ref(),
631 /// _ => return internal_err!("PhysicalExprNode is not a TryCastExpr"),
632 /// };
633 /// ```
634 ///
635 /// With this macro that collapses to:
636 ///
637 /// ```ignore
638 /// let try_cast = expect_expr_variant!(
639 /// node,
640 /// protobuf::physical_expr_node::ExprType::TryCast,
641 /// "TryCastExpr",
642 /// );
643 /// ```
644 ///
645 /// Pass the variant as a `::` path so the macro stays agnostic to how
646 /// the caller imports the proto types.
647 #[macro_export]
648 macro_rules! expect_expr_variant {
649 ($node:expr, $variant:path, $expr_name:literal $(,)?) => {{
650 match &$node.expr_type {
651 ::core::option::Option::Some($variant(inner)) => inner,
652 _ => {
653 return ::datafusion_common::internal_err!(concat!(
654 "PhysicalExprNode is not a ",
655 $expr_name
656 ));
657 }
658 }
659 }};
660 }
661 #[doc(inline)]
662 pub use expect_expr_variant;
663
664 /// Decoder context handed to per-expression `try_from_proto` constructors.
665 ///
666 /// Wraps an internal [`PhysicalExprDecode`] trait object plus a borrowed
667 /// schema. The trait stays an implementation detail of `datafusion-proto`;
668 /// expression authors only see this struct.
669 pub struct PhysicalExprDecodeCtx<'a> {
670 schema: &'a Schema,
671 decoder: &'a dyn PhysicalExprDecode,
672 }
673
674 impl<'a> PhysicalExprDecodeCtx<'a> {
675 /// Construct a new decode context. Typically called by
676 /// `datafusion-proto`; expression authors receive
677 /// `&PhysicalExprDecodeCtx`.
678 pub fn new(schema: &'a Schema, decoder: &'a dyn PhysicalExprDecode) -> Self {
679 Self { schema, decoder }
680 }
681
682 /// The schema bound to this decode context. Use it for column lookups,
683 /// data-type resolution, etc.
684 pub fn schema(&self) -> &Schema {
685 self.schema
686 }
687
688 /// Decode an expression node, recursing into child sub-expressions.
689 ///
690 /// Routes built-in `ExprType` variants through `datafusion-proto`'s
691 /// central match and forwards extension nodes to the registered codec
692 /// (today via [`PhysicalExtensionCodec::try_decode_expr`]; later via
693 /// a per-type registry — see #21835).
694 ///
695 /// [`PhysicalExtensionCodec::try_decode_expr`]: https://docs.rs/datafusion-proto/latest/datafusion_proto/physical_plan/trait.PhysicalExtensionCodec.html#method.try_decode_expr
696 pub fn decode(&self, node: &PhysicalExprNode) -> Result<Arc<dyn PhysicalExpr>> {
697 self.decoder.decode(node, self.schema)
698 }
699
700 /// Decode a required child node, erroring if it is absent.
701 ///
702 /// Proto child expressions are encoded as `Option<Box<PhysicalExprNode>>`;
703 /// pass the field directly (e.g. `node.expr.as_deref()`). `expr_name`
704 /// is the expression being decoded (e.g. `"InListExpr"`) and `field`
705 /// the proto field (e.g. `"expr"`); both are woven into the error so
706 /// it names *where* the missing field is, without each author
707 /// hand-rolling the string.
708 pub fn decode_required_expression(
709 &self,
710 node: Option<&PhysicalExprNode>,
711 expr_name: &str,
712 field: &str,
713 ) -> Result<Arc<dyn PhysicalExpr>> {
714 let node = node.ok_or_else(|| {
715 datafusion_common::internal_datafusion_err!(
716 "{expr_name} is missing required field '{field}'"
717 )
718 })?;
719 self.decode(node)
720 }
721
722 /// Decode a sequence of child nodes, preserving order.
723 ///
724 /// Convenience wrapper over [`Self::decode`] for expressions holding a
725 /// `repeated` proto field (e.g. the `list` of an `InList`). The first
726 /// decode error short-circuits.
727 pub fn decode_children_expressions<'b, I>(
728 &self,
729 nodes: I,
730 ) -> Result<Vec<Arc<dyn PhysicalExpr>>>
731 where
732 I: IntoIterator<Item = &'b PhysicalExprNode>,
733 {
734 nodes.into_iter().map(|node| self.decode(node)).collect()
735 }
736 }
737
738 /// Unwrap a required non-expression proto field.
739 ///
740 /// Mirrors [`PhysicalExprDecodeCtx::decode_required_expression`] for proto
741 /// fields that aren't [`PhysicalExprNode`]s — e.g. the `arrow_type` of a
742 /// `PhysicalCastNode` or the `scalar` of a `PhysicalLiteralNode`. Keeps
743 /// the "missing required field" message format identical across
744 /// expressions:
745 ///
746 /// ```ignore
747 /// let arrow_type = require_proto_field(
748 /// cast_expr.arrow_type.as_ref(),
749 /// "CastExpr",
750 /// "arrow_type",
751 /// )?;
752 /// ```
753 pub fn require_proto_field<T>(
754 opt: Option<T>,
755 expr_name: &str,
756 field: &str,
757 ) -> Result<T> {
758 opt.ok_or_else(|| {
759 datafusion_common::internal_datafusion_err!(
760 "{expr_name} is missing required field '{field}'"
761 )
762 })
763 }
764
765 /// Internal dispatch trait. Implementors live in `datafusion-proto`.
766 /// Expression authors should use [`PhysicalExprDecodeCtx`] instead of
767 /// calling this directly.
768 pub trait PhysicalExprDecode {
769 /// Decode a proto node into a concrete `PhysicalExpr`. The schema is
770 /// passed alongside so implementations can support recursive children
771 /// and rebind the context per call (e.g. for nested plans).
772 fn decode(
773 &self,
774 node: &PhysicalExprNode,
775 schema: &Schema,
776 ) -> Result<Arc<dyn PhysicalExpr>>;
777 }
778}
779
780#[deprecated(
781 since = "50.0.0",
782 note = "Use `datafusion_expr_common::dyn_eq` instead"
783)]
784pub use datafusion_expr_common::dyn_eq::{DynEq, DynHash};
785
786impl dyn PhysicalExpr {
787 /// Returns `true` if the expression is of type `T`.
788 ///
789 /// Prefer this over `downcast_ref::<T>().is_some()`. Works correctly when
790 /// called on `Arc<dyn PhysicalExpr>` via auto-deref.
791 pub fn is<T: PhysicalExpr>(&self) -> bool {
792 (self as &dyn Any).is::<T>()
793 }
794
795 /// Attempts to downcast this expression to a concrete type `T`, returning
796 /// `None` if the expression is not of that type.
797 ///
798 /// Works correctly when called on `Arc<dyn PhysicalExpr>` via auto-deref,
799 /// unlike `(&arc as &dyn Any).downcast_ref::<T>()` which would attempt to
800 /// downcast the `Arc` itself.
801 pub fn downcast_ref<T: PhysicalExpr>(&self) -> Option<&T> {
802 (self as &dyn Any).downcast_ref()
803 }
804}
805
806impl PartialEq for dyn PhysicalExpr {
807 fn eq(&self, other: &Self) -> bool {
808 self.dyn_eq(other as &dyn Any)
809 }
810}
811impl Eq for dyn PhysicalExpr {}
812
813impl Hash for dyn PhysicalExpr {
814 fn hash<H: Hasher>(&self, state: &mut H) {
815 self.dyn_hash(state);
816 }
817}
818
819/// Returns a copy of this expr if we change any child according to the pointer comparison.
820/// The size of `children` must be equal to the size of `PhysicalExpr::children()`.
821pub fn with_new_children_if_necessary(
822 expr: Arc<dyn PhysicalExpr>,
823 children: Vec<Arc<dyn PhysicalExpr>>,
824) -> Result<Arc<dyn PhysicalExpr>> {
825 let old_children = expr.children();
826 assert_eq_or_internal_err!(
827 children.len(),
828 old_children.len(),
829 "PhysicalExpr: Wrong number of children"
830 );
831
832 if children.is_empty()
833 || children
834 .iter()
835 .zip(old_children.iter())
836 .any(|(c1, c2)| !Arc::ptr_eq(c1, c2))
837 {
838 Ok(expr.with_new_children(children)?)
839 } else {
840 Ok(expr)
841 }
842}
843
844/// Returns [`Display`] able a list of [`PhysicalExpr`]
845///
846/// Example output: `[a + 1, b]`
847pub fn format_physical_expr_list<T>(exprs: T) -> impl Display
848where
849 T: IntoIterator,
850 T::Item: Display,
851 T::IntoIter: Clone,
852{
853 struct DisplayWrapper<I>(I)
854 where
855 I: Iterator + Clone,
856 I::Item: Display;
857
858 impl<I> Display for DisplayWrapper<I>
859 where
860 I: Iterator + Clone,
861 I::Item: Display,
862 {
863 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
864 let mut iter = self.0.clone();
865 write!(f, "[")?;
866 if let Some(expr) = iter.next() {
867 write!(f, "{expr}")?;
868 }
869 for expr in iter {
870 write!(f, ", {expr}")?;
871 }
872 write!(f, "]")?;
873 Ok(())
874 }
875 }
876
877 DisplayWrapper(exprs.into_iter())
878}
879
880/// Prints a [`PhysicalExpr`] in a SQL-like format
881///
882/// # Example
883/// ```
884/// # // The boilerplate needed to create a `PhysicalExpr` for the example
885/// use std::collections::HashMap;
886/// # use std::fmt::Formatter;
887/// # use std::sync::Arc;
888/// # use arrow::array::RecordBatch;
889/// # use arrow::datatypes::{DataType, Field, FieldRef, Schema};
890/// # use datafusion_common::Result;
891/// # use datafusion_expr_common::columnar_value::ColumnarValue;
892/// # use datafusion_physical_expr_common::physical_expr::{fmt_sql, DynEq, PhysicalExpr};
893/// # #[derive(Debug, PartialEq, Eq, Hash)]
894/// # struct MyExpr {}
895/// # impl PhysicalExpr for MyExpr {
896/// # fn data_type(&self, input_schema: &Schema) -> Result<DataType> { unimplemented!() }
897/// # fn nullable(&self, input_schema: &Schema) -> Result<bool> { unimplemented!() }
898/// # fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> { unimplemented!() }
899/// # fn return_field(&self, input_schema: &Schema) -> Result<FieldRef> { unimplemented!() }
900/// # fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>>{ unimplemented!() }
901/// # fn with_new_children(self: Arc<Self>, children: Vec<Arc<dyn PhysicalExpr>>) -> Result<Arc<dyn PhysicalExpr>> { unimplemented!() }
902/// # fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { write!(f, "CASE a > b THEN 1 ELSE 0 END") }
903/// # }
904/// # impl std::fmt::Display for MyExpr {fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { unimplemented!() } }
905/// # fn make_physical_expr() -> Arc<dyn PhysicalExpr> { Arc::new(MyExpr{}) }
906/// let expr: Arc<dyn PhysicalExpr> = make_physical_expr();
907/// // wrap the expression in `sql_fmt` which can be used with
908/// // `format!`, `to_string()`, etc
909/// let expr_as_sql = fmt_sql(expr.as_ref());
910/// assert_eq!(
911/// "The SQL: CASE a > b THEN 1 ELSE 0 END",
912/// format!("The SQL: {expr_as_sql}")
913/// );
914/// ```
915pub fn fmt_sql(expr: &dyn PhysicalExpr) -> impl Display + '_ {
916 struct Wrapper<'a> {
917 expr: &'a dyn PhysicalExpr,
918 }
919
920 impl Display for Wrapper<'_> {
921 fn fmt(&self, f: &mut Formatter) -> fmt::Result {
922 self.expr.fmt_sql(f)?;
923 Ok(())
924 }
925 }
926
927 Wrapper { expr }
928}
929
930/// Take a snapshot of the given `PhysicalExpr` if it is dynamic.
931///
932/// Take a snapshot of this `PhysicalExpr` if it is dynamic.
933/// This is used to capture the current state of `PhysicalExpr`s that may contain
934/// dynamic references to other operators in order to serialize it over the wire
935/// or treat it via downcast matching.
936///
937/// See the documentation of [`PhysicalExpr::snapshot`] for more details.
938///
939/// # Returns
940///
941/// Returns a snapshot of the `PhysicalExpr` if it is dynamic, otherwise
942/// returns itself.
943pub fn snapshot_physical_expr(
944 expr: Arc<dyn PhysicalExpr>,
945) -> Result<Arc<dyn PhysicalExpr>> {
946 snapshot_physical_expr_opt(expr).data()
947}
948
949/// Take a snapshot of the given `PhysicalExpr` if it is dynamic.
950///
951/// Take a snapshot of this `PhysicalExpr` if it is dynamic.
952/// This is used to capture the current state of `PhysicalExpr`s that may contain
953/// dynamic references to other operators in order to serialize it over the wire
954/// or treat it via downcast matching.
955///
956/// See the documentation of [`PhysicalExpr::snapshot`] for more details.
957///
958/// # Returns
959///
960/// Returns a `[`Transformed`] indicating whether a snapshot was taken,
961/// along with the resulting `PhysicalExpr`.
962pub fn snapshot_physical_expr_opt(
963 expr: Arc<dyn PhysicalExpr>,
964) -> Result<Transformed<Arc<dyn PhysicalExpr>>> {
965 expr.transform_up(|e| {
966 if let Some(snapshot) = e.snapshot()? {
967 Ok(Transformed::yes(snapshot))
968 } else {
969 Ok(Transformed::no(Arc::clone(&e)))
970 }
971 })
972}
973
974/// Check the generation of this `PhysicalExpr`.
975/// Dynamic `PhysicalExpr`s may have a generation that is incremented
976/// every time the state of the `PhysicalExpr` changes.
977/// If the generation changes that means this `PhysicalExpr` or one of its children
978/// has changed since the last time it was evaluated.
979///
980/// This algorithm will not produce collisions as long as the structure of the
981/// `PhysicalExpr` does not change and no `PhysicalExpr` decrements its own generation.
982pub fn snapshot_generation(expr: &Arc<dyn PhysicalExpr>) -> u64 {
983 let mut generation = 0u64;
984 expr.apply(|e| {
985 // Add the current generation of the `PhysicalExpr` to our global generation.
986 generation = generation.wrapping_add(e.snapshot_generation());
987 Ok(TreeNodeRecursion::Continue)
988 })
989 .expect("this traversal is infallible");
990
991 generation
992}
993
994/// Check if the given `PhysicalExpr` is dynamic.
995/// Internally this calls [`snapshot_generation`] to check if the generation is non-zero,
996/// any dynamic `PhysicalExpr` should have a non-zero generation.
997#[deprecated(
998 since = "55.0.0",
999 note = "Downcast to `DynamicFilterPhysicalExpr`, or use \
1000 `DynamicFilterTracking::classify(expr).contains_dynamic_filter()` from \
1001 `datafusion_physical_expr`"
1002)]
1003pub fn is_dynamic_physical_expr(expr: &Arc<dyn PhysicalExpr>) -> bool {
1004 // If the generation is non-zero, then this `PhysicalExpr` is dynamic.
1005 snapshot_generation(expr) != 0
1006}
1007
1008/// Returns true if the expression is volatile, i.e. whether it can return different
1009/// results when evaluated multiple times with the same input.
1010///
1011/// For example the function call `RANDOM()` is volatile as each call will
1012/// return a different value.
1013///
1014/// This method recursively checks if any sub-expression is volatile, for example
1015/// `1 + RANDOM()` will return `true`.
1016pub fn is_volatile(expr: &Arc<dyn PhysicalExpr>) -> bool {
1017 if expr.is_volatile_node() {
1018 return true;
1019 }
1020 let mut is_volatile = false;
1021 expr.apply(|e| {
1022 if e.is_volatile_node() {
1023 is_volatile = true;
1024 Ok(TreeNodeRecursion::Stop)
1025 } else {
1026 Ok(TreeNodeRecursion::Continue)
1027 }
1028 })
1029 .expect("infallible closure should not fail");
1030 is_volatile
1031}
1032
1033#[cfg(test)]
1034mod test {
1035 use crate::physical_expr::PhysicalExpr;
1036 use arrow::array::{Array, BooleanArray, Int64Array, RecordBatch};
1037 use arrow::datatypes::{DataType, Schema};
1038 use datafusion_expr_common::columnar_value::ColumnarValue;
1039 use std::fmt::{Display, Formatter};
1040 use std::sync::Arc;
1041
1042 #[derive(Debug, PartialEq, Eq, Hash)]
1043 struct TestExpr {}
1044
1045 impl PhysicalExpr for TestExpr {
1046 fn data_type(&self, _schema: &Schema) -> datafusion_common::Result<DataType> {
1047 Ok(DataType::Int64)
1048 }
1049
1050 fn nullable(&self, _schema: &Schema) -> datafusion_common::Result<bool> {
1051 Ok(false)
1052 }
1053
1054 fn evaluate(
1055 &self,
1056 batch: &RecordBatch,
1057 ) -> datafusion_common::Result<ColumnarValue> {
1058 let data = vec![1; batch.num_rows()];
1059 Ok(ColumnarValue::Array(Arc::new(Int64Array::from(data))))
1060 }
1061
1062 fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
1063 vec![]
1064 }
1065
1066 fn with_new_children(
1067 self: Arc<Self>,
1068 _children: Vec<Arc<dyn PhysicalExpr>>,
1069 ) -> datafusion_common::Result<Arc<dyn PhysicalExpr>> {
1070 Ok(Arc::new(Self {}))
1071 }
1072
1073 fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
1074 f.write_str("TestExpr")
1075 }
1076 }
1077
1078 impl Display for TestExpr {
1079 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
1080 self.fmt_sql(f)
1081 }
1082 }
1083
1084 macro_rules! assert_arrays_eq {
1085 ($EXPECTED: expr, $ACTUAL: expr, $MESSAGE: expr) => {
1086 let expected = $EXPECTED.to_array(1).unwrap();
1087 let actual = $ACTUAL;
1088
1089 let actual_array = actual.to_array(expected.len()).unwrap();
1090 let actual_ref = actual_array.as_ref();
1091 let expected_ref = expected.as_ref();
1092 assert!(
1093 actual_ref == expected_ref,
1094 "{}: expected: {:?}, actual: {:?}",
1095 $MESSAGE,
1096 $EXPECTED,
1097 actual_ref
1098 );
1099 };
1100 }
1101
1102 fn test_evaluate_selection(
1103 batch: &RecordBatch,
1104 selection: &BooleanArray,
1105 expected: &ColumnarValue,
1106 ) {
1107 let expr = TestExpr {};
1108
1109 // First check that the `evaluate_selection` is the expected one
1110 let selection_result = expr.evaluate_selection(batch, selection).unwrap();
1111 assert_eq!(
1112 expected.to_array(1).unwrap().len(),
1113 selection_result.to_array(1).unwrap().len(),
1114 "evaluate_selection should output row count should match input record batch"
1115 );
1116 assert_arrays_eq!(
1117 expected,
1118 &selection_result,
1119 "evaluate_selection returned unexpected value"
1120 );
1121
1122 // If we're selecting all rows, the result should be the same as calling `evaluate`
1123 // with the full record batch.
1124 if (0..batch.num_rows())
1125 .all(|row_idx| row_idx < selection.len() && selection.value(row_idx))
1126 {
1127 let empty_result = expr.evaluate(batch).unwrap();
1128
1129 assert_arrays_eq!(
1130 empty_result,
1131 &selection_result,
1132 "evaluate_selection does not match unfiltered evaluate result"
1133 );
1134 }
1135 }
1136
1137 fn test_evaluate_selection_error(batch: &RecordBatch, selection: &BooleanArray) {
1138 let expr = TestExpr {};
1139
1140 // First check that the `evaluate_selection` is the expected one
1141 let selection_result = expr.evaluate_selection(batch, selection);
1142 assert!(selection_result.is_err(), "evaluate_selection should fail");
1143 }
1144
1145 #[test]
1146 pub fn test_evaluate_selection_with_empty_record_batch() {
1147 test_evaluate_selection(
1148 &RecordBatch::new_empty(Arc::new(Schema::empty())),
1149 &BooleanArray::from(vec![false; 0]),
1150 &ColumnarValue::Array(Arc::new(Int64Array::new_null(0))),
1151 );
1152 }
1153
1154 #[test]
1155 pub fn test_evaluate_selection_with_empty_record_batch_with_larger_false_selection() {
1156 test_evaluate_selection_error(
1157 &RecordBatch::new_empty(Arc::new(Schema::empty())),
1158 &BooleanArray::from(vec![false; 10]),
1159 );
1160 }
1161
1162 #[test]
1163 pub fn test_evaluate_selection_with_empty_record_batch_with_larger_true_selection() {
1164 test_evaluate_selection_error(
1165 &RecordBatch::new_empty(Arc::new(Schema::empty())),
1166 &BooleanArray::from(vec![true; 10]),
1167 );
1168 }
1169
1170 #[test]
1171 pub fn test_evaluate_selection_with_non_empty_record_batch() {
1172 test_evaluate_selection(
1173 &unsafe { RecordBatch::new_unchecked(Arc::new(Schema::empty()), vec![], 10) },
1174 &BooleanArray::from(vec![true; 10]),
1175 &ColumnarValue::Array(Arc::new(Int64Array::from(vec![1; 10]))),
1176 );
1177 }
1178
1179 #[test]
1180 pub fn test_evaluate_selection_with_non_empty_record_batch_with_larger_false_selection()
1181 {
1182 test_evaluate_selection_error(
1183 &unsafe { RecordBatch::new_unchecked(Arc::new(Schema::empty()), vec![], 10) },
1184 &BooleanArray::from(vec![false; 20]),
1185 );
1186 }
1187
1188 #[test]
1189 pub fn test_evaluate_selection_with_non_empty_record_batch_with_larger_true_selection()
1190 {
1191 test_evaluate_selection_error(
1192 &unsafe { RecordBatch::new_unchecked(Arc::new(Schema::empty()), vec![], 10) },
1193 &BooleanArray::from(vec![true; 20]),
1194 );
1195 }
1196
1197 #[test]
1198 pub fn test_evaluate_selection_with_non_empty_record_batch_with_smaller_false_selection()
1199 {
1200 test_evaluate_selection_error(
1201 &unsafe { RecordBatch::new_unchecked(Arc::new(Schema::empty()), vec![], 10) },
1202 &BooleanArray::from(vec![false; 5]),
1203 );
1204 }
1205
1206 #[test]
1207 pub fn test_evaluate_selection_with_non_empty_record_batch_with_smaller_true_selection()
1208 {
1209 test_evaluate_selection_error(
1210 &unsafe { RecordBatch::new_unchecked(Arc::new(Schema::empty()), vec![], 10) },
1211 &BooleanArray::from(vec![true; 5]),
1212 );
1213 }
1214}
1215
1216#[cfg(all(test, feature = "proto"))]
1217mod proto_helper_tests {
1218 use datafusion_common::DataFusionError;
1219 use datafusion_proto_models::protobuf::{
1220 self, PhysicalColumn, PhysicalExprNode, physical_expr_node,
1221 };
1222
1223 use crate::expect_expr_variant;
1224 use crate::physical_expr::proto_decode::require_proto_field;
1225
1226 fn column_node() -> PhysicalExprNode {
1227 PhysicalExprNode {
1228 expr_id: None,
1229 expr_type: Some(physical_expr_node::ExprType::Column(PhysicalColumn {
1230 name: "a".to_string(),
1231 index: 0,
1232 })),
1233 }
1234 }
1235
1236 #[test]
1237 fn require_proto_field_returns_inner() {
1238 let v = require_proto_field(Some(7_u32), "FooExpr", "answer").unwrap();
1239 assert_eq!(v, 7);
1240 }
1241
1242 #[test]
1243 fn require_proto_field_reports_missing() {
1244 let err = require_proto_field::<u32>(None, "FooExpr", "answer").unwrap_err();
1245 assert!(matches!(
1246 err,
1247 DataFusionError::Internal(msg)
1248 if msg.contains("FooExpr is missing required field 'answer'")
1249 ));
1250 }
1251
1252 fn expect_column(
1253 node: &PhysicalExprNode,
1254 ) -> Result<&PhysicalColumn, DataFusionError> {
1255 let inner =
1256 expect_expr_variant!(node, physical_expr_node::ExprType::Column, "Column",);
1257 Ok(inner)
1258 }
1259
1260 #[test]
1261 fn expect_expr_variant_returns_inner_payload() {
1262 let node = column_node();
1263 let col = expect_column(&node).unwrap();
1264 assert_eq!(col.name, "a");
1265 }
1266
1267 #[test]
1268 fn expect_expr_variant_rejects_wrong_variant() {
1269 let node = PhysicalExprNode {
1270 expr_id: None,
1271 expr_type: Some(physical_expr_node::ExprType::Negative(Box::new(
1272 protobuf::PhysicalNegativeNode { expr: None },
1273 ))),
1274 };
1275 let err = expect_column(&node).unwrap_err();
1276 assert!(matches!(
1277 err,
1278 DataFusionError::Internal(msg) if msg.contains("PhysicalExprNode is not a Column")
1279 ));
1280 }
1281
1282 #[test]
1283 fn expect_expr_variant_rejects_missing_expr_type() {
1284 let node = PhysicalExprNode {
1285 expr_id: None,
1286 expr_type: None,
1287 };
1288 let err = expect_column(&node).unwrap_err();
1289 assert!(matches!(
1290 err,
1291 DataFusionError::Internal(msg) if msg.contains("PhysicalExprNode is not a Column")
1292 ));
1293 }
1294}