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

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