Skip to main content

densa_data/
filter.rs

1//! Row filters for views, with SQL's three-valued logic: a comparison with a null is
2//! unknown, and only rows whose filter is true are shown.
3
4use std::cmp::Ordering;
5
6use arrow_array::{Array, RecordBatch};
7use arrow_schema::{DataType, Schema};
8
9use crate::value::{compare_literal, literal_fits, value_reader, Kind, Value, ValueFn};
10use crate::DataError;
11
12/// A literal value a column is compared with.
13#[derive(Debug, Clone, PartialEq)]
14#[non_exhaustive]
15pub enum Scalar {
16    /// A signed integer; compares with any numeric column (decimals at their scale).
17    Int(i64),
18    /// An unsigned integer; compares with any numeric column.
19    UInt(u64),
20    /// A float; compares with any numeric column.
21    Float(f64),
22    /// A boolean.
23    Bool(bool),
24    /// Text; compares with string and dictionary-of-string columns.
25    Str(String),
26}
27
28/// How a column compares with a literal.
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
30#[non_exhaustive]
31pub enum CompareOp {
32    /// Equal.
33    Eq,
34    /// Not equal. Like every comparison, unknown for null and for NaN.
35    Ne,
36    /// Less than.
37    Lt,
38    /// Less than or equal.
39    Le,
40    /// Greater than.
41    Gt,
42    /// Greater than or equal.
43    Ge,
44}
45
46impl CompareOp {
47    fn holds(self, order: Ordering) -> bool {
48        match self {
49            Self::Eq => order.is_eq(),
50            Self::Ne => order.is_ne(),
51            Self::Lt => order.is_lt(),
52            Self::Le => order.is_le(),
53            Self::Gt => order.is_gt(),
54            Self::Ge => order.is_ge(),
55        }
56    }
57}
58
59/// Which rows a view shows. Dates, times, timestamps and durations compare as their integer
60/// count of units (milliseconds since the epoch for `Timestamp(Millisecond)`).
61#[derive(Debug, Clone, PartialEq)]
62#[non_exhaustive]
63pub enum Filter {
64    /// `column op value`; unknown when the column is null.
65    Compare {
66        /// Column name.
67        column: String,
68        /// Comparison.
69        op: CompareOp,
70        /// Literal.
71        value: Scalar,
72    },
73    /// The column is null. Works on columns of any type.
74    IsNull(String),
75    /// The string column contains `text` (case-sensitive); unknown when null.
76    Contains {
77        /// Column name.
78        column: String,
79        /// Text searched for.
80        text: String,
81    },
82    /// Every filter holds (true for none).
83    All(Vec<Filter>),
84    /// Some filter holds (false for none).
85    Any(Vec<Filter>),
86    /// The filter does not hold (unknown stays unknown).
87    Not(Box<Filter>),
88}
89
90impl Filter {
91    /// `column op value`.
92    pub fn compare(column: impl Into<String>, op: CompareOp, value: Scalar) -> Self {
93        Self::Compare {
94            column: column.into(),
95            op,
96            value,
97        }
98    }
99
100    /// The column is null.
101    pub fn is_null(column: impl Into<String>) -> Self {
102        Self::IsNull(column.into())
103    }
104
105    /// The column is not null.
106    pub fn is_not_null(column: impl Into<String>) -> Self {
107        Self::not(Self::is_null(column))
108    }
109
110    /// The string column contains `text`.
111    pub fn contains(column: impl Into<String>, text: impl Into<String>) -> Self {
112        Self::Contains {
113            column: column.into(),
114            text: text.into(),
115        }
116    }
117
118    /// Every one of `filters` holds.
119    pub fn all(filters: impl IntoIterator<Item = Filter>) -> Self {
120        Self::All(filters.into_iter().collect())
121    }
122
123    /// At least one of `filters` holds.
124    pub fn any(filters: impl IntoIterator<Item = Filter>) -> Self {
125        Self::Any(filters.into_iter().collect())
126    }
127
128    /// `filter` does not hold.
129    #[expect(
130        clippy::should_implement_trait,
131        reason = "a constructor, not negation of a value"
132    )]
133    pub fn not(filter: Filter) -> Self {
134        Self::Not(Box::new(filter))
135    }
136}
137
138/// A filter bound to column indices and readers.
139#[derive(Debug, Clone)]
140pub(crate) enum Predicate {
141    Compare {
142        column: usize,
143        kind: Kind,
144        read: ValueFn,
145        op: CompareOp,
146        value: Scalar,
147    },
148    IsNull {
149        column: usize,
150        /// The value reader, for types views read (dictionaries: a null value counts);
151        /// other types use the array's own nulls.
152        read: Option<ValueFn>,
153    },
154    Contains {
155        column: usize,
156        read: ValueFn,
157        text: String,
158    },
159    All(Vec<Predicate>),
160    Any(Vec<Predicate>),
161    Not(Box<Predicate>),
162}
163
164/// Resolves column `name` in `schema` to its index, kind and reader.
165pub(crate) fn resolve(schema: &Schema, name: &str) -> Result<(usize, Kind, ValueFn), DataError> {
166    let (index, field) = schema
167        .column_with_name(name)
168        .ok_or_else(|| DataError::UnknownColumn(name.to_owned()))?;
169    let (kind, read) = value_reader(field.data_type())
170        .ok_or_else(|| DataError::UnsupportedColumn(name.to_owned()))?;
171    Ok((index, kind, read))
172}
173
174impl Predicate {
175    pub(crate) fn compile(filter: &Filter, schema: &Schema) -> Result<Self, DataError> {
176        Ok(match filter {
177            Filter::Compare { column, op, value } => {
178                let (index, kind, read) = resolve(schema, column)?;
179                if !literal_fits(kind, value) {
180                    return Err(DataError::LiteralMismatch(column.clone()));
181                }
182                Self::Compare {
183                    column: index,
184                    kind,
185                    read,
186                    op: *op,
187                    value: value.clone(),
188                }
189            }
190            Filter::IsNull(column) => {
191                let (index, field) = schema
192                    .column_with_name(column)
193                    .ok_or_else(|| DataError::UnknownColumn(column.clone()))?;
194                Self::IsNull {
195                    column: index,
196                    read: value_reader(field.data_type()).map(|(_, read)| read),
197                }
198            }
199            Filter::Contains { column, text } => {
200                let (index, kind, read) = resolve(schema, column)?;
201                if kind != Kind::Str {
202                    return Err(DataError::LiteralMismatch(column.clone()));
203                }
204                Self::Contains {
205                    column: index,
206                    read,
207                    text: text.clone(),
208                }
209            }
210            Filter::All(filters) => Self::All(compile_all(filters, schema)?),
211            Filter::Any(filters) => Self::Any(compile_all(filters, schema)?),
212            Filter::Not(filter) => Self::Not(Box::new(Self::compile(filter, schema)?)),
213        })
214    }
215
216    /// The filter's value at `row` of `batch`: `None` for unknown.
217    pub(crate) fn eval(&self, batch: &RecordBatch, row: usize) -> Option<bool> {
218        let value = |column: usize, read: ValueFn| read(batch.column(column).as_ref(), row);
219        match self {
220            Self::Compare {
221                column,
222                kind,
223                read,
224                op,
225                value: literal,
226            } => compare_literal(*kind, value(*column, *read), literal).map(|o| op.holds(o)),
227            Self::IsNull { column, read } => Some(match read {
228                Some(read) => value(*column, *read) == Value::Null,
229                None => {
230                    let array = batch.column(*column);
231                    // A `Null` column has no validity buffer: every value is null.
232                    *array.data_type() == DataType::Null || array.is_null(row)
233                }
234            }),
235            Self::Contains { column, read, text } => match value(*column, *read) {
236                Value::Str(s) => Some(s.contains(text.as_str())),
237                _ => None,
238            },
239            Self::All(predicates) => kleene_all(predicates, batch, row),
240            Self::Any(predicates) => kleene_any(predicates, batch, row),
241            Self::Not(predicate) => predicate.eval(batch, row).map(|b| !b),
242        }
243    }
244}
245
246fn compile_all(filters: &[Filter], schema: &Schema) -> Result<Vec<Predicate>, DataError> {
247    filters
248        .iter()
249        .map(|f| Predicate::compile(f, schema))
250        .collect()
251}
252
253/// False if any is false, else unknown if any is unknown, else true.
254fn kleene_all(predicates: &[Predicate], batch: &RecordBatch, row: usize) -> Option<bool> {
255    let mut unknown = false;
256    for predicate in predicates {
257        match predicate.eval(batch, row) {
258            Some(false) => return Some(false),
259            None => unknown = true,
260            Some(true) => {}
261        }
262    }
263    (!unknown).then_some(true)
264}
265
266/// True if any is true, else unknown if any is unknown, else false.
267fn kleene_any(predicates: &[Predicate], batch: &RecordBatch, row: usize) -> Option<bool> {
268    let mut unknown = false;
269    for predicate in predicates {
270        match predicate.eval(batch, row) {
271            Some(true) => return Some(true),
272            None => unknown = true,
273            Some(false) => {}
274        }
275    }
276    (!unknown).then_some(false)
277}