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::HashSet;
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: HashSet::with_capacity(inlist.list.len()),
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 super::*;
57    use crate::error::Error;
58    use crate::sst::index::inverted_index::applier::builder::tests::{
59        encoded_string, field_column, int64_lit, nonexistent_column, string_lit, tag_column,
60        test_object_store, test_region_metadata,
61    };
62    use crate::sst::index::puffin_manager::PuffinManagerFactory;
63
64    #[test]
65    fn test_collect_in_list_basic() {
66        let (_d, facotry) = PuffinManagerFactory::new_for_test_block("test_collect_in_list_basic_");
67        let metadata = test_region_metadata();
68        let mut builder = InvertedIndexApplierBuilder::new(
69            "test".to_string(),
70            test_object_store(),
71            &metadata,
72            HashSet::from_iter([1, 2, 3]),
73            facotry,
74        );
75
76        let in_list = InList {
77            expr: Box::new(tag_column()),
78            list: vec![string_lit("foo"), string_lit("bar")],
79            negated: false,
80        };
81
82        builder.collect_inlist(&in_list).unwrap();
83
84        let predicates = builder.output.get(&1).unwrap();
85        assert_eq!(predicates.len(), 1);
86        assert_eq!(
87            predicates[0],
88            Predicate::InList(InListPredicate {
89                list: HashSet::from_iter([encoded_string("foo"), encoded_string("bar")])
90            })
91        );
92    }
93
94    #[test]
95    fn test_collect_in_list_negated() {
96        let (_d, facotry) =
97            PuffinManagerFactory::new_for_test_block("test_collect_in_list_negated_");
98        let metadata = test_region_metadata();
99        let mut builder = InvertedIndexApplierBuilder::new(
100            "test".to_string(),
101            test_object_store(),
102            &metadata,
103            HashSet::from_iter([1, 2, 3]),
104            facotry,
105        );
106
107        let in_list = InList {
108            expr: Box::new(tag_column()),
109            list: vec![string_lit("foo"), string_lit("bar")],
110            negated: true,
111        };
112
113        builder.collect_inlist(&in_list).unwrap();
114        assert!(builder.output.is_empty());
115    }
116
117    #[test]
118    fn test_collect_in_list_field_column() {
119        let (_d, facotry) =
120            PuffinManagerFactory::new_for_test_block("test_collect_in_list_field_column_");
121        let metadata = test_region_metadata();
122        let mut builder = InvertedIndexApplierBuilder::new(
123            "test".to_string(),
124            test_object_store(),
125            &metadata,
126            HashSet::from_iter([1, 2, 3]),
127            facotry,
128        );
129
130        let in_list = InList {
131            expr: Box::new(field_column()),
132            list: vec![string_lit("foo"), string_lit("bar")],
133            negated: false,
134        };
135
136        builder.collect_inlist(&in_list).unwrap();
137
138        let predicates = builder.output.get(&3).unwrap();
139        assert_eq!(predicates.len(), 1);
140        assert_eq!(
141            predicates[0],
142            Predicate::InList(InListPredicate {
143                list: HashSet::from_iter([encoded_string("foo"), encoded_string("bar")])
144            })
145        );
146    }
147
148    #[test]
149    fn test_collect_in_list_type_mismatch() {
150        let (_d, facotry) =
151            PuffinManagerFactory::new_for_test_block("test_collect_in_list_type_mismatch_");
152        let metadata = test_region_metadata();
153        let mut builder = InvertedIndexApplierBuilder::new(
154            "test".to_string(),
155            test_object_store(),
156            &metadata,
157            HashSet::from_iter([1, 2, 3]),
158            facotry,
159        );
160
161        let in_list = InList {
162            expr: Box::new(tag_column()),
163            list: vec![int64_lit(123), int64_lit(456)],
164            negated: false,
165        };
166
167        let res = builder.collect_inlist(&in_list);
168        assert!(matches!(res, Err(Error::FieldTypeMismatch { .. })));
169        assert!(builder.output.is_empty());
170    }
171
172    #[test]
173    fn test_collect_in_list_nonexistent_column() {
174        let (_d, facotry) =
175            PuffinManagerFactory::new_for_test_block("test_collect_in_list_nonexistent_column_");
176
177        let metadata = test_region_metadata();
178        let mut builder = InvertedIndexApplierBuilder::new(
179            "test".to_string(),
180            test_object_store(),
181            &metadata,
182            HashSet::from_iter([1, 2, 3]),
183            facotry,
184        );
185
186        let in_list = InList {
187            expr: Box::new(nonexistent_column()),
188            list: vec![string_lit("foo"), string_lit("bar")],
189            negated: false,
190        };
191
192        let res = builder.collect_inlist(&in_list);
193        assert!(matches!(res, Err(Error::ColumnNotFound { .. })));
194        assert!(builder.output.is_empty());
195    }
196}
mito2/sst/index/inverted_index/applier/builder/in_list.rs

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