## 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/invariants.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 datafusion_common::{
    19
        assert_or_internal_err, plan_err,
    20
        tree_node::{TreeNode, TreeNodeRecursion},
    21
        DFSchemaRef, DataFusionError, Result,
    22
    };
    23
    24
    use crate::{
    25
        expr::{Exists, InSubquery},
    26
        expr_rewriter::strip_outer_reference,
    27
        utils::{collect_subquery_cols, split_conjunction},
    28
        Aggregate, Expr, Filter, Join, JoinType, LogicalPlan, Window,
    29
    };
    30
    31
    use super::Extension;
    32
    33
    #[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Hash)]
    34
    pub enum InvariantLevel {
    35
        /// Invariants that are always true in DataFusion `LogicalPlan`s
    36
        /// such as the number of expected children and no duplicated output fields
    37
        Always,
    38
        /// Invariants that must hold true for the plan to be "executable"
    39
        /// such as the type and number of function arguments are correct and
    40
        /// that wildcards have been expanded
    41
        ///
    42
        /// To ensure a LogicalPlan satisfies the `Executable` invariants, run the
    43
        /// `Analyzer`
    44
        Executable,
    45
    }
    46
    47
    /// Apply the [`InvariantLevel::Always`] check at the current plan node only.
    48
    ///
    49
    /// This does not recurs to any child nodes.
    50
    0
    pub fn assert_always_invariants_at_current_node(plan: &LogicalPlan) -> Result<()> {
    51
        // Refer to <https://datafusion.apache.org/contributor-guide/specification/invariants.html#relation-name-tuples-in-logical-fields-and-logical-columns-are-unique>
    52
    0
        assert_unique_field_names(plan)?;
    53
    54
    0
        Ok(())
    55
    0
    }
    56
    57
    /// Visit the plan nodes, and confirm the [`InvariantLevel::Executable`]
    58
    /// as well as the less stringent [`InvariantLevel::Always`] checks.
    59
    0
    pub fn assert_executable_invariants(plan: &LogicalPlan) -> Result<()> {
    60
        // Always invariants
    61
    0
        assert_always_invariants_at_current_node(plan)?;
    62
    0
        assert_valid_extension_nodes(plan, InvariantLevel::Always)?;
    63
    64
        // Executable invariants
    65
    0
        assert_valid_extension_nodes(plan, InvariantLevel::Executable)?;
    66
    0
        assert_valid_semantic_plan(plan)?;
    67
    0
        Ok(())
    68
    0
    }
    69
    70
    /// Asserts that the query plan, and subplan, extension nodes have valid invariants.
    71
    ///
    72
    /// Refer to [`UserDefinedLogicalNode::check_invariants`](super::UserDefinedLogicalNode)
    73
    /// for more details of user-provided extension node invariants.
    74
    0
    fn assert_valid_extension_nodes(plan: &LogicalPlan, check: InvariantLevel) -> Result<()> {
    75
    0
        plan.apply_with_subqueries(|plan: &LogicalPlan| {
    76
    0
            if let LogicalPlan::Extension(Extension { node }) = plan {
    77
    0
                node.check_invariants(check)?;
    78
    0
            }
    79
    0
            plan.apply_expressions(|expr| {
    80
                // recursively look for subqueries
    81
    0
                expr.apply(|expr| {
    82
    0
                    match expr {
    83
    0
                        Expr::Exists(Exists { subquery, .. })
    84
    0
                        | Expr::InSubquery(InSubquery { subquery, .. })
    85
    0
                        | Expr::ScalarSubquery(subquery) => {
    86
    0
                            assert_valid_extension_nodes(&subquery.subquery, check)?;
    87
                        }
    88
    0
                        _ => {}
    89
                    };
    90
    0
                    Ok(TreeNodeRecursion::Continue)
    91
    0
                })
    92
    0
            })
    93
    0
        })
    94
    0
        .map(|_| ())
    95
    0
    }
    96
    97
    /// Returns an error if plan, and subplans, do not have unique fields.
    98
    ///
    99
    /// This invariant is subject to change.
    100
    /// refer: <https://github.com/apache/datafusion/issues/13525#issuecomment-2494046463>
    101
    0
    fn assert_unique_field_names(plan: &LogicalPlan) -> Result<()> {
    102
    0
        plan.schema().check_names()
    103
    0
    }
    104
    105
    /// Returns an error if the plan is not semantically valid.
    106
    0
    fn assert_valid_semantic_plan(plan: &LogicalPlan) -> Result<()> {
    107
    0
        assert_subqueries_are_valid(plan)?;
    108
    109
    0
        Ok(())
    110
    0
    }
    111
    112
    /// Returns an error if the plan does not have the expected schema.
    113
    /// Ignores metadata and nullability.
    114
    0
    pub fn assert_expected_schema(schema: &DFSchemaRef, plan: &LogicalPlan) -> Result<()> {
    115
    0
        let compatible = plan.schema().logically_equivalent_names_and_types(schema);
    116
    117
    0
        assert_or_internal_err!(
    118
    0
            compatible,
    119
    0
            "Failed due to a difference in schemas: original schema: {:?}, new schema: {:?}",
    120
            schema,
    121
    0
            plan.schema()
    122
        );
    123
    0
        Ok(())
    124
    0
    }
    125
    126
    /// Asserts that the subqueries are structured properly with valid node placement.
    127
    ///
    128
    /// Refer to [`check_subquery_expr`] for more details of the internal invariants.
    129
    0
    fn assert_subqueries_are_valid(plan: &LogicalPlan) -> Result<()> {
    130
    0
        plan.apply_with_subqueries(|plan: &LogicalPlan| {
    131
    0
            plan.apply_expressions(|expr| {
    132
                // recursively look for subqueries
    133
    0
                expr.apply(|expr| {
    134
    0
                    match expr {
    135
    0
                        Expr::Exists(Exists { subquery, .. })
    136
    0
                        | Expr::InSubquery(InSubquery { subquery, .. })
    137
    0
                        | Expr::ScalarSubquery(subquery) => {
    138
    0
                            check_subquery_expr(plan, &subquery.subquery, expr)?;
    139
                        }
    140
    0
                        _ => {}
    141
                    };
    142
    0
                    Ok(TreeNodeRecursion::Continue)
    143
    0
                })
    144
    0
            })
    145
    0
        })
    146
    0
        .map(|_| ())
    147
    0
    }
    148
    149
    /// Do necessary check on subquery expressions and fail the invalid plan
    150
    /// 1) Check whether the outer plan is in the allowed outer plans list to use subquery expressions,
    151
    ///    the allowed while list: [Projection, Filter, Window, Aggregate, Join].
    152
    /// 2) Check whether the inner plan is in the allowed inner plans list to use correlated(outer) expressions.
    153
    /// 3) Check and validate unsupported cases to use the correlated(outer) expressions inside the subquery(inner) plans/inner expressions.
    154
    ///    For example, we do not want to support to use correlated expressions as the Join conditions in the subquery plan when the Join
    155
    ///    is a Full Out Join
    156
    0
    pub fn check_subquery_expr(
    157
    0
        outer_plan: &LogicalPlan,
    158
    0
        inner_plan: &LogicalPlan,
    159
    0
        expr: &Expr,
    160
    0
    ) -> Result<()> {
    161
    0
        assert_subqueries_are_valid(inner_plan)?;
    162
    0
        if let Expr::ScalarSubquery(subquery) = expr {
    163
            // Scalar subquery should only return one column
    164
    0
            if subquery.subquery.schema().fields().len() > 1 {
    165
    0
                return plan_err!(
    166
    0
                    "Scalar subquery should only return one column, but found {}: {}",
    167
    0
                    subquery.subquery.schema().fields().len(),
    168
    0
                    subquery.subquery.schema().field_names().join(", ")
    169
                );
    170
    0
            }
    171
            // Correlated scalar subquery must be aggregated to return at most one row
    172
    0
            if !subquery.outer_ref_columns.is_empty() {
    173
    0
                match strip_inner_query(inner_plan) {
    174
    0
                    LogicalPlan::Aggregate(agg) => {
    175
    0
                        check_aggregation_in_scalar_subquery(inner_plan, agg)
    176
                    }
    177
    0
                    LogicalPlan::Filter(Filter { input, .. })
    178
    0
                        if matches!(input.as_ref(), LogicalPlan::Aggregate(_)) =>
    179
                    {
    180
    0
                        if let LogicalPlan::Aggregate(agg) = input.as_ref() {
    181
    0
                            check_aggregation_in_scalar_subquery(inner_plan, agg)
    182
                        } else {
    183
    0
                            Ok(())
    184
                        }
    185
                    }
    186
                    _ => {
    187
    0
                        if inner_plan
    188
    0
                            .max_rows()
    189
    0
                            .filter(|max_row| *max_row <= 1)
    190
    0
                            .is_some()
    191
                        {
    192
    0
                            Ok(())
    193
                        } else {
    194
    0
                            plan_err!(
    195
    0
                                "Correlated scalar subquery must be aggregated to return at most one row"
    196
                            )
    197
                        }
    198
                    }
    199
    0
                }?;
    200
    0
                match outer_plan {
    201
                    LogicalPlan::Projection(_)
    202
    0
                    | LogicalPlan::Filter(_) => Ok(()),
    203
    0
                    LogicalPlan::Aggregate(Aggregate { group_expr, aggr_expr, .. }) => {
    204
    0
                        if group_expr.contains(expr) && !aggr_expr.contains(expr) {
    205
                            // TODO revisit this validation logic
    206
    0
                            plan_err!(
    207
    0
                                "Correlated scalar subquery in the GROUP BY clause must also be in the aggregate expressions"
    208
                            )
    209
                        } else {
    210
    0
                            Ok(())
    211
                        }
    212
                    }
    213
    0
                    _ => plan_err!(
    214
    0
                        "Correlated scalar subquery can only be used in Projection, Filter, Aggregate plan nodes"
    215
                    )
    216
    0
                }?;
    217
    0
            }
    218
    0
            check_correlations_in_subquery(inner_plan)
    219
        } else {
    220
    0
            if let Expr::InSubquery(subquery) = expr {
    221
                // InSubquery should only return one column
    222
    0
                if subquery.subquery.subquery.schema().fields().len() > 1 {
    223
    0
                    return plan_err!(
    224
    0
                        "InSubquery should only return one column, but found {}: {}",
    225
    0
                        subquery.subquery.subquery.schema().fields().len(),
    226
    0
                        subquery.subquery.subquery.schema().field_names().join(", ")
    227
                    );
    228
    0
                }
    229
    0
            }
    230
    0
            match outer_plan {
    231
                LogicalPlan::Projection(_)
    232
                | LogicalPlan::Filter(_)
    233
                | LogicalPlan::TableScan(_)
    234
                | LogicalPlan::Window(_)
    235
                | LogicalPlan::Aggregate(_)
    236
    0
                | LogicalPlan::Join(_) => Ok(()),
    237
    0
                _ => plan_err!(
    238
    0
                    "In/Exist subquery can only be used in \
    239
    0
                    Projection, Filter, TableScan, Window functions, Aggregate and Join plan nodes, \
    240
    0
                    but was used in [{}]",
    241
    0
                    outer_plan.display()
    242
                ),
    243
    0
            }?;
    244
    0
            check_correlations_in_subquery(inner_plan)
    245
        }
    246
    0
    }
    247
    248
    // Recursively check the unsupported outer references in the sub query plan.
    249
    0
    fn check_correlations_in_subquery(inner_plan: &LogicalPlan) -> Result<()> {
    250
    0
        check_inner_plan(inner_plan)
    251
    0
    }
    252
    253
    // Recursively check the unsupported outer references in the sub query plan.
    254
    #[cfg_attr(feature = "recursive_protection", recursive::recursive)]
    255
    0
    fn check_inner_plan(inner_plan: &LogicalPlan) -> Result<()> {
    256
        // We want to support as many operators as possible inside the correlated subquery
    257
    0
        match inner_plan {
    258
            LogicalPlan::Aggregate(_) => {
    259
    0
                inner_plan.apply_children(|plan| {
    260
    0
                    check_inner_plan(plan)?;
    261
    0
                    Ok(TreeNodeRecursion::Continue)
    262
    0
                })?;
    263
    0
                Ok(())
    264
            }
    265
    0
            LogicalPlan::Filter(Filter { input, .. }) => check_inner_plan(input),
    266
    0
            LogicalPlan::Window(window) => {
    267
    0
                check_mixed_out_refer_in_window(window)?;
    268
    0
                inner_plan.apply_children(|plan| {
    269
    0
                    check_inner_plan(plan)?;
    270
    0
                    Ok(TreeNodeRecursion::Continue)
    271
    0
                })?;
    272
    0
                Ok(())
    273
            }
    274
            LogicalPlan::Projection(_)
    275
            | LogicalPlan::Distinct(_)
    276
            | LogicalPlan::Sort(_)
    277
            | LogicalPlan::Union(_)
    278
            | LogicalPlan::TableScan(_)
    279
            | LogicalPlan::EmptyRelation(_)
    280
            | LogicalPlan::Limit(_)
    281
            | LogicalPlan::Values(_)
    282
            | LogicalPlan::Subquery(_)
    283
            | LogicalPlan::SubqueryAlias(_)
    284
            | LogicalPlan::Unnest(_) => {
    285
    0
                inner_plan.apply_children(|plan| {
    286
    0
                    check_inner_plan(plan)?;
    287
    0
                    Ok(TreeNodeRecursion::Continue)
    288
    0
                })?;
    289
    0
                Ok(())
    290
            }
    291
            LogicalPlan::Join(Join {
    292
    0
                left,
    293
    0
                right,
    294
    0
                join_type,
    295
                ..
    296
    0
            }) => match join_type {
    297
                JoinType::Inner => {
    298
    0
                    inner_plan.apply_children(|plan| {
    299
    0
                        check_inner_plan(plan)?;
    300
    0
                        Ok(TreeNodeRecursion::Continue)
    301
    0
                    })?;
    302
    0
                    Ok(())
    303
                }
    304
                JoinType::Left
    305
                | JoinType::LeftSemi
    306
                | JoinType::LeftAnti
    307
                | JoinType::LeftMark => {
    308
    0
                    check_inner_plan(left)?;
    309
    0
                    check_no_outer_references(right)
    310
                }
    311
                JoinType::Right
    312
                | JoinType::RightSemi
    313
                | JoinType::RightAnti
    314
                | JoinType::RightMark => {
    315
    0
                    check_no_outer_references(left)?;
    316
    0
                    check_inner_plan(right)
    317
                }
    318
                JoinType::Full => {
    319
    0
                    inner_plan.apply_children(|plan| {
    320
    0
                        check_no_outer_references(plan)?;
    321
    0
                        Ok(TreeNodeRecursion::Continue)
    322
    0
                    })?;
    323
    0
                    Ok(())
    324
                }
    325
            },
    326
    0
            LogicalPlan::Extension(_) => Ok(()),
    327
    0
            plan => check_no_outer_references(plan),
    328
        }
    329
    0
    }
    330
    331
    0
    fn check_no_outer_references(inner_plan: &LogicalPlan) -> Result<()> {
    332
    0
        if inner_plan.contains_outer_reference() {
    333
    0
            plan_err!(
    334
    0
                "Accessing outer reference columns is not allowed in the plan: {}",
    335
    0
                inner_plan.display()
    336
            )
    337
        } else {
    338
    0
            Ok(())
    339
        }
    340
    0
    }
    341
    342
    0
    fn check_aggregation_in_scalar_subquery(
    343
    0
        inner_plan: &LogicalPlan,
    344
    0
        agg: &Aggregate,
    345
    0
    ) -> Result<()> {
    346
    0
        if agg.aggr_expr.is_empty() {
    347
    0
            return plan_err!(
    348
    0
                "Correlated scalar subquery must be aggregated to return at most one row"
    349
            );
    350
    0
        }
    351
    0
        if !agg.group_expr.is_empty() {
    352
    0
            let correlated_exprs = get_correlated_expressions(inner_plan)?;
    353
    0
            let inner_subquery_cols =
    354
    0
                collect_subquery_cols(&correlated_exprs, agg.input.schema())?;
    355
    0
            let mut group_columns = agg
    356
    0
                .group_expr
    357
    0
                .iter()
    358
    0
                .map(|group| Ok(group.column_refs().into_iter().cloned().collect::<Vec<_>>()))
    359
    0
                .collect::<Result<Vec<_>>>()?
    360
    0
                .into_iter()
    361
    0
                .flatten();
    362
    363
    0
            if !group_columns.all(|group| inner_subquery_cols.contains(&group)) {
    364
                // Group BY columns must be a subset of columns in the correlated expressions
    365
    0
                return plan_err!(
    366
    0
                    "A GROUP BY clause in a scalar correlated subquery cannot contain non-correlated columns"
    367
                );
    368
    0
            }
    369
    0
        }
    370
    0
        Ok(())
    371
    0
    }
    372
    373
    0
    fn strip_inner_query(inner_plan: &LogicalPlan) -> &LogicalPlan {
    374
    0
        match inner_plan {
    375
    0
            LogicalPlan::Projection(projection) => {
    376
    0
                strip_inner_query(projection.input.as_ref())
    377
            }
    378
    0
            LogicalPlan::SubqueryAlias(alias) => strip_inner_query(alias.input.as_ref()),
    379
    0
            other => other,
    380
        }
    381
    0
    }
    382
    383
    0
    fn get_correlated_expressions(inner_plan: &LogicalPlan) -> Result<Vec<Expr>> {
    384
    0
        let mut exprs = vec![];
    385
    0
        inner_plan.apply_with_subqueries(|plan| {
    386
    0
            if let LogicalPlan::Filter(Filter { predicate, .. }) = plan {
    387
    0
                let (correlated, _): (Vec<_>, Vec<_>) = split_conjunction(predicate)
    388
    0
                    .into_iter()
    389
    0
                    .partition(|e| e.contains_outer());
    390
    391
    0
                for expr in correlated {
    392
    0
                    exprs.push(strip_outer_reference(expr.clone()));
    393
    0
                }
    394
    0
            }
    395
    0
            Ok(TreeNodeRecursion::Continue)
    396
    0
        })?;
    397
    0
        Ok(exprs)
    398
    0
    }
    399
    400
    /// Check whether the window expressions contain a mixture of out reference columns and inner columns
    401
    0
    fn check_mixed_out_refer_in_window(window: &Window) -> Result<()> {
    402
    0
        let mixed = window
    403
    0
            .window_expr
    404
    0
            .iter()
    405
    0
            .any(|win_expr| win_expr.contains_outer() && win_expr.any_column_refs());
    406
    0
        if mixed {
    407
    0
            plan_err!(
    408
    0
                "Window expressions should not contain a mixed of outer references and inner columns"
    409
            )
    410
        } else {
    411
    0
            Ok(())
    412
        }
    413
    0
    }
    414
    415
    #[cfg(test)]
    416
    mod test {
    417
        use std::cmp::Ordering;
    418
        use std::sync::Arc;
    419
    420
        use crate::{Extension, UserDefinedLogicalNodeCore};
    421
        use datafusion_common::{DFSchema, DFSchemaRef};
    422
    423
        use super::*;
    424
    425
        #[derive(Debug, PartialEq, Eq, Hash)]
    426
        struct MockUserDefinedLogicalPlan {
    427
            empty_schema: DFSchemaRef,
    428
        }
    429
    430
        impl PartialOrd for MockUserDefinedLogicalPlan {
    431
            fn partial_cmp(&self, _other: &Self) -> Option<Ordering> {
    432
                None
    433
            }
    434
        }
    435
    436
        impl UserDefinedLogicalNodeCore for MockUserDefinedLogicalPlan {
    437
            fn name(&self) -> &str {
    438
                "MockUserDefinedLogicalPlan"
    439
            }
    440
    441
            fn inputs(&self) -> Vec<&LogicalPlan> {
    442
                vec![]
    443
            }
    444
    445
            fn schema(&self) -> &DFSchemaRef {
    446
                &self.empty_schema
    447
            }
    448
    449
            fn expressions(&self) -> Vec<Expr> {
    450
                vec![]
    451
            }
    452
    453
            fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
    454
                write!(f, "MockUserDefinedLogicalPlan")
    455
            }
    456
    457
            fn with_exprs_and_inputs(
    458
                &self,
    459
                _exprs: Vec<Expr>,
    460
                _inputs: Vec<LogicalPlan>,
    461
            ) -> Result<Self> {
    462
                Ok(Self {
    463
                    empty_schema: Arc::clone(&self.empty_schema),
    464
                })
    465
            }
    466
    467
            fn supports_limit_pushdown(&self) -> bool {
    468
                false // Disallow limit push-down by default
    469
            }
    470
        }
    471
    472
        #[test]
    473
        fn wont_fail_extension_plan() {
    474
            let plan = LogicalPlan::Extension(Extension {
    475
                node: Arc::new(MockUserDefinedLogicalPlan {
    476
                    empty_schema: DFSchemaRef::new(DFSchema::empty()),
    477
                }),
    478
            });
    479
    480
            check_inner_plan(&plan).unwrap();
    481
        }
    482
    }