## 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/expr/src/logical_plan/extension.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
    //! This module defines the interface for logical nodes
    19
    use crate::{Expr, LogicalPlan};
    20
    use datafusion_common::{DFSchema, DFSchemaRef, Result};
    21
    use std::cmp::Ordering;
    22
    use std::hash::{Hash, Hasher};
    23
    use std::{any::Any, collections::HashSet, fmt, sync::Arc};
    24
    25
    use super::InvariantLevel;
    26
    27
    /// This defines the interface for [`LogicalPlan`] nodes that can be
    28
    /// used to extend DataFusion with custom relational operators.
    29
    ///
    30
    /// The [`UserDefinedLogicalNodeCore`] trait is *the recommended way to implement*
    31
    /// this trait and avoids having implementing some required boiler plate code.
    32
    pub trait UserDefinedLogicalNode: fmt::Debug + Send + Sync {
    33
        /// Return a reference to self as Any, to support dynamic downcasting
    34
        ///
    35
        /// Typically this will look like:
    36
        ///
    37
        /// ```
    38
        /// # use std::any::Any;
    39
        /// # struct Dummy { }
    40
        ///
    41
        /// # impl Dummy {
    42
        /// // canonical boiler plate
    43
        /// fn as_any(&self) -> &dyn Any {
    44
        ///     self
    45
        /// }
    46
        /// # }
    47
        /// ```
    48
        fn as_any(&self) -> &dyn Any;
    49
    50
        /// Return the plan's name.
    51
        fn name(&self) -> &str;
    52
    53
        /// Return the logical plan's inputs.
    54
        fn inputs(&self) -> Vec<&LogicalPlan>;
    55
    56
        /// Return the output schema of this logical plan node.
    57
        fn schema(&self) -> &DFSchemaRef;
    58
    59
        /// Perform check of invariants for the extension node.
    60
        fn check_invariants(&self, check: InvariantLevel) -> Result<()>;
    61
    62
        /// Returns all expressions in the current logical plan node. This should
    63
        /// not include expressions of any inputs (aka non-recursively).
    64
        ///
    65
        /// These expressions are used for optimizer
    66
        /// passes and rewrites. See [`LogicalPlan::expressions`] for more details.
    67
        fn expressions(&self) -> Vec<Expr>;
    68
    69
        /// A list of output columns (e.g. the names of columns in
    70
        /// self.schema()) for which predicates can not be pushed below
    71
        /// this node without changing the output.
    72
        ///
    73
        /// By default, this returns all columns and thus prevents any
    74
        /// predicates from being pushed below this node.
    75
    0
        fn prevent_predicate_push_down_columns(&self) -> HashSet<String> {
    76
            // default (safe) is all columns in the schema.
    77
    0
            get_all_columns_from_schema(self.schema())
    78
    0
        }
    79
    80
        /// Write a single line, human readable string to `f` for use in explain plan.
    81
        ///
    82
        /// For example: `TopK: k=10`
    83
        fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result;
    84
    85
        /// Create a new `UserDefinedLogicalNode` with the specified children
    86
        /// and expressions. This function is used during optimization
    87
        /// when the plan is being rewritten and a new instance of the
    88
        /// `UserDefinedLogicalNode` must be created.
    89
        ///
    90
        /// Note that exprs and inputs are in the same order as the result
    91
        /// of self.inputs and self.exprs.
    92
        ///
    93
        /// So, `self.with_exprs_and_inputs(exprs, ..).expressions() == exprs
    94
        fn with_exprs_and_inputs(
    95
            &self,
    96
            exprs: Vec<Expr>,
    97
            inputs: Vec<LogicalPlan>,
    98
        ) -> Result<Arc<dyn UserDefinedLogicalNode>>;
    99
    100
        /// Returns the necessary input columns for this node required to compute
    101
        /// the columns in the output schema
    102
        ///
    103
        /// This is used for projection push-down when DataFusion has determined that
    104
        /// only a subset of the output columns of this node are needed by its parents.
    105
        /// This API is used to tell DataFusion which, if any, of the input columns are no longer
    106
        /// needed.
    107
        ///
    108
        /// Return `None`, the default, if this information can not be determined.
    109
        /// Returns `Some(_)` with the column indices for each child of this node that are
    110
        /// needed to compute `output_columns`
    111
    0
        fn necessary_children_exprs(
    112
    0
            &self,
    113
    0
            _output_columns: &[usize],
    114
    0
        ) -> Option<Vec<Vec<usize>>> {
    115
    0
            None
    116
    0
        }
    117
    118
        /// Update the hash `state` with this node requirements from
    119
        /// [`Hash`].
    120
        ///
    121
        /// Note: consider using [`UserDefinedLogicalNodeCore`] instead of
    122
        /// [`UserDefinedLogicalNode`] directly.
    123
        ///
    124
        /// This method is required to support hashing [`LogicalPlan`]s.  To
    125
        /// implement it, typically the type implementing
    126
        /// [`UserDefinedLogicalNode`] typically implements [`Hash`] and
    127
        /// then the following boiler plate is used:
    128
        ///
    129
        /// # Example:
    130
        /// ```
    131
        /// // User defined node that derives Hash
    132
        /// #[derive(Hash, Debug, PartialEq, Eq)]
    133
        /// struct MyNode {
    134
        ///     val: u64,
    135
        /// }
    136
        ///
    137
        /// // impl UserDefinedLogicalNode {
    138
        /// // ...
    139
        /// # impl MyNode {
    140
        /// // Boiler plate to call the derived Hash impl
    141
        /// fn dyn_hash(&self, state: &mut dyn std::hash::Hasher) {
    142
        ///     use std::hash::Hash;
    143
        ///     let mut s = state;
    144
        ///     self.hash(&mut s);
    145
        /// }
    146
        /// // }
    147
        /// # }
    148
        /// ```
    149
        /// Note: [`UserDefinedLogicalNode`] is not constrained by [`Hash`]
    150
        /// directly because it must remain object safe.
    151
        fn dyn_hash(&self, state: &mut dyn Hasher);
    152
    153
        /// Compare `other`, respecting requirements from [Eq].
    154
        ///
    155
        /// Note: consider using [`UserDefinedLogicalNodeCore`] instead of
    156
        /// [`UserDefinedLogicalNode`] directly.
    157
        ///
    158
        /// When `other` has an another type than `self`, then the values
    159
        /// are *not* equal.
    160
        ///
    161
        /// This method is required to support Eq on [`LogicalPlan`]s.  To
    162
        /// implement it, typically the type implementing
    163
        /// [`UserDefinedLogicalNode`] typically implements [`Eq`] and
    164
        /// then the following boiler plate is used:
    165
        ///
    166
        /// # Example:
    167
        /// ```
    168
        /// # use datafusion_expr::UserDefinedLogicalNode;
    169
        /// // User defined node that derives Eq
    170
        /// #[derive(Hash, Debug, PartialEq, Eq)]
    171
        /// struct MyNode {
    172
        ///     val: u64,
    173
        /// }
    174
        ///
    175
        /// // impl UserDefinedLogicalNode {
    176
        /// // ...
    177
        /// # impl MyNode {
    178
        /// // Boiler plate to call the derived Eq impl
    179
        /// fn dyn_eq(&self, other: &dyn UserDefinedLogicalNode) -> bool {
    180
        ///     match other.as_any().downcast_ref::<Self>() {
    181
        ///         Some(o) => self == o,
    182
        ///         None => false,
    183
        ///     }
    184
        /// }
    185
        /// // }
    186
        /// # }
    187
        /// ```
    188
        /// Note: [`UserDefinedLogicalNode`] is not constrained by [`Eq`]
    189
        /// directly because it must remain object safe.
    190
        fn dyn_eq(&self, other: &dyn UserDefinedLogicalNode) -> bool;
    191
    192
        /// Compare `other`, respecting requirements from [PartialOrd].
    193
        /// Must return `Some(Equal)` if and only if `self.dyn_eq(other)`.
    194
        fn dyn_ord(&self, other: &dyn UserDefinedLogicalNode) -> Option<Ordering>;
    195
    196
        /// Returns `true` if a limit can be safely pushed down through this
    197
        /// `UserDefinedLogicalNode` node.
    198
        ///
    199
        /// If this method returns `true`, and the query plan contains a limit at
    200
        /// the output of this node, DataFusion will push the limit to the input
    201
        /// of this node.
    202
    0
        fn supports_limit_pushdown(&self) -> bool {
    203
    0
            false
    204
    0
        }
    205
    }
    206
    207
    impl Hash for dyn UserDefinedLogicalNode {
    208
    0
        fn hash<H: Hasher>(&self, state: &mut H) {
    209
    0
            self.dyn_hash(state);
    210
    0
        }
    211
    }
    212
    213
    impl PartialEq for dyn UserDefinedLogicalNode {
    214
    0
        fn eq(&self, other: &Self) -> bool {
    215
    0
            self.dyn_eq(other)
    216
    0
        }
    217
    }
    218
    219
    impl PartialOrd for dyn UserDefinedLogicalNode {
    220
    0
        fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
    221
    0
            self.dyn_ord(other)
    222
    0
        }
    223
    }
    224
    225
    impl Eq for dyn UserDefinedLogicalNode {}
    226
    227
    /// This trait facilitates implementation of the [`UserDefinedLogicalNode`].
    228
    ///
    229
    /// See the example in
    230
    /// [user_defined_plan.rs](https://github.com/apache/datafusion/blob/main/datafusion/core/tests/user_defined/user_defined_plan.rs)
    231
    /// file for an example of how to use this extension API.
    232
    pub trait UserDefinedLogicalNodeCore:
    233
        fmt::Debug + Eq + PartialOrd + Hash + Sized + Send + Sync + 'static
    234
    {
    235
        /// Return the plan's name.
    236
        fn name(&self) -> &str;
    237
    238
        /// Return the logical plan's inputs.
    239
        fn inputs(&self) -> Vec<&LogicalPlan>;
    240
    241
        /// Return the output schema of this logical plan node.
    242
        fn schema(&self) -> &DFSchemaRef;
    243
    244
        /// Perform check of invariants for the extension node.
    245
        ///
    246
        /// This is the default implementation for extension nodes.
    247
    0
        fn check_invariants(&self, _check: InvariantLevel) -> Result<()> {
    248
    0
            Ok(())
    249
    0
        }
    250
    251
        /// Returns all expressions in the current logical plan node. This
    252
        /// should not include expressions of any inputs (aka
    253
        /// non-recursively). These expressions are used for optimizer
    254
        /// passes and rewrites.
    255
        fn expressions(&self) -> Vec<Expr>;
    256
    257
        /// A list of output columns (e.g. the names of columns in
    258
        /// self.schema()) for which predicates can not be pushed below
    259
        /// this node without changing the output.
    260
        ///
    261
        /// By default, this returns all columns and thus prevents any
    262
        /// predicates from being pushed below this node.
    263
    0
        fn prevent_predicate_push_down_columns(&self) -> HashSet<String> {
    264
            // default (safe) is all columns in the schema.
    265
    0
            get_all_columns_from_schema(self.schema())
    266
    0
        }
    267
    268
        /// Write a single line, human readable string to `f` for use in explain plan.
    269
        ///
    270
        /// For example: `TopK: k=10`
    271
        fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result;
    272
    273
        /// Create a new `UserDefinedLogicalNode` with the specified children
    274
        /// and expressions. This function is used during optimization
    275
        /// when the plan is being rewritten and a new instance of the
    276
        /// `UserDefinedLogicalNode` must be created.
    277
        ///
    278
        /// Note that exprs and inputs are in the same order as the result
    279
        /// of self.inputs and self.exprs.
    280
        ///
    281
        /// So, `self.with_exprs_and_inputs(exprs, ..).expressions() == exprs
    282
        fn with_exprs_and_inputs(
    283
            &self,
    284
            exprs: Vec<Expr>,
    285
            inputs: Vec<LogicalPlan>,
    286
        ) -> Result<Self>;
    287
    288
        /// Returns the necessary input columns for this node required to compute
    289
        /// the columns in the output schema
    290
        ///
    291
        /// This is used for projection push-down when DataFusion has determined that
    292
        /// only a subset of the output columns of this node are needed by its parents.
    293
        /// This API is used to tell DataFusion which, if any, of the input columns are no longer
    294
        /// needed.
    295
        ///
    296
        /// Return `None`, the default, if this information can not be determined.
    297
        /// Returns `Some(_)` with the column indices for each child of this node that are
    298
        /// needed to compute `output_columns`
    299
    0
        fn necessary_children_exprs(
    300
    0
            &self,
    301
    0
            _output_columns: &[usize],
    302
    0
        ) -> Option<Vec<Vec<usize>>> {
    303
    0
            None
    304
    0
        }
    305
    306
        /// Returns `true` if a limit can be safely pushed down through this
    307
        /// `UserDefinedLogicalNode` node.
    308
        ///
    309
        /// If this method returns `true`, and the query plan contains a limit at
    310
        /// the output of this node, DataFusion will push the limit to the input
    311
        /// of this node.
    312
    0
        fn supports_limit_pushdown(&self) -> bool {
    313
    0
            false // Disallow limit push-down by default
    314
    0
        }
    315
    }
    316
    317
    /// Automatically derive UserDefinedLogicalNode to `UserDefinedLogicalNode`
    318
    /// to avoid boiler plate for implementing `as_any`, `Hash`, `PartialEq` and `PartialOrd`.
    319
    impl<T: UserDefinedLogicalNodeCore> UserDefinedLogicalNode for T {
    320
    0
        fn as_any(&self) -> &dyn Any {
    321
    0
            self
    322
    0
        }
    323
    324
    0
        fn name(&self) -> &str {
    325
    0
            self.name()
    326
    0
        }
    327
    328
    0
        fn inputs(&self) -> Vec<&LogicalPlan> {
    329
    0
            self.inputs()
    330
    0
        }
    331
    332
    0
        fn schema(&self) -> &DFSchemaRef {
    333
    0
            self.schema()
    334
    0
        }
    335
    336
    0
        fn check_invariants(&self, check: InvariantLevel) -> Result<()> {
    337
    0
            self.check_invariants(check)
    338
    0
        }
    339
    340
    0
        fn expressions(&self) -> Vec<Expr> {
    341
    0
            self.expressions()
    342
    0
        }
    343
    344
    0
        fn prevent_predicate_push_down_columns(&self) -> HashSet<String> {
    345
    0
            self.prevent_predicate_push_down_columns()
    346
    0
        }
    347
    348
    0
        fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
    349
    0
            self.fmt_for_explain(f)
    350
    0
        }
    351
    352
    0
        fn with_exprs_and_inputs(
    353
    0
            &self,
    354
    0
            exprs: Vec<Expr>,
    355
    0
            inputs: Vec<LogicalPlan>,
    356
    0
        ) -> Result<Arc<dyn UserDefinedLogicalNode>> {
    357
    0
            Ok(Arc::new(self.with_exprs_and_inputs(exprs, inputs)?))
    358
    0
        }
    359
    360
    0
        fn necessary_children_exprs(
    361
    0
            &self,
    362
    0
            output_columns: &[usize],
    363
    0
        ) -> Option<Vec<Vec<usize>>> {
    364
    0
            self.necessary_children_exprs(output_columns)
    365
    0
        }
    366
    367
    0
        fn dyn_hash(&self, state: &mut dyn Hasher) {
    368
    0
            let mut s = state;
    369
    0
            self.hash(&mut s);
    370
    0
        }
    371
    372
    0
        fn dyn_eq(&self, other: &dyn UserDefinedLogicalNode) -> bool {
    373
    0
            match other.as_any().downcast_ref::<Self>() {
    374
    0
                Some(o) => self == o,
    375
    0
                None => false,
    376
            }
    377
    0
        }
    378
    379
    0
        fn dyn_ord(&self, other: &dyn UserDefinedLogicalNode) -> Option<Ordering> {
    380
    0
            other
    381
    0
                .as_any()
    382
    0
                .downcast_ref::<Self>()
    383
    0
                .and_then(|other| self.partial_cmp(other))
    384
    0
        }
    385
    386
    0
        fn supports_limit_pushdown(&self) -> bool {
    387
    0
            self.supports_limit_pushdown()
    388
    0
        }
    389
    }
    390
    391
    0
    fn get_all_columns_from_schema(schema: &DFSchema) -> HashSet<String> {
    392
    0
        schema.fields().iter().map(|f| f.name().clone()).collect()
    393
    0
    }