Coverage Report

Created: 2024-10-13 08:39

/Users/andrewlamb/Software/datafusion/datafusion/expr-common/src/sort_properties.rs
Line
Count
Source (jump to first uncovered line)
1
// Licensed to the Apache Software Foundation (ASF) under one
2
// or more contributor license agreements.  See the NOTICE file
3
// distributed with this work for additional information
4
// regarding copyright ownership.  The ASF licenses this file
5
// to you under the Apache License, Version 2.0 (the
6
// "License"); you may not use this file except in compliance
7
// with the License.  You may obtain a copy of the License at
8
//
9
//   http://www.apache.org/licenses/LICENSE-2.0
10
//
11
// Unless required by applicable law or agreed to in writing,
12
// software distributed under the License is distributed on an
13
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14
// KIND, either express or implied.  See the License for the
15
// specific language governing permissions and limitations
16
// under the License.
17
18
use std::ops::Neg;
19
20
use crate::interval_arithmetic::Interval;
21
22
use arrow::compute::SortOptions;
23
use arrow::datatypes::DataType;
24
25
/// To propagate [`SortOptions`] across the `PhysicalExpr`, it is insufficient
26
/// to simply use `Option<SortOptions>`: There must be a differentiation between
27
/// unordered columns and literal values, since literals may not break the ordering
28
/// when they are used as a child of some binary expression when the other child has
29
/// some ordering. On the other hand, unordered columns cannot maintain ordering when
30
/// they take part in such operations.
31
///
32
/// Example: ((a_ordered + b_unordered) + c_ordered) expression cannot end up with
33
/// sorted data; however the ((a_ordered + 999) + c_ordered) expression can. Therefore,
34
/// we need two different variants for literals and unordered columns as literals are
35
/// often more ordering-friendly under most mathematical operations.
36
#[derive(PartialEq, Debug, Clone, Copy, Default)]
37
pub enum SortProperties {
38
    /// Use the ordinary [`SortOptions`] struct to represent ordered data:
39
    Ordered(SortOptions),
40
    // This alternative represents unordered data:
41
    #[default]
42
    Unordered,
43
    // Singleton is used for single-valued literal numbers:
44
    Singleton,
45
}
46
47
impl SortProperties {
48
0
    pub fn add(&self, rhs: &Self) -> Self {
49
0
        match (self, rhs) {
50
0
            (Self::Singleton, _) => *rhs,
51
0
            (_, Self::Singleton) => *self,
52
0
            (Self::Ordered(lhs), Self::Ordered(rhs))
53
0
                if lhs.descending == rhs.descending =>
54
0
            {
55
0
                Self::Ordered(SortOptions {
56
0
                    descending: lhs.descending,
57
0
                    nulls_first: lhs.nulls_first || rhs.nulls_first,
58
                })
59
            }
60
0
            _ => Self::Unordered,
61
        }
62
0
    }
63
64
0
    pub fn sub(&self, rhs: &Self) -> Self {
65
0
        match (self, rhs) {
66
0
            (Self::Singleton, Self::Singleton) => Self::Singleton,
67
0
            (Self::Singleton, Self::Ordered(rhs)) => Self::Ordered(SortOptions {
68
0
                descending: !rhs.descending,
69
0
                nulls_first: rhs.nulls_first,
70
0
            }),
71
0
            (_, Self::Singleton) => *self,
72
0
            (Self::Ordered(lhs), Self::Ordered(rhs))
73
0
                if lhs.descending != rhs.descending =>
74
0
            {
75
0
                Self::Ordered(SortOptions {
76
0
                    descending: lhs.descending,
77
0
                    nulls_first: lhs.nulls_first || rhs.nulls_first,
78
                })
79
            }
80
0
            _ => Self::Unordered,
81
        }
82
0
    }
83
84
0
    pub fn gt_or_gteq(&self, rhs: &Self) -> Self {
85
0
        match (self, rhs) {
86
0
            (Self::Singleton, Self::Ordered(rhs)) => Self::Ordered(SortOptions {
87
0
                descending: !rhs.descending,
88
0
                nulls_first: rhs.nulls_first,
89
0
            }),
90
0
            (_, Self::Singleton) => *self,
91
0
            (Self::Ordered(lhs), Self::Ordered(rhs))
92
0
                if lhs.descending != rhs.descending =>
93
0
            {
94
0
                *self
95
            }
96
0
            _ => Self::Unordered,
97
        }
98
0
    }
99
100
0
    pub fn and_or(&self, rhs: &Self) -> Self {
101
0
        match (self, rhs) {
102
0
            (Self::Ordered(lhs), Self::Ordered(rhs))
103
0
                if lhs.descending == rhs.descending =>
104
0
            {
105
0
                Self::Ordered(SortOptions {
106
0
                    descending: lhs.descending,
107
0
                    nulls_first: lhs.nulls_first || rhs.nulls_first,
108
                })
109
            }
110
0
            (Self::Ordered(opt), Self::Singleton)
111
0
            | (Self::Singleton, Self::Ordered(opt)) => Self::Ordered(SortOptions {
112
0
                descending: opt.descending,
113
0
                nulls_first: opt.nulls_first,
114
0
            }),
115
0
            (Self::Singleton, Self::Singleton) => Self::Singleton,
116
0
            _ => Self::Unordered,
117
        }
118
0
    }
119
}
120
121
impl Neg for SortProperties {
122
    type Output = Self;
123
124
0
    fn neg(mut self) -> Self::Output {
125
0
        if let SortProperties::Ordered(SortOptions { descending, .. }) = &mut self {
126
0
            *descending = !*descending;
127
0
        }
128
0
        self
129
0
    }
130
}
131
132
/// Represents the properties of a `PhysicalExpr`, including its sorting and range attributes.
133
#[derive(Debug, Clone)]
134
pub struct ExprProperties {
135
    pub sort_properties: SortProperties,
136
    pub range: Interval,
137
}
138
139
impl ExprProperties {
140
    /// Creates a new `ExprProperties` instance with unknown sort properties and unknown range.
141
934
    pub fn new_unknown() -> Self {
142
934
        Self {
143
934
            sort_properties: SortProperties::default(),
144
934
            range: Interval::make_unbounded(&DataType::Null).unwrap(),
145
934
        }
146
934
    }
147
148
    /// Sets the sorting properties of the expression and returns the modified instance.
149
0
    pub fn with_order(mut self, order: SortProperties) -> Self {
150
0
        self.sort_properties = order;
151
0
        self
152
0
    }
153
154
    /// Sets the range of the expression and returns the modified instance.
155
0
    pub fn with_range(mut self, range: Interval) -> Self {
156
0
        self.range = range;
157
0
        self
158
0
    }
159
}