From 577de012b2e2817932e26ea8e51876d375c95590 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 14 Sep 2026 09:01:09 +0000 Subject: [PATCH] perf(servers): group Prometheus response rows by label runs (#9092) * perf(servers): group Prometheus response rows by label runs Query output tends to be clustered by series, but a matrix response read the same label values out of the tag columns once per row, and allocated a key vector per row to look the series up. Use `arrow::compute::partition` to find the runs of rows that share their labels and build the series key once per run. The key buffer is hoisted out of the row loop and handed to the map only when the series is new, through the raw entry API so the key is hashed once either way. Partitioning does not pay off when rows are not clustered, so a few adjacent row pairs are probed first to pick between the run path and the row-by-row path. Both paths produce the same series. Drop the per-row "same labels as the previous row" check from #8815. Runs cover the clustered case it was written for, and it now costs more than it saves: 10% on a result with one row per series, 1-3% on clustered ones. Signed-off-by: Dennis Zhuang * refactor(servers): find label runs with cmp::distinct `arrow::compute::partition` computes the same ranges on the same kernel, but its contract takes lexicographically sorted columns, and query output is not sorted: range queries run without the plan's output sort since #9090, `sort`/`topk` order by value, and the tag column order in the schema does not have to match any sort key. An implementation that exploited the precondition would merge `a, b, a` into one run and attribute one series' samples to another, without failing. `cmp::distinct` is element-wise, so it holds for any row order, and its null handling is the one a series key needs: a null label and an empty one are distinct, two nulls are not. Building the ranges from the boundary bitmask also folds away the `tag_columns.is_empty()` case, since no columns means no boundaries means a single run. Same kernel, so the benchmark does not move: -1.4% to +1.4% across shapes with no consistent sign, against +-3% run-to-run drift. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- .../src/http/result/prometheus_resp.rs | 398 ++++++++++++++---- 1 file changed, 325 insertions(+), 73 deletions(-) diff --git a/src/servers/src/http/result/prometheus_resp.rs b/src/servers/src/http/result/prometheus_resp.rs index 3ebec961d4..eca3833871 100644 --- a/src/servers/src/http/result/prometheus_resp.rs +++ b/src/servers/src/http/result/prometheus_resp.rs @@ -15,8 +15,12 @@ //! prom supply the prometheus HTTP API Server compliance use std::cmp::Ordering; use std::collections::{BTreeMap, HashMap}; +use std::hash::BuildHasher; +use std::ops::Range; -use arrow::array::{Array, AsArray, StructArray}; +use arrow::array::{Array, ArrayRef, AsArray, StructArray}; +use arrow::buffer::BooleanBuffer; +use arrow::compute::kernels::cmp::distinct; use arrow::datatypes::{Float64Type, TimestampMillisecondType}; use arrow_schema::DataType; use axum::Json; @@ -36,6 +40,9 @@ use common_recordbatch::RecordBatches; use datatypes::arrow_array::string_array_value_at_index; use datatypes::prelude::ConcreteDataType; use indexmap::IndexMap; +use indexmap::map::RawEntryApiV1; +use indexmap::map::raw_entry_v1::RawEntryMut; +use itertools::Either; use promql_parser::label::METRIC_NAME; use promql_parser::parser::value::ValueType; use ryu::Buffer; @@ -356,13 +363,8 @@ impl PrometheusJsonResponse { // Tag order matters, e.g., after sorc and sort_desc, the output order must be kept. let mut buffer = IndexMap::, PromSeriesSamples>::new(); - // Consecutive rows often belong to the same series: instant query plans - // keep their output sort, and range query plans, which no longer do, still - // tend to emit a series' rows together. Remember the index of the previous - // row's entry in `buffer` and reuse it when the tags are unchanged, so the - // label vector is not rebuilt and rehashed per row. Unclustered rows only - // cost one `Vec` comparison before falling back to the map lookup. - let mut last_entry_index = None; + // Only a series that is new to `buffer` needs its own key vector. + let mut tags = Vec::with_capacity(num_label_columns + 1); let schema = batches.schema(); for batch in batches.iter() { @@ -398,72 +400,95 @@ impl PrometheusJsonResponse { }) .transpose()?; + // Read the labels once per run of rows that share them instead of + // once per row. `label_runs` finds every boundary, so the probe only + // decides whether looking for runs is worth its cost. + let label_runs = if prefer_label_runs(&tag_columns, batch.num_rows()) { + Either::Left(label_runs(&tag_columns, batch.num_rows())?.into_iter()) + } else { + Either::Right((0..batch.num_rows()).map(|row| row..row + 1)) + }; + // assemble rows - for row_index in 0..batch.num_rows() { - let value = field_column.and_then(|field_column| { - if !field_column.is_valid(row_index) { - return None; + for run in label_runs { + let mut run_entry_index = None; + for row_index in run { + let value = field_column.and_then(|field_column| { + if !field_column.is_valid(row_index) { + return None; + } + let value = field_column.value(row_index); + (!is_prometheus_stale_nan(value)) + .then_some((timestamp_column.value(row_index), value)) + }); + let histogram = native_histogram_column + .and_then(|column| { + read_histogram(column, row_index) + .context(DataFusionSnafu) + .transpose() + }) + .transpose()? + .filter(|histogram| !is_prometheus_stale_nan(histogram.sum)) + .map(|histogram| { + prometheus_native_histogram(&histogram) + .map(|histogram| (timestamp_column.value(row_index), histogram)) + }) + .transpose()?; + + if value.is_none() && histogram.is_none() { + continue; } - let value = field_column.value(row_index); - (!is_prometheus_stale_nan(value)) - .then_some((timestamp_column.value(row_index), value)) - }); - let histogram = native_histogram_column - .and_then(|column| { - read_histogram(column, row_index) - .context(DataFusionSnafu) - .transpose() - }) - .transpose()? - .filter(|histogram| !is_prometheus_stale_nan(histogram.sum)) - .map(|histogram| { - prometheus_native_histogram(&histogram) - .map(|histogram| (timestamp_column.value(row_index), histogram)) - }) - .transpose()?; - if value.is_none() && histogram.is_none() { - continue; - } + let entry_index = match run_entry_index { + Some(index) => index, + None => { + // retrieve tags + tags.clear(); + if let Some(metric_name) = &metric_name { + tags.push((METRIC_NAME, metric_name.as_str())); + } + for (tag_column, tag_name) in tag_columns.iter().zip(tag_names.iter()) { + if let Some(tag_value) = + string_array_value_at_index(tag_column, row_index) + { + tags.push((tag_name, tag_value)); + } + } - // retrieve tags - let mut tags = Vec::with_capacity(num_label_columns + 1); - if let Some(metric_name) = &metric_name { - tags.push((METRIC_NAME, metric_name.as_str())); - } - for (tag_column, tag_name) in tag_columns.iter().zip(tag_names.iter()) { - if let Some(tag_value) = string_array_value_at_index(tag_column, row_index) { - tags.push((tag_name, tag_value)); + let hash = buffer.hasher().hash_one(&tags); + let entry = buffer + .raw_entry_mut_v1() + .from_key_hashed_nocheck(hash, &tags); + let index = entry.index(); + if let RawEntryMut::Vacant(entry) = entry { + // Hand the key buffer over with the hash that was + // just computed for the lookup. + let key = std::mem::replace( + &mut tags, + Vec::with_capacity(num_label_columns + 1), + ); + entry.insert_hashed_nocheck( + hash, + key, + PromSeriesSamples::default(), + ); + } + run_entry_index = Some(index); + index + } + }; + let samples = &mut buffer[entry_index]; + if let Some((timestamp_millis, histogram)) = histogram { + samples + .histograms + .push((timestamp_millis as f64 / 1000.0, histogram)); + } else if let Some((timestamp_millis, value)) = value { + samples.values.push(( + timestamp_millis as f64 / 1000.0, + PromSampleValue::Number(value), + )); } } - - let reuse = last_entry_index.filter(|index| { - buffer - .get_index(*index) - .is_some_and(|(key, _)| key == &tags) - }); - let samples = if let Some(index) = reuse { - buffer - .get_index_mut(index) - .map(|(_, samples)| samples) - .with_context(|| UnexpectedResultSnafu { - reason: "reused series entry must exist", - })? - } else { - let entry = buffer.entry(tags); - last_entry_index = Some(entry.index()); - entry.or_default() - }; - if let Some((timestamp_millis, histogram)) = histogram { - samples - .histograms - .push((timestamp_millis as f64 / 1000.0, histogram)); - } else if let Some((timestamp_millis, value)) = value { - samples.values.push(( - timestamp_millis as f64 / 1000.0, - PromSampleValue::Number(value), - )); - } } } @@ -545,6 +570,71 @@ impl PrometheusJsonResponse { } } +/// Ranges of consecutive rows whose label values are all equal. +/// +/// `arrow::compute::partition` computes the same ranges, but its contract takes +/// lexicographically sorted columns, which query output is not: range queries +/// run without the plan's output sort, and `sort`/`topk` order by value. +/// `distinct` is element-wise, so it holds for any row order, and its null +/// handling is the one a series key needs: a null label and an empty one are +/// distinct, and two nulls are not. +fn label_runs(columns: &[&ArrayRef], rows: usize) -> Result>> { + let mut boundaries: Option = None; + if rows >= 2 { + for column in columns { + let changed = distinct(&column.slice(0, rows - 1), &column.slice(1, rows - 1)) + .context(ArrowSnafu)?; + boundaries = Some(match boundaries { + Some(accumulated) => &accumulated | changed.values(), + None => changed.values().clone(), + }); + } + } + + let mut runs = Vec::new(); + let mut start = 0; + for boundary in boundaries.iter().flat_map(BooleanBuffer::set_indices) { + runs.push(start..boundary + 1); + start = boundary + 1; + } + runs.push(start..rows); + Ok(runs) +} + +/// Decides whether to group the rows of a batch into runs that share their +/// label values. +/// +/// Sampling adjacent row pairs keeps the decision independent of the batch +/// size. Probe positions come from a xorshift sequence instead of a fixed +/// stride, which would alias with periodic series layouts. A wrong guess only +/// costs time: `partition` still validates every boundary, and the row-by-row +/// path builds the labels of every row. +fn prefer_label_runs(columns: &[&ArrayRef], rows: usize) -> bool { + if columns.is_empty() { + return true; + } + if rows < 2 { + return false; + } + let mut position = 0x9e37_79b9_u32; + let mut changes = 0; + for _ in 0..8 { + position ^= position << 13; + position ^= position >> 17; + position ^= position << 5; + let row = position as usize % (rows - 1); + if columns.iter().any(|column| { + string_array_value_at_index(column, row) != string_array_value_at_index(column, row + 1) + }) { + changes += 1; + if changes > 2 { + return false; + } + } + } + true +} + fn merge_annotations(target: &mut Option>, source: Option>) { let Some(source) = source else { return; @@ -953,11 +1043,10 @@ mod tests { } #[test] - fn record_batches_to_data_reuses_entries_for_clustered_series() { - // Rows are clustered by series (a, b, a, a, b, c): consecutive rows of - // the same series exercise the entry-reuse fast path, while series - // transitions fall back to the map lookup. The result must keep the - // first-occurrence order and accumulate values per series as before. + fn record_batches_to_data_groups_clustered_series() { + // Rows are clustered by series (a, b, a, a, b, c) and `a` is revisited + // after `b`. The result must keep the first-occurrence order and + // accumulate values per series. let schema = Arc::new(Schema::new(vec![ ColumnSchema::new( "timestamp", @@ -1023,6 +1112,169 @@ mod tests { assert_eq!(series[2].value, Some((6.0, "6.0".to_string()))); } + #[test] + fn label_strategy_switches_preserve_all_rows_when_probes_miss_changes() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ColumnSchema::new("host", ConcreteDataType::string_datatype(), false), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), false), + ])); + let mut batches = Vec::new(); + let mut expected = BTreeMap::>::new(); + for layout in 0..3 { + let labels = (0..1024) + .map(|row| { + let alternating = match layout { + 0 => true, + 1 => row < 64, + _ => row >= 64, + }; + if alternating && row % 2 != 0 { + "b" + } else { + "a" + } + }) + .collect::>(); + for (row, label) in labels.iter().enumerate() { + let value = (layout * 1024 + row) as f64; + expected + .entry((*label).to_string()) + .or_default() + .push((value, PromSampleValue::Number(value))); + } + let batch = RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_values( + (0..1024).map(|row| (layout * 1024 + row) as i64 * 1000), + )) as _, + Arc::new(StringVector::from( + labels.into_iter().map(Some).collect::>(), + )) as _, + Arc::new(Float64Vector::from( + (0..1024) + .map(|row| Some((layout * 1024 + row) as f64)) + .collect::>(), + )) as _, + ], + ) + .unwrap(); + // The middle batch alternates only over a prefix the probes miss, + // so it takes the run path even though most rows are not runs. + assert_eq!( + prefer_label_runs(&[batch.column(1)], batch.num_rows()), + layout == 1 + ); + batches.push(batch); + } + let response = PrometheusJsonResponse::record_batches_to_data( + RecordBatches::try_new(schema, batches).unwrap(), + None, + ValueType::Matrix, + ) + .unwrap(); + let PrometheusResponse::PromData(PromData { + result: PromQueryResult::Matrix(series), + .. + }) = response + else { + panic!("expected matrix response"); + }; + assert_eq!(series.len(), expected.len()); + for series in series { + assert_eq!(series.metric.len(), 1); + assert!(series.histograms.is_empty()); + assert_eq!( + series.values, + expected.remove(&series.metric["host"]).unwrap() + ); + } + assert!(expected.is_empty()); + } + + #[test] + fn label_runs_keep_null_and_empty_labels_apart_across_batches() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("rack", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), false), + ])); + // Two batches of two runs each, with only `rack` changing: a null label + // and an empty one must not share a run, and the run opening the second + // batch continues the series that ended the first one. + let mut batches = Vec::new(); + let mut expected = vec![Vec::new(), Vec::new()]; + for (batch_index, leading_null) in [true, false].into_iter().enumerate() { + let mut racks = Vec::new(); + let mut values = Vec::new(); + for row in 0..1024 { + let value = (batch_index * 1024 + row) as f64; + let null_rack = (row < 512) == leading_null; + racks.push((!null_rack).then_some("")); + values.push(Some(value)); + expected[usize::from(!null_rack)].push((value, PromSampleValue::Number(value))); + } + batches.push( + RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_values( + values.iter().map(|value| value.unwrap() as i64 * 1000), + )) as _, + Arc::new(StringVector::from(vec![Some("a"); 1024])) as _, + Arc::new(StringVector::from(racks)) as _, + Arc::new(Float64Vector::from(values)) as _, + ], + ) + .unwrap(), + ); + } + for batch in &batches { + assert!(prefer_label_runs( + &[batch.column(1), batch.column(2)], + batch.num_rows() + )); + } + for series in &mut expected { + series.sort_by(|left, right| left.0.total_cmp(&right.0)); + } + + let response = PrometheusJsonResponse::record_batches_to_data( + RecordBatches::try_new(schema, batches).unwrap(), + None, + ValueType::Matrix, + ) + .unwrap(); + let PrometheusResponse::PromData(PromData { + result: PromQueryResult::Matrix(series), + .. + }) = response + else { + panic!("expected matrix response"); + }; + assert_eq!(series.len(), 2); + assert_eq!( + series[0].metric, + BTreeMap::from([("host".into(), "a".into())]) + ); + assert_eq!(series[0].values, expected[0]); + assert_eq!( + series[1].metric, + BTreeMap::from([("host".into(), "a".into()), ("rack".into(), "".into())]) + ); + assert_eq!(series[1].values, expected[1]); + } + #[test] fn matrix_response_is_independent_of_input_row_order() { // Range queries run without the plan's output sort, so this function sees