diff --git a/src/promql/src/binary/mod.rs b/src/promql/src/binary/mod.rs index 83da443402..8b64ec6d01 100644 --- a/src/promql/src/binary/mod.rs +++ b/src/promql/src/binary/mod.rs @@ -17,4 +17,6 @@ mod scalar_operations; mod vector_operations; pub(crate) use scalar_operations::scalar_binary_operations; -pub(crate) use vector_operations::{vector_bin_op, vector_scalar_bin_op}; +pub(crate) use vector_operations::{ + vector_bin_op, vector_scalar_bin_op, vector_step_scalar_bin_op, +}; diff --git a/src/promql/src/binary/vector_operations.rs b/src/promql/src/binary/vector_operations.rs index 44c1f3a2f7..4a413ce34b 100644 --- a/src/promql/src/binary/vector_operations.rs +++ b/src/promql/src/binary/vector_operations.rs @@ -46,6 +46,28 @@ pub fn vector_scalar_bin_op( left: Vec, right: f64, swapped_lhs_rhs: bool, +) -> Result { + vector_scalar_op(expr, left, swapped_lhs_rhs, |_| Some(right)) +} + +/// Implement the operation between a matrix and a scalar that holds one value per step. +pub fn vector_step_scalar_bin_op( + expr: &BinaryExpr, + left: Vec, + right: &[Sample], + swapped_lhs_rhs: bool, +) -> Result { + let right: HashMap = right.iter().map(|s| (s.timestamp, s.value)).collect(); + vector_scalar_op(expr, left, swapped_lhs_rhs, |timestamp| { + right.get(×tamp).copied() + }) +} + +fn vector_scalar_op( + expr: &BinaryExpr, + left: Vec, + swapped_lhs_rhs: bool, + scalar_at: impl Fn(i64) -> Option + Sync, ) -> Result { let is_comparison_operator = expr.op.is_comparison_operator(); let return_bool = expr.return_bool(); @@ -56,6 +78,7 @@ pub fn vector_scalar_bin_op( .samples .into_iter() .filter_map(|sample| { + let right = scalar_at(sample.timestamp)?; let (lhs, rhs) = if swapped_lhs_rhs { (right, sample.value) } else { diff --git a/src/promql/src/engine/mod.rs b/src/promql/src/engine/mod.rs index b5ed88c8f3..261f718763 100644 --- a/src/promql/src/engine/mod.rs +++ b/src/promql/src/engine/mod.rs @@ -116,12 +116,25 @@ impl Engine { let lhs = scalar_operand(lhs, &expr.lhs, &self.eval_ctx); let rhs = scalar_operand(rhs, &expr.rhs, &self.eval_ctx); + let lhs_scalar = expr.lhs.value_type() == ValueType::Scalar; + let rhs_scalar = expr.rhs.value_type() == ValueType::Scalar; match (lhs, rhs) { (Value::Float(left), Value::Float(right)) => { let value = binary::scalar_binary_operations(token, left, right, return_bool, op)?; Value::Float(value) } + // a range-query scalar is one label-less series that label matching would drop + (Value::Matrix(left), Value::Matrix(right)) + if rhs_scalar && !lhs_scalar && right.len() == 1 => + { + binary::vector_step_scalar_bin_op(expr, left, &right[0].samples, false)? + } + (Value::Matrix(left), Value::Matrix(right)) + if lhs_scalar && !rhs_scalar && left.len() == 1 => + { + binary::vector_step_scalar_bin_op(expr, right, &left[0].samples, true)? + } (Value::Matrix(left), Value::Matrix(right)) => { binary::vector_bin_op(expr, left, right)? } @@ -768,10 +781,17 @@ pub(crate) mod tests { ] { let series = matrix(eval_on_empty(&query, 3).await.unwrap()); assert_eq!(series.len(), 1, "{query}"); - assert_eq!(series[0].samples.len(), 1, "{query}"); - assert_eq!(series[0].samples[0].timestamp, timestamp, "{query}"); + // a scalar holds a value at every step, so a rejected step is NaN, not absent + assert_eq!(series[0].samples.len(), 3, "{query}"); + let matched: Vec<_> = series[0] + .samples + .iter() + .filter(|sample| !sample.value.is_nan()) + .collect(); + assert_eq!(matched.len(), 1, "{query}"); + assert_eq!(matched[0].timestamp, timestamp, "{query}"); assert_eq!( - series[0].samples[0].value, + matched[0].value, timestamp as f64 / 1_000_000.0 + 1.0, "{query}" ); @@ -802,6 +822,24 @@ pub(crate) mod tests { } } + #[tokio::test] + async fn test_range_scalar_operand_broadcasts_to_labelled_series() { + let labelled = r#"label_replace(vector(6), "job", "x", "", "")"#; + for (query, expected) in [ + (format!("{labelled} / scalar(vector(3))"), 2.0), + (format!("scalar(vector(3)) / {labelled}"), 0.5), + (format!("{labelled} > scalar(vector(3))"), 6.0), + (format!("scalar(vector(3)) < {labelled}"), 6.0), + ] { + let series = matrix(eval_on_empty(&query, 3).await.unwrap()); + assert_eq!(series.len(), 1, "{query}"); + assert_eq!(series[0].labels.len(), 1, "{query}"); + assert_eq!(series[0].labels[0].value, "x", "{query}"); + let values: Vec = series[0].samples.iter().map(|s| s.value).collect(); + assert_eq!(values, vec![expected; 3], "{query}"); + } + } + #[test] fn test_scalar_operand_keeps_single_step_range_timestamp() { let eval_ctx = EvalContext::new(1_000_000, 1_500_000, 1_000_000, "test".into()); diff --git a/src/promql/src/functions/scalar.rs b/src/promql/src/functions/scalar.rs index ac0027c537..2372457217 100644 --- a/src/promql/src/functions/scalar.rs +++ b/src/promql/src/functions/scalar.rs @@ -13,20 +13,53 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use config::meta::promql::value::{EvalContext, Value}; +use config::meta::promql::value::{EvalContext, Labels, RangeValue, Sample, Value}; use datafusion::error::{DataFusionError, Result}; +/// https://prometheus.io/docs/prometheus/latest/querying/functions/#scalar pub(crate) fn scalar(data: Value, eval_ctx: &EvalContext) -> Result { - match data { - Value::Float(_) | Value::Matrix(_) => Ok(data), - Value::None => { - // Generate samples using timestamps from eval_ctx - super::vector(Value::Float(f64::NAN), eval_ctx) + let matrix = match data { + Value::Float(_) => return Ok(data), + Value::Matrix(matrix) => matrix, + Value::None => Vec::new(), + _ => { + return Err(DataFusionError::Plan( + "Unexpected input. Expected: \"vector(s scalar)\"".into(), + )); + } + }; + + let timestamps = eval_ctx.timestamps(); + let mut counts = vec![0u32; timestamps.len()]; + let mut values = vec![f64::NAN; timestamps.len()]; + for sample in matrix.iter().flat_map(|series| &series.samples) { + if let Ok(step) = timestamps.binary_search(&sample.timestamp) { + counts[step] += 1; + values[step] = sample.value; } - _ => Err(DataFusionError::Plan( - "Unexpected input. Expected: \"vector(s scalar)\"".into(), - )), } + // a series can report at some steps and not others, so the element count belongs to a step + for (value, count) in values.iter_mut().zip(counts) { + if count != 1 { + *value = f64::NAN; + } + } + + // exec and the binary operators only treat a Float as a scalar + if eval_ctx.is_instant() { + return Ok(Value::Float(values[0])); + } + + Ok(Value::Matrix(vec![RangeValue { + labels: Labels::default(), + samples: timestamps + .into_iter() + .zip(values) + .map(|(timestamp, value)| Sample::new(timestamp, value)) + .collect(), + exemplars: None, + time_window: None, + }])) } #[cfg(test)] @@ -61,11 +94,62 @@ mod tests { assert!(ranges[0].samples.iter().all(|s| s.value.is_nan())); } + fn series(samples: Vec) -> RangeValue { + RangeValue { + labels: Labels::default(), + samples, + exemplars: None, + time_window: None, + } + } + #[test] - fn test_scalar_matrix_passthrough() { - let matrix = Value::Matrix(vec![]); - let result = scalar(matrix, &instant_ctx()).unwrap(); - assert!(matches!(result, Value::Matrix(_))); + fn test_scalar_of_one_element_is_that_value() { + let one = Value::Matrix(vec![series(vec![Sample::new(1_000_000, 42.0)])]); + let result = scalar(one, &instant_ctx()).unwrap(); + assert!(matches!(result, Value::Float(f) if f == 42.0)); + } + + #[test] + fn test_scalar_of_several_elements_is_nan() { + let several = Value::Matrix(vec![ + series(vec![Sample::new(1_000_000, 42.0)]), + series(vec![Sample::new(1_000_000, 7.0)]), + ]); + let result = scalar(several, &instant_ctx()).unwrap(); + assert!( + matches!(result, Value::Float(f) if f.is_nan()), + "more than one element is not a scalar" + ); + } + + #[test] + fn test_scalar_of_no_elements_is_nan() { + let result = scalar(Value::Matrix(vec![]), &instant_ctx()).unwrap(); + assert!(matches!(result, Value::Float(f) if f.is_nan())); + } + + #[test] + fn test_scalar_counts_elements_at_each_step_of_a_range_query() { + // Two series report at the middle step only, so that step alone is ambiguous. + let matrix = Value::Matrix(vec![ + series(vec![ + Sample::new(1_000_000, 42.0), + Sample::new(2_000_000, 42.0), + Sample::new(3_000_000, 42.0), + ]), + series(vec![Sample::new(2_000_000, 7.0)]), + ]); + let result = scalar(matrix, &range_ctx()).unwrap(); + + let Value::Matrix(ranges) = result else { + panic!("expected Matrix"); + }; + assert_eq!(ranges.len(), 1); + let values: Vec = ranges[0].samples.iter().map(|s| s.value).collect(); + assert_eq!(values[0], 42.0); + assert!(values[1].is_nan(), "two elements at this step"); + assert_eq!(values[2], 42.0); } #[test] diff --git a/src/promql_service/src/search/mod.rs b/src/promql_service/src/search/mod.rs index 721041cf5b..c6f2e83520 100644 --- a/src/promql_service/src/search/mod.rs +++ b/src/promql_service/src/search/mod.rs @@ -479,9 +479,8 @@ async fn search_in_cluster( series_data.push(series); }); - // with cache maybe we only get the last point from original data, then the result_type will - // return as vector, but if the query is range query, the result_type should be matrix - if result_type == "vector" && original_start != end { + // a worker left with a lone point answers as an instant query, the range is still a matrix + if (result_type == "vector" || result_type == "scalar") && original_start != end { result_type = "matrix".to_string(); }