mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-09-05 21:18:57 +00:00
feat(flow): add generic delta merge for incremental aggregates (#8938)
* feat(function): add internal delta merge aggregates Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(function): cover delta merge aggregates in SQL Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * feat(function): add Welford delta merge aggregate Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
This commit is contained in:
discord9
@@ -65,6 +65,12 @@ pub fn aggr_merge_func_name(aggr_name: &str) -> String {
|
||||
format!("__{}_merge", aggr_name)
|
||||
}
|
||||
|
||||
/// Returns the globally registered name used to merge a delta state with a
|
||||
/// persisted state.
|
||||
pub fn aggr_delta_merge_func_name(state_aggregate_name: &str) -> String {
|
||||
format!("__{}_delta_merge", state_aggregate_name)
|
||||
}
|
||||
|
||||
/// Check if the given aggregate expression is steppable.
|
||||
/// As in if it can be split into multiple steps:
|
||||
/// i.e. on datanode first call `state(input)` then
|
||||
@@ -685,6 +691,222 @@ impl Accumulator for StateAccum {
|
||||
}
|
||||
}
|
||||
|
||||
/// A globally registerable wrapper for a state-family merge UDAF.
|
||||
///
|
||||
/// The wrapped merge function has the family contract `P..., State`. This
|
||||
/// wrapper exposes `P..., delta_state, persisted_state` and forwards the two
|
||||
/// state columns to the existing accumulator in that order. State values are
|
||||
/// opaque to this adapter; null is the inner merge family's identity value.
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct DeltaMergeWrapper {
|
||||
inner: AggregateUDF,
|
||||
name: String,
|
||||
signature: Signature,
|
||||
inner_types: Vec<DataType>,
|
||||
state_type: DataType,
|
||||
}
|
||||
|
||||
impl DeltaMergeWrapper {
|
||||
/// Build the wrapper for one of the explicitly supported exact merge UDAFs.
|
||||
///
|
||||
/// The caller supplies the known merge signature, so construction cannot
|
||||
/// fail while inspecting an arbitrary UDAF signature.
|
||||
pub(crate) fn new(
|
||||
inner: AggregateUDF,
|
||||
state_name: &str,
|
||||
inner_types: Vec<DataType>,
|
||||
state_type: DataType,
|
||||
) -> Self {
|
||||
let mut wrapper_types = inner_types.clone();
|
||||
wrapper_types.push(state_type.clone());
|
||||
Self {
|
||||
name: aggr_delta_merge_func_name(state_name),
|
||||
signature: Signature::exact(wrapper_types, inner.signature().volatility),
|
||||
inner,
|
||||
inner_types,
|
||||
state_type,
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_inner_args(
|
||||
&self,
|
||||
input_fields: &[FieldRef],
|
||||
) -> datafusion_common::Result<Vec<FieldRef>> {
|
||||
if input_fields.len() != self.inner_types.len() + 1 {
|
||||
return Err(datafusion_common::DataFusionError::Plan(
|
||||
"delta merge requires parameters, delta state, and persisted state".to_string(),
|
||||
));
|
||||
}
|
||||
for (field, expected_type) in input_fields[..self.inner_types.len()]
|
||||
.iter()
|
||||
.zip(&self.inner_types)
|
||||
{
|
||||
if field.data_type() != expected_type {
|
||||
return Err(datafusion_common::DataFusionError::Plan(format!(
|
||||
"delta merge argument type does not match its exact signature: {:?} != {expected_type:?}",
|
||||
field.data_type()
|
||||
)));
|
||||
}
|
||||
}
|
||||
let persisted = &input_fields[self.inner_types.len()];
|
||||
if persisted.data_type() != &self.state_type && persisted.data_type() != &DataType::Null {
|
||||
return Err(datafusion_common::DataFusionError::Plan(format!(
|
||||
"persisted state type does not match the exact state type: {:?} != {:?}",
|
||||
persisted.data_type(),
|
||||
self.state_type
|
||||
)));
|
||||
}
|
||||
Ok(input_fields[..self.inner_types.len()].to_vec())
|
||||
}
|
||||
}
|
||||
|
||||
impl AggregateUDFImpl for DeltaMergeWrapper {
|
||||
fn accumulator<'a, 'b>(
|
||||
&'a self,
|
||||
acc_args: datafusion_expr::function::AccumulatorArgs<'b>,
|
||||
) -> datafusion_common::Result<Box<dyn Accumulator>> {
|
||||
if acc_args.exprs.len() != acc_args.expr_fields.len() {
|
||||
return Err(datafusion_common::DataFusionError::Plan(
|
||||
"delta merge expression and field arities differ".to_string(),
|
||||
));
|
||||
}
|
||||
let inner_fields = self.resolve_inner_args(acc_args.expr_fields)?;
|
||||
for (expr, expected_type) in acc_args.exprs.iter().zip(
|
||||
self.inner_types
|
||||
.iter()
|
||||
.chain(std::iter::once(&self.state_type)),
|
||||
) {
|
||||
if expr.data_type(acc_args.schema)? != *expected_type {
|
||||
return Err(datafusion_common::DataFusionError::Internal(
|
||||
"delta merge physical expression type is not resolved".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
let state_index = self.inner_types.len() - 1;
|
||||
let inner_args = datafusion_expr::function::AccumulatorArgs {
|
||||
return_field: acc_args.return_field,
|
||||
schema: acc_args.schema,
|
||||
ignore_nulls: acc_args.ignore_nulls,
|
||||
order_bys: acc_args.order_bys,
|
||||
is_reversed: acc_args.is_reversed,
|
||||
name: self.inner.name(),
|
||||
is_distinct: acc_args.is_distinct,
|
||||
exprs: &acc_args.exprs[..=state_index],
|
||||
expr_fields: &inner_fields,
|
||||
};
|
||||
Ok(Box::new(DeltaMergeAccum {
|
||||
inner: self.inner.accumulator(inner_args)?,
|
||||
params: state_index,
|
||||
}))
|
||||
}
|
||||
|
||||
fn as_any(&self) -> &dyn std::any::Any {
|
||||
self
|
||||
}
|
||||
|
||||
fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
|
||||
fn is_nullable(&self) -> bool {
|
||||
self.inner.is_nullable()
|
||||
}
|
||||
|
||||
fn return_type(&self, arg_types: &[DataType]) -> datafusion_common::Result<DataType> {
|
||||
let fields = arg_types
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, data_type)| {
|
||||
Arc::new(Field::new(index.to_string(), data_type.clone(), true))
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let inner_fields = self.resolve_inner_args(&fields)?;
|
||||
self.inner.return_type(
|
||||
&inner_fields
|
||||
.iter()
|
||||
.map(|field| field.data_type().clone())
|
||||
.collect::<Vec<_>>(),
|
||||
)
|
||||
}
|
||||
|
||||
fn return_field(&self, arg_fields: &[FieldRef]) -> datafusion_common::Result<FieldRef> {
|
||||
let inner_fields = self.resolve_inner_args(arg_fields)?;
|
||||
self.inner.return_field(&inner_fields)
|
||||
}
|
||||
|
||||
fn signature(&self) -> &Signature {
|
||||
&self.signature
|
||||
}
|
||||
|
||||
fn state_fields(
|
||||
&self,
|
||||
args: datafusion_expr::function::StateFieldsArgs,
|
||||
) -> datafusion_common::Result<Vec<FieldRef>> {
|
||||
let inner_fields = self.resolve_inner_args(args.input_fields)?;
|
||||
self.inner
|
||||
.state_fields(datafusion_expr::function::StateFieldsArgs {
|
||||
name: args.name,
|
||||
input_fields: &inner_fields,
|
||||
return_field: args.return_field,
|
||||
ordering_fields: args.ordering_fields,
|
||||
is_distinct: args.is_distinct,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialEq for DeltaMergeWrapper {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
self.name == other.name && self.inner == other.inner
|
||||
}
|
||||
}
|
||||
impl Eq for DeltaMergeWrapper {}
|
||||
impl Hash for DeltaMergeWrapper {
|
||||
fn hash<H: Hasher>(&self, state: &mut H) {
|
||||
self.name.hash(state);
|
||||
self.inner.hash(state);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct DeltaMergeAccum {
|
||||
inner: Box<dyn Accumulator>,
|
||||
params: usize,
|
||||
}
|
||||
|
||||
impl Accumulator for DeltaMergeAccum {
|
||||
fn evaluate(&mut self) -> datafusion_common::Result<ScalarValue> {
|
||||
self.inner.evaluate()
|
||||
}
|
||||
|
||||
fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion_common::Result<()> {
|
||||
if values.len() != self.params + 2 {
|
||||
return Err(datafusion_common::DataFusionError::Plan(format!(
|
||||
"delta merge expected {} arguments, got {}",
|
||||
self.params + 2,
|
||||
values.len()
|
||||
)));
|
||||
}
|
||||
// Null-state identity is an inner family precondition. The wrapper
|
||||
// forwards opaque states and never decodes or filters them.
|
||||
let mut inner_values = values[..self.params + 1].to_vec();
|
||||
self.inner.update_batch(&inner_values)?;
|
||||
inner_values[self.params] = values[self.params + 1].clone();
|
||||
self.inner.update_batch(&inner_values)
|
||||
}
|
||||
|
||||
fn merge_batch(&mut self, states: &[ArrayRef]) -> datafusion_common::Result<()> {
|
||||
self.inner.merge_batch(states)
|
||||
}
|
||||
|
||||
fn size(&self) -> usize {
|
||||
self.inner.size()
|
||||
}
|
||||
|
||||
fn state(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
|
||||
self.inner.state()
|
||||
}
|
||||
}
|
||||
|
||||
/// TODO(discord9): mark this function as non-ser/de able
|
||||
///
|
||||
/// This wrapper shouldn't be register as a udaf, as it contain extra data that is not serializable.
|
||||
|
||||
@@ -18,7 +18,8 @@ use std::sync::{Arc, Mutex};
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
use arrow::array::{
|
||||
ArrayRef, BooleanArray, Float64Array, Int64Array, TimestampMillisecondArray, UInt64Array,
|
||||
ArrayRef, BinaryArray, BooleanArray, Float64Array, Int64Array, StringArray,
|
||||
TimestampMillisecondArray, UInt64Array,
|
||||
};
|
||||
use arrow::record_batch::RecordBatch;
|
||||
use arrow_schema::SchemaRef;
|
||||
@@ -43,20 +44,23 @@ use datafusion_expr::expr::{AggregateFunction, NullTreatment};
|
||||
use datafusion_expr::function::AccumulatorArgs;
|
||||
use datafusion_expr::{
|
||||
Aggregate, AggregateUDFImpl, ColumnarValue, Expr, LogicalPlan, ScalarFunctionArgs, SortExpr,
|
||||
TableScan, lit,
|
||||
TableScan, TypeSignature, lit,
|
||||
};
|
||||
use datafusion_physical_expr::aggregate::AggregateExprBuilder;
|
||||
use datafusion_physical_expr::expressions::col;
|
||||
use datafusion_physical_expr::expressions::{Column as PhysicalColumn, col, lit as physical_lit};
|
||||
use datafusion_physical_expr::{EquivalenceProperties, Partitioning};
|
||||
use datatypes::arrow_array::StringArray;
|
||||
use futures::{Stream, StreamExt as _};
|
||||
use hyperloglogplus::HyperLogLog;
|
||||
use pretty_assertions::assert_eq;
|
||||
use uddsketch::UddSketchRef;
|
||||
|
||||
use super::*;
|
||||
use crate::aggrs::approximate::hll::HllState;
|
||||
use crate::aggrs::approximate::hll::{HllState, HllStateType};
|
||||
use crate::aggrs::approximate::uddsketch::UddSketchState;
|
||||
use crate::aggrs::approximate::welford::{WelfordAccumulator, WelfordState};
|
||||
use crate::aggrs::count_hash::CountHash;
|
||||
use crate::function::Function as _;
|
||||
use crate::function_registry::FUNCTION_REGISTRY;
|
||||
use crate::scalars::hll_count::HllCalcFunction;
|
||||
use crate::scalars::uddsketch_calc::UddSketchCalcFunction;
|
||||
|
||||
@@ -920,9 +924,204 @@ fn test_avg_state_groups_accumulator_state_merge_evaluate() {
|
||||
);
|
||||
}
|
||||
|
||||
/// For testing whether the UDAF state fields are correctly implemented.
|
||||
/// esp. for our own custom UDAF's state fields.
|
||||
/// By compare eval results before and after split to state/merge functions.
|
||||
#[test]
|
||||
fn test_registered_hll_delta_merge_semantics() {
|
||||
let hll = FUNCTION_REGISTRY
|
||||
.get_aggr_func(&aggr_delta_merge_func_name("hll"))
|
||||
.expect("global approximate functions register hll delta merge");
|
||||
assert_eq!(hll.name(), "__hll_delta_merge");
|
||||
assert_eq!(
|
||||
hll.signature().type_signature,
|
||||
TypeSignature::Exact(vec![DataType::Binary, DataType::Binary])
|
||||
);
|
||||
|
||||
let schema = Arc::new(arrow_schema::Schema::new(vec![
|
||||
Field::new("delta", DataType::Binary, true),
|
||||
Field::new("persisted", DataType::Binary, true),
|
||||
]));
|
||||
let expr = AggregateExprBuilder::new(
|
||||
Arc::new(hll),
|
||||
vec![
|
||||
Arc::new(PhysicalColumn::new("delta", 0)),
|
||||
Arc::new(PhysicalColumn::new("persisted", 1)),
|
||||
],
|
||||
)
|
||||
.schema(schema)
|
||||
.alias("hll_delta_merge")
|
||||
.build()
|
||||
.unwrap();
|
||||
let mut accum = expr.create_accumulator().unwrap();
|
||||
|
||||
let mut delta = HllState::new();
|
||||
delta
|
||||
.update_batch(&[Arc::new(StringArray::from(vec![Some("delta")]))])
|
||||
.unwrap();
|
||||
let mut persisted = HllState::new();
|
||||
persisted
|
||||
.update_batch(&[Arc::new(StringArray::from(vec![Some("persisted")]))])
|
||||
.unwrap();
|
||||
let ScalarValue::Binary(Some(delta)) = delta.evaluate().unwrap() else {
|
||||
panic!("HLL state must be binary");
|
||||
};
|
||||
let ScalarValue::Binary(Some(persisted)) = persisted.evaluate().unwrap() else {
|
||||
panic!("HLL state must be binary");
|
||||
};
|
||||
|
||||
accum
|
||||
.update_batch(&[
|
||||
Arc::new(BinaryArray::from(vec![Some(delta.as_slice())])),
|
||||
Arc::new(BinaryArray::from(vec![Some(persisted.as_slice())])),
|
||||
])
|
||||
.unwrap();
|
||||
let ScalarValue::Binary(Some(merged)) = accum.evaluate().unwrap() else {
|
||||
panic!("HLL delta merge state must be binary");
|
||||
};
|
||||
let mut merged: HllStateType = bincode::deserialize(&merged).unwrap();
|
||||
assert_eq!(merged.count().trunc() as u32, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_registered_welford_delta_merge_semantics() {
|
||||
let welford = FUNCTION_REGISTRY
|
||||
.get_aggr_func(&aggr_delta_merge_func_name("stddev_pop_state"))
|
||||
.expect("global approximate functions register Welford delta merge");
|
||||
assert_eq!(welford.name(), "__stddev_pop_state_delta_merge");
|
||||
assert_eq!(
|
||||
welford.signature().type_signature,
|
||||
TypeSignature::Exact(vec![DataType::Binary, DataType::Binary])
|
||||
);
|
||||
|
||||
let schema = Arc::new(arrow_schema::Schema::new(vec![
|
||||
Field::new("delta", DataType::Binary, true),
|
||||
Field::new("persisted", DataType::Binary, true),
|
||||
]));
|
||||
let expr = AggregateExprBuilder::new(
|
||||
Arc::new(welford),
|
||||
vec![
|
||||
Arc::new(PhysicalColumn::new("delta", 0)),
|
||||
Arc::new(PhysicalColumn::new("persisted", 1)),
|
||||
],
|
||||
)
|
||||
.schema(schema)
|
||||
.alias("stddev_pop_state_delta_merge")
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
fn state(values: &[f64]) -> Vec<u8> {
|
||||
let mut accumulator = WelfordAccumulator::default();
|
||||
accumulator
|
||||
.update_batch(&[Arc::new(Float64Array::from(
|
||||
values.iter().copied().map(Some).collect::<Vec<_>>(),
|
||||
)) as ArrayRef])
|
||||
.unwrap();
|
||||
let ScalarValue::Binary(Some(state)) = accumulator.evaluate().unwrap() else {
|
||||
panic!("Welford state must be binary");
|
||||
};
|
||||
state
|
||||
}
|
||||
|
||||
let delta = state(&[1.0, 2.0]);
|
||||
let persisted = state(&[3.0, 4.0]);
|
||||
let expected = state(&[1.0, 2.0, 3.0, 4.0]);
|
||||
|
||||
for (delta, persisted, expected) in [
|
||||
(Some(delta.as_slice()), None, state(&[1.0, 2.0])),
|
||||
(None, Some(persisted.as_slice()), state(&[3.0, 4.0])),
|
||||
(None, None, state(&[])),
|
||||
(Some(delta.as_slice()), Some(persisted.as_slice()), expected),
|
||||
] {
|
||||
let mut accum = expr.create_accumulator().unwrap();
|
||||
accum
|
||||
.update_batch(&[
|
||||
Arc::new(BinaryArray::from(vec![delta])),
|
||||
Arc::new(BinaryArray::from(vec![persisted])),
|
||||
])
|
||||
.unwrap();
|
||||
let ScalarValue::Binary(Some(merged)) = accum.evaluate().unwrap() else {
|
||||
panic!("Welford delta merge state must be binary");
|
||||
};
|
||||
assert_eq!(
|
||||
WelfordState::decode(&merged).unwrap(),
|
||||
WelfordState::decode(&expected).unwrap()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_registered_uddsketch_delta_merge_semantics() {
|
||||
let uddsketch = FUNCTION_REGISTRY
|
||||
.get_aggr_func(&aggr_delta_merge_func_name("uddsketch_state"))
|
||||
.expect("global approximate functions register uddsketch delta merge");
|
||||
assert_eq!(uddsketch.name(), "__uddsketch_state_delta_merge");
|
||||
assert_eq!(
|
||||
uddsketch.signature().type_signature,
|
||||
TypeSignature::Exact(vec![
|
||||
DataType::Int64,
|
||||
DataType::Float64,
|
||||
DataType::Binary,
|
||||
DataType::Binary,
|
||||
])
|
||||
);
|
||||
|
||||
let schema = Arc::new(arrow_schema::Schema::new(vec![
|
||||
Field::new("delta", DataType::Binary, true),
|
||||
Field::new("persisted", DataType::Binary, true),
|
||||
]));
|
||||
let expr = AggregateExprBuilder::new(
|
||||
Arc::new(uddsketch),
|
||||
vec![
|
||||
physical_lit(10_i64),
|
||||
physical_lit(0.01_f64),
|
||||
Arc::new(PhysicalColumn::new("delta", 0)),
|
||||
Arc::new(PhysicalColumn::new("persisted", 1)),
|
||||
],
|
||||
)
|
||||
.schema(schema)
|
||||
.alias("uddsketch_delta_merge")
|
||||
.build()
|
||||
.unwrap();
|
||||
let mut accum = expr.create_accumulator().unwrap();
|
||||
|
||||
let mut delta = UddSketchState::new(10, 0.01).unwrap();
|
||||
delta
|
||||
.update_batch(&[
|
||||
Arc::new(Int64Array::from(vec![10])) as ArrayRef,
|
||||
Arc::new(Float64Array::from(vec![0.01])) as ArrayRef,
|
||||
Arc::new(Float64Array::from(vec![2.0])) as ArrayRef,
|
||||
])
|
||||
.unwrap();
|
||||
let mut persisted = UddSketchState::new(10, 0.01).unwrap();
|
||||
persisted
|
||||
.update_batch(&[
|
||||
Arc::new(Int64Array::from(vec![10])) as ArrayRef,
|
||||
Arc::new(Float64Array::from(vec![0.01])) as ArrayRef,
|
||||
Arc::new(Float64Array::from(vec![1.0])) as ArrayRef,
|
||||
])
|
||||
.unwrap();
|
||||
let ScalarValue::Binary(Some(delta)) = delta.evaluate().unwrap() else {
|
||||
panic!("UDDSketch state must be binary");
|
||||
};
|
||||
let ScalarValue::Binary(Some(persisted)) = persisted.evaluate().unwrap() else {
|
||||
panic!("UDDSketch state must be binary");
|
||||
};
|
||||
|
||||
accum
|
||||
.update_batch(&[
|
||||
Arc::new(Int64Array::from(vec![10])) as ArrayRef,
|
||||
Arc::new(Float64Array::from(vec![0.01])) as ArrayRef,
|
||||
Arc::new(BinaryArray::from(vec![Some(delta.as_slice())])),
|
||||
Arc::new(BinaryArray::from(vec![Some(persisted.as_slice())])),
|
||||
])
|
||||
.unwrap();
|
||||
let ScalarValue::Binary(Some(merged)) = accum.evaluate().unwrap() else {
|
||||
panic!("UDDSketch delta merge state must be binary");
|
||||
};
|
||||
let merged = UddSketchRef::parse(&merged).unwrap();
|
||||
assert_eq!(merged.count(), 2);
|
||||
let median = merged.quantile(0.5).unwrap().unwrap();
|
||||
assert!((1.0..=2.0).contains(&median));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_udaf_correct_eval_result() {
|
||||
struct TestCase {
|
||||
|
||||
@@ -12,6 +12,10 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use datafusion_expr::AggregateUDF;
|
||||
use datatypes::arrow::datatypes::DataType;
|
||||
|
||||
use crate::aggrs::aggr_wrapper::DeltaMergeWrapper;
|
||||
use crate::function_registry::FunctionRegistry;
|
||||
|
||||
pub mod hll;
|
||||
@@ -22,16 +26,42 @@ pub(crate) struct ApproximateFunction;
|
||||
|
||||
impl ApproximateFunction {
|
||||
pub fn register(registry: &FunctionRegistry) {
|
||||
// uddsketch
|
||||
registry.register_aggr(uddsketch::UddSketchState::state_udf_impl());
|
||||
registry.register_aggr(uddsketch::UddSketchState::merge_udf_impl());
|
||||
let uddsketch_state = uddsketch::UddSketchState::state_udf_impl();
|
||||
let uddsketch_merge = uddsketch::UddSketchState::merge_udf_impl();
|
||||
let uddsketch_delta = AggregateUDF::new_from_impl(DeltaMergeWrapper::new(
|
||||
uddsketch_merge.clone(),
|
||||
uddsketch::UDDSKETCH_STATE_NAME,
|
||||
vec![DataType::Int64, DataType::Float64, DataType::Binary],
|
||||
DataType::Binary,
|
||||
));
|
||||
registry.register_aggr(uddsketch_state);
|
||||
registry.register_aggr(uddsketch_merge);
|
||||
registry.register_aggr(uddsketch_delta);
|
||||
|
||||
// hll
|
||||
registry.register_aggr(hll::HllState::state_udf_impl());
|
||||
registry.register_aggr(hll::HllState::merge_udf_impl());
|
||||
let hll_state = hll::HllState::state_udf_impl();
|
||||
let hll_merge = hll::HllState::merge_udf_impl();
|
||||
let hll_delta = AggregateUDF::new_from_impl(DeltaMergeWrapper::new(
|
||||
hll_merge.clone(),
|
||||
hll::HLL_NAME,
|
||||
vec![DataType::Binary],
|
||||
DataType::Binary,
|
||||
));
|
||||
registry.register_aggr(hll_state);
|
||||
registry.register_aggr(hll_merge);
|
||||
registry.register_aggr(hll_delta);
|
||||
|
||||
// welford
|
||||
registry.register_aggr(welford::WelfordAccumulator::state_udf_impl());
|
||||
registry.register_aggr(welford::WelfordAccumulator::merge_udf_impl());
|
||||
let welford_state = welford::WelfordAccumulator::state_udf_impl();
|
||||
let welford_merge = welford::WelfordAccumulator::merge_udf_impl();
|
||||
let welford_delta = AggregateUDF::new_from_impl(DeltaMergeWrapper::new(
|
||||
welford_merge.clone(),
|
||||
welford::STDDEV_POP_STATE_NAME,
|
||||
vec![DataType::Binary],
|
||||
DataType::Binary,
|
||||
));
|
||||
registry.register_aggr(welford_state);
|
||||
registry.register_aggr(welford_merge);
|
||||
registry.register_aggr(welford_delta);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -191,6 +191,11 @@ impl FunctionRegistry {
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Returns a registered aggregate function by name.
|
||||
pub fn get_aggr_func(&self, name: &str) -> Option<AggregateUDF> {
|
||||
self.aggregate_functions.read().unwrap().get(name).cloned()
|
||||
}
|
||||
|
||||
/// Returns true if an aggregate function with the given name exists in the registry.
|
||||
pub fn is_aggr_func_exist(&self, name: &str) -> bool {
|
||||
self.aggregate_functions.read().unwrap().contains_key(name)
|
||||
|
||||
@@ -74,6 +74,36 @@ select hll_count(hll_merge(`state`)) from test_hll_merge;
|
||||
| 12 |
|
||||
+--------------------------------------------+
|
||||
|
||||
CREATE TABLE test_hll_delta_merge (
|
||||
`id` INT PRIMARY KEY,
|
||||
`delta` BINARY,
|
||||
`persisted` BINARY,
|
||||
`ts` timestamp time index default now()
|
||||
);
|
||||
|
||||
Affected Rows: 0
|
||||
|
||||
INSERT INTO test_hll_delta_merge (`id`, `delta`, `persisted`)
|
||||
SELECT 1, hll(`value`), NULL FROM test_hll WHERE `id` <= 15;
|
||||
|
||||
Affected Rows: 1
|
||||
|
||||
INSERT INTO test_hll_delta_merge (`id`, `delta`, `persisted`)
|
||||
SELECT 2, NULL, hll(`value`) FROM test_hll WHERE `id` > 15;
|
||||
|
||||
Affected Rows: 1
|
||||
|
||||
SELECT
|
||||
hll_count(__hll_delta_merge(`delta`, `persisted`)) AS merged_count,
|
||||
(SELECT hll_count(hll(`value`)) FROM test_hll) AS direct_count
|
||||
FROM test_hll_delta_merge;
|
||||
|
||||
+--------------+--------------+
|
||||
| merged_count | direct_count |
|
||||
+--------------+--------------+
|
||||
| 12 | 12 |
|
||||
+--------------+--------------+
|
||||
|
||||
drop table test_hll;
|
||||
|
||||
Affected Rows: 0
|
||||
@@ -82,3 +112,7 @@ drop table test_hll_merge;
|
||||
|
||||
Affected Rows: 0
|
||||
|
||||
drop table test_hll_delta_merge;
|
||||
|
||||
Affected Rows: 0
|
||||
|
||||
|
||||
@@ -44,6 +44,26 @@ select 2, hll(`value`) from test_hll;
|
||||
|
||||
select hll_count(hll_merge(`state`)) from test_hll_merge;
|
||||
|
||||
CREATE TABLE test_hll_delta_merge (
|
||||
`id` INT PRIMARY KEY,
|
||||
`delta` BINARY,
|
||||
`persisted` BINARY,
|
||||
`ts` timestamp time index default now()
|
||||
);
|
||||
|
||||
INSERT INTO test_hll_delta_merge (`id`, `delta`, `persisted`)
|
||||
SELECT 1, hll(`value`), NULL FROM test_hll WHERE `id` <= 15;
|
||||
|
||||
INSERT INTO test_hll_delta_merge (`id`, `delta`, `persisted`)
|
||||
SELECT 2, NULL, hll(`value`) FROM test_hll WHERE `id` > 15;
|
||||
|
||||
SELECT
|
||||
hll_count(__hll_delta_merge(`delta`, `persisted`)) AS merged_count,
|
||||
(SELECT hll_count(hll(`value`)) FROM test_hll) AS direct_count
|
||||
FROM test_hll_delta_merge;
|
||||
|
||||
drop table test_hll;
|
||||
|
||||
drop table test_hll_merge;
|
||||
|
||||
drop table test_hll_delta_merge;
|
||||
|
||||
@@ -98,6 +98,38 @@ FROM grouped_uddsketch;
|
||||
| 0.5 |
|
||||
+-------------+
|
||||
|
||||
CREATE TABLE delta_uddsketch (
|
||||
`id` INT PRIMARY KEY,
|
||||
`delta` BINARY,
|
||||
`persisted` BINARY,
|
||||
`ts` timestamp time index default now()
|
||||
);
|
||||
|
||||
Affected Rows: 0
|
||||
|
||||
INSERT INTO delta_uddsketch (`id`, `delta`, `persisted`)
|
||||
SELECT 1, uddsketch_state(128, 0.01, `value`), NULL
|
||||
FROM test_uddsketch WHERE `id` <= 5;
|
||||
|
||||
Affected Rows: 1
|
||||
|
||||
INSERT INTO delta_uddsketch (`id`, `delta`, `persisted`)
|
||||
SELECT 2, NULL, uddsketch_state(128, 0.01, `value`)
|
||||
FROM test_uddsketch WHERE `id` > 5;
|
||||
|
||||
Affected Rows: 1
|
||||
|
||||
SELECT
|
||||
uddsketch_calc(0.5, __uddsketch_state_delta_merge(128, 0.01, `delta`, `persisted`)) AS merged_p50,
|
||||
(SELECT uddsketch_calc(0.5, uddsketch_state(128, 0.01, `value`)) FROM test_uddsketch) AS direct_p50
|
||||
FROM delta_uddsketch;
|
||||
|
||||
+--------------------+--------------------+
|
||||
| merged_p50 | direct_p50 |
|
||||
+--------------------+--------------------+
|
||||
| 59.745049810145126 | 59.745049810145126 |
|
||||
+--------------------+--------------------+
|
||||
|
||||
-- should fail
|
||||
SELECT uddsketch_calc(0.1, uddsketch_merge(128, 0.1, `state`)) FROM grouped_uddsketch;
|
||||
|
||||
@@ -126,3 +158,7 @@ drop table grouped_uddsketch;
|
||||
|
||||
Affected Rows: 0
|
||||
|
||||
drop table delta_uddsketch;
|
||||
|
||||
Affected Rows: 0
|
||||
|
||||
|
||||
@@ -48,6 +48,26 @@ SELECT uddsketch_calc(0.1, uddsketch_merge(128, 0.01, `state`)) FROM grouped_udd
|
||||
SELECT ROUND(uddsketch_rank(55, uddsketch_merge(128, 0.01, `state`)), 4) AS merged_rank
|
||||
FROM grouped_uddsketch;
|
||||
|
||||
CREATE TABLE delta_uddsketch (
|
||||
`id` INT PRIMARY KEY,
|
||||
`delta` BINARY,
|
||||
`persisted` BINARY,
|
||||
`ts` timestamp time index default now()
|
||||
);
|
||||
|
||||
INSERT INTO delta_uddsketch (`id`, `delta`, `persisted`)
|
||||
SELECT 1, uddsketch_state(128, 0.01, `value`), NULL
|
||||
FROM test_uddsketch WHERE `id` <= 5;
|
||||
|
||||
INSERT INTO delta_uddsketch (`id`, `delta`, `persisted`)
|
||||
SELECT 2, NULL, uddsketch_state(128, 0.01, `value`)
|
||||
FROM test_uddsketch WHERE `id` > 5;
|
||||
|
||||
SELECT
|
||||
uddsketch_calc(0.5, __uddsketch_state_delta_merge(128, 0.01, `delta`, `persisted`)) AS merged_p50,
|
||||
(SELECT uddsketch_calc(0.5, uddsketch_state(128, 0.01, `value`)) FROM test_uddsketch) AS direct_p50
|
||||
FROM delta_uddsketch;
|
||||
|
||||
-- should fail
|
||||
SELECT uddsketch_calc(0.1, uddsketch_merge(128, 0.1, `state`)) FROM grouped_uddsketch;
|
||||
|
||||
@@ -61,3 +81,4 @@ WHERE `id` < 0;
|
||||
drop table test_uddsketch;
|
||||
|
||||
drop table grouped_uddsketch;
|
||||
drop table delta_uddsketch;
|
||||
|
||||
@@ -87,6 +87,64 @@ FROM (
|
||||
|
||||
Error: 1001(Unsupported), This feature is not implemented: Welford DISTINCT aggregations are not available
|
||||
|
||||
CREATE TABLE welford_delta_merge (
|
||||
`id` INT PRIMARY KEY,
|
||||
`delta` BINARY,
|
||||
`persisted` BINARY,
|
||||
`ts` TIMESTAMP TIME INDEX DEFAULT now()
|
||||
);
|
||||
|
||||
Affected Rows: 0
|
||||
|
||||
INSERT INTO welford_delta_merge (`id`, `delta`, `persisted`)
|
||||
SELECT 1, stddev_pop_state(`value`), NULL
|
||||
FROM test_welford WHERE id <= 5;
|
||||
|
||||
Affected Rows: 1
|
||||
|
||||
INSERT INTO welford_delta_merge (`id`, `delta`, `persisted`)
|
||||
SELECT 2, NULL, stddev_pop_state(`value`)
|
||||
FROM test_welford WHERE id > 5;
|
||||
|
||||
Affected Rows: 1
|
||||
|
||||
INSERT INTO welford_delta_merge (`id`, `delta`, `persisted`)
|
||||
VALUES (3, NULL, NULL);
|
||||
|
||||
Affected Rows: 1
|
||||
|
||||
-- Delta-only, persisted-only, and both-null (empty-state) identity behavior.
|
||||
SELECT
|
||||
`id`,
|
||||
stddev_pop_calc(__stddev_pop_state_delta_merge(`delta`, `persisted`)) AS merged_stddev
|
||||
FROM welford_delta_merge
|
||||
GROUP BY `id`
|
||||
ORDER BY `id`;
|
||||
|
||||
+----+--------------------+
|
||||
| id | merged_stddev |
|
||||
+----+--------------------+
|
||||
| 1 | 14.142135623730951 |
|
||||
| 2 | 14.142135623730951 |
|
||||
| 3 | |
|
||||
+----+--------------------+
|
||||
|
||||
-- Merging delta and persisted states must equal direct stddev_pop state calculation.
|
||||
SELECT
|
||||
stddev_pop_calc(__stddev_pop_state_delta_merge(`delta`, `persisted`)) AS merged_stddev,
|
||||
(SELECT stddev_pop_calc(stddev_pop_state(`value`)) FROM test_welford) AS direct_stddev
|
||||
FROM welford_delta_merge;
|
||||
|
||||
+--------------------+--------------------+
|
||||
| merged_stddev | direct_stddev |
|
||||
+--------------------+--------------------+
|
||||
| 28.722813232690143 | 28.722813232690143 |
|
||||
+--------------------+--------------------+
|
||||
|
||||
DROP TABLE welford_delta_merge;
|
||||
|
||||
Affected Rows: 0
|
||||
|
||||
DROP TABLE grouped_welford;
|
||||
|
||||
Affected Rows: 0
|
||||
|
||||
@@ -55,6 +55,39 @@ FROM (
|
||||
SELECT 3 AS `id`, `state` FROM grouped_welford WHERE id = 1
|
||||
) AS duplicated_states;
|
||||
|
||||
CREATE TABLE welford_delta_merge (
|
||||
`id` INT PRIMARY KEY,
|
||||
`delta` BINARY,
|
||||
`persisted` BINARY,
|
||||
`ts` TIMESTAMP TIME INDEX DEFAULT now()
|
||||
);
|
||||
|
||||
INSERT INTO welford_delta_merge (`id`, `delta`, `persisted`)
|
||||
SELECT 1, stddev_pop_state(`value`), NULL
|
||||
FROM test_welford WHERE id <= 5;
|
||||
|
||||
INSERT INTO welford_delta_merge (`id`, `delta`, `persisted`)
|
||||
SELECT 2, NULL, stddev_pop_state(`value`)
|
||||
FROM test_welford WHERE id > 5;
|
||||
|
||||
INSERT INTO welford_delta_merge (`id`, `delta`, `persisted`)
|
||||
VALUES (3, NULL, NULL);
|
||||
|
||||
-- Delta-only, persisted-only, and both-null (empty-state) identity behavior.
|
||||
SELECT
|
||||
`id`,
|
||||
stddev_pop_calc(__stddev_pop_state_delta_merge(`delta`, `persisted`)) AS merged_stddev
|
||||
FROM welford_delta_merge
|
||||
GROUP BY `id`
|
||||
ORDER BY `id`;
|
||||
|
||||
-- Merging delta and persisted states must equal direct stddev_pop state calculation.
|
||||
SELECT
|
||||
stddev_pop_calc(__stddev_pop_state_delta_merge(`delta`, `persisted`)) AS merged_stddev,
|
||||
(SELECT stddev_pop_calc(stddev_pop_state(`value`)) FROM test_welford) AS direct_stddev
|
||||
FROM welford_delta_merge;
|
||||
|
||||
DROP TABLE welford_delta_merge;
|
||||
DROP TABLE grouped_welford;
|
||||
DROP TABLE test_welford;
|
||||
|
||||
|
||||
Reference in New Issue
Block a user