diff --git a/src/datatypes/src/series.rs b/src/datatypes/src/series.rs index 11595b4de22..56369e93dd8 100644 --- a/src/datatypes/src/series.rs +++ b/src/datatypes/src/series.rs @@ -22,7 +22,11 @@ pub mod time; #[cfg(test)] mod tests; -use std::sync::Arc; +use std::any::Any; +use std::sync::{Arc, OnceLock}; + +use arrow::array::{Array, StringArray}; +use arrow::record_batch::RecordBatch; use crate::error::{InvalidSeriesSnafu, Result}; @@ -67,10 +71,10 @@ pub struct Label { /// The metric name is only the `__name__` value (see [`Self::name`]); this type is /// the whole label set a PromQL series shows. It is neither a canonical /// matching/grouping identity nor a storage identity (read key / primary key). -#[derive(Clone, Debug, PartialEq, Eq, Hash)] +#[derive(Clone, Debug)] pub struct SeriesLabels { name: Option>, - tags: Arc<[Label]>, + tags: LabelStorage, } impl serde::Serialize for SeriesLabels { @@ -78,7 +82,7 @@ impl serde::Serialize for SeriesLabels { &self, serializer: S, ) -> std::result::Result { - (&self.name, self.tags.as_ref()).serialize(serializer) + (&self.name, self.tags()).serialize(serializer) } } impl<'de> serde::Deserialize<'de> for SeriesLabels { @@ -91,6 +95,129 @@ impl<'de> serde::Deserialize<'de> for SeriesLabels { } } +/// Shared, bounded series metadata and the owner of its memory reservation. +/// Columns map logical label names to UTF-8 arrays; null labels are absent. +pub struct SeriesLabelBatch { + batch: RecordBatch, + columns: Vec<(Arc, usize)>, + _owner: Arc, +} + +impl std::fmt::Debug for SeriesLabelBatch { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("SeriesLabelBatch") + .field("rows", &self.batch.num_rows()) + .field("columns", &self.columns) + .finish_non_exhaustive() + } +} + +impl SeriesLabelBatch { + /// Validates mappings once. The owner must account for the retained backing buffers. + pub fn try_new( + batch: RecordBatch, + mut columns: Vec<(Arc, usize)>, + owner: Arc, + ) -> Result { + columns.sort_unstable_by(|a, b| a.0.cmp(&b.0)); + if columns.iter().any(|(name, index)| { + name.as_ref() == "__name__" + || batch + .columns() + .get(*index) + .is_none_or(|a| !a.as_any().is::()) + }) || columns.windows(2).any(|p| p[0].0 == p[1].0) + { + return invalid_series("invalid batch label mapping"); + } + Ok(Self { + batch, + columns, + _owner: owner, + }) + } + + /// Estimates retained buffers and one row view per batch row, including lazy + /// owned-label materialization. Shared buffers are conservatively counted per + /// array. Additional independent views of a row need their own reservation. + /// This is a reservation estimate, not an allocator-exact accounting contract. + pub fn estimated_size(&self) -> usize { + let arc_header = 2 * std::mem::size_of::(); + let row_views = self + .batch + .num_rows() + .saturating_mul(std::mem::size_of::() + arc_header); + let mut bytes = self + .batch + .get_array_memory_size() + .saturating_add(std::mem::size_of::() + arc_header) + .saturating_add( + self.columns + .capacity() + .saturating_mul(std::mem::size_of::<(Arc, usize)>()), + ) + .saturating_add(row_views); + for (name, column) in &self.columns { + bytes = bytes.saturating_add(name.len() + arc_header); + for row in 0..self.batch.num_rows() { + if let Some(value) = self.value(*column, row) { + // Vec-to-Arc conversion may temporarily retain both label arrays. + bytes = bytes + .saturating_add(2 * std::mem::size_of::