mito2/sst/index/inverted_index/applier/builder/
in_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::expr::InList;
18use index::inverted_index::search::predicate::{InListPredicate, Predicate};
19
20use crate::error::Result;
21use crate::sst::index::inverted_index::applier::builder::InvertedIndexApplierBuilder;
22
23impl InvertedIndexApplierBuilder<'_> {
24    /// Collects an in list expression in the form of `column IN (lit, lit, ...)`.
25    pub(crate) fn collect_inlist(&mut self, inlist: &InList) -> Result<()> {
26        if inlist.negated {
27            return Ok(());
28        }
29        let Some(column_name) = Self::column_name(&inlist.expr) else {
30            return Ok(());
31        };
32        let Some((column_id, data_type)) = self.column_id_and_type(column_name)? else {
33            return Ok(());
34        };
35
36        let mut predicate = InListPredicate {
37            list: BTreeSet::new(),
38        };
39        for lit in &inlist.list {
40            let Some(lit) = Self::nonnull_lit(lit) else {
41                return Ok(());
42            };
43
44            predicate
45                .list
46                .insert(Self::encode_lit(lit, data_type.clone())?);
47        }
48
49        self.add_predicate(column_id, Predicate::InList(predicate));
50        Ok(())
51    }
52}
53
54#[cfg(test)]
55mod tests {
56    use std::collections::HashSet;
57
58    use super::*;
59    use crate::error::Error;
60    use crate::sst::index::inverted_index::applier::builder::tests::{
61        encoded_string, field_column, int64_lit, nonexistent_column, string_lit, tag_column,
62        test_object_store, test_region_metadata,
63    };
64    use crate::sst::index::puffin_manager::PuffinManagerFactory;
65
66    #[test]
67    fn test_collect_in_list_basic() {
68        let (_d, facotry) = PuffinManagerFactory::new_for_test_block("test_collect_in_list_basic_");
69        let metadata = test_region_metadata();
70        let mut builder = InvertedIndexApplierBuilder::new(
71            "test".to_string(),
72            test_object_store(),
73            &metadata,
74            HashSet::from_iter([1, 2, 3]),
75            facotry,
76        );
77
78        let in_list = InList {
79            expr: Box::new(tag_column()),
80            list: vec![string_lit("foo"), string_lit("bar")],
81            negated: false,
82        };
83
84        builder.collect_inlist(&in_list).unwrap();
85
86        let predicates = builder.output.get(&1).unwrap();
87        assert_eq!(predicates.len(), 1);
88        assert_eq!(
89            predicates[0],
90            Predicate::InList(InListPredicate {
91                list: BTreeSet::from_iter([encoded_string("foo"), encoded_string("bar")])
92            })
93        );
94    }
95
96    #[test]
97    fn test_collect_in_list_negated() {
98        let (_d, facotry) =
99            PuffinManagerFactory::new_for_test_block("test_collect_in_list_negated_");
100        let metadata = test_region_metadata();
101        let mut builder = InvertedIndexApplierBuilder::new(
102            "test".to_string(),
103            test_object_store(),
104            &metadata,
105            HashSet::from_iter([1, 2, 3]),
106            facotry,
107        );
108
109        let in_list = InList {
110            expr: Box::new(tag_column()),
111            list: vec![string_lit("foo"), string_lit("bar")],
112            negated: true,
113        };
114
115        builder.collect_inlist(&in_list).unwrap();
116        assert!(builder.output.is_empty());
117    }
118
119    #[test]
120    fn test_collect_in_list_field_column() {
121        let (_d, facotry) =
122            PuffinManagerFactory::new_for_test_block("test_collect_in_list_field_column_");
123        let metadata = test_region_metadata();
124        let mut builder = InvertedIndexApplierBuilder::new(
125            "test".to_string(),
126            test_object_store(),
127            &metadata,
128            HashSet::from_iter([1, 2, 3]),
129            facotry,
130        );
131
132        let in_list = InList {
133            expr: Box::new(field_column()),
134            list: vec![string_lit("foo"), string_lit("bar")],
135            negated: false,
136        };
137
138        builder.collect_inlist(&in_list).unwrap();
139
140        let predicates = builder.output.get(&3).unwrap();
141        assert_eq!(predicates.len(), 1);
142        assert_eq!(
143            predicates[0],
144            Predicate::InList(InListPredicate {
145                list: BTreeSet::from_iter([encoded_string("foo"), encoded_string("bar")])
146            })
147        );
148    }
149
150    #[test]
151    fn test_collect_in_list_type_mismatch() {
152        let (_d, facotry) =
153            PuffinManagerFactory::new_for_test_block("test_collect_in_list_type_mismatch_");
154        let metadata = test_region_metadata();
155        let mut builder = InvertedIndexApplierBuilder::new(
156            "test".to_string(),
157            test_object_store(),
158            &metadata,
159            HashSet::from_iter([1, 2, 3]),
160            facotry,
161        );
162
163        let in_list = InList {
164            expr: Box::new(tag_column()),
165            list: vec![int64_lit(123), int64_lit(456)],
166            negated: false,
167        };
168
169        let res = builder.collect_inlist(&in_list);
170        assert!(matches!(res, Err(Error::FieldTypeMismatch { .. })));
171        assert!(builder.output.is_empty());
172    }
173
174    #[test]
175    fn test_collect_in_list_nonexistent_column() {
176        let (_d, facotry) =
177            PuffinManagerFactory::new_for_test_block("test_collect_in_list_nonexistent_column_");
178
179        let metadata = test_region_metadata();
180        let mut builder = InvertedIndexApplierBuilder::new(
181            "test".to_string(),
182            test_object_store(),
183            &metadata,
184            HashSet::from_iter([1, 2, 3]),
185            facotry,
186        );
187
188        let in_list = InList {
189            expr: Box::new(nonexistent_column()),
190            list: vec![string_lit("foo"), string_lit("bar")],
191            negated: false,
192        };
193
194        let res = builder.collect_inlist(&in_list);
195        assert!(matches!(res, Err(Error::ColumnNotFound { .. })));
196        assert!(builder.output.is_empty());
197    }
198}
mito2/sst/index/inverted_index/applier/builder/in_list.rs

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