## Coverage Report

#### Created: 2025-11-22 06:00

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /Users/andrewlamb/Software/datafusion/datafusion/physical-expr-common/src/physical_expr.rs

    Line
    Count
    Source
    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
    18
    use std::any::Any;
    19
    use std::fmt;
    20
    use std::fmt::{Debug, Display, Formatter};
    21
    use std::hash::{Hash, Hasher};
    22
    use std::sync::Arc;
    23
    24
    use crate::utils::scatter;
    25
    26
    use arrow::array::{new_empty_array, ArrayRef, BooleanArray};
    27
    use arrow::compute::filter_record_batch;
    28
    use arrow::datatypes::{DataType, Field, FieldRef, Schema};
    29
    use arrow::record_batch::RecordBatch;
    30
    use datafusion_common::tree_node::{
    31
        Transformed, TransformedResult, TreeNode, TreeNodeRecursion,
    32
    };
    33
    use datafusion_common::{
    34
        assert_eq_or_internal_err, exec_err, not_impl_err, DataFusionError, Result,
    35
        ScalarValue,
    36
    };
    37
    use datafusion_expr_common::columnar_value::ColumnarValue;
    38
    use datafusion_expr_common::interval_arithmetic::Interval;
    39
    use datafusion_expr_common::sort_properties::ExprProperties;
    40
    use datafusion_expr_common::statistics::Distribution;
    41
    42
    use itertools::izip;
    43
    44
    /// Shared [`PhysicalExpr`].
    45
    pub type PhysicalExprRef = Arc<dyn PhysicalExpr>;
    46
    47
    /// [`PhysicalExpr`]s represent expressions such as `A + 1` or `CAST(c1 AS int)`.
    48
    ///
    49
    /// `PhysicalExpr` knows its type, nullability and can be evaluated directly on
    50
    /// a [`RecordBatch`] (see [`Self::evaluate`]).
    51
    ///
    52
    /// `PhysicalExpr` are the physical counterpart to [`Expr`] used in logical
    53
    /// planning. They are typically created from [`Expr`] by a [`PhysicalPlanner`]
    54
    /// invoked from a higher level API
    55
    ///
    56
    /// Some important examples of `PhysicalExpr` are:
    57
    /// * [`Column`]: Represents a column at a given index in a RecordBatch
    58
    ///
    59
    /// To create `PhysicalExpr` from  `Expr`, see
    60
    /// * [`SessionContext::create_physical_expr`]: A high level API
    61
    /// * [`create_physical_expr`]: A low level API
    62
    ///
    63
    /// # Formatting `PhysicalExpr` as strings
    64
    /// There are three ways to format `PhysicalExpr` as a string:
    65
    /// * [`Debug`]: Standard Rust debugging format (e.g. `Constant { value: ... }`)
    66
    /// * [`Display`]: Detailed SQL-like format that shows expression structure (e.g. (`Utf8 ("foobar")`). This is often used for debugging and tests
    67
    /// * [`Self::fmt_sql`]: SQL-like human readable format (e.g. ('foobar')`), See also [`sql_fmt`]
    68
    ///
    69
    /// [`SessionContext::create_physical_expr`]: https://docs.rs/datafusion/latest/datafusion/execution/context/struct.SessionContext.html#method.create_physical_expr
    70
    /// [`PhysicalPlanner`]: https://docs.rs/datafusion/latest/datafusion/physical_planner/trait.PhysicalPlanner.html
    71
    /// [`Expr`]: https://docs.rs/datafusion/latest/datafusion/logical_expr/enum.Expr.html
    72
    /// [`create_physical_expr`]: https://docs.rs/datafusion/latest/datafusion/physical_expr/fn.create_physical_expr.html
    73
    /// [`Column`]: https://docs.rs/datafusion/latest/datafusion/physical_expr/expressions/struct.Column.html
    74
    pub trait PhysicalExpr: Any + Send + Sync + Display + Debug + DynEq + DynHash {
    75
        /// Returns the physical expression as [`Any`] so that it can be
    76
        /// downcast to a specific implementation.
    77
        fn as_any(&self) -> &dyn Any;
    78
        /// Get the data type of this expression, given the schema of the input
    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
    0
        fn return_field(&self, input_schema: &Schema) -> Result<FieldRef> {
    90
    0
            Ok(Arc::new(Field::new(
    91
    0
                format!("{self}"),
    92
    0
                self.data_type(input_schema)?,
    93
    0
                self.nullable(input_schema)?,
    94
            )))
    95
    0
        }
    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
    0
        fn evaluate_selection(
    103
    0
            &self,
    104
    0
            batch: &RecordBatch,
    105
    0
            selection: &BooleanArray,
    106
    0
        ) -> Result<ColumnarValue> {
    107
    0
            let row_count = batch.num_rows();
    108
    0
            if row_count != selection.len() {
    109
    0
                return exec_err!("Selection array length does not match batch row count: {} != {row_count}", selection.len());
    110
    0
            }
    111
    112
    0
            let selection_count = selection.true_count();
    113
    114
            // First, check if we can avoid filtering altogether.
    115
    0
            if selection_count == row_count {
    116
                // All values from the `selection` filter are true and match the input batch.
    117
                // No need to perform any filtering.
    118
    0
                return self.evaluate(batch);
    119
    0
            }
    120
    121
            // Next, prepare the result array for each 'true' row in the selection vector.
    122
    0
            let filtered_result = if selection_count == 0 {
    123
                // Do not call `evaluate` when the selection is empty.
    124
                // `evaluate_selection` is used to conditionally evaluate expressions.
    125
                // When the expression in question is fallible, evaluating it with an empty
    126
                // record batch may trigger a runtime error (e.g. division by zero).
    127
                //
    128
                // Instead, create an empty array matching the expected return type.
    129
    0
                let datatype = self.data_type(batch.schema_ref().as_ref())?;
    130
    0
                ColumnarValue::Array(new_empty_array(&datatype))
    131
            } else {
    132
                // If we reach this point, there's no other option than to filter the batch.
    133
                // This is a fairly costly operation since it requires creating partial copies
    134
                // (worst case of length `row_count - 1`) of all the arrays in the record batch.
    135
                // The resulting `filtered_batch` will contain `selection_count` rows.
    136
    0
                let filtered_batch = filter_record_batch(batch, selection)?;
    137
    0
                self.evaluate(&filtered_batch)?
    138
            };
    139
    140
            // Finally, scatter the filtered result array so that the indices match the input rows again.
    141
    0
            match &filtered_result {
    142
    0
                ColumnarValue::Array(a) => {
    143
    0
                    scatter(selection, a.as_ref()).map(ColumnarValue::Array)
    144
                }
    145
    0
                ColumnarValue::Scalar(ScalarValue::Boolean(value)) => {
    146
                    // When the scalar is true or false, skip the scatter process
    147
    0
                    if let Some(v) = value {
    148
    0
                        if *v {
    149
    0
                            Ok(ColumnarValue::from(Arc::new(selection.clone()) as ArrayRef))
    150
                        } else {
    151
    0
                            Ok(filtered_result)
    152
                        }
    153
                    } else {
    154
    0
                        let array = BooleanArray::from(vec![None; row_count]);
    155
    0
                        scatter(selection, &array).map(ColumnarValue::Array)
    156
                    }
    157
                }
    158
    0
                ColumnarValue::Scalar(_) => Ok(filtered_result),
    159
            }
    160
    0
        }
    161
    162
        /// Get a list of child PhysicalExpr that provide the input for this expr.
    163
        fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>>;
    164
    165
        /// Returns a new PhysicalExpr where all children were replaced by new exprs.
    166
        fn with_new_children(
    167
            self: Arc<Self>,
    168
            children: Vec<Arc<dyn PhysicalExpr>>,
    169
        ) -> Result<Arc<dyn PhysicalExpr>>;
    170
    171
        /// Computes the output interval for the expression, given the input
    172
        /// intervals.
    173
        ///
    174
        /// # Parameters
    175
        ///
    176
        /// * `children` are the intervals for the children (inputs) of this
    177
        ///   expression.
    178
        ///
    179
        /// # Returns
    180
        ///
    181
        /// A `Result` containing the output interval for the expression in
    182
        /// case of success, or an error object in case of failure.
    183
        ///
    184
        /// # Example
    185
        ///
    186
        /// If the expression is `a + b`, and the input intervals are `a: [1, 2]`
    187
        /// and `b: [3, 4]`, then the output interval would be `[4, 6]`.
    188
    0
        fn evaluate_bounds(&self, _children: &[&Interval]) -> Result<Interval> {
    189
    0
            not_impl_err!("Not implemented for {self}")
    190
    0
        }
    191
    192
        /// Updates bounds for child expressions, given a known interval for this
    193
        /// expression.
    194
        ///
    195
        /// This is used to propagate constraints down through an expression tree.
    196
        ///
    197
        /// # Parameters
    198
        ///
    199
        /// * `interval` is the currently known interval for this expression.
    200
        /// * `children` are the current intervals for the children of this expression.
    201
        ///
    202
        /// # Returns
    203
        ///
    204
        /// A `Result` containing a `Vec` of new intervals for the children (in order)
    205
        /// in case of success, or an error object in case of failure.
    206
        ///
    207
        /// If constraint propagation reveals an infeasibility for any child, returns
    208
        /// [`None`]. If none of the children intervals change as a result of
    209
        /// propagation, may return an empty vector instead of cloning `children`.
    210
        /// This is the default (and conservative) return value.
    211
        ///
    212
        /// # Example
    213
        ///
    214
        /// If the expression is `a + b`, the current `interval` is `[4, 5]` and the
    215
        /// inputs `a` and `b` are respectively given as `[0, 2]` and `[-∞, 4]`, then
    216
        /// propagation would return `[0, 2]` and `[2, 4]` as `b` must be at least
    217
        /// `2` to make the output at least `4`.
    218
    0
        fn propagate_constraints(
    219
    0
            &self,
    220
    0
            _interval: &Interval,
    221
    0
            _children: &[&Interval],
    222
    0
        ) -> Result<Option<Vec<Interval>>> {
    223
    0
            Ok(Some(vec![]))
    224
    0
        }
    225
    226
        /// Computes the output statistics for the expression, given the input
    227
        /// statistics.
    228
        ///
    229
        /// # Parameters
    230
        ///
    231
        /// * `children` are the statistics for the children (inputs) of this
    232
        ///   expression.
    233
        ///
    234
        /// # Returns
    235
        ///
    236
        /// A `Result` containing the output statistics for the expression in
    237
        /// case of success, or an error object in case of failure.
    238
        ///
    239
        /// Expressions (should) implement this function and utilize the independence
    240
        /// assumption, match on children distribution types and compute the output
    241
        /// statistics accordingly. The default implementation simply creates an
    242
        /// unknown output distribution by combining input ranges. This logic loses
    243
        /// distribution information, but is a safe default.
    244
    0
        fn evaluate_statistics(&self, children: &[&Distribution]) -> Result<Distribution> {
    245
    0
            let children_ranges = children
    246
    0
                .iter()
    247
    0
                .map(|c| c.range())
    248
    0
                .collect::<Result<Vec<_>>>()?;
    249
    0
            let children_ranges_refs = children_ranges.iter().collect::<Vec<_>>();
    250
    0
            let output_interval = self.evaluate_bounds(children_ranges_refs.as_slice())?;
    251
    0
            let dt = output_interval.data_type();
    252
    0
            if dt.eq(&DataType::Boolean) {
    253
    0
                let p = if output_interval.eq(&Interval::TRUE) {
    254
    0
                    ScalarValue::new_one(&dt)
    255
    0
                } else if output_interval.eq(&Interval::FALSE) {
    256
    0
                    ScalarValue::new_zero(&dt)
    257
                } else {
    258
    0
                    ScalarValue::try_from(&dt)
    259
    0
                }?;
    260
    0
                Distribution::new_bernoulli(p)
    261
            } else {
    262
    0
                Distribution::new_from_interval(output_interval)
    263
            }
    264
    0
        }
    265
    266
        /// Updates children statistics using the given parent statistic for this
    267
        /// expression.
    268
        ///
    269
        /// This is used to propagate statistics down through an expression tree.
    270
        ///
    271
        /// # Parameters
    272
        ///
    273
        /// * `parent` is the currently known statistics for this expression.
    274
        /// * `children` are the current statistics for the children of this expression.
    275
        ///
    276
        /// # Returns
    277
        ///
    278
        /// A `Result` containing a `Vec` of new statistics for the children (in order)
    279
        /// in case of success, or an error object in case of failure.
    280
        ///
    281
        /// If statistics propagation reveals an infeasibility for any child, returns
    282
        /// [`None`]. If none of the children statistics change as a result of
    283
        /// propagation, may return an empty vector instead of cloning `children`.
    284
        /// This is the default (and conservative) return value.
    285
        ///
    286
        /// Expressions (should) implement this function and apply Bayes rule to
    287
        /// reconcile and update parent/children statistics. This involves utilizing
    288
        /// the independence assumption, and matching on distribution types. The
    289
        /// default implementation simply creates an unknown distribution if it can
    290
        /// narrow the range by propagating ranges. This logic loses distribution
    291
        /// information, but is a safe default.
    292
    0
        fn propagate_statistics(
    293
    0
            &self,
    294
    0
            parent: &Distribution,
    295
    0
            children: &[&Distribution],
    296
    0
        ) -> Result<Option<Vec<Distribution>>> {
    297
    0
            let children_ranges = children
    298
    0
                .iter()
    299
    0
                .map(|c| c.range())
    300
    0
                .collect::<Result<Vec<_>>>()?;
    301
    0
            let children_ranges_refs = children_ranges.iter().collect::<Vec<_>>();
    302
    0
            let parent_range = parent.range()?;
    303
    0
            let Some(propagated_children) =
    304
    0
                self.propagate_constraints(&parent_range, children_ranges_refs.as_slice())?
    305
            else {
    306
    0
                return Ok(None);
    307
            };
    308
    0
            izip!(propagated_children.into_iter(), children_ranges, children)
    309
    0
                .map(|(new_interval, old_interval, child)| {
    310
    0
                    if new_interval == old_interval {
    311
                        // We weren't able to narrow the range, preserve the old statistics.
    312
    0
                        Ok((*child).clone())
    313
    0
                    } else if new_interval.data_type().eq(&DataType::Boolean) {
    314
    0
                        let dt = old_interval.data_type();
    315
    0
                        let p = if new_interval.eq(&Interval::TRUE) {
    316
    0
                            ScalarValue::new_one(&dt)
    317
    0
                        } else if new_interval.eq(&Interval::FALSE) {
    318
    0
                            ScalarValue::new_zero(&dt)
    319
                        } else {
    320
    0
                            unreachable!("Given that we have a range reduction for a boolean interval, we should have certainty")
    321
    0
                        }?;
    322
    0
                        Distribution::new_bernoulli(p)
    323
                    } else {
    324
    0
                        Distribution::new_from_interval(new_interval)
    325
                    }
    326
    0
                })
    327
    0
                .collect::<Result<_>>()
    328
    0
                .map(Some)
    329
    0
        }
    330
    331
        /// Calculates the properties of this [`PhysicalExpr`] based on its
    332
        /// children's properties (i.e. order and range), recursively aggregating
    333
        /// the information from its children. In cases where the [`PhysicalExpr`]
    334
        /// has no children (e.g., `Literal` or `Column`), these properties should
    335
        /// be specified externally, as the function defaults to unknown properties.
    336
    0
        fn get_properties(&self, _children: &[ExprProperties]) -> Result<ExprProperties> {
    337
    0
            Ok(ExprProperties::new_unknown())
    338
    0
        }
    339
    340
        /// Format this `PhysicalExpr` in nice human readable "SQL" format
    341
        ///
    342
        /// Specifically, this format is designed to be readable by humans, at the
    343
        /// expense of details. Use `Display` or `Debug` for more detailed
    344
        /// representation.
    345
        ///
    346
        /// See the [`fmt_sql`] function for an example of printing `PhysicalExpr`s as SQL.
    347
        fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result;
    348
    349
        /// Take a snapshot of this `PhysicalExpr`, if it is dynamic.
    350
        ///
    351
        /// "Dynamic" in this case means containing references to structures that may change
    352
        /// during plan execution, such as hash tables.
    353
        ///
    354
        /// This method is used to capture the current state of `PhysicalExpr`s that may contain
    355
        /// dynamic references to other operators in order to serialize it over the wire
    356
        /// or treat it via downcast matching.
    357
        ///
    358
        /// You should not call this method directly as it does not handle recursion.
    359
        /// Instead use [`snapshot_physical_expr`] to handle recursion and capture the
    360
        /// full state of the `PhysicalExpr`.
    361
        ///
    362
        /// This is expected to return "simple" expressions that do not have mutable state
    363
        /// and are composed of DataFusion's built-in `PhysicalExpr` implementations.
    364
        /// Callers however should *not* assume anything about the returned expressions
    365
        /// since callers and implementers may not agree on what "simple" or "built-in"
    366
        /// means.
    367
        /// In other words, if you need to serialize a `PhysicalExpr` across the wire
    368
        /// you should call this method and then try to serialize the result,
    369
        /// but you should handle unknown or unexpected `PhysicalExpr` implementations gracefully
    370
        /// just as if you had not called this method at all.
    371
        ///
    372
        /// In particular, consider:
    373
        /// * A `PhysicalExpr` that references the current state of a `datafusion::physical_plan::TopK`
    374
        ///   that is involved in a query with `SELECT * FROM t1 ORDER BY a LIMIT 10`.
    375
        ///   This function may return something like `a >= 12`.
    376
        /// * A `PhysicalExpr` that references the current state of a `datafusion::physical_plan::joins::HashJoinExec`
    377
        ///   from a query such as `SELECT * FROM t1 JOIN t2 ON t1.a = t2.b`.
    378
        ///   This function may return something like `t2.b IN (1, 5, 7)`.
    379
        ///
    380
        /// A system or function that can only deal with a hardcoded set of `PhysicalExpr` implementations
    381
        /// or needs to serialize this state to bytes may not be able to handle these dynamic references.
    382
        /// In such cases, we should return a simplified version of the `PhysicalExpr` that does not
    383
        /// contain these dynamic references.
    384
        ///
    385
        /// Systems that implement remote execution of plans, e.g. serialize a portion of the query plan
    386
        /// and send it across the wire to a remote executor may want to call this method after
    387
        /// every batch on the source side and broadcast / update the current snapshot to the remote executor.
    388
        ///
    389
        /// Note for implementers: this method should *not* handle recursion.
    390
        /// Recursion is handled in [`snapshot_physical_expr`].
    391
    0
        fn snapshot(&self) -> Result<Option<Arc<dyn PhysicalExpr>>> {
    392
            // By default, we return None to indicate that this PhysicalExpr does not
    393
            // have any dynamic references or state.
    394
            // This is a safe default behavior.
    395
    0
            Ok(None)
    396
    0
        }
    397
    398
        /// Returns the generation of this `PhysicalExpr` for snapshotting purposes.
    399
        /// The generation is an arbitrary u64 that can be used to track changes
    400
        /// in the state of the `PhysicalExpr` over time without having to do an exhaustive comparison.
    401
        /// This is useful to avoid unnecessary computation or serialization if there are no changes to the expression.
    402
        /// In particular, dynamic expressions that may change over time; this allows cheap checks for changes.
    403
        /// Static expressions that do not change over time should return 0, as does the default implementation.
    404
        /// You should not call this method directly as it does not handle recursion.
    405
        /// Instead use [`snapshot_generation`] to handle recursion and capture the
    406
        /// full state of the `PhysicalExpr`.
    407
    0
        fn snapshot_generation(&self) -> u64 {
    408
            // By default, we return 0 to indicate that this PhysicalExpr does not
    409
            // have any dynamic references or state.
    410
            // Since the recursive algorithm XORs the generations of all children the overall
    411
            // generation will be 0 if no children have a non-zero generation, meaning that
    412
            // static expressions will always return 0.
    413
    0
            0
    414
    0
        }
    415
    416
        /// Returns true if the expression node is volatile, i.e. whether it can return
    417
        /// different results when evaluated multiple times with the same input.
    418
        ///
    419
        /// Note: unlike [`is_volatile`], this function does not consider inputs:
    420
        /// - `random()` returns `true`,
    421
        /// - `a + random()` returns `false` (because the operation `+` itself is not volatile.)
    422
        ///
    423
        /// The default to this function was set to `false` when it was created
    424
        /// to avoid imposing API churn on implementers, but this is not a safe default in general.
    425
        /// It is highly recommended that volatile expressions implement this method and return `true`.
    426
        /// This default may be removed in the future if it causes problems or we decide to
    427
        /// eat the cost of the breaking change and require all implementers to make a choice.
    428
    0
        fn is_volatile_node(&self) -> bool {
    429
    0
            false
    430
    0
        }
    431
    }
    432
    433
    #[deprecated(
    434
        since = "50.0.0",
    435
        note = "Use `datafusion_expr_common::dyn_eq` instead"
    436
    )]
    437
    pub use datafusion_expr_common::dyn_eq::{DynEq, DynHash};
    438
    439
    impl PartialEq for dyn PhysicalExpr {
    440
    3
        fn eq(&self, other: &Self) -> bool {
    441
    3
            self.dyn_eq(other.as_any())
    442
    3
        }
    443
    }
    444
    impl Eq for dyn PhysicalExpr {}
    445
    446
    impl Hash for dyn PhysicalExpr {
    447
    3
        fn hash<H: Hasher>(&self, state: &mut H) {
    448
    3
            self.dyn_hash(state);
    449
    3
        }
    450
    }
    451
    452
    /// Returns a copy of this expr if we change any child according to the pointer comparison.
    453
    /// The size of `children` must be equal to the size of `PhysicalExpr::children()`.
    454
    1
    pub fn with_new_children_if_necessary(
    455
    1
        expr: Arc<dyn PhysicalExpr>,
    456
    1
        children: Vec<Arc<dyn PhysicalExpr>>,
    457
    1
    ) -> Result<Arc<dyn PhysicalExpr>> {
    458
    1
        let old_children = expr.children();
    459
    1
        assert_eq_or_internal_err!(
    460
    1
            children.len(),
    461
    1
            old_children.len(),
    462
    0
            "PhysicalExpr: Wrong number of children"
    463
        );
    464
    465
    1
        if children.is_empty()
    466
    1
            || children
    467
    1
                .iter()
    468
    1
                .zip(old_children.iter())
    469
    1
                .any(|(c1, c2)| !Arc::ptr_eq(c1, c2))
    470
        {
    471
    1
            Ok(expr.with_new_children(children)?0)
    472
        } else {
    473
    0
            Ok(expr)
    474
        }
    475
    1
    }
    476
    477
    /// Returns [`Display`] able a list of [`PhysicalExpr`]
    478
    ///
    479
    /// Example output: `[a + 1, b]`
    480
    0
    pub fn format_physical_expr_list<T>(exprs: T) -> impl Display
    481
    0
    where
    482
    0
        T: IntoIterator,
    483
    0
        T::Item: Display,
    484
    0
        T::IntoIter: Clone,
    485
    {
    486
        struct DisplayWrapper<I>(I)
    487
        where
    488
            I: Iterator + Clone,
    489
            I::Item: Display;
    490
    491
        impl<I> Display for DisplayWrapper<I>
    492
        where
    493
            I: Iterator + Clone,
    494
            I::Item: Display,
    495
        {
    496
    0
            fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
    497
    0
                let mut iter = self.0.clone();
    498
    0
                write!(f, "[")?;
    499
    0
                if let Some(expr) = iter.next() {
    500
    0
                    write!(f, "{expr}")?;
    501
    0
                }
    502
    0
                for expr in iter {
    503
    0
                    write!(f, ", {expr}")?;
    504
                }
    505
    0
                write!(f, "]")?;
    506
    0
                Ok(())
    507
    0
            }
    508
        }
    509
    510
    0
        DisplayWrapper(exprs.into_iter())
    511
    0
    }
    512
    513
    /// Prints a [`PhysicalExpr`] in a SQL-like format
    514
    ///
    515
    /// # Example
    516
    /// ```
    517
    /// # // The boilerplate needed to create a `PhysicalExpr` for the example
    518
    /// # use std::any::Any;
    519
    /// use std::collections::HashMap;
    520
    /// # use std::fmt::Formatter;
    521
    /// # use std::sync::Arc;
    522
    /// # use arrow::array::RecordBatch;
    523
    /// # use arrow::datatypes::{DataType, Field, FieldRef, Schema};
    524
    /// # use datafusion_common::Result;
    525
    /// # use datafusion_expr_common::columnar_value::ColumnarValue;
    526
    /// # use datafusion_physical_expr_common::physical_expr::{fmt_sql, DynEq, PhysicalExpr};
    527
    /// # #[derive(Debug, PartialEq, Eq, Hash)]
    528
    /// # struct MyExpr {}
    529
    /// # impl PhysicalExpr for MyExpr {fn as_any(&self) -> &dyn Any { unimplemented!() }
    530
    /// # fn data_type(&self, input_schema: &Schema) -> Result<DataType> { unimplemented!() }
    531
    /// # fn nullable(&self, input_schema: &Schema) -> Result<bool> { unimplemented!() }
    532
    /// # fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> { unimplemented!() }
    533
    /// # fn return_field(&self, input_schema: &Schema) -> Result<FieldRef> { unimplemented!() }
    534
    /// # fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>>{ unimplemented!() }
    535
    /// # fn with_new_children(self: Arc<Self>, children: Vec<Arc<dyn PhysicalExpr>>) -> Result<Arc<dyn PhysicalExpr>> { unimplemented!() }
    536
    /// # fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { write!(f, "CASE a > b THEN 1 ELSE 0 END") }
    537
    /// # }
    538
    /// # impl std::fmt::Display for MyExpr {fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { unimplemented!() } }
    539
    /// # fn make_physical_expr() -> Arc<dyn PhysicalExpr> { Arc::new(MyExpr{}) }
    540
    /// let expr: Arc<dyn PhysicalExpr> = make_physical_expr();
    541
    /// // wrap the expression in `sql_fmt` which can be used with
    542
    /// // `format!`, `to_string()`, etc
    543
    /// let expr_as_sql = fmt_sql(expr.as_ref());
    544
    /// assert_eq!(
    545
    ///   "The SQL: CASE a > b THEN 1 ELSE 0 END",
    546
    ///   format!("The SQL: {expr_as_sql}")
    547
    /// );
    548
    /// ```
    549
    4
    pub fn fmt_sql(expr: &dyn PhysicalExpr) -> impl Display + '_ {
    550
        struct Wrapper<'a> {
    551
            expr: &'a dyn PhysicalExpr,
    552
        }
    553
    554
        impl Display for Wrapper<'_> {
    555
    4
            fn fmt(&self, f: &mut Formatter) -> fmt::Result {
    556
    4
                self.expr.fmt_sql(f)?0;
    557
    4
                Ok(())
    558
    4
            }
    559
        }
    560
    561
    4
        Wrapper { expr }
    562
    4
    }
    563
    564
    /// Take a snapshot of the given `PhysicalExpr` if it is dynamic.
    565
    ///
    566
    /// Take a snapshot of this `PhysicalExpr` if it is dynamic.
    567
    /// This is used to capture the current state of `PhysicalExpr`s that may contain
    568
    /// dynamic references to other operators in order to serialize it over the wire
    569
    /// or treat it via downcast matching.
    570
    ///
    571
    /// See the documentation of [`PhysicalExpr::snapshot`] for more details.
    572
    ///
    573
    /// # Returns
    574
    ///
    575
    /// Returns a snapshot of the `PhysicalExpr` if it is dynamic, otherwise
    576
    /// returns itself.
    577
    0
    pub fn snapshot_physical_expr(
    578
    0
        expr: Arc<dyn PhysicalExpr>,
    579
    0
    ) -> Result<Arc<dyn PhysicalExpr>> {
    580
    0
        expr.transform_up(|e| {
    581
    0
            if let Some(snapshot) = e.snapshot()? {
    582
    0
                Ok(Transformed::yes(snapshot))
    583
            } else {
    584
    0
                Ok(Transformed::no(Arc::clone(&e)))
    585
            }
    586
    0
        })
    587
    0
        .data()
    588
    0
    }
    589
    590
    /// Check the generation of this `PhysicalExpr`.
    591
    /// Dynamic `PhysicalExpr`s may have a generation that is incremented
    592
    /// every time the state of the `PhysicalExpr` changes.
    593
    /// If the generation changes that means this `PhysicalExpr` or one of its children
    594
    /// has changed since the last time it was evaluated.
    595
    ///
    596
    /// This algorithm will not produce collisions as long as the structure of the
    597
    /// `PhysicalExpr` does not change and no `PhysicalExpr` decrements its own generation.
    598
    0
    pub fn snapshot_generation(expr: &Arc<dyn PhysicalExpr>) -> u64 {
    599
    0
        let mut generation = 0u64;
    600
    0
        expr.apply(|e| {
    601
            // Add the current generation of the `PhysicalExpr` to our global generation.
    602
    0
            generation = generation.wrapping_add(e.snapshot_generation());
    603
    0
            Ok(TreeNodeRecursion::Continue)
    604
    0
        })
    605
    0
        .expect("this traversal is infallible");
    606
    607
    0
        generation
    608
    0
    }
    609
    610
    /// Check if the given `PhysicalExpr` is dynamic.
    611
    /// Internally this calls [`snapshot_generation`] to check if the generation is non-zero,
    612
    /// any dynamic `PhysicalExpr` should have a non-zero generation.
    613
    0
    pub fn is_dynamic_physical_expr(expr: &Arc<dyn PhysicalExpr>) -> bool {
    614
        // If the generation is non-zero, then this `PhysicalExpr` is dynamic.
    615
    0
        snapshot_generation(expr) != 0
    616
    0
    }
    617
    618
    /// Returns true if the expression is volatile, i.e. whether it can return different
    619
    /// results when evaluated multiple times with the same input.
    620
    ///
    621
    /// For example the function call `RANDOM()` is volatile as each call will
    622
    /// return a different value.
    623
    ///
    624
    /// This method recursively checks if any sub-expression is volatile, for example
    625
    /// `1 + RANDOM()` will return `true`.
    626
    0
    pub fn is_volatile(expr: &Arc<dyn PhysicalExpr>) -> bool {
    627
    0
        if expr.is_volatile_node() {
    628
    0
            return true;
    629
    0
        }
    630
    0
        let mut is_volatile = false;
    631
    0
        expr.apply(|e| {
    632
    0
            if e.is_volatile_node() {
    633
    0
                is_volatile = true;
    634
    0
                Ok(TreeNodeRecursion::Stop)
    635
            } else {
    636
    0
                Ok(TreeNodeRecursion::Continue)
    637
            }
    638
    0
        })
    639
    0
        .expect("infallible closure should not fail");
    640
    0
        is_volatile
    641
    0
    }
    642
    643
    #[cfg(test)]
    644
    mod test {
    645
        use crate::physical_expr::PhysicalExpr;
    646
        use arrow::array::{Array, BooleanArray, Int64Array, RecordBatch};
    647
        use arrow::datatypes::{DataType, Schema};
    648
        use datafusion_expr_common::columnar_value::ColumnarValue;
    649
        use std::fmt::{Display, Formatter};
    650
        use std::sync::Arc;
    651
    652
        #[derive(Debug, PartialEq, Eq, Hash)]
    653
        struct TestExpr {}
    654
    655
        impl PhysicalExpr for TestExpr {
    656
            fn as_any(&self) -> &dyn std::any::Any {
    657
                self
    658
            }
    659
    660
            fn data_type(&self, _schema: &Schema) -> datafusion_common::Result<DataType> {
    661
                Ok(DataType::Int64)
    662
            }
    663
    664
            fn nullable(&self, _schema: &Schema) -> datafusion_common::Result<bool> {
    665
                Ok(false)
    666
            }
    667
    668
            fn evaluate(
    669
                &self,
    670
                batch: &RecordBatch,
    671
            ) -> datafusion_common::Result<ColumnarValue> {
    672
                let data = vec![1; batch.num_rows()];
    673
                Ok(ColumnarValue::Array(Arc::new(Int64Array::from(data))))
    674
            }
    675
    676
            fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
    677
                vec![]
    678
            }
    679
    680
            fn with_new_children(
    681
                self: Arc<Self>,
    682
                _children: Vec<Arc<dyn PhysicalExpr>>,
    683
            ) -> datafusion_common::Result<Arc<dyn PhysicalExpr>> {
    684
                Ok(Arc::new(Self {}))
    685
            }
    686
    687
            fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
    688
                f.write_str("TestExpr")
    689
            }
    690
        }
    691
    692
        impl Display for TestExpr {
    693
            fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
    694
                self.fmt_sql(f)
    695
            }
    696
        }
    697
    698
        macro_rules! assert_arrays_eq {
    699
            ($EXPECTED: expr, $ACTUAL: expr, $MESSAGE: expr) => {
    700
                let expected = $EXPECTED.to_array(1).unwrap();
    701
                let actual = $ACTUAL;
    702
    703
                let actual_array = actual.to_array(expected.len()).unwrap();
    704
                let actual_ref = actual_array.as_ref();
    705
                let expected_ref = expected.as_ref();
    706
                assert!(
    707
                    actual_ref == expected_ref,
    708
                    "{}: expected: {:?}, actual: {:?}",
    709
                    $MESSAGE,
    710
                    $EXPECTED,
    711
                    actual_ref
    712
                );
    713
            };
    714
        }
    715
    716
        fn test_evaluate_selection(
    717
            batch: &RecordBatch,
    718
            selection: &BooleanArray,
    719
            expected: &ColumnarValue,
    720
        ) {
    721
            let expr = TestExpr {};
    722
    723
            // First check that the `evaluate_selection` is the expected one
    724
            let selection_result = expr.evaluate_selection(batch, selection).unwrap();
    725
            assert_eq!(
    726
                expected.to_array(1).unwrap().len(),
    727
                selection_result.to_array(1).unwrap().len(),
    728
                "evaluate_selection should output row count should match input record batch"
    729
            );
    730
            assert_arrays_eq!(
    731
                expected,
    732
                &selection_result,
    733
                "evaluate_selection returned unexpected value"
    734
            );
    735
    736
            // If we're selecting all rows, the result should be the same as calling `evaluate`
    737
            // with the full record batch.
    738
            if (0..batch.num_rows())
    739
                .all(|row_idx| row_idx < selection.len() && selection.value(row_idx))
    740
            {
    741
                let empty_result = expr.evaluate(batch).unwrap();
    742
    743
                assert_arrays_eq!(
    744
                    empty_result,
    745
                    &selection_result,
    746
                    "evaluate_selection does not match unfiltered evaluate result"
    747
                );
    748
            }
    749
        }
    750
    751
        fn test_evaluate_selection_error(batch: &RecordBatch, selection: &BooleanArray) {
    752
            let expr = TestExpr {};
    753
    754
            // First check that the `evaluate_selection` is the expected one
    755
            let selection_result = expr.evaluate_selection(batch, selection);
    756
            assert!(selection_result.is_err(), "evaluate_selection should fail");
    757
        }
    758
    759
        #[test]
    760
        pub fn test_evaluate_selection_with_empty_record_batch() {
    761
            test_evaluate_selection(
    762
                &RecordBatch::new_empty(Arc::new(Schema::empty())),
    763
                &BooleanArray::from(vec![false; 0]),
    764
                &ColumnarValue::Array(Arc::new(Int64Array::new_null(0))),
    765
            );
    766
        }
    767
    768
        #[test]
    769
        pub fn test_evaluate_selection_with_empty_record_batch_with_larger_false_selection() {
    770
            test_evaluate_selection_error(
    771
                &RecordBatch::new_empty(Arc::new(Schema::empty())),
    772
                &BooleanArray::from(vec![false; 10]),
    773
            );
    774
        }
    775
    776
        #[test]
    777
        pub fn test_evaluate_selection_with_empty_record_batch_with_larger_true_selection() {
    778
            test_evaluate_selection_error(
    779
                &RecordBatch::new_empty(Arc::new(Schema::empty())),
    780
                &BooleanArray::from(vec![true; 10]),
    781
            );
    782
        }
    783
    784
        #[test]
    785
        pub fn test_evaluate_selection_with_non_empty_record_batch() {
    786
            test_evaluate_selection(
    787
                unsafe { &RecordBatch::new_unchecked(Arc::new(Schema::empty()), vec![], 10) },
    788
                &BooleanArray::from(vec![true; 10]),
    789
                &ColumnarValue::Array(Arc::new(Int64Array::from(vec![1; 10]))),
    790
            );
    791
        }
    792
    793
        #[test]
    794
        pub fn test_evaluate_selection_with_non_empty_record_batch_with_larger_false_selection(
    795
        ) {
    796
            test_evaluate_selection_error(
    797
                unsafe { &RecordBatch::new_unchecked(Arc::new(Schema::empty()), vec![], 10) },
    798
                &BooleanArray::from(vec![false; 20]),
    799
            );
    800
        }
    801
    802
        #[test]
    803
        pub fn test_evaluate_selection_with_non_empty_record_batch_with_larger_true_selection(
    804
        ) {
    805
            test_evaluate_selection_error(
    806
                unsafe { &RecordBatch::new_unchecked(Arc::new(Schema::empty()), vec![], 10) },
    807
                &BooleanArray::from(vec![true; 20]),
    808
            );
    809
        }
    810
    811
        #[test]
    812
        pub fn test_evaluate_selection_with_non_empty_record_batch_with_smaller_false_selection(
    813
        ) {
    814
            test_evaluate_selection_error(
    815
                unsafe { &RecordBatch::new_unchecked(Arc::new(Schema::empty()), vec![], 10) },
    816
                &BooleanArray::from(vec![false; 5]),
    817
            );
    818
        }
    819
    820
        #[test]
    821
        pub fn test_evaluate_selection_with_non_empty_record_batch_with_smaller_true_selection(
    822
        ) {
    823
            test_evaluate_selection_error(
    824
                unsafe { &RecordBatch::new_unchecked(Arc::new(Schema::empty()), vec![], 10) },
    825
                &BooleanArray::from(vec![true; 5]),
    826
            );
    827
        }
    828
    }