mito2/sst/index/inverted_index/applier/builder/
comparison.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::{Expr as DfExpr, Operator};
16use index::Bytes;
17use index::inverted_index::search::predicate::{Bound, Predicate, Range, RangePredicate};
18
19use crate::error::Result;
20use crate::sst::index::inverted_index::applier::builder::InvertedIndexApplierBuilder;
21
22impl InvertedIndexApplierBuilder<'_> {
23    /// Collects `<`, `<=`, `>`, or `>=` comparisons between an indexed column or
24    /// JSON path expression and a non-null literal.
25    ///
26    /// The literal may appear on either side; comparisons with the literal on
27    /// the left are reversed before building the index predicate.
28    ///
29    /// For example, `42 < col` becomes `col > 42`, and
30    /// `42 < json_get(attrs, '$.value')` becomes `json_get(attrs, '$.value') > 42`.
31    pub(crate) fn collect_comparison_expr(
32        &mut self,
33        left: &DfExpr,
34        op: &Operator,
35        right: &DfExpr,
36    ) -> Result<()> {
37        match op {
38            Operator::Lt => {
39                if Self::nonnull_lit(left).is_some() {
40                    self.collect_column_gt_lit(right, left)
41                } else {
42                    self.collect_column_lt_lit(left, right)
43                }
44            }
45            Operator::LtEq => {
46                if Self::nonnull_lit(left).is_some() {
47                    self.collect_column_ge_lit(right, left)
48                } else {
49                    self.collect_column_le_lit(left, right)
50                }
51            }
52            Operator::Gt => {
53                if Self::nonnull_lit(left).is_some() {
54                    self.collect_column_lt_lit(right, left)
55                } else {
56                    self.collect_column_gt_lit(left, right)
57                }
58            }
59            Operator::GtEq => {
60                if Self::nonnull_lit(left).is_some() {
61                    self.collect_column_le_lit(right, left)
62                } else {
63                    self.collect_column_ge_lit(left, right)
64                }
65            }
66            _ => Ok(()),
67        }
68    }
69
70    fn collect_column_lt_lit(&mut self, left: &DfExpr, right: &DfExpr) -> Result<()> {
71        self.collect_column_cmp_lit(left, right, |value| Range {
72            lower: None,
73            upper: Some(Bound {
74                inclusive: false,
75                value,
76            }),
77        })
78    }
79
80    fn collect_column_gt_lit(&mut self, left: &DfExpr, right: &DfExpr) -> Result<()> {
81        self.collect_column_cmp_lit(left, right, |value| Range {
82            lower: Some(Bound {
83                inclusive: false,
84                value,
85            }),
86            upper: None,
87        })
88    }
89
90    fn collect_column_le_lit(&mut self, left: &DfExpr, right: &DfExpr) -> Result<()> {
91        self.collect_column_cmp_lit(left, right, |value| Range {
92            lower: None,
93            upper: Some(Bound {
94                inclusive: true,
95                value,
96            }),
97        })
98    }
99
100    fn collect_column_ge_lit(&mut self, left: &DfExpr, right: &DfExpr) -> Result<()> {
101        self.collect_column_cmp_lit(left, right, |value| Range {
102            lower: Some(Bound {
103                inclusive: true,
104                value,
105            }),
106            upper: None,
107        })
108    }
109
110    fn collect_column_cmp_lit(
111        &mut self,
112        column: &DfExpr,
113        literal: &DfExpr,
114        range_builder: impl FnOnce(Bytes) -> Range,
115    ) -> Result<()> {
116        let Some(lit) = Self::nonnull_lit(literal) else {
117            return Ok(());
118        };
119        let Some((target, data_type)) = self.expr_to_index_target(column)? else {
120            return Ok(());
121        };
122
123        let predicate = Predicate::Range(RangePredicate {
124            range: range_builder(Self::encode_lit(lit, data_type)?),
125        });
126
127        self.add_predicate(target, predicate);
128        Ok(())
129    }
130}
131
132#[cfg(test)]
133mod tests {
134
135    use std::collections::HashSet;
136
137    use store_api::region_request::PathType;
138
139    use super::*;
140    use crate::error::Error;
141    use crate::sst::index::inverted_index::applier::builder::tests::{
142        encoded_string, field_column, int64_lit, nonexistent_column, string_lit, tag_column,
143        test_object_store, test_region_metadata,
144    };
145    use crate::sst::index::puffin_manager::PuffinManagerFactory;
146
147    #[test]
148    fn test_collect_comparison_basic() {
149        let cases = [
150            (
151                (&tag_column(), &Operator::Lt, &string_lit("123")),
152                Range {
153                    lower: None,
154                    upper: Some(Bound {
155                        inclusive: false,
156                        value: encoded_string("123"),
157                    }),
158                },
159            ),
160            (
161                (&string_lit("123"), &Operator::Lt, &tag_column()),
162                Range {
163                    lower: Some(Bound {
164                        inclusive: false,
165                        value: encoded_string("123"),
166                    }),
167                    upper: None,
168                },
169            ),
170            (
171                (&tag_column(), &Operator::LtEq, &string_lit("123")),
172                Range {
173                    lower: None,
174                    upper: Some(Bound {
175                        inclusive: true,
176                        value: encoded_string("123"),
177                    }),
178                },
179            ),
180            (
181                (&string_lit("123"), &Operator::LtEq, &tag_column()),
182                Range {
183                    lower: Some(Bound {
184                        inclusive: true,
185                        value: encoded_string("123"),
186                    }),
187                    upper: None,
188                },
189            ),
190            (
191                (&tag_column(), &Operator::Gt, &string_lit("123")),
192                Range {
193                    lower: Some(Bound {
194                        inclusive: false,
195                        value: encoded_string("123"),
196                    }),
197                    upper: None,
198                },
199            ),
200            (
201                (&string_lit("123"), &Operator::Gt, &tag_column()),
202                Range {
203                    lower: None,
204                    upper: Some(Bound {
205                        inclusive: false,
206                        value: encoded_string("123"),
207                    }),
208                },
209            ),
210            (
211                (&tag_column(), &Operator::GtEq, &string_lit("123")),
212                Range {
213                    lower: Some(Bound {
214                        inclusive: true,
215                        value: encoded_string("123"),
216                    }),
217                    upper: None,
218                },
219            ),
220            (
221                (&string_lit("123"), &Operator::GtEq, &tag_column()),
222                Range {
223                    lower: None,
224                    upper: Some(Bound {
225                        inclusive: true,
226                        value: encoded_string("123"),
227                    }),
228                },
229            ),
230        ];
231
232        let (_d, factory) =
233            PuffinManagerFactory::new_for_test_block("test_collect_comparison_basic_");
234        let metadata = test_region_metadata();
235        let mut builder = InvertedIndexApplierBuilder::new(
236            "test".to_string(),
237            PathType::Bare,
238            test_object_store(),
239            &metadata,
240            HashSet::from_iter([1, 2, 3]),
241            factory,
242        );
243
244        for ((left, op, right), _) in &cases {
245            builder.collect_comparison_expr(left, op, right).unwrap();
246        }
247
248        let predicates = builder
249            .output
250            .get(&index::target::IndexTarget::ColumnId(1))
251            .unwrap();
252        assert_eq!(predicates.len(), cases.len());
253        for ((_, expected), actual) in cases.into_iter().zip(predicates) {
254            assert_eq!(
255                actual,
256                &Predicate::Range(RangePredicate { range: expected })
257            );
258        }
259    }
260
261    #[test]
262    fn test_collect_comparison_type_mismatch() {
263        let (_d, factory) =
264            PuffinManagerFactory::new_for_test_block("test_collect_comparison_type_mismatch_");
265        let metadata = test_region_metadata();
266        let mut builder = InvertedIndexApplierBuilder::new(
267            "test".to_string(),
268            PathType::Bare,
269            test_object_store(),
270            &metadata,
271            HashSet::from_iter([1, 2, 3]),
272            factory,
273        );
274
275        let res = builder.collect_comparison_expr(&tag_column(), &Operator::Lt, &int64_lit(10));
276        assert!(matches!(res, Err(Error::Encode { .. })));
277        assert!(builder.output.is_empty());
278    }
279
280    #[test]
281    fn test_collect_comparison_field_column() {
282        let (_d, factory) =
283            PuffinManagerFactory::new_for_test_block("test_collect_comparison_field_column_");
284        let metadata = test_region_metadata();
285        let mut builder = InvertedIndexApplierBuilder::new(
286            "test".to_string(),
287            PathType::Bare,
288            test_object_store(),
289            &metadata,
290            HashSet::from_iter([1, 2, 3]),
291            factory,
292        );
293
294        builder
295            .collect_comparison_expr(&field_column(), &Operator::Lt, &string_lit("abc"))
296            .unwrap();
297
298        let predicates = builder
299            .output
300            .get(&index::target::IndexTarget::ColumnId(3))
301            .unwrap();
302        assert_eq!(predicates.len(), 1);
303        assert_eq!(
304            predicates[0],
305            Predicate::Range(RangePredicate {
306                range: Range {
307                    lower: None,
308                    upper: Some(Bound {
309                        inclusive: false,
310                        value: encoded_string("abc"),
311                    }),
312                }
313            })
314        );
315    }
316
317    #[test]
318    fn test_collect_comparison_nonexistent_column() {
319        let (_d, factory) =
320            PuffinManagerFactory::new_for_test_block("test_collect_comparison_nonexistent_column_");
321        let metadata = test_region_metadata();
322        let mut builder = InvertedIndexApplierBuilder::new(
323            "test".to_string(),
324            PathType::Bare,
325            test_object_store(),
326            &metadata,
327            HashSet::from_iter([1, 2, 3]),
328            factory,
329        );
330
331        let res = builder.collect_comparison_expr(
332            &nonexistent_column(),
333            &Operator::Lt,
334            &string_lit("abc"),
335        );
336        assert!(matches!(res, Err(Error::ColumnNotFound { .. })));
337        assert!(builder.output.is_empty());
338    }
339}
mito2/sst/index/inverted_index/applier/builder/comparison.rs

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