From 2933f95672be7bd0798690a781e4024f689ea764 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Thu, 10 Sep 2026 04:25:48 +0000 Subject: [PATCH] perf(servers): drop the output sort for Prometheus range queries (#9090) * perf(servers): drop the output sort for Prometheus range queries Matrix responses sort samples by timestamp and series by labels while serializing, so the sort the PromQL planner puts at the root of a range query plan never reaches the client. Mark the query as not requiring output ordering in `do_range_query`, and let the frontend strip the root sort (through projections, and only when it carries no fetch) before execution. Instant queries are unchanged, so `sort()` and `sort_desc()` keep their observable order. Signed-off-by: Dennis Zhuang * test(servers): pin the matrix response against input row order Signed-off-by: Dennis Zhuang * test(servers): drop redundant response assertions in the ordering test Signed-off-by: Dennis Zhuang * perf(servers): drop the output sort for gRPC gateway range queries Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/frontend/src/instance.rs | 7 + src/frontend/src/instance/promql.rs | 93 +++++++++++++ src/servers/src/grpc/prom_query_gateway.rs | 10 +- src/servers/src/http/prometheus.rs | 69 +++++++++- .../src/http/result/prometheus_resp.rs | 130 +++++++++++++++++- src/servers/src/prometheus_handler.rs | 19 ++- 6 files changed, 316 insertions(+), 12 deletions(-) diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index dd89a1f3f9..f834722cca 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -1329,6 +1329,7 @@ impl PrometheusHandler for Instance { self.check_query_target_permission(targets, &query_ctx) .await?; + let requires_output_ordering = query.requires_output_ordering(); let (query, stmt) = query.into_parts(); let QueryStatement::Promql(eval_stmt, _) = &stmt else { @@ -1342,6 +1343,12 @@ impl PrometheusHandler for Instance { .map_err(BoxedError::new) .context(ExecuteQuerySnafu)?; + let plan = if requires_output_ordering { + plan + } else { + promql::remove_output_sort(plan) + }; + interceptor.pre_execute(&query, &eval_stmt.expr, Some(&plan), query_ctx.clone())?; // Take the EvalStmt from the original QueryStatement and use it to create the CatalogQueryStatement. diff --git a/src/frontend/src/instance/promql.rs b/src/frontend/src/instance/promql.rs index 0f6121441f..3df7022470 100644 --- a/src/frontend/src/instance/promql.rs +++ b/src/frontend/src/instance/promql.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; use std::time::SystemTime; use auth::PermissionTableTarget; @@ -21,6 +22,7 @@ use common_catalog::consts::INFORMATION_SCHEMA_NAME; use common_catalog::format_full_table_name; use common_recordbatch::util; use common_telemetry::tracing; +use datafusion_expr::LogicalPlan; use promql_parser::label::{Matcher, Matchers}; use query::promql; use query::promql::planner::PromPlanner; @@ -35,6 +37,21 @@ use crate::error::{ }; use crate::instance::Instance; +/// Strips the output sort a PromQL plan ends with, keeping the plan schema intact. +/// +/// Only the sort the caller would observe is removed: recursion stops at any other +/// node, so ordering consumed by windows, limits or PromQL extension nodes stays. +pub(super) fn remove_output_sort(plan: LogicalPlan) -> LogicalPlan { + match plan { + LogicalPlan::Sort(sort) if sort.fetch.is_none() => Arc::unwrap_or_clone(sort.input), + LogicalPlan::Projection(mut projection) => { + projection.input = Arc::new(remove_output_sort(Arc::unwrap_or_clone(projection.input))); + LogicalPlan::Projection(projection) + } + plan => plan, + } +} + impl Instance { /// Handles metric names query request, returns the names. #[tracing::instrument(skip_all)] @@ -180,3 +197,79 @@ impl Instance { Ok(results) } } + +#[cfg(test)] +mod tests { + use datafusion::arrow::array::Int64Array; + use datafusion::prelude::{SessionConfig, SessionContext}; + use datafusion_expr::{LogicalPlanBuilder, Sort, col, lit}; + + use super::*; + + #[tokio::test] + async fn remove_output_sort_keeps_schema_and_row_selection() { + let input = + LogicalPlanBuilder::values(vec![vec![lit(3_i64)], vec![lit(1_i64)], vec![lit(2_i64)]]) + .unwrap() + .build() + .unwrap(); + let sort = Sort { + expr: vec![col("column1").sort(true, false)], + input: Arc::new(input), + fetch: None, + }; + let projected = LogicalPlanBuilder::from(LogicalPlan::Sort(sort.clone())) + .project(vec![col("column1").alias("sample")]) + .unwrap() + .project(vec![col("sample")]) + .unwrap() + .build() + .unwrap(); + // The root sort is removed, but the sort feeding the limit selects the rows. + let limited = LogicalPlanBuilder::from(LogicalPlan::Sort(sort.clone())) + .limit(0, Some(2)) + .unwrap() + .sort(vec![col("column1").sort(false, false)]) + .unwrap() + .build() + .unwrap(); + let fetched = LogicalPlan::Sort(Sort { + fetch: Some(2), + ..sort.clone() + }); + + let context = + SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); + for (name, plan, expected) in [ + ("root", LogicalPlan::Sort(sort), vec![3, 1, 2]), + ("projection", projected, vec![3, 1, 2]), + ("sort below limit", limited, vec![1, 2]), + ("fetch", fetched, vec![1, 2]), + ] { + let schema = plan.schema().clone(); + let plan = remove_output_sort(plan); + assert_eq!(plan.schema(), &schema, "{name}"); + let output = context + .execute_logical_plan(plan) + .await + .unwrap() + .collect() + .await + .unwrap(); + let values = output + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + assert_eq!(values, expected, "{name}"); + } + } +} diff --git a/src/servers/src/grpc/prom_query_gateway.rs b/src/servers/src/grpc/prom_query_gateway.rs index 99940ad9a2..92bd5d722e 100644 --- a/src/servers/src/grpc/prom_query_gateway.rs +++ b/src/servers/src/grpc/prom_query_gateway.rs @@ -135,11 +135,15 @@ impl PrometheusGatewayService { }; let (metric_name, mut result_type) = retrieve_metric_name_and_result_type(query.expr()); let query_id = ctx.remote_query_id().map(str::to_string); - let result = self.handler.do_query_parsed(query, ctx).await; - // range query only returns matrix - if is_range_query { + // A range query only returns a matrix, and matrix serialization sorts + // samples and series, so execution order never reaches the response. + let query = if is_range_query { result_type = ValueType::Matrix; + query.with_unordered_output() + } else { + query }; + let result = self.handler.do_query_parsed(query, ctx).await; PrometheusJsonResponse::from_query_result( result, diff --git a/src/servers/src/http/prometheus.rs b/src/servers/src/http/prometheus.rs index 2c9d8267ee..1f9988ede4 100644 --- a/src/servers/src/http/prometheus.rs +++ b/src/servers/src/http/prometheus.rs @@ -624,7 +624,11 @@ async fn do_range_query( ) -> PrometheusJsonResponse { let (metric_name, _) = retrieve_metric_name_and_result_type(prom_query.expr()); let query_id = query_ctx.remote_query_id().map(str::to_string); - let result = handler.do_query_parsed(prom_query, query_ctx).await; + // Matrix serialization sorts samples and series, so execution order never + // reaches the response. + let result = handler + .do_query_parsed(prom_query.with_unordered_output(), query_ctx) + .await; PrometheusJsonResponse::from_query_result( result, metric_name, @@ -2404,6 +2408,7 @@ mod tests { denied_table: Option<&'static str>, metric_names: Vec, queries: Mutex>, + ordered_outputs: Mutex>, } #[async_trait::async_trait] @@ -2417,6 +2422,10 @@ mod tests { query: ParsedPromQuery, _: QueryContextRef, ) -> Result { + self.ordered_outputs + .lock() + .unwrap() + .push(query.requires_output_ordering()); self.queries .lock() .unwrap() @@ -2531,6 +2540,59 @@ mod tests { ); } + #[tokio::test] + async fn range_query_does_not_require_execution_output_ordering() { + let handler = Arc::new(TestPrometheusHandler { + catalog_manager: MemoryCatalogManager::new(), + deny_operation: false, + denied_table: None, + metric_names: Vec::new(), + queries: Mutex::new(Vec::new()), + ordered_outputs: Mutex::new(Vec::new()), + }); + let state: PrometheusHandlerRef = handler.clone(); + instant_query( + State(state.clone()), + Query(InstantQuery { + query: Some("sort(vector(1))".to_string()), + time: Some("0".to_string()), + ..Default::default() + }), + Extension(QueryContext::with( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + )), + Form(InstantQuery::default()), + ) + .await; + + // Both a single-point and a multi-step range query take the same path. + for end in ["0", "1"] { + range_query( + State(state.clone()), + Query(RangeQuery { + query: Some("sort(vector(1))".to_string()), + start: Some("0".to_string()), + end: Some(end.to_string()), + step: Some("1s".to_string()), + ..Default::default() + }), + Extension(QueryContext::with( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + )), + Form(RangeQuery::default()), + ) + .await; + } + + // `sort()` stays observable for instant queries, but not for range queries. + assert_eq!( + *handler.ordered_outputs.lock().unwrap(), + vec![true, false, false] + ); + } + #[tokio::test] async fn test_promql_timer_records_parse_errors() { let handler: PrometheusHandlerRef = Arc::new(TestPrometheusHandler { @@ -2539,6 +2601,7 @@ mod tests { denied_table: None, metric_names: Vec::new(), queries: Mutex::new(Vec::new()), + ordered_outputs: Mutex::new(Vec::new()), }); let query_ctx = QueryContext::with("promql_timer_test", "parse_error"); let db = query_ctx.get_db_string(); @@ -2622,6 +2685,7 @@ mod tests { denied_table: Some("denied"), metric_names: Vec::new(), queries: Mutex::new(Vec::new()), + ordered_outputs: Mutex::new(Vec::new()), })), Path(FIELD_NAME_LABEL.to_string()), Extension(query_ctx), @@ -2665,6 +2729,7 @@ mod tests { denied_table: None, metric_names: vec!["cpu_user".to_string(), "cpu_system".to_string()], queries: Mutex::new(Vec::new()), + ordered_outputs: Mutex::new(Vec::new()), }); let state: PrometheusHandlerRef = handler.clone(); let query_ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME); @@ -3571,6 +3636,7 @@ mod tests { denied_table: None, metric_names: Vec::new(), queries: Mutex::new(Vec::new()), + ordered_outputs: Mutex::new(Vec::new()), })), Query(MetadataQuery::default()), Extension(query_ctx.clone()), @@ -3584,6 +3650,7 @@ mod tests { denied_table: Some("denied"), metric_names: Vec::new(), queries: Mutex::new(Vec::new()), + ordered_outputs: Mutex::new(Vec::new()), }); let response = metadata_query( State(handler.clone()), diff --git a/src/servers/src/http/result/prometheus_resp.rs b/src/servers/src/http/result/prometheus_resp.rs index 7ec15a8c1e..3ebec961d4 100644 --- a/src/servers/src/http/result/prometheus_resp.rs +++ b/src/servers/src/http/result/prometheus_resp.rs @@ -356,13 +356,12 @@ impl PrometheusJsonResponse { // Tag order matters, e.g., after sorc and sort_desc, the output order must be kept. let mut buffer = IndexMap::, PromSeriesSamples>::new(); - // Query output is clustered by series (the range plan sorts by series - // key + timestamp), so consecutive rows usually belong to the same - // series. Remember the index of the previous row's entry in `buffer`, - // and reuse it directly when its tags are unchanged. This avoids - // building and hashing the label vector on every row; the worst case - // adds one `Vec` comparison per series transition before falling back - // to the map lookup. + // 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; let schema = batches.schema(); @@ -1024,6 +1023,123 @@ mod tests { assert_eq!(series[2].value, Some((6.0, "6.0".to_string()))); } + #[test] + fn matrix_response_is_independent_of_input_row_order() { + // Range queries run without the plan's output sort, so this function sees + // series interleaved across batches with timestamps out of order. The + // serialized matrix must be the same either way. + 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(), true), + ColumnSchema::new("histogram", native_histogram_value_type().clone(), true), + ])); + let histogram = |sum: f64| NativeHistogram { + sum, + ..sample_histogram() + }; + // timestamp, host, rack, float value, histogram value + type Row = ( + i64, + Option<&'static str>, + Option<&'static str>, + Option, + Option, + ); + let rows: Vec = vec![ + (1_000, Some("a"), Some("r"), Some(1.0), None), + (3_000, Some("a"), Some("r"), Some(3.0), None), + (2_000, Some("a"), Some("r"), Some(2.0), None), + (5_000, Some("a"), None, Some(5.0), None), + (4_000, Some("a"), None, Some(4.0), None), + (7_000, Some(""), None, Some(7.0), None), + (8_000, None, None, Some(8.0), None), + (6_000, None, None, Some(6.0), None), + (2_000, Some("h"), None, None, Some(histogram(20.0))), + (1_000, Some("h"), None, None, Some(histogram(10.0))), + ]; + let matrix = |order: &[usize], splits: &[usize]| { + let batch = RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_vec( + order.iter().map(|&row| rows[row].0).collect(), + )) as _, + Arc::new(StringVector::from( + order.iter().map(|&row| rows[row].1).collect::>(), + )) as _, + Arc::new(StringVector::from( + order.iter().map(|&row| rows[row].2).collect::>(), + )) as _, + Arc::new(Float64Vector::from( + order.iter().map(|&row| rows[row].3).collect::>(), + )) as _, + histogram_vector( + &order + .iter() + .map(|&row| rows[row].4.clone()) + .collect::>(), + ), + ], + ) + .unwrap(); + let mut batches = Vec::new(); + let mut start = 0; + for &end in splits.iter().chain(std::iter::once(&order.len())) { + batches.push(batch.slice(start, end - start).unwrap()); + start = end; + } + let response = PrometheusJsonResponse::record_batches_to_data( + RecordBatches::try_new(schema.clone(), batches).unwrap(), + Some("metric".to_string()), + ValueType::Matrix, + ) + .unwrap(); + let PrometheusResponse::PromData(PromData { + result: PromQueryResult::Matrix(series), + .. + }) = response + else { + panic!("expected matrix response"); + }; + series + }; + + let clustered = matrix(&[0, 2, 1, 4, 3, 5, 7, 6, 9, 8], &[5]); + let interleaved = matrix(&[8, 5, 1, 6, 0, 3, 9, 2, 7, 4], &[3, 6]); + assert_eq!( + serde_json::to_value(&interleaved).unwrap(), + serde_json::to_value(&clustered).unwrap() + ); + + // Pin the canonical arrangement itself, not only its stability. + assert_eq!( + serde_json::to_value(&clustered[..4]).unwrap(), + serde_json::json!([ + {"metric": {"__name__": "metric"}, "values": [[6.0, "6.0"], [8.0, "8.0"]]}, + {"metric": {"__name__": "metric", "host": ""}, "values": [[7.0, "7.0"]]}, + {"metric": {"__name__": "metric", "host": "a"}, + "values": [[4.0, "4.0"], [5.0, "5.0"]]}, + {"metric": {"__name__": "metric", "host": "a", "rack": "r"}, + "values": [[1.0, "1.0"], [2.0, "2.0"], [3.0, "3.0"]]}, + ]) + ); + assert_eq!(clustered[4].metric["host"], "h"); + assert_eq!( + clustered[4] + .histograms + .iter() + .map(|(timestamp, histogram)| (*timestamp, histogram.sum.as_str())) + .collect::>(), + vec![(1.0, "10"), (2.0, "20")] + ); + } + #[test] fn record_batches_to_data_preserves_mixed_float_and_histogram_rows() { let schema = Arc::new(Schema::new(vec![ diff --git a/src/servers/src/prometheus_handler.rs b/src/servers/src/prometheus_handler.rs index d56668d016..4f4b7c5d19 100644 --- a/src/servers/src/prometheus_handler.rs +++ b/src/servers/src/prometheus_handler.rs @@ -40,6 +40,7 @@ pub type PrometheusHandlerRef = Arc; pub struct ParsedPromQuery { query: PromQuery, statement: QueryStatement, + requires_output_ordering: bool, } impl ParsedPromQuery { @@ -51,7 +52,11 @@ impl ParsedPromQuery { query: query.clone(), } })?; - Ok(Self { query, statement }) + Ok(Self { + query, + statement, + requires_output_ordering: true, + }) } /// Returns the original query parameters. @@ -64,6 +69,18 @@ impl ParsedPromQuery { &self.statement } + /// Returns whether the caller observes the query output in execution order. + /// When it does not, the executor may drop the plan's output sort. + pub fn requires_output_ordering(&self) -> bool { + self.requires_output_ordering + } + + /// Marks the output order as irrelevant to the caller. + pub(crate) fn with_unordered_output(mut self) -> Self { + self.requires_output_ordering = false; + self + } + /// Returns the parsed PromQL expression. pub fn expr(&self) -> &PromqlExpr { let QueryStatement::Promql(eval_stmt, _) = &self.statement else {