diff --git a/CHANGELOG.md b/CHANGELOG.md index 7086e819d..e015ea372 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/Cargo.lock b/Cargo.lock index b48c2df21..de80edb3f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1975,6 +1975,7 @@ dependencies = [ "tower-test", "tracing", "tracing-subscriber", + "uuid", "zstd", ] diff --git a/lading/Cargo.toml b/lading/Cargo.toml index a2ecd3e6c..b76ca90f4 100644 --- a/lading/Cargo.toml +++ b/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 } diff --git a/lading/src/bin/captool/analyze/parquet.rs b/lading/src/bin/captool/analyze/parquet.rs index 371f3c247..30c8246a6 100644 --- a/lading/src/bin/captool/analyze/parquet.rs +++ b/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>( // 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::() - .ok_or_else(|| Error::InvalidColumnType("Label keys not String".to_string()))?; - let value_array = labels_slice - .column(1) - .as_any() - .downcast_ref::() - .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>( 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 { + 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 = (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 = ["env:prod".to_string(), "service:api".to_string()].into(); + let staging_key: BTreeSet = ["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); + } +} diff --git a/lading_capture/src/formats.rs b/lading_capture/src/formats.rs index 60da6b7a0..a9b8f5913 100644 --- a/lading_capture/src/formats.rs +++ b/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 = [ + ("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 = (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> = 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 = [("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::() + .expect("fetch_index is UInt64Array"); + let labels_array = batch + .column_by_name("labels") + .expect("labels column") + .as_any() + .downcast_ref::() + .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 = (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, Box> { 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::() - .expect("label keys are strings"); - let values = labels_slice - .column(1) - .as_any() - .downcast_ref::() - .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() { diff --git a/lading_capture/src/formats/parquet.rs b/lading_capture/src/formats/parquet.rs index d26b4f739..df2b10526 100644 --- a/lading_capture/src/formats/parquet.rs +++ b/lading_capture/src/formats/parquet.rs @@ -10,12 +10,13 @@ use std::{ }; use arrow_array::{ - ArrayRef, BinaryArray, Float64Array, MapArray, RecordBatch, StringArray, StructArray, - TimestampMillisecondArray, UInt64Array, + Array, ArrayRef, BinaryArray, DictionaryArray, Float64Array, MapArray, RecordBatch, + StringArray, StructArray, TimestampMillisecondArray, TypedDictionaryArray, UInt64Array, + builder::StringDictionaryBuilder, types::Int32Type, }; use arrow_buffer::OffsetBuffer; use arrow_schema::{ArrowError, DataType, Field, Fields, Schema}; -use lading_capture_schema::{capture_schema, columns}; +use lading_capture_schema::{capture_schema, columns, label_dictionary_type}; use parquet::{ arrow::ArrowWriter, basic::{Compression, ZstdLevel}, @@ -25,6 +26,29 @@ use parquet::{ use crate::line; +/// Resolve a label map column typed as [`label_dictionary_type`] to a typed +/// dictionary view of its `Utf8` values. +/// +/// The single decode path for both the `key` and `value` columns, shared across +/// crates (e.g. `captool`) so a schema change surfaces everywhere at once. +/// `column_name` prefixes the error so callers can tell the two columns apart. +/// +/// # Errors +/// +/// Errors if `column` is not a `Dictionary(Int32, _)` or its values are not +/// `Utf8`. +pub fn resolve_label_dictionary<'a>( + column: &'a dyn Array, + column_name: &str, +) -> Result, String> { + column + .as_any() + .downcast_ref::>() + .ok_or_else(|| format!("{column_name} are not Dictionary(Int32,Utf8)"))? + .downcast_dict::() + .ok_or_else(|| format!("{column_name} dictionary values are not Utf8")) +} + /// Parquet format errors #[derive(thiserror::Error, Debug)] pub enum Error { @@ -210,16 +234,33 @@ impl Format { label_offsets.push(0i32); label_offsets.extend_from_slice(&self.buffers.label_offsets); - // Build the labels map array using pre-allocated buffers - let keys_array = Arc::new(StringArray::from(self.buffers.label_keys.clone())); - let values_array = Arc::new(StringArray::from(self.buffers.label_values.clone())); + // Keys and values are dictionary-encoded (Int32,Utf8) so readers + // materialize each distinct string once; labels dominate capture memory. + let mut keys_builder = StringDictionaryBuilder::::new(); + for key in &self.buffers.label_keys { + keys_builder.append_value(key); + } + let keys_array = Arc::new(keys_builder.finish()); + let mut values_builder = StringDictionaryBuilder::::new(); + for value in &self.buffers.label_values { + values_builder.append_value(value); + } + let values_array = Arc::new(values_builder.finish()); let struct_array = StructArray::from(vec![ ( - Arc::new(Field::new(columns::LABEL_KEY, DataType::Utf8, false)), + Arc::new(Field::new( + columns::LABEL_KEY, + label_dictionary_type(), + false, + )), keys_array as ArrayRef, ), ( - Arc::new(Field::new(columns::LABEL_VALUE, DataType::Utf8, false)), + Arc::new(Field::new( + columns::LABEL_VALUE, + label_dictionary_type(), + false, + )), values_array as ArrayRef, ), ]); @@ -227,8 +268,8 @@ impl Format { let field = Arc::new(Field::new( columns::LABEL_ENTRIES, DataType::Struct(Fields::from(vec![ - Field::new(columns::LABEL_KEY, DataType::Utf8, false), - Field::new(columns::LABEL_VALUE, DataType::Utf8, false), + Field::new(columns::LABEL_KEY, label_dictionary_type(), false), + Field::new(columns::LABEL_VALUE, label_dictionary_type(), false), ])), false, )); diff --git a/lading_capture/src/validate/parquet.rs b/lading_capture/src/validate/parquet.rs index 23331a3b6..53da3db89 100644 --- a/lading_capture/src/validate/parquet.rs +++ b/lading_capture/src/validate/parquet.rs @@ -13,12 +13,13 @@ use std::hash::{BuildHasher, Hasher}; use std::path::Path; use arrow_array::{ - Array, MapArray, StringArray, StructArray, TimestampMillisecondArray, UInt64Array, + Array, ArrayAccessor, MapArray, StringArray, StructArray, TimestampMillisecondArray, + UInt64Array, }; use lading_capture_schema::columns; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; -use crate::validate::ValidationResult; +use crate::{formats::parquet::resolve_label_dictionary, validate::ValidationResult}; /// Errors for parquet validation #[derive(thiserror::Error, Debug)] @@ -181,20 +182,10 @@ pub fn validate_parquet>( } let labels_slice: StructArray = labels_array.value(row); - let key_array = labels_slice - .column(0) - .as_any() - .downcast_ref::() - .ok_or_else(|| { - Error::InvalidColumnType("Labels keys are not StringArray".to_string()) - })?; - let value_array = labels_slice - .column(1) - .as_any() - .downcast_ref::() - .ok_or_else(|| { - Error::InvalidColumnType("Labels values are not StringArray".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)?; let mut sorted_labels: BTreeSet = BTreeSet::new(); for i in 0..key_array.len() { diff --git a/lading_capture_schema/src/lib.rs b/lading_capture_schema/src/lib.rs index 0562a9e69..76979bac0 100644 --- a/lading_capture_schema/src/lib.rs +++ b/lading_capture_schema/src/lib.rs @@ -68,9 +68,20 @@ fn labels_entry_field() -> Arc { Arc::new(Field::new( columns::LABEL_ENTRIES, DataType::Struct(Fields::from(vec![ - Field::new(columns::LABEL_KEY, DataType::Utf8, false), - Field::new(columns::LABEL_VALUE, DataType::Utf8, false), + Field::new(columns::LABEL_KEY, label_dictionary_type(), false), + Field::new(columns::LABEL_VALUE, label_dictionary_type(), false), ])), false, )) } + +/// Arrow logical type for label keys and values. +/// +/// Label strings repeat heavily across rows and dominate capture memory. +/// `Dictionary(Int32, Utf8)` lets readers hold each distinct string once, +/// cutting resident memory (not just on-disk size, which parquet already +/// dictionary-encoded). +#[must_use] +pub fn label_dictionary_type() -> DataType { + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)) +}