mito2/sst/index/inverted_index/applier/builder/
eq_list.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 std::collections::BTreeSet;
16
17use datafusion_expr::{BinaryExpr, Expr as DfExpr, Operator};
18use datatypes::data_type::ConcreteDataType;
19use index::Bytes;
20use index::inverted_index::search::predicate::{InListPredicate, Predicate};
21use index::target::IndexTarget;
22
23use crate::error::Result;
24use crate::sst::index::inverted_index::applier::builder::InvertedIndexApplierBuilder;
25
26impl InvertedIndexApplierBuilder<'_> {
27    /// Collects an eq expression in the form of `column = lit`.
28    pub(crate) fn collect_eq(&mut self, left: &DfExpr, right: &DfExpr) -> Result<()> {
29        let resolved = match self.expr_to_index_target(left)? {
30            Some(target) => Some(target),
31            None => self.expr_to_index_target(right)?,
32        };
33        let Some((target, data_type)) = resolved else {
34            return Ok(());
35        };
36        let Some(lit) = Self::nonnull_lit(right).or_else(|| Self::nonnull_lit(left)) else {
37            return Ok(());
38        };
39
40        let predicate = Predicate::InList(InListPredicate {
41            list: BTreeSet::from_iter([Self::encode_lit(lit, data_type)?]),
42        });
43        self.add_predicate(target, predicate);
44        Ok(())
45    }
46
47    /// Collects eq list in the form of `column = lit OR column = lit OR ...`.
48    pub(crate) fn collect_or_eq_list(&mut self, eq_expr: &DfExpr, or_list: &DfExpr) -> Result<()> {
49        let DfExpr::BinaryExpr(BinaryExpr {
50            left,
51            op: Operator::Eq,
52            right,
53        }) = eq_expr
54        else {
55            return Ok(());
56        };
57
58        let resolved = match self.expr_to_index_target(left)? {
59            Some(target) => Some(target),
60            None => self.expr_to_index_target(right)?,
61        };
62        let Some((target, data_type)) = resolved else {
63            return Ok(());
64        };
65        let Some(lit) = Self::nonnull_lit(right).or_else(|| Self::nonnull_lit(left)) else {
66            return Ok(());
67        };
68
69        let bytes = Self::encode_lit(lit, data_type.clone())?;
70        let mut inlist = BTreeSet::from_iter([bytes]);
71
72        if self.collect_eq_list_inner(&target, &data_type, or_list, &mut inlist)? {
73            let predicate = Predicate::InList(InListPredicate { list: inlist });
74            self.add_predicate(target, predicate);
75        }
76
77        Ok(())
78    }
79
80    /// Recursively collects eq list.
81    ///
82    /// Returns false if the expression doesn't match the form then
83    /// caller can safely ignore the expression.
84    fn collect_eq_list_inner(
85        &self,
86        target: &IndexTarget,
87        data_type: &ConcreteDataType,
88        expr: &DfExpr,
89        inlist: &mut BTreeSet<Bytes>,
90    ) -> Result<bool> {
91        let DfExpr::BinaryExpr(BinaryExpr {
92            left,
93            op: op @ (Operator::Eq | Operator::Or),
94            right,
95        }) = expr
96        else {
97            return Ok(false);
98        };
99
100        if op == &Operator::Or {
101            let r = self
102                .collect_eq_list_inner(target, data_type, left, inlist)?
103                .then(|| self.collect_eq_list_inner(target, data_type, right, inlist))
104                .transpose()?
105                .unwrap_or(false);
106            return Ok(r);
107        }
108
109        if op == &Operator::Eq {
110            let resolved = match self.expr_to_index_target(left)? {
111                Some(target) => Some(target),
112                None => self.expr_to_index_target(right)?,
113            };
114            let Some((candidate, _)) = resolved else {
115                return Ok(false);
116            };
117            if target != &candidate {
118                return Ok(false);
119            }
120            let Some(lit) = Self::nonnull_lit(right).or_else(|| Self::nonnull_lit(left)) else {
121                return Ok(false);
122            };
123
124            inlist.insert(Self::encode_lit(lit, data_type.clone())?);
125            return Ok(true);
126        }
127
128        Ok(false)
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use std::collections::HashSet;
135
136    use store_api::region_request::PathType;
137
138    use super::*;
139    use crate::error::Error;
140    use crate::sst::index::inverted_index::applier::builder::tests::{
141        encoded_string, field_column, int64_lit, nonexistent_column, string_lit, tag_column,
142        tag_column2, test_object_store, test_region_metadata,
143    };
144    use crate::sst::index::puffin_manager::PuffinManagerFactory;
145
146    #[test]
147    fn test_collect_eq_basic() {
148        let (_d, factory) = PuffinManagerFactory::new_for_test_block("test_collect_eq_basic_");
149        let metadata = test_region_metadata();
150        let mut builder = InvertedIndexApplierBuilder::new(
151            "test".to_string(),
152            PathType::Bare,
153            test_object_store(),
154            &metadata,
155            HashSet::from_iter([1, 2, 3]),
156            factory,
157        );
158
159        builder
160            .collect_eq(&tag_column(), &string_lit("foo"))
161            .unwrap();
162        builder
163            .collect_eq(&string_lit("bar"), &tag_column())
164            .unwrap();
165
166        let predicates = builder
167            .output
168            .get(&index::target::IndexTarget::ColumnId(1))
169            .unwrap();
170        assert_eq!(predicates.len(), 2);
171        assert_eq!(
172            predicates[0],
173            Predicate::InList(InListPredicate {
174                list: BTreeSet::from_iter([encoded_string("foo")])
175            })
176        );
177        assert_eq!(
178            predicates[1],
179            Predicate::InList(InListPredicate {
180                list: BTreeSet::from_iter([encoded_string("bar")])
181            })
182        );
183    }
184
185    #[test]
186    fn test_collect_eq_field_column() {
187        let (_d, factory) =
188            PuffinManagerFactory::new_for_test_block("test_collect_eq_field_column_");
189        let metadata = test_region_metadata();
190        let mut builder = InvertedIndexApplierBuilder::new(
191            "test".to_string(),
192            PathType::Bare,
193            test_object_store(),
194            &metadata,
195            HashSet::from_iter([1, 2, 3]),
196            factory,
197        );
198
199        builder
200            .collect_eq(&field_column(), &string_lit("abc"))
201            .unwrap();
202
203        let predicates = builder
204            .output
205            .get(&index::target::IndexTarget::ColumnId(3))
206            .unwrap();
207        assert_eq!(predicates.len(), 1);
208        assert_eq!(
209            predicates[0],
210            Predicate::InList(InListPredicate {
211                list: BTreeSet::from_iter([encoded_string("abc")])
212            })
213        );
214    }
215
216    #[test]
217    fn test_collect_eq_nonexistent_column() {
218        let (_d, factory) =
219            PuffinManagerFactory::new_for_test_block("test_collect_eq_nonexistent_column_");
220        let metadata = test_region_metadata();
221        let mut builder = InvertedIndexApplierBuilder::new(
222            "test".to_string(),
223            PathType::Bare,
224            test_object_store(),
225            &metadata,
226            HashSet::from_iter([1, 2, 3]),
227            factory,
228        );
229
230        let res = builder.collect_eq(&nonexistent_column(), &string_lit("abc"));
231        assert!(matches!(res, Err(Error::ColumnNotFound { .. })));
232        assert!(builder.output.is_empty());
233    }
234
235    #[test]
236    fn test_collect_eq_type_mismatch() {
237        let (_d, factory) =
238            PuffinManagerFactory::new_for_test_block("test_collect_eq_type_mismatch_");
239        let metadata = test_region_metadata();
240        let mut builder = InvertedIndexApplierBuilder::new(
241            "test".to_string(),
242            PathType::Bare,
243            test_object_store(),
244            &metadata,
245            HashSet::from_iter([1, 2, 3]),
246            factory,
247        );
248
249        let res = builder.collect_eq(&tag_column(), &int64_lit(1));
250        assert!(matches!(res, Err(Error::Encode { .. })));
251        assert!(builder.output.is_empty());
252    }
253
254    #[test]
255    fn test_collect_or_eq_list_basic() {
256        let (_d, factory) =
257            PuffinManagerFactory::new_for_test_block("test_collect_or_eq_list_basic_");
258        let metadata = test_region_metadata();
259        let mut builder = InvertedIndexApplierBuilder::new(
260            "test".to_string(),
261            PathType::Bare,
262            test_object_store(),
263            &metadata,
264            HashSet::from_iter([1, 2, 3]),
265            factory,
266        );
267
268        let eq_expr = DfExpr::BinaryExpr(BinaryExpr {
269            left: Box::new(tag_column()),
270            op: Operator::Eq,
271            right: Box::new(string_lit("abc")),
272        });
273        let or_eq_list = DfExpr::BinaryExpr(BinaryExpr {
274            left: Box::new(DfExpr::BinaryExpr(BinaryExpr {
275                left: Box::new(tag_column()),
276                op: Operator::Eq,
277                right: Box::new(string_lit("foo")),
278            })),
279            op: Operator::Or,
280            right: Box::new(DfExpr::BinaryExpr(BinaryExpr {
281                left: Box::new(DfExpr::BinaryExpr(BinaryExpr {
282                    left: Box::new(tag_column()),
283                    op: Operator::Eq,
284                    right: Box::new(string_lit("bar")),
285                })),
286                op: Operator::Or,
287                right: Box::new(DfExpr::BinaryExpr(BinaryExpr {
288                    left: Box::new(tag_column()),
289                    op: Operator::Eq,
290                    right: Box::new(string_lit("baz")),
291                })),
292            })),
293        });
294
295        builder.collect_or_eq_list(&eq_expr, &or_eq_list).unwrap();
296
297        let predicates = builder
298            .output
299            .get(&index::target::IndexTarget::ColumnId(1))
300            .unwrap();
301        assert_eq!(predicates.len(), 1);
302        assert_eq!(
303            predicates[0],
304            Predicate::InList(InListPredicate {
305                list: BTreeSet::from_iter([
306                    encoded_string("abc"),
307                    encoded_string("foo"),
308                    encoded_string("bar"),
309                    encoded_string("baz")
310                ])
311            })
312        );
313    }
314
315    #[test]
316    fn test_collect_or_eq_list_invalid_op() {
317        let (_d, factory) =
318            PuffinManagerFactory::new_for_test_block("test_collect_or_eq_list_invalid_op_");
319        let metadata = test_region_metadata();
320        let mut builder = InvertedIndexApplierBuilder::new(
321            "test".to_string(),
322            PathType::Bare,
323            test_object_store(),
324            &metadata,
325            HashSet::from_iter([1, 2, 3]),
326            factory,
327        );
328
329        let eq_expr = DfExpr::BinaryExpr(BinaryExpr {
330            left: Box::new(tag_column()),
331            op: Operator::Eq,
332            right: Box::new(string_lit("abc")),
333        });
334        let or_eq_list = DfExpr::BinaryExpr(BinaryExpr {
335            left: Box::new(DfExpr::BinaryExpr(BinaryExpr {
336                left: Box::new(tag_column()),
337                op: Operator::Eq,
338                right: Box::new(string_lit("foo")),
339            })),
340            op: Operator::Or,
341            right: Box::new(DfExpr::BinaryExpr(BinaryExpr {
342                left: Box::new(tag_column()),
343                op: Operator::Gt, // invalid op
344                right: Box::new(string_lit("foo")),
345            })),
346        });
347
348        builder.collect_or_eq_list(&eq_expr, &or_eq_list).unwrap();
349        assert!(builder.output.is_empty());
350    }
351
352    #[test]
353    fn test_collect_or_eq_list_multiple_columns() {
354        let (_d, factory) =
355            PuffinManagerFactory::new_for_test_block("test_collect_or_eq_list_multiple_columns_");
356        let metadata = test_region_metadata();
357        let mut builder = InvertedIndexApplierBuilder::new(
358            "test".to_string(),
359            PathType::Bare,
360            test_object_store(),
361            &metadata,
362            HashSet::from_iter([1, 2, 3]),
363            factory,
364        );
365
366        let eq_expr = DfExpr::BinaryExpr(BinaryExpr {
367            left: Box::new(tag_column()),
368            op: Operator::Eq,
369            right: Box::new(string_lit("abc")),
370        });
371        let or_eq_list = DfExpr::BinaryExpr(BinaryExpr {
372            left: Box::new(DfExpr::BinaryExpr(BinaryExpr {
373                left: Box::new(tag_column()),
374                op: Operator::Eq,
375                right: Box::new(string_lit("foo")),
376            })),
377            op: Operator::Or,
378            right: Box::new(DfExpr::BinaryExpr(BinaryExpr {
379                left: Box::new(tag_column2()), // different column
380                op: Operator::Eq,
381                right: Box::new(string_lit("bar")),
382            })),
383        });
384
385        builder.collect_or_eq_list(&eq_expr, &or_eq_list).unwrap();
386        assert!(builder.output.is_empty());
387    }
388}
mito2/sst/index/inverted_index/applier/builder/eq_list.rs

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