Type capture label keys and values as Dictionary(Int32,Utf8)

CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## Unreleased
 
 ## Changed
+- **Breaking Change**: Capture parquet files now type the `labels` map `key` and
+  `value` as `Dictionary(Int32, Utf8)` instead of `Utf8`. Labels dominate capture
+  memory, and dictionary typing lets readers materialize each distinct string
+  once. Consumers must read the label key/value columns as dictionary arrays; the
+  old and new schemas are not interchangeable.
 - **Breaking Change**: Datadog blackhole now records received series under the
   `target/` prefix, matching the prometheus and expvar target metrics
   collectors. The `record` policy still matches on the unprefixed series name.

lading/Cargo.toml

@@ -107,6 +107,7 @@ async-pidfd = { version = "0.1" }
 tempfile = { workspace = true }
 proptest = { workspace = true }
 tower-test = { version = "0.4" }
+uuid = { workspace = true }
 
 [build-dependencies]
 prost-build = { workspace = true }

lading/src/bin/captool/analyze/parquet.rs

@@ -7,7 +7,8 @@ use std::collections::BTreeSet;
 use std::fs::File;
 use std::path::Path;
 
-use arrow_array::{Array, Float64Array, MapArray, StringArray, UInt64Array};
+use arrow_array::{Array, ArrayAccessor, Float64Array, MapArray, StringArray, UInt64Array};
+use lading_capture::formats::parquet::resolve_label_dictionary;
 use lading_capture_schema::columns;
 use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
 use rustc_hash::FxHashMap;
@@ -170,16 +171,10 @@ pub(crate) fn analyze_metric<P: AsRef<Path>>(
             // Extract labels
             let mut sorted_labels = BTreeSet::new();
             let labels_slice = labels_array.value(row);
-            let key_array = labels_slice
-                .column(0)
-                .as_any()
-                .downcast_ref::<StringArray>()
-                .ok_or_else(|| Error::InvalidColumnType("Label keys not String".to_string()))?;
-            let value_array = labels_slice
-                .column(1)
-                .as_any()
-                .downcast_ref::<StringArray>()
-                .ok_or_else(|| Error::InvalidColumnType("Label values not String".to_string()))?;
+            let key_array = resolve_label_dictionary(labels_slice.column(0), "Labels keys")
+                .map_err(Error::InvalidColumnType)?;
+            let value_array = resolve_label_dictionary(labels_slice.column(1), "Labels values")
+                .map_err(Error::InvalidColumnType)?;
 
             for i in 0..key_array.len() {
                 let key = key_array.value(i);
@@ -231,3 +226,100 @@ pub(crate) fn analyze_metric<P: AsRef<Path>>(
 
     Ok(results)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use lading_capture::formats::parquet::Format;
+    use lading_capture::line::{Line, LineValue, MetricKind};
+    use rustc_hash::FxHashMap;
+    use tempfile::NamedTempFile;
+    use uuid::Uuid;
+
+    fn labels(pairs: &[(&str, &str)]) -> FxHashMap<String, String> {
+        pairs
+            .iter()
+            .map(|(k, v)| ((*k).to_string(), (*v).to_string()))
+            .collect()
+    }
+
+    /// Write the given lines to a parquet capture file using the production
+    /// [`Format`] writer and return the backing temp file.
+    fn write_capture(lines: &[Line]) -> NamedTempFile {
+        let tmp = NamedTempFile::new().expect("create temp file");
+        let file = tmp.reopen().expect("reopen temp file");
+        let mut writer = Format::new(file, 3).expect("create parquet writer");
+        for line in lines {
+            writer.write_metric(line).expect("write metric");
+        }
+        writer.flush().expect("flush");
+        writer.close().expect("close");
+        tmp
+    }
+
+    /// `analyze_metric` reads captures written by the production [`Format`] writer.
+    ///
+    /// Binds `captool analyze` to the writer's on-disk schema: a label-encoding
+    /// change the analyzer fails to track breaks this test rather than shipping
+    /// silently broken.
+    #[test]
+    #[expect(
+        clippy::float_cmp,
+        reason = "stats are computed from exact integer-valued inputs"
+    )]
+    fn analyze_metric_reads_writer_output() {
+        let run_id = Uuid::from_u128(0x1234_5678_9abc_def0_1234_5678_9abc_def0);
+        let prod = labels(&[("env", "prod"), ("service", "api")]);
+        let staging = labels(&[("env", "staging")]);
+
+        let mut lines: Vec<Line> = (0u64..3)
+            .map(|i| Line {
+                run_id,
+                time: 1_000 + u128::from(i),
+                fetch_index: i,
+                metric_name: "requests".to_string(),
+                metric_kind: MetricKind::Counter,
+                value: LineValue::Int((i + 1) * 10),
+                labels: prod.clone(),
+                value_histogram: Vec::new(),
+            })
+            .collect();
+        lines.extend((0u64..2).map(|i| Line {
+            run_id,
+            time: 2_000 + u128::from(i),
+            fetch_index: i,
+            metric_name: "requests".to_string(),
+            metric_kind: MetricKind::Counter,
+            value: LineValue::Int((i + 1) * 100),
+            labels: staging.clone(),
+            value_histogram: Vec::new(),
+        }));
+
+        let tmp = write_capture(&lines);
+
+        let metrics = list_metrics(tmp.path()).expect("list metrics");
+        assert!(
+            metrics.iter().any(|m| m.name == "requests"),
+            "requests metric should be listed, got {metrics:?}"
+        );
+
+        let series = analyze_metric(tmp.path(), "requests").expect("analyze requests");
+        assert_eq!(series.len(), 2, "two distinct label sets expected");
+
+        let prod_key: BTreeSet<String> = ["env:prod".to_string(), "service:api".to_string()].into();
+        let staging_key: BTreeSet<String> = ["env:staging".to_string()].into();
+
+        let prod_stats = series.get(&prod_key).expect("prod series present");
+        assert_eq!(prod_stats.min, 10.0);
+        assert_eq!(prod_stats.max, 30.0);
+        assert_eq!(prod_stats.mean, 20.0);
+        assert!(prod_stats.is_monotonic);
+
+        let staging_stats = series.get(&staging_key).expect("staging series present");
+        assert_eq!(staging_stats.min, 100.0);
+        assert_eq!(staging_stats.max, 200.0);
+        assert_eq!(staging_stats.mean, 150.0);
+        assert!(staging_stats.is_monotonic);
+    }
+}

lading_capture/src/formats.rs

@@ -68,10 +68,11 @@ pub trait OutputFormat {
 #[cfg(test)]
 mod tests {
     use super::jsonl;
+    use crate::formats::parquet::resolve_label_dictionary;
     use crate::line::{Line, LineValue, MetricKind};
     use approx::relative_eq;
     use arrow_array::{
-        Array, BinaryArray, Float64Array, MapArray, StringArray, StructArray,
+        Array, ArrayAccessor, BinaryArray, Float64Array, MapArray, StringArray, StructArray,
         TimestampMillisecondArray, UInt64Array,
     };
     use bytes::Bytes;
@@ -399,6 +400,171 @@ mod tests {
         }
     }
 
+    /// Roundtrip many rows sharing identical label strings.
+    ///
+    /// The property test's random labels only repeat by chance; this pins the
+    /// production pattern where a few distinct strings (`env=prod`) repeat across
+    /// every row and are dictionary-deduplicated.
+    #[test]
+    fn parquet_round_trip_repeated_labels() {
+        let run_id = Uuid::from_u128(0x1234_5678_9abc_def0_1234_5678_9abc_def0);
+        let shared_labels: FxHashMap<String, String> = [
+            ("env".to_string(), "prod".to_string()),
+            ("service".to_string(), "api".to_string()),
+        ]
+        .into_iter()
+        .collect();
+
+        // 500 rows sharing two labels, plus one distinct row so the dictionary
+        // holds more than one entry.
+        let mut input_lines: Vec<Line> = (0u64..500)
+            .map(|i| Line {
+                run_id,
+                time: 1_000 + u128::from(i),
+                fetch_index: i,
+                metric_name: "requests".to_string(),
+                metric_kind: MetricKind::Counter,
+                value: LineValue::Int(i),
+                labels: shared_labels.clone(),
+                value_histogram: Vec::new(),
+            })
+            .collect();
+        input_lines.push(Line {
+            run_id,
+            time: 2_000,
+            fetch_index: 500,
+            metric_name: "requests".to_string(),
+            metric_kind: MetricKind::Counter,
+            value: LineValue::Int(500),
+            labels: [("env".to_string(), "staging".to_string())]
+                .into_iter()
+                .collect(),
+            value_histogram: Vec::new(),
+        });
+
+        let mut buffer = Cursor::new(Vec::new());
+        {
+            let mut writer = super::parquet::Format::new(&mut buffer, 3).expect("create writer");
+            for line in &input_lines {
+                writer.write_metric(line).expect("write");
+            }
+            writer.flush().expect("flush");
+            writer.close().expect("close");
+        }
+        let bytes = buffer.into_inner();
+
+        let deserialized_lines = read_parquet_lines(&bytes).expect("read parquet");
+
+        assert_eq!(input_lines.len(), deserialized_lines.len());
+        for (input, output) in input_lines.iter().zip(deserialized_lines.iter()) {
+            assert_eq!(input.labels, output.labels);
+            assert_eq!(input.metric_name, output.metric_name);
+            assert_eq!(input.fetch_index, output.fetch_index);
+        }
+    }
+
+    /// Decode `Dictionary(Int32, Utf8)` labels when a file is read as several
+    /// `RecordBatch`es.
+    ///
+    /// The reader emits one `RecordBatch` per `batch_size` rows, so production
+    /// files are read as several batches — a path a single-batch roundtrip misses.
+    /// A small `batch_size` forces that here; each batch's labels are decoded via
+    /// the production `resolve_label_dictionary` and checked, and the test asserts
+    /// more than one batch so the coverage cannot silently lapse.
+    ///
+    /// It does not exercise divergent per-batch dictionaries: the writer emits one
+    /// row group with a single dictionary page, shared by every batch.
+    #[test]
+    fn parquet_round_trip_multiple_read_batches() {
+        let run_id = Uuid::from_u128(0x0fed_cba9_8765_4321_0fed_cba9_8765_4321);
+        let envs = ["prod", "staging", "prod", "staging", "prod", "staging"];
+
+        let mut buffer = Cursor::new(Vec::new());
+        let mut expected_labels: FxHashMap<u64, FxHashMap<String, String>> = FxHashMap::default();
+        {
+            let mut writer = super::parquet::Format::new(&mut buffer, 3).expect("create writer");
+            for (i, env) in envs.iter().enumerate() {
+                let fetch_index = i as u64;
+                let labels: FxHashMap<String, String> = [("env".to_string(), (*env).to_string())]
+                    .into_iter()
+                    .collect();
+                let line = Line {
+                    run_id,
+                    time: 1_000 + u128::from(fetch_index),
+                    fetch_index,
+                    metric_name: "requests".to_string(),
+                    metric_kind: MetricKind::Counter,
+                    value: LineValue::Int(fetch_index),
+                    labels: labels.clone(),
+                    value_histogram: Vec::new(),
+                };
+                writer.write_metric(&line).expect("write");
+                expected_labels.insert(fetch_index, labels);
+            }
+            writer.flush().expect("flush");
+            writer.close().expect("close");
+        }
+        let bytes = buffer.into_inner();
+
+        // A batch size below the row count forces the reader to emit several
+        // batches.
+        let bytes_buf = Bytes::copy_from_slice(&bytes);
+        let reader = ParquetRecordBatchReaderBuilder::try_new(bytes_buf)
+            .expect("reader builder")
+            .with_batch_size(2)
+            .build()
+            .expect("reader");
+
+        let mut batch_count = 0usize;
+        let mut seen = 0usize;
+        for batch_result in reader {
+            let batch = batch_result.expect("batch");
+            if batch.num_rows() == 0 {
+                continue;
+            }
+            batch_count += 1;
+
+            let fetch_index_array = batch
+                .column_by_name("fetch_index")
+                .expect("fetch_index column")
+                .as_any()
+                .downcast_ref::<UInt64Array>()
+                .expect("fetch_index is UInt64Array");
+            let labels_array = batch
+                .column_by_name("labels")
+                .expect("labels column")
+                .as_any()
+                .downcast_ref::<MapArray>()
+                .expect("labels is MapArray");
+
+            for row in 0..batch.num_rows() {
+                let fetch_index = fetch_index_array.value(row);
+                let labels_slice: StructArray = labels_array.value(row);
+                let keys = resolve_label_dictionary(labels_slice.column(0), "Labels keys")
+                    .expect("label keys are Dictionary(Int32,Utf8)");
+                let values = resolve_label_dictionary(labels_slice.column(1), "Labels values")
+                    .expect("label values are Dictionary(Int32,Utf8)");
+                let decoded: FxHashMap<String, String> = (0..keys.len())
+                    .map(|i| (keys.value(i).to_string(), values.value(i).to_string()))
+                    .collect();
+                assert_eq!(
+                    &decoded,
+                    expected_labels
+                        .get(&fetch_index)
+                        .expect("known fetch_index"),
+                    "labels mismatch at fetch_index {fetch_index}"
+                );
+                seen += 1;
+            }
+        }
+
+        assert!(
+            batch_count > 1,
+            "expected multiple read batches, got {batch_count}"
+        );
+        assert_eq!(seen, envs.len(), "every row decoded exactly once");
+    }
+
     #[expect(clippy::too_many_lines)]
     fn read_parquet_lines(bytes: &[u8]) -> Result<Vec<Line>, Box<dyn std::error::Error>> {
         let bytes_buf = Bytes::copy_from_slice(bytes);
@@ -502,16 +668,10 @@ mod tests {
                 };
 
                 let labels_slice: StructArray = labels_array.value(row_idx);
-                let keys = labels_slice
-                    .column(0)
-                    .as_any()
-                    .downcast_ref::<StringArray>()
-                    .expect("label keys are strings");
-                let values = labels_slice
-                    .column(1)
-                    .as_any()
-                    .downcast_ref::<StringArray>()
-                    .expect("label values are strings");
+                let keys = resolve_label_dictionary(labels_slice.column(0), "Labels keys")
+                    .expect("label keys are Dictionary(Int32,Utf8)");
+                let values = resolve_label_dictionary(labels_slice.column(1), "Labels values")
+                    .expect("label values are Dictionary(Int32,Utf8)");
 
                 let mut labels = FxHashMap::default();
                 for i in 0..keys.len() {