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 store_api::region_request::PathType;
59
60    use super::*;
61    use crate::error::Error;
62    use crate::sst::index::inverted_index::applier::builder::tests::{
63        encoded_string, field_column, int64_lit, nonexistent_column, string_lit, tag_column,
64        test_object_store, test_region_metadata,
65    };
66    use crate::sst::index::puffin_manager::PuffinManagerFactory;
67
68    #[test]
69    fn test_collect_in_list_basic() {
70        let (_d, facotry) = PuffinManagerFactory::new_for_test_block("test_collect_in_list_basic_");
71        let metadata = test_region_metadata();
72        let mut builder = InvertedIndexApplierBuilder::new(
73            "test".to_string(),
74            PathType::Bare,
75            test_object_store(),
76            &metadata,
77            HashSet::from_iter([1, 2, 3]),
78            facotry,
79        );
80
81        let in_list = InList {
82            expr: Box::new(tag_column()),
83            list: vec![string_lit("foo"), string_lit("bar")],
84            negated: false,
85        };
86
87        builder.collect_inlist(&in_list).unwrap();
88
89        let predicates = builder.output.get(&1).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, facotry) =
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            facotry,
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, facotry) =
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            facotry,
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.output.get(&3).unwrap();
146        assert_eq!(predicates.len(), 1);
147        assert_eq!(
148            predicates[0],
149            Predicate::InList(InListPredicate {
150                list: BTreeSet::from_iter([encoded_string("foo"), encoded_string("bar")])
151            })
152        );
153    }
154
155    #[test]
156    fn test_collect_in_list_type_mismatch() {
157        let (_d, facotry) =
158            PuffinManagerFactory::new_for_test_block("test_collect_in_list_type_mismatch_");
159        let metadata = test_region_metadata();
160        let mut builder = InvertedIndexApplierBuilder::new(
161            "test".to_string(),
162            PathType::Bare,
163            test_object_store(),
164            &metadata,
165            HashSet::from_iter([1, 2, 3]),
166            facotry,
167        );
168
169        let in_list = InList {
170            expr: Box::new(tag_column()),
171            list: vec![int64_lit(123), int64_lit(456)],
172            negated: false,
173        };
174
175        let res = builder.collect_inlist(&in_list);
176        assert!(matches!(res, Err(Error::Encode { .. })));
177        assert!(builder.output.is_empty());
178    }
179
180    #[test]
181    fn test_collect_in_list_nonexistent_column() {
182        let (_d, facotry) =
183            PuffinManagerFactory::new_for_test_block("test_collect_in_list_nonexistent_column_");
184
185        let metadata = test_region_metadata();
186        let mut builder = InvertedIndexApplierBuilder::new(
187            "test".to_string(),
188            PathType::Bare,
189            test_object_store(),
190            &metadata,
191            HashSet::from_iter([1, 2, 3]),
192            facotry,
193        );
194
195        let in_list = InList {
196            expr: Box::new(nonexistent_column()),
197            list: vec![string_lit("foo"), string_lit("bar")],
198            negated: false,
199        };
200
201        let res = builder.collect_inlist(&in_list);
202        assert!(matches!(res, Err(Error::ColumnNotFound { .. })));
203        assert!(builder.output.is_empty());
204    }
205}
mito2/sst/index/inverted_index/applier/builder/in_list.rs

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