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 {
113 Self::Contains {
114 column: column.into(),
115 text: text.into(),
116 }
117 }
118
119 pub fn all(filters: impl IntoIterator<Item = Filter>) -> Self {
121 Self::All(filters.into_iter().collect())
122 }
123
124 pub fn any(filters: impl IntoIterator<Item = Filter>) -> Self {
126 Self::Any(filters.into_iter().collect())
127 }
128
129 #[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#[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 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
165pub(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 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 *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
254fn 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
267fn 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}