1use akar_common::types::{PhysicalTypeID, Value};
4use akar_common::vector::{DataChunk, ValueVector};
5
6pub(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#[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#[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#[inline]
203pub(crate) fn hash_row(row: &[Value]) -> u64 {
204 use std::hash::{Hash, Hasher};
205 let mut hasher = std::collections::hash_map::DefaultHasher::new();
206 for (i, val) in row.iter().enumerate() {
207 i.hash(&mut hasher);
208 hash_value_into(val, &mut hasher);
209 }
210 hasher.finish()
211}
212
213#[cfg(test)]
214mod tests {
215 use super::*;
216
217 #[test]
218 fn store_value_in_vector_string_overflow_returns_error() {
219 let mut v = ValueVector::new(PhysicalTypeID::String, 1);
220 let err = store_value_in_vector(&mut v, 0, &Value::String("a".repeat(256))).unwrap_err();
221 assert!(err.contains("255"), "err: {err}");
222 }
223
224 #[test]
225 fn store_value_in_vector_string_255_round_trips() {
226 let mut v = ValueVector::new(PhysicalTypeID::String, 1);
227 v.resize(1);
228 let s = "b".repeat(255);
229 store_value_in_vector(&mut v, 0, &Value::String(s.clone())).unwrap();
230 assert_eq!(v.get_value(0), Some(Value::String(s)));
231 }
232}