refactor: centralize native histogram encoding in common-query (#8945)

* refactor: centralize native histogram encoding in common-query

Signed-off-by: shuiyisong <xixing.sys@gmail.com>

* fix: fmt

Signed-off-by: shuiyisong <xixing.sys@gmail.com>

---------

Signed-off-by: shuiyisong <xixing.sys@gmail.com>

src/common/query/src/native_histogram.rs

@@ -19,6 +19,8 @@
 //! [`NativeHistogram`] is the query-time representation and therefore normalizes
 //! integer and floating-point payloads to absolute `f64` counts.
 
+mod encoding;
+
 use std::collections::BTreeMap;
 use std::sync::Arc;
 
@@ -34,6 +36,7 @@ use datafusion::arrow::datatypes::{
 use datafusion_common::{DataFusionError, Result as DfResult};
 use datatypes::data_type::{ConcreteDataType, DataType};
 use datatypes::types::{StructField, StructType};
+pub use encoding::{NativeHistogramError, encode_native_histogram, native_histogram_column_schema};
 use once_cell::sync::Lazy;
 
 use crate::prelude::greptime_native_histogram;

src/common/query/src/native_histogram/encoding.rs

@@ -17,23 +17,27 @@ use api::greptime_proto::io::prometheus::write::v2::{BucketSpan, Histogram};
 use api::helper::ColumnDataTypeWrapper;
 use api::v1::value::ValueData;
 use api::v1::{ColumnSchema, ListValue, SemanticType, Value};
-use common_query::native_histogram::*;
-use common_query::prelude::greptime_native_histogram;
 use snafu::{Snafu, ensure};
 
-const MAX_NATIVE_HISTOGRAM_SCHEMA: i32 = 8;
+use crate::native_histogram::{
+    CUSTOM_BUCKETS_SCHEMA, MAX_EXPONENTIAL_SCHEMA, NATIVE_HISTOGRAM_FIELD_NAMES,
+    exponential_overflow_bucket_index, native_histogram_value_type,
+};
+use crate::prelude::greptime_native_histogram;
+
 const MAX_REDUCIBLE_NATIVE_HISTOGRAM_SCHEMA: i32 = 52;
 
+/// Error returned while validating or encoding a native histogram.
 #[derive(Debug, Snafu)]
 #[snafu(display("{message}"))]
-pub(crate) struct NativeHistogramError {
+pub struct NativeHistogramError {
     message: String,
 }
 
 type Result<T> = std::result::Result<T, NativeHistogramError>;
 
 /// Returns the canonical column schema for a native histogram value.
-pub(crate) fn native_histogram_column_schema() -> Result<ColumnSchema> {
+pub fn native_histogram_column_schema() -> Result<ColumnSchema> {
     let (datatype, datatype_extension) =
         ColumnDataTypeWrapper::try_from(native_histogram_value_type().clone())
             .map_err(|error| NativeHistogramError {
@@ -51,7 +55,7 @@ pub(crate) fn native_histogram_column_schema() -> Result<ColumnSchema> {
 }
 
 /// Validates and encodes a Prometheus histogram into the canonical Struct value.
-pub(crate) fn encode_native_histogram(histogram: &Histogram) -> Result<ValueData> {
+pub fn encode_native_histogram(histogram: &Histogram) -> Result<ValueData> {
     let uses_float_counts = native_histogram_uses_float_counts(histogram)?;
     validate_native_histogram(histogram, uses_float_counts)?;
 
@@ -193,7 +197,7 @@ fn validate_native_histogram_schema(schema: i32) -> Result<Option<i32>> {
         return Ok(Some(overflow_index));
     }
 
-    if (MAX_NATIVE_HISTOGRAM_SCHEMA + 1..=MAX_REDUCIBLE_NATIVE_HISTOGRAM_SCHEMA).contains(&schema) {
+    if (MAX_EXPONENTIAL_SCHEMA + 1..=MAX_REDUCIBLE_NATIVE_HISTOGRAM_SCHEMA).contains(&schema) {
         Err(NativeHistogramError {
             message: format!("native histogram schema {schema} must be reduced before ingestion"),
         })

src/servers/src/lib.rs

@@ -34,7 +34,6 @@ pub mod interceptor;
 pub mod metrics;
 pub mod metrics_handler;
 pub mod mysql;
-pub(crate) mod native_histogram;
 pub mod opentsdb;
 pub mod otel_arrow;
 pub mod otlp;

src/servers/src/otlp/metrics.rs

@@ -20,7 +20,9 @@ use api::greptime_proto::io::prometheus::write::v2::{BucketSpan, Histogram as Pr
 use api::v1::value::ValueData;
 use api::v1::{RowInsertRequests, SemanticType, Value};
 use common_grpc::precision::Precision;
-use common_query::native_histogram::native_histogram_value_type;
+use common_query::native_histogram::{
+    encode_native_histogram, native_histogram_column_schema, native_histogram_value_type,
+};
 use common_query::prelude::{GREPTIME_COUNT, greptime_timestamp, greptime_value};
 use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS;
 use lazy_static::lazy_static;
@@ -34,7 +36,6 @@ use table::requests::{
 };
 
 use crate::error::{self, Result};
-use crate::native_histogram::{encode_native_histogram, native_histogram_column_schema};
 use crate::otlp::trace::{KEY_SERVICE_INSTANCE_ID, KEY_SERVICE_NAME};
 use crate::query_handler::MetricsIngestOutcome;
 use crate::row_writer::{self, MultiTableData, TableData};

src/servers/src/prom_remote_write/v2.rs

@@ -30,9 +30,11 @@ use api::v1::value::ValueData;
 use api::v1::{ColumnDataType, RowInsertRequest, Rows, SemanticType, Value};
 use bytes::{Buf, Bytes};
 use common_grpc::precision::Precision;
-use common_query::native_histogram::NATIVE_HISTOGRAM_FIELD;
 #[cfg(test)]
 use common_query::native_histogram::*;
+use common_query::native_histogram::{
+    NATIVE_HISTOGRAM_FIELD, encode_native_histogram, native_histogram_column_schema,
+};
 use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value};
 use pipeline::{ContextOpt, ContextReq};
 use prost::encoding::{
@@ -46,7 +48,6 @@ use table::requests::{
 };
 
 use crate::error::{self, Result};
-use crate::native_histogram::{encode_native_histogram, native_histogram_column_schema};
 use crate::prom_remote_write::row_builder::PromCtx;
 use crate::prom_remote_write::validation::validate_label_name;
 use crate::prom_remote_write::{REMOTE_WRITE_V2_VERSION, try_decompress};