common_recordbatch/
filter.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
15//! Precise single-column predicates and record batch filtering.
16
17mod predicate;
18#[cfg(test)]
19mod tests;
20mod timestamp;
21
22use std::borrow::Borrow;
23use std::sync::Arc;
24
25use datafusion::error::Result as DfResult;
26use datafusion::logical_expr::{Expr, Literal, Operator, utils};
27use datafusion::physical_plan::PhysicalExpr;
28use datafusion_common::arrow::array::{ArrayRef, Datum};
29use datafusion_common::arrow::buffer::BooleanBuffer;
30use datafusion_common::cast::{as_boolean_array, as_null_array};
31use datafusion_common::{DataFusionError, ScalarValue, internal_err};
32use datatypes::arrow::array::{Array, BooleanArray, BooleanBufferBuilder, RecordBatch};
33use datatypes::arrow::compute::filter_record_batch;
34use datatypes::arrow::datatypes::DataType;
35use datatypes::data_type::{ConcreteDataType, DataType as _};
36use datatypes::value::Value;
37use datatypes::vectors::VectorRef;
38use snafu::ResultExt;
39
40pub use self::predicate::{regexp_is_match_dictionary, regexp_is_match_scalar};
41pub use self::timestamp::timestamp_scalar_value;
42use crate::error::{Result, ToArrowScalarSnafu};
43use crate::filter::predicate::SimplePredicate;
44
45/// Prunes candidates using known metadata columns. Missing columns and unsupported
46/// predicates cannot exclude a candidate. Complete SQL predicates must still run
47/// after pruning.
48#[derive(Debug, Clone, Default)]
49pub struct MetadataFilter {
50    filters: Arc<[SimpleFilterEvaluator]>,
51}
52
53impl MetadataFilter {
54    /// Compiles supported conjuncts once for reuse across scan stages. OR is
55    /// never split: an unknown alternative might match the candidate.
56    pub fn new(expressions: &[Expr]) -> Self {
57        Self {
58            filters: expressions
59                .iter()
60                .flat_map(utils::split_conjunction)
61                .filter_map(SimpleFilterEvaluator::try_new)
62                .collect(),
63        }
64    }
65
66    /// Returns false only when known metadata proves a conjunct cannot match.
67    /// `None` means unavailable metadata, distinct from a known SQL NULL scalar.
68    /// Values may be borrowed so repeated predicates can reuse allocated strings.
69    pub fn matches<T: Borrow<ScalarValue>>(
70        &self,
71        mut value: impl FnMut(&str) -> Option<T>,
72    ) -> bool {
73        self.filters.iter().all(|filter| {
74            value(filter.column_name())
75                .is_none_or(|value| filter.evaluate_scalar(value.borrow()).unwrap_or(true))
76        })
77    }
78}
79
80/// Evaluates supported single-column predicates with precompiled regular expressions.
81/// Compound predicates preserve SQL null semantics until conversion to a scan mask.
82#[derive(Debug, Clone)]
83pub struct SimpleFilterEvaluator {
84    column_name: String,
85    predicate: SimplePredicate,
86}
87
88impl SimpleFilterEvaluator {
89    pub fn new<T: Literal>(column_name: String, lit: T, op: Operator) -> Option<Self> {
90        if !matches!(
91            op,
92            Operator::Eq
93                | Operator::NotEq
94                | Operator::Lt
95                | Operator::LtEq
96                | Operator::Gt
97                | Operator::GtEq
98        ) {
99            return None;
100        }
101        let Expr::Literal(value, _) = lit.lit() else {
102            return None;
103        };
104        Some(Self {
105            column_name,
106            predicate: SimplePredicate::Comparison {
107                literal: value.to_scalar().ok()?,
108                op,
109                regex: None,
110                regex_negative: false,
111            },
112        })
113    }
114
115    pub fn try_new(predicate: &Expr) -> Option<Self> {
116        Self::try_new_with_column_type(predicate, &|_| None)
117    }
118
119    /// Recognizes single-column filters, including lossless string dictionary casts
120    /// when the caller can resolve the column's type.
121    pub fn try_new_with_column_type(
122        predicate: &Expr,
123        column_type: &impl Fn(&str) -> Option<DataType>,
124    ) -> Option<Self> {
125        let columns = predicate.column_refs();
126        if columns.len() != 1 {
127            return None;
128        }
129        let column_name = columns.into_iter().next()?.name.clone();
130        let column_type = column_type(&column_name);
131        Some(Self {
132            column_name,
133            predicate: SimplePredicate::try_new(predicate, column_type.as_ref())?,
134        })
135    }
136
137    /// Returns the referenced column name.
138    pub fn column_name(&self) -> &str {
139        &self.column_name
140    }
141
142    pub fn is_eq(&self) -> bool {
143        matches!(
144            self.predicate,
145            SimplePredicate::Comparison {
146                op: Operator::Eq,
147                ..
148            }
149        )
150    }
151    pub fn is_not_eq(&self) -> bool {
152        matches!(
153            self.predicate,
154            SimplePredicate::Comparison {
155                op: Operator::NotEq,
156                ..
157            }
158        )
159    }
160    pub fn is_lt(&self) -> bool {
161        matches!(
162            self.predicate,
163            SimplePredicate::Comparison {
164                op: Operator::Lt,
165                ..
166            }
167        )
168    }
169    pub fn is_lt_eq(&self) -> bool {
170        matches!(
171            self.predicate,
172            SimplePredicate::Comparison {
173                op: Operator::LtEq,
174                ..
175            }
176        )
177    }
178    pub fn is_gt(&self) -> bool {
179        matches!(
180            self.predicate,
181            SimplePredicate::Comparison {
182                op: Operator::Gt,
183                ..
184            }
185        )
186    }
187    pub fn is_gt_eq(&self) -> bool {
188        matches!(
189            self.predicate,
190            SimplePredicate::Comparison {
191                op: Operator::GtEq,
192                ..
193            }
194        )
195    }
196
197    /// Returns true for an OR chain of equality comparisons.
198    pub fn is_or_eq_chain(&self) -> bool {
199        matches!(self.predicate, SimplePredicate::OrEqChain { .. })
200    }
201
202    /// Returns the literal of a bare comparison, without discarding compound conditions.
203    pub fn literal_value(&self) -> Option<Value> {
204        let SimplePredicate::Comparison { literal, .. } = &self.predicate else {
205            return None;
206        };
207        Value::try_from(ScalarValue::try_from_array(literal.get().0, 0).ok()?).ok()
208    }
209
210    /// Returns equality-OR literals, or None for other predicates or unsupported values.
211    pub fn literal_list_values(&self) -> Option<Vec<Value>> {
212        let SimplePredicate::OrEqChain { literals } = &self.predicate else {
213            return None;
214        };
215        literals
216            .iter()
217            .map(|literal| {
218                Value::try_from(ScalarValue::try_from_array(literal.get().0, 0).ok()?).ok()
219            })
220            .collect()
221    }
222
223    pub fn evaluate_scalar(&self, input: &ScalarValue) -> Result<bool> {
224        let input = input
225            .to_scalar()
226            .with_context(|_| ToArrowScalarSnafu { v: input.clone() })?;
227        Ok(self.evaluate_datum(&input, 1)?.value(0))
228    }
229
230    pub fn evaluate_array(&self, input: &ArrayRef) -> Result<BooleanBuffer> {
231        self.evaluate_datum(input, input.len())
232    }
233
234    pub fn evaluate_vector(&self, input: &VectorRef) -> Result<BooleanBuffer> {
235        self.evaluate_array(&input.to_arrow_array())
236    }
237
238    fn evaluate_datum(&self, input: &impl Datum, len: usize) -> Result<BooleanBuffer> {
239        let result = self.predicate.evaluate(input, len)?;
240        Ok(boolean_array_to_scan_mask(&result).values().clone())
241    }
242
243    /// Adapts timestamp predicates to an older SST's timestamp unit without rounding
244    /// away boundary rows. Nonrepresentable equality literals cannot match a row.
245    pub fn cast_timestamp_unit(&self, target: &ConcreteDataType) -> Option<TimestampUnitCast> {
246        let DataType::Timestamp(target_unit, _) = target.as_arrow_type() else {
247            return None;
248        };
249        Some(match self.predicate.cast_timestamp_unit(target_unit)? {
250            SimplePredicate::Constant(Some(true)) => TimestampUnitCast::Matched,
251            SimplePredicate::Constant(_) => TimestampUnitCast::Pruned,
252            predicate => TimestampUnitCast::Filter(Self {
253                column_name: self.column_name.clone(),
254                predicate,
255            }),
256        })
257    }
258}
259
260/// The result of casting a [`SimpleFilterEvaluator`] to a different timestamp
261/// unit, preserving the predicate's semantics on rows stored in that unit.
262/// See [`SimpleFilterEvaluator::cast_timestamp_unit`].
263#[derive(Debug, Clone)]
264pub enum TimestampUnitCast {
265    /// The cast filter; evaluates column values stored in the target unit.
266    Filter(SimpleFilterEvaluator),
267    /// No value in the target unit satisfies the filter.
268    Pruned,
269    /// Every value in the target unit satisfies the filter.
270    Matched,
271}
272
273/// Evaluate the predicate on the input [RecordBatch], and return a new [RecordBatch].
274/// Copy from datafusion::physical_plan::src::filter.rs
275pub fn batch_filter(
276    batch: &RecordBatch,
277    predicate: &Arc<dyn PhysicalExpr>,
278) -> DfResult<RecordBatch> {
279    predicate
280        .evaluate(batch)
281        .and_then(|v| v.into_array(batch.num_rows()))
282        .and_then(|array| {
283            let filter_array = match as_boolean_array(&array) {
284                Ok(boolean_array) => Ok(boolean_array.clone()),
285                Err(_) => {
286                    let Ok(null_array) = as_null_array(&array) else {
287                        return internal_err!(
288                            "Cannot create filter_array from non-boolean predicates"
289                        );
290                    };
291
292                    // if the predicate is null, then the result is also null
293                    Ok::<BooleanArray, DataFusionError>(BooleanArray::new_null(null_array.len()))
294                }
295            }?;
296            Ok(filter_record_batch(
297                batch,
298                &boolean_array_to_scan_mask(&filter_array),
299            )?)
300        })
301}
302
303/// Converts nullable SQL predicate values to a scan mask, where `NULL` is `false`.
304fn boolean_array_to_scan_mask(array: &BooleanArray) -> BooleanArray {
305    if array.null_count() == 0 {
306        return array.clone();
307    }
308
309    let mut values = BooleanBufferBuilder::new(array.len());
310    for index in 0..array.len() {
311        values.append(array.is_valid(index) && array.value(index));
312    }
313    BooleanArray::new(values.into(), None)
314}
common_recordbatch/filter.rs

common_recordbatch/
filter.rs
