common_query/logical_plan/
expr.rs

1// Copyright 2023 Greptime Team
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 common_time::Timestamp;
16use common_time::range::TimestampRange;
17use common_time::timestamp::TimeUnit;
18use datafusion::arrow::datatypes::DataType as ArrowDataType;
19use datafusion_common::{Column, ScalarValue};
20use datafusion_expr::expr::Expr;
21use datafusion_expr::{Operator, and, binary_expr};
22use datatypes::data_type::DataType;
23use datatypes::schema::ColumnSchema;
24use datatypes::value::Value;
25
26/// Whether a scalar type is in the conservative lossless Substrait transport subset.
27///
28/// In particular, timestamp timezone identifiers are lost by Substrait even when
29/// encoding succeeds. Additional Arrow types need transport coverage first.
30pub fn is_lossless_substrait_type(data_type: &ArrowDataType) -> bool {
31    matches!(
32        data_type,
33        ArrowDataType::Null
34            | ArrowDataType::Boolean
35            | ArrowDataType::Int8
36            | ArrowDataType::Int16
37            | ArrowDataType::Int32
38            | ArrowDataType::Int64
39            | ArrowDataType::UInt8
40            | ArrowDataType::UInt16
41            | ArrowDataType::UInt32
42            | ArrowDataType::UInt64
43            | ArrowDataType::Float32
44            | ArrowDataType::Float64
45            | ArrowDataType::Utf8
46            | ArrowDataType::Binary
47            | ArrowDataType::Date32
48            | ArrowDataType::Timestamp(_, None)
49            | ArrowDataType::Decimal128(_, _)
50    )
51}
52
53/// Builds a filter for a timestamp column with the same type as the timestamp column.
54/// Returns [None] if time range is [None] or full time range.
55pub fn build_same_type_ts_filter(
56    ts_schema: &ColumnSchema,
57    time_range: Option<TimestampRange>,
58) -> Option<Expr> {
59    let time_range = time_range?;
60    let start = time_range
61        .start()
62        .and_then(|start| ts_schema.data_type.try_cast(Value::Timestamp(start)));
63    let end = time_range
64        .end()
65        .and_then(|end| ts_schema.data_type.try_cast(Value::Timestamp(end)));
66
67    let time_range = match (start, end) {
68        (Some(Value::Timestamp(start)), Some(Value::Timestamp(end))) => {
69            TimestampRange::new(start, end)
70        }
71        (Some(Value::Timestamp(start)), None) => Some(TimestampRange::from_start(start)),
72        (None, Some(Value::Timestamp(end))) => Some(TimestampRange::until_end(end, false)),
73        _ => return None,
74    };
75    build_filter_from_timestamp(&ts_schema.name, time_range.as_ref())
76}
77
78/// Builds an `Expr` that filters timestamp column from given timestamp range.
79/// Returns [None] if time range is [None] or full time range.
80pub fn build_filter_from_timestamp(
81    ts_col_name: &str,
82    time_range: Option<&TimestampRange>,
83) -> Option<Expr> {
84    let time_range = time_range?;
85    let ts_col_expr = Expr::Column(Column::from_name(ts_col_name));
86
87    match (time_range.start(), time_range.end()) {
88        (None, None) => None,
89        (Some(start), None) => Some(binary_expr(
90            ts_col_expr,
91            Operator::GtEq,
92            timestamp_to_literal(start),
93        )),
94        (None, Some(end)) => Some(binary_expr(
95            ts_col_expr,
96            Operator::Lt,
97            timestamp_to_literal(end),
98        )),
99        (Some(start), Some(end)) => Some(and(
100            binary_expr(
101                ts_col_expr.clone(),
102                Operator::GtEq,
103                timestamp_to_literal(start),
104            ),
105            binary_expr(ts_col_expr, Operator::Lt, timestamp_to_literal(end)),
106        )),
107    }
108}
109
110/// Converts a [Timestamp] to datafusion literal value.
111fn timestamp_to_literal(timestamp: &Timestamp) -> Expr {
112    let scalar_value = match timestamp.unit() {
113        TimeUnit::Second => ScalarValue::TimestampSecond(Some(timestamp.value()), None),
114        TimeUnit::Millisecond => ScalarValue::TimestampMillisecond(Some(timestamp.value()), None),
115        TimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(Some(timestamp.value()), None),
116        TimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(Some(timestamp.value()), None),
117    };
118    Expr::Literal(scalar_value, None)
119}
120
121#[cfg(test)]
122mod tests {
123    use datafusion_expr::Literal;
124
125    use super::*;
126
127    #[test]
128    fn test_timestamp_to_literal() {
129        let timestamp = Timestamp::new(123456789, TimeUnit::Second);
130        let expected = ScalarValue::TimestampSecond(Some(123456789), None).lit();
131        assert_eq!(timestamp_to_literal(&timestamp), expected);
132
133        let timestamp = Timestamp::new(123456789, TimeUnit::Millisecond);
134        let expected = ScalarValue::TimestampMillisecond(Some(123456789), None).lit();
135        assert_eq!(timestamp_to_literal(&timestamp), expected);
136
137        let timestamp = Timestamp::new(123456789, TimeUnit::Microsecond);
138        let expected = ScalarValue::TimestampMicrosecond(Some(123456789), None).lit();
139        assert_eq!(timestamp_to_literal(&timestamp), expected);
140
141        let timestamp = Timestamp::new(123456789, TimeUnit::Nanosecond);
142        let expected = ScalarValue::TimestampNanosecond(Some(123456789), None).lit();
143        assert_eq!(timestamp_to_literal(&timestamp), expected);
144    }
145}
common_query/logical_plan/expr.rs

common_query/logical_plan/
expr.rs
