oxidelake_core/
params.rs

1//! Operator parameter types and the bounded v1 GPU type coverage.
2//!
3//! These types are deliberately independent of DataFusion's `PhysicalExpr`
4//! tree: the planner lowers eligible DataFusion expressions into them, the
5//! codec serializes them with `postcard`, and every backend executes them.
6
7use datafusion::arrow::datatypes::DataType;
8use serde::{Deserialize, Serialize};
9
10/// Comparison operators supported by the fused filter kernel.
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
12pub enum Comparison {
13    /// `=`
14    Eq,
15    /// `<`
16    Lt,
17    /// `<=`
18    LtEq,
19    /// `>`
20    Gt,
21    /// `>=`
22    GtEq,
23}
24
25impl Comparison {
26    /// SQL spelling, used in `EXPLAIN` output.
27    pub const fn symbol(self) -> &'static str {
28        match self {
29            Comparison::Eq => "=",
30            Comparison::Lt => "<",
31            Comparison::LtEq => "<=",
32            Comparison::Gt => ">",
33            Comparison::GtEq => ">=",
34        }
35    }
36}
37
38/// A literal a column is compared against.
39#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
40pub enum Literal {
41    /// 64-bit signed integer.
42    Int64(i64),
43    /// 64-bit float.
44    Float64(f64),
45}
46
47impl Literal {
48    /// The Arrow type this literal compares against.
49    pub const fn data_type(&self) -> DataType {
50        match self {
51            Literal::Int64(_) => DataType::Int64,
52            Literal::Float64(_) => DataType::Float64,
53        }
54    }
55}
56
57/// A filter predicate in the bounded v1 grammar: column-vs-literal comparisons
58/// combined with `AND`.
59#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
60pub enum Predicate {
61    /// `column <op> literal`
62    Compare {
63        /// Index of the column in the input schema.
64        column: usize,
65        /// The comparison operator.
66        op: Comparison,
67        /// The literal operand.
68        literal: Literal,
69    },
70    /// Logical conjunction.
71    And(Box<Predicate>, Box<Predicate>),
72}
73
74impl Predicate {
75    /// Builds a comparison predicate.
76    pub const fn compare(column: usize, op: Comparison, literal: Literal) -> Self {
77        Self::Compare {
78            column,
79            op,
80            literal,
81        }
82    }
83
84    /// Conjoins two predicates.
85    pub fn and(left: Predicate, right: Predicate) -> Self {
86        Self::And(Box::new(left), Box::new(right))
87    }
88
89    /// Every column index referenced, in evaluation order (duplicates kept).
90    pub fn columns(&self) -> Vec<usize> {
91        let mut out = Vec::new();
92        self.collect_columns(&mut out);
93        out
94    }
95
96    fn collect_columns(&self, out: &mut Vec<usize>) {
97        match self {
98            Predicate::Compare { column, .. } => out.push(*column),
99            Predicate::And(l, r) => {
100                l.collect_columns(out);
101                r.collect_columns(out);
102            }
103        }
104    }
105
106    /// Number of comparison leaves.
107    pub fn leaf_count(&self) -> usize {
108        match self {
109            Predicate::Compare { .. } => 1,
110            Predicate::And(l, r) => l.leaf_count() + r.leaf_count(),
111        }
112    }
113}
114
115/// Aggregate functions supported by the grouped-aggregation kernel.
116#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
117pub enum AggregateFunction {
118    /// `SUM`
119    Sum,
120    /// `COUNT`
121    Count,
122    /// `MIN`
123    Min,
124    /// `MAX`
125    Max,
126}
127
128impl AggregateFunction {
129    /// SQL spelling, used in `EXPLAIN` output and output column names.
130    pub const fn name(self) -> &'static str {
131        match self {
132            AggregateFunction::Sum => "SUM",
133            AggregateFunction::Count => "COUNT",
134            AggregateFunction::Min => "MIN",
135            AggregateFunction::Max => "MAX",
136        }
137    }
138}
139
140/// A grouped aggregation over one `Int64` key.
141#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
142pub struct AggregateSpec {
143    /// Index of the `Int64` group-by column.
144    pub group_by: usize,
145    /// `(function, input column index)` pairs, in output order.
146    pub aggregates: Vec<(AggregateFunction, usize)>,
147}
148
149/// Distance metric for the vector kernel.
150#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
151pub enum DistanceMetric {
152    /// Euclidean distance.
153    L2,
154    /// Cosine distance (`1 - cosine similarity`).
155    Cosine,
156}
157
158impl DistanceMetric {
159    /// Lowercase name used in `EXPLAIN` output and SQL function names.
160    pub const fn name(self) -> &'static str {
161        match self {
162            DistanceMetric::L2 => "l2",
163            DistanceMetric::Cosine => "cosine",
164        }
165    }
166}
167
168/// `true` for the scalar types the v1 GPU kernels accept as filter, join-key
169/// and aggregate inputs.
170pub const fn gpu_eligible_scalar(dt: &DataType) -> bool {
171    matches!(dt, DataType::Int64 | DataType::Float64)
172}
173
174/// For a `FixedSizeList<Float32, n>` column returns `Some(n)`; otherwise `None`.
175pub fn vector_dimension(dt: &DataType) -> Option<usize> {
176    match dt {
177        DataType::FixedSizeList(field, n) if *field.data_type() == DataType::Float32 && *n > 0 => {
178            usize::try_from(*n).ok()
179        }
180        _ => None,
181    }
182}
183
184#[cfg(test)]
185#[allow(clippy::unwrap_used, clippy::expect_used)]
186mod tests {
187    use std::sync::Arc;
188
189    use datafusion::arrow::datatypes::Field;
190
191    use super::*;
192
193    #[test]
194    fn predicate_collects_columns_in_order() {
195        let p = Predicate::and(
196            Predicate::compare(2, Comparison::Gt, Literal::Int64(1)),
197            Predicate::and(
198                Predicate::compare(0, Comparison::LtEq, Literal::Float64(0.5)),
199                Predicate::compare(2, Comparison::Eq, Literal::Int64(9)),
200            ),
201        );
202        assert_eq!(p.columns(), vec![2, 0, 2]);
203        assert_eq!(p.leaf_count(), 3);
204    }
205
206    #[test]
207    fn coverage_rules() {
208        assert!(gpu_eligible_scalar(&DataType::Int64));
209        assert!(gpu_eligible_scalar(&DataType::Float64));
210        assert!(!gpu_eligible_scalar(&DataType::Int32));
211        assert!(!gpu_eligible_scalar(&DataType::Utf8));
212
213        let vec3 =
214            DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, false)), 3);
215        assert_eq!(vector_dimension(&vec3), Some(3));
216        let f64s =
217            DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float64, false)), 3);
218        assert_eq!(vector_dimension(&f64s), None);
219        assert_eq!(vector_dimension(&DataType::Int64), None);
220    }
221
222    #[test]
223    fn names_are_stable_for_explain_output() {
224        assert_eq!(AggregateFunction::Sum.name(), "SUM");
225        assert_eq!(AggregateFunction::Count.name(), "COUNT");
226        assert_eq!(DistanceMetric::Cosine.name(), "cosine");
227        assert_eq!(Comparison::GtEq.symbol(), ">=");
228        assert_eq!(Literal::Float64(1.5).data_type(), DataType::Float64);
229    }
230}
oxidelake_core/params.rs

oxidelake_core/
params.rs
