akar_processor/physical/
common.rs

1//! Common utility functions used across physical operators.
2
3use akar_common::types::{PhysicalTypeID, Value};
4use akar_common::vector::{DataChunk, ValueVector};
5
6/// Materialize a set of global row indices (which may span multiple input
7/// chunks) into a new `DataChunk` via Arrow's `take`.
8///
9/// Unlike the legacy `ValueVector` round-trip, this preserves complex types
10/// (List/Struct/Array), which `store_value_in_vector` would drop to NULL.
11pub(crate) fn take_global_rows(
12    chunks: &[DataChunk],
13    global_indices: &[usize],
14    field_names: Vec<String>,
15) -> Result<DataChunk, String> {
16    let num_fields = chunks.first().map(|c| c.num_fields()).unwrap_or(0);
17    if num_fields == 0 {
18        return Ok(DataChunk::new(Vec::new(), Vec::new()).with_names(field_names));
19    }
20    let indices_arr = arrow::array::UInt32Array::from_iter_values(global_indices.iter().map(|&i| i as u32));
21    let mut fields = Vec::with_capacity(num_fields);
22    for col in 0..num_fields {
23        let field = if chunks.len() == 1 {
24            chunks[0].fields[col].clone()
25        } else {
26            let parts: Vec<arrow::array::ArrayRef> = chunks.iter().map(|c| c.fields[col].clone()).collect();
27            let refs: Vec<&dyn arrow::array::Array> = parts.iter().map(|p| p.as_ref()).collect();
28            arrow::compute::concat(&refs).map_err(|e| e.to_string())?
29        };
30        fields.push(arrow::compute::take(field.as_ref(), &indices_arr, None).map_err(|e| e.to_string())?);
31    }
32    Ok(DataChunk::new(fields, chunks[0].field_types.clone()).with_names(field_names))
33}
34
35#[inline]
36pub(crate) fn store_value_in_vector(v: &mut ValueVector, row: usize, val: &Value) -> Result<(), String> {
37    match val {
38        Value::Null => {
39            v.set_null(row, true);
40        }
41        Value::Bool(x) => {
42            if v.physical_type() == PhysicalTypeID::Bool {
43                v.data_mut()[row] = if *x { 1 } else { 0 };
44                v.set_null(row, false);
45            }
46        }
47        Value::Int64(x) => {
48            let offset = row * 8;
49            if offset + 8 <= v.data().len() {
50                v.data_mut()[offset..offset + 8].copy_from_slice(&x.to_le_bytes());
51                v.set_null(row, false);
52            }
53        }
54        Value::UInt64(x) => {
55            let offset = row * 8;
56            if offset + 8 <= v.data().len() {
57                v.data_mut()[offset..offset + 8].copy_from_slice(&x.to_le_bytes());
58                v.set_null(row, false);
59            }
60        }
61        Value::Int32(x) => {
62            let offset = row * 4;
63            if offset + 4 <= v.data().len() {
64                v.data_mut()[offset..offset + 4].copy_from_slice(&x.to_le_bytes());
65                v.set_null(row, false);
66            }
67        }
68        Value::Double(x) => {
69            let offset = row * 8;
70            if offset + 8 <= v.data().len() {
71                v.data_mut()[offset..offset + 8].copy_from_slice(&x.to_le_bytes());
72                v.set_null(row, false);
73            }
74        }
75        Value::Float(x) => {
76            let offset = row * 4;
77            if offset + 4 <= v.data().len() {
78                v.data_mut()[offset..offset + 4].copy_from_slice(&x.to_le_bytes());
79                v.set_null(row, false);
80            }
81        }
82        Value::String(s) => {
83            let bytes = s.as_bytes();
84            if bytes.len() > 255 {
85                return Err(format!(
86                    "Cannot store string of {} bytes: inline string storage limit is 255 bytes",
87                    bytes.len()
88                ));
89            }
90            let offset = row * 256;
91            if offset < v.data().len() {
92                v.data_mut()[offset] = bytes.len() as u8;
93                if offset + 1 + bytes.len() <= v.data().len() {
94                    v.data_mut()[offset + 1..offset + 1 + bytes.len()].copy_from_slice(bytes);
95                }
96                v.set_null(row, false);
97            }
98        }
99        _ => {
100            v.set_null(row, true);
101        }
102    }
103    Ok(())
104}
105
106#[inline(always)]
107pub(crate) fn value_cmp(a: &Value, b: &Value) -> std::cmp::Ordering {
108    match (a, b) {
109        (Value::Null, Value::Null) => std::cmp::Ordering::Equal,
110        (Value::Null, _) => std::cmp::Ordering::Greater,
111        (_, Value::Null) => std::cmp::Ordering::Less,
112        (Value::Bool(x), Value::Bool(y)) => x.cmp(y),
113        (Value::Int64(x), Value::Int64(y)) => x.cmp(y),
114        (Value::Int32(x), Value::Int32(y)) => x.cmp(y),
115        (Value::Int16(x), Value::Int16(y)) => (*x as i64).cmp(&(*y as i64)),
116        (Value::Int8(x), Value::Int8(y)) => (*x as i64).cmp(&(*y as i64)),
117        (Value::UInt64(x), Value::UInt64(y)) => x.cmp(y),
118        (Value::UInt32(x), Value::UInt32(y)) => (*x as u64).cmp(&(*y as u64)),
119        (Value::Double(x), Value::Double(y)) => double_cmp(*x, *y),
120        (Value::Float(x), Value::Float(y)) => double_cmp(*x as f64, *y as f64),
121        (Value::String(x), Value::String(y)) => x.cmp(y),
122        (Value::Date(x), Value::Date(y)) => x.0.cmp(&y.0),
123        (Value::Timestamp(x), Value::Timestamp(y)) => x.0.cmp(&y.0),
124        _ => std::cmp::Ordering::Equal,
125    }
126}
127
128/// Total-order comparison for floats with a NaN convention: NaN sorts greater
129/// than every finite value, and NaN == NaN. `partial_cmp(...).unwrap_or(Equal)`
130/// would treat NaN as equal to everything, breaking ORDER BY / TOP-K sort order.
131#[inline(always)]
132pub(crate) fn double_cmp(a: f64, b: f64) -> std::cmp::Ordering {
133    if a.is_nan() {
134        if b.is_nan() {
135            std::cmp::Ordering::Equal
136        } else {
137            std::cmp::Ordering::Greater
138        }
139    } else if b.is_nan() {
140        std::cmp::Ordering::Less
141    } else {
142        a.partial_cmp(&b).unwrap_or(std::cmp::Ordering::Equal)
143    }
144}
145
146#[inline]
147pub(crate) fn value_hash(val: &Value) -> u64 {
148    use std::hash::Hasher;
149    let mut hasher = std::collections::hash_map::DefaultHasher::new();
150    hash_value_into(val, &mut hasher);
151    hasher.finish()
152}
153
154/// Write a Value's hash into an arbitrary Hasher.
155#[inline]
156pub(crate) fn hash_value_into(val: &Value, hasher: &mut impl std::hash::Hasher) {
157    use std::hash::Hash;
158    match val {
159        Value::Null => 0u8.hash(hasher),
160        Value::Bool(b) => b.hash(hasher),
161        Value::Int64(i) => i.hash(hasher),
162        Value::Int32(i) => i.hash(hasher),
163        Value::Int16(i) => i.hash(hasher),
164        Value::Int8(i) => i.hash(hasher),
165        Value::UInt64(i) => i.hash(hasher),
166        Value::UInt32(i) => i.hash(hasher),
167        Value::UInt16(i) => (*i as u64).hash(hasher),
168        Value::UInt8(i) => (*i as u64).hash(hasher),
169        Value::Int128(i) => i.hash(hasher),
170        Value::Double(f) => f.to_bits().hash(hasher),
171        Value::Float(f) => f.to_bits().hash(hasher),
172        Value::String(s) => s.hash(hasher),
173        Value::Blob(b) => b.hash(hasher),
174        Value::Date(d) => d.0.hash(hasher),
175        Value::Timestamp(t) => t.0.hash(hasher),
176        Value::TimestampTz(t) => t.0.hash(hasher),
177        Value::TimestampNs(t) => t.0.hash(hasher),
178        Value::TimestampMs(t) => t.0.hash(hasher),
179        Value::TimestampSec(t) => t.0.hash(hasher),
180        Value::Interval(i) => (i.months, i.days, i.micros).hash(hasher),
181        Value::InternalID(id) => id.offset.hash(hasher),
182        Value::UInt128(i) => i.hash(hasher),
183        Value::List(vals) => {
184            for v in vals {
185                hash_value_into(v, hasher);
186            }
187        }
188        Value::Map(kvs) => {
189            for (k, v) in kvs {
190                hash_value_into(k, hasher);
191                hash_value_into(v, hasher);
192            }
193        }
194        Value::Union(_, v) => hash_value_into(v, hasher),
195        _ => std::mem::discriminant(val).hash(hasher),
196    }
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202
203    #[test]
204    fn store_value_in_vector_string_overflow_returns_error() {
205        let mut v = ValueVector::new(PhysicalTypeID::String, 1);
206        let err = store_value_in_vector(&mut v, 0, &Value::String("a".repeat(256))).unwrap_err();
207        assert!(err.contains("255"), "err: {err}");
208    }
209
210    #[test]
211    fn store_value_in_vector_string_255_round_trips() {
212        let mut v = ValueVector::new(PhysicalTypeID::String, 1);
213        v.resize(1);
214        let s = "b".repeat(255);
215        store_value_in_vector(&mut v, 0, &Value::String(s.clone())).unwrap();
216        assert_eq!(v.get_value(0), Some(Value::String(s)));
217    }
218}
akar_processor/physical/common.rs

akar_processor/physical/
common.rs
