Skip to main content

tonbo_predicate/core/
value.rs

1use std::cmp::Ordering;
2
3use typed_arrow_dyn::{DynCell, DynCellRaw, DynCellRef};
4
5/// Literal values accepted by predicate operands, backed by `DynCell`.
6#[derive(Clone, Debug)]
7pub struct ScalarValue {
8    cell: DynCell,
9}
10
11impl ScalarValue {
12    /// Represents SQL/Arrow `NULL`.
13    #[must_use]
14    pub fn null() -> Self {
15        Self {
16            cell: DynCell::Null,
17        }
18    }
19
20    pub(crate) fn from_dyn(cell: DynCell) -> Self {
21        Self { cell }
22    }
23
24    /// Returns true when the literal is `NULL`.
25    #[must_use]
26    pub fn is_null(&self) -> bool {
27        matches!(self.cell, DynCell::Null)
28    }
29
30    /// Returns a borrowed view over this scalar value.
31    #[must_use]
32    pub fn as_ref(&self) -> ScalarValueRef<'_> {
33        let ref_cell = self
34            .cell
35            .as_ref()
36            .expect("ScalarValue should only hold scalar DynCell variants");
37        ScalarValueRef::from_dyn(ref_cell)
38    }
39
40    /// Compares this scalar with another, returning the ordering when both sides are comparable.
41    pub fn compare(&self, other: &Self) -> Option<Ordering> {
42        self.as_ref().compare(&other.as_ref())
43    }
44
45    /// Access the underlying dynamic cell.
46    #[must_use]
47    pub fn as_dyn(&self) -> &DynCell {
48        &self.cell
49    }
50
51    /// Consume this scalar and return the underlying dynamic cell.
52    pub fn into_dyn(self) -> DynCell {
53        self.cell
54    }
55}
56
57impl PartialEq for ScalarValue {
58    fn eq(&self, other: &Self) -> bool {
59        let left = self.as_ref();
60        let right = other.as_ref();
61        match (left.is_null(), right.is_null()) {
62            (true, true) => true,
63            _ => left
64                .compare(&right)
65                .map(|ord| ord == Ordering::Equal)
66                .unwrap_or_else(|| left.eq(&right.as_dyn())),
67        }
68    }
69}
70
71impl From<bool> for ScalarValue {
72    fn from(value: bool) -> Self {
73        ScalarValue::from_dyn(DynCell::Bool(value))
74    }
75}
76
77impl From<i64> for ScalarValue {
78    fn from(value: i64) -> Self {
79        ScalarValue::from_dyn(DynCell::I64(value))
80    }
81}
82
83impl From<u64> for ScalarValue {
84    fn from(value: u64) -> Self {
85        ScalarValue::from_dyn(DynCell::U64(value))
86    }
87}
88
89impl From<f64> for ScalarValue {
90    fn from(value: f64) -> Self {
91        ScalarValue::from_dyn(DynCell::F64(value))
92    }
93}
94
95impl From<String> for ScalarValue {
96    fn from(value: String) -> Self {
97        ScalarValue::from_dyn(DynCell::Str(value))
98    }
99}
100
101impl From<&str> for ScalarValue {
102    fn from(value: &str) -> Self {
103        ScalarValue::from_dyn(DynCell::Str(value.to_owned()))
104    }
105}
106
107impl From<Vec<u8>> for ScalarValue {
108    fn from(value: Vec<u8>) -> Self {
109        ScalarValue::from_dyn(DynCell::Bin(value))
110    }
111}
112
113impl From<&[u8]> for ScalarValue {
114    fn from(value: &[u8]) -> Self {
115        ScalarValue::from_dyn(DynCell::Bin(value.to_vec()))
116    }
117}
118
119/// Borrowed view over a scalar value backed by `DynCellRef`.
120#[derive(Clone, Debug)]
121pub struct ScalarValueRef<'a> {
122    cell: DynCellRef<'a>,
123}
124
125impl PartialEq<DynCellRef<'_>> for ScalarValueRef<'_> {
126    fn eq(&self, other: &DynCellRef<'_>) -> bool {
127        self.cells_equal(&ScalarValueRef::from_dyn(other.clone()))
128    }
129}
130
131impl PartialOrd for ScalarValueRef<'_> {
132    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
133        self.compare(other)
134    }
135}
136
137fn signed_int_to_i128(raw: &DynCellRaw) -> Option<i128> {
138    match raw {
139        DynCellRaw::I8(v) => Some(i128::from(*v)),
140        DynCellRaw::I16(v) => Some(i128::from(*v)),
141        DynCellRaw::I32(v) => Some(i128::from(*v)),
142        DynCellRaw::I64(v) => Some(i128::from(*v)),
143        _ => None,
144    }
145}
146
147fn unsigned_int_to_u128(raw: &DynCellRaw) -> Option<u128> {
148    match raw {
149        DynCellRaw::U8(v) => Some(u128::from(*v)),
150        DynCellRaw::U16(v) => Some(u128::from(*v)),
151        DynCellRaw::U32(v) => Some(u128::from(*v)),
152        DynCellRaw::U64(v) => Some(u128::from(*v)),
153        _ => None,
154    }
155}
156
157impl<'a> ScalarValueRef<'a> {
158    fn cells_equal_option(lhs: Option<DynCellRef<'_>>, rhs: Option<DynCellRef<'_>>) -> bool {
159        match (lhs, rhs) {
160            (None, None) => true,
161            (Some(l), Some(r)) => {
162                let lref = ScalarValueRef::from_dyn(l);
163                let rref = ScalarValueRef::from_dyn(r);
164                lref.cells_equal(&rref)
165            }
166            _ => false,
167        }
168    }
169
170    /// Deep Arrow-semantic equality against another scalar reference.
171    fn cells_equal(&self, rhs: &ScalarValueRef<'_>) -> bool {
172        use DynCellRaw::*;
173        match (self.cell.as_raw(), rhs.cell.as_raw()) {
174            (Null, Null) => true,
175            (Bool(a), Bool(b)) => a == b,
176            (I64(a), I64(b)) => a == b,
177            (U64(a), U64(b)) => a == b,
178            (F64(a), F64(b)) => a.to_bits() == b.to_bits(),
179            (Str { ptr: ap, len: al }, Str { ptr: bp, len: bl }) => unsafe {
180                std::slice::from_raw_parts(ap.as_ptr() as *const u8, *al)
181                    == std::slice::from_raw_parts(bp.as_ptr() as *const u8, *bl)
182            },
183            (Bin { ptr: ap, len: al }, Bin { ptr: bp, len: bl }) => unsafe {
184                std::slice::from_raw_parts(ap.as_ptr() as *const u8, *al)
185                    == std::slice::from_raw_parts(bp.as_ptr() as *const u8, *bl)
186            },
187            _ => {
188                if let (Some(ls), Some(rs)) = (self.cell.as_struct(), rhs.cell.as_struct()) {
189                    if ls.len() != rs.len() {
190                        return false;
191                    }
192                    for idx in 0..ls.len() {
193                        let l = ls.get(idx).ok().flatten();
194                        let r = rs.get(idx).ok().flatten();
195                        if !ScalarValueRef::cells_equal_option(l, r) {
196                            return false;
197                        }
198                    }
199                    return true;
200                }
201                if let (Some(ll), Some(rl)) = (self.cell.as_list(), rhs.cell.as_list()) {
202                    if ll.len() != rl.len() {
203                        return false;
204                    }
205                    for idx in 0..ll.len() {
206                        let l = ll.get(idx).ok().flatten();
207                        let r = rl.get(idx).ok().flatten();
208                        if !ScalarValueRef::cells_equal_option(l, r) {
209                            return false;
210                        }
211                    }
212                    return true;
213                }
214                if let (Some(lf), Some(rf)) = (
215                    self.cell.as_fixed_size_list(),
216                    rhs.cell.as_fixed_size_list(),
217                ) {
218                    if lf.len() != rf.len() {
219                        return false;
220                    }
221                    for idx in 0..lf.len() {
222                        let l = lf.get(idx).ok().flatten();
223                        let r = rf.get(idx).ok().flatten();
224                        if !ScalarValueRef::cells_equal_option(l, r) {
225                            return false;
226                        }
227                    }
228                    return true;
229                }
230                if let (Some(lm), Some(rm)) = (self.cell.as_map(), rhs.cell.as_map()) {
231                    if lm.len() != rm.len() {
232                        return false;
233                    }
234                    for idx in 0..lm.len() {
235                        let l = lm.get(idx).ok();
236                        let r = rm.get(idx).ok();
237                        let (lk, lv) = match l {
238                            Some(pair) => pair,
239                            None => return false,
240                        };
241                        let (rk, rv) = match r {
242                            Some(pair) => pair,
243                            None => return false,
244                        };
245                        if !ScalarValueRef::cells_equal_option(Some(lk), Some(rk))
246                            || !ScalarValueRef::cells_equal_option(lv, rv)
247                        {
248                            return false;
249                        }
250                    }
251                    return true;
252                }
253                if let (Some(lu), Some(ru)) = (self.cell.as_union(), rhs.cell.as_union()) {
254                    if lu.type_id() != ru.type_id() {
255                        return false;
256                    }
257                    let lval = lu.value().ok().flatten();
258                    let rval = ru.value().ok().flatten();
259                    return ScalarValueRef::cells_equal_option(lval, rval);
260                }
261                false
262            }
263        }
264    }
265
266    /// Returns true when the literal is the `Null` variant.
267    #[must_use]
268    pub fn is_null(&self) -> bool {
269        self.cell.is_null()
270    }
271
272    /// Compares this scalar with another, returning the ordering when both sides are comparable.
273    pub fn compare(&self, other: &ScalarValueRef<'_>) -> Option<Ordering> {
274        use DynCellRaw::*;
275        match (self.cell.as_raw(), other.cell.as_raw()) {
276            (Null, _) | (_, Null) => None,
277            (Bool(lhs), Bool(rhs)) => Some(lhs.cmp(rhs)),
278            (I8(lhs), I8(rhs)) => Some(lhs.cmp(rhs)),
279            (I16(lhs), I16(rhs)) => Some(lhs.cmp(rhs)),
280            (I32(lhs), I32(rhs)) => Some(lhs.cmp(rhs)),
281            (I64(lhs), I64(rhs)) => Some(lhs.cmp(rhs)),
282            (U8(lhs), U8(rhs)) => Some(lhs.cmp(rhs)),
283            (U16(lhs), U16(rhs)) => Some(lhs.cmp(rhs)),
284            (U32(lhs), U32(rhs)) => Some(lhs.cmp(rhs)),
285            (U64(lhs), U64(rhs)) => Some(lhs.cmp(rhs)),
286            (F32(lhs), F32(rhs)) => lhs.partial_cmp(rhs),
287            (F64(lhs), F64(rhs)) => lhs.partial_cmp(rhs),
288            (Str { ptr: lp, len: ll }, Str { ptr: rp, len: rl }) => {
289                let l = unsafe { std::slice::from_raw_parts(lp.as_ptr() as *const u8, *ll) };
290                let r = unsafe { std::slice::from_raw_parts(rp.as_ptr() as *const u8, *rl) };
291                Some(l.cmp(r))
292            }
293            (Bin { ptr: lp, len: ll }, Bin { ptr: rp, len: rl }) => {
294                let l = unsafe { std::slice::from_raw_parts(lp.as_ptr() as *const u8, *ll) };
295                let r = unsafe { std::slice::from_raw_parts(rp.as_ptr() as *const u8, *rl) };
296                Some(l.cmp(r))
297            }
298            _ => {
299                // Allow mixed-width numeric comparisons when both sides are ints of the same sign.
300                if let (Some(lhs), Some(rhs)) = (
301                    signed_int_to_i128(self.cell.as_raw()),
302                    signed_int_to_i128(other.cell.as_raw()),
303                ) {
304                    return Some(lhs.cmp(&rhs));
305                }
306                if let (Some(lhs), Some(rhs)) = (
307                    unsigned_int_to_u128(self.cell.as_raw()),
308                    unsigned_int_to_u128(other.cell.as_raw()),
309                ) {
310                    return Some(lhs.cmp(&rhs));
311                }
312                self.cells_equal(other).then_some(Ordering::Equal)
313            }
314        }
315    }
316
317    /// Extract as `bool` when possible.
318    pub fn as_bool(&self) -> Option<bool> {
319        self.cell.as_bool()
320    }
321
322    /// Extract as signed integer across supported widths.
323    pub fn as_int_i128(&self) -> Option<i128> {
324        match self.cell.as_raw() {
325            DynCellRaw::I8(v) => Some(i128::from(*v)),
326            DynCellRaw::I16(v) => Some(i128::from(*v)),
327            DynCellRaw::I32(v) => Some(i128::from(*v)),
328            DynCellRaw::I64(v) => Some(i128::from(*v)),
329            _ => None,
330        }
331    }
332
333    /// Extract as unsigned integer across supported widths.
334    pub fn as_uint_u128(&self) -> Option<u128> {
335        match self.cell.as_raw() {
336            DynCellRaw::U8(v) => Some(u128::from(*v)),
337            DynCellRaw::U16(v) => Some(u128::from(*v)),
338            DynCellRaw::U32(v) => Some(u128::from(*v)),
339            DynCellRaw::U64(v) => Some(u128::from(*v)),
340            _ => None,
341        }
342    }
343
344    /// Extract as 64-bit floating point.
345    pub fn as_f64(&self) -> Option<f64> {
346        match self.cell.as_raw() {
347            DynCellRaw::F32(value) => Some(f64::from(*value)),
348            DynCellRaw::F64(value) => Some(*value),
349            _ => None,
350        }
351    }
352
353    /// Extract as string slice.
354    pub fn as_utf8(&self) -> Option<&'a str> {
355        self.cell.as_str()
356    }
357
358    /// Extract as binary slice.
359    pub fn as_binary(&self) -> Option<&'a [u8]> {
360        self.cell.as_bin()
361    }
362
363    /// Access the underlying dynamic cell reference.
364    #[must_use]
365    pub fn as_dyn(&self) -> DynCellRef<'a> {
366        self.cell.clone()
367    }
368
369    /// Construct from a dynamic cell reference.
370    pub fn from_dyn(cell: DynCellRef<'a>) -> Self {
371        Self { cell }
372    }
373}
374
375impl<'a> PartialEq for ScalarValueRef<'a> {
376    fn eq(&self, other: &Self) -> bool {
377        match (self.is_null(), other.is_null()) {
378            (true, true) => true,
379            _ => self
380                .compare(other)
381                .map(|ord| ord == Ordering::Equal)
382                .unwrap_or(false),
383        }
384    }
385}