Use collision-free interned group keys in DAG execution

data_plane/benches/sketch_db.rs

@@ -456,12 +456,52 @@ fn bench_reconcile_per_batch(c: &mut Criterion) {
     g.finish();
 }
 
+fn bench_group_key_projection(c: &mut Criterion) {
+    use data_plane::precompute_engine::group_key::intern_pairs;
+
+    let cardinality = 100_000usize;
+    let labels = (0..cardinality)
+        .map(|index| (format!("region;{index}"), format!("service={index}")))
+        .collect::<Vec<_>>();
+    let mut group = c.benchmark_group("group_key_projection");
+    group.throughput(Throughput::Elements(cardinality as u64));
+    group.bench_function("cold_high_cardinality", |b| {
+        b.iter(|| {
+            for (region, service) in &labels {
+                black_box(intern_pairs([
+                    ("region", region.as_str()),
+                    ("service", service.as_str()),
+                ]));
+            }
+        });
+    });
+    // Warm the bounded interner, then measure shared-DAG reuse.
+    for (region, service) in &labels {
+        black_box(intern_pairs([
+            ("region", region.as_str()),
+            ("service", service.as_str()),
+        ]));
+    }
+    group.bench_function("warm_shared_projection", |b| {
+        b.iter(|| {
+            for (region, service) in &labels {
+                black_box(intern_pairs([
+                    ("region", region.as_str()),
+                    ("service", service.as_str()),
+                ]));
+            }
+        });
+    });
+    group.finish();
+}
+
 criterion_group!(
     benches,
     bench_append_sample,
     bench_append_precompute,
     bench_query_range,
     bench_query_precomputes_by_agg,
     bench_reconcile_per_batch,
+    bench_group_key_projection,
 );
 criterion_main!(benches);

data_plane/src/drivers/ingest/otel.rs

@@ -660,7 +660,11 @@ async fn route_otlp_to_precompute(
     // alongside the sid in the WorkerMessage so the worker can resolve
     // the source config and render emit-time labels without consulting
     // the sid → attrs reverse mapping.
-    type BucketTuple = (u64, asap_types::PolicyFingerprint, String); // (sid, policy_fp, group_key)
+    type BucketTuple = (
+        u64,
+        asap_types::PolicyFingerprint,
+        Arc<crate::precompute_engine::group_key::GroupKey>,
+    ); // (sid, policy_fp, group_key)
     type SampleTuple = (String, i64, f64);
     let mut by_bucket: HashMap<u64, (BucketTuple, Vec<SampleTuple>)> = HashMap::new();
     let mut raw_matched = 0usize;
@@ -4601,7 +4605,7 @@ mod sid_bucketing_tests {
         let groups: Vec<(
             u64,
             asap_types::PolicyFingerprint,
-            String,
+            Arc<crate::precompute_engine::group_key::GroupKey>,
             Vec<(String, i64, f64)>,
         )> = messages
             .into_iter()

data_plane/src/drivers/ingest/prometheus_remote_write.rs

@@ -524,7 +524,11 @@ fn route_messages(
     ingest: &Arc<IngestState>,
     physical_plan: &crate::storage_engines::types::ActivePhysicalPlan,
 ) -> Vec<WorkerMessage> {
-    type Bucket = (u64, asap_types::PolicyFingerprint, String);
+    type Bucket = (
+        u64,
+        asap_types::PolicyFingerprint,
+        Arc<crate::precompute_engine::group_key::GroupKey>,
+    );
     type RoutedSample = (String, i64, f64);
     let snapshot = physical_plan.runtime_config.clone();
     let _ = crate::storage_engines::sketch_db::lifecycle::reconcile_if_config_changed(
@@ -579,7 +583,14 @@ fn route_messages(
                     .collect()
             };
             let group_key = if series_scoped {
-                sample.population_key.to_string()
+                let mut names = sample.labels.keys().map(String::as_str).collect::<Vec<_>>();
+                names.sort_unstable();
+                crate::precompute_engine::group_key::intern_pairs(names.into_iter().map(|name| {
+                    (
+                        name,
+                        sample.labels.get(name).map(String::as_str).unwrap_or(""),
+                    )
+                }))
             } else {
                 group_key
             };
@@ -1229,7 +1240,7 @@ mod tests {
         else {
             panic!("expected GroupSamples");
         };
-        assert_eq!(group_key, "api");
+        assert_eq!(group_key.values().labels, vec!["api"]);
         assert_eq!(
             samples,
             vec![("requests_total{job=\"api\"}".into(), 100, 4.0)]

data_plane/src/precompute_engine/group_key.rs

@@ -0,0 +1,142 @@
+use std::collections::HashMap;
+use std::sync::{Arc, OnceLock};
+
+use moka::sync::Cache;
+
+use crate::storage_engines::types::KeyByLabelValues;
+
+/// Collision-free, positional identity for a physical GROUP BY partition.
+/// Label names are retained so different DAG projections cannot alias merely
+/// because their values happen to be equal.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct GroupKey {
+    labels: Arc<[(Arc<str>, Arc<str>)]>,
+    canonical: Arc<[u8]>,
+}
+
+impl GroupKey {
+    pub fn project<'a>(
+        names: impl IntoIterator<Item = &'a str>,
+        labels: &HashMap<String, String>,
+    ) -> Self {
+        Self::new(
+            names
+                .into_iter()
+                .map(|name| (name, labels.get(name).map(String::as_str).unwrap_or(""))),
+        )
+    }
+
+    pub fn new<'a>(pairs: impl IntoIterator<Item = (&'a str, &'a str)>) -> Self {
+        let labels: Arc<[(Arc<str>, Arc<str>)]> = pairs
+            .into_iter()
+            .map(|(name, value)| (Arc::from(name), Arc::from(value)))
+            .collect::<Vec<_>>()
+            .into();
+        let mut canonical = Vec::new();
+        canonical.extend_from_slice(b"ASAPGK\x01");
+        canonical.extend_from_slice(&(labels.len() as u32).to_be_bytes());
+        for (name, value) in labels.iter() {
+            put_component(&mut canonical, name.as_bytes());
+            put_component(&mut canonical, value.as_bytes());
+        }
+        Self {
+            labels,
+            canonical: canonical.into(),
+        }
+    }
+
+    pub fn canonical_bytes(&self) -> &[u8] {
+        &self.canonical
+    }
+
+    pub fn values(&self) -> KeyByLabelValues {
+        KeyByLabelValues::new_with_labels(
+            self.labels
+                .iter()
+                .map(|(_, value)| value.to_string())
+                .collect(),
+        )
+    }
+
+    pub fn as_population_labels(&self) -> std::collections::BTreeMap<String, String> {
+        self.labels
+            .iter()
+            .map(|(name, value)| (name.to_string(), value.to_string()))
+            .collect()
+    }
+}
+
+impl std::fmt::Display for GroupKey {
+    fn fmt(&self, output: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        output.write_str("{")?;
+        for (index, (name, value)) in self.labels.iter().enumerate() {
+            if index != 0 {
+                output.write_str(",")?;
+            }
+            write!(output, "{name}={value:?}")?;
+        }
+        output.write_str("}")
+    }
+}
+
+fn put_component(target: &mut Vec<u8>, value: &[u8]) {
+    target.extend_from_slice(&(value.len() as u32).to_be_bytes());
+    target.extend_from_slice(value);
+}
+
+/// Bounded interner shared by ingress adapters. Repeated DAG consumers with
+/// the same projection carry one immutable key allocation into workers.
+pub fn intern(key: GroupKey) -> Arc<GroupKey> {
+    intern_pairs(
+        key.labels
+            .iter()
+            .map(|(name, value)| (name.as_ref(), value.as_ref())),
+    )
+}
+
+/// Project and intern from borrowed labels. The hot path builds only the
+/// compact lookup bytes on a cache hit; label/name strings are allocated once
+/// when a new high-cardinality group first appears.
+pub fn intern_pairs<'a>(pairs: impl IntoIterator<Item = (&'a str, &'a str)>) -> Arc<GroupKey> {
+    let pairs = pairs.into_iter().collect::<Vec<_>>();
+    let mut encoded = Vec::with_capacity(
+        11 + pairs
+            .iter()
+            .map(|(name, value)| name.len() + value.len() + 8)
+            .sum::<usize>(),
+    );
+    encoded.extend_from_slice(b"ASAPGK\x01");
+    encoded.extend_from_slice(&(pairs.len() as u32).to_be_bytes());
+    for (name, value) in &pairs {
+        put_component(&mut encoded, name.as_bytes());
+        put_component(&mut encoded, value.as_bytes());
+    }
+    static INTERNER: OnceLock<Cache<Vec<u8>, Arc<GroupKey>>> = OnceLock::new();
+    let interner = INTERNER.get_or_init(|| Cache::builder().max_capacity(131_072).build());
+    interner.get_with(encoded, || Arc::new(GroupKey::new(pairs)))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn delimiters_names_order_and_missing_values_cannot_alias() {
+        let a = GroupKey::new([("x", "a;b"), ("y", "c")]);
+        let b = GroupKey::new([("x", "a"), ("y", "b;c")]);
+        let reordered = GroupKey::new([("y", "c"), ("x", "a;b")]);
+        let other_names = GroupKey::new([("p", "a;b"), ("q", "c")]);
+        let missing = GroupKey::new([("x", ""), ("y", "c")]);
+        for other in [&b, &reordered, &other_names, &missing] {
+            assert_ne!(a.canonical_bytes(), other.canonical_bytes());
+        }
+        assert_eq!(a.values().labels, vec!["a;b", "c"]);
+    }
+
+    #[test]
+    fn interner_reuses_equal_projection() {
+        let left = intern(GroupKey::new([("region", "us-east"), ("job", "api")]));
+        let right = intern(GroupKey::new([("region", "us-east"), ("job", "api")]));
+        assert!(Arc::ptr_eq(&left, &right));
+    }
+}

data_plane/src/precompute_engine/ingest_handler.rs

@@ -260,7 +260,10 @@ impl IngestState {
     /// Extract the group key for a series key against a given aggregation
     /// config. Re-exports the module-private helper so that out-of-module
     /// ingest sources (e.g. OTLP) can reuse it.
-    pub fn extract_group_key_for(series_key: &str, config: &AggregationConfig) -> String {
+    pub fn extract_group_key_for(
+        series_key: &str,
+        config: &AggregationConfig,
+    ) -> Arc<crate::precompute_engine::group_key::GroupKey> {
         extract_group_key(series_key, config)
     }
 
@@ -274,33 +277,33 @@ impl IngestState {
     pub fn extract_group_key_from_labels(
         labels: &std::collections::HashMap<String, String>,
         config: &AggregationConfig,
-    ) -> String {
-        let mut values = Vec::with_capacity(config.grouping_labels.labels.len());
-        for label_name in &config.grouping_labels.labels {
-            values.push(
-                labels
-                    .get(label_name.as_str())
-                    .map(|s| s.as_str())
-                    .unwrap_or(""),
-            );
-        }
-        values.join(";")
+    ) -> Arc<crate::precompute_engine::group_key::GroupKey> {
+        crate::precompute_engine::group_key::intern_pairs(config.grouping_labels.labels.iter().map(
+            |name| {
+                (
+                    name.as_str(),
+                    labels.get(name).map(String::as_str).unwrap_or(""),
+                )
+            },
+        ))
     }
 }
 
 /// Extract the group key (grouping label values joined by semicolons)
 /// for a given series key and aggregation config.
-fn extract_group_key(series_key: &str, config: &AggregationConfig) -> String {
+fn extract_group_key(
+    series_key: &str,
+    config: &AggregationConfig,
+) -> Arc<crate::precompute_engine::group_key::GroupKey> {
     let labels = parse_labels_from_series_key(series_key);
-    let mut values = Vec::new();
-    for label_name in &config.grouping_labels.labels {
-        if let Some(val) = labels.get(label_name.as_str()) {
-            values.push(*val);
-        } else {
-            values.push("");
-        }
-    }
-    values.join(";")
+    crate::precompute_engine::group_key::intern_pairs(config.grouping_labels.labels.iter().map(
+        |name| {
+            (
+                name.as_str(),
+                labels.get(name.as_str()).copied().unwrap_or(""),
+            )
+        },
+    ))
 }
 
 #[cfg(test)]

data_plane/src/precompute_engine/mod.rs

@@ -2,6 +2,7 @@ pub mod accumulator_factory;
 pub mod config;
 mod engine;
 pub mod frame_lineage;
+pub mod group_key;
 pub mod ingest_handler;
 pub(crate) mod metrics;
 pub mod operators;

data_plane/src/precompute_engine/series_router.rs

@@ -1,8 +1,10 @@
+use crate::precompute_engine::group_key::GroupKey;
 use crate::storage_engines::types::AggregateCore;
 use asap_types::PolicyFingerprint;
 use futures::future::try_join_all;
 use std::collections::HashMap;
 use std::fmt;
+use std::sync::Arc;
 use std::time::Instant;
 use tokio::sync::mpsc;
 use xxhash_rust::xxh64::xxh64;
@@ -48,7 +50,7 @@ pub enum WorkerMessage {
         /// Grouping label values joined by semicolons (e.g. "constant").
         /// Empty string if the aggregation has no grouping labels. Used
         /// at emit time to render the output's `KeyByLabelValues`.
-        group_key: String,
+        group_key: Arc<GroupKey>,
         /// Each entry: (series_key, timestamp_ms, value).
         /// series_key is needed for keyed (MultipleSubpopulation) accumulators
         /// to extract the aggregated-label key.
@@ -75,7 +77,7 @@ pub enum WorkerMessage {
         /// Grouping label values joined by semicolons, matching the
         /// format produced by `IngestState::extract_group_key_for`.
         /// Used at emit time to render the output's `KeyByLabelValues`.
-        group_key: String,
+        group_key: Arc<GroupKey>,
         /// Wall-clock timestamp the sketch refers to (millis since epoch).
         /// Used to place the sketch into the correct pane.
         timestamp_ms: i64,