diff --git a/Cargo.lock b/Cargo.lock index ece381b7b8..0f473d4e03 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2477,6 +2477,7 @@ dependencies = [ "common-telemetry", "common-time", "common-version", + "criterion 0.7.0", "datafusion", "datafusion-common", "datafusion-expr", @@ -2511,7 +2512,8 @@ dependencies = [ "store-api", "table", "tokio", - "uddsketch", + "uddsketch-rs", + "vu128", "wkt", ] @@ -15823,11 +15825,13 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971" [[package]] -name = "uddsketch" +name = "uddsketch-rs" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/timescaledb-toolkit.git?rev=84828fe8fb494a6a61412a3da96517fc80f7bb20#84828fe8fb494a6a61412a3da96517fc80f7bb20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f0495d3e02c0e9b074e5631de7af419b9f4fdbfec4b1ca1b3180c8b3a926cb5" dependencies = [ - "serde", + "thiserror 2.0.17", + "vu128", ] [[package]] @@ -16216,6 +16220,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "vu128" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b18da3bd753c6f4373511e5f025423986560dfe4a5e7d642cc9a0266847f9fdd" + [[package]] name = "walkdir" version = "2.5.0" diff --git a/Cargo.toml b/Cargo.toml index bfd6153bf5..daee1bd409 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -266,8 +266,10 @@ tracing-appender = "0.2" tracing-opentelemetry = "0.33.0" tracing-subscriber = { version = "0.3", features = ["env-filter", "json", "fmt"] } typetag = "0.2" +uddsketch-rs = "0.1.0" uuid = { version = "1.17", features = ["serde", "v4", "v7", "fast-rng"] } vrl = "0.33" +vu128 = "1.1" zstd = "0.13" # DO_NOT_REMOVE_THIS: END_OF_EXTERNAL_DEPENDENCIES diff --git a/src/common/function/Cargo.toml b/src/common/function/Cargo.toml index 622f3ddfd2..a1ea31bc32 100644 --- a/src/common/function/Cargo.toml +++ b/src/common/function/Cargo.toml @@ -65,12 +65,18 @@ snafu.workspace = true sql.workspace = true store-api.workspace = true table.workspace = true -uddsketch = { git = "https://github.com/GreptimeTeam/timescaledb-toolkit.git", rev = "84828fe8fb494a6a61412a3da96517fc80f7bb20" } +uddsketch-rs.workspace = true +vu128.workspace = true wkt = { version = "0.11", optional = true } [dev-dependencies] approx = "0.5" +criterion.workspace = true futures.workspace = true pretty_assertions.workspace = true serde = { version = "1.0", features = ["derive"] } tokio.workspace = true + +[[bench]] +name = "uddsketch" +harness = false diff --git a/src/common/function/benches/uddsketch.rs b/src/common/function/benches/uddsketch.rs new file mode 100644 index 0000000000..6c8ad8779c --- /dev/null +++ b/src/common/function/benches/uddsketch.rs @@ -0,0 +1,201 @@ +// 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::hint::black_box; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Float64Array, Int64Array}; +use arrow::datatypes::{DataType, Field, Schema}; +use common_function::aggrs::approximate::uddsketch::UddSketchState; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use datafusion::common::ScalarValue; +use datafusion::logical_expr::Accumulator; +use datafusion::logical_expr::function::AccumulatorArgs; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_expr::expressions::{Column, Literal}; + +const BATCH_SIZES: [usize; 4] = [128, 256, 1024, 2048]; +const BUCKET_SIZE: i64 = 128; +const ERROR_RATE: f64 = 0.01; + +struct AccumulatorFactory { + udf: datafusion::logical_expr::AggregateUDF, + schema: Schema, + exprs: Vec>, + expr_fields: Vec>, + return_field: Arc, +} + +impl AccumulatorFactory { + fn new() -> Self { + let udf = UddSketchState::state_udf_impl(); + let schema = Schema::new(vec![ + Field::new("bucket_size", DataType::Int64, false), + Field::new("error", DataType::Float64, false), + Field::new("value", DataType::Float64, false), + ]); + let exprs: Vec> = vec![ + Arc::new(Literal::new(ScalarValue::Int64(Some(BUCKET_SIZE)))), + Arc::new(Literal::new(ScalarValue::Float64(Some(ERROR_RATE)))), + Arc::new(Column::new("value", 2)), + ]; + let expr_fields = exprs + .iter() + .map(|expr| expr.return_field(&schema).unwrap()) + .collect::>(); + let return_type = udf + .return_type(&[DataType::Int64, DataType::Float64, DataType::Float64]) + .unwrap(); + + Self { + udf, + schema, + exprs, + expr_fields, + return_field: Arc::new(Field::new("uddsketch_state", return_type, true)), + } + } + + fn create(&self) -> Box { + self.udf + .accumulator(AccumulatorArgs { + return_field: Arc::clone(&self.return_field), + schema: &self.schema, + ignore_nulls: false, + order_bys: &[], + is_reversed: false, + name: "uddsketch_state", + is_distinct: false, + exprs: &self.exprs, + expr_fields: &self.expr_fields, + }) + .unwrap() + } +} + +fn input_arrays(batch_size: usize) -> Vec { + let bucket_sizes = Arc::new(Int64Array::from_value(BUCKET_SIZE, batch_size)) as ArrayRef; + let errors = Arc::new(Float64Array::from_value(ERROR_RATE, batch_size)) as ArrayRef; + let mut state = 0x9e37_79b9_7f4a_7c15_u64; + let values = (0..batch_size) + .map(|index| { + state = state + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407) + .wrapping_add(index as u64); + let unit = (state >> 11) as f64 * (1.0 / (1_u64 << 53) as f64); + let magnitude = 10_f64.powf(-9.0 + 18.0 * unit); + if state & 1 == 0 { + magnitude + } else { + -magnitude + } + }) + .collect::>(); + let values = Arc::new(Float64Array::from(values)) as ArrayRef; + + vec![bucket_sizes, errors, values] +} + +fn validate(factory: &AccumulatorFactory, values: &[ArrayRef]) { + let mut accumulator = factory.create(); + accumulator.update_batch(values).unwrap(); + match accumulator.evaluate().unwrap() { + ScalarValue::Binary(Some(encoded)) => assert!(!encoded.is_empty()), + encoded => panic!("expected non-empty Binary, got {encoded:?}"), + } +} + +fn bench_uddsketch(c: &mut Criterion) { + let factory = AccumulatorFactory::new(); + let inputs = BATCH_SIZES + .into_iter() + .map(|batch_size| { + let values = input_arrays(batch_size); + validate(&factory, &values); + (batch_size, values) + }) + .collect::>(); + + let mut group = c.benchmark_group("uddsketch/ingest/fresh"); + for (batch_size, values) in &inputs { + group.throughput(Throughput::Elements(*batch_size as u64)); + group.bench_with_input( + BenchmarkId::new("batch_size", batch_size), + values, + |b, values| { + b.iter(|| { + let mut accumulator = factory.create(); + accumulator.update_batch(black_box(values)).unwrap(); + black_box(accumulator); + }); + }, + ); + } + group.finish(); + + let mut group = c.benchmark_group("uddsketch/ingest/reused"); + for (batch_size, values) in &inputs { + group.throughput(Throughput::Elements(*batch_size as u64)); + group.bench_with_input( + BenchmarkId::new("batch_size", batch_size), + values, + |b, values| { + let mut accumulator = factory.create(); + b.iter(|| { + accumulator.update_batch(black_box(values)).unwrap(); + black_box(&mut accumulator); + }); + }, + ); + } + group.finish(); + + let mut group = c.benchmark_group("uddsketch/ingest_evaluate/fresh"); + for (batch_size, values) in &inputs { + group.throughput(Throughput::Elements(*batch_size as u64)); + group.bench_with_input( + BenchmarkId::new("batch_size", batch_size), + values, + |b, values| { + b.iter(|| { + let mut accumulator = factory.create(); + accumulator.update_batch(black_box(values)).unwrap(); + black_box(accumulator.evaluate().unwrap()); + }); + }, + ); + } + group.finish(); + + let mut group = c.benchmark_group("uddsketch/ingest_evaluate/reused"); + for (batch_size, values) in &inputs { + group.throughput(Throughput::Elements(*batch_size as u64)); + group.bench_with_input( + BenchmarkId::new("batch_size", batch_size), + values, + |b, values| { + let mut accumulator = factory.create(); + b.iter(|| { + accumulator.update_batch(black_box(values)).unwrap(); + black_box(accumulator.evaluate().unwrap()); + }); + }, + ); + } + group.finish(); +} + +criterion_group!(benches, bench_uddsketch); +criterion_main!(benches); diff --git a/src/common/function/src/aggrs/approximate/uddsketch.rs b/src/common/function/src/aggrs/approximate/uddsketch.rs index 24b5c5455c..f7d1558d13 100644 --- a/src/common/function/src/aggrs/approximate/uddsketch.rs +++ b/src/common/function/src/aggrs/approximate/uddsketch.rs @@ -21,7 +21,6 @@ use std::sync::Arc; use common_query::prelude::*; -use common_telemetry::trace; use datafusion::common::cast::{as_binary_array, as_primitive_array}; use datafusion::common::not_impl_err; use datafusion::error::{DataFusionError, Result as DfResult}; @@ -29,27 +28,40 @@ use datafusion::logical_expr::function::AccumulatorArgs; use datafusion::logical_expr::{Accumulator as DfAccumulator, AggregateUDF, Volatility}; use datafusion::physical_plan::expressions::Literal; use datafusion::prelude::create_udaf; -use datatypes::arrow::array::ArrayRef; +use datatypes::arrow::array::{Array, ArrayRef}; use datatypes::arrow::datatypes::{DataType, Float64Type}; -use serde::{Deserialize, Serialize}; -use uddsketch::{SketchHashKey, UDDSketch}; +use uddsketch::{BatchWorkspace, UddSketch}; + +use crate::uddsketch_compat; pub const UDDSKETCH_STATE_NAME: &str = "uddsketch_state"; pub const UDDSKETCH_MERGE_NAME: &str = "uddsketch_merge"; -#[derive(Debug, Serialize, Deserialize)] +const MAX_BUCKETS: u32 = 1_000_000; + +#[derive(Debug)] pub struct UddSketchState { - uddsketch: UDDSketch, - error_rate: f64, + uddsketch: UddSketch, + workspace: BatchWorkspace, + values: Vec, } impl UddSketchState { - pub fn new(bucket_size: u64, error_rate: f64) -> Self { - Self { - uddsketch: UDDSketch::new(bucket_size, error_rate), - error_rate, + pub fn new(bucket_size: u32, error_rate: f64) -> DfResult { + if bucket_size > MAX_BUCKETS { + return Err(DataFusionError::Plan(format!( + "UDDSketch bucket size exceeds the maximum of {}", + MAX_BUCKETS + ))); } + let uddsketch = UddSketch::new(bucket_size, error_rate) + .map_err(|e| DataFusionError::Plan(e.to_string()))?; + Ok(Self { + uddsketch, + workspace: BatchWorkspace::default(), + values: Vec::new(), + }) } pub fn state_udf_impl() -> AggregateUDF { @@ -60,7 +72,7 @@ impl UddSketchState { Volatility::Immutable, Arc::new(|args| { let (bucket_size, error_rate) = downcast_accumulator_args(args)?; - Ok(Box::new(UddSketchState::new(bucket_size, error_rate))) + Ok(Box::new(UddSketchState::new(bucket_size, error_rate)?)) }), Arc::new(vec![DataType::Binary]), ) @@ -80,51 +92,44 @@ impl UddSketchState { Volatility::Immutable, Arc::new(|args| { let (bucket_size, error_rate) = downcast_accumulator_args(args)?; - Ok(Box::new(UddSketchState::new(bucket_size, error_rate))) + Ok(Box::new(UddSketchState::new(bucket_size, error_rate)?)) }), Arc::new(vec![DataType::Binary]), ) } - fn update(&mut self, value: f64) { - self.uddsketch.add_value(value); - } - fn merge(&mut self, raw: &[u8]) -> DfResult<()> { - if let Ok(uddsketch) = bincode::deserialize::(raw) { - if uddsketch.uddsketch.count() != 0 { - if self.uddsketch.max_allowed_buckets() != uddsketch.uddsketch.max_allowed_buckets() - || (self.error_rate - uddsketch.error_rate).abs() >= 1e-9 - { - return Err(DataFusionError::Plan(format!( - "Merging UDDSketch with different parameters: arguments={:?} vs actual input={:?}", - (self.uddsketch.max_allowed_buckets(), self.error_rate), - ( - uddsketch.uddsketch.max_allowed_buckets(), - uddsketch.error_rate - ) - ))); - } - self.uddsketch.merge_sketch(&uddsketch.uddsketch); - } - } else { - trace!("Warning: Failed to deserialize UDDSketch from {:?}", raw); - return Err(DataFusionError::Plan( - "Failed to deserialize UDDSketch from binary".to_string(), - )); + let uddsketch = uddsketch_compat::decode(raw).map_err(|e| { + common_telemetry::trace!("Failed to deserialize UDDSketch: {}", e); + DataFusionError::Plan("Failed to deserialize UDDSketch from binary".to_string()) + })?; + if uddsketch.count() == 0 { + return Ok(()); } - - Ok(()) + if self.uddsketch.max_buckets() != uddsketch.max_buckets() + || self.uddsketch.initial_error().to_bits() != uddsketch.initial_error().to_bits() + { + return Err(DataFusionError::Plan(format!( + "Merging UDDSketch with different parameters: arguments={:?} vs actual input={:?}", + (self.uddsketch.max_buckets(), self.uddsketch.initial_error()), + (uddsketch.max_buckets(), uddsketch.initial_error()) + ))); + } + self.uddsketch + .merge(&uddsketch) + .map_err(|e| DataFusionError::Plan(e.to_string())) } } -fn downcast_accumulator_args(args: AccumulatorArgs) -> DfResult<(u64, f64)> { +fn downcast_accumulator_args(args: AccumulatorArgs) -> DfResult<(u32, f64)> { let bucket_size = match args.exprs[0] .as_any() .downcast_ref::() .map(|lit| lit.value()) { - Some(ScalarValue::Int64(Some(value))) => *value as u64, + Some(ScalarValue::Int64(Some(value))) => u32::try_from(*value).map_err(|_| { + DataFusionError::Plan(format!("Invalid UDDSketch bucket size: {}", value)) + })?, _ => { return not_impl_err!( "{} not supported for bucket size: {}", @@ -158,9 +163,16 @@ impl DfAccumulator for UddSketchState { match array.data_type() { DataType::Float64 => { let f64_array = as_primitive_array::(array)?; - for v in f64_array.iter().flatten() { - self.update(v); - } + let values: &[f64] = if f64_array.null_count() == 0 { + f64_array.values().as_ref() + } else { + self.values.clear(); + self.values.extend(f64_array.iter().flatten()); + self.values.as_slice() + }; + self.uddsketch + .add_batch_with_workspace(values, &mut self.workspace) + .map_err(|e| DataFusionError::Execution(e.to_string()))?; } // meaning instantiate as `uddsketch_merge` DataType::Binary => self.merge_batch(std::slice::from_ref(array))?, @@ -176,35 +188,21 @@ impl DfAccumulator for UddSketchState { } fn evaluate(&mut self) -> DfResult { - Ok(ScalarValue::Binary(Some( - bincode::serialize(&self).map_err(|e| { - DataFusionError::Internal(format!("Failed to serialize UDDSketch: {}", e)) - })?, - ))) + Ok(ScalarValue::Binary(Some(self.uddsketch.encode().map_err( + |e| DataFusionError::Internal(format!("Failed to serialize UDDSketch: {}", e)), + )?))) } fn size(&self) -> usize { - // Base size of UDDSketch struct fields - let mut total_size = std::mem::size_of::() * 3 + // alpha, gamma, values_sum - std::mem::size_of::() + // compactions - std::mem::size_of::() * 2; // max_buckets, num_values - - // Size of buckets (SketchHashMap) - // Each bucket entry contains: - // - SketchHashKey (enum with i64/Zero/Invalid variants) - // - SketchHashEntry (count: u64, next: SketchHashKey) - let bucket_entry_size = std::mem::size_of::() + // key - std::mem::size_of::() + // count - std::mem::size_of::(); // next - - total_size += self.uddsketch.current_buckets_count() * bucket_entry_size; - - total_size + std::mem::size_of::() - std::mem::size_of::() + + self.uddsketch.allocated_size() + + self.workspace.allocated_size() + + self.values.capacity() * std::mem::size_of::() } fn state(&mut self) -> DfResult> { Ok(vec![ScalarValue::Binary(Some( - bincode::serialize(&self).map_err(|e| { + self.uddsketch.encode().map_err(|e| { DataFusionError::Internal(format!("Failed to serialize UDDSketch: {}", e)) })?, ))]) @@ -224,20 +222,21 @@ impl DfAccumulator for UddSketchState { #[cfg(test)] mod tests { use datafusion::arrow::array::{BinaryArray, Float64Array}; + use uddsketch::UddSketchRef; use super::*; #[test] fn test_uddsketch_state_basic() { - let mut state = UddSketchState::new(10, 0.01); - state.update(1.0); - state.update(2.0); - state.update(3.0); + let mut state = UddSketchState::new(10, 0.01).unwrap(); + state.uddsketch.add(1.0).unwrap(); + state.uddsketch.add(2.0).unwrap(); + state.uddsketch.add(3.0).unwrap(); let result = state.evaluate().unwrap(); if let ScalarValue::Binary(Some(bytes)) = result { - let deserialized: UddSketchState = bincode::deserialize(&bytes).unwrap(); - assert_eq!(deserialized.uddsketch.count(), 3); + let encoded = UddSketchRef::parse(&bytes).unwrap(); + assert_eq!(encoded.count(), 3); } else { panic!("Expected binary scalar value"); } @@ -245,39 +244,38 @@ mod tests { #[test] fn test_uddsketch_state_roundtrip() { - let mut state = UddSketchState::new(10, 0.01); - state.update(1.0); - state.update(2.0); + let mut state = UddSketchState::new(10, 0.01).unwrap(); + state.uddsketch.add(1.0).unwrap(); + state.uddsketch.add(2.0).unwrap(); // Serialize let serialized = state.evaluate().unwrap(); // Create new state and merge the serialized data - let mut new_state = UddSketchState::new(10, 0.01); + let mut new_state = UddSketchState::new(10, 0.01).unwrap(); if let ScalarValue::Binary(Some(bytes)) = &serialized { new_state.merge(bytes).unwrap(); - // Verify the merged state matches original by comparing deserialized values - let original_sketch: UddSketchState = bincode::deserialize(bytes).unwrap(); - let original_sketch = original_sketch.uddsketch; + let original_sketch = UddSketchRef::parse(bytes).unwrap(); let new_result = new_state.evaluate().unwrap(); if let ScalarValue::Binary(Some(new_bytes)) = new_result { - let new_sketch: UddSketchState = bincode::deserialize(&new_bytes).unwrap(); - let new_sketch = new_sketch.uddsketch; + let new_sketch = UddSketchRef::parse(&new_bytes).unwrap(); assert_eq!(original_sketch.count(), new_sketch.count()); assert_eq!(original_sketch.sum(), new_sketch.sum()); - assert_eq!(original_sketch.mean(), new_sketch.mean()); - assert_eq!(original_sketch.max_error(), new_sketch.max_error()); + assert_eq!( + original_sketch.max_error().unwrap(), + new_sketch.max_error().unwrap() + ); // Compare a few quantiles to ensure statistical equivalence for q in [0.1, 0.5, 0.9].iter() { + let original = original_sketch.quantile(*q).unwrap().unwrap(); + let merged = new_sketch.quantile(*q).unwrap().unwrap(); assert!( - (original_sketch.estimate_quantile(*q) - new_sketch.estimate_quantile(*q)) - .abs() - < 1e-10, + (original - merged).abs() < 1e-10, "Quantile {} mismatch: original={}, new={}", q, - original_sketch.estimate_quantile(*q), - new_sketch.estimate_quantile(*q) + original, + merged ); } } else { @@ -288,10 +286,41 @@ mod tests { } } + #[test] + fn test_uddsketch_state_merges_legacy_state() { + let mut state = UddSketchState::new(128, 0.01).unwrap(); + + state.merge(uddsketch_compat::LEGACY_STATE).unwrap(); + + let ScalarValue::Binary(Some(encoded)) = state.evaluate().unwrap() else { + panic!("Expected binary scalar value"); + }; + let sketch = UddSketchRef::parse(&encoded).unwrap(); + assert_eq!(sketch.count(), 4); + assert_eq!(sketch.sum(), 1.0); + assert_eq!(sketch.quantile(0.5).unwrap(), Some(0.9900000000000001)); + } + + #[test] + fn test_uddsketch_state_merges_compacted_legacy_state() { + let mut legacy_state = uddsketch_compat::COMPACTED_LEGACY_SKETCH.to_vec(); + legacy_state.extend_from_slice(&0.01_f64.to_le_bytes()); + let mut state = UddSketchState::new(7, 0.01).unwrap(); + + state.merge(&legacy_state).unwrap(); + + let ScalarValue::Binary(Some(encoded)) = state.evaluate().unwrap() else { + panic!("Expected binary scalar value"); + }; + let sketch = UddSketchRef::parse(&encoded).unwrap(); + assert_eq!(sketch.count(), 201); + assert_eq!(sketch.times_compacted(), 12); + } + #[test] fn test_uddsketch_state_batch_update() { - let mut state = UddSketchState::new(10, 0.01); - let values = vec![1.0f64, 2.0, 3.0]; + let mut state = UddSketchState::new(10, 0.01).unwrap(); + let values = vec![Some(1.0f64), None, Some(2.0), Some(3.0)]; let array = Arc::new(Float64Array::from(values)) as ArrayRef; state @@ -300,25 +329,52 @@ mod tests { let result = state.evaluate().unwrap(); if let ScalarValue::Binary(Some(bytes)) = result { - let deserialized: UddSketchState = bincode::deserialize(&bytes).unwrap(); - let deserialized = deserialized.uddsketch; - assert_eq!(deserialized.count(), 3); + let encoded = UddSketchRef::parse(&bytes).unwrap(); + assert_eq!(encoded.count(), 3); } else { panic!("Expected binary scalar value"); } } + #[test] + fn test_uddsketch_state_non_null_batch_avoids_values_buffer() { + let mut state = UddSketchState::new(10, 0.01).unwrap(); + let array = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])) as ArrayRef; + + state + .update_batch(&[array.clone(), array.clone(), array]) + .unwrap(); + + assert_eq!(state.uddsketch.count(), 3); + assert_eq!(state.values.capacity(), 0); + } + + #[test] + fn test_uddsketch_state_non_null_sliced_batch() { + let mut state = UddSketchState::new(10, 0.01).unwrap(); + let array = Float64Array::from(vec![1.0, 2.0, 3.0, 4.0]); + let array = array.slice(1, 2); + let array = Arc::new(array) as ArrayRef; + + state + .update_batch(&[array.clone(), array.clone(), array]) + .unwrap(); + + assert_eq!(state.uddsketch.count(), 2); + assert_eq!(state.uddsketch.sum(), 5.0); + } + #[test] fn test_uddsketch_state_merge_batch() { - let mut state1 = UddSketchState::new(10, 0.01); - state1.update(1.0); + let mut state1 = UddSketchState::new(10, 0.01).unwrap(); + state1.uddsketch.add(1.0).unwrap(); let state1_binary = state1.evaluate().unwrap(); - let mut state2 = UddSketchState::new(10, 0.01); - state2.update(2.0); + let mut state2 = UddSketchState::new(10, 0.01).unwrap(); + state2.uddsketch.add(2.0).unwrap(); let state2_binary = state2.evaluate().unwrap(); - let mut merged_state = UddSketchState::new(10, 0.01); + let mut merged_state = UddSketchState::new(10, 0.01).unwrap(); if let (ScalarValue::Binary(Some(bytes1)), ScalarValue::Binary(Some(bytes2))) = (&state1_binary, &state2_binary) { @@ -330,9 +386,8 @@ mod tests { let result = merged_state.evaluate().unwrap(); if let ScalarValue::Binary(Some(bytes)) = result { - let deserialized: UddSketchState = bincode::deserialize(&bytes).unwrap(); - let deserialized = deserialized.uddsketch; - assert_eq!(deserialized.count(), 2); + let encoded = UddSketchRef::parse(&bytes).unwrap(); + assert_eq!(encoded.count(), 2); } else { panic!("Expected binary scalar value"); } @@ -343,13 +398,14 @@ mod tests { #[test] fn test_uddsketch_state_size() { - let mut state = UddSketchState::new(10, 0.01); + let mut state = UddSketchState::new(10, 0.01).unwrap(); let initial_size = state.size(); // Add some values to create buckets - state.update(1.0); - state.update(2.0); - state.update(3.0); + let array = Arc::new(Float64Array::from_iter_values((0..64).map(f64::from))) as ArrayRef; + state + .update_batch(&[array.clone(), array.clone(), array]) + .unwrap(); let size_with_values = state.size(); assert!( @@ -358,14 +414,29 @@ mod tests { initial_size, size_with_values ); + } - // Verify size increases with more buckets - state.update(10.0); // This should create a new bucket - assert!( - state.size() > size_with_values, - "Size should increase after adding new bucket: prev={}, new={}", - size_with_values, - state.size() - ); + #[test] + fn test_uddsketch_state_rejects_invalid_config() { + assert!(UddSketchState::new(6, 0.01).is_err()); + assert!(UddSketchState::new(10, 1.0).is_err()); + + let mut maximum = UddSketchState::new(1_000_000, 0.01).unwrap(); + let ScalarValue::Binary(Some(encoded)) = maximum.evaluate().unwrap() else { + panic!("Expected binary scalar value"); + }; + UddSketchRef::parse(&encoded).unwrap(); + assert!(UddSketchState::new(1_000_001, 0.01).is_err()); + } + + #[test] + fn test_uddsketch_state_rejects_nan_batch() { + let mut state = UddSketchState::new(10, 0.01).unwrap(); + let array = Arc::new(Float64Array::from(vec![1.0, f64::NAN])) as ArrayRef; + + let error = state + .update_batch(&[array.clone(), array.clone(), array]) + .unwrap_err(); + assert!(error.to_string().contains("NaN values are not supported")); } } diff --git a/src/common/function/src/lib.rs b/src/common/function/src/lib.rs index 7abd595367..336f8d489c 100644 --- a/src/common/function/src/lib.rs +++ b/src/common/function/src/lib.rs @@ -18,6 +18,7 @@ mod admin; mod flush_flow; mod macros; mod system; +mod uddsketch_compat; pub mod aggrs; pub mod function; diff --git a/src/common/function/src/scalars/uddsketch_calc.rs b/src/common/function/src/scalars/uddsketch_calc.rs index d5d91e7edb..d28bb78d1b 100644 --- a/src/common/function/src/scalars/uddsketch_calc.rs +++ b/src/common/function/src/scalars/uddsketch_calc.rs @@ -22,10 +22,10 @@ use datafusion_common::DataFusionError; use datafusion_common::arrow::array::{Array, AsArray, Float64Builder}; use datafusion_common::arrow::datatypes::{DataType, Float64Type}; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, Volatility}; -use uddsketch::UDDSketch; use crate::function::{Function, extract_args}; use crate::function_registry::FunctionRegistry; +use crate::uddsketch_compat; const NAME: &str = "uddsketch_calc"; @@ -114,26 +114,19 @@ impl Function for UddSketchCalcFunction { let sketch_bytes = sketch_opt.unwrap(); let perc = perc_opt.unwrap(); - // Deserialize the UDDSketch from its bincode representation - let sketch: UDDSketch = match bincode::deserialize(sketch_bytes) { - Ok(s) => s, + let value = match uddsketch_compat::quantile(sketch_bytes, perc) { + Ok(value) => value, Err(e) => { - common_telemetry::trace!("Failed to deserialize UDDSketch: {}", e); + common_telemetry::trace!("Failed to parse UDDSketch: {}", e); builder.append_null(); continue; } }; - // Check if the sketch is empty, if so, return null - // This is important to avoid panics when calling estimate_quantile on an empty sketch - // In practice, this will happen if input is all null - if sketch.bucket_iter().count() == 0 { - builder.append_null(); - continue; + match value { + Some(value) => builder.append_value(value), + None => builder.append_null(), } - // Compute the estimated quantile from the sketch - let result = sketch.estimate_quantile(perc); - builder.append_value(result); } Ok(ColumnarValue::Array(Arc::new(builder.finish()))) @@ -146,6 +139,7 @@ mod tests { use arrow_schema::Field; use datafusion_common::arrow::array::{BinaryArray, Float64Array}; + use uddsketch::UddSketch; use super::*; @@ -159,24 +153,17 @@ mod tests { ); // Create a test sketch - let mut sketch = UDDSketch::new(128, 0.01); - sketch.add_value(10.0); - sketch.add_value(20.0); - sketch.add_value(30.0); - sketch.add_value(40.0); - sketch.add_value(50.0); - sketch.add_value(60.0); - sketch.add_value(70.0); - sketch.add_value(80.0); - sketch.add_value(90.0); - sketch.add_value(100.0); + let mut sketch = UddSketch::new(128, 0.01).unwrap(); + sketch + .add_batch(&[10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0]) + .unwrap(); // Get expected values directly from the sketch - let expected_p50 = sketch.estimate_quantile(0.5); - let expected_p90 = sketch.estimate_quantile(0.9); - let expected_p95 = sketch.estimate_quantile(0.95); + let expected_p50 = sketch.quantile(0.5).unwrap().unwrap(); + let expected_p90 = sketch.quantile(0.9).unwrap().unwrap(); + let expected_p95 = sketch.quantile(0.95).unwrap().unwrap(); - let serialized = bincode::serialize(&sketch).unwrap(); + let serialized = sketch.encode().unwrap(); let percentiles = vec![0.5, 0.9, 0.95]; let args = vec![ @@ -207,6 +194,34 @@ mod tests { assert!((result.value(2) - expected_p95).abs() < 1e-10); } + #[test] + fn test_uddsketch_calc_function_reads_legacy_state() { + let function = UddSketchCalcFunction::default(); + let args = vec![ + ColumnarValue::Array(Arc::new(Float64Array::from(vec![0.5]))), + ColumnarValue::Array(Arc::new(BinaryArray::from_iter_values(vec![ + uddsketch_compat::LEGACY_STATE, + ]))), + ]; + + let result = function + .invoke_with_args(ScalarFunctionArgs { + args, + arg_fields: vec![], + number_rows: 1, + return_field: Arc::new(Field::new("x", DataType::Float64, false)), + config_options: Arc::new(Default::default()), + }) + .unwrap(); + let ColumnarValue::Array(result) = result else { + unreachable!() + }; + let result = result.as_primitive::(); + assert_eq!(result.len(), 1); + assert!(!result.is_null(0)); + assert_eq!(result.value(0), 0.9900000000000001); + } + #[test] fn test_uddsketch_calc_function_errors() { let function = UddSketchCalcFunction::default(); @@ -249,5 +264,31 @@ mod tests { let result = result.as_primitive::(); assert_eq!(result.len(), 1); assert!(result.is_null(0)); + + let empty = UddSketch::new(128, 0.01).unwrap().encode().unwrap(); + let mut populated = UddSketch::new(128, 0.01).unwrap(); + populated.add(1.0).unwrap(); + let populated = populated.encode().unwrap(); + let args = vec![ + ColumnarValue::Array(Arc::new(Float64Array::from(vec![0.5, -0.1, f64::NAN]))), + ColumnarValue::Array(Arc::new(BinaryArray::from_iter_values(vec![ + empty, + populated.clone(), + populated, + ]))), + ]; + let result = function + .invoke_with_args(ScalarFunctionArgs { + args, + arg_fields: vec![], + number_rows: 3, + return_field: Arc::new(Field::new("x", DataType::Float64, false)), + config_options: Arc::new(Default::default()), + }) + .unwrap(); + let ColumnarValue::Array(result) = result else { + unreachable!() + }; + assert_eq!(result.as_primitive::().null_count(), 3); } } diff --git a/src/common/function/src/uddsketch_compat.rs b/src/common/function/src/uddsketch_compat.rs new file mode 100644 index 0000000000..733b4ab043 --- /dev/null +++ b/src/common/function/src/uddsketch_compat.rs @@ -0,0 +1,619 @@ +// 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. + +//! Compatibility decoder for UDDSketch states written before the canonical v1 format. + +use std::collections::{HashMap, HashSet}; + +use bincode::Options; +use serde::Deserialize; +use uddsketch::{UddSketch, UddSketchRef}; + +const MAX_BYTES: usize = 64 * 1024 * 1024; +const MAX_BUCKETS: usize = 1_000_000; +const HEADER_LEN: usize = 48; + +#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, PartialEq)] +enum LegacyBucketKey { + Negative(i64), + Zero, + Positive(i64), + Invalid, +} + +#[derive(Debug, Deserialize)] +struct LegacyBucket { + count: u64, + next: LegacyBucketKey, +} + +#[derive(Debug, Deserialize)] +struct LegacyBucketStore { + map: HashMap, + head: LegacyBucketKey, +} + +#[derive(Debug, Deserialize)] +struct LegacyUddSketch { + buckets: LegacyBucketStore, + alpha: f64, + gamma: f64, + compactions: u32, + max_buckets: u64, + count: u64, + sum: f64, +} + +#[derive(Debug, Deserialize)] +struct LegacyUddSketchState { + uddsketch: LegacyUddSketch, + initial_error: f64, +} + +pub(crate) fn decode(raw: &[u8]) -> Result { + match UddSketch::decode(raw) { + Ok(sketch) => Ok(sketch), + Err(current_error) => decode_legacy_state(raw).map_err(|legacy_error| { + format!( + "canonical decode failed: {current_error}; legacy decode failed: {legacy_error}" + ) + }), + } +} + +pub(crate) fn quantile(raw: &[u8], quantile: f64) -> Result, String> { + match UddSketchRef::parse(raw) { + Ok(sketch) => sketch.quantile(quantile).map_err(|error| error.to_string()), + Err(current_error) => decode_legacy_sketch(raw) + .and_then(|sketch| sketch.quantile(quantile)) + .map_err(|legacy_error| { + format!( + "canonical decode failed: {current_error}; legacy decode failed: {legacy_error}" + ) + }), + } +} + +fn validate_legacy_input(raw: &[u8]) -> Result<(), String> { + if raw.len() > MAX_BYTES { + return Err("input exceeds the legacy decode byte limit".to_string()); + } + let map_len = raw + .get(..8) + .and_then(|bytes| bytes.try_into().ok()) + .map(u64::from_le_bytes) + .ok_or_else(|| "legacy input is truncated before the bucket count".to_string())?; + if map_len > MAX_BUCKETS as u64 { + return Err("legacy populated bucket count exceeds decode limit".to_string()); + } + Ok(()) +} + +fn legacy_options() -> impl Options { + bincode::DefaultOptions::new() + .with_fixint_encoding() + .with_limit(MAX_BYTES as u64) + .reject_trailing_bytes() +} + +fn decode_legacy_state(raw: &[u8]) -> Result { + validate_legacy_input(raw)?; + let state = legacy_options() + .deserialize::(raw) + .map_err(|error| error.to_string())?; + let encoded = state.into_canonical()?; + UddSketch::decode(&encoded).map_err(|error| error.to_string()) +} + +fn decode_legacy_sketch(raw: &[u8]) -> Result { + validate_legacy_input(raw)?; + legacy_options() + .deserialize::(raw) + .map(|state| state.uddsketch) + .or_else(|state_error| { + legacy_options() + .deserialize::(raw) + .map_err(|sketch_error| { + format!( + "legacy state decode failed: {state_error}; legacy sketch decode failed: {sketch_error}" + ) + }) + }) +} + +impl LegacyUddSketchState { + fn into_canonical(self) -> Result, String> { + let sketch = self.uddsketch; + let max_buckets = u32::try_from(sketch.max_buckets) + .map_err(|_| "legacy maximum bucket count exceeds u32".to_string())?; + if !(7..=MAX_BUCKETS as u32).contains(&max_buckets) { + return Err("legacy maximum bucket count is outside supported limits".to_string()); + } + let compactions = u8::try_from(sketch.compactions) + .map_err(|_| "legacy compaction count exceeds u8".to_string())?; + if compactions > 63 { + return Err("legacy compaction count exceeds 63".to_string()); + } + + let (expected_alpha, expected_gamma) = mapping(self.initial_error, compactions)?; + if sketch.alpha.to_bits() != expected_alpha.to_bits() + || sketch.gamma.to_bits() != expected_gamma.to_bits() + { + return Err("legacy mapping metadata is inconsistent".to_string()); + } + + let buckets = sketch.buckets.ordered()?; + if buckets.len() > max_buckets as usize { + return Err("legacy populated bucket count exceeds its configured limit".to_string()); + } + let decoded_count = buckets.iter().try_fold(0_u64, |total, (_, count)| { + total + .checked_add(*count) + .ok_or_else(|| "legacy bucket count sum overflows u64".to_string()) + })?; + if decoded_count != sketch.count { + return Err("legacy bucket counts do not match the value count".to_string()); + } + if sketch.count == 0 && sketch.sum.to_bits() != 0.0_f64.to_bits() { + return Err("legacy empty sketch sum is not positive zero".to_string()); + } + + encode_canonical( + max_buckets, + self.initial_error, + compactions, + sketch.count, + sketch.sum, + &buckets, + ) + } +} + +impl LegacyUddSketch { + fn quantile(self, quantile: f64) -> Result, String> { + if !quantile.is_finite() || !(0.0..=1.0).contains(&quantile) { + return Err("invalid quantile".to_string()); + } + if self.compactions >= 64 { + return Err("legacy compaction count must be below 64".to_string()); + } + validate_current_mapping(self.alpha, self.gamma, self.compactions)?; + if self.max_buckets == 0 || self.max_buckets > MAX_BUCKETS as u64 { + return Err("legacy maximum bucket count is outside supported limits".to_string()); + } + + let count = self.count; + let alpha = self.alpha; + let gamma = self.gamma; + let buckets = self.buckets.ordered()?; + if buckets.len() > self.max_buckets as usize { + return Err("legacy populated bucket count exceeds its configured limit".to_string()); + } + for (key, _) in &buckets { + if !legacy_bucket_key_is_attainable(self.gamma, *key) { + return Err("legacy bucket index cannot represent an f64".to_string()); + } + } + let decoded_count = buckets.iter().try_fold(0_u64, |total, (_, count)| { + total + .checked_add(*count) + .ok_or_else(|| "legacy bucket count sum overflows u64".to_string()) + })?; + if decoded_count != count { + return Err("legacy bucket counts do not match the value count".to_string()); + } + if count == 0 { + return Ok(None); + } + + let target = if quantile == 1.0 { + count + } else { + ((count as f64 * quantile) as u64) + .saturating_add(1) + .min(count) + }; + let mut seen = 0_u64; + for (key, bucket_count) in buckets { + seen += bucket_count; + if seen >= target { + return Ok(Some(legacy_bucket_value(alpha, gamma, key)?)); + } + } + Err("legacy bucket counts do not cover the quantile rank".to_string()) + } +} + +fn validate_current_mapping( + mut alpha: f64, + mut gamma: f64, + compactions: u32, +) -> Result<(), String> { + if !alpha.is_finite() || !(0.0..=1.0).contains(&alpha) { + return Err("legacy current error is outside [0, 1]".to_string()); + } + if !gamma.is_finite() || gamma <= 1.0 { + return (alpha == 1.0 && gamma == f64::INFINITY && compactions >= 5) + .then_some(()) + .ok_or_else(|| "legacy gamma must be greater than one".to_string()); + } + if alpha == 1.0 { + return (compactions > 0 && 1.0 - 2.0 / (gamma + 1.0) == 1.0) + .then_some(()) + .ok_or_else(|| "legacy saturated mapping metadata is inconsistent".to_string()); + } + + for _ in 0..compactions { + alpha /= 1.0 + (1.0 - alpha * alpha).sqrt(); + gamma = gamma.sqrt(); + } + let expected_gamma = (1.0 + alpha) / (1.0 - alpha); + let relative_difference = (gamma - expected_gamma).abs() / expected_gamma; + if relative_difference > 1e-10 { + return Err("legacy mapping metadata is inconsistent".to_string()); + } + Ok(()) +} + +fn legacy_bucket_key_is_attainable(gamma: f64, key: LegacyBucketKey) -> bool { + let index = match key { + LegacyBucketKey::Zero => return true, + LegacyBucketKey::Negative(index) | LegacyBucketKey::Positive(index) => index, + LegacyBucketKey::Invalid => return false, + }; + if index == i64::MAX { + return true; + } + + let minimum = f64::from_bits(1).log(gamma).ceil() as i64; + let maximum = f64::MAX.log(gamma).ceil() as i64; + (minimum..=maximum).contains(&index) +} + +fn legacy_bucket_value(alpha: f64, gamma: f64, key: LegacyBucketKey) -> Result { + let magnitude = |index: i64| gamma.powf(index as f64 - 1.0) * (1.0 + alpha); + match key { + LegacyBucketKey::Negative(index) => Ok(-magnitude(index)), + LegacyBucketKey::Zero => Ok(0.0), + LegacyBucketKey::Positive(index) => Ok(magnitude(index)), + LegacyBucketKey::Invalid => Err("legacy bucket chain contains the end marker".to_string()), + } +} + +impl LegacyBucketStore { + fn ordered(self) -> Result, String> { + if self.map.len() > MAX_BUCKETS { + return Err("legacy populated bucket count exceeds decode limit".to_string()); + } + if self.map.is_empty() { + if self.head != LegacyBucketKey::Invalid { + return Err("legacy empty bucket store has a nonempty head".to_string()); + } + return Ok(Vec::new()); + } + + let mut buckets = Vec::with_capacity(self.map.len()); + let mut visited = HashSet::with_capacity(self.map.len()); + let mut key = self.head; + while key != LegacyBucketKey::Invalid { + if !visited.insert(key) { + return Err("legacy bucket chain contains a cycle".to_string()); + } + let bucket = self + .map + .get(&key) + .ok_or_else(|| "legacy bucket chain references a missing bucket".to_string())?; + if bucket.count == 0 { + return Err("legacy bucket has a zero count".to_string()); + } + buckets.push((key, bucket.count)); + key = bucket.next; + } + if buckets.len() != self.map.len() || self.map.contains_key(&LegacyBucketKey::Invalid) { + return Err("legacy bucket store contains unreachable buckets".to_string()); + } + if !buckets.windows(2).all(|pair| key_lt(pair[0].0, pair[1].0)) { + return Err("legacy bucket chain is not strictly ordered".to_string()); + } + Ok(buckets) + } +} + +fn mapping(initial_error: f64, compactions: u8) -> Result<(f64, f64), String> { + if !initial_error.is_finite() || !(1e-12..1.0).contains(&initial_error) { + return Err("legacy initial error is outside [1e-12, 1)".to_string()); + } + let mut alpha = initial_error; + let mut gamma = (1.0 + initial_error) / (1.0 - initial_error); + for _ in 0..compactions { + gamma *= gamma; + alpha = 2.0 * alpha / (1.0 + alpha.powi(2)); + } + Ok((alpha, gamma)) +} + +fn encode_canonical( + max_buckets: u32, + initial_error: f64, + compactions: u8, + count: u64, + sum: f64, + buckets: &[(LegacyBucketKey, u64)], +) -> Result, String> { + let negative = buckets + .iter() + .filter_map(|(key, count)| match key { + LegacyBucketKey::Negative(index) => Some((*index, *count)), + _ => None, + }) + .collect::>(); + let zero_count = buckets + .iter() + .find_map(|(key, count)| (*key == LegacyBucketKey::Zero).then_some(*count)) + .unwrap_or(0); + let positive = buckets + .iter() + .filter_map(|(key, count)| match key { + LegacyBucketKey::Positive(index) => Some((*index, *count)), + _ => None, + }) + .collect::>(); + + let mut encoded = vec![0; HEADER_LEN]; + put_varint(&mut encoded, negative.len() as u64); + put_varint(&mut encoded, zero_count); + put_varint(&mut encoded, positive.len() as u64); + encode_section(&mut encoded, &negative, true)?; + encode_section(&mut encoded, &positive, false)?; + + let payload_len = u32::try_from(encoded.len() - HEADER_LEN) + .map_err(|_| "legacy canonical payload exceeds u32".to_string())?; + encoded[0..4].copy_from_slice(b"UDDS"); + encoded[4] = 1; + encoded[6] = compactions; + encoded[8..12].copy_from_slice(&max_buckets.to_le_bytes()); + encoded[12..16].copy_from_slice(&(buckets.len() as u32).to_le_bytes()); + encoded[16..24].copy_from_slice(&initial_error.to_bits().to_le_bytes()); + encoded[24..32].copy_from_slice(&count.to_le_bytes()); + encoded[32..40].copy_from_slice(&sum.to_bits().to_le_bytes()); + encoded[40..44].copy_from_slice(&payload_len.to_le_bytes()); + Ok(encoded) +} + +fn encode_section( + output: &mut Vec, + buckets: &[(i64, u64)], + descending: bool, +) -> Result<(), String> { + let mut previous = None; + for &(index, count) in buckets { + let encoded_index = match previous { + None => zigzag(index), + Some((previous_index, _)) => { + let delta = if descending { + previous_index as i128 - index as i128 + } else { + index as i128 - previous_index as i128 + }; + u64::try_from(delta) + .ok() + .filter(|delta| *delta != 0) + .ok_or_else(|| "legacy bucket indices are not strictly ordered".to_string())? + } + }; + let encoded_count = match previous { + None => count, + Some((_, previous_count)) => zigzag(count.wrapping_sub(previous_count) as i64), + }; + put_varint(output, encoded_index); + put_varint(output, encoded_count); + previous = Some((index, count)); + } + Ok(()) +} + +fn put_varint(output: &mut Vec, value: u64) { + let mut buffer = [0; 9]; + let len = vu128::encode_u64(&mut buffer, value); + output.extend_from_slice(&buffer[..len]); +} + +const fn zigzag(value: i64) -> u64 { + (value.wrapping_shl(1) ^ (value >> 63)) as u64 +} + +fn key_lt(left: LegacyBucketKey, right: LegacyBucketKey) -> bool { + use LegacyBucketKey::*; + match (left, right) { + (Negative(left), Negative(right)) => left > right, + (Negative(_), Zero | Positive(_)) | (Zero, Positive(_)) => true, + (Positive(left), Positive(right)) => left < right, + _ => false, + } +} + +#[cfg(test)] +pub(crate) const LEGACY_STATE: &[u8] = &[ + 4, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 116, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 116, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 116, + 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 116, 0, 0, 0, 0, 0, 0, 0, + 123, 20, 174, 71, 225, 122, 132, 63, 253, 74, 129, 90, 191, 82, 240, 63, 0, 0, 0, 0, 128, 0, 0, + 0, 0, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 240, 63, 123, 20, 174, 71, 225, 122, + 132, 63, +]; + +#[cfg(test)] +pub(crate) const COMPACTED_LEGACY_SKETCH: &[u8] = &[ + 6, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 255, 255, 255, 255, 255, 255, 255, 255, 36, 0, 0, 0, 0, 0, + 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 36, 0, 0, 0, 0, + 0, 0, 0, 2, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 254, 255, 255, 255, 255, 255, 255, + 255, 29, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 255, 255, 255, 255, 255, 255, 255, 255, 2, 0, 0, 0, + 2, 0, 0, 0, 0, 0, 0, 0, 36, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, + 0, 3, 0, 0, 0, 0, 0, 0, 0, 29, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 2, 0, 0, 0, 1, 0, 0, 0, 0, 0, + 0, 0, 35, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 254, 255, 255, + 255, 255, 255, 255, 255, 0, 0, 0, 0, 0, 0, 240, 63, 169, 137, 186, 120, 1, 63, 82, 71, 12, 0, + 0, 0, 7, 0, 0, 0, 0, 0, 0, 0, 201, 0, 0, 0, 0, 0, 0, 0, 112, 103, 108, 212, 220, 81, 180, 84, +]; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn current_format_is_decoded_without_legacy_fallback() { + let mut sketch = UddSketch::new(128, 0.01).unwrap(); + sketch.add(42.0).unwrap(); + let expected = sketch.quantile(0.5).unwrap(); + let encoded = sketch.encode().unwrap(); + + assert!(decode_legacy_state(&encoded).is_err()); + assert_eq!(decode(&encoded).unwrap(), sketch); + assert_eq!(quantile(&encoded, 0.5).unwrap(), expected); + } + + #[test] + fn legacy_decoder_rejects_oversized_bucket_count_before_deserializing() { + let mut encoded = LEGACY_STATE.to_vec(); + encoded[..8].copy_from_slice(&((MAX_BUCKETS as u64) + 1).to_le_bytes()); + + assert_eq!( + decode_legacy_state(&encoded).unwrap_err(), + "legacy populated bucket count exceeds decode limit" + ); + } + + #[test] + fn legacy_decoder_reads_bare_sketch_from_scalar_callers() { + let bare_sketch = &LEGACY_STATE[..LEGACY_STATE.len() - std::mem::size_of::()]; + + assert_eq!( + quantile(bare_sketch, 0.5).unwrap(), + Some(0.9900000000000001) + ); + } + + #[test] + fn legacy_decoder_reads_compacted_bare_sketch() { + assert!(decode_legacy_state(COMPACTED_LEGACY_SKETCH).is_err()); + assert!(quantile(COMPACTED_LEGACY_SKETCH, 0.5).unwrap().is_some()); + } + + #[test] + fn legacy_quantile_rejects_invalid_mapping() { + let bare_len = LEGACY_STATE.len() - std::mem::size_of::(); + let mut invalid = LEGACY_STATE[..bare_len].to_vec(); + let alpha_offset = bare_len - 44; + invalid[alpha_offset..alpha_offset + 8].copy_from_slice(&f64::NAN.to_le_bytes()); + + assert!(quantile(&invalid, 0.5).is_err()); + + let mut sketch = decode_legacy_sketch(&LEGACY_STATE[..bare_len]).unwrap(); + sketch.gamma = 2.0; + assert!(sketch.quantile(0.5).is_err()); + + let mut sketch = decode_legacy_sketch(&LEGACY_STATE[..bare_len]).unwrap(); + sketch.alpha = 1.0; + sketch.gamma = 2.0; + assert!(sketch.quantile(0.5).is_err()); + + let mut sketch = decode_legacy_sketch(&LEGACY_STATE[..bare_len]).unwrap(); + sketch.alpha = 1e-12; + sketch.gamma = 1.0 + 1e-9; + sketch.compactions = 0; + assert!(sketch.quantile(0.5).is_err()); + + let mut sketch = decode_legacy_sketch(&LEGACY_STATE[..bare_len]).unwrap(); + sketch.alpha = 1.0; + sketch.gamma = 1e20; + sketch.compactions = 0; + assert!(sketch.quantile(0.5).is_err()); + + let mut sketch = decode_legacy_sketch(&LEGACY_STATE[..bare_len]).unwrap(); + sketch.alpha = 1.0; + sketch.gamma = f64::INFINITY; + sketch.compactions = 1; + assert!(validate_current_mapping(sketch.alpha, sketch.gamma, sketch.compactions).is_err()); + + let mut sketch = decode_legacy_sketch(&LEGACY_STATE[..bare_len]).unwrap(); + sketch.alpha = 0.99999999; + sketch.gamma = 2.0; + sketch.compactions = 0; + assert!(sketch.quantile(0.5).is_err()); + + let mut sketch = decode_legacy_sketch(&LEGACY_STATE[..bare_len]).unwrap(); + let LegacyBucketKey::Negative(index) = sketch.buckets.head else { + panic!("Expected negative head bucket"); + }; + let bucket = sketch + .buckets + .map + .remove(&LegacyBucketKey::Negative(index)) + .unwrap(); + sketch + .buckets + .map + .insert(LegacyBucketKey::Negative(i64::MIN), bucket); + sketch.buckets.head = LegacyBucketKey::Negative(i64::MIN); + assert!(sketch.quantile(0.5).is_err()); + } + + #[test] + fn legacy_quantile_accepts_saturated_mapping() { + let sketch = LegacyUddSketch { + buckets: LegacyBucketStore { + map: HashMap::from([( + LegacyBucketKey::Positive(0), + LegacyBucket { + count: 1, + next: LegacyBucketKey::Invalid, + }, + )]), + head: LegacyBucketKey::Positive(0), + }, + alpha: 1.0, + gamma: f64::INFINITY, + compactions: 63, + max_buckets: 7, + count: 1, + sum: f64::INFINITY, + }; + + assert_eq!(sketch.quantile(0.5).unwrap(), Some(0.0)); + } + + #[test] + fn legacy_quantile_handles_maximum_count_at_one() { + let sketch = LegacyUddSketch { + buckets: LegacyBucketStore { + map: HashMap::from([( + LegacyBucketKey::Zero, + LegacyBucket { + count: u64::MAX, + next: LegacyBucketKey::Invalid, + }, + )]), + head: LegacyBucketKey::Zero, + }, + alpha: 0.01, + gamma: (1.0 + 0.01) / (1.0 - 0.01), + compactions: 0, + max_buckets: 128, + count: u64::MAX, + sum: 0.0, + }; + + assert_eq!(sketch.quantile(1.0).unwrap(), Some(0.0)); + } +} diff --git a/tests/compatibility/cases/uddsketch_legacy_state/case.toml b/tests/compatibility/cases/uddsketch_legacy_state/case.toml new file mode 100644 index 0000000000..b22670037d --- /dev/null +++ b/tests/compatibility/cases/uddsketch_legacy_state/case.toml @@ -0,0 +1,9 @@ +name = "uddsketch_legacy_state" +reason = "Verify UDDSketch state bytes persisted by old binaries (wrapped legacy bincode UddSketchState) are read directly by uddsketch_calc and uddsketch_merge on the new binary, merged into canonical v1, and merged together with newly generated canonical states." +introduced_by = "PR #8867" +topologies = ["distributed", "standalone"] +from_range = [">=v1.0.0"] +to_range = ["*"] +features = ["table", "query", "udf", "aggregate"] +owner = "query" +namespace = "uddsketch_legacy_state" diff --git a/tests/compatibility/cases/uddsketch_legacy_state/setup.sql b/tests/compatibility/cases/uddsketch_legacy_state/setup.sql new file mode 100644 index 0000000000..abe681a7b9 --- /dev/null +++ b/tests/compatibility/cases/uddsketch_legacy_state/setup.sql @@ -0,0 +1,31 @@ +CREATE TABLE uddsketch_values ( + seq_id INT PRIMARY KEY, + val DOUBLE, + ts TIMESTAMP TIME INDEX DEFAULT now() +); + +INSERT INTO uddsketch_values (seq_id, val) 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); + +CREATE TABLE uddsketch_states ( + state BINARY, + grp INT PRIMARY KEY, + ts TIMESTAMP TIME INDEX DEFAULT now() +); + +INSERT INTO uddsketch_states (state, grp) +SELECT uddsketch_state(128, 0.01, val), seq_id / 5 * 5 AS grp +FROM uddsketch_values +GROUP BY grp; + +ADMIN FLUSH_TABLE('uddsketch_values'); +ADMIN FLUSH_TABLE('uddsketch_states'); diff --git a/tests/compatibility/cases/uddsketch_legacy_state/verify.result b/tests/compatibility/cases/uddsketch_legacy_state/verify.result new file mode 100644 index 0000000000..b4c8a3d156 --- /dev/null +++ b/tests/compatibility/cases/uddsketch_legacy_state/verify.result @@ -0,0 +1,63 @@ +-- (1) direct uddsketch_calc of each persisted legacy row +SELECT grp, uddsketch_calc(0.5, state) AS p50 +FROM uddsketch_states +ORDER BY grp; + ++-----+--------------------+ +| grp | p50 | ++-----+--------------------+ +| 0 | 30.267171338721894 | +| 5 | 70.11183939140265 | +| 10 | 100.49456770856492 | ++-----+--------------------+ + +-- (2) uddsketch_merge over the persisted legacy rows, then calc. +-- Succeeds only if every legacy row decodes to the exact declared +-- parameters (128, 0.01); the merge parameter check is the guard. +SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 +FROM uddsketch_states; + ++--------------------+ +| p50 | ++--------------------+ +| 59.745049810145126 | ++--------------------+ + +-- (3) generate fresh canonical states with the new binary and merge mixed legacy + canonical rows +CREATE TABLE uddsketch_new_states ( + state BINARY, + grp INT PRIMARY KEY, + ts TIMESTAMP TIME INDEX DEFAULT now() +); + +Affected Rows: 0 + +INSERT INTO uddsketch_new_states (state, grp) +SELECT uddsketch_state(128, 0.01, val), seq_id / 5 * 5 AS grp +FROM uddsketch_values +GROUP BY grp; + +Affected Rows: 3 + +SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 +FROM ( + SELECT state FROM uddsketch_states + UNION ALL + SELECT state FROM uddsketch_new_states +) AS all_states; + ++--------------------+ +| p50 | ++--------------------+ +| 59.745049810145126 | ++--------------------+ + +-- (4) canonical-only control +SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 +FROM uddsketch_new_states; + ++--------------------+ +| p50 | ++--------------------+ +| 59.745049810145126 | ++--------------------+ diff --git a/tests/compatibility/cases/uddsketch_legacy_state/verify.sql b/tests/compatibility/cases/uddsketch_legacy_state/verify.sql new file mode 100644 index 0000000000..8cf84cfbbc --- /dev/null +++ b/tests/compatibility/cases/uddsketch_legacy_state/verify.sql @@ -0,0 +1,33 @@ +-- (1) direct uddsketch_calc of each persisted legacy row +SELECT grp, uddsketch_calc(0.5, state) AS p50 +FROM uddsketch_states +ORDER BY grp; + +-- (2) uddsketch_merge over the persisted legacy rows, then calc. +-- Succeeds only if every legacy row decodes to the exact declared +-- parameters (128, 0.01); the merge parameter check is the guard. +SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 +FROM uddsketch_states; + +-- (3) generate fresh canonical states with the new binary and merge mixed legacy + canonical rows +CREATE TABLE uddsketch_new_states ( + state BINARY, + grp INT PRIMARY KEY, + ts TIMESTAMP TIME INDEX DEFAULT now() +); + +INSERT INTO uddsketch_new_states (state, grp) +SELECT uddsketch_state(128, 0.01, val), seq_id / 5 * 5 AS grp +FROM uddsketch_values +GROUP BY grp; + +SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 +FROM ( + SELECT state FROM uddsketch_states + UNION ALL + SELECT state FROM uddsketch_new_states +) AS all_states; + +-- (4) canonical-only control +SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 +FROM uddsketch_new_states;