From 283f5d91d104577e49faedcc736f3a6cfa5b6108 Mon Sep 17 00:00:00 2001 From: shuiyisong Date: Mon, 6 Jul 2026 16:11:33 +0800 Subject: [PATCH] feat: implement native histogram read Signed-off-by: shuiyisong --- Cargo.lock | 1 + src/common/query/src/lib.rs | 39 +- src/common/query/src/native_histogram.rs | 1464 +++++++++++- src/common/query/src/promql_annotations.rs | 170 ++ src/promql/Cargo.toml | 1 + src/promql/src/functions.rs | 13 + src/promql/src/functions/native_histogram.rs | 2061 +++++++++++++++++ src/query/src/datafusion.rs | 19 +- src/query/src/planner.rs | 24 +- src/query/src/promql/planner.rs | 691 +++++- src/servers/src/http/prometheus.rs | 14 + .../src/http/result/prometheus_resp.rs | 448 +++- src/servers/src/prom_remote_write/v2.rs | 402 +++- tests-integration/tests/grpc.rs | 5 + tests-integration/tests/http.rs | 51 +- 15 files changed, 5299 insertions(+), 104 deletions(-) create mode 100644 src/common/query/src/promql_annotations.rs create mode 100644 src/promql/src/functions/native_histogram.rs diff --git a/Cargo.lock b/Cargo.lock index 62e91c0de85..d38680c11b6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11021,6 +11021,7 @@ dependencies = [ "bytemuck", "common-error", "common-macro", + "common-query", "common-recordbatch", "common-telemetry", "criterion 0.7.0", diff --git a/src/common/query/src/lib.rs b/src/common/query/src/lib.rs index bdba49eaf70..4028bc690f4 100644 --- a/src/common/query/src/lib.rs +++ b/src/common/query/src/lib.rs @@ -17,6 +17,7 @@ pub mod error; pub mod logical_plan; pub mod native_histogram; pub mod prelude; +pub mod promql_annotations; pub mod request; pub mod stream; #[cfg(any(test, feature = "testing"))] @@ -32,6 +33,8 @@ use datafusion::physical_plan::ExecutionPlan; use serde::{Deserialize, Serialize}; use sqlparser_derive::{Visit, VisitMut}; +use crate::promql_annotations::PromqlAnnotationCollector; + /// new Output struct with output data(previously Output) and output meta #[derive(Debug)] pub struct Output { @@ -69,6 +72,9 @@ pub struct OutputMeta { /// May exist for query output. One can retrieve execution metrics from this plan. pub plan: Option>, pub cost: OutputCost, + pub warnings: Vec, + pub infos: Vec, + pub promql_annotations: Option, } impl Output { @@ -121,18 +127,47 @@ impl Debug for OutputData { impl OutputMeta { pub fn new(plan: Option>, cost: usize) -> Self { - Self { plan, cost } + Self { + plan, + cost, + warnings: Vec::new(), + infos: Vec::new(), + promql_annotations: None, + } } pub fn new_with_plan(plan: Arc) -> Self { Self { plan: Some(plan), cost: 0, + warnings: Vec::new(), + infos: Vec::new(), + promql_annotations: None, } } pub fn new_with_cost(cost: usize) -> Self { - Self { plan: None, cost } + Self { + plan: None, + cost, + warnings: Vec::new(), + infos: Vec::new(), + promql_annotations: None, + } + } + + pub fn with_promql_annotations( + mut self, + promql_annotations: Option, + ) -> Self { + self.promql_annotations = promql_annotations; + self + } + + pub fn collect_promql_annotations(&mut self) { + if let Some(collector) = &self.promql_annotations { + collector.append_to(&mut self.warnings, &mut self.infos); + } } } diff --git a/src/common/query/src/native_histogram.rs b/src/common/query/src/native_histogram.rs index b10c78104a5..3520c64505c 100644 --- a/src/common/query/src/native_histogram.rs +++ b/src/common/query/src/native_histogram.rs @@ -17,9 +17,20 @@ //! Prom remote-write v2 stores these names as children of one Struct field, while //! metric-engine uses the same contract to recognize native histogram tables. +use std::collections::BTreeMap; use std::sync::Arc; -use datatypes::data_type::ConcreteDataType; +use datafusion::arrow::array::{ + Array, ArrayRef, Float64Array, Int32Array, ListArray, PrimitiveArray, StructArray, + TimestampMillisecondArray, UInt64Array, +}; +use datafusion::arrow::buffer::NullBuffer; +use datafusion::arrow::datatypes::{ + ArrowPrimitiveType, DataType as ArrowDataType, Float64Type, Int32Type, Int64Type, + TimestampMillisecondType, UInt32Type, UInt64Type, +}; +use datafusion_common::{DataFusionError, Result as DfResult}; +use datatypes::data_type::{ConcreteDataType, DataType}; use datatypes::types::{StructField, StructType}; use once_cell::sync::Lazy; @@ -106,6 +117,1455 @@ pub fn native_histogram_value_type() -> &'static ConcreteDataType { &NATIVE_HISTOGRAM_VALUE_TYPE } +pub fn is_native_histogram_value_type(data_type: &ConcreteDataType) -> bool { + data_type == &*NATIVE_HISTOGRAM_VALUE_TYPE +} + pub fn is_native_histogram_value_schema(name: &str, data_type: &ConcreteDataType) -> bool { - name == NATIVE_HISTOGRAM_FIELD && data_type == native_histogram_value_type() + name == NATIVE_HISTOGRAM_FIELD && is_native_histogram_value_type(data_type) +} + +pub const CUSTOM_BUCKETS_SCHEMA: i32 = -53; +const MIN_EXPONENTIAL_SCHEMA: i32 = -4; +const MAX_EXPONENTIAL_SCHEMA: i32 = 8; +pub const UNKNOWN_COUNTER_RESET_HINT: i32 = 0; +pub const COUNTER_RESET_HINT: i32 = 1; +pub const NOT_COUNTER_RESET_HINT: i32 = 2; +pub const GAUGE_RESET_HINT: i32 = 3; + +#[derive(Clone, Debug, PartialEq)] +pub struct Span { + pub offset: i32, + pub length: u32, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum BoundaryRule { + OpenLeft = 0, + OpenRight = 1, + OpenBoth = 2, + ClosedBoth = 3, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct Bucket { + pub lower: f64, + pub upper: f64, + pub count: f64, + pub boundary_rule: BoundaryRule, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct NativeHistogram { + pub schema: i32, + pub zero_threshold: f64, + pub sum: f64, + pub reset_hint: i32, + pub start_timestamp: Option, + pub custom_values: Vec, + pub positive_spans: Vec, + pub negative_spans: Vec, + pub count: f64, + pub zero_count: f64, + pub positive_buckets: Vec, + pub negative_buckets: Vec, +} + +impl NativeHistogram { + pub fn uses_custom_buckets(&self) -> bool { + self.schema == CUSTOM_BUCKETS_SCHEMA + } + + pub fn compatible_with(&self, other: &Self) -> bool { + self.schema == other.schema + && self.zero_threshold == other.zero_threshold + && self.custom_values == other.custom_values + } + + pub fn zero_like(&self) -> Self { + let mut result = self.clone(); + result.count = 0.0; + result.zero_count = 0.0; + result.sum = 0.0; + result.positive_buckets.fill(0.0); + result.negative_buckets.fill(0.0); + result + } + + pub fn add_exact(&self, other: &Self) -> Option { + if !self.compatible_with(other) { + return None; + } + + let mut result = self.clone(); + result.count += other.count; + result.zero_count += other.zero_count; + result.sum += other.sum; + result.reset_hint = add_reset_hint(self.reset_hint, other.reset_hint); + (result.positive_spans, result.positive_buckets) = merge_side( + &self.positive_spans, + &self.positive_buckets, + &other.positive_spans, + &other.positive_buckets, + |left, right| left + right, + )?; + (result.negative_spans, result.negative_buckets) = merge_side( + &self.negative_spans, + &self.negative_buckets, + &other.negative_spans, + &other.negative_buckets, + |left, right| left + right, + )?; + Some(result) + } + + pub fn sub_exact(&self, other: &Self) -> Option { + if !self.compatible_with(other) { + return None; + } + + let mut result = self.clone(); + result.count -= other.count; + result.zero_count -= other.zero_count; + result.sum -= other.sum; + result.reset_hint = GAUGE_RESET_HINT; + (result.positive_spans, result.positive_buckets) = merge_side( + &self.positive_spans, + &self.positive_buckets, + &other.positive_spans, + &other.positive_buckets, + |left, right| left - right, + )?; + (result.negative_spans, result.negative_buckets) = merge_side( + &self.negative_spans, + &self.negative_buckets, + &other.negative_spans, + &other.negative_buckets, + |left, right| left - right, + )?; + Some(result) + } + + pub fn add(&self, other: &Self) -> Option { + if self.is_empty_payload() { + return Some(other.clone().compact()); + } + if other.is_empty_payload() { + return Some(self.clone().compact()); + } + + let (left, right) = self.reconcile(other)?; + left.add_exact(&right).map(Self::compact) + } + + pub fn sub(&self, other: &Self) -> Option { + if self.is_empty_payload() { + return Some(other.clone().negated().compact()); + } + if other.is_empty_payload() { + return Some(self.clone().into_gauge().compact()); + } + + let (left, right) = self.reconcile(other)?; + left.sub_exact(&right).map(Self::compact) + } + + pub fn negated(self) -> Self { + self.scale(-1.0) + } + + pub fn into_gauge(mut self) -> Self { + self.reset_hint = GAUGE_RESET_HINT; + self + } + + pub fn counter_reset_hints_contradict(&self, other: &Self) -> bool { + matches!( + (self.reset_hint, other.reset_hint), + (COUNTER_RESET_HINT, NOT_COUNTER_RESET_HINT) + | (NOT_COUNTER_RESET_HINT, COUNTER_RESET_HINT) + ) + } + + pub fn needs_custom_reconciliation(&self, other: &Self) -> bool { + self.uses_custom_buckets() + && other.uses_custom_buckets() + && self.custom_values != other.custom_values + } + + pub fn promql_eq(&self, other: &Self) -> bool { + self.schema == other.schema + && self.zero_threshold == other.zero_threshold + && self.custom_values == other.custom_values + && self.count == other.count + && self.zero_count == other.zero_count + && self.sum == other.sum + && side_counts(&self.positive_spans, &self.positive_buckets) + == side_counts(&other.positive_spans, &other.positive_buckets) + && side_counts(&self.negative_spans, &self.negative_buckets) + == side_counts(&other.negative_spans, &other.negative_buckets) + } + + pub fn promql_string(&self) -> String { + let mut parts = vec![format!("count:{}", self.count), format!("sum:{}", self.sum)]; + if let Some(buckets) = self.all_buckets() { + parts.extend( + buckets + .into_iter() + .filter(|bucket| bucket.count != 0.0) + .map(|bucket| { + let (left, right) = match bucket.boundary_rule { + BoundaryRule::OpenLeft => ("(", "]"), + BoundaryRule::OpenRight => ("[", ")"), + BoundaryRule::OpenBoth => ("(", ")"), + BoundaryRule::ClosedBoth => ("[", "]"), + }; + format!( + "{}{},{}{}:{}", + left, bucket.lower, bucket.upper, right, bucket.count + ) + }), + ); + } + format!("{{{}}}", parts.join(", ")) + } + + pub fn estimated_stdvar(&self) -> f64 { + if self.count == 0.0 { + return f64::NAN; + } + let mean = self.sum / self.count; + let Some(buckets) = self.all_buckets() else { + return f64::NAN; + }; + buckets + .into_iter() + .map(|bucket| { + let midpoint = self.bucket_midpoint(&bucket); + bucket.count * (midpoint - mean).powi(2) + }) + .sum::() + / self.count + } + + pub fn estimated_stddev(&self) -> f64 { + self.estimated_stdvar().sqrt() + } + + pub fn scale(mut self, factor: f64) -> Self { + self.count *= factor; + self.zero_count *= factor; + self.sum *= factor; + for count in &mut self.positive_buckets { + *count *= factor; + } + for count in &mut self.negative_buckets { + *count *= factor; + } + if factor < 0.0 { + self.reset_hint = GAUGE_RESET_HINT; + } + self + } + + pub fn divide_by(mut self, divisor: f64) -> Self { + self.count /= divisor; + self.zero_count /= divisor; + self.sum /= divisor; + if divisor == 0.0 { + self.positive_spans.clear(); + self.positive_buckets.clear(); + self.negative_spans.clear(); + self.negative_buckets.clear(); + } else { + for count in &mut self.positive_buckets { + *count /= divisor; + } + for count in &mut self.negative_buckets { + *count /= divisor; + } + } + if divisor < 0.0 { + self.reset_hint = GAUGE_RESET_HINT; + } + self + } + + pub fn compact(mut self) -> Self { + let (spans, buckets) = compact_side(&self.positive_spans, &self.positive_buckets); + self.positive_spans = spans; + self.positive_buckets = buckets; + let (spans, buckets) = compact_side(&self.negative_spans, &self.negative_buckets); + self.negative_spans = spans; + self.negative_buckets = buckets; + self + } + + pub fn detect_reset(&self, previous: &Self) -> bool { + match self.reset_hint { + COUNTER_RESET_HINT => return true, + NOT_COUNTER_RESET_HINT => return false, + _ => {} + } + + let Some((current, previous)) = self.reconcile(previous) else { + return true; + }; + if current.count < previous.count || current.zero_count < previous.zero_count { + return true; + } + current.side_has_reset(true, &previous) || current.side_has_reset(false, &previous) + } + + pub fn detect_start_timestamp_reset( + &self, + previous: &Self, + previous_ts: i64, + current_ts: i64, + ) -> bool { + let current_start = self.start_timestamp.unwrap_or_default(); + if current_start == 0 || current_start >= current_ts || current_start < previous_ts { + return false; + } + if current_start > previous_ts { + return true; + } + + let previous_start = previous.start_timestamp.unwrap_or_default(); + previous_start <= previous_ts && previous_start != 0 && previous_start != previous_ts + } + + pub fn detect_counter_reset(&self, previous: &Self, previous_ts: i64, current_ts: i64) -> bool { + self.detect_start_timestamp_reset(previous, previous_ts, current_ts) + || self.detect_reset(previous) + } + + pub fn all_buckets(&self) -> Option> { + let mut buckets = self.side_buckets(false)?; + buckets.reverse(); + if self.zero_count != 0.0 { + buckets.push(Bucket { + lower: -self.zero_threshold, + upper: self.zero_threshold, + count: self.zero_count, + boundary_rule: BoundaryRule::ClosedBoth, + }); + } + buckets.extend(self.side_buckets(true)?); + Some(buckets) + } + + pub fn side_buckets(&self, positive: bool) -> Option> { + let (spans, counts) = if positive { + (&self.positive_spans, &self.positive_buckets) + } else { + (&self.negative_spans, &self.negative_buckets) + }; + + let mut result = Vec::with_capacity(counts.len()); + for (idx, count) in side_bucket_indices(spans)?.into_iter().zip(counts) { + let upper = get_bound(idx, self.schema, &self.custom_values)?; + let lower = get_bound(idx - 1, self.schema, &self.custom_values)?; + if positive { + result.push(Bucket { + lower, + upper, + count: *count, + boundary_rule: BoundaryRule::OpenLeft, + }); + } else { + result.push(Bucket { + lower: -upper, + upper: -lower, + count: *count, + boundary_rule: BoundaryRule::OpenRight, + }); + } + } + Some(result) + } + + pub fn quantile(&self, q: f64) -> f64 { + if q < 0.0 { + return f64::NEG_INFINITY; + } + if q > 1.0 { + return f64::INFINITY; + } + if self.count == 0.0 || q.is_nan() { + return f64::NAN; + } + + let Some(mut buckets) = self.all_buckets() else { + return f64::NAN; + }; + let rank = q * self.count; + let mut count = 0.0; + for bucket in &mut buckets { + if bucket.count == 0.0 { + continue; + } + count += bucket.count; + if count < rank { + continue; + } + + if !self.uses_custom_buckets() && bucket.lower < 0.0 && bucket.upper > 0.0 { + if self.negative_buckets.is_empty() && !self.positive_buckets.is_empty() { + bucket.lower = 0.0; + } else if self.positive_buckets.is_empty() && !self.negative_buckets.is_empty() { + bucket.upper = 0.0; + } + } else if self.uses_custom_buckets() { + if bucket.lower == f64::NEG_INFINITY { + if bucket.upper <= 0.0 { + return bucket.upper; + } + bucket.lower = 0.0; + } else if bucket.upper == f64::INFINITY { + return bucket.lower; + } + } + + let rank_in_bucket = rank - (count - bucket.count); + let fraction = rank_in_bucket / bucket.count; + if self.uses_custom_buckets() || (bucket.lower <= 0.0 && bucket.upper >= 0.0) { + return bucket.lower + (bucket.upper - bucket.lower) * fraction; + } + + let log_lower = bucket.lower.abs().log2(); + let log_upper = bucket.upper.abs().log2(); + if bucket.lower > 0.0 { + return 2.0_f64.powf(log_lower + (log_upper - log_lower) * fraction); + } + return -2.0_f64.powf(log_upper + (log_lower - log_upper) * (1.0 - fraction)); + } + + f64::NAN + } + + pub fn fraction(&self, lower: f64, upper: f64) -> f64 { + if self.count == 0.0 || lower.is_nan() || upper.is_nan() { + return f64::NAN; + } + if lower >= upper { + return 0.0; + } + + let Some(mut buckets) = self.all_buckets() else { + return f64::NAN; + }; + + let mut rank = 0.0; + let mut lower_rank = 0.0; + let mut upper_rank = 0.0; + let mut lower_set = false; + let mut upper_set = false; + + for bucket in &mut buckets { + let zero_bucket = bucket.lower <= 0.0 && bucket.upper >= 0.0; + if zero_bucket { + if self.negative_buckets.is_empty() && !self.positive_buckets.is_empty() { + bucket.lower = 0.0; + } else if self.positive_buckets.is_empty() && !self.negative_buckets.is_empty() { + bucket.upper = 0.0; + } + } + + if !lower_set && bucket.lower >= lower { + lower_rank = rank; + lower_set = true; + } + if !upper_set && bucket.lower >= upper { + upper_rank = rank; + upper_set = true; + } + if lower_set && upper_set { + break; + } + if !lower_set && bucket.lower < lower && bucket.upper > lower { + lower_rank = self.interpolate_rank(bucket, rank, lower, zero_bucket); + lower_set = true; + } + if !upper_set && bucket.lower < upper && bucket.upper > upper { + upper_rank = self.interpolate_rank(bucket, rank, upper, zero_bucket); + upper_set = true; + } + if lower_set && upper_set { + break; + } + rank += bucket.count; + } + + if !lower_set || lower_rank > self.count { + lower_rank = self.count; + } + if !upper_set || upper_rank > self.count { + upper_rank = self.count; + } + + (upper_rank - lower_rank) / self.count + } + + pub fn to_prometheus_buckets(&self) -> Option> { + Some( + self.all_buckets()? + .into_iter() + .filter(|bucket| bucket.count != 0.0) + .map(|bucket| { + ( + bucket.boundary_rule as u8, + bucket.lower.to_string(), + bucket.upper.to_string(), + bucket.count.to_string(), + ) + }) + .collect(), + ) + } + + fn interpolate_rank(&self, bucket: &Bucket, rank: f64, value: f64, zero_bucket: bool) -> f64 { + if self.uses_custom_buckets() || zero_bucket { + if bucket.lower == f64::NEG_INFINITY { + return bucket.count; + } + return rank + bucket.count * (value - bucket.lower) / (bucket.upper - bucket.lower); + } + + let log_lower = bucket.lower.abs().log2(); + let log_upper = bucket.upper.abs().log2(); + let log_value = value.abs().log2(); + let fraction = if value > 0.0 { + (log_value - log_lower) / (log_upper - log_lower) + } else { + 1.0 - ((log_value - log_upper) / (log_lower - log_upper)) + }; + rank + bucket.count * fraction + } + + fn bucket_midpoint(&self, bucket: &Bucket) -> f64 { + if bucket.lower == f64::NEG_INFINITY && bucket.upper == f64::INFINITY { + return 0.0; + } + if bucket.lower == f64::NEG_INFINITY { + return bucket.upper; + } + if bucket.upper == f64::INFINITY { + return bucket.lower; + } + if self.uses_custom_buckets() || (bucket.lower <= 0.0 && bucket.upper >= 0.0) { + return (bucket.lower + bucket.upper) / 2.0; + } + if bucket.upper < 0.0 { + -((bucket.lower.abs() * bucket.upper.abs()).sqrt()) + } else { + (bucket.lower * bucket.upper).sqrt() + } + } + + fn side_has_reset(&self, positive: bool, previous: &Self) -> bool { + let (current_spans, current_buckets, previous_spans, previous_buckets) = if positive { + ( + &self.positive_spans, + &self.positive_buckets, + &previous.positive_spans, + &previous.positive_buckets, + ) + } else { + ( + &self.negative_spans, + &self.negative_buckets, + &previous.negative_spans, + &previous.negative_buckets, + ) + }; + let Some(current) = side_counts(current_spans, current_buckets) else { + return true; + }; + let Some(previous) = side_counts(previous_spans, previous_buckets) else { + return true; + }; + previous.keys().chain(current.keys()).any(|idx| { + current.get(idx).copied().unwrap_or_default() + < previous.get(idx).copied().unwrap_or_default() + }) + } + + fn reconcile(&self, other: &Self) -> Option<(Self, Self)> { + match (self.uses_custom_buckets(), other.uses_custom_buckets()) { + (true, true) => reconcile_custom(self, other), + (false, false) => reconcile_exponential(self, other), + _ => None, + } + } + + fn is_empty_payload(&self) -> bool { + self.count == 0.0 + && self.zero_count == 0.0 + && self.sum == 0.0 + && self.positive_buckets.iter().all(|count| *count == 0.0) + && self.negative_buckets.iter().all(|count| *count == 0.0) + } +} + +fn reconcile_exponential( + left: &NativeHistogram, + right: &NativeHistogram, +) -> Option<(NativeHistogram, NativeHistogram)> { + let schema = left.schema.min(right.schema); + let mut left = left.copy_to_schema(schema)?; + let mut right = right.copy_to_schema(schema)?; + let mut zero_threshold = left.zero_threshold.max(right.zero_threshold); + loop { + let expanded = left + .expanded_zero_threshold(zero_threshold)? + .max(right.expanded_zero_threshold(zero_threshold)?); + if expanded == zero_threshold { + break; + } + zero_threshold = expanded; + } + left.grow_zero_threshold(zero_threshold)?; + right.grow_zero_threshold(zero_threshold)?; + Some((left.compact(), right.compact())) +} + +fn reconcile_custom( + left: &NativeHistogram, + right: &NativeHistogram, +) -> Option<(NativeHistogram, NativeHistogram)> { + let custom_values = if left.custom_values == right.custom_values { + left.custom_values.clone() + } else { + left.custom_values + .iter() + .copied() + .filter(|value| right.custom_values.contains(value)) + .collect() + }; + + Some(( + left.copy_to_custom_values(custom_values.clone())?.compact(), + right.copy_to_custom_values(custom_values)?.compact(), + )) +} + +impl NativeHistogram { + fn copy_to_schema(&self, target_schema: i32) -> Option { + if self.uses_custom_buckets() + || !(MIN_EXPONENTIAL_SCHEMA..=MAX_EXPONENTIAL_SCHEMA).contains(&target_schema) + || target_schema > self.schema + { + return None; + } + if target_schema == self.schema { + return Some(self.clone()); + } + + let mut result = self.clone(); + result.schema = target_schema; + (result.positive_spans, result.positive_buckets) = reduce_side( + &self.positive_spans, + &self.positive_buckets, + self.schema, + target_schema, + )?; + (result.negative_spans, result.negative_buckets) = reduce_side( + &self.negative_spans, + &self.negative_buckets, + self.schema, + target_schema, + )?; + Some(result) + } + + fn copy_to_custom_values(&self, custom_values: Vec) -> Option { + if !self.uses_custom_buckets() { + return None; + } + if self.custom_values == custom_values { + return Some(self.clone()); + } + + let mut result = self.clone(); + result.custom_values = custom_values.clone(); + result.negative_spans.clear(); + result.negative_buckets.clear(); + (result.positive_spans, result.positive_buckets) = map_custom_side( + &self.positive_spans, + &self.positive_buckets, + &self.custom_values, + &custom_values, + )?; + Some(result) + } + + fn grow_zero_threshold(&mut self, zero_threshold: f64) -> Option<()> { + if self.uses_custom_buckets() || zero_threshold == self.zero_threshold { + self.zero_threshold = zero_threshold; + return Some(()); + } + + let (spans, buckets, zero_count) = fold_zero_side( + &self.positive_spans, + &self.positive_buckets, + self.schema, + zero_threshold, + )?; + self.positive_spans = spans; + self.positive_buckets = buckets; + self.zero_count += zero_count; + + let (spans, buckets, zero_count) = fold_zero_side( + &self.negative_spans, + &self.negative_buckets, + self.schema, + zero_threshold, + )?; + self.negative_spans = spans; + self.negative_buckets = buckets; + self.zero_count += zero_count; + self.zero_threshold = zero_threshold; + Some(()) + } + + fn expanded_zero_threshold(&self, mut zero_threshold: f64) -> Option { + if self.uses_custom_buckets() { + return Some(zero_threshold); + } + zero_threshold = expand_zero_threshold_side( + &self.positive_spans, + &self.positive_buckets, + self.schema, + zero_threshold, + )?; + expand_zero_threshold_side( + &self.negative_spans, + &self.negative_buckets, + self.schema, + zero_threshold, + ) + } +} + +fn add_reset_hint(left: i32, right: i32) -> i32 { + if left == GAUGE_RESET_HINT || right == GAUGE_RESET_HINT { + GAUGE_RESET_HINT + } else if left == right { + left + } else { + UNKNOWN_COUNTER_RESET_HINT + } +} + +fn side_bucket_indices(spans: &[Span]) -> Option> { + let mut indices = Vec::new(); + let mut current_index = 0i32; + let mut first = true; + for span in spans { + if first { + current_index = span.offset; + first = false; + } else { + current_index = current_index.checked_add(span.offset)?; + } + for _ in 0..span.length { + indices.push(current_index); + current_index = current_index.checked_add(1)?; + } + } + Some(indices) +} + +fn spans_from_indices_counts(values: Vec<(i32, f64)>) -> (Vec, Vec) { + let mut spans = Vec::new(); + let mut buckets = Vec::new(); + let mut current_span_start = None::; + let mut previous_index = 0i32; + + for (idx, count) in values { + if count == 0.0 { + continue; + } + match current_span_start { + None => { + spans.push(Span { + offset: idx, + length: 1, + }); + current_span_start = Some(idx); + } + Some(_) if idx == previous_index + 1 => { + spans.last_mut().expect("span exists").length += 1; + } + Some(_) => { + spans.push(Span { + offset: idx - previous_index - 1, + length: 1, + }); + } + } + buckets.push(count); + previous_index = idx; + } + + (spans, buckets) +} + +fn side_counts(spans: &[Span], buckets: &[f64]) -> Option> { + let mut values = BTreeMap::new(); + for (idx, count) in side_bucket_indices(spans)?.into_iter().zip(buckets) { + if *count != 0.0 { + values.insert(idx, *count); + } + } + Some(values) +} + +fn merge_side( + left_spans: &[Span], + left_buckets: &[f64], + right_spans: &[Span], + right_buckets: &[f64], + op: impl Fn(f64, f64) -> f64, +) -> Option<(Vec, Vec)> { + let left = side_counts(left_spans, left_buckets)?; + let right = side_counts(right_spans, right_buckets)?; + let mut values = BTreeMap::new(); + for idx in left.keys().chain(right.keys()) { + values.insert( + *idx, + op( + left.get(idx).copied().unwrap_or_default(), + right.get(idx).copied().unwrap_or_default(), + ), + ); + } + Some(spans_from_indices_counts(values.into_iter().collect())) +} + +fn reduce_side( + spans: &[Span], + buckets: &[f64], + schema: i32, + target_schema: i32, +) -> Option<(Vec, Vec)> { + let factor = 1_i32.checked_shl((schema - target_schema) as u32)?; + let mut values = std::collections::BTreeMap::::new(); + for (idx, count) in side_bucket_indices(spans)?.into_iter().zip(buckets) { + let target_idx = ceil_div(idx, factor); + *values.entry(target_idx).or_default() += *count; + } + Some(spans_from_indices_counts(values.into_iter().collect())) +} + +fn compact_side(spans: &[Span], buckets: &[f64]) -> (Vec, Vec) { + let Some(indices) = side_bucket_indices(spans) else { + return (Vec::new(), Vec::new()); + }; + spans_from_indices_counts(indices.into_iter().zip(buckets.iter().copied()).collect()) +} + +fn map_custom_side( + spans: &[Span], + buckets: &[f64], + old_values: &[f64], + new_values: &[f64], +) -> Option<(Vec, Vec)> { + let mut values = std::collections::BTreeMap::::new(); + for (idx, count) in side_bucket_indices(spans)?.into_iter().zip(buckets) { + let upper = get_bound(idx, CUSTOM_BUCKETS_SCHEMA, old_values)?; + let target_idx = new_values + .iter() + .position(|value| *value >= upper) + .unwrap_or(new_values.len()) as i32; + *values.entry(target_idx).or_default() += *count; + } + Some(spans_from_indices_counts(values.into_iter().collect())) +} + +fn fold_zero_side( + spans: &[Span], + buckets: &[f64], + schema: i32, + zero_threshold: f64, +) -> Option<(Vec, Vec, f64)> { + let mut kept = Vec::new(); + let mut zero_count = 0.0; + for (idx, count) in side_bucket_indices(spans)?.into_iter().zip(buckets) { + if get_bound(idx, schema, &[])? <= zero_threshold { + zero_count += *count; + } else { + kept.push((idx, *count)); + } + } + let (spans, buckets) = spans_from_indices_counts(kept); + Some((spans, buckets, zero_count)) +} + +fn expand_zero_threshold_side( + spans: &[Span], + buckets: &[f64], + schema: i32, + mut zero_threshold: f64, +) -> Option { + for (idx, count) in side_bucket_indices(spans)?.into_iter().zip(buckets) { + if *count == 0.0 { + continue; + } + let lower = get_bound(idx - 1, schema, &[])?; + let upper = get_bound(idx, schema, &[])?; + if lower < zero_threshold && zero_threshold < upper { + zero_threshold = upper; + } + } + Some(zero_threshold) +} + +fn ceil_div(value: i32, divisor: i32) -> i32 { + -((-value).div_euclid(divisor)) +} + +pub fn get_bound(idx: i32, schema: i32, custom_values: &[f64]) -> Option { + if schema == CUSTOM_BUCKETS_SCHEMA { + return match idx { + -1 => Some(f64::NEG_INFINITY), + idx if idx == custom_values.len() as i32 => Some(f64::INFINITY), + idx if idx >= 0 && (idx as usize) < custom_values.len() => { + Some(custom_values[idx as usize]) + } + _ => None, + }; + } + + if !(MIN_EXPONENTIAL_SCHEMA..=MAX_EXPONENTIAL_SCHEMA).contains(&schema) { + return None; + } + if schema < 0 { + return Some(2.0_f64.powi(idx.checked_shl((-schema) as u32)?)); + } + Some(2.0_f64.powf(idx as f64 / (1u32 << schema) as f64)) +} + +pub fn native_histogram_arrow_type() -> ArrowDataType { + native_histogram_value_type().as_arrow_type() +} + +fn native_histogram_fields() -> datafusion::arrow::datatypes::Fields { + match native_histogram_arrow_type() { + ArrowDataType::Struct(fields) => fields, + _ => unreachable!("native histogram must be a struct"), + } +} + +fn struct_child<'a>(array: &'a StructArray, name: &str) -> DfResult<&'a ArrayRef> { + let index = array + .fields() + .iter() + .position(|field| field.name() == name) + .ok_or_else(|| { + DataFusionError::Execution(format!("native histogram missing field {name}")) + })?; + Ok(array.column(index)) +} + +fn primitive_child<'a, T>(array: &'a StructArray, name: &str) -> DfResult<&'a PrimitiveArray> +where + T: ArrowPrimitiveType, +{ + let child = struct_child(array, name)?; + child + .as_any() + .downcast_ref::>() + .ok_or_else(|| { + DataFusionError::Execution(format!( + "native histogram field {name} has invalid type {}", + child.data_type() + )) + }) +} + +fn list_child<'a>(array: &'a StructArray, name: &str) -> DfResult<&'a ListArray> { + let child = struct_child(array, name)?; + child.as_any().downcast_ref::().ok_or_else(|| { + DataFusionError::Execution(format!( + "native histogram field {name} has invalid type {}", + child.data_type() + )) + }) +} + +fn required_primitive(array: &StructArray, name: &str, row: usize) -> DfResult +where + T: ArrowPrimitiveType, +{ + let values = primitive_child::(array, name)?; + if values.is_null(row) { + return Err(DataFusionError::Execution(format!( + "native histogram field {name} is null" + ))); + } + Ok(values.value(row)) +} + +fn optional_primitive(array: &StructArray, name: &str, row: usize) -> DfResult> +where + T: ArrowPrimitiveType, +{ + let values = primitive_child::(array, name)?; + Ok((!values.is_null(row)).then(|| values.value(row))) +} + +fn list_values(array: &StructArray, name: &str, row: usize) -> DfResult> +where + T: ArrowPrimitiveType, +{ + let list = list_child(array, name)?; + if list.is_null(row) { + return Ok(Vec::new()); + } + + let values = list.value(row); + let values = values + .as_any() + .downcast_ref::>() + .ok_or_else(|| { + DataFusionError::Execution(format!( + "native histogram list field {name} has invalid value type {}", + values.data_type() + )) + })?; + + values + .iter() + .map(|value| { + value.ok_or_else(|| { + DataFusionError::Execution(format!( + "native histogram list field {name} contains null" + )) + }) + }) + .collect() +} + +fn read_spans(offsets: Vec, lengths: Vec, name: &str) -> DfResult> { + if offsets.len() != lengths.len() { + return Err(DataFusionError::Execution(format!( + "native histogram {name} span offsets and lengths mismatch: {} vs {}", + offsets.len(), + lengths.len() + ))); + } + Ok(offsets + .into_iter() + .zip(lengths) + .map(|(offset, length)| Span { offset, length }) + .collect()) +} + +fn check_span_bucket_count(spans: &[Span], buckets: usize, name: &str) -> DfResult<()> { + let span_len = spans + .iter() + .try_fold(0usize, |sum, span| sum.checked_add(span.length as usize)) + .ok_or_else(|| { + DataFusionError::Execution(format!("native histogram {name} spans overflow")) + })?; + if span_len != buckets { + return Err(DataFusionError::Execution(format!( + "native histogram {name} spans describe {span_len} buckets, found {buckets}" + ))); + } + Ok(()) +} + +pub fn read_histogram(array: &StructArray, row: usize) -> DfResult> { + if array.is_null(row) { + return Ok(None); + } + + let schema = required_primitive::(array, SCHEMA_FIELD, row)?; + let positive_spans = read_spans( + list_values::(array, POSITIVE_SPAN_OFFSETS_FIELD, row)?, + list_values::(array, POSITIVE_SPAN_LENGTHS_FIELD, row)?, + "positive", + )?; + let negative_spans = read_spans( + list_values::(array, NEGATIVE_SPAN_OFFSETS_FIELD, row)?, + list_values::(array, NEGATIVE_SPAN_LENGTHS_FIELD, row)?, + "negative", + )?; + + let (count, zero_count, positive_buckets, negative_buckets) = + if let Some(count) = optional_primitive::(array, COUNT_F64_FIELD, row)? { + ( + count, + optional_primitive::(array, ZERO_COUNT_F64_FIELD, row)? + .unwrap_or_default(), + list_values::(array, POSITIVE_BUCKETS_F64_FIELD, row)?, + list_values::(array, NEGATIVE_BUCKETS_F64_FIELD, row)?, + ) + } else { + ( + required_primitive::(array, COUNT_U64_FIELD, row)? as f64, + optional_primitive::(array, ZERO_COUNT_U64_FIELD, row)? + .unwrap_or_default() as f64, + list_values::(array, POSITIVE_BUCKETS_I64_FIELD, row)? + .into_iter() + .map(|value| value as f64) + .collect(), + list_values::(array, NEGATIVE_BUCKETS_I64_FIELD, row)? + .into_iter() + .map(|value| value as f64) + .collect(), + ) + }; + + check_span_bucket_count(&positive_spans, positive_buckets.len(), "positive")?; + check_span_bucket_count(&negative_spans, negative_buckets.len(), "negative")?; + + Ok(Some(NativeHistogram { + schema, + zero_threshold: required_primitive::(array, ZERO_THRESHOLD_FIELD, row)?, + sum: required_primitive::(array, SUM_FIELD, row)?, + reset_hint: required_primitive::(array, RESET_HINT_FIELD, row)?, + start_timestamp: optional_primitive::( + array, + START_TIMESTAMP_FIELD, + row, + )?, + custom_values: list_values::(array, CUSTOM_VALUES_FIELD, row)?, + positive_spans, + negative_spans, + count, + zero_count, + positive_buckets, + negative_buckets, + })) +} + +fn list_opt(values: Vec) -> Option>> { + Some(values.into_iter().map(Some).collect()) +} + +pub fn build_histogram_array(values: &[Option]) -> ArrayRef { + let mut schemas = Vec::with_capacity(values.len()); + let mut zero_thresholds = Vec::with_capacity(values.len()); + let mut sums = Vec::with_capacity(values.len()); + let mut reset_hints = Vec::with_capacity(values.len()); + let mut start_timestamps = Vec::with_capacity(values.len()); + let mut custom_values = Vec::with_capacity(values.len()); + let mut positive_span_offsets = Vec::with_capacity(values.len()); + let mut positive_span_lengths = Vec::with_capacity(values.len()); + let mut negative_span_offsets = Vec::with_capacity(values.len()); + let mut negative_span_lengths = Vec::with_capacity(values.len()); + let mut count_u64 = Vec::>::with_capacity(values.len()); + let mut zero_count_u64 = Vec::>::with_capacity(values.len()); + let mut positive_buckets_i64 = Vec::with_capacity(values.len()); + let mut negative_buckets_i64 = Vec::with_capacity(values.len()); + let mut count_f64 = Vec::with_capacity(values.len()); + let mut zero_count_f64 = Vec::with_capacity(values.len()); + let mut positive_buckets_f64 = Vec::with_capacity(values.len()); + let mut negative_buckets_f64 = Vec::with_capacity(values.len()); + let mut validity = Vec::with_capacity(values.len()); + + for value in values { + validity.push(value.is_some()); + if let Some(histogram) = value { + schemas.push(Some(histogram.schema)); + zero_thresholds.push(Some(histogram.zero_threshold)); + sums.push(Some(histogram.sum)); + reset_hints.push(Some(histogram.reset_hint)); + start_timestamps.push(histogram.start_timestamp); + custom_values.push(list_opt(histogram.custom_values.clone())); + positive_span_offsets.push(list_opt( + histogram + .positive_spans + .iter() + .map(|span| span.offset) + .collect(), + )); + positive_span_lengths.push(list_opt( + histogram + .positive_spans + .iter() + .map(|span| span.length) + .collect(), + )); + negative_span_offsets.push(list_opt( + histogram + .negative_spans + .iter() + .map(|span| span.offset) + .collect(), + )); + negative_span_lengths.push(list_opt( + histogram + .negative_spans + .iter() + .map(|span| span.length) + .collect(), + )); + count_u64.push(None); + zero_count_u64.push(None); + positive_buckets_i64.push(list_opt(Vec::::new())); + negative_buckets_i64.push(list_opt(Vec::::new())); + count_f64.push(Some(histogram.count)); + zero_count_f64.push(Some(histogram.zero_count)); + positive_buckets_f64.push(list_opt(histogram.positive_buckets.clone())); + negative_buckets_f64.push(list_opt(histogram.negative_buckets.clone())); + } else { + schemas.push(None); + zero_thresholds.push(None); + sums.push(None); + reset_hints.push(None); + start_timestamps.push(None); + custom_values.push(None); + positive_span_offsets.push(None); + positive_span_lengths.push(None); + negative_span_offsets.push(None); + negative_span_lengths.push(None); + count_u64.push(None); + zero_count_u64.push(None); + positive_buckets_i64.push(None); + negative_buckets_i64.push(None); + count_f64.push(None); + zero_count_f64.push(None); + positive_buckets_f64.push(None); + negative_buckets_f64.push(None); + } + } + + let arrays: Vec = vec![ + Arc::new(Int32Array::from(schemas)), + Arc::new(Float64Array::from(zero_thresholds)), + Arc::new(Float64Array::from(sums)), + Arc::new(Int32Array::from(reset_hints)), + Arc::new(TimestampMillisecondArray::from_iter(start_timestamps)), + Arc::new(ListArray::from_iter_primitive::( + custom_values, + )), + Arc::new(ListArray::from_iter_primitive::( + positive_span_offsets, + )), + Arc::new(ListArray::from_iter_primitive::( + positive_span_lengths, + )), + Arc::new(ListArray::from_iter_primitive::( + negative_span_offsets, + )), + Arc::new(ListArray::from_iter_primitive::( + negative_span_lengths, + )), + Arc::new(UInt64Array::from(count_u64)), + Arc::new(UInt64Array::from(zero_count_u64)), + Arc::new(ListArray::from_iter_primitive::( + positive_buckets_i64, + )), + Arc::new(ListArray::from_iter_primitive::( + negative_buckets_i64, + )), + Arc::new(Float64Array::from(count_f64)), + Arc::new(Float64Array::from(zero_count_f64)), + Arc::new(ListArray::from_iter_primitive::( + positive_buckets_f64, + )), + Arc::new(ListArray::from_iter_primitive::( + negative_buckets_f64, + )), + ]; + + Arc::new(StructArray::new( + native_histogram_fields(), + arrays, + Some(NullBuffer::from(validity)), + )) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn histogram(positive_spans: Vec, positive_buckets: Vec) -> NativeHistogram { + let count = positive_buckets.iter().sum(); + NativeHistogram { + schema: 0, + zero_threshold: 0.0, + sum: count, + reset_hint: COUNTER_RESET_HINT, + start_timestamp: None, + custom_values: Vec::new(), + positive_spans, + negative_spans: Vec::new(), + count, + zero_count: 0.0, + positive_buckets, + negative_buckets: Vec::new(), + } + } + + #[test] + fn add_uses_sparse_bucket_union() { + let left = histogram( + vec![Span { + offset: 0, + length: 1, + }], + vec![1.0], + ); + let right = histogram( + vec![Span { + offset: 2, + length: 1, + }], + vec![2.0], + ); + + let result = left.add(&right).unwrap(); + assert_eq!( + result.positive_spans, + vec![ + Span { + offset: 0, + length: 1, + }, + Span { + offset: 1, + length: 1, + }, + ] + ); + assert_eq!(result.positive_buckets, vec![1.0, 2.0]); + } + + #[test] + fn zero_threshold_expands_to_populated_bucket_boundary() { + let left = NativeHistogram { + zero_threshold: 0.5, + ..histogram( + vec![Span { + offset: 0, + length: 1, + }], + vec![1.0], + ) + }; + let right = NativeHistogram { + zero_threshold: 0.75, + ..histogram( + vec![Span { + offset: 1, + length: 1, + }], + vec![1.0], + ) + }; + + let result = left.add(&right).unwrap(); + assert_eq!(result.zero_threshold, 1.0); + assert_eq!(result.zero_count, 1.0); + assert_eq!(result.positive_buckets, vec![1.0]); + } + + #[test] + fn custom_reconciliation_without_shared_bounds_uses_overflow_bucket() { + let mut left = histogram( + vec![Span { + offset: 0, + length: 1, + }], + vec![1.0], + ); + left.schema = CUSTOM_BUCKETS_SCHEMA; + left.custom_values = vec![1.0]; + + let mut right = histogram( + vec![Span { + offset: 0, + length: 1, + }], + vec![2.0], + ); + right.schema = CUSTOM_BUCKETS_SCHEMA; + right.custom_values = vec![2.0]; + + let result = left.add(&right).unwrap(); + assert!(result.custom_values.is_empty()); + assert_eq!( + result.positive_spans, + vec![Span { + offset: 0, + length: 1, + }] + ); + assert_eq!(result.positive_buckets, vec![3.0]); + } + + #[test] + fn promql_eq_ignores_reset_hint_start_timestamp_and_sparse_zeros() { + let mut left = histogram( + vec![Span { + offset: 0, + length: 2, + }], + vec![1.0, 0.0], + ); + left.reset_hint = COUNTER_RESET_HINT; + left.start_timestamp = Some(1000); + + let mut right = histogram( + vec![Span { + offset: 0, + length: 1, + }], + vec![1.0], + ); + right.reset_hint = NOT_COUNTER_RESET_HINT; + right.start_timestamp = Some(2000); + + assert!(left.promql_eq(&right)); + } + + #[test] + fn subtraction_returns_gauge_histogram() { + let left = histogram( + vec![Span { + offset: 0, + length: 1, + }], + vec![3.0], + ); + let right = histogram( + vec![Span { + offset: 0, + length: 1, + }], + vec![1.0], + ); + + let result = left.sub(&right).unwrap(); + assert_eq!(result.reset_hint, GAUGE_RESET_HINT); + assert_eq!(result.positive_buckets, vec![2.0]); + } + + #[test] + fn subtraction_treats_empty_left_as_zero() { + let left = histogram(vec![], vec![]); + let mut right = histogram( + vec![Span { + offset: 0, + length: 1, + }], + vec![2.0], + ); + right.schema = CUSTOM_BUCKETS_SCHEMA; + right.custom_values = vec![1.0]; + + let result = left.sub(&right).unwrap(); + + assert_eq!(result.schema, CUSTOM_BUCKETS_SCHEMA); + assert_eq!(result.reset_hint, GAUGE_RESET_HINT); + assert_eq!(result.count, -2.0); + assert_eq!(result.sum, -2.0); + assert_eq!(result.positive_buckets, vec![-2.0]); + } } diff --git a/src/common/query/src/promql_annotations.rs b/src/common/query/src/promql_annotations.rs new file mode 100644 index 00000000000..70f4db0af70 --- /dev/null +++ b/src/common/query/src/promql_annotations.rs @@ -0,0 +1,170 @@ +// 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::collections::{BTreeSet, HashMap}; +use std::sync::{Arc, Mutex, Weak}; + +use datafusion::config::{ConfigEntry, ConfigExtension, ExtensionOptions}; +use datafusion_common::{DataFusionError, Result as DfResult}; +use once_cell::sync::Lazy; + +type PromqlAnnotationStateRef = Mutex; + +static PROMQL_ANNOTATION_REGISTRY: Lazy>>> = + Lazy::new(|| Mutex::new(HashMap::new())); + +#[derive(Debug, Default)] +struct PromqlAnnotationState { + warnings: BTreeSet, + infos: BTreeSet, +} + +#[derive(Clone, Debug)] +pub struct PromqlAnnotationCollector { + inner: Arc, +} + +impl Default for PromqlAnnotationCollector { + fn default() -> Self { + Self { + inner: Arc::new(Mutex::new(PromqlAnnotationState::default())), + } + } +} + +impl PromqlAnnotationCollector { + pub fn record_warning>(&self, warning: S) { + self.inner + .lock() + .expect("promql annotation collector poisoned") + .warnings + .insert(warning.into()); + } + + pub fn record_info>(&self, info: S) { + self.inner + .lock() + .expect("promql annotation collector poisoned") + .infos + .insert(info.into()); + } + + pub fn append_to(&self, warnings: &mut Vec, infos: &mut Vec) { + let inner = self + .inner + .lock() + .expect("promql annotation collector poisoned"); + append_unique(warnings, inner.warnings.iter().cloned()); + append_unique(infos, inner.infos.iter().cloned()); + } +} + +impl ConfigExtension for PromqlAnnotationCollector { + const PREFIX: &'static str = "greptime_promql_annotations"; +} + +impl ExtensionOptions for PromqlAnnotationCollector { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn cloned(&self) -> Box { + Box::new(self.clone()) + } + + fn set(&mut self, key: &str, value: &str) -> DfResult<()> { + Err(DataFusionError::NotImplemented(format!( + "PromqlAnnotationCollector does not support set key: {key} with value: {value}" + ))) + } + + fn entries(&self) -> Vec { + vec![] + } +} + +pub fn promql_annotation_collector(query_id: &str) -> PromqlAnnotationCollector { + let mut registry = PROMQL_ANNOTATION_REGISTRY + .lock() + .expect("promql annotation registry poisoned"); + registry.retain(|_, collector| collector.strong_count() > 0); + + if let Some(inner) = registry.get(query_id).and_then(Weak::upgrade) { + return PromqlAnnotationCollector { inner }; + } + + let collector = PromqlAnnotationCollector::default(); + registry.insert(query_id.to_string(), Arc::downgrade(&collector.inner)); + collector +} + +pub fn get_promql_annotation_collector(query_id: &str) -> Option { + let mut registry = PROMQL_ANNOTATION_REGISTRY + .lock() + .expect("promql annotation registry poisoned"); + registry.retain(|_, collector| collector.strong_count() > 0); + registry + .get(query_id) + .and_then(Weak::upgrade) + .map(|inner| PromqlAnnotationCollector { inner }) +} + +fn append_unique(target: &mut Vec, values: impl Iterator) { + let mut seen = target.iter().cloned().collect::>(); + for value in values { + if seen.insert(value.clone()) { + target.push(value); + } + } + target.sort(); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn collector_deduplicates_and_orders_annotations() { + let collector = PromqlAnnotationCollector::default(); + collector.record_info("z-info"); + collector.record_info("a-info"); + collector.record_info("z-info"); + collector.record_warning("b-warning"); + collector.record_warning("a-warning"); + + let mut warnings = vec!["c-warning".to_string()]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + + assert_eq!(warnings, vec!["a-warning", "b-warning", "c-warning"]); + assert_eq!(infos, vec!["a-info", "z-info"]); + } + + #[test] + fn registry_prunes_dead_collectors() { + let key = "registry_prunes_dead_collectors"; + { + let collector = promql_annotation_collector(key); + assert!(get_promql_annotation_collector(key).is_some()); + drop(collector); + } + + assert!(get_promql_annotation_collector(key).is_none()); + assert!(!PROMQL_ANNOTATION_REGISTRY.lock().unwrap().contains_key(key)); + } +} diff --git a/src/promql/Cargo.toml b/src/promql/Cargo.toml index 460be8ddd95..36ac213fef4 100644 --- a/src/promql/Cargo.toml +++ b/src/promql/Cargo.toml @@ -13,6 +13,7 @@ async-trait.workspace = true bytemuck.workspace = true common-error.workspace = true common-macro.workspace = true +common-query.workspace = true common-telemetry.workspace = true datafusion.workspace = true datafusion-common.workspace = true diff --git a/src/promql/src/functions.rs b/src/promql/src/functions.rs index 7c7452566a6..3c8b64fad6d 100644 --- a/src/promql/src/functions.rs +++ b/src/promql/src/functions.rs @@ -18,6 +18,7 @@ mod deriv; mod double_exponential_smoothing; mod extrapolate_rate; mod idelta; +mod native_histogram; mod predict_linear; mod quantile; mod quantile_aggr; @@ -42,6 +43,18 @@ pub use deriv::Deriv; pub use double_exponential_smoothing::DoubleExponentialSmoothing; pub use extrapolate_rate::{Delta, Increase, Rate}; pub use idelta::IDelta; +pub use native_histogram::{ + NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAggAvg, + NativeHistogramAggSum, NativeHistogramAvg, NativeHistogramAvgOverTime, NativeHistogramChanges, + NativeHistogramCount, NativeHistogramCountOverTime, NativeHistogramDelta, + NativeHistogramDivScalar, NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction, + NativeHistogramIDelta, NativeHistogramIRate, NativeHistogramIncrease, + NativeHistogramLastOverTime, NativeHistogramMulScalar, NativeHistogramNeg, + NativeHistogramNotEq, NativeHistogramPresentOverTime, NativeHistogramQuantile, + NativeHistogramRate, NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev, + NativeHistogramStdvar, NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime, + NativeHistogramToString, +}; pub use predict_linear::PredictLinear; pub use quantile::QuantileOverTime; pub use quantile_aggr::{QUANTILE_NAME, quantile_udaf}; diff --git a/src/promql/src/functions/native_histogram.rs b/src/promql/src/functions/native_histogram.rs new file mode 100644 index 00000000000..b8cc9b0cc33 --- /dev/null +++ b/src/promql/src/functions/native_histogram.rs @@ -0,0 +1,2061 @@ +// 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. + +//! Native histogram PromQL helpers. + +use std::any::Any; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use common_query::native_histogram::*; +use common_query::promql_annotations::PromqlAnnotationCollector; +use datafusion::arrow::array::{ + Array, ArrayRef, BooleanArray, DictionaryArray, Float64Array, Float64Builder, Int64Array, + StringBuilder, StructArray, TimestampMillisecondArray, UInt64Array, +}; +use datafusion::arrow::datatypes::{DataType, Int64Type, TimeUnit}; +use datafusion::common::{DataFusionError, Result as DfResult}; +use datafusion::logical_expr::{Accumulator as DfAccumulator, AggregateUDF, ScalarUDF, Volatility}; +use datafusion::physical_plan::ColumnarValue; +use datafusion_common::ScalarValue; +use datafusion_expr::function::AccumulatorArgs; +use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature, create_udaf, create_udf}; + +use crate::functions::extract_array; +use crate::range_array::{RangeArray, unpack}; + +fn extract_histogram_array(value: &ColumnarValue, func_name: &str) -> DfResult { + let array = extract_array(value)?; + if array.data_type() != &native_histogram_arrow_type() { + return Err(DataFusionError::Execution(format!( + "{func_name}: expected native histogram struct, found {}", + array.data_type() + ))); + } + Ok(array) +} + +fn read_scalar_f64_arg( + value: &ColumnarValue, + row: usize, + len: usize, + func_name: &str, +) -> DfResult { + match value { + ColumnarValue::Scalar(ScalarValue::Float64(value)) => Ok(value.unwrap_or(f64::NAN)), + ColumnarValue::Array(array) => { + let array = array + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Execution(format!( + "{func_name}: expected Float64 argument, found {}", + array.data_type() + )) + })?; + if array.len() != len { + return Err(DataFusionError::Execution(format!( + "{func_name}: Float64 argument length mismatch: {} vs {len}", + array.len() + ))); + } + Ok(if array.is_null(row) { + f64::NAN + } else { + array.value(row) + }) + } + other => Err(DataFusionError::Execution(format!( + "{func_name}: expected Float64 argument, found {}", + other.data_type() + ))), + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +enum AnnotationReturn { + FloatNull, + BooleanFalse, +} + +impl AnnotationReturn { + fn data_type(self) -> DataType { + match self { + Self::FloatNull => DataType::Float64, + Self::BooleanFalse => DataType::Boolean, + } + } + + fn scalar_value(self) -> ScalarValue { + match self { + Self::FloatNull => ScalarValue::Float64(None), + Self::BooleanFalse => ScalarValue::Boolean(Some(false)), + } + } +} + +#[derive(Debug, Clone)] +struct NativeHistogramAnnotationUdf { + name: &'static str, + signature: Signature, + return_kind: AnnotationReturn, + message: String, + collector: Option, +} + +impl NativeHistogramAnnotationUdf { + fn new( + name: &'static str, + return_kind: AnnotationReturn, + message: String, + collector: Option, + ) -> Self { + Self { + name, + signature: Signature::variadic_any(Volatility::Volatile), + return_kind, + message, + collector, + } + } +} + +impl PartialEq for NativeHistogramAnnotationUdf { + fn eq(&self, other: &Self) -> bool { + self.name == other.name + && self.return_kind == other.return_kind + && self.message == other.message + } +} + +impl Eq for NativeHistogramAnnotationUdf {} + +impl Hash for NativeHistogramAnnotationUdf { + fn hash(&self, state: &mut H) { + self.name.hash(state); + self.return_kind.hash(state); + self.message.hash(state); + } +} + +impl ScalarUDFImpl for NativeHistogramAnnotationUdf { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> DfResult { + Ok(self.return_kind.data_type()) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DfResult { + if let Some(collector) = args + .config_options + .extensions + .get::() + .cloned() + .or_else(|| self.collector.clone()) + { + collector.record_info(self.message.clone()); + } + Ok(ColumnarValue::Scalar(self.return_kind.scalar_value())) + } +} + +pub struct NativeHistogramDrop; + +impl NativeHistogramDrop { + const fn float_null_name() -> &'static str { + "prom_native_histogram_drop_float" + } + + const fn bool_false_name() -> &'static str { + "prom_native_histogram_drop_bool" + } + + pub fn float_null_udf( + message: String, + collector: Option, + ) -> ScalarUDF { + ScalarUDF::new_from_impl(NativeHistogramAnnotationUdf::new( + Self::float_null_name(), + AnnotationReturn::FloatNull, + message, + collector, + )) + } + + pub fn bool_false_udf( + message: String, + collector: Option, + ) -> ScalarUDF { + ScalarUDF::new_from_impl(NativeHistogramAnnotationUdf::new( + Self::bool_false_name(), + AnnotationReturn::BooleanFalse, + message, + collector, + )) + } +} + +// Synthetic drop UDFs only cover info-level invalid PromQL combinations. +// Warning annotations stay in the histogram UDFs that actually drop samples. +fn record_info(collector: &Option, message: impl Into) { + if let Some(collector) = collector { + collector.record_info(message); + } +} + +fn record_warning(collector: &Option, message: impl Into) { + if let Some(collector) = collector { + collector.record_warning(message); + } +} + +fn record_custom_reconciliation( + collector: &Option, + name: &'static str, + lhs: &NativeHistogram, + rhs: &NativeHistogram, +) { + if lhs.needs_custom_reconciliation(rhs) { + record_info( + collector, + format!("{name}: reconciled native histograms with different custom buckets"), + ); + } +} + +fn record_counter_reset_contradiction( + collector: &Option, + name: &'static str, + lhs: &NativeHistogram, + rhs: &NativeHistogram, +) { + if lhs.counter_reset_hints_contradict(rhs) { + record_warning( + collector, + format!("{name}: native histogram counter reset hints contradict"), + ); + } +} + +fn scalar_histogram_udf( + name: &'static str, + extra_input_types: Vec, + calc: fn(&NativeHistogram, &[ColumnarValue], usize, usize, &'static str) -> DfResult, +) -> ScalarUDF { + let mut input_types = vec![native_histogram_arrow_type()]; + input_types.extend(extra_input_types); + create_udf( + name, + input_types, + DataType::Float64, + Volatility::Volatile, + Arc::new(move |input: &[ColumnarValue]| { + if input.is_empty() { + return Err(DataFusionError::Plan(format!( + "{name} requires a native histogram argument" + ))); + } + let histograms = extract_histogram_array(&input[0], name)?; + let histograms = histograms + .as_any() + .downcast_ref::() + .expect("validated native histogram struct"); + let mut result = Float64Builder::with_capacity(histograms.len()); + for row in 0..histograms.len() { + match read_histogram(histograms, row)? { + Some(histogram) => { + result.append_value(calc(&histogram, input, row, histograms.len(), name)?) + } + None => result.append_null(), + } + } + Ok(ColumnarValue::Array(Arc::new(result.finish()))) + }) as _, + ) +} + +fn histogram_pair_udf( + name: &'static str, + op: fn(&NativeHistogram, &NativeHistogram) -> Option, +) -> ScalarUDF { + histogram_pair_udf_with_collector(name, op, None) +} + +fn histogram_pair_udf_with_collector( + name: &'static str, + op: fn(&NativeHistogram, &NativeHistogram) -> Option, + collector: Option, +) -> ScalarUDF { + create_udf( + name, + vec![native_histogram_arrow_type(), native_histogram_arrow_type()], + native_histogram_arrow_type(), + Volatility::Volatile, + Arc::new(move |input: &[ColumnarValue]| { + let lhs = extract_histogram_array(&input[0], name)?; + let rhs = extract_histogram_array(&input[1], name)?; + if lhs.len() != rhs.len() { + return Err(DataFusionError::Execution(format!( + "{name}: native histogram argument length mismatch: {} vs {}", + lhs.len(), + rhs.len() + ))); + } + + let lhs = lhs + .as_any() + .downcast_ref::() + .expect("validated native histogram struct"); + let rhs = rhs + .as_any() + .downcast_ref::() + .expect("validated native histogram struct"); + let mut result = Vec::with_capacity(lhs.len()); + for row in 0..lhs.len() { + result.push( + match (read_histogram(lhs, row)?, read_histogram(rhs, row)?) { + (Some(lhs), Some(rhs)) => { + record_custom_reconciliation(&collector, name, &lhs, &rhs); + let result = op(&lhs, &rhs); + if result.is_some() && name == NativeHistogramAdd::name() { + record_counter_reset_contradiction(&collector, name, &lhs, &rhs); + } + if result.is_none() { + record_warning( + &collector, + format!( + "{name}: dropped native histogram sample with incompatible schemas" + ), + ); + } + result + } + _ => None, + }, + ); + } + Ok(ColumnarValue::Array(build_histogram_array(&result))) + }) as _, + ) +} + +fn histogram_transform_udf( + name: &'static str, + op: fn(NativeHistogram) -> NativeHistogram, +) -> ScalarUDF { + create_udf( + name, + vec![native_histogram_arrow_type()], + native_histogram_arrow_type(), + Volatility::Volatile, + Arc::new(move |input: &[ColumnarValue]| { + let histograms = extract_histogram_array(&input[0], name)?; + let histograms = histograms + .as_any() + .downcast_ref::() + .expect("validated native histogram struct"); + let mut result = Vec::with_capacity(histograms.len()); + for row in 0..histograms.len() { + result.push(read_histogram(histograms, row)?.map(op)); + } + Ok(ColumnarValue::Array(build_histogram_array(&result))) + }) as _, + ) +} + +fn histogram_string_udf(name: &'static str) -> ScalarUDF { + create_udf( + name, + vec![native_histogram_arrow_type()], + DataType::Utf8, + Volatility::Volatile, + Arc::new(move |input: &[ColumnarValue]| { + let histograms = extract_histogram_array(&input[0], name)?; + let histograms = histograms + .as_any() + .downcast_ref::() + .expect("validated native histogram struct"); + let mut result = StringBuilder::with_capacity(histograms.len(), histograms.len() * 32); + for row in 0..histograms.len() { + match read_histogram(histograms, row)? { + Some(histogram) => result.append_value(histogram.promql_string()), + None => result.append_null(), + } + } + Ok(ColumnarValue::Array(Arc::new(result.finish()))) + }) as _, + ) +} + +fn histogram_scalar_udf( + name: &'static str, + input_types: Vec, + histogram_index: usize, + scalar_index: usize, + op: fn(NativeHistogram, f64) -> Option, +) -> ScalarUDF { + create_udf( + name, + input_types, + native_histogram_arrow_type(), + Volatility::Volatile, + Arc::new(move |input: &[ColumnarValue]| { + let histograms = extract_histogram_array(&input[histogram_index], name)?; + let histograms = histograms + .as_any() + .downcast_ref::() + .expect("validated native histogram struct"); + let mut result = Vec::with_capacity(histograms.len()); + for row in 0..histograms.len() { + result.push(match read_histogram(histograms, row)? { + Some(histogram) => { + let scalar = + read_scalar_f64_arg(&input[scalar_index], row, histograms.len(), name)?; + op(histogram, scalar) + } + None => None, + }); + } + Ok(ColumnarValue::Array(build_histogram_array(&result))) + }) as _, + ) +} + +fn histogram_compare_udf( + name: &'static str, + op: fn(&NativeHistogram, &NativeHistogram) -> bool, +) -> ScalarUDF { + create_udf( + name, + vec![native_histogram_arrow_type(), native_histogram_arrow_type()], + DataType::Boolean, + Volatility::Volatile, + Arc::new(move |input: &[ColumnarValue]| { + let lhs = extract_histogram_array(&input[0], name)?; + let rhs = extract_histogram_array(&input[1], name)?; + if lhs.len() != rhs.len() { + return Err(DataFusionError::Execution(format!( + "{name}: native histogram argument length mismatch: {} vs {}", + lhs.len(), + rhs.len() + ))); + } + + let lhs = lhs + .as_any() + .downcast_ref::() + .expect("validated native histogram struct"); + let rhs = rhs + .as_any() + .downcast_ref::() + .expect("validated native histogram struct"); + let mut result = Vec::with_capacity(lhs.len()); + for row in 0..lhs.len() { + result.push( + match (read_histogram(lhs, row)?, read_histogram(rhs, row)?) { + (Some(lhs), Some(rhs)) => Some(op(&lhs, &rhs)), + _ => None, + }, + ); + } + Ok(ColumnarValue::Array(Arc::new(BooleanArray::from(result)))) + }) as _, + ) +} + +pub struct NativeHistogramAdd; + +impl NativeHistogramAdd { + pub const fn name() -> &'static str { + "prom_native_histogram_add" + } + + pub fn scalar_udf() -> ScalarUDF { + histogram_pair_udf(Self::name(), NativeHistogram::add) + } + + pub fn scalar_udf_with_collector(collector: Option) -> ScalarUDF { + histogram_pair_udf_with_collector(Self::name(), NativeHistogram::add, collector) + } +} + +pub struct NativeHistogramSub; + +impl NativeHistogramSub { + pub const fn name() -> &'static str { + "prom_native_histogram_sub" + } + + pub fn scalar_udf() -> ScalarUDF { + histogram_pair_udf(Self::name(), NativeHistogram::sub) + } + + pub fn scalar_udf_with_collector(collector: Option) -> ScalarUDF { + histogram_pair_udf_with_collector(Self::name(), NativeHistogram::sub, collector) + } +} + +pub struct NativeHistogramMulScalar; + +impl NativeHistogramMulScalar { + pub const fn name() -> &'static str { + "prom_native_histogram_mul_scalar" + } + + pub fn scalar_udf() -> ScalarUDF { + histogram_scalar_udf( + Self::name(), + vec![native_histogram_arrow_type(), DataType::Float64], + 0, + 1, + |histogram, scalar| Some(histogram.scale(scalar)), + ) + } +} + +pub struct NativeHistogramScalarMul; + +impl NativeHistogramScalarMul { + pub const fn name() -> &'static str { + "prom_native_histogram_scalar_mul" + } + + pub fn scalar_udf() -> ScalarUDF { + histogram_scalar_udf( + Self::name(), + vec![DataType::Float64, native_histogram_arrow_type()], + 1, + 0, + |histogram, scalar| Some(histogram.scale(scalar)), + ) + } +} + +pub struct NativeHistogramDivScalar; + +impl NativeHistogramDivScalar { + pub const fn name() -> &'static str { + "prom_native_histogram_div_scalar" + } + + pub fn scalar_udf() -> ScalarUDF { + histogram_scalar_udf( + Self::name(), + vec![native_histogram_arrow_type(), DataType::Float64], + 0, + 1, + |histogram, scalar| Some(histogram.divide_by(scalar)), + ) + } +} + +pub struct NativeHistogramNeg; + +impl NativeHistogramNeg { + pub const fn name() -> &'static str { + "prom_native_histogram_neg" + } + + pub fn scalar_udf() -> ScalarUDF { + histogram_transform_udf(Self::name(), NativeHistogram::negated) + } +} + +pub struct NativeHistogramEq; + +impl NativeHistogramEq { + pub const fn name() -> &'static str { + "prom_native_histogram_eq" + } + + pub fn scalar_udf() -> ScalarUDF { + histogram_compare_udf(Self::name(), NativeHistogram::promql_eq) + } +} + +pub struct NativeHistogramNotEq; + +impl NativeHistogramNotEq { + pub const fn name() -> &'static str { + "prom_native_histogram_not_eq" + } + + pub fn scalar_udf() -> ScalarUDF { + histogram_compare_udf(Self::name(), |lhs, rhs| !lhs.promql_eq(rhs)) + } +} + +pub struct NativeHistogramCount; + +impl NativeHistogramCount { + pub const fn name() -> &'static str { + "prom_native_histogram_count" + } + + pub fn scalar_udf() -> ScalarUDF { + scalar_histogram_udf(Self::name(), vec![], |histogram, _, _, _, _| { + Ok(histogram.count) + }) + } +} + +pub struct NativeHistogramSum; + +impl NativeHistogramSum { + pub const fn name() -> &'static str { + "prom_native_histogram_sum" + } + + pub fn scalar_udf() -> ScalarUDF { + scalar_histogram_udf(Self::name(), vec![], |histogram, _, _, _, _| { + Ok(histogram.sum) + }) + } +} + +pub struct NativeHistogramAvg; + +impl NativeHistogramAvg { + pub const fn name() -> &'static str { + "prom_native_histogram_avg" + } + + pub fn scalar_udf() -> ScalarUDF { + scalar_histogram_udf(Self::name(), vec![], |histogram, _, _, _, _| { + Ok(histogram.sum / histogram.count) + }) + } +} + +pub struct NativeHistogramStddev; + +impl NativeHistogramStddev { + pub const fn name() -> &'static str { + "prom_native_histogram_stddev" + } + + pub fn scalar_udf() -> ScalarUDF { + scalar_histogram_udf(Self::name(), vec![], |histogram, _, _, _, _| { + Ok(histogram.estimated_stddev()) + }) + } +} + +pub struct NativeHistogramStdvar; + +impl NativeHistogramStdvar { + pub const fn name() -> &'static str { + "prom_native_histogram_stdvar" + } + + pub fn scalar_udf() -> ScalarUDF { + scalar_histogram_udf(Self::name(), vec![], |histogram, _, _, _, _| { + Ok(histogram.estimated_stdvar()) + }) + } +} + +pub struct NativeHistogramToString; + +impl NativeHistogramToString { + pub const fn name() -> &'static str { + "prom_native_histogram_to_string" + } + + pub fn scalar_udf() -> ScalarUDF { + histogram_string_udf(Self::name()) + } +} + +pub struct NativeHistogramQuantile; + +impl NativeHistogramQuantile { + pub const fn name() -> &'static str { + "prom_native_histogram_quantile" + } + + pub fn scalar_udf() -> ScalarUDF { + scalar_histogram_udf( + Self::name(), + vec![DataType::Float64], + |histogram, input, row, len, name| { + Ok(histogram.quantile(read_scalar_f64_arg(&input[1], row, len, name)?)) + }, + ) + } +} + +pub struct NativeHistogramFraction; + +impl NativeHistogramFraction { + pub const fn name() -> &'static str { + "prom_native_histogram_fraction" + } + + pub fn scalar_udf() -> ScalarUDF { + scalar_histogram_udf( + Self::name(), + vec![DataType::Float64, DataType::Float64], + |histogram, input, row, len, name| { + let lower = read_scalar_f64_arg(&input[1], row, len, name)?; + let upper = read_scalar_f64_arg(&input[2], row, len, name)?; + Ok(histogram.fraction(lower, upper)) + }, + ) + } +} + +#[derive(Debug, Clone, Copy)] +enum NativeHistogramAggregateKind { + Sum, + Avg, +} + +impl NativeHistogramAggregateKind { + const fn name(self) -> &'static str { + match self { + Self::Sum => NativeHistogramAggSum::name(), + Self::Avg => NativeHistogramAggAvg::name(), + } + } + + const fn needs_count(self) -> bool { + matches!(self, Self::Avg) + } +} + +pub struct NativeHistogramAggSum; + +impl NativeHistogramAggSum { + pub const fn name() -> &'static str { + "prom_native_histogram_agg_sum" + } + + pub fn aggregate_udf() -> AggregateUDF { + native_histogram_aggregate_udf(NativeHistogramAggregateKind::Sum, None) + } + + pub fn aggregate_udf_with_collector( + collector: Option, + ) -> AggregateUDF { + native_histogram_aggregate_udf(NativeHistogramAggregateKind::Sum, collector) + } +} + +pub struct NativeHistogramAggAvg; + +impl NativeHistogramAggAvg { + pub const fn name() -> &'static str { + "prom_native_histogram_agg_avg" + } + + pub fn aggregate_udf() -> AggregateUDF { + native_histogram_aggregate_udf(NativeHistogramAggregateKind::Avg, None) + } + + pub fn aggregate_udf_with_collector( + collector: Option, + ) -> AggregateUDF { + native_histogram_aggregate_udf(NativeHistogramAggregateKind::Avg, collector) + } +} + +#[derive(Debug)] +struct NativeHistogramAggregateAccumulator { + kind: NativeHistogramAggregateKind, + value: Option, + count: u64, + dropped_incompatible: bool, + collector: Option, +} + +impl NativeHistogramAggregateAccumulator { + fn new( + kind: NativeHistogramAggregateKind, + collector: Option, + ) -> Self { + Self { + kind, + value: None, + count: 0, + dropped_incompatible: false, + collector, + } + } + + fn from_args( + kind: NativeHistogramAggregateKind, + collector: Option, + _args: AccumulatorArgs, + ) -> DfResult> { + Ok(Box::new(Self::new(kind, collector))) + } + + fn push_histogram(&mut self, histogram: NativeHistogram, count: u64) { + if self.dropped_incompatible { + return; + } + + let value = match self.value.take() { + Some(value) => { + record_custom_reconciliation(&self.collector, self.kind.name(), &value, &histogram); + record_counter_reset_contradiction( + &self.collector, + self.kind.name(), + &value, + &histogram, + ); + match value.add(&histogram) { + Some(value) => Some(value), + None => { + self.record_incompatible(); + None + } + } + } + None => Some(histogram), + }; + if !self.dropped_incompatible { + self.value = value; + self.count += count; + } + } + + fn mark_incompatible(&mut self) { + self.value = None; + self.count = 0; + self.dropped_incompatible = true; + } + + fn record_incompatible(&mut self) { + self.mark_incompatible(); + record_warning( + &self.collector, + format!( + "{}: dropped native histogram aggregate with incompatible schemas", + self.kind.name() + ), + ); + } +} + +fn range_sum_histograms( + samples: Vec, + name: &'static str, + collector: &Option, +) -> Option { + let mut value = None; + for histogram in samples { + value = match value { + Some(value) => { + record_custom_reconciliation(collector, name, &value, &histogram); + record_counter_reset_contradiction(collector, name, &value, &histogram); + match value.add(&histogram) { + Some(value) => Some(value), + None => { + record_warning( + collector, + format!( + "{name}: dropped native histogram range with incompatible schemas" + ), + ); + return None; + } + } + } + None => Some(histogram), + }; + } + value +} + +#[derive(Debug, Clone, Copy)] +enum NativeHistogramRangeHistogramKind { + Sum, + Avg, + Last, +} + +#[derive(Debug, Clone, Copy)] +enum NativeHistogramRangeFloatKind { + Absent, + Count, + Present, + Changes, + Resets, +} + +fn collect_window_histograms( + histograms: &StructArray, + offset: usize, + length: usize, +) -> DfResult>> { + let mut samples = Vec::with_capacity(length); + for row in offset..offset + length { + let Some(histogram) = read_histogram(histograms, row)? else { + return Ok(None); + }; + samples.push(histogram); + } + Ok(Some(samples)) +} + +fn native_histogram_range_histogram( + input: &[ColumnarValue], + kind: NativeHistogramRangeHistogramKind, + func_name: &'static str, + collector: Option, +) -> DfResult { + if input.len() != 2 { + return Err(DataFusionError::Plan(format!( + "{func_name} function should have 2 inputs" + ))); + } + + let ts_range = extract_range_dict( + &input[0], + func_name, + "timestamp range vector", + &DataType::Timestamp(TimeUnit::Millisecond, None), + )?; + let value_range = extract_range_dict( + &input[1], + func_name, + "value range vector", + &native_histogram_arrow_type(), + )?; + if ts_range.keys().values() != value_range.keys().values() { + return Err(DataFusionError::Execution(format!( + "{func_name}: timestamp and value ranges should have the same window layout" + ))); + } + + let histograms = value_range + .values() + .as_any() + .downcast_ref::() + .expect("validated native histogram range"); + let mut result = Vec::with_capacity(value_range.keys().len()); + for key in value_range.keys().values() { + let (offset, length) = unpack(*key); + let offset = offset as usize; + let length = length as usize; + if length == 0 { + result.push(None); + continue; + } + let Some(samples) = collect_window_histograms(histograms, offset, length)? else { + result.push(None); + continue; + }; + let histogram = match kind { + NativeHistogramRangeHistogramKind::Sum => { + range_sum_histograms(samples, func_name, &collector) + .map(NativeHistogram::into_gauge) + } + NativeHistogramRangeHistogramKind::Avg => { + let count = samples.len(); + range_sum_histograms(samples, func_name, &collector) + .map(|histogram| histogram.divide_by(count as f64).into_gauge()) + } + NativeHistogramRangeHistogramKind::Last => { + samples.last().cloned().map(NativeHistogram::into_gauge) + } + }; + result.push(histogram); + } + + Ok(ColumnarValue::Array(build_histogram_array(&result))) +} + +fn native_histogram_range_float( + input: &[ColumnarValue], + kind: NativeHistogramRangeFloatKind, + func_name: &'static str, +) -> DfResult { + if input.len() != 2 { + return Err(DataFusionError::Plan(format!( + "{func_name} function should have 2 inputs" + ))); + } + + let ts_range = extract_range_dict( + &input[0], + func_name, + "timestamp range vector", + &DataType::Timestamp(TimeUnit::Millisecond, None), + )?; + let value_range = extract_range_dict( + &input[1], + func_name, + "value range vector", + &native_histogram_arrow_type(), + )?; + if ts_range.keys().values() != value_range.keys().values() { + return Err(DataFusionError::Execution(format!( + "{func_name}: timestamp and value ranges should have the same window layout" + ))); + } + + let timestamps = ts_range + .values() + .as_any() + .downcast_ref::() + .expect("validated timestamp range") + .values(); + let histograms = value_range + .values() + .as_any() + .downcast_ref::() + .expect("validated native histogram range"); + let mut result = Float64Builder::with_capacity(value_range.keys().len()); + for key in value_range.keys().values() { + let (offset, length) = unpack(*key); + let offset = offset as usize; + let length = length as usize; + if length == 0 { + match kind { + NativeHistogramRangeFloatKind::Absent => result.append_value(1.0), + _ => result.append_null(), + } + continue; + } + let Some(samples) = collect_window_histograms(histograms, offset, length)? else { + result.append_null(); + continue; + }; + let value = match kind { + NativeHistogramRangeFloatKind::Absent => { + result.append_null(); + continue; + } + NativeHistogramRangeFloatKind::Count => length as f64, + NativeHistogramRangeFloatKind::Present => 1.0, + NativeHistogramRangeFloatKind::Changes => samples + .windows(2) + .filter(|pair| !pair[0].promql_eq(&pair[1])) + .count() as f64, + NativeHistogramRangeFloatKind::Resets => samples + .windows(2) + .zip(timestamps[offset..offset + length].windows(2)) + .filter(|(pair, ts_pair)| { + pair[1].detect_counter_reset(&pair[0], ts_pair[0], ts_pair[1]) + }) + .count() as f64, + }; + result.append_value(value); + } + + Ok(ColumnarValue::Array(Arc::new(result.finish()))) +} + +fn create_native_range_histogram_udf( + name: &'static str, + kind: NativeHistogramRangeHistogramKind, + collector: Option, +) -> ScalarUDF { + create_udf( + name, + vec![ + RangeArray::convert_data_type(DataType::Timestamp(TimeUnit::Millisecond, None)), + RangeArray::convert_data_type(native_histogram_arrow_type()), + ], + native_histogram_arrow_type(), + Volatility::Volatile, + Arc::new(move |input: &[ColumnarValue]| { + native_histogram_range_histogram(input, kind, name, collector.clone()) + }) as _, + ) +} + +fn create_native_range_float_udf( + name: &'static str, + kind: NativeHistogramRangeFloatKind, +) -> ScalarUDF { + create_udf( + name, + vec![ + RangeArray::convert_data_type(DataType::Timestamp(TimeUnit::Millisecond, None)), + RangeArray::convert_data_type(native_histogram_arrow_type()), + ], + DataType::Float64, + Volatility::Volatile, + Arc::new(move |input: &[ColumnarValue]| native_histogram_range_float(input, kind, name)) + as _, + ) +} + +pub struct NativeHistogramSumOverTime; +pub struct NativeHistogramAvgOverTime; +pub struct NativeHistogramAbsentOverTime; +pub struct NativeHistogramCountOverTime; +pub struct NativeHistogramLastOverTime; +pub struct NativeHistogramPresentOverTime; +pub struct NativeHistogramChanges; +pub struct NativeHistogramResets; + +impl NativeHistogramSumOverTime { + pub const fn name() -> &'static str { + "prom_native_histogram_sum_over_time" + } + + pub fn scalar_udf() -> ScalarUDF { + Self::scalar_udf_with_collector(None) + } + + pub fn scalar_udf_with_collector(collector: Option) -> ScalarUDF { + create_native_range_histogram_udf( + Self::name(), + NativeHistogramRangeHistogramKind::Sum, + collector, + ) + } +} + +impl NativeHistogramAvgOverTime { + pub const fn name() -> &'static str { + "prom_native_histogram_avg_over_time" + } + + pub fn scalar_udf() -> ScalarUDF { + Self::scalar_udf_with_collector(None) + } + + pub fn scalar_udf_with_collector(collector: Option) -> ScalarUDF { + create_native_range_histogram_udf( + Self::name(), + NativeHistogramRangeHistogramKind::Avg, + collector, + ) + } +} + +impl NativeHistogramAbsentOverTime { + pub const fn name() -> &'static str { + "prom_native_histogram_absent_over_time" + } + + pub fn scalar_udf() -> ScalarUDF { + create_native_range_float_udf(Self::name(), NativeHistogramRangeFloatKind::Absent) + } +} + +impl NativeHistogramCountOverTime { + pub const fn name() -> &'static str { + "prom_native_histogram_count_over_time" + } + + pub fn scalar_udf() -> ScalarUDF { + create_native_range_float_udf(Self::name(), NativeHistogramRangeFloatKind::Count) + } +} + +impl NativeHistogramLastOverTime { + pub const fn name() -> &'static str { + "prom_native_histogram_last_over_time" + } + + pub fn scalar_udf() -> ScalarUDF { + create_native_range_histogram_udf( + Self::name(), + NativeHistogramRangeHistogramKind::Last, + None, + ) + } +} + +impl NativeHistogramPresentOverTime { + pub const fn name() -> &'static str { + "prom_native_histogram_present_over_time" + } + + pub fn scalar_udf() -> ScalarUDF { + create_native_range_float_udf(Self::name(), NativeHistogramRangeFloatKind::Present) + } +} + +impl NativeHistogramChanges { + pub const fn name() -> &'static str { + "prom_native_histogram_changes" + } + + pub fn scalar_udf() -> ScalarUDF { + create_native_range_float_udf(Self::name(), NativeHistogramRangeFloatKind::Changes) + } +} + +impl NativeHistogramResets { + pub const fn name() -> &'static str { + "prom_native_histogram_resets" + } + + pub fn scalar_udf() -> ScalarUDF { + create_native_range_float_udf(Self::name(), NativeHistogramRangeFloatKind::Resets) + } +} + +fn native_histogram_scalar(histogram: Option) -> ScalarValue { + let array = build_histogram_array(&[histogram]); + let histogram = array + .as_any() + .downcast_ref::() + .expect("native histogram array is a StructArray") + .clone(); + ScalarValue::Struct(Arc::new(histogram)) +} + +fn native_histogram_aggregate_udf( + kind: NativeHistogramAggregateKind, + collector: Option, +) -> AggregateUDF { + let state_types = if kind.needs_count() { + vec![ + native_histogram_arrow_type(), + DataType::UInt64, + DataType::Boolean, + ] + } else { + vec![native_histogram_arrow_type(), DataType::Boolean] + }; + + create_udaf( + kind.name(), + vec![native_histogram_arrow_type()], + Arc::new(native_histogram_arrow_type()), + Volatility::Volatile, + Arc::new(move |args| { + NativeHistogramAggregateAccumulator::from_args(kind, collector.clone(), args) + }), + Arc::new(state_types), + ) +} + +impl DfAccumulator for NativeHistogramAggregateAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> DfResult<()> { + let histograms = values + .first() + .and_then(|array| array.as_any().downcast_ref::()) + .ok_or_else(|| { + DataFusionError::Execution(format!( + "{}: expected native histogram struct input", + self.kind.name() + )) + })?; + + for row in 0..histograms.len() { + let Some(histogram) = read_histogram(histograms, row)? else { + continue; + }; + self.push_histogram(histogram, 1); + } + + Ok(()) + } + + fn evaluate(&mut self) -> DfResult { + let histogram = match (self.kind, self.dropped_incompatible, self.value.clone()) { + (_, true, _) => None, + (_, false, None) => None, + (NativeHistogramAggregateKind::Sum, false, value) => value, + (NativeHistogramAggregateKind::Avg, false, Some(value)) if self.count > 0 => { + Some(value.divide_by(self.count as f64)) + } + (NativeHistogramAggregateKind::Avg, _, _) => None, + }; + + Ok(native_histogram_scalar(histogram)) + } + + fn size(&self) -> usize { + std::mem::size_of::() + } + + fn state(&mut self) -> DfResult> { + let mut state = vec![native_histogram_scalar(self.value.clone())]; + if self.kind.needs_count() { + state.push(ScalarValue::UInt64(Some(self.count))); + } + state.push(ScalarValue::Boolean(Some(self.dropped_incompatible))); + Ok(state) + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> DfResult<()> { + if states.is_empty() { + return Ok(()); + } + + let histograms = states[0] + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Execution(format!( + "{}: expected native histogram struct state", + self.kind.name() + )) + })?; + let counts = if self.kind.needs_count() { + Some( + states + .get(1) + .and_then(|array| array.as_any().downcast_ref::()) + .ok_or_else(|| { + DataFusionError::Execution(format!( + "{}: expected UInt64 count state", + self.kind.name() + )) + })?, + ) + } else { + None + }; + let dropped_index = if self.kind.needs_count() { 2 } else { 1 }; + let dropped = states + .get(dropped_index) + .and_then(|array| array.as_any().downcast_ref::()) + .ok_or_else(|| { + DataFusionError::Execution(format!( + "{}: expected Boolean dropped state", + self.kind.name() + )) + })?; + + for row in 0..histograms.len() { + if dropped.value(row) { + self.mark_incompatible(); + continue; + } + let Some(histogram) = read_histogram(histograms, row)? else { + continue; + }; + let count = counts.map(|counts| counts.value(row)).unwrap_or(1); + self.push_histogram(histogram, count); + } + + Ok(()) + } +} + +fn histogram_delta( + samples: &[NativeHistogram], + timestamps: &[i64], + is_counter: bool, +) -> Option { + if samples.len() < 2 || samples.len() != timestamps.len() { + return None; + } + + if !is_counter { + return samples + .last()? + .sub(samples.first()?) + .map(NativeHistogram::into_gauge); + } + + let mut result = samples.first()?.zero_like(); + let mut segment_start = samples.first()?; + for (sample_pair, ts_pair) in samples.windows(2).zip(timestamps.windows(2)) { + let previous_ts = ts_pair[0]; + let current_ts = ts_pair[1]; + let previous = &sample_pair[0]; + let current = &sample_pair[1]; + if current.detect_counter_reset(previous, previous_ts, current_ts) { + result = result.add(&previous.sub(segment_start)?)?; + result = result.add(current)?; + segment_start = current; + } + } + Some( + result + .add(&samples.last()?.sub(segment_start)?)? + .into_gauge(), + ) +} + +fn idelta_value( + samples: &[NativeHistogram], + is_rate: bool, + previous_ts: i64, + current_ts: i64, + sampled_interval_secs: f64, +) -> Option { + if samples.len() < 2 { + return None; + } + let previous = &samples[samples.len() - 2]; + let current = samples.last()?; + let result = if is_rate && current.detect_counter_reset(previous, previous_ts, current_ts) { + current.clone() + } else { + current.sub(previous)? + }; + Some( + if is_rate { + result.scale(1.0 / sampled_interval_secs) + } else { + result + } + .into_gauge(), + ) +} + +fn extract_range_dict( + columnar_value: &ColumnarValue, + func_name: &str, + arg_name: &str, + expected_value_type: &DataType, +) -> DfResult> { + let array = extract_array(columnar_value)?; + let dict = array + .as_any() + .downcast_ref::>() + .ok_or_else(|| { + DataFusionError::Execution(format!( + "{func_name}: expect {arg_name} as DictionaryArray, found {}", + array.data_type() + )) + })? + .clone(); + + if &dict.value_type() != expected_value_type { + return Err(DataFusionError::Execution(format!( + "{func_name}: expect {arg_name} values of type {expected_value_type}, found {}", + dict.value_type() + ))); + } + + RangeArray::try_new(dict.clone()).map_err(DataFusionError::from)?; + Ok(dict) +} + +fn native_extrapolated_rate( + input: &[ColumnarValue], + range_length: i64, + func_name: &'static str, + collector: Option, +) -> DfResult { + if input.len() != 4 { + return Err(DataFusionError::Plan(format!( + "{func_name} function should have 4 inputs" + ))); + } + + let ts_dict = extract_range_dict( + &input[0], + func_name, + "timestamp range vector", + &DataType::Timestamp(TimeUnit::Millisecond, None), + )?; + let value_dict = extract_range_dict( + &input[1], + func_name, + "value range vector", + &native_histogram_arrow_type(), + )?; + let eval_ts = extract_array(&input[2])?; + let eval_ts = eval_ts + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Execution(format!( + "{func_name}: expect evaluation timestamp vector as Timestamp(Millisecond), found {}", + eval_ts.data_type() + )) + })?; + + let keys = ts_dict.keys().values(); + if value_dict.keys().values() != keys || eval_ts.len() != keys.len() { + return Err(DataFusionError::Execution(format!( + "{func_name}: timestamp, value, and evaluation ranges should have the same layout" + ))); + } + + let all_timestamps = ts_dict + .values() + .as_any() + .downcast_ref::() + .expect("validated timestamp range") + .values(); + let all_histograms = value_dict + .values() + .as_any() + .downcast_ref::() + .expect("validated native histogram range"); + let range_length_secs = range_length as f64 / 1000.0; + let mut result = Vec::with_capacity(keys.len()); + + for index in 0..keys.len() { + let (raw_offset, raw_length) = unpack(keys[index]); + let offset = raw_offset as usize; + let length = raw_length as usize; + if length < 2 { + result.push(None); + continue; + } + + let mut samples = Vec::with_capacity(length); + let mut has_null = false; + for row in offset..offset + length { + let Some(histogram) = read_histogram(all_histograms, row)? else { + has_null = true; + break; + }; + samples.push(histogram); + } + if has_null { + result.push(None); + continue; + } + + let timestamps = &all_timestamps[offset..offset + length]; + for pair in samples.windows(2) { + record_custom_reconciliation(&collector, func_name, &pair[0], &pair[1]); + } + let Some(mut histogram) = histogram_delta(&samples, timestamps, IS_COUNTER) else { + record_warning( + &collector, + format!("{func_name}: dropped native histogram range with incompatible schemas"), + ); + result.push(None); + continue; + }; + + let first_ts = all_timestamps[offset]; + let last_ts = all_timestamps[offset + length - 1]; + let range_end = eval_ts.value(index); + let range_start = range_end - range_length; + let sampled_interval_ms = (last_ts - first_ts) as f64; + if sampled_interval_ms <= 0.0 { + result.push(None); + continue; + } + let average_interval_ms = sampled_interval_ms / (length - 1) as f64; + let mut duration_to_start_ms = (first_ts - range_start) as f64; + let duration_to_end_ms = (range_end - last_ts) as f64; + + if IS_COUNTER && histogram.count > 0.0 && samples[0].count >= 0.0 { + let duration_to_zero = sampled_interval_ms * (samples[0].count / histogram.count); + if duration_to_zero < duration_to_start_ms { + duration_to_start_ms = duration_to_zero; + } + } + + let extrapolation_threshold = average_interval_ms * 1.1; + let mut extrapolated_interval_ms = sampled_interval_ms; + if duration_to_start_ms < extrapolation_threshold { + extrapolated_interval_ms += duration_to_start_ms; + } else { + extrapolated_interval_ms += average_interval_ms / 2.0; + } + if duration_to_end_ms < extrapolation_threshold { + extrapolated_interval_ms += duration_to_end_ms; + } else { + extrapolated_interval_ms += average_interval_ms / 2.0; + } + + let mut factor = extrapolated_interval_ms / sampled_interval_ms; + if IS_RATE { + factor /= range_length_secs; + } + histogram = histogram.scale(factor).into_gauge(); + result.push(Some(histogram)); + } + + Ok(ColumnarValue::Array(build_histogram_array(&result))) +} + +fn create_native_extrapolated_udf( + name: &'static str, + collector: Option, +) -> ScalarUDF { + let input_types = vec![ + RangeArray::convert_data_type(DataType::Timestamp(TimeUnit::Millisecond, None)), + RangeArray::convert_data_type(native_histogram_arrow_type()), + DataType::Timestamp(TimeUnit::Millisecond, None), + DataType::Int64, + ]; + create_udf( + name, + input_types, + native_histogram_arrow_type(), + Volatility::Volatile, + Arc::new(move |input: &[ColumnarValue]| { + let range_length = extract_array(&input[3])?; + let range_length = range_length + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Execution(format!( + "{name}: expect Int64 as range length type, found {}", + range_length.data_type() + )) + })?; + if range_length.is_empty() || range_length.is_null(0) { + return Err(DataFusionError::Execution(format!( + "{name}: range length must contain a non-null Int64 value" + ))); + } + native_extrapolated_rate::( + input, + range_length.value(0), + name, + collector.clone(), + ) + }) as _, + ) +} + +pub struct NativeHistogramDelta; +pub struct NativeHistogramRate; +pub struct NativeHistogramIncrease; + +impl NativeHistogramDelta { + pub const fn name() -> &'static str { + "prom_native_histogram_delta" + } + + pub fn scalar_udf() -> ScalarUDF { + Self::scalar_udf_with_collector(None) + } + + pub fn scalar_udf_with_collector(collector: Option) -> ScalarUDF { + create_native_extrapolated_udf::(Self::name(), collector) + } +} + +impl NativeHistogramRate { + pub const fn name() -> &'static str { + "prom_native_histogram_rate" + } + + pub fn scalar_udf() -> ScalarUDF { + Self::scalar_udf_with_collector(None) + } + + pub fn scalar_udf_with_collector(collector: Option) -> ScalarUDF { + create_native_extrapolated_udf::(Self::name(), collector) + } +} + +impl NativeHistogramIncrease { + pub const fn name() -> &'static str { + "prom_native_histogram_increase" + } + + pub fn scalar_udf() -> ScalarUDF { + Self::scalar_udf_with_collector(None) + } + + pub fn scalar_udf_with_collector(collector: Option) -> ScalarUDF { + create_native_extrapolated_udf::(Self::name(), collector) + } +} + +fn native_idelta( + input: &[ColumnarValue], + func_name: &'static str, + collector: Option, +) -> DfResult { + if input.len() != 2 { + return Err(DataFusionError::Plan(format!( + "{func_name} function should have 2 inputs" + ))); + } + + let ts_range = extract_range_dict( + &input[0], + func_name, + "timestamp range vector", + &DataType::Timestamp(TimeUnit::Millisecond, None), + )?; + let value_range = extract_range_dict( + &input[1], + func_name, + "value range vector", + &native_histogram_arrow_type(), + )?; + + if ts_range.keys().values() != value_range.keys().values() { + return Err(DataFusionError::Execution(format!( + "{func_name}: timestamp and value ranges should have the same window layout" + ))); + } + + let ts_values = ts_range + .values() + .as_any() + .downcast_ref::() + .expect("validated timestamp range") + .values(); + let histograms = value_range + .values() + .as_any() + .downcast_ref::() + .expect("validated native histogram range"); + let mut result = Vec::with_capacity(ts_range.keys().len()); + + for key in ts_range.keys().values() { + let (offset, length) = unpack(*key); + let offset = offset as usize; + let length = length as usize; + if length < 2 { + result.push(None); + continue; + } + + let mut samples = Vec::with_capacity(2); + let mut has_null = false; + for row in offset + length - 2..offset + length { + let Some(histogram) = read_histogram(histograms, row)? else { + has_null = true; + break; + }; + samples.push(histogram); + } + if has_null { + result.push(None); + continue; + } + + let sampled_interval_secs = + (ts_values[offset + length - 1] - ts_values[offset + length - 2]) as f64 / 1000.0; + if sampled_interval_secs <= 0.0 { + result.push(None); + continue; + } + record_custom_reconciliation(&collector, func_name, &samples[0], &samples[1]); + let value = idelta_value( + &samples, + IS_RATE, + ts_values[offset + length - 2], + ts_values[offset + length - 1], + sampled_interval_secs, + ); + if value.is_none() { + record_warning( + &collector, + format!("{func_name}: dropped native histogram range with incompatible schemas"), + ); + } + result.push(value); + } + + Ok(ColumnarValue::Array(build_histogram_array(&result))) +} + +fn create_native_idelta_udf( + name: &'static str, + collector: Option, +) -> ScalarUDF { + create_udf( + name, + vec![ + RangeArray::convert_data_type(DataType::Timestamp(TimeUnit::Millisecond, None)), + RangeArray::convert_data_type(native_histogram_arrow_type()), + ], + native_histogram_arrow_type(), + Volatility::Volatile, + Arc::new(move |input: &[ColumnarValue]| { + native_idelta::(input, name, collector.clone()) + }) as _, + ) +} + +pub struct NativeHistogramIDelta; +pub struct NativeHistogramIRate; + +impl NativeHistogramIDelta { + pub const fn name() -> &'static str { + "prom_native_histogram_idelta" + } + + pub fn scalar_udf() -> ScalarUDF { + Self::scalar_udf_with_collector(None) + } + + pub fn scalar_udf_with_collector(collector: Option) -> ScalarUDF { + create_native_idelta_udf::(Self::name(), collector) + } +} + +impl NativeHistogramIRate { + pub const fn name() -> &'static str { + "prom_native_histogram_irate" + } + + pub fn scalar_udf() -> ScalarUDF { + Self::scalar_udf_with_collector(None) + } + + pub fn scalar_udf_with_collector(collector: Option) -> ScalarUDF { + create_native_idelta_udf::(Self::name(), collector) + } +} + +#[cfg(test)] +mod tests { + use datafusion::arrow::datatypes::Field; + use datafusion_common::config::ConfigOptions; + use datafusion_expr::ScalarFunctionArgs; + + use super::*; + + fn sample_histogram(count: f64, sum: f64, positive_buckets: Vec) -> NativeHistogram { + NativeHistogram { + schema: 0, + zero_threshold: 0.0, + sum, + reset_hint: 0, + start_timestamp: None, + custom_values: Vec::new(), + positive_spans: vec![Span { + offset: 0, + length: positive_buckets.len() as u32, + }], + negative_spans: Vec::new(), + count, + zero_count: 0.0, + positive_buckets, + negative_buckets: Vec::new(), + } + } + + fn run_scalar_udf(udf: ScalarUDF, input: Vec) -> f64 { + let result = run_udf(udf, input, DataType::Float64); + extract_array(&result) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + } + + fn run_udf(udf: ScalarUDF, input: Vec, return_type: DataType) -> ColumnarValue { + let arg_fields = input + .iter() + .enumerate() + .map(|(idx, input)| Arc::new(Field::new(format!("arg_{idx}"), input.data_type(), true))) + .collect(); + let args = ScalarFunctionArgs { + args: input, + arg_fields, + number_rows: 1, + return_field: Arc::new(Field::new("result", return_type, true)), + config_options: Arc::new(ConfigOptions::default()), + }; + + udf.invoke_with_args(args).unwrap() + } + + fn run_histogram_udf(udf: ScalarUDF, input: Vec) -> NativeHistogram { + let result = run_udf(udf, input, native_histogram_arrow_type()); + let array = extract_array(&result) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + read_histogram(&array, 0).unwrap().unwrap() + } + + fn run_histogram_range_udf( + udf: ScalarUDF, + histograms: Vec, + ) -> NativeHistogram { + let values = histograms.into_iter().map(Some).collect::>(); + let timestamps = Arc::new(TimestampMillisecondArray::from_iter( + (0..values.len()).map(|idx| Some((idx as i64 + 1) * 1000)), + )); + let histograms = build_histogram_array(&values); + let range = [(0, values.len() as u32)]; + let ts_range = RangeArray::from_ranges(timestamps, range).unwrap(); + let value_range = RangeArray::from_ranges(histograms, range).unwrap(); + + run_histogram_udf( + udf, + vec![ + ColumnarValue::Array(Arc::new(ts_range.into_dict())), + ColumnarValue::Array(Arc::new(value_range.into_dict())), + ], + ) + } + + #[test] + fn count_sum_and_avg_read_struct() { + let histograms = vec![Some(sample_histogram(6.0, 10.0, vec![2.0, 4.0]))]; + let array = build_histogram_array(&histograms); + let input = vec![ColumnarValue::Array(array)]; + + let count = run_scalar_udf(NativeHistogramCount::scalar_udf(), input.clone()); + assert_eq!(count, 6.0); + + let sum = run_scalar_udf(NativeHistogramSum::scalar_udf(), input.clone()); + assert_eq!(sum, 10.0); + + let avg = run_scalar_udf(NativeHistogramAvg::scalar_udf(), input); + assert_eq!(avg, 10.0 / 6.0); + } + + #[test] + fn quantile_uses_bucket_bounds() { + let histogram = sample_histogram(6.0, 10.0, vec![2.0, 4.0]); + assert_eq!(histogram.quantile(0.0), 0.5); + assert!(histogram.quantile(0.5) > 1.0); + assert!(histogram.quantile(0.5) < 2.0); + } + + #[test] + fn comparison_uses_promql_histogram_equality() { + let mut left = sample_histogram(1.0, 1.0, vec![1.0, 0.0]); + left.reset_hint = COUNTER_RESET_HINT; + left.start_timestamp = Some(1000); + let mut right = sample_histogram(1.0, 1.0, vec![1.0]); + right.reset_hint = NOT_COUNTER_RESET_HINT; + right.start_timestamp = Some(2000); + + let result = run_udf( + NativeHistogramEq::scalar_udf(), + vec![ + ColumnarValue::Array(build_histogram_array(&[Some(left)])), + ColumnarValue::Array(build_histogram_array(&[Some(right)])), + ], + DataType::Boolean, + ); + let values = extract_array(&result).unwrap(); + let values = values.as_any().downcast_ref::().unwrap(); + assert!(values.value(0)); + } + + #[test] + fn unary_minus_returns_gauge_histogram() { + let result = run_histogram_udf( + NativeHistogramNeg::scalar_udf(), + vec![ColumnarValue::Array(build_histogram_array(&[Some( + sample_histogram(2.0, 3.0, vec![2.0]), + )]))], + ); + + assert_eq!(result.reset_hint, GAUGE_RESET_HINT); + assert_eq!(result.count, -2.0); + assert_eq!(result.sum, -3.0); + assert_eq!(result.positive_buckets, vec![-2.0]); + } + + #[test] + fn histogram_over_time_functions_return_gauge_histograms() { + let mut first = sample_histogram(1.0, 1.0, vec![1.0]); + first.reset_hint = COUNTER_RESET_HINT; + let mut second = sample_histogram(2.0, 2.0, vec![2.0]); + second.reset_hint = COUNTER_RESET_HINT; + + let result = run_histogram_range_udf( + NativeHistogramSumOverTime::scalar_udf(), + vec![first.clone(), second.clone()], + ); + assert_eq!(result.reset_hint, GAUGE_RESET_HINT); + assert_eq!(result.count, 3.0); + assert_eq!(result.sum, 3.0); + assert_eq!(result.positive_buckets, vec![3.0]); + + let result = run_histogram_range_udf( + NativeHistogramAvgOverTime::scalar_udf(), + vec![first.clone(), second.clone()], + ); + assert_eq!(result.reset_hint, GAUGE_RESET_HINT); + assert_eq!(result.count, 1.5); + assert_eq!(result.sum, 1.5); + assert_eq!(result.positive_buckets, vec![1.5]); + + let result = run_histogram_range_udf( + NativeHistogramLastOverTime::scalar_udf(), + vec![first, second], + ); + assert_eq!(result.reset_hint, GAUGE_RESET_HINT); + assert_eq!(result.count, 2.0); + assert_eq!(result.sum, 2.0); + assert_eq!(result.positive_buckets, vec![2.0]); + } + + #[test] + fn absent_over_time_handles_histogram_ranges() { + let values = vec![Some(sample_histogram(1.0, 1.0, vec![1.0]))]; + let timestamps = Arc::new(TimestampMillisecondArray::from_iter([Some(1000)])); + let histograms = build_histogram_array(&values); + let ranges = [(0, 1), (0, 0)]; + let ts_range = RangeArray::from_ranges(timestamps, ranges).unwrap(); + let value_range = RangeArray::from_ranges(histograms, ranges).unwrap(); + + let result = run_udf( + NativeHistogramAbsentOverTime::scalar_udf(), + vec![ + ColumnarValue::Array(Arc::new(ts_range.into_dict())), + ColumnarValue::Array(Arc::new(value_range.into_dict())), + ], + DataType::Float64, + ); + let result = extract_array(&result).unwrap(); + let result = result.as_any().downcast_ref::().unwrap(); + + assert!(result.is_null(0)); + assert_eq!(result.value(1), 1.0); + } + + #[test] + fn delta_requires_exact_layout() { + let first = sample_histogram(2.0, 3.0, vec![1.0, 1.0]); + let last = sample_histogram(5.0, 8.0, vec![2.0, 3.0]); + let delta = histogram_delta(&[first, last], &[0, 1], false).unwrap(); + assert_eq!(delta.count, 3.0); + assert_eq!(delta.sum, 5.0); + assert_eq!(delta.reset_hint, GAUGE_RESET_HINT); + assert_eq!(delta.positive_buckets, vec![1.0, 2.0]); + } + + #[test] + fn reset_hint_shortcuts_detection() { + let previous = sample_histogram(6.0, 10.0, vec![2.0, 4.0]); + + let mut current = sample_histogram(7.0, 12.0, vec![3.0, 4.0]); + current.reset_hint = COUNTER_RESET_HINT; + assert!(current.detect_reset(&previous)); + + let mut current = sample_histogram(5.0, 8.0, vec![1.0, 4.0]); + current.reset_hint = NOT_COUNTER_RESET_HINT; + assert!(!current.detect_reset(&previous)); + } + + #[test] + fn start_timestamp_detects_counter_reset() { + let first = sample_histogram(6.0, 10.0, vec![2.0, 4.0]); + let mut last = sample_histogram(7.0, 12.0, vec![3.0, 4.0]); + last.start_timestamp = Some(1500); + + let delta = histogram_delta(&[first.clone(), last.clone()], &[1000, 2000], true).unwrap(); + assert_eq!(delta.count, 7.0); + assert_eq!(delta.sum, 12.0); + + let idelta = idelta_value(&[first, last], true, 1000, 2000, 1.0).unwrap(); + assert_eq!(idelta.count, 7.0); + assert_eq!(idelta.sum, 12.0); + } +} diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index 08a401e0483..5beabbddd98 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -28,6 +28,7 @@ use common_catalog::consts::is_readonly_schema; use common_error::ext::BoxedError; use common_function::function::FunctionContext; use common_function::function_factory::ScalarFunctionFactory; +use common_query::promql_annotations::get_promql_annotation_collector; use common_query::{Output, OutputData, OutputMeta}; use common_recordbatch::adapter::{RecordBatchStreamAdapter, RegionQueryStatCounters}; use common_recordbatch::{EmptyRecordBatchStream, SendableRecordBatchStream}; @@ -167,16 +168,19 @@ impl DatafusionQueryEngine { let physical_plan = self.create_physical_plan(&mut ctx, &plan).await?; let physical_plan = self.optimize_physical_plan(&mut ctx, physical_plan)?; let physical_plan = if let Some(wrapper) = self.plugins.get::() { - wrapper.wrap(physical_plan, query_ctx) + wrapper.wrap(physical_plan, query_ctx.clone()) } else { physical_plan }; let stream = self.execute_stream(&ctx, &physical_plan)?; + let promql_annotations = query_ctx + .remote_query_id() + .and_then(get_promql_annotation_collector); Ok(Output::new( OutputData::Stream(stream), - OutputMeta::new_with_plan(physical_plan), + OutputMeta::new_with_plan(physical_plan).with_promql_annotations(promql_annotations), )) } @@ -597,6 +601,17 @@ impl QueryEngine for DatafusionQueryEngine { state: self.engine_state().function_state(), }); + if let Some(promql_annotations) = query_ctx + .remote_query_id() + .and_then(get_promql_annotation_collector) + { + state + .config_mut() + .options_mut() + .extensions + .insert(promql_annotations); + } + // Carry scheduled Flow time through ConfigOptions.extensions so that // the distributed plan analyzer can read it during expression // simplification (preventing wall-clock constant-folding of `now()`). diff --git a/src/query/src/planner.rs b/src/query/src/planner.rs index 44b3a4604dc..cd0366cfac4 100644 --- a/src/query/src/planner.rs +++ b/src/query/src/planner.rs @@ -22,6 +22,7 @@ use arrow_schema::DataType; use async_trait::async_trait; use catalog::table_source::DfTableSourceProvider; use common_error::ext::BoxedError; +use common_query::promql_annotations::promql_annotation_collector; use common_telemetry::tracing; use datafusion::common::{DFSchema, plan_err}; use datafusion::execution::SessionStateBuilder; @@ -299,7 +300,15 @@ impl DfLogicalPlanner { #[tracing::instrument(skip_all)] async fn plan_pql(&self, stmt: &EvalStmt, query_ctx: QueryContextRef) -> Result { - let scheduled_state = self.derive_session_state_with_scheduled_time(&query_ctx)?; + let mut scheduled_state = self.derive_session_state_with_scheduled_time(&query_ctx)?; + let promql_annotations = query_ctx.remote_query_id().map(promql_annotation_collector); + if let Some(collector) = &promql_annotations { + scheduled_state + .config_mut() + .options_mut() + .extensions + .insert(collector.clone()); + } let plan_decoder = Arc::new(DefaultPlanDecoder::new( scheduled_state.clone(), &query_ctx, @@ -314,10 +323,15 @@ impl DfLogicalPlanner { .sql_parser .enable_ident_normalization, ); - let plan = PromPlanner::stmt_to_plan(table_provider, stmt, &self.engine_state) - .await - .map_err(BoxedError::new) - .context(QueryPlanSnafu)?; + let plan = PromPlanner::stmt_to_plan_with_annotations( + table_provider, + stmt, + &self.engine_state, + promql_annotations, + ) + .await + .map_err(BoxedError::new) + .context(QueryPlanSnafu)?; let context = QueryEngineContext::new(scheduled_state, query_ctx); Ok(self diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index c2cd6316ada..1b4085f66b3 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -22,7 +22,9 @@ use catalog::table_source::DfTableSourceProvider; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_function::function::FunctionContext; +use common_query::native_histogram::native_histogram_value_type; use common_query::prelude::greptime_value; +use common_query::promql_annotations::PromqlAnnotationCollector; use datafusion::common::DFSchemaRef; use datafusion::datasource::DefaultTableSource; use datafusion::functions_aggregate::average::avg_udaf; @@ -51,7 +53,7 @@ use datafusion_expr::{ ExprSchemable, Literal, Projection, SortExpr, TableScan, TableSource, col, lit, }; use datatypes::arrow::datatypes::{DataType as ArrowDataType, TimeUnit as ArrowTimeUnit}; -use datatypes::data_type::ConcreteDataType; +use datatypes::data_type::{ConcreteDataType, DataType as GreptimeDataType}; use itertools::Itertools; use once_cell::sync::Lazy; use promql::extension_plan::{ @@ -60,9 +62,18 @@ use promql::extension_plan::{ }; use promql::functions::{ AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing, - IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, PredictLinear, PresentOverTime, - QuantileOverTime, Rate, Resets, Round, StddevOverTime, StdvarOverTime, SumOverTime, - quantile_udaf, + IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, NativeHistogramAbsentOverTime, + NativeHistogramAdd, NativeHistogramAggAvg, NativeHistogramAggSum, NativeHistogramAvg, + NativeHistogramAvgOverTime, NativeHistogramChanges, NativeHistogramCount, + NativeHistogramCountOverTime, NativeHistogramDelta, NativeHistogramDivScalar, + NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction, NativeHistogramIDelta, + NativeHistogramIRate, NativeHistogramIncrease, NativeHistogramLastOverTime, + NativeHistogramMulScalar, NativeHistogramNeg, NativeHistogramNotEq, + NativeHistogramPresentOverTime, NativeHistogramQuantile, NativeHistogramRate, + NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev, NativeHistogramStdvar, + NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime, NativeHistogramToString, + PredictLinear, PresentOverTime, QuantileOverTime, Rate, Resets, Round, StddevOverTime, + StdvarOverTime, SumOverTime, quantile_udaf, }; use promql_parser::label::{METRIC_NAME, MatchOp, Matcher, Matchers}; use promql_parser::parser::token::TokenType; @@ -374,6 +385,7 @@ impl PromPlannerContext { pub struct PromPlanner { table_provider: DfTableSourceProvider, ctx: PromPlannerContext, + promql_annotations: Option, } impl PromPlanner { @@ -381,10 +393,20 @@ impl PromPlanner { table_provider: DfTableSourceProvider, stmt: &EvalStmt, query_engine_state: &QueryEngineState, + ) -> Result { + Self::stmt_to_plan_with_annotations(table_provider, stmt, query_engine_state, None).await + } + + pub async fn stmt_to_plan_with_annotations( + table_provider: DfTableSourceProvider, + stmt: &EvalStmt, + query_engine_state: &QueryEngineState, + promql_annotations: Option, ) -> Result { let mut planner = Self { table_provider, ctx: PromPlannerContext::from_eval_stmt(stmt), + promql_annotations, }; let plan = planner @@ -699,6 +721,22 @@ impl PromPlanner { let group_exprs = self.agg_modifier_to_col(input.schema(), modifier, false)?; + if self.all_field_columns_are_native_histograms(input.schema()) { + let promql_annotations = self.promql_annotations.clone(); + return self.projection_for_each_field_column(input, |col| { + Ok(DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramDrop::float_null_udf( + format!( + "{}: dropped native histogram samples because this aggregation is not supported for native histograms", + op + ), + promql_annotations.clone(), + )), + args: vec![DfExpr::Column(Column::from_name(col))], + })) + }); + } + let val = Self::get_param_as_literal_expr(param, Some(*op), Some(ArrowDataType::Float64))?; // convert op and value columns to window exprs. @@ -775,6 +813,14 @@ impl PromPlanner { let UnaryExpr { expr } = unary_expr; // Unary Expr in PromQL implys the `-` operator let input = self.prom_expr_to_plan(expr, query_engine_state).await?; + if self.all_field_columns_are_native_histograms(input.schema()) { + return self.projection_for_each_field_column(input, |col| { + Ok(DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramNeg::scalar_udf()), + args: vec![DfExpr::Column(col.into())], + })) + }); + } self.projection_for_each_field_column(input, |col| { Ok(DfExpr::Negative(Box::new(DfExpr::Column(col.into())))) }) @@ -818,6 +864,16 @@ impl PromPlanner { }); } + if planned_leaves.iter().any(|leaf| { + leaf.ctx + .field_columns + .iter() + .any(|field| Self::field_column_is_native_histogram(leaf.plan.schema(), field)) + }) { + self.ctx = original_ctx; + return Ok(None); + } + if !Self::binary_island_join_contexts_supported(&planned_leaves) { self.ctx = original_ctx; return Ok(None); @@ -1154,10 +1210,25 @@ impl PromPlanner { if let Some(time_expr) = self.try_build_special_time_expr_with_context(lhs) { expr = time_expr } + let input_schema = input.schema().clone(); + let promql_annotations = self.promql_annotations.clone(); let bin_expr_builder = |col: &String| { let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; - let mut binary_expr = - binary_expr_builder(expr.clone(), DfExpr::Column(col.into()))?; + let rhs_is_histogram = + Self::field_column_is_native_histogram(&input_schema, col); + let rhs = DfExpr::Column(col.into()); + let mut binary_expr = match Self::native_histogram_binary_expr( + *op, + expr.clone(), + false, + rhs.clone(), + rhs_is_histogram, + is_comparison_op && !should_return_bool, + promql_annotations.clone(), + )? { + Some(expr) => expr, + None => binary_expr_builder(expr.clone(), rhs)?, + }; if is_comparison_op && should_return_bool { binary_expr = DfExpr::Cast(Cast { @@ -1180,10 +1251,25 @@ impl PromPlanner { if let Some(time_expr) = self.try_build_special_time_expr_with_context(rhs) { expr = time_expr } + let input_schema = input.schema().clone(); + let promql_annotations = self.promql_annotations.clone(); let bin_expr_builder = |col: &String| { let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; - let mut binary_expr = - binary_expr_builder(DfExpr::Column(col.into()), expr.clone())?; + let lhs_is_histogram = + Self::field_column_is_native_histogram(&input_schema, col); + let lhs = DfExpr::Column(col.into()); + let mut binary_expr = match Self::native_histogram_binary_expr( + *op, + lhs.clone(), + lhs_is_histogram, + expr.clone(), + false, + is_comparison_op && !should_return_bool, + promql_annotations.clone(), + )? { + Some(expr) => expr, + None => binary_expr_builder(lhs, expr.clone())?, + }; if is_comparison_op && should_return_bool { binary_expr = DfExpr::Cast(Cast { @@ -1279,21 +1365,38 @@ impl PromPlanner { &right_context, )?; let join_plan_schema = join_plan.schema().clone(); + let promql_annotations = self.promql_annotations.clone(); let bin_expr_builder = |_: &String| { let (left_col_name, right_col_name) = field_columns.next().unwrap(); - let left_col = join_plan_schema + let left_field = join_plan_schema .qualified_field_with_name(Some(&left_table_ref), left_col_name) - .context(DataFusionPlanningSnafu)? - .into(); - let right_col = join_plan_schema + .context(DataFusionPlanningSnafu)?; + let right_field = join_plan_schema .qualified_field_with_name(Some(&right_table_ref), right_col_name) - .context(DataFusionPlanningSnafu)? - .into(); + .context(DataFusionPlanningSnafu)?; + let left_is_histogram = + left_field.1.data_type() == &Self::native_histogram_arrow_type(); + let right_is_histogram = + right_field.1.data_type() == &Self::native_histogram_arrow_type(); + let left_col = left_field.into(); + let right_col = right_field.into(); let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; - let mut binary_expr = - binary_expr_builder(DfExpr::Column(left_col), DfExpr::Column(right_col))?; + let lhs = DfExpr::Column(left_col); + let rhs = DfExpr::Column(right_col); + let mut binary_expr = match Self::native_histogram_binary_expr( + *op, + lhs.clone(), + left_is_histogram, + rhs.clone(), + right_is_histogram, + is_comparison_op && !should_return_bool, + promql_annotations.clone(), + )? { + Some(expr) => expr, + None => binary_expr_builder(lhs, rhs)?, + }; if is_comparison_op && should_return_bool { binary_expr = DfExpr::Cast(Cast { expr: Box::new(binary_expr), @@ -1613,8 +1716,12 @@ impl PromPlanner { ), }) }; - let (mut func_exprs, new_tags) = - self.create_function_expr(func, args.literals.clone(), query_engine_state)?; + let (mut func_exprs, new_tags) = self.create_function_expr( + func, + args.literals.clone(), + input.schema(), + query_engine_state, + )?; func_exprs.insert(0, self.create_time_index_column_expr()?); func_exprs.extend_from_slice(&self.create_tag_column_exprs()?); if let Some(tsid_col) = @@ -2738,6 +2845,7 @@ impl PromPlanner { &mut self, func: &Function, other_input_exprs: Vec, + input_schema: &DFSchemaRef, query_engine_state: &QueryEngineState, ) -> Result<(Vec, Vec)> { // TODO(ruihang): check function args list @@ -2748,35 +2856,136 @@ impl PromPlanner { let mut exprs = Vec::with_capacity(self.ctx.field_columns.len()); // New labels after executing the function, e.g. `label_replace` etc. let mut new_tags = vec![]; + let promql_annotations = self.promql_annotations.clone(); + let native_histogram_drop_udf = |name: &str| { + Arc::new(NativeHistogramDrop::float_null_udf( + format!( + "{name}: dropped native histogram samples because this function is not supported for native histograms" + ), + promql_annotations.clone(), + )) + }; let scalar_func = match func.name { + "increase" if self.all_field_columns_are_native_histogram_ranges(input_schema) => { + ScalarFunc::ExtrapolateUdf( + Arc::new(NativeHistogramIncrease::scalar_udf_with_collector( + self.promql_annotations.clone(), + )), + self.ctx.range.context(ExpectRangeSelectorSnafu)?, + ) + } "increase" => ScalarFunc::ExtrapolateUdf( Arc::new(Increase::scalar_udf()), self.ctx.range.context(ExpectRangeSelectorSnafu)?, ), + "rate" if self.all_field_columns_are_native_histogram_ranges(input_schema) => { + ScalarFunc::ExtrapolateUdf( + Arc::new(NativeHistogramRate::scalar_udf_with_collector( + self.promql_annotations.clone(), + )), + self.ctx.range.context(ExpectRangeSelectorSnafu)?, + ) + } "rate" => ScalarFunc::ExtrapolateUdf( Arc::new(Rate::scalar_udf()), self.ctx.range.context(ExpectRangeSelectorSnafu)?, ), + "delta" if self.all_field_columns_are_native_histogram_ranges(input_schema) => { + ScalarFunc::ExtrapolateUdf( + Arc::new(NativeHistogramDelta::scalar_udf_with_collector( + self.promql_annotations.clone(), + )), + self.ctx.range.context(ExpectRangeSelectorSnafu)?, + ) + } "delta" => ScalarFunc::ExtrapolateUdf( Arc::new(Delta::scalar_udf()), self.ctx.range.context(ExpectRangeSelectorSnafu)?, ), + "idelta" if self.all_field_columns_are_native_histogram_ranges(input_schema) => { + ScalarFunc::Udf(Arc::new(NativeHistogramIDelta::scalar_udf_with_collector( + self.promql_annotations.clone(), + ))) + } "idelta" => ScalarFunc::Udf(Arc::new(IDelta::::scalar_udf())), + "irate" if self.all_field_columns_are_native_histogram_ranges(input_schema) => { + ScalarFunc::Udf(Arc::new(NativeHistogramIRate::scalar_udf_with_collector( + self.promql_annotations.clone(), + ))) + } "irate" => ScalarFunc::Udf(Arc::new(IDelta::::scalar_udf())), + "resets" if self.all_field_columns_are_native_histogram_ranges(input_schema) => { + ScalarFunc::Udf(Arc::new(NativeHistogramResets::scalar_udf())) + } "resets" => ScalarFunc::Udf(Arc::new(Resets::scalar_udf())), + "changes" if self.all_field_columns_are_native_histogram_ranges(input_schema) => { + ScalarFunc::Udf(Arc::new(NativeHistogramChanges::scalar_udf())) + } "changes" => ScalarFunc::Udf(Arc::new(Changes::scalar_udf())), + "deriv" if self.all_field_columns_are_native_histogram_ranges(input_schema) => { + ScalarFunc::Udf(native_histogram_drop_udf(func.name)) + } "deriv" => ScalarFunc::Udf(Arc::new(Deriv::scalar_udf())), + "avg_over_time" if self.all_field_columns_are_native_histogram_ranges(input_schema) => { + ScalarFunc::Udf(Arc::new( + NativeHistogramAvgOverTime::scalar_udf_with_collector( + self.promql_annotations.clone(), + ), + )) + } "avg_over_time" => ScalarFunc::Udf(Arc::new(AvgOverTime::scalar_udf())), + "min_over_time" | "max_over_time" + if self.all_field_columns_are_native_histogram_ranges(input_schema) => + { + ScalarFunc::Udf(native_histogram_drop_udf(func.name)) + } "min_over_time" => ScalarFunc::Udf(Arc::new(MinOverTime::scalar_udf())), "max_over_time" => ScalarFunc::Udf(Arc::new(MaxOverTime::scalar_udf())), + "sum_over_time" if self.all_field_columns_are_native_histogram_ranges(input_schema) => { + ScalarFunc::Udf(Arc::new( + NativeHistogramSumOverTime::scalar_udf_with_collector( + self.promql_annotations.clone(), + ), + )) + } "sum_over_time" => ScalarFunc::Udf(Arc::new(SumOverTime::scalar_udf())), + "count_over_time" + if self.all_field_columns_are_native_histogram_ranges(input_schema) => + { + ScalarFunc::Udf(Arc::new(NativeHistogramCountOverTime::scalar_udf())) + } "count_over_time" => ScalarFunc::Udf(Arc::new(CountOverTime::scalar_udf())), + "last_over_time" + if self.all_field_columns_are_native_histogram_ranges(input_schema) => + { + ScalarFunc::Udf(Arc::new(NativeHistogramLastOverTime::scalar_udf())) + } "last_over_time" => ScalarFunc::Udf(Arc::new(LastOverTime::scalar_udf())), + "absent_over_time" + if self.all_field_columns_are_native_histogram_ranges(input_schema) => + { + ScalarFunc::Udf(Arc::new(NativeHistogramAbsentOverTime::scalar_udf())) + } "absent_over_time" => ScalarFunc::Udf(Arc::new(AbsentOverTime::scalar_udf())), + "present_over_time" + if self.all_field_columns_are_native_histogram_ranges(input_schema) => + { + ScalarFunc::Udf(Arc::new(NativeHistogramPresentOverTime::scalar_udf())) + } "present_over_time" => ScalarFunc::Udf(Arc::new(PresentOverTime::scalar_udf())), + "stddev_over_time" | "stdvar_over_time" | "quantile_over_time" + if self.all_field_columns_are_native_histogram_ranges(input_schema) => + { + ScalarFunc::Udf(native_histogram_drop_udf(func.name)) + } "stddev_over_time" => ScalarFunc::Udf(Arc::new(StddevOverTime::scalar_udf())), "stdvar_over_time" => ScalarFunc::Udf(Arc::new(StdvarOverTime::scalar_udf())), "quantile_over_time" => ScalarFunc::Udf(Arc::new(QuantileOverTime::scalar_udf())), + "predict_linear" + if self.all_field_columns_are_native_histogram_ranges(input_schema) => + { + ScalarFunc::Udf(native_histogram_drop_udf(func.name)) + } "predict_linear" => { other_input_exprs[0] = DfExpr::Cast(Cast { expr: Box::new(other_input_exprs[0].clone()), @@ -2784,9 +2993,32 @@ impl PromPlanner { }); ScalarFunc::Udf(Arc::new(PredictLinear::scalar_udf())) } + "double_exponential_smoothing" | "holt_winters" + if self.all_field_columns_are_native_histogram_ranges(input_schema) => + { + ScalarFunc::Udf(native_histogram_drop_udf(func.name)) + } "double_exponential_smoothing" | "holt_winters" => { ScalarFunc::Udf(Arc::new(DoubleExponentialSmoothing::scalar_udf())) } + "histogram_count" => { + ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramCount::scalar_udf())) + } + "histogram_sum" => { + ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramSum::scalar_udf())) + } + "histogram_avg" => { + ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramAvg::scalar_udf())) + } + "histogram_stddev" => { + ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramStddev::scalar_udf())) + } + "histogram_stdvar" => { + ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramStdvar::scalar_udf())) + } + "histogram_fraction" => { + ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramFraction::scalar_udf())) + } "time" => { exprs.push(build_special_time_expr( self.ctx.time_index_column.as_ref().unwrap(), @@ -2935,12 +3167,21 @@ impl PromPlanner { ScalarFunc::GeneratedExpr } + "round" if self.all_field_columns_are_native_histograms(input_schema) => { + if other_input_exprs.is_empty() { + other_input_exprs.push_front(0.0f64.lit()); + } + ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name)) + } "round" => { if other_input_exprs.is_empty() { other_input_exprs.push_front(0.0f64.lit()); } ScalarFunc::DataFusionUdf(Arc::new(Round::scalar_udf())) } + "rad" | "deg" | "sgn" if self.all_field_columns_are_native_histograms(input_schema) => { + ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name)) + } "rad" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::radians()), "deg" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::degrees()), "sgn" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::signum()), @@ -2960,20 +3201,32 @@ impl PromPlanner { .scalar_functions() .get(func.name) { - ScalarFunc::DataFusionBuiltin(f.clone()) + if self.all_field_columns_are_native_histograms(input_schema) { + ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name)) + } else { + ScalarFunc::DataFusionBuiltin(f.clone()) + } } else if let Some(factory) = query_engine_state.scalar_function(func.name) { - let func_state = query_engine_state.function_state(); - let query_ctx = self.table_provider.query_ctx(); + if self.all_field_columns_are_native_histograms(input_schema) { + ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name)) + } else { + let func_state = query_engine_state.function_state(); + let query_ctx = self.table_provider.query_ctx(); - ScalarFunc::DataFusionUdf(Arc::new(factory.provide(FunctionContext { - state: func_state, - query_ctx: query_ctx.clone(), - }))) + ScalarFunc::DataFusionUdf(Arc::new(factory.provide(FunctionContext { + state: func_state, + query_ctx: query_ctx.clone(), + }))) + } } else if let Some(f) = datafusion_functions::math::functions() .iter() .find(|f| f.name() == func.name) { - ScalarFunc::DataFusionUdf(f.clone()) + if self.all_field_columns_are_native_histograms(input_schema) { + ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name)) + } else { + ScalarFunc::DataFusionUdf(f.clone()) + } } else { return UnsupportedExprSnafu { name: func.name.to_string(), @@ -3005,6 +3258,25 @@ impl PromPlanner { .collect_vec(); exprs.push(DfExpr::ScalarFunction(ScalarFunction { func, args })) } + ScalarFunc::NativeHistogramUdf(func) => { + if Self::field_column_is_native_histogram(input_schema, value) { + let args = itertools::chain!( + other_input_exprs.iter().take(field_column_pos).cloned(), + std::iter::once(col_expr), + other_input_exprs.iter().skip(field_column_pos).cloned() + ) + .collect_vec(); + exprs.push(DfExpr::ScalarFunction(ScalarFunction { func, args })); + } else { + exprs.push( + DfExpr::Literal(ScalarValue::Float64(None), None).alias(format!( + "{}_{}", + func.name(), + value + )), + ); + } + } ScalarFunc::Udf(func) => { let ts_range_expr = DfExpr::Column(Column::from_name( RangeManipulate::build_timestamp_range_name( @@ -3391,6 +3663,11 @@ impl PromPlanner { } ); } + + if self.all_field_columns_are_native_histograms(input_plan.schema()) { + return self.create_native_histogram_aggregate_exprs(op, input_plan); + } + let aggr = match op.id() { token::T_SUM => sum_udaf(), token::T_QUANTILE => { @@ -3466,6 +3743,85 @@ impl PromPlanner { Ok((exprs, prev_field_exprs)) } + fn create_native_histogram_aggregate_exprs( + &mut self, + op: TokenType, + input_plan: &LogicalPlan, + ) -> Result<(Vec, Vec)> { + let prev_field_exprs = if op.id() == token::T_COUNT_VALUES { + ensure!( + self.ctx.field_columns.len() == 1, + UnsupportedExprSnafu { + name: "count_values on multi-value input" + } + ); + self.ctx + .field_columns + .iter() + .map(|col| { + DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramToString::scalar_udf()), + args: vec![DfExpr::Column(Column::from_name(col))], + }) + }) + .collect::>() + } else { + vec![] + }; + + let exprs = self + .ctx + .field_columns + .iter() + .map(|col| { + let input = DfExpr::Column(Column::from_name(col)); + let expr = match op.id() { + token::T_SUM => Arc::new(NativeHistogramAggSum::aggregate_udf_with_collector( + self.promql_annotations.clone(), + )) + .call(vec![input]) + .alias(col), + token::T_AVG => Arc::new(NativeHistogramAggAvg::aggregate_udf_with_collector( + self.promql_annotations.clone(), + )) + .call(vec![input]) + .alias(col), + token::T_COUNT_VALUES | token::T_COUNT => count_udaf().call(vec![input]).alias(col), + token::T_GROUP => max_udaf().call(vec![lit(1_f64)]).alias(col), + token::T_MIN + | token::T_MAX + | token::T_STDDEV + | token::T_STDVAR + | token::T_QUANTILE + | token::T_TOPK + | token::T_BOTTOMK => sum_udaf() + .call(vec![DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramDrop::float_null_udf( + format!( + "{}: dropped native histogram samples because this aggregation is not supported for native histograms", + op + ), + self.promql_annotations.clone(), + )), + args: vec![input], + })]) + .alias(col), + _ => return UnexpectedTokenSnafu { token: op }.fail(), + }; + Ok(expr) + }) + .collect::>>()?; + + let normalized_exprs = + normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?; + self.ctx.field_columns = normalized_exprs + .into_iter() + .map(|expr| expr.schema_name().to_string()) + .collect(); + + Ok((exprs, prev_field_exprs)) + } + fn get_param_value_as_str(op: TokenType, param: &Option>) -> Result<&str> { let param = param .as_deref() @@ -3634,6 +3990,10 @@ impl PromPlanner { let input_plan = self.strip_tsid_column(input_plan)?; self.ctx.use_tsid = false; + if self.all_field_columns_are_native_histograms(input_plan.schema()) { + return self.create_native_histogram_quantile_plan(phi, input_plan); + } + if !self.ctx.has_le_tag() { // Return empty result instead of error when 'le' column is not found // This handles the case when histogram metrics don't exist @@ -3677,6 +4037,43 @@ impl PromPlanner { })) } + fn create_native_histogram_quantile_plan( + &mut self, + phi: f64, + input_plan: LogicalPlan, + ) -> Result { + ensure!( + self.ctx.field_columns.len() == 1, + MultiFieldsNotSupportedSnafu { + operator: SPECIAL_HISTOGRAM_QUANTILE + }, + ); + + let field_column = self.ctx.field_columns[0].clone(); + let quantile_expr = DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramQuantile::scalar_udf()), + args: vec![ + DfExpr::Column(Column::from_name(field_column)), + DfExpr::Literal(ScalarValue::Float64(Some(phi)), None), + ], + }); + let display_name = quantile_expr.schema_name().to_string(); + self.ctx.field_columns = vec![display_name.clone()]; + + let project_exprs = std::iter::once(self.create_time_index_column_expr()?) + .chain(std::iter::once(quantile_expr.alias(display_name))) + .chain(self.create_tag_column_exprs()?) + .collect::>(); + + LogicalPlanBuilder::from(input_plan) + .project(project_exprs) + .context(DataFusionPlanningSnafu)? + .filter(self.create_empty_values_filter_expr()?) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu) + } + /// Create a [SPECIAL_VECTOR_FUNCTION] plan async fn create_vector_plan(&mut self, args: &PromFunctionArgs) -> Result { if args.args.len() != 1 { @@ -3884,6 +4281,69 @@ impl PromPlanner { } } + fn native_histogram_binary_expr( + token: TokenType, + lhs: DfExpr, + lhs_is_histogram: bool, + rhs: DfExpr, + rhs_is_histogram: bool, + filter_context: bool, + promql_annotations: Option, + ) -> Result> { + if !lhs_is_histogram && !rhs_is_histogram { + return Ok(None); + } + + let scalar_fn = |func: ScalarUdfDef, args| { + DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(func), + args, + }) + }; + let invalid_expr = || { + let message = format!( + "{}: dropped native histogram samples because this binary operation is not supported for native histograms", + token + ); + let func = if filter_context { + NativeHistogramDrop::bool_false_udf(message, promql_annotations.clone()) + } else { + NativeHistogramDrop::float_null_udf(message, promql_annotations.clone()) + }; + let args = vec![lhs.clone(), rhs.clone()]; + scalar_fn(func, args) + }; + + let expr = match (token.id(), lhs_is_histogram, rhs_is_histogram) { + (token::T_ADD, true, true) => scalar_fn( + NativeHistogramAdd::scalar_udf_with_collector(promql_annotations.clone()), + vec![lhs, rhs], + ), + (token::T_SUB, true, true) => scalar_fn( + NativeHistogramSub::scalar_udf_with_collector(promql_annotations.clone()), + vec![lhs, rhs], + ), + (token::T_MUL, true, false) => { + scalar_fn(NativeHistogramMulScalar::scalar_udf(), vec![lhs, rhs]) + } + (token::T_MUL, false, true) => { + scalar_fn(NativeHistogramScalarMul::scalar_udf(), vec![lhs, rhs]) + } + (token::T_DIV, true, false) => { + scalar_fn(NativeHistogramDivScalar::scalar_udf(), vec![lhs, rhs]) + } + (token::T_EQLC, true, true) => { + scalar_fn(NativeHistogramEq::scalar_udf(), vec![lhs, rhs]) + } + (token::T_NEQ, true, true) => { + scalar_fn(NativeHistogramNotEq::scalar_udf(), vec![lhs, rhs]) + } + _ => invalid_expr(), + }; + + Ok(Some(expr)) + } + /// Return a lambda to build binary expression from token. /// Because some binary operator are function in DataFusion like `atan2` or `^`. #[allow(clippy::type_complexity)] @@ -3990,6 +4450,52 @@ impl PromPlanner { .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME) } + fn native_histogram_arrow_type() -> ArrowDataType { + native_histogram_value_type().as_arrow_type() + } + + fn field_column_type<'a>( + schema: &'a DFSchemaRef, + field_column: &str, + ) -> Option<&'a ArrowDataType> { + schema + .index_of_column_by_name(None, field_column) + .map(|idx| schema.field(idx).data_type()) + } + + fn field_column_is_native_histogram(schema: &DFSchemaRef, field_column: &str) -> bool { + Self::field_column_type(schema, field_column) + .is_some_and(|data_type| data_type == &Self::native_histogram_arrow_type()) + } + + fn field_column_is_native_histogram_range(schema: &DFSchemaRef, field_column: &str) -> bool { + Self::field_column_type(schema, field_column).is_some_and(|data_type| { + matches!( + data_type, + ArrowDataType::Dictionary(_, value_type) + if value_type.as_ref() == &Self::native_histogram_arrow_type() + ) + }) + } + + fn all_field_columns_are_native_histograms(&self, schema: &DFSchemaRef) -> bool { + !self.ctx.field_columns.is_empty() + && self + .ctx + .field_columns + .iter() + .all(|field| Self::field_column_is_native_histogram(schema, field)) + } + + fn all_field_columns_are_native_histogram_ranges(&self, schema: &DFSchemaRef) -> bool { + !self.ctx.field_columns.is_empty() + && self + .ctx + .field_columns + .iter() + .all(|field| Self::field_column_is_native_histogram_range(schema, field)) + } + fn optional_tsid_projection( schema: &DFSchemaRef, table_ref: Option<&TableReference>, @@ -4719,6 +5225,9 @@ enum ScalarFunc { /// Similar to DataFusionBuiltin but for custom functions not built into DataFusion. /// Processing: Direct pass-through with argument positioning. DataFusionUdf(Arc), + /// Native histogram helper UDFs. Non-histogram inputs are projected as NULL + /// so the normal PromQL empty-value filter drops them. + NativeHistogramUdf(Arc), /// PromQL-specific functions that operate on time series data with temporal context. /// These functions require both timestamp ranges and values to perform calculations. /// Processing: Automatically injects timestamp_range and value columns as first arguments. @@ -4745,6 +5254,7 @@ mod test { use catalog::memory::{MemoryCatalogManager, new_memory_catalog_manager}; use common_base::Plugins; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; + use common_query::native_histogram::NATIVE_HISTOGRAM_FIELD; use common_query::prelude::greptime_timestamp; use common_query::test_util::DummyDecoder; use datafusion::arrow::datatypes::Schema as ArrowSchema; @@ -4925,6 +5435,62 @@ mod test { ) } + async fn build_test_native_histogram_table_provider(table_name: &str) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + let columns = vec![ + ColumnSchema::new( + "tag_0".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + NATIVE_HISTOGRAM_FIELD.to_string(), + native_histogram_value_type().clone(), + true, + ), + ]; + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![0]) + .value_indices(vec![2]) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name(table_name) + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + table_id: 1024, + table, + }) + .is_ok() + ); + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) + } + async fn build_test_table_provider_with_tsid( table_name_tuples: &[(String, String)], num_tag: usize, @@ -6278,6 +6844,7 @@ mod test { let mut planner = PromPlanner { table_provider, ctx: PromPlannerContext::from_eval_stmt(&eval_stmt), + promql_annotations: None, }; let plan = planner .prom_expr_to_plan(&eval_stmt.expr, &build_query_engine_state()) @@ -7132,6 +7699,78 @@ mod test { indie_query_plan_compare(query, expected).await; } + async fn native_histogram_plan(query: &str) -> String { + let table_provider = build_test_native_histogram_table_provider("some_metric").await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt(query), + &build_query_engine_state(), + ) + .await + .unwrap(); + plan.display_indent_schema().to_string() + } + + #[tokio::test] + async fn native_histogram_count_uses_native_udf() { + let plan = native_histogram_plan("histogram_count(some_metric)").await; + + assert!(plan.contains("prom_native_histogram_count"), "{plan}"); + assert!(!plan.contains("PromHistogramFold"), "{plan}"); + } + + #[tokio::test] + async fn native_histogram_rate_can_feed_count() { + let plan = native_histogram_plan("histogram_count(rate(some_metric[5m]))").await; + + assert!(plan.contains("prom_native_histogram_rate"), "{plan}"); + assert!(plan.contains("prom_native_histogram_count"), "{plan}"); + } + + #[tokio::test] + async fn native_histogram_quantile_skips_classic_fold() { + let plan = native_histogram_plan("histogram_quantile(0.9, some_metric)").await; + + assert!(plan.contains("prom_native_histogram_quantile"), "{plan}"); + assert!(!plan.contains("PromHistogramFold"), "{plan}"); + } + + #[tokio::test] + async fn native_histogram_topk_uses_drop_udf() { + let plan = native_histogram_plan("topk(1, some_metric)").await; + + assert!(plan.contains("prom_native_histogram_drop_float"), "{plan}"); + } + + #[tokio::test] + async fn unsupported_native_histogram_functions_use_drop_udf() { + for query in [ + "deriv(some_metric[5m])", + "min_over_time(some_metric[5m])", + "quantile_over_time(0.9, some_metric[5m])", + "predict_linear(some_metric[5m], 60)", + "round(some_metric)", + "abs(some_metric)", + ] { + let plan = native_histogram_plan(query).await; + + assert!( + plan.contains("prom_native_histogram_drop_float"), + "{query}\n{plan}" + ); + } + } + + #[tokio::test] + async fn native_histogram_absent_over_time_uses_native_udf() { + let plan = native_histogram_plan("absent_over_time(some_metric[5m])").await; + + assert!( + plan.contains("prom_native_histogram_absent_over_time"), + "{plan}" + ); + } + #[tokio::test] async fn less_filter_on_value() { let query = "some_metric < 1.2345"; diff --git a/src/servers/src/http/prometheus.rs b/src/servers/src/http/prometheus.rs index 9d17c59b784..2b4c6613069 100644 --- a/src/servers/src/http/prometheus.rs +++ b/src/servers/src/http/prometheus.rs @@ -74,17 +74,31 @@ use crate::prometheus_handler::PrometheusHandlerRef; /// For [ValueType::Vector] result type #[derive(Debug, Default, Serialize, Deserialize, PartialEq)] +#[serde(deny_unknown_fields)] pub struct PromSeriesVector { pub metric: BTreeMap, #[serde(skip_serializing_if = "Option::is_none")] pub value: Option<(f64, String)>, + #[serde(skip_serializing_if = "Option::is_none")] + pub histogram: Option<(f64, PromNativeHistogram)>, } /// For [ValueType::Matrix] result type #[derive(Debug, Default, Serialize, Deserialize, PartialEq)] +#[serde(deny_unknown_fields)] pub struct PromSeriesMatrix { pub metric: BTreeMap, + #[serde(skip_serializing_if = "Vec::is_empty", default)] pub values: Vec<(f64, String)>, + #[serde(skip_serializing_if = "Vec::is_empty", default)] + pub histograms: Vec<(f64, PromNativeHistogram)>, +} + +#[derive(Debug, Default, Serialize, Deserialize, PartialEq)] +pub struct PromNativeHistogram { + pub count: String, + pub sum: String, + pub buckets: Vec<(u8, String, String, String)>, } /// Variants corresponding to [ValueType] diff --git a/src/servers/src/http/result/prometheus_resp.rs b/src/servers/src/http/result/prometheus_resp.rs index 6f2b1156862..1795595839d 100644 --- a/src/servers/src/http/result/prometheus_resp.rs +++ b/src/servers/src/http/result/prometheus_resp.rs @@ -16,7 +16,7 @@ use std::cmp::Ordering; use std::collections::{BTreeMap, HashMap}; -use arrow::array::{Array, AsArray}; +use arrow::array::{Array, AsArray, StructArray}; use arrow::datatypes::{Float64Type, TimestampMillisecondType}; use arrow_schema::DataType; use axum::Json; @@ -24,6 +24,9 @@ use axum::http::HeaderValue; use axum::response::{IntoResponse, Response}; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; +use common_query::native_histogram::{ + NativeHistogram, is_native_histogram_value_type, read_histogram, +}; use common_query::{Output, OutputData}; use common_recordbatch::RecordBatches; use datatypes::prelude::ConcreteDataType; @@ -35,13 +38,37 @@ use serde_json::Value; use snafu::{OptionExt, ResultExt}; use crate::error::{ - ArrowSnafu, CollectRecordbatchSnafu, Result, UnexpectedResultSnafu, status_code_to_http_status, + ArrowSnafu, CollectRecordbatchSnafu, DataFusionSnafu, Result, UnexpectedResultSnafu, + status_code_to_http_status, }; use crate::http::header::{GREPTIME_DB_HEADER_METRICS, collect_plan_metrics}; use crate::http::prometheus::{ - PromData, PromQueryResult, PromSeriesMatrix, PromSeriesVector, PrometheusResponse, + PromData, PromNativeHistogram, PromQueryResult, PromSeriesMatrix, PromSeriesVector, + PrometheusResponse, }; +#[derive(Default)] +struct PromSeriesSamples { + values: Vec<(f64, String)>, + histograms: Vec<(f64, PromNativeHistogram)>, +} + +fn prometheus_native_histogram(histogram: &NativeHistogram) -> Option { + Some(PromNativeHistogram { + count: histogram.count.to_string(), + sum: histogram.sum.to_string(), + buckets: histogram.to_prometheus_buckets()?, + }) +} + +fn prometheus_native_histogram_or_error( + histogram: &NativeHistogram, +) -> Result { + prometheus_native_histogram(histogram).context(UnexpectedResultSnafu { + reason: "native histogram cannot be converted to Prometheus buckets".to_string(), + }) +} + #[derive(Debug, Default, Serialize, Deserialize, PartialEq)] pub struct PrometheusJsonResponse { pub status: String, @@ -55,6 +82,8 @@ pub struct PrometheusJsonResponse { pub error_type: Option, #[serde(skip_serializing_if = "Option::is_none")] pub warnings: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub infos: Option>, #[serde(skip)] pub status_code: Option, @@ -99,6 +128,7 @@ impl PrometheusJsonResponse { error: Some(reason.into()), error_type: Some(error_type.to_string()), warnings: None, + infos: None, resp_metrics: Default::default(), status_code: Some(error_type), } @@ -111,6 +141,7 @@ impl PrometheusJsonResponse { error: None, error_type: None, warnings: None, + infos: None, resp_metrics: Default::default(), status_code: None, } @@ -124,6 +155,7 @@ impl PrometheusJsonResponse { ) -> Self { let response: Result = try { let result = result?; + let mut meta = result.meta; let mut resp = match result.data { OutputData::RecordBatches(batches) => Self::success( @@ -144,8 +176,16 @@ impl PrometheusJsonResponse { "expected data result, but got affected rows", ), }; + meta.collect_promql_annotations(); - if let Some(physical_plan) = result.meta.plan { + if !meta.warnings.is_empty() { + resp.warnings = Some(meta.warnings); + } + if !meta.infos.is_empty() { + resp.infos = Some(meta.infos); + } + + if let Some(physical_plan) = meta.plan { let mut result_map = HashMap::new(); let mut tmp = vec![&mut result_map]; collect_plan_metrics(&physical_plan, &mut tmp); @@ -199,6 +239,7 @@ impl PrometheusJsonResponse { let mut timestamp_column_index = None; let mut tag_column_indices = Vec::new(); let mut first_field_column_index = None; + let mut native_histogram_column_index = None; let mut num_label_columns = 0; @@ -224,6 +265,11 @@ impl PrometheusJsonResponse { { first_field_column_index = Some(i); } + _ if native_histogram_column_index.is_none() + && is_native_histogram_value_type(&column.data_type) => + { + native_histogram_column_index = Some(i); + } ConcreteDataType::String(_) => { tag_column_indices.push(i); num_label_columns += 1; @@ -235,13 +281,16 @@ impl PrometheusJsonResponse { let timestamp_column_index = timestamp_column_index.context(UnexpectedResultSnafu { reason: "no timestamp column found".to_string(), })?; - let first_field_column_index = first_field_column_index.context(UnexpectedResultSnafu { - reason: "no value column found".to_string(), - })?; + if first_field_column_index.is_none() && native_histogram_column_index.is_none() { + return UnexpectedResultSnafu { + reason: "no value column found".to_string(), + } + .fail(); + } // Preserves the order of output tags. // Tag order matters, e.g., after sorc and sort_desc, the output order must be kept. - let mut buffer = IndexMap::, Vec<(f64, String)>>::new(); + let mut buffer = IndexMap::, PromSeriesSamples>::new(); let schema = batches.schema(); for batch in batches.iter() { @@ -258,42 +307,73 @@ impl PrometheusJsonResponse { .column(timestamp_column_index) .as_primitive::(); - let array = - arrow::compute::cast(batch.column(first_field_column_index), &DataType::Float64) - .context(ArrowSnafu)?; - let field_column = array.as_primitive::(); + let field_array = first_field_column_index + .map(|index| arrow::compute::cast(batch.column(index), &DataType::Float64)) + .transpose() + .context(ArrowSnafu)?; + let field_column = field_array + .as_ref() + .map(|array| array.as_primitive::()); + let native_histogram_column = native_histogram_column_index + .map(|index| { + batch + .column(index) + .as_any() + .downcast_ref::() + .with_context(|| UnexpectedResultSnafu { + reason: "native histogram column is not a struct array".to_string(), + }) + }) + .transpose()?; // assemble rows for row_index in 0..batch.num_rows() { - // retrieve value - if field_column.is_valid(row_index) { - let v = field_column.value(row_index); - // ignore all NaN values to reduce the amount of data to be sent. - if v.is_nan() { - continue; + let value = field_column.and_then(|field_column| { + if !field_column.is_valid(row_index) { + return None; } + let value = field_column.value(row_index); + (!value.is_nan()).then_some((timestamp_column.value(row_index), value)) + }); + let histogram = native_histogram_column + .and_then(|column| { + read_histogram(column, row_index) + .context(DataFusionSnafu) + .transpose() + }) + .transpose()? + .map(|histogram| { + prometheus_native_histogram_or_error(&histogram) + .map(|histogram| (timestamp_column.value(row_index), histogram)) + }) + .transpose()?; - // retrieve tags - let mut tags = Vec::with_capacity(num_label_columns + 1); - if let Some(metric_name) = &metric_name { - tags.push((METRIC_NAME, metric_name.as_str())); - } - for (tag_column, tag_name) in tag_columns.iter().zip(tag_names.iter()) { - // TODO(ruihang): add test for NULL tag - if tag_column.is_valid(row_index) { - tags.push((tag_name, tag_column.value(row_index))); - } + if value.is_none() && histogram.is_none() { + continue; + } + + // retrieve tags + let mut tags = Vec::with_capacity(num_label_columns + 1); + if let Some(metric_name) = &metric_name { + tags.push((METRIC_NAME, metric_name.as_str())); + } + for (tag_column, tag_name) in tag_columns.iter().zip(tag_names.iter()) { + // TODO(ruihang): add test for NULL tag + if tag_column.is_valid(row_index) { + tags.push((tag_name, tag_column.value(row_index))); } + } - // retrieve timestamp - let timestamp_millis = timestamp_column.value(row_index); - let timestamp = timestamp_millis as f64 / 1000.0; - - buffer - .entry(tags) - .or_default() - .push((timestamp, Into::::into(v).to_string())); - }; + let entry = buffer.entry(tags).or_default(); + if let Some((timestamp_millis, histogram)) = histogram { + entry + .histograms + .push((timestamp_millis as f64 / 1000.0, histogram)); + } else if let Some((timestamp_millis, value)) = value { + entry + .values + .push((timestamp_millis as f64 / 1000.0, value.to_string())); + } } } @@ -306,28 +386,46 @@ impl PrometheusJsonResponse { }; // accumulate data into result - buffer.into_iter().for_each(|(tags, mut values)| { + buffer.into_iter().for_each(|(tags, mut samples)| { let metric = tags .into_iter() .map(|(k, v)| (k.to_string(), v.to_string())) .collect::>(); match result { PromQueryResult::Vector(ref mut v) => { + let histogram = samples.histograms.pop(); + let value = if histogram.is_none() { + samples.values.pop() + } else { + None + }; v.push(PromSeriesVector { metric, - value: values.pop(), + value, + histogram, }); } PromQueryResult::Matrix(ref mut v) => { // sort values by timestamp - if !values.is_sorted_by(|a, b| a.0 <= b.0) { - values.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(Ordering::Equal)); + if !samples.values.is_sorted_by(|a, b| a.0 <= b.0) { + samples + .values + .sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(Ordering::Equal)); + } + if !samples.histograms.is_sorted_by(|a, b| a.0 <= b.0) { + samples + .histograms + .sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(Ordering::Equal)); } - v.push(PromSeriesMatrix { metric, values }); + v.push(PromSeriesMatrix { + metric, + values: samples.values, + histograms: samples.histograms, + }); } PromQueryResult::Scalar(ref mut v) => { - *v = values.pop(); + *v = samples.values.pop(); } PromQueryResult::String(ref mut _v) => { // TODO(ruihang): Not supported yet @@ -350,3 +448,267 @@ impl PrometheusJsonResponse { Ok(data) } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use common_query::native_histogram::{ + NATIVE_HISTOGRAM_FIELD, NativeHistogram, Span, build_histogram_array, + native_histogram_value_type, + }; + use common_query::{Output, OutputData, OutputMeta}; + use common_recordbatch::{RecordBatch, RecordBatches}; + use datatypes::data_type::ConcreteDataType; + use datatypes::schema::{ColumnSchema, Schema}; + use datatypes::vectors::{ + Float64Vector, StringVector, StructVector, TimestampMillisecondVector, VectorRef, + }; + + use super::*; + + fn sample_histogram() -> NativeHistogram { + NativeHistogram { + schema: 0, + zero_threshold: 0.001, + sum: 3.0, + reset_hint: 0, + start_timestamp: Some(0), + custom_values: vec![], + positive_spans: vec![Span { + offset: 0, + length: 1, + }], + negative_spans: vec![], + count: 2.0, + zero_count: 1.0, + positive_buckets: vec![1.0], + negative_buckets: vec![], + } + } + + fn histogram_vector(values: &[Option]) -> VectorRef { + let histogram_array = build_histogram_array(values); + let histogram_array = histogram_array + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + let ConcreteDataType::Struct(histogram_type) = native_histogram_value_type().clone() else { + unreachable!("native histogram type must be a struct") + }; + Arc::new(StructVector::try_new(histogram_type, histogram_array).unwrap()) + } + + #[tokio::test] + async fn from_query_result_preserves_warnings_and_infos() { + let recordbatches = RecordBatches::empty(); + let mut meta = OutputMeta::default(); + meta.warnings.push("warn".to_string()); + meta.infos.push("info".to_string()); + let response = PrometheusJsonResponse::from_query_result( + Ok(Output::new(OutputData::RecordBatches(recordbatches), meta)), + None, + ValueType::Vector, + ) + .await; + + assert_eq!(response.warnings, Some(vec!["warn".to_string()])); + assert_eq!(response.infos, Some(vec!["info".to_string()])); + } + + #[test] + fn record_batches_to_data_serializes_aliased_native_histogram_vector() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "greptime_timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new("job", ConcreteDataType::string_datatype(), false), + ColumnSchema::new( + "prom_native_histogram_rate(greptime_native_histogram)", + native_histogram_value_type().clone(), + true, + ), + ])); + let columns: Vec = vec![ + Arc::new(TimestampMillisecondVector::from_values([1000])), + Arc::new(StringVector::from(vec![Some("api")])), + histogram_vector(&[Some(sample_histogram())]), + ]; + let recordbatch = RecordBatch::new(schema.clone(), columns).unwrap(); + let recordbatches = RecordBatches::try_new(schema, vec![recordbatch]).unwrap(); + + let response = PrometheusJsonResponse::record_batches_to_data( + recordbatches, + Some("http_requests".to_string()), + ValueType::Vector, + ) + .unwrap(); + + let PrometheusResponse::PromData(PromData { + result: PromQueryResult::Vector(series), + .. + }) = response + else { + panic!("expected vector response"); + }; + assert_eq!(series.len(), 1); + assert_eq!(series[0].metric["__name__"], "http_requests"); + assert_eq!(series[0].metric["job"], "api"); + assert!(series[0].value.is_none()); + + let (timestamp, histogram) = series[0].histogram.as_ref().unwrap(); + assert_eq!(*timestamp, 1.0); + assert_eq!(histogram.count, "2"); + assert_eq!(histogram.sum, "3"); + assert_eq!( + histogram.buckets, + vec![ + ( + 3, + "-0.001".to_string(), + "0.001".to_string(), + "1".to_string() + ), + (0, "0.5".to_string(), "1".to_string(), "1".to_string()), + ] + ); + } + + #[test] + fn record_batches_to_data_prefers_histogram_for_instant_vector() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "greptime_timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new("job", ConcreteDataType::string_datatype(), false), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), true), + ColumnSchema::new( + NATIVE_HISTOGRAM_FIELD, + native_histogram_value_type().clone(), + true, + ), + ])); + let columns: Vec = vec![ + Arc::new(TimestampMillisecondVector::from_values([1000])), + Arc::new(StringVector::from(vec![Some("api")])), + Arc::new(Float64Vector::from_slice([42.0])), + histogram_vector(&[Some(sample_histogram())]), + ]; + let recordbatch = RecordBatch::new(schema.clone(), columns).unwrap(); + let recordbatches = RecordBatches::try_new(schema, vec![recordbatch]).unwrap(); + + let response = PrometheusJsonResponse::record_batches_to_data( + recordbatches, + Some("http_requests".to_string()), + ValueType::Vector, + ) + .unwrap(); + + let PrometheusResponse::PromData(PromData { + result: PromQueryResult::Vector(series), + .. + }) = response + else { + panic!("expected vector response"); + }; + + assert_eq!(series.len(), 1); + assert!(series[0].value.is_none()); + assert!(series[0].histogram.is_some()); + } + + #[test] + fn record_batches_to_data_prefers_histogram_for_matrix_row() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "greptime_timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new("job", ConcreteDataType::string_datatype(), false), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), true), + ColumnSchema::new( + NATIVE_HISTOGRAM_FIELD, + native_histogram_value_type().clone(), + true, + ), + ])); + let columns: Vec = vec![ + Arc::new(TimestampMillisecondVector::from_values([1000, 2000])), + Arc::new(StringVector::from(vec![Some("api"), Some("api")])), + Arc::new(Float64Vector::from_slice([42.0, 7.0])), + histogram_vector(&[Some(sample_histogram()), None]), + ]; + let recordbatch = RecordBatch::new(schema.clone(), columns).unwrap(); + let recordbatches = RecordBatches::try_new(schema, vec![recordbatch]).unwrap(); + + let response = PrometheusJsonResponse::record_batches_to_data( + recordbatches, + Some("http_requests".to_string()), + ValueType::Matrix, + ) + .unwrap(); + + let PrometheusResponse::PromData(PromData { + result: PromQueryResult::Matrix(series), + .. + }) = response + else { + panic!("expected matrix response"); + }; + + assert_eq!(series.len(), 1); + assert_eq!(series[0].values, vec![(2.0, "7".to_string())]); + assert_eq!(series[0].histograms.len(), 1); + assert_eq!(series[0].histograms[0].0, 1.0); + } + + #[test] + fn record_batches_to_data_errors_on_unrenderable_native_histogram() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "greptime_timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new("job", ConcreteDataType::string_datatype(), false), + ColumnSchema::new( + NATIVE_HISTOGRAM_FIELD, + native_histogram_value_type().clone(), + true, + ), + ])); + let mut histogram = sample_histogram(); + histogram.schema = 100; + let columns: Vec = vec![ + Arc::new(TimestampMillisecondVector::from_values([1000])), + Arc::new(StringVector::from(vec![Some("api")])), + histogram_vector(&[Some(histogram)]), + ]; + let recordbatch = RecordBatch::new(schema.clone(), columns).unwrap(); + let recordbatches = RecordBatches::try_new(schema, vec![recordbatch]).unwrap(); + + let error = PrometheusJsonResponse::record_batches_to_data( + recordbatches, + Some("http_requests".to_string()), + ValueType::Vector, + ) + .unwrap_err(); + + assert!( + error + .to_string() + .contains("native histogram cannot be converted to Prometheus buckets"), + "{error}" + ); + } +} diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index b4128994cef..2bb381a0d3d 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -16,9 +16,11 @@ use std::collections::hash_map::Entry; use ahash::{HashMap, HashMapExt, HashSet, HashSetExt}; use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; +use api::greptime_proto::io::prometheus::write::v2::{ + BucketSpan, Histogram, Request, Sample, TimeSeries, +}; #[cfg(test)] use api::greptime_proto::io::prometheus::write::v2::{Exemplar, Metadata, metadata}; -use api::greptime_proto::io::prometheus::write::v2::{Histogram, Request, Sample, TimeSeries}; use api::helper::ColumnDataTypeWrapper; use api::v1::value::ValueData; use api::v1::{ColumnSchema, ListValue, RowInsertRequest, Rows, SemanticType, Value}; @@ -43,6 +45,9 @@ use crate::row_writer::{self, TableData}; type PromTags = Vec<(String, String)>; type ResolvedSeriesLabels = (PromCtx, String, PromTags); +const MIN_REMOTE_WRITE_V2_SCHEMA: i32 = -4; +const MAX_REMOTE_WRITE_V2_SCHEMA: i32 = 8; +const MAX_REDUCIBLE_REMOTE_WRITE_V2_SCHEMA: i32 = 52; pub(crate) fn decode_remote_write_v2_request(is_zstd: bool, body: Bytes) -> Result { let _timer = crate::metrics::METRIC_HTTP_PROM_STORE_DECODE_ELAPSED.start_timer(); @@ -279,6 +284,9 @@ fn native_histogram_column_schema() -> ColumnSchema { } fn native_histogram_struct_value(histogram: &Histogram) -> Result { + let uses_float_counts = native_histogram_uses_float_counts(histogram)?; + validate_native_histogram(histogram, uses_float_counts)?; + let mut items = Vec::with_capacity(NATIVE_HISTOGRAM_FIELD_NAMES.len()); items.extend([ pb_value(ValueData::I32Value(histogram.schema)), @@ -295,39 +303,39 @@ fn native_histogram_struct_value(histogram: &Histogram) -> Result { u32_list_value(histogram.negative_spans.iter().map(|span| span.length)), ]); - let int_counts = match (&histogram.count, &histogram.zero_count) { - (Some(Count::CountInt(count)), Some(ZeroCount::ZeroCountInt(zero_count))) => { - (*count, *zero_count) - } - (Some(Count::CountInt(count)), _) => (*count, 0), - (_, Some(ZeroCount::ZeroCountInt(zero_count))) => (0, *zero_count), - _ => (0, 0), - }; - let float_counts = match (&histogram.count, &histogram.zero_count) { - (Some(Count::CountFloat(count)), Some(ZeroCount::ZeroCountFloat(zero_count))) => { - Some((*count, *zero_count)) - } - (Some(Count::CountFloat(count)), _) => Some((*count, 0.0)), - _ => None, - }; - - if let Some(counts) = float_counts { + if uses_float_counts { + let count = match histogram.count.as_ref() { + Some(Count::CountFloat(count)) => *count, + _ => 0.0, + }; + let zero_count = match histogram.zero_count.as_ref() { + Some(ZeroCount::ZeroCountFloat(zero_count)) => *zero_count, + _ => 0.0, + }; items.extend([ null_pb_value(), null_pb_value(), i64_list_value(std::iter::empty()), i64_list_value(std::iter::empty()), - pb_value(ValueData::F64Value(counts.0)), - pb_value(ValueData::F64Value(counts.1)), + pb_value(ValueData::F64Value(count)), + pb_value(ValueData::F64Value(zero_count)), f64_list_value(histogram.positive_counts.iter().copied()), f64_list_value(histogram.negative_counts.iter().copied()), ]); } else { + let count = match histogram.count.as_ref() { + Some(Count::CountInt(count)) => *count, + _ => 0, + }; + let zero_count = match histogram.zero_count.as_ref() { + Some(ZeroCount::ZeroCountInt(zero_count)) => *zero_count, + _ => 0, + }; let positive_buckets = bucket_counts_from_deltas(&histogram.positive_deltas)?; let negative_buckets = bucket_counts_from_deltas(&histogram.negative_deltas)?; items.extend([ - pb_value(ValueData::U64Value(int_counts.0)), - pb_value(ValueData::U64Value(int_counts.1)), + pb_value(ValueData::U64Value(count)), + pb_value(ValueData::U64Value(zero_count)), i64_list_value(positive_buckets.iter().copied()), i64_list_value(negative_buckets.iter().copied()), null_pb_value(), @@ -340,6 +348,229 @@ fn native_histogram_struct_value(histogram: &Histogram) -> Result { Ok(ValueData::StructValue(api::v1::StructValue { items })) } +fn validate_native_histogram(histogram: &Histogram, uses_float_counts: bool) -> Result<()> { + validate_native_histogram_schema(histogram.schema)?; + validate_native_histogram_custom_values(histogram)?; + + let (positive_buckets, negative_buckets) = if uses_float_counts { + ( + histogram.positive_counts.len(), + histogram.negative_counts.len(), + ) + } else { + ( + histogram.positive_deltas.len(), + histogram.negative_deltas.len(), + ) + }; + let custom_max_index = + (histogram.schema == CUSTOM_BUCKETS_SCHEMA).then_some(histogram.custom_values.len() as i32); + validate_native_histogram_spans( + "positive", + &histogram.positive_spans, + positive_buckets, + custom_max_index, + )?; + validate_native_histogram_spans( + "negative", + &histogram.negative_spans, + negative_buckets, + None, + )?; + + if histogram.schema == CUSTOM_BUCKETS_SCHEMA { + ensure!( + histogram.zero_threshold == 0.0 && native_histogram_zero_count_is_zero(histogram), + error::InvalidPromRemoteRequestSnafu { + msg: "remote write v2 custom native histogram must not use a zero bucket" + .to_string(), + } + ); + ensure!( + histogram.negative_spans.is_empty() + && histogram.negative_deltas.is_empty() + && histogram.negative_counts.is_empty(), + error::InvalidPromRemoteRequestSnafu { + msg: "remote write v2 custom native histogram must not use negative buckets" + .to_string(), + } + ); + } + + Ok(()) +} + +fn validate_native_histogram_schema(schema: i32) -> Result<()> { + if schema == CUSTOM_BUCKETS_SCHEMA + || (MIN_REMOTE_WRITE_V2_SCHEMA..=MAX_REMOTE_WRITE_V2_SCHEMA).contains(&schema) + { + return Ok(()); + } + + if (MAX_REMOTE_WRITE_V2_SCHEMA + 1..=MAX_REDUCIBLE_REMOTE_WRITE_V2_SCHEMA).contains(&schema) { + error::InvalidPromRemoteRequestSnafu { + msg: format!( + "remote write v2 native histogram schema {schema} must be reduced before ingestion" + ), + } + .fail() + } else { + error::InvalidPromRemoteRequestSnafu { + msg: format!("remote write v2 native histogram schema {schema} is unsupported"), + } + .fail() + } +} + +fn validate_native_histogram_custom_values(histogram: &Histogram) -> Result<()> { + if histogram.schema != CUSTOM_BUCKETS_SCHEMA { + ensure!( + histogram.custom_values.is_empty(), + error::InvalidPromRemoteRequestSnafu { + msg: "remote write v2 standard native histogram must not use custom_values" + .to_string(), + } + ); + return Ok(()); + } + + for value in &histogram.custom_values { + ensure!( + !value.is_nan() && *value != f64::INFINITY, + error::InvalidPromRemoteRequestSnafu { + msg: "remote write v2 custom native histogram custom_values must not contain +Inf or NaN" + .to_string(), + } + ); + } + for values in histogram.custom_values.windows(2) { + ensure!( + values[0] < values[1], + error::InvalidPromRemoteRequestSnafu { + msg: "remote write v2 custom native histogram custom_values must be sorted" + .to_string(), + } + ); + } + + Ok(()) +} + +fn validate_native_histogram_spans( + name: &str, + spans: &[BucketSpan], + bucket_count: usize, + custom_max_index: Option, +) -> Result<()> { + let span_len = spans + .iter() + .try_fold(0usize, |sum, span| sum.checked_add(span.length as usize)) + .context(error::InvalidPromRemoteRequestSnafu { + msg: format!("remote write v2 native histogram {name} spans overflow"), + })?; + ensure!( + span_len == bucket_count, + error::InvalidPromRemoteRequestSnafu { + msg: format!( + "remote write v2 native histogram {name} spans describe {span_len} buckets, found {bucket_count}" + ), + } + ); + + let mut current_index = 0i32; + let mut first = true; + for span in spans { + current_index = if first { + first = false; + span.offset + } else { + current_index.checked_add(span.offset).context( + error::InvalidPromRemoteRequestSnafu { + msg: format!( + "remote write v2 native histogram {name} span index overflows i32" + ), + }, + )? + }; + + for _ in 0..span.length { + if let Some(max_index) = custom_max_index { + ensure!( + (0..=max_index).contains(¤t_index), + error::InvalidPromRemoteRequestSnafu { + msg: format!( + "remote write v2 custom native histogram {name} bucket index {current_index} is out of range" + ), + } + ); + } + current_index = + current_index + .checked_add(1) + .context(error::InvalidPromRemoteRequestSnafu { + msg: format!( + "remote write v2 native histogram {name} span index overflows i32" + ), + })?; + } + } + + Ok(()) +} + +fn native_histogram_zero_count_is_zero(histogram: &Histogram) -> bool { + match histogram.zero_count.as_ref() { + Some(ZeroCount::ZeroCountInt(zero_count)) => *zero_count == 0, + Some(ZeroCount::ZeroCountFloat(zero_count)) => *zero_count == 0.0, + None => true, + } +} + +fn native_histogram_uses_float_counts(histogram: &Histogram) -> Result { + let uses_float_count = matches!(histogram.count, Some(Count::CountFloat(_))) + || matches!(histogram.zero_count, Some(ZeroCount::ZeroCountFloat(_))); + let uses_int_count = matches!(histogram.count, Some(Count::CountInt(_))) + || matches!(histogram.zero_count, Some(ZeroCount::ZeroCountInt(_))); + let uses_float_buckets = + !histogram.positive_counts.is_empty() || !histogram.negative_counts.is_empty(); + let uses_int_buckets = + !histogram.positive_deltas.is_empty() || !histogram.negative_deltas.is_empty(); + + if matches!( + (&histogram.count, &histogram.zero_count), + (Some(Count::CountInt(_)), Some(ZeroCount::ZeroCountFloat(_))) + | (Some(Count::CountFloat(_)), Some(ZeroCount::ZeroCountInt(_))) + ) { + return error::InvalidPromRemoteRequestSnafu { + msg: "remote write v2 native histogram count and zero_count must use the same integer or float family".to_string(), + } + .fail(); + } + + ensure!( + !(uses_float_buckets && uses_int_buckets), + error::InvalidPromRemoteRequestSnafu { + msg: "remote write v2 native histogram bucket counts must use either integer deltas or float counts".to_string(), + } + ); + ensure!( + !(uses_float_count && uses_int_buckets), + error::InvalidPromRemoteRequestSnafu { + msg: "remote write v2 float native histogram must not use integer bucket deltas" + .to_string(), + } + ); + ensure!( + !(uses_int_count && uses_float_buckets), + error::InvalidPromRemoteRequestSnafu { + msg: "remote write v2 integer native histogram must not use float bucket counts" + .to_string(), + } + ); + + Ok(uses_float_count || uses_float_buckets) +} + fn pb_value(value_data: ValueData) -> Value { optional_pb_value(Some(value_data)) } @@ -859,6 +1090,126 @@ mod tests { "conflicts with an internal native histogram label", )); + cases.push(( + "int count with float zero count", + request_with_histogram(Histogram { + count: Some(Count::CountInt(1)), + zero_count: Some(ZeroCount::ZeroCountFloat(0.5)), + ..Default::default() + }), + "count and zero_count must use the same integer or float family", + )); + + cases.push(( + "float count with int zero count", + request_with_histogram(Histogram { + count: Some(Count::CountFloat(1.0)), + zero_count: Some(ZeroCount::ZeroCountInt(1)), + ..Default::default() + }), + "count and zero_count must use the same integer or float family", + )); + + cases.push(( + "reducible schema", + request_with_histogram(Histogram { + schema: 9, + ..Default::default() + }), + "schema 9 must be reduced before ingestion", + )); + + cases.push(( + "unsupported schema", + request_with_histogram(Histogram { + schema: 53, + ..Default::default() + }), + "schema 53 is unsupported", + )); + + cases.push(( + "standard schema with custom values", + request_with_histogram(Histogram { + schema: 1, + custom_values: vec![1.0], + ..Default::default() + }), + "standard native histogram must not use custom_values", + )); + + cases.push(( + "custom values with inf", + request_with_histogram(Histogram { + schema: CUSTOM_BUCKETS_SCHEMA, + custom_values: vec![f64::INFINITY], + ..Default::default() + }), + "custom_values must not contain +Inf or NaN", + )); + + cases.push(( + "custom values not sorted", + request_with_histogram(Histogram { + schema: CUSTOM_BUCKETS_SCHEMA, + custom_values: vec![2.0, 1.0], + ..Default::default() + }), + "custom_values must be sorted", + )); + + cases.push(( + "custom schema with zero bucket", + request_with_histogram(Histogram { + schema: CUSTOM_BUCKETS_SCHEMA, + zero_count: Some(ZeroCount::ZeroCountInt(1)), + ..Default::default() + }), + "custom native histogram must not use a zero bucket", + )); + + cases.push(( + "custom schema with negative buckets", + request_with_histogram(Histogram { + schema: CUSTOM_BUCKETS_SCHEMA, + negative_spans: vec![BucketSpan { + offset: -1, + length: 1, + }], + negative_deltas: vec![1], + ..Default::default() + }), + "custom native histogram must not use negative buckets", + )); + + cases.push(( + "span count mismatch", + request_with_histogram(Histogram { + positive_spans: vec![BucketSpan { + offset: 0, + length: 2, + }], + positive_deltas: vec![1], + ..Default::default() + }), + "positive spans describe 2 buckets, found 1", + )); + + cases.push(( + "custom span index out of range", + request_with_histogram(Histogram { + schema: CUSTOM_BUCKETS_SCHEMA, + custom_values: vec![1.0], + positive_spans: vec![BucketSpan { + offset: 2, + length: 1, + }], + positive_deltas: vec![1], + ..Default::default() + }), + "positive bucket index 2 is out of range", + )); + for (name, request, expected) in cases { assert_invalid(name, request, expected); } @@ -1115,6 +1466,13 @@ mod tests { ) } + fn request_with_histogram(histogram: Histogram) -> Request { + test_util::request_with_labels_and_histograms( + vec![(METRIC_NAME_LABEL, "metric")], + vec![histogram], + ) + } + fn assert_invalid(name: &str, request: Request, expected: &str) { let err = match into_write_requests(request) { Ok(_) => panic!("{name}: expected invalid request error"), diff --git a/tests-integration/tests/grpc.rs b/tests-integration/tests/grpc.rs index d6b96484af0..aaed03b66fa 100644 --- a/tests-integration/tests/grpc.rs +++ b/tests-integration/tests/grpc.rs @@ -872,6 +872,7 @@ pub async fn test_prom_gateway_query(store_type: StorageType) { .into_iter() .collect(), value: Some((5.0, "1".to_string())), + ..Default::default() }, PromSeriesVector { metric: [ @@ -881,6 +882,7 @@ pub async fn test_prom_gateway_query(store_type: StorageType) { .into_iter() .collect(), value: Some((5.0, "2".to_string())), + ..Default::default() }, ] ); @@ -932,6 +934,7 @@ pub async fn test_prom_gateway_query(store_type: StorageType) { .into_iter() .collect(), values: vec![(5.0, "1".to_string()), (10.0, "1".to_string())], + ..Default::default() }, PromSeriesMatrix { metric: [ @@ -941,6 +944,7 @@ pub async fn test_prom_gateway_query(store_type: StorageType) { .into_iter() .collect(), values: vec![(5.0, "2".to_string()), (10.0, "2".to_string())], + ..Default::default() }, ] ); @@ -973,6 +977,7 @@ pub async fn test_prom_gateway_query(store_type: StorageType) { error: None, error_type: None, warnings: None, + infos: None, resp_metrics: Default::default(), status_code: None, }; diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 8b6dcd27aff..888e7434796 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2455,7 +2455,6 @@ pub async fn test_prometheus_remote_write_v2_native_histogram(store_type: Storag reset_hint: 2, timestamp: 3000, start_timestamp: 1500, - custom_values: vec![0.5, 1.5], ..Default::default() }, Histogram { @@ -2508,10 +2507,58 @@ pub async fn test_prometheus_remote_write_v2_native_histogram(store_type: Storag "prometheus_remote_write_v2_native_histogram_rows", &client, "select greptime_timestamp, greptime_native_histogram, job, instance from remote_write_v2_latency_seconds order by greptime_timestamp;", - "[[3000,{\"count_f64\":null,\"count_u64\":4,\"custom_values\":[0.5,1.5],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[1],\"negative_span_lengths\":[1],\"negative_span_offsets\":[-2],\"positive_buckets_f64\":[],\"positive_buckets_i64\":[1,3,2],\"positive_span_lengths\":[3],\"positive_span_offsets\":[0],\"reset_hint\":2,\"schema\":1,\"start_timestamp\":1500,\"sum\":10.0,\"zero_count_f64\":null,\"zero_count_u64\":1,\"zero_threshold\":0.001},\"api\",\"localhost:9090\"],[4000,{\"count_f64\":3.5,\"count_u64\":null,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[],\"negative_span_lengths\":[],\"negative_span_offsets\":[],\"positive_buckets_f64\":[2.0,3.5],\"positive_buckets_i64\":[],\"positive_span_lengths\":[2],\"positive_span_offsets\":[3],\"reset_hint\":3,\"schema\":2,\"start_timestamp\":2500,\"sum\":20.0,\"zero_count_f64\":0.5,\"zero_count_u64\":null,\"zero_threshold\":0.002},\"api\",\"localhost:9090\"]]", + "[[3000,{\"count_f64\":null,\"count_u64\":4,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[1],\"negative_span_lengths\":[1],\"negative_span_offsets\":[-2],\"positive_buckets_f64\":[],\"positive_buckets_i64\":[1,3,2],\"positive_span_lengths\":[3],\"positive_span_offsets\":[0],\"reset_hint\":2,\"schema\":1,\"start_timestamp\":1500,\"sum\":10.0,\"zero_count_f64\":null,\"zero_count_u64\":1,\"zero_threshold\":0.001},\"api\",\"localhost:9090\"],[4000,{\"count_f64\":3.5,\"count_u64\":null,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[],\"negative_span_lengths\":[],\"negative_span_offsets\":[],\"positive_buckets_f64\":[2.0,3.5],\"positive_buckets_i64\":[],\"positive_span_lengths\":[2],\"positive_span_offsets\":[3],\"reset_hint\":3,\"schema\":2,\"start_timestamp\":2500,\"sum\":20.0,\"zero_count_f64\":0.5,\"zero_count_u64\":null,\"zero_threshold\":0.002},\"api\",\"localhost:9090\"]]", ) .await; + let res = client + .get("/v1/prometheus/api/v1/query?query=histogram_count(remote_write_v2_latency_seconds)&time=4") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = serde_json::from_str::(&res.text().await).unwrap(); + assert_eq!(body["status"], "success"); + assert_eq!(body["data"]["resultType"], "vector"); + let series = body["data"]["result"].as_array().unwrap(); + assert_eq!(series.len(), 1); + let value = series[0]["value"].as_array().unwrap(); + assert_eq!(value[0].as_f64(), Some(4.0)); + assert_eq!(value[1], "3.5"); + + let res = client + .get("/v1/prometheus/api/v1/query?query=remote_write_v2_latency_seconds&time=4") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = serde_json::from_str::(&res.text().await).unwrap(); + assert_eq!(body["status"], "success"); + assert_eq!(body["data"]["resultType"], "vector"); + let series = body["data"]["result"].as_array().unwrap(); + assert_eq!(series.len(), 1); + assert!(series[0].get("value").is_none()); + let histogram = series[0]["histogram"].as_array().unwrap(); + assert_eq!(histogram[0].as_f64(), Some(4.0)); + assert_eq!(histogram[1]["count"], "3.5"); + assert_eq!(histogram[1]["sum"], "20"); + + let res = client + .get("/v1/prometheus/api/v1/query_range?query=remote_write_v2_latency_seconds&start=3&end=4&step=1") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = serde_json::from_str::(&res.text().await).unwrap(); + assert_eq!(body["status"], "success"); + assert_eq!(body["data"]["resultType"], "matrix"); + let series = body["data"]["result"].as_array().unwrap(); + assert_eq!(series.len(), 1); + assert!(series[0].get("values").is_none()); + let histograms = series[0]["histograms"].as_array().unwrap(); + assert_eq!(histograms.len(), 2); + assert_eq!(histograms[0][0].as_f64(), Some(3.0)); + assert_eq!(histograms[0][1]["count"], "4"); + assert_eq!(histograms[1][0].as_f64(), Some(4.0)); + assert_eq!(histograms[1][1]["count"], "3.5"); + validate_data( "prometheus_remote_write_v2_native_histogram_table_created", &client,