## Coverage Report

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

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

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

    Line
    Count
    Source
    1
    // Licensed to the Apache Software Foundation (ASF) under one
    2
    // or more contributor license agreements.  See the NOTICE file
    3
    // distributed with this work for additional information
    4
    // regarding copyright ownership.  The ASF licenses this file
    5
    // to you under the Apache License, Version 2.0 (the
    6
    // "License"); you may not use this file except in compliance
    7
    // with the License.  You may obtain a copy of the License at
    8
    //
    9
    //   http://www.apache.org/licenses/LICENSE-2.0
    10
    //
    11
    // Unless required by applicable law or agreed to in writing,
    12
    // software distributed under the License is distributed on an
    13
    // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
    14
    // KIND, either express or implied.  See the License for the
    15
    // specific language governing permissions and limitations
    16
    // under the License.
    17
    18
    //! Sort expressions
    19
    20
    use std::cmp::Ordering;
    21
    use std::fmt::{self, Display, Formatter};
    22
    use std::hash::{Hash, Hasher};
    23
    use std::ops::{Deref, DerefMut};
    24
    use std::sync::Arc;
    25
    use std::vec::IntoIter;
    26
    27
    use crate::physical_expr::{fmt_sql, PhysicalExpr};
    28
    29
    use arrow::compute::kernels::sort::{SortColumn, SortOptions};
    30
    use arrow::datatypes::Schema;
    31
    use arrow::record_batch::RecordBatch;
    32
    use datafusion_common::{HashSet, Result};
    33
    use datafusion_expr_common::columnar_value::ColumnarValue;
    34
    35
    /// Represents Sort operation for a column in a RecordBatch
    36
    ///
    37
    /// Example:
    38
    /// ```
    39
    /// # use std::any::Any;
    40
    /// # use std::collections::HashMap;
    41
    /// # use std::fmt::{Display, Formatter};
    42
    /// # use std::hash::Hasher;
    43
    /// # use std::sync::Arc;
    44
    /// # use arrow::array::RecordBatch;
    45
    /// # use datafusion_common::Result;
    46
    /// # use arrow::compute::SortOptions;
    47
    /// # use arrow::datatypes::{DataType, Field, FieldRef, Schema};
    48
    /// # use datafusion_expr_common::columnar_value::ColumnarValue;
    49
    /// # use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
    50
    /// # use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
    51
    /// # // this crate doesn't have a physical expression implementation
    52
    /// # // so make a really simple one
    53
    /// # #[derive(Clone, Debug, PartialEq, Eq, Hash)]
    54
    /// # struct MyPhysicalExpr;
    55
    /// # impl PhysicalExpr for MyPhysicalExpr {
    56
    /// #  fn as_any(&self) -> &dyn Any {todo!() }
    57
    /// #  fn data_type(&self, input_schema: &Schema) -> Result<DataType> {todo!()}
    58
    /// #  fn nullable(&self, input_schema: &Schema) -> Result<bool> {todo!() }
    59
    /// #  fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {todo!() }
    60
    /// #  fn return_field(&self, input_schema: &Schema) -> Result<FieldRef> { unimplemented!() }
    61
    /// #  fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {todo!()}
    62
    /// #  fn with_new_children(self: Arc<Self>, children: Vec<Arc<dyn PhysicalExpr>>) -> Result<Arc<dyn PhysicalExpr>> {todo!()}
    63
    /// # fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { todo!() }
    64
    /// # }
    65
    /// # impl Display for MyPhysicalExpr {
    66
    /// #    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { write!(f, "a") }
    67
    /// # }
    68
    /// # fn col(name: &str) -> Arc<dyn PhysicalExpr> { Arc::new(MyPhysicalExpr) }
    69
    /// // Sort by a ASC
    70
    /// let options = SortOptions::default();
    71
    /// let sort_expr = PhysicalSortExpr::new(col("a"), options);
    72
    /// assert_eq!(sort_expr.to_string(), "a ASC");
    73
    ///
    74
    /// // Sort by a DESC NULLS LAST
    75
    /// let sort_expr = PhysicalSortExpr::new_default(col("a"))
    76
    ///   .desc()
    77
    ///   .nulls_last();
    78
    /// assert_eq!(sort_expr.to_string(), "a DESC NULLS LAST");
    79
    /// ```
    80
    #[derive(Clone, Debug, Eq)]
    81
    pub struct PhysicalSortExpr {
    82
        /// Physical expression representing the column to sort
    83
        pub expr: Arc<dyn PhysicalExpr>,
    84
        /// Option to specify how the given column should be sorted
    85
        pub options: SortOptions,
    86
    }
    87
    88
    impl PhysicalSortExpr {
    89
        /// Create a new PhysicalSortExpr
    90
    0
        pub fn new(expr: Arc<dyn PhysicalExpr>, options: SortOptions) -> Self {
    91
    0
            Self { expr, options }
    92
    0
        }
    93
    94
        /// Create a new PhysicalSortExpr with default [`SortOptions`]
    95
    0
        pub fn new_default(expr: Arc<dyn PhysicalExpr>) -> Self {
    96
    0
            Self::new(expr, SortOptions::default())
    97
    0
        }
    98
    99
        /// Reverses the sort expression. For instance, `[a ASC NULLS LAST]` turns
    100
        /// into `[a DESC NULLS FIRST]`. Such reversals are useful in planning, e.g.
    101
        /// when constructing equivalent window expressions.
    102
    0
        pub fn reverse(&self) -> Self {
    103
    0
            let mut result = self.clone();
    104
    0
            result.options = !result.options;
    105
    0
            result
    106
    0
        }
    107
    108
        /// Set the sort sort options to ASC
    109
    0
        pub fn asc(mut self) -> Self {
    110
    0
            self.options.descending = false;
    111
    0
            self
    112
    0
        }
    113
    114
        /// Set the sort sort options to DESC
    115
    0
        pub fn desc(mut self) -> Self {
    116
    0
            self.options.descending = true;
    117
    0
            self
    118
    0
        }
    119
    120
        /// Set the sort sort options to NULLS FIRST
    121
    0
        pub fn nulls_first(mut self) -> Self {
    122
    0
            self.options.nulls_first = true;
    123
    0
            self
    124
    0
        }
    125
    126
        /// Set the sort sort options to NULLS LAST
    127
    0
        pub fn nulls_last(mut self) -> Self {
    128
    0
            self.options.nulls_first = false;
    129
    0
            self
    130
    0
        }
    131
    132
        /// Like [`PhysicalExpr::fmt_sql`] prints a [`PhysicalSortExpr`] in a SQL-like format.
    133
    0
        pub fn fmt_sql(&self, f: &mut Formatter) -> fmt::Result {
    134
    0
            write!(
    135
    0
                f,
    136
    0
                "{} {}",
    137
    0
                fmt_sql(self.expr.as_ref()),
    138
    0
                to_str(&self.options)
    139
            )
    140
    0
        }
    141
    142
        /// Evaluates the sort expression into a `SortColumn` that can be passed
    143
        /// into the arrow sort kernel.
    144
    0
        pub fn evaluate_to_sort_column(&self, batch: &RecordBatch) -> Result<SortColumn> {
    145
    0
            let array_to_sort = match self.expr.evaluate(batch)? {
    146
    0
                ColumnarValue::Array(array) => array,
    147
    0
                ColumnarValue::Scalar(scalar) => scalar.to_array_of_size(batch.num_rows())?,
    148
            };
    149
    0
            Ok(SortColumn {
    150
    0
                values: array_to_sort,
    151
    0
                options: Some(self.options),
    152
    0
            })
    153
    0
        }
    154
    155
        /// Checks whether this sort expression satisfies the given `requirement`.
    156
        /// If sort options are unspecified in `requirement`, only expressions are
    157
        /// compared for inequality. See [`options_compatible`] for details on
    158
        /// how sort options compare with one another.
    159
    0
        pub fn satisfy(
    160
    0
            &self,
    161
    0
            requirement: &PhysicalSortRequirement,
    162
    0
            schema: &Schema,
    163
    0
        ) -> bool {
    164
    0
            self.expr.eq(&requirement.expr)
    165
    0
                && requirement.options.is_none_or(|opts| {
    166
    0
                    options_compatible(
    167
    0
                        &self.options,
    168
    0
                        &opts,
    169
    0
                        self.expr.nullable(schema).unwrap_or(true),
    170
                    )
    171
    0
                })
    172
    0
        }
    173
    174
        /// Checks whether this sort expression satisfies the given `sort_expr`.
    175
        /// See [`options_compatible`] for details on how sort options compare with
    176
        /// one another.
    177
    0
        pub fn satisfy_expr(&self, sort_expr: &Self, schema: &Schema) -> bool {
    178
    0
            self.expr.eq(&sort_expr.expr)
    179
    0
                && options_compatible(
    180
    0
                    &self.options,
    181
    0
                    &sort_expr.options,
    182
    0
                    self.expr.nullable(schema).unwrap_or(true),
    183
                )
    184
    0
        }
    185
    }
    186
    187
    impl PartialEq for PhysicalSortExpr {
    188
    0
        fn eq(&self, other: &Self) -> bool {
    189
    0
            self.options == other.options && self.expr.eq(&other.expr)
    190
    0
        }
    191
    }
    192
    193
    impl Hash for PhysicalSortExpr {
    194
    0
        fn hash<H: Hasher>(&self, state: &mut H) {
    195
    0
            self.expr.hash(state);
    196
    0
            self.options.hash(state);
    197
    0
        }
    198
    }
    199
    200
    impl Display for PhysicalSortExpr {
    201
    0
        fn fmt(&self, f: &mut Formatter) -> fmt::Result {
    202
    0
            write!(f, "{} {}", self.expr, to_str(&self.options))
    203
    0
        }
    204
    }
    205
    206
    /// Returns whether the given two [`SortOptions`] are compatible. Here,
    207
    /// compatibility means that they are either exactly equal, or they differ only
    208
    /// in whether NULL values come in first/last, which is immaterial because the
    209
    /// column in question is not nullable (specified by the `nullable` parameter).
    210
    0
    pub fn options_compatible(
    211
    0
        options_lhs: &SortOptions,
    212
    0
        options_rhs: &SortOptions,
    213
    0
        nullable: bool,
    214
    0
    ) -> bool {
    215
    0
        if nullable {
    216
    0
            options_lhs == options_rhs
    217
        } else {
    218
            // If the column is not nullable, NULLS FIRST/LAST is not important.
    219
    0
            options_lhs.descending == options_rhs.descending
    220
        }
    221
    0
    }
    222
    223
    /// Represents sort requirement associated with a plan
    224
    ///
    225
    /// If the requirement includes [`SortOptions`] then both the
    226
    /// expression *and* the sort options must match.
    227
    ///
    228
    /// If the requirement does not include [`SortOptions`]) then only the
    229
    /// expressions must match.
    230
    ///
    231
    /// # Examples
    232
    ///
    233
    /// With sort options (`A`, `DESC NULLS FIRST`):
    234
    /// * `ORDER BY A DESC NULLS FIRST` matches
    235
    /// * `ORDER BY A ASC  NULLS FIRST` does not match (`ASC` vs `DESC`)
    236
    /// * `ORDER BY B DESC NULLS FIRST` does not match (different expr)
    237
    ///
    238
    /// Without sort options (`A`, None):
    239
    /// * `ORDER BY A DESC NULLS FIRST` matches
    240
    /// * `ORDER BY A ASC  NULLS FIRST` matches (`ASC` and `NULL` options ignored)
    241
    /// * `ORDER BY B DESC NULLS FIRST` does not match  (different expr)
    242
    #[derive(Clone, Debug)]
    243
    pub struct PhysicalSortRequirement {
    244
        /// Physical expression representing the column to sort
    245
        pub expr: Arc<dyn PhysicalExpr>,
    246
        /// Option to specify how the given column should be sorted.
    247
        /// If unspecified, there are no constraints on sort options.
    248
        pub options: Option<SortOptions>,
    249
    }
    250
    251
    impl PartialEq for PhysicalSortRequirement {
    252
    0
        fn eq(&self, other: &Self) -> bool {
    253
    0
            self.options == other.options && self.expr.eq(&other.expr)
    254
    0
        }
    255
    }
    256
    257
    impl Display for PhysicalSortRequirement {
    258
    0
        fn fmt(&self, f: &mut Formatter) -> fmt::Result {
    259
    0
            let opts_string = self.options.as_ref().map_or("NA", to_str);
    260
    0
            write!(f, "{} {}", self.expr, opts_string)
    261
    0
        }
    262
    }
    263
    264
    /// Writes a list of [`PhysicalSortRequirement`]s to a `std::fmt::Formatter`.
    265
    ///
    266
    /// Example output: `[a + 1, b]`
    267
    0
    pub fn format_physical_sort_requirement_list(
    268
    0
        exprs: &[PhysicalSortRequirement],
    269
    0
    ) -> impl Display + '_ {
    270
        struct DisplayWrapper<'a>(&'a [PhysicalSortRequirement]);
    271
        impl Display for DisplayWrapper<'_> {
    272
    0
            fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
    273
    0
                let mut iter = self.0.iter();
    274
    0
                write!(f, "[")?;
    275
    0
                if let Some(expr) = iter.next() {
    276
    0
                    write!(f, "{expr}")?;
    277
    0
                }
    278
    0
                for expr in iter {
    279
    0
                    write!(f, ", {expr}")?;
    280
                }
    281
    0
                write!(f, "]")?;
    282
    0
                Ok(())
    283
    0
            }
    284
        }
    285
    0
        DisplayWrapper(exprs)
    286
    0
    }
    287
    288
    impl PhysicalSortRequirement {
    289
        /// Creates a new requirement.
    290
        ///
    291
        /// If `options` is `Some(..)`, creates an `exact` requirement,
    292
        /// which must match both `options` and `expr`.
    293
        ///
    294
        /// If `options` is `None`, Creates a new `expr_only` requirement,
    295
        /// which must match only `expr`.
    296
        ///
    297
        /// See [`PhysicalSortRequirement`] for examples.
    298
    0
        pub fn new(expr: Arc<dyn PhysicalExpr>, options: Option<SortOptions>) -> Self {
    299
    0
            Self { expr, options }
    300
    0
        }
    301
    302
        /// Returns whether this requirement is equal or more specific than `other`.
    303
    0
        pub fn compatible(&self, other: &Self) -> bool {
    304
    0
            self.expr.eq(&other.expr)
    305
    0
                && other
    306
    0
                    .options
    307
    0
                    .is_none_or(|other_opts| self.options == Some(other_opts))
    308
    0
        }
    309
    }
    310
    311
    /// Returns the SQL string representation of the given [`SortOptions`] object.
    312
    #[inline]
    313
    0
    fn to_str(options: &SortOptions) -> &str {
    314
    0
        match (options.descending, options.nulls_first) {
    315
    0
            (true, true) => "DESC",
    316
    0
            (true, false) => "DESC NULLS LAST",
    317
    0
            (false, true) => "ASC",
    318
    0
            (false, false) => "ASC NULLS LAST",
    319
        }
    320
    0
    }
    321
    322
    // Cross-conversion utilities between `PhysicalSortExpr` and `PhysicalSortRequirement`
    323
    impl From<PhysicalSortExpr> for PhysicalSortRequirement {
    324
    0
        fn from(value: PhysicalSortExpr) -> Self {
    325
    0
            Self::new(value.expr, Some(value.options))
    326
    0
        }
    327
    }
    328
    329
    impl From<PhysicalSortRequirement> for PhysicalSortExpr {
    330
        /// The default sort options `ASC, NULLS LAST` when the requirement does
    331
        /// not specify sort options. This default is consistent with PostgreSQL.
    332
        ///
    333
        /// Reference: <https://www.postgresql.org/docs/current/queries-order.html>
    334
    0
        fn from(value: PhysicalSortRequirement) -> Self {
    335
    0
            let options = value
    336
    0
                .options
    337
    0
                .unwrap_or_else(|| SortOptions::new(false, false));
    338
    0
            Self::new(value.expr, options)
    339
    0
        }
    340
    }
    341
    342
    /// This object represents a lexicographical ordering and contains a vector
    343
    /// of `PhysicalSortExpr` objects.
    344
    ///
    345
    /// For example, a `vec![a ASC, b DESC]` represents a lexicographical ordering
    346
    /// that first sorts by column `a` in ascending order, then by column `b` in
    347
    /// descending order.
    348
    ///
    349
    /// # Invariants
    350
    ///
    351
    /// The following always hold true for a `LexOrdering`:
    352
    ///
    353
    /// 1. It is non-degenerate, meaning it contains at least one element.
    354
    /// 2. It is duplicate-free, meaning it does not contain multiple entries for
    355
    ///    the same column.
    356
    #[derive(Debug, Clone)]
    357
    pub struct LexOrdering {
    358
        /// Vector of sort expressions representing the lexicographical ordering.
    359
        exprs: Vec<PhysicalSortExpr>,
    360
        /// Set of expressions in the lexicographical ordering, used to ensure
    361
        /// that the ordering is duplicate-free. Note that the elements in this
    362
        /// set are the same underlying physical expressions as in `exprs`.
    363
        set: HashSet<Arc<dyn PhysicalExpr>>,
    364
    }
    365
    366
    impl LexOrdering {
    367
        /// Creates a new [`LexOrdering`] from the given vector of sort expressions.
    368
        /// If the vector is empty, returns `None`.
    369
    0
        pub fn new(exprs: impl IntoIterator<Item = PhysicalSortExpr>) -> Option<Self> {
    370
    0
            let exprs = exprs.into_iter();
    371
    0
            let mut candidate = Self {
    372
    0
                // not valid yet; valid publicly-returned instance must be non-empty
    373
    0
                exprs: Vec::new(),
    374
    0
                set: HashSet::new(),
    375
    0
            };
    376
    0
            for expr in exprs {
    377
    0
                candidate.push(expr);
    378
    0
            }
    379
    0
            if candidate.exprs.is_empty() {
    380
    0
                None
    381
            } else {
    382
    0
                Some(candidate)
    383
            }
    384
    0
        }
    385
    386
        /// Appends an element to the back of the `LexOrdering`.
    387
    0
        pub fn push(&mut self, sort_expr: PhysicalSortExpr) {
    388
    0
            if self.set.insert(Arc::clone(&sort_expr.expr)) {
    389
    0
                self.exprs.push(sort_expr);
    390
    0
            }
    391
    0
        }
    392
    393
        /// Add all elements from `iter` to the `LexOrdering`.
    394
    0
        pub fn extend(&mut self, sort_exprs: impl IntoIterator<Item = PhysicalSortExpr>) {
    395
    0
            for sort_expr in sort_exprs {
    396
    0
                self.push(sort_expr);
    397
    0
            }
    398
    0
        }
    399
    400
        /// Returns the leading `PhysicalSortExpr` of the `LexOrdering`. Note that
    401
        /// this function does not return an `Option`, as a `LexOrdering` is always
    402
        /// non-degenerate (i.e. it contains at least one element).
    403
    0
        pub fn first(&self) -> &PhysicalSortExpr {
    404
            // Can safely `unwrap` because `LexOrdering` is non-degenerate:
    405
    0
            self.exprs.first().unwrap()
    406
    0
        }
    407
    408
        /// Returns the number of elements that can be stored in the `LexOrdering`
    409
        /// without reallocating.
    410
    0
        pub fn capacity(&self) -> usize {
    411
    0
            self.exprs.capacity()
    412
    0
        }
    413
    414
        /// Truncates the `LexOrdering`, keeping only the first `len` elements.
    415
        /// Returns `true` if truncation made a change, `false` otherwise. Negative
    416
        /// cases happen in two scenarios: (1) When `len` is greater than or equal
    417
        /// to the number of expressions inside this `LexOrdering`, making truncation
    418
        /// a no-op, or (2) when `len` is `0`, making truncation impossible.
    419
    0
        pub fn truncate(&mut self, len: usize) -> bool {
    420
    0
            if len == 0 || len >= self.exprs.len() {
    421
    0
                return false;
    422
    0
            }
    423
    0
            for PhysicalSortExpr { expr, .. } in self.exprs[len..].iter() {
    424
    0
                self.set.remove(expr);
    425
    0
            }
    426
    0
            self.exprs.truncate(len);
    427
    0
            true
    428
    0
        }
    429
    }
    430
    431
    impl PartialEq for LexOrdering {
    432
    0
        fn eq(&self, other: &Self) -> bool {
    433
            let Self {
    434
    0
                exprs,
    435
                set: _, // derived from `exprs`
    436
    0
            } = self;
    437
            // PartialEq must be consistent with PartialOrd
    438
    0
            exprs == &other.exprs
    439
    0
        }
    440
    }
    441
    impl Eq for LexOrdering {}
    442
    impl PartialOrd for LexOrdering {
    443
        /// There is a partial ordering among `LexOrdering` objects. For example, the
    444
        /// ordering `[a ASC]` is coarser (less) than ordering `[a ASC, b ASC]`.
    445
        /// If two orderings do not share a prefix, they are incomparable.
    446
    0
        fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
    447
            // PartialEq must be consistent with PartialOrd
    448
    0
            self.exprs
    449
    0
                .iter()
    450
    0
                .zip(other.exprs.iter())
    451
    0
                .all(|(lhs, rhs)| lhs == rhs)
    452
    0
                .then(|| self.len().cmp(&other.len()))
    453
    0
        }
    454
    }
    455
    456
    impl<const N: usize> From<[PhysicalSortExpr; N]> for LexOrdering {
    457
    0
        fn from(value: [PhysicalSortExpr; N]) -> Self {
    458
            // TODO: Replace this assertion with a condition on the generic parameter
    459
            //       when Rust supports it.
    460
    0
            assert!(N > 0);
    461
    0
            Self::new(value)
    462
    0
                .expect("A LexOrdering from non-empty array must be non-degenerate")
    463
    0
        }
    464
    }
    465
    466
    impl Deref for LexOrdering {
    467
        type Target = [PhysicalSortExpr];
    468
    469
    0
        fn deref(&self) -> &Self::Target {
    470
    0
            self.exprs.as_slice()
    471
    0
        }
    472
    }
    473
    474
    impl Display for LexOrdering {
    475
    0
        fn fmt(&self, f: &mut Formatter) -> fmt::Result {
    476
    0
            let mut first = true;
    477
    0
            for sort_expr in &self.exprs {
    478
    0
                if first {
    479
    0
                    first = false;
    480
    0
                } else {
    481
    0
                    write!(f, ", ")?;
    482
                }
    483
    0
                write!(f, "{sort_expr}")?;
    484
            }
    485
    0
            Ok(())
    486
    0
        }
    487
    }
    488
    489
    impl IntoIterator for LexOrdering {
    490
        type Item = PhysicalSortExpr;
    491
        type IntoIter = IntoIter<Self::Item>;
    492
    493
    0
        fn into_iter(self) -> Self::IntoIter {
    494
    0
            self.exprs.into_iter()
    495
    0
        }
    496
    }
    497
    498
    impl<'a> IntoIterator for &'a LexOrdering {
    499
        type Item = &'a PhysicalSortExpr;
    500
        type IntoIter = std::slice::Iter<'a, PhysicalSortExpr>;
    501
    502
    0
        fn into_iter(self) -> Self::IntoIter {
    503
    0
            self.exprs.iter()
    504
    0
        }
    505
    }
    506
    507
    impl From<LexOrdering> for Vec<PhysicalSortExpr> {
    508
    0
        fn from(ordering: LexOrdering) -> Self {
    509
    0
            ordering.exprs
    510
    0
        }
    511
    }
    512
    513
    /// This object represents a lexicographical ordering requirement and contains
    514
    /// a vector of `PhysicalSortRequirement` objects.
    515
    ///
    516
    /// For example, a `vec![a Some(ASC), b None]` represents a lexicographical
    517
    /// requirement that firsts imposes an ordering by column `a` in ascending
    518
    /// order, then by column `b` in *any* (ascending or descending) order. The
    519
    /// ordering is non-degenerate, meaning it contains at least one element, and
    520
    /// it is duplicate-free, meaning it does not contain multiple entries for the
    521
    /// same column.
    522
    ///
    523
    /// Note that a `LexRequirement` need not enforce the uniqueness of its sort
    524
    /// expressions after construction like a `LexOrdering` does, because it provides
    525
    /// no mutation methods. If such methods become necessary, we will need to
    526
    /// enforce uniqueness like the latter object.
    527
    #[derive(Debug, Clone, PartialEq)]
    528
    pub struct LexRequirement {
    529
        reqs: Vec<PhysicalSortRequirement>,
    530
    }
    531
    532
    impl LexRequirement {
    533
        /// Creates a new [`LexRequirement`] from the given vector of sort expressions.
    534
        /// If the vector is empty, returns `None`.
    535
    0
        pub fn new(reqs: impl IntoIterator<Item = PhysicalSortRequirement>) -> Option<Self> {
    536
    0
            let (non_empty, requirements) = Self::construct(reqs);
    537
    0
            non_empty.then_some(requirements)
    538
    0
        }
    539
    540
        /// Returns the leading `PhysicalSortRequirement` of the `LexRequirement`.
    541
        /// Note that this function does not return an `Option`, as a `LexRequirement`
    542
        /// is always non-degenerate (i.e. it contains at least one element).
    543
    0
        pub fn first(&self) -> &PhysicalSortRequirement {
    544
            // Can safely `unwrap` because `LexRequirement` is non-degenerate:
    545
    0
            self.reqs.first().unwrap()
    546
    0
        }
    547
    548
        /// Constructs a new `LexRequirement` from the given sort requirements w/o
    549
        /// enforcing non-degeneracy. This function is used internally and is not
    550
        /// meant (or safe) for external use.
    551
    0
        fn construct(
    552
    0
            reqs: impl IntoIterator<Item = PhysicalSortRequirement>,
    553
    0
        ) -> (bool, Self) {
    554
    0
            let mut set = HashSet::new();
    555
    0
            let reqs = reqs
    556
    0
                .into_iter()
    557
    0
                .filter_map(|r| set.insert(Arc::clone(&r.expr)).then_some(r))
    558
    0
                .collect();
    559
    0
            (!set.is_empty(), Self { reqs })
    560
    0
        }
    561
    }
    562
    563
    impl<const N: usize> From<[PhysicalSortRequirement; N]> for LexRequirement {
    564
    0
        fn from(value: [PhysicalSortRequirement; N]) -> Self {
    565
            // TODO: Replace this assertion with a condition on the generic parameter
    566
            //       when Rust supports it.
    567
    0
            assert!(N > 0);
    568
    0
            let (non_empty, requirement) = Self::construct(value);
    569
    0
            debug_assert!(non_empty);
    570
    0
            requirement
    571
    0
        }
    572
    }
    573
    574
    impl Deref for LexRequirement {
    575
        type Target = [PhysicalSortRequirement];
    576
    577
    0
        fn deref(&self) -> &Self::Target {
    578
    0
            self.reqs.as_slice()
    579
    0
        }
    580
    }
    581
    582
    impl IntoIterator for LexRequirement {
    583
        type Item = PhysicalSortRequirement;
    584
        type IntoIter = IntoIter<Self::Item>;
    585
    586
    0
        fn into_iter(self) -> Self::IntoIter {
    587
    0
            self.reqs.into_iter()
    588
    0
        }
    589
    }
    590
    591
    impl<'a> IntoIterator for &'a LexRequirement {
    592
        type Item = &'a PhysicalSortRequirement;
    593
        type IntoIter = std::slice::Iter<'a, PhysicalSortRequirement>;
    594
    595
    0
        fn into_iter(self) -> Self::IntoIter {
    596
    0
            self.reqs.iter()
    597
    0
        }
    598
    }
    599
    600
    impl From<LexRequirement> for Vec<PhysicalSortRequirement> {
    601
    0
        fn from(requirement: LexRequirement) -> Self {
    602
    0
            requirement.reqs
    603
    0
        }
    604
    }
    605
    606
    // Cross-conversion utilities between `LexOrdering` and `LexRequirement`
    607
    impl From<LexOrdering> for LexRequirement {
    608
    0
        fn from(value: LexOrdering) -> Self {
    609
            // Can construct directly as `value` is non-degenerate:
    610
    0
            let (non_empty, requirements) =
    611
    0
                Self::construct(value.into_iter().map(Into::into));
    612
    0
            debug_assert!(non_empty);
    613
    0
            requirements
    614
    0
        }
    615
    }
    616
    617
    impl From<LexRequirement> for LexOrdering {
    618
    0
        fn from(value: LexRequirement) -> Self {
    619
            // Can construct directly as `value` is non-degenerate
    620
    0
            Self::new(value.into_iter().map(Into::into))
    621
    0
                .expect("A LexOrdering from LexRequirement must be non-degenerate")
    622
    0
        }
    623
    }
    624
    625
    /// Represents a plan's input ordering requirements. Vector elements represent
    626
    /// alternative ordering requirements in the order of preference. The list of
    627
    /// alternatives can be either hard or soft, depending on whether the operator
    628
    /// can work without an input ordering.
    629
    ///
    630
    /// # Invariants
    631
    ///
    632
    /// The following always hold true for a `OrderingRequirements`:
    633
    ///
    634
    /// 1. It is non-degenerate, meaning it contains at least one ordering. The
    635
    ///    absence of an input ordering requirement is represented by a `None` value
    636
    ///    in `ExecutionPlan` APIs, which return an `Option<OrderingRequirements>`.
    637
    #[derive(Debug, Clone, PartialEq)]
    638
    pub enum OrderingRequirements {
    639
        /// The operator is not able to work without one of these requirements.
    640
        Hard(Vec<LexRequirement>),
    641
        /// The operator can benefit from these input orderings when available,
    642
        /// but can still work in the absence of any input ordering.
    643
        Soft(Vec<LexRequirement>),
    644
    }
    645
    646
    impl OrderingRequirements {
    647
        /// Creates a new instance from the given alternatives. If an empty list of
    648
        /// alternatives are given, returns `None`.
    649
        pub fn new_alternatives(
    650
            alternatives: impl IntoIterator<Item = LexRequirement>,
    651
            soft: bool,
    652
        ) -> Option<Self> {
    653
            let alternatives = alternatives.into_iter().collect::<Vec<_>>();
    654
            (!alternatives.is_empty()).then(|| {
    655
                if soft {
    656
                    Self::Soft(alternatives)
    657
                } else {
    658
                    Self::Hard(alternatives)
    659
                }
    660
            })
    661
        }
    662
    663
        /// Creates a new instance with a single hard requirement.
    664
    0
        pub fn new(requirement: LexRequirement) -> Self {
    665
    0
            Self::Hard(vec![requirement])
    666
    0
        }
    667
    668
        /// Creates a new instance with a single soft requirement.
    669
    0
        pub fn new_soft(requirement: LexRequirement) -> Self {
    670
    0
            Self::Soft(vec![requirement])
    671
    0
        }
    672
    673
        /// Adds an alternative requirement to the list of alternatives.
    674
    0
        pub fn add_alternative(&mut self, requirement: LexRequirement) {
    675
    0
            match self {
    676
    0
                Self::Hard(alts) | Self::Soft(alts) => alts.push(requirement),
    677
            }
    678
    0
        }
    679
    680
        /// Returns the first (i.e. most preferred) `LexRequirement` among
    681
        /// alternative requirements.
    682
    0
        pub fn into_single(self) -> LexRequirement {
    683
    0
            match self {
    684
    0
                Self::Hard(mut alts) | Self::Soft(mut alts) => alts.swap_remove(0),
    685
            }
    686
    0
        }
    687
    688
        /// Returns a reference to the first (i.e. most preferred) `LexRequirement`
    689
        /// among alternative requirements.
    690
    0
        pub fn first(&self) -> &LexRequirement {
    691
    0
            match self {
    692
    0
                Self::Hard(alts) | Self::Soft(alts) => &alts[0],
    693
            }
    694
    0
        }
    695
    696
        /// Returns all alternatives as a vector of `LexRequirement` objects and a
    697
        /// boolean value indicating softness/hardness of the requirements.
    698
    0
        pub fn into_alternatives(self) -> (Vec<LexRequirement>, bool) {
    699
    0
            match self {
    700
    0
                Self::Hard(alts) => (alts, false),
    701
    0
                Self::Soft(alts) => (alts, true),
    702
            }
    703
    0
        }
    704
    }
    705
    706
    impl From<LexRequirement> for OrderingRequirements {
    707
    0
        fn from(requirement: LexRequirement) -> Self {
    708
    0
            Self::new(requirement)
    709
    0
        }
    710
    }
    711
    712
    impl From<LexOrdering> for OrderingRequirements {
    713
    0
        fn from(ordering: LexOrdering) -> Self {
    714
    0
            Self::new(ordering.into())
    715
    0
        }
    716
    }
    717
    718
    impl Deref for OrderingRequirements {
    719
        type Target = [LexRequirement];
    720
    721
    0
        fn deref(&self) -> &Self::Target {
    722
    0
            match &self {
    723
    0
                Self::Hard(alts) | Self::Soft(alts) => alts.as_slice(),
    724
            }
    725
    0
        }
    726
    }
    727
    728
    impl DerefMut for OrderingRequirements {
    729
    0
        fn deref_mut(&mut self) -> &mut Self::Target {
    730
    0
            match self {
    731
    0
                Self::Hard(alts) | Self::Soft(alts) => alts.as_mut_slice(),
    732
            }
    733
    0
        }
    734
    }