From c6b10bfbb95bd04983e0934d16743f46e4b31c68 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Mon, 31 Aug 2026 14:15:33 +0000 Subject: [PATCH] feat(function): add mergeable stddev_pop state functions (#8972) * feat(function): add Welford stddev functions Signed-off-by: Lei, HUANG * test(function): cover merged Welford time windows Signed-off-by: Lei, HUANG * refactor(function): use stddev_pop SQL names Signed-off-by: Lei, HUANG * fix(function): make Welford arithmetic partition-stable Signed-off-by: Lei, HUANG * fix(function): reject invalid singleton Welford states Signed-off-by: Lei, HUANG * test(function): pin Welford state compatibility Signed-off-by: Lei, HUANG * refactor(function): remove unreachable variance clamp Signed-off-by: Lei, HUANG * fix(function): reject DISTINCT Welford aggregates Signed-off-by: Lei, HUANG * test(compat): bound Welford downgrade targets Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- src/common/function/src/aggrs/approximate.rs | 5 + .../function/src/aggrs/approximate/welford.rs | 528 ++++++++++++++++++ src/common/function/src/function_registry.rs | 2 + src/common/function/src/scalars.rs | 1 + .../function/src/scalars/welford_stddev.rs | 199 +++++++ .../common/aggregate/welford.result | 192 +++++++ .../standalone/common/aggregate/welford.sql | 134 +++++ .../cases/stddev_pop_state/case.toml | 9 + .../cases/stddev_pop_state/setup.sql | 26 + .../cases/stddev_pop_state/verify.result | 21 + .../cases/stddev_pop_state/verify.sql | 8 + 11 files changed, 1125 insertions(+) create mode 100644 src/common/function/src/aggrs/approximate/welford.rs create mode 100644 src/common/function/src/scalars/welford_stddev.rs create mode 100644 tests/cases/standalone/common/aggregate/welford.result create mode 100644 tests/cases/standalone/common/aggregate/welford.sql create mode 100644 tests/compatibility/cases/stddev_pop_state/case.toml create mode 100644 tests/compatibility/cases/stddev_pop_state/setup.sql create mode 100644 tests/compatibility/cases/stddev_pop_state/verify.result create mode 100644 tests/compatibility/cases/stddev_pop_state/verify.sql diff --git a/src/common/function/src/aggrs/approximate.rs b/src/common/function/src/aggrs/approximate.rs index 11b906b971..e4fbc3bab1 100644 --- a/src/common/function/src/aggrs/approximate.rs +++ b/src/common/function/src/aggrs/approximate.rs @@ -16,6 +16,7 @@ use crate::function_registry::FunctionRegistry; pub mod hll; pub mod uddsketch; +pub mod welford; pub(crate) struct ApproximateFunction; @@ -28,5 +29,9 @@ impl ApproximateFunction { // hll registry.register_aggr(hll::HllState::state_udf_impl()); registry.register_aggr(hll::HllState::merge_udf_impl()); + + // welford + registry.register_aggr(welford::WelfordAccumulator::state_udf_impl()); + registry.register_aggr(welford::WelfordAccumulator::merge_udf_impl()); } } diff --git a/src/common/function/src/aggrs/approximate/welford.rs b/src/common/function/src/aggrs/approximate/welford.rs new file mode 100644 index 0000000000..5d8bb11374 --- /dev/null +++ b/src/common/function/src/aggrs/approximate/welford.rs @@ -0,0 +1,528 @@ +// 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. + +//! Mergeable Welford state for population standard deviation. +//! +//! Input samples and intermediate states must contain only finite values. + +use std::sync::Arc; + +use datafusion::arrow::array::ArrayRef; +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 STDDEV_POP_STATE_NAME: &str = "stddev_pop_state"; +pub const STDDEV_POP_MERGE_NAME: &str = "stddev_pop_merge"; + +const ENCODED_LEN: usize = 28; +const MAGIC: &[u8; 4] = b"WLF1"; + +#[derive(Debug, Clone, Copy, PartialEq)] +pub(crate) struct WelfordState { + pub(crate) count: u64, + pub(crate) mean: f64, + pub(crate) m2: f64, +} + +impl Default for WelfordState { + fn default() -> Self { + Self { + count: 0, + mean: 0.0, + m2: 0.0, + } + } +} + +impl WelfordState { + 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.mean.to_bits().to_le_bytes()); + encoded[20..28].copy_from_slice(&self.m2.to_bits().to_le_bytes()); + encoded + } + + pub(crate) fn decode(encoded: &[u8]) -> DfResult { + if encoded.len() != ENCODED_LEN || &encoded[..4] != MAGIC { + return Err(invalid_state()); + } + + let state = Self { + count: decode_u64(encoded, 4), + mean: decode_f64(encoded, 12), + m2: decode_f64(encoded, 20), + }; + if !state.is_valid() { + return Err(invalid_state()); + } + + Ok(state) + } + + fn is_valid(&self) -> bool { + match self.count { + 0 => self.mean.to_bits() == 0 && self.m2.to_bits() == 0, + 1 => self.mean.is_finite() && self.m2.to_bits() == 0, + _ => self.mean.is_finite() && self.m2.is_finite() && self.m2 >= 0.0, + } + } + + fn update(&mut self, sample: f64) -> DfResult<()> { + if !sample.is_finite() { + return Err(non_finite_input()); + } + + let count = self + .count + .checked_add(1) + .ok_or_else(|| DataFusionError::Execution("Welford count overflow".to_string()))?; + let candidate_state = if self.count == 0 { + Self { + count, + mean: sample, + m2: 0.0, + } + } else { + let delta = sample - self.mean; + if !delta.is_finite() { + return Err(non_finite_arithmetic()); + } + let mean = self.mean + delta / count as f64; + let delta2 = sample - mean; + Self { + count, + mean, + m2: self.m2 + delta * delta2, + } + }; + self.replace_with_candidate(candidate_state) + } + + fn merge(&mut self, other: &Self) -> DfResult<()> { + if other.count == 0 { + return Ok(()); + } + if self.count == 0 { + return self.replace_with_candidate(*other); + } + + let count = self + .count + .checked_add(other.count) + .ok_or_else(|| DataFusionError::Execution("Welford count overflow".to_string()))?; + let delta = other.mean - self.mean; + if !delta.is_finite() { + return Err(non_finite_arithmetic()); + } + let self_count = self.count as f64; + let other_count = other.count as f64; + let count_f64 = count as f64; + let mean_delta = if delta.abs() <= f64::MAX / other_count { + delta * other_count / count_f64 + } else { + delta * (other_count / count_f64) + }; + let weighted_count = self_count * other_count / count_f64; + let candidate_state = Self { + count, + mean: self.mean + mean_delta, + m2: self.m2 + other.m2 + checked_weighted_square(delta, weighted_count)?, + }; + self.replace_with_candidate(candidate_state) + } + + fn replace_with_candidate(&mut self, candidate_state: Self) -> DfResult<()> { + if !candidate_state.is_valid() { + return Err(non_finite_arithmetic()); + } + + *self = candidate_state; + Ok(()) + } + + pub(crate) fn population_stddev(&self) -> Option { + if self.count == 0 { + return None; + } + + Some((self.m2 / self.count as f64).sqrt()) + } +} + +fn checked_weighted_square(delta: f64, weight: f64) -> DfResult { + if delta.abs() <= f64::MAX.sqrt() { + return Ok(delta * delta * weight); + } + if weight <= 1.0 { + // Applying the weight first avoids overflow when the weighted square is representable. + return Ok(delta * weight * delta); + } + + Err(non_finite_arithmetic()) +} + +/// Accumulates and merges versioned Welford states. +#[derive(Debug, Default)] +pub struct WelfordAccumulator { + state: WelfordState, +} + +impl WelfordAccumulator { + /// Creates the `stddev_pop_state` aggregate function. + pub fn state_udf_impl() -> AggregateUDF { + create_udaf( + STDDEV_POP_STATE_NAME, + vec![DataType::Float64], + Arc::new(DataType::Binary), + Volatility::Immutable, + Arc::new(Self::create_accumulator), + Arc::new(vec![DataType::Binary]), + ) + } + + /// Creates the `stddev_pop_merge` aggregate function. + pub fn merge_udf_impl() -> AggregateUDF { + create_udaf( + STDDEV_POP_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!("Welford DISTINCT aggregations are not available"); + } + Ok(Box::new(Self::default())) + } +} + +impl DfAccumulator for WelfordAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> DfResult<()> { + let array = &values[0]; + match array.data_type() { + DataType::Float64 => { + for sample in as_primitive_array::(array)?.iter().flatten() { + self.state.update(sample)?; + } + } + DataType::Binary => self.merge_batch(std::slice::from_ref(array))?, + other => { + return not_impl_err!("Welford functions do not support data type: {other}"); + } + } + Ok(()) + } + + 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<()> { + let array = as_binary_array(&states[0])?; + for encoded in array.iter().flatten() { + self.state.merge(&WelfordState::decode(encoded)?)?; + } + Ok(()) + } +} + +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 decode_f64(encoded: &[u8], offset: usize) -> f64 { + f64::from_bits(decode_u64(encoded, offset)) +} + +fn invalid_state() -> DataFusionError { + DataFusionError::Execution("Invalid Welford state".to_string()) +} + +fn non_finite_input() -> DataFusionError { + DataFusionError::Execution("Welford state requires finite input values".to_string()) +} + +fn non_finite_arithmetic() -> DataFusionError { + DataFusionError::Execution("Welford arithmetic produced a non-finite state".to_string()) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datafusion::arrow::array::{ArrayRef, BinaryArray, Float64Array}; + use datafusion_common::ScalarValue; + + use super::*; + + fn state_from_values(values: &[f64]) -> WelfordState { + let mut state = WelfordState::default(); + for value in values { + state.update(*value).unwrap(); + } + state + } + + #[test] + fn test_welford_state_encoding_contract() { + let state = WelfordState { + count: 3, + mean: 2.0, + m2: 6.0, + }; + + let encoded = state.encode(); + + assert_eq!( + encoded, + [ + b'W', b'L', b'F', b'1', // magic + 3, 0, 0, 0, 0, 0, 0, 0, // count + 0, 0, 0, 0, 0, 0, 0, 64, // mean + 0, 0, 0, 0, 0, 0, 24, 64, // m2 + ] + ); + assert_eq!(WelfordState::decode(&encoded).unwrap(), state); + } + + #[test] + fn test_welford_state_online_update() { + let mut state = WelfordState::default(); + for value in [1.0, 2.0, 3.0, 4.0] { + state.update(value).unwrap(); + } + + assert_eq!(state.count, 4); + assert_eq!(state.mean, 2.5); + assert_eq!(state.m2, 5.0); + assert_eq!(state.population_stddev(), Some(1.25_f64.sqrt())); + } + + #[test] + fn test_welford_state_empty_and_single_value() { + let mut state = WelfordState::default(); + assert_eq!(state.population_stddev(), None); + + state.update(42.0).unwrap(); + assert_eq!(state.population_stddev(), Some(0.0)); + } + + #[test] + fn test_welford_non_finite_values_fail_independent_of_partitioning() { + for sample in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] { + let mut one_pass = WelfordState::default(); + let update_failed = one_pass.update(sample).is_err(); + + let mut merged = WelfordState::default(); + let non_finite_state = WelfordState { + count: 1, + mean: sample, + m2: 0.0, + }; + let merge_failed = merged.merge(&non_finite_state).is_err(); + + assert_eq!((update_failed, merge_failed), (true, true)); + assert_eq!(one_pass, WelfordState::default()); + assert_eq!(merged, WelfordState::default()); + } + } + + #[test] + fn test_welford_state_rejects_malformed_encoding() { + assert!(WelfordState::decode(b"").is_err()); + assert!(WelfordState::decode(&[0; 28]).is_err()); + + let mut encoded = WelfordState::default().encode().to_vec(); + encoded.push(0); + assert!(WelfordState::decode(&encoded).is_err()); + + let noncanonical_empty = WelfordState { + count: 0, + mean: 1.0, + m2: 0.0, + }; + assert!(WelfordState::decode(&noncanonical_empty.encode()).is_err()); + + let negative_m2 = WelfordState { + count: 2, + mean: 1.0, + m2: -1.0, + }; + assert!(WelfordState::decode(&negative_m2.encode()).is_err()); + + for m2 in [1.0, -0.0] { + let noncanonical_singleton = WelfordState { + count: 1, + mean: 0.0, + m2, + }; + assert!(WelfordState::decode(&noncanonical_singleton.encode()).is_err()); + } + + for (mean, m2) in [ + (f64::NAN, 0.0), + (f64::INFINITY, 0.0), + (f64::NEG_INFINITY, 0.0), + (0.0, f64::NAN), + (0.0, f64::INFINITY), + (0.0, f64::NEG_INFINITY), + ] { + let non_finite = WelfordState { count: 1, mean, m2 }; + assert!(WelfordState::decode(&non_finite.encode()).is_err()); + } + } + + #[test] + fn test_welford_state_merge_matches_one_pass_update() { + let mut merged = state_from_values(&[1.0, 2.0]); + merged.merge(&state_from_values(&[3.0, 4.0])).unwrap(); + + assert_eq!(merged, state_from_values(&[1.0, 2.0, 3.0, 4.0])); + } + + #[test] + fn test_welford_large_finite_variance_matches_partitioned_merge() { + let large_sample = f64::MAX.sqrt() * 1.1; + let one_pass = state_from_values(&[0.0, large_sample]); + let mut merged = state_from_values(&[0.0]); + + merged.merge(&state_from_values(&[large_sample])).unwrap(); + + assert_eq!(merged, one_pass); + } + + #[test] + fn test_welford_extreme_values_fail_independent_of_partitioning() { + for values in [[f64::MAX, -f64::MAX], [-f64::MAX, f64::MAX]] { + let mut one_pass = state_from_values(&values[..1]); + let original_one_pass = one_pass; + let update_failed = one_pass.update(values[1]).is_err(); + + let mut merged = state_from_values(&values[..1]); + let original_merged = merged; + let merge_failed = merged.merge(&state_from_values(&values[1..])).is_err(); + + assert_eq!((update_failed, merge_failed), (true, true)); + assert_eq!(one_pass, original_one_pass); + assert_eq!(merged, original_merged); + } + } + + #[test] + fn test_welford_state_empty_merge_identity() { + let populated = state_from_values(&[1.0, 2.0]); + let mut left = WelfordState::default(); + left.merge(&populated).unwrap(); + assert_eq!(left, populated); + + let mut right = populated; + right.merge(&WelfordState::default()).unwrap(); + assert_eq!(right, populated); + } + + #[test] + fn test_welford_state_merge_rejects_count_overflow() { + let mut state = WelfordState { + count: u64::MAX, + mean: 1.0, + m2: 0.0, + }; + let other = WelfordState { + count: 1, + mean: 1.0, + m2: 0.0, + }; + + assert!(state.merge(&other).is_err()); + } + + #[test] + fn test_welford_accumulator_ignores_nulls() { + let mut accumulator = WelfordAccumulator::default(); + let array = Arc::new(Float64Array::from(vec![Some(1.0), None, Some(3.0)])) as ArrayRef; + + accumulator.update_batch(&[array]).unwrap(); + + let ScalarValue::Binary(Some(encoded)) = accumulator.evaluate().unwrap() else { + panic!("Expected binary scalar value"); + }; + assert_eq!( + WelfordState::decode(&encoded).unwrap(), + state_from_values(&[1.0, 3.0]) + ); + } + + #[test] + fn test_welford_accumulator_merges_binary_states() { + let first = state_from_values(&[1.0, 2.0]).encode(); + let second = state_from_values(&[3.0, 4.0]).encode(); + let states = Arc::new(BinaryArray::from(vec![ + Some(first.as_slice()), + None, + Some(second.as_slice()), + ])) as ArrayRef; + let mut accumulator = WelfordAccumulator::default(); + + accumulator.merge_batch(&[states]).unwrap(); + + let ScalarValue::Binary(Some(encoded)) = accumulator.state().unwrap().remove(0) else { + panic!("Expected binary scalar value"); + }; + assert_eq!( + WelfordState::decode(&encoded).unwrap(), + state_from_values(&[1.0, 2.0, 3.0, 4.0]) + ); + } + + #[test] + fn test_welford_accumulator_rejects_malformed_state() { + let noncanonical_singleton = WelfordState { + count: 1, + mean: 0.0, + m2: 1.0, + } + .encode(); + + for encoded in [b"invalid".to_vec(), noncanonical_singleton.to_vec()] { + let states = Arc::new(BinaryArray::from(vec![Some(encoded.as_slice())])) as ArrayRef; + let mut accumulator = WelfordAccumulator::default(); + + assert!(accumulator.merge_batch(&[states]).is_err()); + } + } +} diff --git a/src/common/function/src/function_registry.rs b/src/common/function/src/function_registry.rs index 15721ccabc..6ff585b075 100644 --- a/src/common/function/src/function_registry.rs +++ b/src/common/function/src/function_registry.rs @@ -43,6 +43,7 @@ use crate::scalars::timestamp::TimestampFunction; use crate::scalars::uddsketch_calc::UddSketchCalcFunction; use crate::scalars::uddsketch_rank::UddSketchRankFunction; use crate::scalars::vector::VectorFunction as VectorScalarFunction; +use crate::scalars::welford_stddev::WelfordStddevFunction; use crate::system::SystemFunction; #[derive(Default)] @@ -212,6 +213,7 @@ pub static FUNCTION_REGISTRY: LazyLock> = LazyLock::new(|| UddSketchCalcFunction::register(&function_registry); UddSketchRankFunction::register(&function_registry); HllCalcFunction::register(&function_registry); + WelfordStddevFunction::register(&function_registry); DecodePrimaryKeyFunction::register(&function_registry); // Full text search function diff --git a/src/common/function/src/scalars.rs b/src/common/function/src/scalars.rs index 75dec84ea6..1a8c22e2d6 100644 --- a/src/common/function/src/scalars.rs +++ b/src/common/function/src/scalars.rs @@ -33,3 +33,4 @@ pub(crate) mod timestamp; pub(crate) mod uddsketch_calc; pub(crate) mod uddsketch_rank; pub mod udf; +pub(crate) mod welford_stddev; diff --git a/src/common/function/src/scalars/welford_stddev.rs b/src/common/function/src/scalars/welford_stddev.rs new file mode 100644 index 0000000000..7463feac2b --- /dev/null +++ b/src/common/function/src/scalars/welford_stddev.rs @@ -0,0 +1,199 @@ +// 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 `stddev_pop_calc`. + +use std::fmt; +use std::fmt::Display; +use std::sync::Arc; + +use datafusion_common::DataFusionError; +use datafusion_common::arrow::array::{Array, AsArray, Float64Builder}; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, Volatility}; +use datatypes::arrow::datatypes::DataType; + +use crate::aggrs::approximate::welford::WelfordState; +use crate::function::{Function, extract_args}; +use crate::function_registry::FunctionRegistry; + +const NAME: &str = "stddev_pop_calc"; + +/// Calculates population standard deviation from a serialized Welford state. +#[derive(Debug)] +pub(crate) struct WelfordStddevFunction { + signature: Signature, +} + +impl WelfordStddevFunction { + pub fn register(registry: &FunctionRegistry) { + registry.register_scalar(Self::default()); + } +} + +impl Default for WelfordStddevFunction { + fn default() -> Self { + Self { + signature: Signature::exact(vec![DataType::Binary], Volatility::Immutable), + } + } +} + +impl Display for WelfordStddevFunction { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "{}", NAME.to_ascii_uppercase()) + } +} + +impl Function for WelfordStddevFunction { + 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] = extract_args(self.name(), &args)?; + let Some(states) = arg.as_binary_opt::() else { + return Err(DataFusionError::Execution(format!( + "'{}' expects argument to be Binary datatype, got {}", + self.name(), + arg.data_type() + ))); + }; + let mut builder = Float64Builder::with_capacity(states.len()); + for state in states.iter() { + match state.and_then(decode_population_stddev) { + Some(stddev) => builder.append_value(stddev), + None => builder.append_null(), + } + } + + Ok(ColumnarValue::Array(Arc::new(builder.finish()))) + } +} + +fn decode_population_stddev(encoded: &[u8]) -> Option { + match WelfordState::decode(encoded) { + Ok(state) => state.population_stddev(), + Err(error) => { + common_telemetry::trace!("Failed to decode Welford state: {}", error); + None + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_schema::Field; + use datafusion_common::arrow::array::{Array, AsArray, BinaryArray}; + use datafusion_common::arrow::datatypes::Float64Type; + use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; + use datatypes::arrow::datatypes::DataType; + + use super::*; + use crate::aggrs::approximate::welford::WelfordState; + use crate::function::Function; + + fn invoke(states: BinaryArray) -> ColumnarValue { + WelfordStddevFunction::default() + .invoke_with_args(ScalarFunctionArgs { + number_rows: states.len(), + args: vec![ColumnarValue::Array(Arc::new(states))], + arg_fields: vec![], + return_field: Arc::new(Field::new("x", DataType::Float64, true)), + config_options: Arc::new(Default::default()), + }) + .unwrap() + } + + #[test] + fn test_populated_welford_state_returns_population_stddev() { + let populated = WelfordState { + count: 4, + mean: 2.5, + m2: 5.0, + } + .encode(); + + let ColumnarValue::Array(output) = + invoke(BinaryArray::from(vec![Some(populated.as_slice())])) + else { + panic!("Expected array result"); + }; + let output = output.as_primitive::(); + assert!((output.value(0) - 1.25_f64.sqrt()).abs() < 1e-12); + } + + #[test] + fn test_empty_malformed_and_null_states_return_null() { + let empty = WelfordState::default().encode(); + let noncanonical_singleton = WelfordState { + count: 1, + mean: 0.0, + m2: 1.0, + } + .encode(); + + let ColumnarValue::Array(output) = invoke(BinaryArray::from(vec![ + Some(empty.as_slice()), + Some(b"invalid".as_slice()), + Some(noncanonical_singleton.as_slice()), + None, + ])) else { + panic!("Expected array result"); + }; + assert_eq!(output.null_count(), 4); + } + + #[test] + fn test_stddev_pop_calc_metadata() { + let function = WelfordStddevFunction::default(); + + assert_eq!(function.name(), "stddev_pop_calc"); + assert_eq!( + function.return_type(&[DataType::Binary]).unwrap(), + DataType::Float64 + ); + } + + #[test] + fn test_stddev_pop_calc_rejects_wrong_argument_count() { + let error = WelfordStddevFunction::default() + .invoke_with_args(ScalarFunctionArgs { + args: vec![], + arg_fields: vec![], + number_rows: 0, + return_field: Arc::new(Field::new("x", DataType::Float64, true)), + config_options: Arc::new(Default::default()), + }) + .unwrap_err(); + + assert!( + error + .to_string() + .contains("stddev_pop_calc function requires 1 argument, got 0") + ); + } +} diff --git a/tests/cases/standalone/common/aggregate/welford.result b/tests/cases/standalone/common/aggregate/welford.result new file mode 100644 index 0000000000..4840d7288c --- /dev/null +++ b/tests/cases/standalone/common/aggregate/welford.result @@ -0,0 +1,192 @@ +CREATE TABLE test_welford ( + `id` INT PRIMARY KEY, + `value` DOUBLE, + `ts` TIMESTAMP TIME INDEX DEFAULT now() +); + +Affected Rows: 0 + +INSERT INTO test_welford (`id`, `value`) VALUES + (1, 10.0), + (2, 20.0), + (3, 30.0), + (4, 40.0), + (5, 50.0), + (6, 60.0), + (7, 70.0), + (8, 80.0), + (9, 90.0), + (10, 100.0), + (11, NULL); + +Affected Rows: 11 + +SELECT stddev_pop_calc(stddev_pop_state(`value`)) FROM test_welford; + ++-------------------------------------------------------+ +| stddev_pop_calc(stddev_pop_state(test_welford.value)) | ++-------------------------------------------------------+ +| 28.722813232690143 | ++-------------------------------------------------------+ + +-- A second DISTINCT argument set prevents DataFusion from rewriting this to GROUP BY. +SELECT + stddev_pop_calc(stddev_pop_state(DISTINCT `value`)), + count(DISTINCT `id`) +FROM ( + SELECT `id`, `value` FROM test_welford + UNION ALL + SELECT 12 AS `id`, 100.0 AS `value` +) AS duplicated_welford; + +Error: 1001(Unsupported), This feature is not implemented: Welford DISTINCT aggregations are not available + +SELECT stddev_pop_calc(stddev_pop_state(`value`)) FROM test_welford WHERE false; + ++-------------------------------------------------------+ +| stddev_pop_calc(stddev_pop_state(test_welford.value)) | ++-------------------------------------------------------+ +| | ++-------------------------------------------------------+ + +CREATE TABLE grouped_welford ( + `id` INT PRIMARY KEY, + `state` BINARY, + `ts` TIMESTAMP TIME INDEX DEFAULT now() +); + +Affected Rows: 0 + +INSERT INTO grouped_welford (`id`, `state`) +SELECT 1, stddev_pop_state(`value`) FROM test_welford WHERE id <= 5; + +Affected Rows: 1 + +INSERT INTO grouped_welford (`id`, `state`) +SELECT 2, stddev_pop_state(`value`) FROM test_welford WHERE id > 5; + +Affected Rows: 1 + +SELECT stddev_pop_calc(stddev_pop_merge(`state`)) FROM grouped_welford; + ++----------------------------------------------------------+ +| stddev_pop_calc(stddev_pop_merge(grouped_welford.state)) | ++----------------------------------------------------------+ +| 28.722813232690143 | ++----------------------------------------------------------+ + +-- A second DISTINCT argument set prevents DataFusion from rewriting this to GROUP BY. +SELECT + stddev_pop_calc(stddev_pop_merge(DISTINCT `state`)), + count(DISTINCT `id`) +FROM ( + SELECT `id`, `state` FROM grouped_welford + UNION ALL + SELECT 3 AS `id`, `state` FROM grouped_welford WHERE id = 1 +) AS duplicated_states; + +Error: 1001(Unsupported), This feature is not implemented: Welford DISTINCT aggregations are not available + +DROP TABLE grouped_welford; + +Affected Rows: 0 + +DROP TABLE test_welford; + +Affected Rows: 0 + +CREATE TABLE welford_window_raw ( + `id` INT PRIMARY KEY, + `value` DOUBLE, + `ts` TIMESTAMP TIME INDEX +); + +Affected Rows: 0 + +INSERT INTO welford_window_raw VALUES + (1, 1.0, '2024-01-01 00:01:05'), + (2, 2.0, '2024-01-01 00:01:20'), + (3, 3.0, '2024-01-01 00:01:50'), + (4, 10.0, '2024-01-01 00:02:10'), + (5, 20.0, '2024-01-01 00:02:40'), + (6, 4.0, '2024-01-01 00:03:05'), + (7, 8.0, '2024-01-01 00:03:15'), + (8, 12.0, '2024-01-01 00:03:35'), + (9, 16.0, '2024-01-01 00:03:55'), + (10, 100.0, '2024-01-01 00:04:05'), + (11, 200.0, '2024-01-01 00:04:25'), + (12, 300.0, '2024-01-01 00:04:45'); + +Affected Rows: 12 + +CREATE TABLE welford_minute_states ( + `minute_ts` TIMESTAMP TIME INDEX, + `state` BINARY +); + +Affected Rows: 0 + +INSERT INTO welford_minute_states (`minute_ts`, `state`) +SELECT + date_bin(INTERVAL '1 minute', `ts`) AS minute_ts, + stddev_pop_state(`value`) AS state +FROM welford_window_raw +GROUP BY minute_ts; + +Affected Rows: 4 + +-- Merging persisted minute states must reproduce aggregation over the raw samples. +WITH ranges AS ( + SELECT + '1-3' AS range_name, + CAST('2024-01-01 00:01:00' AS TIMESTAMP) AS start_ts, + CAST('2024-01-01 00:04:00' AS TIMESTAMP) AS end_ts + UNION ALL + SELECT + '2-4' AS range_name, + CAST('2024-01-01 00:02:00' AS TIMESTAMP) AS start_ts, + CAST('2024-01-01 00:05:00' AS TIMESTAMP) AS end_ts +), direct AS ( + SELECT + ranges.range_name, + count(*) AS sample_count, + stddev_pop(raw.`value`) AS stddev + FROM ranges CROSS JOIN welford_window_raw AS raw + WHERE raw.`ts` >= ranges.start_ts + AND raw.`ts` < ranges.end_ts + GROUP BY ranges.range_name +), merged AS ( + SELECT + ranges.range_name, + count(*) AS state_count, + stddev_pop_calc(stddev_pop_merge(states.`state`)) AS stddev + FROM ranges CROSS JOIN welford_minute_states AS states + WHERE states.minute_ts >= ranges.start_ts + AND states.minute_ts < ranges.end_ts + GROUP BY ranges.range_name +) +SELECT + direct.range_name, + direct.sample_count, + merged.state_count, + direct.stddev AS direct_stddev, + merged.stddev AS merged_stddev, + abs(direct.stddev - merged.stddev) AS difference +FROM direct JOIN merged ON direct.range_name = merged.range_name +ORDER BY direct.range_name; + ++------------+--------------+-------------+--------------------+--------------------+------------+ +| range_name | sample_count | state_count | direct_stddev | merged_stddev | difference | ++------------+--------------+-------------+--------------------+--------------------+------------+ +| 1-3 | 9 | 3 | 6.25586144900411 | 6.25586144900411 | 0.0 | +| 2-4 | 9 | 3 | 100.61048223620871 | 100.61048223620871 | 0.0 | ++------------+--------------+-------------+--------------------+--------------------+------------+ + +DROP TABLE welford_minute_states; + +Affected Rows: 0 + +DROP TABLE welford_window_raw; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/aggregate/welford.sql b/tests/cases/standalone/common/aggregate/welford.sql new file mode 100644 index 0000000000..edaed6ab69 --- /dev/null +++ b/tests/cases/standalone/common/aggregate/welford.sql @@ -0,0 +1,134 @@ +CREATE TABLE test_welford ( + `id` INT PRIMARY KEY, + `value` DOUBLE, + `ts` TIMESTAMP TIME INDEX DEFAULT now() +); + +INSERT INTO test_welford (`id`, `value`) VALUES + (1, 10.0), + (2, 20.0), + (3, 30.0), + (4, 40.0), + (5, 50.0), + (6, 60.0), + (7, 70.0), + (8, 80.0), + (9, 90.0), + (10, 100.0), + (11, NULL); + +SELECT stddev_pop_calc(stddev_pop_state(`value`)) FROM test_welford; + +-- A second DISTINCT argument set prevents DataFusion from rewriting this to GROUP BY. +SELECT + stddev_pop_calc(stddev_pop_state(DISTINCT `value`)), + count(DISTINCT `id`) +FROM ( + SELECT `id`, `value` FROM test_welford + UNION ALL + SELECT 12 AS `id`, 100.0 AS `value` +) AS duplicated_welford; + +SELECT stddev_pop_calc(stddev_pop_state(`value`)) FROM test_welford WHERE false; + +CREATE TABLE grouped_welford ( + `id` INT PRIMARY KEY, + `state` BINARY, + `ts` TIMESTAMP TIME INDEX DEFAULT now() +); + +INSERT INTO grouped_welford (`id`, `state`) +SELECT 1, stddev_pop_state(`value`) FROM test_welford WHERE id <= 5; + +INSERT INTO grouped_welford (`id`, `state`) +SELECT 2, stddev_pop_state(`value`) FROM test_welford WHERE id > 5; + +SELECT stddev_pop_calc(stddev_pop_merge(`state`)) FROM grouped_welford; + +-- A second DISTINCT argument set prevents DataFusion from rewriting this to GROUP BY. +SELECT + stddev_pop_calc(stddev_pop_merge(DISTINCT `state`)), + count(DISTINCT `id`) +FROM ( + SELECT `id`, `state` FROM grouped_welford + UNION ALL + SELECT 3 AS `id`, `state` FROM grouped_welford WHERE id = 1 +) AS duplicated_states; + +DROP TABLE grouped_welford; +DROP TABLE test_welford; + +CREATE TABLE welford_window_raw ( + `id` INT PRIMARY KEY, + `value` DOUBLE, + `ts` TIMESTAMP TIME INDEX +); + +INSERT INTO welford_window_raw VALUES + (1, 1.0, '2024-01-01 00:01:05'), + (2, 2.0, '2024-01-01 00:01:20'), + (3, 3.0, '2024-01-01 00:01:50'), + (4, 10.0, '2024-01-01 00:02:10'), + (5, 20.0, '2024-01-01 00:02:40'), + (6, 4.0, '2024-01-01 00:03:05'), + (7, 8.0, '2024-01-01 00:03:15'), + (8, 12.0, '2024-01-01 00:03:35'), + (9, 16.0, '2024-01-01 00:03:55'), + (10, 100.0, '2024-01-01 00:04:05'), + (11, 200.0, '2024-01-01 00:04:25'), + (12, 300.0, '2024-01-01 00:04:45'); + +CREATE TABLE welford_minute_states ( + `minute_ts` TIMESTAMP TIME INDEX, + `state` BINARY +); + +INSERT INTO welford_minute_states (`minute_ts`, `state`) +SELECT + date_bin(INTERVAL '1 minute', `ts`) AS minute_ts, + stddev_pop_state(`value`) AS state +FROM welford_window_raw +GROUP BY minute_ts; + +-- Merging persisted minute states must reproduce aggregation over the raw samples. +WITH ranges AS ( + SELECT + '1-3' AS range_name, + CAST('2024-01-01 00:01:00' AS TIMESTAMP) AS start_ts, + CAST('2024-01-01 00:04:00' AS TIMESTAMP) AS end_ts + UNION ALL + SELECT + '2-4' AS range_name, + CAST('2024-01-01 00:02:00' AS TIMESTAMP) AS start_ts, + CAST('2024-01-01 00:05:00' AS TIMESTAMP) AS end_ts +), direct AS ( + SELECT + ranges.range_name, + count(*) AS sample_count, + stddev_pop(raw.`value`) AS stddev + FROM ranges CROSS JOIN welford_window_raw AS raw + WHERE raw.`ts` >= ranges.start_ts + AND raw.`ts` < ranges.end_ts + GROUP BY ranges.range_name +), merged AS ( + SELECT + ranges.range_name, + count(*) AS state_count, + stddev_pop_calc(stddev_pop_merge(states.`state`)) AS stddev + FROM ranges CROSS JOIN welford_minute_states AS states + WHERE states.minute_ts >= ranges.start_ts + AND states.minute_ts < ranges.end_ts + GROUP BY ranges.range_name +) +SELECT + direct.range_name, + direct.sample_count, + merged.state_count, + direct.stddev AS direct_stddev, + merged.stddev AS merged_stddev, + abs(direct.stddev - merged.stddev) AS difference +FROM direct JOIN merged ON direct.range_name = merged.range_name +ORDER BY direct.range_name; + +DROP TABLE welford_minute_states; +DROP TABLE welford_window_raw; diff --git a/tests/compatibility/cases/stddev_pop_state/case.toml b/tests/compatibility/cases/stddev_pop_state/case.toml new file mode 100644 index 0000000000..717e73b3e7 --- /dev/null +++ b/tests/compatibility/cases/stddev_pop_state/case.toml @@ -0,0 +1,9 @@ +name = "stddev_pop_state" +reason = "Verify WLF1 states persisted by an old binary retain their binary layout and are readable directly and through stddev_pop_merge on the new binary." +introduced_by = "PR #8972" +topologies = ["distributed", "standalone"] +from_range = [">=v1.3.0"] +to_range = [">=v1.3.0"] +features = ["table", "query", "aggregate"] +owner = "query" +namespace = "stddev_pop_state" diff --git a/tests/compatibility/cases/stddev_pop_state/setup.sql b/tests/compatibility/cases/stddev_pop_state/setup.sql new file mode 100644 index 0000000000..c6f6b21a8e --- /dev/null +++ b/tests/compatibility/cases/stddev_pop_state/setup.sql @@ -0,0 +1,26 @@ +CREATE TABLE stddev_values ( + seq_id INT PRIMARY KEY, + grp INT, + val DOUBLE, + ts TIMESTAMP TIME INDEX DEFAULT now() +); + +INSERT INTO stddev_values (seq_id, grp, val) VALUES + (1, 0, 1.0), + (2, 0, 2.0), + (3, 1, 3.0), + (4, 1, 4.0); + +CREATE TABLE stddev_states ( + grp INT PRIMARY KEY, + state BINARY, + ts TIMESTAMP TIME INDEX DEFAULT now() +); + +INSERT INTO stddev_states (grp, state) +SELECT grp, stddev_pop_state(val) +FROM stddev_values +GROUP BY grp; + +ADMIN FLUSH_TABLE('stddev_values'); +ADMIN FLUSH_TABLE('stddev_states'); diff --git a/tests/compatibility/cases/stddev_pop_state/verify.result b/tests/compatibility/cases/stddev_pop_state/verify.result new file mode 100644 index 0000000000..6da57639a6 --- /dev/null +++ b/tests/compatibility/cases/stddev_pop_state/verify.result @@ -0,0 +1,21 @@ +-- Direct calculation from each state persisted by the old binary. +SELECT grp, stddev_pop_calc(state) AS stddev +FROM stddev_states +ORDER BY grp; + ++-----+--------+ +| grp | stddev | ++-----+--------+ +| 0 | 0.5 | +| 1 | 0.5 | ++-----+--------+ + +-- Merge all persisted states with the new binary before calculating. +SELECT stddev_pop_calc(stddev_pop_merge(state)) AS stddev +FROM stddev_states; + ++-------------------+ +| stddev | ++-------------------+ +| 1.118033988749895 | ++-------------------+ diff --git a/tests/compatibility/cases/stddev_pop_state/verify.sql b/tests/compatibility/cases/stddev_pop_state/verify.sql new file mode 100644 index 0000000000..70bc82d143 --- /dev/null +++ b/tests/compatibility/cases/stddev_pop_state/verify.sql @@ -0,0 +1,8 @@ +-- Direct calculation from each state persisted by the old binary. +SELECT grp, stddev_pop_calc(state) AS stddev +FROM stddev_states +ORDER BY grp; + +-- Merge all persisted states with the new binary before calculating. +SELECT stddev_pop_calc(stddev_pop_merge(state)) AS stddev +FROM stddev_states;