risingwave_common/util/
sort_util.rs

1// Copyright 2022 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::cmp::Ordering;
16use std::fmt;
17use std::sync::Arc;
18
19use parse_display::Display;
20use risingwave_common_estimate_size::EstimateSize;
21use risingwave_pb::common::{PbColumnOrder, PbDirection, PbNullsAre, PbOrderType};
22
23use super::iter_util::ZipEqDebug;
24use crate::array::{Array, DataChunk};
25use crate::catalog::{FieldDisplay, Schema};
26use crate::dispatch_array_variants;
27use crate::row::Row;
28use crate::types::{DefaultOrdered, ToDatumRef};
29
30/// Sort direction, ascending/descending.
31#[derive(PartialEq, Eq, Hash, Copy, Clone, Debug, Display, Default)]
32pub enum Direction {
33    #[default]
34    #[display("ASC")]
35    Ascending,
36    #[display("DESC")]
37    Descending,
38}
39
40impl Direction {
41    fn from_protobuf(direction: &PbDirection) -> Self {
42        match direction {
43            PbDirection::Ascending => Self::Ascending,
44            PbDirection::Descending => Self::Descending,
45            PbDirection::Unspecified => unreachable!(),
46        }
47    }
48
49    fn to_protobuf(self) -> PbDirection {
50        match self {
51            Self::Ascending => PbDirection::Ascending,
52            Self::Descending => PbDirection::Descending,
53        }
54    }
55}
56
57impl Direction {
58    fn reverse(self) -> Self {
59        match self {
60            Self::Ascending => Self::Descending,
61            Self::Descending => Self::Ascending,
62        }
63    }
64}
65
66/// Nulls are largest/smallest.
67#[derive(PartialEq, Eq, Hash, Copy, Clone, Debug, Display, Default)]
68enum NullsAre {
69    #[default]
70    #[display("LARGEST")]
71    Largest,
72    #[display("SMALLEST")]
73    Smallest,
74}
75
76impl NullsAre {
77    fn from_protobuf(nulls_are: &PbNullsAre) -> Self {
78        match nulls_are {
79            PbNullsAre::Largest => Self::Largest,
80            PbNullsAre::Smallest => Self::Smallest,
81            PbNullsAre::Unspecified => unreachable!(),
82        }
83    }
84
85    fn to_protobuf(self) -> PbNullsAre {
86        match self {
87            Self::Largest => PbNullsAre::Largest,
88            Self::Smallest => PbNullsAre::Smallest,
89        }
90    }
91}
92
93/// Order type of a column.
94#[derive(PartialEq, Eq, Hash, Copy, Clone, Debug, Default)]
95pub struct OrderType {
96    direction: Direction,
97    nulls_are: NullsAre,
98}
99
100impl OrderType {
101    pub fn from_protobuf(order_type: &PbOrderType) -> OrderType {
102        OrderType {
103            direction: Direction::from_protobuf(&order_type.direction()),
104            nulls_are: NullsAre::from_protobuf(&order_type.nulls_are()),
105        }
106    }
107
108    pub fn to_protobuf(self) -> PbOrderType {
109        PbOrderType {
110            direction: self.direction.to_protobuf() as _,
111            nulls_are: self.nulls_are.to_protobuf() as _,
112        }
113    }
114}
115
116impl OrderType {
117    fn new(direction: Direction, nulls_are: NullsAre) -> Self {
118        Self {
119            direction,
120            nulls_are,
121        }
122    }
123
124    fn nulls_first(direction: Direction) -> Self {
125        match direction {
126            Direction::Ascending => Self::new(direction, NullsAre::Smallest),
127            Direction::Descending => Self::new(direction, NullsAre::Largest),
128        }
129    }
130
131    fn nulls_last(direction: Direction) -> Self {
132        match direction {
133            Direction::Ascending => Self::new(direction, NullsAre::Largest),
134            Direction::Descending => Self::new(direction, NullsAre::Smallest),
135        }
136    }
137
138    pub fn from_bools(asc: Option<bool>, nulls_first: Option<bool>) -> Self {
139        let direction = match asc {
140            None => Direction::default(),
141            Some(true) => Direction::Ascending,
142            Some(false) => Direction::Descending,
143        };
144        match nulls_first {
145            None => Self::new(direction, NullsAre::default()),
146            Some(true) => Self::nulls_first(direction),
147            Some(false) => Self::nulls_last(direction),
148        }
149    }
150
151    // TODO(rc): Many places that call `ascending` should've call `default`.
152    /// Create an `ASC` order type.
153    pub fn ascending() -> Self {
154        Self {
155            direction: Direction::Ascending,
156            nulls_are: NullsAre::default(),
157        }
158    }
159
160    /// Create a `DESC` order type.
161    pub fn descending() -> Self {
162        Self {
163            direction: Direction::Descending,
164            nulls_are: NullsAre::default(),
165        }
166    }
167
168    /// Create an `ASC NULLS FIRST` order type.
169    pub fn ascending_nulls_first() -> Self {
170        Self::nulls_first(Direction::Ascending)
171    }
172
173    /// Create an `ASC NULLS LAST` order type.
174    pub fn ascending_nulls_last() -> Self {
175        Self::nulls_last(Direction::Ascending)
176    }
177
178    /// Create a `DESC NULLS FIRST` order type.
179    pub fn descending_nulls_first() -> Self {
180        Self::nulls_first(Direction::Descending)
181    }
182
183    /// Create a `DESC NULLS LAST` order type.
184    pub fn descending_nulls_last() -> Self {
185        Self::nulls_last(Direction::Descending)
186    }
187
188    pub fn direction(&self) -> Direction {
189        self.direction
190    }
191
192    pub fn is_ascending(&self) -> bool {
193        self.direction == Direction::Ascending
194    }
195
196    pub fn is_descending(&self) -> bool {
197        self.direction == Direction::Descending
198    }
199
200    pub fn nulls_are_largest(&self) -> bool {
201        self.nulls_are == NullsAre::Largest
202    }
203
204    pub fn nulls_are_smallest(&self) -> bool {
205        self.nulls_are == NullsAre::Smallest
206    }
207
208    pub fn nulls_are_first(&self) -> bool {
209        self.is_ascending() && self.nulls_are_smallest()
210            || self.is_descending() && self.nulls_are_largest()
211    }
212
213    pub fn nulls_are_last(&self) -> bool {
214        !self.nulls_are_first()
215    }
216
217    pub fn reverse(self) -> Self {
218        Self::new(self.direction.reverse(), self.nulls_are)
219    }
220
221    pub fn all() -> Vec<Self> {
222        vec![
223            Self::ascending_nulls_first(),
224            Self::ascending_nulls_last(),
225            Self::descending_nulls_first(),
226            Self::descending_nulls_last(),
227        ]
228    }
229}
230
231impl fmt::Display for OrderType {
232    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
233        write!(f, "{}", self.direction)?;
234        if self.nulls_are != NullsAre::default() {
235            write!(
236                f,
237                " NULLS {}",
238                if self.nulls_are_first() {
239                    "FIRST"
240                } else {
241                    "LAST"
242                }
243            )?;
244        }
245        Ok(())
246    }
247}
248
249/// Column index with an order type (ASC or DESC). Used to represent a sort key
250/// (`Vec<ColumnOrder>`).
251///
252/// Corresponds to protobuf [`PbColumnOrder`].
253#[derive(Clone, PartialEq, Eq, Hash, Copy)]
254pub struct ColumnOrder {
255    pub column_index: usize,
256    pub order_type: OrderType,
257}
258
259impl ColumnOrder {
260    pub fn new(column_index: usize, order_type: OrderType) -> Self {
261        Self {
262            column_index,
263            order_type,
264        }
265    }
266
267    /// Shift the column index with offset.
268    pub fn shift_with_offset(&mut self, offset: isize) {
269        self.column_index = (self.column_index as isize + offset) as usize;
270    }
271}
272
273impl ColumnOrder {
274    pub fn from_protobuf(column_order: &PbColumnOrder) -> Self {
275        ColumnOrder {
276            column_index: column_order.column_index as _,
277            order_type: OrderType::from_protobuf(column_order.get_order_type().unwrap()),
278        }
279    }
280
281    pub fn to_protobuf(self) -> PbColumnOrder {
282        PbColumnOrder {
283            column_index: self.column_index as _,
284            order_type: Some(self.order_type.to_protobuf()),
285        }
286    }
287}
288
289impl fmt::Display for ColumnOrder {
290    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
291        write!(f, "${} {}", self.column_index, self.order_type)
292    }
293}
294
295impl fmt::Debug for ColumnOrder {
296    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
297        write!(f, "{}", self)
298    }
299}
300
301/// Returns the index of the first `ORDER BY` column of a streaming TopN operator if watermarks on
302/// that column can be forwarded by the operator, i.e. when the column is ordered `ASC NULLS LAST`.
303///
304/// A watermark `wm` on a column promises that no row with the column value `< wm` will ever arrive
305/// again. To forward it, the operator must guarantee that it will never emit any change (insertion,
306/// deletion or update) to rows with the column value `< wm` afterwards. A row arriving at a TopN
307/// operator can only change its own output and the output of rows ordered *after* it, by evicting
308/// them from or restoring them into the top N. When the first `ORDER BY` column is ordered
309/// `ASC NULLS LAST`, all such rows are not smaller than the arriving row in that column, which is
310/// `>= wm` (or NULL, treated as the largest) by the watermark guarantee. With `DESC` or
311/// `ASC NULLS FIRST` ordering, an arriving row may evict rows with values `< wm` instead, so
312/// watermarks on the column must be dropped.
313///
314/// The rule is shared by the optimizer and the executors so that they stay in sync.
315pub fn topn_watermark_forwardable_order_key(order_by: &[ColumnOrder]) -> Option<usize> {
316    let first = order_by.first()?;
317    (first.order_type.is_ascending() && first.order_type.nulls_are_largest())
318        .then_some(first.column_index)
319}
320
321pub struct ColumnOrderDisplay<'a> {
322    pub column_order: &'a ColumnOrder,
323    pub input_schema: &'a Schema,
324}
325
326impl fmt::Display for ColumnOrderDisplay<'_> {
327    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
328        let that = self.column_order;
329        write!(
330            f,
331            "{} {}",
332            FieldDisplay(self.input_schema.fields.get(that.column_index).unwrap()),
333            that.order_type
334        )
335    }
336}
337
338#[derive(Clone, Debug)]
339pub struct HeapElem {
340    column_orders: Arc<Vec<ColumnOrder>>,
341    chunk: DataChunk,
342    chunk_idx: usize,
343    elem_idx: usize,
344    /// `DataChunk` can be encoded to accelerate the comparison.
345    /// Use `risingwave_common::util::encoding_for_comparison::encode_chunk`
346    /// to perform encoding, otherwise the comparison will be performed
347    /// column by column.
348    encoded_chunk: Option<Arc<Vec<Vec<u8>>>>,
349    estimated_size: usize,
350}
351
352impl HeapElem {
353    pub fn new(
354        column_orders: Arc<Vec<ColumnOrder>>,
355        chunk: DataChunk,
356        chunk_idx: usize,
357        elem_idx: usize,
358        encoded_chunk: Option<Arc<Vec<Vec<u8>>>>,
359    ) -> Self {
360        let estimated_size = encoded_chunk
361            .as_ref()
362            .map(|inner| inner.iter().map(|i| i.capacity()).sum())
363            .unwrap_or(0);
364
365        Self {
366            column_orders,
367            chunk,
368            chunk_idx,
369            elem_idx,
370            encoded_chunk,
371            estimated_size,
372        }
373    }
374
375    #[inline(always)]
376    pub fn chunk_idx(&self) -> usize {
377        self.chunk_idx
378    }
379
380    #[inline(always)]
381    pub fn elem_idx(&self) -> usize {
382        self.elem_idx
383    }
384
385    pub fn chunk(&self) -> &DataChunk {
386        &self.chunk
387    }
388}
389
390impl Ord for HeapElem {
391    fn cmp(&self, other: &Self) -> Ordering {
392        let ord = if let (Some(lhs_encoded_chunk), Some(rhs_encoded_chunk)) =
393            (self.encoded_chunk.as_ref(), other.encoded_chunk.as_ref())
394        {
395            lhs_encoded_chunk[self.elem_idx]
396                .as_slice()
397                .cmp(rhs_encoded_chunk[other.elem_idx].as_slice())
398        } else {
399            compare_rows_in_chunk(
400                &self.chunk,
401                self.elem_idx,
402                &other.chunk,
403                other.elem_idx,
404                self.column_orders.as_ref(),
405            )
406        };
407        ord.reverse()
408    }
409}
410
411impl EstimateSize for HeapElem {
412    fn estimated_heap_size(&self) -> usize {
413        self.estimated_size
414    }
415}
416
417impl PartialOrd for HeapElem {
418    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
419        Some(self.cmp(other))
420    }
421}
422
423impl PartialEq for HeapElem {
424    fn eq(&self, other: &Self) -> bool {
425        self.cmp(other) == Ordering::Equal
426    }
427}
428
429impl Eq for HeapElem {}
430
431fn generic_partial_cmp<T: PartialOrd>(
432    lhs: Option<&T>,
433    rhs: Option<&T>,
434    order_type: OrderType,
435) -> Option<Ordering> {
436    let ord = match (lhs, rhs, order_type.nulls_are) {
437        (Some(l), Some(r), _) => l.partial_cmp(r),
438        (None, None, _) => Some(Ordering::Equal),
439        (Some(_), None, NullsAre::Largest) => Some(Ordering::Less),
440        (Some(_), None, NullsAre::Smallest) => Some(Ordering::Greater),
441        (None, Some(_), NullsAre::Largest) => Some(Ordering::Greater),
442        (None, Some(_), NullsAre::Smallest) => Some(Ordering::Less),
443    };
444    ord.map(|o| {
445        if order_type.is_descending() {
446            o.reverse()
447        } else {
448            o
449        }
450    })
451}
452
453fn compare_values_in_array<'a, T>(
454    lhs_array: &'a T,
455    lhs_idx: usize,
456    rhs_array: &'a T,
457    rhs_idx: usize,
458    order_type: OrderType,
459) -> Ordering
460where
461    T: Array,
462    <T as Array>::RefItem<'a>: PartialOrd,
463{
464    generic_partial_cmp(
465        lhs_array.value_at(lhs_idx).as_ref(),
466        rhs_array.value_at(rhs_idx).as_ref(),
467        order_type,
468    )
469    .expect("items in the same `Array` type should be able to compare")
470}
471
472fn compare_rows_in_chunk(
473    lhs_data_chunk: &DataChunk,
474    lhs_idx: usize,
475    rhs_data_chunk: &DataChunk,
476    rhs_idx: usize,
477    column_orders: &[ColumnOrder],
478) -> Ordering {
479    for column_order in column_orders {
480        let lhs_array = lhs_data_chunk.column_at(column_order.column_index);
481        let rhs_array = rhs_data_chunk.column_at(column_order.column_index);
482
483        let res = dispatch_array_variants!(&**lhs_array, lhs_inner, {
484            #[expect(
485                clippy::unnecessary_fallible_conversions,
486                reason = "FIXME: Array's `into` is not safe. We should revise it."
487            )]
488            let rhs_inner = (&**rhs_array).try_into().unwrap_or_else(|_| {
489                panic!(
490                    "Unmatched array types, lhs array is: {}, rhs array is: {}",
491                    lhs_array.get_ident(),
492                    rhs_array.get_ident(),
493                )
494            });
495            compare_values_in_array(
496                lhs_inner,
497                lhs_idx,
498                rhs_inner,
499                rhs_idx,
500                column_order.order_type,
501            )
502        });
503
504        if res != Ordering::Equal {
505            return res;
506        }
507    }
508    Ordering::Equal
509}
510
511/// Partial compare two `Datum`s with specified order type.
512pub fn partial_cmp_datum(
513    lhs: impl ToDatumRef,
514    rhs: impl ToDatumRef,
515    order_type: OrderType,
516) -> Option<Ordering> {
517    let lhs = lhs.to_datum_ref().map(DefaultOrdered);
518    let rhs = rhs.to_datum_ref().map(DefaultOrdered);
519    generic_partial_cmp(lhs.as_ref(), rhs.as_ref(), order_type)
520}
521
522/// Compare two `Datum`s with specified order type.
523///
524/// # Panics
525///
526/// Panics if the data types of `lhs` and `rhs` are not matched.
527pub fn cmp_datum(lhs: impl ToDatumRef, rhs: impl ToDatumRef, order_type: OrderType) -> Ordering {
528    let lhs = lhs.to_datum_ref();
529    let rhs = rhs.to_datum_ref();
530    partial_cmp_datum(lhs, rhs, order_type)
531        .unwrap_or_else(|| panic!("cannot compare {lhs:?} with {rhs:?}"))
532}
533
534/// Compare two `Datum` iterators with specified order types.
535pub fn partial_cmp_datum_iter(
536    lhs: impl IntoIterator<Item = impl ToDatumRef>,
537    rhs: impl IntoIterator<Item = impl ToDatumRef>,
538    order_types: impl IntoIterator<Item = OrderType>,
539) -> Option<Ordering> {
540    let mut order_types_iter = order_types.into_iter();
541    lhs.into_iter().partial_cmp_by(rhs, |x, y| {
542        let order_type = order_types_iter.next()?;
543        partial_cmp_datum(x, y, order_type)
544    })
545}
546
547/// Compare two `Datum` iterators with specified order types.
548///
549/// # Panics
550///
551/// Panics if the number of `OrderType`s is smaller than the number of `Datum`s,
552/// or if the data types of `lhs` and `rhs` are not matched.
553pub fn cmp_datum_iter(
554    lhs: impl IntoIterator<Item = impl ToDatumRef>,
555    rhs: impl IntoIterator<Item = impl ToDatumRef>,
556    order_types: impl IntoIterator<Item = OrderType>,
557) -> Ordering {
558    let mut order_types_iter = order_types.into_iter();
559    lhs.into_iter().cmp_by(rhs, |x, y| {
560        let order_type = order_types_iter
561            .next()
562            .expect("number of `OrderType`s is not enough");
563        cmp_datum(x, y, order_type)
564    })
565}
566
567/// Partial compare two `Row`s with specified order types.
568///
569/// NOTE: This function returns `None` if two rows have different schema.
570pub fn partial_cmp_rows(
571    lhs: impl Row,
572    rhs: impl Row,
573    order_types: &[OrderType],
574) -> Option<Ordering> {
575    if lhs.len() != rhs.len() {
576        return None;
577    }
578    lhs.iter()
579        .zip_eq_debug(rhs.iter())
580        .zip_eq_debug(order_types)
581        .try_fold(Ordering::Equal, |acc, ((l, r), order_type)| match acc {
582            Ordering::Equal => partial_cmp_datum(l, r, *order_type),
583            acc => Some(acc),
584        })
585}
586
587/// Compare two `Row`s with specified order types.
588///
589/// # Panics
590///
591/// Panics if the length of `lhs`, `rhs` and `order_types` are not equal,
592/// or, if the schemas of `lhs` and `rhs` are not matched.
593pub fn cmp_rows(lhs: impl Row, rhs: impl Row, order_types: &[OrderType]) -> Ordering {
594    assert_eq!(lhs.len(), rhs.len());
595    lhs.iter()
596        .zip_eq_debug(rhs.iter())
597        .zip_eq_debug(order_types)
598        .fold(Ordering::Equal, |acc, ((l, r), order_type)| match acc {
599            Ordering::Equal => cmp_datum(l, r, *order_type),
600            acc => acc,
601        })
602}
603
604/// Compare two rows column-by-column, assuming all columns are in ascending order.
605/// This function avoids the allocation of `order_types` before each call.
606///
607/// # Panics
608///
609/// See [`cmp_rows`]
610pub fn cmp_rows_ascending(lhs: impl Row, rhs: impl Row) -> Ordering {
611    assert_eq!(lhs.len(), rhs.len());
612    let order_type = OrderType::ascending();
613    lhs.iter()
614        .zip_eq_debug(rhs.iter())
615        .fold(Ordering::Equal, |acc, (l, r)| match acc {
616            Ordering::Equal => cmp_datum(l, r, order_type),
617            acc => acc,
618        })
619}
620
621#[cfg(test)]
622mod tests {
623    use itertools::Itertools;
624
625    use super::*;
626    use crate::array::{ListValue, StructValue};
627    use crate::row::OwnedRow;
628    use crate::types::{DataType, Datum, ScalarImpl, StructType};
629
630    #[test]
631    fn test_order_type() {
632        assert_eq!(OrderType::default(), OrderType::ascending());
633        assert_eq!(
634            OrderType::default(),
635            OrderType::new(Direction::Ascending, NullsAre::Largest)
636        );
637        assert_eq!(
638            OrderType::default(),
639            OrderType::from_bools(Some(true), Some(false))
640        );
641        assert_eq!(OrderType::default(), OrderType::from_bools(None, None));
642
643        assert!(OrderType::ascending().is_ascending());
644        assert!(OrderType::ascending().nulls_are_largest());
645        assert!(OrderType::ascending().nulls_are_last());
646
647        assert!(OrderType::descending().is_descending());
648        assert!(OrderType::descending().nulls_are_largest());
649        assert!(OrderType::descending().nulls_are_first());
650
651        assert!(OrderType::ascending_nulls_first().is_ascending());
652        assert!(OrderType::ascending_nulls_first().nulls_are_smallest());
653        assert!(OrderType::ascending_nulls_first().nulls_are_first());
654
655        assert!(OrderType::ascending_nulls_last().is_ascending());
656        assert!(OrderType::ascending_nulls_last().nulls_are_largest());
657        assert!(OrderType::ascending_nulls_last().nulls_are_last());
658
659        assert!(OrderType::descending_nulls_first().is_descending());
660        assert!(OrderType::descending_nulls_first().nulls_are_largest());
661        assert!(OrderType::descending_nulls_first().nulls_are_first());
662
663        assert!(OrderType::descending_nulls_last().is_descending());
664        assert!(OrderType::descending_nulls_last().nulls_are_smallest());
665        assert!(OrderType::descending_nulls_last().nulls_are_last());
666
667        assert_eq!(OrderType::ascending().reverse(), OrderType::descending());
668        assert_eq!(OrderType::descending().reverse(), OrderType::ascending());
669        assert_eq!(
670            OrderType::ascending_nulls_first().reverse(),
671            OrderType::descending_nulls_last()
672        );
673        assert_eq!(
674            OrderType::ascending_nulls_last().reverse(),
675            OrderType::descending_nulls_first()
676        );
677    }
678
679    #[test]
680    fn test_compare_rows_in_chunk() {
681        let v10 = Some(ScalarImpl::Int32(42));
682        let v11 = Some(ScalarImpl::Utf8("hello".into()));
683        let v12 = Some(ScalarImpl::Float32(4.0.into()));
684        let v20 = Some(ScalarImpl::Int32(42));
685        let v21 = Some(ScalarImpl::Utf8("hell".into()));
686        let v22 = Some(ScalarImpl::Float32(3.0.into()));
687
688        let row1 = OwnedRow::new(vec![v10, v11, v12]);
689        let row2 = OwnedRow::new(vec![v20, v21, v22]);
690        let chunk = DataChunk::from_rows(
691            &[row1, row2],
692            &[DataType::Int32, DataType::Varchar, DataType::Float32],
693        );
694        let column_orders = vec![
695            ColumnOrder::new(0, OrderType::ascending()),
696            ColumnOrder::new(1, OrderType::descending()),
697        ];
698
699        assert_eq!(
700            Ordering::Equal,
701            compare_rows_in_chunk(&chunk, 0, &chunk, 0, &column_orders)
702        );
703        assert_eq!(
704            Ordering::Less,
705            compare_rows_in_chunk(&chunk, 0, &chunk, 1, &column_orders)
706        );
707    }
708
709    #[test]
710    fn test_compare_all_types() {
711        let row1 = OwnedRow::new(vec![
712            Some(ScalarImpl::Int16(16)),
713            Some(ScalarImpl::Int32(32)),
714            Some(ScalarImpl::Int64(64)),
715            Some(ScalarImpl::Float32(3.2.into())),
716            Some(ScalarImpl::Float64(6.4.into())),
717            Some(ScalarImpl::Utf8("hello".into())),
718            Some(ScalarImpl::Bool(true)),
719            Some(ScalarImpl::Decimal(10.into())),
720            Some(ScalarImpl::Interval(Default::default())),
721            Some(ScalarImpl::Date(Default::default())),
722            Some(ScalarImpl::Timestamp(Default::default())),
723            Some(ScalarImpl::Time(Default::default())),
724            Some(ScalarImpl::Struct(StructValue::new(vec![
725                Some(ScalarImpl::Int32(1)),
726                Some(ScalarImpl::Float32(3.0.into())),
727            ]))),
728            Some(ScalarImpl::List(ListValue::from_iter([1, 2]))),
729        ]);
730        let row2 = OwnedRow::new(vec![
731            Some(ScalarImpl::Int16(16)),
732            Some(ScalarImpl::Int32(32)),
733            Some(ScalarImpl::Int64(64)),
734            Some(ScalarImpl::Float32(3.2.into())),
735            Some(ScalarImpl::Float64(6.4.into())),
736            Some(ScalarImpl::Utf8("hello".into())),
737            Some(ScalarImpl::Bool(true)),
738            Some(ScalarImpl::Decimal(10.into())),
739            Some(ScalarImpl::Interval(Default::default())),
740            Some(ScalarImpl::Date(Default::default())),
741            Some(ScalarImpl::Timestamp(Default::default())),
742            Some(ScalarImpl::Time(Default::default())),
743            Some(ScalarImpl::Struct(StructValue::new(vec![
744                Some(ScalarImpl::Int32(1)),
745                Some(ScalarImpl::Float32(33333.0.into())), // larger than row1
746            ]))),
747            Some(ScalarImpl::List(ListValue::from_iter([1, 2]))),
748        ]);
749
750        let column_orders = (0..row1.len())
751            .map(|i| ColumnOrder::new(i, OrderType::ascending()))
752            .collect_vec();
753
754        let chunk = DataChunk::from_rows(
755            &[row1, row2],
756            &[
757                DataType::Int16,
758                DataType::Int32,
759                DataType::Int64,
760                DataType::Float32,
761                DataType::Float64,
762                DataType::Varchar,
763                DataType::Boolean,
764                DataType::Decimal,
765                DataType::Interval,
766                DataType::Date,
767                DataType::Timestamp,
768                DataType::Time,
769                StructType::unnamed(vec![DataType::Int32, DataType::Float32]).into(),
770                DataType::Int32.list(),
771            ],
772        );
773        assert_eq!(
774            Ordering::Equal,
775            compare_rows_in_chunk(&chunk, 0, &chunk, 0, &column_orders)
776        );
777        assert_eq!(
778            Ordering::Less,
779            compare_rows_in_chunk(&chunk, 0, &chunk, 1, &column_orders)
780        );
781    }
782
783    fn common_compare_datum<CmpFn>(compare: CmpFn)
784    where
785        CmpFn: Fn(Datum, Datum, OrderType) -> Ordering,
786    {
787        assert_eq!(
788            Ordering::Equal,
789            compare(Some(42.into()), Some(42.into()), OrderType::default(),)
790        );
791        assert_eq!(Ordering::Equal, compare(None, None, OrderType::default(),));
792        assert_eq!(
793            Ordering::Less,
794            compare(Some(42.into()), Some(100.into()), OrderType::ascending(),)
795        );
796        assert_eq!(
797            Ordering::Greater,
798            compare(Some(42.into()), None, OrderType::ascending_nulls_first(),)
799        );
800        assert_eq!(
801            Ordering::Less,
802            compare(Some(42.into()), None, OrderType::ascending_nulls_last(),)
803        );
804        assert_eq!(
805            Ordering::Greater,
806            compare(Some(42.into()), None, OrderType::descending_nulls_first(),)
807        );
808        assert_eq!(
809            Ordering::Less,
810            compare(Some(42.into()), None, OrderType::descending_nulls_last(),)
811        );
812    }
813
814    fn common_compare_rows<CmpFn>(compare: CmpFn)
815    where
816        CmpFn: Fn(OwnedRow, OwnedRow, &[OrderType]) -> Ordering,
817    {
818        assert_eq!(
819            Ordering::Equal,
820            compare(
821                OwnedRow::new(vec![Some(42.into()), Some(42.into())]),
822                OwnedRow::new(vec![Some(42.into()), Some(42.into())]),
823                &[OrderType::ascending(), OrderType::ascending()],
824            )
825        );
826
827        assert_eq!(
828            Ordering::Greater,
829            compare(
830                OwnedRow::new(vec![Some(42.into()), Some(42.into())]),
831                OwnedRow::new(vec![Some(42.into()), Some(100.into())]),
832                &[OrderType::ascending(), OrderType::descending()],
833            )
834        );
835
836        assert_eq!(
837            Ordering::Less,
838            compare(
839                OwnedRow::new(vec![Some(42.into()), Some(42.into())]),
840                OwnedRow::new(vec![Some(42.into()), Some(100.into())]),
841                &[OrderType::ascending(), OrderType::ascending()],
842            )
843        );
844    }
845
846    #[test]
847    fn test_compare_datum() {
848        common_compare_datum(cmp_datum);
849    }
850
851    #[test]
852    fn test_compare_rows() {
853        common_compare_rows(cmp_rows);
854    }
855
856    #[test]
857    fn test_partial_compare_datum() {
858        common_compare_datum(|lhs, rhs, order| partial_cmp_datum(lhs, rhs, order).unwrap());
859
860        assert_eq!(
861            None,
862            partial_cmp_datum(
863                Some(ScalarImpl::from(42)),
864                Some(ScalarImpl::from("abc")),
865                OrderType::default()
866            )
867        )
868    }
869
870    #[test]
871    fn test_partial_compare_rows() {
872        common_compare_rows(|lhs, rhs, orders| partial_cmp_rows(lhs, rhs, orders).unwrap());
873
874        assert_eq!(
875            None,
876            partial_cmp_rows(
877                OwnedRow::new(vec![Some(42.into())]),
878                OwnedRow::new(vec![Some("abc".into())]),
879                &[OrderType::default()]
880            )
881        )
882    }
883}
risingwave_common/util/sort_util.rs

risingwave_common/util/
sort_util.rs
