From 775094165ccfe53ef4b4907cedfb292fe749cc35 Mon Sep 17 00:00:00 2001 From: Huaijin Date: Wed, 26 Aug 2026 14:35:33 +0000 Subject: [PATCH] perf(promql): preallocate hash-sorted series samples (#13971) --- src/promql/src/load_series/load_labels.rs | 2 +- src/promql/src/load_series/mod.rs | 213 ++++++++---------- src/promql/src/load_series/series_capacity.rs | 108 +++++++++ 3 files changed, 207 insertions(+), 116 deletions(-) create mode 100644 src/promql/src/load_series/series_capacity.rs diff --git a/src/promql/src/load_series/load_labels.rs b/src/promql/src/load_series/load_labels.rs index bea4073240..cfa30938fe 100644 --- a/src/promql/src/load_series/load_labels.rs +++ b/src/promql/src/load_series/load_labels.rs @@ -503,7 +503,7 @@ mod tests { .unwrap(); let sample_df = ctx.read_batch(sample_batch).unwrap(); let (metrics, _) = - load_samples_from_datafusion("test", &DataType::UInt64, sample_df, false) + load_samples_from_datafusion("test", &DataType::UInt64, sample_df, false, 1, 0) .await .unwrap(); let label_df = ctx.read_batch(label_batch).unwrap(); diff --git a/src/promql/src/load_series/mod.rs b/src/promql/src/load_series/mod.rs index 369d4b7def..49555c46e8 100644 --- a/src/promql/src/load_series/mod.rs +++ b/src/promql/src/load_series/mod.rs @@ -18,8 +18,9 @@ mod label_cache; mod labels; mod load_labels; +mod series_capacity; -use std::sync::Arc; +use std::{borrow::Cow, sync::Arc}; use config::{ TIMESTAMP_COL_NAME, @@ -34,8 +35,8 @@ use config::{ }; use datafusion::{ arrow::{ - array::{Float64Array, Int64Array, StringArray, UInt64Array}, - datatypes::{DataType, Schema}, + array::{Array, AsArray, RecordBatch}, + datatypes::{DataType, Float64Type, Int64Type, Schema, UInt64Type}, }, error::{DataFusionError, Result}, logical_expr::utils::disjunction, @@ -45,10 +46,13 @@ use datafusion::{ prelude::{DataFrame, Expr, SessionContext, col, lit}, }; use futures::TryStreamExt; -use hashbrown::{HashMap, HashSet}; +use hashbrown::{HashMap, HashSet, hash_map::Entry}; use promql_parser::parser::VectorSelector; -use self::labels::load_series_labels; +use self::{ + labels::load_series_labels, + series_capacity::{batch_run_len, initial_series_capacity, series_fragment_hint}, +}; use super::utils::{apply_label_selector, apply_matchers}; pub(super) type PartitionedMetrics = Vec>; @@ -189,11 +193,14 @@ pub(super) async fn selector_load_data_from_datafusion( ) .await? } else { + let query_duration = end.saturating_sub(start.saturating_sub(lookback)); load_samples_from_datafusion( &query_ctx.trace_id, hash_field_type, df_group.clone(), !skip_labels, + series_fragment_hint(query_duration), + query_duration, ) .await? }; @@ -246,6 +253,8 @@ pub(super) async fn load_samples_from_datafusion( hash_field_type: &DataType, df: DataFrame, collect_timestamps: bool, + fragment_hint: usize, + query_duration: i64, ) -> Result<(PartitionedMetrics, HashSet)> { let ctx = Arc::new(df.task_ctx()); let target_partitions = ctx.session_config().target_partitions(); @@ -278,60 +287,18 @@ pub(super) async fn load_samples_from_datafusion( loop { match stream.try_next().await { Ok(Some(batch)) => { - let time_values = batch - .column_by_name(TIMESTAMP_COL_NAME) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - let value_values = batch - .column_by_name(VALUE_LABEL) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); + let time_values = batch[TIMESTAMP_COL_NAME].as_primitive::(); + let value_values = batch[VALUE_LABEL].as_primitive::(); - if hash_field_type == DataType::UInt64 { - let hash_values = batch - .column_by_name(HASH_LABEL) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - for i in 0..batch.num_rows() { - let timestamp = time_values.value(i); - let hash: u64 = hash_values.value(i); - let entry = metrics.entry(hash).or_insert_with(|| RangeValue { - labels: vec![], - samples: vec![], - exemplars: None, - time_window: None, - }); - entry - .samples - .push(Sample::new(timestamp, value_values.value(i))); - } - } else { - let hash_values = batch - .column_by_name(HASH_LABEL) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - for i in 0..batch.num_rows() { - let timestamp = time_values.value(i); - let hash: u64 = gxhash::new().sum64(hash_values.value(i)); - let entry = metrics.entry(hash).or_insert_with(|| RangeValue { - labels: vec![], - samples: vec![], - exemplars: None, - time_window: None, - }); - entry - .samples - .push(Sample::new(timestamp, value_values.value(i))); - } - } + let hashes = batch_hash_values(&batch, &hash_field_type); + append_batch_samples( + &mut metrics, + &hashes, + time_values.values(), + value_values.values(), + fragment_hint, + query_duration, + ); } Ok(None) => break, Err(e) => { @@ -368,6 +335,62 @@ pub(super) async fn load_samples_from_datafusion( Ok((metrics, all_unique_timestamps)) } +fn append_batch_samples( + metrics: &mut HashMap, + hashes: &[u64], + timestamps: &[i64], + values: &[f64], + fragment_hint: usize, + query_duration: i64, +) { + let mut i = 0; + while i < hashes.len() { + let run_len = batch_run_len(hashes, i); + let entry = match metrics.entry(hashes[i]) { + Entry::Occupied(entry) => entry.into_mut(), + Entry::Vacant(entry) => { + let capacity = initial_series_capacity( + run_len, + fragment_hint, + timestamps[i], + timestamps[i + run_len - 1], + query_duration, + ); + entry.insert(RangeValue { + labels: vec![], + samples: Vec::with_capacity(capacity), + exemplars: None, + time_window: None, + }) + } + }; + entry.samples.extend( + timestamps[i..i + run_len] + .iter() + .zip(&values[i..i + run_len]) + .map(|(×tamp, &value)| Sample::new(timestamp, value)), + ); + i += run_len; + } +} + +/// The hash column as u64 values: zero-copy for UInt64, hashed per row for +/// Utf8. +fn batch_hash_values<'a>(batch: &'a RecordBatch, hash_field_type: &DataType) -> Cow<'a, [u64]> { + if *hash_field_type == DataType::UInt64 { + let hash_values = batch[HASH_LABEL].as_primitive::(); + Cow::Borrowed(hash_values.values().as_ref()) + } else { + let hash_values = batch[HASH_LABEL].as_string::(); + let mut hasher = gxhash::new(); + Cow::Owned( + (0..hash_values.len()) + .map(|i| hasher.sum64(hash_values.value(i))) + .collect(), + ) + } +} + async fn load_exemplars_from_datafusion( trace_id: &str, hash_field_type: &DataType, @@ -406,61 +429,21 @@ async fn load_exemplars_from_datafusion( loop { match stream.try_next().await { Ok(Some(batch)) => { - let exemplars_values = batch - .column_by_name(EXEMPLARS_LABEL) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - if hash_field_type == DataType::UInt64 { - let hash_values = batch - .column_by_name(HASH_LABEL) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - for i in 0..batch.num_rows() { - let hash: u64 = hash_values.value(i); - let exemplar = exemplars_values.value(i); - if let Ok(exemplars) = json::from_str::>(exemplar) - { - let entry = metrics.entry(hash).or_insert_with(|| RangeValue { - labels: vec![], - samples: vec![], - exemplars: Some(vec![]), - time_window: None, - }); - let entry = entry.exemplars.as_mut().unwrap(); - for exemplar in exemplars { - if let Some(exemplar) = exemplar.as_object() { - entry.push(Arc::new(Exemplar::from(exemplar))); - } - } - } - } - } else { - let hash_values = batch - .column_by_name(HASH_LABEL) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - for i in 0..batch.num_rows() { - let hash: u64 = gxhash::new().sum64(hash_values.value(i)); - let exemplar = exemplars_values.value(i); - if let Ok(exemplars) = json::from_str::>(exemplar) - { - let entry = metrics.entry(hash).or_insert_with(|| RangeValue { - labels: vec![], - samples: vec![], - exemplars: Some(vec![]), - time_window: None, - }); - let entry = entry.exemplars.as_mut().unwrap(); - for exemplar in exemplars { - if let Some(exemplar) = exemplar.as_object() { - entry.push(Arc::new(Exemplar::from(exemplar))); - } + let exemplars_values = batch[EXEMPLARS_LABEL].as_string::(); + let hashes = batch_hash_values(&batch, &hash_field_type); + for (i, &hash) in hashes.iter().enumerate() { + let exemplar = exemplars_values.value(i); + if let Ok(exemplars) = json::from_str::>(exemplar) { + let entry = metrics.entry(hash).or_insert_with(|| RangeValue { + labels: vec![], + samples: vec![], + exemplars: Some(vec![]), + time_window: None, + }); + let entry = entry.exemplars.as_mut().unwrap(); + for exemplar in exemplars { + if let Some(exemplar) = exemplar.as_object() { + entry.push(Arc::new(Exemplar::from(exemplar))); } } } @@ -569,7 +552,7 @@ mod tests { let df = ctx.read_batch(batch).unwrap(); let (metrics_without_timestamps, skipped_timestamps) = - load_samples_from_datafusion("test", &DataType::UInt64, df.clone(), false) + load_samples_from_datafusion("test", &DataType::UInt64, df.clone(), false, 1, 0) .await .unwrap(); assert!(skipped_timestamps.is_empty()); @@ -578,7 +561,7 @@ mod tests { assert_eq!(metrics_without_timestamps[&11].samples.len(), 2); let (metrics, timestamps) = - load_samples_from_datafusion("test", &DataType::UInt64, df, true) + load_samples_from_datafusion("test", &DataType::UInt64, df, true, 1, 0) .await .unwrap(); let metrics = merge_partitioned_metrics(metrics); diff --git a/src/promql/src/load_series/series_capacity.rs b/src/promql/src/load_series/series_capacity.rs new file mode 100644 index 0000000000..177a8f7588 --- /dev/null +++ b/src/promql/src/load_series/series_capacity.rs @@ -0,0 +1,108 @@ +// Copyright 2026 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +//! Sample-capacity estimation for hash-sorted series materialization. + +const STORAGE_HOUR_MICROS: i64 = 60 * 60 * 1_000_000; +const MAX_SERIES_FRAGMENT_HINT: usize = 24; +const MAX_INITIAL_SERIES_CAPACITY: usize = 2048; + +/// Length of the contiguous run of equal hashes starting at `start`. +pub(super) fn batch_run_len(hashes: &[u64], start: usize) -> usize { + let hash = hashes[start]; + let mut end = start + 1; + while end < hashes.len() && hashes[end] == hash { + end += 1; + } + end - start +} + +/// Estimate the total sample count of a series from its first contiguous run +/// of `first_run_len` rows spanning `[run_first_ts, run_last_ts]`. +/// +/// Two estimates cover the common shapes: `run × fragments` fits a series +/// whose runs arrive whole, `duration / interval` recovers a first run +/// truncated by a batch boundary. Overestimating (short-lived series, the +/// sparse-window optimization path) only wastes capacity, bounded by +/// `MAX_INITIAL_SERIES_CAPACITY`. +pub(super) fn initial_series_capacity( + first_run_len: usize, + fragment_hint: usize, + run_first_ts: i64, + run_last_ts: i64, + query_duration: i64, +) -> usize { + let run_based = first_run_len.saturating_mul(fragment_hint); + // At least two intervals, so a single anomalous gap (adjacent + // near-duplicate rows in time-sorted input) cannot dictate the estimate. + let interval_based = if first_run_len >= 3 && run_last_ts > run_first_ts { + let sample_interval = + ((run_last_ts - run_first_ts) / (first_run_len - 1) as i64).max(1) as u64; + let samples = (query_duration.max(0) as u64) + .div_ceil(sample_interval) + .saturating_add(1); + usize::try_from(samples).unwrap_or(MAX_INITIAL_SERIES_CAPACITY) + } else { + 0 + }; + // The current batch already proves first_run_len samples exist, so the + // cap never allocates below that. + let cap = MAX_INITIAL_SERIES_CAPACITY.max(first_run_len); + run_based.max(interval_based).min(cap) +} + +/// Hash-sorted parquet is written per storage hour, so a series arrives as +/// roughly one contiguous run per hour fragment of the query span. +pub(super) fn series_fragment_hint(query_duration: i64) -> usize { + let hourly_fragments = (query_duration.max(0) as u64).div_ceil(STORAGE_HOUR_MICROS as u64); + hourly_fragments.clamp(1, MAX_SERIES_FRAGMENT_HINT as u64) as usize +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_batch_run_len() { + let hashes = [7u64, 7, 7, 9, 9, 1]; + assert_eq!(batch_run_len(&hashes, 0), 3); + assert_eq!(batch_run_len(&hashes, 3), 2); + assert_eq!(batch_run_len(&hashes, 5), 1); + } + + #[test] + fn test_initial_series_capacity_is_bounded() { + assert_eq!(initial_series_capacity(160, 4, 0, 0, 0), 640); + assert_eq!( + initial_series_capacity(100, 4, 0, 99 * 15 * 1_000_000, 11_100 * 1_000_000), + 741, + ); + assert_eq!(initial_series_capacity(1024, 4, 0, 0, 0), 2048); + assert_eq!(initial_series_capacity(4096, 4, 0, 0, 0), 4096); + // A 2-row run must not infer an interval from its single gap. + assert_eq!( + initial_series_capacity(2, 4, 0, 1_000, 11_100 * 1_000_000), + 8 + ); + } + + #[test] + fn test_series_fragment_hint_is_bounded() { + let hour = STORAGE_HOUR_MICROS; + assert_eq!(series_fragment_hint(3 * hour + 5 * 60 * 1_000_000), 4); + assert_eq!(series_fragment_hint(0), 1); + assert_eq!(series_fragment_hint(100 * hour), MAX_SERIES_FRAGMENT_HINT); + } +}