Skip to main content

datafusion_physical_expr_common/
sort_expr.rs

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
20use std::cmp::Ordering;
21use std::fmt::{self, Display, Formatter};
22use std::hash::{Hash, Hasher};
23use std::ops::{Deref, DerefMut};
24use std::sync::Arc;
25use std::vec::IntoIter;
26
27use crate::physical_expr::{PhysicalExpr, fmt_sql};
28
29use arrow::compute::kernels::sort::{SortColumn, SortOptions};
30use arrow::datatypes::Schema;
31use arrow::record_batch::RecordBatch;
32use datafusion_common::{HashSet, Result};
33use datafusion_expr_common::columnar_value::ColumnarValue;
34use indexmap::IndexSet;
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 data_type(&self, input_schema: &Schema) -> Result<DataType> {todo!()}
57/// #  fn nullable(&self, input_schema: &Schema) -> Result<bool> {todo!() }
58/// #  fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {todo!() }
59/// #  fn return_field(&self, input_schema: &Schema) -> Result<FieldRef> { unimplemented!() }
60/// #  fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {todo!()}
61/// #  fn with_new_children(self: Arc<Self>, children: Vec<Arc<dyn PhysicalExpr>>) -> Result<Arc<dyn PhysicalExpr>> {todo!()}
62/// # fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { todo!() }
63/// # }
64/// # impl Display for MyPhysicalExpr {
65/// #    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { write!(f, "a") }
66/// # }
67/// # fn col(name: &str) -> Arc<dyn PhysicalExpr> { Arc::new(MyPhysicalExpr) }
68/// // Sort by a ASC
69/// let options = SortOptions::default();
70/// let sort_expr = PhysicalSortExpr::new(col("a"), options);
71/// assert_eq!(sort_expr.to_string(), "a ASC");
72///
73/// // Sort by a DESC NULLS LAST
74/// let sort_expr = PhysicalSortExpr::new_default(col("a"))
75///   .desc()
76///   .nulls_last();
77/// assert_eq!(sort_expr.to_string(), "a DESC NULLS LAST");
78/// ```
79#[derive(Clone, Debug, Eq)]
80pub struct PhysicalSortExpr {
81    /// Physical expression representing the column to sort
82    pub expr: Arc<dyn PhysicalExpr>,
83    /// Option to specify how the given column should be sorted
84    pub options: SortOptions,
85}
86
87impl PhysicalSortExpr {
88    /// Create a new PhysicalSortExpr
89    pub fn new(expr: Arc<dyn PhysicalExpr>, options: SortOptions) -> Self {
90        Self { expr, options }
91    }
92
93    /// Create a new PhysicalSortExpr with default [`SortOptions`]
94    pub fn new_default(expr: Arc<dyn PhysicalExpr>) -> Self {
95        Self::new(expr, SortOptions::default())
96    }
97
98    /// Reverses the sort expression. For instance, `[a ASC NULLS LAST]` turns
99    /// into `[a DESC NULLS FIRST]`. Such reversals are useful in planning, e.g.
100    /// when constructing equivalent window expressions.
101    pub fn reverse(&self) -> Self {
102        let mut result = self.clone();
103        result.options = !result.options;
104        result
105    }
106
107    /// Set the sort sort options to ASC
108    pub fn asc(mut self) -> Self {
109        self.options.descending = false;
110        self
111    }
112
113    /// Set the sort sort options to DESC
114    pub fn desc(mut self) -> Self {
115        self.options.descending = true;
116        self
117    }
118
119    /// Set the sort sort options to NULLS FIRST
120    pub fn nulls_first(mut self) -> Self {
121        self.options.nulls_first = true;
122        self
123    }
124
125    /// Set the sort sort options to NULLS LAST
126    pub fn nulls_last(mut self) -> Self {
127        self.options.nulls_first = false;
128        self
129    }
130
131    /// Like [`PhysicalExpr::fmt_sql`] prints a [`PhysicalSortExpr`] in a SQL-like format.
132    pub fn fmt_sql(&self, f: &mut Formatter) -> fmt::Result {
133        write!(
134            f,
135            "{} {}",
136            fmt_sql(self.expr.as_ref()),
137            to_str(&self.options)
138        )
139    }
140
141    /// Evaluates the sort expression into a `SortColumn` that can be passed
142    /// into the arrow sort kernel.
143    pub fn evaluate_to_sort_column(&self, batch: &RecordBatch) -> Result<SortColumn> {
144        let array_to_sort = match self.expr.evaluate(batch)? {
145            ColumnarValue::Array(array) => array,
146            ColumnarValue::Scalar(scalar) => scalar.to_array_of_size(batch.num_rows())?,
147        };
148        Ok(SortColumn {
149            values: array_to_sort,
150            options: Some(self.options),
151        })
152    }
153
154    /// Checks whether this sort expression satisfies the given `requirement`.
155    /// If sort options are unspecified in `requirement`, only expressions are
156    /// compared for inequality. See [`options_compatible`] for details on
157    /// how sort options compare with one another.
158    pub fn satisfy(
159        &self,
160        requirement: &PhysicalSortRequirement,
161        schema: &Schema,
162    ) -> bool {
163        self.expr.eq(&requirement.expr)
164            && requirement.options.is_none_or(|opts| {
165                options_compatible(
166                    &self.options,
167                    &opts,
168                    self.expr.nullable(schema).unwrap_or(true),
169                )
170            })
171    }
172
173    /// Checks whether this sort expression satisfies the given `sort_expr`.
174    /// See [`options_compatible`] for details on how sort options compare with
175    /// one another.
176    pub fn satisfy_expr(&self, sort_expr: &Self, schema: &Schema) -> bool {
177        self.expr.eq(&sort_expr.expr)
178            && options_compatible(
179                &self.options,
180                &sort_expr.options,
181                self.expr.nullable(schema).unwrap_or(true),
182            )
183    }
184}
185
186/// Protobuf conversions for [`PhysicalSortExpr`].
187///
188/// This is the flat [`PhysicalSortExprNode`] representation used wherever the
189/// wire format stores an ordering (scan output orderings, range partitioning,
190/// window frames, …). It is *not* the `PhysicalExprNode::Sort` wrapping that
191/// `SortExec` uses for its own `expr` field.
192///
193/// [`PhysicalSortExprNode`]: datafusion_proto_models::protobuf::PhysicalSortExprNode
194#[cfg(feature = "proto")]
195impl PhysicalSortExpr {
196    /// Serialize this sort expression, encoding its child expression through
197    /// `ctx`.
198    pub fn try_to_proto(
199        &self,
200        ctx: &crate::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>,
201    ) -> Result<datafusion_proto_models::protobuf::PhysicalSortExprNode> {
202        Ok(datafusion_proto_models::protobuf::PhysicalSortExprNode {
203            expr: Some(Box::new(ctx.encode_child(&self.expr)?)),
204            asc: !self.options.descending,
205            nulls_first: self.options.nulls_first,
206        })
207    }
208
209    /// Reconstruct a [`PhysicalSortExpr`] from its protobuf representation.
210    pub fn try_from_proto(
211        node: &datafusion_proto_models::protobuf::PhysicalSortExprNode,
212        ctx: &crate::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>,
213    ) -> Result<Self> {
214        let expr = ctx.decode_required_expression(
215            node.expr.as_deref(),
216            "PhysicalSortExpr",
217            "expr",
218        )?;
219        Ok(PhysicalSortExpr {
220            expr,
221            options: SortOptions {
222                descending: !node.asc,
223                nulls_first: node.nulls_first,
224            },
225        })
226    }
227}
228
229/// Serialize a sequence of sort expressions into the flat
230/// [`PhysicalSortExprNode`] list the wire format uses for an ordering.
231///
232/// Accepts anything that yields [`PhysicalSortExpr`]s by value or by reference,
233/// so a [`LexOrdering`], a `&[PhysicalSortExpr]`, or a [`LexRequirement`]
234/// mapped through [`PhysicalSortExpr::from`] all work:
235///
236/// ```ignore
237/// let nodes = sort_exprs_try_to_proto(ordering.iter(), ctx)?;
238/// let nodes = sort_exprs_try_to_proto(
239///     requirement.iter().map(|req| PhysicalSortExpr::from(req.clone())),
240///     ctx,
241/// )?;
242/// ```
243///
244/// The `PhysicalSortExprNodeCollection` message some plans use is just this
245/// list in a wrapper, so those callers wrap the result themselves rather than
246/// this function guessing which shape they mean.
247///
248/// [`PhysicalSortExprNode`]: datafusion_proto_models::protobuf::PhysicalSortExprNode
249#[cfg(feature = "proto")]
250pub fn sort_exprs_try_to_proto<E: std::borrow::Borrow<PhysicalSortExpr>>(
251    exprs: impl IntoIterator<Item = E>,
252    ctx: &crate::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>,
253) -> Result<Vec<datafusion_proto_models::protobuf::PhysicalSortExprNode>> {
254    exprs
255        .into_iter()
256        .map(|expr| expr.borrow().try_to_proto(ctx))
257        .collect()
258}
259
260/// Reconstruct a sequence of sort expressions from the flat
261/// [`PhysicalSortExprNode`] list, the counterpart of
262/// [`sort_exprs_try_to_proto`].
263///
264/// Returns the expressions rather than a [`LexOrdering`] or a
265/// [`LexRequirement`], because callers differ in what an empty list means:
266/// `LexOrdering::new` / `LexRequirement::new` return `None` for it, which is
267/// "no ordering declared" for a scan and an error for an operator that requires
268/// one. Callers with the former convention can use
269/// [`optional_ordering_try_from_proto`] instead.
270///
271/// [`PhysicalSortExprNode`]: datafusion_proto_models::protobuf::PhysicalSortExprNode
272#[cfg(feature = "proto")]
273pub fn sort_exprs_try_from_proto(
274    nodes: &[datafusion_proto_models::protobuf::PhysicalSortExprNode],
275    ctx: &crate::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>,
276) -> Result<Vec<PhysicalSortExpr>> {
277    nodes
278        .iter()
279        .map(|node| PhysicalSortExpr::try_from_proto(node, ctx))
280        .collect()
281}
282
283/// Serialize an optional [`LexOrdering`], encoding `None` as an empty list.
284#[cfg(feature = "proto")]
285pub fn optional_ordering_try_to_proto(
286    ordering: Option<&LexOrdering>,
287    ctx: &crate::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>,
288) -> Result<Vec<datafusion_proto_models::protobuf::PhysicalSortExprNode>> {
289    sort_exprs_try_to_proto(ordering.into_iter().flatten(), ctx)
290}
291
292/// Counterpart of [`optional_ordering_try_to_proto`]: an empty list decodes
293/// as `None`.
294#[cfg(feature = "proto")]
295pub fn optional_ordering_try_from_proto(
296    nodes: &[datafusion_proto_models::protobuf::PhysicalSortExprNode],
297    ctx: &crate::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>,
298) -> Result<Option<LexOrdering>> {
299    Ok(LexOrdering::new(sort_exprs_try_from_proto(nodes, ctx)?))
300}
301
302impl PartialEq for PhysicalSortExpr {
303    fn eq(&self, other: &Self) -> bool {
304        self.options == other.options && self.expr.eq(&other.expr)
305    }
306}
307
308impl Hash for PhysicalSortExpr {
309    fn hash<H: Hasher>(&self, state: &mut H) {
310        self.expr.hash(state);
311        self.options.hash(state);
312    }
313}
314
315impl Display for PhysicalSortExpr {
316    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
317        write!(f, "{} {}", self.expr, to_str(&self.options))
318    }
319}
320
321/// Returns whether the given two [`SortOptions`] are compatible. Here,
322/// compatibility means that they are either exactly equal, or they differ only
323/// in whether NULL values come in first/last, which is immaterial because the
324/// column in question is not nullable (specified by the `nullable` parameter).
325pub fn options_compatible(
326    options_lhs: &SortOptions,
327    options_rhs: &SortOptions,
328    nullable: bool,
329) -> bool {
330    if nullable {
331        options_lhs == options_rhs
332    } else {
333        // If the column is not nullable, NULLS FIRST/LAST is not important.
334        options_lhs.descending == options_rhs.descending
335    }
336}
337
338/// Represents sort requirement associated with a plan
339///
340/// If the requirement includes [`SortOptions`] then both the
341/// expression *and* the sort options must match.
342///
343/// If the requirement does not include [`SortOptions`]) then only the
344/// expressions must match.
345///
346/// # Examples
347///
348/// With sort options (`A`, `DESC NULLS FIRST`):
349/// * `ORDER BY A DESC NULLS FIRST` matches
350/// * `ORDER BY A ASC  NULLS FIRST` does not match (`ASC` vs `DESC`)
351/// * `ORDER BY B DESC NULLS FIRST` does not match (different expr)
352///
353/// Without sort options (`A`, None):
354/// * `ORDER BY A DESC NULLS FIRST` matches
355/// * `ORDER BY A ASC  NULLS FIRST` matches (`ASC` and `NULL` options ignored)
356/// * `ORDER BY B DESC NULLS FIRST` does not match  (different expr)
357#[derive(Clone, Debug)]
358pub struct PhysicalSortRequirement {
359    /// Physical expression representing the column to sort
360    pub expr: Arc<dyn PhysicalExpr>,
361    /// Option to specify how the given column should be sorted.
362    /// If unspecified, there are no constraints on sort options.
363    pub options: Option<SortOptions>,
364}
365
366impl PartialEq for PhysicalSortRequirement {
367    fn eq(&self, other: &Self) -> bool {
368        self.options == other.options && self.expr.eq(&other.expr)
369    }
370}
371
372impl Display for PhysicalSortRequirement {
373    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
374        let opts_string = self.options.as_ref().map_or("NA", to_str);
375        write!(f, "{} {}", self.expr, opts_string)
376    }
377}
378
379/// Writes a list of [`PhysicalSortRequirement`]s to a `std::fmt::Formatter`.
380///
381/// Example output: `[a + 1, b]`
382pub fn format_physical_sort_requirement_list(
383    exprs: &[PhysicalSortRequirement],
384) -> impl Display + '_ {
385    struct DisplayWrapper<'a>(&'a [PhysicalSortRequirement]);
386    impl Display for DisplayWrapper<'_> {
387        fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
388            let mut iter = self.0.iter();
389            write!(f, "[")?;
390            if let Some(expr) = iter.next() {
391                write!(f, "{expr}")?;
392            }
393            for expr in iter {
394                write!(f, ", {expr}")?;
395            }
396            write!(f, "]")?;
397            Ok(())
398        }
399    }
400    DisplayWrapper(exprs)
401}
402
403impl PhysicalSortRequirement {
404    /// Creates a new requirement.
405    ///
406    /// If `options` is `Some(..)`, creates an `exact` requirement,
407    /// which must match both `options` and `expr`.
408    ///
409    /// If `options` is `None`, Creates a new `expr_only` requirement,
410    /// which must match only `expr`.
411    ///
412    /// See [`PhysicalSortRequirement`] for examples.
413    pub fn new(expr: Arc<dyn PhysicalExpr>, options: Option<SortOptions>) -> Self {
414        Self { expr, options }
415    }
416
417    /// Returns whether this requirement is equal or more specific than `other`.
418    pub fn compatible(&self, other: &Self) -> bool {
419        self.expr.eq(&other.expr)
420            && other
421                .options
422                .is_none_or(|other_opts| self.options == Some(other_opts))
423    }
424}
425
426/// Returns the SQL string representation of the given [`SortOptions`] object.
427#[inline]
428fn to_str(options: &SortOptions) -> &str {
429    match (options.descending, options.nulls_first) {
430        (true, true) => "DESC",
431        (true, false) => "DESC NULLS LAST",
432        (false, true) => "ASC",
433        (false, false) => "ASC NULLS LAST",
434    }
435}
436
437// Cross-conversion utilities between `PhysicalSortExpr` and `PhysicalSortRequirement`
438impl From<PhysicalSortExpr> for PhysicalSortRequirement {
439    fn from(value: PhysicalSortExpr) -> Self {
440        Self::new(value.expr, Some(value.options))
441    }
442}
443
444impl From<PhysicalSortRequirement> for PhysicalSortExpr {
445    /// The default sort options `ASC, NULLS LAST` when the requirement does
446    /// not specify sort options. This default is consistent with PostgreSQL.
447    ///
448    /// Reference: <https://www.postgresql.org/docs/current/queries-order.html>
449    fn from(value: PhysicalSortRequirement) -> Self {
450        let options = value
451            .options
452            .unwrap_or_else(|| SortOptions::new(false, false));
453        Self::new(value.expr, options)
454    }
455}
456
457/// This object represents a lexicographical ordering and contains a vector
458/// of `PhysicalSortExpr` objects.
459///
460/// For example, a `vec![a ASC, b DESC]` represents a lexicographical ordering
461/// that first sorts by column `a` in ascending order, then by column `b` in
462/// descending order.
463///
464/// # Invariants
465///
466/// The following always hold true for a `LexOrdering`:
467///
468/// 1. It is non-degenerate, meaning it contains at least one element.
469/// 2. It is duplicate-free, meaning it does not contain multiple entries for
470///    the same column.
471#[derive(Clone, Debug)]
472pub struct LexOrdering {
473    /// Vector of sort expressions representing the lexicographical ordering.
474    exprs: Vec<PhysicalSortExpr>,
475    /// Set of expressions in the lexicographical ordering, used to ensure
476    /// that the ordering is duplicate-free. Note that the elements in this
477    /// set are the same underlying physical expressions as in `exprs`.
478    set: IndexSet<Arc<dyn PhysicalExpr>>,
479}
480
481impl LexOrdering {
482    /// Creates a new [`LexOrdering`] from the given vector of sort expressions.
483    /// If the vector is empty, returns `None`.
484    pub fn new(exprs: impl IntoIterator<Item = PhysicalSortExpr>) -> Option<Self> {
485        let exprs = exprs.into_iter();
486        let mut candidate = Self {
487            // not valid yet; valid publicly-returned instance must be non-empty
488            exprs: Vec::new(),
489            set: IndexSet::new(),
490        };
491        for expr in exprs {
492            candidate.push(expr);
493        }
494        if candidate.exprs.is_empty() {
495            None
496        } else {
497            Some(candidate)
498        }
499    }
500
501    /// Appends an element to the back of the `LexOrdering`.
502    pub fn push(&mut self, sort_expr: PhysicalSortExpr) {
503        if self.set.insert(Arc::clone(&sort_expr.expr)) {
504            self.exprs.push(sort_expr);
505        }
506    }
507
508    /// Add all elements from `iter` to the `LexOrdering`.
509    pub fn extend(&mut self, sort_exprs: impl IntoIterator<Item = PhysicalSortExpr>) {
510        for sort_expr in sort_exprs {
511            self.push(sort_expr);
512        }
513    }
514
515    /// Returns the leading `PhysicalSortExpr` of the `LexOrdering`. Note that
516    /// this function does not return an `Option`, as a `LexOrdering` is always
517    /// non-degenerate (i.e. it contains at least one element).
518    pub fn first(&self) -> &PhysicalSortExpr {
519        // Can safely `unwrap` because `LexOrdering` is non-degenerate:
520        self.exprs.first().unwrap()
521    }
522
523    /// Returns the number of elements that can be stored in the `LexOrdering`
524    /// without reallocating.
525    pub fn capacity(&self) -> usize {
526        self.exprs.capacity()
527    }
528
529    /// Truncates the `LexOrdering`, keeping only the first `len` elements.
530    /// Returns `true` if truncation made a change, `false` otherwise. Negative
531    /// cases happen in two scenarios: (1) When `len` is greater than or equal
532    /// to the number of expressions inside this `LexOrdering`, making truncation
533    /// a no-op, or (2) when `len` is `0`, making truncation impossible.
534    pub fn truncate(&mut self, len: usize) -> bool {
535        if len == 0 || len >= self.exprs.len() {
536            return false;
537        }
538        for PhysicalSortExpr { expr, .. } in self.exprs[len..].iter() {
539            self.set.swap_remove(expr);
540        }
541        self.exprs.truncate(len);
542        true
543    }
544
545    /// Check if reversing this ordering would satisfy another ordering requirement.
546    ///
547    /// This supports **prefix matching**: if this ordering is `[A DESC, B ASC]`
548    /// and `other` is `[A ASC]`, reversing this gives `[A ASC, B DESC]`, which
549    /// satisfies `other` since `[A ASC]` is a prefix.
550    ///
551    /// # Arguments
552    /// * `other` - The ordering requirement to check against
553    ///
554    /// # Returns
555    /// `true` if reversing this ordering would satisfy `other`
556    ///
557    /// # Example
558    /// ```text
559    /// self:  [number DESC, letter ASC]
560    /// other: [number ASC]
561    /// After reversing self: [number ASC, letter DESC]  ✓ Prefix match!
562    /// ```
563    pub fn is_reverse(&self, other: &LexOrdering) -> bool {
564        let self_exprs = self.as_ref();
565        let other_exprs = other.as_ref();
566
567        if other_exprs.len() > self_exprs.len() {
568            return false;
569        }
570
571        other_exprs.iter().zip(self_exprs.iter()).all(|(req, cur)| {
572            req.expr.eq(&cur.expr) && is_reversed_sort_options(&req.options, &cur.options)
573        })
574    }
575
576    /// Returns the sort options for the given expression if one is defined in this `LexOrdering`.
577    pub fn get_sort_options(&self, expr: &dyn PhysicalExpr) -> Option<SortOptions> {
578        for e in self {
579            if e.expr.as_ref().dyn_eq(expr) {
580                return Some(e.options);
581            }
582        }
583
584        None
585    }
586}
587
588/// Check if two SortOptions represent reversed orderings.
589///
590/// Returns `true` if both `descending` and `nulls_first` are opposite.
591///
592/// # Example
593/// ```
594/// use arrow::compute::SortOptions;
595/// # use datafusion_physical_expr_common::sort_expr::is_reversed_sort_options;
596///
597/// let asc_nulls_last = SortOptions {
598///     descending: false,
599///     nulls_first: false,
600/// };
601/// let desc_nulls_first = SortOptions {
602///     descending: true,
603///     nulls_first: true,
604/// };
605///
606/// assert!(is_reversed_sort_options(&asc_nulls_last, &desc_nulls_first));
607/// assert!(is_reversed_sort_options(&desc_nulls_first, &asc_nulls_last));
608/// ```
609pub fn is_reversed_sort_options(lhs: &SortOptions, rhs: &SortOptions) -> bool {
610    lhs.descending != rhs.descending && lhs.nulls_first != rhs.nulls_first
611}
612
613impl PartialEq for LexOrdering {
614    fn eq(&self, other: &Self) -> bool {
615        let Self {
616            exprs,
617            set: _, // derived from `exprs`
618        } = self;
619        // PartialEq must be consistent with PartialOrd
620        exprs == &other.exprs
621    }
622}
623impl Eq for LexOrdering {}
624impl PartialOrd for LexOrdering {
625    /// There is a partial ordering among `LexOrdering` objects. For example, the
626    /// ordering `[a ASC]` is coarser (less) than ordering `[a ASC, b ASC]`.
627    /// If two orderings do not share a prefix, they are incomparable.
628    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
629        // PartialEq must be consistent with PartialOrd
630        self.exprs
631            .iter()
632            .zip(other.exprs.iter())
633            .all(|(lhs, rhs)| lhs == rhs)
634            .then(|| self.len().cmp(&other.len()))
635    }
636}
637
638impl<const N: usize> From<[PhysicalSortExpr; N]> for LexOrdering {
639    fn from(value: [PhysicalSortExpr; N]) -> Self {
640        // TODO: Replace this assertion with a condition on the generic parameter
641        //       when Rust supports it.
642        assert!(N > 0);
643        Self::new(value)
644            .expect("A LexOrdering from non-empty array must be non-degenerate")
645    }
646}
647
648impl Deref for LexOrdering {
649    type Target = [PhysicalSortExpr];
650
651    fn deref(&self) -> &Self::Target {
652        self.exprs.as_slice()
653    }
654}
655
656impl Display for LexOrdering {
657    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
658        let mut first = true;
659        for sort_expr in &self.exprs {
660            if first {
661                first = false;
662            } else {
663                write!(f, ", ")?;
664            }
665            write!(f, "{sort_expr}")?;
666        }
667        Ok(())
668    }
669}
670
671impl IntoIterator for LexOrdering {
672    type Item = PhysicalSortExpr;
673    type IntoIter = IntoIter<Self::Item>;
674
675    fn into_iter(self) -> Self::IntoIter {
676        self.exprs.into_iter()
677    }
678}
679
680impl<'a> IntoIterator for &'a LexOrdering {
681    type Item = &'a PhysicalSortExpr;
682    type IntoIter = std::slice::Iter<'a, PhysicalSortExpr>;
683
684    fn into_iter(self) -> Self::IntoIter {
685        self.exprs.iter()
686    }
687}
688
689impl From<LexOrdering> for Vec<PhysicalSortExpr> {
690    fn from(ordering: LexOrdering) -> Self {
691        ordering.exprs
692    }
693}
694
695/// This object represents a lexicographical ordering requirement and contains
696/// a vector of `PhysicalSortRequirement` objects.
697///
698/// For example, a `vec![a Some(ASC), b None]` represents a lexicographical
699/// requirement that firsts imposes an ordering by column `a` in ascending
700/// order, then by column `b` in *any* (ascending or descending) order. The
701/// ordering is non-degenerate, meaning it contains at least one element, and
702/// it is duplicate-free, meaning it does not contain multiple entries for the
703/// same column.
704///
705/// Note that a `LexRequirement` need not enforce the uniqueness of its sort
706/// expressions after construction like a `LexOrdering` does, because it provides
707/// no mutation methods. If such methods become necessary, we will need to
708/// enforce uniqueness like the latter object.
709#[derive(Debug, Clone, PartialEq)]
710pub struct LexRequirement {
711    reqs: Vec<PhysicalSortRequirement>,
712}
713
714impl LexRequirement {
715    /// Creates a new [`LexRequirement`] from the given vector of sort expressions.
716    /// If the vector is empty, returns `None`.
717    pub fn new(reqs: impl IntoIterator<Item = PhysicalSortRequirement>) -> Option<Self> {
718        let (non_empty, requirements) = Self::construct(reqs);
719        non_empty.then_some(requirements)
720    }
721
722    /// Returns the leading `PhysicalSortRequirement` of the `LexRequirement`.
723    /// Note that this function does not return an `Option`, as a `LexRequirement`
724    /// is always non-degenerate (i.e. it contains at least one element).
725    pub fn first(&self) -> &PhysicalSortRequirement {
726        // Can safely `unwrap` because `LexRequirement` is non-degenerate:
727        self.reqs.first().unwrap()
728    }
729
730    /// Constructs a new `LexRequirement` from the given sort requirements w/o
731    /// enforcing non-degeneracy. This function is used internally and is not
732    /// meant (or safe) for external use.
733    fn construct(
734        reqs: impl IntoIterator<Item = PhysicalSortRequirement>,
735    ) -> (bool, Self) {
736        let mut set = HashSet::new();
737        let reqs = reqs
738            .into_iter()
739            .filter_map(|r| set.insert(Arc::clone(&r.expr)).then_some(r))
740            .collect();
741        (!set.is_empty(), Self { reqs })
742    }
743}
744
745impl<const N: usize> From<[PhysicalSortRequirement; N]> for LexRequirement {
746    fn from(value: [PhysicalSortRequirement; N]) -> Self {
747        // TODO: Replace this assertion with a condition on the generic parameter
748        //       when Rust supports it.
749        assert!(N > 0);
750        let (non_empty, requirement) = Self::construct(value);
751        debug_assert!(non_empty);
752        requirement
753    }
754}
755
756impl Deref for LexRequirement {
757    type Target = [PhysicalSortRequirement];
758
759    fn deref(&self) -> &Self::Target {
760        self.reqs.as_slice()
761    }
762}
763
764impl IntoIterator for LexRequirement {
765    type Item = PhysicalSortRequirement;
766    type IntoIter = IntoIter<Self::Item>;
767
768    fn into_iter(self) -> Self::IntoIter {
769        self.reqs.into_iter()
770    }
771}
772
773impl<'a> IntoIterator for &'a LexRequirement {
774    type Item = &'a PhysicalSortRequirement;
775    type IntoIter = std::slice::Iter<'a, PhysicalSortRequirement>;
776
777    fn into_iter(self) -> Self::IntoIter {
778        self.reqs.iter()
779    }
780}
781
782impl From<LexRequirement> for Vec<PhysicalSortRequirement> {
783    fn from(requirement: LexRequirement) -> Self {
784        requirement.reqs
785    }
786}
787
788// Cross-conversion utilities between `LexOrdering` and `LexRequirement`
789impl From<LexOrdering> for LexRequirement {
790    fn from(value: LexOrdering) -> Self {
791        // Can construct directly as `value` is non-degenerate:
792        let (non_empty, requirements) =
793            Self::construct(value.into_iter().map(Into::into));
794        debug_assert!(non_empty);
795        requirements
796    }
797}
798
799impl From<LexRequirement> for LexOrdering {
800    fn from(value: LexRequirement) -> Self {
801        // Can construct directly as `value` is non-degenerate
802        Self::new(value.into_iter().map(Into::into))
803            .expect("A LexOrdering from LexRequirement must be non-degenerate")
804    }
805}
806
807/// Represents a plan's input ordering requirements. Vector elements represent
808/// alternative ordering requirements in the order of preference. The list of
809/// alternatives can be either hard or soft, depending on whether the operator
810/// can work without an input ordering.
811///
812/// # Invariants
813///
814/// The following always hold true for a `OrderingRequirements`:
815///
816/// 1. It is non-degenerate, meaning it contains at least one ordering. The
817///    absence of an input ordering requirement is represented by a `None` value
818///    in `ExecutionPlan` APIs, which return an `Option<OrderingRequirements>`.
819#[derive(Debug, Clone, PartialEq)]
820pub enum OrderingRequirements {
821    /// The operator is not able to work without one of these requirements.
822    Hard(Vec<LexRequirement>),
823    /// The operator can benefit from these input orderings when available,
824    /// but can still work in the absence of any input ordering.
825    Soft(Vec<LexRequirement>),
826}
827
828impl OrderingRequirements {
829    /// Creates a new instance from the given alternatives. If an empty list of
830    /// alternatives are given, returns `None`.
831    pub fn new_alternatives(
832        alternatives: impl IntoIterator<Item = LexRequirement>,
833        soft: bool,
834    ) -> Option<Self> {
835        let alternatives = alternatives.into_iter().collect::<Vec<_>>();
836        (!alternatives.is_empty()).then(|| {
837            if soft {
838                Self::Soft(alternatives)
839            } else {
840                Self::Hard(alternatives)
841            }
842        })
843    }
844
845    /// Creates a new instance with a single hard requirement.
846    pub fn new(requirement: LexRequirement) -> Self {
847        Self::Hard(vec![requirement])
848    }
849
850    /// Creates a new instance with a single soft requirement.
851    pub fn new_soft(requirement: LexRequirement) -> Self {
852        Self::Soft(vec![requirement])
853    }
854
855    /// Adds an alternative requirement to the list of alternatives.
856    pub fn add_alternative(&mut self, requirement: LexRequirement) {
857        match self {
858            Self::Hard(alts) | Self::Soft(alts) => alts.push(requirement),
859        }
860    }
861
862    /// Returns the first (i.e. most preferred) `LexRequirement` among
863    /// alternative requirements.
864    pub fn into_single(self) -> LexRequirement {
865        match self {
866            Self::Hard(mut alts) | Self::Soft(mut alts) => alts.swap_remove(0),
867        }
868    }
869
870    /// Returns a reference to the first (i.e. most preferred) `LexRequirement`
871    /// among alternative requirements.
872    pub fn first(&self) -> &LexRequirement {
873        match self {
874            Self::Hard(alts) | Self::Soft(alts) => &alts[0],
875        }
876    }
877
878    /// Returns all alternatives as a vector of `LexRequirement` objects and a
879    /// boolean value indicating softness/hardness of the requirements.
880    pub fn into_alternatives(self) -> (Vec<LexRequirement>, bool) {
881        match self {
882            Self::Hard(alts) => (alts, false),
883            Self::Soft(alts) => (alts, true),
884        }
885    }
886}
887
888impl From<LexRequirement> for OrderingRequirements {
889    fn from(requirement: LexRequirement) -> Self {
890        Self::new(requirement)
891    }
892}
893
894impl From<LexOrdering> for OrderingRequirements {
895    fn from(ordering: LexOrdering) -> Self {
896        Self::new(ordering.into())
897    }
898}
899
900impl Deref for OrderingRequirements {
901    type Target = [LexRequirement];
902
903    fn deref(&self) -> &Self::Target {
904        match &self {
905            Self::Hard(alts) | Self::Soft(alts) => alts.as_slice(),
906        }
907    }
908}
909
910impl DerefMut for OrderingRequirements {
911    fn deref_mut(&mut self) -> &mut Self::Target {
912        match self {
913            Self::Hard(alts) | Self::Soft(alts) => alts.as_mut_slice(),
914        }
915    }
916}
917
918#[cfg(test)]
919mod tests {
920    use super::*;
921
922    #[test]
923    fn test_is_reversed_sort_options() {
924        // Test basic reversal: ASC NULLS LAST ↔ DESC NULLS FIRST
925        let asc_nulls_last = SortOptions {
926            descending: false,
927            nulls_first: false,
928        };
929        let desc_nulls_first = SortOptions {
930            descending: true,
931            nulls_first: true,
932        };
933        assert!(is_reversed_sort_options(&asc_nulls_last, &desc_nulls_first));
934        assert!(is_reversed_sort_options(&desc_nulls_first, &asc_nulls_last));
935
936        // Test another reversal: ASC NULLS FIRST ↔ DESC NULLS LAST
937        let asc_nulls_first = SortOptions {
938            descending: false,
939            nulls_first: true,
940        };
941        let desc_nulls_last = SortOptions {
942            descending: true,
943            nulls_first: false,
944        };
945        assert!(is_reversed_sort_options(&asc_nulls_first, &desc_nulls_last));
946        assert!(is_reversed_sort_options(&desc_nulls_last, &asc_nulls_first));
947
948        // Test non-reversal: same options
949        assert!(!is_reversed_sort_options(&asc_nulls_last, &asc_nulls_last));
950        assert!(!is_reversed_sort_options(
951            &desc_nulls_first,
952            &desc_nulls_first
953        ));
954
955        // Test non-reversal: only descending differs
956        assert!(!is_reversed_sort_options(&asc_nulls_last, &desc_nulls_last));
957        assert!(!is_reversed_sort_options(&desc_nulls_last, &asc_nulls_last));
958
959        // Test non-reversal: only nulls_first differs
960        assert!(!is_reversed_sort_options(&asc_nulls_last, &asc_nulls_first));
961        assert!(!is_reversed_sort_options(&asc_nulls_first, &asc_nulls_last));
962    }
963}