akar-processor 0.1.4

//! Common utility functions used across physical operators.

use akar_common::types::{PhysicalTypeID, Value};
use akar_common::vector::{DataChunk, ValueVector};

/// Materialize a set of global row indices (which may span multiple input
/// chunks) into a new `DataChunk` via Arrow's `take`.
///
/// Unlike the legacy `ValueVector` round-trip, this preserves complex types
/// (List/Struct/Array), which `store_value_in_vector` would drop to NULL.
pub(crate) fn take_global_rows(
    chunks: &[DataChunk],
    global_indices: &[usize],
    field_names: Vec<String>,
) -> Result<DataChunk, String> {
    let num_fields = chunks.first().map(|c| c.num_fields()).unwrap_or(0);
    if num_fields == 0 {
        return Ok(DataChunk::new(Vec::new(), Vec::new()).with_names(field_names));
    }
    let indices_arr = arrow::array::UInt32Array::from_iter_values(global_indices.iter().map(|&i| i as u32));
    let mut fields = Vec::with_capacity(num_fields);
    for col in 0..num_fields {
        let field = if chunks.len() == 1 {
            chunks[0].fields[col].clone()
        } else {
            let parts: Vec<arrow::array::ArrayRef> = chunks.iter().map(|c| c.fields[col].clone()).collect();
            let refs: Vec<&dyn arrow::array::Array> = parts.iter().map(|p| p.as_ref()).collect();
            arrow::compute::concat(&refs).map_err(|e| e.to_string())?
        };
        fields.push(arrow::compute::take(field.as_ref(), &indices_arr, None).map_err(|e| e.to_string())?);
    }
    Ok(DataChunk::new(fields, chunks[0].field_types.clone()).with_names(field_names))
}

#[inline]
pub(crate) fn store_value_in_vector(v: &mut ValueVector, row: usize, val: &Value) -> Result<(), String> {
    match val {
        Value::Null => {
            v.set_null(row, true);
        }
        Value::Bool(x) => {
            if v.physical_type() == PhysicalTypeID::Bool {
                v.data_mut()[row] = if *x { 1 } else { 0 };
                v.set_null(row, false);
            }
        }
        Value::Int64(x) => {
            let offset = row * 8;
            if offset + 8 <= v.data().len() {
                v.data_mut()[offset..offset + 8].copy_from_slice(&x.to_le_bytes());
                v.set_null(row, false);
            }
        }
        Value::UInt64(x) => {
            let offset = row * 8;
            if offset + 8 <= v.data().len() {
                v.data_mut()[offset..offset + 8].copy_from_slice(&x.to_le_bytes());
                v.set_null(row, false);
            }
        }
        Value::Int32(x) => {
            let offset = row * 4;
            if offset + 4 <= v.data().len() {
                v.data_mut()[offset..offset + 4].copy_from_slice(&x.to_le_bytes());
                v.set_null(row, false);
            }
        }
        Value::Double(x) => {
            let offset = row * 8;
            if offset + 8 <= v.data().len() {
                v.data_mut()[offset..offset + 8].copy_from_slice(&x.to_le_bytes());
                v.set_null(row, false);
            }
        }
        Value::Float(x) => {
            let offset = row * 4;
            if offset + 4 <= v.data().len() {
                v.data_mut()[offset..offset + 4].copy_from_slice(&x.to_le_bytes());
                v.set_null(row, false);
            }
        }
        Value::String(s) => {
            let bytes = s.as_bytes();
            if bytes.len() > 255 {
                return Err(format!(
                    "Cannot store string of {} bytes: inline string storage limit is 255 bytes",
                    bytes.len()
                ));
            }
            let offset = row * 256;
            if offset < v.data().len() {
                v.data_mut()[offset] = bytes.len() as u8;
                if offset + 1 + bytes.len() <= v.data().len() {
                    v.data_mut()[offset + 1..offset + 1 + bytes.len()].copy_from_slice(bytes);
                }
                v.set_null(row, false);
            }
        }
        _ => {
            v.set_null(row, true);
        }
    }
    Ok(())
}

#[inline(always)]
pub(crate) fn value_cmp(a: &Value, b: &Value) -> std::cmp::Ordering {
    match (a, b) {
        (Value::Null, Value::Null) => std::cmp::Ordering::Equal,
        (Value::Null, _) => std::cmp::Ordering::Greater,
        (_, Value::Null) => std::cmp::Ordering::Less,
        (Value::Bool(x), Value::Bool(y)) => x.cmp(y),
        (Value::Int64(x), Value::Int64(y)) => x.cmp(y),
        (Value::Int32(x), Value::Int32(y)) => x.cmp(y),
        (Value::Int16(x), Value::Int16(y)) => (*x as i64).cmp(&(*y as i64)),
        (Value::Int8(x), Value::Int8(y)) => (*x as i64).cmp(&(*y as i64)),
        (Value::UInt64(x), Value::UInt64(y)) => x.cmp(y),
        (Value::UInt32(x), Value::UInt32(y)) => (*x as u64).cmp(&(*y as u64)),
        (Value::Double(x), Value::Double(y)) => double_cmp(*x, *y),
        (Value::Float(x), Value::Float(y)) => double_cmp(*x as f64, *y as f64),
        (Value::String(x), Value::String(y)) => x.cmp(y),
        (Value::Date(x), Value::Date(y)) => x.0.cmp(&y.0),
        (Value::Timestamp(x), Value::Timestamp(y)) => x.0.cmp(&y.0),
        _ => std::cmp::Ordering::Equal,
    }
}

/// Total-order comparison for floats with a NaN convention: NaN sorts greater
/// than every finite value, and NaN == NaN. `partial_cmp(...).unwrap_or(Equal)`
/// would treat NaN as equal to everything, breaking ORDER BY / TOP-K sort order.
#[inline(always)]
pub(crate) fn double_cmp(a: f64, b: f64) -> std::cmp::Ordering {
    if a.is_nan() {
        if b.is_nan() {
            std::cmp::Ordering::Equal
        } else {
            std::cmp::Ordering::Greater
        }
    } else if b.is_nan() {
        std::cmp::Ordering::Less
    } else {
        a.partial_cmp(&b).unwrap_or(std::cmp::Ordering::Equal)
    }
}

#[inline]
pub(crate) fn value_hash(val: &Value) -> u64 {
    use std::hash::Hasher;
    let mut hasher = std::collections::hash_map::DefaultHasher::new();
    hash_value_into(val, &mut hasher);
    hasher.finish()
}

/// Write a Value's hash into an arbitrary Hasher.
#[inline]
pub(crate) fn hash_value_into(val: &Value, hasher: &mut impl std::hash::Hasher) {
    use std::hash::Hash;
    match val {
        Value::Null => 0u8.hash(hasher),
        Value::Bool(b) => b.hash(hasher),
        Value::Int64(i) => i.hash(hasher),
        Value::Int32(i) => i.hash(hasher),
        Value::Int16(i) => i.hash(hasher),
        Value::Int8(i) => i.hash(hasher),
        Value::UInt64(i) => i.hash(hasher),
        Value::UInt32(i) => i.hash(hasher),
        Value::UInt16(i) => (*i as u64).hash(hasher),
        Value::UInt8(i) => (*i as u64).hash(hasher),
        Value::Int128(i) => i.hash(hasher),
        Value::Double(f) => f.to_bits().hash(hasher),
        Value::Float(f) => f.to_bits().hash(hasher),
        Value::String(s) => s.hash(hasher),
        Value::Blob(b) => b.hash(hasher),
        Value::Date(d) => d.0.hash(hasher),
        Value::Timestamp(t) => t.0.hash(hasher),
        Value::TimestampTz(t) => t.0.hash(hasher),
        Value::TimestampNs(t) => t.0.hash(hasher),
        Value::TimestampMs(t) => t.0.hash(hasher),
        Value::TimestampSec(t) => t.0.hash(hasher),
        Value::Interval(i) => (i.months, i.days, i.micros).hash(hasher),
        Value::InternalID(id) => id.offset.hash(hasher),
        Value::UInt128(i) => i.hash(hasher),
        Value::List(vals) => {
            for v in vals {
                hash_value_into(v, hasher);
            }
        }
        Value::Map(kvs) => {
            for (k, v) in kvs {
                hash_value_into(k, hasher);
                hash_value_into(v, hasher);
            }
        }
        Value::Union(_, v) => hash_value_into(v, hasher),
        _ => std::mem::discriminant(val).hash(hasher),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn store_value_in_vector_string_overflow_returns_error() {
        let mut v = ValueVector::new(PhysicalTypeID::String, 1);
        let err = store_value_in_vector(&mut v, 0, &Value::String("a".repeat(256))).unwrap_err();
        assert!(err.contains("255"), "err: {err}");
    }

    #[test]
    fn store_value_in_vector_string_255_round_trips() {
        let mut v = ValueVector::new(PhysicalTypeID::String, 1);
        v.resize(1);
        let s = "b".repeat(255);
        store_value_in_vector(&mut v, 0, &Value::String(s.clone())).unwrap();
        assert_eq!(v.get_value(0), Some(Value::String(s)));
    }
}