## 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/src/projection.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::ops::Deref;
    19
    use std::sync::Arc;
    20
    21
    use crate::expressions::Column;
    22
    use crate::utils::collect_columns;
    23
    use crate::PhysicalExpr;
    24
    25
    use arrow::array::{RecordBatch, RecordBatchOptions};
    26
    use arrow::datatypes::{Field, Schema, SchemaRef};
    27
    use datafusion_common::stats::{ColumnStatistics, Precision};
    28
    use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
    29
    use datafusion_common::{
    30
        assert_or_internal_err, internal_datafusion_err, plan_err, DataFusionError, Result,
    31
    };
    32
    33
    use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
    34
    use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays;
    35
    use indexmap::IndexMap;
    36
    use itertools::Itertools;
    37
    38
    /// A projection expression as used by projection operations.
    39
    ///
    40
    /// The expression is evaluated and the result is stored in a column
    41
    /// with the name specified by `alias`.
    42
    ///
    43
    /// For example, the SQL expression `a + b AS sum_ab` would be represented
    44
    /// as a `ProjectionExpr` where `expr` is the expression `a + b`
    45
    /// and `alias` is the string `sum_ab`.
    46
    #[derive(Debug, Clone)]
    47
    pub struct ProjectionExpr {
    48
        /// The expression that will be evaluated.
    49
        pub expr: Arc<dyn PhysicalExpr>,
    50
        /// The name of the output column for use an output schema.
    51
        pub alias: String,
    52
    }
    53
    54
    impl PartialEq for ProjectionExpr {
    55
    0
        fn eq(&self, other: &Self) -> bool {
    56
    0
            let ProjectionExpr { expr, alias } = self;
    57
    0
            expr.eq(&other.expr) && *alias == other.alias
    58
    0
        }
    59
    }
    60
    61
    impl Eq for ProjectionExpr {}
    62
    63
    impl std::fmt::Display for ProjectionExpr {
    64
    0
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
    65
    0
            if self.expr.to_string() == self.alias {
    66
    0
                write!(f, "{}", self.alias)
    67
            } else {
    68
    0
                write!(f, "{} AS {}", self.expr, self.alias)
    69
            }
    70
    0
        }
    71
    }
    72
    73
    impl ProjectionExpr {
    74
        /// Create a new projection expression
    75
    0
        pub fn new(expr: Arc<dyn PhysicalExpr>, alias: String) -> Self {
    76
    0
            Self { expr, alias }
    77
    0
        }
    78
    79
        /// Create a new projection expression from an expression and a schema using the expression's output field name as alias.
    80
    0
        pub fn new_from_expression(
    81
    0
            expr: Arc<dyn PhysicalExpr>,
    82
    0
            schema: &Schema,
    83
    0
        ) -> Result<Self> {
    84
    0
            let field = expr.return_field(schema)?;
    85
    0
            Ok(Self {
    86
    0
                expr,
    87
    0
                alias: field.name().to_string(),
    88
    0
            })
    89
    0
        }
    90
    }
    91
    92
    impl From<(Arc<dyn PhysicalExpr>, String)> for ProjectionExpr {
    93
    0
        fn from(value: (Arc<dyn PhysicalExpr>, String)) -> Self {
    94
    0
            Self::new(value.0, value.1)
    95
    0
        }
    96
    }
    97
    98
    impl From<&(Arc<dyn PhysicalExpr>, String)> for ProjectionExpr {
    99
    0
        fn from(value: &(Arc<dyn PhysicalExpr>, String)) -> Self {
    100
    0
            Self::new(Arc::clone(&value.0), value.1.clone())
    101
    0
        }
    102
    }
    103
    104
    impl From<ProjectionExpr> for (Arc<dyn PhysicalExpr>, String) {
    105
    0
        fn from(value: ProjectionExpr) -> Self {
    106
    0
            (value.expr, value.alias)
    107
    0
        }
    108
    }
    109
    110
    /// A collection of projection expressions.
    111
    ///
    112
    /// This struct encapsulates multiple `ProjectionExpr` instances,
    113
    /// representing a complete projection operation and provides
    114
    /// methods to manipulate and analyze the projection as a whole.
    115
    #[derive(Debug, Clone, PartialEq, Eq)]
    116
    pub struct ProjectionExprs {
    117
        exprs: Vec<ProjectionExpr>,
    118
    }
    119
    120
    impl std::fmt::Display for ProjectionExprs {
    121
    0
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
    122
    0
            let exprs: Vec<String> = self.exprs.iter().map(|e| e.to_string()).collect();
    123
    0
            write!(f, "Projection[{}]", exprs.join(", "))
    124
    0
        }
    125
    }
    126
    127
    impl From<Vec<ProjectionExpr>> for ProjectionExprs {
    128
    0
        fn from(value: Vec<ProjectionExpr>) -> Self {
    129
    0
            Self { exprs: value }
    130
    0
        }
    131
    }
    132
    133
    impl From<&[ProjectionExpr]> for ProjectionExprs {
    134
    0
        fn from(value: &[ProjectionExpr]) -> Self {
    135
    0
            Self {
    136
    0
                exprs: value.to_vec(),
    137
    0
            }
    138
    0
        }
    139
    }
    140
    141
    impl FromIterator<ProjectionExpr> for ProjectionExprs {
    142
    0
        fn from_iter<T: IntoIterator<Item = ProjectionExpr>>(exprs: T) -> Self {
    143
    0
            Self {
    144
    0
                exprs: exprs.into_iter().collect::<Vec<_>>(),
    145
    0
            }
    146
    0
        }
    147
    }
    148
    149
    impl AsRef<[ProjectionExpr]> for ProjectionExprs {
    150
    0
        fn as_ref(&self) -> &[ProjectionExpr] {
    151
    0
            &self.exprs
    152
    0
        }
    153
    }
    154
    155
    impl ProjectionExprs {
    156
    0
        pub fn new<I>(exprs: I) -> Self
    157
    0
        where
    158
    0
            I: IntoIterator<Item = ProjectionExpr>,
    159
        {
    160
    0
            Self {
    161
    0
                exprs: exprs.into_iter().collect::<Vec<_>>(),
    162
    0
            }
    163
    0
        }
    164
    165
        /// Creates a [`ProjectionExpr`] from a list of column indices.
    166
        ///
    167
        /// This is a convenience method for creating simple column-only projections, where each projection expression is a reference to a column
    168
        /// in the input schema.
    169
        ///
    170
        /// # Behavior
    171
        /// - Ordering: the output projection preserves the exact order of indices provided in the input slice
    172
        ///   For example, `[2, 0, 1]` will produce projections for columns 2, 0, then 1 in that order
    173
        /// - Duplicates: Duplicate indices are allowed and will create multiple projection expressions referencing the same source column
    174
        ///   For example, `[0, 0]` creates 2 separate projections both referencing column 0
    175
        ///
    176
        /// # Panics
    177
        /// Panics if any index in `indices` is out of bounds for the provided schema.
    178
        ///
    179
        /// # Example
    180
        ///
    181
        /// ```rust
    182
        /// use arrow::datatypes::{DataType, Field, Schema};
    183
        /// use datafusion_physical_expr::projection::ProjectionExprs;
    184
        /// use std::sync::Arc;
    185
        ///
    186
        /// // Create a schema with three columns
    187
        /// let schema = Arc::new(Schema::new(vec![
    188
        ///     Field::new("a", DataType::Int32, false),
    189
        ///     Field::new("b", DataType::Utf8, false),
    190
        ///     Field::new("c", DataType::Float64, false),
    191
        /// ]));
    192
        ///
    193
        /// // Project columns at indices 2 and 0 (c and a) - ordering is preserved
    194
        /// let projection = ProjectionExprs::from_indices(&[2, 0], &schema);
    195
        ///
    196
        /// // This creates: SELECT c@2 AS c, a@0 AS a
    197
        /// assert_eq!(projection.as_ref().len(), 2);
    198
        /// assert_eq!(projection.as_ref()[0].alias, "c");
    199
        /// assert_eq!(projection.as_ref()[1].alias, "a");
    200
        ///
    201
        /// // Duplicate indices are allowed
    202
        /// let projection_with_dups = ProjectionExprs::from_indices(&[0, 0, 1], &schema);
    203
        /// assert_eq!(projection_with_dups.as_ref().len(), 3);
    204
        /// assert_eq!(projection_with_dups.as_ref()[0].alias, "a");
    205
        /// assert_eq!(projection_with_dups.as_ref()[1].alias, "a"); // duplicate
    206
        /// assert_eq!(projection_with_dups.as_ref()[2].alias, "b");
    207
        /// ```
    208
    0
        pub fn from_indices(indices: &[usize], schema: &Schema) -> Self {
    209
    0
            let projection_exprs = indices.iter().map(|&i| {
    210
    0
                let field = schema.field(i);
    211
    0
                ProjectionExpr {
    212
    0
                    expr: Arc::new(Column::new(field.name(), i)),
    213
    0
                    alias: field.name().clone(),
    214
    0
                }
    215
    0
            });
    216
    217
    0
            Self::from_iter(projection_exprs)
    218
    0
        }
    219
    220
        /// Returns an iterator over the projection expressions
    221
    0
        pub fn iter(&self) -> impl Iterator<Item = &ProjectionExpr> {
    222
    0
            self.exprs.iter()
    223
    0
        }
    224
    225
        /// Creates a ProjectionMapping from this projection
    226
    0
        pub fn projection_mapping(
    227
    0
            &self,
    228
    0
            input_schema: &SchemaRef,
    229
    0
        ) -> Result<ProjectionMapping> {
    230
    0
            ProjectionMapping::try_new(
    231
    0
                self.exprs
    232
    0
                    .iter()
    233
    0
                    .map(|p| (Arc::clone(&p.expr), p.alias.clone())),
    234
    0
                input_schema,
    235
            )
    236
    0
        }
    237
    238
        /// Iterate over a clone of the projection expressions.
    239
    0
        pub fn expr_iter(&self) -> impl Iterator<Item = Arc<dyn PhysicalExpr>> + '_ {
    240
    0
            self.exprs.iter().map(|e| Arc::clone(&e.expr))
    241
    0
        }
    242
    243
        /// Apply another projection on top of this projection, returning the combined projection.
    244
        /// For example, if this projection is `SELECT c@2 AS x, b@1 AS y, a@0 as z` and the other projection is `SELECT x@0 + 1 AS c1, y@1 + z@2 as c2`,
    245
        /// we return a projection equivalent to `SELECT c@2 + 1 AS c1, b@1 + a@0 as c2`.
    246
        ///
    247
        /// # Example
    248
        ///
    249
        /// ```rust
    250
        /// use datafusion_common::{Result, ScalarValue};
    251
        /// use datafusion_expr::Operator;
    252
        /// use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal};
    253
        /// use datafusion_physical_expr::projection::{ProjectionExpr, ProjectionExprs};
    254
        /// use std::sync::Arc;
    255
        ///
    256
        /// fn main() -> Result<()> {
    257
        ///     // Example from the docstring:
    258
        ///     // Base projection: SELECT c@2 AS x, b@1 AS y, a@0 AS z
    259
        ///     let base = ProjectionExprs::new(vec![
    260
        ///         ProjectionExpr {
    261
        ///             expr: Arc::new(Column::new("c", 2)),
    262
        ///             alias: "x".to_string(),
    263
        ///         },
    264
        ///         ProjectionExpr {
    265
        ///             expr: Arc::new(Column::new("b", 1)),
    266
        ///             alias: "y".to_string(),
    267
        ///         },
    268
        ///         ProjectionExpr {
    269
        ///             expr: Arc::new(Column::new("a", 0)),
    270
        ///             alias: "z".to_string(),
    271
        ///         },
    272
        ///     ]);
    273
        ///
    274
        ///     // Top projection: SELECT x@0 + 1 AS c1, y@1 + z@2 AS c2
    275
        ///     let top = ProjectionExprs::new(vec![
    276
        ///         ProjectionExpr {
    277
        ///             expr: Arc::new(BinaryExpr::new(
    278
        ///                 Arc::new(Column::new("x", 0)),
    279
        ///                 Operator::Plus,
    280
        ///                 Arc::new(Literal::new(ScalarValue::Int32(Some(1)))),
    281
        ///             )),
    282
        ///             alias: "c1".to_string(),
    283
        ///         },
    284
        ///         ProjectionExpr {
    285
        ///             expr: Arc::new(BinaryExpr::new(
    286
        ///                 Arc::new(Column::new("y", 1)),
    287
        ///                 Operator::Plus,
    288
        ///                 Arc::new(Column::new("z", 2)),
    289
        ///             )),
    290
        ///             alias: "c2".to_string(),
    291
        ///         },
    292
        ///     ]);
    293
        ///
    294
        ///     // Expected result: SELECT c@2 + 1 AS c1, b@1 + a@0 AS c2
    295
        ///     let result = base.try_merge(&top)?;
    296
        ///
    297
        ///     assert_eq!(result.as_ref().len(), 2);
    298
        ///     assert_eq!(result.as_ref()[0].alias, "c1");
    299
        ///     assert_eq!(result.as_ref()[1].alias, "c2");
    300
        ///
    301
        ///     Ok(())
    302
        /// }
    303
        /// ```
    304
        ///
    305
        /// # Errors
    306
        /// This function returns an error if any expression in the `other` projection cannot be
    307
        /// applied on top of this projection.
    308
    0
        pub fn try_merge(&self, other: &ProjectionExprs) -> Result<ProjectionExprs> {
    309
    0
            let mut new_exprs = Vec::with_capacity(other.exprs.len());
    310
    0
            for proj_expr in &other.exprs {
    311
    0
                let new_expr = update_expr(&proj_expr.expr, &self.exprs, true)?
    312
    0
                    .ok_or_else(|| {
    313
    0
                        internal_datafusion_err!(
    314
    0
                            "Failed to combine projections: expression {} could not be applied on top of existing projections {}",
    315
                            proj_expr.expr,
    316
    0
                            self.exprs.iter().map(|e| format!("{e}")).join(", ")
    317
                        )
    318
    0
                    })?;
    319
    0
                new_exprs.push(ProjectionExpr {
    320
    0
                    expr: new_expr,
    321
    0
                    alias: proj_expr.alias.clone(),
    322
    0
                });
    323
            }
    324
    0
            Ok(ProjectionExprs::new(new_exprs))
    325
    0
        }
    326
    327
        /// Extract the column indices used in this projection.
    328
        /// For example, for a projection `SELECT a AS x, b + 1 AS y`, where `a` is at index 0 and `b` is at index 1,
    329
        /// this function would return `[0, 1]`.
    330
        /// Repeated indices are returned only once, and the order is ascending.
    331
    0
        pub fn column_indices(&self) -> Vec<usize> {
    332
    0
            self.exprs
    333
    0
                .iter()
    334
    0
                .flat_map(|e| collect_columns(&e.expr).into_iter().map(|col| col.index()))
    335
    0
                .sorted_unstable()
    336
    0
                .dedup()
    337
    0
                .collect_vec()
    338
    0
        }
    339
    340
        /// Extract the ordered column indices for a column-only projection.
    341
        ///
    342
        /// This function assumes that all expressions in the projection are simple column references.
    343
        /// It returns the column indices in the order they appear in the projection.
    344
        ///
    345
        /// # Panics
    346
        ///
    347
        /// Panics if any expression in the projection is not a simple column reference. This includes:
    348
        /// - Computed expressions (e.g., `a + 1`, `CAST(a AS INT)`)
    349
        /// - Function calls (e.g., `UPPER(name)`, `SUM(amount)`)
    350
        /// - Literals (e.g., `42`, `'hello'`)
    351
        /// - Complex nested expressions (e.g., `CASE WHEN ... THEN ... END`)
    352
        ///
    353
        /// # Returns
    354
        ///
    355
        /// A vector of column indices in projection order. Unlike [`column_indices()`](Self::column_indices),
    356
        /// this function:
    357
        /// - Preserves the projection order (does not sort)
    358
        /// - Preserves duplicates (does not deduplicate)
    359
        ///
    360
        /// # Example
    361
        ///
    362
        /// For a projection `SELECT c, a, c` where `a` is at index 0 and `c` is at index 2,
    363
        /// this function would return `[2, 0, 2]`.
    364
        ///
    365
        /// Use [`column_indices()`](Self::column_indices) instead if the projection may contain
    366
        /// non-column expressions or if you need a deduplicated sorted list.
    367
    0
        pub fn ordered_column_indices(&self) -> Vec<usize> {
    368
    0
            self.exprs
    369
    0
                .iter()
    370
    0
                .map(|e| {
    371
    0
                    e.expr
    372
    0
                        .as_any()
    373
    0
                        .downcast_ref::<Column>()
    374
    0
                        .expect("Expected column reference in projection")
    375
    0
                        .index()
    376
    0
                })
    377
    0
                .collect()
    378
    0
        }
    379
    380
        /// Project a schema according to this projection.
    381
        /// For example, for a projection `SELECT a AS x, b + 1 AS y`, where `a` is at index 0 and `b` is at index 1,
    382
        /// if the input schema is `[a: Int32, b: Int32, c: Int32]`, the output schema would be `[x: Int32, y: Int32]`.
    383
        /// Fields' metadata are preserved from the input schema.
    384
    0
        pub fn project_schema(&self, input_schema: &Schema) -> Result<Schema> {
    385
    0
            let fields: Result<Vec<Field>> = self
    386
    0
                .exprs
    387
    0
                .iter()
    388
    0
                .map(|proj_expr| {
    389
    0
                    let metadata = proj_expr
    390
    0
                        .expr
    391
    0
                        .return_field(input_schema)?
    392
    0
                        .metadata()
    393
    0
                        .clone();
    394
    395
    0
                    let field = Field::new(
    396
    0
                        &proj_expr.alias,
    397
    0
                        proj_expr.expr.data_type(input_schema)?,
    398
    0
                        proj_expr.expr.nullable(input_schema)?,
    399
                    )
    400
    0
                    .with_metadata(metadata);
    401
    402
    0
                    Ok(field)
    403
    0
                })
    404
    0
                .collect();
    405
    406
    0
            Ok(Schema::new_with_metadata(
    407
    0
                fields?,
    408
    0
                input_schema.metadata().clone(),
    409
            ))
    410
    0
        }
    411
    412
        /// Create a new [`Projector`] from this projection and an input schema.
    413
        ///
    414
        /// A [`Projector`] can be used to apply this projection to record batches.
    415
        ///
    416
        /// # Errors
    417
        /// This function returns an error if the output schema cannot be constructed from the input schema
    418
        /// with the given projection expressions.
    419
        /// For example, if an expression only works with integer columns but the input schema has a string column at that index.
    420
    0
        pub fn make_projector(&self, input_schema: &Schema) -> Result<Projector> {
    421
    0
            let output_schema = Arc::new(self.project_schema(input_schema)?);
    422
    0
            Ok(Projector {
    423
    0
                projection: self.clone(),
    424
    0
                output_schema,
    425
    0
            })
    426
    0
        }
    427
    428
        /// Project statistics according to this projection.
    429
        /// For example, for a projection `SELECT a AS x, b + 1 AS y`, where `a` is at index 0 and `b` is at index 1,
    430
        /// if the input statistics has column statistics for columns `a`, `b`, and `c`, the output statistics would have column statistics for columns `x` and `y`.
    431
    0
        pub fn project_statistics(
    432
    0
            &self,
    433
    0
            mut stats: datafusion_common::Statistics,
    434
    0
            input_schema: &Schema,
    435
    0
        ) -> Result<datafusion_common::Statistics> {
    436
    0
            let mut primitive_row_size = 0;
    437
    0
            let mut primitive_row_size_possible = true;
    438
    0
            let mut column_statistics = vec![];
    439
    440
    0
            for proj_expr in &self.exprs {
    441
    0
                let expr = &proj_expr.expr;
    442
    0
                let col_stats = if let Some(col) = expr.as_any().downcast_ref::<Column>() {
    443
    0
                    stats.column_statistics[col.index()].clone()
    444
                } else {
    445
                    // TODO stats: estimate more statistics from expressions
    446
                    // (expressions should compute their statistics themselves)
    447
    0
                    ColumnStatistics::new_unknown()
    448
                };
    449
    0
                column_statistics.push(col_stats);
    450
    0
                let data_type = expr.data_type(input_schema)?;
    451
    0
                if let Some(value) = data_type.primitive_width() {
    452
    0
                    primitive_row_size += value;
    453
    0
                    continue;
    454
    0
                }
    455
    0
                primitive_row_size_possible = false;
    456
            }
    457
    458
    0
            if primitive_row_size_possible {
    459
    0
                stats.total_byte_size =
    460
    0
                    Precision::Exact(primitive_row_size).multiply(&stats.num_rows);
    461
    0
            }
    462
    0
            stats.column_statistics = column_statistics;
    463
    0
            Ok(stats)
    464
    0
        }
    465
    }
    466
    467
    impl<'a> IntoIterator for &'a ProjectionExprs {
    468
        type Item = &'a ProjectionExpr;
    469
        type IntoIter = std::slice::Iter<'a, ProjectionExpr>;
    470
    471
    0
        fn into_iter(self) -> Self::IntoIter {
    472
    0
            self.exprs.iter()
    473
    0
        }
    474
    }
    475
    476
    /// Applies a projection to record batches.
    477
    ///
    478
    /// A [`Projector`] uses a set of projection expressions to transform
    479
    /// and a pre-computed output schema to project record batches accordingly.
    480
    ///
    481
    /// The main reason to use a `Projector` is to avoid repeatedly computing
    482
    /// the output schema for each batch, which can be costly if the projection
    483
    /// expressions are complex.
    484
    #[derive(Clone, Debug)]
    485
    pub struct Projector {
    486
        projection: ProjectionExprs,
    487
        output_schema: SchemaRef,
    488
    }
    489
    490
    impl Projector {
    491
        /// Project a record batch according to this projector's expressions.
    492
        ///
    493
        /// # Errors
    494
        /// This function returns an error if any expression evaluation fails
    495
        /// or if the output schema of the resulting record batch does not match
    496
        /// the pre-computed output schema of the projector.
    497
    0
        pub fn project_batch(&self, batch: &RecordBatch) -> Result<RecordBatch> {
    498
    0
            let arrays = evaluate_expressions_to_arrays(
    499
    0
                self.projection.exprs.iter().map(|p| &p.expr),
    500
    0
                batch,
    501
    0
            )?;
    502
    503
    0
            if arrays.is_empty() {
    504
    0
                let options =
    505
    0
                    RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
    506
    0
                RecordBatch::try_new_with_options(
    507
    0
                    Arc::clone(&self.output_schema),
    508
    0
                    arrays,
    509
    0
                    &options,
    510
                )
    511
    0
                .map_err(Into::into)
    512
            } else {
    513
    0
                RecordBatch::try_new(Arc::clone(&self.output_schema), arrays)
    514
    0
                    .map_err(Into::into)
    515
            }
    516
    0
        }
    517
    518
    0
        pub fn output_schema(&self) -> &SchemaRef {
    519
    0
            &self.output_schema
    520
    0
        }
    521
    522
    0
        pub fn projection(&self) -> &ProjectionExprs {
    523
    0
            &self.projection
    524
    0
        }
    525
    }
    526
    527
    impl IntoIterator for ProjectionExprs {
    528
        type Item = ProjectionExpr;
    529
        type IntoIter = std::vec::IntoIter<ProjectionExpr>;
    530
    531
    0
        fn into_iter(self) -> Self::IntoIter {
    532
    0
            self.exprs.into_iter()
    533
    0
        }
    534
    }
    535
    536
    /// The function operates in two modes:
    537
    ///
    538
    /// 1) When `sync_with_child` is `true`:
    539
    ///
    540
    ///    The function updates the indices of `expr` if the expression resides
    541
    ///    in the input plan. For instance, given the expressions `a@1 + b@2`
    542
    ///    and `c@0` with the input schema `c@2, a@0, b@1`, the expressions are
    543
    ///    updated to `a@0 + b@1` and `c@2`.
    544
    ///
    545
    /// 2) When `sync_with_child` is `false`:
    546
    ///
    547
    ///    The function determines how the expression would be updated if a projection
    548
    ///    was placed before the plan associated with the expression. If the expression
    549
    ///    cannot be rewritten after the projection, it returns `None`. For example,
    550
    ///    given the expressions `c@0`, `a@1` and `b@2`, and the projection with
    551
    ///    an output schema of `a, c_new`, then `c@0` becomes `c_new@1`, `a@1` becomes
    552
    ///    `a@0`, but `b@2` results in `None` since the projection does not include `b`.
    553
    ///
    554
    /// # Errors
    555
    /// This function returns an error if `sync_with_child` is `true` and if any expression references
    556
    /// an index that is out of bounds for `projected_exprs`.
    557
    /// For example:
    558
    ///
    559
    /// - `expr` is `a@3`
    560
    /// - `projected_exprs` is \[`a@0`, `b@1`\]
    561
    ///
    562
    /// In this case, `a@3` references index 3, which is out of bounds for `projected_exprs` (which has length 2).
    563
    0
    pub fn update_expr(
    564
    0
        expr: &Arc<dyn PhysicalExpr>,
    565
    0
        projected_exprs: &[ProjectionExpr],
    566
    0
        sync_with_child: bool,
    567
    0
    ) -> Result<Option<Arc<dyn PhysicalExpr>>> {
    568
        #[derive(Debug, PartialEq)]
    569
        enum RewriteState {
    570
            /// The expression is unchanged.
    571
            Unchanged,
    572
            /// Some part of the expression has been rewritten
    573
            RewrittenValid,
    574
            /// Some part of the expression has been rewritten, but some column
    575
            /// references could not be.
    576
            RewrittenInvalid,
    577
        }
    578
    579
    0
        let mut state = RewriteState::Unchanged;
    580
    581
    0
        let new_expr = Arc::clone(expr)
    582
    0
            .transform_up(|expr| {
    583
    0
                if state == RewriteState::RewrittenInvalid {
    584
    0
                    return Ok(Transformed::no(expr));
    585
    0
                }
    586
    587
    0
                let Some(column) = expr.as_any().downcast_ref::<Column>() else {
    588
    0
                    return Ok(Transformed::no(expr));
    589
                };
    590
    0
                if sync_with_child {
    591
    0
                    state = RewriteState::RewrittenValid;
    592
                    // Update the index of `column`:
    593
    0
                    let projected_expr = projected_exprs.get(column.index()).ok_or_else(|| {
    594
    0
                        internal_datafusion_err!(
    595
    0
                            "Column index {} out of bounds for projected expressions of length {}",
    596
    0
                            column.index(),
    597
    0
                            projected_exprs.len()
    598
                        )
    599
    0
                    })?;
    600
    0
                    Ok(Transformed::yes(Arc::clone(&projected_expr.expr)))
    601
                } else {
    602
                    // default to invalid, in case we can't find the relevant column
    603
    0
                    state = RewriteState::RewrittenInvalid;
    604
                    // Determine how to update `column` to accommodate `projected_exprs`
    605
    0
                    projected_exprs
    606
    0
                        .iter()
    607
    0
                        .enumerate()
    608
    0
                        .find_map(|(index, proj_expr)| {
    609
    0
                            proj_expr.expr.as_any().downcast_ref::<Column>().and_then(
    610
    0
                                |projected_column| {
    611
    0
                                    (column.name().eq(projected_column.name())
    612
    0
                                        && column.index() == projected_column.index())
    613
    0
                                    .then(|| {
    614
    0
                                        state = RewriteState::RewrittenValid;
    615
    0
                                        Arc::new(Column::new(&proj_expr.alias, index)) as _
    616
    0
                                    })
    617
    0
                                },
    618
                            )
    619
    0
                        })
    620
    0
                        .map_or_else(
    621
    0
                            || Ok(Transformed::no(expr)),
    622
    0
                            |c| Ok(Transformed::yes(c)),
    623
                        )
    624
                }
    625
    0
            })
    626
    0
            .data()?;
    627
    628
    0
        match state {
    629
    0
            RewriteState::RewrittenInvalid => Ok(None),
    630
            // Both Unchanged and RewrittenValid are valid:
    631
            // - Unchanged means no columns to rewrite (e.g., literals)
    632
            // - RewrittenValid means columns were successfully rewritten
    633
    0
            RewriteState::Unchanged | RewriteState::RewrittenValid => Ok(Some(new_expr)),
    634
        }
    635
    0
    }
    636
    637
    /// Stores target expressions, along with their indices, that associate with a
    638
    /// source expression in a projection mapping.
    639
    #[derive(Clone, Debug, Default)]
    640
    pub struct ProjectionTargets {
    641
        /// A non-empty vector of pairs of target expressions and their indices.
    642
        /// Consider using a special non-empty collection type in the future (e.g.
    643
        /// if Rust provides one in the standard library).
    644
        exprs_indices: Vec<(Arc<dyn PhysicalExpr>, usize)>,
    645
    }
    646
    647
    impl ProjectionTargets {
    648
        /// Returns the first target expression and its index.
    649
    0
        pub fn first(&self) -> &(Arc<dyn PhysicalExpr>, usize) {
    650
            // Since the vector is non-empty, we can safely unwrap:
    651
    0
            self.exprs_indices.first().unwrap()
    652
    0
        }
    653
    654
        /// Adds a target expression and its index to the list of targets.
    655
    0
        pub fn push(&mut self, target: (Arc<dyn PhysicalExpr>, usize)) {
    656
    0
            self.exprs_indices.push(target);
    657
    0
        }
    658
    }
    659
    660
    impl Deref for ProjectionTargets {
    661
        type Target = [(Arc<dyn PhysicalExpr>, usize)];
    662
    663
    0
        fn deref(&self) -> &Self::Target {
    664
    0
            &self.exprs_indices
    665
    0
        }
    666
    }
    667
    668
    impl From<Vec<(Arc<dyn PhysicalExpr>, usize)>> for ProjectionTargets {
    669
    0
        fn from(exprs_indices: Vec<(Arc<dyn PhysicalExpr>, usize)>) -> Self {
    670
    0
            Self { exprs_indices }
    671
    0
        }
    672
    }
    673
    674
    /// Stores the mapping between source expressions and target expressions for a
    675
    /// projection.
    676
    #[derive(Clone, Debug)]
    677
    pub struct ProjectionMapping {
    678
        /// Mapping between source expressions and target expressions.
    679
        /// Vector indices correspond to the indices after projection.
    680
        map: IndexMap<Arc<dyn PhysicalExpr>, ProjectionTargets>,
    681
    }
    682
    683
    impl ProjectionMapping {
    684
        /// Constructs the mapping between a projection's input and output
    685
        /// expressions.
    686
        ///
    687
        /// For example, given the input projection expressions (`a + b`, `c + d`)
    688
        /// and an output schema with two columns `"c + d"` and `"a + b"`, the
    689
        /// projection mapping would be:
    690
        ///
    691
        /// ```text
    692
        ///  [0]: (c + d, [(col("c + d"), 0)])
    693
        ///  [1]: (a + b, [(col("a + b"), 1)])
    694
        /// ```
    695
        ///
    696
        /// where `col("c + d")` means the column named `"c + d"`.
    697
    0
        pub fn try_new(
    698
    0
            expr: impl IntoIterator<Item = (Arc<dyn PhysicalExpr>, String)>,
    699
    0
            input_schema: &SchemaRef,
    700
    0
        ) -> Result<Self> {
    701
            // Construct a map from the input expressions to the output expression of the projection:
    702
    0
            let mut map = IndexMap::<_, ProjectionTargets>::new();
    703
    0
            for (expr_idx, (expr, name)) in expr.into_iter().enumerate() {
    704
    0
                let target_expr = Arc::new(Column::new(&name, expr_idx)) as _;
    705
    0
                let source_expr = expr.transform_down(|e| match e.as_any().downcast_ref::<Column>() {
    706
    0
                    Some(col) => {
    707
                        // Sometimes, an expression and its name in the input_schema
    708
                        // doesn't match. This can cause problems, so we make sure
    709
                        // that the expression name matches with the name in `input_schema`.
    710
                        // Conceptually, `source_expr` and `expression` should be the same.
    711
    0
                        let idx = col.index();
    712
    0
                        let matching_field = input_schema.field(idx);
    713
    0
                        let matching_name = matching_field.name();
    714
    0
                        assert_or_internal_err!(
    715
    0
                            col.name() == matching_name,
    716
    0
                            "Input field name {matching_name} does not match with the projection expression {}",
    717
    0
                            col.name()
    718
                        );
    719
    0
                        let matching_column = Column::new(matching_name, idx);
    720
    0
                        Ok(Transformed::yes(Arc::new(matching_column)))
    721
                    }
    722
    0
                    None => Ok(Transformed::no(e)),
    723
    0
                })
    724
    0
                .data()?;
    725
    0
                map.entry(source_expr)
    726
    0
                    .or_default()
    727
    0
                    .push((target_expr, expr_idx));
    728
            }
    729
    0
            Ok(Self { map })
    730
    0
        }
    731
    732
        /// Constructs a subset mapping using the provided indices.
    733
        ///
    734
        /// This is used when the output is a subset of the input without any
    735
        /// other transformations. The indices are for columns in the schema.
    736
    0
        pub fn from_indices(indices: &[usize], schema: &SchemaRef) -> Result<Self> {
    737
    0
            let projection_exprs = indices.iter().map(|index| {
    738
    0
                let field = schema.field(*index);
    739
    0
                let column = Arc::new(Column::new(field.name(), *index));
    740
    0
                (column as _, field.name().clone())
    741
    0
            });
    742
    0
            ProjectionMapping::try_new(projection_exprs, schema)
    743
    0
        }
    744
    }
    745
    746
    impl Deref for ProjectionMapping {
    747
        type Target = IndexMap<Arc<dyn PhysicalExpr>, ProjectionTargets>;
    748
    749
    0
        fn deref(&self) -> &Self::Target {
    750
    0
            &self.map
    751
    0
        }
    752
    }
    753
    754
    impl FromIterator<(Arc<dyn PhysicalExpr>, ProjectionTargets)> for ProjectionMapping {
    755
    0
        fn from_iter<T: IntoIterator<Item = (Arc<dyn PhysicalExpr>, ProjectionTargets)>>(
    756
    0
            iter: T,
    757
    0
        ) -> Self {
    758
    0
            Self {
    759
    0
                map: IndexMap::from_iter(iter),
    760
    0
            }
    761
    0
        }
    762
    }
    763
    764
    /// Projects a slice of [LexOrdering]s onto the given schema.
    765
    ///
    766
    /// This is a convenience wrapper that applies [project_ordering] to each
    767
    /// input ordering and collects the successful projections:
    768
    /// - For each input ordering, the result of [project_ordering] is appended to
    769
    ///   the output if it is `Some(...)`.
    770
    /// - Order is preserved and no deduplication is attempted.
    771
    /// - If none of the input orderings can be projected, an empty `Vec` is
    772
    ///   returned.
    773
    ///
    774
    /// See [project_ordering] for the semantics of projecting a single
    775
    /// [LexOrdering].
    776
    0
    pub fn project_orderings(
    777
    0
        orderings: &[LexOrdering],
    778
    0
        schema: &SchemaRef,
    779
    0
    ) -> Vec<LexOrdering> {
    780
    0
        let mut projected_orderings = vec![];
    781
    782
    0
        for ordering in orderings {
    783
    0
            projected_orderings.extend(project_ordering(ordering, schema));
    784
    0
        }
    785
    786
    0
        projected_orderings
    787
    0
    }
    788
    789
    /// Projects a single [LexOrdering] onto the given schema.
    790
    ///
    791
    /// This function attempts to rewrite every [PhysicalSortExpr] in the provided
    792
    /// [LexOrdering] so that any [Column] expressions point at the correct field
    793
    /// indices in `schema`.
    794
    ///
    795
    /// Key details:
    796
    /// - Columns are matched by name, not by index. The index of each matched
    797
    ///   column is looked up with [Schema::column_with_name](arrow::datatypes::Schema::column_with_name) and a new
    798
    ///   [Column] with the correct [index](Column::index) is substituted.
    799
    /// - If an expression references a column name that does not exist in
    800
    ///   `schema`, projection of the current ordering stops and only the already
    801
    ///   rewritten prefix is kept. This models the fact that a lexicographical
    802
    ///   ordering remains valid for any leading prefix whose expressions are
    803
    ///   present in the projected schema.
    804
    /// - If no expressions can be projected (i.e. the first one is missing), the
    805
    ///   function returns `None`.
    806
    ///
    807
    /// Return value:
    808
    /// - `Some(LexOrdering)` if at least one sort expression could be projected.
    809
    ///   The returned ordering may be a strict prefix of the input ordering.
    810
    /// - `None` if no part of the ordering can be projected onto `schema`.
    811
    ///
    812
    /// Example
    813
    ///
    814
    /// Suppose we have an input ordering `[col("a@0"), col("b@1")]` but the projected
    815
    /// schema only contains b and not a. The result will be `Some([col("a@0")])`. In other
    816
    /// words, the column reference is reindexed to match the projected schema.
    817
    /// If neither a nor b is present, the result will be None.
    818
    0
    pub fn project_ordering(
    819
    0
        ordering: &LexOrdering,
    820
    0
        schema: &SchemaRef,
    821
    0
    ) -> Option<LexOrdering> {
    822
    0
        let mut projected_exprs = vec![];
    823
    0
        for PhysicalSortExpr { expr, options } in ordering.iter() {
    824
    0
            let transformed = Arc::clone(expr).transform_up(|expr| {
    825
    0
                let Some(col) = expr.as_any().downcast_ref::<Column>() else {
    826
    0
                    return Ok(Transformed::no(expr));
    827
                };
    828
    829
    0
                let name = col.name();
    830
    0
                if let Some((idx, _)) = schema.column_with_name(name) {
    831
                    // Compute the new column expression (with correct index) after projection:
    832
    0
                    Ok(Transformed::yes(Arc::new(Column::new(name, idx))))
    833
                } else {
    834
                    // Cannot find expression in the projected_schema,
    835
                    // signal this using an Err result
    836
    0
                    plan_err!("")
    837
                }
    838
    0
            });
    839
    840
    0
            match transformed {
    841
    0
                Ok(transformed) => {
    842
    0
                    projected_exprs.push(PhysicalSortExpr::new(transformed.data, *options));
    843
    0
                }
    844
                Err(_) => {
    845
                    // Err result indicates an expression could not be found in the
    846
                    // projected_schema, stop iterating since rest of the orderings are violated
    847
    0
                    break;
    848
                }
    849
            }
    850
        }
    851
    852
    0
        LexOrdering::new(projected_exprs)
    853
    0
    }
    854
    855
    #[cfg(test)]
    856
    pub(crate) mod tests {
    857
        use std::collections::HashMap;
    858
    859
        use super::*;
    860
        use crate::equivalence::{convert_to_orderings, EquivalenceProperties};
    861
        use crate::expressions::{col, BinaryExpr, Literal};
    862
        use crate::utils::tests::TestScalarUDF;
    863
        use crate::{PhysicalExprRef, ScalarFunctionExpr};
    864
    865
        use arrow::compute::SortOptions;
    866
        use arrow::datatypes::{DataType, Field, Schema, TimeUnit};
    867
        use datafusion_common::config::ConfigOptions;
    868
        use datafusion_common::{ScalarValue, Statistics};
    869
        use datafusion_expr::{Operator, ScalarUDF};
    870
        use insta::assert_snapshot;
    871
    872
    0
        pub(crate) fn output_schema(
    873
    0
            mapping: &ProjectionMapping,
    874
    0
            input_schema: &Arc<Schema>,
    875
    0
        ) -> Result<SchemaRef> {
    876
            // Calculate output schema:
    877
    0
            let mut fields = vec![];
    878
    0
            for (source, targets) in mapping.iter() {
    879
    0
                let data_type = source.data_type(input_schema)?;
    880
    0
                let nullable = source.nullable(input_schema)?;
    881
    0
                for (target, _) in targets.iter() {
    882
    0
                    let Some(column) = target.as_any().downcast_ref::<Column>() else {
    883
    0
                        return plan_err!("Expects to have column");
    884
                    };
    885
    0
                    fields.push(Field::new(column.name(), data_type.clone(), nullable));
    886
                }
    887
            }
    888
    889
    0
            let output_schema = Arc::new(Schema::new_with_metadata(
    890
    0
                fields,
    891
    0
                input_schema.metadata().clone(),
    892
            ));
    893
    894
    0
            Ok(output_schema)
    895
    0
        }
    896
    897
        #[test]
    898
    0
        fn project_orderings() -> Result<()> {
    899
    0
            let schema = Arc::new(Schema::new(vec![
    900
    0
                Field::new("a", DataType::Int32, true),
    901
    0
                Field::new("b", DataType::Int32, true),
    902
    0
                Field::new("c", DataType::Int32, true),
    903
    0
                Field::new("d", DataType::Int32, true),
    904
    0
                Field::new("e", DataType::Int32, true),
    905
    0
                Field::new("ts", DataType::Timestamp(TimeUnit::Nanosecond, None), true),
    906
            ]));
    907
    0
            let col_a = &col("a", &schema)?;
    908
    0
            let col_b = &col("b", &schema)?;
    909
    0
            let col_c = &col("c", &schema)?;
    910
    0
            let col_d = &col("d", &schema)?;
    911
    0
            let col_e = &col("e", &schema)?;
    912
    0
            let col_ts = &col("ts", &schema)?;
    913
    0
            let a_plus_b = Arc::new(BinaryExpr::new(
    914
    0
                Arc::clone(col_a),
    915
    0
                Operator::Plus,
    916
    0
                Arc::clone(col_b),
    917
    0
            )) as Arc<dyn PhysicalExpr>;
    918
    0
            let b_plus_d = Arc::new(BinaryExpr::new(
    919
    0
                Arc::clone(col_b),
    920
    0
                Operator::Plus,
    921
    0
                Arc::clone(col_d),
    922
    0
            )) as Arc<dyn PhysicalExpr>;
    923
    0
            let b_plus_e = Arc::new(BinaryExpr::new(
    924
    0
                Arc::clone(col_b),
    925
    0
                Operator::Plus,
    926
    0
                Arc::clone(col_e),
    927
    0
            )) as Arc<dyn PhysicalExpr>;
    928
    0
            let c_plus_d = Arc::new(BinaryExpr::new(
    929
    0
                Arc::clone(col_c),
    930
    0
                Operator::Plus,
    931
    0
                Arc::clone(col_d),
    932
    0
            )) as Arc<dyn PhysicalExpr>;
    933
    934
    0
            let option_asc = SortOptions {
    935
    0
                descending: false,
    936
    0
                nulls_first: false,
    937
    0
            };
    938
    0
            let option_desc = SortOptions {
    939
    0
                descending: true,
    940
    0
                nulls_first: true,
    941
    0
            };
    942
    943
    0
            let test_cases = vec![
    944
                // ---------- TEST CASE 1 ------------
    945
    0
                (
    946
    0
                    // orderings
    947
    0
                    vec![
    948
    0
                        // [b ASC]
    949
    0
                        vec![(col_b, option_asc)],
    950
    0
                    ],
    951
    0
                    // projection exprs
    952
    0
                    vec![(col_b, "b_new".to_string()), (col_a, "a_new".to_string())],
    953
    0
                    // expected
    954
    0
                    vec![
    955
    0
                        // [b_new ASC]
    956
    0
                        vec![("b_new", option_asc)],
    957
    0
                    ],
    958
    0
                ),
    959
                // ---------- TEST CASE 2 ------------
    960
    0
                (
    961
    0
                    // orderings
    962
    0
                    vec![
    963
    0
                        // empty ordering
    964
    0
                    ],
    965
    0
                    // projection exprs
    966
    0
                    vec![(col_c, "c_new".to_string()), (col_b, "b_new".to_string())],
    967
    0
                    // expected
    968
    0
                    vec![
    969
    0
                        // no ordering at the output
    970
    0
                    ],
    971
    0
                ),
    972
                // ---------- TEST CASE 3 ------------
    973
    0
                (
    974
    0
                    // orderings
    975
    0
                    vec![
    976
    0
                        // [ts ASC]
    977
    0
                        vec![(col_ts, option_asc)],
    978
    0
                    ],
    979
    0
                    // projection exprs
    980
    0
                    vec![
    981
    0
                        (col_b, "b_new".to_string()),
    982
    0
                        (col_a, "a_new".to_string()),
    983
    0
                        (col_ts, "ts_new".to_string()),
    984
    0
                    ],
    985
    0
                    // expected
    986
    0
                    vec![
    987
    0
                        // [ts_new ASC]
    988
    0
                        vec![("ts_new", option_asc)],
    989
    0
                    ],
    990
    0
                ),
    991
                // ---------- TEST CASE 4 ------------
    992
    0
                (
    993
    0
                    // orderings
    994
    0
                    vec![
    995
    0
                        // [a ASC, ts ASC]
    996
    0
                        vec![(col_a, option_asc), (col_ts, option_asc)],
    997
    0
                        // [b ASC, ts ASC]
    998
    0
                        vec![(col_b, option_asc), (col_ts, option_asc)],
    999
    0
                    ],
    1000
    0
                    // projection exprs
    1001
    0
                    vec![
    1002
    0
                        (col_b, "b_new".to_string()),
    1003
    0
                        (col_a, "a_new".to_string()),
    1004
    0
                        (col_ts, "ts_new".to_string()),
    1005
    0
                    ],
    1006
    0
                    // expected
    1007
    0
                    vec![
    1008
    0
                        // [a_new ASC, ts_new ASC]
    1009
    0
                        vec![("a_new", option_asc), ("ts_new", option_asc)],
    1010
    0
                        // [b_new ASC, ts_new ASC]
    1011
    0
                        vec![("b_new", option_asc), ("ts_new", option_asc)],
    1012
    0
                    ],
    1013
    0
                ),
    1014
                // ---------- TEST CASE 5 ------------
    1015
    0
                (
    1016
    0
                    // orderings
    1017
    0
                    vec![
    1018
    0
                        // [a + b ASC]
    1019
    0
                        vec![(&a_plus_b, option_asc)],
    1020
    0
                    ],
    1021
    0
                    // projection exprs
    1022
    0
                    vec![
    1023
    0
                        (col_b, "b_new".to_string()),
    1024
    0
                        (col_a, "a_new".to_string()),
    1025
    0
                        (&a_plus_b, "a+b".to_string()),
    1026
    0
                    ],
    1027
    0
                    // expected
    1028
    0
                    vec![
    1029
    0
                        // [a + b ASC]
    1030
    0
                        vec![("a+b", option_asc)],
    1031
    0
                    ],
    1032
    0
                ),
    1033
                // ---------- TEST CASE 6 ------------
    1034
    0
                (
    1035
    0
                    // orderings
    1036
    0
                    vec![
    1037
    0
                        // [a + b ASC, c ASC]
    1038
    0
                        vec![(&a_plus_b, option_asc), (col_c, option_asc)],
    1039
    0
                    ],
    1040
    0
                    // projection exprs
    1041
    0
                    vec![
    1042
    0
                        (col_b, "b_new".to_string()),
    1043
    0
                        (col_a, "a_new".to_string()),
    1044
    0
                        (col_c, "c_new".to_string()),
    1045
    0
                        (&a_plus_b, "a+b".to_string()),
    1046
    0
                    ],
    1047
    0
                    // expected
    1048
    0
                    vec![
    1049
    0
                        // [a + b ASC, c_new ASC]
    1050
    0
                        vec![("a+b", option_asc), ("c_new", option_asc)],
    1051
    0
                    ],
    1052
    0
                ),
    1053
                // ------- TEST CASE 7 ----------
    1054
    0
                (
    1055
    0
                    vec![
    1056
    0
                        // [a ASC, b ASC, c ASC]
    1057
    0
                        vec![(col_a, option_asc), (col_b, option_asc)],
    1058
    0
                        // [a ASC, d ASC]
    1059
    0
                        vec![(col_a, option_asc), (col_d, option_asc)],
    1060
    0
                    ],
    1061
    0
                    // b as b_new, a as a_new, d as d_new b+d
    1062
    0
                    vec![
    1063
    0
                        (col_b, "b_new".to_string()),
    1064
    0
                        (col_a, "a_new".to_string()),
    1065
    0
                        (col_d, "d_new".to_string()),
    1066
    0
                        (&b_plus_d, "b+d".to_string()),
    1067
    0
                    ],
    1068
    0
                    // expected
    1069
    0
                    vec![
    1070
    0
                        // [a_new ASC, b_new ASC]
    1071
    0
                        vec![("a_new", option_asc), ("b_new", option_asc)],
    1072
    0
                        // [a_new ASC, d_new ASC]
    1073
    0
                        vec![("a_new", option_asc), ("d_new", option_asc)],
    1074
    0
                        // [a_new ASC, b+d ASC]
    1075
    0
                        vec![("a_new", option_asc), ("b+d", option_asc)],
    1076
    0
                    ],
    1077
    0
                ),
    1078
                // ------- TEST CASE 8 ----------
    1079
    0
                (
    1080
    0
                    // orderings
    1081
    0
                    vec![
    1082
    0
                        // [b+d ASC]
    1083
    0
                        vec![(&b_plus_d, option_asc)],
    1084
    0
                    ],
    1085
    0
                    // proj exprs
    1086
    0
                    vec![
    1087
    0
                        (col_b, "b_new".to_string()),
    1088
    0
                        (col_a, "a_new".to_string()),
    1089
    0
                        (col_d, "d_new".to_string()),
    1090
    0
                        (&b_plus_d, "b+d".to_string()),
    1091
    0
                    ],
    1092
    0
                    // expected
    1093
    0
                    vec![
    1094
    0
                        // [b+d ASC]
    1095
    0
                        vec![("b+d", option_asc)],
    1096
    0
                    ],
    1097
    0
                ),
    1098
                // ------- TEST CASE 9 ----------
    1099
    0
                (
    1100
    0
                    // orderings
    1101
    0
                    vec![
    1102
    0
                        // [a ASC, d ASC, b ASC]
    1103
    0
                        vec![
    1104
    0
                            (col_a, option_asc),
    1105
    0
                            (col_d, option_asc),
    1106
    0
                            (col_b, option_asc),
    1107
    0
                        ],
    1108
    0
                        // [c ASC]
    1109
    0
                        vec![(col_c, option_asc)],
    1110
    0
                    ],
    1111
    0
                    // proj exprs
    1112
    0
                    vec![
    1113
    0
                        (col_b, "b_new".to_string()),
    1114
    0
                        (col_a, "a_new".to_string()),
    1115
    0
                        (col_d, "d_new".to_string()),
    1116
    0
                        (col_c, "c_new".to_string()),
    1117
    0
                    ],
    1118
    0
                    // expected
    1119
    0
                    vec![
    1120
    0
                        // [a_new ASC, d_new ASC, b_new ASC]
    1121
    0
                        vec![
    1122
    0
                            ("a_new", option_asc),
    1123
    0
                            ("d_new", option_asc),
    1124
    0
                            ("b_new", option_asc),
    1125
    0
                        ],
    1126
    0
                        // [c_new ASC],
    1127
    0
                        vec![("c_new", option_asc)],
    1128
    0
                    ],
    1129
    0
                ),
    1130
                // ------- TEST CASE 10 ----------
    1131
    0
                (
    1132
    0
                    vec![
    1133
    0
                        // [a ASC, b ASC, c ASC]
    1134
    0
                        vec![
    1135
    0
                            (col_a, option_asc),
    1136
    0
                            (col_b, option_asc),
    1137
    0
                            (col_c, option_asc),
    1138
    0
                        ],
    1139
    0
                        // [a ASC, d ASC]
    1140
    0
                        vec![(col_a, option_asc), (col_d, option_asc)],
    1141
    0
                    ],
    1142
    0
                    // proj exprs
    1143
    0
                    vec![
    1144
    0
                        (col_b, "b_new".to_string()),
    1145
    0
                        (col_a, "a_new".to_string()),
    1146
    0
                        (col_c, "c_new".to_string()),
    1147
    0
                        (&c_plus_d, "c+d".to_string()),
    1148
    0
                    ],
    1149
    0
                    // expected
    1150
    0
                    vec![
    1151
    0
                        // [a_new ASC, b_new ASC, c_new ASC]
    1152
    0
                        vec![
    1153
    0
                            ("a_new", option_asc),
    1154
    0
                            ("b_new", option_asc),
    1155
    0
                            ("c_new", option_asc),
    1156
    0
                        ],
    1157
    0
                        // [a_new ASC, b_new ASC, c+d ASC]
    1158
    0
                        vec![
    1159
    0
                            ("a_new", option_asc),
    1160
    0
                            ("b_new", option_asc),
    1161
    0
                            ("c+d", option_asc),
    1162
    0
                        ],
    1163
    0
                    ],
    1164
    0
                ),
    1165
                // ------- TEST CASE 11 ----------
    1166
    0
                (
    1167
    0
                    // orderings
    1168
    0
                    vec![
    1169
    0
                        // [a ASC, b ASC]
    1170
    0
                        vec![(col_a, option_asc), (col_b, option_asc)],
    1171
    0
                        // [a ASC, d ASC]
    1172
    0
                        vec![(col_a, option_asc), (col_d, option_asc)],
    1173
    0
                    ],
    1174
    0
                    // proj exprs
    1175
    0
                    vec![
    1176
    0
                        (col_b, "b_new".to_string()),
    1177
    0
                        (col_a, "a_new".to_string()),
    1178
    0
                        (&b_plus_d, "b+d".to_string()),
    1179
    0
                    ],
    1180
    0
                    // expected
    1181
    0
                    vec![
    1182
    0
                        // [a_new ASC, b_new ASC]
    1183
    0
                        vec![("a_new", option_asc), ("b_new", option_asc)],
    1184
    0
                        // [a_new ASC, b + d ASC]
    1185
    0
                        vec![("a_new", option_asc), ("b+d", option_asc)],
    1186
    0
                    ],
    1187
    0
                ),
    1188
                // ------- TEST CASE 12 ----------
    1189
    0
                (
    1190
    0
                    // orderings
    1191
    0
                    vec![
    1192
    0
                        // [a ASC, b ASC, c ASC]
    1193
    0
                        vec![
    1194
    0
                            (col_a, option_asc),
    1195
    0
                            (col_b, option_asc),
    1196
    0
                            (col_c, option_asc),
    1197
    0
                        ],
    1198
    0
                    ],
    1199
    0
                    // proj exprs
    1200
    0
                    vec![(col_c, "c_new".to_string()), (col_a, "a_new".to_string())],
    1201
    0
                    // expected
    1202
    0
                    vec![
    1203
    0
                        // [a_new ASC]
    1204
    0
                        vec![("a_new", option_asc)],
    1205
    0
                    ],
    1206
    0
                ),
    1207
                // ------- TEST CASE 13 ----------
    1208
    0
                (
    1209
    0
                    // orderings
    1210
    0
                    vec![
    1211
    0
                        // [a ASC, b ASC, c ASC]
    1212
    0
                        vec![
    1213
    0
                            (col_a, option_asc),
    1214
    0
                            (col_b, option_asc),
    1215
    0
                            (col_c, option_asc),
    1216
    0
                        ],
    1217
    0
                        // [a ASC, a + b ASC, c ASC]
    1218
    0
                        vec![
    1219
    0
                            (col_a, option_asc),
    1220
    0
                            (&a_plus_b, option_asc),
    1221
    0
                            (col_c, option_asc),
    1222
    0
                        ],
    1223
    0
                    ],
    1224
    0
                    // proj exprs
    1225
    0
                    vec![
    1226
    0
                        (col_c, "c_new".to_string()),
    1227
    0
                        (col_b, "b_new".to_string()),
    1228
    0
                        (col_a, "a_new".to_string()),
    1229
    0
                        (&a_plus_b, "a+b".to_string()),
    1230
    0
                    ],
    1231
    0
                    // expected
    1232
    0
                    vec![
    1233
    0
                        // [a_new ASC, b_new ASC, c_new ASC]
    1234
    0
                        vec![
    1235
    0
                            ("a_new", option_asc),
    1236
    0
                            ("b_new", option_asc),
    1237
    0
                            ("c_new", option_asc),
    1238
    0
                        ],
    1239
    0
                        // [a_new ASC, a+b ASC, c_new ASC]
    1240
    0
                        vec![
    1241
    0
                            ("a_new", option_asc),
    1242
    0
                            ("a+b", option_asc),
    1243
    0
                            ("c_new", option_asc),
    1244
    0
                        ],
    1245
    0
                    ],
    1246
    0
                ),
    1247
                // ------- TEST CASE 14 ----------
    1248
    0
                (
    1249
    0
                    // orderings
    1250
    0
                    vec![
    1251
    0
                        // [a ASC, b ASC]
    1252
    0
                        vec![(col_a, option_asc), (col_b, option_asc)],
    1253
    0
                        // [c ASC, b ASC]
    1254
    0
                        vec![(col_c, option_asc), (col_b, option_asc)],
    1255
    0
                        // [d ASC, e ASC]
    1256
    0
                        vec![(col_d, option_asc), (col_e, option_asc)],
    1257
    0
                    ],
    1258
    0
                    // proj exprs
    1259
    0
                    vec![
    1260
    0
                        (col_c, "c_new".to_string()),
    1261
    0
                        (col_d, "d_new".to_string()),
    1262
    0
                        (col_a, "a_new".to_string()),
    1263
    0
                        (&b_plus_e, "b+e".to_string()),
    1264
    0
                    ],
    1265
    0
                    // expected
    1266
    0
                    vec![
    1267
    0
                        // [a_new ASC, d_new ASC, b+e ASC]
    1268
    0
                        vec![
    1269
    0
                            ("a_new", option_asc),
    1270
    0
                            ("d_new", option_asc),
    1271
    0
                            ("b+e", option_asc),
    1272
    0
                        ],
    1273
    0
                        // [d_new ASC, a_new ASC, b+e ASC]
    1274
    0
                        vec![
    1275
    0
                            ("d_new", option_asc),
    1276
    0
                            ("a_new", option_asc),
    1277
    0
                            ("b+e", option_asc),
    1278
    0
                        ],
    1279
    0
                        // [c_new ASC, d_new ASC, b+e ASC]
    1280
    0
                        vec![
    1281
    0
                            ("c_new", option_asc),
    1282
    0
                            ("d_new", option_asc),
    1283
    0
                            ("b+e", option_asc),
    1284
    0
                        ],
    1285
    0
                        // [d_new ASC, c_new ASC, b+e ASC]
    1286
    0
                        vec![
    1287
    0
                            ("d_new", option_asc),
    1288
    0
                            ("c_new", option_asc),
    1289
    0
                            ("b+e", option_asc),
    1290
    0
                        ],
    1291
    0
                    ],
    1292
    0
                ),
    1293
                // ------- TEST CASE 15 ----------
    1294
    0
                (
    1295
    0
                    // orderings
    1296
    0
                    vec![
    1297
    0
                        // [a ASC, c ASC, b ASC]
    1298
    0
                        vec![
    1299
    0
                            (col_a, option_asc),
    1300
    0
                            (col_c, option_asc),
    1301
    0
                            (col_b, option_asc),
    1302
    0
                        ],
    1303
    0
                    ],
    1304
    0
                    // proj exprs
    1305
    0
                    vec![
    1306
    0
                        (col_c, "c_new".to_string()),
    1307
    0
                        (col_a, "a_new".to_string()),
    1308
    0
                        (&a_plus_b, "a+b".to_string()),
    1309
    0
                    ],
    1310
    0
                    // expected
    1311
    0
                    vec![
    1312
    0
                        // [a_new ASC, d_new ASC, b+e ASC]
    1313
    0
                        vec![
    1314
    0
                            ("a_new", option_asc),
    1315
    0
                            ("c_new", option_asc),
    1316
    0
                            ("a+b", option_asc),
    1317
    0
                        ],
    1318
    0
                    ],
    1319
    0
                ),
    1320
                // ------- TEST CASE 16 ----------
    1321
    0
                (
    1322
    0
                    // orderings
    1323
    0
                    vec![
    1324
    0
                        // [a ASC, b ASC]
    1325
    0
                        vec![(col_a, option_asc), (col_b, option_asc)],
    1326
    0
                        // [c ASC, b DESC]
    1327
    0
                        vec![(col_c, option_asc), (col_b, option_desc)],
    1328
    0
                        // [e ASC]
    1329
    0
                        vec![(col_e, option_asc)],
    1330
    0
                    ],
    1331
    0
                    // proj exprs
    1332
    0
                    vec![
    1333
    0
                        (col_c, "c_new".to_string()),
    1334
    0
                        (col_a, "a_new".to_string()),
    1335
    0
                        (col_b, "b_new".to_string()),
    1336
    0
                        (&b_plus_e, "b+e".to_string()),
    1337
    0
                    ],
    1338
    0
                    // expected
    1339
    0
                    vec![
    1340
    0
                        // [a_new ASC, b_new ASC]
    1341
    0
                        vec![("a_new", option_asc), ("b_new", option_asc)],
    1342
    0
                        // [a_new ASC, b_new ASC]
    1343
    0
                        vec![("a_new", option_asc), ("b+e", option_asc)],
    1344
    0
                        // [c_new ASC, b_new DESC]
    1345
    0
                        vec![("c_new", option_asc), ("b_new", option_desc)],
    1346
    0
                    ],
    1347
    0
                ),
    1348
            ];
    1349
    1350
    0
            for (idx, (orderings, proj_exprs, expected)) in test_cases.into_iter().enumerate()
    1351
            {
    1352
    0
                let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema));
    1353
    1354
    0
                let orderings = convert_to_orderings(&orderings);
    1355
    0
                eq_properties.add_orderings(orderings);
    1356
    1357
    0
                let proj_exprs = proj_exprs
    1358
    0
                    .into_iter()
    1359
    0
                    .map(|(expr, name)| (Arc::clone(expr), name));
    1360
    0
                let projection_mapping = ProjectionMapping::try_new(proj_exprs, &schema)?;
    1361
    0
                let output_schema = output_schema(&projection_mapping, &schema)?;
    1362
    1363
    0
                let expected = expected
    1364
    0
                    .into_iter()
    1365
    0
                    .map(|ordering| {
    1366
    0
                        ordering
    1367
    0
                            .into_iter()
    1368
    0
                            .map(|(name, options)| {
    1369
    0
                                (col(name, &output_schema).unwrap(), options)
    1370
    0
                            })
    1371
    0
                            .collect::<Vec<_>>()
    1372
    0
                    })
    1373
    0
                    .collect::<Vec<_>>();
    1374
    0
                let expected = convert_to_orderings(&expected);
    1375
    1376
    0
                let projected_eq = eq_properties.project(&projection_mapping, output_schema);
    1377
    0
                let orderings = projected_eq.oeq_class();
    1378
    1379
    0
                let err_msg = format!(
    1380
    0
                    "test_idx: {idx:?}, actual: {orderings:?}, expected: {expected:?}, projection_mapping: {projection_mapping:?}"
    1381
                );
    1382
    1383
    0
                assert_eq!(orderings.len(), expected.len(), "{err_msg}");
    1384
    0
                for expected_ordering in &expected {
    1385
    0
                    assert!(orderings.contains(expected_ordering), "{}", err_msg)
    1386
                }
    1387
            }
    1388
    1389
    0
            Ok(())
    1390
    0
        }
    1391
    1392
        #[test]
    1393
    0
        fn project_orderings2() -> Result<()> {
    1394
    0
            let schema = Arc::new(Schema::new(vec![
    1395
    0
                Field::new("a", DataType::Int32, true),
    1396
    0
                Field::new("b", DataType::Int32, true),
    1397
    0
                Field::new("c", DataType::Int32, true),
    1398
    0
                Field::new("d", DataType::Int32, true),
    1399
    0
                Field::new("ts", DataType::Timestamp(TimeUnit::Nanosecond, None), true),
    1400
            ]));
    1401
    0
            let col_a = &col("a", &schema)?;
    1402
    0
            let col_b = &col("b", &schema)?;
    1403
    0
            let col_c = &col("c", &schema)?;
    1404
    0
            let col_ts = &col("ts", &schema)?;
    1405
    0
            let a_plus_b = Arc::new(BinaryExpr::new(
    1406
    0
                Arc::clone(col_a),
    1407
    0
                Operator::Plus,
    1408
    0
                Arc::clone(col_b),
    1409
    0
            )) as Arc<dyn PhysicalExpr>;
    1410
    1411
    0
            let test_fun = Arc::new(ScalarUDF::new_from_impl(TestScalarUDF::new()));
    1412
    1413
    0
            let round_c = Arc::new(ScalarFunctionExpr::try_new(
    1414
    0
                test_fun,
    1415
    0
                vec![Arc::clone(col_c)],
    1416
    0
                &schema,
    1417
    0
                Arc::new(ConfigOptions::default()),
    1418
    0
            )?) as PhysicalExprRef;
    1419
    1420
    0
            let option_asc = SortOptions {
    1421
    0
                descending: false,
    1422
    0
                nulls_first: false,
    1423
    0
            };
    1424
    1425
    0
            let proj_exprs = vec![
    1426
    0
                (col_b, "b_new".to_string()),
    1427
    0
                (col_a, "a_new".to_string()),
    1428
    0
                (col_c, "c_new".to_string()),
    1429
    0
                (&round_c, "round_c_res".to_string()),
    1430
            ];
    1431
    0
            let proj_exprs = proj_exprs
    1432
    0
                .into_iter()
    1433
    0
                .map(|(expr, name)| (Arc::clone(expr), name));
    1434
    0
            let projection_mapping = ProjectionMapping::try_new(proj_exprs, &schema)?;
    1435
    0
            let output_schema = output_schema(&projection_mapping, &schema)?;
    1436
    1437
    0
            let col_a_new = &col("a_new", &output_schema)?;
    1438
    0
            let col_b_new = &col("b_new", &output_schema)?;
    1439
    0
            let col_c_new = &col("c_new", &output_schema)?;
    1440
    0
            let col_round_c_res = &col("round_c_res", &output_schema)?;
    1441
    0
            let a_new_plus_b_new = Arc::new(BinaryExpr::new(
    1442
    0
                Arc::clone(col_a_new),
    1443
    0
                Operator::Plus,
    1444
    0
                Arc::clone(col_b_new),
    1445
    0
            )) as Arc<dyn PhysicalExpr>;
    1446
    1447
    0
            let test_cases = [
    1448
    0
                // ---------- TEST CASE 1 ------------
    1449
    0
                (
    1450
    0
                    // orderings
    1451
    0
                    vec![
    1452
    0
                        // [a ASC]
    1453
    0
                        vec![(col_a, option_asc)],
    1454
    0
                    ],
    1455
    0
                    // expected
    1456
    0
                    vec![
    1457
    0
                        // [b_new ASC]
    1458
    0
                        vec![(col_a_new, option_asc)],
    1459
    0
                    ],
    1460
    0
                ),
    1461
    0
                // ---------- TEST CASE 2 ------------
    1462
    0
                (
    1463
    0
                    // orderings
    1464
    0
                    vec![
    1465
    0
                        // [a+b ASC]
    1466
    0
                        vec![(&a_plus_b, option_asc)],
    1467
    0
                    ],
    1468
    0
                    // expected
    1469
    0
                    vec![
    1470
    0
                        // [b_new ASC]
    1471
    0
                        vec![(&a_new_plus_b_new, option_asc)],
    1472
    0
                    ],
    1473
    0
                ),
    1474
    0
                // ---------- TEST CASE 3 ------------
    1475
    0
                (
    1476
    0
                    // orderings
    1477
    0
                    vec![
    1478
    0
                        // [a ASC, ts ASC]
    1479
    0
                        vec![(col_a, option_asc), (col_ts, option_asc)],
    1480
    0
                    ],
    1481
    0
                    // expected
    1482
    0
                    vec![
    1483
    0
                        // [a_new ASC, date_bin_res ASC]
    1484
    0
                        vec![(col_a_new, option_asc)],
    1485
    0
                    ],
    1486
    0
                ),
    1487
    0
                // ---------- TEST CASE 4 ------------
    1488
    0
                (
    1489
    0
                    // orderings
    1490
    0
                    vec![
    1491
    0
                        // [a ASC, ts ASC, b ASC]
    1492
    0
                        vec![
    1493
    0
                            (col_a, option_asc),
    1494
    0
                            (col_ts, option_asc),
    1495
    0
                            (col_b, option_asc),
    1496
    0
                        ],
    1497
    0
                    ],
    1498
    0
                    // expected
    1499
    0
                    vec![
    1500
    0
                        // [a_new ASC, date_bin_res ASC]
    1501
    0
                        vec![(col_a_new, option_asc)],
    1502
    0
                    ],
    1503
    0
                ),
    1504
    0
                // ---------- TEST CASE 5 ------------
    1505
    0
                (
    1506
    0
                    // orderings
    1507
    0
                    vec![
    1508
    0
                        // [a ASC, c ASC]
    1509
    0
                        vec![(col_a, option_asc), (col_c, option_asc)],
    1510
    0
                    ],
    1511
    0
                    // expected
    1512
    0
                    vec![
    1513
    0
                        // [a_new ASC, round_c_res ASC, c_new ASC]
    1514
    0
                        vec![(col_a_new, option_asc), (col_round_c_res, option_asc)],
    1515
    0
                        // [a_new ASC, c_new ASC]
    1516
    0
                        vec![(col_a_new, option_asc), (col_c_new, option_asc)],
    1517
    0
                    ],
    1518
    0
                ),
    1519
    0
                // ---------- TEST CASE 6 ------------
    1520
    0
                (
    1521
    0
                    // orderings
    1522
    0
                    vec![
    1523
    0
                        // [c ASC, b ASC]
    1524
    0
                        vec![(col_c, option_asc), (col_b, option_asc)],
    1525
    0
                    ],
    1526
    0
                    // expected
    1527
    0
                    vec![
    1528
    0
                        // [round_c_res ASC]
    1529
    0
                        vec![(col_round_c_res, option_asc)],
    1530
    0
                        // [c_new ASC, b_new ASC]
    1531
    0
                        vec![(col_c_new, option_asc), (col_b_new, option_asc)],
    1532
    0
                    ],
    1533
    0
                ),
    1534
    0
                // ---------- TEST CASE 7 ------------
    1535
    0
                (
    1536
    0
                    // orderings
    1537
    0
                    vec![
    1538
    0
                        // [a+b ASC, c ASC]
    1539
    0
                        vec![(&a_plus_b, option_asc), (col_c, option_asc)],
    1540
    0
                    ],
    1541
    0
                    // expected
    1542
    0
                    vec![
    1543
    0
                        // [a+b ASC, round(c) ASC, c_new ASC]
    1544
    0
                        vec![
    1545
    0
                            (&a_new_plus_b_new, option_asc),
    1546
    0
                            (col_round_c_res, option_asc),
    1547
    0
                        ],
    1548
    0
                        // [a+b ASC, c_new ASC]
    1549
    0
                        vec![(&a_new_plus_b_new, option_asc), (col_c_new, option_asc)],
    1550
    0
                    ],
    1551
    0
                ),
    1552
    0
            ];
    1553
    1554
    0
            for (idx, (orderings, expected)) in test_cases.iter().enumerate() {
    1555
    0
                let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema));
    1556
    1557
    0
                let orderings = convert_to_orderings(orderings);
    1558
    0
                eq_properties.add_orderings(orderings);
    1559
    1560
    0
                let expected = convert_to_orderings(expected);
    1561
    1562
    0
                let projected_eq =
    1563
    0
                    eq_properties.project(&projection_mapping, Arc::clone(&output_schema));
    1564
    0
                let orderings = projected_eq.oeq_class();
    1565
    1566
    0
                let err_msg = format!(
    1567
    0
                    "test idx: {idx:?}, actual: {orderings:?}, expected: {expected:?}, projection_mapping: {projection_mapping:?}"
    1568
                );
    1569
    1570
    0
                assert_eq!(orderings.len(), expected.len(), "{err_msg}");
    1571
    0
                for expected_ordering in &expected {
    1572
    0
                    assert!(orderings.contains(expected_ordering), "{}", err_msg)
    1573
                }
    1574
            }
    1575
    0
            Ok(())
    1576
    0
        }
    1577
    1578
        #[test]
    1579
    0
        fn project_orderings3() -> Result<()> {
    1580
    0
            let schema = Arc::new(Schema::new(vec![
    1581
    0
                Field::new("a", DataType::Int32, true),
    1582
    0
                Field::new("b", DataType::Int32, true),
    1583
    0
                Field::new("c", DataType::Int32, true),
    1584
    0
                Field::new("d", DataType::Int32, true),
    1585
    0
                Field::new("e", DataType::Int32, true),
    1586
    0
                Field::new("f", DataType::Int32, true),
    1587
            ]));
    1588
    0
            let col_a = &col("a", &schema)?;
    1589
    0
            let col_b = &col("b", &schema)?;
    1590
    0
            let col_c = &col("c", &schema)?;
    1591
    0
            let col_d = &col("d", &schema)?;
    1592
    0
            let col_e = &col("e", &schema)?;
    1593
    0
            let col_f = &col("f", &schema)?;
    1594
    0
            let a_plus_b = Arc::new(BinaryExpr::new(
    1595
    0
                Arc::clone(col_a),
    1596
    0
                Operator::Plus,
    1597
    0
                Arc::clone(col_b),
    1598
    0
            )) as Arc<dyn PhysicalExpr>;
    1599
    1600
    0
            let option_asc = SortOptions {
    1601
    0
                descending: false,
    1602
    0
                nulls_first: false,
    1603
    0
            };
    1604
    1605
    0
            let proj_exprs = vec![
    1606
    0
                (col_c, "c_new".to_string()),
    1607
    0
                (col_d, "d_new".to_string()),
    1608
    0
                (&a_plus_b, "a+b".to_string()),
    1609
            ];
    1610
    0
            let proj_exprs = proj_exprs
    1611
    0
                .into_iter()
    1612
    0
                .map(|(expr, name)| (Arc::clone(expr), name));
    1613
    0
            let projection_mapping = ProjectionMapping::try_new(proj_exprs, &schema)?;
    1614
    0
            let output_schema = output_schema(&projection_mapping, &schema)?;
    1615
    1616
    0
            let col_a_plus_b_new = &col("a+b", &output_schema)?;
    1617
    0
            let col_c_new = &col("c_new", &output_schema)?;
    1618
    0
            let col_d_new = &col("d_new", &output_schema)?;
    1619
    1620
    0
            let test_cases = vec![
    1621
                // ---------- TEST CASE 1 ------------
    1622
    0
                (
    1623
    0
                    // orderings
    1624
    0
                    vec![
    1625
    0
                        // [d ASC, b ASC]
    1626
    0
                        vec![(col_d, option_asc), (col_b, option_asc)],
    1627
    0
                        // [c ASC, a ASC]
    1628
    0
                        vec![(col_c, option_asc), (col_a, option_asc)],
    1629
    0
                    ],
    1630
    0
                    // equal conditions
    1631
    0
                    vec![],
    1632
    0
                    // expected
    1633
    0
                    vec![
    1634
    0
                        // [d_new ASC, c_new ASC, a+b ASC]
    1635
    0
                        vec![
    1636
    0
                            (col_d_new, option_asc),
    1637
    0
                            (col_c_new, option_asc),
    1638
    0
                            (col_a_plus_b_new, option_asc),
    1639
    0
                        ],
    1640
    0
                        // [c_new ASC, d_new ASC, a+b ASC]
    1641
    0
                        vec![
    1642
    0
                            (col_c_new, option_asc),
    1643
    0
                            (col_d_new, option_asc),
    1644
    0
                            (col_a_plus_b_new, option_asc),
    1645
    0
                        ],
    1646
    0
                    ],
    1647
    0
                ),
    1648
                // ---------- TEST CASE 2 ------------
    1649
    0
                (
    1650
    0
                    // orderings
    1651
    0
                    vec![
    1652
    0
                        // [d ASC, b ASC]
    1653
    0
                        vec![(col_d, option_asc), (col_b, option_asc)],
    1654
    0
                        // [c ASC, e ASC], Please note that a=e
    1655
    0
                        vec![(col_c, option_asc), (col_e, option_asc)],
    1656
    0
                    ],
    1657
    0
                    // equal conditions
    1658
    0
                    vec![(col_e, col_a)],
    1659
    0
                    // expected
    1660
    0
                    vec![
    1661
    0
                        // [d_new ASC, c_new ASC, a+b ASC]
    1662
    0
                        vec![
    1663
    0
                            (col_d_new, option_asc),
    1664
    0
                            (col_c_new, option_asc),
    1665
    0
                            (col_a_plus_b_new, option_asc),
    1666
    0
                        ],
    1667
    0
                        // [c_new ASC, d_new ASC, a+b ASC]
    1668
    0
                        vec![
    1669
    0
                            (col_c_new, option_asc),
    1670
    0
                            (col_d_new, option_asc),
    1671
    0
                            (col_a_plus_b_new, option_asc),
    1672
    0
                        ],
    1673
    0
                    ],
    1674
    0
                ),
    1675
                // ---------- TEST CASE 3 ------------
    1676
    0
                (
    1677
    0
                    // orderings
    1678
    0
                    vec![
    1679
    0
                        // [d ASC, b ASC]
    1680
    0
                        vec![(col_d, option_asc), (col_b, option_asc)],
    1681
    0
                        // [c ASC, e ASC], Please note that a=f
    1682
    0
                        vec![(col_c, option_asc), (col_e, option_asc)],
    1683
    0
                    ],
    1684
    0
                    // equal conditions
    1685
    0
                    vec![(col_a, col_f)],
    1686
    0
                    // expected
    1687
    0
                    vec![
    1688
    0
                        // [d_new ASC]
    1689
    0
                        vec![(col_d_new, option_asc)],
    1690
    0
                        // [c_new ASC]
    1691
    0
                        vec![(col_c_new, option_asc)],
    1692
    0
                    ],
    1693
    0
                ),
    1694
            ];
    1695
    0
            for (orderings, equal_columns, expected) in test_cases {
    1696
    0
                let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema));
    1697
    0
                for (lhs, rhs) in equal_columns {
    1698
    0
                    eq_properties.add_equal_conditions(Arc::clone(lhs), Arc::clone(rhs))?;
    1699
                }
    1700
    1701
    0
                let orderings = convert_to_orderings(&orderings);
    1702
    0
                eq_properties.add_orderings(orderings);
    1703
    1704
    0
                let expected = convert_to_orderings(&expected);
    1705
    1706
    0
                let projected_eq =
    1707
    0
                    eq_properties.project(&projection_mapping, Arc::clone(&output_schema));
    1708
    0
                let orderings = projected_eq.oeq_class();
    1709
    1710
    0
                let err_msg = format!(
    1711
    0
                    "actual: {orderings:?}, expected: {expected:?}, projection_mapping: {projection_mapping:?}"
    1712
                );
    1713
    1714
    0
                assert_eq!(orderings.len(), expected.len(), "{err_msg}");
    1715
    0
                for expected_ordering in &expected {
    1716
    0
                    assert!(orderings.contains(expected_ordering), "{}", err_msg)
    1717
                }
    1718
            }
    1719
    1720
    0
            Ok(())
    1721
    0
        }
    1722
    1723
    0
        fn get_stats() -> Statistics {
    1724
    0
            Statistics {
    1725
    0
                num_rows: Precision::Exact(5),
    1726
    0
                total_byte_size: Precision::Exact(23),
    1727
    0
                column_statistics: vec![
    1728
    0
                    ColumnStatistics {
    1729
    0
                        distinct_count: Precision::Exact(5),
    1730
    0
                        max_value: Precision::Exact(ScalarValue::Int64(Some(21))),
    1731
    0
                        min_value: Precision::Exact(ScalarValue::Int64(Some(-4))),
    1732
    0
                        sum_value: Precision::Exact(ScalarValue::Int64(Some(42))),
    1733
    0
                        null_count: Precision::Exact(0),
    1734
    0
                    },
    1735
    0
                    ColumnStatistics {
    1736
    0
                        distinct_count: Precision::Exact(1),
    1737
    0
                        max_value: Precision::Exact(ScalarValue::from("x")),
    1738
    0
                        min_value: Precision::Exact(ScalarValue::from("a")),
    1739
    0
                        sum_value: Precision::Absent,
    1740
    0
                        null_count: Precision::Exact(3),
    1741
    0
                    },
    1742
    0
                    ColumnStatistics {
    1743
    0
                        distinct_count: Precision::Absent,
    1744
    0
                        max_value: Precision::Exact(ScalarValue::Float32(Some(1.1))),
    1745
    0
                        min_value: Precision::Exact(ScalarValue::Float32(Some(0.1))),
    1746
    0
                        sum_value: Precision::Exact(ScalarValue::Float32(Some(5.5))),
    1747
    0
                        null_count: Precision::Absent,
    1748
    0
                    },
    1749
    0
                ],
    1750
    0
            }
    1751
    0
        }
    1752
    1753
    0
        fn get_schema() -> Schema {
    1754
    0
            let field_0 = Field::new("col0", DataType::Int64, false);
    1755
    0
            let field_1 = Field::new("col1", DataType::Utf8, false);
    1756
    0
            let field_2 = Field::new("col2", DataType::Float32, false);
    1757
    0
            Schema::new(vec![field_0, field_1, field_2])
    1758
    0
        }
    1759
    1760
        #[test]
    1761
    0
        fn test_stats_projection_columns_only() {
    1762
    0
            let source = get_stats();
    1763
    0
            let schema = get_schema();
    1764
    1765
    0
            let projection = ProjectionExprs::new(vec![
    1766
    0
                ProjectionExpr {
    1767
    0
                    expr: Arc::new(Column::new("col1", 1)),
    1768
    0
                    alias: "col1".to_string(),
    1769
    0
                },
    1770
    0
                ProjectionExpr {
    1771
    0
                    expr: Arc::new(Column::new("col0", 0)),
    1772
    0
                    alias: "col0".to_string(),
    1773
    0
                },
    1774
            ]);
    1775
    1776
    0
            let result = projection.project_statistics(source, &schema).unwrap();
    1777
    1778
    0
            let expected = Statistics {
    1779
    0
                num_rows: Precision::Exact(5),
    1780
    0
                total_byte_size: Precision::Exact(23),
    1781
    0
                column_statistics: vec![
    1782
    0
                    ColumnStatistics {
    1783
    0
                        distinct_count: Precision::Exact(1),
    1784
    0
                        max_value: Precision::Exact(ScalarValue::from("x")),
    1785
    0
                        min_value: Precision::Exact(ScalarValue::from("a")),
    1786
    0
                        sum_value: Precision::Absent,
    1787
    0
                        null_count: Precision::Exact(3),
    1788
    0
                    },
    1789
    0
                    ColumnStatistics {
    1790
    0
                        distinct_count: Precision::Exact(5),
    1791
    0
                        max_value: Precision::Exact(ScalarValue::Int64(Some(21))),
    1792
    0
                        min_value: Precision::Exact(ScalarValue::Int64(Some(-4))),
    1793
    0
                        sum_value: Precision::Exact(ScalarValue::Int64(Some(42))),
    1794
    0
                        null_count: Precision::Exact(0),
    1795
    0
                    },
    1796
    0
                ],
    1797
    0
            };
    1798
    1799
    0
            assert_eq!(result, expected);
    1800
    0
        }
    1801
    1802
        #[test]
    1803
    0
        fn test_stats_projection_column_with_primitive_width_only() {
    1804
    0
            let source = get_stats();
    1805
    0
            let schema = get_schema();
    1806
    1807
    0
            let projection = ProjectionExprs::new(vec![
    1808
    0
                ProjectionExpr {
    1809
    0
                    expr: Arc::new(Column::new("col2", 2)),
    1810
    0
                    alias: "col2".to_string(),
    1811
    0
                },
    1812
    0
                ProjectionExpr {
    1813
    0
                    expr: Arc::new(Column::new("col0", 0)),
    1814
    0
                    alias: "col0".to_string(),
    1815
    0
                },
    1816
            ]);
    1817
    1818
    0
            let result = projection.project_statistics(source, &schema).unwrap();
    1819
    1820
    0
            let expected = Statistics {
    1821
    0
                num_rows: Precision::Exact(5),
    1822
    0
                total_byte_size: Precision::Exact(60),
    1823
    0
                column_statistics: vec![
    1824
    0
                    ColumnStatistics {
    1825
    0
                        distinct_count: Precision::Absent,
    1826
    0
                        max_value: Precision::Exact(ScalarValue::Float32(Some(1.1))),
    1827
    0
                        min_value: Precision::Exact(ScalarValue::Float32(Some(0.1))),
    1828
    0
                        sum_value: Precision::Exact(ScalarValue::Float32(Some(5.5))),
    1829
    0
                        null_count: Precision::Absent,
    1830
    0
                    },
    1831
    0
                    ColumnStatistics {
    1832
    0
                        distinct_count: Precision::Exact(5),
    1833
    0
                        max_value: Precision::Exact(ScalarValue::Int64(Some(21))),
    1834
    0
                        min_value: Precision::Exact(ScalarValue::Int64(Some(-4))),
    1835
    0
                        sum_value: Precision::Exact(ScalarValue::Int64(Some(42))),
    1836
    0
                        null_count: Precision::Exact(0),
    1837
    0
                    },
    1838
    0
                ],
    1839
    0
            };
    1840
    1841
    0
            assert_eq!(result, expected);
    1842
    0
        }
    1843
    1844
        // Tests for Projection struct
    1845
    1846
        #[test]
    1847
    0
        fn test_projection_new() -> Result<()> {
    1848
    0
            let exprs = vec![
    1849
    0
                ProjectionExpr {
    1850
    0
                    expr: Arc::new(Column::new("a", 0)),
    1851
    0
                    alias: "a".to_string(),
    1852
    0
                },
    1853
    0
                ProjectionExpr {
    1854
    0
                    expr: Arc::new(Column::new("b", 1)),
    1855
    0
                    alias: "b".to_string(),
    1856
    0
                },
    1857
            ];
    1858
    0
            let projection = ProjectionExprs::new(exprs.clone());
    1859
    0
            assert_eq!(projection.as_ref().len(), 2);
    1860
    0
            Ok(())
    1861
    0
        }
    1862
    1863
        #[test]
    1864
    0
        fn test_projection_from_vec() -> Result<()> {
    1865
    0
            let exprs = vec![ProjectionExpr {
    1866
    0
                expr: Arc::new(Column::new("x", 0)),
    1867
    0
                alias: "x".to_string(),
    1868
    0
            }];
    1869
    0
            let projection: ProjectionExprs = exprs.clone().into();
    1870
    0
            assert_eq!(projection.as_ref().len(), 1);
    1871
    0
            Ok(())
    1872
    0
        }
    1873
    1874
        #[test]
    1875
    0
        fn test_projection_as_ref() -> Result<()> {
    1876
    0
            let exprs = vec![
    1877
    0
                ProjectionExpr {
    1878
    0
                    expr: Arc::new(Column::new("col1", 0)),
    1879
    0
                    alias: "col1".to_string(),
    1880
    0
                },
    1881
    0
                ProjectionExpr {
    1882
    0
                    expr: Arc::new(Column::new("col2", 1)),
    1883
    0
                    alias: "col2".to_string(),
    1884
    0
                },
    1885
            ];
    1886
    0
            let projection = ProjectionExprs::new(exprs);
    1887
    0
            let as_ref: &[ProjectionExpr] = projection.as_ref();
    1888
    0
            assert_eq!(as_ref.len(), 2);
    1889
    0
            Ok(())
    1890
    0
        }
    1891
    1892
        #[test]
    1893
    0
        fn test_column_indices_multiple_columns() -> Result<()> {
    1894
            // Test with reversed column order to ensure proper reordering
    1895
    0
            let projection = ProjectionExprs::new(vec![
    1896
    0
                ProjectionExpr {
    1897
    0
                    expr: Arc::new(Column::new("c", 5)),
    1898
    0
                    alias: "c".to_string(),
    1899
    0
                },
    1900
    0
                ProjectionExpr {
    1901
    0
                    expr: Arc::new(Column::new("b", 2)),
    1902
    0
                    alias: "b".to_string(),
    1903
    0
                },
    1904
    0
                ProjectionExpr {
    1905
    0
                    expr: Arc::new(Column::new("a", 0)),
    1906
    0
                    alias: "a".to_string(),
    1907
    0
                },
    1908
            ]);
    1909
            // Should return sorted indices regardless of projection order
    1910
    0
            assert_eq!(projection.column_indices(), vec![0, 2, 5]);
    1911
    0
            Ok(())
    1912
    0
        }
    1913
    1914
        #[test]
    1915
    0
        fn test_column_indices_duplicates() -> Result<()> {
    1916
            // Test that duplicate column indices appear only once
    1917
    0
            let projection = ProjectionExprs::new(vec![
    1918
    0
                ProjectionExpr {
    1919
    0
                    expr: Arc::new(Column::new("a", 1)),
    1920
    0
                    alias: "a".to_string(),
    1921
    0
                },
    1922
    0
                ProjectionExpr {
    1923
    0
                    expr: Arc::new(Column::new("b", 3)),
    1924
    0
                    alias: "b".to_string(),
    1925
    0
                },
    1926
    0
                ProjectionExpr {
    1927
    0
                    expr: Arc::new(Column::new("a2", 1)), // duplicate index
    1928
    0
                    alias: "a2".to_string(),
    1929
    0
                },
    1930
            ]);
    1931
    0
            assert_eq!(projection.column_indices(), vec![1, 3]);
    1932
    0
            Ok(())
    1933
    0
        }
    1934
    1935
        #[test]
    1936
    0
        fn test_column_indices_unsorted() -> Result<()> {
    1937
            // Test that column indices are sorted in the output
    1938
    0
            let projection = ProjectionExprs::new(vec![
    1939
    0
                ProjectionExpr {
    1940
    0
                    expr: Arc::new(Column::new("c", 5)),
    1941
    0
                    alias: "c".to_string(),
    1942
    0
                },
    1943
    0
                ProjectionExpr {
    1944
    0
                    expr: Arc::new(Column::new("a", 1)),
    1945
    0
                    alias: "a".to_string(),
    1946
    0
                },
    1947
    0
                ProjectionExpr {
    1948
    0
                    expr: Arc::new(Column::new("b", 3)),
    1949
    0
                    alias: "b".to_string(),
    1950
    0
                },
    1951
            ]);
    1952
    0
            assert_eq!(projection.column_indices(), vec![1, 3, 5]);
    1953
    0
            Ok(())
    1954
    0
        }
    1955
    1956
        #[test]
    1957
    0
        fn test_column_indices_complex_expr() -> Result<()> {
    1958
            // Test with complex expressions containing multiple columns
    1959
    0
            let expr = Arc::new(BinaryExpr::new(
    1960
    0
                Arc::new(Column::new("a", 1)),
    1961
    0
                Operator::Plus,
    1962
    0
                Arc::new(Column::new("b", 4)),
    1963
            ));
    1964
    0
            let projection = ProjectionExprs::new(vec![
    1965
    0
                ProjectionExpr {
    1966
    0
                    expr,
    1967
    0
                    alias: "sum".to_string(),
    1968
    0
                },
    1969
    0
                ProjectionExpr {
    1970
    0
                    expr: Arc::new(Column::new("c", 2)),
    1971
    0
                    alias: "c".to_string(),
    1972
    0
                },
    1973
            ]);
    1974
            // Should return [1, 2, 4] - all columns used, sorted and deduplicated
    1975
    0
            assert_eq!(projection.column_indices(), vec![1, 2, 4]);
    1976
    0
            Ok(())
    1977
    0
        }
    1978
    1979
        #[test]
    1980
    0
        fn test_column_indices_empty() -> Result<()> {
    1981
    0
            let projection = ProjectionExprs::new(vec![]);
    1982
    0
            assert_eq!(projection.column_indices(), Vec::<usize>::new());
    1983
    0
            Ok(())
    1984
    0
        }
    1985
    1986
        #[test]
    1987
    0
        fn test_merge_simple_columns() -> Result<()> {
    1988
            // First projection: SELECT c@2 AS x, b@1 AS y, a@0 AS z
    1989
    0
            let base_projection = ProjectionExprs::new(vec![
    1990
    0
                ProjectionExpr {
    1991
    0
                    expr: Arc::new(Column::new("c", 2)),
    1992
    0
                    alias: "x".to_string(),
    1993
    0
                },
    1994
    0
                ProjectionExpr {
    1995
    0
                    expr: Arc::new(Column::new("b", 1)),
    1996
    0
                    alias: "y".to_string(),
    1997
    0
                },
    1998
    0
                ProjectionExpr {
    1999
    0
                    expr: Arc::new(Column::new("a", 0)),
    2000
    0
                    alias: "z".to_string(),
    2001
    0
                },
    2002
            ]);
    2003
    2004
            // Second projection: SELECT y@1 AS col2, x@0 AS col1
    2005
    0
            let top_projection = ProjectionExprs::new(vec![
    2006
    0
                ProjectionExpr {
    2007
    0
                    expr: Arc::new(Column::new("y", 1)),
    2008
    0
                    alias: "col2".to_string(),
    2009
    0
                },
    2010
    0
                ProjectionExpr {
    2011
    0
                    expr: Arc::new(Column::new("x", 0)),
    2012
    0
                    alias: "col1".to_string(),
    2013
    0
                },
    2014
            ]);
    2015
    2016
            // Merge should produce: SELECT b@1 AS col2, c@2 AS col1
    2017
    0
            let merged = base_projection.try_merge(&top_projection)?;
    2018
    0
            assert_snapshot!(format!("{merged}"), @"Projection[b@1 AS col2, c@2 AS col1]");
    2019
    2020
    0
            Ok(())
    2021
    0
        }
    2022
    2023
        #[test]
    2024
    0
        fn test_merge_with_expressions() -> Result<()> {
    2025
            // First projection: SELECT c@2 AS x, b@1 AS y, a@0 AS z
    2026
    0
            let base_projection = ProjectionExprs::new(vec![
    2027
    0
                ProjectionExpr {
    2028
    0
                    expr: Arc::new(Column::new("c", 2)),
    2029
    0
                    alias: "x".to_string(),
    2030
    0
                },
    2031
    0
                ProjectionExpr {
    2032
    0
                    expr: Arc::new(Column::new("b", 1)),
    2033
    0
                    alias: "y".to_string(),
    2034
    0
                },
    2035
    0
                ProjectionExpr {
    2036
    0
                    expr: Arc::new(Column::new("a", 0)),
    2037
    0
                    alias: "z".to_string(),
    2038
    0
                },
    2039
            ]);
    2040
    2041
            // Second projection: SELECT y@1 + z@2 AS c2, x@0 + 1 AS c1
    2042
    0
            let top_projection = ProjectionExprs::new(vec![
    2043
    0
                ProjectionExpr {
    2044
    0
                    expr: Arc::new(BinaryExpr::new(
    2045
    0
                        Arc::new(Column::new("y", 1)),
    2046
    0
                        Operator::Plus,
    2047
    0
                        Arc::new(Column::new("z", 2)),
    2048
    0
                    )),
    2049
    0
                    alias: "c2".to_string(),
    2050
    0
                },
    2051
    0
                ProjectionExpr {
    2052
    0
                    expr: Arc::new(BinaryExpr::new(
    2053
    0
                        Arc::new(Column::new("x", 0)),
    2054
    0
                        Operator::Plus,
    2055
    0
                        Arc::new(Literal::new(ScalarValue::Int32(Some(1)))),
    2056
    0
                    )),
    2057
    0
                    alias: "c1".to_string(),
    2058
    0
                },
    2059
            ]);
    2060
    2061
            // Merge should produce: SELECT b@1 + a@0 AS c2, c@2 + 1 AS c1
    2062
    0
            let merged = base_projection.try_merge(&top_projection)?;
    2063
    0
            assert_snapshot!(format!("{merged}"), @"Projection[b@1 + a@0 AS c2, c@2 + 1 AS c1]");
    2064
    2065
    0
            Ok(())
    2066
    0
        }
    2067
    2068
        #[test]
    2069
    0
        fn try_merge_error() {
    2070
            // Create a base projection
    2071
    0
            let base = ProjectionExprs::new(vec![
    2072
    0
                ProjectionExpr {
    2073
    0
                    expr: Arc::new(Column::new("a", 0)),
    2074
    0
                    alias: "x".to_string(),
    2075
    0
                },
    2076
    0
                ProjectionExpr {
    2077
    0
                    expr: Arc::new(Column::new("b", 1)),
    2078
    0
                    alias: "y".to_string(),
    2079
    0
                },
    2080
            ]);
    2081
    2082
            // Create a top projection that references a non-existent column index
    2083
    0
            let top = ProjectionExprs::new(vec![ProjectionExpr {
    2084
    0
                expr: Arc::new(Column::new("z", 5)), // Invalid index
    2085
    0
                alias: "result".to_string(),
    2086
    0
            }]);
    2087
    2088
            // Attempt to merge and expect an error
    2089
    0
            let err_msg = base.try_merge(&top).unwrap_err().to_string();
    2090
    0
            assert!(
    2091
    0
                err_msg.contains("Internal error: Column index 5 out of bounds for projected expressions of length 2"),
    2092
    0
                "Unexpected error message: {err_msg}",
    2093
            );
    2094
    0
        }
    2095
    2096
        #[test]
    2097
    0
        fn test_merge_empty_projection_with_literal() -> Result<()> {
    2098
            // This test reproduces the issue from roundtrip_empty_projection test
    2099
            // Query like: SELECT 1 FROM table
    2100
            // where the file scan needs no columns (empty projection)
    2101
            // but we project a literal on top
    2102
    2103
            // Empty base projection (no columns needed from file)
    2104
    0
            let base_projection = ProjectionExprs::new(vec![]);
    2105
    2106
            // Top projection with a literal expression: SELECT 1
    2107
    0
            let top_projection = ProjectionExprs::new(vec![ProjectionExpr {
    2108
    0
                expr: Arc::new(Literal::new(ScalarValue::Int64(Some(1)))),
    2109
    0
                alias: "Int64(1)".to_string(),
    2110
    0
            }]);
    2111
    2112
            // This should succeed - literals don't reference columns so they should
    2113
            // pass through unchanged when merged with an empty projection
    2114
    0
            let merged = base_projection.try_merge(&top_projection)?;
    2115
    0
            assert_snapshot!(format!("{merged}"), @"Projection[1 AS Int64(1)]");
    2116
    2117
    0
            Ok(())
    2118
    0
        }
    2119
    2120
        #[test]
    2121
    0
        fn test_update_expr_with_literal() -> Result<()> {
    2122
            // Test that update_expr correctly handles expressions without column references
    2123
    0
            let literal_expr: Arc<dyn PhysicalExpr> =
    2124
    0
                Arc::new(Literal::new(ScalarValue::Int64(Some(42))));
    2125
    0
            let empty_projection: Vec<ProjectionExpr> = vec![];
    2126
    2127
            // Updating a literal with an empty projection should return the literal unchanged
    2128
    0
            let result = update_expr(&literal_expr, &empty_projection, true)?;
    2129
    0
            assert!(result.is_some(), "Literal expression should be valid");
    2130
    2131
    0
            let result_expr = result.unwrap();
    2132
    0
            assert_eq!(
    2133
    0
                result_expr
    2134
    0
                    .as_any()
    2135
    0
                    .downcast_ref::<Literal>()
    2136
    0
                    .unwrap()
    2137
    0
                    .value(),
    2138
                &ScalarValue::Int64(Some(42))
    2139
            );
    2140
    2141
    0
            Ok(())
    2142
    0
        }
    2143
    2144
        #[test]
    2145
    0
        fn test_update_expr_with_complex_literal_expr() -> Result<()> {
    2146
            // Test update_expr with an expression containing both literals and a column
    2147
            // This tests the case where we have: literal + column
    2148
    0
            let expr: Arc<dyn PhysicalExpr> = Arc::new(BinaryExpr::new(
    2149
    0
                Arc::new(Literal::new(ScalarValue::Int64(Some(10)))),
    2150
    0
                Operator::Plus,
    2151
    0
                Arc::new(Column::new("x", 0)),
    2152
    0
            ));
    2153
    2154
            // Base projection that maps column 0 to a different expression
    2155
    0
            let base_projection = vec![ProjectionExpr {
    2156
    0
                expr: Arc::new(Column::new("a", 5)),
    2157
    0
                alias: "x".to_string(),
    2158
    0
            }];
    2159
    2160
            // The expression should be updated: 10 + x@0 becomes 10 + a@5
    2161
    0
            let result = update_expr(&expr, &base_projection, true)?;
    2162
    0
            assert!(result.is_some(), "Expression should be valid");
    2163
    2164
    0
            let result_expr = result.unwrap();
    2165
    0
            let binary = result_expr
    2166
    0
                .as_any()
    2167
    0
                .downcast_ref::<BinaryExpr>()
    2168
    0
                .expect("Should be a BinaryExpr");
    2169
    2170
            // Left side should still be the literal
    2171
    0
            assert!(binary.left().as_any().downcast_ref::<Literal>().is_some());
    2172
    2173
            // Right side should be updated to reference column at index 5
    2174
    0
            let right_col = binary
    2175
    0
                .right()
    2176
    0
                .as_any()
    2177
    0
                .downcast_ref::<Column>()
    2178
    0
                .expect("Right should be a Column");
    2179
    0
            assert_eq!(right_col.index(), 5);
    2180
    2181
    0
            Ok(())
    2182
    0
        }
    2183
    2184
        #[test]
    2185
    0
        fn test_project_schema_simple_columns() -> Result<()> {
    2186
            // Input schema: [col0: Int64, col1: Utf8, col2: Float32]
    2187
    0
            let input_schema = get_schema();
    2188
    2189
            // Projection: SELECT col2 AS c, col0 AS a
    2190
    0
            let projection = ProjectionExprs::new(vec![
    2191
    0
                ProjectionExpr {
    2192
    0
                    expr: Arc::new(Column::new("col2", 2)),
    2193
    0
                    alias: "c".to_string(),
    2194
    0
                },
    2195
    0
                ProjectionExpr {
    2196
    0
                    expr: Arc::new(Column::new("col0", 0)),
    2197
    0
                    alias: "a".to_string(),
    2198
    0
                },
    2199
            ]);
    2200
    2201
    0
            let output_schema = projection.project_schema(&input_schema)?;
    2202
    2203
            // Should have 2 fields
    2204
    0
            assert_eq!(output_schema.fields().len(), 2);
    2205
    2206
            // First field should be "c" with Float32 type
    2207
    0
            assert_eq!(output_schema.field(0).name(), "c");
    2208
    0
            assert_eq!(output_schema.field(0).data_type(), &DataType::Float32);
    2209
    2210
            // Second field should be "a" with Int64 type
    2211
    0
            assert_eq!(output_schema.field(1).name(), "a");
    2212
    0
            assert_eq!(output_schema.field(1).data_type(), &DataType::Int64);
    2213
    2214
    0
            Ok(())
    2215
    0
        }
    2216
    2217
        #[test]
    2218
    0
        fn test_project_schema_with_expressions() -> Result<()> {
    2219
            // Input schema: [col0: Int64, col1: Utf8, col2: Float32]
    2220
    0
            let input_schema = get_schema();
    2221
    2222
            // Projection: SELECT col0 + 1 AS incremented
    2223
    0
            let projection = ProjectionExprs::new(vec![ProjectionExpr {
    2224
    0
                expr: Arc::new(BinaryExpr::new(
    2225
    0
                    Arc::new(Column::new("col0", 0)),
    2226
    0
                    Operator::Plus,
    2227
    0
                    Arc::new(Literal::new(ScalarValue::Int64(Some(1)))),
    2228
    0
                )),
    2229
    0
                alias: "incremented".to_string(),
    2230
    0
            }]);
    2231
    2232
    0
            let output_schema = projection.project_schema(&input_schema)?;
    2233
    2234
            // Should have 1 field
    2235
    0
            assert_eq!(output_schema.fields().len(), 1);
    2236
    2237
            // Field should be "incremented" with Int64 type
    2238
    0
            assert_eq!(output_schema.field(0).name(), "incremented");
    2239
    0
            assert_eq!(output_schema.field(0).data_type(), &DataType::Int64);
    2240
    2241
    0
            Ok(())
    2242
    0
        }
    2243
    2244
        #[test]
    2245
    0
        fn test_project_schema_preserves_metadata() -> Result<()> {
    2246
            // Create schema with metadata
    2247
    0
            let mut metadata = HashMap::new();
    2248
    0
            metadata.insert("key".to_string(), "value".to_string());
    2249
    0
            let field_with_metadata =
    2250
    0
                Field::new("col0", DataType::Int64, false).with_metadata(metadata.clone());
    2251
    0
            let input_schema = Schema::new(vec![
    2252
    0
                field_with_metadata,
    2253
    0
                Field::new("col1", DataType::Utf8, false),
    2254
            ]);
    2255
    2256
            // Projection: SELECT col0 AS renamed
    2257
    0
            let projection = ProjectionExprs::new(vec![ProjectionExpr {
    2258
    0
                expr: Arc::new(Column::new("col0", 0)),
    2259
    0
                alias: "renamed".to_string(),
    2260
    0
            }]);
    2261
    2262
    0
            let output_schema = projection.project_schema(&input_schema)?;
    2263
    2264
            // Should have 1 field
    2265
    0
            assert_eq!(output_schema.fields().len(), 1);
    2266
    2267
            // Field should be "renamed" with metadata preserved
    2268
    0
            assert_eq!(output_schema.field(0).name(), "renamed");
    2269
    0
            assert_eq!(output_schema.field(0).metadata(), &metadata);
    2270
    2271
    0
            Ok(())
    2272
    0
        }
    2273
    2274
        #[test]
    2275
    0
        fn test_project_schema_empty() -> Result<()> {
    2276
    0
            let input_schema = get_schema();
    2277
    0
            let projection = ProjectionExprs::new(vec![]);
    2278
    2279
    0
            let output_schema = projection.project_schema(&input_schema)?;
    2280
    2281
    0
            assert_eq!(output_schema.fields().len(), 0);
    2282
    2283
    0
            Ok(())
    2284
    0
        }
    2285
    2286
        #[test]
    2287
    0
        fn test_project_statistics_columns_only() -> Result<()> {
    2288
    0
            let input_stats = get_stats();
    2289
    0
            let input_schema = get_schema();
    2290
    2291
            // Projection: SELECT col1 AS text, col0 AS num
    2292
    0
            let projection = ProjectionExprs::new(vec![
    2293
    0
                ProjectionExpr {
    2294
    0
                    expr: Arc::new(Column::new("col1", 1)),
    2295
    0
                    alias: "text".to_string(),
    2296
    0
                },
    2297
    0
                ProjectionExpr {
    2298
    0
                    expr: Arc::new(Column::new("col0", 0)),
    2299
    0
                    alias: "num".to_string(),
    2300
    0
                },
    2301
            ]);
    2302
    2303
    0
            let output_stats = projection.project_statistics(input_stats, &input_schema)?;
    2304
    2305
            // Row count should be preserved
    2306
    0
            assert_eq!(output_stats.num_rows, Precision::Exact(5));
    2307
    2308
            // Should have 2 column statistics (reordered from input)
    2309
    0
            assert_eq!(output_stats.column_statistics.len(), 2);
    2310
    2311
            // First column (col1 from input)
    2312
    0
            assert_eq!(
    2313
    0
                output_stats.column_statistics[0].distinct_count,
    2314
                Precision::Exact(1)
    2315
            );
    2316
    0
            assert_eq!(
    2317
    0
                output_stats.column_statistics[0].max_value,
    2318
    0
                Precision::Exact(ScalarValue::from("x"))
    2319
            );
    2320
    2321
            // Second column (col0 from input)
    2322
    0
            assert_eq!(
    2323
    0
                output_stats.column_statistics[1].distinct_count,
    2324
                Precision::Exact(5)
    2325
            );
    2326
    0
            assert_eq!(
    2327
    0
                output_stats.column_statistics[1].max_value,
    2328
                Precision::Exact(ScalarValue::Int64(Some(21)))
    2329
            );
    2330
    2331
    0
            Ok(())
    2332
    0
        }
    2333
    2334
        #[test]
    2335
    0
        fn test_project_statistics_with_expressions() -> Result<()> {
    2336
    0
            let input_stats = get_stats();
    2337
    0
            let input_schema = get_schema();
    2338
    2339
            // Projection with expression: SELECT col0 + 1 AS incremented, col1 AS text
    2340
    0
            let projection = ProjectionExprs::new(vec![
    2341
    0
                ProjectionExpr {
    2342
    0
                    expr: Arc::new(BinaryExpr::new(
    2343
    0
                        Arc::new(Column::new("col0", 0)),
    2344
    0
                        Operator::Plus,
    2345
    0
                        Arc::new(Literal::new(ScalarValue::Int64(Some(1)))),
    2346
    0
                    )),
    2347
    0
                    alias: "incremented".to_string(),
    2348
    0
                },
    2349
    0
                ProjectionExpr {
    2350
    0
                    expr: Arc::new(Column::new("col1", 1)),
    2351
    0
                    alias: "text".to_string(),
    2352
    0
                },
    2353
            ]);
    2354
    2355
    0
            let output_stats = projection.project_statistics(input_stats, &input_schema)?;
    2356
    2357
            // Row count should be preserved
    2358
    0
            assert_eq!(output_stats.num_rows, Precision::Exact(5));
    2359
    2360
            // Should have 2 column statistics
    2361
    0
            assert_eq!(output_stats.column_statistics.len(), 2);
    2362
    2363
            // First column (expression) should have unknown statistics
    2364
    0
            assert_eq!(
    2365
    0
                output_stats.column_statistics[0].distinct_count,
    2366
                Precision::Absent
    2367
            );
    2368
    0
            assert_eq!(
    2369
    0
                output_stats.column_statistics[0].max_value,
    2370
                Precision::Absent
    2371
            );
    2372
    2373
            // Second column (col1) should preserve statistics
    2374
    0
            assert_eq!(
    2375
    0
                output_stats.column_statistics[1].distinct_count,
    2376
                Precision::Exact(1)
    2377
            );
    2378
    2379
    0
            Ok(())
    2380
    0
        }
    2381
    2382
        #[test]
    2383
    0
        fn test_project_statistics_primitive_width_only() -> Result<()> {
    2384
    0
            let input_stats = get_stats();
    2385
    0
            let input_schema = get_schema();
    2386
    2387
            // Projection with only primitive width columns: SELECT col2 AS f, col0 AS i
    2388
    0
            let projection = ProjectionExprs::new(vec![
    2389
    0
                ProjectionExpr {
    2390
    0
                    expr: Arc::new(Column::new("col2", 2)),
    2391
    0
                    alias: "f".to_string(),
    2392
    0
                },
    2393
    0
                ProjectionExpr {
    2394
    0
                    expr: Arc::new(Column::new("col0", 0)),
    2395
    0
                    alias: "i".to_string(),
    2396
    0
                },
    2397
            ]);
    2398
    2399
    0
            let output_stats = projection.project_statistics(input_stats, &input_schema)?;
    2400
    2401
            // Row count should be preserved
    2402
    0
            assert_eq!(output_stats.num_rows, Precision::Exact(5));
    2403
    2404
            // Total byte size should be recalculated for primitive types
    2405
            // Float32 (4 bytes) + Int64 (8 bytes) = 12 bytes per row, 5 rows = 60 bytes
    2406
    0
            assert_eq!(output_stats.total_byte_size, Precision::Exact(60));
    2407
    2408
            // Should have 2 column statistics
    2409
    0
            assert_eq!(output_stats.column_statistics.len(), 2);
    2410
    2411
    0
            Ok(())
    2412
    0
        }
    2413
    2414
        #[test]
    2415
    0
        fn test_project_statistics_empty() -> Result<()> {
    2416
    0
            let input_stats = get_stats();
    2417
    0
            let input_schema = get_schema();
    2418
    2419
    0
            let projection = ProjectionExprs::new(vec![]);
    2420
    2421
    0
            let output_stats = projection.project_statistics(input_stats, &input_schema)?;
    2422
    2423
            // Row count should be preserved
    2424
    0
            assert_eq!(output_stats.num_rows, Precision::Exact(5));
    2425
    2426
            // Should have no column statistics
    2427
    0
            assert_eq!(output_stats.column_statistics.len(), 0);
    2428
    2429
            // Total byte size should be 0 for empty projection
    2430
    0
            assert_eq!(output_stats.total_byte_size, Precision::Exact(0));
    2431
    2432
    0
            Ok(())
    2433
    0
        }
    2434
    }