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`, case included: `"AAPL"` contains `"AA"` but not
111    /// `"aa"`. To ignore case, keep the column in one case and change the text to match.
112    pub fn contains(column: impl Into<String>, text: impl Into<String>) -> Self {
113        Self::Contains {
114            column: column.into(),
115            text: text.into(),
116        }
117    }
118
119    /// Every one of `filters` holds.
120    pub fn all(filters: impl IntoIterator<Item = Filter>) -> Self {
121        Self::All(filters.into_iter().collect())
122    }
123
124    /// At least one of `filters` holds.
125    pub fn any(filters: impl IntoIterator<Item = Filter>) -> Self {
126        Self::Any(filters.into_iter().collect())
127    }
128
129    /// `filter` does not hold.
130    #[expect(
131        clippy::should_implement_trait,
132        reason = "a constructor, not negation of a value"
133    )]
134    pub fn not(filter: Filter) -> Self {
135        Self::Not(Box::new(filter))
136    }
137}
138
139/// A filter bound to column indices and readers.
140#[derive(Debug, Clone)]
141pub(crate) enum Predicate {
142    Compare {
143        column: usize,
144        kind: Kind,
145        read: ValueFn,
146        op: CompareOp,
147        value: Scalar,
148    },
149    IsNull {
150        column: usize,
151        /// The value reader, for types views read (dictionaries: a null value counts);
152        /// other types use the array's own nulls.
153        read: Option<ValueFn>,
154    },
155    Contains {
156        column: usize,
157        read: ValueFn,
158        text: String,
159    },
160    All(Vec<Predicate>),
161    Any(Vec<Predicate>),
162    Not(Box<Predicate>),
163}
164
165/// Resolves column `name` in `schema` to its index, kind and reader.
166pub(crate) fn resolve(schema: &Schema, name: &str) -> Result<(usize, Kind, ValueFn), DataError> {
167    let (index, field) = schema
168        .column_with_name(name)
169        .ok_or_else(|| DataError::UnknownColumn(name.to_owned()))?;
170    let (kind, read) = value_reader(field.data_type())
171        .ok_or_else(|| DataError::UnsupportedColumn(name.to_owned()))?;
172    Ok((index, kind, read))
173}
174
175impl Predicate {
176    pub(crate) fn compile(filter: &Filter, schema: &Schema) -> Result<Self, DataError> {
177        Ok(match filter {
178            Filter::Compare { column, op, value } => {
179                let (index, kind, read) = resolve(schema, column)?;
180                if !literal_fits(kind, value) {
181                    return Err(DataError::LiteralMismatch(column.clone()));
182                }
183                Self::Compare {
184                    column: index,
185                    kind,
186                    read,
187                    op: *op,
188                    value: value.clone(),
189                }
190            }
191            Filter::IsNull(column) => {
192                let (index, field) = schema
193                    .column_with_name(column)
194                    .ok_or_else(|| DataError::UnknownColumn(column.clone()))?;
195                Self::IsNull {
196                    column: index,
197                    read: value_reader(field.data_type()).map(|(_, read)| read),
198                }
199            }
200            Filter::Contains { column, text } => {
201                let (index, kind, read) = resolve(schema, column)?;
202                if kind != Kind::Str {
203                    return Err(DataError::LiteralMismatch(column.clone()));
204                }
205                Self::Contains {
206                    column: index,
207                    read,
208                    text: text.clone(),
209                }
210            }
211            Filter::All(filters) => Self::All(compile_all(filters, schema)?),
212            Filter::Any(filters) => Self::Any(compile_all(filters, schema)?),
213            Filter::Not(filter) => Self::Not(Box::new(Self::compile(filter, schema)?)),
214        })
215    }
216
217    /// The filter's value at `row` of `batch`: `None` for unknown.
218    pub(crate) fn eval(&self, batch: &RecordBatch, row: usize) -> Option<bool> {
219        let value = |column: usize, read: ValueFn| read(batch.column(column).as_ref(), row);
220        match self {
221            Self::Compare {
222                column,
223                kind,
224                read,
225                op,
226                value: literal,
227            } => compare_literal(*kind, value(*column, *read), literal).map(|o| op.holds(o)),
228            Self::IsNull { column, read } => Some(match read {
229                Some(read) => value(*column, *read) == Value::Null,
230                None => {
231                    let array = batch.column(*column);
232                    // A `Null` column has no validity buffer: every value is null.
233                    *array.data_type() == DataType::Null || array.is_null(row)
234                }
235            }),
236            Self::Contains { column, read, text } => match value(*column, *read) {
237                Value::Str(s) => Some(s.contains(text.as_str())),
238                _ => None,
239            },
240            Self::All(predicates) => kleene_all(predicates, batch, row),
241            Self::Any(predicates) => kleene_any(predicates, batch, row),
242            Self::Not(predicate) => predicate.eval(batch, row).map(|b| !b),
243        }
244    }
245}
246
247fn compile_all(filters: &[Filter], schema: &Schema) -> Result<Vec<Predicate>, DataError> {
248    filters
249        .iter()
250        .map(|f| Predicate::compile(f, schema))
251        .collect()
252}
253
254/// False if any is false, else unknown if any is unknown, else true.
255fn kleene_all(predicates: &[Predicate], batch: &RecordBatch, row: usize) -> Option<bool> {
256    let mut unknown = false;
257    for predicate in predicates {
258        match predicate.eval(batch, row) {
259            Some(false) => return Some(false),
260            None => unknown = true,
261            Some(true) => {}
262        }
263    }
264    (!unknown).then_some(true)
265}
266
267/// True if any is true, else unknown if any is unknown, else false.
268fn kleene_any(predicates: &[Predicate], batch: &RecordBatch, row: usize) -> Option<bool> {
269    let mut unknown = false;
270    for predicate in predicates {
271        match predicate.eval(batch, row) {
272            Some(true) => return Some(true),
273            None => unknown = true,
274            Some(false) => {}
275        }
276    }
277    (!unknown).then_some(false)
278}