diff --git a/src/servers/src/http/result/prometheus_resp.rs b/src/servers/src/http/result/prometheus_resp.rs index eca38338719..c2ce10a0db3 100644 --- a/src/servers/src/http/result/prometheus_resp.rs +++ b/src/servers/src/http/result/prometheus_resp.rs @@ -36,12 +36,14 @@ use common_query::promql_annotations::{ PromqlAnnotationCollector, get_promql_annotation_collector, }; use common_query::{Output, OutputData}; -use common_recordbatch::RecordBatches; +use common_recordbatch::{RecordBatch, RecordBatches, SendableRecordBatchStream}; use datatypes::arrow_array::string_array_value_at_index; use datatypes::prelude::ConcreteDataType; -use indexmap::IndexMap; +use datatypes::schema::Schema; +use futures::TryStreamExt; use indexmap::map::RawEntryApiV1; use indexmap::map::raw_entry_v1::RawEntryMut; +use indexmap::{Equivalent, IndexMap}; use itertools::Either; use promql_parser::label::METRIC_NAME; use promql_parser::parser::value::ValueType; @@ -236,16 +238,9 @@ impl PrometheusJsonResponse { OutputData::RecordBatches(batches) => Self::success( Self::record_batches_to_data(batches, metric_name, result_type)?, ), - OutputData::Stream(stream) => { - let record_batches = RecordBatches::try_collect(stream) - .await - .context(CollectRecordbatchSnafu)?; - Self::success(Self::record_batches_to_data( - record_batches, - metric_name, - result_type, - )?) - } + OutputData::Stream(stream) => Self::success( + Self::consume_stream_to_data(stream, metric_name, result_type).await?, + ), OutputData::AffectedRows(_) => Self::error( StatusCode::Unexpected, "expected data result, but got affected rows", @@ -299,22 +294,83 @@ impl PrometheusJsonResponse { ) -> Result { // Return empty result if no batches if batches.iter().next().is_none() { - return Ok(PrometheusResponse::PromData(PromData { - result_type: result_type.to_string(), - ..Default::default() - })); + return Ok(empty_data(result_type)); } - // infer semantic type of each column from schema. - // TODO(ruihang): wish there is a better way to do this. + let layout = ColumnLayout::infer(&batches.schema())?; + + // Preserves the order of output tags. + // Tag order matters, e.g., after sorc and sort_desc, the output order must be kept. + let mut buffer = IndexMap::, PromSeriesSamples>::new(); + + for batch in batches.iter() { + merge_batch(&mut buffer, &layout, batch, metric_name.as_deref())?; + } + + samples_to_data(buffer, result_type) + } + + /// Consume a streaming query result into [PromData]. + /// + /// Every batch is merged into the per-series buffer and dropped before the + /// next one is polled, so the whole query result and the assembled series + /// never coexist in memory. + async fn consume_stream_to_data( + mut stream: SendableRecordBatchStream, + metric_name: Option, + result_type: ValueType, + ) -> Result { + // An empty stream carries no batch to infer the column roles from, and + // reads as an empty result, like empty [RecordBatches]. + let Some(first_batch) = stream.try_next().await.context(CollectRecordbatchSnafu)? else { + return Ok(empty_data(result_type)); + }; + + let layout = ColumnLayout::infer(&stream.schema())?; + + // Preserves the order of output tags. + // Tag order matters, e.g., after sorc and sort_desc, the output order must be kept. + let mut buffer = IndexMap::, PromSeriesSamples>::new(); + + merge_batch(&mut buffer, &layout, &first_batch, metric_name.as_deref())?; + // Release the batch before polling the next one. + drop(first_batch); + + while let Some(batch) = stream.try_next().await.context(CollectRecordbatchSnafu)? { + merge_batch(&mut buffer, &layout, &batch, metric_name.as_deref())?; + } + + samples_to_data(buffer, result_type) + } +} + +/// The empty data of a Prometheus HTTP API response of `result_type`. +fn empty_data(result_type: ValueType) -> PrometheusResponse { + PrometheusResponse::PromData(PromData { + result_type: result_type.to_string(), + ..Default::default() + }) +} + +/// The role of every column of a query result. +struct ColumnLayout { + timestamp_column_index: usize, + tag_column_indices: Vec, + tag_names: Vec, + first_field_column_index: Option, + native_histogram_column_index: Option, +} + +impl ColumnLayout { + /// Infer the role of every column from `schema`. + // TODO(ruihang): wish there is a better way to do this. + fn infer(schema: &Schema) -> Result { let mut timestamp_column_index = None; let mut tag_column_indices = Vec::new(); let mut first_field_column_index = None; let mut native_histogram_column_index = None; - let mut num_label_columns = 0; - - for (i, column) in batches.schema().column_schemas().iter().enumerate() { + for (i, column) in schema.column_schemas().iter().enumerate() { match column.data_type { ConcreteDataType::Timestamp(datatypes::types::TimestampType::Millisecond(_)) if timestamp_column_index.is_none() => @@ -343,7 +399,6 @@ impl PrometheusJsonResponse { } ConcreteDataType::String(_) => { tag_column_indices.push(i); - num_label_columns += 1; } _ => {} } @@ -359,215 +414,261 @@ impl PrometheusJsonResponse { .fail(); } - // Preserves the order of output tags. - // Tag order matters, e.g., after sorc and sort_desc, the output order must be kept. - let mut buffer = IndexMap::, PromSeriesSamples>::new(); + let tag_names = tag_column_indices + .iter() + .map(|index| schema.column_name_by_index(*index).to_string()) + .collect::>(); - // Only a series that is new to `buffer` needs its own key vector. - let mut tags = Vec::with_capacity(num_label_columns + 1); + Ok(Self { + timestamp_column_index, + tag_column_indices, + tag_names, + first_field_column_index, + native_histogram_column_index, + }) + } - let schema = batches.schema(); - for batch in batches.iter() { - // prepare things... - let tag_columns = tag_column_indices + /// Capacity of a scratch label buffer: one slot per label plus the metric + /// name. + fn tags_capacity(&self) -> usize { + self.tag_column_indices.len() + 1 + } +} + +/// The lookup key of a series in the per-series sample buffer. +/// +/// The blanket `Equivalent` implementation compares through `Borrow`, which +/// `Vec<(String, String)>` cannot provide for labels borrowed from a batch, so +/// this wrapper compares the labels itself. It is built for a single lookup and +/// hashes like the owned key it looks up. +struct SeriesKeyLookup<'a>(&'a [(&'a str, &'a str)]); + +impl Equivalent> for SeriesKeyLookup<'_> { + fn equivalent(&self, key: &Vec<(String, String)>) -> bool { + self.0.len() == key.len() + && self + .0 .iter() - .map(|i| batch.column(*i)) - .collect::>(); - let tag_names = tag_column_indices - .iter() - .map(|c| schema.column_name_by_index(*c)) - .collect::>(); - let timestamp_column = batch - .column(timestamp_column_index) - .as_primitive::(); + .zip(key) + .all(|((name, value), (key_name, key_value))| { + *name == key_name.as_str() && *value == key_value.as_str() + }) + } +} - let field_array = first_field_column_index - .map(|index| arrow::compute::cast(batch.column(index), &DataType::Float64)) - .transpose() - .context(ArrowSnafu)?; - let field_column = field_array - .as_ref() - .map(|array| array.as_primitive::()); - let native_histogram_column = native_histogram_column_index - .map(|index| { - batch - .column(index) - .as_any() - .downcast_ref::() - .with_context(|| UnexpectedResultSnafu { - reason: "native histogram column is not a struct array", - }) +/// Merge the rows of one batch into the per-series sample buffer. +fn merge_batch( + buffer: &mut IndexMap, PromSeriesSamples>, + layout: &ColumnLayout, + batch: &RecordBatch, + metric_name: Option<&str>, +) -> Result<()> { + // Scratch buffer of the labels of the series being read, borrowed from + // `batch` and `layout` and refilled for every series that is not in + // `buffer` yet, so looking a series up allocates nothing and only a series + // that is new allocates its labels. + let mut tags: Vec<(&str, &str)> = Vec::with_capacity(layout.tags_capacity()); + + // prepare things... + let tag_columns = layout + .tag_column_indices + .iter() + .map(|i| batch.column(*i)) + .collect::>(); + let timestamp_column = batch + .column(layout.timestamp_column_index) + .as_primitive::(); + + let field_array = layout + .first_field_column_index + .map(|index| arrow::compute::cast(batch.column(index), &DataType::Float64)) + .transpose() + .context(ArrowSnafu)?; + let field_column = field_array + .as_ref() + .map(|array| array.as_primitive::()); + let native_histogram_column = layout + .native_histogram_column_index + .map(|index| { + batch + .column(index) + .as_any() + .downcast_ref::() + .with_context(|| UnexpectedResultSnafu { + reason: "native histogram column is not a struct array", + }) + }) + .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 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()?; - // 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)) - }; + if value.is_none() && histogram.is_none() { + continue; + } - // assemble rows - 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 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)); + } + for (tag_column, tag_name) in tag_columns.iter().zip(layout.tag_names.iter()) { + if let Some(tag_value) = string_array_value_at_index(tag_column, row_index) + { + tags.push((tag_name.as_str(), tag_value)); } - 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)); - } - } - - 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), - )); + // Borrowed and owned labels hash the same: both + // hash their `str` contents, so this hash serves the + // lookup and the insertion of a new series below. + let hash = buffer.hasher().hash_one(&tags); + let entry = buffer + .raw_entry_mut_v1() + .from_key_hashed_nocheck(hash, &SeriesKeyLookup(tags.as_slice())); + let index = entry.index(); + if let RawEntryMut::Vacant(entry) = entry { + // A series that is new to the buffer takes + // ownership of its labels. + let key = tags + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect::>(); + 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), + )); } } + } + Ok(()) +} - // initialize result to return - let mut result = match result_type { - ValueType::Vector => PromQueryResult::Vector(vec![]), - ValueType::Matrix => PromQueryResult::Matrix(vec![]), - ValueType::Scalar => PromQueryResult::Scalar(None), - ValueType::String => PromQueryResult::String(None), - }; +/// Assemble the per-series samples of `buffer` into the data of a Prometheus +/// HTTP API response of `result_type`. +fn samples_to_data( + buffer: IndexMap, PromSeriesSamples>, + result_type: ValueType, +) -> Result { + // initialize result to return + let mut result = match result_type { + ValueType::Vector => PromQueryResult::Vector(vec![]), + ValueType::Matrix => PromQueryResult::Matrix(vec![]), + ValueType::Scalar => PromQueryResult::Scalar(None), + ValueType::String => PromQueryResult::String(None), + }; - // accumulate data into result - buffer.into_iter().for_each(|(tags, mut samples)| { - let metric = tags - .into_iter() - .map(|(k, v)| (k.to_string(), v.to_string())) - .collect::>(); - match result { - PromQueryResult::Vector(ref mut v) => { - let histogram = samples.histograms.pop(); - let value = if histogram.is_none() { - samples - .values - .pop() - .map(|(timestamp, value)| (timestamp, value.into_string())) - } else { - None - }; - v.push(PromSeriesVector { - metric, - value, - histogram, - }); - } - PromQueryResult::Matrix(ref mut v) => { - // sort values by timestamp - if !samples.values.is_sorted_by(|a, b| a.0 <= b.0) { - samples - .values - .sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(Ordering::Equal)); - } - if !samples.histograms.is_sorted_by(|a, b| a.0 <= b.0) { - samples - .histograms - .sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(Ordering::Equal)); - } - - v.push(PromSeriesMatrix { - metric, - values: samples.values, - histograms: samples.histograms, - }); - } - PromQueryResult::Scalar(ref mut v) => { - *v = samples + // accumulate data into result + buffer.into_iter().for_each(|(tags, mut samples)| { + let metric = tags.into_iter().collect::>(); + match result { + PromQueryResult::Vector(ref mut v) => { + let histogram = samples.histograms.pop(); + let value = if histogram.is_none() { + samples .values .pop() - .map(|(timestamp, value)| (timestamp, value.into_string())); - } - PromQueryResult::String(ref mut _v) => { - // TODO(ruihang): Not supported yet - } + .map(|(timestamp, value)| (timestamp, value.into_string())) + } else { + None + }; + v.push(PromSeriesVector { + metric, + value, + histogram, + }); } - }); + PromQueryResult::Matrix(ref mut v) => { + // sort values by timestamp + if !samples.values.is_sorted_by(|a, b| a.0 <= b.0) { + samples + .values + .sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(Ordering::Equal)); + } + if !samples.histograms.is_sorted_by(|a, b| a.0 <= b.0) { + samples + .histograms + .sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(Ordering::Equal)); + } - // sort matrix by metric - // see: https://prometheus.io/docs/prometheus/3.5/querying/api/#range-vectors - if let PromQueryResult::Matrix(ref mut v) = result { - v.sort_by(|a, b| a.metric.cmp(&b.metric)); + v.push(PromSeriesMatrix { + metric, + values: samples.values, + histograms: samples.histograms, + }); + } + PromQueryResult::Scalar(ref mut v) => { + *v = samples + .values + .pop() + .map(|(timestamp, value)| (timestamp, value.into_string())); + } + PromQueryResult::String(ref mut _v) => { + // TODO(ruihang): Not supported yet + } } + }); - let result_type_string = result_type.to_string(); - let data = PrometheusResponse::PromData(PromData { - result_type: result_type_string, - result, - }); - - Ok(data) + // sort matrix by metric + // see: https://prometheus.io/docs/prometheus/3.5/querying/api/#range-vectors + if let PromQueryResult::Matrix(ref mut v) = result { + v.sort_by(|a, b| a.metric.cmp(&b.metric)); } + + let result_type_string = result_type.to_string(); + let data = PrometheusResponse::PromData(PromData { + result_type: result_type_string, + result, + }); + + Ok(data) } /// Ranges of consecutive rows whose label values are all equal. @@ -667,6 +768,163 @@ mod tests { use super::*; + #[tokio::test] + async fn empty_stream_reads_as_empty_data() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), false), + ])); + let stream = RecordBatches::try_new(schema, vec![]).unwrap().as_stream(); + + let response = + PrometheusJsonResponse::consume_stream_to_data(stream, None, ValueType::Matrix).await; + + assert!(matches!( + response.unwrap(), + PrometheusResponse::PromData(PromData { + result: PromQueryResult::Matrix(series), + .. + }) if series.is_empty() + )); + } + + #[tokio::test] + async fn stream_result_matches_record_batches_result() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), false), + ])); + // Two batches with the same series split across them, so streaming + // merges the samples of one series across a batch boundary. + let mut batches = Vec::new(); + for (batch_index, host) in [(0, "a"), (1, "b")].into_iter() { + let timestamps = (0..4).map(|i| (batch_index * 4 + i) as i64 * 1000); + let values = (0..4).map(|i| (batch_index * 4 + i) as f64); + batches.push( + RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_values(timestamps)) as _, + Arc::new(StringVector::from(vec![Some(host); 4])) as _, + Arc::new(Float64Vector::from_values(values)) as _, + ], + ) + .unwrap(), + ); + } + + let streamed = PrometheusJsonResponse::consume_stream_to_data( + RecordBatches::try_new(schema.clone(), batches.clone()) + .unwrap() + .as_stream(), + Some("metric".to_string()), + ValueType::Matrix, + ) + .await + .unwrap(); + let eager = PrometheusJsonResponse::record_batches_to_data( + RecordBatches::try_new(schema, batches).unwrap(), + Some("metric".to_string()), + ValueType::Matrix, + ) + .unwrap(); + + assert_eq!( + serde_json::to_value(&streamed).unwrap(), + serde_json::to_value(&eager).unwrap() + ); + let PrometheusResponse::PromData(PromData { + result: PromQueryResult::Matrix(series), + .. + }) = streamed + else { + panic!("expected matrix response"); + }; + assert_eq!(series.len(), 2); + } + + #[tokio::test] + async fn stream_mixed_histogram_result_matches_record_batches_result() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ColumnSchema::new("kind", ConcreteDataType::string_datatype(), false), + ColumnSchema::new("float", ConcreteDataType::float64_datatype(), true), + ColumnSchema::new("histogram", native_histogram_value_type().clone(), true), + ])); + // The float series splits across the batch boundary so streaming must + // merge its samples like the eager path; the histogram row only arrives + // in the second batch. + let first = RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_values([1_000, 2_000])) as _, + Arc::new(StringVector::from(vec![Some("float"); 2])) as _, + Arc::new(Float64Vector::from(vec![Some(1.25), Some(0.5)])) as _, + histogram_vector(&[None, None]), + ], + ) + .unwrap(); + let second = RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_values([3_000, 3_000])) as _, + Arc::new(StringVector::from(vec![Some("float"), Some("histogram")])) as _, + Arc::new(Float64Vector::from(vec![Some(2.0), None])) as _, + histogram_vector(&[None, Some(sample_histogram())]), + ], + ) + .unwrap(); + + let streamed = PrometheusJsonResponse::consume_stream_to_data( + RecordBatches::try_new(schema.clone(), vec![first.clone(), second.clone()]) + .unwrap() + .as_stream(), + Some("mixed_metric".to_string()), + ValueType::Vector, + ) + .await + .unwrap(); + let eager = PrometheusJsonResponse::record_batches_to_data( + RecordBatches::try_new(schema, vec![first, second]).unwrap(), + Some("mixed_metric".to_string()), + ValueType::Vector, + ) + .unwrap(); + + assert_eq!( + serde_json::to_value(&streamed).unwrap(), + serde_json::to_value(&eager).unwrap() + ); + let PrometheusResponse::PromData(PromData { + result: PromQueryResult::Vector(series), + .. + }) = streamed + else { + panic!("expected vector response"); + }; + assert_eq!(series.len(), 2); + let histogram = series + .iter() + .find(|series| series.metric["kind"] == "histogram") + .unwrap(); + assert!(histogram.histogram.is_some()); + assert!(histogram.value.is_none()); + } + #[tokio::test] async fn query_response_preserves_and_merges_promql_annotations() { let query_id = "query_response_preserves_and_merges_promql_annotations";