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}