vortex-array 0.79.0

Vortex in memory columnar data format
Documentation
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use std::sync::Arc;

use vortex_error::VortexResult;
use vortex_session::VortexSession;

use crate::aggregate_fn::AggregateFn;
use crate::aggregate_fn::AggregateFnId;
use crate::aggregate_fn::AggregateFnRef;
use crate::aggregate_fn::AggregateFnVTable;
use crate::dtype::DType;

/// Reference-counted pointer to an aggregate function plugin.
pub type AggregateFnPluginRef = Arc<dyn AggregateFnPlugin>;

/// Registry trait for ID-based deserialization of aggregate functions.
///
/// Plugins are registered in the session by their [`AggregateFnId`]. When a serialized aggregate
/// function is encountered, the session resolves the ID to the plugin and calls [`deserialize`]
/// to reconstruct the value as an [`AggregateFnRef`].
///
/// [`deserialize`]: AggregateFnPlugin::deserialize
pub trait AggregateFnPlugin: 'static + Send + Sync {
    /// Returns the ID for this aggregate function.
    fn id(&self) -> AggregateFnId;

    /// Deserialize an aggregate function from serialized metadata.
    fn deserialize(&self, metadata: &[u8], session: &VortexSession)
    -> VortexResult<AggregateFnRef>;

    /// The default zone statistic (per-chunk) for a column of `input_dtype`, or `None` if the
    /// dtype is not supported (or this aggregate is not a zone statistic at all).
    ///
    /// This is how a registered aggregate volunteers itself as a per-chunk statistic: when a
    /// zoned layout writer opens a column, it collects every plugin's default via
    /// [`AggregateFnSession::zone_stat_defaults`], so new statistics can be added without the
    /// writer knowing about them.
    ///
    /// [`AggregateFnSession::zone_stat_defaults`]: crate::aggregate_fn::session::AggregateFnSession::zone_stat_defaults
    fn zone_stat_default(&self, _input_dtype: &DType) -> Option<AggregateFnRef> {
        None
    }
}

impl std::fmt::Debug for dyn AggregateFnPlugin {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_tuple("AggregateFnPlugin")
            .field(&self.id())
            .finish()
    }
}

impl<V: AggregateFnVTable> AggregateFnPlugin for V {
    fn id(&self) -> AggregateFnId {
        AggregateFnVTable::id(self)
    }

    fn deserialize(
        &self,
        metadata: &[u8],
        session: &VortexSession,
    ) -> VortexResult<AggregateFnRef> {
        let options = AggregateFnVTable::deserialize(self, metadata, session)?;
        Ok(AggregateFn::new(self.clone(), options).erased())
    }

    fn zone_stat_default(&self, input_dtype: &DType) -> Option<AggregateFnRef> {
        AggregateFnVTable::zone_stat_default(self, input_dtype)
    }
}