feat(function): add mergeable binary average states (#9062)

* feat(function): add mergeable binary average states Signed-off-by: discord9 <discord9@163.com> * feat(function): expose avg_calc as an OSS scalar Signed-off-by: discord9 <discord9@163.com> * fix(function): borrow invalid AVG scalar argument type Signed-off-by: discord9 <discord9@163.com> * test(function): verify public AVG state SQL finalization Signed-off-by: discord9 <discord9@163.com> * fix(function): return canonical AVG1 state for empty window frames DataFusion's plain-aggregate window executor bypasses the accumulator and calls default_value() directly when a window frame contains no rows. create_udaf's SimpleAggregateUDF derives default_value from the return type, which yields SQL NULL for Binary output instead of the canonical AVG1 empty state, so empty frames and frames over only NULL inputs became observable differently (IS NULL, direct state saves, state comparisons). Register avg_state/avg_merge through a small AggregateUDFImpl (AvgUdaf) that keeps the existing accumulator and declares the canonical AVG1 empty state as default_value, restoring the documented contract. Add a unit test asserting default_value equals the empty accumulator's evaluate() and an sqlness case covering the empty-frame window. Signed-off-by: discord9 <discord9@outlook.com> --------- Signed-off-by: discord9 <discord9@163.com> Signed-off-by: discord9 <discord9@outlook.com> Co-authored-by: discord9 <discord9@outlook.com>
2026-09-24 14:15:49 +00:00 · 2026-09-18 07:02:10 +00:00
parent 8d8ebd3cc5
commit 5ef8a46e3f
11 changed files with 1050 additions and 0 deletions
@@ -18,6 +18,7 @@ use datatypes::arrow::datatypes::DataType;
 use crate::aggrs::aggr_wrapper::DeltaMergeWrapper;
 use crate::function_registry::FunctionRegistry;

+pub mod avg;
 pub mod hll;
 pub mod uddsketch;
 pub mod welford;
@@ -26,6 +27,16 @@ pub(crate) struct ApproximateFunction;

 impl ApproximateFunction {
    pub fn register(registry: &FunctionRegistry) {
+        let avg_merge = avg::AvgAccumulator::merge_udf_impl();
+        registry.register_aggr(avg::AvgAccumulator::state_udf_impl());
+        registry.register_aggr(avg_merge.clone());
+        registry.register_aggr(AggregateUDF::new_from_impl(DeltaMergeWrapper::new(
+            avg_merge.clone(),
+            avg::AVG_STATE_NAME,
+            vec![DataType::Binary],
+            DataType::Binary,
+        )));
+
        let uddsketch_state = uddsketch::UddSketchState::state_udf_impl();
        let uddsketch_merge = uddsketch::UddSketchState::merge_udf_impl();
        let uddsketch_delta = AggregateUDF::new_from_impl(DeltaMergeWrapper::new(
@@ -0,0 +1,594 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use datafusion::arrow::array::{ArrayRef, Float64Array};
+use datafusion::arrow::compute::sum;
+use datafusion::common::cast::{as_binary_array, as_primitive_array};
+use datafusion::common::not_impl_err;
+use datafusion::error::{DataFusionError, Result as DfResult};
+use datafusion::logical_expr::function::AccumulatorArgs;
+use datafusion::logical_expr::{
+    Accumulator as DfAccumulator, AggregateUDF, AggregateUDFImpl, Signature,
+};
+use datafusion_common::ScalarValue;
+use datatypes::arrow::datatypes::{DataType, Float64Type};
+
+pub const AVG_STATE_NAME: &str = "avg_state";
+pub const AVG_MERGE_NAME: &str = "avg_merge";
+
+const ENCODED_LEN: usize = 20;
+const MAGIC: &[u8; 4] = b"AVG1";
+
+/// The portable state used by the Float64 average aggregate functions.
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub struct AvgState {
+    count: u64,
+    sum: f64,
+}
+
+impl Default for AvgState {
+    fn default() -> Self {
+        Self { count: 0, sum: 0.0 }
+    }
+}
+
+impl AvgState {
+    /// Returns the exact AVG1 representation of this state.
+    pub(crate) fn encode(&self) -> [u8; ENCODED_LEN] {
+        let mut encoded = [0; ENCODED_LEN];
+        encoded[..4].copy_from_slice(MAGIC);
+        encoded[4..12].copy_from_slice(&self.count.to_le_bytes());
+        encoded[12..20].copy_from_slice(&self.sum.to_bits().to_le_bytes());
+        encoded
+    }
+
+    /// Decodes and validates an AVG1 state.
+    pub fn decode(encoded: &[u8]) -> DfResult<Self> {
+        if encoded.len() != ENCODED_LEN || &encoded[..4] != MAGIC {
+            return Err(invalid_state());
+        }
+        let count = decode_u64(encoded, 4);
+        let sum = f64::from_bits(decode_u64(encoded, 12));
+        if count == 0 && sum.to_bits() != 0 {
+            return Err(invalid_state());
+        }
+        Ok(Self { count, sum })
+    }
+
+    /// Returns the number of non-null input values in this state.
+    pub(crate) fn count(&self) -> u64 {
+        self.count
+    }
+
+    /// Returns the average, or `None` for the canonical empty state.
+    pub fn average(&self) -> Option<f64> {
+        (self.count() != 0).then(|| self.sum / self.count() as f64)
+    }
+}
+
+fn decode_u64(encoded: &[u8], offset: usize) -> u64 {
+    let mut bytes = [0; 8];
+    bytes.copy_from_slice(&encoded[offset..offset + 8]);
+    u64::from_le_bytes(bytes)
+}
+
+fn invalid_state() -> DataFusionError {
+    DataFusionError::Execution("Invalid AVG1 state".to_string())
+}
+
+fn count_overflow() -> DataFusionError {
+    DataFusionError::Execution("AVG count overflow".to_string())
+}
+
+/// The `avg_state` / `avg_merge` aggregate UDF.
+///
+/// Declares a canonical AVG1 empty state as `default_value` so window frames
+/// without rows observe the same contract as the accumulator's `evaluate`.
+#[derive(Debug, Clone, Eq, PartialEq, Hash)]
+struct AvgUdaf {
+    name: &'static str,
+    signature: Signature,
+    input: InputKind,
+}
+
+impl AggregateUDFImpl for AvgUdaf {
+    fn name(&self) -> &str {
+        self.name
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, _arg_types: &[DataType]) -> DfResult<DataType> {
+        Ok(DataType::Binary)
+    }
+
+    fn accumulator(&self, acc_args: AccumulatorArgs) -> DfResult<Box<dyn DfAccumulator>> {
+        if acc_args.is_distinct {
+            return not_impl_err!("AVG DISTINCT aggregations are not available");
+        }
+        let input = match acc_args.exprs[0].data_type(acc_args.schema)? {
+            DataType::Float64 => InputKind::Float64,
+            DataType::Binary => InputKind::Binary,
+            data_type => return not_impl_err!("AVG functions do not support {data_type:?}"),
+        };
+        Ok(Box::new(AvgAccumulator {
+            state: AvgState::default(),
+            input,
+        }))
+    }
+
+    fn default_value(&self, _data_type: &DataType) -> DfResult<ScalarValue> {
+        Ok(ScalarValue::Binary(Some(
+            AvgState::default().encode().to_vec(),
+        )))
+    }
+}
+
+#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
+enum InputKind {
+    Float64,
+    Binary,
+}
+
+/// Accumulates and merges AVG1 states.
+#[derive(Debug)]
+pub(crate) struct AvgAccumulator {
+    state: AvgState,
+    input: InputKind,
+}
+
+impl Default for AvgAccumulator {
+    fn default() -> Self {
+        Self {
+            state: AvgState::default(),
+            input: InputKind::Float64,
+        }
+    }
+}
+
+impl AvgAccumulator {
+    pub fn state_udf_impl() -> AggregateUDF {
+        AggregateUDF::new_from_impl(AvgUdaf {
+            name: AVG_STATE_NAME,
+            signature: Signature::exact(
+                vec![DataType::Float64],
+                datafusion::logical_expr::Volatility::Immutable,
+            ),
+            input: InputKind::Float64,
+        })
+    }
+
+    pub fn merge_udf_impl() -> AggregateUDF {
+        AggregateUDF::new_from_impl(AvgUdaf {
+            name: AVG_MERGE_NAME,
+            signature: Signature::exact(
+                vec![DataType::Binary],
+                datafusion::logical_expr::Volatility::Immutable,
+            ),
+            input: InputKind::Binary,
+        })
+    }
+
+    fn update_float64(&mut self, array: &ArrayRef) -> DfResult<()> {
+        let array = as_primitive_array::<Float64Type>(array)?;
+        let mut count = self.state.count;
+        for _ in array.iter().flatten() {
+            count = count.checked_add(1).ok_or_else(count_overflow)?;
+        }
+        let sum = sum(array)
+            .map(|batch_sum| self.state.sum + batch_sum)
+            .unwrap_or(self.state.sum);
+        self.state = AvgState { count, sum };
+        Ok(())
+    }
+
+    fn merge_states(&mut self, array: &ArrayRef) -> DfResult<()> {
+        let array = as_binary_array(array)?;
+        let states = array
+            .iter()
+            .flatten()
+            .map(AvgState::decode)
+            .collect::<DfResult<Vec<_>>>()?;
+        let count = states.iter().try_fold(self.state.count, |count, state| {
+            count.checked_add(state.count).ok_or_else(count_overflow)
+        })?;
+        let sums = states
+            .iter()
+            .filter(|state| state.count != 0)
+            .map(|state| Some(state.sum))
+            .collect::<Vec<_>>();
+        let sum = sum(&Float64Array::from(sums))
+            .map(|batch_sum| self.state.sum + batch_sum)
+            .unwrap_or(self.state.sum);
+        self.state = AvgState { count, sum };
+        Ok(())
+    }
+}
+
+impl DfAccumulator for AvgAccumulator {
+    fn update_batch(&mut self, values: &[ArrayRef]) -> DfResult<()> {
+        let array = &values[0];
+        match (self.input, array.data_type()) {
+            (InputKind::Float64, DataType::Float64) => self.update_float64(array),
+            (InputKind::Binary, DataType::Binary) => self.merge_states(array),
+            (_, data_type) => not_impl_err!("AVG input type does not match: {data_type:?}"),
+        }
+    }
+
+    fn evaluate(&mut self) -> DfResult<ScalarValue> {
+        Ok(ScalarValue::Binary(Some(self.state.encode().to_vec())))
+    }
+
+    fn size(&self) -> usize {
+        std::mem::size_of::<Self>()
+    }
+
+    fn state(&mut self) -> DfResult<Vec<ScalarValue>> {
+        Ok(vec![ScalarValue::Binary(Some(
+            self.state.encode().to_vec(),
+        ))])
+    }
+
+    fn merge_batch(&mut self, states: &[ArrayRef]) -> DfResult<()> {
+        self.merge_states(&states[0])
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use arrow::array::{BinaryArray, Float64Array};
+    use datafusion_common::ScalarValue;
+    use datafusion_common::arrow::datatypes::DataType;
+    use datafusion_expr::TypeSignature;
+    use datafusion_physical_expr::aggregate::AggregateExprBuilder;
+    use datafusion_physical_expr::expressions::{Column, lit as physical_lit};
+
+    use super::*;
+    use crate::aggrs::aggr_wrapper::{aggr_delta_merge_func_name, aggr_state_func_name};
+    use crate::function_registry::FUNCTION_REGISTRY;
+
+    fn state(count: u64, sum: f64) -> Vec<u8> {
+        AvgState { count, sum }.encode().to_vec()
+    }
+
+    #[test]
+    fn codec_golden_and_roundtrip() {
+        let empty = AvgState::default().encode();
+        assert_eq!(empty.len(), ENCODED_LEN);
+        assert_eq!(empty.as_slice(), b"AVG1\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0");
+        let mut accumulator = AvgAccumulator::default();
+        accumulator
+            .update_batch(&[Arc::new(Float64Array::from(vec![Some(1.5)]))])
+            .unwrap();
+        let one = accumulator.state.encode();
+        assert_eq!(one.len(), ENCODED_LEN);
+        assert_eq!(
+            one.as_slice(),
+            &[
+                b'A', b'V', b'G', b'1', 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xf8, 0x3f,
+            ]
+        );
+        assert_eq!(&one[12..20], &1.5f64.to_bits().to_le_bytes());
+        assert_eq!(AvgState::decode(&empty).unwrap().encode(), empty);
+        assert_eq!(AvgState::decode(&one).unwrap().encode(), one);
+    }
+
+    #[test]
+    fn codec_rejects_malformed_states() {
+        assert!(AvgState::decode(b"").is_err());
+        assert!(AvgState::decode(&[0; 19]).is_err());
+        assert!(AvgState::decode(&[0; 21]).is_err());
+        let mut avg2 = AvgState::default().encode();
+        avg2[..4].copy_from_slice(b"AVG2");
+        assert!(AvgState::decode(&avg2).is_err());
+        let mut wrong_magic = AvgState::default().encode();
+        wrong_magic[0] = b'X';
+        assert!(AvgState::decode(&wrong_magic).is_err());
+        for sum in [1.0, -0.0] {
+            assert!(AvgState::decode(&state(0, sum)).is_err());
+        }
+        let mut count = AvgState {
+            count: 0x0102_0304_0506_0708,
+            sum: 0.0,
+        }
+        .encode();
+        assert_eq!(&count[4..12], &0x0102_0304_0506_0708u64.to_le_bytes());
+        count[4..12].reverse();
+        assert_ne!(
+            AvgState::decode(&count).unwrap().count(),
+            0x0102_0304_0506_0708
+        );
+        let mut sum = AvgState { count: 1, sum: 1.5 }.encode();
+        assert_eq!(&sum[12..20], &1.5f64.to_bits().to_le_bytes());
+        sum[12..20].reverse();
+        assert_ne!(AvgState::decode(&sum).unwrap().average(), Some(1.5));
+    }
+
+    #[test]
+    fn codec_preserves_populated_float_bits() {
+        for bits in [
+            0.0f64.to_bits(),
+            (-0.0f64).to_bits(),
+            f64::INFINITY.to_bits(),
+            f64::NEG_INFINITY.to_bits(),
+            0x7ff8_0000_0000_0001,
+            0x7ff0_0000_0000_0001,
+        ] {
+            let encoded = state(1, f64::from_bits(bits));
+            assert_eq!(
+                AvgState::decode(&encoded).unwrap().encode().as_slice(),
+                encoded
+            );
+        }
+    }
+
+    #[test]
+    fn distinct_is_rejected() {
+        let udf = AvgAccumulator::state_udf_impl();
+        let schema = arrow_schema::Schema::empty();
+        let expr = physical_lit(1.0f64);
+        let field = Arc::new(arrow_schema::Field::new("in", DataType::Float64, true));
+        let args = AccumulatorArgs {
+            return_field: Arc::new(arrow_schema::Field::new("out", DataType::Binary, true)),
+            schema: &schema,
+            ignore_nulls: false,
+            order_bys: &[],
+            is_reversed: false,
+            name: AVG_STATE_NAME,
+            is_distinct: true,
+            exprs: std::slice::from_ref(&expr),
+            expr_fields: std::slice::from_ref(&field),
+        };
+        assert!(udf.accumulator(args).is_err());
+    }
+
+    #[test]
+    fn state_counts_nulls_and_empty_is_canonical() {
+        let mut accumulator = AvgAccumulator::default();
+        accumulator
+            .update_batch(&[Arc::new(Float64Array::from(vec![None, None]))])
+            .unwrap();
+        assert_eq!(accumulator.state.encode(), AvgState::default().encode());
+        accumulator
+            .update_batch(&[Arc::new(Float64Array::from(vec![
+                Some(1.0),
+                None,
+                Some(3.0),
+                Some(8.0),
+            ]))])
+            .unwrap();
+        assert_eq!(accumulator.state.count(), 3);
+        assert_eq!(accumulator.state.average(), Some(4.0));
+    }
+
+    #[test]
+    fn default_value_matches_empty_accumulator_evaluate() {
+        for udf in [
+            AvgAccumulator::state_udf_impl(),
+            AvgAccumulator::merge_udf_impl(),
+        ] {
+            let default = udf.default_value(&DataType::Binary).unwrap();
+            let expected = ScalarValue::Binary(Some(AvgState::default().encode().to_vec()));
+            assert_eq!(default, expected);
+            let mut empty = AvgAccumulator::default();
+            assert_eq!(
+                default,
+                empty.evaluate().unwrap(),
+                "{} default_value must equal the empty accumulator evaluate",
+                udf.name()
+            );
+        }
+    }
+
+    #[test]
+    fn merge_preserves_populated_negative_zero_for_empty_input() {
+        let mut accumulator = AvgAccumulator {
+            state: AvgState {
+                count: 1,
+                sum: -0.0,
+            },
+            input: InputKind::Binary,
+        };
+        let expected = accumulator.state.encode();
+        accumulator
+            .update_batch(&[Arc::new(BinaryArray::from(vec![
+                None,
+                Some(AvgState::default().encode().as_slice()),
+            ]))])
+            .unwrap();
+        assert_eq!(accumulator.state.encode(), expected);
+    }
+
+    #[test]
+    fn merge_ignores_nulls_and_merges_weighted_states() {
+        let mut accumulator = AvgAccumulator {
+            state: AvgState::default(),
+            input: InputKind::Binary,
+        };
+        accumulator
+            .update_batch(&[Arc::new(BinaryArray::from(vec![
+                Some(state(2, 4.0).as_slice()),
+                None,
+                Some(state(3, 15.0).as_slice()),
+            ]))])
+            .unwrap();
+        assert_eq!(accumulator.state.count(), 5);
+        assert_eq!(accumulator.state.average(), Some(19.0 / 5.0));
+        let before = accumulator.state;
+        assert!(
+            accumulator
+                .update_batch(&[Arc::new(BinaryArray::from(vec![Some(&[][..])]))])
+                .is_err()
+        );
+        assert_eq!(accumulator.state, before);
+    }
+
+    #[test]
+    fn overflow_does_not_mutate_update_or_merge() {
+        let mut update = AvgAccumulator {
+            state: AvgState {
+                count: u64::MAX,
+                sum: 1.0,
+            },
+            input: InputKind::Float64,
+        };
+        let before = update.state;
+        assert!(
+            update
+                .update_batch(&[Arc::new(Float64Array::from(vec![Some(2.0)]))])
+                .is_err()
+        );
+        assert_eq!(update.state, before);
+
+        let mut merge = AvgAccumulator {
+            state: AvgState {
+                count: u64::MAX,
+                sum: 1.0,
+            },
+            input: InputKind::Binary,
+        };
+        let before = merge.state;
+        assert!(
+            merge
+                .update_batch(&[Arc::new(BinaryArray::from(vec![Some(
+                    state(1, 2.0).as_slice()
+                )]))])
+                .is_err()
+        );
+        assert_eq!(merge.state, before);
+    }
+
+    #[test]
+    fn registered_delta_merge_has_four_way_and_malformed_behavior() {
+        let udf = FUNCTION_REGISTRY
+            .get_aggr_func(&aggr_delta_merge_func_name(AVG_STATE_NAME))
+            .unwrap();
+        assert_eq!(udf.name(), "__avg_state_delta_merge");
+        assert_eq!(
+            udf.signature().type_signature,
+            TypeSignature::Exact(vec![DataType::Binary, DataType::Binary])
+        );
+        let schema = Arc::new(arrow_schema::Schema::new(vec![
+            arrow_schema::Field::new("delta", DataType::Binary, true),
+            arrow_schema::Field::new("persisted", DataType::Binary, true),
+        ]));
+        let expr = AggregateExprBuilder::new(
+            Arc::new(udf),
+            vec![
+                Arc::new(Column::new("delta", 0)),
+                Arc::new(Column::new("persisted", 1)),
+            ],
+        )
+        .schema(schema)
+        .alias("avg_delta_merge")
+        .build()
+        .unwrap();
+        let delta = state(2, 3.0);
+        let persisted = state(2, 7.0);
+        for (left, right, expected) in [
+            (
+                Some(delta.as_slice()),
+                None,
+                AvgState { count: 2, sum: 3.0 }.encode(),
+            ),
+            (
+                None,
+                Some(persisted.as_slice()),
+                AvgState { count: 2, sum: 7.0 }.encode(),
+            ),
+            (None, None, AvgState::default().encode()),
+            (
+                Some(delta.as_slice()),
+                Some(persisted.as_slice()),
+                AvgState {
+                    count: 4,
+                    sum: 10.0,
+                }
+                .encode(),
+            ),
+        ] {
+            let mut accumulator = expr.create_accumulator().unwrap();
+            accumulator
+                .update_batch(&[
+                    Arc::new(BinaryArray::from(vec![left])),
+                    Arc::new(BinaryArray::from(vec![right])),
+                ])
+                .unwrap();
+            let ScalarValue::Binary(Some(actual)) = accumulator.evaluate().unwrap() else {
+                panic!("AVG delta merge state must be binary");
+            };
+            assert_eq!(actual.as_slice(), expected.as_slice());
+        }
+        let mut accumulator = expr.create_accumulator().unwrap();
+        assert!(
+            accumulator
+                .update_batch(&[
+                    Arc::new(BinaryArray::from(vec![Some(&[][..])])),
+                    Arc::new(BinaryArray::from(vec![None])),
+                ])
+                .is_err()
+        );
+        let mut accumulator = expr.create_accumulator().unwrap();
+        assert!(
+            accumulator
+                .update_batch(&[
+                    Arc::new(BinaryArray::from(vec![None])),
+                    Arc::new(BinaryArray::from(vec![Some(&[][..])])),
+                ])
+                .is_err()
+        );
+        let mut accumulator = expr.create_accumulator().unwrap();
+        accumulator
+            .update_batch(&[
+                Arc::new(BinaryArray::from(vec![Some(delta.as_slice())])),
+                Arc::new(BinaryArray::from(vec![None])),
+            ])
+            .unwrap();
+        assert!(
+            accumulator
+                .update_batch(&[
+                    Arc::new(BinaryArray::from(vec![Some(delta.as_slice())])),
+                    Arc::new(BinaryArray::from(vec![Some(
+                        AvgState {
+                            count: u64::MAX,
+                            sum: 1.0
+                        }
+                        .encode()
+                        .as_slice()
+                    )])),
+                ])
+                .is_err()
+        );
+    }
+
+    #[test]
+    fn avg_registry_does_not_replace_native_state_registry() {
+        let avg = FUNCTION_REGISTRY.get_aggr_func(AVG_STATE_NAME).unwrap();
+        let native = FUNCTION_REGISTRY
+            .get_aggr_func(&aggr_state_func_name("avg"))
+            .unwrap();
+        assert_eq!(
+            avg.return_type(&[DataType::Float64]).unwrap(),
+            DataType::Binary
+        );
+        assert!(matches!(
+            native.return_type(&[DataType::Float64]).unwrap(),
+            DataType::Struct(_)
+        ));
+    }
+}
@@ -29,6 +29,7 @@ use crate::aggrs::vector::VectorFunction as VectorAggrFunction;
 use crate::function::{Function, FunctionRef};
 use crate::function_factory::ScalarFunctionFactory;
 use crate::scalars::anomaly::AnomalyFunction;
+use crate::scalars::avg_calc::AvgCalcFunction;
 use crate::scalars::date::DateFunction;
 use crate::scalars::expression::ExpressionFunction;
 use crate::scalars::hll_count::HllCalcFunction;
@@ -215,6 +216,7 @@ pub static FUNCTION_REGISTRY: LazyLock<Arc<FunctionRegistry>> = LazyLock::new(||
    TimestampFunction::register(&function_registry);
    DateFunction::register(&function_registry);
    ExpressionFunction::register(&function_registry);
+    AvgCalcFunction::register(&function_registry);
    UddSketchCalcFunction::register(&function_registry);
    UddSketchRankFunction::register(&function_registry);
    HllCalcFunction::register(&function_registry);
@@ -25,6 +25,7 @@ pub mod primary_key;
 pub(crate) mod string;
 pub mod vector;

+pub(crate) mod avg_calc;
 pub(crate) mod hll_count;
 pub mod ip;
 #[cfg(test)]
@@ -0,0 +1,266 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Implementation of the scalar function `avg_calc`.
+
+use std::fmt;
+use std::fmt::Display;
+use std::sync::Arc;
+
+use datafusion_common::arrow::array::{Array, AsArray, Float64Builder};
+use datafusion_common::{DataFusionError, ScalarValue};
+use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, Volatility};
+use datatypes::arrow::datatypes::DataType;
+
+use crate::aggrs::approximate::avg::AvgState;
+use crate::function::Function;
+use crate::function_registry::FunctionRegistry;
+
+const NAME: &str = "avg_calc";
+
+/// Calculates an average from a serialized AVG1 state.
+#[derive(Debug)]
+pub(crate) struct AvgCalcFunction {
+    signature: Signature,
+}
+
+impl AvgCalcFunction {
+    pub fn register(registry: &FunctionRegistry) {
+        registry.register_scalar(Self::default());
+    }
+}
+
+impl Default for AvgCalcFunction {
+    fn default() -> Self {
+        Self {
+            signature: Signature::exact(vec![DataType::Binary], Volatility::Immutable),
+        }
+    }
+}
+
+impl Display for AvgCalcFunction {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{}", NAME.to_ascii_uppercase())
+    }
+}
+
+impl Function for AvgCalcFunction {
+    fn name(&self) -> &str {
+        NAME
+    }
+
+    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
+        Ok(DataType::Float64)
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn invoke_with_args(
+        &self,
+        args: ScalarFunctionArgs,
+    ) -> datafusion_common::Result<ColumnarValue> {
+        let [arg] = datafusion_common::utils::take_function_args(self.name(), &args.args)?;
+        match arg {
+            ColumnarValue::Scalar(ScalarValue::Binary(state)) => {
+                Ok(ColumnarValue::Scalar(ScalarValue::Float64(
+                    state
+                        .as_deref()
+                        .map(AvgState::decode)
+                        .transpose()?
+                        .and_then(|state| state.average()),
+                )))
+            }
+            ColumnarValue::Scalar(ScalarValue::Null) => {
+                Ok(ColumnarValue::Scalar(ScalarValue::Float64(None)))
+            }
+            ColumnarValue::Array(states) => {
+                let Some(states) = states.as_binary_opt::<i32>() else {
+                    return Err(invalid_type(self.name(), states.data_type()));
+                };
+                let mut builder = Float64Builder::with_capacity(states.len());
+                for state in states.iter() {
+                    builder.append_option(match state {
+                        Some(state) => AvgState::decode(state)?.average(),
+                        None => None,
+                    });
+                }
+                Ok(ColumnarValue::Array(Arc::new(builder.finish())))
+            }
+            _ => Err(invalid_type(self.name(), &arg.data_type())),
+        }
+    }
+}
+
+fn invalid_type(name: &str, data_type: &DataType) -> DataFusionError {
+    DataFusionError::Execution(format!(
+        "'{name}' expects argument to be Binary datatype, got {data_type}"
+    ))
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use arrow_schema::Field;
+    use datafusion::arrow::array::{Array, AsArray, BinaryArray, Float64Array};
+    use datafusion::logical_expr::Accumulator;
+    use datafusion::prelude::SessionContext;
+    use datafusion_common::arrow::datatypes::Float64Type;
+    use datafusion_expr::{ColumnarValue, ScalarFunctionArgs};
+
+    use super::*;
+    use crate::aggrs::approximate::avg::AvgAccumulator;
+    use crate::function::{Function, FunctionContext};
+    use crate::function_registry::FUNCTION_REGISTRY;
+
+    fn produce_state(values: Vec<Option<f64>>) -> Vec<u8> {
+        let mut accumulator = AvgAccumulator::default();
+        accumulator
+            .update_batch(&[Arc::new(Float64Array::from(values))])
+            .unwrap();
+        let ScalarValue::Binary(Some(state)) = accumulator.evaluate().unwrap() else {
+            panic!("AVG state must be binary");
+        };
+        state
+    }
+
+    fn invoke(arg: ColumnarValue, number_rows: usize) -> datafusion_common::Result<ColumnarValue> {
+        AvgCalcFunction::default().invoke_with_args(ScalarFunctionArgs {
+            args: vec![arg],
+            arg_fields: vec![],
+            number_rows,
+            return_field: Arc::new(Field::new("x", DataType::Float64, true)),
+            config_options: Arc::new(Default::default()),
+        })
+    }
+
+    #[test]
+    fn scalar_and_array_states_decode_to_averages() {
+        let state = produce_state(vec![Some(1.0), Some(2.0), Some(6.0)]);
+        let ColumnarValue::Scalar(ScalarValue::Float64(Some(value))) =
+            invoke(ColumnarValue::Scalar(ScalarValue::Binary(Some(state))), 1).unwrap()
+        else {
+            panic!("Expected Float64 scalar");
+        };
+        assert_eq!(value, 3.0);
+
+        let ColumnarValue::Scalar(ScalarValue::Float64(None)) =
+            invoke(ColumnarValue::Scalar(ScalarValue::Binary(None)), 1).unwrap()
+        else {
+            panic!("Expected NULL Float64 scalar");
+        };
+        let empty = produce_state(vec![None]);
+        let ColumnarValue::Scalar(ScalarValue::Float64(None)) = invoke(
+            ColumnarValue::Scalar(ScalarValue::Binary(Some(empty.clone()))),
+            1,
+        )
+        .unwrap() else {
+            panic!("Expected NULL Float64 scalar");
+        };
+
+        let infinity = produce_state(vec![Some(f64::INFINITY)]);
+        let nan = produce_state(vec![Some(f64::NAN)]);
+        let ColumnarValue::Array(result) = invoke(
+            ColumnarValue::Array(Arc::new(BinaryArray::from(vec![
+                Some(empty.as_slice()),
+                None,
+                Some(infinity.as_slice()),
+                Some(nan.as_slice()),
+            ]))),
+            4,
+        )
+        .unwrap() else {
+            panic!("Expected Float64 array");
+        };
+        let result = result.as_primitive::<Float64Type>();
+        assert!(result.is_null(0));
+        assert!(result.is_null(1));
+        assert_eq!(result.value(2), f64::INFINITY);
+        assert!(result.value(3).is_nan());
+    }
+
+    #[test]
+    fn malformed_and_unknown_version_states_fail_the_whole_batch() {
+        let valid = produce_state(vec![Some(3.0)]);
+        let mut unknown_version = valid.clone();
+        unknown_version[..4].copy_from_slice(b"AVG2");
+
+        assert!(
+            invoke(
+                ColumnarValue::Scalar(ScalarValue::Binary(Some(unknown_version.clone()))),
+                1,
+            )
+            .is_err()
+        );
+        assert!(
+            invoke(
+                ColumnarValue::Array(Arc::new(BinaryArray::from(vec![
+                    Some(valid.as_slice()),
+                    Some(b"malformed".as_slice()),
+                    Some(unknown_version.as_slice()),
+                ]))),
+                3,
+            )
+            .is_err()
+        );
+    }
+
+    #[tokio::test]
+    async fn registry_query_decodes_avg_state_and_weighted_avg_merge() {
+        let ctx = SessionContext::new();
+        let avg_calc = FUNCTION_REGISTRY
+            .get_function(NAME)
+            .expect("avg_calc must be registered")
+            .provide(FunctionContext::default());
+        ctx.register_udf(avg_calc);
+        for name in ["avg_state", "avg_merge"] {
+            ctx.register_udaf(
+                FUNCTION_REGISTRY
+                    .get_aggr_func(name)
+                    .expect("AVG aggregate must be registered"),
+            );
+        }
+
+        let batches = ctx
+            .sql(
+                "SELECT avg_calc(avg_state(CAST(value AS DOUBLE))) FROM \
+                 (VALUES (1.0), (2.0), (6.0)) AS values_table(value)",
+            )
+            .await
+            .unwrap()
+            .collect()
+            .await
+            .unwrap();
+        let result = batches[0].column(0).as_primitive::<Float64Type>();
+        assert_eq!(result.value(0), 3.0);
+
+        let batches = ctx
+            .sql(
+                "WITH states AS (\
+                 SELECT avg_state(CAST(value AS DOUBLE)) AS state FROM (VALUES (1.0), (3.0)) AS left_values(value) \
+                 UNION ALL \
+                 SELECT avg_state(CAST(value AS DOUBLE)) AS state FROM (VALUES (6.0)) AS right_values(value)\
+                 ) SELECT avg_calc(avg_merge(state)) FROM states",
+            )
+            .await
+            .unwrap()
+            .collect()
+            .await
+            .unwrap();
+        let result = batches[0].column(0).as_primitive::<Float64Type>();
+        assert_eq!(result.value(0), 10.0 / 3.0);
+    }
+}
@@ -89,6 +89,76 @@ SELECT AVG(i), AVG(j) FROM vals;
 -- FIXME(dennis): AVG(DISTINCT) not supported
 -- https://github.com/apache/datafusion/issues/2408
 -- SELECT AVG(DISTINCT i), AVG(DISTINCT j) FROM vals;
+-- AVG1 state scalar calculation
+SELECT avg_calc(avg_state(CAST(i AS DOUBLE))) FROM integers;
+
+---------------------------------+
+| avg_calc(avg_state(integers.i)) |
+---------------------------------+
+| 2.0                             |
+---------------------------------+
+
+-- Merged states retain their unequal group weights.
+WITH states AS (
+  SELECT avg_state(CAST(i AS DOUBLE)) AS state FROM integers WHERE i < 3
+  UNION ALL
+  SELECT avg_state(CAST(i AS DOUBLE)) AS state FROM integers WHERE i >= 3
+)
+SELECT avg_calc(avg_merge(state)) FROM states;
+
+-----------------------------------+
+| avg_calc(avg_merge(states.state)) |
+-----------------------------------+
+| 2.0                               |
+-----------------------------------+
+
+-- Empty, null-only, and null binary states calculate to NULL.
+SELECT avg_calc(avg_state(CAST(i AS DOUBLE))) FROM integers WHERE i > 100;
+
+---------------------------------+
+| avg_calc(avg_state(integers.i)) |
+---------------------------------+
+|                                 |
+---------------------------------+
+
+SELECT avg_calc(avg_state(NULL::DOUBLE));
+
+---------------------------+
+| avg_calc(avg_state(NULL)) |
+---------------------------+
+|                           |
+---------------------------+
+
+SELECT avg_calc(NULL::BYTEA);
+
+----------------+
+| avg_calc(NULL) |
+----------------+
+|                |
+----------------+
+
+-- Empty window frames return the canonical AVG1 state, not SQL NULL.
+SELECT
+    i,
+    avg_state(CAST(i AS DOUBLE)) OVER (
+        ORDER BY i
+        ROWS BETWEEN UNBOUNDED PRECEDING AND 1 PRECEDING
+    ) AS state
+FROM integers;
+
+---+------------------------------------------+
+| i | state                                    |
+---+------------------------------------------+
+| 1 | 4156473100000000000000000000000000000000 |
+| 2 | 415647310100000000000000000000000000f03f |
+| 3 | 4156473102000000000000000000000000000840 |
+---+------------------------------------------+
+
+-- Invalid AVG1 state propagates an error.
+SELECT avg_calc(X'00');
+
+Error: 3001(EngineExecuteQuery), Execution error: Invalid AVG1 state
+
 -- cleanup
 DROP TABLE integers;

@@ -40,6 +40,34 @@ SELECT AVG(i), AVG(j) FROM vals;
 -- https://github.com/apache/datafusion/issues/2408
 -- SELECT AVG(DISTINCT i), AVG(DISTINCT j) FROM vals;

+-- AVG1 state scalar calculation
+SELECT avg_calc(avg_state(CAST(i AS DOUBLE))) FROM integers;
+
+-- Merged states retain their unequal group weights.
+WITH states AS (
+  SELECT avg_state(CAST(i AS DOUBLE)) AS state FROM integers WHERE i < 3
+  UNION ALL
+  SELECT avg_state(CAST(i AS DOUBLE)) AS state FROM integers WHERE i >= 3
+)
+SELECT avg_calc(avg_merge(state)) FROM states;
+
+-- Empty, null-only, and null binary states calculate to NULL.
+SELECT avg_calc(avg_state(CAST(i AS DOUBLE))) FROM integers WHERE i > 100;
+SELECT avg_calc(avg_state(NULL::DOUBLE));
+SELECT avg_calc(NULL::BYTEA);
+
+-- Empty window frames return the canonical AVG1 state, not SQL NULL.
+SELECT
+    i,
+    avg_state(CAST(i AS DOUBLE)) OVER (
+        ORDER BY i
+        ROWS BETWEEN UNBOUNDED PRECEDING AND 1 PRECEDING
+    ) AS state
+FROM integers;
+
+-- Invalid AVG1 state propagates an error.
+SELECT avg_calc(X'00');
+
 -- cleanup
 DROP TABLE integers;

@@ -0,0 +1,10 @@
+name = "avg_state_binary"
+reason = "Verify persisted AVG1 binary states are decoded, merged, and reproduced exactly after upgrade."
+introduced_by = "PR #9035"
+topologies = ["distributed", "standalone"]
+from_range = ["*"]
+# The runner compares core versions only, so this also admits other 1.3.0 prereleases.
+to_range = [">=v1.3.0-alpha.1"]
+features = ["table", "query", "aggregate"]
+owner = "query"
+namespace = "avg_state_binary"
@@ -0,0 +1,14 @@
+CREATE TABLE avg1_states (
+  seq_id INT PRIMARY KEY,
+  state BINARY,
+  ts TIMESTAMP TIME INDEX
+);
+
+-- AVG1: magic (4 bytes), little-endian u64 count, little-endian f64 sum.
+INSERT INTO avg1_states (seq_id, state, ts) VALUES
+  (1, X'4156473102000000000000000000000000000840', '2026-01-01 00:00:00'),
+  (2, X'4156473101000000000000000000000000001840', '2026-01-01 00:00:01'),
+  (3, NULL, '2026-01-01 00:00:02'),
+  (4, X'4156473100000000000000000000000000000000', '2026-01-01 00:00:03');
+
+ADMIN FLUSH_TABLE('avg1_states');
@@ -0,0 +1,36 @@
+-- The persisted states have counts 2 and 1 and sums 3.0 and 6.0.
+-- Null and canonical empty states do not change the merged AVG1 state.
+SELECT avg_merge(state) = X'4156473103000000000000000000000000002240' AS merged_state_matches
+FROM avg1_states;
+
+----------------------+
+| merged_state_matches |
+----------------------+
+| true                 |
+----------------------+
+
+-- A new state over the equivalent Float64 values has the exact same AVG1 bytes.
+WITH generated_values AS (
+  SELECT CAST(1.0 AS DOUBLE) AS value
+  UNION ALL SELECT CAST(2.0 AS DOUBLE)
+  UNION ALL SELECT CAST(6.0 AS DOUBLE)
+)
+SELECT avg_state(value) = X'4156473103000000000000000000000000002240' AS generated_state_matches
+FROM generated_values;
+
+-------------------------+
+| generated_state_matches |
+-------------------------+
+| true                    |
+-------------------------+
+
+-- A null-only merge is the canonical empty AVG1 state.
+SELECT avg_merge(state) = X'4156473100000000000000000000000000000000' AS null_state_is_empty
+FROM avg1_states
+WHERE state IS NULL;
+
+---------------------+
+| null_state_is_empty |
+---------------------+
+| true                |
+---------------------+
@@ -0,0 +1,18 @@
+-- The persisted states have counts 2 and 1 and sums 3.0 and 6.0.
+-- Null and canonical empty states do not change the merged AVG1 state.
+SELECT avg_merge(state) = X'4156473103000000000000000000000000002240' AS merged_state_matches
+FROM avg1_states;
+
+-- A new state over the equivalent Float64 values has the exact same AVG1 bytes.
+WITH generated_values AS (
+  SELECT CAST(1.0 AS DOUBLE) AS value
+  UNION ALL SELECT CAST(2.0 AS DOUBLE)
+  UNION ALL SELECT CAST(6.0 AS DOUBLE)
+)
+SELECT avg_state(value) = X'4156473103000000000000000000000000002240' AS generated_state_matches
+FROM generated_values;
+
+-- A null-only merge is the canonical empty AVG1 state.
+SELECT avg_merge(state) = X'4156473100000000000000000000000000000000' AS null_state_is_empty
+FROM avg1_states
+WHERE state IS NULL;