diff --git a/src/common/query/src/native_histogram.rs b/src/common/query/src/native_histogram.rs index 3fdec9b309a..11f616d4a2e 100644 --- a/src/common/query/src/native_histogram.rs +++ b/src/common/query/src/native_histogram.rs @@ -23,13 +23,13 @@ use std::collections::BTreeMap; use std::sync::Arc; use datafusion::arrow::array::{ - Array, ArrayRef, Float64Array, Int32Array, ListArray, PrimitiveArray, StructArray, - TimestampMillisecondArray, UInt64Array, + Array, ArrayRef, Float64Array, Int32Array, Int64Array, ListArray, PrimitiveArray, StructArray, + TimestampMillisecondArray, }; use datafusion::arrow::buffer::NullBuffer; use datafusion::arrow::datatypes::{ ArrowPrimitiveType, DataType as ArrowDataType, Field, Float64Type, Int32Type, Int64Type, - TimestampMillisecondType, UInt32Type, UInt64Type, + TimestampMillisecondType, }; use datafusion_common::{DataFusionError, Result as DfResult}; use datatypes::data_type::{ConcreteDataType, DataType}; @@ -50,8 +50,8 @@ pub const POSITIVE_SPAN_OFFSETS_FIELD: &str = "positive_span_offsets"; pub const POSITIVE_SPAN_LENGTHS_FIELD: &str = "positive_span_lengths"; pub const NEGATIVE_SPAN_OFFSETS_FIELD: &str = "negative_span_offsets"; pub const NEGATIVE_SPAN_LENGTHS_FIELD: &str = "negative_span_lengths"; -pub const COUNT_U64_FIELD: &str = "count_u64"; -pub const ZERO_COUNT_U64_FIELD: &str = "zero_count_u64"; +pub const COUNT_I64_FIELD: &str = "count_i64"; +pub const ZERO_COUNT_I64_FIELD: &str = "zero_count_i64"; pub const POSITIVE_BUCKETS_I64_FIELD: &str = "positive_buckets_i64"; pub const NEGATIVE_BUCKETS_I64_FIELD: &str = "negative_buckets_i64"; pub const COUNT_F64_FIELD: &str = "count_f64"; @@ -72,8 +72,8 @@ pub const NATIVE_HISTOGRAM_FIELD_NAMES: &[&str] = &[ POSITIVE_SPAN_LENGTHS_FIELD, NEGATIVE_SPAN_OFFSETS_FIELD, NEGATIVE_SPAN_LENGTHS_FIELD, - COUNT_U64_FIELD, - ZERO_COUNT_U64_FIELD, + COUNT_I64_FIELD, + ZERO_COUNT_I64_FIELD, POSITIVE_BUCKETS_I64_FIELD, NEGATIVE_BUCKETS_I64_FIELD, COUNT_F64_FIELD, @@ -108,9 +108,9 @@ pub fn native_histogram_field_type(name: &str) -> Option { ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::int32_datatype())), ), POSITIVE_SPAN_LENGTHS_FIELD | NEGATIVE_SPAN_LENGTHS_FIELD => Some( - ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::uint32_datatype())), + ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::int32_datatype())), ), - COUNT_U64_FIELD | ZERO_COUNT_U64_FIELD => Some(ConcreteDataType::uint64_datatype()), + COUNT_I64_FIELD | ZERO_COUNT_I64_FIELD => Some(ConcreteDataType::int64_datatype()), POSITIVE_BUCKETS_I64_FIELD | NEGATIVE_BUCKETS_I64_FIELD => Some( ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::int64_datatype())), ), @@ -195,7 +195,7 @@ pub struct Span { /// The first bucket index, or the gap after the preceding span. pub offset: i32, /// Number of consecutive buckets in the span. - pub length: u32, + pub length: i32, } /// Inclusion rules for a materialized bucket's lower and upper bounds. @@ -1320,7 +1320,7 @@ where .collect() } -fn read_spans(offsets: Vec, lengths: Vec, name: &str) -> DfResult> { +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 {}", @@ -1328,11 +1328,18 @@ fn read_spans(offsets: Vec, lengths: Vec, name: &str) -> DfResult DfResult<()> { @@ -1362,12 +1369,12 @@ pub fn read_histogram(array: &StructArray, row: usize) -> DfResult(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)?, + 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)?, + list_values::(array, NEGATIVE_SPAN_LENGTHS_FIELD, row)?, "negative", )?; @@ -1382,8 +1389,8 @@ pub fn read_histogram(array: &StructArray, row: usize) -> DfResult(array, COUNT_U64_FIELD, row)? as f64, - optional_primitive::(array, ZERO_COUNT_U64_FIELD, row)? + required_primitive::(array, COUNT_I64_FIELD, row)? as f64, + optional_primitive::(array, ZERO_COUNT_I64_FIELD, row)? .unwrap_or_default() as f64, list_values::(array, POSITIVE_BUCKETS_I64_FIELD, row)? .into_iter() @@ -1439,8 +1446,8 @@ pub fn build_histogram_array(values: &[Option]) -> ArrayRef { 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 count_i64 = Vec::>::with_capacity(values.len()); + let mut zero_count_i64 = 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()); @@ -1486,8 +1493,8 @@ pub fn build_histogram_array(values: &[Option]) -> ArrayRef { .map(|span| span.length) .collect(), )); - count_u64.push(None); - zero_count_u64.push(None); + count_i64.push(None); + zero_count_i64.push(None); positive_buckets_i64.push(list_opt(Vec::::new())); negative_buckets_i64.push(list_opt(Vec::::new())); count_f64.push(Some(histogram.count)); @@ -1505,8 +1512,8 @@ pub fn build_histogram_array(values: &[Option]) -> ArrayRef { positive_span_lengths.push(None); negative_span_offsets.push(None); negative_span_lengths.push(None); - count_u64.push(None); - zero_count_u64.push(None); + count_i64.push(None); + zero_count_i64.push(None); positive_buckets_i64.push(None); negative_buckets_i64.push(None); count_f64.push(None); @@ -1542,7 +1549,7 @@ pub fn build_histogram_array(values: &[Option]) -> ArrayRef { ), ( POSITIVE_SPAN_LENGTHS_FIELD, - Arc::new(ListArray::from_iter_primitive::( + Arc::new(ListArray::from_iter_primitive::( positive_span_lengths, )), ), @@ -1554,14 +1561,14 @@ pub fn build_histogram_array(values: &[Option]) -> ArrayRef { ), ( NEGATIVE_SPAN_LENGTHS_FIELD, - Arc::new(ListArray::from_iter_primitive::( + Arc::new(ListArray::from_iter_primitive::( negative_span_lengths, )), ), - (COUNT_U64_FIELD, Arc::new(UInt64Array::from(count_u64))), + (COUNT_I64_FIELD, Arc::new(Int64Array::from(count_i64))), ( - ZERO_COUNT_U64_FIELD, - Arc::new(UInt64Array::from(zero_count_u64)), + ZERO_COUNT_I64_FIELD, + Arc::new(Int64Array::from(zero_count_i64)), ), ( POSITIVE_BUCKETS_I64_FIELD, @@ -1672,7 +1679,7 @@ mod tests { assert_eq!(read_histogram(array, 0).unwrap(), Some(expected)); assert!( - primitive_child::(array, COUNT_U64_FIELD) + primitive_child::(array, COUNT_I64_FIELD) .unwrap() .is_null(0) ); diff --git a/src/promql/src/functions/native_histogram.rs b/src/promql/src/functions/native_histogram.rs index 5208db8c6b1..5074a5b6ec7 100644 --- a/src/promql/src/functions/native_histogram.rs +++ b/src/promql/src/functions/native_histogram.rs @@ -1922,7 +1922,7 @@ mod tests { custom_values: Vec::new(), positive_spans: vec![Span { offset: 0, - length: positive_buckets.len() as u32, + length: positive_buckets.len() as i32, }], negative_spans: Vec::new(), count, diff --git a/src/servers/src/prom_remote_write/README.md b/src/servers/src/prom_remote_write/README.md index 1f15f8288a7..d21eead1901 100644 --- a/src/servers/src/prom_remote_write/README.md +++ b/src/servers/src/prom_remote_write/README.md @@ -47,7 +47,7 @@ Each histogram row stores one Struct field named - common scalar children: `schema`, `zero_threshold`, `sum`, `reset_hint`, `start_timestamp`; -- count children: `count_u64` / `zero_count_u64` or `count_f64` / `zero_count_f64`; +- count children: `count_i64` / `zero_count_i64` or `count_f64` / `zero_count_f64`; - list children for custom values, spans, and positive/negative buckets; - original Prometheus labels as Greptime tags. diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index 923f860be38..1010b915953 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -288,6 +288,8 @@ fn native_histogram_struct_value(histogram: &Histogram) -> Result { validate_native_histogram(histogram, uses_float_counts)?; let mut items = Vec::with_capacity(NATIVE_HISTOGRAM_FIELD_NAMES.len()); + let positive_span_lengths = i32_span_lengths("positive", &histogram.positive_spans)?; + let negative_span_lengths = i32_span_lengths("negative", &histogram.negative_spans)?; items.extend([ pb_value(ValueData::I32Value(histogram.schema)), pb_value(ValueData::F64Value(histogram.zero_threshold)), @@ -298,9 +300,9 @@ fn native_histogram_struct_value(histogram: &Histogram) -> Result { )), f64_list_value(histogram.custom_values.iter().copied()), i32_list_value(histogram.positive_spans.iter().map(|span| span.offset)), - u32_list_value(histogram.positive_spans.iter().map(|span| span.length)), + i32_list_value(positive_span_lengths.iter().copied()), i32_list_value(histogram.negative_spans.iter().map(|span| span.offset)), - u32_list_value(histogram.negative_spans.iter().map(|span| span.length)), + i32_list_value(negative_span_lengths.iter().copied()), ]); if uses_float_counts { @@ -335,9 +337,23 @@ fn native_histogram_struct_value(histogram: &Histogram) -> Result { let positive_buckets = bucket_counts_from_deltas(&histogram.positive_deltas)?; let negative_buckets = bucket_counts_from_deltas(&histogram.negative_deltas)?; validate_integer_native_histogram_counts(histogram, &positive_buckets, &negative_buckets)?; + let count = i64::try_from(count) + .ok() + .context(error::InvalidPromRemoteRequestSnafu { + msg: format!( + "remote write v2 native histogram integer count {count} overflows i64" + ), + })?; + let zero_count = i64::try_from(zero_count).ok().context( + error::InvalidPromRemoteRequestSnafu { + msg: format!( + "remote write v2 native histogram integer zero_count {zero_count} overflows i64" + ), + }, + )?; items.extend([ - pb_value(ValueData::U64Value(count)), - pb_value(ValueData::U64Value(zero_count)), + pb_value(ValueData::I64Value(count)), + pb_value(ValueData::I64Value(zero_count)), i64_list_value(positive_buckets.iter().copied()), i64_list_value(negative_buckets.iter().copied()), null_pb_value(), @@ -695,8 +711,20 @@ fn i32_list_value(values: impl IntoIterator) -> Value { list_value(values.into_iter().map(ValueData::I32Value)) } -fn u32_list_value(values: impl IntoIterator) -> Value { - list_value(values.into_iter().map(ValueData::U32Value)) +fn i32_span_lengths(name: &str, spans: &[BucketSpan]) -> Result> { + spans + .iter() + .map(|span| { + i32::try_from(span.length) + .ok() + .context(error::InvalidPromRemoteRequestSnafu { + msg: format!( + "remote write v2 native histogram {name} span length {} overflows i32", + span.length + ), + }) + }) + .collect() } fn i64_list_value(values: impl IntoIterator) -> Value { @@ -1524,8 +1552,8 @@ mod tests { Some(ValueData::I32Value(0)) ); assert_eq!( - histogram_field_value(&rows, 0, COUNT_U64_FIELD), - Some(ValueData::U64Value(0)) + histogram_field_value(&rows, 0, COUNT_I64_FIELD), + Some(ValueData::I64Value(0)) ); assert_eq!(histogram_field_value(&rows, 0, COUNT_F64_FIELD), None); } @@ -1641,8 +1669,8 @@ mod tests { ); assert_eq!( - histogram_field_value(&rows, 0, COUNT_U64_FIELD), - Some(ValueData::U64Value(0)) + histogram_field_value(&rows, 0, COUNT_I64_FIELD), + Some(ValueData::I64Value(0)) ); assert_eq!(histogram_field_value(&rows, 0, COUNT_F64_FIELD), None); assert!(matches!( @@ -1655,7 +1683,7 @@ mod tests { POSITIVE_BUCKETS_F64_FIELD ))); - assert_eq!(histogram_field_value(&rows, 1, COUNT_U64_FIELD), None); + assert_eq!(histogram_field_value(&rows, 1, COUNT_I64_FIELD), None); assert_eq!( histogram_field_value(&rows, 1, COUNT_F64_FIELD), Some(ValueData::F64Value(6.0)) diff --git a/src/servers/tests/prom_remote_write_v2_test.rs b/src/servers/tests/prom_remote_write_v2_test.rs index a44583dccd1..94d260d3fac 100644 --- a/src/servers/tests/prom_remote_write_v2_test.rs +++ b/src/servers/tests/prom_remote_write_v2_test.rs @@ -18,7 +18,7 @@ use api::v1::value::ValueData; use api::v1::{ColumnSchema, Rows}; use bytes::Bytes; use common_query::native_histogram::{ - COUNT_U64_FIELD, NATIVE_HISTOGRAM_FIELD_NAMES, POSITIVE_BUCKETS_F64_FIELD, + COUNT_I64_FIELD, NATIVE_HISTOGRAM_FIELD_NAMES, POSITIVE_BUCKETS_F64_FIELD, POSITIVE_BUCKETS_I64_FIELD, POSITIVE_SPAN_OFFSETS_FIELD, SCHEMA_FIELD, }; use common_query::prelude::greptime_native_histogram; @@ -100,8 +100,8 @@ fn test_decode_remote_write_v2_native_histogram_dump() { Some(ValueData::I32Value(3)) ); assert_eq!( - histogram_field_value(rows, 0, COUNT_U64_FIELD), - Some(ValueData::U64Value(24)) + histogram_field_value(rows, 0, COUNT_I64_FIELD), + Some(ValueData::I64Value(24)) ); assert_eq!( list_i32_values(histogram_field_value(rows, 0, POSITIVE_SPAN_OFFSETS_FIELD)), diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index dbf085d970e..d524c96328a 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2811,7 +2811,7 @@ 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\":8,\"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\":6.0,\"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_i64\":8,\"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_i64\":1,\"zero_threshold\":0.001},\"api\",\"localhost:9090\"],[4000,{\"count_f64\":6.0,\"count_i64\":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_i64\":null,\"zero_threshold\":0.002},\"api\",\"localhost:9090\"]]", ) .await;