diff --git a/src/promql/src/extension_plan/instant_manipulate.rs b/src/promql/src/extension_plan/instant_manipulate.rs index 9abec5001f1..b73a968d973 100644 --- a/src/promql/src/extension_plan/instant_manipulate.rs +++ b/src/promql/src/extension_plan/instant_manipulate.rs @@ -331,6 +331,7 @@ impl Stream for InstantManipulateStream { impl InstantManipulateStream { // refer to Go version: https://github.com/prometheus/prometheus/blob/e934d0f01158a1d55fa0ebb035346b195fcc1260/promql/engine.go#L1571 + // and the function `vectorSelectorSingle` pub fn manipulate(&self, input: RecordBatch) -> DataFusionResult { let mut take_indices = Vec::with_capacity(input.num_rows()); // TODO(ruihang): maybe the input is not timestamp millisecond array @@ -373,28 +374,36 @@ impl InstantManipulateStream { if cursor == ts_column.len() { cursor -= 1; } - // then, search backward to lookback - loop { - let curr = ts_column.value(cursor); - if let Some(field_column) = &field_column && field_column.value(cursor).is_nan() { - // if the newest value is NaN, it means the value is stale, so we should not use it + + // then examine the value + let curr_ts = ts_column.value(cursor); + if curr_ts + self.lookback_delta < expected_ts { + take_indices.push(None); + continue; + } + if curr_ts > expected_ts { + // exceeds current expected timestamp, examine the previous value + if let Some(prev_cursor) = cursor.checked_sub(1) { + let prev_ts = ts_column.value(prev_cursor); + if prev_ts + self.lookback_delta < expected_ts { + // not found in lookback, leave this field blank. + take_indices.push(None); + } else if let Some(field_column) = &field_column && field_column.value(prev_cursor).is_nan() { + // if the newest value is NaN, it means the value is stale, so we should not use it + take_indices.push(None); + } else { + // use this point + take_indices.push(Some(prev_cursor as u64)); + } + } else { take_indices.push(None); - break; } - if curr + self.lookback_delta < expected_ts { - // not found in lookback, leave this field blank. - take_indices.push(None); - break; - } else if curr < expected_ts && curr + self.lookback_delta >= expected_ts { - // find the expected value, push and break - take_indices.push(Some(cursor as u64)); - break; - } else if cursor == 0 { - // reach the first value and not found in lookback, leave this field blank - take_indices.push(None); - break; - } - cursor -= 1; + } else if let Some(field_column) = &field_column && field_column.value(cursor).is_nan() { + // if the newest value is NaN, it means the value is stale, so we should not use it + take_indices.push(None); + } else { + // use this point + take_indices.push(Some(cursor as u64)); } } @@ -478,15 +487,20 @@ mod test { lookback_delta: Millisecond, interval: Millisecond, expected: String, + contains_nan: bool, ) { - let memory_exec = Arc::new(prepare_test_data()); + let memory_exec = if contains_nan { + Arc::new(prepare_test_data_with_nan()) + } else { + Arc::new(prepare_test_data()) + }; let normalize_exec = Arc::new(InstantManipulateExec { start, end, lookback_delta, interval, time_index_column: TIME_INDEX_COLUMN.to_string(), - field_column: None, + field_column: Some("value".to_string()), input: memory_exec, metric: ExecutionPlanMetricsSet::new(), }); @@ -517,7 +531,7 @@ mod test { \n| 1970-01-01T00:05:00 | 1.0 | foo |\ \n+---------------------+-------+------+", ); - do_normalize_test(0, 310_000, 10_000, 30_000, expected).await; + do_normalize_test(0, 310_000, 10_000, 30_000, expected, false).await; } #[tokio::test] @@ -544,7 +558,7 @@ mod test { \n| 1970-01-01T00:05:00 | 1.0 | foo |\ \n+---------------------+-------+------+", ); - do_normalize_test(0, 300_000, 10_000, 10_000, expected).await; + do_normalize_test(0, 300_000, 10_000, 10_000, expected, false).await; } #[tokio::test] @@ -566,7 +580,7 @@ mod test { \n| 1970-01-01T00:05:00 | 1.0 | foo |\ \n+---------------------+-------+------+", ); - do_normalize_test(0, 300_000, 30_000, 30_000, expected).await; + do_normalize_test(0, 300_000, 30_000, 30_000, expected, false).await; } #[tokio::test] @@ -604,7 +618,7 @@ mod test { \n| 1970-01-01T00:05:00 | 1.0 | foo |\ \n+---------------------+-------+------+", ); - do_normalize_test(0, 300_000, 30_000, 10_000, expected).await; + do_normalize_test(0, 300_000, 30_000, 10_000, expected, false).await; } #[tokio::test] @@ -646,7 +660,7 @@ mod test { \n| 1970-01-01T00:05:00 | 1.0 | foo |\ \n+---------------------+-------+------+", ); - do_normalize_test(0, 300_000, 60_000, 10_000, expected).await; + do_normalize_test(0, 300_000, 60_000, 10_000, expected, false).await; } #[tokio::test] @@ -668,7 +682,7 @@ mod test { \n| 1970-01-01T00:05:00 | 1.0 | foo |\ \n+---------------------+-------+------+", ); - do_normalize_test(0, 300_000, 60_000, 30_000, expected).await; + do_normalize_test(0, 300_000, 60_000, 30_000, expected, false).await; } #[tokio::test] @@ -681,7 +695,7 @@ mod test { \n| 1970-01-01T00:04:01 | 1.0 | foo |\ \n+---------------------+-------+------+", ); - do_normalize_test(230_000, 245_000, 0, 1_000, expected).await; + do_normalize_test(230_000, 245_000, 0, 1_000, expected, false).await; } #[tokio::test] @@ -695,7 +709,7 @@ mod test { \n| 1970-01-01T00:00:30 | 1.0 | foo |\ \n+---------------------+-------+------+", ); - do_normalize_test(0, 30_000, 10_000, 10_000, expected).await; + do_normalize_test(0, 30_000, 10_000, 10_000, expected, false).await; } #[tokio::test] @@ -712,7 +726,7 @@ mod test { \n| 1970-01-01T00:05:00 | 1.0 | foo |\ \n+---------------------+-------+------+", ); - do_normalize_test(-900_000, 900_000, 30_000, 60_000, expected).await; + do_normalize_test(-900_000, 900_000, 30_000, 60_000, expected, false).await; } #[tokio::test] @@ -733,6 +747,58 @@ mod test { \n| 1970-01-01T00:05:00 | 1.0 | foo |\ \n+---------------------+-------+------+", ); - do_normalize_test(190_000, 300_000, 30_000, 10_000, expected).await; + do_normalize_test(190_000, 300_000, 30_000, 10_000, expected, false).await; + } + + fn prepare_test_data_with_nan() -> MemoryExec { + let schema = Arc::new(Schema::new(vec![ + Field::new(TIME_INDEX_COLUMN, TimestampMillisecondType::DATA_TYPE, true), + Field::new("value", DataType::Float64, true), + ])); + let timestamp_column = Arc::new(TimestampMillisecondArray::from_slice([ + 0, 30_000, 60_000, 90_000, 120_000, // every 30s + ])) as _; + let field_column = Arc::new(Float64Array::from_slice([ + 0.0, + f64::NAN, + 6.0, + f64::NAN, + 12.0, + ])) as _; + let data = + RecordBatch::try_new(schema.clone(), vec![timestamp_column, field_column]).unwrap(); + + MemoryExec::try_new(&[vec![data]], schema, None).unwrap() + } + + #[tokio::test] + async fn lookback_10s_interval_10s_with_nan() { + let expected = String::from( + "+---------------------+-------+\ + \n| timestamp | value |\ + \n+---------------------+-------+\ + \n| 1970-01-01T00:00:00 | 0.0 |\ + \n| 1970-01-01T00:00:10 | 0.0 |\ + \n| 1970-01-01T00:01:00 | 6.0 |\ + \n| 1970-01-01T00:01:10 | 6.0 |\ + \n| 1970-01-01T00:02:00 | 12.0 |\ + \n| 1970-01-01T00:02:10 | 12.0 |\ + \n+---------------------+-------+", + ); + do_normalize_test(0, 300_000, 10_000, 10_000, expected, true).await; + } + + #[tokio::test] + async fn lookback_10s_interval_10s_with_nan_unaligned() { + let expected = String::from( + "+-------------------------+-------+\ + \n| timestamp | value |\ + \n+-------------------------+-------+\ + \n| 1970-01-01T00:00:00.001 | 0.0 |\ + \n| 1970-01-01T00:01:00.001 | 6.0 |\ + \n| 1970-01-01T00:02:00.001 | 12.0 |\ + \n+-------------------------+-------+", + ); + do_normalize_test(1, 300_001, 10_000, 10_000, expected, true).await; } } diff --git a/src/promql/src/planner.rs b/src/promql/src/planner.rs index 95bb81c2efa..dcaaf891245 100644 --- a/src/promql/src/planner.rs +++ b/src/promql/src/planner.rs @@ -184,11 +184,8 @@ impl PromPlanner { self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string()); self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()]; self.ctx.table_name = Some(String::new()); - let field_expr = DfExpr::BinaryExpr(BinaryExpr { - left: Box::new(lhs), - op: Self::prom_token_to_binary_op(*op)?, - right: Box::new(rhs), - }); + let field_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; + let field_expr = field_expr_builder(lhs, rhs)?; LogicalPlan::Extension(Extension { node: Arc::new( @@ -208,11 +205,10 @@ impl PromPlanner { (Some(expr), None) => { let input = self.prom_expr_to_plan(*rhs.clone()).await?; let bin_expr_builder = |col: &String| { - let mut binary_expr = DfExpr::BinaryExpr(BinaryExpr { - left: Box::new(expr.clone()), - op: Self::prom_token_to_binary_op(*op)?, - right: Box::new(DfExpr::Column(col.into())), - }); + let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; + let mut binary_expr = + binary_expr_builder(expr.clone(), DfExpr::Column(col.into()))?; + if is_comparison_op && should_return_bool { binary_expr = DfExpr::Cast(Cast { expr: Box::new(binary_expr), @@ -231,11 +227,10 @@ impl PromPlanner { (None, Some(expr)) => { let input = self.prom_expr_to_plan(*lhs.clone()).await?; let bin_expr_builder = |col: &String| { - let mut binary_expr = DfExpr::BinaryExpr(BinaryExpr { - left: Box::new(DfExpr::Column(col.into())), - op: Self::prom_token_to_binary_op(*op)?, - right: Box::new(expr.clone()), - }); + let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; + let mut binary_expr = + binary_expr_builder(DfExpr::Column(col.into()), expr.clone())?; + if is_comparison_op && should_return_bool { binary_expr = DfExpr::Cast(Cast { expr: Box::new(binary_expr), @@ -275,11 +270,11 @@ impl PromPlanner { .context(DataFusionPlanningSnafu)? .qualified_column(); - let mut binary_expr = DfExpr::BinaryExpr(BinaryExpr { - left: Box::new(DfExpr::Column(left_col)), - op: Self::prom_token_to_binary_op(*op)?, - right: Box::new(DfExpr::Column(right_col)), - }); + let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; + let mut binary_expr = binary_expr_builder( + DfExpr::Column(left_col), + DfExpr::Column(right_col), + )?; if is_comparison_op && should_return_bool { binary_expr = DfExpr::Cast(Cast { expr: Box::new(binary_expr), @@ -1103,35 +1098,65 @@ impl PromPlanner { PromExpr::Paren(ParenExpr { expr }) => Self::try_build_literal_expr(expr), // TODO(ruihang): support Unary operator PromExpr::Unary(UnaryExpr { expr, .. }) => Self::try_build_literal_expr(expr), - PromExpr::Binary(PromBinaryExpr { lhs, rhs, op, .. }) => { + PromExpr::Binary(PromBinaryExpr { + lhs, + rhs, + op, + modifier, + }) => { let lhs = Self::try_build_literal_expr(lhs)?; let rhs = Self::try_build_literal_expr(rhs)?; - let op = Self::prom_token_to_binary_op(*op).ok()?; - Some(DfExpr::BinaryExpr(BinaryExpr { - left: Box::new(lhs), - op, - right: Box::new(rhs), - })) + let is_comparison_op = Self::is_token_a_comparison_op(*op); + let expr_builder = Self::prom_token_to_binary_expr_builder(*op).ok()?; + let expr = expr_builder(lhs, rhs).ok()?; + + let should_return_bool = if let Some(m) = modifier { + m.return_bool + } else { + false + }; + if is_comparison_op && should_return_bool { + Some(DfExpr::Cast(Cast { + expr: Box::new(expr), + data_type: ArrowDataType::Float64, + })) + } else { + Some(expr) + } } } } - fn prom_token_to_binary_op(token: TokenType) -> Result { + /// Return a lambda to build binary expression from token. + /// Because some binary operator are function in DataFusion like `atan2` or `^`. + #[allow(clippy::type_complexity)] + fn prom_token_to_binary_expr_builder( + token: TokenType, + ) -> Result Result>> { match token.id() { - token::T_ADD => Ok(Operator::Plus), - token::T_SUB => Ok(Operator::Minus), - token::T_MUL => Ok(Operator::Multiply), - token::T_DIV => Ok(Operator::Divide), - token::T_MOD => Ok(Operator::Modulo), - token::T_EQLC => Ok(Operator::Eq), - token::T_NEQ => Ok(Operator::NotEq), - token::T_GTR => Ok(Operator::Gt), - token::T_LSS => Ok(Operator::Lt), - token::T_GTE => Ok(Operator::GtEq), - token::T_LTE => Ok(Operator::LtEq), - // TODO(ruihang): support these two operators - // token::T_POW => Ok(Operator::Power), - // token::T_ATAN2 => Ok(Operator::Atan2), + token::T_ADD => Ok(Box::new(|lhs, rhs| Ok(lhs + rhs))), + token::T_SUB => Ok(Box::new(|lhs, rhs| Ok(lhs - rhs))), + token::T_MUL => Ok(Box::new(|lhs, rhs| Ok(lhs * rhs))), + token::T_DIV => Ok(Box::new(|lhs, rhs| Ok(lhs / rhs))), + token::T_MOD => Ok(Box::new(|lhs: DfExpr, rhs| Ok(lhs % rhs))), + token::T_EQLC => Ok(Box::new(|lhs, rhs| Ok(lhs.eq(rhs)))), + token::T_NEQ => Ok(Box::new(|lhs, rhs| Ok(lhs.not_eq(rhs)))), + token::T_GTR => Ok(Box::new(|lhs, rhs| Ok(lhs.gt(rhs)))), + token::T_LSS => Ok(Box::new(|lhs, rhs| Ok(lhs.lt(rhs)))), + token::T_GTE => Ok(Box::new(|lhs, rhs| Ok(lhs.gt_eq(rhs)))), + token::T_LTE => Ok(Box::new(|lhs, rhs| Ok(lhs.lt_eq(rhs)))), + token::T_POW => Ok(Box::new(|lhs, rhs| { + Ok(DfExpr::ScalarFunction(ScalarFunction { + fun: BuiltinScalarFunction::Power, + args: vec![lhs, rhs], + })) + })), + token::T_ATAN2 => Ok(Box::new(|lhs, rhs| { + Ok(DfExpr::ScalarFunction(ScalarFunction { + fun: BuiltinScalarFunction::Atan2, + args: vec![lhs, rhs], + })) + })), _ => UnexpectedTokenSnafu { token }.fail(), } } diff --git a/tests/cases/standalone/common/tql/binary_operator.result b/tests/cases/standalone/common/tql/binary_operator.result new file mode 100644 index 00000000000..bce78de11c1 --- /dev/null +++ b/tests/cases/standalone/common/tql/binary_operator.result @@ -0,0 +1,39 @@ +create table data (ts timestamp(3) time index, val double); + +Affected Rows: 0 + +insert into data values (0, 1), (10000, 2), (20000, 3); + +Affected Rows: 3 + +tql eval (0, 30, '10s'), data < 1; + +++ +++ + +tql eval (0, 30, '10s'), data + (1 < bool 2); + ++---------------------+-------------------------------+ +| ts | val + Float64(1) < Float64(2) | ++---------------------+-------------------------------+ +| 1970-01-01T00:00:00 | 2.0 | +| 1970-01-01T00:00:10 | 3.0 | +| 1970-01-01T00:00:20 | 4.0 | +| 1970-01-01T00:00:30 | 4.0 | ++---------------------+-------------------------------+ + +tql eval (0, 30, '10s'), data + (1 > bool 2); + ++---------------------+-------------------------------+ +| ts | val + Float64(1) > Float64(2) | ++---------------------+-------------------------------+ +| 1970-01-01T00:00:00 | 1.0 | +| 1970-01-01T00:00:10 | 2.0 | +| 1970-01-01T00:00:20 | 3.0 | +| 1970-01-01T00:00:30 | 3.0 | ++---------------------+-------------------------------+ + +drop table data; + +Affected Rows: 1 + diff --git a/tests/cases/standalone/common/tql/binary_operator.sql b/tests/cases/standalone/common/tql/binary_operator.sql new file mode 100644 index 00000000000..82c2682fbfb --- /dev/null +++ b/tests/cases/standalone/common/tql/binary_operator.sql @@ -0,0 +1,11 @@ +create table data (ts timestamp(3) time index, val double); + +insert into data values (0, 1), (10000, 2), (20000, 3); + +tql eval (0, 30, '10s'), data < 1; + +tql eval (0, 30, '10s'), data + (1 < bool 2); + +tql eval (0, 30, '10s'), data + (1 > bool 2); + +drop table data; diff --git a/tests/cases/standalone/common/tql/literal_only.result b/tests/cases/standalone/common/tql/literal_only.result index f30b53fda42..a21aa97b08f 100644 --- a/tests/cases/standalone/common/tql/literal_only.result +++ b/tests/cases/standalone/common/tql/literal_only.result @@ -28,3 +28,13 @@ TQL EVAL (0, 10, '5s') "1+1"; | 1970-01-01T00:00:10 | 1+1 | +---------------------+-------+ +TQL EVAL (0, 10, '5s') 1 * 2 + 4 / 6 - 10 % 2 ^ 2; + ++---------------------+--------------------+ +| time | value | ++---------------------+--------------------+ +| 1970-01-01T00:00:00 | 0.6666666666666665 | +| 1970-01-01T00:00:05 | 0.6666666666666665 | +| 1970-01-01T00:00:10 | 0.6666666666666665 | ++---------------------+--------------------+ + diff --git a/tests/cases/standalone/common/tql/literal_only.sql b/tests/cases/standalone/common/tql/literal_only.sql index dc353146f61..0dd34f10909 100644 --- a/tests/cases/standalone/common/tql/literal_only.sql +++ b/tests/cases/standalone/common/tql/literal_only.sql @@ -3,3 +3,5 @@ TQL EVAL (0, 10, '5s') 1; TQL EVAL (0, 10, '5s') 1+1; TQL EVAL (0, 10, '5s') "1+1"; + +TQL EVAL (0, 10, '5s') 1 * 2 + 4 / 6 - 10 % 2 ^ 2; diff --git a/tests/cases/standalone/common/tql/operator.result b/tests/cases/standalone/common/tql/operator.result new file mode 100644 index 00000000000..f5a4d66675e --- /dev/null +++ b/tests/cases/standalone/common/tql/operator.result @@ -0,0 +1,53 @@ +create table trigy (ts timestamp(3) time index, val double); + +Affected Rows: 0 + +insert into trigy values (0,10); + +Affected Rows: 1 + +create table trigx (ts timestamp(3) time index, val double); + +Affected Rows: 0 + +insert into trigx values (0,20); + +Affected Rows: 1 + +create table trignan (ts timestamp(3) time index, val double); + +Affected Rows: 0 + +insert into trignan values (0, 'NaN'::double); + +Affected Rows: 1 + +tql eval (300, 300, '1s') trigy atan2 trigx; + ++---------------------+----------------------------+ +| ts | atan2(trigy.val,trigx.val) | ++---------------------+----------------------------+ +| 1970-01-01T00:05:00 | 0.4636476090008061 | ++---------------------+----------------------------+ + +tql eval (300, 300, '1s') trigy atan2 trignan; + +++ +++ + +tql eval (300, 300, '1s') 10 atan2 20; + ++---------------------+--------------------+ +| time | value | ++---------------------+--------------------+ +| 1970-01-01T00:05:00 | 0.4636476090008061 | ++---------------------+--------------------+ + +tql eval (300, 300, '1s') 10 atan2 NaN; + ++---------------------+-------+ +| time | value | ++---------------------+-------+ +| 1970-01-01T00:05:00 | NaN | ++---------------------+-------+ + diff --git a/tests/cases/standalone/common/tql/operator.sql b/tests/cases/standalone/common/tql/operator.sql new file mode 100644 index 00000000000..6a7efc37c49 --- /dev/null +++ b/tests/cases/standalone/common/tql/operator.sql @@ -0,0 +1,35 @@ +-- Port from operators.test L607 - L630, commit 001ee2620e094970e5657ce39275b2fccdbd1359 +-- Include atan2 + +-- load 5m +-- trigy{} 10 +-- trigx{} 20 +-- trigNaN{} NaN +create table trigy (ts timestamp(3) time index, val double); + +insert into trigy values (0,10); + +create table trigx (ts timestamp(3) time index, val double); + +insert into trigx values (0,20); + +create table trignan (ts timestamp(3) time index, val double); + +insert into trignan values (0, 'NaN'::double); + +-- eval instant at 5m trigy atan2 trigx +-- trigy{} 0.4636476090008061 +tql eval (300, 300, '1s') trigy atan2 trigx; + +-- eval instant at 5m trigy atan2 trigNaN +-- trigy{} NaN +-- This query doesn't have result because `trignan` is NaN and will be filtered out. +tql eval (300, 300, '1s') trigy atan2 trignan; + +-- eval instant at 5m 10 atan2 20 +-- 0.4636476090008061 +tql eval (300, 300, '1s') 10 atan2 20; + +-- eval instant at 5m 10 atan2 NaN +-- NaN +tql eval (300, 300, '1s') 10 atan2 NaN;