From 8c2a501bf33b8b4c58f243d4b7499a194c6dd5ec Mon Sep 17 00:00:00 2001 From: discord9 Date: Tue, 8 Sep 2026 17:21:00 +0800 Subject: [PATCH] refactor(flow): separate AVG prerequisite from independent Flow changes Signed-off-by: discord9 --- src/common/function/src/aggrs/approximate.rs | 11 - .../function/src/aggrs/approximate/avg.rs | 545 ------------------ src/flow/src/batching_mode/utils.rs | 21 +- src/flow/src/batching_mode/utils/test.rs | 108 +--- .../cases/avg_state_binary/case.toml | 10 - .../cases/avg_state_binary/setup.sql | 14 - .../cases/avg_state_binary/verify.result | 36 -- .../cases/avg_state_binary/verify.sql | 18 - 8 files changed, 4 insertions(+), 759 deletions(-) delete mode 100644 src/common/function/src/aggrs/approximate/avg.rs delete mode 100644 tests/compatibility/cases/avg_state_binary/case.toml delete mode 100644 tests/compatibility/cases/avg_state_binary/setup.sql delete mode 100644 tests/compatibility/cases/avg_state_binary/verify.result delete mode 100644 tests/compatibility/cases/avg_state_binary/verify.sql diff --git a/src/common/function/src/aggrs/approximate.rs b/src/common/function/src/aggrs/approximate.rs index 562723f672..bfab225755 100644 --- a/src/common/function/src/aggrs/approximate.rs +++ b/src/common/function/src/aggrs/approximate.rs @@ -18,7 +18,6 @@ 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; @@ -27,16 +26,6 @@ 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 deleted file mode 100644 index 49b0ed3274..0000000000 --- a/src/common/function/src/aggrs/approximate/avg.rs +++ /dev/null @@ -1,545 +0,0 @@ -// 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 std::sync::Arc; - -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, Volatility}; -use datafusion::prelude::create_udaf; -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()) -} - -#[derive(Debug, Clone, Copy)] -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 { - create_udaf( - AVG_STATE_NAME, - vec![DataType::Float64], - Arc::new(DataType::Binary), - Volatility::Immutable, - Arc::new(Self::create_accumulator), - Arc::new(vec![DataType::Binary]), - ) - } - - pub fn merge_udf_impl() -> AggregateUDF { - create_udaf( - AVG_MERGE_NAME, - vec![DataType::Binary], - Arc::new(DataType::Binary), - Volatility::Immutable, - Arc::new(Self::create_accumulator), - Arc::new(vec![DataType::Binary]), - ) - } - - fn create_accumulator(args: AccumulatorArgs) -> DfResult> { - if args.is_distinct { - return not_impl_err!("AVG DISTINCT aggregations are not available"); - } - let input = match args.exprs[0].data_type(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(Self { - state: AvgState::default(), - input, - })) - } - - 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 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/flow/src/batching_mode/utils.rs b/src/flow/src/batching_mode/utils.rs index 0bb79bd196..454abf1182 100644 --- a/src/flow/src/batching_mode/utils.rs +++ b/src/flow/src/batching_mode/utils.rs @@ -21,7 +21,6 @@ use catalog::CatalogManagerRef; use common_error::ext::BoxedError; use common_function::aggrs::aggr_wrapper::get_aggr_func; use common_telemetry::debug; -use datafusion::arrow::datatypes::DataType as ArrowDataType; use datafusion::datasource::DefaultTableSource; use datafusion::error::Result as DfResult; use datafusion::logical_expr::Expr; @@ -327,10 +326,7 @@ fn is_literal_or_cast_literal(expr: &Expr) -> bool { } } -fn merge_op_for_aggregate_expr( - aggr_expr: &Expr, - input_schema: &DFSchema, -) -> Result { +fn merge_op_for_aggregate_expr(aggr_expr: &Expr) -> Result { let Some(aggr_func) = get_aggr_func(aggr_expr) else { return Err(aggr_expr.to_string()); }; @@ -350,8 +346,6 @@ fn merge_op_for_aggregate_expr( params, }) }; - let is_type = |expr: &Expr, data_type| expr.get_type(input_schema).ok() == Some(data_type); - match aggr_func.func.name().to_ascii_lowercase().as_str() { "sum" | "count" => Ok(IncrementalAggregateMergeOp::Sum), "min" => Ok(IncrementalAggregateMergeOp::Min), @@ -362,9 +356,6 @@ fn merge_op_for_aggregate_expr( "bit_or" => Ok(IncrementalAggregateMergeOp::BitOr), "bit_xor" => Ok(IncrementalAggregateMergeOp::BitXor), // Preserve state-family parameters; value coercion is handled by the aggregate. - "avg_state" if aggr_func.params.args.len() == 1 => { - state_delta_merge("__avg_state_delta_merge", vec![]) - } "hll" if aggr_func.params.args.len() == 1 => state_delta_merge("__hll_delta_merge", vec![]), "stddev_pop_state" if aggr_func.params.args.len() == 1 => { state_delta_merge("__stddev_pop_state_delta_merge", vec![]) @@ -383,14 +374,6 @@ fn merge_op_for_aggregate_expr( vec![bucket_size.clone(), error_rate.clone()], ) } - // AVG's binary merge form is admitted because its state argument is - // already the aggregate result stored by the sink. - "avg_merge" - if aggr_func.params.args.len() == 1 - && is_type(&aggr_func.params.args[0], ArrowDataType::Binary) => - { - state_delta_merge("__avg_state_delta_merge", vec![]) - } _ => Err(aggr_expr.to_string()), } } @@ -537,7 +520,7 @@ pub fn analyze_incremental_aggregate_plan( &group_key_names, )); for aggr_expr in aggr_exprs { - let merge_op = match merge_op_for_aggregate_expr(&aggr_expr, aggregate.input.schema()) { + let merge_op = match merge_op_for_aggregate_expr(&aggr_expr) { Ok(merge_op) => merge_op, Err(reason) => { unsupported_exprs.push(reason); diff --git a/src/flow/src/batching_mode/utils/test.rs b/src/flow/src/batching_mode/utils/test.rs index dc5d92f06d..5479874f26 100644 --- a/src/flow/src/batching_mode/utils/test.rs +++ b/src/flow/src/batching_mode/utils/test.rs @@ -1766,114 +1766,10 @@ async fn test_analyze_incremental_aggregate_plan_rejects_avg() { assert!(!analysis.unsupported_exprs.is_empty()); } -#[tokio::test] -async fn test_analyze_incremental_aggregate_plan_supports_avg_state() { - let query_engine = create_test_query_engine(); - let ctx = QueryContext::arc(); - let sql = "SELECT avg_state(number) AS avg_num, ts FROM numbers_with_ts GROUP BY ts"; - let plan = sql_to_df_plan(ctx, query_engine, sql, false).await.unwrap(); - - let analysis = analyze_incremental_aggregate_plan(&plan).unwrap().unwrap(); - assert!( - analysis.unsupported_exprs.is_empty(), - "avg_state should be supported: {:?}", - analysis.unsupported_exprs - ); - assert_eq!(analysis.merge_columns.len(), 1); - assert_eq!(analysis.merge_columns[0].output_field_name, "avg_num"); - assert!(matches!( - &analysis.merge_columns[0].merge_op, - IncrementalAggregateMergeOp::StateDeltaMerge { - function_name: "__avg_state_delta_merge", - params, - } if params.is_empty() - )); -} - -#[tokio::test] -async fn test_analyze_incremental_aggregate_plan_supports_avg_merge() { - let query_engine = create_test_query_engine(); - let ctx = QueryContext::arc(); - let sql = "SELECT avg_merge(avg_state(number)) AS avg_num, ts FROM numbers_with_ts GROUP BY ts"; - let plan = sql_to_df_plan(ctx, query_engine, sql, false).await.unwrap(); - - let analysis = analyze_incremental_aggregate_plan(&plan).unwrap().unwrap(); - assert!( - analysis.unsupported_exprs.is_empty(), - "avg_merge should be supported: {:?}", - analysis.unsupported_exprs - ); - assert_eq!(analysis.merge_columns.len(), 1); - assert!(matches!( - &analysis.merge_columns[0].merge_op, - IncrementalAggregateMergeOp::StateDeltaMerge { - function_name: "__avg_state_delta_merge", - params, - } if params.is_empty() - )); -} - -#[tokio::test] -async fn test_analyze_incremental_aggregate_plan_supports_duplicate_avg_projections() { - let analysis = analyze_test_sql( - "SELECT avg_state(number) AS avg_num, avg_state(number + 1) AS avg_num_plus, ts FROM numbers_with_ts GROUP BY ts", - ) - .await; - - assert!(analysis.unsupported_exprs.is_empty()); - assert_eq!(analysis.merge_columns.len(), 2); - assert!(analysis.merge_columns.iter().all(|column| { - matches!( - &column.merge_op, - IncrementalAggregateMergeOp::StateDeltaMerge { - function_name: "__avg_state_delta_merge", - params, - } if params.is_empty() - ) - })); - assert!( - analysis - .merge_columns - .iter() - .any(|column| column.output_field_name == "avg_num") - ); - assert!( - analysis - .merge_columns - .iter() - .any(|column| column.output_field_name == "avg_num_plus") - ); -} - -#[tokio::test] -async fn test_analyze_incremental_aggregate_plan_supports_avg_with_native_aggregate() { - let analysis = analyze_test_sql( - "SELECT avg_state(number) AS avg_num, sum(number) AS total, ts FROM numbers_with_ts GROUP BY ts", - ) - .await; - - assert!(analysis.unsupported_exprs.is_empty()); - assert_eq!(analysis.merge_columns.len(), 2); - assert!(analysis.merge_columns.iter().any(|column| { - column.output_field_name == "avg_num" - && matches!( - &column.merge_op, - IncrementalAggregateMergeOp::StateDeltaMerge { - function_name: "__avg_state_delta_merge", - params, - } if params.is_empty() - ) - })); - assert!(analysis.merge_columns.iter().any(|column| { - column.output_field_name == "total" && column.merge_op == IncrementalAggregateMergeOp::Sum - })); -} - #[tokio::test] async fn test_analyze_incremental_aggregate_plan_supports_mixed_state_families() { let analysis = analyze_test_sql( - "SELECT avg_state(number) AS avg_num, \ - hll(CAST(number AS VARCHAR)) AS hll_a, \ + "SELECT hll(CAST(number AS VARCHAR)) AS hll_a, \ hll(CAST(number AS VARCHAR)) AS hll_b, \ uddsketch_state(128, 0.01, CAST(number AS DOUBLE)) AS percentile_a, \ uddsketch_state(256, 0.02, number) AS percentile_b, \ @@ -1887,7 +1783,7 @@ async fn test_analyze_incremental_aggregate_plan_supports_mixed_state_families() "mixed state aggregate should be supported: {:?}", analysis.unsupported_exprs ); - assert_eq!(analysis.merge_columns.len(), 7); + assert_eq!(analysis.merge_columns.len(), 6); assert!(analysis.merge_columns.iter().any(|column| { column.output_field_name == "hll_a" && column.merge_op diff --git a/tests/compatibility/cases/avg_state_binary/case.toml b/tests/compatibility/cases/avg_state_binary/case.toml deleted file mode 100644 index 0049fca8c3..0000000000 --- a/tests/compatibility/cases/avg_state_binary/case.toml +++ /dev/null @@ -1,10 +0,0 @@ -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 deleted file mode 100644 index 2bb2db577e..0000000000 --- a/tests/compatibility/cases/avg_state_binary/setup.sql +++ /dev/null @@ -1,14 +0,0 @@ -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 deleted file mode 100644 index c3eba781a4..0000000000 --- a/tests/compatibility/cases/avg_state_binary/verify.result +++ /dev/null @@ -1,36 +0,0 @@ --- 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 deleted file mode 100644 index 30cade275b..0000000000 --- a/tests/compatibility/cases/avg_state_binary/verify.sql +++ /dev/null @@ -1,18 +0,0 @@ --- 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;