mito2/sst/index/inverted_index/applier/builder/
between.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 datafusion_expr::Between;
16use index::inverted_index::search::predicate::{Bound, Predicate, Range, RangePredicate};
17
18use crate::error::Result;
19use crate::sst::index::inverted_index::applier::builder::InvertedIndexApplierBuilder;
20
21impl InvertedIndexApplierBuilder<'_> {
22    /// Collects a `BETWEEN` expression in the form of `column BETWEEN lit AND lit`.
23    pub(crate) fn collect_between(&mut self, between: &Between) -> Result<()> {
24        if between.negated {
25            return Ok(());
26        }
27
28        let Some(column_name) = Self::column_name(&between.expr) else {
29            return Ok(());
30        };
31        let Some((column_id, data_type)) = self.column_id_and_type(column_name)? else {
32            return Ok(());
33        };
34        let Some(low) = Self::nonnull_lit(&between.low) else {
35            return Ok(());
36        };
37        let Some(high) = Self::nonnull_lit(&between.high) else {
38            return Ok(());
39        };
40
41        let predicate = Predicate::Range(RangePredicate {
42            range: Range {
43                lower: Some(Bound {
44                    inclusive: true,
45                    value: Self::encode_lit(low, data_type.clone())?,
46                }),
47                upper: Some(Bound {
48                    inclusive: true,
49                    value: Self::encode_lit(high, data_type)?,
50                }),
51            },
52        });
53
54        self.add_predicate(column_id, predicate);
55        Ok(())
56    }
57}
58
59#[cfg(test)]
60mod tests {
61    use std::collections::HashSet;
62
63    use super::*;
64    use crate::error::Error;
65    use crate::sst::index::inverted_index::applier::builder::tests::{
66        encoded_string, field_column, int64_lit, nonexistent_column, string_lit, tag_column,
67        test_object_store, test_region_metadata,
68    };
69    use crate::sst::index::puffin_manager::PuffinManagerFactory;
70
71    #[test]
72    fn test_collect_between_basic() {
73        let (_d, facotry) = PuffinManagerFactory::new_for_test_block("test_collect_between_basic_");
74        let metadata = test_region_metadata();
75        let mut builder = InvertedIndexApplierBuilder::new(
76            "test".to_string(),
77            test_object_store(),
78            &metadata,
79            HashSet::from_iter([1, 2, 3]),
80            facotry,
81        );
82
83        let between = Between {
84            negated: false,
85            expr: Box::new(tag_column()),
86            low: Box::new(string_lit("abc")),
87            high: Box::new(string_lit("def")),
88        };
89
90        builder.collect_between(&between).unwrap();
91
92        let predicates = builder.output.get(&1).unwrap();
93        assert_eq!(predicates.len(), 1);
94        assert_eq!(
95            predicates[0],
96            Predicate::Range(RangePredicate {
97                range: Range {
98                    lower: Some(Bound {
99                        inclusive: true,
100                        value: encoded_string("abc"),
101                    }),
102                    upper: Some(Bound {
103                        inclusive: true,
104                        value: encoded_string("def"),
105                    }),
106                }
107            })
108        );
109    }
110
111    #[test]
112    fn test_collect_between_negated() {
113        let (_d, facotry) =
114            PuffinManagerFactory::new_for_test_block("test_collect_between_negated_");
115        let metadata = test_region_metadata();
116        let mut builder = InvertedIndexApplierBuilder::new(
117            "test".to_string(),
118            test_object_store(),
119            &metadata,
120            HashSet::from_iter([1, 2, 3]),
121            facotry,
122        );
123
124        let between = Between {
125            negated: true,
126            expr: Box::new(tag_column()),
127            low: Box::new(string_lit("abc")),
128            high: Box::new(string_lit("def")),
129        };
130
131        builder.collect_between(&between).unwrap();
132        assert!(builder.output.is_empty());
133    }
134
135    #[test]
136    fn test_collect_between_field_column() {
137        let (_d, facotry) =
138            PuffinManagerFactory::new_for_test_block("test_collect_between_field_column_");
139        let metadata = test_region_metadata();
140        let mut builder = InvertedIndexApplierBuilder::new(
141            "test".to_string(),
142            test_object_store(),
143            &metadata,
144            HashSet::from_iter([1, 2, 3]),
145            facotry,
146        );
147
148        let between = Between {
149            negated: false,
150            expr: Box::new(field_column()),
151            low: Box::new(string_lit("abc")),
152            high: Box::new(string_lit("def")),
153        };
154
155        builder.collect_between(&between).unwrap();
156
157        let predicates = builder.output.get(&3).unwrap();
158        assert_eq!(predicates.len(), 1);
159        assert_eq!(
160            predicates[0],
161            Predicate::Range(RangePredicate {
162                range: Range {
163                    lower: Some(Bound {
164                        inclusive: true,
165                        value: encoded_string("abc"),
166                    }),
167                    upper: Some(Bound {
168                        inclusive: true,
169                        value: encoded_string("def"),
170                    }),
171                }
172            })
173        );
174    }
175
176    #[test]
177    fn test_collect_between_type_mismatch() {
178        let (_d, facotry) =
179            PuffinManagerFactory::new_for_test_block("test_collect_between_type_mismatch_");
180        let metadata = test_region_metadata();
181        let mut builder = InvertedIndexApplierBuilder::new(
182            "test".to_string(),
183            test_object_store(),
184            &metadata,
185            HashSet::from_iter([1, 2, 3]),
186            facotry,
187        );
188
189        let between = Between {
190            negated: false,
191            expr: Box::new(tag_column()),
192            low: Box::new(int64_lit(123)),
193            high: Box::new(int64_lit(456)),
194        };
195
196        let res = builder.collect_between(&between);
197        assert!(matches!(res, Err(Error::FieldTypeMismatch { .. })));
198        assert!(builder.output.is_empty());
199    }
200
201    #[test]
202    fn test_collect_between_nonexistent_column() {
203        let (_d, facotry) =
204            PuffinManagerFactory::new_for_test_block("test_collect_between_nonexistent_column_");
205        let metadata = test_region_metadata();
206        let mut builder = InvertedIndexApplierBuilder::new(
207            "test".to_string(),
208            test_object_store(),
209            &metadata,
210            HashSet::from_iter([1, 2, 3]),
211            facotry,
212        );
213
214        let between = Between {
215            negated: false,
216            expr: Box::new(nonexistent_column()),
217            low: Box::new(string_lit("abc")),
218            high: Box::new(string_lit("def")),
219        };
220
221        let res = builder.collect_between(&between);
222        assert!(matches!(res, Err(Error::ColumnNotFound { .. })));
223        assert!(builder.output.is_empty());
224    }
225}
mito2/sst/index/inverted_index/applier/builder/between.rs

mito2/sst/index/inverted_index/applier/builder/
between.rs
