spiral-sdk 0.1.8

//! Top-k and aggregate emits: what each split of a scan returns in place of its rows, and how
//! clients merge the splits' results. Partial results and their merges live here, so a client
//! merges without the native runtime.

mod aggregate;
mod top_k;

use std::sync::LazyLock;

use arrow_array::RecordBatch;
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use vortex_session::VortexSession;

use self::aggregate::AggregateMerger;
pub use self::aggregate::BoundAggregate;
pub use self::top_k::BoundTopK;
pub use super::api::{
    Aggregate, AggregateFunction, ColumnAggregate, CountAggregate, Emit, HistogramAggregate,
    NamedAggregate, SortKey, TopK,
};
use super::error::ScandalError;

/// The most rows a top-k keeps. Larger orderings belong in the client's own engine.
pub const MAX_TOP_K: u32 = 10_000;
/// The groups an aggregate keeps unless it sets `maxGroups`.
pub const DEFAULT_MAX_GROUPS: u32 = 1_000;
/// The most groups an aggregate may keep.
pub const MAX_GROUPS: u32 = 100_000;
/// The most bins a histogram counts.
pub const MAX_HISTOGRAM_BINS: u32 = 4_096;
/// The aggregate output column holding the rows each group covers.
pub const ROWS_COLUMN: &str = "scandal.rows";
/// The aggregate output column that is true on the row folding the groups past the cap.
pub const OTHER_COLUMN: &str = "scandal.other";

/// An emit checked against its partial results' schema.
#[derive(Clone, Debug)]
pub enum BoundEmit {
    TopK(BoundTopK),
    Aggregate(BoundAggregate),
}

/// The session that merges results: Vortex's arrays, dtypes and aggregate functions.
static MERGE_SESSION: LazyLock<VortexSession> = LazyLock::new(vortex_array::array_session);

impl Emit {
    /// Check this emit against the schema of its results, as a client does to merge them
    /// without the rows they came from.
    pub fn bind_results(&self, results: &SchemaRef) -> Result<BoundEmit, ScandalError> {
        match self {
            Self::TopK(top_k) => BoundTopK::bind(top_k, results).map(BoundEmit::TopK),
            Self::Aggregate(aggregate) => {
                BoundAggregate::bind_results(aggregate, results, &MERGE_SESSION)
                    .map(BoundEmit::Aggregate)
            }
            Self::Unknown { tag, .. } => Err(invalid(format!("unknown emit `{tag}`"))),
        }
    }
}

impl BoundEmit {
    /// The schema of each split's results and of their merge.
    pub fn schema(&self) -> &SchemaRef {
        match self {
            Self::TopK(top_k) => top_k.schema(),
            Self::Aggregate(aggregate) => aggregate.schema(),
        }
    }

    /// Merge the splits' results, in split order, into the scan's result.
    pub fn merge(&self, results: &[RecordBatch]) -> Result<RecordBatch, ScandalError> {
        match self {
            Self::TopK(top_k) => top_k.merge(results),
            Self::Aggregate(aggregate) => {
                let mut table = AggregateMerger::new(aggregate)?;
                for result in results {
                    table.merge(result)?;
                }
                table.finish()
            }
        }
    }
}

/// The index and field of the output column `name`, which an emit's `role` names.
fn column<'a>(
    rows: &'a Schema,
    name: &str,
    role: &str,
) -> Result<(usize, &'a Field), ScandalError> {
    rows.column_with_name(name)
        .ok_or_else(|| invalid(format!("{role} `{name}` is not an output column")))
}

/// Refuse a column whose values do not order as scalars: lists, structs, maps and unions.
fn require_orderable(field: &Field, role: &str) -> Result<(), ScandalError> {
    let orderable = matches!(
        field.data_type(),
        DataType::Boolean
            | DataType::Int8
            | DataType::Int16
            | DataType::Int32
            | DataType::Int64
            | DataType::UInt8
            | DataType::UInt16
            | DataType::UInt32
            | DataType::UInt64
            | DataType::Float16
            | DataType::Float32
            | DataType::Float64
            | DataType::Decimal32(..)
            | DataType::Decimal64(..)
            | DataType::Decimal128(..)
            | DataType::Decimal256(..)
            | DataType::Date32
            | DataType::Date64
            | DataType::Time32(_)
            | DataType::Time64(_)
            | DataType::Timestamp(..)
            | DataType::Duration(_)
            | DataType::Utf8
            | DataType::LargeUtf8
            | DataType::Binary
            | DataType::LargeBinary
            | DataType::FixedSizeBinary(_)
    );
    if orderable {
        Ok(())
    } else {
        Err(invalid(format!(
            "{role} `{}` has type {}, which does not order",
            field.name(),
            field.data_type()
        )))
    }
}

fn invalid(message: impl Into<String>) -> ScandalError {
    ScandalError::InvalidArgument {
        message: message.into(),
    }
}

fn internal(error: impl ToString) -> ScandalError {
    ScandalError::Internal {
        message: error.to_string(),
    }
}