diff --git a/src/common/function/src/aggrs/approximate.rs b/src/common/function/src/aggrs/approximate.rs index bfab225755a..562723f6727 100644 --- a/src/common/function/src/aggrs/approximate.rs +++ b/src/common/function/src/aggrs/approximate.rs @@ -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( diff --git a/src/common/function/src/aggrs/approximate/avg.rs b/src/common/function/src/aggrs/approximate/avg.rs new file mode 100644 index 00000000000..9ca952edb7c --- /dev/null +++ b/src/common/function/src/aggrs/approximate/avg.rs @@ -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 { + 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 { + (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 { + Ok(DataType::Binary) + } + + fn accumulator(&self, acc_args: AccumulatorArgs) -> DfResult> { + 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 { + 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::(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::>>()?; + 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::>(); + 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 { + Ok(ScalarValue::Binary(Some(self.state.encode().to_vec()))) + } + + fn size(&self) -> usize { + std::mem::size_of::() + } + + fn state(&mut self) -> DfResult> { + 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 { + 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(_) + )); + } +} diff --git a/src/common/function/src/function_registry.rs b/src/common/function/src/function_registry.rs index 80c34eeb91f..c8e9874abc9 100644 --- a/src/common/function/src/function_registry.rs +++ b/src/common/function/src/function_registry.rs @@ -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> = 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); diff --git a/src/common/function/src/scalars.rs b/src/common/function/src/scalars.rs index 1a8c22e2d6f..306d1cd17e9 100644 --- a/src/common/function/src/scalars.rs +++ b/src/common/function/src/scalars.rs @@ -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)] diff --git a/src/common/function/src/scalars/avg_calc.rs b/src/common/function/src/scalars/avg_calc.rs new file mode 100644 index 00000000000..cae76cb2b1f --- /dev/null +++ b/src/common/function/src/scalars/avg_calc.rs @@ -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 { + Ok(DataType::Float64) + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn invoke_with_args( + &self, + args: ScalarFunctionArgs, + ) -> datafusion_common::Result { + 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::() 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>) -> Vec { + 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 { + 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::(); + 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::(); + 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::(); + assert_eq!(result.value(0), 10.0 / 3.0); + } +} diff --git a/tests/cases/standalone/common/aggregate/avg.result b/tests/cases/standalone/common/aggregate/avg.result index 99fd8a09e4c..686c7df9e5b 100644 --- a/tests/cases/standalone/common/aggregate/avg.result +++ b/tests/cases/standalone/common/aggregate/avg.result @@ -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; diff --git a/tests/cases/standalone/common/aggregate/avg.sql b/tests/cases/standalone/common/aggregate/avg.sql index cbb12edcbee..0667bbcfde8 100644 --- a/tests/cases/standalone/common/aggregate/avg.sql +++ b/tests/cases/standalone/common/aggregate/avg.sql @@ -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; diff --git a/tests/compatibility/cases/avg_state_binary/case.toml b/tests/compatibility/cases/avg_state_binary/case.toml new file mode 100644 index 00000000000..0049fca8c39 --- /dev/null +++ b/tests/compatibility/cases/avg_state_binary/case.toml @@ -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" diff --git a/tests/compatibility/cases/avg_state_binary/setup.sql b/tests/compatibility/cases/avg_state_binary/setup.sql new file mode 100644 index 00000000000..2bb2db577ee --- /dev/null +++ b/tests/compatibility/cases/avg_state_binary/setup.sql @@ -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'); diff --git a/tests/compatibility/cases/avg_state_binary/verify.result b/tests/compatibility/cases/avg_state_binary/verify.result new file mode 100644 index 00000000000..c3eba781a49 --- /dev/null +++ b/tests/compatibility/cases/avg_state_binary/verify.result @@ -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 | ++---------------------+ diff --git a/tests/compatibility/cases/avg_state_binary/verify.sql b/tests/compatibility/cases/avg_state_binary/verify.sql new file mode 100644 index 00000000000..30cade275b6 --- /dev/null +++ b/tests/compatibility/cases/avg_state_binary/verify.sql @@ -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;