spiral-sdk 0.1.12

//! `spql.agg.histogram`: counts of a numeric column's values in equal-width bins, as a Vortex
//! aggregate function whose partial counts merge by adding.

use std::fmt::{self, Display, Formatter};
use std::hash::{Hash, Hasher};
use std::sync::Arc;

use vortex_array::aggregate_fn::{AggregateArgs, AggregateFnId, AggregateFnVTable};
use vortex_array::arrays::PrimitiveArray;
use vortex_array::builtins::ArrayBuiltins as _;
use vortex_array::dtype::{DType, Nullability, PType};
use vortex_array::scalar::Scalar;
use vortex_array::{ArrayRef, Canonical, Columnar, ExecutionCtx, IntoArray as _};
use vortex_error::{VortexResult, vortex_bail, vortex_err};
use vortex_session::VortexSession;
use vortex_session::registry::CachedId;

/// Counts of a numeric column's values in `bins` bins of `width` from `start`: values below
/// `start`, each bin in order, then values at or past the last bin's end. Nulls and NaN are not
/// counted. The partial and the result are a non-null fixed-size list of `bins + 2` `u64`s.
#[derive(Clone, Debug)]
pub struct Histogram;

/// Where a histogram's bins start, how wide each is and how many there are.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct HistogramOptions {
    pub start: f64,
    pub width: f64,
    pub bins: u32,
}

impl HistogramOptions {
    /// Counts per histogram: one per bin, and one each below and past them.
    pub fn slots(&self) -> usize {
        self.bins as usize + 2
    }

    /// The slot that `value`, a non-NaN number, counts in.
    fn slot(&self, value: f64) -> usize {
        if value < self.start {
            return 0;
        }
        let bin = ((value - self.start) / self.width).floor();
        if bin >= f64::from(self.bins) {
            self.slots() - 1
        } else {
            bin as usize + 1
        }
    }

    fn dtype(&self) -> DType {
        DType::FixedSizeList(
            Arc::new(DType::Primitive(PType::U64, Nullability::NonNullable)),
            self.bins + 2,
            Nullability::NonNullable,
        )
    }
}

impl Eq for HistogramOptions {}

impl Hash for HistogramOptions {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.start.to_bits().hash(state);
        self.width.to_bits().hash(state);
        self.bins.hash(state);
    }
}

impl Display for HistogramOptions {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "start={}, width={}, bins={}",
            self.start, self.width, self.bins
        )
    }
}

impl AggregateFnVTable for Histogram {
    type Options = HistogramOptions;
    type Partial = Vec<u64>;

    fn id(&self) -> AggregateFnId {
        static ID: CachedId = CachedId::new("spql.agg.histogram");
        *ID
    }

    fn serialize(&self, options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
        let mut bytes = Vec::with_capacity(20);
        bytes.extend_from_slice(&options.start.to_le_bytes());
        bytes.extend_from_slice(&options.width.to_le_bytes());
        bytes.extend_from_slice(&options.bins.to_le_bytes());
        Ok(Some(bytes))
    }

    fn deserialize(
        &self,
        metadata: &[u8],
        _session: &VortexSession,
    ) -> VortexResult<Self::Options> {
        let bytes = |range: std::ops::Range<usize>| -> VortexResult<[u8; 8]> {
            let mut bytes = [0; 8];
            let field = metadata
                .get(range)
                .ok_or_else(|| vortex_err!("histogram options are truncated"))?;
            bytes[..field.len()].copy_from_slice(field);
            Ok(bytes)
        };
        let bins = bytes(16..20)?;
        Ok(HistogramOptions {
            start: f64::from_le_bytes(bytes(0..8)?),
            width: f64::from_le_bytes(bytes(8..16)?),
            bins: u32::from_le_bytes([bins[0], bins[1], bins[2], bins[3]]),
        })
    }

    fn return_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
        matches!(input_dtype, DType::Primitive(..)).then(|| options.dtype())
    }

    fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
        self.return_dtype(options, input_dtype)
    }

    fn empty_partial(&self, args: AggregateArgs<'_, Self::Options>) -> VortexResult<Vec<u64>> {
        Ok(vec![0; args.options.slots()])
    }

    fn partial_from_scalar(
        &self,
        args: AggregateArgs<'_, Self::Options>,
        scalar: &Scalar,
    ) -> VortexResult<Vec<u64>> {
        let counts = scalar
            .as_list()
            .elements()
            .ok_or_else(|| vortex_err!("histogram partial is null"))?
            .iter()
            .map(|count| {
                count
                    .as_primitive()
                    .typed_value::<u64>()
                    .ok_or_else(|| vortex_err!("histogram count is null"))
            })
            .collect::<VortexResult<Vec<_>>>()?;
        if counts.len() != args.options.slots() {
            vortex_bail!("histogram partial has {} counts", counts.len());
        }
        Ok(counts)
    }

    fn merge_partials(
        &self,
        _args: AggregateArgs<'_, Self::Options>,
        mut first: Vec<u64>,
        second: Vec<u64>,
    ) -> VortexResult<Vec<u64>> {
        for (count, other) in first.iter_mut().zip(second) {
            *count += other;
        }
        Ok(first)
    }

    fn to_scalar(
        &self,
        _args: AggregateArgs<'_, Self::Options>,
        partial: &Vec<u64>,
    ) -> VortexResult<Scalar> {
        Ok(Scalar::fixed_size_list(
            Arc::new(DType::Primitive(PType::U64, Nullability::NonNullable)),
            partial
                .iter()
                .map(|count| Scalar::primitive(*count, Nullability::NonNullable))
                .collect(),
            Nullability::NonNullable,
        ))
    }

    fn is_saturated(&self, _args: AggregateArgs<'_, Self::Options>, _partial: &Vec<u64>) -> bool {
        false
    }

    fn accumulate(
        &self,
        args: AggregateArgs<'_, Self::Options>,
        partial: &mut Vec<u64>,
        batch: &Columnar,
        ctx: &mut ExecutionCtx,
    ) -> VortexResult<()> {
        let float = DType::Primitive(PType::F64, Nullability::Nullable);
        match batch {
            Columnar::Constant(constant) => {
                let value = constant.scalar().cast(&float)?;
                if let Some(value) = value.as_primitive().typed_value::<f64>()
                    && !value.is_nan()
                {
                    partial[args.options.slot(value)] += constant.len() as u64;
                }
            }
            Columnar::Canonical(Canonical::Primitive(values)) => {
                let values = values
                    .clone()
                    .into_array()
                    .cast(float)?
                    .execute::<PrimitiveArray>(ctx)?;
                let valid = values.validity()?.execute_mask(values.len(), ctx)?;
                let numbers = values.as_slice::<f64>();
                let mut count = |value: f64| {
                    if !value.is_nan() {
                        partial[args.options.slot(value)] += 1;
                    }
                };
                if valid.all_true() {
                    numbers.iter().copied().for_each(&mut count);
                } else if !valid.all_false() {
                    numbers
                        .iter()
                        .zip(valid.iter())
                        .filter_map(|(value, valid)| valid.then_some(*value))
                        .for_each(&mut count);
                }
            }
            Columnar::Canonical(_) => {
                vortex_bail!("a histogram counts numbers, not {}", batch.dtype())
            }
        }
        Ok(())
    }

    fn finalize(
        &self,
        _args: AggregateArgs<'_, Self::Options>,
        partials: ArrayRef,
    ) -> VortexResult<ArrayRef> {
        Ok(partials)
    }

    fn finalize_scalar(
        &self,
        args: AggregateArgs<'_, Self::Options>,
        partial: &Vec<u64>,
    ) -> VortexResult<Scalar> {
        self.to_scalar(args, partial)
    }
}