compute timeseries_id for sketches (#6348)

* compute timeseries_id for sketches

* address comment

* linter

Author: mattmkim

quickwit/quickwit-opentelemetry/src/otlp/arrow_metrics.rs

@@ -37,6 +37,8 @@ use quickwit_proto::types::DocUid;
 
 use super::otel_metrics::{MetricDataPoint, MetricType};
 
+const COMPUTED_TIMESERIES_ID_FIELD: &str = "timeseries_id";
+
 /// Builder for creating Arrow RecordBatch from MetricDataPoints.
 ///
 /// Accumulates data points and discovers the schema dynamically at `finish()`
@@ -70,7 +72,12 @@ impl ArrowMetricsBatchBuilder {
         let mut tag_keys: BTreeSet<&str> = BTreeSet::new();
         for dp in &self.data_points {
             for key in dp.tags.keys() {
-                tag_keys.insert(key.as_str());
+                // Exclude any user-provided tag named "timeseries_id" to prevent collision with the
+                // computed column. TODO: if user sets "timeseries_id" as a tag, we
+                // are excluding it.
+                if key != COMPUTED_TIMESERIES_ID_FIELD {
+                    tag_keys.insert(key.as_str());
+                }
             }
         }
         let sorted_tag_keys: Vec<&str> = tag_keys.into_iter().collect();
@@ -86,9 +93,11 @@ impl ArrowMetricsBatchBuilder {
         fields.push(Field::new("metric_type", DataType::UInt8, false));
         fields.push(Field::new("timestamp_secs", DataType::UInt64, false));
         fields.push(Field::new("value", DataType::Float64, false));
-        // TODO: customer could submit a timeseries_id tag, and I don't think we want to explicitly
-        // reserve it.
-        fields.push(Field::new("timeseries_id", DataType::Int64, false));
+        fields.push(Field::new(
+            COMPUTED_TIMESERIES_ID_FIELD,
+            DataType::Int64,
+            false,
+        ));
 
         for &tag_key in &sorted_tag_keys {
             fields.push(Field::new(
@@ -132,11 +141,9 @@ impl ArrowMetricsBatchBuilder {
             // TODO: can we not have to compute the timeseries for every point? especially with
             // compaction, there may be many points with the same tags, in the same
             // batch.
-            timeseries_id_builder.append_value(compute_timeseries_id(
-                &dp.metric_name,
-                dp.metric_type as u8,
-                &dp.tags,
-            ));
+            let timeseries_id =
+                compute_metric_timeseries_id(&dp.metric_name, dp.metric_type, &dp.tags);
+            timeseries_id_builder.append_value(timeseries_id);
 
             // Only touch builders for tags this data point has.
             for (tag_key, tag_val) in &dp.tags {
@@ -185,6 +192,22 @@ impl ArrowMetricsBatchBuilder {
     }
 }
 
+// Computes timeseries_id for a metric data point. If the user provided a timeseries_id tag, we will
+// remove it.
+fn compute_metric_timeseries_id(
+    metric_name: &str,
+    metric_type: MetricType,
+    tags: &HashMap<String, String>,
+) -> i64 {
+    if !tags.contains_key(COMPUTED_TIMESERIES_ID_FIELD) {
+        return compute_timeseries_id(metric_name, metric_type as u8, tags);
+    }
+
+    let mut filtered_tags = tags.clone();
+    filtered_tags.remove(COMPUTED_TIMESERIES_ID_FIELD);
+    compute_timeseries_id(metric_name, metric_type as u8, &filtered_tags)
+}
+
 /// Error type for Arrow IPC operations.
 #[derive(Debug, thiserror::Error)]
 pub enum ArrowIpcError {
@@ -411,6 +434,35 @@ mod tests {
         assert_eq!(batch.num_columns(), 14);
     }
 
+    #[test]
+    fn test_arrow_batch_filters_user_timeseries_id_tag() {
+        let mut dp = make_test_data_point();
+        dp.tags
+            .insert("timeseries_id".to_string(), "user-provided".to_string());
+
+        let mut builder = ArrowMetricsBatchBuilder::with_capacity(1);
+        builder.append(dp);
+
+        let batch = builder.finish();
+        let schema = batch.schema();
+        assert_eq!(
+            schema
+                .fields()
+                .iter()
+                .filter(|field| field.name().as_str() == "timeseries_id")
+                .count(),
+            1
+        );
+        assert_eq!(schema.index_of("timeseries_id").unwrap(), 4);
+        assert!(schema.index_of("service").is_ok());
+
+        quickwit_parquet_engine::sorted_series::append_sorted_series_column(
+            quickwit_parquet_engine::table_config::ProductType::Metrics.default_sort_fields(),
+            &batch,
+        )
+        .unwrap();
+    }
+
     #[test]
     fn test_arrow_batch_builder_multiple_rows() {
         let mut builder = ArrowMetricsBatchBuilder::with_capacity(100);

quickwit/quickwit-parquet-engine/src/ingest/arrow_sketches.rs

@@ -21,11 +21,21 @@ use std::collections::{BTreeSet, HashMap};
 use std::sync::Arc;
 
 use arrow::array::{
-    ArrayRef, Float64Builder, Int16Builder, ListBuilder, RecordBatch, StringDictionaryBuilder,
-    UInt32Builder, UInt64Builder,
+    ArrayRef, Float64Builder, Int16Builder, Int64Builder, ListBuilder, RecordBatch,
+    StringDictionaryBuilder, UInt32Builder, UInt64Builder,
 };
 use arrow::datatypes::{DataType, Field, Int32Type, Schema as ArrowSchema};
 
+use crate::timeseries_id::compute_timeseries_id;
+
+/// `metric_type` discriminant fed into [`compute_timeseries_id`] for sketch
+/// data points. Mirrors `MetricType::Summary as u8` in
+/// `quickwit-opentelemetry::otlp::otel_metrics` — chosen because DDSketches
+/// are distribution summaries and because it does not collide with the
+/// discriminants used on the metrics path (`Gauge = 0`, `Sum = 1`).
+const SKETCH_METRIC_TYPE_DISCRIMINANT: u8 = 4;
+const COMPUTED_TIMESERIES_ID_FIELD: &str = "timeseries_id";
+
 /// A single DDSketch data point with tags.
 #[derive(Debug, Clone)]
 pub struct SketchDataPoint {
@@ -75,14 +85,21 @@ impl ArrowSketchBatchBuilder {
         let mut tag_keys: BTreeSet<&str> = BTreeSet::new();
         for dp in &self.data_points {
             for key in dp.tags.keys() {
-                tag_keys.insert(key.as_str());
+                // Exclude any user-provided tag named "timeseries_id" to prevent collision with the
+                // computed column. TODO: if user sets "timeseries_id" as a tag, we
+                // are excluding it.
+                if key != COMPUTED_TIMESERIES_ID_FIELD {
+                    tag_keys.insert(key.as_str());
+                }
             }
         }
         let sorted_tag_keys: Vec<String> = tag_keys.into_iter().map(str::to_owned).collect();
 
-        // Build the Arrow schema dynamically
+        // Build the Arrow schema dynamically.
+        // 10 fixed columns: metric_name, timestamp_secs, count, sum, min, max,
+        // flags, keys, counts, timeseries_id.
         let dict_type = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
-        let mut fields = Vec::with_capacity(9 + sorted_tag_keys.len());
+        let mut fields = Vec::with_capacity(10 + sorted_tag_keys.len());
         fields.push(Field::new("metric_name", dict_type.clone(), false));
         fields.push(Field::new("timestamp_secs", DataType::UInt64, false));
         fields.push(Field::new("count", DataType::UInt64, false));
@@ -100,6 +117,11 @@ impl ArrowSketchBatchBuilder {
             DataType::List(Arc::new(Field::new("item", DataType::UInt64, false))),
             false,
         ));
+        fields.push(Field::new(
+            COMPUTED_TIMESERIES_ID_FIELD,
+            DataType::Int64,
+            false,
+        ));
 
         for tag_key in &sorted_tag_keys {
             fields.push(Field::new(tag_key, dict_type.clone(), true));
@@ -126,6 +148,7 @@ impl ArrowSketchBatchBuilder {
             DataType::UInt64,
             false,
         ));
+        let mut timeseries_id_builder = Int64Builder::with_capacity(num_rows);
 
         let mut tag_builders: Vec<StringDictionaryBuilder<Int32Type>> = sorted_tag_keys
             .iter()
@@ -153,6 +176,9 @@ impl ArrowSketchBatchBuilder {
             }
             counts_builder.append(true);
 
+            let timeseries_id = compute_sketch_timeseries_id(&dp.metric_name, &dp.tags);
+            timeseries_id_builder.append_value(timeseries_id);
+
             for (tag_idx, tag_key) in sorted_tag_keys.iter().enumerate() {
                 match dp.tags.get(tag_key) {
                     Some(tag_val) => tag_builders[tag_idx].append_value(tag_val),
@@ -161,7 +187,7 @@ impl ArrowSketchBatchBuilder {
             }
         }
 
-        let mut arrays: Vec<ArrayRef> = Vec::with_capacity(9 + sorted_tag_keys.len());
+        let mut arrays: Vec<ArrayRef> = Vec::with_capacity(10 + sorted_tag_keys.len());
         arrays.push(Arc::new(metric_name_builder.finish()));
         arrays.push(Arc::new(timestamp_secs_builder.finish()));
         arrays.push(Arc::new(count_builder.finish()));
@@ -171,6 +197,7 @@ impl ArrowSketchBatchBuilder {
         arrays.push(Arc::new(flags_builder.finish()));
         arrays.push(Arc::new(keys_builder.finish()));
         arrays.push(Arc::new(counts_builder.finish()));
+        arrays.push(Arc::new(timeseries_id_builder.finish()));
 
         for tag_builder in &mut tag_builders {
             arrays.push(Arc::new(tag_builder.finish()));
@@ -193,6 +220,17 @@ impl ArrowSketchBatchBuilder {
     }
 }
 
+// Computes timeseries_id for a sketch. If the user provided a timeseries_id tag, we will remove it.
+fn compute_sketch_timeseries_id(metric_name: &str, tags: &HashMap<String, String>) -> i64 {
+    if !tags.contains_key(COMPUTED_TIMESERIES_ID_FIELD) {
+        return compute_timeseries_id(metric_name, SKETCH_METRIC_TYPE_DISCRIMINANT, tags);
+    }
+
+    let mut filtered_tags = tags.clone();
+    filtered_tags.remove(COMPUTED_TIMESERIES_ID_FIELD);
+    compute_timeseries_id(metric_name, SKETCH_METRIC_TYPE_DISCRIMINANT, &filtered_tags)
+}
+
 #[cfg(test)]
 mod tests {
     use std::collections::HashMap;
@@ -230,8 +268,8 @@ mod tests {
 
         let batch = builder.finish();
         assert_eq!(batch.num_rows(), 1);
-        // 9 fixed columns + 3 tag columns (env, host, service)
-        assert_eq!(batch.num_columns(), 12);
+        // 10 fixed columns + 3 tag columns (env, host, service)
+        assert_eq!(batch.num_columns(), 13);
 
         // Verify schema field names
         let schema = batch.schema();
@@ -248,13 +286,43 @@ mod tests {
                 "flags",
                 "keys",
                 "counts",
+                "timeseries_id",
                 "env",
                 "host",
                 "service",
             ]
         );
     }
 
+    #[test]
+    fn test_sketch_batch_filters_user_timeseries_id_tag() {
+        let mut dp = make_test_sketch_point();
+        dp.tags
+            .insert("timeseries_id".to_string(), "user-provided".to_string());
+
+        let mut builder = ArrowSketchBatchBuilder::with_capacity(1);
+        builder.append(dp);
+
+        let batch = builder.finish();
+        let schema = batch.schema();
+        assert_eq!(
+            schema
+                .fields()
+                .iter()
+                .filter(|field| field.name().as_str() == "timeseries_id")
+                .count(),
+            1
+        );
+        assert_eq!(schema.index_of("timeseries_id").unwrap(), 9);
+        assert!(schema.index_of("service").is_ok());
+
+        crate::sorted_series::append_sorted_series_column(
+            crate::table_config::ProductType::Sketches.default_sort_fields(),
+            &batch,
+        )
+        .unwrap();
+    }
+
     #[test]
     fn test_sketch_batch_builder_multiple_rows() {
         let mut builder = ArrowSketchBatchBuilder::with_capacity(100);
@@ -281,8 +349,8 @@ mod tests {
 
         let batch = builder.finish();
         assert_eq!(batch.num_rows(), 100);
-        // 9 fixed + 2 tags
-        assert_eq!(batch.num_columns(), 11);
+        // 10 fixed + 2 tags
+        assert_eq!(batch.num_columns(), 12);
     }
 
     #[test]
@@ -325,14 +393,14 @@ mod tests {
 
         let batch = builder.finish();
         assert_eq!(batch.num_rows(), 2);
-        // 9 fixed + 3 tags (env, host, region)
-        assert_eq!(batch.num_columns(), 12);
+        // 10 fixed + 3 tags (env, host, region)
+        assert_eq!(batch.num_columns(), 13);
 
         let schema = batch.schema();
         let tag_names: Vec<&str> = schema
             .fields()
             .iter()
-            .skip(9)
+            .skip(10)
             .map(|f| f.name().as_str())
             .collect();
         assert_eq!(tag_names, vec!["env", "host", "region"]);
@@ -345,8 +413,8 @@ mod tests {
 
         let batch = builder.finish();
         assert_eq!(batch.num_rows(), 0);
-        // 9 fixed columns, no tags
-        assert_eq!(batch.num_columns(), 9);
+        // 10 fixed columns, no tags
+        assert_eq!(batch.num_columns(), 10);
     }
 
     #[test]