diff --git a/src/catalog/src/kvbackend/client.rs b/src/catalog/src/kvbackend/client.rs index 69b31d62322..402af1d9865 100644 --- a/src/catalog/src/kvbackend/client.rs +++ b/src/catalog/src/kvbackend/client.rs @@ -192,6 +192,10 @@ impl KvBackend for CachedKvBackend { } } + if miss_keys.is_empty() { + return Ok(BatchGetResponse { kvs }); + } + let batch_get_req = BatchGetRequest::new().with_keys(miss_keys.clone()); let pre_version = self.version(); @@ -503,6 +507,7 @@ mod tests { pub struct SimpleKvBackend { inner_map: DashMap, Vec>, get_execute_times: Arc, + batch_get_execute_times: Arc, } impl TxnService for SimpleKvBackend { @@ -520,6 +525,7 @@ mod tests { } async fn batch_get(&self, req: BatchGetRequest) -> Result { + self.batch_get_execute_times.fetch_add(1, Ordering::SeqCst); let mut kvs = Vec::with_capacity(req.keys.len()); for key in req.keys.iter() { if let Some(kv) = self.get(key).await? { @@ -571,6 +577,7 @@ mod tests { async fn test_cached_kv_backend() { let simple_kv = Arc::new(SimpleKvBackend::default()); let get_execute_times = simple_kv.get_execute_times.clone(); + let batch_get_execute_times = simple_kv.batch_get_execute_times.clone(); let cached_kv = CachedKvBackend::wrap(simple_kv); add_some_vals(&cached_kv).await; @@ -582,9 +589,11 @@ mod tests { assert_eq!(get_execute_times.load(Ordering::SeqCst), 0); for _ in 0..10 { - let _batch_get_resp = cached_kv.batch_get(batch_get_req.clone()).await.unwrap(); + let batch_get_resp = cached_kv.batch_get(batch_get_req.clone()).await.unwrap(); + assert_eq!(batch_get_resp.kvs.len(), 2); assert_eq!(get_execute_times.load(Ordering::SeqCst), 2); + assert_eq!(batch_get_execute_times.load(Ordering::SeqCst), 1); } let batch_get_req = BatchGetRequest { @@ -594,14 +603,29 @@ mod tests { let _batch_get_resp = cached_kv.batch_get(batch_get_req.clone()).await.unwrap(); assert_eq!(get_execute_times.load(Ordering::SeqCst), 3); + assert_eq!(batch_get_execute_times.load(Ordering::SeqCst), 2); for _ in 0..10 { - let _batch_get_resp = cached_kv.batch_get(batch_get_req.clone()).await.unwrap(); + let batch_get_resp = cached_kv.batch_get(batch_get_req.clone()).await.unwrap(); + assert_eq!(batch_get_resp.kvs.len(), 3); assert_eq!(get_execute_times.load(Ordering::SeqCst), 3); + assert_eq!(batch_get_execute_times.load(Ordering::SeqCst), 2); } } + #[tokio::test] + async fn test_cached_kv_backend_empty_batch_get() { + let simple_kv = Arc::new(SimpleKvBackend::default()); + let cached_kv = CachedKvBackend::wrap(simple_kv.clone()); + + let response = cached_kv.batch_get(BatchGetRequest::new()).await.unwrap(); + + assert!(response.kvs.is_empty()); + assert_eq!(simple_kv.batch_get_execute_times.load(Ordering::SeqCst), 0); + assert_eq!(simple_kv.get_execute_times.load(Ordering::SeqCst), 0); + } + #[tokio::test] async fn test_cached_kv_backend_rejects_writes_with_read_only_inner() { let inner = Arc::new(MemoryKvBackend::::new()); diff --git a/src/common/function/src/aggrs/aggr_wrapper.rs b/src/common/function/src/aggrs/aggr_wrapper.rs index c2e8b567caa..4e638756509 100644 --- a/src/common/function/src/aggrs/aggr_wrapper.rs +++ b/src/common/function/src/aggrs/aggr_wrapper.rs @@ -25,7 +25,7 @@ use std::hash::{Hash, Hasher}; use std::sync::Arc; -use arrow::array::{ArrayRef, BooleanArray, StructArray}; +use arrow::array::{ArrayData, ArrayRef, BooleanArray, StructArray, make_array}; use arrow_schema::{FieldRef, Fields}; use common_telemetry::debug; use datafusion::functions_aggregate::all_default_aggregate_functions; @@ -84,6 +84,19 @@ pub fn is_all_aggr_exprs_steppable(aggr_exprs: &[Expr]) -> bool { return false; } + // DataFusion only sorts the input of an aggregate with a hard ordering requirement + // when the requirement is already satisfied or the aggregate has a reverse + // expression (apache/datafusion#25676). The state wrapper has none, so e.g. + // `nth_value(.. ORDER BY ..)` would read unsorted input on datanodes. Ordered-set + // aggregates like `approx_percentile_cont(..) WITHIN GROUP (ORDER BY ..)` are + // exempt: their ORDER BY names the value, and they don't need sorted input. + if !aggr_func.params.order_by.is_empty() + && aggr_func.func.order_sensitivity().hard_requires() + && !aggr_func.func.supports_within_group_clause() + { + return false; + } + // whether the corresponding state function exists in the registry FUNCTION_REGISTRY.is_aggr_func_exist(&aggr_state_func_name(aggr_func.func.name())) } else { @@ -516,43 +529,7 @@ impl StateGroupsAccum { } fn wrap_state_arrays(&self, arrays: Vec) -> datafusion_common::Result { - let array_type = arrays - .iter() - .map(|array| array.data_type().clone()) - .collect::>(); - let expected_type = self - .state_fields - .iter() - .map(|field| field.data_type().clone()) - .collect::>(); - if array_type != expected_type { - debug!( - "State mismatch, expected: {}, got: {} for expected fields: {:?} and given array types: {:?}", - self.state_fields.len(), - arrays.len(), - self.state_fields, - array_type, - ); - let guess_schema = arrays - .iter() - .enumerate() - .map(|(index, array)| { - Field::new( - format!("col_{index}[mismatch_state]").as_str(), - array.data_type().clone(), - true, - ) - }) - .collect::(); - let array = StructArray::try_new(guess_schema, arrays, None)?; - return Ok(Arc::new(array)); - } - - Ok(Arc::new(StructArray::try_new( - self.state_fields.clone(), - arrays, - None, - )?)) + Ok(Arc::new(state_struct_array(&self.state_fields, arrays)?)) } } @@ -605,6 +582,81 @@ impl GroupsAccumulator for StateGroupsAccum { } } +/// Wraps the state arrays of an accumulator into a struct of the declared state fields. +/// +/// The declared state type is derived from logical expressions, while the accumulator names +/// nested fields after physical expressions. For example, `array_agg(v ORDER BY ts)` declares +/// its orderings as `List(Struct("ts": ..))` but produces `List(Struct("ts@0": ..))`. Arrays that +/// differ only in nested field names are cast to the declared type; any other difference is an +/// error. +fn state_struct_array( + state_fields: &Fields, + arrays: Vec, +) -> datafusion_common::Result { + if arrays.len() != state_fields.len() { + return Err(datafusion_common::DataFusionError::Internal(format!( + "Expected {} state arrays for fields {:?}, got {}", + state_fields.len(), + state_fields, + arrays.len() + ))); + } + let arrays = arrays + .into_iter() + .zip(state_fields.iter()) + .map(|(array, field)| { + let expected = field.data_type(); + if array.data_type() == expected { + Ok(array) + } else if array.data_type().equals_datatype(expected) { + Ok(make_array(relabel_nested_fields( + array.to_data(), + expected, + )?)) + } else { + Err(datafusion_common::DataFusionError::Internal(format!( + "State field `{}` expects type {expected}, but the accumulator produced {}", + field.name(), + array.data_type() + ))) + } + }) + .collect::>>()?; + Ok(StructArray::try_new(state_fields.clone(), arrays, None)?) +} + +/// Rebuilds `data` with the type `target`, which must match it position by position apart +/// from nested field names and metadata. Unlike a cast, children are never matched by name. +fn relabel_nested_fields( + data: ArrayData, + target: &DataType, +) -> datafusion_common::Result { + let child_types = match target { + DataType::List(field) + | DataType::LargeList(field) + | DataType::FixedSizeList(field, _) + | DataType::Map(field, _) => vec![field.data_type()], + DataType::Struct(fields) => fields.iter().map(|f| f.data_type()).collect(), + _ if data.child_data().is_empty() => vec![], + _ => { + return Err(datafusion_common::DataFusionError::NotImplemented(format!( + "Relabeling nested fields of {target}" + ))); + } + }; + let children = data + .child_data() + .iter() + .zip(child_types) + .map(|(child, child_type)| relabel_nested_fields(child.clone(), child_type)) + .collect::>>()?; + Ok(data + .into_builder() + .data_type(target.clone()) + .child_data(children) + .build()?) +} + impl StateAccum { pub fn new( inner: Box, @@ -631,40 +683,7 @@ impl Accumulator for StateAccum { .iter() .map(|s| s.to_array()) .collect::, _>>()?; - let array_type = array - .iter() - .map(|a| a.data_type().clone()) - .collect::>(); - let expected_type: Vec<_> = self - .state_fields - .iter() - .map(|f| f.data_type().clone()) - .collect(); - if array_type != expected_type { - debug!( - "State mismatch, expected: {}, got: {} for expected fields: {:?} and given array types: {:?}", - self.state_fields.len(), - array.len(), - self.state_fields, - array_type, - ); - let guess_schema = array - .iter() - .enumerate() - .map(|(index, array)| { - Field::new( - format!("col_{index}[mismatch_state]").as_str(), - array.data_type().clone(), - true, - ) - }) - .collect::(); - let arr = StructArray::try_new(guess_schema, array, None)?; - - return Ok(ScalarValue::Struct(Arc::new(arr))); - } - - let struct_array = StructArray::try_new(self.state_fields.clone(), array, None)?; + let struct_array = state_struct_array(&self.state_fields, array)?; Ok(ScalarValue::Struct(Arc::new(struct_array))) } diff --git a/src/common/function/src/aggrs/aggr_wrapper/fix_order.rs b/src/common/function/src/aggrs/aggr_wrapper/fix_order.rs index 35aa35b5837..2294f8be031 100644 --- a/src/common/function/src/aggrs/aggr_wrapper/fix_order.rs +++ b/src/common/function/src/aggrs/aggr_wrapper/fix_order.rs @@ -154,13 +154,17 @@ fn rewrite_expr( if is_fix { // then always fix the ordering field&distinct flag and more let order_by = aggregate_function.params.order_by.clone(); + // DataFusion always makes ordering fields nullable when building the physical + // aggregate (`ordering_fields` in `datafusion-functions-aggregate-common`), and some + // accumulators (e.g. `array_agg`) nest these fields in their state type. Keep the + // declared state type in line with what the accumulator produces. let ordering_fields: Vec<_> = order_by .iter() .map(|sort_expr| { sort_expr .expr .to_field(&aggregate_input.schema()) - .map(|(_, f)| f) + .map(|(_, f)| Arc::new(f.as_ref().clone().with_nullable(true))) }) .collect::>>()?; let distinct = aggregate_function.params.distinct; diff --git a/src/common/function/src/aggrs/aggr_wrapper/tests.rs b/src/common/function/src/aggrs/aggr_wrapper/tests.rs index 3f59139f05b..e51caacc35b 100644 --- a/src/common/function/src/aggrs/aggr_wrapper/tests.rs +++ b/src/common/function/src/aggrs/aggr_wrapper/tests.rs @@ -27,8 +27,11 @@ use common_telemetry::init_default_ut_logging; use datafusion::catalog::{Session, TableProvider}; use datafusion::datasource::DefaultTableSource; use datafusion::execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext}; +use datafusion::functions_aggregate::approx_percentile_cont::approx_percentile_cont_udaf; +use datafusion::functions_aggregate::array_agg::array_agg_udaf; use datafusion::functions_aggregate::average::avg_udaf; use datafusion::functions_aggregate::count::count_udaf; +use datafusion::functions_aggregate::nth_value::nth_value_udaf; use datafusion::functions_aggregate::sum::sum_udaf; use datafusion::optimizer::AnalyzerRule; use datafusion::optimizer::analyzer::type_coercion::TypeCoercion; @@ -1312,6 +1315,40 @@ async fn test_udaf_correct_eval_result() { order_by: vec![], null_treatment: None, }, + // The ordering state of `array_agg` nests the ORDER BY fields in `List(Struct(..))`. + TestCase { + func: array_agg_udaf(), + input_schema: Arc::new(arrow_schema::Schema::new(vec![ + Field::new("number", DataType::Float64, true), + Field::new( + "ts", + DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, None), + false, + ), + ])), + args: vec![Expr::Column(Column::new_unqualified("number"))], + input: vec![ + Arc::new(Float64Array::from(vec![Some(3.), Some(1.), Some(2.)])), + Arc::new(TimestampMillisecondArray::from(vec![3000, 1000, 2000])), + ], + expected_output: Some(ScalarValue::List(ScalarValue::new_list_nullable( + &[ + ScalarValue::Float64(Some(1.)), + ScalarValue::Float64(Some(2.)), + ScalarValue::Float64(Some(3.)), + ], + &DataType::Float64, + ))), + expected_fn: None, + distinct: false, + filter: None, + order_by: vec![SortExpr::new( + Expr::Column(Column::new_unqualified("ts")), + true, + true, + )], + null_treatment: None, + }, // TODO(discord9): udd_merge/hll_merge/geo_path/quantile_aggr tests ]; let test_table_ref = TableReference::bare("TestTable"); @@ -1421,3 +1458,89 @@ async fn execute_phy_plan( } Ok(batches) } + +#[test] +fn test_state_struct_array_rejects_mismatched_state() { + let fields = Fields::from(vec![Field::new("sum", DataType::Int64, true)]); + let arrays: Vec = vec![Arc::new(Float64Array::from(vec![1.0]))]; + let err = state_struct_array(&fields, arrays).unwrap_err(); + assert!( + err.to_string() + .contains("State field `sum` expects type Int64, but the accumulator produced Float64"), + "{err}" + ); +} + +#[test] +fn test_state_struct_array_keeps_child_order() { + let int_field = |name: &str| Field::new(name, DataType::Int64, true); + // Same child types, names crossed: children must stay in place, not be matched by name. + let produced = StructArray::from(vec![ + ( + Arc::new(int_field("a")), + Arc::new(Int64Array::from(vec![10])) as ArrayRef, + ), + ( + Arc::new(int_field("b")), + Arc::new(Int64Array::from(vec![20])) as ArrayRef, + ), + ]); + let declared = Fields::from(vec![Field::new( + "state", + DataType::Struct(Fields::from(vec![int_field("b"), int_field("a")])), + true, + )]); + + let state = state_struct_array(&declared, vec![Arc::new(produced)]).unwrap(); + let state = state.column(0).as_struct(); + assert_eq!( + state + .column_by_name("b") + .unwrap() + .as_primitive::() + .value(0), + 10 + ); + assert_eq!( + state + .column_by_name("a") + .unwrap() + .as_primitive::() + .value(0), + 20 + ); +} + +#[test] +fn test_hard_ordered_aggr_not_steppable() { + let order_by = vec![SortExpr::new( + Expr::Column(Column::new_unqualified("ts")), + true, + true, + )]; + let aggr = |func: Arc, args: Vec| { + Expr::AggregateFunction(AggregateFunction::new_udf( + func, + args, + false, + None, + order_by.clone(), + None, + )) + }; + let number = Expr::Column(Column::new_unqualified("number")); + + assert!(!is_all_aggr_exprs_steppable(&[aggr( + nth_value_udaf(), + vec![number.clone(), lit(2i64)], + )])); + assert!(is_all_aggr_exprs_steppable(&[aggr( + array_agg_udaf(), + vec![number.clone()] + )])); + // WITHIN GROUP (ORDER BY number) + assert!(is_all_aggr_exprs_steppable(&[aggr( + approx_percentile_cont_udaf(), + vec![number, lit(0.5f64)] + )])); +} diff --git a/src/common/function/src/aggrs/geo/encoding.rs b/src/common/function/src/aggrs/geo/encoding.rs index 35a59d1ecfb..67dfd493a77 100644 --- a/src/common/function/src/aggrs/geo/encoding.rs +++ b/src/common/function/src/aggrs/geo/encoding.rs @@ -221,18 +221,24 @@ impl DfAccumulator for JsonEncodePathAccumulator { ))); } - for state in states { - let state = as_struct_array(state)?; - let lat_list = as_list_array(state.column(0))?.value(0); - let lat_array = as_primitive_array::(&lat_list)?; - let lng_list = as_list_array(state.column(1))?.value(0); - let lng_array = as_primitive_array::(&lng_list)?; - let ts_list = as_list_array(state.column(2))?.value(0); - let ts_array = as_primitive_array::(&ts_list)?; + let state = as_struct_array(&states[0])?; + let lat_lists = as_list_array(state.column(0))?; + let lng_lists = as_list_array(state.column(1))?; + let ts_lists = as_list_array(state.column(2))?; + for row in 0..state.len() { + if state.is_null(row) { + continue; + } + let lat_list = lat_lists.value(row); + let lng_list = lng_lists.value(row); + let ts_list = ts_lists.value(row); - self.lat.extend(lat_array); - self.lng.extend(lng_array); - self.timestamp.extend(ts_array); + self.lat + .extend(as_primitive_array::(&lat_list)?); + self.lng + .extend(as_primitive_array::(&lng_list)?); + self.timestamp + .extend(as_primitive_array::(&ts_list)?); } Ok(()) @@ -331,9 +337,9 @@ mod tests { _ => panic!("Expected Struct scalar value"), }; - // Merge state arrays - merged.merge_batch(&[state_array1]).unwrap(); - merged.merge_batch(&[state_array2]).unwrap(); + // States from different partitions arrive as rows of one batch + let states = compute::concat(&[state_array1.as_ref(), state_array2.as_ref()]).unwrap(); + merged.merge_batch(&[states]).unwrap(); // Evaluate merged result let result = merged.evaluate().unwrap(); diff --git a/src/common/function/src/aggrs/geo/geo_path.rs b/src/common/function/src/aggrs/geo/geo_path.rs index b09745eaeb7..c4e366c83b0 100644 --- a/src/common/function/src/aggrs/geo/geo_path.rs +++ b/src/common/function/src/aggrs/geo/geo_path.rs @@ -250,18 +250,24 @@ impl DfAccumulator for GeoPathAccumulator { ))); } - for state in states { - let state = as_struct_array(state)?; - let lat_list = as_list_array(state.column(0))?.value(0); - let lat_array = as_primitive_array::(&lat_list)?; - let lng_list = as_list_array(state.column(1))?.value(0); - let lng_array = as_primitive_array::(&lng_list)?; - let ts_list = as_list_array(state.column(2))?.value(0); - let ts_array = as_primitive_array::(&ts_list)?; + let state = as_struct_array(&states[0])?; + let lat_lists = as_list_array(state.column(0))?; + let lng_lists = as_list_array(state.column(1))?; + let ts_lists = as_list_array(state.column(2))?; + for row in 0..state.len() { + if state.is_null(row) { + continue; + } + let lat_list = lat_lists.value(row); + let lng_list = lng_lists.value(row); + let ts_list = ts_lists.value(row); - self.lat.extend(lat_array); - self.lng.extend(lng_array); - self.timestamp.extend(ts_array); + self.lat + .extend(as_primitive_array::(&lat_list)?); + self.lng + .extend(as_primitive_array::(&lng_list)?); + self.timestamp + .extend(as_primitive_array::(&ts_list)?); } Ok(()) @@ -403,9 +409,9 @@ mod tests { _ => panic!("Expected Struct scalar value"), }; - // Merge state arrays - merged.merge_batch(&[state_array1]).unwrap(); - merged.merge_batch(&[state_array2]).unwrap(); + // States from different partitions arrive as rows of one batch + let states = compute::concat(&[state_array1.as_ref(), state_array2.as_ref()]).unwrap(); + merged.merge_batch(&[states]).unwrap(); // Evaluate merged result let result = merged.evaluate().unwrap(); diff --git a/src/frontend/src/error.rs b/src/frontend/src/error.rs index 556852caa0c..451040c884a 100644 --- a/src/frontend/src/error.rs +++ b/src/frontend/src/error.rs @@ -174,6 +174,15 @@ pub enum Error { location: Location, }, + #[snafu(display("Unexpected type {data_type} for column '{column}' of table '{table_name}'"))] + UnexpectedColumnType { + table_name: String, + column: String, + data_type: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to collect recordbatch"))] CollectRecordbatch { #[snafu(implicit)] @@ -422,7 +431,9 @@ impl ErrorExt for Error { Error::RequestQuery { source, .. } => source.status_code(), - Error::CacheRequired { .. } => StatusCode::Internal, + Error::CacheRequired { .. } | Error::UnexpectedColumnType { .. } => { + StatusCode::Internal + } Error::TableNotFound { .. } => StatusCode::TableNotFound, diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index f834722cca5..609e6879a9a 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -1509,6 +1509,20 @@ impl PrometheusHandler for Instance { .context(ExecuteQuerySnafu) } + async fn query_metric_names_by_labels( + &self, + matchers: Vec, + schema: &str, + start: SystemTime, + end: SystemTime, + ctx: &QueryContextRef, + ) -> server_error::Result> { + self.handle_query_metric_names_by_labels(matchers, schema, start, end, ctx) + .await + .map_err(BoxedError::new) + .context(ExecuteQuerySnafu) + } + async fn query_label_values( &self, metric: String, diff --git a/src/frontend/src/instance/otlp/trace_types.rs b/src/frontend/src/instance/otlp/trace_types.rs index 0295d12e685..6f37f14e34a 100644 --- a/src/frontend/src/instance/otlp/trace_types.rs +++ b/src/frontend/src/instance/otlp/trace_types.rs @@ -519,31 +519,34 @@ mod tests { ); } - /// The PHP instrumentation case: an empty string must be distinguishable - /// from any other unparsable value in the reported diagnostics. + /// The PHP instrumentation case: `getHeaderLine('Content-Length')` sends "" + /// when the header is absent. The span must be kept with a null attribute. #[test] - fn test_prepare_trace_column_rewrites_reports_empty_string_value() { - let rows = vec![Row { - values: vec![Value { - value_data: Some(ValueData::StringValue(String::new())), + fn test_prepare_trace_column_rewrites_nulls_empty_string_value() { + let mut rows = Rows { + schema: vec![ColumnSchema { + column_name: "span_attributes.http.response.body.size".to_string(), + datatype: ColumnDataType::String as i32, + ..Default::default() }], - }]; + rows: vec![Row { + values: vec![Value { + value_data: Some(ValueData::StringValue(String::new())), + }], + }], + }; let pending_rewrites = vec![PendingTraceColumnRewrite { col_idx: 0, target_type: ColumnDataType::Int64, column_name: "span_attributes.http.response.body.size".to_string(), }]; - let err = prepare_trace_column_rewrites(&rows, pending_rewrites, "opentelemetry_traces") - .unwrap_err(); - assert!( - err.error.to_string().contains( - "'span_attributes.http.response.body.size' in table 'opentelemetry_traces' \ - from String(\"\") to Int64" - ), - "unexpected error message: {}", - err.error - ); + prepare_trace_column_rewrites(&rows.rows, pending_rewrites, "opentelemetry_traces") + .unwrap() + .apply(&mut rows); + + assert_eq!(rows.schema[0].datatype, ColumnDataType::Int64 as i32); + assert_eq!(rows.rows[0].values[0].value_data, None); } #[test] diff --git a/src/frontend/src/instance/promql.rs b/src/frontend/src/instance/promql.rs index 3df70224703..3ad9ae217ae 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::collections::HashSet; use std::sync::Arc; use std::time::SystemTime; @@ -20,20 +21,25 @@ use catalog::information_schema::TABLES; use client::OutputData; use common_catalog::consts::INFORMATION_SCHEMA_NAME; use common_catalog::format_full_table_name; -use common_recordbatch::util; +use common_recordbatch::{RecordBatch, util}; use common_telemetry::tracing; use datafusion_expr::LogicalPlan; +use datatypes::arrow::array::{Array, UInt32Array}; +use futures::StreamExt; use promql_parser::label::{Matcher, Matchers}; use query::promql; use query::promql::planner::PromPlanner; use servers::prometheus; use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt}; +use store_api::metric_engine_consts::DATA_SCHEMA_TABLE_ID_COLUMN_NAME; +use store_api::storage::TableId; +use table::TableRef; use crate::error::{ CatalogSnafu, CollectRecordbatchSnafu, ExecLogicalPlanSnafu, PrometheusLabelValuesQueryPlanSnafu, PrometheusMetricNamesQueryPlanSnafu, ReadTableSnafu, - Result, TableNotFoundSnafu, TableSnafu, + Result, TableNotFoundSnafu, TableSnafu, UnexpectedColumnTypeSnafu, }; use crate::instance::Instance; @@ -117,6 +123,173 @@ impl Instance { Ok(results) } + /// Handles a metric names query constrained by matchers on ordinary labels. + /// + /// [`Instance::handle_query_metric_names`] answers from table metadata, which + /// cannot resolve a label matcher: whether a metric carries `pod="abc"` is a + /// property of its data. The metric engine multiplexes the logical tables of + /// one physical table into a region holding the union of their label columns + /// alongside `__table_id`, so a single distinct scan per physical table + /// resolves the matchers for all of its logical tables at once. + /// + /// Only metric engine tables are covered. Tables on other engines share no + /// column space to scan, and a scan per table does not scale to the table + /// counts this API is expected to answer over. + #[tracing::instrument(skip_all)] + pub(crate) async fn handle_query_metric_names_by_labels( + &self, + matchers: Vec, + schema: &str, + start: SystemTime, + end: SystemTime, + ctx: &QueryContextRef, + ) -> Result> { + let _timer = crate::metrics::PROMQL_QUERY_METRICS_ELAPSED + .with_label_values(&[ctx.get_db_string().as_str()]) + .start_timer(); + + let catalog = ctx.current_catalog(); + let mut table_ids = HashSet::new(); + for physical in self.physical_metric_tables(catalog, schema, ctx).await? { + table_ids.extend( + self.scan_matching_table_ids(physical, &matchers, start, end, ctx) + .await?, + ); + } + + // Batch-resolve only the ids the scan produced. An id dropped between the + // scan and here simply has no entry. + let table_ids = table_ids.into_iter().collect::>(); + let mut names = self + .catalog_manager + .tables_by_ids(catalog, schema, &table_ids) + .await + .context(CatalogSnafu)? + .into_iter() + .map(|table| table.table_info().name.clone()) + .collect::>(); + names.sort_unstable(); + Ok(names) + } + + /// Returns the metric engine physical tables of a schema. + /// + /// Their data regions carry the union of their logical tables' label columns, + /// so scanning these covers every metric of the schema. + async fn physical_metric_tables( + &self, + catalog: &str, + schema: &str, + ctx: &QueryContextRef, + ) -> Result> { + let mut tables = self.catalog_manager.tables(catalog, schema, Some(ctx)); + let mut physical_tables = Vec::new(); + + while let Some(table) = tables.next().await { + let table = table.context(CatalogSnafu)?; + if table.table_info().is_physical_table() { + physical_tables.push(table); + } + } + + Ok(physical_tables) + } + + /// Scans `table` for the distinct values of `column` that match `matchers` + /// within the time range. + /// + /// Callers look the table up and decode the batches; the plan between is the + /// same whether the projected column is a label or `__table_id`. + async fn scan_distinct_column( + &self, + table: TableRef, + matchers: Vec, + column: String, + start: SystemTime, + end: SystemTime, + ctx: &QueryContextRef, + ) -> Result> { + let dataframe = self + .query_engine + .read_table(table.clone()) + .with_context(|_| ReadTableSnafu { + table_name: table.table_info().full_table_name(), + })?; + + let scan_plan = dataframe.into_unoptimized_plan(); + let conditions = PromPlanner::matchers_to_expr(Matchers::new(matchers), scan_plan.schema()) + .context(PrometheusLabelValuesQueryPlanSnafu)?; + let logical_plan = promql::label_values::rewrite_label_values_query( + table, scan_plan, conditions, column, start, end, + ) + .context(PrometheusLabelValuesQueryPlanSnafu)?; + + let results = self + .query_engine + .execute(logical_plan, ctx.clone()) + .await + .context(ExecLogicalPlanSnafu)?; + + match results.data { + OutputData::Stream(stream) => { + util::collect(stream).await.context(CollectRecordbatchSnafu) + } + OutputData::RecordBatches(rbs) => Ok(rbs.take()), + _ => unreachable!("should not happen"), + } + } + + /// Returns the `__table_id`s of `physical` carrying a row that matches every + /// matcher within the time range. + async fn scan_matching_table_ids( + &self, + physical: TableRef, + matchers: &[Matcher], + start: SystemTime, + end: SystemTime, + ctx: &QueryContextRef, + ) -> Result> { + // `__table_id` attributes a row to its logical table, and a physical + // table that never took a column from one does not expose it. Skipping + // such a table can only miss a metric carrying no labels at all. + if physical + .schema() + .column_schema_by_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME) + .is_none() + { + return Ok(Vec::new()); + } + + let table_name = physical.table_info().full_table_name(); + let batches = self + .scan_distinct_column( + physical, + matchers.to_vec(), + DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(), + start, + end, + ctx, + ) + .await?; + + let mut table_ids = Vec::new(); + for batch in batches { + // Only one column in results, ensured by `rewrite_label_values_query`. + let column = batch.column(0); + let ids = column + .as_any() + .downcast_ref::() + .with_context(|| UnexpectedColumnTypeSnafu { + table_name: table_name.clone(), + column: DATA_SCHEMA_TABLE_ID_COLUMN_NAME, + data_type: column.data_type().to_string(), + })?; + table_ids.extend(ids.iter().flatten()); + } + + Ok(table_ids) + } + /// Handles label values query request, returns the values. #[tracing::instrument(skip_all)] pub(crate) async fn handle_query_label_values( @@ -150,40 +323,9 @@ impl Instance { .context(TableSnafu); } - let dataframe = self - .query_engine - .read_table(table.clone()) - .with_context(|_| ReadTableSnafu { - table_name: full_table_name, - })?; - - let scan_plan = dataframe.into_unoptimized_plan(); - let filter_conditions = - PromPlanner::matchers_to_expr(Matchers::new(matchers), scan_plan.schema()) - .context(PrometheusLabelValuesQueryPlanSnafu)?; - let logical_plan = promql::label_values::rewrite_label_values_query( - table, - scan_plan, - filter_conditions, - label_name, - start, - end, - ) - .context(PrometheusLabelValuesQueryPlanSnafu)?; - - let results = self - .query_engine - .execute(logical_plan, ctx.clone()) - .await - .context(ExecLogicalPlanSnafu)?; - - let batches = match results.data { - OutputData::Stream(stream) => util::collect(stream) - .await - .context(CollectRecordbatchSnafu)?, - OutputData::RecordBatches(rbs) => rbs.take(), - _ => unreachable!("should not happen"), - }; + let batches = self + .scan_distinct_column(table, matchers, label_name, start, end, ctx) + .await?; let mut results = Vec::with_capacity(batches.iter().map(|b| b.num_rows()).sum()); for batch in batches { diff --git a/src/meta-srv/AGENTS.md b/src/meta-srv/AGENTS.md index e1e23581525..5f9f573978c 100644 --- a/src/meta-srv/AGENTS.md +++ b/src/meta-srv/AGENTS.md @@ -38,7 +38,11 @@ abstraction, election, key encoding, and the DDL manager live in `common-meta`; - **Heartbeat**: `service/heartbeat.rs` receives datanode/frontend heartbeats and runs the handler chain in `handler/`; responses carry mailbox messages (DDL - results, cache invalidation). + results, cache invalidation). Node information is refreshed in memory on each + heartbeat. Datanode statistics collection and flownode remapping share + `handler/utils.rs` for per-node address persistence. Datanode address + updates retain the statistics handler's per-node ordering. Address epochs are + confirmed only after a successful save; cache invalidation remains best effort. - **DDL**: `service/procedure.rs` (`ddl`) is leader-only and hands off to `common-meta`'s `DdlManager`, executed via `common-procedure`. - **Region migration**: `procedure/region_migration/manager.rs` plus per-step diff --git a/src/meta-srv/src/handler.rs b/src/meta-srv/src/handler.rs index b2fd83b9e96..d974e4f2580 100644 --- a/src/meta-srv/src/handler.rs +++ b/src/meta-srv/src/handler.rs @@ -81,6 +81,7 @@ pub mod publish_heartbeat_handler; pub mod region_lease_handler; pub mod remap_flow_peer_handler; pub mod response_header_handler; +mod utils; #[cfg(test)] pub mod test_utils; diff --git a/src/meta-srv/src/handler/collect_cluster_info_handler.rs b/src/meta-srv/src/handler/collect_cluster_info_handler.rs index 6bee86ad711..12fee395136 100644 --- a/src/meta-srv/src/handler/collect_cluster_info_handler.rs +++ b/src/meta-srv/src/handler/collect_cluster_info_handler.rs @@ -46,30 +46,21 @@ impl HeartbeatHandler for CollectFrontendClusterInfoHandler { ctx: &mut Context, _acc: &mut HeartbeatAccumulator, ) -> Result { - let Some((key, peer, info, env_vars)) = extract_base_info(req) else { + let Some(update) = NodeInfoUpdate::from_request(req) else { return Ok(HandleControl::Continue); }; let frontend_workloads = get_frontend_workloads(req.node_workloads.as_ref()); - let value = NodeInfo { - peer, - last_activity_ts: common_time::util::current_time_millis(), - status: NodeStatus::Frontend(FrontendStatus { - workloads: frontend_workloads, - }), - version: info.version, - git_commit: info.git_commit, - start_time_ms: info.start_time_ms, - total_cpu_millicores: info.total_cpu_millicores, - total_memory_bytes: info.total_memory_bytes, - cpu_usage_millicores: info.cpu_usage_millicores, - memory_usage_bytes: info.memory_usage_bytes, - hostname: info.hostname, - env_vars, - }; - - put_into_memory_store(ctx, key, value).await?; + update + .save_to_memory( + ctx, + common_time::util::current_time_millis(), + NodeStatus::Frontend(FrontendStatus { + workloads: frontend_workloads, + }), + ) + .await?; Ok(HandleControl::Continue) } @@ -89,29 +80,20 @@ impl HeartbeatHandler for CollectFlownodeClusterInfoHandler { ctx: &mut Context, _acc: &mut HeartbeatAccumulator, ) -> Result { - let Some((key, peer, info, env_vars)) = extract_base_info(req) else { + let Some(update) = NodeInfoUpdate::from_request(req) else { return Ok(HandleControl::Continue); }; let flownode_workloads = get_flownode_workloads(req.node_workloads.as_ref()); - let value = NodeInfo { - peer, - last_activity_ts: common_time::util::current_time_millis(), - status: NodeStatus::Flownode(FlownodeStatus { - workloads: flownode_workloads, - }), - version: info.version, - git_commit: info.git_commit, - start_time_ms: info.start_time_ms, - total_cpu_millicores: info.total_cpu_millicores, - total_memory_bytes: info.total_memory_bytes, - cpu_usage_millicores: info.cpu_usage_millicores, - memory_usage_bytes: info.memory_usage_bytes, - hostname: info.hostname, - env_vars, - }; - - put_into_memory_store(ctx, key, value).await?; + update + .save_to_memory( + ctx, + common_time::util::current_time_millis(), + NodeStatus::Flownode(FlownodeStatus { + workloads: flownode_workloads, + }), + ) + .await?; Ok(HandleControl::Continue) } @@ -132,7 +114,7 @@ impl HeartbeatHandler for CollectDatanodeClusterInfoHandler { ctx: &mut Context, acc: &mut HeartbeatAccumulator, ) -> Result { - let Some((key, peer, info, env_vars)) = extract_base_info(req) else { + let Some(update) = NodeInfoUpdate::from_request(req) else { return Ok(HandleControl::Continue); }; @@ -147,16 +129,64 @@ impl HeartbeatHandler for CollectDatanodeClusterInfoHandler { .count(); let follower_regions = stat.region_stats.len() - leader_regions; - let value = NodeInfo { + update + .save_to_memory( + ctx, + stat.timestamp_millis, + NodeStatus::Datanode(DatanodeStatus { + rcus: stat.rcus, + wcus: stat.wcus, + leader_regions, + follower_regions, + workloads: stat.datanode_workloads.clone(), + }), + ) + .await?; + + Ok(HandleControl::Continue) + } +} + +/// Common fields collected for one heartbeat's in-memory node information +/// update. Persistent routing addresses have a separate lifecycle. +struct NodeInfoUpdate { + key: NodeInfoKey, + peer: Peer, + info: PbNodeInfo, + env_vars: HashMap, +} + +impl NodeInfoUpdate { + fn from_request(request: &HeartbeatRequest) -> Option { + let key = NodeInfoKey::new(request)?; + let peer = request.peer.clone()?; + let info = request.info.clone()?; + let env_vars = EnvVars::from_extensions(&request.extensions) + .inspect_err(|e| { + warn!(e; "Failed to deserialize __env_vars from heartbeat extensions, peer: {}", peer); + }) + .unwrap_or_default() + .map(|e| e.vars) + .unwrap_or_default(); + Some(Self { + key, peer, - last_activity_ts: stat.timestamp_millis, - status: NodeStatus::Datanode(DatanodeStatus { - rcus: stat.rcus, - wcus: stat.wcus, - leader_regions, - follower_regions, - workloads: stat.datanode_workloads.clone(), - }), + info, + env_vars, + }) + } + + async fn save_to_memory( + self, + ctx: &Context, + last_activity_ts: i64, + status: NodeStatus, + ) -> Result<()> { + let info = self.info; + let value = NodeInfo { + peer: self.peer, + last_activity_ts, + status, version: info.version, git_commit: info.git_commit, start_time_ms: info.start_time_ms, @@ -165,53 +195,88 @@ impl HeartbeatHandler for CollectDatanodeClusterInfoHandler { cpu_usage_millicores: info.cpu_usage_millicores, memory_usage_bytes: info.memory_usage_bytes, hostname: info.hostname, - env_vars, + env_vars: self.env_vars, }; - - put_into_memory_store(ctx, key, value).await?; - - Ok(HandleControl::Continue) + let value = value.try_into().context(InvalidClusterInfoFormatSnafu)?; + ctx.in_memory + .put(PutRequest { + key: (&self.key).into(), + value, + ..Default::default() + }) + .await + .context(SaveClusterInfoSnafu)?; + Ok(()) } } -fn extract_base_info( - request: &HeartbeatRequest, -) -> Option<(NodeInfoKey, Peer, PbNodeInfo, HashMap)> { - let HeartbeatRequest { peer, info, .. } = request; - let key = NodeInfoKey::new(request)?; - let Some(peer) = &peer else { - return None; - }; - let Some(info) = &info else { - return None; - }; +#[cfg(test)] +mod tests { + use api::v1::meta::RequestHeader; + use common_meta::datanode::Stat; - let env_vars = EnvVars::from_extensions(&request.extensions) - .inspect_err(|e| { - warn!(e; - "Failed to deserialize __env_vars from heartbeat extensions, peer: {}", peer + use super::*; + use crate::handler::test_utils::TestEnv; + + #[tokio::test] + async fn test_node_info_refreshes_without_epoch_change() { + let mut ctx = TestEnv::new().ctx(); + let handlers: Vec<(Role, Box)> = vec![ + (Role::Frontend, Box::new(CollectFrontendClusterInfoHandler)), + (Role::Datanode, Box::new(CollectDatanodeClusterInfoHandler)), + (Role::Flownode, Box::new(CollectFlownodeClusterInfoHandler)), + ]; + for (role, handler) in handlers { + let mut req = HeartbeatRequest { + header: Some(RequestHeader { + role: role as i32, + ..Default::default() + }), + peer: Some(Peer { + id: 7, + addr: "node:3001".into(), + }), + info: Some(PbNodeInfo { + version: "v1".into(), + hostname: "host".into(), + ..Default::default() + }), + node_epoch: 10, + ..Default::default() + }; + let mut acc = HeartbeatAccumulator { + stat: Some(Stat { + timestamp_millis: 123, + rcus: 5, + wcus: 6, + ..Default::default() + }), + ..Default::default() + }; + handler.handle(&req, &mut ctx, &mut acc).await.unwrap(); + let key: Vec = (&NodeInfoKey::new(&req).unwrap()).into(); + let first = + NodeInfo::try_from(ctx.in_memory.get(&key).await.unwrap().unwrap().value).unwrap(); + assert_eq!(first.peer, req.peer.clone().unwrap()); + assert_eq!(first.version, "v1"); + assert_eq!(first.hostname, "host"); + req.info.as_mut().unwrap().version = "v2".into(); + req.info.as_mut().unwrap().memory_usage_bytes = 42; + acc.stat.as_mut().unwrap().timestamp_millis = 456; + handler.handle(&req, &mut ctx, &mut acc).await.unwrap(); + let second = + NodeInfo::try_from(ctx.in_memory.get(&key).await.unwrap().unwrap().value).unwrap(); + assert_eq!(second.version, "v2"); + assert_eq!(second.memory_usage_bytes, 42); + assert_eq!( + serde_json::to_value(&second.status).unwrap(), + serde_json::to_value(&first.status).unwrap() ); - }) - .unwrap_or_default() - .map(|e| e.vars) - .unwrap_or_default(); - - Some((key, peer.clone(), info.clone(), env_vars)) -} - -async fn put_into_memory_store(ctx: &mut Context, key: NodeInfoKey, value: NodeInfo) -> Result<()> { - let key = (&key).into(); - let value = value.try_into().context(InvalidClusterInfoFormatSnafu)?; - let put_req = PutRequest { - key, - value, - ..Default::default() - }; - - ctx.in_memory - .put(put_req) - .await - .context(SaveClusterInfoSnafu)?; - - Ok(()) + if role == Role::Datanode { + assert_eq!(second.last_activity_ts, 456); + } + // Node information must not be persisted to the configured backend. + assert!(ctx.kv_backend.get(&key).await.unwrap().is_none()); + } + } } diff --git a/src/meta-srv/src/handler/collect_stats_handler.rs b/src/meta-srv/src/handler/collect_stats_handler.rs index e3ae7ebe915..5e63c49cfe9 100644 --- a/src/meta-srv/src/handler/collect_stats_handler.rs +++ b/src/meta-srv/src/handler/collect_stats_handler.rs @@ -15,19 +15,19 @@ use std::cmp::Ordering; use std::sync::Arc; -use api::v1::meta::{HeartbeatRequest, Role}; +use api::v1::meta::{HeartbeatRequest, Peer, Role}; use common_meta::datanode::{DatanodeStatKey, DatanodeStatValue, Stat}; use common_meta::instruction::CacheIdent; -use common_meta::key::node_address::{NodeAddressKey, NodeAddressValue}; -use common_meta::key::{MetadataKey, MetadataValue}; -use common_meta::peer::Peer; +use common_meta::key::MetadataKey; +use common_meta::key::node_address::NodeAddressKey; use common_meta::rpc::store::PutRequest; -use common_telemetry::{error, info, warn}; +use common_telemetry::{error, warn}; use dashmap::DashMap; use snafu::ResultExt; use tokio::sync::Mutex; use crate::error::{self, Result}; +use crate::handler::utils::{self, NodeAddressUpdater}; use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler}; use crate::metasrv::Context; @@ -74,6 +74,7 @@ const DEFAULT_FLUSH_STATS_FACTOR: usize = 3; pub struct CollectStatsHandler { stats_cache: DashMap>>, flush_stats_factor: usize, + address_updater: NodeAddressUpdater, } impl Default for CollectStatsHandler { @@ -87,6 +88,7 @@ impl CollectStatsHandler { Self { flush_stats_factor: flush_stats_factor.unwrap_or(DEFAULT_FLUSH_STATS_FACTOR), stats_cache: DashMap::default(), + address_updater: NodeAddressUpdater::default(), } } } @@ -119,6 +121,7 @@ impl HeartbeatHandler for CollectStatsHandler { let key: Vec = key.into(); + let current_epoch = current_stat.node_epoch; let refresh = if let Some(epoch) = epoch_stats.epoch() { match current_stat.node_epoch.cmp(&epoch) { Ordering::Greater => { @@ -146,11 +149,13 @@ impl HeartbeatHandler for CollectStatsHandler { true }; - // Need to refresh the [datanode -> address] mapping - if refresh { - // Safety: `epoch_stats.stats` is not empty - let last = epoch_stats.stats.last().unwrap(); - rewrite_node_address(ctx, last).await; + // Keep address updates and statistics in the original per-node order. + // Address acknowledgement is independent of the statistics epoch so a + // failed save can be retried by the next heartbeat of the same epoch. + if epoch_stats.epoch() == Some(current_epoch) + && let Some(stat) = epoch_stats.stats.last() + { + update_datanode_address(&self.address_updater, ctx, stat).await; } if !refresh && epoch_stats.len() < self.flush_stats_factor { @@ -178,48 +183,26 @@ impl HeartbeatHandler for CollectStatsHandler { } } -async fn rewrite_node_address(ctx: &Context, stat: &Stat) { - let peer = Peer { - id: stat.id, - addr: stat.addr.clone(), - }; - let key = NodeAddressKey::with_datanode(peer.id).to_bytes(); - if let Ok(value) = NodeAddressValue::new(peer.clone()).try_as_raw_value() { - let put = PutRequest { - key, - value, - prev_kv: false, - }; - - match ctx.leader_cached_kv_backend.put(put).await { - Ok(_) => { - info!( - "Successfully updated datanode `NodeAddressValue`: {:?}", - peer - ); - // broadcast invalidating cache - let cache_idents = stat - .table_ids() - .into_iter() - .map(CacheIdent::TableId) - .collect::>(); - if let Err(e) = ctx - .cache_invalidator - .invalidate(&Default::default(), &cache_idents) - .await - { - error!(e; "Failed to invalidate {} `NodeAddressKey` cache, peer: {:?}", cache_idents.len(), peer); - } - } - Err(e) => { - error!(e; "Failed to update datanode `NodeAddressValue`: {:?}", peer); - } - } - } else { - warn!( - "Failed to serialize datanode `NodeAddressValue`: {:?}", - peer - ); +/// Applies datanode-specific address and table-cache updates. The caller keeps +/// its statistics lock so notifications and statistics retain their ordering. +async fn update_datanode_address(updater: &NodeAddressUpdater, ctx: &Context, stat: &Stat) { + let updated = updater + .update_if_needed((Role::Datanode, stat.id), stat.node_epoch, || async { + let peer = Peer { id: stat.id, addr: stat.addr.clone() }; + let key = NodeAddressKey::with_datanode(stat.id).to_bytes(); + utils::save_node_address(ctx, key, peer).await + }) + .await + .inspect_err(|e| { + error!(e; "Failed to update datanode address, node: {}, address: {}", stat.id, stat.addr); + }); + if let Ok(true) = updated { + let cache_idents = stat + .table_ids() + .into_iter() + .map(CacheIdent::TableId) + .collect::>(); + utils::invalidate_address_caches(ctx, stat.id, &cache_idents).await; } } @@ -553,4 +536,127 @@ mod tests { handler.handle(&req, &mut ctx, &mut acc).await.unwrap(); } } + + #[tokio::test] + async fn test_persistence_failure_retries_but_notification_failure_does_not() { + use api::v1::meta::HeartbeatRequest; + use common_meta::key::MetadataValue; + use common_meta::key::node_address::NodeAddressValue; + use common_meta::kv_backend::test_util::MockKvBackendBuilder; + use common_meta::rpc::store::PutResponse; + + use crate::handler::collect_stats_handler::CollectStatsHandler; + use crate::handler::test_utils::FailingCacheInvalidator; + use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler}; + use crate::service::store::cached_kv::LeaderCachedKvBackend; + let mut ctx = TestEnv::new().ctx(); + let puts = Arc::new(std::sync::Mutex::new(Vec::::new())); + let backend = MockKvBackendBuilder::default() + .put_fn(Arc::new({ + let puts = puts.clone(); + move |req: PutRequest| { + let mut puts = puts.lock().unwrap(); + puts.push(req); + if puts.len() == 1 { + return common_meta::error::UnexpectedSnafu { + err_msg: "injected persistence failure", + } + .fail(); + } + Ok(PutResponse::default()) + } + }) as _) + .build() + .unwrap(); + ctx.leader_cached_kv_backend = + Arc::new(LeaderCachedKvBackend::with_always_leader(Arc::new(backend))); + let invalidator = Arc::new(FailingCacheInvalidator::default()); + ctx.cache_invalidator = invalidator.clone(); + let handler = CollectStatsHandler::default(); + let mut stat = Stat { + id: 101, + addr: "dn-original".into(), + node_epoch: 7, + ..Default::default() + }; + let req = HeartbeatRequest::default(); + + // A backend error must not interrupt subsequent heartbeat handlers. + assert!(matches!( + handler + .handle( + &req, + &mut ctx, + &mut HeartbeatAccumulator { + stat: Some(stat.clone()), + ..Default::default() + } + ) + .await + .unwrap(), + HandleControl::Continue + )); + assert_eq!(puts.lock().unwrap().len(), 1); + assert_eq!(invalidator.attempts(), 0); + + let stats_key: Vec = common_meta::datanode::DatanodeStatKey { node_id: 101 }.into(); + assert!(ctx.in_memory.get(&stats_key).await.unwrap().is_some()); + + // Same epoch retries persistence using the current heartbeat's address. + stat.addr = "dn-retry".into(); + handler + .handle( + &req, + &mut ctx, + &mut HeartbeatAccumulator { + stat: Some(stat.clone()), + ..Default::default() + }, + ) + .await + .unwrap(); + { + let puts = puts.lock().unwrap(); + assert_eq!(puts.len(), 2); + assert_eq!(puts[1].key, NodeAddressKey::with_datanode(101).to_bytes()); + assert_eq!( + NodeAddressValue::try_from_raw_value(&puts[1].value) + .unwrap() + .peer + .addr, + "dn-retry" + ); + } + assert_eq!(invalidator.attempts(), 1); + + // Notification remains best-effort: its failure neither rewrites nor retries. + handler + .handle( + &req, + &mut ctx, + &mut HeartbeatAccumulator { + stat: Some(stat.clone()), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(puts.lock().unwrap().len(), 2); + assert_eq!(invalidator.attempts(), 1); + + stat.node_epoch = 8; + handler + .handle( + &req, + &mut ctx, + &mut HeartbeatAccumulator { + stat: Some(stat.clone()), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(puts.lock().unwrap().len(), 3); + assert_eq!(invalidator.attempts(), 2); + } } diff --git a/src/meta-srv/src/handler/remap_flow_peer_handler.rs b/src/meta-srv/src/handler/remap_flow_peer_handler.rs index f6b3c67a9e5..6a2e7040e6a 100644 --- a/src/meta-srv/src/handler/remap_flow_peer_handler.rs +++ b/src/meta-srv/src/handler/remap_flow_peer_handler.rs @@ -12,22 +12,20 @@ // See the License for the specific language governing permissions and // limitations under the License. -use api::v1::meta::{HeartbeatRequest, Peer, Role}; +use api::v1::meta::{HeartbeatRequest, Role}; use common_meta::instruction::CacheIdent; -use common_meta::key::node_address::{NodeAddressKey, NodeAddressValue}; -use common_meta::key::{MetadataKey, MetadataValue}; -use common_meta::rpc::store::PutRequest; -use common_telemetry::{error, info, warn}; -use dashmap::DashMap; +use common_meta::key::MetadataKey; +use common_meta::key::node_address::NodeAddressKey; +use common_telemetry::error; use crate::Result; +use crate::handler::utils::{self, NodeAddressUpdater}; use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler}; use crate::metasrv::Context; #[derive(Debug, Default)] pub struct RemapFlowPeerHandler { - /// flow_node_id -> epoch - epoch_cache: DashMap, + address_updater: NodeAddressUpdater, } #[async_trait::async_trait] @@ -45,61 +43,85 @@ impl HeartbeatHandler for RemapFlowPeerHandler { let Some(peer) = req.peer.as_ref() else { return Ok(HandleControl::Continue); }; - - let current_epoch = req.node_epoch; - let flow_node_id = peer.id; - - let refresh = if let Some(mut epoch) = self.epoch_cache.get_mut(&flow_node_id) { - if current_epoch > *epoch.value() { - *epoch.value_mut() = current_epoch; - true - } else { - false - } - } else { - self.epoch_cache.insert(flow_node_id, current_epoch); - true - }; - - if refresh { - rewrite_node_address(ctx, peer).await; + let updated = self + .address_updater + .update_if_needed((Role::Flownode, peer.id), req.node_epoch, || { + let key = NodeAddressKey::with_flownode(peer.id).to_bytes(); + utils::save_node_address(ctx, key, peer.clone()) + }) + .await + .inspect_err(|e| { + // Address errors must not prevent later heartbeat handlers from running. + error!(e; "Failed to update flownode address, peer: {:?}", peer); + }); + if let Ok(true) = updated { + let cache_idents = [CacheIdent::FlowNodeAddressChange(peer.id)]; + utils::invalidate_address_caches(ctx, peer.id, &cache_idents).await; } - Ok(HandleControl::Continue) } } -async fn rewrite_node_address(ctx: &mut Context, peer: &Peer) { - let key = NodeAddressKey::with_flownode(peer.id).to_bytes(); - if let Ok(value) = NodeAddressValue::new(peer.clone()).try_as_raw_value() { - let put = PutRequest { - key, - value, - prev_kv: false, - }; +#[cfg(test)] +mod tests { + use std::sync::Arc; - match ctx.leader_cached_kv_backend.put(put).await { - Ok(_) => { - info!("Successfully updated flow `NodeAddressValue`: {:?}", peer); - // broadcast invalidating cache to all frontends - let cache_idents = vec![CacheIdent::FlowNodeAddressChange(peer.id)]; - info!( - "Invalidate flow node cache for new address with cache idents: {:?}", - cache_idents - ); - if let Err(e) = ctx - .cache_invalidator - .invalidate(&Default::default(), &cache_idents) - .await - { - error!(e; "Failed to invalidate {} `NodeAddressKey` cache, peer: {:?}", cache_idents.len(), peer); - } - } - Err(e) => { - error!(e; "Failed to update flow `NodeAddressValue`: {:?}", peer); - } - } - } else { - warn!("Failed to serialize flow `NodeAddressValue`: {:?}", peer); + use api::v1::meta::Peer; + + use super::*; + use crate::handler::test_utils::TestEnv; + #[tokio::test] + async fn test_notification_failure_keeps_persisted_epoch() { + use api::v1::meta::HeartbeatRequest; + use common_meta::key::MetadataValue; + use common_meta::key::node_address::NodeAddressValue; + + use crate::handler::test_utils::FailingCacheInvalidator; + use crate::handler::{HeartbeatAccumulator, HeartbeatHandler}; + let mut ctx = TestEnv::new().ctx(); + let invalidator = Arc::new(FailingCacheInvalidator::default()); + ctx.cache_invalidator = invalidator.clone(); + let handler = RemapFlowPeerHandler::default(); + let mut req = HeartbeatRequest { + peer: Some(Peer { + id: 1, + addr: "first".into(), + }), + node_epoch: 1, + ..Default::default() + }; + let key = NodeAddressKey::with_flownode(1).to_bytes(); + handler + .handle(&req, &mut ctx, &mut HeartbeatAccumulator::default()) + .await + .unwrap(); + req.peer.as_mut().unwrap().addr = "same-epoch".into(); + handler + .handle(&req, &mut ctx, &mut HeartbeatAccumulator::default()) + .await + .unwrap(); + assert_eq!(invalidator.attempts(), 1); + let kv = ctx.kv_backend.get(&key).await.unwrap().unwrap(); + assert_eq!( + NodeAddressValue::try_from_raw_value(&kv.value) + .unwrap() + .peer + .addr, + "first" + ); + req.node_epoch = 2; + handler + .handle(&req, &mut ctx, &mut HeartbeatAccumulator::default()) + .await + .unwrap(); + assert_eq!(invalidator.attempts(), 2); + let kv = ctx.kv_backend.get(&key).await.unwrap().unwrap(); + assert_eq!( + NodeAddressValue::try_from_raw_value(&kv.value) + .unwrap() + .peer + .addr, + "same-epoch" + ); } } diff --git a/src/meta-srv/src/handler/test_utils.rs b/src/meta-srv/src/handler/test_utils.rs index d8241c785ea..c2afb3cb14a 100644 --- a/src/meta-srv/src/handler/test_utils.rs +++ b/src/meta-srv/src/handler/test_utils.rs @@ -13,9 +13,13 @@ // limitations under the License. use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; -use common_meta::cache_invalidator::{CacheInvalidatorRef, DummyCacheInvalidator}; +use common_meta::cache_invalidator::{ + CacheInvalidator, CacheInvalidatorRef, Context as InvalidationContext, DummyCacheInvalidator, +}; use common_meta::distributed_time_constants::BASE_HEARTBEAT_INTERVAL; +use common_meta::instruction::CacheIdent; use common_meta::key::{TableMetadataManager, TableMetadataManagerRef}; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::kv_backend::{KvBackendRef, ResettableKvBackendRef}; @@ -97,3 +101,34 @@ impl TestEnv { } } } + +/// Counts failed notification attempts without requiring heartbeat connections. +#[derive(Default)] +pub(crate) struct FailingCacheInvalidator { + attempts: AtomicUsize, +} + +impl FailingCacheInvalidator { + pub(crate) fn attempts(&self) -> usize { + self.attempts.load(Ordering::Relaxed) + } +} + +#[async_trait::async_trait] +impl CacheInvalidator for FailingCacheInvalidator { + async fn invalidate( + &self, + _: &InvalidationContext, + _: &[CacheIdent], + ) -> common_meta::error::Result<()> { + self.attempts.fetch_add(1, Ordering::Relaxed); + common_meta::error::UnexpectedSnafu { + err_msg: "injected notification failure", + } + .fail() + } + + fn invalidate_all(&self) -> common_meta::error::Result<()> { + Ok(()) + } +} diff --git a/src/meta-srv/src/handler/utils.rs b/src/meta-srv/src/handler/utils.rs new file mode 100644 index 00000000000..96624d87306 --- /dev/null +++ b/src/meta-srv/src/handler/utils.rs @@ -0,0 +1,281 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::future::Future; +use std::sync::Arc; + +use api::v1::meta::{Peer, Role}; +use common_meta::instruction::CacheIdent; +use common_meta::key::MetadataValue; +use common_meta::key::node_address::NodeAddressValue; +use common_meta::rpc::store::PutRequest; +use common_telemetry::{error, info}; +use dashmap::DashMap; +use snafu::ResultExt; +use tokio::sync::Mutex; + +use crate::Result; +use crate::error::{InvalidNodeInfoFormatSnafu, KvBackendSnafu}; +use crate::metasrv::Context; + +/// Serializes address changes per node and acknowledges only successful updates. +/// Node information refreshed on every heartbeat has a separate lifecycle. +#[derive(Debug, Default)] +pub(crate) struct NodeAddressUpdater { + states: DashMap<(Role, u64), Arc>>, +} + +#[derive(Debug, Default)] +struct NodeAddressEpochs { + /// Rejects stale heartbeats even after a newer epoch failed to save. + seen_epoch: Option, + /// Suppresses retries only after persistence succeeds. + saved_epoch: Option, +} + +impl NodeAddressUpdater { + /// Saves a new or previously failed epoch. Returns true only when this call + /// saved the address, or false when the heartbeat is stale or already saved. + /// Cache invalidation is separate and does not affect acknowledgement. + pub(crate) async fn update_if_needed( + &self, + node: (Role, u64), + epoch: u64, + save: F, + ) -> Result + where + F: FnOnce() -> Fut, + Fut: Future>, + { + // Never retain a DashMap guard across await, even while waiting for a + // different node: its entry may live in the same shard. + let state = { + let entry = self.states.entry(node).or_default(); + Arc::clone(entry.value()) + }; + let mut state = state.lock().await; + if state.seen_epoch.is_some_and(|seen| epoch < seen) || state.saved_epoch == Some(epoch) { + return Ok(false); + } + + // Seeing an epoch rejects stale heartbeats, but only a successful save + // suppresses retries of that same epoch. Cancellation also leaves it retryable. + state.seen_epoch = Some(epoch); + save().await?; + state.saved_epoch = Some(epoch); + Ok(true) + } +} + +/// Persists an address through the leader cache. +pub(crate) async fn save_node_address(ctx: &Context, key: Vec, peer: Peer) -> Result<()> { + let address = NodeAddressValue::new(peer); + let value = address + .try_as_raw_value() + .context(InvalidNodeInfoFormatSnafu)?; + ctx.leader_cached_kv_backend + .put(PutRequest { + key, + value, + prev_kv: false, + }) + .await + .context(KvBackendSnafu)?; + info!("Successfully updated node address: {:?}", address.peer); + Ok(()) +} + +/// Preserves the existing best-effort notification behavior. Notification +/// failures do not undo a persisted address or cause same-epoch rewrites. +pub(crate) async fn invalidate_address_caches( + ctx: &Context, + node_id: u64, + cache_idents: &[CacheIdent], +) { + if let Err(e) = ctx + .cache_invalidator + .invalidate(&Default::default(), cache_idents) + .await + { + error!(e; "Failed to invalidate address caches for node {}: {:?}", node_id, cache_idents); + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use tokio::sync::Notify; + use tokio::time::timeout; + + use super::*; + + #[tokio::test] + async fn test_failed_save_retries_same_epoch_and_rejects_stale_heartbeats() { + let updater = NodeAddressUpdater::default(); + let key = (Role::Datanode, 1); + assert!( + updater + .update_if_needed(key, 2, || async { + crate::error::UnexpectedSnafu { + violated: "injected save failure", + } + .fail() + }) + .await + .is_err() + ); + assert!( + !updater + .update_if_needed(key, 1, || async { + panic!("stale heartbeat must not save"); + }) + .await + .unwrap() + ); + assert!( + updater + .update_if_needed(key, 2, || async { Ok(()) }) + .await + .unwrap() + ); + assert!( + !updater + .update_if_needed(key, 2, || async { + panic!("successful epoch must not save again"); + }) + .await + .unwrap() + ); + assert!( + updater + .update_if_needed(key, 3, || async { Ok(()) }) + .await + .unwrap() + ); + } + + #[tokio::test] + async fn test_cancelled_save_can_be_retried() { + let updater = NodeAddressUpdater::default(); + let key = (Role::Datanode, 1); + let entered = Notify::new(); + { + let first = updater.update_if_needed(key, 0, || async { + entered.notify_one(); + std::future::pending::>().await + }); + tokio::pin!(first); + tokio::select! { + _ = entered.notified() => {}, + _ = &mut first => panic!("save must remain pending"), + } + } + assert!( + timeout( + Duration::from_secs(2), + updater.update_if_needed(key, 0, || async { Ok(()) }) + ) + .await + .unwrap() + .unwrap() + ); + } + + #[test] + fn test_same_node_serialized_other_nodes_progress() { + // A watchdog outside the runtime detects synchronous DashMap deadlocks; + // a Tokio timeout alone cannot fire if the only worker is blocked. + let (progress_tx, progress_rx) = std::sync::mpsc::sync_channel(1); + let release = Arc::new(Notify::new()); + let finish = Arc::new(Notify::new()); + let worker = std::thread::spawn({ + let release = release.clone(); + let finish = finish.clone(); + move || { + tokio::runtime::Builder::new_current_thread() + .enable_time() + .build() + .unwrap() + .block_on(async move { + let updater = Arc::new(NodeAddressUpdater::default()); + let entered = Arc::new(Notify::new()); + + let first = tokio::spawn({ + let updater = updater.clone(); + let entered = entered.clone(); + let release = release.clone(); + async move { + updater + .update_if_needed((Role::Datanode, 1), 1, || async { + entered.notify_one(); + release.notified().await; + Ok(()) + }) + .await + } + }); + entered.notified().await; + let second = tokio::spawn({ + let updater = updater.clone(); + async move { + updater + .update_if_needed((Role::Datanode, 1), 2, || async { Ok(()) }) + .await + } + }); + tokio::task::yield_now().await; + assert!(!second.is_finished()); + // Exercise shared DashMap shards on a single-thread runtime while the + // first node is suspended in external I/O. + timeout(Duration::from_secs(2), async { + for id in 0_u64..128 { + assert!( + updater + .update_if_needed((Role::Datanode, id + 2), 1, || async { + Ok(()) + }) + .await + .unwrap() + ); + } + }) + .await + .unwrap(); + progress_tx.send(()).unwrap(); + finish.notified().await; + timeout(Duration::from_secs(2), async { + assert!(first.await.unwrap().unwrap()); + assert!(second.await.unwrap().unwrap()); + }) + .await + .unwrap(); + assert!( + !updater + .update_if_needed((Role::Datanode, 1), 1, || async { + panic!("old epoch must not overwrite a newer one"); + }) + .await + .unwrap() + ); + }); + } + }); + let progress = progress_rx.recv_timeout(Duration::from_secs(5)); + release.notify_one(); + finish.notify_one(); + progress.expect("other node updates must progress while one node save is blocked"); + worker.join().unwrap(); + } +} diff --git a/src/meta-srv/src/procedure/region_migration/manager.rs b/src/meta-srv/src/procedure/region_migration/manager.rs index c21de506243..8a115cd9f8e 100644 --- a/src/meta-srv/src/procedure/region_migration/manager.rs +++ b/src/meta-srv/src/procedure/region_migration/manager.rs @@ -510,17 +510,14 @@ impl RegionMigrationManager { ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); let procedure_id = procedure_with_id.id; info!("Starting region migration procedure {procedure_id} for {task}"); - let procedure_manager = self.procedure_manager.clone(); + let mut watcher = self + .procedure_manager + .submit(procedure_with_id) + .await + .context(error::SubmitProcedureSnafu)?; let num_region = task.region_ids.len(); common_runtime::spawn_global(async move { - let watcher = &mut match procedure_manager.submit(procedure_with_id).await { - Ok(watcher) => watcher, - Err(e) => { - error!(e; "Failed to submit region migration procedure {procedure_id} for {task}"); - return; - } - }; METRIC_META_REGION_MIGRATION_DATANODES .with_label_values(&["src", &task.from_peer.id.to_string()]) .inc_by(num_region as u64); @@ -528,7 +525,7 @@ impl RegionMigrationManager { .with_label_values(&["desc", &task.to_peer.id.to_string()]) .inc_by(num_region as u64); - if let Err(e) = watcher::wait(watcher).await { + if let Err(e) = watcher::wait(&mut watcher).await { error!(e; "Failed to wait region migration procedure {procedure_id} for {task}"); METRIC_META_REGION_MIGRATION_FAIL.inc(); return; @@ -613,15 +610,12 @@ impl RegionMigrationManager { ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); let procedure_id = procedure_with_id.id; info!("Starting region migration procedure {procedure_id} for {task}"); - let procedure_manager = self.procedure_manager.clone(); + let mut watcher = self + .procedure_manager + .submit(procedure_with_id) + .await + .context(error::SubmitProcedureSnafu)?; common_runtime::spawn_global(async move { - let watcher = &mut match procedure_manager.submit(procedure_with_id).await { - Ok(watcher) => watcher, - Err(e) => { - error!(e; "Failed to submit region migration procedure {procedure_id} for {task}"); - return; - } - }; METRIC_META_REGION_MIGRATION_DATANODES .with_label_values(&["src", &task.from_peer.id.to_string()]) .inc(); @@ -629,7 +623,7 @@ impl RegionMigrationManager { .with_label_values(&["desc", &task.to_peer.id.to_string()]) .inc(); - if let Err(e) = watcher::wait(watcher).await { + if let Err(e) = watcher::wait(&mut watcher).await { error!(e; "Failed to wait region migration procedure {procedure_id} for {task}"); METRIC_META_REGION_MIGRATION_FAIL.inc(); return; @@ -649,10 +643,148 @@ mod test { use common_meta::key::table_route::LogicalTableRouteValue; use common_meta::key::test_utils::new_test_table_info; use common_meta::rpc::router::Region; + use common_meta::state_store::KvStateStore; + use common_procedure::ProcedureManager; + use common_procedure::local::{LocalManager, ManagerConfig, PauseAware}; + use common_procedure::test_util::InMemoryPoisonStore; + use tokio::sync::Notify; use super::*; use crate::procedure::region_migration::test_util::TestingEnv; + #[derive(Default)] + struct SubmissionGate { + entered: Notify, + release: Notify, + } + + #[async_trait::async_trait] + impl PauseAware for SubmissionGate { + async fn is_paused(&self) -> std::result::Result { + self.entered.notify_one(); + self.release.notified().await; + Ok(false) + } + } + + async fn check_submission_registration(batch: bool) { + let env = TestingEnv::new(); + let gate = Arc::new(SubmissionGate::default()); + let procedure_manager = Arc::new(LocalManager::new( + ManagerConfig::default(), + Arc::new(KvStateStore::new(env.kv_backend())), + Arc::new(InMemoryPoisonStore::default()), + Some(gate.clone()), + None, + )); + let manager = RegionMigrationManager::new(procedure_manager.clone(), env.context_factory()); + let region_ids = if batch { + vec![RegionId::new(1024, 1), RegionId::new(1024, 2)] + } else { + vec![RegionId::new(1024, 1)] + }; + let region_routes = region_ids + .iter() + .map(|region_id| RegionRoute { + region: Region::new_test(*region_id), + leader_peer: Some(Peer::empty(1)), + ..Default::default() + }) + .collect(); + env.create_physical_table_metadata(new_test_table_info(1024), region_routes) + .await; + + // A stopped manager must reject submission and release all migration guards. + // Starting it then allows the same regions to be submitted successfully. + for started in [false, true] { + if started { + procedure_manager.start().await.unwrap(); + } + let submission = async { + if batch { + let result = manager + .submit_region_migration_task(RegionMigrationTaskBatch { + region_ids: region_ids.clone(), + from_peer: Peer::empty(1), + to_peer: Peer::empty(2), + timeout: Duration::from_secs(10), + trigger_reason: RegionMigrationTriggerReason::Manual, + }) + .await?; + assert_eq!(result.submitted, region_ids); + Ok(result.procedure_id) + } else { + manager + .submit_procedure( + ProcedureContext::default(), + RegionMigrationProcedureTask::new( + region_ids[0], + Peer::empty(1), + Peer::empty(2), + Duration::from_secs(10), + RegionMigrationTriggerReason::Manual, + ), + ) + .await + } + }; + tokio::pin!(submission); + tokio::select! { + biased; + result = &mut submission => panic!("Submission returned before registration: {result:?}"), + _ = gate.entered.notified() => {} + } + assert!(futures::poll!(&mut submission).is_pending()); + assert!( + procedure_manager + .list_procedures() + .await + .unwrap() + .is_empty() + ); + for region_id in ®ion_ids { + assert!(manager.tracker.contains(*region_id)); + } + + gate.release.notify_one(); + let result = tokio::time::timeout(Duration::from_secs(5), &mut submission) + .await + .expect("Submission should return after registration"); + if started { + let procedure_id = result.unwrap().unwrap(); + assert!( + procedure_manager + .procedure_state(procedure_id) + .await + .unwrap() + .is_some() + ); + } else { + assert_matches!( + result.unwrap_err(), + error::Error::SubmitProcedure { + source: common_procedure::Error::ManagerNotStart { .. }, + .. + } + ); + for region_id in ®ion_ids { + assert!(!manager.tracker.contains(*region_id)); + } + } + } + procedure_manager.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_submit_procedure_registration() { + check_submission_registration(false).await; + } + + #[tokio::test] + async fn test_submit_region_migration_task_registration() { + check_submission_registration(true).await; + } + #[tokio::test] async fn test_insert_running_procedure() { let env = TestingEnv::new(); @@ -1039,6 +1171,7 @@ mod test { #[tokio::test] async fn test_running_regions() { let env = TestingEnv::new(); + env.procedure_manager().start().await.unwrap(); let context_factory = env.context_factory(); let manager = RegionMigrationManager::new(env.procedure_manager().clone(), context_factory); let region_id = RegionId::new(1024, 1); @@ -1072,5 +1205,6 @@ mod test { assert_eq!(result.migrating, vec![region_id]); assert_eq!(result.submitted, vec![RegionId::new(1024, 2)]); assert!(result.procedure_id.is_some()); + env.procedure_manager().stop().await.unwrap(); } } diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index 9a7bd71b2ff..ef150dfb070 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -101,6 +101,10 @@ pub(crate) struct RangeResultMemoryLimiter { semaphore: Arc, permit_bytes: usize, total_permits: usize, + /// Number of acquisitions that found too few permits and parked. Test-only + /// signal to synchronize with a caller waiting inside [Self::acquire]. + #[cfg(test)] + waited_acquires: std::sync::atomic::AtomicUsize, } impl Default for RangeResultMemoryLimiter { @@ -122,6 +126,8 @@ impl RangeResultMemoryLimiter { semaphore: Arc::new(tokio::sync::Semaphore::new(total_permits)), permit_bytes, total_permits, + #[cfg(test)] + waited_acquires: std::sync::atomic::AtomicUsize::new(0), } } @@ -130,6 +136,12 @@ impl RangeResultMemoryLimiter { self.permit_bytes } + #[cfg(test)] + pub(crate) fn waited_acquires(&self) -> usize { + self.waited_acquires + .load(std::sync::atomic::Ordering::Acquire) + } + #[cfg(test)] pub(crate) fn available_permits(&self) -> usize { self.semaphore.available_permits() @@ -146,6 +158,14 @@ impl RangeResultMemoryLimiter { } .fail(); } + // Nothing awaits between this check and the parking below, so an observed + // increment means the caller is about to wait for the missing permits. + #[cfg(test)] + if self.semaphore.available_permits() < permits { + self.waited_acquires + .fetch_add(1, std::sync::atomic::Ordering::Release); + } + self.semaphore .acquire_many(permits as u32) .await diff --git a/src/mito2/src/compaction/twcs.rs b/src/mito2/src/compaction/twcs.rs index b2e3274b062..4546f339c5d 100644 --- a/src/mito2/src/compaction/twcs.rs +++ b/src/mito2/src/compaction/twcs.rs @@ -892,6 +892,7 @@ struct Window { start: Timestamp, end: Timestamp, files: Vec, + /// Start of the half-open window, in seconds. time_window: i64, primary_key_range: Option<(bytes::Bytes, bytes::Bytes)>, } @@ -942,12 +943,7 @@ fn assign_to_windows<'a>( continue; } let (_, end) = f.time_range(); - let time_window = end - .convert_to(TimeUnit::Second) - .unwrap() - .value() - .align_to_ceil_by_bucket(time_window_size) - .unwrap_or(i64::MIN); + let time_window = time_window_start(end, time_window_size); match windows.entry(time_window) { Entry::Occupied(mut e) => { @@ -963,37 +959,41 @@ fn assign_to_windows<'a>( windows.into_iter().collect() } +/// Uses the same half-open windows as memtables, including timestamps in the first +/// second after a boundary. An unrepresentable window start is clamped to `i64::MIN`. +fn time_window_start(timestamp: Timestamp, time_window_size: i64) -> i64 { + // Conversion to seconds cannot overflow. + timestamp + .convert_to(TimeUnit::Second) + .and_then(|ts| ts.value().align_by_bucket(time_window_size)) + .unwrap_or(i64::MIN) +} + fn time_window_intersects_range( - window_end: i64, + window_start: i64, time_window_size: i64, time_range: &TimestampRange, ) -> bool { - let first_window = match time_range.start() { - None => i64::MIN, - Some(start) => { - let Some(first_window) = start - .convert_to(TimeUnit::Second) - .and_then(|timestamp| timestamp.value().align_to_ceil_by_bucket(time_window_size)) - else { - return false; - }; - first_window - } - }; + if time_range.is_empty() { + return false; + } + let first_window = time_range + .start() + .map_or(i64::MIN, |start| time_window_start(start, time_window_size)); let last_window = match time_range.end() { None => i64::MAX, Some(end) => { - let Some(last_window) = end + // The range end is exclusive; keep the last second that intersects it. + let Some(last_second) = end .convert_to_ceil(TimeUnit::Second) .and_then(|timestamp| timestamp.value().checked_sub(1)) - .and_then(|timestamp| timestamp.align_to_ceil_by_bucket(time_window_size)) else { return false; }; - last_window + time_window_start(Timestamp::new_second(last_second), time_window_size) } }; - (first_window..=last_window).contains(&window_end) + (first_window..=last_window).contains(&window_start) } fn window_has_overlap(this: &Window, windows: &BTreeMap) -> bool { @@ -1011,7 +1011,7 @@ fn window_has_overlap(this: &Window, windows: &BTreeMap) -> bool { } /// Finds the latest active writing window among all files. -/// Returns `None` when there are no files or all files are corrupted. +/// Returns `None` when there are no files. fn find_latest_window_in_seconds<'a>( files: impl Iterator, time_window_size: i64, @@ -1027,9 +1027,7 @@ fn find_latest_window_in_seconds<'a>( latest_timestamp = Some(end); } } - latest_timestamp - .and_then(|ts| ts.convert_to_ceil(TimeUnit::Second)) - .and_then(|ts| ts.value().align_to_ceil_by_bucket(time_window_size)) + latest_timestamp.map(|ts| time_window_start(ts, time_window_size)) } /// Finds the active window from the file with the highest sequence number. @@ -1039,9 +1037,7 @@ fn find_latest_window_in_seconds<'a>( /// survives the transient state where level 0 is empty right after its files /// were compacted away. Returns `None` when no file carries a sequence. /// -/// The window key follows the same convention as [`assign_to_windows`] -/// (truncate to seconds, then align up), since it is compared against the -/// window keys produced there. +/// The window key is its inclusive start, as in [`assign_to_windows`]. fn find_active_window_by_sequence<'a>( files: impl Iterator, time_window_size: i64, @@ -1049,12 +1045,7 @@ fn find_active_window_by_sequence<'a>( files .filter(|f| f.meta_ref().sequence.is_some()) .max_by_key(|f| f.meta_ref().sequence) - .and_then(|f| { - f.time_range() - .1 - .convert_to(TimeUnit::Second) - .and_then(|ts| ts.value().align_to_ceil_by_bucket(time_window_size)) - }) + .map(|f| time_window_start(f.time_range().1, time_window_size)) } #[cfg(test)] @@ -1347,7 +1338,7 @@ mod tests { new_file_handle_with_sequence(FileId::random(), 0, 999, 1, 0), new_file_handle_with_sequence(FileId::random(), 2000, 2999, 0, 0), ]; - assert_eq!(Some(3), active_window_of(&files)); + assert_eq!(Some(2), active_window_of(&files)); // No sequence and no L0: None, exactly as before. let files = [new_file_handle_with_sequence( @@ -1363,7 +1354,7 @@ mod tests { #[test] fn test_get_latest_window_in_seconds() { assert_eq!( - Some(1), + Some(0), find_latest_window_in_seconds([new_file_handle(FileId::random(), 0, 999, 0)].iter(), 1) ); assert_eq!( @@ -1375,7 +1366,7 @@ mod tests { ); assert_eq!( - Some(-9223372036854000), + Some(-9223372036857600), find_latest_window_in_seconds( [new_file_handle(FileId::random(), i64::MIN, i64::MIN + 1, 0)].iter(), 3600, @@ -1383,7 +1374,7 @@ mod tests { ); assert_eq!( - (i64::MAX / 10000000 + 1) * 10000, + (i64::MAX / 10000000) * 10000, find_latest_window_in_seconds( [new_file_handle(FileId::random(), i64::MIN, i64::MAX, 0)].iter(), 10000, @@ -1392,7 +1383,7 @@ mod tests { ); assert_eq!( - Some((i64::MAX / 3600000 + 1) * 3600), + Some((i64::MAX / 3600000) * 3600), find_latest_window_in_seconds( [ new_file_handle(FileId::random(), i64::MIN, i64::MAX, 0), @@ -1404,6 +1395,114 @@ mod tests { ); } + #[test] + fn test_window_assignment_and_activity_at_boundaries() { + // Before, at, and after positive/negative boundaries, including the first + // fractional second: all picker paths must agree with half-open partitions. + for (unit, ticks_per_second) in [ + (TimeUnit::Second, 1), + (TimeUnit::Millisecond, 1_000), + (TimeUnit::Microsecond, 1_000_000), + (TimeUnit::Nanosecond, 1_000_000_000), + ] { + for window_size in [1, 3, 86_400] { + for boundary in [-window_size, 0, window_size] { + for (offset, expected_window) in [ + (-1, boundary - window_size), + (0, boundary), + (ticks_per_second - 1, boundary), + (ticks_per_second, boundary + i64::from(window_size == 1)), + ] { + let timestamp = Timestamp::new(boundary * ticks_per_second + offset, unit); + let file = FileHandle::new( + FileMeta { + time_range: (timestamp, timestamp), + sequence: NonZeroU64::new(1), + ..Default::default() + }, + crate::test_util::new_noop_file_purger(), + ); + let files = [file]; + let windows = assign_to_windows(files.iter(), window_size); + assert_eq!( + vec![expected_window], + windows.keys().copied().collect::>(), + "{timestamp:?}, window size {window_size}" + ); + assert_eq!( + Some(expected_window), + find_active_window_by_sequence(files.iter(), window_size) + ); + assert_eq!( + Some(expected_window), + find_latest_window_in_seconds(files.iter(), window_size) + ); + } + } + } + } + } + + #[test] + fn test_window_assignment_at_timestamp_limits() { + for (seconds, expected_window) in [(i64::MIN, i64::MIN), (i64::MAX, i64::MAX - 1)] { + let timestamp = Timestamp::new_second(seconds); + let files = [FileHandle::new( + FileMeta { + time_range: (timestamp, timestamp), + sequence: NonZeroU64::new(1), + ..Default::default() + }, + crate::test_util::new_noop_file_purger(), + )]; + let windows = assign_to_windows(files.iter(), 3); + assert_eq!( + vec![expected_window], + windows.keys().copied().collect::>() + ); + assert_eq!( + Some(expected_window), + find_active_window_by_sequence(files.iter(), 3) + ); + assert_eq!( + Some(expected_window), + find_latest_window_in_seconds(files.iter(), 3) + ); + assert!(time_window_intersects_range( + expected_window, + 3, + &TimestampRange::single(timestamp) + )); + } + } + + #[tokio::test] + async fn test_midnight_file_is_not_compacted_with_previous_day() { + // Regression for #9441: flush already separates the midnight-only SST. + let midnight = 86_400_000; + CompactionPickerTestCase { + window_size: 86_400, + input_files: [ + (midnight - 240_000, midnight - 180_000), + (midnight - 180_000, midnight - 90_000), + (midnight - 90_000, midnight - 30_000), + (midnight, midnight), + ] + .into_iter() + .enumerate() + .map(|(index, (start, end))| { + new_file_handle_with_sequence(FileId::random(), start, end, 0, index as u64 + 1) + }) + .collect(), + expected_outputs: vec![ExpectedOutput { + input_files: vec![0, 1, 2], + output_level: 1, + }], + } + .check() + .await; + } + #[test] fn test_assign_to_windows() { let windows = assign_to_windows( @@ -1433,7 +1532,7 @@ mod tests { assert_eq!( files[0], windows - .get(&0) + .get(&-3) .unwrap() .files() .next() @@ -1444,7 +1543,7 @@ mod tests { assert_eq!( files[1], windows - .get(&3) + .get(&0) .unwrap() .files() .next() @@ -1455,7 +1554,7 @@ mod tests { assert_eq!( files[2], windows - .get(&12) + .get(&9) .unwrap() .files() .next() @@ -1550,6 +1649,7 @@ mod tests { .collect::>(); let windows = assign_to_windows(file_handles.iter(), time_window); + assert_eq!(expected_files.len(), windows.len()); for (expected_window, overlapping, window_files) in expected_files { let actual_window = windows.get(expected_window).unwrap(); @@ -1574,8 +1674,8 @@ mod tests { &[(0, 999), (1000, 1999), (2000, 2999)], 2, &[ - (0, false, vec![(0, 999)]), - (2, false, vec![(1000, 1999), (2000, 2999)]), + (0, false, vec![(0, 999), (1000, 1999)]), + (2, false, vec![(2000, 2999)]), ], ); @@ -1592,9 +1692,8 @@ mod tests { &[(0, 999), (1000, 1999), (2000, 2999), (3000, 3999)], 2, &[ - (0, false, vec![(0, 999)]), - (2, false, vec![(1000, 1999), (2000, 2999)]), - (4, false, vec![(3000, 3999)]), + (0, false, vec![(0, 999), (1000, 1999)]), + (2, false, vec![(2000, 2999), (3000, 3999)]), ], ); @@ -1608,9 +1707,8 @@ mod tests { ], 2, &[ - (0, true, vec![(0, 999)]), - (2, true, vec![(1000, 1999), (2000, 2999)]), - (4, true, vec![(0, 3999), (3000, 3999)]), + (0, true, vec![(0, 999), (1000, 1999)]), + (2, true, vec![(0, 3999), (2000, 2999), (3000, 3999)]), ], ); @@ -1624,43 +1722,39 @@ mod tests { ], 2, &[ - (0, false, vec![(0, 999)]), - (2, true, vec![(1000, 1999), (2000, 2999)]), - (4, true, vec![(1999, 3999), (3000, 3999)]), + (0, true, vec![(0, 999), (1000, 1999)]), + (2, true, vec![(1999, 3999), (2000, 2999), (3000, 3999)]), ], ); check_assign_to_windows_with_overlapping( &[ (0, 999), // window 0 - (1000, 1999), // window 2 + (1000, 1999), // window 0 (2000, 2999), // window 2 - (3000, 3999), // window 4 - (2999, 3999), // window 4 + (3000, 3999), // window 2 + (2999, 3999), // window 2 ], 2, &[ - // window 2 overlaps with window 4 - (0, false, vec![(0, 999)]), - (2, true, vec![(1000, 1999), (2000, 2999)]), - (4, true, vec![(2999, 3999), (3000, 3999)]), + // Overlap within a window does not imply cross-window overlap. + (0, false, vec![(0, 999), (1000, 1999)]), + (2, false, vec![(2000, 2999), (2999, 3999), (3000, 3999)]), ], ); check_assign_to_windows_with_overlapping( &[ (0, 999), // window 0 - (1000, 1999), // window 2 + (1000, 1999), // window 0 (2000, 2999), // window 2 - (3000, 3999), // window 4 - (0, 1000), // // window 2 + (3000, 3999), // window 2 + (0, 2000), // window 2 ], 2, &[ - // only window 0 overlaps with window 2. - (0, true, vec![(0, 999)]), - (2, true, vec![(0, 1000), (1000, 1999), (2000, 2999)]), - (4, false, vec![(3000, 3999)]), + (0, true, vec![(0, 999), (1000, 1999)]), + (2, true, vec![(0, 2000), (2000, 2999), (3000, 3999)]), ], ); } @@ -1680,7 +1774,7 @@ mod tests { new_file_handle_with_size_sequence_and_primary_key_range( FileId::random(), 500, - 1999, + 2999, 0, 2, 10, @@ -1690,7 +1784,8 @@ mod tests { let windows = assign_to_windows(files.iter(), 2); - let overlapping = window_has_overlap(windows.get(&2).unwrap(), &windows); + assert_eq!(2, windows.len()); + let overlapping = window_has_overlap(windows.get(&0).unwrap(), &windows); assert!(!overlapping); } @@ -2410,7 +2505,7 @@ mod tests { }; let output = picker - .build_output_with_time_range(RegionId::from_u64(123), windows, Some(9), None) + .build_output_with_time_range(RegionId::from_u64(123), windows, Some(6), None) .await .unwrap(); @@ -2427,27 +2522,60 @@ mod tests { #[test] fn test_filter_time_windows_by_time_range() { - let time_range = TimestampRange::new( - Timestamp::new_millisecond(1_200), - Timestamp::new_millisecond(1_800), - ) - .unwrap(); + // Interior, exclusive boundary ends, fractional boundary ends, negative + // timestamps, and empty ranges must all use the same half-open windows. + for (start, end, expected_windows) in [ + (1_200, 1_800, vec![0]), + (0, 3_000, vec![0]), + (3_000, 6_000, vec![3]), + (2_999, 3_000, vec![0]), + (2_999, 3_001, vec![0, 3]), + (3_000, 3_001, vec![3]), + (3_001, 3_002, vec![3]), + (-3_000, 0, vec![-3]), + (-3_001, -2_999, vec![-6, -3]), + (0, 0, vec![]), + (1_500, 1_500, vec![]), + ] { + let time_range = TimestampRange::with_unit(start, end, TimeUnit::Millisecond).unwrap(); + let selected = [-6, -3, 0, 3, 6] + .into_iter() + .filter(|window| time_window_intersects_range(*window, 3, &time_range)) + .collect::>(); + assert_eq!(expected_windows, selected, "{time_range}"); + } + } - assert!(time_window_intersects_range(3, 3, &time_range)); - assert!(!time_window_intersects_range(9, 3, &time_range)); + #[test] + fn test_filter_time_windows_with_unbounded_and_extreme_ranges() { + for (range, expected_windows) in [ + (TimestampRange::min_to_max(), vec![-3, 0, 3]), + ( + TimestampRange::from_start(Timestamp::new_second(0)), + vec![0, 3], + ), + ( + TimestampRange::until_end(Timestamp::new_second(0), false), + vec![-3], + ), + ] { + let selected = [-3, 0, 3] + .into_iter() + .filter(|window| time_window_intersects_range(*window, 3, &range)) + .collect::>(); + assert_eq!(expected_windows, selected, "{range}"); + } - let boundary_range = - TimestampRange::new(Timestamp::new_second(0), Timestamp::new_second(3)).unwrap(); - assert!(time_window_intersects_range(0, 3, &boundary_range)); - assert!(time_window_intersects_range(3, 3, &boundary_range)); - assert!(!time_window_intersects_range(6, 3, &boundary_range)); - - let overflowing_range = TimestampRange::new( + let upper_range = TimestampRange::new( Timestamp::new_second(i64::MAX - 1), Timestamp::new_second(i64::MAX), ) .unwrap(); - assert!(!time_window_intersects_range(0, 4, &overflowing_range)); + assert!(!time_window_intersects_range(0, 4, &upper_range)); + assert!(time_window_intersects_range(i64::MAX - 3, 4, &upper_range)); + + let lower_range = TimestampRange::until_end(Timestamp::new_second(i64::MIN), false); + assert!(!time_window_intersects_range(i64::MIN, 3, &lower_range)); } #[tokio::test] @@ -2477,7 +2605,7 @@ mod tests { }; let output = picker - .build_output_with_time_range(RegionId::from_u64(123), windows, Some(9), Some(3)) + .build_output_with_time_range(RegionId::from_u64(123), windows, Some(6), Some(3)) .await .unwrap(); diff --git a/src/mito2/src/read/range_cache.rs b/src/mito2/src/read/range_cache.rs index ddf8c3cc356..1108b074697 100644 --- a/src/mito2/src/read/range_cache.rs +++ b/src/mito2/src/read/range_cache.rs @@ -736,26 +736,33 @@ struct CacheBatchBuffer { buffered_rows: usize, buffered_size: usize, sender: Option>, + /// Handle of the background concat task. Cleared once the task owns a finish + /// command, so only unfinished scans cancel it. + concat_task: Option>, } impl CacheBatchBuffer { fn new(cache_strategy: &CacheStrategy) -> Self { - let sender = cache_strategy.range_result_memory_limiter().map(|limiter| { - let skip_threshold_bytes = cache_strategy.range_result_cache_size().unwrap_or(0); - let (tx, rx) = mpsc::unbounded_channel(); - common_runtime::spawn_query(run_cache_concat_task( - rx, - limiter.clone(), - skip_threshold_bytes, - )); - tx - }); + let (sender, concat_task) = cache_strategy + .range_result_memory_limiter() + .map(|limiter| { + let skip_threshold_bytes = cache_strategy.range_result_cache_size().unwrap_or(0); + let (tx, rx) = mpsc::unbounded_channel(); + let task = common_runtime::spawn_query(run_cache_concat_task( + rx, + limiter.clone(), + skip_threshold_bytes, + )); + (tx, task) + }) + .unzip(); Self { buffered_batches: Vec::new(), buffered_rows: 0, buffered_size: 0, sender, + concat_task, } } @@ -814,9 +821,22 @@ impl CacheBatchBuffer { part_metrics, result_tx, }) - .is_err() + .is_ok() { - self.sender = None; + // The task now owns the finish command, so it may keep populating the + // cache after the scan stream is dropped. + self.concat_task = None; + } + } +} + +impl Drop for CacheBatchBuffer { + fn drop(&mut self) { + // Still holding the handle means no finish command was sent: the scan was + // cancelled or failed and the queued batches can never reach the cache. + // Aborting releases them even while the task waits for a memory permit. + if let Some(task) = &self.concat_task { + task.abort(); } } } @@ -1570,6 +1590,43 @@ mod tests { assert_eq!(replayed[1].num_rows(), 1); } + #[tokio::test] + async fn dropped_cache_buffer_releases_batches_while_waiting_for_memory() { + let strategy = test_cache_strategy(); + let limiter = strategy.range_result_memory_limiter().unwrap(); + // Take every permit so the concat task blocks before it can compact. + let _permit = limiter + .acquire(limiter.available_permits() * limiter.permit_bytes()) + .await + .unwrap(); + let batch = make_batch(&vec![1; DEFAULT_READ_BATCH_SIZE / 2 + 1]); + let weak = Arc::downgrade(batch.column(0)); + let mut buffer = CacheBatchBuffer::new(&strategy); + buffer.push(batch.clone()).unwrap(); + buffer.push(batch).unwrap(); + // Both batches were handed to the task instead of staying in the buffer. + assert!(buffer.buffered_batches.is_empty()); + + // Cancel only once the task is parked on the permit, otherwise the abort + // could land on a task that was never polled. + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while limiter.waited_acquires() == 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("concat task must wait for a memory permit"); + drop(buffer); + + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while weak.upgrade().is_some() { + tokio::task::yield_now().await; + } + }) + .await + .expect("cancelled cache work must release input without waiting for a permit"); + } + #[tokio::test] async fn cache_batch_buffer_finishes_pending_batches() { let strategy = test_cache_strategy(); diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index aea8aa19347..0751e16d5cd 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -1674,7 +1674,15 @@ impl ScanInput { }; match maybe_batch { Some(Ok(batch)) => { - let _ = sender.send(Ok(batch)).await; + // The receiver is gone when the query is cancelled or finishes early, + // so stop reading the source. + if let Err(e) = sender.send(Ok(batch)).await { + debug!( + "Stop parallel scan task, receiver dropped, region_id: {}, error: {}", + region_id, e + ); + break; + } } Some(Err(e)) => { let _ = sender.send(Err(e)).await; @@ -3229,4 +3237,59 @@ mod tests { .build(); assert!(exact_sequence_range(&request, &version).is_err()); } + + #[tokio::test] + async fn test_parallel_flat_source_stops_after_receiver_dropped() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use datatypes::arrow::array::Int64Array; + + const SOURCE_BATCHES: usize = 10_000; + + let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false)); + let mapper = FlatProjectionMapper::new(&metadata, [0, 2, 3]).unwrap(); + let env = SchedulerEnv::new().await; + let input = ScanInput::builder(env.access_layer.clone(), mapper).build(); + + let schema = Arc::new(ArrowSchema::new(vec![Field::new( + "v", + ArrowDataType::Int64, + false, + )])); + let batch = + RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(vec![1]))]).unwrap(); + let pulled = Arc::new(AtomicUsize::new(0)); + let mut released = Vec::new(); + let sources = (0..2) + .map(|_| { + let (release_tx, release_rx) = tokio::sync::oneshot::channel::<()>(); + released.push(release_rx); + let batch = batch.clone(); + let pulled = pulled.clone(); + futures::stream::iter(0..SOURCE_BATCHES) + .map(move |_| { + // Dropped together with the source stream. + let _ = &release_tx; + pulled.fetch_add(1, Ordering::Relaxed); + Ok(batch.clone()) + }) + .boxed() + }) + .collect(); + + let streams = input + .create_parallel_flat_sources(sources, Arc::new(Semaphore::new(2)), 1) + .unwrap(); + drop(streams); + + for release_rx in released { + tokio::time::timeout(std::time::Duration::from_secs(10), release_rx) + .await + .unwrap() + .unwrap_err(); + } + // Each task may have read a few batches ahead into its channel before the + // receiver went away, but must not drain the rest of its source. + assert!(pulled.load(Ordering::Relaxed) < 100); + } } diff --git a/src/promql/src/extension_plan/histogram_fold.rs b/src/promql/src/extension_plan/histogram_fold.rs index 5a5f2f8d466..aae9df85e40 100644 --- a/src/promql/src/extension_plan/histogram_fold.rs +++ b/src/promql/src/extension_plan/histogram_fold.rs @@ -833,6 +833,8 @@ impl HistogramFoldStream { self.find_first_complete_bucket(&batch) } + /// Returns the first group's candidate bucket count once a group boundary is + /// observed. Optimistic validation checks its bounds before folding. fn find_first_complete_bucket(&self, batch: &RecordBatch) -> DataFusionResult> { if batch.num_rows() == 0 { return Ok(None); @@ -840,21 +842,14 @@ impl HistogramFoldStream { let vectors = Helper::try_into_vectors(batch.columns()) .map_err(|e| DataFusionError::Execution(e.to_string()))?; - let le_array = batch.column(self.le_column_index); - let mut tag_values_buf = Vec::with_capacity(self.normal_indices.len()); self.collect_tag_values(&vectors, 0, &mut tag_values_buf); - let mut group_start = 0usize; - for row in 0..batch.num_rows() { + for row in 1..batch.num_rows() { if !self.is_same_group(&vectors, row, &tag_values_buf) { - // new group begins - self.collect_tag_values(&vectors, row, &mut tag_values_buf); - group_start = row; - } - - if Self::is_positive_infinity(le_array, row) { - return Ok(Some(row - group_start + 1)); + // A new series/timestamp proves that the previous group is + // complete, including any numerically equal +Inf boundaries. + return Ok(Some(row)); } } @@ -875,7 +870,10 @@ impl HistogramFoldStream { let field_array = field_array.as_primitive::(); let mut tag_values_buf = Vec::with_capacity(self.normal_indices.len()); - while remaining_rows >= bucket_num && self.mode == FoldMode::Optimistic { + // Keep one row of lookahead: +Inf alone does not prove completion. + // A group ending at the batch boundary stays buffered until more input + // arrives, or flush_remaining completes it at EOF. + while remaining_rows > bucket_num && self.mode == FoldMode::Optimistic { self.collect_tag_values(&vectors, cursor, &mut tag_values_buf); if !self.validate_optimistic_group( &vectors, @@ -974,6 +972,9 @@ impl HistogramFoldStream { tag_values: &[ValueRef<'_>], ) -> bool { let inf_index = cursor + bucket_num - 1; + if self.is_same_group(vectors, inf_index + 1, tag_values) { + return false; + } if !Self::is_positive_infinity(le_array, inf_index) { return false; } @@ -1261,8 +1262,31 @@ impl HistogramFoldStream { } } - // check input value - if !bucket.windows(2).all(|w| w[0] <= w[1]) { + // Detect equal numeric boundaries during validation so unique buckets + // need neither another scan nor an allocation. + let has_duplicates = !bucket.windows(2).all(|w| w[0] < w[1]); + if has_duplicates && !bucket.windows(2).all(|w| w[0] <= w[1]) { + return Ok(f64::NAN); + } + // Counts must be coalesced before repairing monotonicity. + let coalesced = has_duplicates.then(|| { + let mut bounds = Vec::with_capacity(bucket.len()); + let mut counts = Vec::with_capacity(counter.len()); + for (&bound, &count) in bucket.iter().zip(counter) { + if bounds.last() == Some(&bound) { + *counts.last_mut().unwrap() += count; + } else { + bounds.push(bound); + counts.push(count); + } + } + (bounds, counts) + }); + let (bucket, counter) = match &coalesced { + Some((bounds, counts)) => (bounds.as_slice(), counts.as_slice()), + None => (bucket, counter), + }; + if matches!(operation, HistogramFoldOperation::Quantile(_)) && bucket.len() < 2 { return Ok(f64::NAN); } let counter = match operation { @@ -1326,26 +1350,6 @@ impl HistogramFoldStream { } fn evaluate_fraction(lower: f64, upper: f64, bucket: &[f64], counter: &[f64]) -> f64 { - let coalesced = bucket - .windows(2) - .any(|bounds| bounds[0] == bounds[1]) - .then(|| { - let mut bounds = Vec::with_capacity(bucket.len()); - let mut counts = Vec::with_capacity(counter.len()); - for (&bound, &count) in bucket.iter().zip(counter) { - if bounds.last() == Some(&bound) { - *counts.last_mut().unwrap() += count; - } else { - bounds.push(bound); - counts.push(count); - } - } - (bounds, counts) - }); - let (bucket, counter) = match &coalesced { - Some((bounds, counts)) => (bounds.as_slice(), counts.as_slice()), - None => (bucket, counter), - }; let total = *counter.last().unwrap(); if total == 0.0 || lower.is_nan() || upper.is_nan() { return f64::NAN; @@ -1453,8 +1457,8 @@ mod test { // 12 items let host_column_1 = Arc::new(StringArray::from(vec![ - "host_1", "host_1", "host_1", "host_1", "host_1", "host_1", "host_1", "host_1", - "host_1", "host_1", "host_1", "host_1", + "host_1", "host_1", "host_1", "host_1", "host_1", "host_2", "host_2", "host_2", + "host_2", "host_2", "host_3", "host_3", ])) as _; let le_column_1 = Arc::new(StringArray::from(vec![ "0.001", "0.1", "10", "1000", "+Inf", "0.001", "0.1", "10", "1000", "+inf", "0.001", @@ -1465,14 +1469,14 @@ mod test { ])) as _; // 2 items - let host_column_2 = Arc::new(StringArray::from(vec!["host_1", "host_1"])) as _; + let host_column_2 = Arc::new(StringArray::from(vec!["host_3", "host_3"])) as _; let le_column_2 = Arc::new(StringArray::from(vec!["10", "1000"])) as _; let val_column_2 = Arc::new(Float64Array::from(vec![1.0, 1.0])) as _; // 11 items let host_column_3 = Arc::new(StringArray::from(vec![ - "host_1", "host_2", "host_2", "host_2", "host_2", "host_2", "host_2", "host_2", - "host_2", "host_2", "host_2", + "host_3", "host_4", "host_4", "host_4", "host_4", "host_4", "host_5", "host_5", + "host_5", "host_5", "host_5", ])) as _; let le_column_3 = Arc::new(StringArray::from(vec![ "+INF", "0.001", "0.1", "10", "1000", "+iNf", "0.001", "0.1", "10", "1000", "+Inf", @@ -1638,10 +1642,10 @@ mod test { | host | val | +--------+-------------------+ | host_1 | 257.5 | -| host_1 | 5.05 | -| host_1 | 0.0004 | -| host_2 | NaN | -| host_2 | 6.040000000000001 | +| host_2 | 5.05 | +| host_3 | 0.0004 | +| host_4 | NaN | +| host_5 | 6.040000000000001 | +--------+-------------------+", ); assert_eq!(result_literal, expected); @@ -1818,7 +1822,7 @@ mod test { Field::new("le", DataType::Utf8, true), Field::new("val", DataType::Float64, true), ])); - let host_column = Arc::new(StringArray::from(vec!["a", "a", "a", "a", "b", "b"])) as _; + let host_column = Arc::new(StringArray::from(vec!["a", "a", "b", "b", "c", "c"])) as _; let le_column = Arc::new(StringArray::from(vec![ "0.1", "+Inf", "0.1", "1.0", "0.1", "+Inf", ])) as _; @@ -1839,8 +1843,8 @@ mod test { | host | val | +------+-----+ | a | 0.1 | -| a | NaN | -| b | 0.1 | +| b | NaN | +| c | 0.1 | +------+-----+", ); assert_eq!(result_literal, expected); @@ -1917,6 +1921,79 @@ mod test { ); } + #[tokio::test] + async fn trailing_invalid_bounds_fall_back_before_eof() { + let schema = Arc::new(Schema::new(vec![ + Field::new("host", DataType::Utf8, false), + Field::new("le", DataType::Utf8, true), + Field::new("val", DataType::Float64, false), + ])); + for trailing_bound in [Some("bad"), None] { + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec![ + "a", "a", "a", "b", "b", "b", "c", "c", "c", + ])), + Arc::new(StringArray::from( + [Some("1"), Some("+Inf"), trailing_bound].repeat(3), + )), + Arc::new(Float64Array::from([2.0, 4.0, 99.0].repeat(3))), + ], + ) + .unwrap(); + let fold = build_fold_exec_from_batches(vec![batch.clone()], schema.clone(), 0.5, 0); + let mut stream = HistogramFoldStream { + le_column_index: 1, + field_column_index: 2, + histogram_column_index: None, + operation: fold.operation, + normal_indices: vec![0], + bucket_size: None, + batch_size: 1, + output_schema: fold.output_schema.clone(), + input_schema: schema.clone(), + mode: FoldMode::Optimistic, + safe_group: None, + input_buffer: vec![], + input_buffered_rows: 0, + output_buffer: HistogramFoldStream::empty_output_buffer(&schema).unwrap(), + output_buffered_rows: 0, + input: fold + .input + .execute(0, SessionContext::default().task_ctx()) + .unwrap(), + metric: BaselineMetrics::new(&fold.metric, 0), + }; + + // Neither +Inf nor a batch boundary alone completes the group. + assert!(stream.fold_input(batch.slice(0, 2)).unwrap().is_none()); + assert!(stream.fold_input(batch.slice(2, 1)).unwrap().is_none()); + assert_eq!(stream.mode, FoldMode::Optimistic); + + for (offset, host) in [(3, "a"), (6, "b")] { + let output = stream + .fold_input(batch.slice(offset, 3)) + .unwrap() + .expect("a completed group should produce output before EOF") + .unwrap(); + assert_eq!(stream.mode, FoldMode::Safe); + assert_eq!(stream.input_buffered_rows, 0); + assert!(stream.input_buffer.is_empty()); + assert_eq!(output.num_rows(), 1); + assert_eq!(string_array_value_at_index(output.column(0), 0), Some(host)); + assert_eq!(output.column(1).as_primitive::().value(0), 1.0); + } + + stream.flush_remaining().unwrap(); + let output = stream.take_output_buf().unwrap().unwrap(); + assert_eq!(output.num_rows(), 1); + assert_eq!(string_array_value_at_index(output.column(0), 0), Some("c")); + assert_eq!(output.column(1).as_primitive::().value(0), 1.0); + assert!(stream.take_output_buf().unwrap().is_none()); + } + } + #[tokio::test] async fn safe_mode_handles_misaligned_groups() { let schema = Arc::new(Schema::new(vec![ @@ -2032,6 +2109,171 @@ mod test { assert!((values[1] - 0.55).abs() < 1e-10, "{values:?}"); } + #[test] + fn evaluate_quantile_duplicate_bounds() { + let cases = [ + // Issue #9443: coalesce before repairing the count at 1000. + ( + vec![10.0, 10.0, 100.0, 100.0, 1000.0, f64::INFINITY], + vec![40.0, 40.0, 50.0, 50.0, 50.0, 100.0], + 0.95, + 77.5, + ), + // Repairing the split counts first would incorrectly inflate them. + ( + vec![10.0, 10.0, 100.0, 100.0, f64::INFINITY], + vec![30.0, 10.0, 40.0, 20.0, 60.0], + 0.5, + 7.5, + ), + ( + vec![10.0, f64::INFINITY, f64::INFINITY], + vec![40.0, 20.0, 30.0], + 0.5, + 6.25, + ), + ( + vec![f64::INFINITY, f64::INFINITY], + vec![20.0, 30.0], + 0.5, + f64::NAN, + ), + ( + vec![-0.0, 0.0, 10.0, f64::INFINITY], + vec![10.0, 10.0, 40.0, 40.0], + 0.75, + 5.0, + ), + ]; + for (bounds, counts, quantile, expected) in cases { + let actual = HistogramFoldStream::evaluate_row( + HistogramFoldOperation::Quantile(quantile.into()), + &bounds, + &counts, + ) + .unwrap(); + if expected.is_nan() { + assert!(actual.is_nan()); + } else { + assert!((actual - expected).abs() < 1e-10, "{actual} != {expected}"); + } + } + } + + #[tokio::test] + async fn duplicate_infinity_bounds_across_batches() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("le", DataType::Utf8, false), + Field::new("val", DataType::Float64, false), + ])); + let duplicate = [("10", 20.0), ("10.0", 20.0), ("+Inf", 25.0), ("+inf", 25.0)]; + let regular = [("10", 20.0), ("10.0", 20.0), ("+Inf", 50.0)]; + // Cover discovery, an already established optimistic bucket count, + // fallback, and the final group at EOF. + for groups in [ + vec![duplicate.as_slice()], + vec![duplicate.as_slice(), regular.as_slice()], + vec![regular.as_slice(), duplicate.as_slice(), regular.as_slice()], + ] { + let mut timestamps = Vec::new(); + let mut bounds = Vec::new(); + let mut counts = Vec::new(); + for (timestamp, group) in groups.iter().enumerate() { + for &(bound, count) in *group { + timestamps.push(timestamp as i64); + bounds.push(bound); + counts.push(count); + } + } + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(timestamps)), + Arc::new(StringArray::from(bounds)), + Arc::new(Float64Array::from(counts)), + ], + ) + .unwrap(); + for batch_size in 1..=batch.num_rows() { + let batches = (0..batch.num_rows()) + .step_by(batch_size) + .map(|offset| batch.slice(offset, batch_size.min(batch.num_rows() - offset))) + .collect(); + let fold = build_fold_exec_from_batches(batches, schema.clone(), 0.5, 0); + let result = + datafusion::physical_plan::collect(fold, SessionContext::default().task_ctx()) + .await + .unwrap(); + let values: Vec<_> = result + .iter() + .flat_map(|batch| { + batch + .column(1) + .as_primitive::() + .values() + .iter() + .copied() + }) + .collect(); + assert_eq!(values, vec![6.25; groups.len()], "batch size {batch_size}"); + } + } + } + + #[tokio::test] + async fn duplicate_bounds_across_batches_and_safe_fallback() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("le", DataType::Utf8, false), + Field::new("val", DataType::Float64, false), + ])); + // The first group establishes six buckets for optimistic folding. The + // shorter second group forces safe mode; both have duplicate bounds. + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![ + 1000, 1000, 1000, 1000, 1000, 1000, 2000, 2000, 2000, 2000, 2000, + ])), + Arc::new(StringArray::from(vec![ + "10", "10.0", "100", "100.0", "1000", "+Inf", "10", "10.0", "100", "100.0", + "+Inf", + ])), + Arc::new(Float64Array::from(vec![ + 40.0, 40.0, 50.0, 50.0, 50.0, 100.0, 40.0, 40.0, 50.0, 50.0, 100.0, + ])), + ], + ) + .unwrap(); + // Split each pair of equal boundaries across an input batch boundary. + let batches = (0..batch.num_rows()).map(|i| batch.slice(i, 1)).collect(); + let fold = build_fold_exec_from_batches(batches, schema, 0.95, 0); + let result = datafusion::physical_plan::collect(fold, SessionContext::default().task_ctx()) + .await + .unwrap(); + let values: Vec<_> = result + .iter() + .flat_map(|batch| { + batch + .column(1) + .as_primitive::() + .values() + .iter() + .copied() + }) + .collect(); + assert_eq!(values, vec![77.5, 77.5]); + } + #[test] fn evaluate_row_normal_case() { let bucket = [0.0, 1.0, 2.0, 3.0, 4.0, f64::INFINITY]; diff --git a/src/promql/src/extension_plan/series_divide.rs b/src/promql/src/extension_plan/series_divide.rs index 9e6a46d991a..090ead28c3e 100644 --- a/src/promql/src/extension_plan/series_divide.rs +++ b/src/promql/src/extension_plan/series_divide.rs @@ -554,7 +554,9 @@ impl Stream for SeriesDivideStream { } error => return Poll::Ready(error), }; - self.buffer.push(batch); + if batch.num_rows() != 0 { + self.buffer.push(batch); + } continue; } } @@ -624,7 +626,8 @@ impl SeriesDivideStream { #[cfg(test)] mod test { use datafusion::arrow::array::{ - DictionaryArray, Int32Array, LargeStringArray, StringArray, StringViewArray, UInt32Array, + DictionaryArray, Int32Array, Int64Array, LargeStringArray, StringArray, StringViewArray, + UInt32Array, }; use datafusion::arrow::datatypes::{DataType, Field, Int32Type, Schema, UInt32Type}; use datafusion::common::ToDFSchema; @@ -1245,4 +1248,216 @@ mod test { assert!(tsid_array.iter().all(|v| v == Some(expected_tsid))); } } + + /// The input stream may legally start with a batch that has zero rows. + /// `SeriesDivideStream` must skip it and keep working on the real rows that + /// follow (here: two batches of the very same tagged series). + #[tokio::test] + async fn test_stream_leading_empty_batch_then_same_tags() { + let schema = Arc::new(Schema::new(vec![ + Field::new("host", DataType::Utf8, false), + Field::new("path", DataType::Utf8, false), + Field::new( + "time_index", + DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + ])); + + let make_batch = |hosts: Vec<&str>, paths: Vec<&str>, times: Vec| { + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(hosts)), + Arc::new(StringArray::from(paths)), + Arc::new(datafusion::arrow::array::TimestampMillisecondArray::from( + times, + )), + ], + ) + .unwrap() + }; + + let memory_exec: Arc = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new( + &[vec![ + RecordBatch::new_empty(schema.clone()), + make_batch( + vec!["server1", "server1"], + vec!["/var/log", "/var/log"], + vec![1000, 2000], + ), + make_batch(vec!["server1"], vec!["/var/log"], vec![3000]), + ]], + schema.clone(), + None, + ) + .unwrap(), + ))); + + let divide_exec = Arc::new(SeriesDivideExec { + tag_columns: vec!["host".to_string(), "path".to_string()], + time_index_column: "time_index".to_string(), + input: memory_exec, + metric: ExecutionPlanMetricsSet::new(), + }); + + let mut stream = divide_exec + .execute(0, SessionContext::default().task_ctx()) + .unwrap(); + + let batch = stream.next().await.unwrap().unwrap(); + assert_eq!(batch.num_rows(), 3); + let expected = [ + ("server1", "/var/log", 1000i64), + ("server1", "/var/log", 2000), + ("server1", "/var/log", 3000), + ]; + for (row, (host, path, time)) in expected.into_iter().enumerate() { + assert_eq!( + string_array_value_at_index(batch.column(0), row), + Some(host) + ); + assert_eq!( + string_array_value_at_index(batch.column(1), row), + Some(path) + ); + assert_eq!( + batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .value(row), + time + ); + } + assert!(stream.next().await.is_none()); + } + + /// Several leading empty batches (and one empty batch in the middle of a + /// series) must not disturb series division: the two series are still + /// emitted separately, including the last one flushed at end of input. + #[tokio::test] + async fn test_stream_repeated_empty_batches_and_two_series() { + let schema = Arc::new(Schema::new(vec![ + Field::new("host", DataType::Utf8, false), + Field::new("value", DataType::Int64, false), + Field::new( + "time_index", + DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + ])); + + let make_batch = |hosts: Vec<&str>, values: Vec| { + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(hosts)), + Arc::new(Int64Array::from(values.clone())), + Arc::new(datafusion::arrow::array::TimestampMillisecondArray::from( + values.iter().map(|v| v * 1000).collect::>(), + )), + ], + ) + .unwrap() + }; + + let memory_exec: Arc = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new( + &[vec![ + RecordBatch::new_empty(schema.clone()), + RecordBatch::new_empty(schema.clone()), + make_batch(vec!["server1", "server1"], vec![1, 2]), + RecordBatch::new_empty(schema.clone()), + make_batch(vec!["server1"], vec![3]), + make_batch(vec!["server2", "server2"], vec![4, 5]), + ]], + schema.clone(), + None, + ) + .unwrap(), + ))); + + let divide_exec = Arc::new(SeriesDivideExec { + tag_columns: vec!["host".to_string()], + time_index_column: "time_index".to_string(), + input: memory_exec, + metric: ExecutionPlanMetricsSet::new(), + }); + + let mut stream = divide_exec + .execute(0, SessionContext::default().task_ctx()) + .unwrap(); + + for (expected_host, expected_values) in + [("server1", vec![1i64, 2, 3]), ("server2", vec![4, 5])] + { + let batch = stream.next().await.unwrap().unwrap(); + assert_eq!(batch.num_rows(), expected_values.len()); + for (row, value) in expected_values.into_iter().enumerate() { + assert_eq!( + string_array_value_at_index(batch.column(0), row), + Some(expected_host) + ); + assert_eq!( + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(row), + value + ); + } + } + assert!(stream.next().await.is_none()); + } + + /// A stream that only ever yields zero-row batches must terminate with no + /// batches instead of failing. + #[tokio::test] + async fn test_stream_all_empty_batches_yields_nothing() { + let schema = Arc::new(Schema::new(vec![ + Field::new("host", DataType::Utf8, false), + Field::new("value", DataType::Int64, false), + Field::new( + "time_index", + DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + ])); + + let memory_exec: Arc = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new( + &[vec![ + RecordBatch::new_empty(schema.clone()), + RecordBatch::new_empty(schema.clone()), + RecordBatch::new_empty(schema.clone()), + ]], + schema.clone(), + None, + ) + .unwrap(), + ))); + + let divide_exec = Arc::new(SeriesDivideExec { + tag_columns: vec!["host".to_string()], + time_index_column: "time_index".to_string(), + input: memory_exec, + metric: ExecutionPlanMetricsSet::new(), + }); + + let mut stream = divide_exec + .execute(0, SessionContext::default().task_ctx()) + .unwrap(); + assert!(stream.next().await.is_none()); + + let collected = + datafusion::physical_plan::collect(divide_exec, SessionContext::default().task_ctx()) + .await + .unwrap(); + assert!(collected.is_empty()); + } } diff --git a/src/query/src/optimizer/constant_term.rs b/src/query/src/optimizer/constant_term.rs index 47f3faf59d7..3f7e987b96e 100644 --- a/src/query/src/optimizer/constant_term.rs +++ b/src/query/src/optimizer/constant_term.rs @@ -205,10 +205,14 @@ impl PhysicalOptimizerRule for MatchesConstantTermOptimizer { })?; if new_pred.transformed { - let exec = FilterExecBuilder::new(new_pred.data, filter.input().clone()) - .with_default_selectivity(filter.default_selectivity()) - .apply_projection_by_ref(filter.projection().as_ref()) - .and_then(|x| x.build())?; + // `FilterExecBuilder::from(&FilterExec)` already carries over the + // input, projection, default selectivity, batch size and fetch, so + // only the predicate has to be replaced. Re-applying the projection + // is not only pointless but wrong: `apply_projection` composes with + // the projection copied from the builder. + let exec = FilterExecBuilder::from(filter) + .with_predicate(new_pred.data) + .build()?; return Ok(Transformed::yes(Arc::new(exec) as _)); } } diff --git a/src/query/src/promql/error.rs b/src/query/src/promql/error.rs index dba55c8df8b..45e08349623 100644 --- a/src/query/src/promql/error.rs +++ b/src/query/src/promql/error.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::any::Any; +use std::time::SystemTime; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; @@ -202,6 +203,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Time out of the representable millisecond range: {:?}", time))] + SystemTimeOutOfRange { + time: SystemTime, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("vector cannot contain metrics with the same labelset"))] SameLabelSet { #[snafu(implicit)] @@ -243,6 +251,7 @@ impl ErrorExt for Error { | UnsupportedMatcherOp { .. } | SameLabelSet { .. } | TimestampOutOfRange { .. } + | SystemTimeOutOfRange { .. } | InvalidRegularExpression { .. } | InvalidDestinationLabelName { .. } => StatusCode::InvalidArguments, diff --git a/src/query/src/promql/label_values.rs b/src/query/src/promql/label_values.rs index 4f6ea42a60b..44c173477bb 100644 --- a/src/query/src/promql/label_values.rs +++ b/src/query/src/promql/label_values.rs @@ -23,9 +23,31 @@ use snafu::{OptionExt, ResultExt}; use table::TableRef; use crate::promql::error::{ - DataFusionPlanningSnafu, Result, TimeIndexNotFoundSnafu, TimestampOutOfRangeSnafu, + DataFusionPlanningSnafu, Result, SystemTimeOutOfRangeSnafu, TimeIndexNotFoundSnafu, + TimestampOutOfRangeSnafu, }; +/// Converts a [`SystemTime`] to a millisecond [`Timestamp`]. +/// +/// `duration_since` reports an instant before the epoch as an error rather than +/// a negative duration, and an RFC3339 query parameter can name one, so the sign +/// is recovered here instead of unwrapping. +fn millis_since_epoch(time: SystemTime) -> Result { + let (millis, before_epoch) = match time.duration_since(UNIX_EPOCH) { + Ok(duration) => (duration.as_millis(), false), + Err(earlier) => (earlier.duration().as_millis(), true), + }; + let millis = i64::try_from(millis) + .ok() + .with_context(|| SystemTimeOutOfRangeSnafu { time })?; + + Ok(Timestamp::new_millisecond(if before_epoch { + -millis + } else { + millis + })) +} + fn build_time_filter(time_index_expr: Expr, start: Timestamp, end: Timestamp) -> Expr { time_index_expr .clone() @@ -67,14 +89,12 @@ pub fn rewrite_label_values_query( })?; // We only support millisecond precision at most. - let start = - Timestamp::new_millisecond(start.duration_since(UNIX_EPOCH).unwrap().as_millis() as i64); + let start = millis_since_epoch(start)?; let start = start.convert_to(unit).context(TimestampOutOfRangeSnafu { timestamp: start.value(), unit, })?; - let end = - Timestamp::new_millisecond(end.duration_since(UNIX_EPOCH).unwrap().as_millis() as i64); + let end = millis_since_epoch(end)?; let end = end.convert_to(unit).context(TimestampOutOfRangeSnafu { timestamp: end.value(), unit, @@ -98,3 +118,28 @@ pub fn rewrite_label_values_query( Ok(logical_plan) } + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use super::*; + + #[test] + fn millis_before_the_epoch_are_negative() { + // `SystemTime::duration_since` reports these as an error; unwrapping it + // panicked on any request whose RFC3339 start named a pre-epoch instant. + let time = UNIX_EPOCH - Duration::from_millis(1); + assert_eq!(millis_since_epoch(time).unwrap().value(), -1); + assert_eq!(millis_since_epoch(UNIX_EPOCH).unwrap().value(), 0); + + let time = UNIX_EPOCH + Duration::from_millis(1); + assert_eq!(millis_since_epoch(time).unwrap().value(), 1); + } + + #[test] + fn millis_beyond_i64_are_rejected() { + let time = UNIX_EPOCH + Duration::from_secs(1 << 60); + assert!(millis_since_epoch(time).is_err()); + } +} diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 7f1d5d35fbe..d1d98330b97 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -1554,6 +1554,18 @@ impl PromPlanner { } } else if right_is_empty_metric && !left_is_empty_metric { self.ctx = left_context.clone(); + } else if rhs.value_type() == ValueType::Scalar + && lhs.value_type() == ValueType::Vector + { + // A computed scalar on the right (`time()`, `scalar(...)`) has no labels; + // the result keeps the labels of the vector on the left. + self.ctx = left_context.clone(); + } + // An operand without a table name, such as a comparison filter's projection, is + // joined under the `""` qualifier; resolve the projected columns against it. + // Otherwise a time index present on both sides is ambiguous. + if self.ctx.table_name.is_none() { + self.ctx.table_name = Some(String::new()); } // Computed scalars reach this join path instead of the literal projection paths. // Broadcast them for arithmetic in the same way as literal scalars. @@ -2596,13 +2608,17 @@ impl PromPlanner { let accepts_empty = matcher.is_match(""); let column_name = Self::find_case_sensitive_column(table_schema, matcher.name.as_str()); + // Prometheus reads a label a series does not carry as the empty + // string. A row can miss a label two ways: the table has no column + // for it, or the column exists but is NULL on that row — the latter + // is the norm for logical metrics sharing a physical table, which + // holds the union of their label columns. let col = if let Some(column_name) = column_name { let column = DfExpr::Column(Column::from_name(&column_name)); let field = table_schema .index_of_column_by_name(None, &column_name) .map(|index| table_schema.field(index)); if accepts_empty - && column_name == OTLP_AGGREGATION_TEMPORALITY_LABEL && let Some(data_type) = field .filter(|field| { field.is_nullable() @@ -2621,7 +2637,6 @@ impl PromPlanner { } } else { DfExpr::Literal(ScalarValue::Utf8(Some(String::new())), None) - .alias(matcher.name.clone()) }; let lit = DfExpr::Literal(ScalarValue::Utf8(Some(matcher.value)), None); let expr = match matcher.op { diff --git a/src/query/src/promql/planner/test/delta.rs b/src/query/src/promql/planner/test/delta.rs index 93e214d45b5..abfb7b35c03 100644 --- a/src/query/src/promql/planner/test/delta.rs +++ b/src/query/src/promql/planner/test/delta.rs @@ -521,7 +521,7 @@ async fn delta_mixed_ranges_drop_and_float_ranges_sum() { } #[tokio::test] -async fn temporality_matchers_treat_null_as_absent() { +async fn matchers_read_absent_labels_as_empty() { let marker = OTLP_AGGREGATION_TEMPORALITY_LABEL; let schema = Arc::new(ArrowSchema::new(vec![Field::new( marker, @@ -593,31 +593,82 @@ async fn temporality_matchers_treat_null_as_absent() { ); } - let ordinary_schema = Arc::new(ArrowSchema::new(vec![Field::new( + // The rule is about NULL, not about the marker: any nullable label column + // can be NULL where the series does not carry the label. A non-nullable one + // has nothing to normalize. + for (nullable, wants_coalesce) in [(true, true), (false, false)] { + let ordinary_schema = Arc::new(ArrowSchema::new(vec![Field::new( + "label", + ArrowDataType::Utf8, + nullable, + )])); + let ordinary_scan = LogicalPlanBuilder::scan( + "ordinary_labels", + provider_as_source(Arc::new( + MemTable::try_new(ordinary_schema, vec![vec![]]).unwrap(), + )), + None, + ) + .unwrap() + .build() + .unwrap(); + let PromExpr::VectorSelector(selector) = + parser::parse(r#"metric{label!="delta"}"#).unwrap() + else { + unreachable!() + }; + let expressions = PromPlanner::matchers_to_expr(selector.matchers, ordinary_scan.schema()) + .unwrap() + .iter() + .map(ToString::to_string) + .join(" AND "); + assert_eq!( + wants_coalesce, + expressions.contains("coalesce"), + "nullable={nullable}: {expressions}" + ); + } + + // The other way a label goes absent is having no column for it at all. The + // literal standing in for the label has to be usable as a predicate. + let schema = Arc::new(ArrowSchema::new(vec![Field::new( "label", ArrowDataType::Utf8, true, )])); - let ordinary_scan = LogicalPlanBuilder::scan( + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(StringArray::from(vec![Some("value")]))], + ) + .unwrap(); + let scan = LogicalPlanBuilder::scan( "ordinary_labels", provider_as_source(Arc::new( - MemTable::try_new(ordinary_schema, vec![vec![]]).unwrap(), + MemTable::try_new(schema, vec![vec![batch]]).unwrap(), )), None, ) .unwrap() .build() .unwrap(); - let PromExpr::VectorSelector(selector) = parser::parse(r#"metric{label!="delta"}"#).unwrap() + let PromExpr::VectorSelector(selector) = parser::parse(r#"metric{absent!="delta"}"#).unwrap() else { unreachable!() }; - let expressions = PromPlanner::matchers_to_expr(selector.matchers, ordinary_scan.schema()) + let expressions = PromPlanner::matchers_to_expr(selector.matchers, scan.schema()).unwrap(); + assert_eq!( + r#"Utf8("") != Utf8("delta")"#, + expressions.iter().map(ToString::to_string).join(" AND ") + ); + + let plan = LogicalPlanBuilder::from(scan) + .filter(conjunction(expressions).unwrap()) .unwrap() - .iter() - .map(ToString::to_string) - .join(" AND "); - assert!(!expressions.contains("coalesce"), "{expressions}"); + .build() + .unwrap(); + let (_, batches) = execute(plan, &build_query_engine_state()).await; + let rows = batches.iter().map(|batch| batch.num_rows()).sum::(); + assert_eq!(1, rows); } #[tokio::test] diff --git a/src/servers/src/addrs.rs b/src/servers/src/addrs.rs index 07422028f97..d68568726fa 100644 --- a/src/servers/src/addrs.rs +++ b/src/servers/src/addrs.rs @@ -25,7 +25,7 @@ pub fn resolve_addr(bind_addr: &str, server_addr: Option<&str>) -> String { } else { // otherwise, resolve port from bind_addr // should be safe to unwrap here because bind_addr is already validated - let port = bind_addr.split(':').nth(1).unwrap(); + let port = bind_addr.rsplit_once(':').unwrap().1; format!("{server_addr}:{port}") } } @@ -50,6 +50,16 @@ mod tests { super::resolve_addr("127.0.0.1:3001", Some("tomcat:3002")) ); + assert_eq!( + "tomcat:3001", + super::resolve_addr("[::]:3001", Some("tomcat")) + ); + + assert_eq!( + "tomcat:3001", + super::resolve_addr("[2001:db8::10]:3001", Some("tomcat")) + ); + assert_eq!( "127.0.0.1:3001", super::resolve_addr("127.0.0.1:3001", None) diff --git a/src/servers/src/http/prometheus.rs b/src/servers/src/http/prometheus.rs index 1f9988ede4c..da89d2a7b6c 100644 --- a/src/servers/src/http/prometheus.rs +++ b/src/servers/src/http/prometheus.rs @@ -49,13 +49,13 @@ use futures::future::join_all; use futures::{StreamExt, TryStreamExt}; use itertools::Itertools; use promql_parser::label::{METRIC_NAME, MatchOp, Matcher, Matchers}; -use promql_parser::parser::token::{self}; +use promql_parser::parser::token::TokenType; use promql_parser::parser::value::ValueType; use promql_parser::parser::{ - AggregateExpr, BinaryExpr, Call, Expr as PromqlExpr, LabelModifier, MatrixSelector, ParenExpr, - SubqueryExpr, UnaryExpr, VectorSelector, + AggregateExpr, BinModifier, BinaryExpr, Call, Expr as PromqlExpr, LabelModifier, + MatrixSelector, ParenExpr, SubqueryExpr, UnaryExpr, VectorMatchCardinality, VectorSelector, }; -use query::parser::{DEFAULT_LOOKBACK_STRING, PromQuery, QueryStatement}; +use query::parser::{DEFAULT_LOOKBACK_STRING, PromQuery, QueryLanguageParser, QueryStatement}; use serde::de::{self, MapAccess, Visitor}; use serde::{Deserialize, Serialize}; use serde_json::Value; @@ -74,7 +74,7 @@ use table::requests::{ pub use super::result::prometheus_resp::{PromSampleValue, PrometheusJsonResponse}; use crate::error::{ CollectRecordbatchSnafu, ConvertScalarValueSnafu, DataFusionSnafu, Error, InvalidQuerySnafu, - NotSupportedSnafu, Result, TableNotFoundSnafu, UnexpectedResultSnafu, + NotSupportedSnafu, ParseTimestampSnafu, Result, TableNotFoundSnafu, UnexpectedResultSnafu, }; use crate::http::header::collect_plan_metrics; use crate::otlp::metrics::ucum_to_openmetrics_unit; @@ -1373,6 +1373,72 @@ fn promql_expr_to_metric_name(expr: &PromqlExpr) -> Option { } } +/// Whether the metric name survives a binary operation; see +/// [`binary_metric_name_side`] for the operand that supplies it. +/// +/// Follows Prometheus' `shouldDropMetricName` and `resultMetric`: set operators +/// and comparisons keep the name, while arithmetic operators drop it. A +/// comparison also drops it when it returns a bool or is a one-to-one +/// `on(...)` match that does not keep `__name__` or an `ignoring(...)` match +/// that excludes it. Prometheus keeps only the `on(...)` labels or deletes the +/// `ignoring(...)` labels for one-to-one matches. +fn binary_keeps_metric_name(op: &TokenType, modifier: Option<&BinModifier>) -> bool { + if op.is_set_operator() { + return true; + } + + if !op.is_comparison_operator() { + return false; + } + + let Some(m) = modifier else { + return true; + }; + + if m.return_bool { + return false; + } + + if m.card == VectorMatchCardinality::OneToOne { + match &m.matching { + Some(LabelModifier::Include(labels)) + if !labels.labels.contains(&METRIC_NAME.to_string()) => + { + return false; + } + Some(LabelModifier::Exclude(labels)) + if labels.labels.contains(&METRIC_NAME.to_string()) => + { + return false; + } + _ => {} + } + } + + true +} + +/// The operand whose labels supply the result metric name of a binary +/// operation. +/// +/// Prometheus builds the result from the side whose labels it carries, i.e. the +/// "many" side of the match: `group_right` makes the right-hand operand that +/// side, and a scalar left-hand operand (for example `0.5 < a`) leaves the +/// vector on the right. Every other case takes the left-hand operand. +fn binary_metric_name_side<'a>( + lhs: &'a PromqlExpr, + rhs: &'a PromqlExpr, + modifier: Option<&BinModifier>, +) -> &'a PromqlExpr { + let group_right = + modifier.is_some_and(|m| matches!(m.card, VectorMatchCardinality::OneToMany(_))); + if group_right || matches!(lhs.value_type(), ValueType::Scalar) { + rhs + } else { + lhs + } +} + /// Recursively collect all metric names from a PromQL expression fn collect_metric_names(expr: &PromqlExpr, metric_names: &mut HashSet) { match expr { @@ -1395,14 +1461,16 @@ fn collect_metric_names(expr: &PromqlExpr, metric_names: &mut HashSet) { collect_metric_names(expr, metric_names) } PromqlExpr::Unary(UnaryExpr { .. }) => metric_names.clear(), - PromqlExpr::Binary(BinaryExpr { lhs, op, .. }) => { - if matches!( - op.id(), - token::T_LAND // INTERSECT - | token::T_LOR // UNION - | token::T_LUNLESS // EXCEPT - ) { - collect_metric_names(lhs, metric_names) + PromqlExpr::Binary(BinaryExpr { + lhs, + op, + rhs, + modifier, + .. + }) => { + if binary_keeps_metric_name(op, modifier.as_ref()) { + let side = binary_metric_name_side(lhs, rhs, modifier.as_ref()); + collect_metric_names(side, metric_names) } else { metric_names.clear() } @@ -1711,9 +1779,35 @@ pub async fn label_values_query( ); let catalog_manager = handler.catalog_manager(); - let mut table_names = try_call_return_response!( - retrieve_table_names(&query_ctx, catalog_manager, matches).await - ); + // An empty `match[]` enumerates every metric; otherwise only the + // selectors answerable from metadata go down that path. + let enumerate_all = matches.is_empty(); + let (metadata_selectors, label_selectors) = + try_call_return_response!(split_selectors_by_label_use(&matches)); + + let mut table_names = if enumerate_all || !metadata_selectors.is_empty() { + try_call_return_response!( + retrieve_table_names(&query_ctx, catalog_manager, metadata_selectors).await + ) + } else { + Vec::new() + }; + + if !label_selectors.is_empty() { + table_names.extend(try_call_return_response!( + retrieve_table_names_by_labels( + &handler, + label_selectors, + params.start.as_deref(), + params.end.as_deref(), + &query_ctx, + ) + .await + )); + table_names.sort_unstable(); + table_names.dedup(); + } + table_names = try_call_return_response!( handler .filter_metadata_metric_names( @@ -1897,6 +1991,131 @@ fn take_metric_name(selector: &mut VectorSelector) -> Option { Some(name) } +/// Removes every `__name__` matcher from the selector and returns them, so the +/// rest can be planned as column predicates. A name given as `VectorSelector::name` +/// comes back as an equality matcher, making both spellings filter alike. +fn take_metric_name_matchers(selector: &mut VectorSelector) -> Vec { + let mut taken = Vec::new(); + if let Some(name) = selector.name.take() { + taken.push(Matcher::new(MatchOp::Equal, METRIC_NAME_LABEL, &name)); + } + + let (name_matchers, rest) = std::mem::take(&mut selector.matchers.matchers) + .into_iter() + .partition(|matcher| matcher.name == METRIC_NAME_LABEL); + selector.matchers.matchers = rest; + taken.extend(name_matchers); + + taken +} + +/// Whether a metric name satisfies every `__name__` matcher of one selector. +/// +/// Negated matchers are honoured here, unlike in [`retrieve_table_names`] where +/// they keep every table so the caller authorizes the full candidate set: the +/// names reaching this point are already narrowed by the data, so filtering them +/// can only remove names, never widen what a caller gets to see. +fn metric_name_matches(table_name: &str, matchers: &[Matcher]) -> bool { + matchers.iter().all(|matcher| match &matcher.op { + MatchOp::Equal => table_name == matcher.value, + MatchOp::NotEqual => table_name != matcher.value, + MatchOp::Re(re) => re.is_match(table_name), + MatchOp::NotRe(re) => !re.is_match(table_name), + }) +} + +/// Whether a matcher constrains an ordinary label. The others name the metric, +/// the database or the field, none of which is a column to scan. +fn is_ordinary_label_matcher(matcher: &Matcher) -> bool { + matcher.name != METRIC_NAME_LABEL + && matcher.name != FIELD_NAME_LABEL + && !is_database_selection_label(&matcher.name) +} + +/// Splits `match[]` selectors by whether they constrain an ordinary label. The +/// first group is answerable from table metadata; the second needs the data read +/// and is returned as parsed selectors. +/// +/// `or` matchers stay in the metadata group, which ignores them, rather than +/// being silently dropped from a data scan that cannot express them. +fn split_selectors_by_label_use(matches: &[String]) -> Result<(Vec, Vec)> { + let mut metadata_only = Vec::new(); + let mut with_labels = Vec::new(); + + for selector in matches { + let expr = promql_parser::parser::parse(selector) + .map_err(|reason| InvalidQuerySnafu { reason }.build())?; + let PromqlExpr::VectorSelector(vector_selector) = expr else { + return InvalidQuerySnafu { + reason: "expected vector selector".to_string(), + } + .fail(); + }; + + let constrains_labels = vector_selector.matchers.or_matchers.is_empty() + && vector_selector + .matchers + .matchers + .iter() + .any(is_ordinary_label_matcher); + if constrains_labels { + with_labels.push(vector_selector); + } else { + metadata_only.push(selector.clone()); + } + } + + Ok((metadata_only, with_labels)) +} + +/// Resolves selectors constraining ordinary labels into metric names: the data +/// answers the label matchers, then each selector's `__name__` matchers narrow +/// the names it found. +async fn retrieve_table_names_by_labels( + handler: &PrometheusHandlerRef, + selectors: Vec, + start: Option<&str>, + end: Option<&str>, + query_ctx: &QueryContextRef, +) -> Result> { + let start_arg = start.map(str::to_string).unwrap_or_else(yesterday_rfc3339); + let end_arg = end.map(str::to_string).unwrap_or_else(current_time_rfc3339); + let start = QueryLanguageParser::parse_promql_timestamp(&start_arg).with_context(|_| { + ParseTimestampSnafu { + timestamp: start_arg.clone(), + } + })?; + let end = QueryLanguageParser::parse_promql_timestamp(&end_arg).with_context(|_| { + ParseTimestampSnafu { + timestamp: end_arg.clone(), + } + })?; + + let schema = query_ctx.current_schema(); + let mut table_names = Vec::new(); + for mut selector in selectors { + let name_matchers = take_metric_name_matchers(&mut selector); + // The database and field matchers name no column, and the metadata path + // ignores them too. + let label_matchers = selector + .matchers + .matchers + .into_iter() + .filter(is_ordinary_label_matcher) + .collect(); + let matched = handler + .query_metric_names_by_labels(label_matchers, &schema, start, end, query_ctx) + .await?; + table_names.extend( + matched + .into_iter() + .filter(|name| metric_name_matches(name, &name_matchers)), + ); + } + + Ok(table_names) +} + async fn retrieve_table_names( query_ctx: &QueryContext, catalog_manager: CatalogManagerRef, @@ -2407,6 +2626,10 @@ mod tests { deny_operation: bool, denied_table: Option<&'static str>, metric_names: Vec, + /// Names the label-matcher path resolves, kept apart from `metric_names` + /// so a test can tell which path answered. + label_metric_names: Vec, + label_lookups: Mutex>>, queries: Mutex>, ordered_outputs: Mutex>, } @@ -2482,6 +2705,18 @@ mod tests { Ok(self.metric_names.clone()) } + async fn query_metric_names_by_labels( + &self, + matchers: Vec, + _: &str, + _: std::time::SystemTime, + _: std::time::SystemTime, + _: &QueryContextRef, + ) -> Result> { + self.label_lookups.lock().unwrap().push(matchers); + Ok(self.label_metric_names.clone()) + } + async fn query_label_values( &self, _: String, @@ -2547,6 +2782,8 @@ mod tests { deny_operation: false, denied_table: None, metric_names: Vec::new(), + label_metric_names: Vec::new(), + label_lookups: Mutex::new(Vec::new()), queries: Mutex::new(Vec::new()), ordered_outputs: Mutex::new(Vec::new()), }); @@ -2600,6 +2837,8 @@ mod tests { deny_operation: false, denied_table: None, metric_names: Vec::new(), + label_metric_names: Vec::new(), + label_lookups: Mutex::new(Vec::new()), queries: Mutex::new(Vec::new()), ordered_outputs: Mutex::new(Vec::new()), }); @@ -2684,6 +2923,8 @@ mod tests { deny_operation: false, denied_table: Some("denied"), metric_names: Vec::new(), + label_metric_names: Vec::new(), + label_lookups: Mutex::new(Vec::new()), queries: Mutex::new(Vec::new()), ordered_outputs: Mutex::new(Vec::new()), })), @@ -2700,6 +2941,153 @@ mod tests { ); } + /// A handler over the logical metric tables `cpu_user` and `cpu_system`, + /// which is what the metadata path enumerates, with the label-matcher path + /// answering `label_metric_names`. + fn label_values_handler(label_metric_names: Vec<&str>) -> Arc { + let mut cpu_user = test_table_info( + 1024, + "cpu_user", + DEFAULT_SCHEMA_NAME, + DEFAULT_CATALOG_NAME, + Arc::new(Schema::new(vec![])), + ); + cpu_user.meta.options.extra_options.insert( + LOGICAL_TABLE_METADATA_KEY.to_string(), + "physical_metrics".to_string(), + ); + let manager = MemoryCatalogManager::new_with_table(EmptyTable::from_table_info(&cpu_user)); + let mut cpu_system = cpu_user.clone(); + cpu_system.ident.table_id = 1025; + cpu_system.name = "cpu_system".to_string(); + manager + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: cpu_system.name.clone(), + table_id: cpu_system.table_id(), + table: EmptyTable::from_table_info(&cpu_system), + }) + .unwrap(); + + Arc::new(TestPrometheusHandler { + catalog_manager: manager, + deny_operation: false, + denied_table: None, + metric_names: Vec::new(), + label_metric_names: label_metric_names.into_iter().map(String::from).collect(), + label_lookups: Mutex::new(Vec::new()), + queries: Mutex::new(Vec::new()), + ordered_outputs: Mutex::new(Vec::new()), + }) + } + + async fn query_metric_name_values( + handler: Arc, + matches: Vec<&str>, + ) -> Vec { + let state: PrometheusHandlerRef = handler; + let response = label_values_query( + State(state), + Path(METRIC_NAME_LABEL.to_string()), + Extension(QueryContext::with( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + )), + Query(LabelValueQuery { + matches: Matches(matches.into_iter().map(String::from).collect()), + ..Default::default() + }), + ) + .await; + + assert!( + response.status_code.is_none(), + "status={:?}, error={:?}", + response.status_code, + response.error + ); + match response.data { + PrometheusResponse::LabelValues(values) => values, + other => panic!("expected label values, got {other:?}"), + } + } + + #[tokio::test] + async fn label_matchers_resolve_metric_names_from_data() { + let handler = label_values_handler(vec!["cpu_user"]); + let values = query_metric_name_values(handler.clone(), vec![r#"{pod="abc"}"#]).await; + + // The metadata path would have enumerated both metrics. + assert_eq!(vec!["cpu_user".to_string()], values); + + let lookups = handler.label_lookups.lock().unwrap(); + assert_eq!(1, lookups.len()); + assert_eq!( + vec!["pod".to_string()], + lookups[0] + .iter() + .map(|matcher| matcher.name.clone()) + .collect::>() + ); + } + + #[tokio::test] + async fn metric_name_matchers_narrow_data_resolved_names() { + let handler = label_values_handler(vec!["cpu_user", "cpu_system"]); + let values = + query_metric_name_values(handler, vec![r#"{__name__=~"cpu_u.*", pod="abc"}"#]).await; + + assert_eq!(vec!["cpu_user".to_string()], values); + } + + #[tokio::test] + async fn special_matchers_are_stripped_before_the_data_lookup() { + let handler = label_values_handler(vec!["cpu_user"]); + let values = query_metric_name_values( + handler.clone(), + vec![r#"{pod="abc", __field__="value", __database__="public"}"#], + ) + .await; + + assert_eq!(vec!["cpu_user".to_string()], values); + let lookups = handler.label_lookups.lock().unwrap(); + assert_eq!( + vec!["pod".to_string()], + lookups[0] + .iter() + .map(|matcher| matcher.name.clone()) + .collect::>() + ); + } + + #[tokio::test] + async fn database_and_field_matchers_stay_on_the_metadata_path() { + let handler = label_values_handler(vec!["never_returned"]); + let values = query_metric_name_values( + handler.clone(), + vec![r#"{__name__=~"cpu_.*", __field__="value", __database__="other"}"#], + ) + .await; + + assert_eq!( + vec!["cpu_system".to_string(), "cpu_user".to_string()], + values + ); + assert!(handler.label_lookups.lock().unwrap().is_empty()); + } + + #[tokio::test] + async fn empty_match_still_enumerates_every_metric() { + let handler = label_values_handler(Vec::new()); + let values = query_metric_name_values(handler, Vec::new()).await; + + assert_eq!( + vec!["cpu_system".to_string(), "cpu_user".to_string()], + values + ); + } + #[tokio::test] async fn test_series_query_expands_metric_name_regex() { let cpu_user = test_table_info( @@ -2728,6 +3116,8 @@ mod tests { deny_operation: false, denied_table: None, metric_names: vec!["cpu_user".to_string(), "cpu_system".to_string()], + label_metric_names: Vec::new(), + label_lookups: Mutex::new(Vec::new()), queries: Mutex::new(Vec::new()), ordered_outputs: Mutex::new(Vec::new()), }); @@ -2889,6 +3279,134 @@ mod tests { expected_type: ValueType::Vector, should_error: false, }, + // Arithmetic operators always drop the metric name. + TestCase { + name: "arithmetic between metrics", + promql: "a + b", + expected_metric: None, + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "arithmetic with scalar", + promql: "a * 2", + expected_metric: None, + expected_type: ValueType::Vector, + should_error: false, + }, + // Comparisons keep the metric name from the side whose labels the + // result carries, i.e. the "many" side of the match: normally the + // left operand, but `group_right` takes it from the right operand and + // a scalar left operand leaves the vector on the right. The name is + // dropped when the comparison returns a bool or is a one-to-one + // `on(...)` match that does not keep `__name__` or an `ignoring(...)` + // match that excludes it. A `group_left`/`group_right` modifier does + // not reduce the labels, so the name is kept. + TestCase { + name: "bool comparison between metrics", + promql: "a > bool b", + expected_metric: None, + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "bool comparison with scalar", + promql: "a > bool 0.5", + expected_metric: None, + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with scalar", + promql: "a > 0.5", + expected_metric: Some("a"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with scalar on the left", + promql: "0.5 < a", + expected_metric: Some("a"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "equality comparison with scalar", + promql: "up == 0", + expected_metric: Some("up"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison between metrics", + promql: "a > b", + expected_metric: Some("a"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with ignoring", + promql: "a > ignoring(x) b", + expected_metric: Some("a"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with ignoring __name__", + promql: "a > ignoring(__name__) b", + expected_metric: None, + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with ignoring __name__ and group_left", + promql: "a > ignoring(__name__) group_left b", + expected_metric: Some("a"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with ignoring __name__ and group_right", + promql: "a > ignoring(__name__) group_right b", + expected_metric: Some("b"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "unless with ignoring __name__", + promql: "a unless ignoring(__name__) b", + expected_metric: Some("a"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with on", + promql: "a > on(x) b", + expected_metric: None, + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with on keeping __name__", + promql: "a > on(__name__) b", + expected_metric: Some("a"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with on and group_left", + promql: "a > on(x) group_left b", + expected_metric: Some("a"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with on and group_right", + promql: "a > on(x) group_right b", + expected_metric: Some("b"), + expected_type: ValueType::Vector, + should_error: false, + }, // Subquery cases TestCase { name: "basic subquery", @@ -3635,6 +4153,8 @@ mod tests { deny_operation: true, denied_table: None, metric_names: Vec::new(), + label_metric_names: Vec::new(), + label_lookups: Mutex::new(Vec::new()), queries: Mutex::new(Vec::new()), ordered_outputs: Mutex::new(Vec::new()), })), @@ -3649,6 +4169,8 @@ mod tests { deny_operation: false, denied_table: Some("denied"), metric_names: Vec::new(), + label_metric_names: Vec::new(), + label_lookups: Mutex::new(Vec::new()), queries: Mutex::new(Vec::new()), ordered_outputs: Mutex::new(Vec::new()), }); diff --git a/src/servers/src/otlp/coerce.rs b/src/servers/src/otlp/coerce.rs index 2e7dd36ed7a..2f8d25247cf 100644 --- a/src/servers/src/otlp/coerce.rs +++ b/src/servers/src/otlp/coerce.rs @@ -30,6 +30,9 @@ pub enum TraceCoerceError { // - String to Int64 // - String to Float64 // - String to Boolean +// The empty string and "-" coerce to null instead of failing: instrumentations +// fill numeric attributes from absent headers this way (PSR-7 `getHeaderLine` +// returns ""), and rejecting the value would reject the whole span. // // Lossless signed-to-unsigned integer casts. These let the built-in data // models move from unsigned to signed integers while existing unsigned tables @@ -86,6 +89,13 @@ pub fn coerce_value_data( return Ok(None); }; + if let ValueData::StringValue(s) = v + && matches!(s.as_str(), "" | "-") + && is_supported_trace_coercion(request_type, target) + { + return Ok(None); + } + let Some(value) = coerce_non_null_value(target, request_type, v) else { return Err(TraceCoerceError::Unsupported); }; @@ -273,6 +283,24 @@ mod tests { assert_eq!(result, Err(TraceCoerceError::Unsupported)); } + #[test] + fn test_coerce_absent_value_strings_to_null() { + for s in ["", "-"] { + for target in [ + ColumnDataType::Int64, + ColumnDataType::Float64, + ColumnDataType::Boolean, + ] { + let result = coerce_value_data( + &Some(ValueData::StringValue(s.to_string())), + target, + ColumnDataType::String, + ); + assert_eq!(result, Ok(None), "{s:?} to {target:?}"); + } + } + } + #[test] fn test_coerce_float64_to_int64_not_supported() { let result = coerce_value_data( diff --git a/src/servers/src/prom_store.rs b/src/servers/src/prom_store.rs index 980504515a8..9182f64ad00 100644 --- a/src/servers/src/prom_store.rs +++ b/src/servers/src/prom_store.rs @@ -362,22 +362,37 @@ pub fn recordbatches_to_timeseries( value_column_name: &str, recordbatches: RecordBatches, ) -> Result> { - Ok(recordbatches - .take() - .into_iter() - .map(|x| recordbatch_to_timeseries(table_name, timestamp_column_name, value_column_name, x)) - .collect::>>()? - .into_iter() - .flatten() - .collect()) + let mut timeseries: Vec = Vec::new(); + let mut timeseries_by_hash: HashMap> = HashMap::new(); + + for recordbatch in recordbatches.take() { + append_recordbatch_to_timeseries( + table_name, + timestamp_column_name, + value_column_name, + recordbatch, + &mut timeseries, + &mut timeseries_by_hash, + )?; + } + + for ts in &mut timeseries { + ts.samples.sort_unstable_by_key(|s| s.timestamp); + } + + timeseries + .sort_unstable_by(|left, right| compare_timeseries_labels(&left.labels, &right.labels)); + Ok(timeseries) } -fn recordbatch_to_timeseries( +fn append_recordbatch_to_timeseries( table: &str, timestamp_column_name: &str, value_column_name: &str, recordbatch: RecordBatch, -) -> Result> { + timeseries: &mut Vec, + timeseries_by_hash: &mut HashMap>, +) -> Result<()> { let ts_column = recordbatch .column_by_name(timestamp_column_name) .with_context(|| error::InvalidPromRemoteReadQueryResultSnafu { @@ -409,8 +424,6 @@ fn recordbatch_to_timeseries( })?; let columns = label_columns(&recordbatch, timestamp_column_name, value_column_name)?; - let mut timeseries: Vec = Vec::new(); - let mut timeseries_by_hash: HashMap> = HashMap::new(); let mut previous_timeseries: Option = None; for row in 0..recordbatch.num_rows() { @@ -447,9 +460,7 @@ fn recordbatch_to_timeseries( timeseries[timeseries_index].samples.push(sample); } - timeseries - .sort_unstable_by(|left, right| compare_timeseries_labels(&left.labels, &right.labels)); - Ok(timeseries) + Ok(()) } pub fn to_grpc_row_insert_requests(request: &WriteRequest) -> Result<(RowInsertRequests, usize)> { @@ -1122,6 +1133,121 @@ mod tests { ); } + #[test] + fn test_recordbatches_to_timeseries_merges_across_batches() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + greptime_timestamp(), + ConcreteDataType::timestamp_millisecond_datatype(), + true, + ), + ColumnSchema::new(greptime_value(), ConcreteDataType::float64_datatype(), true), + ColumnSchema::new("instance", ConcreteDataType::string_datatype(), true), + ])); + + let recordbatches = RecordBatches::try_new( + schema.clone(), + vec![ + RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_vec(vec![3000, 1500])) as _, + Arc::new(Float64Vector::from_vec(vec![30.0, 15.0])) as _, + Arc::new(StringVector::from(vec!["host1", "host2"])) as _, + ], + ) + .unwrap(), + RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_vec(vec![1000])) as _, + Arc::new(Float64Vector::from_vec(vec![10.0])) as _, + Arc::new(StringVector::from(vec!["host1"])) as _, + ], + ) + .unwrap(), + RecordBatch::new( + schema, + vec![ + Arc::new(TimestampMillisecondVector::from_vec(vec![2000, 2500])) as _, + Arc::new(Float64Vector::from_vec(vec![20.0, 25.0])) as _, + Arc::new(StringVector::from(vec!["host1", "host2"])) as _, + ], + ) + .unwrap(), + ], + ) + .unwrap(); + + let timeseries = recordbatches_to_timeseries( + "cpu_usage", + greptime_timestamp(), + greptime_value(), + recordbatches, + ) + .unwrap(); + + assert_eq!(2, timeseries.len()); + + assert_eq!( + vec![ + Label { + name: METRIC_NAME_LABEL.to_string(), + value: "cpu_usage".to_string(), + }, + Label { + name: "instance".to_string(), + value: "host1".to_string(), + }, + ], + timeseries[0].labels + ); + assert_eq!( + vec![ + Sample { + value: 10.0, + timestamp: 1000, + }, + Sample { + value: 20.0, + timestamp: 2000, + }, + Sample { + value: 30.0, + timestamp: 3000, + }, + ], + timeseries[0].samples + ); + + assert_eq!( + vec![ + Label { + name: METRIC_NAME_LABEL.to_string(), + value: "cpu_usage".to_string(), + }, + Label { + name: "instance".to_string(), + value: "host2".to_string(), + }, + ], + timeseries[1].labels + ); + assert_eq!( + vec![ + Sample { + value: 15.0, + timestamp: 1500, + }, + Sample { + value: 25.0, + timestamp: 2500, + }, + ], + timeseries[1].samples + ); + } + #[test] fn test_recordbatches_to_timeseries_borrows_and_groups_dictionary_labels() { let arrow_schema = Arc::new(ArrowSchema::new(vec![ @@ -1212,7 +1338,7 @@ mod tests { } #[test] - fn test_recordbatch_to_timeseries_groups_non_contiguous_series() { + fn test_recordbatches_to_timeseries_groups_non_contiguous_series() { let schema = Arc::new(Schema::new(vec![ ColumnSchema::new( greptime_timestamp(), @@ -1223,7 +1349,7 @@ mod tests { ColumnSchema::new("instance", ConcreteDataType::string_datatype(), true), ])); let recordbatch = RecordBatch::new( - schema, + schema.clone(), vec![ Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as _, Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as _, @@ -1232,11 +1358,12 @@ mod tests { ) .unwrap(); - let timeseries = recordbatch_to_timeseries( + let recordbatches = RecordBatches::try_new(schema, vec![recordbatch]).unwrap(); + let timeseries = recordbatches_to_timeseries( "metric1", greptime_timestamp(), greptime_value(), - recordbatch, + recordbatches, ) .unwrap(); @@ -1259,7 +1386,7 @@ mod tests { } #[test] - fn test_recordbatch_to_timeseries_arrow_label_types_and_nulls() { + fn test_recordbatches_to_timeseries_arrow_label_types_and_nulls() { let schema = Arc::new(Schema::new(vec![ ColumnSchema::new( greptime_timestamp(), @@ -1272,7 +1399,7 @@ mod tests { ColumnSchema::new("shard", ConcreteDataType::int32_datatype(), true), ])); let recordbatch = RecordBatch::new( - schema, + schema.clone(), vec![ Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as _, Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as _, @@ -1289,11 +1416,12 @@ mod tests { ) .unwrap(); - let timeseries = recordbatch_to_timeseries( + let recordbatches = RecordBatches::try_new(schema, vec![recordbatch]).unwrap(); + let timeseries = recordbatches_to_timeseries( "metric1", greptime_timestamp(), greptime_value(), - recordbatch, + recordbatches, ) .unwrap(); diff --git a/src/servers/src/prometheus_handler.rs b/src/servers/src/prometheus_handler.rs index 4f4b7c5d19c..1e211b501ab 100644 --- a/src/servers/src/prometheus_handler.rs +++ b/src/servers/src/prometheus_handler.rs @@ -184,6 +184,21 @@ pub trait PrometheusHandler { ctx: &QueryContextRef, ) -> Result>; + /// Query metric table names that carry data matching `matchers` in the time + /// range. `matchers` must hold only ordinary label matchers: `__name__` + /// names a table and the database and field matchers name no column, so the + /// caller resolves all three before calling. + /// + /// Only metric engine tables are covered. + async fn query_metric_names_by_labels( + &self, + matchers: Vec, + schema: &str, + start: SystemTime, + end: SystemTime, + ctx: &QueryContextRef, + ) -> Result>; + async fn query_label_values( &self, metric: String, diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index db1df3b339e..6536269c748 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -704,6 +704,125 @@ LEFT JOIN (SELECT k, count(*) n FROM join_repro WHERE kind = 'child' AND k < '8' } } +/// Regression test for `MatchesConstantTermOptimizer` dropping `FilterExec`'s `fetch`. +/// +/// On a single-partition plan the `LIMIT` is enforced by the `fetch` that the physical +/// `LimitPushdown` rule leaves on the `FilterExec`, so a rule rebuilding that `FilterExec` +/// has to carry the `fetch` over. While it does not, the limit is silently lifted for +/// `matches_term` filters and +/// `SELECT count(*) FROM (SELECT ts, msg FROM t WHERE matches_term(msg, 'foo') LIMIT 1) s` +/// returns the number of matching rows (2) instead of 1. +/// +/// `query_parallelism=4` is expected to stay correct, because the limit is then materialized +/// above the scan instead of being delegated to the `FilterExec`; the test pins down that +/// contrast. The metric-engine logical table covers the same filter on a shared physical +/// region. +/// +/// Both matching rows share `ts` and `msg`, so the limited output does not depend on the +/// order in which the scan returns them. +#[apply(both_instances_cases)] +async fn test_matches_term_subquery_limit_count(instance: Arc) { + let frontend = instance.frontend(); + + execute_sql( + &frontend, + r#" +CREATE TABLE f07_mito_msgs ( + ts TIMESTAMP TIME INDEX, + host STRING PRIMARY KEY, + msg STRING, + val DOUBLE +) +ENGINE = mito +"#, + ) + .await; + execute_sql( + &frontend, + r#"INSERT INTO f07_mito_msgs (ts, host, msg, val) VALUES + (0, 'host1', 'foo bar', 1.0), + (0, 'host2', 'foo bar', 2.0), + (0, 'host3', 'other baz', 3.0)"#, + ) + .await; + + execute_sql( + &frontend, + r#" +CREATE TABLE f07_metric_phy ( + ts TIMESTAMP TIME INDEX, + val DOUBLE, + host STRING, + msg STRING, + PRIMARY KEY (host, msg) +) +ENGINE = metric +WITH ('physical_metric_table' = '') +"#, + ) + .await; + execute_sql( + &frontend, + r#" +CREATE TABLE f07_metric_msgs ( + ts TIMESTAMP TIME INDEX, + val DOUBLE, + host STRING, + msg STRING, + PRIMARY KEY (host, msg) +) +ENGINE = metric +WITH ('on_physical_table' = 'f07_metric_phy') +"#, + ) + .await; + execute_sql( + &frontend, + r#"INSERT INTO f07_metric_msgs (ts, host, msg, val) VALUES + (0, 'host1', 'foo bar', 1.0), + (0, 'host2', 'foo bar', 2.0), + (0, 'host3', 'other baz', 3.0)"#, + ) + .await; + + for table in ["f07_mito_msgs", "f07_metric_msgs"] { + let limited = format!("SELECT ts, msg FROM {table} WHERE matches_term(msg, 'foo') LIMIT 1"); + let count = format!("SELECT count(*) FROM ({limited}) s"); + + let limit_rows = execute_sql_with_query_parallelism(&frontend, &limited, 1) + .await + .data + .pretty_print() + .await; + assert_eq!( + limit_rows, + r#"+---------------------+---------+ +| ts | msg | ++---------------------+---------+ +| 1970-01-01T00:00:00 | foo bar | ++---------------------+---------+"#, + "query_parallelism=1: `{limited}` must return exactly the single limited row" + ); + + for parallelism in [1, 4] { + let count_rows = execute_sql_with_query_parallelism(&frontend, &count, parallelism) + .await + .data + .pretty_print() + .await; + assert_eq!( + count_rows, + r#"+----------+ +| count(*) | ++----------+ +| 1 | ++----------+"#, + "query_parallelism={parallelism}: the LIMIT 1 of the subquery must cap `{count}` to one row" + ); + } + } +} + #[apply(both_instances_cases)] async fn test_partitioned_merge_scan_semi_join_repartitions(instance: Arc) { let frontend = instance.frontend(); diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index c8f18be5d66..fb4b8ddb468 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -1516,6 +1516,117 @@ pub async fn test_prom_http_api(store_type: StorageType) { .unwrap() ); + // query `__name__` by a matcher on an ordinary label: the metric engine + // physical tables are scanned, so only metrics carrying the label value are + // returned. `demo_metrics` shares `phy` with `demo` but has no `host` value. + let res = client + .get("/v1/prometheus/api/v1/label/__name__/values?match[]={host=\"host1\"}&start=0&end=600") + .send() + .await; + let status = res.status(); + let text = res.text().await; + assert_eq!(status, StatusCode::OK, "{text}"); + let prom_resp = serde_json::from_str::(&text).unwrap(); + assert_eq!(prom_resp.status, "success"); + assert!(prom_resp.error.is_none()); + assert_eq!( + prom_resp.data, + serde_json::from_value::(json!(["demo", "multi_labels"])).unwrap() + ); + + // `__name__` matchers narrow the names the data resolved: `multi_labels` + // also carries `idc="idc1"` but its name does not match. + let res = client + .get("/v1/prometheus/api/v1/label/__name__/values?match[]={__name__=~\"demo.*\", idc=\"idc1\"}&start=0&end=600") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let prom_resp = res.json::().await; + assert_eq!(prom_resp.status, "success"); + assert!(prom_resp.error.is_none()); + assert_eq!( + prom_resp.data, + serde_json::from_value::(json!([ + "demo_metrics", + "demo_metrics_with_nanos", + ])) + .unwrap() + ); + + // The time range selects the series: `demo` carries `host="host2"` only at + // t=600, so narrowing the range drops it while `multi_labels` at t=0 stays. + let res = client + .get("/v1/prometheus/api/v1/label/__name__/values?match[]={host=\"host2\"}&start=0&end=600") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let prom_resp = res.json::().await; + assert_eq!( + prom_resp.data, + serde_json::from_value::(json!(["demo", "multi_labels"])).unwrap() + ); + let res = client + .get("/v1/prometheus/api/v1/label/__name__/values?match[]={host=\"host2\"}&start=0&end=300") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let prom_resp = res.json::().await; + assert_eq!( + prom_resp.data, + serde_json::from_value::(json!(["multi_labels"])).unwrap() + ); + + // Logical metrics sharing a physical table have NULL in the label columns + // they don't use. Prometheus reads a label a series doesn't carry as the + // empty string, so `demo_metrics` and `demo_metrics_with_nanos` — neither of + // which has a `host` label — match both of these. + // + // `.%2B` is `.+`; a bare `+` decodes to a space in a query string. Grafana + // encodes it the same way. + let res = client + .get("/v1/prometheus/api/v1/label/__name__/values?match[]={__name__=~\".%2B\", host=\"\"}&start=0&end=600") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let prom_resp = res.json::().await; + assert_eq!( + prom_resp.data, + serde_json::from_value::(json!([ + "demo_metrics", + "demo_metrics_with_nanos", + ])) + .unwrap() + ); + + let res = client + .get("/v1/prometheus/api/v1/label/__name__/values?match[]={__name__=~\".%2B\", host!=\"host1\"}&start=0&end=600") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let prom_resp = res.json::().await; + assert_eq!( + prom_resp.data, + serde_json::from_value::(json!([ + "demo", + "demo_metrics", + "demo_metrics_with_nanos", + "multi_labels", + ])) + .unwrap() + ); + + // A pre-epoch RFC3339 bound is a valid range, not a panic. + let res = client + .get("/v1/prometheus/api/v1/label/__name__/values?match[]={host=\"host1\"}&start=1969-12-31T23:59:59Z&end=600") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let prom_resp = res.json::().await; + assert_eq!( + prom_resp.data, + serde_json::from_value::(json!(["demo", "multi_labels"])).unwrap() + ); + // buildinfo let res = client .get("/v1/prometheus/api/v1/status/buildinfo") diff --git a/tests/cases/standalone/common/aggregate/array_agg.result b/tests/cases/standalone/common/aggregate/array_agg.result index 5a73c289899..e1ea29beac5 100644 --- a/tests/cases/standalone/common/aggregate/array_agg.result +++ b/tests/cases/standalone/common/aggregate/array_agg.result @@ -144,3 +144,94 @@ DROP TABLE doubles; Affected Rows: 0 +-- On partitioned tables the aggregate is split into partial state and merge +CREATE TABLE array_agg_partitioned ( + ts TIMESTAMP TIME INDEX, + k INT, + lat DOUBLE, + PRIMARY KEY(k) +) +PARTITION ON COLUMNS (k) (k < 10, k >= 10 AND k < 20, k >= 20); + +Affected Rows: 0 + +INSERT INTO array_agg_partitioned VALUES + (1000, 1, 1), + (2000, 11, 2), + (3000, 21, 3), + (4000, 2, 4), + (5000, 12, 5), + (6000, 22, 6), + (7000, 3, 7); + +Affected Rows: 7 + +SELECT array_agg(lat ORDER BY ts) FROM array_agg_partitioned; + ++-----------------------------------------------------------------------------------------+ +| array_agg(array_agg_partitioned.lat) ORDER BY [array_agg_partitioned.ts ASC NULLS LAST] | ++-----------------------------------------------------------------------------------------+ +| [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0] | ++-----------------------------------------------------------------------------------------+ + +SELECT lat > 3 AS g, array_agg(lat ORDER BY ts DESC) FROM array_agg_partitioned GROUP BY g ORDER BY g; + ++-------+-------------------------------------------------------------------------------------------+ +| g | array_agg(array_agg_partitioned.lat) ORDER BY [array_agg_partitioned.ts DESC NULLS FIRST] | ++-------+-------------------------------------------------------------------------------------------+ +| false | [3.0, 2.0, 1.0] | +| true | [7.0, 6.0, 5.0, 4.0] | ++-------+-------------------------------------------------------------------------------------------+ + +SELECT array_agg(k ORDER BY k % 10, ts DESC) FROM array_agg_partitioned; + ++---------------------------------------------------------------------------------------------------------------------------------------------+ +| array_agg(array_agg_partitioned.k) ORDER BY [array_agg_partitioned.k % Int64(10) ASC NULLS LAST, array_agg_partitioned.ts DESC NULLS FIRST] | ++---------------------------------------------------------------------------------------------------------------------------------------------+ +| [21, 11, 1, 22, 12, 2, 3] | ++---------------------------------------------------------------------------------------------------------------------------------------------+ + +-- nth_value needs sorted input, so it isn't split into partial state and merge +SELECT nth_value(lat, 2 ORDER BY ts), nth_value(lat, 3 ORDER BY ts DESC) FROM array_agg_partitioned; + ++--------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------+ +| nth_value(array_agg_partitioned.lat,Int64(2)) ORDER BY [array_agg_partitioned.ts ASC NULLS LAST] | nth_value(array_agg_partitioned.lat,Int64(3)) ORDER BY [array_agg_partitioned.ts DESC NULLS FIRST] | ++--------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------+ +| 2.0 | 5.0 | ++--------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------+ + +-- WITHIN GROUP aggregates don't need sorted input and are still split +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +EXPLAIN SELECT sum(lat), approx_percentile_cont(0.75) WITHIN GROUP (ORDER BY lat DESC) FROM array_agg_partitioned; + ++-+-+ +| plan_type_| plan_| ++-+-+ +| logical_plan_| Aggregate: groupBy=[[]], aggr=[[__sum_merge(__sum_state(array_agg_partitioned.lat)) AS sum(array_agg_partitioned.lat), __approx_percentile_cont_merge(__approx_percentile_cont_state(array_agg_partitioned.lat,Float64(0.75)) ORDER BY [array_agg_partitioned.lat DESC NULLS FIRST]) AS approx_percentile_cont(Float64(0.75)) WITHIN GROUP [array_agg_partitioned.lat DESC NULLS FIRST]]] | +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| Aggregate: groupBy=[[]], aggr=[[__sum_state(array_agg_partitioned.lat), __approx_percentile_cont_state(array_agg_partitioned.lat, Float64(0.75)) ORDER BY [array_agg_partitioned.lat DESC NULLS FIRST]]]_| +|_|_TableScan: array_agg_partitioned_| +|_| ]]_| +| physical_plan | AggregateExec: mode=Final, gby=[], aggr=[sum(array_agg_partitioned.lat), approx_percentile_cont(Float64(0.75)) WITHIN GROUP [array_agg_partitioned.lat DESC NULLS FIRST]]_| +|_|_CoalescePartitionsExec_| +|_|_AggregateExec: mode=Partial, gby=[], aggr=[sum(array_agg_partitioned.lat), approx_percentile_cont(Float64(0.75)) WITHIN GROUP [array_agg_partitioned.lat DESC NULLS FIRST]]_| +|_|_RepartitionExec: partitioning=REDACTED +|_|_MergeScanExec: REDACTED +|_|_| ++-+-+ + +SELECT sum(lat), approx_percentile_cont(0.75) WITHIN GROUP (ORDER BY lat DESC) FROM array_agg_partitioned; + ++--------------------------------+-------------------------------------------------------------------------------------------------+ +| sum(array_agg_partitioned.lat) | approx_percentile_cont(Float64(0.75)) WITHIN GROUP [array_agg_partitioned.lat DESC NULLS FIRST] | ++--------------------------------+-------------------------------------------------------------------------------------------------+ +| 28.0 | 2.25 | ++--------------------------------+-------------------------------------------------------------------------------------------------+ + +DROP TABLE array_agg_partitioned; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/aggregate/array_agg.sql b/tests/cases/standalone/common/aggregate/array_agg.sql index dedabf1a18f..d0d7caa574f 100644 --- a/tests/cases/standalone/common/aggregate/array_agg.sql +++ b/tests/cases/standalone/common/aggregate/array_agg.sql @@ -54,3 +54,41 @@ DROP TABLE integers; DROP TABLE strings; DROP TABLE doubles; + +-- On partitioned tables the aggregate is split into partial state and merge +CREATE TABLE array_agg_partitioned ( + ts TIMESTAMP TIME INDEX, + k INT, + lat DOUBLE, + PRIMARY KEY(k) +) +PARTITION ON COLUMNS (k) (k < 10, k >= 10 AND k < 20, k >= 20); + +INSERT INTO array_agg_partitioned VALUES + (1000, 1, 1), + (2000, 11, 2), + (3000, 21, 3), + (4000, 2, 4), + (5000, 12, 5), + (6000, 22, 6), + (7000, 3, 7); + +SELECT array_agg(lat ORDER BY ts) FROM array_agg_partitioned; + +SELECT lat > 3 AS g, array_agg(lat ORDER BY ts DESC) FROM array_agg_partitioned GROUP BY g ORDER BY g; + +SELECT array_agg(k ORDER BY k % 10, ts DESC) FROM array_agg_partitioned; + +-- nth_value needs sorted input, so it isn't split into partial state and merge +SELECT nth_value(lat, 2 ORDER BY ts), nth_value(lat, 3 ORDER BY ts DESC) FROM array_agg_partitioned; + +-- WITHIN GROUP aggregates don't need sorted input and are still split +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +EXPLAIN SELECT sum(lat), approx_percentile_cont(0.75) WITHIN GROUP (ORDER BY lat DESC) FROM array_agg_partitioned; + +SELECT sum(lat), approx_percentile_cont(0.75) WITHIN GROUP (ORDER BY lat DESC) FROM array_agg_partitioned; + +DROP TABLE array_agg_partitioned; diff --git a/tests/cases/standalone/common/function/geo.result b/tests/cases/standalone/common/function/geo.result index 65ba42e6ced..f3ae45550be 100644 --- a/tests/cases/standalone/common/function/geo.result +++ b/tests/cases/standalone/common/function/geo.result @@ -437,3 +437,47 @@ FROM | true | false | true | false | false | true | +--------------------------+--------------------------+------------------------+------------------------+----------------------------------+----------------------------------+ +CREATE TABLE geo_path_partitioned ( + ts TIMESTAMP TIME INDEX, + k INT, + lat DOUBLE, + lon DOUBLE, + PRIMARY KEY(k) +) +PARTITION ON COLUMNS (k) (k < 10, k >= 10 AND k < 20, k >= 20); + +Affected Rows: 0 + +INSERT INTO geo_path_partitioned VALUES + (1000, 1, 1, 11), + (2000, 11, 2, 12), + (3000, 21, 3, 13), + (4000, 2, 4, 14), + (5000, 12, 5, 15), + (6000, 22, 6, 16), + (7000, 3, 7, 17); + +Affected Rows: 7 + +SELECT lat > 3 AS g, geo_path(lat, lon, ts) FROM geo_path_partitioned GROUP BY g ORDER BY g; + ++-------+-------------------------------------------------------------------------------------+ +| g | geo_path(geo_path_partitioned.lat,geo_path_partitioned.lon,geo_path_partitioned.ts) | ++-------+-------------------------------------------------------------------------------------+ +| false | {lat: [1.0, 2.0, 3.0], lng: [11.0, 12.0, 13.0]} | +| true | {lat: [4.0, 5.0, 6.0, 7.0], lng: [14.0, 15.0, 16.0, 17.0]} | ++-------+-------------------------------------------------------------------------------------+ + +SELECT lat > 3 AS g, json_encode_path(lat, lon, ts) FROM geo_path_partitioned GROUP BY g ORDER BY g; + ++-------+---------------------------------------------------------------------------------------------+ +| g | json_encode_path(geo_path_partitioned.lat,geo_path_partitioned.lon,geo_path_partitioned.ts) | ++-------+---------------------------------------------------------------------------------------------+ +| false | [[11.0,1.0],[12.0,2.0],[13.0,3.0]] | +| true | [[14.0,4.0],[15.0,5.0],[16.0,6.0],[17.0,7.0]] | ++-------+---------------------------------------------------------------------------------------------+ + +DROP TABLE geo_path_partitioned; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/function/geo.sql b/tests/cases/standalone/common/function/geo.sql index 43025cc196b..1de7d7c82aa 100644 --- a/tests/cases/standalone/common/function/geo.sql +++ b/tests/cases/standalone/common/function/geo.sql @@ -180,3 +180,27 @@ FROM 'POLYGON ((-121.491698 38.653343, -121.582353 38.556757, -121.469721 38.449287, -121.315883 38.541721, -121.491698 38.653343))' AS polygon2, 'POLYGON ((-122.089628 37.450332, -122.20535 37.378342, -122.093062 37.36088, -122.044301 37.372886, -122.089628 37.450332))' AS polygon3, ); + +CREATE TABLE geo_path_partitioned ( + ts TIMESTAMP TIME INDEX, + k INT, + lat DOUBLE, + lon DOUBLE, + PRIMARY KEY(k) +) +PARTITION ON COLUMNS (k) (k < 10, k >= 10 AND k < 20, k >= 20); + +INSERT INTO geo_path_partitioned VALUES + (1000, 1, 1, 11), + (2000, 11, 2, 12), + (3000, 21, 3, 13), + (4000, 2, 4, 14), + (5000, 12, 5, 15), + (6000, 22, 6, 16), + (7000, 3, 7, 17); + +SELECT lat > 3 AS g, geo_path(lat, lon, ts) FROM geo_path_partitioned GROUP BY g ORDER BY g; + +SELECT lat > 3 AS g, json_encode_path(lat, lon, ts) FROM geo_path_partitioned GROUP BY g ORDER BY g; + +DROP TABLE geo_path_partitioned; diff --git a/tests/cases/standalone/common/promql/computed_scalar_rhs.result b/tests/cases/standalone/common/promql/computed_scalar_rhs.result new file mode 100644 index 00000000000..96b925d1cf5 --- /dev/null +++ b/tests/cases/standalone/common/promql/computed_scalar_rhs.result @@ -0,0 +1,97 @@ +-- A vector combined with a computed scalar on the right keeps the vector's +-- labels, like a scalar on the left does. Expected results are from Prometheus 3. +create table computed_scalar_lhs ( + ts timestamp time index, + val double, + host string primary key +); + +Affected Rows: 0 + +create table computed_scalar_rhs ( + ts timestamp time index, + val double, + host string primary key +); + +Affected Rows: 0 + +insert into computed_scalar_lhs values + (0, 1.0, 'a'), + (0, 2.0, 'b'); + +Affected Rows: 2 + +insert into computed_scalar_rhs values + (0, 10.0, 'x'); + +Affected Rows: 1 + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') computed_scalar_lhs + time(); + ++------+---------------------+-------------------------------------------------+ +| host | ts | computed_scalar_lhs.val + .time / Float64(1000) | ++------+---------------------+-------------------------------------------------+ +| a | 1970-01-01T00:01:00 | 61.0 | +| b | 1970-01-01T00:01:00 | 62.0 | ++------+---------------------+-------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') time() + computed_scalar_lhs; + ++------+---------------------+-------------------------------------------------+ +| host | ts | .time / Float64(1000) + computed_scalar_lhs.val | ++------+---------------------+-------------------------------------------------+ +| a | 1970-01-01T00:01:00 | 61.0 | +| b | 1970-01-01T00:01:00 | 62.0 | ++------+---------------------+-------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') computed_scalar_lhs * scalar(computed_scalar_rhs); + ++------+---------------------+-----------------------------------------------------------+ +| host | ts | computed_scalar_lhs.val * computed_scalar_rhs.scalar(val) | ++------+---------------------+-----------------------------------------------------------+ +| a | 1970-01-01T00:01:00 | 10.0 | +| b | 1970-01-01T00:01:00 | 20.0 | ++------+---------------------+-----------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') computed_scalar_lhs < bool (time() * 2); + ++------+---------------------+--------------------------------------------------------------+ +| host | ts | computed_scalar_lhs.val < .time / Float64(1000) * Float64(2) | ++------+---------------------+--------------------------------------------------------------+ +| a | 1970-01-01T00:01:00 | 1.0 | +| b | 1970-01-01T00:01:00 | 1.0 | ++------+---------------------+--------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') (time() >= computed_scalar_lhs) ^ scalar(computed_scalar_rhs); + ++------+---------------------+---------------------------------------------+ +| host | ts | power(.val,computed_scalar_rhs.scalar(val)) | ++------+---------------------+---------------------------------------------+ +| a | 1970-01-01T00:01:00 | 1.0 | +| b | 1970-01-01T00:01:00 | 1024.0 | ++------+---------------------+---------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') scalar(computed_scalar_rhs) * (time() >= computed_scalar_lhs); + ++------+---------------------+----------------------------------------+ +| host | ts | computed_scalar_rhs.scalar(val) * .val | ++------+---------------------+----------------------------------------+ +| a | 1970-01-01T00:01:00 | 10.0 | +| b | 1970-01-01T00:01:00 | 20.0 | ++------+---------------------+----------------------------------------+ + +drop table computed_scalar_lhs; + +Affected Rows: 0 + +drop table computed_scalar_rhs; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/computed_scalar_rhs.sql b/tests/cases/standalone/common/promql/computed_scalar_rhs.sql new file mode 100644 index 00000000000..57cc431c1f9 --- /dev/null +++ b/tests/cases/standalone/common/promql/computed_scalar_rhs.sql @@ -0,0 +1,43 @@ +-- A vector combined with a computed scalar on the right keeps the vector's +-- labels, like a scalar on the left does. Expected results are from Prometheus 3. + +create table computed_scalar_lhs ( + ts timestamp time index, + val double, + host string primary key +); + +create table computed_scalar_rhs ( + ts timestamp time index, + val double, + host string primary key +); + +insert into computed_scalar_lhs values + (0, 1.0, 'a'), + (0, 2.0, 'b'); + +insert into computed_scalar_rhs values + (0, 10.0, 'x'); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') computed_scalar_lhs + time(); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') time() + computed_scalar_lhs; + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') computed_scalar_lhs * scalar(computed_scalar_rhs); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') computed_scalar_lhs < bool (time() * 2); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') (time() >= computed_scalar_lhs) ^ scalar(computed_scalar_rhs); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') scalar(computed_scalar_rhs) * (time() >= computed_scalar_lhs); + +drop table computed_scalar_lhs; + +drop table computed_scalar_rhs; diff --git a/tests/cases/standalone/common/promql/math.result b/tests/cases/standalone/common/promql/math.result index a6271c158cb..e1b884c50f1 100644 --- a/tests/cases/standalone/common/promql/math.result +++ b/tests/cases/standalone/common/promql/math.result @@ -55,15 +55,15 @@ TQL EVAL (0, 20, '5s') pi() + 1; -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 20, '5s') angles{unit="radians"} * pi(); -+---------------------+--------------------+ -| time | angles.val * .pi() | -+---------------------+--------------------+ -| 1970-01-01T00:00:00 | 0.0 | -| 1970-01-01T00:00:05 | 4.934813740258847 | -| 1970-01-01T00:00:10 | 9.869627480517694 | -| 1970-01-01T00:00:15 | 14.80444122077654 | -| 1970-01-01T00:00:20 | 19.73925496103539 | -+---------------------+--------------------+ ++---------+---------------------+--------------------+ +| unit | ts | angles.val * .pi() | ++---------+---------------------+--------------------+ +| radians | 1970-01-01T00:00:00 | 0.0 | +| radians | 1970-01-01T00:00:05 | 4.934813740258847 | +| radians | 1970-01-01T00:00:10 | 9.869627480517694 | +| radians | 1970-01-01T00:00:15 | 14.80444122077654 | +| radians | 1970-01-01T00:00:20 | 19.73925496103539 | ++---------+---------------------+--------------------+ -- Test rad() function (converts degrees to radians) -- SQLNESS SORT_RESULT 3 1 @@ -310,18 +310,18 @@ TQL EVAL (25, 35, '5s') sgn(angles - 42.5); -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 35, '5s') sgn(angles{unit="radians"} - pi()); -+---------------------+----------------------------+ -| time | signum(angles.val - .pi()) | -+---------------------+----------------------------+ -| 1970-01-01T00:00:00 | -1.0 | -| 1970-01-01T00:00:05 | -1.0 | -| 1970-01-01T00:00:10 | 1.0 | -| 1970-01-01T00:00:15 | 1.0 | -| 1970-01-01T00:00:20 | 1.0 | -| 1970-01-01T00:00:25 | 1.0 | -| 1970-01-01T00:00:30 | 1.0 | -| 1970-01-01T00:00:35 | 1.0 | -+---------------------+----------------------------+ ++---------------------+----------------------------+---------+ +| ts | signum(angles.val - .pi()) | unit | ++---------------------+----------------------------+---------+ +| 1970-01-01T00:00:00 | -1.0 | radians | +| 1970-01-01T00:00:05 | -1.0 | radians | +| 1970-01-01T00:00:10 | 1.0 | radians | +| 1970-01-01T00:00:15 | 1.0 | radians | +| 1970-01-01T00:00:20 | 1.0 | radians | +| 1970-01-01T00:00:25 | 1.0 | radians | +| 1970-01-01T00:00:30 | 1.0 | radians | +| 1970-01-01T00:00:35 | 1.0 | radians | ++---------------------+----------------------------+---------+ -- SQLNESS SORT_RESULT 3 1 TQL EVAL (25, 35, '5s') sgn(angles) * angles; diff --git a/tests/cases/standalone/common/promql/regex.result b/tests/cases/standalone/common/promql/regex.result index bccbd0f3402..8905694d1ad 100644 --- a/tests/cases/standalone/common/promql/regex.result +++ b/tests/cases/standalone/common/promql/regex.result @@ -144,8 +144,9 @@ TQL ANALYZE VERBOSE (0, 0, '1s') test{host!~".+"}; |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..0], lookback=[300000], interval=[1000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["host"] REDACTED +|_|_|_FilterExec: CASE WHEN host@1 IS NOT NULL THEN host@1 ELSE_END = REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "projection": ["ts", "host", "val"], "filters": ["host = Dictionary(UInt32, Utf8(\"\"))", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(0, None)"], "REDACTED +|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["ts", "host", "val"], "filters": ["CASE WHEN host IS NOT NULL THEN host ELSE Dictionary(UInt32, Utf8(\"\")) END = Dictionary(UInt32, Utf8(\"\"))", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(0, None)"], "REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ diff --git a/tests/cases/standalone/common/promql/simple_histogram.result b/tests/cases/standalone/common/promql/simple_histogram.result index e02e138f80e..e53d3977867 100644 --- a/tests/cases/standalone/common/promql/simple_histogram.result +++ b/tests/cases/standalone/common/promql/simple_histogram.result @@ -428,3 +428,92 @@ drop table histogram5_bucket; Affected Rows: 0 +-- Issue #9443: string labels stay distinct during aggregation, but equal +-- numeric bucket boundaries must be coalesced before monotonicity repair. +CREATE TABLE equivalent_bounds_bucket ( + ts TIMESTAMP TIME INDEX, + instance STRING, + le STRING, + val DOUBLE, + PRIMARY KEY (instance, le) +); + +Affected Rows: 0 + +INSERT INTO equivalent_bounds_bucket VALUES + (3000000, 'a', '10', 40), + (3000000, 'a', '100', 50), + (3000000, 'a', '+Inf', 50), + (3000000, 'b', '10.0', 40), + (3000000, 'b', '100.0', 50), + (3000000, 'b', '1000.0', 50), + (3000000, 'b', '+Inf', 50); + +Affected Rows: 7 + +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') sum by (le) (equivalent_bounds_bucket); + ++--------+---------------------+-----------------------------------+ +| le | ts | sum(equivalent_bounds_bucket.val) | ++--------+---------------------+-----------------------------------+ +| +Inf | 1970-01-01T00:50:00 | 100.0 | +| 10 | 1970-01-01T00:50:00 | 40.0 | +| 10.0 | 1970-01-01T00:50:00 | 40.0 | +| 100 | 1970-01-01T00:50:00 | 50.0 | +| 100.0 | 1970-01-01T00:50:00 | 50.0 | +| 1000.0 | 1970-01-01T00:50:00 | 50.0 | ++--------+---------------------+-----------------------------------+ + +tql eval (3000, 3000, '1s') histogram_quantile(0.95, sum by (le) (equivalent_bounds_bucket)); + ++---------------------+-----------------------------------+ +| ts | sum(equivalent_bounds_bucket.val) | ++---------------------+-----------------------------------+ +| 1970-01-01T00:50:00 | 77.5 | ++---------------------+-----------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') histogram_quantile(0.95, sum by (instance, le) (equivalent_bounds_bucket)); + ++----------+---------------------+-----------------------------------+ +| instance | ts | sum(equivalent_bounds_bucket.val) | ++----------+---------------------+-----------------------------------+ +| a | 1970-01-01T00:50:00 | 77.5 | +| b | 1970-01-01T00:50:00 | 77.5 | ++----------+---------------------+-----------------------------------+ + +DROP TABLE equivalent_bounds_bucket; + +Affected Rows: 0 + +-- Equivalent infinity labels belong to the same histogram until the group ends. +CREATE TABLE infinity_bounds_bucket ( + ts TIMESTAMP TIME INDEX, + le STRING, + val DOUBLE, + PRIMARY KEY (le) +); + +Affected Rows: 0 + +INSERT INTO infinity_bounds_bucket VALUES + (3000000, '10', 20), + (3000000, '10.0', 20), + (3000000, '+Inf', 25), + (3000000, '+inf', 25); + +Affected Rows: 4 + +tql eval (3000, 3000, '1s') histogram_quantile(0.5, infinity_bounds_bucket); + ++---------------------+------+ +| ts | val | ++---------------------+------+ +| 1970-01-01T00:50:00 | 6.25 | ++---------------------+------+ + +DROP TABLE infinity_bounds_bucket; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/simple_histogram.sql b/tests/cases/standalone/common/promql/simple_histogram.sql index ec0bf9b1439..6f73325a61f 100644 --- a/tests/cases/standalone/common/promql/simple_histogram.sql +++ b/tests/cases/standalone/common/promql/simple_histogram.sql @@ -253,3 +253,50 @@ insert into histogram5_bucket values tql eval (3000, 3015, '3s') histogram_quantile(0.5, histogram5_bucket); drop table histogram5_bucket; + +-- Issue #9443: string labels stay distinct during aggregation, but equal +-- numeric bucket boundaries must be coalesced before monotonicity repair. +CREATE TABLE equivalent_bounds_bucket ( + ts TIMESTAMP TIME INDEX, + instance STRING, + le STRING, + val DOUBLE, + PRIMARY KEY (instance, le) +); + +INSERT INTO equivalent_bounds_bucket VALUES + (3000000, 'a', '10', 40), + (3000000, 'a', '100', 50), + (3000000, 'a', '+Inf', 50), + (3000000, 'b', '10.0', 40), + (3000000, 'b', '100.0', 50), + (3000000, 'b', '1000.0', 50), + (3000000, 'b', '+Inf', 50); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') sum by (le) (equivalent_bounds_bucket); + +tql eval (3000, 3000, '1s') histogram_quantile(0.95, sum by (le) (equivalent_bounds_bucket)); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') histogram_quantile(0.95, sum by (instance, le) (equivalent_bounds_bucket)); + +DROP TABLE equivalent_bounds_bucket; + +-- Equivalent infinity labels belong to the same histogram until the group ends. +CREATE TABLE infinity_bounds_bucket ( + ts TIMESTAMP TIME INDEX, + le STRING, + val DOUBLE, + PRIMARY KEY (le) +); + +INSERT INTO infinity_bounds_bucket VALUES + (3000000, '10', 20), + (3000000, '10.0', 20), + (3000000, '+Inf', 25), + (3000000, '+inf', 25); + +tql eval (3000, 3000, '1s') histogram_quantile(0.5, infinity_bounds_bucket); + +DROP TABLE infinity_bounds_bucket; diff --git a/tests/cases/standalone/common/promql/time_fn.result b/tests/cases/standalone/common/promql/time_fn.result index 92ee594150c..793876f867e 100644 --- a/tests/cases/standalone/common/promql/time_fn.result +++ b/tests/cases/standalone/common/promql/time_fn.result @@ -150,7 +150,7 @@ tql eval (1, 2, '1s') time() == bool metrics; tql eval (1, 2, '1s') metrics + time(); +---------------------+-------------------------------------+ -| time | metrics.val + .time / Float64(1000) | +| ts | metrics.val + .time / Float64(1000) | +---------------------+-------------------------------------+ | 1970-01-01T00:00:01 | 2.0 | | 1970-01-01T00:00:02 | 4.0 | @@ -170,7 +170,7 @@ tql eval (1, 2, '1s') metrics == time(); tql eval (1, 2, '1s') metrics == bool time(); +---------------------+-------------------------------------+ -| time | metrics.val = .time / Float64(1000) | +| ts | metrics.val = .time / Float64(1000) | +---------------------+-------------------------------------+ | 1970-01-01T00:00:01 | 1.0 | | 1970-01-01T00:00:02 | 1.0 | diff --git a/tests/cases/standalone/common/promql/timestamp_fn.result b/tests/cases/standalone/common/promql/timestamp_fn.result index 7cc2b24484e..a7d7007a776 100644 --- a/tests/cases/standalone/common/promql/timestamp_fn.result +++ b/tests/cases/standalone/common/promql/timestamp_fn.result @@ -95,7 +95,7 @@ tql eval (0, 60, '30s') timestamp(timestamp_test) > bool 30; tql eval (0, 60, '30s') timestamp(timestamp_test) - time(); +---------------------+----------------------------------------------+ -| time | timestamp_test.value - .time / Float64(1000) | +| ts | timestamp_test.value - .time / Float64(1000) | +---------------------+----------------------------------------------+ | 1970-01-01T00:00:00 | 0.0 | | 1970-01-01T00:00:30 | -29.0 |