vortex-layout 0.84.0

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

//! Shared helpers for the zoned layout's auxiliary stats-table schema.

use std::sync::Arc;

use vortex_array::aggregate_fn::AggregateFnId;
use vortex_array::aggregate_fn::AggregateFnRef;
use vortex_array::aggregate_fn::session::AggregateFnSessionExt;
use vortex_array::dtype::DType;
use vortex_array::dtype::Nullability;
use vortex_array::dtype::StructFields;
use vortex_array::expr::stats::Stat;
use vortex_error::VortexResult;
use vortex_error::vortex_bail;
use vortex_error::vortex_err;
use vortex_session::VortexSession;

pub const MAX_IS_TRUNCATED: &str = "max_is_truncated";
pub const MIN_IS_TRUNCATED: &str = "min_is_truncated";

#[derive(Clone, PartialEq, Eq, ::prost::Message)]
pub(crate) struct AggregateSpecProto {
    #[prost(string, tag = "1")]
    id: String,
    #[prost(bytes, tag = "2")]
    options: Vec<u8>,
}

impl AggregateSpecProto {
    pub(crate) fn try_from_aggregate_fn(aggregate_fn: &AggregateFnRef) -> VortexResult<Self> {
        let options = aggregate_fn.options().serialize()?.ok_or_else(|| {
            vortex_err!(
                "Aggregate function '{}' is not serializable",
                aggregate_fn.id()
            )
        })?;

        Ok(Self {
            id: aggregate_fn.id().to_string(),
            options,
        })
    }

    /// Resolve this spec to its aggregate function, returning `Ok(None)` for an unknown aggregate
    /// ID when the session allows unknown plugins.
    ///
    /// A `None` result means the zone map cannot be reconstructed for this aggregate, so callers
    /// should disable zoned pruning rather than fail the read.
    fn to_aggregate_fn_opt(&self, session: &VortexSession) -> VortexResult<Option<AggregateFnRef>> {
        #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")]
        let aggregate_fn_id = AggregateFnId::new(self.id.as_str());

        let Some(plugin) = session.aggregate_fns().find_plugin(&aggregate_fn_id) else {
            if session.allows_unknown() {
                return Ok(None);
            }

            vortex_bail!("unknown aggregate function id: {}", self.id);
        };

        let aggregate_fn = plugin.deserialize(&self.options, session)?;
        if aggregate_fn.id() != aggregate_fn_id {
            vortex_bail!(
                "Aggregate function ID mismatch: expected {}, got {}",
                aggregate_fn_id,
                aggregate_fn.id()
            );
        }

        Ok(Some(aggregate_fn))
    }

    /// Construct a spec referencing an aggregate ID that no session resolves, for exercising the
    /// unknown-aggregate read paths.
    #[cfg(test)]
    pub(crate) fn new_unknown(id: &str) -> Self {
        Self {
            id: id.to_string(),
            options: vec![],
        }
    }
}

/// Return the auxiliary stats-table schema for a zoned layout.
pub(crate) fn aggregate_stats_table_dtype(
    column_dtype: &DType,
    aggregate_fns: &[AggregateFnRef],
) -> DType {
    DType::Struct(
        StructFields::from_iter(aggregate_fns.iter().filter_map(|aggregate_fn| {
            aggregate_state_dtype(column_dtype, aggregate_fn)
                .map(|dtype| (aggregate_fn.to_string(), dtype.as_nullable()))
        })),
        Nullability::NonNullable,
    )
}

pub(crate) fn legacy_stats_table_dtype(column_dtype: &DType, present_stats: &[Stat]) -> DType {
    assert!(present_stats.is_sorted(), "Stats must be sorted");
    DType::Struct(
        StructFields::from_iter(
            present_stats
                .iter()
                .filter_map(|stat| {
                    stat.dtype(column_dtype)
                        .or_else(|| {
                            // Backward compat: older files may have stored stats (e.g. Sum)
                            // for extension types by resolving through the storage dtype.
                            if let DType::Extension(ext) = column_dtype {
                                stat.dtype(ext.storage_dtype())
                            } else {
                                None
                            }
                        })
                        .map(|dtype| (stat, dtype.as_nullable()))
                })
                .flat_map(|(stat, dtype)| match stat {
                    Stat::Max => vec![
                        (stat.name(), dtype),
                        (MAX_IS_TRUNCATED, DType::Bool(Nullability::NonNullable)),
                    ],
                    Stat::Min => vec![
                        (stat.name(), dtype),
                        (MIN_IS_TRUNCATED, DType::Bool(Nullability::NonNullable)),
                    ],
                    _ => vec![(stat.name(), dtype)],
                }),
        ),
        Nullability::NonNullable,
    )
}

/// Serializes each live [`AggregateFnRef`] into an [`AggregateSpecProto`] for storage in zoned
/// metadata. The inverse of [`try_aggregate_fns_from_specs`].
///
/// Errors if any aggregate is not serializable.
pub(crate) fn aggregate_specs_from_fns(
    aggregate_fns: &[AggregateFnRef],
) -> VortexResult<Arc<[AggregateSpecProto]>> {
    aggregate_fns
        .iter()
        .map(AggregateSpecProto::try_from_aggregate_fn)
        .collect::<VortexResult<Vec<_>>>()
        .map(Into::into)
}

/// Resolves each [`AggregateSpecProto`] into a live [`AggregateFnRef`], tolerating aggregates this
/// session does not recognize. Rebuilding the `zones` child needs every spec resolved, since its
/// dtype is derived from them.
///
/// Returns:
///
/// - `Ok(Some(fns))`: all resolved, so the zone map can be rebuilt.
/// - `Ok(None)`: an aggregate is unknown but unknown plugins are allowed (see
///   [`VortexSession::allows_unknown`]). The caller should disable pruning and scan only the `data`
///   child, which still returns correct results.
/// - `Err(..)`: any other failure (an unknown aggregate while disallowed, or undeserializable
///   options).
pub(crate) fn try_aggregate_fns_from_specs(
    aggregate_specs: &[AggregateSpecProto],
    session: &VortexSession,
) -> VortexResult<Option<Arc<[AggregateFnRef]>>> {
    aggregate_specs
        .iter()
        .map(|aggregate_spec| aggregate_spec.to_aggregate_fn_opt(session))
        .collect::<VortexResult<Option<Vec<_>>>>()
        .map(|aggregate_fns| aggregate_fns.map(Arc::from))
}

pub(crate) fn aggregate_state_dtype(
    column_dtype: &DType,
    aggregate_fn: &AggregateFnRef,
) -> Option<DType> {
    aggregate_fn.state_dtype(column_dtype).or_else(|| {
        if let DType::Extension(ext) = column_dtype {
            aggregate_fn.state_dtype(ext.storage_dtype())
        } else {
            None
        }
    })
}

pub(crate) fn default_bounded_stat_max_bytes() -> std::num::NonZeroUsize {
    // SAFETY: 64 is non-zero.
    unsafe { std::num::NonZeroUsize::new_unchecked(64) }
}

#[cfg(test)]
mod tests {
    use vortex_array::aggregate_fn::AggregateFnVTableExt;
    use vortex_array::aggregate_fn::NumericalAggregateOpts;
    use vortex_array::aggregate_fn::fns::max::Max;
    use vortex_array::aggregate_fn::fns::min::Min;
    use vortex_array::aggregate_fn::fns::sum::Sum;
    use vortex_array::dtype::DType;
    use vortex_array::dtype::Nullability;
    use vortex_array::dtype::PType;
    use vortex_array::extension::datetime::Date;
    use vortex_array::extension::datetime::TimeUnit;

    use super::*;

    #[test]
    fn unknown_aggregate_errors_without_allow_unknown() {
        let session = vortex_array::array_session();
        let specs = [AggregateSpecProto::new_unknown("vortex.test.unknown")];
        assert!(try_aggregate_fns_from_specs(&specs, &session).is_err());
    }

    #[test]
    fn unknown_aggregate_is_none_with_allow_unknown() -> VortexResult<()> {
        let session = vortex_array::array_session();
        session.allow_unknown();
        let specs = [AggregateSpecProto::new_unknown("vortex.test.unknown")];
        assert!(try_aggregate_fns_from_specs(&specs, &session)?.is_none());
        Ok(())
    }

    #[test]
    fn stats_table_dtype_adds_truncation_flags() {
        let dtype = legacy_stats_table_dtype(
            &DType::Primitive(PType::I32, Nullability::NonNullable),
            &[Stat::Max, Stat::Min, Stat::Sum],
        );

        assert_eq!(
            dtype.as_struct_fields().names().as_ref(),
            &[
                Stat::Max.name(),
                MAX_IS_TRUNCATED,
                Stat::Min.name(),
                MIN_IS_TRUNCATED,
                Stat::Sum.name(),
            ]
        );
    }

    #[test]
    fn stats_table_dtype_uses_storage_dtype_for_extensions() {
        let dtype = DType::Extension(Date::new(TimeUnit::Days, Nullability::NonNullable).erased());
        let stats_dtype = legacy_stats_table_dtype(&dtype, &[Stat::Max, Stat::Min]);

        assert_eq!(
            stats_dtype.as_struct_fields().names().as_ref(),
            &[
                Stat::Max.name(),
                MAX_IS_TRUNCATED,
                Stat::Min.name(),
                MIN_IS_TRUNCATED,
            ]
        );
    }

    #[test]
    fn aggregate_stats_table_dtype_uses_display_names() {
        let dtype = aggregate_stats_table_dtype(
            &DType::Primitive(PType::I32, Nullability::NonNullable),
            &[
                Max.bind(NumericalAggregateOpts::skip_nans()),
                Min.bind(NumericalAggregateOpts::skip_nans()),
                Sum.bind(NumericalAggregateOpts::skip_nans()),
            ],
        );

        assert_eq!(
            dtype.as_struct_fields().names().as_ref(),
            &[
                Max.bind(NumericalAggregateOpts::skip_nans()).to_string(),
                Min.bind(NumericalAggregateOpts::skip_nans()).to_string(),
                Sum.bind(NumericalAggregateOpts::skip_nans()).to_string(),
            ]
        );
    }
}