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 {