diff --git a/src/promql/src/select.rs b/src/promql/src/select.rs index a9d4035d98e..c61000f3486 100644 --- a/src/promql/src/select.rs +++ b/src/promql/src/select.rs @@ -70,12 +70,20 @@ pub fn topk( ) -> Result> { single_field(schema, "topk/bottomk")?; let points = grid_points(schema, series)?; - // Prometheus converts k with int64() and caps it at the number of input series. - let k = if k.is_nan() || k < 1.0 { - 0 - } else { - (k.trunc().min(series.len() as f64)) as usize - }; + // Prometheus returns nothing for k below 1, rejects NaN and values out of int64, then + // truncates k and caps it at the number of input series. + if k < 1.0 { + return Ok(Vec::new()); + } + if k.is_nan() { + return Err(DataFusionError::Execution("Parameter value is NaN".into())); + } + if k >= i64::MAX as f64 { + return Err(DataFusionError::Execution(format!( + "Scalar value {k} overflows int64" + ))); + } + let k = k.trunc().min(series.len() as f64) as usize; if k == 0 { return Ok(Vec::new()); } @@ -162,28 +170,28 @@ pub fn absent( labels: MetricName, ) -> Result> { let points = grid_points(schema, series)?; - let mut present = vec![false; points]; + let fields = schema.fields.len(); + // Each field is absent on its own. + let mut present = vec![false; points * fields]; for item in series { let TimeAxis::Grid(range) = &item.time else { unreachable!("checked by grid_points"); }; - let points = &mut present[range.start()..range.start() + range.len()]; - for (row, present) in points.iter_mut().enumerate() { - if item - .fields - .iter() - .any(|field| !matches!(field.sample(row), SampleRef::Missing)) - { - *present = true; + for (field, samples) in item.fields.iter().enumerate() { + for row in 0..range.len() { + if !matches!(samples.sample(row), SampleRef::Missing) { + present[(range.start() + row) * fields + field] = true; + } } } } - let ones = vec![Some(1.0); schema.fields.len()]; let mut builder = GridSeriesBuilder::try_new(Arc::new(schema.clone()))?; - for (point, present) in present.into_iter().enumerate() { - if !present { - builder.push(point, &ones)?; + let mut row = vec![None; fields]; + for point in 0..points { + for (field, value) in row.iter_mut().enumerate() { + *value = (!present[point * fields + field]).then_some(1.0); } + builder.push(point, &row)?; } builder.finish(labels) } diff --git a/src/promql/src/select/tests.rs b/src/promql/src/select/tests.rs index 376da22062d..786e0285777 100644 --- a/src/promql/src/select/tests.rs +++ b/src/promql/src/select/tests.rs @@ -271,3 +271,31 @@ fn rejects_value_operators_over_several_fields() { Err(DataFusionError::NotImplemented(_)) )); } + +/// Prometheus returns nothing for `k` below 1 and rejects NaN and values out of int64. +#[test] +fn topk_validates_k_like_prometheus() { + let schema = schema(1, 1); + let series = vec![build(&schema, metric("__name__=a"), &vec![vec![Some(1.0)]]).unwrap()]; + for k in [f64::NAN, f64::INFINITY, 2f64.powi(63)] { + assert!( + topk(k, true, &Grouping::By(Vec::new()), &schema, &series).is_err(), + "{k}" + ); + } + for k in [0.5, f64::NEG_INFINITY] { + let result = topk(k, false, &Grouping::By(Vec::new()), &schema, &series).unwrap(); + assert!(result.is_empty(), "{k}"); + } +} + +/// Each field of `absent` is 1 where that field has no sample. +#[test] +fn absent_evaluates_fields_independently() { + let schema = schema(2, 2); + let input = vec![vec![Some(3.0), None], vec![None, Some(4.0)]]; + let series = vec![build(&schema, metric("__name__=a"), &input).unwrap()]; + let result = absent(&schema, &series, metric("-")).unwrap(); + let expected = vec![vec![None, Some(1.0)], vec![Some(1.0), None]]; + assert_series_eq(result.as_ref(), &expected, &metric("-")); +}