vortex_array/aggregate_fn/

proto.rs

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

use vortex_error::VortexResult;
use vortex_error::vortex_bail;
use vortex_error::vortex_err;
use vortex_proto::expr as pb;
use vortex_session::VortexSession;

use crate::aggregate_fn::AggregateFnId;
use crate::aggregate_fn::AggregateFnRef;
use crate::aggregate_fn::new_foreign_aggregate_fn;
use crate::aggregate_fn::session::AggregateFnSessionExt;

impl AggregateFnRef {
    /// Serialize this aggregate function to its protobuf representation.
    ///
    /// Note: the serialization format is not stable and may change between versions.
    pub fn serialize_proto(&self) -> VortexResult<pb::AggregateFn> {
        let metadata = self
            .options()
            .serialize()?
            .ok_or_else(|| vortex_err!("Aggregate function '{}' is not serializable", self.id()))?;

        Ok(pb::AggregateFn {
            id: self.id().to_string(),
            metadata: Some(metadata),
        })
    }

    /// Deserialize an aggregate function from its protobuf representation.
    ///
    /// Looks up the aggregate function plugin by ID in the session's registry
    /// and delegates deserialization to it.
    ///
    /// Note: the serialization format is not stable and may change between versions.
    pub fn from_proto(proto: &pb::AggregateFn, session: &VortexSession) -> VortexResult<Self> {
        #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")]
        let agg_fn_id: AggregateFnId = AggregateFnId::new(proto.id.as_str());
        let agg_fn = if let Some(plugin) = session.aggregate_fns().find_plugin(&agg_fn_id) {
            plugin.deserialize(proto.metadata(), session)?
        } else if session.allows_unknown() {
            new_foreign_aggregate_fn(agg_fn_id, proto.metadata().to_vec())
        } else {
            return Err(vortex_err!("unknown aggregate function id: {}", proto.id));
        };

        if agg_fn.id() != agg_fn_id {
            vortex_bail!(
                "Aggregate function ID mismatch: expected {}, got {}",
                agg_fn_id,
                agg_fn.id()
            );
        }

        Ok(agg_fn)
    }
}

#[cfg(test)]
mod tests {
    use prost::Message;
    use rstest::rstest;
    use vortex_error::VortexResult;
    use vortex_error::vortex_panic;
    use vortex_proto::expr as pb;
    use vortex_session::VortexSession;

    use crate::ArrayRef;
    use crate::Columnar;
    use crate::ExecutionCtx;
    use crate::aggregate_fn::AggregateFnId;
    use crate::aggregate_fn::AggregateFnRef;
    use crate::aggregate_fn::AggregateFnVTable;
    use crate::aggregate_fn::AggregateFnVTableExt;
    use crate::aggregate_fn::EmptyOptions;
    use crate::aggregate_fn::NumericalAggregateOpts;
    use crate::aggregate_fn::fns::sum::Sum;
    use crate::aggregate_fn::session::AggregateFnSession;
    use crate::aggregate_fn::session::AggregateFnSessionExt;
    use crate::dtype::DType;
    use crate::scalar::Scalar;

    /// A minimal serializable aggregate function used solely to exercise the serde round-trip.
    #[derive(Clone, Debug)]
    struct TestAgg;

    impl AggregateFnVTable for TestAgg {
        type Options = EmptyOptions;
        type Partial = ();

        #[expect(clippy::disallowed_methods, reason = "test-only id")]
        fn id(&self) -> AggregateFnId {
            AggregateFnId::new("vortex.test.proto")
        }

        fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
            Ok(Some(vec![]))
        }

        fn deserialize(
            &self,
            _metadata: &[u8],
            _session: &VortexSession,
        ) -> VortexResult<Self::Options> {
            Ok(EmptyOptions)
        }

        fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option<DType> {
            Some(input_dtype.clone())
        }

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

        fn empty_partial(
            &self,
            _options: &Self::Options,
            _input_dtype: &DType,
        ) -> VortexResult<Self::Partial> {
            Ok(())
        }

        fn combine_partials(
            &self,
            _partial: &mut Self::Partial,
            _other: Scalar,
        ) -> VortexResult<()> {
            Ok(())
        }

        fn to_scalar(&self, _partial: &Self::Partial) -> VortexResult<Scalar> {
            vortex_panic!("TestAgg is for serde tests only");
        }

        fn reset(&self, _partial: &mut Self::Partial) {}

        fn is_saturated(&self, _partial: &Self::Partial) -> bool {
            true
        }

        fn accumulate(
            &self,
            _state: &mut Self::Partial,
            _batch: &Columnar,
            _ctx: &mut ExecutionCtx,
        ) -> VortexResult<()> {
            Ok(())
        }

        fn finalize(&self, partials: ArrayRef) -> VortexResult<ArrayRef> {
            Ok(partials)
        }

        fn finalize_scalar(&self, _partial: &Self::Partial) -> VortexResult<Scalar> {
            vortex_panic!("TestAgg is for serde tests only");
        }
    }

    #[test]
    fn aggregate_fn_serde() {
        let session = crate::array_session();
        session.aggregate_fns().register(TestAgg);

        let agg_fn = TestAgg.bind(EmptyOptions);

        let serialized = agg_fn.serialize_proto().unwrap();
        let buf = serialized.encode_to_vec();
        let deserialized_proto = pb::AggregateFn::decode(buf.as_slice()).unwrap();
        let deserialized = AggregateFnRef::from_proto(&deserialized_proto, &session).unwrap();

        assert_eq!(deserialized, agg_fn);
    }

    /// The `skip_nans` option must survive a protobuf serialize/deserialize round-trip for the
    /// numeric aggregates, including the non-default NaN-including configuration.
    #[rstest]
    #[case(NumericalAggregateOpts::skip_nans())]
    #[case(NumericalAggregateOpts::include_nans())]
    fn numeric_aggregate_options_round_trip(
        #[case] options: NumericalAggregateOpts,
    ) -> VortexResult<()> {
        let session = crate::array_session();
        let agg_fn = Sum.bind(options);

        let proto = agg_fn.serialize_proto()?;
        let buf = proto.encode_to_vec();
        let decoded = pb::AggregateFn::decode(buf.as_slice())?;
        let round_tripped = AggregateFnRef::from_proto(&decoded, &session)?;

        assert_eq!(round_tripped, agg_fn);
        Ok(())
    }

    #[test]
    fn unknown_aggregate_fn_id_allow_unknown() {
        let session = VortexSession::empty()
            .with::<AggregateFnSession>()
            .allow_unknown();

        let proto = pb::AggregateFn {
            id: "vortex.test.foreign_aggregate".to_string(),
            metadata: Some(vec![7, 8, 9]),
        };

        let agg_fn = AggregateFnRef::from_proto(&proto, &session).unwrap();
        assert_eq!(agg_fn.id().as_ref(), "vortex.test.foreign_aggregate");

        let roundtrip = agg_fn.serialize_proto().unwrap();
        assert_eq!(roundtrip.id, proto.id);
        assert_eq!(roundtrip.metadata(), proto.metadata());
    }
}