starred/openobserve
1
0
Fork 0
mirror of https://github.com/openobserve/openobserve.git synced 2026-11-05 16:10:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf(promql): preallocate hash-sorted series samples (#13971)

Browse Source
This commit is contained in:
Huaijin
2026-08-26 14:35:33 +00:00
committed by GitHub
parent 94a0a58866
commit 775094165c
3 changed files with 207 additions and 116 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/promql/src/load_series/load_labels.rs
+1 -1
View File
@@ -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();
src/promql/src/load_series/mod.rs
+98 -115
View File
@@ -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<HashMap<u64, RangeValue>>;
@@ -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<i64>)> {
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::<Int64Array>()
.unwrap();
let value_values = batch
.column_by_name(VALUE_LABEL)
.unwrap()
.as_any()
.downcast_ref::<Float64Array>()
.unwrap();
let time_values = batch[TIMESTAMP_COL_NAME].as_primitive::<Int64Type>();
let value_values = batch[VALUE_LABEL].as_primitive::<Float64Type>();
if hash_field_type == DataType::UInt64 {
let hash_values = batch
.column_by_name(HASH_LABEL)
.unwrap()
.as_any()
.downcast_ref::<UInt64Array>()
.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::<StringArray>()
.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<u64, RangeValue>,
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(|(&timestamp, &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::<UInt64Type>();
Cow::Borrowed(hash_values.values().as_ref())
} else {
let hash_values = batch[HASH_LABEL].as_string::<i32>();
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::<StringArray>()
.unwrap();
if hash_field_type == DataType::UInt64 {
let hash_values = batch
.column_by_name(HASH_LABEL)
.unwrap()
.as_any()
.downcast_ref::<UInt64Array>()
.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::<Vec<json::Value>>(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::<StringArray>()
.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::<Vec<json::Value>>(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::<i32>();
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::<Vec<json::Value>>(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);
src/promql/src/load_series/series_capacity.rs
+108
View File
@@ -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 <http://www.gnu.org/licenses/>.
//! 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);
}
}