rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-09-06 13:38:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(udaf): replace uddsketch implementation (#8867)

Browse Source 
* refactor(function): replace uddsketch implementation

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

* bench(function): compare uddsketch batch ingestion

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

* perf(function): avoid copying non-null uddsketch batches

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

* fix(function): decode legacy uddsketch states

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

* fix(function): harden legacy uddsketch validation

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

* format: taplo

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

* test: add compatibility tests for uddsketch functions

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

---------

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>
This commit is contained in:
Lei, HUANG
2026-08-14 12:20:04 +00:00
committed by GitHub
parent cb30837cd5
commit a8924bb95c
12 changed files with 1235 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/common/function/Cargo.toml
+7 -1
View File
@@ -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
src/common/function/benches/uddsketch.rs
+201
View File
@@ -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<Arc<dyn PhysicalExpr>>,
expr_fields: Vec<Arc<Field>>,
return_field: Arc<Field>,
}
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<Arc<dyn PhysicalExpr>> = 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::<Vec<_>>();
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<dyn Accumulator> {
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<ArrayRef> {
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::<Vec<_>>();
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::<Vec<_>>();
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);
src/common/function/src/aggrs/approximate/uddsketch.rs
+185 -114
View File
@@ -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<f64>,
}
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<Self> {
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::<Self>(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::<Literal>()
.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::<Float64Type>(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<ScalarValue> {
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::<f64>() * 3 + // alpha, gamma, values_sum
std::mem::size_of::<u32>() + // compactions
std::mem::size_of::<u64>() * 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::<SketchHashKey>() + // key
std::mem::size_of::<u64>() + // count
std::mem::size_of::<SketchHashKey>(); // next
total_size += self.uddsketch.current_buckets_count() * bucket_entry_size;
total_size
std::mem::size_of::<Self>() - std::mem::size_of::<UddSketch>()
+ self.uddsketch.allocated_size()
+ self.workspace.allocated_size()
+ self.values.capacity() * std::mem::size_of::<f64>()
}
fn state(&mut self) -> DfResult<Vec<ScalarValue>> {
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"));
}
}
src/common/function/src/lib.rs
+1
View File
@@ -18,6 +18,7 @@ mod admin;
mod flush_flow;
mod macros;
mod system;
mod uddsketch_compat;
pub mod aggrs;
pub mod function;
src/common/function/src/scalars/uddsketch_calc.rs
+70 -29
View File
@@ -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::<Float64Type>();
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::<Float64Type>();
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::<Float64Type>().null_count(), 3);
}
}
src/common/function/src/uddsketch_compat.rs
+619
View File
@@ -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<LegacyBucketKey, LegacyBucket>,
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<UddSketch, String> {
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<Option<f64>, 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<UddSketch, String> {
validate_legacy_input(raw)?;
let state = legacy_options()
.deserialize::<LegacyUddSketchState>(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<LegacyUddSketch, String> {
validate_legacy_input(raw)?;
legacy_options()
.deserialize::<LegacyUddSketchState>(raw)
.map(|state| state.uddsketch)
.or_else(|state_error| {
legacy_options()
.deserialize::<LegacyUddSketch>(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<Vec<u8>, 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<Option<f64>, 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<f64, String> {
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<Vec<(LegacyBucketKey, u64)>, 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<Vec<u8>, String> {
let negative = buckets
.iter()
.filter_map(|(key, count)| match key {
LegacyBucketKey::Negative(index) => Some((*index, *count)),
_ => None,
})
.collect::<Vec<_>>();
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::<Vec<_>>();
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<u8>,
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<u8>, 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::<f64>()];
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::<f64>();
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));
}
}