fix: make scalar() reduce to a scalar (#14619)

Co-authored-by: Huaijin <haohuaijin@gmail.com>

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,
+};

src/promql/src/binary/vector_operations.rs

@@ -46,6 +46,28 @@ pub fn vector_scalar_bin_op(
     left: Vec<RangeValue>,
     right: f64,
     swapped_lhs_rhs: bool,
+) -> Result<Value> {
+    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<RangeValue>,
+    right: &[Sample],
+    swapped_lhs_rhs: bool,
+) -> Result<Value> {
+    let right: HashMap<i64, f64> = right.iter().map(|s| (s.timestamp, s.value)).collect();
+    vector_scalar_op(expr, left, swapped_lhs_rhs, |timestamp| {
+        right.get(&timestamp).copied()
+    })
+}
+
+fn vector_scalar_op(
+    expr: &BinaryExpr,
+    left: Vec<RangeValue>,
+    swapped_lhs_rhs: bool,
+    scalar_at: impl Fn(i64) -> Option<f64> + Sync,
 ) -> Result<Value> {
     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 {

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<f64> = 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());

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 <http://www.gnu.org/licenses/>.
 
-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<Value> {
-    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<Sample>) -> 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<f64> = 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]

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();
     }