1use 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#[derive(Debug, Clone, PartialEq)]
14#[non_exhaustive]
15pub enum Scalar {
16 Int(i64),
18 UInt(u64),
20 Float(f64),
22 Bool(bool),
24 Str(String),
26}
27
28#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
30#[non_exhaustive]
31pub enum CompareOp {
32 Eq,
34 Ne,
36 Lt,
38 Le,
40 Gt,
42 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#[derive(Debug, Clone, PartialEq)]
62#[non_exhaustive]
63pub enum Filter {
64 Compare {
66 column: String,
68 op: CompareOp,
70 value: Scalar,
72 },
73 IsNull(String),
75 Contains {
77 column: String,
79 text: String,
81 },
82 All(Vec<Filter>),
84 Any(Vec<Filter>),
86 Not(Box<Filter>),
88}
89
90impl Filter {
91 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 pub fn is_null(column: impl Into<String>) -> Self {
102 Self::IsNull(column.into())
103 }
104
105 pub fn is_not_null(column: impl Into<String>) -> Self {
107 Self::not(Self::is_null(column))
108 }
109
110 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 pub fn all(filters: impl IntoIterator<Item = Filter>) -> Self {
120 Self::All(filters.into_iter().collect())
121 }
122
123 pub fn any(filters: impl IntoIterator<Item = Filter>) -> Self {
125 Self::Any(filters.into_iter().collect())
126 }
127
128 #[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#[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 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
164pub(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 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 *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
253fn 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
266fn 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}