From f956f4da30a9702097e5825a91eaba60a2178645 Mon Sep 17 00:00:00 2001 From: discord9 Date: Mon, 28 Sep 2026 05:30:33 +0000 Subject: [PATCH] refactor(promql): split planner into focused modules (#9354) * refactor(promql): move planner tests into separate module Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(promql): move function-specific planner methods into module Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(promql): move OR planner method into module Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(promql): move binary island planner into module Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(promql): keep binary result labels in planner root Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(promql): limit planner helper visibility to parent module Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(promql): group set operator planning and localize imports Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): use crate-rooted planner imports Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/promql/planner.rs | 9415 +---------------- .../src/promql/planner/function_plans.rs | 473 + src/query/src/promql/planner/island.rs | 536 + src/query/src/promql/planner/set_operator.rs | 832 ++ src/query/src/promql/planner/test.rs | 7585 +++++++++++++ 5 files changed, 9451 insertions(+), 9390 deletions(-) create mode 100644 src/query/src/promql/planner/function_plans.rs create mode 100644 src/query/src/promql/planner/island.rs create mode 100644 src/query/src/promql/planner/set_operator.rs create mode 100644 src/query/src/promql/planner/test.rs diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 73332c237c8..08e7005f9fb 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -12,9 +12,15 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod function_plans; + +mod island; + mod matching_filters; -use std::collections::{BTreeSet, HashMap, HashSet, VecDeque}; +mod set_operator; + +use std::collections::{BTreeSet, HashSet, VecDeque}; use std::sync::Arc; use std::time::UNIX_EPOCH; @@ -46,14 +52,12 @@ use datafusion::logical_expr::{ BinaryExpr, Cast, Extension, LogicalPlan, LogicalPlanBuilder, Operator, ScalarUDF as ScalarUdfDef, WindowFrame, WindowFunctionDefinition, }; -use datafusion::optimizer::simplify_expressions::ExprSimplifier; use datafusion::prelude as df_prelude; use datafusion::prelude::{Column, Expr as DfExpr, JoinType}; use datafusion::scalar::ScalarValue; use datafusion_common::{DFSchema, NullEquality, TableReference}; use datafusion_expr::expr::WindowFunctionParams; use datafusion_expr::expr_fn::when; -use datafusion_expr::simplify::SimplifyContext; use datafusion_expr::utils::{conjunction, disjunction}; use datafusion_expr::{ExprSchemable, Literal, SortExpr, TableSource, col, lit}; use datafusion_functions::core::coalesce; @@ -62,8 +66,7 @@ use datatypes::data_type::{ConcreteDataType, DataType as GreptimeDataType}; use itertools::Itertools; use once_cell::sync::Lazy; use promql::extension_plan::{ - Absent, EmptyMetric, HistogramFold, HistogramFoldOperation, InstantManipulate, Millisecond, - RangeManipulate, ScalarCalculate, SeriesDivide, SeriesNormalize, UnionDistinctOn, + EmptyMetric, InstantManipulate, Millisecond, RangeManipulate, SeriesDivide, SeriesNormalize, build_special_time_expr, }; use promql::functions::{ @@ -72,24 +75,22 @@ use promql::functions::{ NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAggAvg, NativeHistogramAggSum, NativeHistogramAvg, NativeHistogramAvgOverTime, NativeHistogramChanges, NativeHistogramCount, NativeHistogramCountOverTime, NativeHistogramDelta, - NativeHistogramDivScalar, NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction, - NativeHistogramIDelta, NativeHistogramIRate, NativeHistogramIncrease, - NativeHistogramLastOverTime, NativeHistogramMulScalar, NativeHistogramNeg, - NativeHistogramNotEq, NativeHistogramPresentOverTime, NativeHistogramQuantile, - NativeHistogramRate, NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev, - NativeHistogramStdvar, NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime, - NativeHistogramToString, PredictLinear, PresentOverTime, PromqlFloatToString, QuantileOverTime, - Rate, Resets, Round, StddevOverTime, StdvarOverTime, SumOverTime, UniqueMatchGroup, - quantile_udaf, + NativeHistogramDivScalar, NativeHistogramDrop, NativeHistogramEq, NativeHistogramIDelta, + NativeHistogramIRate, NativeHistogramIncrease, NativeHistogramLastOverTime, + NativeHistogramMulScalar, NativeHistogramNeg, NativeHistogramNotEq, + NativeHistogramPresentOverTime, NativeHistogramRate, NativeHistogramResets, + NativeHistogramScalarMul, NativeHistogramStddev, NativeHistogramStdvar, NativeHistogramSub, + NativeHistogramSum, NativeHistogramSumOverTime, NativeHistogramToString, PredictLinear, + PresentOverTime, PromqlFloatToString, QuantileOverTime, Rate, Resets, Round, StddevOverTime, + StdvarOverTime, SumOverTime, UniqueMatchGroup, quantile_udaf, }; use promql_parser::label::{METRIC_NAME, MatchOp, Matcher, Matchers}; use promql_parser::parser::token::TokenType; use promql_parser::parser::value::ValueType; use promql_parser::parser::{ AggregateExpr, BinModifier, BinaryExpr as PromBinaryExpr, Call, EvalStmt, Expr as PromExpr, - Function, FunctionArgs as PromFunctionArgs, LabelModifier, MatrixSelector, NumberLiteral, - Offset, ParenExpr, StringLiteral, SubqueryExpr, UnaryExpr, VectorMatchCardinality, - VectorSelector, token, + Function, LabelModifier, MatrixSelector, NumberLiteral, Offset, ParenExpr, StringLiteral, + SubqueryExpr, UnaryExpr, VectorMatchCardinality, VectorSelector, token, }; use regex::{self, Regex}; use snafu::{OptionExt, ResultExt, ensure}; @@ -104,14 +105,12 @@ use crate::parser::{ EXPLAIN_VERBOSE_NODE_NAME, }; use crate::promql::error::{ - CatalogSnafu, ColumnNotFoundSnafu, CombineTableColumnMismatchSnafu, DataFusionPlanningSnafu, - ExpectRangeSelectorSnafu, FunctionInvalidArgumentSnafu, InvalidDestinationLabelNameSnafu, - InvalidRegularExpressionSnafu, InvalidTimeRangeSnafu, MultiFieldsNotSupportedSnafu, - MultipleMetricMatchersSnafu, MultipleVectorSnafu, NoMetricMatcherSnafu, PromqlPlanNodeSnafu, - Result, SameLabelSetSnafu, TableNameNotFoundSnafu, TimeIndexNotFoundSnafu, - UnexpectedPlanExprSnafu, UnexpectedTokenSnafu, UnknownTableSnafu, UnsupportedExprSnafu, - UnsupportedMatcherOpSnafu, UnsupportedVectorMatchSnafu, ValueNotFoundSnafu, - ZeroRangeSelectorSnafu, + CatalogSnafu, ColumnNotFoundSnafu, DataFusionPlanningSnafu, ExpectRangeSelectorSnafu, + FunctionInvalidArgumentSnafu, InvalidDestinationLabelNameSnafu, InvalidRegularExpressionSnafu, + InvalidTimeRangeSnafu, MultiFieldsNotSupportedSnafu, MultipleMetricMatchersSnafu, + MultipleVectorSnafu, NoMetricMatcherSnafu, Result, SameLabelSetSnafu, TableNameNotFoundSnafu, + TimeIndexNotFoundSnafu, UnexpectedPlanExprSnafu, UnexpectedTokenSnafu, UnknownTableSnafu, + UnsupportedExprSnafu, UnsupportedMatcherOpSnafu, ValueNotFoundSnafu, ZeroRangeSelectorSnafu, }; use crate::query_engine::QueryEngineState; @@ -147,8 +146,6 @@ const FIELD_COLUMN_MATCHER: &str = "__field__"; const SCHEMA_COLUMN_MATCHER: &str = "__schema__"; const DB_COLUMN_MATCHER: &str = "__database__"; -/// Prefix for generated binary island leaf aliases. -const BINARY_ISLAND_LEAF_ALIAS_PREFIX: &str = "__prom_v"; const OR_FLOAT_FIELD_PREFIX: &str = "__promql_or_float_"; const OR_HISTOGRAM_FIELD_PREFIX: &str = "__promql_or_histogram_"; const TIMESTAMP_VALUE_PREFIX: &str = "__promql_timestamp_value_"; @@ -198,103 +195,6 @@ struct PromPlannerContext { range: Option, } -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -struct VectorLeafKey { - metric_name: String, - matchers: Vec<(String, String, String)>, - or_matchers: Vec>, - offset_ms: i128, - at: String, -} - -#[derive(Debug, Clone)] -struct IslandLeaf { - selector: VectorSelector, - display_table: String, -} - -#[derive(Debug, Clone)] -enum IslandExpr { - VectorLeaf(usize), - Scalar(DfExpr), - Unary { - input: Box, - }, - Binary { - op: TokenType, - lhs: Box, - rhs: Box, - }, -} - -impl IslandExpr { - fn try_new(expr: &PromExpr, env: &mut IslandCollectEnv) -> Option { - if let Some(expr) = PromPlanner::try_build_literal_expr(expr) { - return Some(Self::Scalar(expr)); - } - - match expr { - PromExpr::Paren(ParenExpr { expr }) => Self::try_new(expr, env), - PromExpr::VectorSelector(selector) => { - let leaf = env.intern_leaf(selector)?; - Some(Self::VectorLeaf(leaf)) - } - PromExpr::Unary(UnaryExpr { expr }) => { - let input = Self::try_new(expr, env)?; - Some(Self::Unary { - input: Box::new(input), - }) - } - PromExpr::Binary(PromBinaryExpr { - lhs, - rhs, - op, - modifier, - }) if matches!( - op.id(), - token::T_ADD - | token::T_SUB - | token::T_MUL - | token::T_DIV - | token::T_MOD - | token::T_POW - | token::T_ATAN2 - ) && modifier.as_ref().is_none_or(|modifier| { - !modifier.return_bool - && modifier.matching.is_none() - && matches!(modifier.card, VectorMatchCardinality::OneToOne) - && modifier.fill_values.lhs.is_none() - && modifier.fill_values.rhs.is_none() - }) => - { - let lhs = Self::try_new(lhs, env)?; - let rhs = Self::try_new(rhs, env)?; - Some(Self::Binary { - op: *op, - lhs: Box::new(lhs), - rhs: Box::new(rhs), - }) - } - _ => None, - } - } -} - -#[derive(Debug, Default)] -struct IslandCollectEnv { - leaf_by_key: HashMap, - leaves: Vec, - vector_occurrences: usize, -} - -#[derive(Debug)] -struct PlannedIslandLeaf { - plan: LogicalPlan, - ctx: PromPlannerContext, - alias: TableReference, - display_table: String, -} - /// Result labels a vector-vector binary operation derives from its matching modifier, projected /// from the operand each label belongs to. #[derive(Debug)] @@ -324,81 +224,6 @@ impl BinaryResultLabels { } } -#[derive(Debug)] -struct IslandFieldExprs { - exprs: Vec, - names: Vec, - scalar: bool, -} - -impl VectorLeafKey { - fn from_selector(selector: &VectorSelector) -> Option { - let mut metric_name = selector.name.clone(); - let mut matchers = Vec::with_capacity(selector.matchers.matchers.len()); - let matcher_key = |matcher: &Matcher| { - ( - matcher.name.clone(), - matcher.op.to_string(), - matcher.value.clone(), - ) - }; - - for matcher in &selector.matchers.matchers { - if matcher.name == METRIC_NAME { - if matcher.op != MatchOp::Equal || metric_name.is_some() { - return None; - } - metric_name = Some(matcher.value.clone()); - } else { - matchers.push(matcher_key(matcher)); - } - } - matchers.sort(); - - let mut or_matchers = selector - .matchers - .or_matchers - .iter() - .map(|group| { - let mut group = group.iter().map(matcher_key).collect::>(); - group.sort(); - group - }) - .collect::>(); - or_matchers.sort(); - - Some(Self { - metric_name: metric_name?, - matchers, - or_matchers, - offset_ms: match &selector.offset { - Some(Offset::Pos(duration)) => duration.as_millis() as i128, - Some(Offset::Neg(duration)) => -(duration.as_millis() as i128), - None => 0, - }, - at: format!("{:?}", selector.at), - }) - } -} - -impl IslandCollectEnv { - fn intern_leaf(&mut self, selector: &VectorSelector) -> Option { - self.vector_occurrences += 1; - let key = VectorLeafKey::from_selector(selector)?; - if let Some(id) = self.leaf_by_key.get(&key) { - return Some(*id); - } - - let id = self.leaves.len(); - self.leaves.push(IslandLeaf { - selector: selector.clone(), - display_table: key.metric_name.clone(), - }); - self.leaf_by_key.insert(key, id); - Some(id) - } -} - impl PromPlannerContext { fn from_eval_stmt(stmt: &EvalStmt) -> Self { Self { @@ -1039,330 +864,6 @@ impl PromPlanner { }) } - async fn try_plan_binary_island( - &mut self, - binary_expr: &PromBinaryExpr, - ) -> Result> { - let original_ctx = self.ctx.clone(); - let mut collect_env = IslandCollectEnv::default(); - let Some(island_expr) = - IslandExpr::try_new(&PromExpr::Binary(binary_expr.clone()), &mut collect_env) - else { - return Ok(None); - }; - - if collect_env.leaves.is_empty() - || collect_env.vector_occurrences <= collect_env.leaves.len() - { - return Ok(None); - } - - let mut planned_leaves = Vec::with_capacity(collect_env.leaves.len()); - for (idx, leaf) in collect_env.leaves.iter().enumerate() { - let plan = self - .prom_vector_selector_to_plan(&leaf.selector, false) - .await?; - let ctx = self.ctx.clone(); - let alias = TableReference::bare(format!("{BINARY_ISLAND_LEAF_ALIAS_PREFIX}{idx}")); - let plan = LogicalPlanBuilder::from(plan) - .alias(alias.clone()) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - planned_leaves.push(PlannedIslandLeaf { - plan, - ctx, - alias, - display_table: leaf.display_table.clone(), - }); - } - - if planned_leaves.iter().any(|leaf| { - Self::field_columns_contain_native_histogram( - leaf.plan.schema(), - &leaf.ctx.field_columns, - ) - }) { - self.ctx = original_ctx; - return Ok(None); - } - - if !Self::binary_island_join_contexts_supported(&planned_leaves) { - self.ctx = original_ctx; - return Ok(None); - } - - let mut input = planned_leaves[0].plan.clone(); - for right_idx in 1..planned_leaves.len() { - input = self.join_binary_island_leaf( - input, - &planned_leaves[0], - &planned_leaves[right_idx], - )?; - } - - let field_exprs = - Self::build_binary_island_field_exprs(&island_expr, &planned_leaves, input.schema())?; - if field_exprs.scalar || field_exprs.exprs.is_empty() { - self.ctx = original_ctx; - return Ok(None); - } - - let plan = self.project_binary_island( - input, - &planned_leaves[0].alias, - &planned_leaves[0].ctx, - field_exprs, - )?; - Ok(Some(plan)) - } - - fn binary_island_join_contexts_supported(leaves: &[PlannedIslandLeaf]) -> bool { - if leaves - .iter() - .any(|leaf| leaf.ctx.time_index_column.is_none()) - { - return false; - } - - if leaves.len() <= 1 { - return true; - } - - let first_tags = leaves[0].ctx.tag_columns.iter().collect::>(); - - leaves.iter().skip(1).all(|leaf| { - (Self::plan_has_tsid_column(&leaves[0].plan) && Self::plan_has_tsid_column(&leaf.plan)) - || leaf.ctx.tag_columns.iter().collect::>() == first_tags - }) - } - - fn join_binary_island_leaf( - &self, - left: LogicalPlan, - first_leaf: &PlannedIslandLeaf, - right_leaf: &PlannedIslandLeaf, - ) -> Result { - let only_join_time_index = (first_leaf.ctx.tag_columns.is_empty() - || right_leaf.ctx.tag_columns.is_empty()) - && !first_leaf - .ctx - .tag_columns - .iter() - .chain(&right_leaf.ctx.tag_columns) - .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL); - let (mut left_keys, mut right_keys, force_empty_join) = self.binary_join_key_columns( - left.schema(), - right_leaf.plan.schema(), - &first_leaf.ctx, - &right_leaf.ctx, - only_join_time_index, - &None, - )?; - - if let (Some(left_time_index_column), Some(right_time_index_column)) = ( - first_leaf.ctx.time_index_column.clone(), - right_leaf.ctx.time_index_column.clone(), - ) { - left_keys.insert(left_time_index_column); - right_keys.insert(right_time_index_column); - } - - LogicalPlanBuilder::from(left) - .join_detailed( - right_leaf.plan.clone(), - JoinType::Inner, - ( - left_keys - .into_iter() - .map(|name| Column::new(Some(first_leaf.alias.clone()), name)) - .collect::>(), - right_keys - .into_iter() - .map(|name| Column::new(Some(right_leaf.alias.clone()), name)) - .collect::>(), - ), - force_empty_join.then_some(lit(false)), - NullEquality::NullEqualsNull, - ) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu) - } - - fn build_binary_island_field_exprs( - expr: &IslandExpr, - leaves: &[PlannedIslandLeaf], - schema: &DFSchemaRef, - ) -> Result { - match expr { - IslandExpr::VectorLeaf(id) => { - let leaf = &leaves[*id]; - let exprs = leaf - .ctx - .field_columns - .iter() - .map(|field| { - schema - .qualified_field_with_name(Some(&leaf.alias), field) - .context(DataFusionPlanningSnafu) - .map(|field| DfExpr::Column(field.into())) - }) - .collect::>>()?; - let names = leaf - .ctx - .field_columns - .iter() - .map(|field| format!("{}.{}", leaf.display_table, field)) - .collect(); - Ok(IslandFieldExprs { - exprs, - names, - scalar: false, - }) - } - IslandExpr::Scalar(expr) => Ok(IslandFieldExprs { - exprs: vec![expr.clone()], - names: vec![expr.schema_name().to_string()], - scalar: true, - }), - IslandExpr::Unary { input } => { - let input = Self::build_binary_island_field_exprs(input, leaves, schema)?; - let mut exprs = Vec::with_capacity(input.exprs.len()); - let mut names = Vec::with_capacity(input.names.len()); - for (expr, name) in input.exprs.into_iter().zip(input.names) { - exprs.push(DfExpr::Negative(Box::new(expr))); - names.push(format!("-{name}")); - } - Ok(IslandFieldExprs { - exprs, - names, - scalar: input.scalar, - }) - } - IslandExpr::Binary { op, lhs, rhs } => { - let same_leaf = match (&**lhs, &**rhs) { - (IslandExpr::VectorLeaf(left), IslandExpr::VectorLeaf(right)) - if left == right => - { - Some(*left) - } - _ => None, - }; - let lhs = Self::build_binary_island_field_exprs(lhs, leaves, schema)?; - let rhs = Self::build_binary_island_field_exprs(rhs, leaves, schema)?; - let expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; - let scalar = lhs.scalar && rhs.scalar; - let op = op.to_string(); - - let (exprs, names) = match (lhs.scalar, rhs.scalar) { - (true, true) => { - let expr = expr_builder(lhs.exprs[0].clone(), rhs.exprs[0].clone())?; - let name = format!("{} {op} {}", lhs.names[0], rhs.names[0]); - (vec![expr], vec![name]) - } - (true, false) => { - let mut exprs = Vec::with_capacity(rhs.exprs.len()); - let mut names = Vec::with_capacity(rhs.names.len()); - for (rhs_expr, rhs_name) in rhs.exprs.into_iter().zip(rhs.names) { - exprs.push(expr_builder(lhs.exprs[0].clone(), rhs_expr)?); - names.push(format!("{} {op} {rhs_name}", lhs.names[0])); - } - (exprs, names) - } - (false, true) => { - let mut exprs = Vec::with_capacity(lhs.exprs.len()); - let mut names = Vec::with_capacity(lhs.names.len()); - for (lhs_expr, lhs_name) in lhs.exprs.into_iter().zip(lhs.names) { - exprs.push(expr_builder(lhs_expr, rhs.exprs[0].clone())?); - names.push(format!("{lhs_name} {op} {}", rhs.names[0])); - } - (exprs, names) - } - (false, false) => { - let mut exprs = Vec::new(); - let mut names = Vec::new(); - for (idx, ((lhs_expr, rhs_expr), (mut lhs_name, mut rhs_name))) in lhs - .exprs - .into_iter() - .zip(rhs.exprs) - .zip(lhs.names.into_iter().zip(rhs.names)) - .enumerate() - { - if let Some(leaf) = same_leaf { - let field = leaves[leaf] - .ctx - .field_columns - .get(idx) - .cloned() - .unwrap_or_else(|| lhs_name.clone()); - lhs_name = format!("lhs.{field}"); - rhs_name = format!("rhs.{field}"); - } - exprs.push(expr_builder(lhs_expr, rhs_expr)?); - names.push(format!("{lhs_name} {op} {rhs_name}")); - } - (exprs, names) - } - }; - - Ok(IslandFieldExprs { - exprs, - names, - scalar, - }) - } - } - } - - fn project_binary_island( - &mut self, - input: LogicalPlan, - base_alias: &TableReference, - base_ctx: &PromPlannerContext, - field_exprs: IslandFieldExprs, - ) -> Result { - self.ctx = base_ctx.clone(); - - let schema = input.schema(); - let non_field_exprs = base_ctx - .tag_columns - .iter() - .chain(base_ctx.time_index_column.iter()) - .map(|column| { - schema - .qualified_field_with_name(Some(base_alias), column) - .context(DataFusionPlanningSnafu) - .map(|field| DfExpr::Column(field.into())) - }); - let tsid_expr = Self::optional_tsid_projection(schema, Some(base_alias), base_ctx.use_tsid) - .into_iter() - .map(Ok); - - self.ctx.field_columns = field_exprs.names; - let field_exprs = field_exprs - .exprs - .into_iter() - .zip(self.ctx.field_columns.iter()) - .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::, name)))); - - let project_exprs = non_field_exprs - .chain(tsid_expr) - .chain(field_exprs) - .collect::>>()?; - - let plan = LogicalPlanBuilder::from(input) - .project(project_exprs) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - - self.ctx.table_name = None; - self.ctx.schema_name = None; - - Ok(plan) - } - async fn prom_binary_expr_to_plan( &mut self, query_engine_state: &QueryEngineState, @@ -5000,423 +4501,6 @@ impl PromPlanner { Ok(normalized_exprs) } - /// Create a classic, native, or mixed histogram helper plan. - async fn create_histogram_plan( - &mut self, - function_name: &str, - args: &PromFunctionArgs, - query_engine_state: &QueryEngineState, - ) -> Result { - let float_literal = |param: &PromExpr| -> Result { - let value = (|| { - let expr = Self::get_param_as_literal_expr( - Some(param), - None, - Some(ArrowDataType::Float64), - ) - .ok()?; - let simplifier = ExprSimplifier::new(SimplifyContext::default()); - let expr = simplifier.coerce(expr, &DFSchema::empty()).ok()?; - let DfExpr::Literal(value, _) = simplifier.simplify(expr).ok()? else { - return None; - }; - let ScalarValue::Float64(Some(value)) = - value.cast_to(&ArrowDataType::Float64).ok()? - else { - return None; - }; - Some(value) - })() - .with_context(|| FunctionInvalidArgumentSnafu { - fn_name: function_name.to_string(), - })?; - Ok(value) - }; - let (function, input) = match (function_name, args.args.as_slice()) { - (SPECIAL_HISTOGRAM_QUANTILE, [quantile, input]) => ( - HistogramFoldOperation::Quantile(float_literal(quantile)?.into()), - input.as_ref().clone(), - ), - (SPECIAL_HISTOGRAM_FRACTION, [lower, upper, input]) => ( - HistogramFoldOperation::Fraction { - lower: float_literal(lower)?.into(), - upper: float_literal(upper)?.into(), - }, - input.as_ref().clone(), - ), - _ => { - return FunctionInvalidArgumentSnafu { - fn_name: function_name.to_string(), - } - .fail(); - } - }; - - let input_plan = self.prom_expr_to_plan(&input, query_engine_state).await?; - // Histogram helpers fold buckets across `le`, so `__tsid` (which includes `le`) is not a - // stable series identifier anymore. HistogramFold must not treat it as a label column. - let input_plan = self.strip_tsid_column(input_plan)?; - self.ctx.use_tsid = false; - - if let Some((float_field, histogram_field)) = - Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns) - .map(|(float, histogram)| (float.to_string(), histogram.to_string())) - { - if self.ctx.has_le_tag() { - return self.create_mixed_histogram_plan( - function, - input_plan, - float_field, - histogram_field, - ); - } - self.ctx.field_columns = vec![histogram_field]; - } - if self.all_field_columns_are_native_histograms(input_plan.schema()) { - return self.create_native_histogram_plan(function, input_plan); - } - - if !self.ctx.has_le_tag() { - // Return empty result instead of error when 'le' column is not found - // This handles the case when histogram metrics don't exist - return Ok(LogicalPlan::EmptyRelation( - datafusion::logical_expr::EmptyRelation { - produce_one_row: false, - schema: input_plan.schema().clone(), - }, - )); - } - let time_index_column = - self.ctx - .time_index_column - .clone() - .with_context(|| TimeIndexNotFoundSnafu { - table: self.ctx.table_name.clone().unwrap_or_default(), - })?; - // FIXME(ruihang): support multi fields - let field_column = self - .ctx - .field_columns - .first() - .with_context(|| FunctionInvalidArgumentSnafu { - fn_name: function.function_name().to_string(), - })? - .clone(); - // remove le column from tag columns - self.ctx.tag_columns.retain(|col| col != LE_COLUMN_NAME); - - let fold = HistogramFold::new_with_operation( - LE_COLUMN_NAME.to_string(), - field_column, - time_index_column, - function, - None, - input_plan, - ) - .context(DataFusionPlanningSnafu)?; - Ok(LogicalPlan::Extension(Extension { - node: Arc::new(fold), - })) - } - - fn create_native_histogram_expr( - &self, - function: HistogramFoldOperation, - field_column: &str, - ) -> DfExpr { - let field = DfExpr::Column(Column::from_name(field_column)); - let (func, args) = match function { - HistogramFoldOperation::Quantile(quantile) => ( - Arc::new(NativeHistogramQuantile::scalar_udf_with_collector( - self.promql_annotations.clone(), - )), - vec![field, lit(f64::from(quantile))], - ), - HistogramFoldOperation::Fraction { lower, upper } => ( - Arc::new(NativeHistogramFraction::scalar_udf_with_collector( - self.promql_annotations.clone(), - )), - vec![field, lit(f64::from(lower)), lit(f64::from(upper))], - ), - }; - DfExpr::ScalarFunction(ScalarFunction { func, args }) - } - - fn create_native_histogram_plan( - &mut self, - function: HistogramFoldOperation, - input_plan: LogicalPlan, - ) -> Result { - ensure!( - self.ctx.field_columns.len() == 1, - MultiFieldsNotSupportedSnafu { - operator: function.function_name() - }, - ); - - let field_column = self.ctx.field_columns[0].clone(); - let function_expr = self.create_native_histogram_expr(function, &field_column); - let display_name = function_expr.schema_name().to_string(); - self.ctx.field_columns = vec![display_name.clone()]; - - let project_exprs = std::iter::once(self.create_time_index_column_expr()?) - .chain(std::iter::once(function_expr.alias(display_name))) - .chain(self.create_tag_column_exprs()?) - .collect::>(); - - LogicalPlanBuilder::from(input_plan) - .project(project_exprs) - .context(DataFusionPlanningSnafu)? - .filter(self.create_empty_values_filter_expr(false)?) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu) - } - - fn create_mixed_histogram_plan( - &mut self, - function: HistogramFoldOperation, - input_plan: LogicalPlan, - float_field: String, - histogram_field: String, - ) -> Result { - let time_index_column = - self.ctx - .time_index_column - .clone() - .with_context(|| TimeIndexNotFoundSnafu { - table: self.ctx.table_name.clone().unwrap_or_default(), - })?; - let tag_columns = self.ctx.tag_columns.clone(); - let folded = HistogramFold::new_with_operation( - LE_COLUMN_NAME.to_string(), - float_field.clone(), - time_index_column.clone(), - function, - Some(histogram_field.clone()), - input_plan, - ) - .context(DataFusionPlanningSnafu)?; - let record_collision = DfExpr::ScalarFunction(ScalarFunction { - func: Arc::new(NativeHistogramDrop::warning_bool_false_udf( - "vector contains a mix of classic and native histograms".to_string(), - self.promql_annotations.clone(), - )), - args: vec![col(&float_field), col(&histogram_field)], - }); - let keep = when( - col(&float_field) - .is_not_null() - .and(col(&histogram_field).is_not_null()), - record_collision, - ) - .otherwise(lit(true)) - .context(DataFusionPlanningSnafu)?; - - let native_expr = self.create_native_histogram_expr(function, &histogram_field); - let output_field = native_expr.schema_name().to_string(); - let value = DfExpr::ScalarFunction(ScalarFunction { - func: coalesce(), - args: vec![col(&float_field), native_expr], - }); - self.ctx.field_columns = vec![output_field.clone()]; - LogicalPlanBuilder::from(LogicalPlan::Extension(Extension { - node: Arc::new(folded), - })) - .filter(keep) - .context(DataFusionPlanningSnafu)? - .project( - std::iter::once(col(&time_index_column)) - .chain(std::iter::once(value.alias(output_field))) - .chain(tag_columns.iter().map(col)), - ) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu) - } - - /// Create a [SPECIAL_VECTOR_FUNCTION] plan - async fn create_vector_plan(&mut self, args: &PromFunctionArgs) -> Result { - if args.args.len() != 1 { - return FunctionInvalidArgumentSnafu { - fn_name: SPECIAL_VECTOR_FUNCTION.to_string(), - } - .fail(); - } - let lit = Self::get_param_as_literal_expr(Some(args.args[0].as_ref()), None, None)?; - - // reuse `SPECIAL_TIME_FUNCTION` as name of time index column - self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string()); - self.ctx.reset_table_name_and_schema(); - self.ctx.tag_columns = vec![]; - self.ctx.aggregation_field_labels.clear(); - self.ctx.field_columns = vec![greptime_value().to_string()]; - Ok(LogicalPlan::Extension(Extension { - node: Arc::new( - EmptyMetric::new( - self.ctx.start, - self.ctx.end, - self.ctx.interval, - SPECIAL_TIME_FUNCTION.to_string(), - greptime_value().to_string(), - Some(lit), - ) - .context(DataFusionPlanningSnafu)?, - ), - })) - } - - /// Create a [SCALAR_FUNCTION] plan - async fn create_scalar_plan( - &mut self, - args: &PromFunctionArgs, - query_engine_state: &QueryEngineState, - ) -> Result { - ensure!( - args.len() == 1, - FunctionInvalidArgumentSnafu { - fn_name: SCALAR_FUNCTION - } - ); - let input = self - .prom_expr_to_plan(&args.args[0], query_engine_state) - .await?; - let input_schema = input.schema().clone(); - let alternative_samples = - Self::field_columns_are_alternative_samples(&input_schema, &self.ctx.field_columns); - let histogram_fields = self - .ctx - .field_columns - .iter() - .filter(|field| Self::field_column_is_native_histogram(&input_schema, field)) - .count(); - ensure!( - self.ctx.field_columns.len() == 1 || alternative_samples, - MultiFieldsNotSupportedSnafu { - operator: SCALAR_FUNCTION - }, - ); - let scalar_field = self - .ctx - .field_columns - .iter() - .find(|field| !Self::field_column_is_native_histogram(&input_schema, field)) - .or_else(|| self.ctx.field_columns.first()) - .cloned() - .with_context(|| FunctionInvalidArgumentSnafu { - fn_name: SCALAR_FUNCTION, - })?; - let input = if histogram_fields == self.ctx.field_columns.len() { - // scalar() ignores histogram samples. An empty input makes ScalarCalculate emit NaN - // for every evaluation timestamp without attempting a Struct-to-Float64 cast. - LogicalPlanBuilder::from(input) - .filter(lit(false)) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)? - } else if histogram_fields > 0 { - // A mixed vector contributes only its float samples to scalar(). - LogicalPlanBuilder::from(input) - .filter(DfExpr::Column(Column::from_name(&scalar_field)).is_not_null()) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)? - } else { - input - }; - let scalar_plan = LogicalPlan::Extension(Extension { - node: Arc::new( - ScalarCalculate::new( - self.ctx.start, - self.ctx.end, - self.ctx.interval, - input, - self.ctx.time_index_column.as_ref().unwrap(), - &self.ctx.tag_columns, - &scalar_field, - self.ctx.table_name.as_deref(), - ) - .context(PromqlPlanNodeSnafu)?, - ), - }); - // scalar plan have no tag columns - self.ctx.tag_columns.clear(); - self.ctx.aggregation_field_labels.clear(); - self.ctx.field_columns.clear(); - self.ctx - .field_columns - .push(scalar_plan.schema().field(1).name().clone()); - Ok(scalar_plan) - } - - /// Create a [SPECIAL_ABSENT_FUNCTION] plan - async fn create_absent_plan( - &mut self, - args: &PromFunctionArgs, - query_engine_state: &QueryEngineState, - ) -> Result { - if args.args.len() != 1 { - return FunctionInvalidArgumentSnafu { - fn_name: SPECIAL_ABSENT_FUNCTION.to_string(), - } - .fail(); - } - let input = self - .prom_expr_to_plan(&args.args[0], query_engine_state) - .await?; - - let time_index_expr = self.create_time_index_column_expr()?; - let first_field_expr = - self.create_field_column_exprs()? - .pop() - .with_context(|| ValueNotFoundSnafu { - table: self.ctx.table_name.clone().unwrap_or_default(), - })?; - let first_value_expr = first_value(first_field_expr, vec![]); - - let ordered_aggregated_input = LogicalPlanBuilder::from(input) - .aggregate( - vec![time_index_expr.clone()], - vec![first_value_expr.clone()], - ) - .context(DataFusionPlanningSnafu)? - .sort(vec![time_index_expr.sort(true, false)]) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - - let fake_labels = self - .ctx - .selector_matcher - .iter() - .filter_map(|matcher| match matcher.op { - MatchOp::Equal => Some((matcher.name.clone(), matcher.value.clone())), - _ => None, - }) - .collect::>(); - - // Create the absent plan - let absent_plan = LogicalPlan::Extension(Extension { - node: Arc::new( - Absent::try_new( - self.ctx.start, - self.ctx.end, - self.ctx.interval, - self.ctx.time_index_column.as_ref().unwrap().clone(), - self.ctx.field_columns[0].clone(), - fake_labels, - ordered_aggregated_input, - ) - .context(DataFusionPlanningSnafu)?, - ), - }); - - // The absent series carries the equality matchers as labels, not the input's - // tags or value fields, so the input's field grouping labels no longer apply. - self.ctx.aggregation_field_labels.clear(); - Ok(absent_plan) - } - /// Try to build a DataFusion Literal Expression from PromQL Expr, return /// `None` if the input is not a literal expression. fn try_build_literal_expr(expr: &PromExpr) -> Option { @@ -6618,203 +5702,6 @@ impl PromPlanner { matches!(plan, LogicalPlan::EmptyRelation(relation) if !relation.produce_one_row) } - /// Build a set operator (AND/OR/UNLESS) - fn set_op_on_non_field_columns( - &mut self, - mut left: LogicalPlan, - mut right: LogicalPlan, - left_context: PromPlannerContext, - right_context: PromPlannerContext, - op: TokenType, - modifier: &Option, - ) -> Result { - let left_tag_col_set = left_context - .tag_columns - .iter() - .cloned() - .collect::>(); - let right_tag_col_set = right_context - .tag_columns - .iter() - .cloned() - .collect::>(); - - if matches!(op.id(), token::T_LOR) { - return self.or_operator( - left, - right, - left_tag_col_set, - right_tag_col_set, - left_context, - right_context, - modifier, - ); - } - - if let Some(modifier) = modifier { - ensure!( - matches!( - modifier.card, - VectorMatchCardinality::OneToOne | VectorMatchCardinality::ManyToMany - ), - UnsupportedVectorMatchSnafu { - name: modifier.card.clone(), - }, - ); - } - - let output_context = left_context.clone(); - let visible_left_schema = left.schema().clone(); - let mut left_context = left_context; - let mut right_context = right_context; - let added_marker_to_left = if Self::only_temporality_match_label_mismatches( - &left_context, - &right_context, - modifier, - ) { - let aligned = Self::align_temporality_match_column( - left, - right, - &mut left_context, - &mut right_context, - )?; - left = aligned.0; - right = aligned.1; - aligned.2 - } else { - false - }; - - let mut left_tag_col_set = left_context - .tag_columns - .iter() - .cloned() - .collect::>(); - let mut right_tag_col_set = right_context - .tag_columns - .iter() - .cloned() - .collect::>(); - if let Some(matching) = modifier - .as_ref() - .and_then(|modifier| modifier.matching.as_ref()) - { - match matching { - LabelModifier::Include(on) => { - let mask = on.labels.iter().cloned().collect::>(); - left_tag_col_set = left_tag_col_set.intersection(&mask).cloned().collect(); - right_tag_col_set = right_tag_col_set.intersection(&mask).cloned().collect(); - } - LabelModifier::Exclude(ignoring) => { - for label in &ignoring.labels { - let _ = left_tag_col_set.remove(label); - let _ = right_tag_col_set.remove(label); - } - } - } - } - ensure!( - left_tag_col_set == right_tag_col_set, - CombineTableColumnMismatchSnafu { - left: left_tag_col_set.iter().cloned().collect::>(), - right: right_tag_col_set.iter().cloned().collect::>(), - } - ); - - let left_time_index = left_context.time_index_column.clone().unwrap(); - let right_time_index = right_context.time_index_column.clone().unwrap(); - - // alias right time index column if necessary - if left_context.time_index_column != right_context.time_index_column { - let right_project_exprs = right - .schema() - .fields() - .iter() - .map(|field| { - if field.name() == &right_time_index { - DfExpr::Column(Column::from_name(&right_time_index)).alias(&left_time_index) - } else { - DfExpr::Column(Column::from_name(field.name())) - } - }) - .collect::>(); - - right = LogicalPlanBuilder::from(right) - .project(right_project_exprs) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - } - - let join_keys = left_tag_col_set - .into_iter() - .chain([left_time_index]) - .collect::>(); - - ensure!( - left_context.field_columns.len() == 1 - || Self::field_columns_are_alternative_samples( - left.schema(), - &left_context.field_columns, - ), - MultiFieldsNotSupportedSnafu { - operator: "AND/UNLESS operator" - } - ); - // Generate join plan. - // All set operations in PromQL are "distinct" - let result = match op.id() { - token::T_LAND => LogicalPlanBuilder::from(left) - .distinct() - .context(DataFusionPlanningSnafu)? - .join_detailed( - right, - JoinType::LeftSemi, - (join_keys.clone(), join_keys), - None, - NullEquality::NullEqualsNull, - ) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu), - token::T_LUNLESS => LogicalPlanBuilder::from(left) - .distinct() - .context(DataFusionPlanningSnafu)? - .join_detailed( - right, - JoinType::LeftAnti, - (join_keys.clone(), join_keys), - None, - NullEquality::NullEqualsNull, - ) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu), - token::T_LOR => { - // OR is handled at the beginning of this function, as it cannot - // be expressed using JOIN like AND and UNLESS. - unreachable!() - } - _ => UnexpectedTokenSnafu { token: op }.fail(), - }?; - let result = if added_marker_to_left { - LogicalPlanBuilder::from(result) - .project(visible_left_schema.iter().map(|(qualifier, field)| { - DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) - })) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)? - } else { - result - }; - - // AND/UNLESS preserve the complete left operand's visible columns and values; encoded - // markers are decoded. - self.ctx = output_context; - Ok(result) - } - fn string_value_data_type(data_type: &ArrowDataType) -> Option<&ArrowDataType> { match data_type { data_type if data_type.is_string() => Some(data_type), @@ -6870,601 +5757,6 @@ impl PromPlanner { } } - // TODO(ruihang): change function name - #[allow(clippy::too_many_arguments)] - fn or_operator( - &mut self, - left: LogicalPlan, - right: LogicalPlan, - left_tag_cols_set: HashSet, - right_tag_cols_set: HashSet, - left_context: PromPlannerContext, - right_context: PromPlannerContext, - modifier: &Option, - ) -> Result { - let left_is_empty = Self::is_zero_row_empty_relation(&left); - let right_is_empty = Self::is_zero_row_empty_relation(&right); - match (left_is_empty, right_is_empty) { - (true, false) => { - self.ctx = right_context; - return Ok(right); - } - (false, true) => { - self.ctx = left_context; - return Ok(left); - } - (true, true) => { - self.ctx = left_context; - return Ok(left); - } - (false, false) => {} - } - - ensure!( - !left.schema().fields().is_empty() && !right.schema().fields().is_empty(), - UnexpectedPlanExprSnafu { - desc: "OR operator input has zero columns", - } - ); - let left_has_alternative_samples = - Self::field_columns_are_alternative_samples(left.schema(), &left_context.field_columns); - let right_has_alternative_samples = Self::field_columns_are_alternative_samples( - right.schema(), - &right_context.field_columns, - ); - ensure!( - left_context.field_columns.len() == 1 || left_has_alternative_samples, - MultiFieldsNotSupportedSnafu { - operator: "OR operator" - } - ); - ensure!( - right_context.field_columns.len() == 1 || right_has_alternative_samples, - MultiFieldsNotSupportedSnafu { - operator: "OR operator" - } - ); - - // prepare hash sets - let all_tags = left_tag_cols_set - .union(&right_tag_cols_set) - .cloned() - .collect::>(); - let left_qualifier = left.schema().qualified_field(0).0.cloned(); - let right_qualifier = right.schema().qualified_field(0).0.cloned(); - let left_qualifier_string = left_qualifier - .as_ref() - .map(|l| l.to_string()) - .unwrap_or_default(); - let right_qualifier_string = right_qualifier - .as_ref() - .map(|r| r.to_string()) - .unwrap_or_default(); - let left_time_index_column = - left_context - .time_index_column - .clone() - .with_context(|| TimeIndexNotFoundSnafu { - table: left_qualifier_string.clone(), - })?; - let right_time_index_column = - right_context - .time_index_column - .clone() - .with_context(|| TimeIndexNotFoundSnafu { - table: right_qualifier_string.clone(), - })?; - let native_histogram_type = Self::native_histogram_arrow_type(); - let is_numeric = |data_type: &ArrowDataType| { - matches!( - data_type, - ArrowDataType::Int8 - | ArrowDataType::Int16 - | ArrowDataType::Int32 - | ArrowDataType::Int64 - | ArrowDataType::UInt8 - | ArrowDataType::UInt16 - | ArrowDataType::UInt32 - | ArrowDataType::UInt64 - | ArrowDataType::Float32 - | ArrowDataType::Float64 - ) - }; - let left_fields = left_context - .field_columns - .iter() - .map(|name| { - left.schema() - .iter() - .find(|(_, field)| field.name() == name) - .map(|(qualifier, field)| { - (name.clone(), qualifier.cloned(), field.data_type().clone()) - }) - .with_context(|| ColumnNotFoundSnafu { col: name.clone() }) - }) - .collect::>>()?; - let right_fields = right_context - .field_columns - .iter() - .map(|name| { - right - .schema() - .iter() - .find(|(_, field)| field.name() == name) - .map(|(qualifier, field)| { - (name.clone(), qualifier.cloned(), field.data_type().clone()) - }) - .with_context(|| ColumnNotFoundSnafu { col: name.clone() }) - }) - .collect::>>()?; - let left_field = &left_fields[0]; - let right_field = &right_fields[0]; - let left_field_col = &left_field.0; - let right_field_col = &right_field.0; - let fields_are_samples = |fields: &[(String, Option, ArrowDataType)]| { - fields.iter().all(|(_, _, data_type)| { - is_numeric(data_type) || data_type == &native_histogram_type - }) - }; - let mixed_sample_types = if left_has_alternative_samples || right_has_alternative_samples { - if !fields_are_samples(&left_fields) || !fields_are_samples(&right_fields) { - return UnexpectedPlanExprSnafu { - desc: format!( - "OR value fields have incompatible types: {:?} and {:?}", - left_fields - .iter() - .map(|(_, _, data_type)| data_type) - .collect::>(), - right_fields - .iter() - .map(|(_, _, data_type)| data_type) - .collect::>() - ), - } - .fail(); - } - true - } else { - (left_field.2 == native_histogram_type && is_numeric(&right_field.2)) - || (right_field.2 == native_histogram_type && is_numeric(&left_field.2)) - }; - let target_field_type = if mixed_sample_types { - // Mixed vectors use the existing response representation: one nullable float column - // and one nullable native-histogram column. - ArrowDataType::Float64 - } else if left_field.2 == right_field.2 { - left_field.2.clone() - } else if is_numeric(&left_field.2) && is_numeric(&right_field.2) { - ArrowDataType::Float64 - } else { - return UnexpectedPlanExprSnafu { - desc: format!( - "OR value fields have incompatible types: {:?} and {:?}", - left_field.2, right_field.2 - ), - } - .fail(); - }; - let (mixed_float_field_col, mixed_histogram_field_col) = if mixed_sample_types { - let mut reserved_names = left - .schema() - .fields() - .iter() - .chain(right.schema().fields().iter()) - .map(|field| field.name().clone()) - .collect::>(); - for (name, _, _) in left_fields.iter().chain(&right_fields) { - reserved_names.remove(name); - } - reserved_names.extend(all_tags.iter().cloned()); - let unique_name = |prefix: &str, reserved_names: &mut HashSet| { - let mut index = 0; - loop { - let name = format!("{prefix}{index}"); - index += 1; - if reserved_names.insert(name.clone()) { - break name; - } - } - }; - let float_field = unique_name(OR_FLOAT_FIELD_PREFIX, &mut reserved_names); - let histogram_field = unique_name(OR_HISTOGRAM_FIELD_PREFIX, &mut reserved_names); - (float_field, histogram_field) - } else { - (left_field_col.clone(), String::new()) - }; - let left_tag_types = left_tag_cols_set - .iter() - .map(|label| { - left.schema() - .fields() - .iter() - .find(|field| field.name() == label) - .map(|field| (label.clone(), field.data_type().clone())) - .with_context(|| ColumnNotFoundSnafu { col: label.clone() }) - }) - .collect::>>()?; - let right_tag_types = right_tag_cols_set - .iter() - .map(|label| { - right - .schema() - .fields() - .iter() - .find(|field| field.name() == label) - .map(|field| (label.clone(), field.data_type().clone())) - .with_context(|| ColumnNotFoundSnafu { col: label.clone() }) - }) - .collect::>>()?; - let mut target_tag_types = HashMap::with_capacity(all_tags.len()); - for label in &all_tags { - let Some(data_type) = - Self::common_label_data_type(left_tag_types.get(label), right_tag_types.get(label)) - else { - return UnexpectedPlanExprSnafu { - desc: format!( - "OR label {label} has incompatible types: {:?} and {:?}", - left_tag_types.get(label), - right_tag_types.get(label) - ), - } - .fail(); - }; - target_tag_types.insert(label.clone(), data_type); - } - let left_has_tsid = left - .schema() - .fields() - .iter() - .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME); - let right_has_tsid = right - .schema() - .fields() - .iter() - .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME); - - // step 0: fill all columns in output schema - let mut all_columns_set = left - .schema() - .fields() - .iter() - .chain(right.schema().fields().iter()) - .map(|field| field.name().clone()) - .collect::>(); - // Keep `__tsid` only when both sides contain it, otherwise it may break schema alignment - // (e.g. `unknown_metric or some_metric`). - if !(left_has_tsid && right_has_tsid) { - all_columns_set.remove(DATA_SCHEMA_TSID_COLUMN_NAME); - } - // remove time index column - all_columns_set.remove(&left_time_index_column); - all_columns_set.remove(&right_time_index_column); - if mixed_sample_types { - for (name, _, _) in left_fields.iter().chain(&right_fields) { - all_columns_set.remove(name); - } - all_columns_set.extend(all_tags.iter().cloned()); - all_columns_set.insert(mixed_float_field_col.clone()); - all_columns_set.insert(mixed_histogram_field_col.clone()); - } else if left_field_col != right_field_col { - // remove field column in the right - all_columns_set.remove(right_field_col); - } - let mut all_columns = all_columns_set.into_iter().collect::>(); - // sort to ensure the generated schema is not volatile - all_columns.sort_unstable(); - // use left time index column name as the result time index column name - all_columns.insert(0, left_time_index_column.clone()); - let mut occupied_column_names = left - .schema() - .fields() - .iter() - .chain(right.schema().fields().iter()) - .map(|field| field.name().clone()) - .collect::>(); - - // step 1: align schema using project, fill non-exist columns with null - let aligned_label_expr = |col: &String, source_types: &HashMap| { - let target_type = &target_tag_types[col]; - if let Some(source_type) = source_types.get(col) { - let expr = DfExpr::Column(Column::new(None::, col)); - if source_type == target_type { - expr - } else { - DfExpr::Cast(Cast::new(Box::new(expr), target_type.clone())).alias(col.clone()) - } - } else { - DfExpr::Literal( - Self::string_scalar_value(target_type, None) - .expect("target label type is a string"), - None, - ) - .alias(col.clone()) - } - }; - let null_histogram = - ScalarValue::try_new_null(&native_histogram_type).context(DataFusionPlanningSnafu)?; - let mixed_value_expr = |fields: &[(String, Option, ArrowDataType)], - output_col: &String| { - if output_col == &mixed_float_field_col { - if let Some((name, qualifier, data_type)) = fields - .iter() - .find(|(_, _, data_type)| is_numeric(data_type)) - { - let expr = DfExpr::Column(Column::new(qualifier.clone(), name)); - if data_type == &ArrowDataType::Float64 { - expr.alias(output_col) - } else { - DfExpr::Cast(Cast::new(Box::new(expr), ArrowDataType::Float64)) - .alias(output_col) - } - } else { - DfExpr::Literal(ScalarValue::Float64(None), None).alias(output_col) - } - } else { - fields - .iter() - .find(|(_, _, data_type)| data_type == &native_histogram_type) - .map(|(name, qualifier, _)| { - DfExpr::Column(Column::new(qualifier.clone(), name)).alias(output_col) - }) - .unwrap_or_else(|| { - DfExpr::Literal(null_histogram.clone(), None).alias(output_col) - }) - } - }; - let left_proj_exprs = all_columns.iter().map(|col| { - if mixed_sample_types - && (col == &mixed_float_field_col || col == &mixed_histogram_field_col) - { - mixed_value_expr(&left_fields, col) - } else if !mixed_sample_types - && col == left_field_col - && left_field.2 != target_field_type - { - DfExpr::Cast(Cast::new( - Box::new(DfExpr::Column(Column::new( - left_field.1.clone(), - left_field_col, - ))), - target_field_type.clone(), - )) - .alias(left_field_col.clone()) - } else if target_tag_types.contains_key(col) { - aligned_label_expr(col, &left_tag_types) - } else { - DfExpr::Column(Column::new(None::, col)) - } - }); - let right_time_index_expr = DfExpr::Column(Column::new( - right_qualifier.clone(), - right_time_index_column, - )) - .alias(left_time_index_column.clone()); - // The field column in right side may not have qualifier (it may be removed by join operation), - // so we need to find it from the schema. - // `skip（1)` to skip the time index column - let right_proj_exprs_without_time_index = all_columns.iter().skip(1).map(|col| { - // expr - if mixed_sample_types - && (col == &mixed_float_field_col || col == &mixed_histogram_field_col) - { - mixed_value_expr(&right_fields, col) - } else if !mixed_sample_types && col == left_field_col { - let expr = DfExpr::Column(Column::new(right_field.1.clone(), right_field_col)); - if right_field.2 != target_field_type { - DfExpr::Cast(Cast::new(Box::new(expr), target_field_type.clone())) - .alias(left_field_col.clone()) - } else if left_field_col != right_field_col { - expr.alias(left_field_col.clone()) - } else { - expr - } - } else if target_tag_types.contains_key(col) { - aligned_label_expr(col, &right_tag_types) - } else { - DfExpr::Column(Column::new(None::, col)) - } - }); - let right_proj_exprs = [right_time_index_expr] - .into_iter() - .chain(right_proj_exprs_without_time_index); - - let left_projected = LogicalPlanBuilder::from(left) - .project(left_proj_exprs) - .context(DataFusionPlanningSnafu)? - .alias(left_qualifier_string.clone()) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - let right_projected = LogicalPlanBuilder::from(right) - .project(right_proj_exprs) - .context(DataFusionPlanningSnafu)? - .alias(right_qualifier_string.clone()) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - - // step 2: compute match columns - let mut match_columns = if let Some(modifier) = modifier - && let Some(matching) = &modifier.matching - { - match matching { - // keeps columns mentioned in `on` - LabelModifier::Include(on) => on.labels.clone(), - // removes columns memtioned in `ignoring` - LabelModifier::Exclude(ignoring) => { - let ignoring = ignoring.labels.iter().cloned().collect::>(); - all_tags.difference(&ignoring).cloned().collect() - } - } - } else { - all_tags.iter().cloned().collect() - }; - // sort to ensure the generated plan is not volatile - match_columns.sort_unstable(); - match_columns.dedup(); - occupied_column_names.extend( - left_projected - .schema() - .fields() - .iter() - .chain(right_projected.schema().fields().iter()) - .map(|field| field.name().clone()), - ); - - let visible_schema = left_projected.schema().clone(); - let visible_left_exprs = left_projected - .schema() - .iter() - .map(|(qualifier, field)| { - DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) - }) - .collect::>(); - let visible_right_exprs = right_projected - .schema() - .iter() - .map(|(qualifier, field)| { - DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) - }) - .collect::>(); - let mut left_match_exprs = Vec::with_capacity(match_columns.len()); - let mut right_match_exprs = Vec::with_capacity(match_columns.len()); - let mut next_internal_column = 0; - - for label in &match_columns { - let left_field = if left_tag_cols_set.contains(label) { - Some( - left_projected - .schema() - .iter() - .find(|(_, field)| field.name() == label) - .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone())) - .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?, - ) - } else { - None - }; - let right_field = if right_tag_cols_set.contains(label) { - Some( - right_projected - .schema() - .iter() - .find(|(_, field)| field.name() == label) - .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone())) - .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?, - ) - } else { - None - }; - let data_type = match (left_field.as_ref(), right_field.as_ref()) { - (Some((_, left_type)), Some((_, right_type))) if left_type == right_type => { - left_type.clone() - } - (Some((_, left_type)), Some((_, right_type))) => { - return UnexpectedPlanExprSnafu { - desc: format!( - "OR match label {label} has incompatible types: {left_type:?} and {right_type:?}" - ), - } - .fail(); - } - (Some((_, data_type)), None) | (None, Some((_, data_type))) => data_type.clone(), - (None, None) => ArrowDataType::Utf8, - }; - let Some(value_type) = Self::string_value_data_type(&data_type).cloned() else { - return UnexpectedPlanExprSnafu { - desc: format!("OR match label {label} must be a string"), - } - .fail(); - }; - let internal_name = loop { - let name = format!("__promql_or_match_{next_internal_column}"); - next_internal_column += 1; - if occupied_column_names.insert(name.clone()) { - break name; - } - }; - left_match_exprs.push(Self::normalized_match_key_expr( - label, - left_field, - &value_type, - &internal_name, - )); - right_match_exprs.push(Self::normalized_match_key_expr( - label, - right_field, - &value_type, - &internal_name, - )); - } - - let left_augmented = LogicalPlanBuilder::from(left_projected) - .project(visible_left_exprs.into_iter().chain(left_match_exprs)) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - let right_augmented = LogicalPlanBuilder::from(right_projected) - .project(visible_right_exprs.into_iter().chain(right_match_exprs)) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - - // step 3: build `UnionDistinctOn` with normalized internal match keys. - let visible_field_count = visible_schema.fields().len(); - let compare_key_indices = - (visible_field_count..visible_field_count + match_columns.len()).collect::>(); - let (time_qualifier, _) = visible_schema - .iter() - .find(|(_, field)| field.name() == &left_time_index_column) - .with_context(|| TimeIndexNotFoundSnafu { - table: left_qualifier_string.clone(), - })?; - let ts_col_idx = left_augmented - .schema() - .iter() - .position(|(qualifier, field)| { - qualifier == time_qualifier && field.name() == &left_time_index_column - }) - .with_context(|| TimeIndexNotFoundSnafu { - table: left_qualifier_string.clone(), - })?; - let union_distinct_on = UnionDistinctOn::try_new( - left_augmented, - right_augmented, - compare_key_indices, - ts_col_idx, - ) - .context(DataFusionPlanningSnafu)?; - let augmented_result = LogicalPlan::Extension(Extension { - node: Arc::new(union_distinct_on), - }); - let result = LogicalPlanBuilder::from(augmented_result) - .project(visible_schema.iter().map(|(qualifier, field)| { - DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) - })) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - - // step 4: update context - let output_field_col = left_field_col.clone(); - let mut output_context = left_context; - let mut visible_tags = all_tags.into_iter().collect::>(); - visible_tags.sort_unstable(); - output_context.time_index_column = Some(left_time_index_column); - output_context.tag_columns = visible_tags; - output_context.field_columns = if mixed_sample_types { - vec![mixed_float_field_col, mixed_histogram_field_col] - } else { - vec![output_field_col] - }; - output_context.use_tsid = left_has_tsid && right_has_tsid; - self.ctx = output_context; - - Ok(result) - } - /// Build a projection that project and perform operation expr for every value columns. /// Non-value columns (tag and timestamp) will be preserved in the projection. /// @@ -7705,7661 +5997,4 @@ enum ScalarFunc { } #[cfg(test)] -mod test { - use std::time::{Duration, UNIX_EPOCH}; - - use catalog::RegisterTableRequest; - use catalog::memory::{MemoryCatalogManager, new_memory_catalog_manager}; - use common_base::Plugins; - use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; - use common_query::native_histogram::{ - CUSTOM_BUCKETS_SCHEMA, CounterResetHint, NativeHistogram, build_histogram_array, - }; - use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value}; - use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; - use common_query::test_util::DummyDecoder; - use common_recordbatch::RecordBatch as GreptimeRecordBatch; - use datafusion::arrow::array::{ - Array, ArrayRef, Float64Array, Int64Array, StringArray, TimestampMillisecondArray, - }; - use datafusion::arrow::datatypes::{Field, Schema as ArrowSchema}; - use datafusion::arrow::record_batch::RecordBatch; - use datafusion::catalog::{CatalogProvider, MemoryCatalogProvider, MemorySchemaProvider}; - use datafusion::datasource::memory::MemorySourceConfig; - use datafusion::datasource::source::DataSourceExec; - use datafusion::datasource::{MemTable, provider_as_source}; - use datafusion::execution::context::SessionContext; - use datafusion::logical_expr::Extension; - use datatypes::prelude::ConcreteDataType; - use datatypes::schema::{ColumnSchema, Schema}; - use promql_parser::label::Labels; - use promql_parser::parser; - use session::context::QueryContext; - use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; - use table::Table; - use table::metadata::{FilterPushDownType, TableInfoBuilder, TableMetaBuilder}; - use table::test_util::{EmptyTable, MemTable as GreptimeMemTable}; - - use super::*; - use crate::QueryEngineContext; - use crate::options::QueryOptions; - use crate::parser::QueryLanguageParser; - use crate::query_engine::DefaultSerializer; - - mod delta; - - fn find_instant_manipulate(plan: &LogicalPlan) -> Option<&InstantManipulate> { - if let LogicalPlan::Extension(Extension { node }) = plan - && let Some(instant_manipulate) = node.as_any().downcast_ref::() - { - return Some(instant_manipulate); - } - - plan.inputs().into_iter().find_map(find_instant_manipulate) - } - - fn build_query_engine_state() -> QueryEngineState { - QueryEngineState::new( - new_memory_catalog_manager().unwrap(), - None, - None, - None, - None, - None, - false, - Plugins::default(), - QueryOptions::default(), - ) - } - - #[test] - fn common_label_type_preserves_only_shared_dictionary_encoding() { - let dictionary = ArrowDataType::Dictionary( - Box::new(ArrowDataType::UInt32), - Box::new(ArrowDataType::Utf8), - ); - let other_dictionary = ArrowDataType::Dictionary( - Box::new(ArrowDataType::Int32), - Box::new(ArrowDataType::Utf8), - ); - - assert_eq!( - Some(dictionary.clone()), - PromPlanner::common_label_data_type(Some(&dictionary), Some(&dictionary)) - ); - assert_eq!( - Some(ArrowDataType::Utf8), - PromPlanner::common_label_data_type(Some(&dictionary), Some(&ArrowDataType::Utf8)) - ); - assert_eq!( - Some(ArrowDataType::Utf8), - PromPlanner::common_label_data_type(Some(&dictionary), Some(&other_dictionary)) - ); - assert_eq!( - Some(ArrowDataType::Utf8), - PromPlanner::common_label_data_type(Some(&dictionary), None) - ); - } - - async fn build_optimized_promql_plan( - table_provider: DfTableSourceProvider, - eval_stmt: &EvalStmt, - ) -> LogicalPlan { - let state = build_query_engine_state(); - let raw_plan = PromPlanner::stmt_to_plan(table_provider, eval_stmt, &state) - .await - .unwrap(); - let context = QueryEngineContext::new(state.session_state(), QueryContext::arc()); - state - .optimize_by_extension_rules(raw_plan, &context) - .unwrap() - } - - async fn build_optimized_tsid_plan( - query: &str, - num_tag: usize, - num_field: usize, - end_secs: u64, - lookback_secs: u64, - ) -> String { - let eval_stmt = EvalStmt { - expr: parser::parse(query).unwrap(), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(end_secs)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(lookback_secs), - }; - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - num_tag, - num_field, - ) - .await; - - build_optimized_promql_plan(table_provider, &eval_stmt) - .await - .display_indent_schema() - .to_string() - } - - async fn assert_nested_count_rewrite_applies(query: &str, expected_outer_agg: &str) { - let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await; - - assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); - assert!(plan_str.contains("Projection: some_metric.timestamp, some_metric.tag_0")); - assert!(plan_str.contains("Distinct:")); - assert!(plan_str.contains(expected_outer_agg), "{plan_str}"); - assert!(!plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]")); - } - - async fn assert_nested_count_rewrite_missing(query: &str, num_tag: usize, lookback_secs: u64) { - let plan_str = build_optimized_tsid_plan(query, num_tag, 1, 100_000, lookback_secs).await; - assert!(!plan_str.contains("Distinct:"), "{plan_str}"); - } - - fn build_eval_stmt(expr: &str) -> EvalStmt { - EvalStmt { - expr: parser::parse(expr).unwrap(), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - } - } - - enum DirectOrValue { - Float64(f64), - Int64(i64), - NativeHistogram(NativeHistogram), - Utf8(&'static str), - } - - impl DirectOrValue { - fn data_type(&self) -> ArrowDataType { - match self { - Self::Float64(_) => ArrowDataType::Float64, - Self::Int64(_) => ArrowDataType::Int64, - Self::NativeHistogram(_) => native_histogram_value_type().as_arrow_type(), - Self::Utf8(_) => ArrowDataType::Utf8, - } - } - fn array(&self) -> Arc { - match self { - Self::Float64(v) => Arc::new(Float64Array::from(vec![*v])), - Self::Int64(v) => Arc::new(Int64Array::from(vec![*v])), - Self::NativeHistogram(v) => build_histogram_array(&[Some(v.clone())]), - Self::Utf8(v) => Arc::new(StringArray::from(vec![*v])), - } - } - } - - fn direct_or_histogram() -> NativeHistogram { - NativeHistogram { - schema: 0, - zero_threshold: 0.0, - sum: 1.0, - reset_hint: CounterResetHint::Unknown, - start_timestamp: None, - custom_values: vec![], - positive_spans: vec![], - negative_spans: vec![], - count: 1.0, - zero_count: 1.0, - positive_buckets: vec![], - negative_buckets: vec![], - } - } - - fn operator_metric_table( - name: &str, - table_id: u32, - tag: &str, - le: Option<&str>, - value: DirectOrValue, - ) -> table::TableRef { - let value_type = match &value { - DirectOrValue::Float64(_) => ConcreteDataType::float64_datatype(), - DirectOrValue::Int64(_) => ConcreteDataType::int64_datatype(), - DirectOrValue::NativeHistogram(_) => native_histogram_value_type().clone(), - DirectOrValue::Utf8(_) => ConcreteDataType::string_datatype(), - }; - let tag_count = 1 + usize::from(le.is_some()); - let mut columns = vec![ColumnSchema::new( - "tag".to_string(), - ConcreteDataType::string_datatype(), - false, - )]; - if le.is_some() { - columns.push(ColumnSchema::new( - LE_COLUMN_NAME.to_string(), - ConcreteDataType::string_datatype(), - false, - )); - } - columns.extend([ - ColumnSchema::new( - "ts".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ColumnSchema::new("v".to_string(), value_type, true), - ]); - let schema = Arc::new(Schema::new(columns)); - let mut arrays = vec![Arc::new(StringArray::from(vec![tag])) as Arc]; - if let Some(le) = le { - arrays.push(Arc::new(StringArray::from(vec![le]))); - } - arrays.extend([ - Arc::new(TimestampMillisecondArray::from(vec![1_000])) as Arc, - value.array(), - ]); - let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap(); - let backing = GreptimeMemTable::new_with_catalog( - name, - GreptimeRecordBatch::from_df_record_batch(schema.clone(), batch), - table_id, - DEFAULT_CATALOG_NAME.to_string(), - DEFAULT_SCHEMA_NAME.to_string(), - ); - let value_index = tag_count + 1; - let meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices((0..tag_count).collect()) - .value_indices(vec![value_index]) - .next_column_id((value_index + 1) as u32) - .build() - .unwrap(); - let info = Arc::new( - TableInfoBuilder::default() - .table_id(table_id) - .name(name) - .meta(meta) - .build() - .unwrap(), - ); - Arc::new(Table::new( - info, - FilterPushDownType::Unsupported, - backing.data_source(), - )) - } - - fn operator_table_provider() -> DfTableSourceProvider { - let catalog = MemoryCatalogManager::with_default_setup(); - let tables = [ - operator_metric_table("lf", 2_001, "a", None, DirectOrValue::Float64(2.0)), - operator_metric_table( - "lh", - 2_002, - "b", - None, - DirectOrValue::NativeHistogram(direct_or_histogram()), - ), - operator_metric_table("rf", 2_003, "b", None, DirectOrValue::Float64(3.0)), - operator_metric_table( - "rh", - 2_004, - "a", - None, - DirectOrValue::NativeHistogram(direct_or_histogram()), - ), - operator_metric_table("fallback", 2_005, "c", None, DirectOrValue::Float64(7.0)), - operator_metric_table( - "bad_classic", - 2_006, - "d", - Some("broken"), - DirectOrValue::Float64(1.0), - ), - operator_metric_table( - "bad_native", - 2_007, - "d", - None, - DirectOrValue::NativeHistogram(direct_or_histogram()), - ), - ]; - for table in tables { - let info = table.table_info(); - catalog - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: info.name.clone(), - table_id: info.ident.table_id, - table, - }) - .unwrap(); - } - DfTableSourceProvider::new( - catalog, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - fn operator_eval_stmt(expr: &str) -> EvalStmt { - let time = UNIX_EPOCH.checked_add(Duration::from_secs(1)).unwrap(); - EvalStmt { - expr: parser::parse(expr).unwrap(), - start: time, - end: time, - interval: Duration::from_secs(1), - lookback_delta: Duration::from_secs(5), - } - } - - struct DirectOrSource { - name: &'static str, - empty: bool, - timestamp: i64, - tags: Vec<(&'static str, Option<&'static str>)>, - value: DirectOrValue, - } - - fn source( - name: &'static str, - empty: bool, - timestamp: i64, - tags: Vec<(&'static str, Option<&'static str>)>, - value: DirectOrValue, - ) -> DirectOrSource { - DirectOrSource { - name, - empty, - timestamp, - tags, - value, - } - } - - fn tagged_source( - name: &'static str, - empty: bool, - tag: (&'static str, Option<&'static str>), - value: DirectOrValue, - ) -> DirectOrSource { - source(name, empty, 1, vec![("job", Some("job")), tag], value) - } - - fn job_source(name: &'static str, value: DirectOrValue) -> DirectOrSource { - source(name, true, 1, vec![("job", Some("job"))], value) - } - - fn table(source: &DirectOrSource) -> Arc { - let mut fields = vec![Field::new( - "ts", - ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), - false, - )]; - fields.extend( - source - .tags - .iter() - .map(|(name, _)| Field::new(*name, ArrowDataType::Utf8, true)), - ); - fields.push(Field::new("v", source.value.data_type(), true)); - let schema = Arc::new(ArrowSchema::new(fields)); - let partitions = if source.empty { - vec![vec![]] - } else { - let mut columns: Vec> = - vec![Arc::new(TimestampMillisecondArray::from(vec![ - source.timestamp, - ]))]; - columns.extend( - source - .tags - .iter() - .map(|(_, value)| Arc::new(StringArray::from(vec![*value])) as Arc), - ); - columns.push(source.value.array()); - vec![vec![RecordBatch::try_new(schema.clone(), columns).unwrap()]] - }; - Arc::new(MemTable::try_new(schema, partitions).unwrap()) - } - - fn scan(source: &DirectOrSource) -> LogicalPlan { - LogicalPlanBuilder::scan(source.name, provider_as_source(table(source)), None) - .unwrap() - .build() - .unwrap() - } - - fn direct_or_context(qualifier: &str, tags: &[&str], field: &str) -> PromPlannerContext { - PromPlannerContext { - table_name: Some(qualifier.to_string()), - time_index_column: Some("ts".to_string()), - field_columns: vec![field.to_string()], - tag_columns: tags.iter().map(|tag| (*tag).to_string()).collect(), - ..Default::default() - } - } - - fn or_modifier(expr: &str) -> Option { - let PromExpr::Binary(expr) = parser::parse(expr).unwrap() else { - unreachable!() - }; - expr.modifier - } - - async fn plan_direct_or( - left: LogicalPlan, - right: LogicalPlan, - left_context: PromPlannerContext, - right_context: PromPlannerContext, - modifier: &Option, - ) -> LogicalPlan { - let table_provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], - &[], - ) - .await; - let mut planner = PromPlanner { - table_provider, - ctx: PromPlannerContext::default(), - promql_annotations: None, - }; - planner - .or_operator( - left, - right, - left_context.tag_columns.iter().cloned().collect(), - right_context.tag_columns.iter().cloned().collect(), - left_context, - right_context, - modifier, - ) - .unwrap() - } - - async fn execute( - plan: LogicalPlan, - state: &QueryEngineState, - ) -> (LogicalPlan, Vec) { - let context = QueryEngineContext::new(state.session_state(), QueryContext::arc()); - let optimized = state.optimize_by_extension_rules(plan, &context).unwrap(); - let physical = state - .session_state() - .create_physical_plan(&optimized) - .await - .unwrap(); - let batches = - datafusion::physical_plan::collect(physical, state.session_state().task_ctx()) - .await - .unwrap(); - (optimized, batches) - } - - async fn run( - left: &DirectOrSource, - right: &DirectOrSource, - left_context: PromPlannerContext, - right_context: PromPlannerContext, - modifier: &Option, - ) -> (LogicalPlan, Vec) { - let plan = plan_direct_or( - scan(left), - scan(right), - left_context, - right_context, - modifier, - ) - .await; - execute(plan, &build_query_engine_state()).await - } - - async fn mixed_direct_or(histogram_on_left: bool) -> (PromPlanner, LogicalPlan) { - let sample = |histogram: bool| { - if histogram { - DirectOrValue::NativeHistogram(direct_or_histogram()) - } else { - DirectOrValue::Float64(1.25) - } - }; - let left = tagged_source( - "lhs", - false, - ( - "k", - Some(if histogram_on_left { - "histogram" - } else { - "float" - }), - ), - sample(histogram_on_left), - ); - let right = tagged_source( - "rhs", - false, - ( - "k", - Some(if histogram_on_left { - "float" - } else { - "histogram" - }), - ), - sample(!histogram_on_left), - ); - let table_provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], - &[], - ) - .await; - let mut planner = PromPlanner { - table_provider, - ctx: PromPlannerContext::default(), - promql_annotations: None, - }; - let left_context = direct_or_context("lhs", &["job", "k"], "v"); - let right_context = direct_or_context("rhs", &["job", "k"], "v"); - let plan = planner - .or_operator( - scan(&left), - scan(&right), - left_context.tag_columns.iter().cloned().collect(), - right_context.tag_columns.iter().cloned().collect(), - left_context, - right_context, - &or_modifier("lhs or on(k) rhs"), - ) - .unwrap(); - (planner, plan) - } - - async fn mixed_aggregate_input(histograms: Vec) -> (PromPlanner, LogicalPlan) { - let float_field = format!("{OR_FLOAT_FIELD_PREFIX}0"); - let histogram_field = format!("{OR_HISTOGRAM_FIELD_PREFIX}0"); - let row_count = histograms.len() + 1; - let schema = Arc::new(ArrowSchema::new(vec![ - Field::new( - "ts", - ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), - false, - ), - Field::new("k", ArrowDataType::Utf8, false), - Field::new(&float_field, ArrowDataType::Float64, true), - Field::new( - &histogram_field, - native_histogram_value_type().as_arrow_type(), - true, - ), - ])); - let mut histogram_values = Vec::with_capacity(row_count); - histogram_values.push(None); - histogram_values.extend(histograms.into_iter().map(Some)); - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(TimestampMillisecondArray::from(vec![1; row_count])), - Arc::new(StringArray::from_iter_values( - (0..row_count).map(|row| format!("kind_{row}")), - )), - Arc::new(Float64Array::from_iter( - (0..row_count).map(|row| (row == 0).then_some(1.25)), - )), - build_histogram_array(&histogram_values), - ], - ) - .unwrap(); - let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); - let plan = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None) - .unwrap() - .build() - .unwrap(); - let table_provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], - &[], - ) - .await; - let planner = PromPlanner { - table_provider, - ctx: PromPlannerContext { - table_name: Some("mixed".to_string()), - time_index_column: Some("ts".to_string()), - field_columns: vec![float_field, histogram_field], - tag_columns: vec!["k".to_string()], - ..Default::default() - }, - promql_annotations: None, - }; - (planner, plan) - } - - fn assert_no_internal_or_keys(schema: &DFSchema) { - assert!( - schema - .fields() - .iter() - .all(|field| !field.name().starts_with("__promql_or_match_")), - "{schema:?}" - ); - } - - fn values(batches: &[RecordBatch], column: &str) -> Vec { - batches - .iter() - .flat_map(|batch| { - batch - .column_by_name(column) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap() - .iter() - .flatten() - }) - .collect() - } - - fn numeric_values(batches: &[RecordBatch], column: &str) -> Vec { - batches - .iter() - .flat_map(|batch| { - let values = datafusion::arrow::compute::cast( - batch.column_by_name(column).unwrap(), - &ArrowDataType::Float64, - ) - .unwrap(); - values - .as_any() - .downcast_ref::() - .unwrap() - .iter() - .flatten() - .collect::>() - }) - .collect() - } - - fn histograms(batches: &[RecordBatch], column: &str) -> Vec { - batches - .iter() - .flat_map(|batch| { - let values = batch - .column_by_name(column) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - (0..values.len()).filter_map(|row| { - common_query::native_histogram::read_histogram(values, row).unwrap() - }) - }) - .collect() - } - - fn rows(batches: &[RecordBatch]) -> Vec<(f64, Option)> { - let mut rows = batches - .iter() - .flat_map(|batch| { - let values = batch - .column_by_name("v") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - let labels = batch - .column_by_name("k") - .map(|column| column.as_any().downcast_ref::().unwrap()); - (0..batch.num_rows()).map(move |i| { - ( - values.value(i), - labels.and_then(|labels| { - (!labels.is_null(i)).then(|| labels.value(i).to_string()) - }), - ) - }) - }) - .collect::>(); - rows.sort_by(|left, right| left.0.total_cmp(&right.0)); - rows - } - - fn matrix_source( - name: &'static str, - k: Option>, - timestamp: i64, - value: f64, - ) -> DirectOrSource { - let mut tags = vec![("job", Some("job"))]; - if let Some(k) = k { - tags.push(("k", k)); - } - source(name, false, timestamp, tags, DirectOrValue::Float64(value)) - } - - fn matrix_context(name: &str, k: Option>) -> PromPlannerContext { - direct_or_context( - name, - if k.is_some() { &["job", "k"] } else { &["job"] }, - "v", - ) - } - - async fn build_missing_le_or_normal_metric_table_provider() -> DfTableSourceProvider { - build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "non_existent_histogram_bucket".to_string(), - ), - (DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string()), - ], - &["pod", "instance"], - ) - .await - } - - fn assert_normal_metric_schema(plan: &LogicalPlan) { - let fields = plan.schema().fields(); - assert_eq!(fields.len(), 4, "{fields:?}"); - assert!( - fields.iter().any(|field| field.name() == "pod"), - "{fields:?}" - ); - assert!( - fields.iter().any(|field| field.name() == "instance"), - "{fields:?}" - ); - assert!( - fields - .iter() - .any(|field| field.name() == greptime_timestamp()), - "{fields:?}" - ); - assert!( - fields.iter().any(|field| { - field.name() == greptime_value() && field.data_type() == &ArrowDataType::Float64 - }), - "{fields:?}" - ); - } - - async fn build_test_table_provider_with_distinct_tags( - table_tags: &[(&str, &[&str])], - ) -> DfTableSourceProvider { - let catalog_list = MemoryCatalogManager::with_default_setup(); - for (table_name, tags) in table_tags { - let mut columns = tags - .iter() - .map(|tag| { - ColumnSchema::new( - (*tag).to_string(), - ConcreteDataType::string_datatype(), - false, - ) - }) - .collect::>(); - columns.push( - ColumnSchema::new( - greptime_timestamp().to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ); - columns.push(ColumnSchema::new( - greptime_value().to_string(), - ConcreteDataType::float64_datatype(), - true, - )); - let table_meta = TableMetaBuilder::empty() - .schema(Arc::new(Schema::new(columns))) - .primary_key_indices((0..tags.len()).collect()) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .name((*table_name).to_string()) - .meta(table_meta) - .build() - .unwrap(); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: (*table_name).to_string(), - table_id: 1024, - table: EmptyTable::from_table_info(&table_info), - }) - .is_ok() - ); - } - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - fn contains_histogram_fold(plan: &LogicalPlan) -> bool { - matches!(plan, LogicalPlan::Extension(Extension { node }) if node.as_any().is::()) - || plan.inputs().into_iter().any(contains_histogram_fold) - } - - async fn build_set_op_context_table_provider() -> DfTableSourceProvider { - build_test_table_provider_with_distinct_tags(&[ - ("bucket_metric", &["job", "le"]), - ("normal_metric", &["job"]), - ("fallback_metric", &["instance"]), - ]) - .await - } - - async fn build_or_context_table_provider() -> DfTableSourceProvider { - build_test_table_provider_with_distinct_tags(&[ - ("normal_metric", &["job"]), - ("other_metric", &["instance"]), - ("non_hist_metric", &["instance"]), - ]) - .await - } - - async fn optimize_and_create_physical_plan( - state: &QueryEngineState, - plan: LogicalPlan, - ) -> ( - LogicalPlan, - Arc, - ) { - let context = QueryEngineContext::new(state.session_state(), QueryContext::arc()); - let optimized = state.optimize_by_extension_rules(plan, &context).unwrap(); - let physical = state - .session_state() - .create_physical_plan(&optimized) - .await - .unwrap(); - (optimized, physical) - } - - async fn build_test_table_provider( - table_name_tuples: &[(String, String)], - num_tag: usize, - num_field: usize, - ) -> DfTableSourceProvider { - let catalog_list = MemoryCatalogManager::with_default_setup(); - for (schema_name, table_name) in table_name_tuples { - let mut columns = vec![]; - for i in 0..num_tag { - columns.push(ColumnSchema::new( - format!("tag_{i}"), - ConcreteDataType::string_datatype(), - false, - )); - } - columns.push( - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ); - for i in 0..num_field { - columns.push(ColumnSchema::new( - format!("field_{i}"), - ConcreteDataType::float64_datatype(), - true, - )); - } - let schema = Arc::new(Schema::new(columns)); - let table_meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices((0..num_tag).collect()) - .value_indices((num_tag + 1..num_tag + 1 + num_field).collect()) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .name(table_name.clone()) - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: schema_name.clone(), - table_name: table_name.clone(), - table_id: 1024, - table, - }) - .is_ok() - ); - } - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - async fn build_test_native_histogram_table_provider(table_name: &str) -> DfTableSourceProvider { - build_test_native_histogram_table_provider_with_marker(table_name, false).await - } - - async fn build_test_native_histogram_table_provider_with_marker( - table_name: &str, - temporality_marker: bool, - ) -> DfTableSourceProvider { - let catalog_list = MemoryCatalogManager::with_default_setup(); - let mut columns = vec![ - ColumnSchema::new( - "tag_0".to_string(), - ConcreteDataType::string_datatype(), - false, - ), - ColumnSchema::new( - LE_COLUMN_NAME.to_string(), - ConcreteDataType::string_datatype(), - true, - ), - ]; - if temporality_marker { - columns.push(ColumnSchema::new( - OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(), - ConcreteDataType::string_datatype(), - true, - )); - } - let tag_count = columns.len(); - columns.extend([ - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ColumnSchema::new( - greptime_native_histogram().to_string(), - native_histogram_value_type().clone(), - true, - ), - ]); - let schema = Arc::new(Schema::new(columns)); - let table_meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices((0..tag_count).collect()) - .value_indices(vec![tag_count + 1]) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .name(table_name) - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: table_name.to_string(), - table_id: 1024, - table, - }) - .is_ok() - ); - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - async fn build_test_multi_histogram_table_provider(table_name: &str) -> DfTableSourceProvider { - let catalog_list = MemoryCatalogManager::with_default_setup(); - let columns = vec![ - ColumnSchema::new( - "tag_0".to_string(), - ConcreteDataType::string_datatype(), - false, - ), - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ColumnSchema::new( - greptime_native_histogram().to_string(), - native_histogram_value_type().clone(), - true, - ), - ColumnSchema::new( - "native_histogram_2".to_string(), - native_histogram_value_type().clone(), - true, - ), - ]; - let schema = Arc::new(Schema::new(columns)); - let table_meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices(vec![0]) - .value_indices(vec![2, 3]) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .name(table_name) - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: table_name.to_string(), - table_id: 1024, - table, - }) - .is_ok() - ); - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - async fn build_test_mixed_native_histogram_table_provider( - table_name: &str, - ) -> DfTableSourceProvider { - build_test_mixed_native_histogram_table_provider_with_marker(table_name, false).await - } - - async fn build_test_mixed_native_histogram_table_provider_with_marker( - table_name: &str, - temporality_marker: bool, - ) -> DfTableSourceProvider { - let catalog_list = MemoryCatalogManager::with_default_setup(); - let mut columns = vec![ColumnSchema::new( - "tag_0".to_string(), - ConcreteDataType::string_datatype(), - false, - )]; - if temporality_marker { - columns.push(ColumnSchema::new( - OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(), - ConcreteDataType::string_datatype(), - true, - )); - } - let tag_count = columns.len(); - columns.extend([ - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ColumnSchema::new( - greptime_native_histogram().to_string(), - native_histogram_value_type().clone(), - true, - ), - ColumnSchema::new( - greptime_value().to_string(), - ConcreteDataType::float64_datatype(), - true, - ), - ]); - let schema = Arc::new(Schema::new(columns)); - let table_meta = TableMetaBuilder::empty() - .schema(schema.clone()) - .primary_key_indices((0..tag_count).collect()) - .value_indices(vec![tag_count + 1, tag_count + 2]) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = Arc::new( - TableInfoBuilder::default() - .name(table_name) - .meta(table_meta) - .build() - .unwrap(), - ); - let mut arrays: Vec> = - vec![Arc::new(StringArray::from(vec!["float", "histogram"]))]; - if temporality_marker { - arrays.push(Arc::new(StringArray::from(vec![ - Some(GREPTIME_TEMPORALITY_DELTA), - Some(GREPTIME_TEMPORALITY_DELTA), - ]))); - } - arrays.extend([ - Arc::new(TimestampMillisecondArray::from(vec![1_000, 1_000])) as Arc, - build_histogram_array(&[None, Some(direct_or_histogram())]), - Arc::new(Float64Array::from(vec![Some(2.0), None])), - ]); - let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap(); - let backing = GreptimeMemTable::new_with_catalog( - table_name, - GreptimeRecordBatch::from_df_record_batch(schema, batch), - 1024, - DEFAULT_CATALOG_NAME.to_string(), - DEFAULT_SCHEMA_NAME.to_string(), - ); - let table = Arc::new(Table::new( - table_info, - FilterPushDownType::Unsupported, - backing.data_source(), - )); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: table_name.to_string(), - table_id: 1024, - table, - }) - .is_ok() - ); - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - fn classic_and_native_histogram_table_provider( - native_tag: &str, - native_le: Option<&str>, - native_histogram: NativeHistogram, - ) -> DfTableSourceProvider { - let table_name = "mixed_histogram"; - let catalog = MemoryCatalogManager::with_default_setup(); - let schema = Arc::new(Schema::new(vec![ - ColumnSchema::new( - "tag".to_string(), - ConcreteDataType::string_datatype(), - false, - ), - ColumnSchema::new( - LE_COLUMN_NAME.to_string(), - ConcreteDataType::string_datatype(), - true, - ), - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ColumnSchema::new( - greptime_native_histogram().to_string(), - native_histogram_value_type().clone(), - true, - ), - ColumnSchema::new( - greptime_value().to_string(), - ConcreteDataType::float64_datatype(), - true, - ), - ])); - let table_meta = TableMetaBuilder::empty() - .schema(schema.clone()) - .primary_key_indices(vec![0, 1]) - .value_indices(vec![3, 4]) - .next_column_id(5) - .build() - .unwrap(); - let table_info = Arc::new( - TableInfoBuilder::default() - .name(table_name) - .meta(table_meta) - .build() - .unwrap(), - ); - let batch = RecordBatch::try_new( - schema.arrow_schema().clone(), - vec![ - Arc::new(StringArray::from(vec![ - "classic", "classic", native_tag, "classic", "classic", native_tag, - ])), - Arc::new(StringArray::from(vec![ - Some("1"), - Some("+Inf"), - native_le, - Some("1"), - Some("+Inf"), - native_le, - ])), - Arc::new(TimestampMillisecondArray::from(vec![ - 1_000, 1_000, 1_000, 2_000, 2_000, 2_000, - ])), - build_histogram_array(&[ - None, - None, - Some(native_histogram.clone()), - None, - None, - Some(native_histogram), - ]), - Arc::new(Float64Array::from(vec![ - Some(2.0), - Some(4.0), - None, - Some(2.0), - Some(4.0), - None, - ])), - ], - ) - .unwrap(); - let backing = GreptimeMemTable::new_with_catalog( - table_name, - GreptimeRecordBatch::from_df_record_batch(schema, batch), - 2_200, - DEFAULT_CATALOG_NAME.to_string(), - DEFAULT_SCHEMA_NAME.to_string(), - ); - let table = Arc::new(Table::new( - table_info, - FilterPushDownType::Unsupported, - backing.data_source(), - )); - catalog - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: table_name.to_string(), - table_id: 2_200, - table, - }) - .unwrap(); - - DfTableSourceProvider::new( - catalog, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - async fn build_test_table_provider_with_tsid( - table_name_tuples: &[(String, String)], - num_tag: usize, - num_field: usize, - ) -> DfTableSourceProvider { - let table_specs = table_name_tuples - .iter() - .map(|(schema_name, table_name)| ((schema_name.clone(), table_name.clone()), num_field)) - .collect::>(); - build_test_table_provider_with_tsid_fields(&table_specs, num_tag).await - } - - async fn build_test_table_provider_with_tsid_fields( - table_specs: &[((String, String), usize)], - num_tag: usize, - ) -> DfTableSourceProvider { - let table_specs = table_specs - .iter() - .map(|(table_name_tuple, num_field)| (table_name_tuple.clone(), num_tag, *num_field)) - .collect::>(); - build_test_table_provider_with_tsid_tag_fields(&table_specs).await - } - - async fn build_test_table_provider_with_tsid_tag_fields( - table_specs: &[((String, String), usize, usize)], - ) -> DfTableSourceProvider { - let catalog_list = MemoryCatalogManager::with_default_setup(); - - let physical_table_name = "phy"; - let physical_table_id = 999u32; - let physical_num_tag = table_specs - .iter() - .map(|(_, num_tag, _)| *num_tag) - .max() - .unwrap_or(0); - let physical_num_field = table_specs - .iter() - .map(|(_, _, num_field)| *num_field) - .max() - .unwrap_or(0); - - // Register a metric engine physical table with internal columns. - { - let mut columns = vec![ - ColumnSchema::new( - DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(), - ConcreteDataType::uint32_datatype(), - false, - ), - ColumnSchema::new( - DATA_SCHEMA_TSID_COLUMN_NAME.to_string(), - ConcreteDataType::uint64_datatype(), - false, - ), - ]; - for i in 0..physical_num_tag { - columns.push(ColumnSchema::new( - format!("tag_{i}"), - ConcreteDataType::string_datatype(), - false, - )); - } - columns.push( - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ); - for i in 0..physical_num_field { - columns.push(ColumnSchema::new( - format!("field_{i}"), - ConcreteDataType::float64_datatype(), - true, - )); - } - - let schema = Arc::new(Schema::new(columns)); - let primary_key_indices = (0..(2 + physical_num_tag)).collect::>(); - let table_meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices(primary_key_indices) - .value_indices( - (2 + physical_num_tag..2 + physical_num_tag + 1 + physical_num_field).collect(), - ) - .engine(METRIC_ENGINE_NAME.to_string()) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .table_id(physical_table_id) - .name(physical_table_name) - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: physical_table_name.to_string(), - table_id: physical_table_id, - table, - }) - .is_ok() - ); - } - - // Register metric engine logical tables without `__tsid`, referencing the physical table. - for (idx, ((schema_name, table_name), num_tag, num_field)) in table_specs.iter().enumerate() - { - let mut columns = vec![]; - for i in 0..*num_tag { - columns.push(ColumnSchema::new( - format!("tag_{i}"), - ConcreteDataType::string_datatype(), - false, - )); - } - columns.push( - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ); - for i in 0..*num_field { - columns.push(ColumnSchema::new( - format!("field_{i}"), - ConcreteDataType::float64_datatype(), - true, - )); - } - - let schema = Arc::new(Schema::new(columns)); - let mut options = table::requests::TableOptions::default(); - options.extra_options.insert( - LOGICAL_TABLE_METADATA_KEY.to_string(), - physical_table_name.to_string(), - ); - let table_id = 1024u32 + idx as u32; - let table_meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices((0..*num_tag).collect()) - .value_indices((*num_tag + 1..*num_tag + 1 + *num_field).collect()) - .engine(METRIC_ENGINE_NAME.to_string()) - .options(options) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .table_id(table_id) - .name(table_name.clone()) - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: schema_name.clone(), - table_name: table_name.clone(), - table_id, - table, - }) - .is_ok() - ); - } - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - async fn build_test_table_provider_with_fields( - table_name_tuples: &[(String, String)], - tags: &[&str], - ) -> DfTableSourceProvider { - let catalog_list = MemoryCatalogManager::with_default_setup(); - for (schema_name, table_name) in table_name_tuples { - let mut columns = vec![]; - let num_tag = tags.len(); - for tag in tags { - columns.push(ColumnSchema::new( - tag.to_string(), - ConcreteDataType::string_datatype(), - false, - )); - } - columns.push( - ColumnSchema::new( - greptime_timestamp().to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ); - columns.push(ColumnSchema::new( - greptime_value().to_string(), - ConcreteDataType::float64_datatype(), - true, - )); - let schema = Arc::new(Schema::new(columns)); - let table_meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices((0..num_tag).collect()) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .name(table_name.clone()) - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: schema_name.clone(), - table_name: table_name.clone(), - table_id: 1024, - table, - }) - .is_ok() - ); - } - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - // { - // input: `abs(some_metric{foo!="bar"})`, - // expected: &Call{ - // Func: MustGetFunction("abs"), - // Args: Expressions{ - // &VectorSelector{ - // Name: "some_metric", - // LabelMatchers: []*labels.Matcher{ - // MustLabelMatcher(labels.MatchNotEqual, "foo", "bar"), - // MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"), - // }, - // }, - // }, - // }, - // }, - async fn do_single_instant_function_call(fn_name: &'static str, plan_name: &str) { - let prom_expr = - parser::parse(&format!("{fn_name}(some_metric{{tag_0!=\"bar\"}})")).unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let expected = String::from( - "Filter: TEMPLATE(field_0) IS NOT NULL [timestamp:Timestamp(ms), TEMPLATE(field_0):Float64;N, tag_0:Utf8]\ - \n Projection: some_metric.timestamp, TEMPLATE(some_metric.field_0) AS TEMPLATE(field_0), some_metric.tag_0 [timestamp:Timestamp(ms), TEMPLATE(field_0):Float64;N, tag_0:Utf8]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]" - ).replace("TEMPLATE", plan_name); - - assert_eq!(plan.display_indent_schema().to_string(), expected); - } - - #[tokio::test] - async fn single_abs() { - do_single_instant_function_call("abs", "abs").await; - } - - #[tokio::test] - #[should_panic] - async fn single_absent() { - do_single_instant_function_call("absent", "").await; - } - - #[tokio::test] - async fn single_ceil() { - do_single_instant_function_call("ceil", "ceil").await; - } - - #[tokio::test] - async fn single_exp() { - do_single_instant_function_call("exp", "exp").await; - } - - #[tokio::test] - async fn single_ln() { - do_single_instant_function_call("ln", "ln").await; - } - - #[tokio::test] - async fn single_log2() { - do_single_instant_function_call("log2", "log2").await; - } - - #[tokio::test] - async fn single_log10() { - do_single_instant_function_call("log10", "log10").await; - } - - #[tokio::test] - #[should_panic] - async fn single_scalar() { - do_single_instant_function_call("scalar", "").await; - } - - #[tokio::test] - #[should_panic] - async fn single_sgn() { - do_single_instant_function_call("sgn", "").await; - } - - #[tokio::test] - #[should_panic] - async fn single_sort() { - do_single_instant_function_call("sort", "").await; - } - - #[tokio::test] - #[should_panic] - async fn single_sort_desc() { - do_single_instant_function_call("sort_desc", "").await; - } - - #[tokio::test] - async fn single_sqrt() { - do_single_instant_function_call("sqrt", "sqrt").await; - } - - #[tokio::test] - async fn single_timestamp_plan_preserves_source_value() { - let eval_stmt = build_eval_stmt(r#"timestamp(some_metric{tag_0!="bar"})"#); - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let expected = String::from( - "Filter: value IS NOT NULL [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\ - \n Projection: some_metric.timestamp, value AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\ - \n Projection: some_metric.timestamp, __promql_timestamp_value_ AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\ - \n Projection: some_metric.tag_0, some_metric.timestamp, some_metric.field_0, CAST(CAST(CAST(CAST(some_metric.timestamp AS Int64) AS Decimal128(19, 0)) * Decimal128(1,1,0) + Decimal128(0,19,0) AS Int64) AS Float64) / Float64(1000) AS __promql_timestamp_value_ [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - assert_eq!(plan.display_indent_schema().to_string(), expected); - } - - #[tokio::test] - async fn single_acos() { - do_single_instant_function_call("acos", "acos").await; - } - - #[tokio::test] - #[should_panic] - async fn single_acosh() { - do_single_instant_function_call("acosh", "").await; - } - - #[tokio::test] - async fn single_asin() { - do_single_instant_function_call("asin", "asin").await; - } - - #[tokio::test] - #[should_panic] - async fn single_asinh() { - do_single_instant_function_call("asinh", "").await; - } - - #[tokio::test] - async fn single_atan() { - do_single_instant_function_call("atan", "atan").await; - } - - #[tokio::test] - #[should_panic] - async fn single_atanh() { - do_single_instant_function_call("atanh", "").await; - } - - #[tokio::test] - async fn single_cos() { - do_single_instant_function_call("cos", "cos").await; - } - - #[tokio::test] - #[should_panic] - async fn single_cosh() { - do_single_instant_function_call("cosh", "").await; - } - - #[tokio::test] - async fn single_sin() { - do_single_instant_function_call("sin", "sin").await; - } - - #[tokio::test] - #[should_panic] - async fn single_sinh() { - do_single_instant_function_call("sinh", "").await; - } - - #[tokio::test] - async fn single_tan() { - do_single_instant_function_call("tan", "tan").await; - } - - #[tokio::test] - #[should_panic] - async fn single_tanh() { - do_single_instant_function_call("tanh", "").await; - } - - #[tokio::test] - #[should_panic] - async fn single_deg() { - do_single_instant_function_call("deg", "").await; - } - - #[tokio::test] - #[should_panic] - async fn single_rad() { - do_single_instant_function_call("rad", "").await; - } - - // { - // input: "avg by (foo)(some_metric)", - // expected: &AggregateExpr{ - // Op: AVG, - // Expr: &VectorSelector{ - // Name: "some_metric", - // LabelMatchers: []*labels.Matcher{ - // MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"), - // }, - // PosRange: PositionRange{ - // Start: 13, - // End: 24, - // }, - // }, - // Grouping: []string{"foo"}, - // PosRange: PositionRange{ - // Start: 0, - // End: 25, - // }, - // }, - // }, - async fn do_aggregate_expr_plan(fn_name: &str, plan_name: &str) { - let prom_expr = parser::parse(&format!( - "{fn_name} by (tag_1)(some_metric{{tag_0!=\"bar\"}})", - )) - .unwrap(); - let mut eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - // test group by - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 2, - 2, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected_no_without = String::from( - "Sort: some_metric.tag_1 ASC NULLS LAST, some_metric.timestamp ASC NULLS LAST [tag_1:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ - \n Aggregate: groupBy=[[some_metric.tag_1, some_metric.timestamp]], aggr=[[TEMPLATE(some_metric.field_0), TEMPLATE(some_metric.field_1)]] [tag_1:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.tag_1 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ - \n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]" - ).replace("TEMPLATE", plan_name); - assert_eq!( - plan.display_indent_schema().to_string(), - expected_no_without - ); - - // test group without - if let PromExpr::Aggregate(AggregateExpr { modifier, .. }) = &mut eval_stmt.expr { - *modifier = Some(LabelModifier::Exclude(Labels { - labels: vec![String::from("tag_1")].into_iter().collect(), - })); - } - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 2, - 2, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected_without = String::from( - "Sort: some_metric.tag_0 ASC NULLS LAST, some_metric.timestamp ASC NULLS LAST [tag_0:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ - \n Aggregate: groupBy=[[some_metric.tag_0, some_metric.timestamp]], aggr=[[TEMPLATE(some_metric.field_0), TEMPLATE(some_metric.field_1)]] [tag_0:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.tag_1 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ - \n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]" - ).replace("TEMPLATE", plan_name); - assert_eq!(plan.display_indent_schema().to_string(), expected_without); - } - - #[tokio::test] - async fn aggregate_sum() { - do_aggregate_expr_plan("sum", "sum").await; - } - - #[tokio::test] - async fn tsid_is_used_for_series_divide_when_available() { - let prom_expr = parser::parse("some_metric").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); - assert!(plan_str.contains("__tsid ASC NULLS FIRST")); - assert!( - !plan - .schema() - .fields() - .iter() - .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME) - ); - - let manipulate = find_instant_manipulate(&plan).unwrap(); - let exec = manipulate.to_execution_plan(Arc::new(DataSourceExec::new(Arc::new( - MemorySourceConfig::try_new( - &[], - Arc::new( - datafusion_expr::UserDefinedLogicalNodeCore::inputs(manipulate)[0] - .schema() - .as_arrow() - .clone(), - ), - None, - ) - .unwrap(), - )))); - assert!(format!("{exec:?}").contains("reuse_tsid_column: true")); - } - - #[tokio::test] - async fn default_binary_join_uses_tsid_when_available() { - let eval_stmt = build_eval_stmt("some_metric / some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!( - plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), - "{plan_str}" - ); - assert!( - !plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"), - "{plan_str}" - ); - } - - #[tokio::test] - async fn reject_binary_fill_modifiers() { - let state = build_query_engine_state(); - - for query in [ - "some_metric + fill(0) some_alt_metric", - "some_metric + fill_left(0) some_alt_metric", - "some_metric + fill_right(0) some_alt_metric", - "(some_metric + fill(0) some_alt_metric) + some_metric", - ] { - let eval_stmt = build_eval_stmt(query); - let table_provider = build_test_table_provider(&[], 0, 0).await; - let err = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &state) - .await - .unwrap_err(); - - assert!( - matches!( - &err, - crate::promql::error::Error::UnsupportedExpr { name, .. } - if name == "PromQL fill modifiers" - ), - "{err}" - ); - } - } - - #[tokio::test] - async fn timestamp_binary_join_falls_back_when_tsid_is_projected_out() { - for query in [ - "timestamp(some_metric) / some_metric", - "some_metric / timestamp(some_metric)", - ] { - let eval_stmt = build_eval_stmt(query); - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(!plan_str.contains("__tsid ="), "{query}: {plan_str}"); - assert!( - plan_str.contains("lhs.tag_0 = rhs.tag_0"), - "{query}: {plan_str}" - ); - assert!( - !plan - .schema() - .fields() - .iter() - .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME), - "{query}: {plan_str}" - ); - } - } - - #[tokio::test] - async fn timestamp_binary_join_rejects_default_matching_on_mismatched_labels() { - let eval_stmt = build_eval_stmt("timestamp(left_host_job) / right_by_job"); - - let table_provider = build_test_table_provider_with_tsid_tag_fields(&[ - ( - (DEFAULT_SCHEMA_NAME.to_string(), "left_host_job".to_string()), - 2, - 1, - ), - ( - (DEFAULT_SCHEMA_NAME.to_string(), "right_by_job".to_string()), - 1, - 1, - ), - ]) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let plan_str = plan.display_indent_schema().to_string(); - - assert!( - plan_str.contains("Boolean(false)") || plan_str.contains("false"), - "{plan_str}" - ); - } - - #[tokio::test] - async fn tsid_is_preserved_for_nested_default_binary_joins() { - let eval_stmt = build_eval_stmt("(some_metric - some_alt_metric) / some_third_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_third_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); - assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); - } - - #[tokio::test] - async fn repeated_tsid_binary_operand_reuses_leaf_plan() { - let eval_stmt = build_eval_stmt("((some_metric - some_alt_metric) / some_metric) * 100"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}"); - assert_eq!( - plan_str - .matches("Filter: phy.__table_id = UInt32(1024)") - .count(), - 1, - "{plan_str}" - ); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 2, - "{plan_str}" - ); - assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); - } - - #[tokio::test] - async fn repeated_tsid_binary_operand_reuses_shorter_field_side() { - let eval_stmt = - build_eval_stmt("((two_field_metric - one_field_metric) / one_field_metric) * 100"); - - let table_provider = build_test_table_provider_with_tsid_fields( - &[ - ( - ( - DEFAULT_SCHEMA_NAME.to_string(), - "two_field_metric".to_string(), - ), - 2, - ), - ( - ( - DEFAULT_SCHEMA_NAME.to_string(), - "one_field_metric".to_string(), - ), - 1, - ), - ], - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let field_names = plan - .schema() - .fields() - .iter() - .map(|field| field.name().clone()) - .collect::>(); - let value_columns = field_names - .iter() - .filter(|name| { - *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME - }) - .count(); - assert_eq!(value_columns, 1, "{field_names:?}"); - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}"); - assert_eq!( - plan_str - .matches("Filter: phy.__table_id = UInt32(1025)") - .count(), - 1, - "{plan_str}" - ); - assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); - } - - #[tokio::test] - async fn binary_island_reuses_self_operand_without_join() { - let eval_stmt = build_eval_stmt("some_metric / some_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 0, "{plan_str}"); - assert_eq!( - plan_str - .matches("Filter: phy.__table_id = UInt32(1024)") - .count(), - 1, - "{plan_str}" - ); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 1, - "{plan_str}" - ); - } - - #[tokio::test] - async fn binary_island_reuses_leaf_across_two_branches() { - let eval_stmt = - build_eval_stmt("(some_metric + some_alt_metric) / (some_metric + third_metric)"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - (DEFAULT_SCHEMA_NAME.to_string(), "third_metric".to_string()), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); - assert_eq!( - plan_str - .matches("Filter: phy.__table_id = UInt32(1024)") - .count(), - 1, - "{plan_str}" - ); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 3, - "{plan_str}" - ); - } - - #[tokio::test] - async fn binary_island_generated_alias_avoids_user_column_names() { - let eval_stmt = build_eval_stmt("(some_metric + some_alt_metric) / some_metric"); - - let table_provider = build_test_table_provider_with_fields( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - &["prom_v0", "__prom_v0"], - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let field_names = plan.schema().field_names(); - assert!(field_names.iter().any(|name| name.ends_with(".prom_v0"))); - assert!(field_names.iter().any(|name| name.ends_with(".__prom_v0"))); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("SubqueryAlias: __prom_v0"), "{plan_str}"); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 2, - "{plan_str}" - ); - } - - #[tokio::test] - async fn binary_island_clears_qualifier_for_nested_unary_projection() { - let eval_stmt = build_eval_stmt("-((some_metric + some_alt_metric) / some_metric)"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}"); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 2, - "{plan_str}" - ); - } - - #[tokio::test] - async fn binary_island_keeps_distinct_matcher_leaves() { - let eval_stmt = build_eval_stmt( - "(some_metric{tag_0=\"foo\"} + some_alt_metric) / some_metric{tag_0=\"bar\"}", - ); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 3, - "{plan_str}" - ); - } - - #[tokio::test] - async fn binary_island_keeps_offset_leaves_distinct() { - let eval_stmt = build_eval_stmt("(some_metric offset 5m + some_alt_metric) / some_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 3, - "{plan_str}" - ); - } - - #[tokio::test] - async fn binary_island_falls_back_for_group_modifier() { - let eval_stmt = build_eval_stmt( - "(some_metric + ignoring(tag_0) group_left some_alt_metric) / some_metric", - ); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 3, - "{plan_str}" - ); - } - - #[tokio::test] - async fn binary_island_falls_back_for_comparison_filter() { - let eval_stmt = build_eval_stmt("(some_metric > some_alt_metric) / some_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 3, - "{plan_str}" - ); - } - - #[tokio::test] - async fn tsid_binary_join_uses_shorter_field_side() { - let eval_stmt = build_eval_stmt("one_field_metric / two_field_metric"); - - let table_provider = build_test_table_provider_with_tsid_fields( - &[ - ( - ( - DEFAULT_SCHEMA_NAME.to_string(), - "one_field_metric".to_string(), - ), - 1, - ), - ( - ( - DEFAULT_SCHEMA_NAME.to_string(), - "two_field_metric".to_string(), - ), - 2, - ), - ], - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let field_names = plan - .schema() - .fields() - .iter() - .map(|field| field.name().clone()) - .collect::>(); - let value_columns = field_names - .iter() - .filter(|name| { - *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME - }) - .count(); - assert_eq!(value_columns, 1, "{field_names:?}"); - } - - #[tokio::test] - async fn comparison_binary_join_uses_shorter_field_side() { - let eval_stmt = build_eval_stmt("two_field_metric > one_field_metric"); - - let table_provider = build_test_table_provider_with_tsid_fields( - &[ - ( - ( - DEFAULT_SCHEMA_NAME.to_string(), - "two_field_metric".to_string(), - ), - 2, - ), - ( - ( - DEFAULT_SCHEMA_NAME.to_string(), - "one_field_metric".to_string(), - ), - 1, - ), - ], - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let field_names = plan - .schema() - .fields() - .iter() - .map(|field| field.name().clone()) - .collect::>(); - assert!( - field_names.iter().any(|name| name == "field_0"), - "{field_names:?}" - ); - assert!( - !field_names.iter().any(|name| name == "field_1"), - "{field_names:?}" - ); - } - - #[tokio::test] - async fn label_matching_modifier_disables_tsid_binary_join() { - let eval_stmt = build_eval_stmt("some_metric / ignoring(tag_0) some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 2, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(!plan_str.contains("__tsid ="), "{plan_str}"); - assert!( - plan_str.contains("some_metric.tag_1 = some_alt_metric.tag_1"), - "{plan_str}" - ); - } - - #[tokio::test] - async fn ignoring_absent_label_keeps_tsid_binary_join() { - let eval_stmt = build_eval_stmt("some_metric / ignoring(missing) some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 2, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!( - plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), - "{plan_str}" - ); - assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); - assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); - } - - #[tokio::test] - async fn range_function_keeps_tsid_for_absent_ignoring_binary_join() { - let eval_stmt = - build_eval_stmt("rate(some_metric[5m]) / ignoring(missing) some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 2, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!( - plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), - "{plan_str}" - ); - assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); - assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); - } - - #[tokio::test] - async fn on_full_label_set_keeps_tsid_binary_join() { - let eval_stmt = build_eval_stmt("some_metric / on(tag_0, tag_1) some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 2, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!( - plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), - "{plan_str}" - ); - assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); - assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); - } - - #[tokio::test] - async fn on_partial_label_set_disables_tsid_binary_join() { - let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 2, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(!plan_str.contains("__tsid ="), "{plan_str}"); - assert!( - plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"), - "{plan_str}" - ); - assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); - } - - #[tokio::test] - async fn on_label_set_must_cover_both_sides_to_use_tsid_binary_join() { - let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid_tag_fields(&[ - ( - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - 2, - 1, - ), - ( - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - 1, - 1, - ), - ]) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(!plan_str.contains("__tsid ="), "{plan_str}"); - assert!( - plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"), - "{plan_str}" - ); - assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); - } - - #[tokio::test] - async fn comparison_binary_join_uses_tsid_and_keeps_it_in_filtered_result() { - let eval_stmt = build_eval_stmt("some_metric > some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 2, - 1, - ) - .await; - let mut planner = PromPlanner { - table_provider, - ctx: PromPlannerContext::from_eval_stmt(&eval_stmt), - promql_annotations: None, - }; - let plan = planner - .prom_expr_to_plan(&eval_stmt.expr, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!( - plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), - "{plan_str}" - ); - assert!( - plan.schema() - .fields() - .iter() - .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME), - "{plan_str}" - ); - assert!(planner.ctx.use_tsid, "{plan_str}"); - } - - #[tokio::test] - async fn comparison_bool_binary_join_uses_tsid_when_available() { - let eval_stmt = build_eval_stmt("some_metric > bool some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 2, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!( - plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), - "{plan_str}" - ); - assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); - assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); - } - - #[tokio::test] - async fn scalar_count_count_range_keeps_full_window() { - let plan_str = build_optimized_tsid_plan( - "scalar(count(count(some_metric) by (tag_0)))", - 1, - 1, - 100_000, - 1, - ) - .await; - assert!(plan_str.contains("ScalarCalculate: tags=[]")); - assert!(plan_str.contains("PromInstantManipulate: range=[0..100000000]")); - assert!(!plan_str.contains("PromInstantManipulate: range=[99999000..99999000]")); - } - - #[tokio::test] - async fn scalar_count_count_rewrite_applies_inside_binary_expr_for_tsid_input() { - let plan_str = build_optimized_tsid_plan( - "sum(irate(some_metric[1h])) / scalar(count(count(some_metric) by (tag_0)))", - 2, - 1, - 10, - 300, - ) - .await; - assert!(plan_str.contains("Distinct:"), "{plan_str}"); - } - - #[tokio::test] - async fn nested_count_rewrite_keeps_full_series_key_with_tsid_input() { - assert_nested_count_rewrite_applies( - "count(count(some_metric) by (tag_0))", - "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(count(some_metric.field_0))]]" - ) - .await; - } - - #[tokio::test] - async fn nested_sum_count_rewrite_keeps_full_series_key_with_tsid_input() { - assert_nested_count_rewrite_applies( - "count(sum(some_metric) by (tag_0))", - "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(sum(some_metric.field_0))]]" - ) - .await; - } - - #[tokio::test] - async fn nested_supported_inner_aggs_rewrite_apply_for_tsid_input() { - for (query, expected_outer_agg) in [ - ( - "count(avg(some_metric) by (tag_0))", - "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(avg(some_metric.field_0))]]", - ), - ( - "count(min(some_metric) by (tag_0))", - "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(min(some_metric.field_0))]]", - ), - ( - "count(max(some_metric) by (tag_0))", - "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(max(some_metric.field_0))]]", - ), - ( - "count(stddev(some_metric) by (tag_0))", - "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(stddev_pop(some_metric.field_0))]]", - ), - ( - "count(stdvar(some_metric) by (tag_0))", - "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(var_pop(some_metric.field_0))]]", - ), - ] { - assert_nested_count_rewrite_applies(query, expected_outer_agg).await; - } - } - - #[tokio::test] - async fn nested_non_count_inner_aggs_rewrite_filter_null_values_for_tsid_input() { - let count_plan = - build_optimized_tsid_plan("count(count(some_metric) by (tag_0))", 2, 1, 100_000, 1) - .await; - assert!( - !count_plan.contains("some_metric.field_0 IS NOT NULL"), - "{count_plan}" - ); - - for query in [ - "count(sum(some_metric) by (tag_0))", - "count(avg(some_metric) by (tag_0))", - "count(min(some_metric) by (tag_0))", - "count(max(some_metric) by (tag_0))", - "count(stddev(some_metric) by (tag_0))", - "count(stdvar(some_metric) by (tag_0))", - ] { - let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await; - assert!( - plan_str.contains("Filter: some_metric.field_0 IS NOT NULL"), - "{query}: {plan_str}" - ); - } - } - - #[tokio::test] - async fn nested_unsupported_or_non_direct_inner_aggs_do_not_rewrite() { - assert_nested_count_rewrite_missing("count(group(some_metric) by (tag_0))", 2, 1).await; - assert_nested_count_rewrite_missing( - "count(sum(irate(some_metric[1h])) by (tag_0))", - 2, - 300, - ) - .await; - } - - #[tokio::test] - async fn physical_table_name_is_not_leaked_in_plan() { - let prom_expr = parser::parse("some_metric").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("TableScan: phy"), "{plan}"); - assert!(plan_str.contains("SubqueryAlias: some_metric")); - assert!(plan_str.contains("Filter: phy.__table_id = UInt32(1024)")); - assert!(!plan_str.contains("TableScan: some_metric")); - } - - #[tokio::test] - async fn sum_without_does_not_group_by_tsid() { - let prom_expr = parser::parse("sum without (tag_0) (some_metric)").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); - - let aggr_line = plan_str - .lines() - .find(|line| line.contains("Aggregate: groupBy=")) - .unwrap(); - assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME)); - } - - #[tokio::test] - async fn topk_without_does_not_partition_by_tsid() { - let prom_expr = parser::parse("topk without (tag_0) (1, some_metric)").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); - - let window_line = plan_str - .lines() - .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()")) - .unwrap(); - let partition_by = window_line - .split("PARTITION BY [") - .nth(1) - .and_then(|s| s.split("] ORDER BY").next()) - .unwrap(); - assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME)); - } - - #[tokio::test] - async fn sum_by_does_not_group_by_tsid() { - let prom_expr = parser::parse("sum by (__tsid) (some_metric)").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); - - let aggr_line = plan_str - .lines() - .find(|line| line.contains("Aggregate: groupBy=")) - .unwrap(); - assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME)); - } - - #[tokio::test] - async fn aggregate_over_binary_time_function_expr() { - for op in ["sum", "min", "max", "avg"] { - let prom_expr = parser::parse(&format!( - "{op} by (tag_0, tag_1, tag_2) (time() - some_metric)" - )) - .unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 3, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - let aggr_line = plan_str - .lines() - .find(|line| line.contains("Aggregate: groupBy=")) - .unwrap(); - assert!(aggr_line.contains(op), "{plan_str}"); - assert!(aggr_line.contains("first_value"), "{plan_str}"); - assert!( - !plan - .schema() - .fields() - .iter() - .any(|field| { field.name() == DATA_SCHEMA_TSID_COLUMN_NAME }) - ); - } - } - - #[tokio::test] - async fn topk_by_does_not_partition_by_tsid() { - let prom_expr = parser::parse("topk by (__tsid) (1, some_metric)").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); - - let window_line = plan_str - .lines() - .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()")) - .unwrap(); - let partition_by = window_line - .split("PARTITION BY [") - .nth(1) - .and_then(|s| s.split("] ORDER BY").next()) - .unwrap(); - assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME)); - } - - #[tokio::test] - async fn selector_matcher_on_tsid_does_not_use_internal_column() { - let prom_expr = parser::parse(r#"some_metric{__tsid="123"}"#).unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - fn collect_filter_cols(plan: &LogicalPlan, out: &mut HashSet) { - if let LogicalPlan::Filter(filter) = plan { - datafusion_expr::utils::expr_to_columns(&filter.predicate, out).unwrap(); - } - for input in plan.inputs() { - collect_filter_cols(input, out); - } - } - - let mut filter_cols = HashSet::new(); - collect_filter_cols(&plan, &mut filter_cols); - assert!( - !filter_cols - .iter() - .any(|c| c.name == DATA_SCHEMA_TSID_COLUMN_NAME) - ); - } - - #[tokio::test] - async fn tsid_is_not_used_when_physical_table_is_missing() { - let prom_expr = parser::parse("some_metric").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let catalog_list = MemoryCatalogManager::with_default_setup(); - - // Register a metric engine logical table referencing a missing physical table. - let mut columns = vec![ColumnSchema::new( - "tag_0".to_string(), - ConcreteDataType::string_datatype(), - false, - )]; - columns.push( - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ); - columns.push(ColumnSchema::new( - "field_0".to_string(), - ConcreteDataType::float64_datatype(), - true, - )); - let schema = Arc::new(Schema::new(columns)); - let mut options = table::requests::TableOptions::default(); - options - .extra_options - .insert(LOGICAL_TABLE_METADATA_KEY.to_string(), "phy".to_string()); - let table_meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices(vec![0]) - .value_indices(vec![2]) - .engine(METRIC_ENGINE_NAME.to_string()) - .options(options) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .table_id(1024) - .name("some_metric") - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: "some_metric".to_string(), - table_id: 1024, - table, - }) - .unwrap(); - - let table_provider = DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ); - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]")); - assert!(!plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); - } - - #[tokio::test] - async fn tsid_is_carried_only_when_aggregate_preserves_label_set() { - let prom_expr = parser::parse("sum by (tag_0) (some_metric)").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("first_value") && plan_str.contains("__tsid")); - assert!( - !plan - .schema() - .fields() - .iter() - .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME) - ); - - // Merging aggregate: label set is reduced, tsid should not be carried. - let prom_expr = parser::parse("sum(some_metric)").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let plan_str = plan.display_indent_schema().to_string(); - assert!(!plan_str.contains("first_value")); - } - - #[tokio::test] - async fn or_operator_with_unknown_metric_does_not_require_tsid() { - let prom_expr = parser::parse("unknown_metric or some_metric").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - assert!( - !plan - .schema() - .fields() - .iter() - .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME) - ); - } - - #[tokio::test] - async fn aggregate_avg() { - do_aggregate_expr_plan("avg", "avg").await; - } - - #[tokio::test] - #[should_panic] // output type doesn't match - async fn aggregate_count() { - do_aggregate_expr_plan("count", "count").await; - } - - #[tokio::test] - async fn aggregate_min() { - do_aggregate_expr_plan("min", "min").await; - } - - #[tokio::test] - async fn aggregate_max() { - do_aggregate_expr_plan("max", "max").await; - } - - #[tokio::test] - async fn aggregate_group() { - // Regression test for `group()` aggregator. - // PromQL: sum(group by (cluster)(kubernetes_build_info{service="kubernetes",job="apiserver"})) - // should be plannable, and `group()` should produce constant 1 for each group. - let prom_expr = parser::parse( - "sum(group by (cluster)(kubernetes_build_info{service=\"kubernetes\",job=\"apiserver\"}))", - ) - .unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_fields( - &[( - DEFAULT_SCHEMA_NAME.to_string(), - "kubernetes_build_info".to_string(), - )], - &["cluster", "service", "job"], - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("max(Float64(1")); - } - - #[tokio::test] - async fn aggregate_stddev() { - do_aggregate_expr_plan("stddev", "stddev_pop").await; - } - - #[tokio::test] - async fn aggregate_stdvar() { - do_aggregate_expr_plan("stdvar", "var_pop").await; - } - - // TODO(ruihang): add range fn tests once exprs are ready. - - // { - // input: "some_metric{tag_0="foo"} + some_metric{tag_0="bar"}", - // expected: &BinaryExpr{ - // Op: ADD, - // LHS: &VectorSelector{ - // Name: "a", - // LabelMatchers: []*labels.Matcher{ - // MustLabelMatcher(labels.MatchEqual, "tag_0", "foo"), - // MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"), - // }, - // }, - // RHS: &VectorSelector{ - // Name: "sum", - // LabelMatchers: []*labels.Matcher{ - // MustLabelMatcher(labels.MatchxEqual, "tag_0", "bar"), - // MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"), - // }, - // }, - // VectorMatching: &VectorMatching{}, - // }, - // }, - #[tokio::test] - async fn binary_op_column_column() { - let prom_expr = - parser::parse(r#"some_metric{tag_0="foo"} + some_metric{tag_0="bar"}"#).unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let expected = String::from( - "Projection: rhs.tag_0, rhs.timestamp, CAST(lhs.field_0 AS Float64) + CAST(rhs.field_0 AS Float64) AS lhs.field_0 + rhs.field_0 [tag_0:Utf8, timestamp:Timestamp(ms), lhs.field_0 + rhs.field_0:Float64;N]\ - \n Inner Join: lhs.tag_0 = rhs.tag_0, lhs.timestamp = rhs.timestamp [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n SubqueryAlias: lhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.tag_0 = Utf8(\"foo\") AND some_metric.tag_0 = Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n SubqueryAlias: rhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.tag_0 = Utf8(\"bar\") AND some_metric.tag_0 = Utf8(\"foo\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - assert_eq!(plan.display_indent_schema().to_string(), expected); - } - - async fn indie_query_plan_compare>(query: &str, expected: T) { - let prom_expr = parser::parse(query).unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - "greptime_private".to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - assert_eq!(plan.display_indent_schema().to_string(), expected.as_ref()); - } - - #[tokio::test] - async fn binary_op_literal_column() { - let query = r#"1 + some_metric{tag_0="bar"}"#; - let expected = String::from( - "Projection: some_metric.tag_0, some_metric.timestamp, Float64(1) + CAST(some_metric.field_0 AS Float64) AS Float64(1) + field_0 [tag_0:Utf8, timestamp:Timestamp(ms), Float64(1) + field_0:Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.tag_0 = Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - } - - #[tokio::test] - async fn binary_op_literal_literal() { - let query = r#"1 + 1"#; - let expected = r#"EmptyMetric: range=[0..100000000], interval=[5000] [time:Timestamp(ms), value:Float64;N] - TableScan: dummy [time:Timestamp(ms), value:Float64;N]"#; - indie_query_plan_compare(query, expected).await; - } - - #[tokio::test] - async fn simple_bool_grammar() { - let query = "some_metric != bool 1.2345"; - let expected = String::from( - "Projection: some_metric.tag_0, some_metric.timestamp, CAST(some_metric.field_0 != Float64(1.2345) AS Float64) AS field_0 != Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0 != Float64(1.2345):Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - } - - #[tokio::test] - async fn bool_with_additional_arithmetic() { - let query = "some_metric + (1 == bool 2)"; - let expected = String::from( - "Projection: some_metric.tag_0, some_metric.timestamp, CAST(some_metric.field_0 AS Float64) + CAST(Float64(1) = Float64(2) AS Float64) AS field_0 + Float64(1) = Float64(2) [tag_0:Utf8, timestamp:Timestamp(ms), field_0 + Float64(1) = Float64(2):Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - } - - #[tokio::test] - async fn simple_unary() { - let query = "-some_metric"; - let expected = String::from( - "Projection: some_metric.tag_0, some_metric.timestamp, (- some_metric.field_0) AS (- field_0) [tag_0:Utf8, timestamp:Timestamp(ms), (- field_0):Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - } - - #[tokio::test] - async fn increase_aggr() { - let query = "increase(some_metric[5m])"; - let expected = String::from( - "Filter: prom_increase(timestamp_range,field_0,timestamp,Int64(300000)) IS NOT NULL [timestamp:Timestamp(ms), prom_increase(timestamp_range,field_0,timestamp,Int64(300000)):Float64;N, tag_0:Utf8]\ - \n Projection: some_metric.timestamp, prom_increase(timestamp_range, field_0, some_metric.timestamp, Int64(300000)) AS prom_increase(timestamp_range,field_0,timestamp,Int64(300000)), some_metric.tag_0 [timestamp:Timestamp(ms), prom_increase(timestamp_range,field_0,timestamp,Int64(300000)):Float64;N, tag_0:Utf8]\ - \n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[300000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\ - \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.timestamp >= TimestampMillisecond(-299999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - } - - async fn native_histogram_plan(query: &str) -> String { - let table_provider = build_test_native_histogram_table_provider("some_metric").await; - let plan = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt(query), - &build_query_engine_state(), - ) - .await - .unwrap(); - plan.display_indent_schema().to_string() - } - - #[tokio::test] - async fn native_histogram_count_uses_native_udf() { - let plan = native_histogram_plan("histogram_count(some_metric)").await; - - assert!(plan.contains("prom_native_histogram_count"), "{plan}"); - assert!(!plan.contains("HistogramFold:"), "{plan}"); - } - - #[tokio::test] - async fn timestamp_filters_native_histogram_stale_marker_before_projection() { - let mut stale = direct_or_histogram(); - stale.sum = f64::from_bits(PROMETHEUS_STALE_NAN_BITS); - let table = operator_metric_table( - "stale_histogram", - 2_100, - "a", - None, - DirectOrValue::NativeHistogram(stale), - ); - let catalog = MemoryCatalogManager::with_default_setup(); - catalog - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: "stale_histogram".to_string(), - table_id: 2_100, - table, - }) - .unwrap(); - let provider = DfTableSourceProvider::new( - catalog, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ); - let state = build_query_engine_state(); - let plan = PromPlanner::stmt_to_plan( - provider, - &operator_eval_stmt("timestamp(stale_histogram)"), - &state, - ) - .await - .unwrap(); - let plan_text = plan.display_indent_schema().to_string(); - assert!(plan_text.contains(TIMESTAMP_VALUE_PREFIX), "{plan_text}"); - - let (_, batches) = execute(plan, &state).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); - } - - #[tokio::test] - async fn timestamp_filters_stale_marker_from_mixed_sample_companion() { - let histograms = build_histogram_array(&[None]); - let schema = Arc::new(ArrowSchema::new(vec![ - Field::new( - "timestamp", - ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), - false, - ), - Field::new( - greptime_native_histogram(), - histograms.data_type().clone(), - true, - ), - Field::new(greptime_value(), ArrowDataType::Float64, true), - ])); - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(TimestampMillisecondArray::from(vec![1_000])), - histograms, - Arc::new(Float64Array::from(vec![f64::from_bits( - PROMETHEUS_STALE_NAN_BITS, - )])), - ], - ) - .unwrap(); - let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); - let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None) - .unwrap() - .build() - .unwrap(); - let input = LogicalPlan::Extension(Extension { - node: Arc::new(SeriesDivide::new( - Vec::new(), - "timestamp".to_string(), - input, - )), - }); - let input = LogicalPlan::Extension(Extension { - node: Arc::new(InstantManipulate::new( - 1_000, - 1_000, - 5_000, - 1_000, - 0, - "timestamp".to_string(), - Vec::new(), - Some(greptime_native_histogram().to_string()), - input, - )), - }); - // Match timestamp()'s parent projection, which otherwise prunes the companion lane. - let plan = LogicalPlanBuilder::from(input) - .project([col("timestamp")]) - .unwrap() - .build() - .unwrap(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); - } - - #[tokio::test] - async fn native_histogram_rate_can_feed_count() { - let plan = native_histogram_plan("histogram_count(rate(some_metric[5m]))").await; - - assert!(plan.contains("prom_native_histogram_rate"), "{plan}"); - assert!(plan.contains("prom_native_histogram_count"), "{plan}"); - } - - #[tokio::test] - async fn native_histogram_quantile_skips_classic_fold() { - let plan = native_histogram_plan("histogram_quantile(0.9, some_metric)").await; - - assert!(plan.contains("prom_native_histogram_quantile"), "{plan}"); - assert!(!plan.contains("HistogramFold:"), "{plan}"); - assert!(plan.contains("some_metric.le"), "{plan}"); - // The phi literal is threaded into the native quantile UDF as its second argument. - assert!(plan.contains("Float64(0.9)"), "{plan}"); - // The empty-values filter drops NULL quantile results so the output is empty - // when all native histogram samples are dropped. - assert!(plan.contains("IS NOT NULL"), "{plan}"); - } - - #[tokio::test] - async fn mixed_native_histogram_quantile_uses_histogram_field() { - let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; - let plan = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt("histogram_quantile(0.9, some_metric)"), - &build_query_engine_state(), - ) - .await - .unwrap() - .display_indent_schema() - .to_string(); - - assert!( - plan.contains("prom_native_histogram_quantile(greptime_native_histogram"), - "{plan}" - ); - assert!(!plan.contains("EmptyRelation"), "{plan}"); - } - - #[tokio::test] - async fn mixed_histogram_helpers_execute_classic_and_native_samples() { - let state = build_query_engine_state(); - for (query, expected) in [ - ( - "histogram_quantile(0.5, mixed_histogram)", - vec![("classic", 1.0), ("native", 0.0)], - ), - ( - "histogram_fraction(-Inf, +Inf, mixed_histogram)", - vec![("classic", 1.0), ("native", 1.0)], - ), - ] { - let plan = PromPlanner::stmt_to_plan( - classic_and_native_histogram_table_provider("native", None, direct_or_histogram()), - &operator_eval_stmt(query), - &state, - ) - .await - .unwrap(); - let plan_text = plan.display_indent_schema().to_string(); - assert!(plan_text.contains("HistogramFold:"), "{plan_text}"); - assert!(plan_text.contains("prom_native_histogram_"), "{plan_text}"); - let value_field = plan - .schema() - .fields() - .iter() - .find(|field| field.data_type() == &ArrowDataType::Float64) - .unwrap() - .name() - .clone(); - - let (_, batches) = execute(plan, &state).await; - let mut actual = batches - .iter() - .flat_map(|batch| { - let tags = batch - .column_by_name("tag") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - let values = batch - .column_by_name(&value_field) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - (0..batch.num_rows()).map(|row| (tags.value(row), values.value(row))) - }) - .collect::>(); - actual.sort_by_key(|(tag, _)| *tag); - assert_eq!(actual, expected, "{query}"); - } - } - - #[tokio::test] - async fn mixed_histogram_helpers_report_annotations() { - let state = build_query_engine_state(); - let mut native_histogram = direct_or_histogram(); - native_histogram.count = 2.0; - native_histogram.sum = f64::NAN; - for (native_tag, expected_rows, expected_warnings, expected_infos) in [ - ( - "classic", - 0, - vec!["vector contains a mix of classic and native histograms"], - vec![], - ), - ( - "native", - 2, - vec![], - vec!["input to histogram_quantile has NaN observations, result is skewed higher"], - ), - ] { - let collector = PromqlAnnotationCollector::default(); - let plan = PromPlanner::stmt_to_plan_with_annotations( - classic_and_native_histogram_table_provider( - native_tag, - None, - native_histogram.clone(), - ), - &operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"), - &state, - Some(collector.clone()), - ) - .await - .unwrap(); - - let (_, batches) = execute(plan, &state).await; - assert_eq!( - batches.iter().map(RecordBatch::num_rows).sum::(), - expected_rows - ); - let mut warnings = vec![]; - let mut infos = vec![]; - collector.append_to(&mut warnings, &mut infos); - assert_eq!(warnings, expected_warnings); - assert_eq!(infos, expected_infos); - } - } - - #[tokio::test] - async fn mixed_histogram_helper_preserves_native_le_and_scans_once() { - let state = build_query_engine_state(); - let mut stmt = operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"); - stmt.end = UNIX_EPOCH.checked_add(Duration::from_secs(2)).unwrap(); - let plan = PromPlanner::stmt_to_plan( - classic_and_native_histogram_table_provider( - "classic", - Some("native"), - direct_or_histogram(), - ), - &stmt, - &state, - ) - .await - .unwrap(); - let plan_text = plan.display_indent_schema().to_string(); - assert_eq!( - plan_text.matches("TableScan: mixed_histogram").count(), - 1, - "{plan_text}" - ); - - let value_field = plan - .schema() - .fields() - .iter() - .find(|field| field.data_type() == &ArrowDataType::Float64) - .unwrap() - .name() - .clone(); - let (_, batches) = execute(plan, &state).await; - let mut actual = batches - .iter() - .flat_map(|batch| { - let le = batch - .column_by_name(LE_COLUMN_NAME) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - let timestamps = batch - .column_by_name("timestamp") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - let values = batch - .column_by_name(&value_field) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - (0..batch.num_rows()).map(|row| { - ( - timestamps.value(row), - (!le.is_null(row)).then(|| le.value(row).to_string()), - values.value(row), - ) - }) - }) - .collect::>(); - actual.sort_by(|lhs, rhs| (lhs.0, &lhs.1).cmp(&(rhs.0, &rhs.1))); - assert_eq!( - actual, - vec![ - (1_000, None, 1.0), - (1_000, Some("native".to_string()), 0.0), - (2_000, None, 1.0), - (2_000, Some("native".to_string()), 0.0), - ] - ); - } - - #[tokio::test] - async fn nested_histogram_helpers_ignore_unparsable_bucket_labels() { - let state = build_query_engine_state(); - for native_le in [None, Some("native")] { - for query in [ - "histogram_quantile(0.5, histogram_quantile(0.5, mixed_histogram))", - "histogram_fraction(-Inf, +Inf, histogram_fraction(-Inf, +Inf, mixed_histogram))", - ] { - let plan = PromPlanner::stmt_to_plan( - classic_and_native_histogram_table_provider( - "native", - native_le, - direct_or_histogram(), - ), - &operator_eval_stmt(query), - &state, - ) - .await - .unwrap(); - - let (_, batches) = execute(plan, &state).await; - assert_eq!( - batches.iter().map(RecordBatch::num_rows).sum::(), - 0, - "native_le={native_le:?}, query={query}" - ); - } - } - } - - #[tokio::test] - async fn native_histogram_quantile_rejects_multi_field_input() { - let table_provider = build_test_multi_histogram_table_provider("some_metric").await; - let result = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt("histogram_quantile(0.9, some_metric)"), - &build_query_engine_state(), - ) - .await; - - let err = result.expect_err("histogram_quantile on two native histogram fields must fail"); - assert!( - err.to_string() - .contains("Multi fields calculation is not supported in histogram_quantile"), - "{err}" - ); - } - - #[tokio::test] - async fn native_histogram_topk_uses_drop_udf() { - let plan = native_histogram_plan("topk(1, some_metric)").await; - - assert!(plan.contains("prom_native_histogram_drop_float"), "{plan}"); - assert!( - plan.contains("Filter: prom_native_histogram_drop_float") - && plan.contains("IS NOT NULL"), - "{plan}" - ); - } - - #[tokio::test] - async fn mixed_or_topk_bottomk_ignore_native_histograms() { - for op in ["topk", "bottomk"] { - let collector = PromqlAnnotationCollector::default(); - let state = build_query_engine_state(); - let plan = PromPlanner::stmt_to_plan_with_annotations( - operator_table_provider(), - &operator_eval_stmt(&format!("{op}(1, lf or on(tag) lh)")), - &state, - Some(collector.clone()), - ) - .await - .unwrap(); - let float_field = plan - .schema() - .fields() - .iter() - .find(|field| field.data_type() == &ArrowDataType::Float64) - .unwrap() - .name() - .clone(); - assert!( - plan.schema() - .fields() - .iter() - .all(|field| field.data_type() != &PromPlanner::native_histogram_arrow_type()), - "{plan:?}" - ); - - let (_, batches) = execute(plan, &state).await; - assert_eq!(values(&batches, &float_field), vec![2.0], "{op}"); - let mut warnings = vec![]; - let mut infos = vec![]; - collector.append_to(&mut warnings, &mut infos); - assert!(warnings.is_empty()); - assert_eq!( - infos, - vec![format!( - "{op}: dropped native histogram samples because this aggregation is not supported for native histograms" - )] - ); - } - } - - #[tokio::test] - async fn native_histogram_scalar_is_ignored_before_scalar_calculate() { - let plan = native_histogram_plan("scalar(some_metric)").await; - - assert!(plan.contains("ScalarCalculate"), "{plan}"); - assert!(plan.contains("Filter: Boolean(false)"), "{plan}"); - assert!(!plan.contains("prom_native_histogram_drop"), "{plan}"); - } - - #[tokio::test] - async fn native_histogram_value_sort_is_empty_but_label_sort_preserves_samples() { - for function in ["sort", "sort_desc"] { - let plan = native_histogram_plan(&format!("{function}(some_metric)")).await; - - assert!(plan.contains("Float64(NULL) IS NOT NULL"), "{plan}"); - assert!( - !plan.contains(&format!("Sort: {}", greptime_native_histogram())), - "{plan}" - ); - assert!(!plan.contains("prom_native_histogram_drop"), "{plan}"); - } - - for (function, direction) in [("sort_by_label", "ASC"), ("sort_by_label_desc", "DESC")] { - let plan = native_histogram_plan(&format!("{function}(some_metric, \"tag_0\")")).await; - - assert!(plan.contains(&format!("tag_0 {direction}")), "{plan}"); - assert!(plan.contains(greptime_native_histogram()), "{plan}"); - assert!(!plan.contains("Float64(NULL) IS NOT NULL"), "{plan}"); - } - } - - #[tokio::test] - async fn unsupported_native_histogram_functions_use_drop_udf() { - for query in [ - "deriv(some_metric[5m])", - "min_over_time(some_metric[5m])", - "quantile_over_time(0.9, some_metric[5m])", - "predict_linear(some_metric[5m], 60)", - "round(some_metric)", - "abs(some_metric)", - ] { - let plan = native_histogram_plan(query).await; - - assert!( - plan.contains("prom_native_histogram_drop_float"), - "{query}\n{plan}" - ); - } - } - - #[tokio::test] - async fn native_histogram_absent_over_time_uses_native_udf() { - let plan = native_histogram_plan("absent_over_time(some_metric[5m])").await; - - assert!( - plan.contains("prom_native_histogram_absent_over_time"), - "{plan}" - ); - } - - #[tokio::test] - async fn native_histogram_all_function_arms_route_correctly() { - // Every native-histogram match arm in `create_function_expr` must route to the - // expected UDF when all field columns are native histograms. `holt_winters` shares - // the `double_exponential_smoothing` arm but is not registered in the promql - // parser (0.10), so it cannot be exercised through a query string. - let cases = [ - // Range functions routed to native histogram UDFs. - ( - "increase(some_metric[5m])", - "prom_native_histogram_increase", - ), - ("rate(some_metric[5m])", "prom_native_histogram_rate"), - ("delta(some_metric[5m])", "prom_native_histogram_delta"), - ("idelta(some_metric[5m])", "prom_native_histogram_idelta"), - ("irate(some_metric[5m])", "prom_native_histogram_irate"), - ("resets(some_metric[5m])", "prom_native_histogram_resets"), - ("changes(some_metric[5m])", "prom_native_histogram_changes"), - ( - "avg_over_time(some_metric[5m])", - "prom_native_histogram_avg_over_time", - ), - ( - "sum_over_time(some_metric[5m])", - "prom_native_histogram_sum_over_time", - ), - ( - "count_over_time(some_metric[5m])", - "prom_native_histogram_count_over_time", - ), - ( - "last_over_time(some_metric[5m])", - "prom_native_histogram_last_over_time", - ), - ( - "present_over_time(some_metric[5m])", - "prom_native_histogram_present_over_time", - ), - // Unsupported functions dropped with the float-null UDF. - ("deriv(some_metric[5m])", "prom_native_histogram_drop_float"), - ( - "min_over_time(some_metric[5m])", - "prom_native_histogram_drop_float", - ), - ( - "max_over_time(some_metric[5m])", - "prom_native_histogram_drop_float", - ), - ( - "stddev_over_time(some_metric[5m])", - "prom_native_histogram_drop_float", - ), - ( - "stdvar_over_time(some_metric[5m])", - "prom_native_histogram_drop_float", - ), - ( - "quantile_over_time(0.9, some_metric[5m])", - "prom_native_histogram_drop_float", - ), - ( - "predict_linear(some_metric[5m], 60)", - "prom_native_histogram_drop_float", - ), - ( - "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)", - "prom_native_histogram_drop_float", - ), - ("round(some_metric)", "prom_native_histogram_drop_float"), - ("rad(some_metric)", "prom_native_histogram_drop_float"), - ("deg(some_metric)", "prom_native_histogram_drop_float"), - ("sgn(some_metric)", "prom_native_histogram_drop_float"), - // Instant helper functions routed to native histogram UDFs. - ( - "histogram_count(some_metric)", - "prom_native_histogram_count", - ), - ("histogram_sum(some_metric)", "prom_native_histogram_sum"), - ("histogram_avg(some_metric)", "prom_native_histogram_avg"), - ( - "histogram_stddev(some_metric)", - "prom_native_histogram_stddev", - ), - ( - "histogram_stdvar(some_metric)", - "prom_native_histogram_stdvar", - ), - ( - "histogram_fraction(-2 + 1, 2 / 2, some_metric)", - "prom_native_histogram_fraction", - ), - ]; - - for (query, expected_udf) in cases { - let plan = native_histogram_plan(query).await; - assert!(plan.contains(expected_udf), "{query}\n{plan}"); - if query.starts_with("histogram_fraction") { - assert!(plan.contains("Float64(-1)"), "{query}\n{plan}"); - } - } - } - - #[tokio::test] - async fn mixed_native_histogram_ranges_use_coordinated_udfs() { - let dual_output = [ - "increase(some_metric[5m])", - "rate(some_metric[5m])", - "delta(some_metric[5m])", - "idelta(some_metric[5m])", - "irate(some_metric[5m])", - "avg_over_time(some_metric[5m])", - "sum_over_time(some_metric[5m])", - "last_over_time(some_metric[5m])", - ]; - let float_output = [ - "resets(some_metric[5m])", - "changes(some_metric[5m])", - "deriv(some_metric[5m])", - "min_over_time(some_metric[5m])", - "max_over_time(some_metric[5m])", - "count_over_time(some_metric[5m])", - "absent_over_time(some_metric[5m])", - "present_over_time(some_metric[5m])", - "stddev_over_time(some_metric[5m])", - "stdvar_over_time(some_metric[5m])", - "quantile_over_time(0.9, some_metric[5m])", - "predict_linear(some_metric[5m], 60)", - "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)", - ]; - - for query in dual_output.iter().chain(float_output.iter()) { - let plan = PromPlanner::stmt_to_plan( - build_test_mixed_native_histogram_table_provider("some_metric").await, - &build_eval_stmt(query), - &build_query_engine_state(), - ) - .await - .unwrap() - .display_indent_schema() - .to_string(); - assert!(plan.contains("prom_mixed_range_float"), "{query}\n{plan}"); - assert_eq!( - plan.contains("prom_mixed_range_histogram"), - dual_output.contains(query), - "{query}\n{plan}" - ); - } - - let plan = PromPlanner::stmt_to_plan( - build_test_mixed_native_histogram_table_provider("some_metric").await, - &build_eval_stmt("sum_over_time(rate(some_metric[5m])[10m:1m])"), - &build_query_engine_state(), - ) - .await - .unwrap() - .display_indent_schema() - .to_string(); - let expected = r#"Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_value, prom_mixed_range_histogram(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=["greptime_value", "greptime_native_histogram"] [timestamp:Timestamp(ms), greptime_value:Dictionary(Int64, Float64);N, greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, tag_0:Utf8, timestamp_range:Dictionary(Int64, Timestamp(ms))] - PromSeriesDivide: tags=["tag_0"] [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_value, prom_mixed_range_histogram(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - PromRangeManipulate: req range=[-540000..100000000], interval=[60000], eval range=[300000], time index=[timestamp], values=["greptime_native_histogram", "greptime_value"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, greptime_value:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))] - PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] - PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] - Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] - Filter: some_metric.timestamp >= TimestampMillisecond(-839999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] - TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]"#; - assert_eq!(plan, expected); - } - - #[tokio::test] - async fn mixed_native_histogram_rate_executes_real_ranges() { - let schema = Arc::new(ArrowSchema::new(vec![ - Field::new( - "timestamp", - ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), - false, - ), - Field::new(greptime_value(), ArrowDataType::Float64, true), - Field::new( - greptime_native_histogram(), - native_histogram_value_type().as_arrow_type(), - true, - ), - ])); - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(TimestampMillisecondArray::from(vec![1000, 2000, 3000])), - Arc::new(Float64Array::from(vec![Some(1.0), None, Some(3.0)])), - build_histogram_array(&[None, Some(direct_or_histogram()), None]), - ], - ) - .unwrap(); - let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); - let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None) - .unwrap() - .build() - .unwrap(); - let collector = PromqlAnnotationCollector::default(); - let mut planner = PromPlanner { - table_provider: build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], - &[], - ) - .await, - ctx: PromPlannerContext { - start: 3000, - end: 3000, - interval: 1000, - range: Some(3000), - time_index_column: Some("timestamp".to_string()), - field_columns: vec![ - greptime_native_histogram().to_string(), - greptime_value().to_string(), - ], - ..Default::default() - }, - promql_annotations: Some(collector.clone()), - }; - let input = LogicalPlan::Extension(Extension { - node: Arc::new( - RangeManipulate::new( - 3000, - 3000, - 1000, - 0, - 3000, - "timestamp".to_string(), - planner.ctx.field_columns.clone(), - input, - ) - .unwrap(), - ), - }); - let PromExpr::Call(call) = parser::parse("rate(mixed[3s])").unwrap() else { - unreachable!() - }; - let preserve_any_value = PromPlanner::field_columns_are_alternative_samples( - input.schema(), - &planner.ctx.field_columns, - ); - let state = build_query_engine_state(); - let (mut exprs, _) = planner - .create_function_expr(&call.func, vec![], input.schema(), &state) - .unwrap(); - exprs.insert(0, planner.create_time_index_column_expr().unwrap()); - let plan = LogicalPlanBuilder::from(input) - .project(exprs) - .unwrap() - .filter( - planner - .create_empty_values_filter_expr(preserve_any_value) - .unwrap(), - ) - .unwrap() - .build() - .unwrap(); - let (_, batches) = execute(plan, &state).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); - let mut warnings = Vec::new(); - collector.append_to(&mut warnings, &mut Vec::new()); - assert!( - warnings - .iter() - .any(|warning| warning.contains("mix of float and native histogram")) - ); - } - - #[tokio::test] - async fn native_histogram_mixed_field_table_behaves() { - // Exercise function planning after float and histogram samples have already been - // represented as alternative nullable fields. Histogram functions must select the - // histogram field without adding a NULL float field that would reject every row. - let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; - let plan = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt("histogram_count(some_metric)"), - &build_query_engine_state(), - ) - .await - .unwrap(); - let plan_str = plan.display_indent_schema().to_string(); - assert!( - plan_str.contains("prom_native_histogram_count"), - "{plan_str}" - ); - assert!(!plan_str.contains("Float64(NULL)"), "{plan_str}"); - assert!( - plan_str.contains("prom_native_histogram_count(greptime_native_histogram) IS NOT NULL"), - "{plan_str}" - ); - - // Value sorting keeps the float column and never sorts by the histogram column. - let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; - let plan = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt("sort(some_metric)"), - &build_query_engine_state(), - ) - .await - .unwrap(); - let plan_str = plan.display_indent_schema().to_string(); - assert!( - plan_str.contains("greptime_value ASC NULLS FIRST"), - "{plan_str}" - ); - assert!( - !plan_str.contains("greptime_native_histogram ASC"), - "{plan_str}" - ); - - // scalar() ignores histogram samples and evaluates only the float field. - let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; - let plan = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt("scalar(some_metric)"), - &build_query_engine_state(), - ) - .await - .unwrap(); - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("ScalarCalculate"), "{plan_str}"); - assert!( - plan_str.contains("greptime_value IS NOT NULL"), - "{plan_str}" - ); - - // Functions that preserve both alternative fields keep rows with either sample type. - let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; - let plan = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt(r#"label_replace(some_metric, "copied", "$1", "tag_0", "(.*)")"#), - &build_query_engine_state(), - ) - .await - .unwrap(); - let plan_str = plan.display_indent_schema().to_string(); - let filter = plan_str.lines().next().unwrap(); - assert!( - filter.starts_with("Filter: ") - && filter.contains("greptime_native_histogram IS NOT NULL") - && filter.contains(" OR ") - && filter.contains("greptime_value IS NOT NULL"), - "{plan_str}" - ); - } - - #[tokio::test] - async fn less_filter_on_value() { - let query = "some_metric < 1.2345"; - let expected = String::from( - "Filter: some_metric.field_0 < Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - } - - #[tokio::test] - async fn count_over_time() { - let query = "count_over_time(some_metric[5m])"; - let expected = String::from( - "Filter: prom_count_over_time(timestamp_range,field_0) IS NOT NULL [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ - \n Projection: some_metric.timestamp, prom_count_over_time(timestamp_range, field_0) AS prom_count_over_time(timestamp_range,field_0), some_metric.tag_0 [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ - \n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[300000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\ - \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.timestamp >= TimestampMillisecond(-299999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - } - - /// The outer `PromRangeManipulate` from a subquery must be preceded by - /// `Sort` + `PromSeriesDivide`. - #[tokio::test] - async fn count_over_time_subquery() { - let query = "count_over_time(some_metric[10m:1m])"; - let expected = String::from( - "Filter: prom_count_over_time(timestamp_range,field_0) IS NOT NULL [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ - \n Projection: some_metric.timestamp, prom_count_over_time(timestamp_range, field_0) AS prom_count_over_time(timestamp_range,field_0), some_metric.tag_0 [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ - \n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromInstantManipulate: range=[-540000..100000000], lookback=[1000], interval=[60000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.timestamp >= TimestampMillisecond(-540999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - indie_query_plan_compare(query, expected).await; - } - - #[tokio::test] - async fn test_hash_join() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let case = r#"http_server_requests_seconds_sum{uri="/accounts/login"} / ignoring(kubernetes_pod_name,kubernetes_namespace) http_server_requests_seconds_count{uri="/accounts/login"}"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "http_server_requests_seconds_sum".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "http_server_requests_seconds_count".to_string(), - ), - ], - &["uri", "kubernetes_namespace", "kubernetes_pod_name"], - ) - .await; - // Should be ok - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected = "Projection: http_server_requests_seconds_sum.uri, http_server_requests_seconds_count.greptime_timestamp, CAST(http_server_requests_seconds_sum.greptime_value AS Float64) / CAST(http_server_requests_seconds_count.greptime_value AS Float64) AS http_server_requests_seconds_sum.greptime_value / http_server_requests_seconds_count.greptime_value\ - \n Projection: http_server_requests_seconds_sum.uri, http_server_requests_seconds_sum.kubernetes_namespace, http_server_requests_seconds_sum.kubernetes_pod_name, http_server_requests_seconds_sum.greptime_timestamp, http_server_requests_seconds_sum.greptime_value, http_server_requests_seconds_count.uri, http_server_requests_seconds_count.kubernetes_namespace, http_server_requests_seconds_count.kubernetes_pod_name, http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_count.greptime_value\ - \n Filter: prom_assert_unique_match_group(__promql_match_group_count, http_server_requests_seconds_sum.uri)\ - \n WindowAggr: windowExpr=[[count(Int64(1)) PARTITION BY [http_server_requests_seconds_sum.uri, http_server_requests_seconds_sum.greptime_timestamp] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS __promql_match_group_count]]\ - \n Inner Join: http_server_requests_seconds_sum.greptime_timestamp = http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_sum.uri = http_server_requests_seconds_count.uri\ - \n SubqueryAlias: http_server_requests_seconds_sum\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\ - \n PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\ - \n Sort: http_server_requests_seconds_sum.uri ASC NULLS FIRST, http_server_requests_seconds_sum.kubernetes_namespace ASC NULLS FIRST, http_server_requests_seconds_sum.kubernetes_pod_name ASC NULLS FIRST, http_server_requests_seconds_sum.greptime_timestamp ASC NULLS FIRST\ - \n Filter: http_server_requests_seconds_sum.uri = Utf8(\"/accounts/login\") AND http_server_requests_seconds_sum.greptime_timestamp >= TimestampMillisecond(-999, None) AND http_server_requests_seconds_sum.greptime_timestamp <= TimestampMillisecond(100000000, None)\ - \n TableScan: http_server_requests_seconds_sum\ - \n SubqueryAlias: http_server_requests_seconds_count\ - \n Projection: http_server_requests_seconds_count.uri, http_server_requests_seconds_count.kubernetes_namespace, http_server_requests_seconds_count.kubernetes_pod_name, http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_count.greptime_value\ - \n Filter: prom_assert_unique_match_group(__promql_match_group_count, http_server_requests_seconds_count.uri)\ - \n WindowAggr: windowExpr=[[count(Int64(1)) PARTITION BY [http_server_requests_seconds_count.uri, http_server_requests_seconds_count.greptime_timestamp] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS __promql_match_group_count]]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\ - \n PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\ - \n Sort: http_server_requests_seconds_count.uri ASC NULLS FIRST, http_server_requests_seconds_count.kubernetes_namespace ASC NULLS FIRST, http_server_requests_seconds_count.kubernetes_pod_name ASC NULLS FIRST, http_server_requests_seconds_count.greptime_timestamp ASC NULLS FIRST\ - \n Filter: http_server_requests_seconds_count.uri = Utf8(\"/accounts/login\") AND http_server_requests_seconds_count.greptime_timestamp >= TimestampMillisecond(-999, None) AND http_server_requests_seconds_count.greptime_timestamp <= TimestampMillisecond(100000000, None)\ - \n TableScan: http_server_requests_seconds_count"; - assert_eq!(plan.to_string(), expected); - } - - #[tokio::test] - async fn test_nested_histogram_quantile() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let case = r#"label_replace(histogram_quantile(0.99, sum by(pod, le, path, code) (rate(greptime_servers_grpc_requests_elapsed_bucket{container="frontend"}[1m0s]))), "pod_new", "$1", "pod", "greptimedb-frontend-[0-9a-z]*-(.*)")"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[( - DEFAULT_SCHEMA_NAME.to_string(), - "greptime_servers_grpc_requests_elapsed_bucket".to_string(), - )], - &["pod", "le", "path", "code", "container"], - ) - .await; - // Should be ok - let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_histogram_quantile_binary_op() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - // Arithmetic applied to a histogram_quantile() result. Regression for #8144: - // HistogramFold used to drop the input column qualifiers, so the binary-op - // projection failed to resolve the qualified tag column. - let case = r#"histogram_quantile(0.5, sum by (le, pod) (rate(http_request_duration_seconds_bucket[5m]))) + 0"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[( - DEFAULT_SCHEMA_NAME.to_string(), - "http_request_duration_seconds_bucket".to_string(), - )], - &["pod", "le"], - ) - .await; - // Should plan without a "No field named ..." error. - let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_parse_and_operator() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let cases = [ - r#"count (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} ) and (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} )) / (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_capacity_bytes{namespace=~".+"} )) >= (80 / 100)) or vector (0)"#, - r#"count (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} ) unless (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} )) / (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_capacity_bytes{namespace=~".+"} )) >= (80 / 100)) or vector (0)"#, - ]; - - for case in cases { - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "kubelet_volume_stats_used_bytes".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "kubelet_volume_stats_capacity_bytes".to_string(), - ), - ], - &["namespace", "persistentvolumeclaim"], - ) - .await; - // Should be ok - let _ = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - } - } - - #[tokio::test] - async fn test_nested_binary_op() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let case = r#"sum(rate(nginx_ingress_controller_requests{job=~".*"}[2m])) - - ( - sum(rate(nginx_ingress_controller_requests{namespace=~".*"}[2m])) - or - vector(0) - )"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[( - DEFAULT_SCHEMA_NAME.to_string(), - "nginx_ingress_controller_requests".to_string(), - )], - &["namespace", "job"], - ) - .await; - // Should be ok - let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_parse_or_operator() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let case = r#" - sum(rate(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}[120s])) by (cluster_name,tenant_name) / - (sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) * 100) - or - 200 * sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) / - sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)"#; - - let table_provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())], - &["tenant_name", "cluster_name"], - ) - .await; - eval_stmt.expr = parser::parse(case).unwrap(); - let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let case = r#"sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) / - (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) + - sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) / - (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0 - or - sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) / - (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0 - or - sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) / - (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0"#; - let table_provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())], - &["tenant_name", "cluster_name"], - ) - .await; - eval_stmt.expr = parser::parse(case).unwrap(); - let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let case = r#"(sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) + - sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or - (sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or - (sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name))"#; - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "background_waitevent_cnt".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "foreground_waitevent_cnt".to_string(), - ), - ], - &["tenant_name", "cluster_name"], - ) - .await; - eval_stmt.expr = parser::parse(case).unwrap(); - let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let case = r#"avg(node_load1{cluster_name=~"cluster1"}) by (cluster_name,host_name) or max(container_cpu_load_average_10s{cluster_name=~"cluster1"}) by (cluster_name,host_name) * 100 / max(container_spec_cpu_quota{cluster_name=~"cluster1"}) by (cluster_name,host_name)"#; - let table_provider = build_test_table_provider_with_fields( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "node_load1".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "container_cpu_load_average_10s".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "container_spec_cpu_quota".to_string(), - ), - ], - &["cluster_name", "host_name"], - ) - .await; - eval_stmt.expr = parser::parse(case).unwrap(); - let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - } - - #[tokio::test] - async fn value_matcher() { - // template - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let cases = [ - // single equal matcher - ( - r#"some_metric{__field__="field_1"}"#, - vec![ - "some_metric.field_1", - "some_metric.tag_0", - "some_metric.tag_1", - "some_metric.tag_2", - "some_metric.timestamp", - ], - ), - // two equal matchers - ( - r#"some_metric{__field__="field_1", __field__="field_0"}"#, - vec![ - "some_metric.field_0", - "some_metric.field_1", - "some_metric.tag_0", - "some_metric.tag_1", - "some_metric.tag_2", - "some_metric.timestamp", - ], - ), - // single not_eq matcher - ( - r#"some_metric{__field__!="field_1"}"#, - vec![ - "some_metric.field_0", - "some_metric.field_2", - "some_metric.tag_0", - "some_metric.tag_1", - "some_metric.tag_2", - "some_metric.timestamp", - ], - ), - // two not_eq matchers - ( - r#"some_metric{__field__!="field_1", __field__!="field_2"}"#, - vec![ - "some_metric.field_0", - "some_metric.tag_0", - "some_metric.tag_1", - "some_metric.tag_2", - "some_metric.timestamp", - ], - ), - // equal and not_eq matchers (no conflict) - ( - r#"some_metric{__field__="field_1", __field__!="field_0"}"#, - vec![ - "some_metric.field_1", - "some_metric.tag_0", - "some_metric.tag_1", - "some_metric.tag_2", - "some_metric.timestamp", - ], - ), - // equal and not_eq matchers (conflict) - ( - r#"some_metric{__field__="field_2", __field__!="field_2"}"#, - vec![ - "some_metric.tag_0", - "some_metric.tag_1", - "some_metric.tag_2", - "some_metric.timestamp", - ], - ), - // single regex eq matcher - ( - r#"some_metric{__field__=~"field_1|field_2"}"#, - vec![ - "some_metric.field_1", - "some_metric.field_2", - "some_metric.tag_0", - "some_metric.tag_1", - "some_metric.tag_2", - "some_metric.timestamp", - ], - ), - // single regex not_eq matcher - ( - r#"some_metric{__field__!~"field_1|field_2"}"#, - vec![ - "some_metric.field_0", - "some_metric.tag_0", - "some_metric.tag_1", - "some_metric.tag_2", - "some_metric.timestamp", - ], - ), - ]; - - for case in cases { - let prom_expr = parser::parse(case.0).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 3, - 3, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let mut fields = plan.schema().field_names(); - let mut expected = case.1.into_iter().map(String::from).collect::>(); - fields.sort(); - expected.sort(); - assert_eq!(fields, expected, "case: {:?}", case.0); - } - - let bad_cases = [ - r#"some_metric{__field__="nonexistent"}"#, - r#"some_metric{__field__!="nonexistent"}"#, - ]; - - for case in bad_cases { - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 3, - 3, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await; - assert!(plan.is_err(), "case: {:?}", case); - } - } - - #[tokio::test] - async fn custom_schema() { - let query = "some_alt_metric{__schema__=\"greptime_private\"}"; - let expected = String::from( - "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: greptime_private.some_alt_metric.tag_0 ASC NULLS FIRST, greptime_private.some_alt_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: greptime_private.some_alt_metric.timestamp >= TimestampMillisecond(-999, None) AND greptime_private.some_alt_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - - let query = "some_alt_metric{__database__=\"greptime_private\"}"; - let expected = String::from( - "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: greptime_private.some_alt_metric.tag_0 ASC NULLS FIRST, greptime_private.some_alt_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: greptime_private.some_alt_metric.timestamp >= TimestampMillisecond(-999, None) AND greptime_private.some_alt_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - - let query = "some_alt_metric{__schema__=\"greptime_private\"} / some_metric"; - let expected = String::from( - "Projection: some_metric.tag_0, some_metric.timestamp, CAST(greptime_private.some_alt_metric.field_0 AS Float64) / CAST(some_metric.field_0 AS Float64) AS greptime_private.some_alt_metric.field_0 / some_metric.field_0 [tag_0:Utf8, timestamp:Timestamp(ms), greptime_private.some_alt_metric.field_0 / some_metric.field_0:Float64;N]\ - \n Inner Join: greptime_private.some_alt_metric.tag_0 = some_metric.tag_0, greptime_private.some_alt_metric.timestamp = some_metric.timestamp [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n SubqueryAlias: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: greptime_private.some_alt_metric.tag_0 ASC NULLS FIRST, greptime_private.some_alt_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: greptime_private.some_alt_metric.timestamp >= TimestampMillisecond(-999, None) AND greptime_private.some_alt_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n SubqueryAlias: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - } - - #[tokio::test] - async fn only_equals_is_supported_for_special_matcher() { - let queries = &[ - "some_alt_metric{__schema__!=\"greptime_private\"}", - "some_alt_metric{__schema__=~\"lalala\"}", - "some_alt_metric{__database__!=\"greptime_private\"}", - "some_alt_metric{__database__=~\"lalala\"}", - ]; - - for query in queries { - let prom_expr = parser::parse(query).unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - "greptime_private".to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await; - assert!(plan.is_err(), "query: {:?}", query); - } - } - - #[tokio::test] - async fn native_scan_bounds_preserve_zero_lookback_and_overflow() { - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let mut planner = PromPlanner { - table_provider, - ctx: PromPlannerContext::from_eval_stmt(&build_eval_stmt("some_metric")), - promql_annotations: None, - }; - planner.ctx.time_index_column = Some("timestamp".to_string()); - planner.ctx.start = 1_000; - planner.ctx.lookback_delta = 0; - let schema = Arc::new( - DFSchema::try_from(ArrowSchema::new(vec![Field::new( - "timestamp", - ArrowDataType::Timestamp(ArrowTimeUnit::Nanosecond, None), - false, - )])) - .unwrap(), - ); - for (end, interval, windows) in [ - (1_000, 1_000, 1), - (2_000, 1_000, 1), - (7_201_000, 7_200_000, 2), - ] { - planner.ctx.end = end; - planner.ctx.interval = interval; - let filter = planner - .build_time_index_filter(0, &schema) - .unwrap() - .unwrap() - .to_string(); - assert_eq!(filter.matches(">=").count(), windows, "{filter}"); - assert!( - filter.contains("TimestampNanosecond(1000000000, None)"), - "{filter}" - ); - } - planner.ctx.end = i64::MAX; - let filter = planner - .build_time_index_filter(0, &schema) - .unwrap() - .unwrap() - .to_string(); - assert!( - filter.contains("timestamp >= TimestampNanosecond(1000000000, None)"), - "{filter}" - ); - - // A lookback subtraction can underflow milliseconds while the upper bound remains - // representable. Keep that upper bound so LastRow cannot select a future sample. - let ms_schema = Arc::new( - DFSchema::try_from(ArrowSchema::new(vec![Field::new( - "timestamp", - ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), - false, - )])) - .unwrap(), - ); - planner.ctx.start = i64::MIN + 100; - planner.ctx.end = planner.ctx.start; - planner.ctx.lookback_delta = 200; - let filter = planner - .build_time_index_filter(0, &ms_schema) - .unwrap() - .unwrap() - .to_string(); - assert_eq!( - filter, - format!( - "timestamp <= TimestampMillisecond({}, None)", - i64::MIN + 100 - ) - ); - - // The lower bound can also overflow while converting milliseconds to native nanoseconds. - // Its representable upper bound still has to reach the scan. - planner.ctx.start = 0; - planner.ctx.end = 0; - planner.ctx.lookback_delta = 300_000; - let filter = planner - .build_time_index_filter(9_223_372_036_854, &schema) - .unwrap() - .unwrap() - .to_string(); - assert_eq!( - filter, - "timestamp <= TimestampNanosecond(-9223372036854000000, None)" - ); - } - - #[tokio::test] - async fn test_non_ms_precision() { - let catalog_list = MemoryCatalogManager::with_default_setup(); - let columns = vec![ - ColumnSchema::new( - "tag".to_string(), - ConcreteDataType::string_datatype(), - false, - ), - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_nanosecond_datatype(), - false, - ) - .with_time_index(true), - ColumnSchema::new( - "field".to_string(), - ConcreteDataType::float64_datatype(), - true, - ), - ]; - let schema = Arc::new(Schema::new(columns)); - let table_meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices(vec![0]) - .value_indices(vec![2]) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .name("metrics".to_string()) - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: "metrics".to_string(), - table_id: 1024, - table, - }) - .is_ok() - ); - - let plan = PromPlanner::stmt_to_plan( - DfTableSourceProvider::new( - catalog_list.clone(), - false, - QueryContext::arc(), - DummyDecoder::arc(), - true, - ), - &EvalStmt { - expr: parser::parse("metrics{tag = \"1\"}").unwrap(), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }, - &build_query_engine_state(), - ) - .await - .unwrap(); - assert_eq!( - plan.display_indent_schema().to_string(), - "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\n PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp > TimestampNanosecond(-1000000000, None) AND metrics.timestamp <= TimestampNanosecond(100000000000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Projection: metrics.field, metrics.tag, metrics.timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]" - ); - let plan = PromPlanner::stmt_to_plan( - DfTableSourceProvider::new( - catalog_list.clone(), - false, - QueryContext::arc(), - DummyDecoder::arc(), - true, - ), - &EvalStmt { - expr: parser::parse("avg_over_time(metrics{tag = \"1\"}[5s])").unwrap(), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }, - &build_query_engine_state(), - ) - .await - .unwrap(); - assert_eq!( - plan.display_indent_schema().to_string(), - "Filter: prom_avg_over_time(timestamp_range,field) IS NOT NULL [timestamp:Timestamp(ms), prom_avg_over_time(timestamp_range,field):Float64;N, tag:Utf8]\n Projection: metrics.timestamp, prom_avg_over_time(timestamp_range, field) AS prom_avg_over_time(timestamp_range,field), metrics.tag [timestamp:Timestamp(ms), prom_avg_over_time(timestamp_range,field):Float64;N, tag:Utf8]\n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[5000], time index=[timestamp], values=[\"field\"] [field:Dictionary(Int64, Float64);N, tag:Utf8, timestamp:Timestamp(ms), timestamp_range:Dictionary(Int64, Timestamp(ms))]\n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp > TimestampNanosecond(-5000000000, None) AND metrics.timestamp <= TimestampNanosecond(100000000000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Projection: metrics.field, metrics.tag, metrics.timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]" - ); - } - - #[tokio::test] - async fn test_nonexistent_label() { - // template - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let case = r#"some_metric{nonexistent="hi"}"#; - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 3, - 3, - ) - .await; - // Should be ok - let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_label_join() { - let prom_expr = parser::parse( - "label_join(up{tag_0='api-server'}, 'foo', ',', 'tag_1', 'tag_2', 'tag_3')", - ) - .unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = - build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 4, 1) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let expected = r#" -Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8] - Projection: up.timestamp, up.field_0, concat_ws(Utf8(","), up.tag_1, up.tag_2, up.tag_3) AS foo, up.tag_0, up.tag_1, up.tag_2, up.tag_3 [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8] - PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] - PromSeriesDivide: tags=["tag_0", "tag_1", "tag_2", "tag_3"] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] - Sort: up.tag_0 ASC NULLS FIRST, up.tag_1 ASC NULLS FIRST, up.tag_2 ASC NULLS FIRST, up.tag_3 ASC NULLS FIRST, up.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] - Filter: up.tag_0 = Utf8("api-server") AND up.timestamp >= TimestampMillisecond(-999, None) AND up.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] - TableScan: up [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#; - - let ret = plan.display_indent_schema().to_string(); - assert_eq!(format!("\n{ret}"), expected, "\n{}", ret); - } - - #[tokio::test] - async fn test_label_replace() { - let prom_expr = parser::parse( - "label_replace(up{tag_0=\"a:c\"}, \"foo\", \"$1\", \"tag_0\", \"(.*):.*\")", - ) - .unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = - build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 1, 1) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let expected = r#" -Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8] - Projection: up.timestamp, up.field_0, regexp_replace(up.tag_0, Utf8("^(?s:(.*):.*)$"), Utf8("$1")) AS foo, up.tag_0 [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8] - PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] - PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] - Sort: up.tag_0 ASC NULLS FIRST, up.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] - Filter: up.tag_0 = Utf8("a:c") AND up.timestamp >= TimestampMillisecond(-999, None) AND up.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] - TableScan: up [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#; - - let ret = plan.display_indent_schema().to_string(); - assert_eq!(format!("\n{ret}"), expected, "\n{}", ret); - } - - #[tokio::test] - async fn label_replace_aggregation_queries_plan_successfully() { - let aggregate = - r#"sum by (foo) (label_replace(some_metric, "foo", "$1", "tag_0", "(.*)"))"#; - let queries = [ - aggregate.to_string(), - format!("{aggregate} <= 10"), - format!("{aggregate} * 0.8"), - format!("0.8 * {aggregate}"), - format!("{aggregate} <= {aggregate} * 0.8"), - ]; - let state = build_query_engine_state(); - let mut failures = Vec::new(); - - for query in queries { - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - if let Err(error) = - PromPlanner::stmt_to_plan(table_provider, &build_eval_stmt(&query), &state).await - { - failures.push(format!("{query}: {error:?}")); - } - } - - assert!(failures.is_empty(), "{}", failures.join("\n")); - } - - #[tokio::test] - async fn test_matchers_to_expr() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - let case = - r#"sum(prometheus_tsdb_head_series{tag_1=~"(10.0.160.237:8080|10.0.160.237:9090)"})"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider( - &[( - DEFAULT_SCHEMA_NAME.to_string(), - "prometheus_tsdb_head_series".to_string(), - )], - 3, - 3, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected = "Sort: prometheus_tsdb_head_series.timestamp ASC NULLS LAST [timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.field_0):Float64;N, sum(prometheus_tsdb_head_series.field_1):Float64;N, sum(prometheus_tsdb_head_series.field_2):Float64;N]\ - \n Aggregate: groupBy=[[prometheus_tsdb_head_series.timestamp]], aggr=[[sum(prometheus_tsdb_head_series.field_0), sum(prometheus_tsdb_head_series.field_1), sum(prometheus_tsdb_head_series.field_2)]] [timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.field_0):Float64;N, sum(prometheus_tsdb_head_series.field_1):Float64;N, sum(prometheus_tsdb_head_series.field_2):Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\", \"tag_2\"] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\ - \n Sort: prometheus_tsdb_head_series.tag_0 ASC NULLS FIRST, prometheus_tsdb_head_series.tag_1 ASC NULLS FIRST, prometheus_tsdb_head_series.tag_2 ASC NULLS FIRST, prometheus_tsdb_head_series.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\ - \n Filter: prometheus_tsdb_head_series.tag_1 ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\ - \n TableScan: prometheus_tsdb_head_series [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]"; - assert_eq!(plan.display_indent_schema().to_string(), expected); - } - - #[tokio::test] - async fn test_topk_expr() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - let case = r#"topk(10, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "prometheus_tsdb_head_series".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "http_server_requests_seconds_count".to_string(), - ), - ], - &["ip"], - ) - .await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected = "Projection: sum(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp [sum(prometheus_tsdb_head_series.greptime_value):Float64;N, ip:Utf8, greptime_timestamp:Timestamp(ms)]\ - \n Sort: prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\ - \n Filter: row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= Float64(10) [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\ - \n WindowAggr: windowExpr=[[row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\ - \n Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ - \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[sum(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"; - - assert_eq!(plan.display_indent_schema().to_string(), expected); - } - - #[tokio::test] - async fn test_count_values_expr() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "prometheus_tsdb_head_series".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "http_server_requests_seconds_count".to_string(), - ), - ], - &["ip"], - ) - .await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected = "Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N]\ - \n Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value) AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N]\ - \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value)]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), prom_float_to_string(prometheus_tsdb_head_series.greptime_value):Utf8;N, count(prometheus_tsdb_head_series.greptime_value):Int64]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"; - - assert_eq!(plan.display_indent_schema().to_string(), expected); - } - - #[tokio::test] - async fn test_value_alias() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - eval_stmt = QueryLanguageParser::apply_alias_extension(eval_stmt, "my_series"); - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "prometheus_tsdb_head_series".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "http_server_requests_seconds_count".to_string(), - ), - ], - &["ip"], - ) - .await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected = r#" -Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.series, prometheus_tsdb_head_series.greptime_timestamp [my_series:Int64, ip:Utf8, series:Utf8;N, greptime_timestamp:Timestamp(ms)] - Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N] - Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value) AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N] - Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value)]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), prom_float_to_string(prometheus_tsdb_head_series.greptime_value):Utf8;N, count(prometheus_tsdb_head_series.greptime_value):Int64] - PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - PromSeriesDivide: tags=["ip"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - Filter: prometheus_tsdb_head_series.ip ~ Utf8("^(?:(10.0.160.237:8080|10.0.160.237:9090))$") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"#; - assert_eq!(format!("\n{}", plan.display_indent_schema()), expected); - } - - #[tokio::test] - async fn test_quantile_expr() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - let case = r#"quantile(0.3, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "prometheus_tsdb_head_series".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "http_server_requests_seconds_count".to_string(), - ), - ], - &["ip"], - ) - .await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected = "Sort: prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [greptime_timestamp:Timestamp(ms), quantile(Float64(0.3),sum(prometheus_tsdb_head_series.greptime_value)):Float64;N]\ - \n Aggregate: groupBy=[[prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[quantile(Float64(0.3), sum(prometheus_tsdb_head_series.greptime_value))]] [greptime_timestamp:Timestamp(ms), quantile(Float64(0.3),sum(prometheus_tsdb_head_series.greptime_value)):Float64;N]\ - \n Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ - \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[sum(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"; - - assert_eq!(plan.display_indent_schema().to_string(), expected); - } - - #[tokio::test] - async fn test_or_not_exists_table_label() { - let state = build_query_engine_state(); - let provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string())], - &["job"], - ) - .await; - let raw = PromPlanner::stmt_to_plan( - provider, - &build_eval_stmt(r#"missing_metric or on(absent_label) normal_metric"#), - &state, - ) - .await - .unwrap(); - assert!( - raw.display_indent_schema() - .to_string() - .contains("__promql_or_match_0@") - ); - let (optimized, batches) = execute(raw, &state).await; - assert_no_internal_or_keys(optimized.schema()); - assert!(batches.iter().all(|batch| { - batch - .schema() - .fields() - .iter() - .all(|field| !field.name().starts_with("__promql_or_match_")) - })); - } - - #[tokio::test] - async fn test_histogram_quantile_missing_le_column() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - // Test case: histogram_quantile with a table that doesn't have 'le' column - let case = r#"histogram_quantile(0.99, sum by(pod,instance,le) (rate(non_existent_histogram_bucket{instance=~"xxx"}[1m])))"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - - // Create a table provider with a table that doesn't have 'le' column - let table_provider = build_test_table_provider_with_fields( - &[( - DEFAULT_SCHEMA_NAME.to_string(), - "non_existent_histogram_bucket".to_string(), - )], - &["pod", "instance"], // Note: no 'le' column - ) - .await; - - // Should return empty result instead of error - let result = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await; - - // This should succeed now (returning empty result) instead of failing with "Cannot find column le" - assert!( - result.is_ok(), - "Expected successful plan creation with empty result, but got error: {:?}", - result.err() - ); - - // Verify that the result is an EmptyRelation - let plan = result.unwrap(); - match plan { - LogicalPlan::EmptyRelation(_) => { - // This is what we expect - } - _ => panic!("Expected EmptyRelation, but got: {:?}", plan), - } - } - - #[tokio::test] - async fn test_direct_or_normalizes_missing_match_labels() { - type Case<'a> = ( - Option>, - Option>, - i64, - i64, - &'a [(f64, Option<&'a str>)], - ); - - let modifier = or_modifier("lhs or on(k) rhs"); - #[rustfmt::skip] - let cases: &[Case<'_>] = &[ - (None, None, 1, 1, &[(1.0, None)]), - (None, Some(Some("")), 1, 1, &[(1.0, None)]), - (Some(Some("")), None, 1, 1, &[(1.0, Some(""))]), - (None, Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]), - (Some(Some("l")), None, 1, 1, &[(1.0, Some("l")), (2.0, None)]), - (Some(None), Some(Some("")), 1, 1, &[(1.0, None)]), - (Some(None), Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]), - (Some(Some("same")), Some(Some("same")), 1, 2, &[(1.0, Some("same")), (2.0, Some("same"))]), - ]; - for &(left, right, left_ts, right_ts, expected) in cases { - let (optimized, batches) = run( - &matrix_source("lhs", left, left_ts, 1.0), - &matrix_source("rhs", right, right_ts, 2.0), - matrix_context("lhs", left), - matrix_context("rhs", right), - &modifier, - ) - .await; - assert_no_internal_or_keys(optimized.schema()); - assert_eq!( - rows(&batches), - expected - .iter() - .map(|(value, label)| (*value, label.map(str::to_string))) - .collect::>() - ); - } - } - - #[tokio::test] - async fn test_direct_or_match_modifiers() { - for (modifier, left, right, expected) in [ - (None, "left", "right", 2), - (or_modifier("lhs or on(k) rhs"), "same", "same", 1), - (or_modifier("lhs or on() rhs"), "left", "right", 1), - (or_modifier("lhs or ignoring(k) rhs"), "left", "right", 1), - ] { - let (_, batches) = run( - &matrix_source("lhs", Some(Some(left)), 1, 1.0), - &matrix_source("rhs", Some(Some(right)), 1, 2.0), - direct_or_context("lhs", &["job", "k"], "v"), - direct_or_context("rhs", &["job", "k"], "v"), - &modifier, - ) - .await; - assert_eq!( - batches.iter().map(RecordBatch::num_rows).sum::(), - expected - ); - } - } - - #[tokio::test] - async fn test_direct_or_nested_projection_uses_left_context() { - let left = matrix_source("lhs", Some(Some("k")), 1, 1.0); - let right = matrix_source("rhs", Some(Some("k")), 1, 2.0); - let raw = plan_direct_or( - scan(&left), - scan(&right), - direct_or_context("lhs", &["job", "k"], "v"), - direct_or_context("rhs", &["job", "k"], "v"), - &or_modifier("lhs or on(k) rhs"), - ) - .await; - assert!(raw.schema().iter().any(|(qualifier, field)| { - qualifier.as_ref().is_some_and(|q| q.to_string() == "lhs") && field.name() == "v" - })); - let nested = LogicalPlanBuilder::from(raw) - .project(vec![ - DfExpr::BinaryExpr(BinaryExpr { - left: Box::new(DfExpr::Column(Column::new( - Some(TableReference::bare("lhs")), - "v", - ))), - op: Operator::Plus, - right: Box::new(lit(1.0)), - }) - .alias("v_plus"), - ]) - .unwrap() - .build() - .unwrap(); - let (_, batches) = execute(nested, &build_query_engine_state()).await; - assert_eq!(values(&batches, "v_plus"), vec![2.0]); - } - - #[tokio::test] - async fn test_direct_or_skips_user_internal_key_name() { - const USER_TAG: &str = "__promql_or_match_0"; - let left = tagged_source( - "lhs", - false, - (USER_TAG, Some("left")), - DirectOrValue::Float64(1.0), - ); - let right = tagged_source( - "rhs", - false, - (USER_TAG, Some("right")), - DirectOrValue::Float64(2.0), - ); - let raw = plan_direct_or( - scan(&left), - scan(&right), - direct_or_context("lhs", &["job", USER_TAG], "v"), - direct_or_context("rhs", &["job", USER_TAG], "v"), - &or_modifier("lhs or on(missing_label) rhs"), - ) - .await; - assert!( - raw.display_indent_schema() - .to_string() - .contains("__promql_or_match_1@") - ); - let (_, batches) = execute(raw, &build_query_engine_state()).await; - assert!( - batches - .iter() - .all(|batch| batch.column_by_name(USER_TAG).is_some()) - ); - } - - #[tokio::test] - async fn test_direct_or_substrait_round_trip_with_normalized_key() { - let state = build_query_engine_state(); - let ctx = SessionContext::new_with_state(state.session_state()); - let catalog = Arc::new(MemoryCatalogProvider::new()); - catalog - .register_schema("public", Arc::new(MemorySchemaProvider::new())) - .unwrap(); - ctx.register_catalog("datafusion", catalog); - let left = matrix_source("lhs", Some(Some("")), 1, 1.0); - let right = matrix_source("rhs", None, 1, 2.0); - ctx.register_table( - TableReference::full("datafusion", "public", "lhs"), - table(&left), - ) - .unwrap(); - ctx.register_table( - TableReference::full("datafusion", "public", "rhs"), - table(&right), - ) - .unwrap(); - let raw = plan_direct_or( - ctx.table("datafusion.public.lhs") - .await - .unwrap() - .into_unoptimized_plan(), - ctx.table("datafusion.public.rhs") - .await - .unwrap() - .into_unoptimized_plan(), - direct_or_context("lhs", &["job", "k"], "v"), - direct_or_context("rhs", &["job"], "v"), - &or_modifier("lhs or on(k) rhs"), - ) - .await; - let decoded = DFLogicalSubstraitConvertor - .decode( - DFLogicalSubstraitConvertor - .encode(&raw, DefaultSerializer) - .unwrap(), - ctx.state(), - ) - .await - .unwrap(); - let (optimized, batches) = execute(decoded, &state).await; - assert_no_internal_or_keys(optimized.schema()); - assert!(batches.iter().all(|batch| { - batch - .schema() - .fields() - .iter() - .all(|field| !field.name().starts_with("__promql_or_match_")) - })); - assert_eq!(values(&batches, "v"), vec![1.0]); - } - - #[tokio::test] - async fn test_direct_or_numeric_value_types() { - let left = tagged_source("lhs", true, ("k", Some("lhs")), DirectOrValue::Int64(0)); - let right = tagged_source( - "rhs", - false, - ("k", Some("rhs")), - DirectOrValue::Float64(0.5), - ); - let (optimized, batches) = run( - &left, - &right, - direct_or_context("lhs", &["job", "k"], "v"), - direct_or_context("rhs", &["job", "k"], "v"), - &or_modifier("lhs or on(k) rhs"), - ) - .await; - assert_eq!( - optimized - .schema() - .field_with_name(None, "v") - .unwrap() - .data_type(), - &ArrowDataType::Float64 - ); - assert_eq!(values(&batches, "v"), vec![0.5]); - let provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], - &[], - ) - .await; - let mut planner = PromPlanner { - table_provider: provider, - ctx: PromPlannerContext::default(), - promql_annotations: None, - }; - let left_context = direct_or_context("lhs", &["job"], "v"); - let right_context = direct_or_context("rhs", &["job"], "v"); - let error = planner - .or_operator( - scan(&job_source("lhs", DirectOrValue::Utf8("x"))), - scan(&job_source("rhs", DirectOrValue::Float64(1.0))), - left_context.tag_columns.iter().cloned().collect(), - right_context.tag_columns.iter().cloned().collect(), - left_context, - right_context, - &or_modifier("lhs or on() rhs"), - ) - .unwrap_err(); - assert!( - error - .to_string() - .contains("OR value fields have incompatible types") - ); - } - - #[tokio::test] - async fn test_or_with_histogram_quantile_missing_le_column() { - let case = r#"histogram_quantile(0.99, non_existent_histogram_bucket) or normal_metric"#; - let eval_stmt = build_eval_stmt(case); - let table_provider = build_missing_le_or_normal_metric_table_provider().await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - assert_normal_metric_schema(&plan); - } - - #[tokio::test] - async fn test_or_with_right_empty_histogram_restores_left_context() { - let eval_stmt = build_eval_stmt( - r#"abs(sum by(instance) (normal_metric) or histogram_quantile(0.99, sum by(pod) (non_existent_histogram_bucket)))"#, - ); - let table_provider = build_missing_le_or_normal_metric_table_provider().await; - - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_or_with_both_empty_histograms() { - let eval_stmt = build_eval_stmt( - r#"histogram_quantile(0.99, sum by(pod) (left_histogram_bucket)) or histogram_quantile(0.99, sum by(instance) (right_histogram_bucket))"#, - ); - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "left_histogram_bucket".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "right_histogram_bucket".to_string(), - ), - ], - &["pod", "instance"], - ) - .await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - match plan { - LogicalPlan::EmptyRelation(relation) => { - assert!(!relation.produce_one_row); - assert!(!relation.schema.fields().is_empty()); - assert!( - relation - .schema - .fields() - .iter() - .any(|field| field.data_type() == &ArrowDataType::Float64) - ); - assert!( - relation - .schema - .fields() - .iter() - .any(|field| field.name() == "pod") - ); - assert!( - !relation - .schema - .fields() - .iter() - .any(|field| field.name() == "instance") - ); - } - _ => panic!("Expected EmptyRelation, but got: {plan:?}"), - } - } - - #[tokio::test] - async fn test_nested_or_with_both_empty_histograms() { - for case in [ - r#"abs(histogram_quantile(0.99, left_histogram_bucket) or histogram_quantile(0.99, right_histogram_bucket))"#, - r#"(histogram_quantile(0.99, left_histogram_bucket) or histogram_quantile(0.99, right_histogram_bucket)) + 1"#, - ] { - let eval_stmt = build_eval_stmt(case); - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "left_histogram_bucket".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "right_histogram_bucket".to_string(), - ), - ], - &["pod", "instance"], - ) - .await; - - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - } - } - - #[tokio::test] - async fn test_or_with_empty_histogram_modifiers() { - for case in [ - r#"histogram_quantile(0.99, non_existent_histogram_bucket) or on(pod) normal_metric"#, - r#"normal_metric or ignoring(instance) histogram_quantile(0.99, non_existent_histogram_bucket)"#, - ] { - let eval_stmt = build_eval_stmt(case); - let table_provider = build_missing_le_or_normal_metric_table_provider().await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - assert_normal_metric_schema(&plan); - } - } - - #[tokio::test] - async fn test_unless_preserves_left_context_for_histogram() { - let eval_stmt = build_eval_stmt( - r#"histogram_quantile(0.99, bucket_metric unless on(job) normal_metric) or fallback_metric"#, - ); - let state = build_query_engine_state(); - let plan = PromPlanner::stmt_to_plan( - build_set_op_context_table_provider().await, - &eval_stmt, - &state, - ) - .await - .unwrap(); - assert!(contains_histogram_fold(&plan), "{plan:?}"); - let (optimized, physical) = optimize_and_create_physical_plan(&state, plan).await; - assert!(contains_histogram_fold(&optimized), "{optimized:?}"); - let batches = - datafusion::physical_plan::collect(physical, state.session_state().task_ctx()) - .await - .unwrap(); - assert!(batches.iter().all(|batch| batch.num_rows() == 0)); - } - - #[tokio::test] - async fn test_and_preserves_left_context_for_histogram() { - let eval_stmt = build_eval_stmt( - r#"histogram_quantile(0.99, bucket_metric and on(job) normal_metric) or fallback_metric"#, - ); - let plan = PromPlanner::stmt_to_plan( - build_set_op_context_table_provider().await, - &eval_stmt, - &build_query_engine_state(), - ) - .await - .unwrap(); - assert!(contains_histogram_fold(&plan), "{plan:?}"); - } - - #[tokio::test] - async fn test_and_preserves_left_context_when_le_is_missing() { - let eval_stmt = - build_eval_stmt(r#"histogram_quantile(0.99, normal_metric and on(job) bucket_metric)"#); - let plan = PromPlanner::stmt_to_plan( - build_set_op_context_table_provider().await, - &eval_stmt, - &build_query_engine_state(), - ) - .await - .unwrap(); - assert!(matches!(&plan, LogicalPlan::EmptyRelation(_)), "{plan:?}"); - assert!(!plan.schema().fields().is_empty()); - assert!(!contains_histogram_fold(&plan), "{plan:?}"); - } - - async fn build_matching_filter_plan(query: &str) -> String { - let table_provider = build_test_table_provider_with_distinct_tags(&[ - ("metric_a", &["host", "device"]), - ("metric_b", &["host", "device"]), - ]) - .await; - let plan = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt(query), - &build_query_engine_state(), - ) - .await - .unwrap(); - plan.display_indent().to_string() - } - - /// [`build_test_table_provider_with_distinct_tags`] plus a `status` string column: a value - /// field that is neither a primary key nor a value column of the metric, so - /// `count by(status) (...)` still reports it among the aggregation's tag columns. - async fn build_test_table_provider_with_string_field( - table_tags: &[(&str, &[&str])], - ) -> DfTableSourceProvider { - let catalog_list = MemoryCatalogManager::with_default_setup(); - for (table_name, tags) in table_tags { - let mut columns = tags - .iter() - .map(|tag| { - ColumnSchema::new( - (*tag).to_string(), - ConcreteDataType::string_datatype(), - false, - ) - }) - .collect::>(); - columns.push( - ColumnSchema::new( - greptime_timestamp().to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ); - columns.push(ColumnSchema::new( - greptime_value().to_string(), - ConcreteDataType::float64_datatype(), - true, - )); - columns.push(ColumnSchema::new( - "status".to_string(), - ConcreteDataType::string_datatype(), - true, - )); - let table_meta = TableMetaBuilder::empty() - .schema(Arc::new(Schema::new(columns))) - .primary_key_indices((0..tags.len()).collect()) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .name((*table_name).to_string()) - .meta(table_meta) - .build() - .unwrap(); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: (*table_name).to_string(), - table_id: 1024, - table: EmptyTable::from_table_info(&table_info), - }) - .is_ok() - ); - } - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - async fn build_matching_filter_plan_with_string_field(query: &str) -> String { - let table_provider = build_test_table_provider_with_string_field(&[ - ("metric_a", &["host", "device"]), - ("metric_b", &["host", "device"]), - ]) - .await; - let plan = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt(query), - &build_query_engine_state(), - ) - .await - .unwrap(); - plan.display_indent().to_string() - } - - #[tokio::test] - async fn binary_matching_label_filter_reaches_both_operands() { - for query in [ - r#"metric_a / metric_b{host="foo"}"#, - r#"metric_a / on(host, device) metric_b{host="foo"}"#, - r#"count_over_time(metric_a[1m]) / on(host) count_over_time(metric_b{host="foo"}[1m])"#, - r#"metric_a / ignoring(device) metric_b{host="foo"}"#, - r#"sum by(host) (metric_a) / on(host) sum by(host) (metric_b{host="foo"})"#, - r#"metric_a / on(host) avg without(device) (metric_b{host="foo"})"#, - ] { - let plan = build_matching_filter_plan(query).await; - assert_eq!( - plan.matches(r#"host = Utf8("foo")"#).count(), - 2, - "{query}\n{plan}" - ); - } - } - - #[tokio::test] - async fn binary_matching_label_filter_reaches_scalar_ranking_and_grouped_operands() { - for query in [ - r#"(8 * metric_a{host="foo"}) / on(host) metric_b"#, - r#"topk(1, metric_a{host="foo"}) / on(host, device) metric_b"#, - r#"(8 * metric_a{host="foo"}) / on(host) group_left topk by(host)(1, max by(host)(metric_b))"#, - ] { - let plan = build_matching_filter_plan(query).await; - assert_eq!( - plan.matches(r#"host = Utf8("foo")"#).count(), - 2, - "{query}\n{plan}" - ); - } - // A global ranking one-side must see every host, so the matcher stays put. - let query = r#"metric_a{host="foo"} / on(host) group_left topk(1, max by(host)(metric_b))"#; - let plan = build_matching_filter_plan(query).await; - assert_eq!( - plan.matches(r#"host = Utf8("foo")"#).count(), - 1, - "{query}\n{plan}" - ); - } - - #[tokio::test] - async fn binary_matching_label_filter_skips_selecting_aggregations() { - // `topk` ranks its input, so filtering before it changes the candidate set. - let query = r#"topk(1, metric_a) / on(host, device) metric_b{host="foo"}"#; - let plan = build_matching_filter_plan(query).await; - assert_eq!(plan.matches("foo").count(), 1, "{query}\n{plan}"); - } - - #[tokio::test] - async fn binary_value_field_matcher_stays_on_its_own_operand() { - let value = greptime_value(); - for query in [ - format!(r#"metric_a / metric_b{{{value}="2"}}"#), - format!(r#"metric_a / on(host, device, {value}) metric_b{{{value}="2"}}"#), - ] { - let plan = build_matching_filter_plan(&query).await; - assert_eq!(plan.matches(r#"Utf8("2")"#).count(), 1, "{query}\n{plan}"); - } - } - - #[tokio::test] - async fn binary_value_field_matcher_is_not_copied_across_aggregations() { - // `status` varies between the samples of one series, so filtering the other operand by it - // would drop the newest sample before sample selection (#9242). - for query in [ - r#"count by(status) (metric_a) / on(status) count by(status) (metric_b{status="ready"})"#, - r#"(8 * count by(status)(metric_a{__field__="status"})) / on(status) topk by(status)(1, count by(status)(metric_b{__field__="status",status="ready"}))"#, - r#"count by(status)(metric_a{__field__="status"}) / on(status) group_left topk by(status)(1, count by(status)(metric_b{__field__="status",status="ready"}))"#, - ] { - let plan = build_matching_filter_plan_with_string_field(query).await; - assert_eq!( - plan.matches(r#"Utf8("ready")"#).count(), - 1, - "{query}\n{plan}" - ); - } - } - - #[tokio::test] - async fn binary_matching_label_filter_reaches_aggregations_grouping_by_tags() { - // `count`, not `sum`: the string value field is not summable. - let query = r#"count by(host) (metric_a) / on(host) count by(host) (metric_b{host="foo"})"#; - let plan = build_matching_filter_plan_with_string_field(query).await; - assert_eq!(plan.matches(r#"Utf8("foo")"#).count(), 2, "{query}\n{plan}"); - } - - #[tokio::test] - async fn binary_matching_label_filter_skips_unproven_expressions() { - for query in [ - r#"metric_a > on(host, device) metric_b{host="foo"}"#, - r#"metric_a / on(host) group_left metric_b{host="foo"}"#, - r#"metric_a / on(host) label_replace(metric_b{host="foo"},"extra","e","host",".*")"#, - r#"metric_a / on(device) metric_b{host="foo"}"#, - ] { - let plan = build_matching_filter_plan(query).await; - assert_eq!(plan.matches("foo").count(), 1, "{query}\n{plan}"); - } - } - - #[tokio::test] - async fn test_or_context_uses_left_qualified_output() { - let case = r#"(normal_metric or other_metric) + 1"#; - let eval_stmt = build_eval_stmt(case); - let state = build_query_engine_state(); - let plan = - PromPlanner::stmt_to_plan(build_or_context_table_provider().await, &eval_stmt, &state) - .await - .unwrap(); - assert!( - plan.schema() - .fields() - .iter() - .any(|field| field.data_type() == &ArrowDataType::Float64), - "{plan:?}" - ); - let (_optimized, _physical) = optimize_and_create_physical_plan(&state, plan).await; - } - - #[tokio::test] - async fn test_or_context_uses_left_qualified_empty_histogram_output() { - let case = r#"(abs(histogram_quantile(0.99, non_hist_metric)) or normal_metric) + 1"#; - let eval_stmt = build_eval_stmt(case); - let plan = PromPlanner::stmt_to_plan( - build_or_context_table_provider().await, - &eval_stmt, - &build_query_engine_state(), - ) - .await - .unwrap(); - assert!( - plan.schema() - .fields() - .iter() - .any(|field| field.data_type() == &ArrowDataType::Float64), - "{plan:?}" - ); - } - - #[tokio::test] - async fn test_direct_or_preserves_float_and_native_histogram_samples() { - for histogram_on_left in [false, true] { - let (planner, plan) = mixed_direct_or(histogram_on_left).await; - - let float_field = &planner.ctx.field_columns[0]; - let histogram_field = &planner.ctx.field_columns[1]; - assert!(float_field.starts_with(OR_FLOAT_FIELD_PREFIX)); - assert!(histogram_field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)); - assert_eq!( - plan.schema() - .field_with_name(None, float_field) - .unwrap() - .data_type(), - &ArrowDataType::Float64 - ); - assert_eq!( - plan.schema() - .field_with_name(None, histogram_field) - .unwrap() - .data_type(), - &native_histogram_value_type().as_arrow_type() - ); - - let (optimized, batches) = execute(plan, &build_query_engine_state()).await; - assert_no_internal_or_keys(optimized.schema()); - let mut sample_kinds = batches - .iter() - .flat_map(|batch| { - let values = batch.column_by_name(float_field).unwrap(); - let histograms = batch.column_by_name(histogram_field).unwrap(); - (0..batch.num_rows()) - .map(|row| (values.is_valid(row), histograms.is_valid(row))) - }) - .collect::>(); - sample_kinds.sort_unstable(); - assert_eq!(sample_kinds, vec![(false, true), (true, false)]); - } - } - - #[tokio::test] - async fn malformed_classic_bucket_does_not_drop_native_histogram() { - let state = build_query_engine_state(); - let collector = PromqlAnnotationCollector::default(); - let plan = PromPlanner::stmt_to_plan_with_annotations( - operator_table_provider(), - &operator_eval_stmt("histogram_quantile(0.5, bad_classic or bad_native)"), - &state, - Some(collector.clone()), - ) - .await - .unwrap(); - let value_field = plan - .schema() - .fields() - .iter() - .find(|field| field.data_type() == &ArrowDataType::Float64) - .unwrap() - .name() - .clone(); - - let (_, batches) = execute(plan, &state).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); - assert_eq!(values(&batches, &value_field), vec![0.0]); - let mut warnings = vec![]; - let mut infos = vec![]; - collector.append_to(&mut warnings, &mut infos); - assert!(warnings.is_empty()); - assert!(infos.is_empty()); - } - - #[tokio::test] - async fn test_mixed_binary_operator_aligns_both_alternative_inputs() { - let state = build_query_engine_state(); - let plan = PromPlanner::stmt_to_plan( - operator_table_provider(), - &operator_eval_stmt("(lf or on(tag) lh) * on(tag) (rf or on(tag) rh)"), - &state, - ) - .await - .unwrap(); - let plan_text = plan.display_indent_schema().to_string(); - assert!( - plan_text.contains("prom_native_histogram_mul_scalar"), - "{plan_text}" - ); - assert!( - plan_text.contains("prom_native_histogram_scalar_mul"), - "{plan_text}" - ); - let float_field = plan - .schema() - .fields() - .iter() - .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX)) - .unwrap() - .name() - .clone(); - let histogram_field = plan - .schema() - .fields() - .iter() - .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX)) - .unwrap() - .name() - .clone(); - - let (_, batches) = execute(plan, &state).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); - assert!(values(&batches, &float_field).is_empty()); - let mut sums = histograms(&batches, &histogram_field) - .into_iter() - .map(|histogram| histogram.sum) - .collect::>(); - sums.sort_by(f64::total_cmp); - assert_eq!(sums, vec![2.0, 3.0]); - } - - #[tokio::test] - async fn test_mixed_binary_operator_reports_only_dropped_samples() { - for (query, expected_rows, expected_infos) in [ - ("(lf or on(tag) lh) + on(tag) (rf or on(tag) rh)", 0, 1), - ("(lf or on(tag) lh) + on(tag) (lf or on(tag) lh)", 2, 0), - ("(lf or on(tag) lh) % on(tag) lh", 0, 1), - ] { - let state = build_query_engine_state(); - let annotations = PromqlAnnotationCollector::default(); - let plan = PromPlanner::stmt_to_plan_with_annotations( - operator_table_provider(), - &operator_eval_stmt(query), - &state, - Some(annotations.clone()), - ) - .await - .unwrap(); - - let (_, batches) = execute(plan, &state).await; - assert_eq!( - batches.iter().map(RecordBatch::num_rows).sum::(), - expected_rows, - "{query}" - ); - let mut warnings = vec![]; - let mut infos = vec![]; - annotations.append_to(&mut warnings, &mut infos); - assert!(warnings.is_empty(), "{query}: {warnings:?}"); - assert_eq!(infos.len(), expected_infos, "{query}: {infos:?}"); - } - } - - #[tokio::test] - async fn test_histogram_only_min_drops_empty_aggregate_group() { - // `min` over native-histogram-only input drops every sample in the group, so the - // NULL-valued aggregate row must be filtered out. Otherwise an outer expression - // like `group()` resurrects the group Prometheus considers unseen. - let state = build_query_engine_state(); - for query in ["min(lh)", "group(min(lh))"] { - let plan = PromPlanner::stmt_to_plan( - operator_table_provider(), - &operator_eval_stmt(query), - &state, - ) - .await - .unwrap(); - let (_, batches) = execute(plan, &state).await; - assert_eq!( - batches.iter().map(RecordBatch::num_rows).sum::(), - 0, - "{query}" - ); - } - } - - #[tokio::test] - async fn test_mixed_min_drops_histogram_only_group() { - // With alternative float/histogram fields, `min by (tag)` keeps float-only groups - // (tag=a from `lf`) and drops histogram-only groups (tag=b from `lh`) instead of - // emitting a NULL-valued row for them. - let state = build_query_engine_state(); - let plan = PromPlanner::stmt_to_plan( - operator_table_provider(), - &operator_eval_stmt("min by (tag) (lf or on(tag) lh)"), - &state, - ) - .await - .unwrap(); - let float_field = plan - .schema() - .fields() - .iter() - .find(|field| field.data_type() == &ArrowDataType::Float64) - .unwrap() - .name() - .clone(); - let (_, batches) = execute(plan, &state).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); - assert_eq!(values(&batches, &float_field), vec![2.0]); - } - - #[tokio::test] - async fn test_mixed_or_can_feed_another_or() { - let state = build_query_engine_state(); - let plan = PromPlanner::stmt_to_plan( - operator_table_provider(), - &operator_eval_stmt("lf or on(tag) lh or on(tag) fallback"), - &state, - ) - .await - .unwrap(); - let float_field = plan - .schema() - .fields() - .iter() - .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX)) - .unwrap() - .name() - .clone(); - let histogram_field = plan - .schema() - .fields() - .iter() - .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX)) - .unwrap() - .name() - .clone(); - - let (_, batches) = execute(plan, &state).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 3); - let mut float_values = values(&batches, &float_field); - float_values.sort_by(f64::total_cmp); - assert_eq!(float_values, vec![2.0, 7.0]); - assert_eq!(histograms(&batches, &histogram_field).len(), 1); - } - - #[tokio::test] - async fn test_mixed_fields_align_with_single_float_vector() { - let (planner, mixed) = mixed_direct_or(false).await; - let scale = tagged_source( - "scale", - false, - ("k", Some("float")), - DirectOrValue::Float64(2.0), - ); - let scale = scan(&scale); - let scale_fields = vec!["v".to_string()]; - let PromExpr::Binary(binary) = parser::parse("lhs * rhs").unwrap() else { - unreachable!() - }; - - let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns( - mixed.schema(), - scale.schema(), - &planner.ctx.field_columns, - &scale_fields, - binary.op, - false, - false, - ); - assert!(invalid_pairs.is_empty()); - assert_eq!( - groups - .iter() - .map(|(output, _)| output.clone()) - .collect::>(), - planner.ctx.field_columns - ); - assert_eq!(groups.len(), 2); - assert!( - groups - .iter() - .flat_map(|(_, pairs)| pairs) - .all(|(_, right)| *right == &scale_fields[0]) - ); - - let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns( - scale.schema(), - mixed.schema(), - &scale_fields, - &planner.ctx.field_columns, - binary.op, - false, - false, - ); - assert!(invalid_pairs.is_empty()); - assert_eq!( - groups - .iter() - .map(|(output, _)| output.clone()) - .collect::>(), - planner.ctx.field_columns - ); - assert_eq!(groups.len(), 2); - assert!( - groups - .iter() - .flat_map(|(_, pairs)| pairs) - .all(|(left, _)| *left == &scale_fields[0]) - ); - } - - #[tokio::test] - async fn test_non_bool_comparison_filters_mixed_sample_lanes() { - let (planner, input) = mixed_direct_or(false).await; - let input_schema = input.schema().clone(); - let plan = planner - .filter_on_field_column(input, |field| { - if PromPlanner::field_column_is_native_histogram(&input_schema, field) { - Ok(lit(false)) - } else { - Ok(col(field).gt(lit(0.0))) - } - }) - .unwrap(); - let float_field = planner.ctx.field_columns[0].clone(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - assert_eq!(values(&batches, &float_field), vec![1.25]); - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); - } - - #[tokio::test] - async fn test_mixed_left_and_unless_preserve_sample_lanes() { - for (expression, expected_sample_kind) in [ - ("lhs and on(k) mask", (false, true)), - ("lhs unless on(k) mask", (true, false)), - ] { - let (mut planner, left) = mixed_direct_or(false).await; - let left_context = planner.ctx.clone(); - let float_field = left_context.field_columns[0].clone(); - let histogram_field = left_context.field_columns[1].clone(); - let mask = tagged_source( - "mask", - false, - ("k", Some("histogram")), - DirectOrValue::Float64(1.0), - ); - let PromExpr::Binary(binary) = parser::parse(expression).unwrap() else { - unreachable!() - }; - let plan = planner - .set_op_on_non_field_columns( - left, - scan(&mask), - left_context, - direct_or_context("mask", &["job", "k"], "v"), - binary.op, - &binary.modifier, - ) - .unwrap(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - let sample_kinds = batches - .iter() - .flat_map(|batch| { - let floats = batch.column_by_name(&float_field).unwrap(); - let histograms = batch.column_by_name(&histogram_field).unwrap(); - (0..batch.num_rows()) - .map(|row| (floats.is_valid(row), histograms.is_valid(row))) - }) - .collect::>(); - assert_eq!(sample_kinds, vec![expected_sample_kind], "{expression}"); - } - } - - #[tokio::test] - async fn test_mixed_fields_arithmetic_broadcasts_computed_scalar() { - let plan = PromPlanner::stmt_to_plan( - build_test_mixed_native_histogram_table_provider("some_metric").await, - &build_eval_stmt("some_metric * scalar(vector(2))"), - &build_query_engine_state(), - ) - .await - .unwrap(); - let schema = plan.schema(); - assert_eq!( - schema - .field_with_unqualified_name(greptime_value()) - .unwrap() - .data_type(), - &ArrowDataType::Float64 - ); - assert_eq!( - schema - .field_with_unqualified_name(greptime_native_histogram()) - .unwrap() - .data_type(), - &native_histogram_value_type().as_arrow_type() - ); - assert!( - plan.display_indent_schema() - .to_string() - .contains("prom_native_histogram_mul_scalar"), - "{plan:?}" - ); - } - - #[tokio::test] - async fn test_unsupported_histogram_binary_does_not_block_or_fallback() { - let state = build_query_engine_state(); - let plan = PromPlanner::stmt_to_plan( - operator_table_provider(), - &operator_eval_stmt("((lf or on(tag) lh) % 2) or on(tag) lh"), - &state, - ) - .await - .unwrap(); - let float_field = plan - .schema() - .fields() - .iter() - .find(|field| field.data_type() == &ArrowDataType::Float64) - .unwrap() - .name() - .clone(); - let histogram_field = plan - .schema() - .fields() - .iter() - .find(|field| field.data_type() == &native_histogram_value_type().as_arrow_type()) - .unwrap() - .name() - .clone(); - - let (_, batches) = execute(plan, &state).await; - assert_eq!(values(&batches, &float_field), vec![0.0]); - assert_eq!(histograms(&batches, &histogram_field).len(), 1); - } - - #[tokio::test] - async fn test_unary_negates_mixed_float_and_native_histogram_samples() { - for histogram_on_left in [false, true] { - let (mut planner, input) = mixed_direct_or(histogram_on_left).await; - let plan = planner.negate_field_columns(input).unwrap(); - assert!(PromPlanner::field_columns_are_alternative_samples( - plan.schema(), - &planner.ctx.field_columns - )); - let float_field = planner - .ctx - .field_columns - .iter() - .find(|field| field.starts_with(OR_FLOAT_FIELD_PREFIX)) - .unwrap(); - let histogram_field = planner - .ctx - .field_columns - .iter() - .find(|field| field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)) - .unwrap(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - assert_eq!(values(&batches, float_field), vec![-1.25]); - let histogram = batches - .iter() - .find_map(|batch| { - let values = batch - .column_by_name(histogram_field) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - (0..values.len()).find_map(|row| { - common_query::native_histogram::read_histogram(values, row).unwrap() - }) - }) - .unwrap(); - assert_eq!(histogram.count, -1.0); - assert_eq!(histogram.sum, -1.0); - assert_eq!(histogram.reset_hint, CounterResetHint::Gauge); - } - } - - #[tokio::test] - async fn test_native_histogram_sum_and_avg_execute_real_batches() { - for op_name in ["sum", "avg"] { - for incompatible in [false, true] { - let mut second = direct_or_histogram(); - if incompatible { - second.schema = CUSTOM_BUCKETS_SCHEMA; - second.custom_values = vec![1.0]; - } - let collector = PromqlAnnotationCollector::default(); - let (mut planner, input) = - mixed_aggregate_input(vec![direct_or_histogram(), second]).await; - planner.promql_annotations = Some(collector.clone()); - let histogram_column = planner.ctx.field_columns[1].clone(); - planner.ctx.field_columns = vec![histogram_column.clone()]; - let input = LogicalPlanBuilder::from(input) - .project([col("ts"), col(&histogram_column)]) - .unwrap() - .build() - .unwrap(); - let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = - parser::parse(&format!("{op_name}(mixed)")).unwrap() - else { - unreachable!() - }; - let (aggregate_exprs, _) = - planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); - let plan = LogicalPlanBuilder::from(input) - .aggregate(vec![col("ts")], aggregate_exprs) - .unwrap() - .filter(planner.create_empty_values_filter_expr(false).unwrap()) - .unwrap() - .build() - .unwrap(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - let mut warnings = vec![]; - let mut infos = vec![]; - collector.append_to(&mut warnings, &mut infos); - assert!(infos.is_empty()); - if incompatible { - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); - assert!(warnings.iter().any(|warning| { - warning - == &format!( - "prom_native_histogram_agg_{op_name}: dropped native histogram aggregate with incompatible schemas" - ) - })); - } else { - let histograms = histograms(&batches, &histogram_column); - assert_eq!(histograms.len(), 1); - let expected = if op_name == "sum" { 2.0 } else { 1.0 }; - assert_eq!(histograms[0].count, expected); - assert_eq!(histograms[0].sum, expected); - assert!(warnings.is_empty()); - } - } - } - } - - #[tokio::test] - async fn test_canonical_mixed_count_group_and_count_values_execute() { - let state = build_query_engine_state(); - for (query, expected) in [ - ("count(some_metric)", vec![2.0]), - ("group(some_metric)", vec![1.0]), - (r#"count_values("sample", some_metric)"#, vec![1.0, 1.0]), - ] { - let plan = PromPlanner::stmt_to_plan( - build_test_mixed_native_histogram_table_provider("some_metric").await, - &operator_eval_stmt(query), - &state, - ) - .await - .unwrap(); - assert!( - plan.schema() - .fields() - .iter() - .all(|field| !field.name().starts_with("__promql_sample_count")), - "{query}: {plan:?}" - ); - let value_fields = plan - .schema() - .fields() - .iter() - .filter(|field| { - matches!( - field.data_type(), - ArrowDataType::Float64 | ArrowDataType::Int64 | ArrowDataType::UInt64 - ) || field.data_type() == &native_histogram_value_type().as_arrow_type() - }) - .collect::>(); - assert_eq!(value_fields.len(), 1, "{query}: {plan:?}"); - assert_ne!( - value_fields[0].data_type(), - &native_histogram_value_type().as_arrow_type(), - "{query}: {plan:?}" - ); - let value_column = value_fields[0].name().clone(); - - let (_, batches) = execute(plan, &state).await; - let mut actual = numeric_values(&batches, &value_column); - actual.sort_by(f64::total_cmp); - assert_eq!(actual, expected, "{query}"); - - if query.starts_with("count_values") { - let mut sample_labels = batches - .iter() - .flat_map(|batch| { - batch - .column_by_name("sample") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap() - .iter() - .flatten() - .map(str::to_string) - }) - .collect::>(); - sample_labels.sort(); - let mut expected_labels = - vec!["2".to_string(), direct_or_histogram().promql_string()]; - expected_labels.sort(); - assert_eq!(sample_labels, expected_labels); - } - } - } - - #[tokio::test] - async fn test_mixed_or_sum_aggregates_each_sample_type() { - let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = - parser::parse("sum(lhs)").unwrap() - else { - unreachable!() - }; - - let collector = PromqlAnnotationCollector::default(); - let (mut planner, input) = mixed_direct_or(false).await; - planner.promql_annotations = Some(collector.clone()); - let float_column = planner.ctx.field_columns[0].clone(); - let histogram_column = planner.ctx.field_columns[1].clone(); - let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); - let plan = LogicalPlanBuilder::from(input) - .aggregate(vec![col("ts"), col("k")], aggregate_exprs) - .unwrap() - .filter( - planner - .mixed_aggregate_filter_expr(op, &float_column, &histogram_column) - .unwrap(), - ) - .unwrap() - .project([ - col(&float_column), - col(&histogram_column), - col("ts"), - col("k"), - ]) - .unwrap() - .build() - .unwrap(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - assert_eq!(values(&batches, &float_column), vec![1.25]); - let histogram = batches - .iter() - .find_map(|batch| { - let values = batch - .column_by_name(&histogram_column)? - .as_any() - .downcast_ref::()?; - (0..values.len()).find_map(|row| { - common_query::native_histogram::read_histogram(values, row).unwrap() - }) - }) - .unwrap(); - assert_eq!(histogram.count, 1.0); - let mut warnings = vec![]; - let mut infos = vec![]; - collector.append_to(&mut warnings, &mut infos); - assert!(warnings.is_empty()); - - let collector = PromqlAnnotationCollector::default(); - let (mut planner, input) = mixed_direct_or(false).await; - planner.promql_annotations = Some(collector.clone()); - let float_column = planner.ctx.field_columns[0].clone(); - let histogram_column = planner.ctx.field_columns[1].clone(); - let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); - let plan = LogicalPlanBuilder::from(input) - .aggregate(vec![col("ts")], aggregate_exprs) - .unwrap() - .filter( - planner - .mixed_aggregate_filter_expr(op, &float_column, &histogram_column) - .unwrap(), - ) - .unwrap() - .project([col(&float_column), col(&histogram_column), col("ts")]) - .unwrap() - .build() - .unwrap(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); - let mut warnings = vec![]; - let mut infos = vec![]; - collector.append_to(&mut warnings, &mut infos); - assert_eq!( - warnings, - vec![ - "sum: dropped aggregation result containing both float and native histogram samples" - ] - ); - } - - #[tokio::test] - async fn test_mixed_or_sum_drops_incompatible_mixed_group() { - let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = - parser::parse("sum(lhs)").unwrap() - else { - unreachable!() - }; - let mut custom = direct_or_histogram(); - custom.schema = CUSTOM_BUCKETS_SCHEMA; - custom.custom_values = vec![1.0]; - let collector = PromqlAnnotationCollector::default(); - let (mut planner, input) = mixed_aggregate_input(vec![direct_or_histogram(), custom]).await; - planner.promql_annotations = Some(collector.clone()); - let float_column = planner.ctx.field_columns[0].clone(); - let histogram_column = planner.ctx.field_columns[1].clone(); - let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); - let plan = LogicalPlanBuilder::from(input) - .aggregate(vec![col("ts")], aggregate_exprs) - .unwrap() - .filter( - planner - .mixed_aggregate_filter_expr(op, &float_column, &histogram_column) - .unwrap(), - ) - .unwrap() - .project([col(&float_column), col(&histogram_column), col("ts")]) - .unwrap() - .build() - .unwrap(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); - let mut warnings = vec![]; - let mut infos = vec![]; - collector.append_to(&mut warnings, &mut infos); - assert!(warnings.iter().any(|warning| { - warning - == "sum: dropped aggregation result containing both float and native histogram samples" - })); - } - - #[tokio::test] - async fn test_mixed_or_min_records_only_present_histograms() { - let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = - parser::parse("min(lhs)").unwrap() - else { - unreachable!() - }; - let expected_info = "min: dropped native histogram samples because this aggregation is not supported for native histograms"; - - for (histograms, expected_infos) in [ - (vec![], vec![]), - (vec![direct_or_histogram()], vec![expected_info]), - ] { - let collector = PromqlAnnotationCollector::default(); - let (mut planner, input) = mixed_aggregate_input(histograms).await; - planner.promql_annotations = Some(collector.clone()); - let float_column = planner.ctx.field_columns[0].clone(); - let histogram_column = planner.ctx.field_columns[1].clone(); - let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); - let plan = LogicalPlanBuilder::from(input) - .aggregate(vec![col("ts")], aggregate_exprs) - .unwrap() - .filter( - planner - .mixed_ignored_histogram_filter_expr(op, &histogram_column) - .unwrap(), - ) - .unwrap() - .project([col(&float_column), col("ts")]) - .unwrap() - .build() - .unwrap(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - assert_eq!(values(&batches, &float_column), vec![1.25]); - let mut warnings = vec![]; - let mut infos = vec![]; - collector.append_to(&mut warnings, &mut infos); - assert!(warnings.is_empty()); - assert_eq!(infos, expected_infos); - } - } - - #[tokio::test] - async fn test_mixed_or_value_aliases_do_not_replace_labels() { - let left = source( - "lhs", - false, - 1, - vec![("job", Some("job")), ("k", Some("float"))], - DirectOrValue::Float64(1.0), - ); - let right = source( - "rhs", - false, - 1, - vec![ - ("job", Some("job")), - ("k", Some("histogram")), - (greptime_value(), Some("value-label")), - ], - DirectOrValue::NativeHistogram(direct_or_histogram()), - ); - let table_provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], - &[], - ) - .await; - let mut planner = PromPlanner { - table_provider, - ctx: PromPlannerContext::default(), - promql_annotations: None, - }; - let left = LogicalPlanBuilder::from(scan(&left)) - .project(vec![ - col("ts"), - col("job"), - col("k"), - col("v").alias(greptime_value()), - ]) - .unwrap() - .build() - .unwrap(); - let left_context = direct_or_context("lhs", &["job", "k"], greptime_value()); - let right_context = direct_or_context("rhs", &["job", "k", greptime_value()], "v"); - let plan = planner - .or_operator( - left, - scan(&right), - left_context.tag_columns.iter().cloned().collect(), - right_context.tag_columns.iter().cloned().collect(), - left_context, - right_context, - &or_modifier("lhs or on(k) rhs"), - ) - .unwrap(); - - assert_eq!( - plan.schema() - .field_with_name(None, greptime_value()) - .unwrap() - .data_type(), - &ArrowDataType::Utf8 - ); - assert!( - planner - .ctx - .field_columns - .iter() - .all(|field| { field != greptime_value() && field != greptime_native_histogram() }) - ); - assert!(PromPlanner::field_columns_are_alternative_samples( - plan.schema(), - &planner.ctx.field_columns - )); - let (_, batches) = execute(plan, &build_query_engine_state()).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); - let labels = batches - .iter() - .flat_map(|batch| { - batch - .column_by_name(greptime_value()) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap() - .iter() - .flatten() - }) - .collect::>(); - assert_eq!(labels, vec!["value-label"]); - } - - #[tokio::test] - async fn test_mixed_or_routes_float_histogram_and_label_functions() { - for (function, expected) in [("abs", 1.25), ("round", 1.0), ("histogram_count", 1.0)] { - let (mut planner, input) = mixed_direct_or(false).await; - let preserve_any_value = PromPlanner::field_columns_are_alternative_samples( - input.schema(), - &planner.ctx.field_columns, - ); - let PromExpr::Call(call) = parser::parse(&format!("{function}(lhs)")).unwrap() else { - unreachable!() - }; - let state = build_query_engine_state(); - let (mut exprs, _) = planner - .create_function_expr(&call.func, vec![], input.schema(), &state) - .unwrap(); - exprs.insert(0, planner.create_time_index_column_expr().unwrap()); - exprs.extend(planner.create_tag_column_exprs().unwrap()); - let plan = LogicalPlanBuilder::from(input) - .project(exprs) - .unwrap() - .filter( - planner - .create_empty_values_filter_expr(preserve_any_value) - .unwrap(), - ) - .unwrap() - .build() - .unwrap(); - let (_, batches) = execute(plan, &state).await; - let values = batches - .iter() - .flat_map(|batch| { - batch - .schema() - .fields() - .iter() - .position(|field| field.data_type() == &ArrowDataType::Float64) - .map(|index| { - batch - .column(index) - .as_any() - .downcast_ref::() - .unwrap() - .iter() - .flatten() - }) - .into_iter() - .flatten() - }) - .collect::>(); - assert_eq!(values, vec![expected], "{function}"); - } - - let (mut planner, input) = mixed_direct_or(false).await; - let preserve_any_value = PromPlanner::field_columns_are_alternative_samples( - input.schema(), - &planner.ctx.field_columns, - ); - let PromExpr::Call(call) = - parser::parse(r#"label_replace(lhs, "copy", "$1", "k", "(.*)")"#).unwrap() - else { - unreachable!() - }; - let args = planner.create_function_args(&call.args.args).unwrap(); - let state = build_query_engine_state(); - let (mut exprs, _) = planner - .create_function_expr(&call.func, args.literals, input.schema(), &state) - .unwrap(); - exprs.insert(0, planner.create_time_index_column_expr().unwrap()); - exprs.extend(planner.create_tag_column_exprs().unwrap()); - let plan = LogicalPlanBuilder::from(input) - .project(exprs) - .unwrap() - .filter( - planner - .create_empty_values_filter_expr(preserve_any_value) - .unwrap(), - ) - .unwrap() - .build() - .unwrap(); - let (_, batches) = execute(plan, &state).await; - let sample_count = batches.iter().map(RecordBatch::num_rows).sum::(); - assert_eq!(sample_count, 2); - } - - /// Table provider with a single metric `cv_metric` holding three series at `ts=1000`: - /// `k="a"` and `k="c"` carry the value 1.0, `k="b"` carries 2.0. - async fn build_count_values_table_provider() -> DfTableSourceProvider { - build_count_values_table_provider_with_values(&[1.0, 2.0, 1.0]).await - } - - /// Table provider with a single metric `cv_metric` holding one series per given value at - /// `ts=1000` (`k` names the series, the sample value is the given value). - async fn build_count_values_table_provider_with_values( - values: &[f64], - ) -> DfTableSourceProvider { - build_count_values_table_provider_with_value_array(Arc::new(Float64Array::from( - values.to_vec(), - ))) - .await - } - - /// Like [`build_count_values_table_provider_with_values`], but with a caller provided value - /// column, so tests can cover value columns that are not `Float64` (e.g. `BIGINT`). - async fn build_count_values_table_provider_with_value_array( - values: ArrayRef, - ) -> DfTableSourceProvider { - let value_data_type = ConcreteDataType::from_arrow_type(values.data_type()); - let catalog_list = MemoryCatalogManager::with_default_setup(); - let columns = vec![ - ColumnSchema::new("k".to_string(), ConcreteDataType::string_datatype(), false), - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ColumnSchema::new(greptime_value().to_string(), value_data_type, true), - ]; - let schema = Arc::new(Schema::new(columns)); - let table_meta = TableMetaBuilder::empty() - .schema(schema.clone()) - .primary_key_indices(vec![0]) - .value_indices(vec![2]) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = Arc::new( - TableInfoBuilder::default() - .table_id(3_001) - .name("cv_metric") - .meta(table_meta) - .build() - .unwrap(), - ); - let batch = RecordBatch::try_new( - schema.arrow_schema().clone(), - vec![ - Arc::new(StringArray::from( - (0..values.len()) - .map(|index| format!("k{index}")) - .collect::>(), - )), - Arc::new(TimestampMillisecondArray::from(vec![1_000; values.len()])), - values.clone(), - ], - ) - .unwrap(); - let backing = GreptimeMemTable::new_with_catalog( - "cv_metric", - GreptimeRecordBatch::from_df_record_batch(schema, batch), - 3_001, - DEFAULT_CATALOG_NAME.to_string(), - DEFAULT_SCHEMA_NAME.to_string(), - ); - let table = Arc::new(Table::new( - table_info, - FilterPushDownType::Unsupported, - backing.data_source(), - )); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: "cv_metric".to_string(), - table_id: 3_001, - table, - }) - .is_ok() - ); - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - /// Collects `(label, value)` pairs of a `count_values` result, where `label` is the - /// PromQL label generated by `count_values` and `value` is the aggregated sample value. - /// - /// The generated label holds the original sample value in PromQL's textual form, so it - /// is asserted as a string: comparing it as a number would not catch formatting bugs - /// (`1.0` instead of `1`, scientific notation, ...). - fn count_values_rows<'a>(batches: &'a [RecordBatch], label: &str) -> Vec<(&'a str, f64)> { - let mut rows = batches - .iter() - .flat_map(|batch| { - // The aggregated value is the only numeric column that is not the generated label. - let value_index = batch - .schema() - .fields() - .iter() - .position(|field| { - field.name() != label - && matches!( - field.data_type(), - ArrowDataType::Float64 - | ArrowDataType::Int64 - | ArrowDataType::UInt64 - ) - }) - .expect("no aggregated value column"); - let labels = batch - .column_by_name(label) - .expect("no generated label column") - .as_any() - .downcast_ref::() - .expect("the generated label must be a string column"); - let values = datafusion::arrow::compute::cast( - batch.column(value_index), - &ArrowDataType::Float64, - ) - .unwrap(); - let values = values.as_any().downcast_ref::().unwrap(); - labels - .iter() - .zip(values.iter()) - .map(|(label, value)| (label.unwrap(), value.unwrap())) - .collect::>() - }) - .collect::>(); - rows.sort_by(|left, right| left.0.cmp(right.0).then(left.1.total_cmp(&right.1))); - rows - } - - /// Asserts that a `count_values` result holds one sample per label set and evaluation - /// timestamp: Prometheus groups by the generated label, so a timestamp must never repeat - /// a label set (that would mean the samples were still grouped by the overwritten input - /// label). - fn assert_unique_label_set_per_timestamp(batches: &[RecordBatch], label: &str) { - let mut seen = HashMap::>::new(); - for batch in batches { - let timestamp_index = batch - .schema() - .fields() - .iter() - .position(|field| matches!(field.data_type(), ArrowDataType::Timestamp(..))) - .expect("no timestamp column"); - let timestamps = batch - .column(timestamp_index) - .as_any() - .downcast_ref::() - .expect("timestamp column is not a millisecond timestamp"); - let labels = batch - .column_by_name(label) - .expect("no generated label column") - .as_any() - .downcast_ref::() - .expect("the generated label must be a string column"); - for (timestamp, label) in timestamps.iter().zip(labels.iter()) { - let timestamp = timestamp.unwrap(); - let label = label.unwrap(); - assert!( - seen.entry(timestamp).or_default().insert(label.to_string()), - "duplicated label set `{label}` at timestamp {timestamp}" - ); - } - } - } - - #[tokio::test] - async fn test_count_values_generated_label_survives_enclosing_expr() { - // https://github.com/GreptimeTeam/greptimedb/issues/9181 - for (case, label) in [ - (r#"count_values("v", prometheus_tsdb_head_series)"#, "v"), - ( - r#"abs(count_values("v", prometheus_tsdb_head_series))"#, - "v", - ), - ( - r#"round(count_values("v", prometheus_tsdb_head_series))"#, - "v", - ), - (r#"count_values("v", prometheus_tsdb_head_series) + 1"#, "v"), - ( - r#"topk(1, count_values("v", prometheus_tsdb_head_series))"#, - "v", - ), - ( - r#"sum by (v) (count_values("v", prometheus_tsdb_head_series))"#, - "v", - ), - ( - r#"label_replace(count_values("v", prometheus_tsdb_head_series), "vcopy", "$1", "v", "(.*)")"#, - "v", - ), - ( - r#"count_values("v", prometheus_tsdb_head_series) by (ip) + 1"#, - "v", - ), - // The generated label overwrites an input label with the same name. - ( - r#"count_values("ip", prometheus_tsdb_head_series) by (ip)"#, - "ip", - ), - ( - r#"count_values("ip", prometheus_tsdb_head_series) by (ip) + 1"#, - "ip", - ), - ] { - let plan = PromPlanner::stmt_to_plan( - build_test_table_provider_with_fields( - &[( - DEFAULT_SCHEMA_NAME.to_string(), - "prometheus_tsdb_head_series".to_string(), - )], - &["ip"], - ) - .await, - &build_eval_stmt(case), - &build_query_engine_state(), - ) - .await - .unwrap(); - - let label_columns = plan - .schema() - .fields() - .iter() - .filter(|field| field.name() == label) - .count(); - assert_eq!( - label_columns, - 1, - "{case}: the `{label}` label must survive: {}", - plan.display_indent() - ); - } - } - - #[tokio::test] - async fn test_count_values_generated_label_in_enclosing_expr_execute() { - // https://github.com/GreptimeTeam/greptimedb/issues/9181 - let state = build_query_engine_state(); - // (query, generated label, expected `(label value, aggregated value)` pairs) - for (query, label, expected) in [ - ( - r#"count_values("v", cv_metric)"#, - "v", - vec![("1", 2.0), ("2", 1.0)], - ), - ( - r#"abs(count_values("v", cv_metric))"#, - "v", - vec![("1", 2.0), ("2", 1.0)], - ), - ( - r#"round(count_values("v", cv_metric))"#, - "v", - vec![("1", 2.0), ("2", 1.0)], - ), - ( - r#"count_values("v", cv_metric) + 1"#, - "v", - vec![("1", 3.0), ("2", 2.0)], - ), - ( - r#"sum by (v) (count_values("v", cv_metric))"#, - "v", - vec![("1", 2.0), ("2", 1.0)], - ), - ( - r#"topk(10, count_values("v", cv_metric))"#, - "v", - vec![("1", 2.0), ("2", 1.0)], - ), - // The generated label overwrites the input label with the same name, and the - // samples are grouped by the generated label only: `{k="1"}` holds the two - // samples of value `1.0` instead of one row per (overwritten label, value). - ( - r#"count_values("k", cv_metric) by (k)"#, - "k", - vec![("1", 2.0), ("2", 1.0)], - ), - ] { - let plan = PromPlanner::stmt_to_plan( - build_count_values_table_provider().await, - &operator_eval_stmt(query), - &state, - ) - .await - .unwrap_or_else(|err| panic!("{query}: {err}")); - - assert_eq!( - plan.schema() - .fields() - .iter() - .filter(|field| field.name() == label) - .count(), - 1, - "{query}: {}", - plan.display_indent() - ); - - let (_, batches) = execute(plan, &state).await; - assert_eq!(count_values_rows(&batches, label), expected, "{query}"); - assert_unique_label_set_per_timestamp(&batches, label); - } - } - - #[tokio::test] - async fn test_count_values_label_is_prometheus_formatted_value() { - // PromQL materializes the generated label with `strconv.FormatFloat(value, 'f', -1, 64)`: - // the shortest decimal form of the sample value without an exponent. The label is a - // label, so it must be a string column holding exactly that text: arrow's - // `Float64 -> Utf8` cast would render `1`/`200`/`1e21` as `1.0`/`200.0`/`1e21`. - let state = build_query_engine_state(); - // Samples are grouped by that formatted text, exactly like Prometheus groups by the - // generated label: `-0.0` and `0.0` are two series (`-0` and `0`), while values that - // round to the same text share one group. `0.0` formats as "0". - let plan = PromPlanner::stmt_to_plan( - build_count_values_table_provider_with_values(&[ - -0.0, 0.0, 1.0, 0.5, 200.0, 1e21, 1e-7, 2.5, - ]) - .await, - &operator_eval_stmt(r#"count_values("v", cv_metric)"#), - &state, - ) - .await - .unwrap(); - - let (_, batches) = execute(plan, &state).await; - assert_unique_label_set_per_timestamp(&batches, "v"); - let mut labels = count_values_rows(&batches, "v") - .into_iter() - .map(|(label, _)| label) - .collect::>(); - labels.sort(); - assert_eq!( - labels, - vec![ - "-0", - "0", - "0.0000001", - "0.5", - "1", - "1000000000000000000000", - "2.5", - "200", - ] - ); - } - - #[tokio::test] - async fn test_count_values_groups_by_formatted_value_for_bigint_input() { - // The grouping key of `count_values` is the formatted sample value, not the raw input - // value. Two `BIGINT` values that differ below the `Float64` precision (`2^53` and - // `2^53 + 1`) cast and format to the same label, so they must share one group and one - // count, exactly like Prometheus, which groups by the generated label text. - let state = build_query_engine_state(); - let plan = PromPlanner::stmt_to_plan( - build_count_values_table_provider_with_value_array(Arc::new(Int64Array::from(vec![ - 9_007_199_254_740_992_i64, - 9_007_199_254_740_993_i64, - ]))) - .await, - &operator_eval_stmt(r#"count_values("v", cv_metric)"#), - &state, - ) - .await - .unwrap(); - - let (_, batches) = execute(plan, &state).await; - // One timestamp must never carry the same label set twice. - assert_unique_label_set_per_timestamp(&batches, "v"); - assert_eq!( - count_values_rows(&batches, "v"), - vec![("9007199254740992", 2.0)] - ); - } -} +mod test; diff --git a/src/query/src/promql/planner/function_plans.rs b/src/query/src/promql/planner/function_plans.rs new file mode 100644 index 00000000000..68a254dd2bc --- /dev/null +++ b/src/query/src/promql/planner/function_plans.rs @@ -0,0 +1,473 @@ +// 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. + +//! Dedicated helper plans for PromQL functions that cannot be expressed as a +//! plain scalar function over the input plan: the histogram helpers, `vector()`, +//! `scalar()`, and `absent()`. + +use std::sync::Arc; + +use common_query::prelude::greptime_value; +use datafusion::functions_aggregate::expr_fn::first_value; +use datafusion::logical_expr::expr::ScalarFunction; +use datafusion::logical_expr::{Extension, LogicalPlan, LogicalPlanBuilder}; +use datafusion::optimizer::simplify_expressions::ExprSimplifier; +use datafusion::prelude::{Column, Expr as DfExpr}; +use datafusion::scalar::ScalarValue; +use datafusion_common::DFSchema; +use datafusion_expr::expr_fn::when; +use datafusion_expr::simplify::SimplifyContext; +use datafusion_expr::{col, lit}; +use datafusion_functions::core::coalesce; +use datatypes::arrow::datatypes::DataType as ArrowDataType; +use promql::extension_plan::{ + Absent, EmptyMetric, HistogramFold, HistogramFoldOperation, ScalarCalculate, +}; +use promql::functions::{NativeHistogramDrop, NativeHistogramFraction, NativeHistogramQuantile}; +use promql_parser::label::MatchOp; +use promql_parser::parser::{Expr as PromExpr, FunctionArgs as PromFunctionArgs}; +use snafu::{OptionExt, ResultExt, ensure}; + +use crate::promql::error::{ + DataFusionPlanningSnafu, FunctionInvalidArgumentSnafu, MultiFieldsNotSupportedSnafu, + PromqlPlanNodeSnafu, Result, TimeIndexNotFoundSnafu, ValueNotFoundSnafu, +}; +use crate::promql::planner::{ + LE_COLUMN_NAME, PromPlanner, SCALAR_FUNCTION, SPECIAL_ABSENT_FUNCTION, + SPECIAL_HISTOGRAM_FRACTION, SPECIAL_HISTOGRAM_QUANTILE, SPECIAL_TIME_FUNCTION, + SPECIAL_VECTOR_FUNCTION, +}; +use crate::query_engine::QueryEngineState; + +impl PromPlanner { + /// Create a classic, native, or mixed histogram helper plan. + pub(super) async fn create_histogram_plan( + &mut self, + function_name: &str, + args: &PromFunctionArgs, + query_engine_state: &QueryEngineState, + ) -> Result { + let float_literal = |param: &PromExpr| -> Result { + let value = (|| { + let expr = Self::get_param_as_literal_expr( + Some(param), + None, + Some(ArrowDataType::Float64), + ) + .ok()?; + let simplifier = ExprSimplifier::new(SimplifyContext::default()); + let expr = simplifier.coerce(expr, &DFSchema::empty()).ok()?; + let DfExpr::Literal(value, _) = simplifier.simplify(expr).ok()? else { + return None; + }; + let ScalarValue::Float64(Some(value)) = + value.cast_to(&ArrowDataType::Float64).ok()? + else { + return None; + }; + Some(value) + })() + .with_context(|| FunctionInvalidArgumentSnafu { + fn_name: function_name.to_string(), + })?; + Ok(value) + }; + let (function, input) = match (function_name, args.args.as_slice()) { + (SPECIAL_HISTOGRAM_QUANTILE, [quantile, input]) => ( + HistogramFoldOperation::Quantile(float_literal(quantile)?.into()), + input.as_ref().clone(), + ), + (SPECIAL_HISTOGRAM_FRACTION, [lower, upper, input]) => ( + HistogramFoldOperation::Fraction { + lower: float_literal(lower)?.into(), + upper: float_literal(upper)?.into(), + }, + input.as_ref().clone(), + ), + _ => { + return FunctionInvalidArgumentSnafu { + fn_name: function_name.to_string(), + } + .fail(); + } + }; + + let input_plan = self.prom_expr_to_plan(&input, query_engine_state).await?; + // Histogram helpers fold buckets across `le`, so `__tsid` (which includes `le`) is not a + // stable series identifier anymore. HistogramFold must not treat it as a label column. + let input_plan = self.strip_tsid_column(input_plan)?; + self.ctx.use_tsid = false; + + if let Some((float_field, histogram_field)) = + Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns) + .map(|(float, histogram)| (float.to_string(), histogram.to_string())) + { + if self.ctx.has_le_tag() { + return self.create_mixed_histogram_plan( + function, + input_plan, + float_field, + histogram_field, + ); + } + self.ctx.field_columns = vec![histogram_field]; + } + if self.all_field_columns_are_native_histograms(input_plan.schema()) { + return self.create_native_histogram_plan(function, input_plan); + } + + if !self.ctx.has_le_tag() { + // Return empty result instead of error when 'le' column is not found + // This handles the case when histogram metrics don't exist + return Ok(LogicalPlan::EmptyRelation( + datafusion::logical_expr::EmptyRelation { + produce_one_row: false, + schema: input_plan.schema().clone(), + }, + )); + } + let time_index_column = + self.ctx + .time_index_column + .clone() + .with_context(|| TimeIndexNotFoundSnafu { + table: self.ctx.table_name.clone().unwrap_or_default(), + })?; + // FIXME(ruihang): support multi fields + let field_column = self + .ctx + .field_columns + .first() + .with_context(|| FunctionInvalidArgumentSnafu { + fn_name: function.function_name().to_string(), + })? + .clone(); + // remove le column from tag columns + self.ctx.tag_columns.retain(|col| col != LE_COLUMN_NAME); + + let fold = HistogramFold::new_with_operation( + LE_COLUMN_NAME.to_string(), + field_column, + time_index_column, + function, + None, + input_plan, + ) + .context(DataFusionPlanningSnafu)?; + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(fold), + })) + } + + fn create_native_histogram_expr( + &self, + function: HistogramFoldOperation, + field_column: &str, + ) -> DfExpr { + let field = DfExpr::Column(Column::from_name(field_column)); + let (func, args) = match function { + HistogramFoldOperation::Quantile(quantile) => ( + Arc::new(NativeHistogramQuantile::scalar_udf_with_collector( + self.promql_annotations.clone(), + )), + vec![field, lit(f64::from(quantile))], + ), + HistogramFoldOperation::Fraction { lower, upper } => ( + Arc::new(NativeHistogramFraction::scalar_udf_with_collector( + self.promql_annotations.clone(), + )), + vec![field, lit(f64::from(lower)), lit(f64::from(upper))], + ), + }; + DfExpr::ScalarFunction(ScalarFunction { func, args }) + } + + fn create_native_histogram_plan( + &mut self, + function: HistogramFoldOperation, + input_plan: LogicalPlan, + ) -> Result { + ensure!( + self.ctx.field_columns.len() == 1, + MultiFieldsNotSupportedSnafu { + operator: function.function_name() + }, + ); + + let field_column = self.ctx.field_columns[0].clone(); + let function_expr = self.create_native_histogram_expr(function, &field_column); + let display_name = function_expr.schema_name().to_string(); + self.ctx.field_columns = vec![display_name.clone()]; + + let project_exprs = std::iter::once(self.create_time_index_column_expr()?) + .chain(std::iter::once(function_expr.alias(display_name))) + .chain(self.create_tag_column_exprs()?) + .collect::>(); + + LogicalPlanBuilder::from(input_plan) + .project(project_exprs) + .context(DataFusionPlanningSnafu)? + .filter(self.create_empty_values_filter_expr(false)?) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu) + } + + fn create_mixed_histogram_plan( + &mut self, + function: HistogramFoldOperation, + input_plan: LogicalPlan, + float_field: String, + histogram_field: String, + ) -> Result { + let time_index_column = + self.ctx + .time_index_column + .clone() + .with_context(|| TimeIndexNotFoundSnafu { + table: self.ctx.table_name.clone().unwrap_or_default(), + })?; + let tag_columns = self.ctx.tag_columns.clone(); + let folded = HistogramFold::new_with_operation( + LE_COLUMN_NAME.to_string(), + float_field.clone(), + time_index_column.clone(), + function, + Some(histogram_field.clone()), + input_plan, + ) + .context(DataFusionPlanningSnafu)?; + let record_collision = DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramDrop::warning_bool_false_udf( + "vector contains a mix of classic and native histograms".to_string(), + self.promql_annotations.clone(), + )), + args: vec![col(&float_field), col(&histogram_field)], + }); + let keep = when( + col(&float_field) + .is_not_null() + .and(col(&histogram_field).is_not_null()), + record_collision, + ) + .otherwise(lit(true)) + .context(DataFusionPlanningSnafu)?; + + let native_expr = self.create_native_histogram_expr(function, &histogram_field); + let output_field = native_expr.schema_name().to_string(); + let value = DfExpr::ScalarFunction(ScalarFunction { + func: coalesce(), + args: vec![col(&float_field), native_expr], + }); + self.ctx.field_columns = vec![output_field.clone()]; + LogicalPlanBuilder::from(LogicalPlan::Extension(Extension { + node: Arc::new(folded), + })) + .filter(keep) + .context(DataFusionPlanningSnafu)? + .project( + std::iter::once(col(&time_index_column)) + .chain(std::iter::once(value.alias(output_field))) + .chain(tag_columns.iter().map(col)), + ) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu) + } + + /// Create a [SPECIAL_VECTOR_FUNCTION] plan + pub(super) async fn create_vector_plan( + &mut self, + args: &PromFunctionArgs, + ) -> Result { + if args.args.len() != 1 { + return FunctionInvalidArgumentSnafu { + fn_name: SPECIAL_VECTOR_FUNCTION.to_string(), + } + .fail(); + } + let lit = Self::get_param_as_literal_expr(Some(args.args[0].as_ref()), None, None)?; + + // reuse `SPECIAL_TIME_FUNCTION` as name of time index column + self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string()); + self.ctx.reset_table_name_and_schema(); + self.ctx.tag_columns = vec![]; + self.ctx.aggregation_field_labels.clear(); + self.ctx.field_columns = vec![greptime_value().to_string()]; + Ok(LogicalPlan::Extension(Extension { + node: Arc::new( + EmptyMetric::new( + self.ctx.start, + self.ctx.end, + self.ctx.interval, + SPECIAL_TIME_FUNCTION.to_string(), + greptime_value().to_string(), + Some(lit), + ) + .context(DataFusionPlanningSnafu)?, + ), + })) + } + + /// Create a [SCALAR_FUNCTION] plan + pub(super) async fn create_scalar_plan( + &mut self, + args: &PromFunctionArgs, + query_engine_state: &QueryEngineState, + ) -> Result { + ensure!( + args.len() == 1, + FunctionInvalidArgumentSnafu { + fn_name: SCALAR_FUNCTION + } + ); + let input = self + .prom_expr_to_plan(&args.args[0], query_engine_state) + .await?; + let input_schema = input.schema().clone(); + let alternative_samples = + Self::field_columns_are_alternative_samples(&input_schema, &self.ctx.field_columns); + let histogram_fields = self + .ctx + .field_columns + .iter() + .filter(|field| Self::field_column_is_native_histogram(&input_schema, field)) + .count(); + ensure!( + self.ctx.field_columns.len() == 1 || alternative_samples, + MultiFieldsNotSupportedSnafu { + operator: SCALAR_FUNCTION + }, + ); + let scalar_field = self + .ctx + .field_columns + .iter() + .find(|field| !Self::field_column_is_native_histogram(&input_schema, field)) + .or_else(|| self.ctx.field_columns.first()) + .cloned() + .with_context(|| FunctionInvalidArgumentSnafu { + fn_name: SCALAR_FUNCTION, + })?; + let input = if histogram_fields == self.ctx.field_columns.len() { + // scalar() ignores histogram samples. An empty input makes ScalarCalculate emit NaN + // for every evaluation timestamp without attempting a Struct-to-Float64 cast. + LogicalPlanBuilder::from(input) + .filter(lit(false)) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)? + } else if histogram_fields > 0 { + // A mixed vector contributes only its float samples to scalar(). + LogicalPlanBuilder::from(input) + .filter(DfExpr::Column(Column::from_name(&scalar_field)).is_not_null()) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)? + } else { + input + }; + let scalar_plan = LogicalPlan::Extension(Extension { + node: Arc::new( + ScalarCalculate::new( + self.ctx.start, + self.ctx.end, + self.ctx.interval, + input, + self.ctx.time_index_column.as_ref().unwrap(), + &self.ctx.tag_columns, + &scalar_field, + self.ctx.table_name.as_deref(), + ) + .context(PromqlPlanNodeSnafu)?, + ), + }); + // scalar plan have no tag columns + self.ctx.tag_columns.clear(); + self.ctx.aggregation_field_labels.clear(); + self.ctx.field_columns.clear(); + self.ctx + .field_columns + .push(scalar_plan.schema().field(1).name().clone()); + Ok(scalar_plan) + } + + /// Create a [SPECIAL_ABSENT_FUNCTION] plan + pub(super) async fn create_absent_plan( + &mut self, + args: &PromFunctionArgs, + query_engine_state: &QueryEngineState, + ) -> Result { + if args.args.len() != 1 { + return FunctionInvalidArgumentSnafu { + fn_name: SPECIAL_ABSENT_FUNCTION.to_string(), + } + .fail(); + } + let input = self + .prom_expr_to_plan(&args.args[0], query_engine_state) + .await?; + + let time_index_expr = self.create_time_index_column_expr()?; + let first_field_expr = + self.create_field_column_exprs()? + .pop() + .with_context(|| ValueNotFoundSnafu { + table: self.ctx.table_name.clone().unwrap_or_default(), + })?; + let first_value_expr = first_value(first_field_expr, vec![]); + + let ordered_aggregated_input = LogicalPlanBuilder::from(input) + .aggregate( + vec![time_index_expr.clone()], + vec![first_value_expr.clone()], + ) + .context(DataFusionPlanningSnafu)? + .sort(vec![time_index_expr.sort(true, false)]) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + + let fake_labels = self + .ctx + .selector_matcher + .iter() + .filter_map(|matcher| match matcher.op { + MatchOp::Equal => Some((matcher.name.clone(), matcher.value.clone())), + _ => None, + }) + .collect::>(); + + // Create the absent plan + let absent_plan = LogicalPlan::Extension(Extension { + node: Arc::new( + Absent::try_new( + self.ctx.start, + self.ctx.end, + self.ctx.interval, + self.ctx.time_index_column.as_ref().unwrap().clone(), + self.ctx.field_columns[0].clone(), + fake_labels, + ordered_aggregated_input, + ) + .context(DataFusionPlanningSnafu)?, + ), + }); + + // The absent series carries the equality matchers as labels, not the input's + // tags or value fields, so the input's field grouping labels no longer apply. + self.ctx.aggregation_field_labels.clear(); + Ok(absent_plan) + } +} diff --git a/src/query/src/promql/planner/island.rs b/src/query/src/promql/planner/island.rs new file mode 100644 index 00000000000..d7c04f75015 --- /dev/null +++ b/src/query/src/promql/planner/island.rs @@ -0,0 +1,536 @@ +// 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. + +//! The binary-island fast path of the PromQL planner. + +use std::collections::{BTreeSet, HashMap}; + +use common_query::prelude::OTLP_AGGREGATION_TEMPORALITY_LABEL; +use datafusion::common::DFSchemaRef; +use datafusion::logical_expr::expr::Alias; +use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder}; +use datafusion::prelude::{Column, Expr as DfExpr, JoinType}; +use datafusion_common::{NullEquality, TableReference}; +use datafusion_expr::lit; +use promql_parser::label::{METRIC_NAME, MatchOp, Matcher}; +use promql_parser::parser::token::{self, TokenType}; +use promql_parser::parser::{ + BinaryExpr as PromBinaryExpr, Expr as PromExpr, Offset, ParenExpr, UnaryExpr, + VectorMatchCardinality, VectorSelector, +}; +use snafu::ResultExt; + +use crate::promql::error::{DataFusionPlanningSnafu, Result}; +use crate::promql::planner::{PromPlanner, PromPlannerContext}; + +/// Prefix for generated binary island leaf aliases. +const BINARY_ISLAND_LEAF_ALIAS_PREFIX: &str = "__prom_v"; + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct VectorLeafKey { + metric_name: String, + matchers: Vec<(String, String, String)>, + or_matchers: Vec>, + offset_ms: i128, + at: String, +} + +#[derive(Debug, Clone)] +struct IslandLeaf { + selector: VectorSelector, + display_table: String, +} + +#[derive(Debug, Clone)] +enum IslandExpr { + VectorLeaf(usize), + Scalar(DfExpr), + Unary { + input: Box, + }, + Binary { + op: TokenType, + lhs: Box, + rhs: Box, + }, +} + +impl IslandExpr { + fn try_new(expr: &PromExpr, env: &mut IslandCollectEnv) -> Option { + if let Some(expr) = PromPlanner::try_build_literal_expr(expr) { + return Some(Self::Scalar(expr)); + } + + match expr { + PromExpr::Paren(ParenExpr { expr }) => Self::try_new(expr, env), + PromExpr::VectorSelector(selector) => { + let leaf = env.intern_leaf(selector)?; + Some(Self::VectorLeaf(leaf)) + } + PromExpr::Unary(UnaryExpr { expr }) => { + let input = Self::try_new(expr, env)?; + Some(Self::Unary { + input: Box::new(input), + }) + } + PromExpr::Binary(PromBinaryExpr { + lhs, + rhs, + op, + modifier, + }) if matches!( + op.id(), + token::T_ADD + | token::T_SUB + | token::T_MUL + | token::T_DIV + | token::T_MOD + | token::T_POW + | token::T_ATAN2 + ) && modifier.as_ref().is_none_or(|modifier| { + !modifier.return_bool + && modifier.matching.is_none() + && matches!(modifier.card, VectorMatchCardinality::OneToOne) + && modifier.fill_values.lhs.is_none() + && modifier.fill_values.rhs.is_none() + }) => + { + let lhs = Self::try_new(lhs, env)?; + let rhs = Self::try_new(rhs, env)?; + Some(Self::Binary { + op: *op, + lhs: Box::new(lhs), + rhs: Box::new(rhs), + }) + } + _ => None, + } + } +} + +#[derive(Debug, Default)] +struct IslandCollectEnv { + leaf_by_key: HashMap, + leaves: Vec, + vector_occurrences: usize, +} + +#[derive(Debug)] +struct PlannedIslandLeaf { + plan: LogicalPlan, + ctx: PromPlannerContext, + alias: TableReference, + display_table: String, +} + +#[derive(Debug)] +struct IslandFieldExprs { + exprs: Vec, + names: Vec, + scalar: bool, +} + +impl VectorLeafKey { + fn from_selector(selector: &VectorSelector) -> Option { + let mut metric_name = selector.name.clone(); + let mut matchers = Vec::with_capacity(selector.matchers.matchers.len()); + let matcher_key = |matcher: &Matcher| { + ( + matcher.name.clone(), + matcher.op.to_string(), + matcher.value.clone(), + ) + }; + + for matcher in &selector.matchers.matchers { + if matcher.name == METRIC_NAME { + if matcher.op != MatchOp::Equal || metric_name.is_some() { + return None; + } + metric_name = Some(matcher.value.clone()); + } else { + matchers.push(matcher_key(matcher)); + } + } + matchers.sort(); + + let mut or_matchers = selector + .matchers + .or_matchers + .iter() + .map(|group| { + let mut group = group.iter().map(matcher_key).collect::>(); + group.sort(); + group + }) + .collect::>(); + or_matchers.sort(); + + Some(Self { + metric_name: metric_name?, + matchers, + or_matchers, + offset_ms: match &selector.offset { + Some(Offset::Pos(duration)) => duration.as_millis() as i128, + Some(Offset::Neg(duration)) => -(duration.as_millis() as i128), + None => 0, + }, + at: format!("{:?}", selector.at), + }) + } +} + +impl IslandCollectEnv { + fn intern_leaf(&mut self, selector: &VectorSelector) -> Option { + self.vector_occurrences += 1; + let key = VectorLeafKey::from_selector(selector)?; + if let Some(id) = self.leaf_by_key.get(&key) { + return Some(*id); + } + + let id = self.leaves.len(); + self.leaves.push(IslandLeaf { + selector: selector.clone(), + display_table: key.metric_name.clone(), + }); + self.leaf_by_key.insert(key, id); + Some(id) + } +} + +impl PromPlanner { + pub(super) async fn try_plan_binary_island( + &mut self, + binary_expr: &PromBinaryExpr, + ) -> Result> { + let original_ctx = self.ctx.clone(); + let mut collect_env = IslandCollectEnv::default(); + let Some(island_expr) = + IslandExpr::try_new(&PromExpr::Binary(binary_expr.clone()), &mut collect_env) + else { + return Ok(None); + }; + + if collect_env.leaves.is_empty() + || collect_env.vector_occurrences <= collect_env.leaves.len() + { + return Ok(None); + } + + let mut planned_leaves = Vec::with_capacity(collect_env.leaves.len()); + for (idx, leaf) in collect_env.leaves.iter().enumerate() { + let plan = self + .prom_vector_selector_to_plan(&leaf.selector, false) + .await?; + let ctx = self.ctx.clone(); + let alias = TableReference::bare(format!("{BINARY_ISLAND_LEAF_ALIAS_PREFIX}{idx}")); + let plan = LogicalPlanBuilder::from(plan) + .alias(alias.clone()) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + planned_leaves.push(PlannedIslandLeaf { + plan, + ctx, + alias, + display_table: leaf.display_table.clone(), + }); + } + + if planned_leaves.iter().any(|leaf| { + Self::field_columns_contain_native_histogram( + leaf.plan.schema(), + &leaf.ctx.field_columns, + ) + }) { + self.ctx = original_ctx; + return Ok(None); + } + + if !Self::binary_island_join_contexts_supported(&planned_leaves) { + self.ctx = original_ctx; + return Ok(None); + } + + let mut input = planned_leaves[0].plan.clone(); + for right_idx in 1..planned_leaves.len() { + input = self.join_binary_island_leaf( + input, + &planned_leaves[0], + &planned_leaves[right_idx], + )?; + } + + let field_exprs = + Self::build_binary_island_field_exprs(&island_expr, &planned_leaves, input.schema())?; + if field_exprs.scalar || field_exprs.exprs.is_empty() { + self.ctx = original_ctx; + return Ok(None); + } + + let plan = self.project_binary_island( + input, + &planned_leaves[0].alias, + &planned_leaves[0].ctx, + field_exprs, + )?; + Ok(Some(plan)) + } + + fn binary_island_join_contexts_supported(leaves: &[PlannedIslandLeaf]) -> bool { + if leaves + .iter() + .any(|leaf| leaf.ctx.time_index_column.is_none()) + { + return false; + } + + if leaves.len() <= 1 { + return true; + } + + let first_tags = leaves[0].ctx.tag_columns.iter().collect::>(); + + leaves.iter().skip(1).all(|leaf| { + (Self::plan_has_tsid_column(&leaves[0].plan) && Self::plan_has_tsid_column(&leaf.plan)) + || leaf.ctx.tag_columns.iter().collect::>() == first_tags + }) + } + + fn join_binary_island_leaf( + &self, + left: LogicalPlan, + first_leaf: &PlannedIslandLeaf, + right_leaf: &PlannedIslandLeaf, + ) -> Result { + let only_join_time_index = (first_leaf.ctx.tag_columns.is_empty() + || right_leaf.ctx.tag_columns.is_empty()) + && !first_leaf + .ctx + .tag_columns + .iter() + .chain(&right_leaf.ctx.tag_columns) + .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL); + let (mut left_keys, mut right_keys, force_empty_join) = self.binary_join_key_columns( + left.schema(), + right_leaf.plan.schema(), + &first_leaf.ctx, + &right_leaf.ctx, + only_join_time_index, + &None, + )?; + + if let (Some(left_time_index_column), Some(right_time_index_column)) = ( + first_leaf.ctx.time_index_column.clone(), + right_leaf.ctx.time_index_column.clone(), + ) { + left_keys.insert(left_time_index_column); + right_keys.insert(right_time_index_column); + } + + LogicalPlanBuilder::from(left) + .join_detailed( + right_leaf.plan.clone(), + JoinType::Inner, + ( + left_keys + .into_iter() + .map(|name| Column::new(Some(first_leaf.alias.clone()), name)) + .collect::>(), + right_keys + .into_iter() + .map(|name| Column::new(Some(right_leaf.alias.clone()), name)) + .collect::>(), + ), + force_empty_join.then_some(lit(false)), + NullEquality::NullEqualsNull, + ) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu) + } + + fn build_binary_island_field_exprs( + expr: &IslandExpr, + leaves: &[PlannedIslandLeaf], + schema: &DFSchemaRef, + ) -> Result { + match expr { + IslandExpr::VectorLeaf(id) => { + let leaf = &leaves[*id]; + let exprs = leaf + .ctx + .field_columns + .iter() + .map(|field| { + schema + .qualified_field_with_name(Some(&leaf.alias), field) + .context(DataFusionPlanningSnafu) + .map(|field| DfExpr::Column(field.into())) + }) + .collect::>>()?; + let names = leaf + .ctx + .field_columns + .iter() + .map(|field| format!("{}.{}", leaf.display_table, field)) + .collect(); + Ok(IslandFieldExprs { + exprs, + names, + scalar: false, + }) + } + IslandExpr::Scalar(expr) => Ok(IslandFieldExprs { + exprs: vec![expr.clone()], + names: vec![expr.schema_name().to_string()], + scalar: true, + }), + IslandExpr::Unary { input } => { + let input = Self::build_binary_island_field_exprs(input, leaves, schema)?; + let mut exprs = Vec::with_capacity(input.exprs.len()); + let mut names = Vec::with_capacity(input.names.len()); + for (expr, name) in input.exprs.into_iter().zip(input.names) { + exprs.push(DfExpr::Negative(Box::new(expr))); + names.push(format!("-{name}")); + } + Ok(IslandFieldExprs { + exprs, + names, + scalar: input.scalar, + }) + } + IslandExpr::Binary { op, lhs, rhs } => { + let same_leaf = match (&**lhs, &**rhs) { + (IslandExpr::VectorLeaf(left), IslandExpr::VectorLeaf(right)) + if left == right => + { + Some(*left) + } + _ => None, + }; + let lhs = Self::build_binary_island_field_exprs(lhs, leaves, schema)?; + let rhs = Self::build_binary_island_field_exprs(rhs, leaves, schema)?; + let expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; + let scalar = lhs.scalar && rhs.scalar; + let op = op.to_string(); + + let (exprs, names) = match (lhs.scalar, rhs.scalar) { + (true, true) => { + let expr = expr_builder(lhs.exprs[0].clone(), rhs.exprs[0].clone())?; + let name = format!("{} {op} {}", lhs.names[0], rhs.names[0]); + (vec![expr], vec![name]) + } + (true, false) => { + let mut exprs = Vec::with_capacity(rhs.exprs.len()); + let mut names = Vec::with_capacity(rhs.names.len()); + for (rhs_expr, rhs_name) in rhs.exprs.into_iter().zip(rhs.names) { + exprs.push(expr_builder(lhs.exprs[0].clone(), rhs_expr)?); + names.push(format!("{} {op} {rhs_name}", lhs.names[0])); + } + (exprs, names) + } + (false, true) => { + let mut exprs = Vec::with_capacity(lhs.exprs.len()); + let mut names = Vec::with_capacity(lhs.names.len()); + for (lhs_expr, lhs_name) in lhs.exprs.into_iter().zip(lhs.names) { + exprs.push(expr_builder(lhs_expr, rhs.exprs[0].clone())?); + names.push(format!("{lhs_name} {op} {}", rhs.names[0])); + } + (exprs, names) + } + (false, false) => { + let mut exprs = Vec::new(); + let mut names = Vec::new(); + for (idx, ((lhs_expr, rhs_expr), (mut lhs_name, mut rhs_name))) in lhs + .exprs + .into_iter() + .zip(rhs.exprs) + .zip(lhs.names.into_iter().zip(rhs.names)) + .enumerate() + { + if let Some(leaf) = same_leaf { + let field = leaves[leaf] + .ctx + .field_columns + .get(idx) + .cloned() + .unwrap_or_else(|| lhs_name.clone()); + lhs_name = format!("lhs.{field}"); + rhs_name = format!("rhs.{field}"); + } + exprs.push(expr_builder(lhs_expr, rhs_expr)?); + names.push(format!("{lhs_name} {op} {rhs_name}")); + } + (exprs, names) + } + }; + + Ok(IslandFieldExprs { + exprs, + names, + scalar, + }) + } + } + } + + fn project_binary_island( + &mut self, + input: LogicalPlan, + base_alias: &TableReference, + base_ctx: &PromPlannerContext, + field_exprs: IslandFieldExprs, + ) -> Result { + self.ctx = base_ctx.clone(); + + let schema = input.schema(); + let non_field_exprs = base_ctx + .tag_columns + .iter() + .chain(base_ctx.time_index_column.iter()) + .map(|column| { + schema + .qualified_field_with_name(Some(base_alias), column) + .context(DataFusionPlanningSnafu) + .map(|field| DfExpr::Column(field.into())) + }); + let tsid_expr = Self::optional_tsid_projection(schema, Some(base_alias), base_ctx.use_tsid) + .into_iter() + .map(Ok); + + self.ctx.field_columns = field_exprs.names; + let field_exprs = field_exprs + .exprs + .into_iter() + .zip(self.ctx.field_columns.iter()) + .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::, name)))); + + let project_exprs = non_field_exprs + .chain(tsid_expr) + .chain(field_exprs) + .collect::>>()?; + + let plan = LogicalPlanBuilder::from(input) + .project(project_exprs) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + + self.ctx.table_name = None; + self.ctx.schema_name = None; + + Ok(plan) + } +} diff --git a/src/query/src/promql/planner/set_operator.rs b/src/query/src/promql/planner/set_operator.rs new file mode 100644 index 00000000000..51cc100408f --- /dev/null +++ b/src/query/src/promql/planner/set_operator.rs @@ -0,0 +1,832 @@ +// 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. + +//! Planning of the PromQL set operators (`and`, `unless`, and `or`). + +use std::collections::{BTreeSet, HashMap, HashSet}; +use std::sync::Arc; + +use datafusion::logical_expr::{Cast, Extension, LogicalPlan, LogicalPlanBuilder}; +use datafusion::prelude::{Column, Expr as DfExpr, JoinType}; +use datafusion::scalar::ScalarValue; +use datafusion_common::{NullEquality, TableReference}; +use datatypes::arrow::datatypes::DataType as ArrowDataType; +use promql::extension_plan::UnionDistinctOn; +use promql_parser::parser::token::{self, TokenType}; +use promql_parser::parser::{BinModifier, LabelModifier, VectorMatchCardinality}; +use snafu::{OptionExt, ResultExt, ensure}; +use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME; + +use crate::promql::error::{ + ColumnNotFoundSnafu, CombineTableColumnMismatchSnafu, DataFusionPlanningSnafu, + MultiFieldsNotSupportedSnafu, Result, TimeIndexNotFoundSnafu, UnexpectedPlanExprSnafu, + UnexpectedTokenSnafu, UnsupportedVectorMatchSnafu, +}; +use crate::promql::planner::{ + OR_FLOAT_FIELD_PREFIX, OR_HISTOGRAM_FIELD_PREFIX, PromPlanner, PromPlannerContext, +}; + +impl PromPlanner { + // TODO(ruihang): change function name + #[allow(clippy::too_many_arguments)] + pub(super) fn or_operator( + &mut self, + left: LogicalPlan, + right: LogicalPlan, + left_tag_cols_set: HashSet, + right_tag_cols_set: HashSet, + left_context: PromPlannerContext, + right_context: PromPlannerContext, + modifier: &Option, + ) -> Result { + let left_is_empty = Self::is_zero_row_empty_relation(&left); + let right_is_empty = Self::is_zero_row_empty_relation(&right); + match (left_is_empty, right_is_empty) { + (true, false) => { + self.ctx = right_context; + return Ok(right); + } + (false, true) => { + self.ctx = left_context; + return Ok(left); + } + (true, true) => { + self.ctx = left_context; + return Ok(left); + } + (false, false) => {} + } + + ensure!( + !left.schema().fields().is_empty() && !right.schema().fields().is_empty(), + UnexpectedPlanExprSnafu { + desc: "OR operator input has zero columns", + } + ); + let left_has_alternative_samples = + Self::field_columns_are_alternative_samples(left.schema(), &left_context.field_columns); + let right_has_alternative_samples = Self::field_columns_are_alternative_samples( + right.schema(), + &right_context.field_columns, + ); + ensure!( + left_context.field_columns.len() == 1 || left_has_alternative_samples, + MultiFieldsNotSupportedSnafu { + operator: "OR operator" + } + ); + ensure!( + right_context.field_columns.len() == 1 || right_has_alternative_samples, + MultiFieldsNotSupportedSnafu { + operator: "OR operator" + } + ); + + // prepare hash sets + let all_tags = left_tag_cols_set + .union(&right_tag_cols_set) + .cloned() + .collect::>(); + let left_qualifier = left.schema().qualified_field(0).0.cloned(); + let right_qualifier = right.schema().qualified_field(0).0.cloned(); + let left_qualifier_string = left_qualifier + .as_ref() + .map(|l| l.to_string()) + .unwrap_or_default(); + let right_qualifier_string = right_qualifier + .as_ref() + .map(|r| r.to_string()) + .unwrap_or_default(); + let left_time_index_column = + left_context + .time_index_column + .clone() + .with_context(|| TimeIndexNotFoundSnafu { + table: left_qualifier_string.clone(), + })?; + let right_time_index_column = + right_context + .time_index_column + .clone() + .with_context(|| TimeIndexNotFoundSnafu { + table: right_qualifier_string.clone(), + })?; + let native_histogram_type = Self::native_histogram_arrow_type(); + let is_numeric = |data_type: &ArrowDataType| { + matches!( + data_type, + ArrowDataType::Int8 + | ArrowDataType::Int16 + | ArrowDataType::Int32 + | ArrowDataType::Int64 + | ArrowDataType::UInt8 + | ArrowDataType::UInt16 + | ArrowDataType::UInt32 + | ArrowDataType::UInt64 + | ArrowDataType::Float32 + | ArrowDataType::Float64 + ) + }; + let left_fields = left_context + .field_columns + .iter() + .map(|name| { + left.schema() + .iter() + .find(|(_, field)| field.name() == name) + .map(|(qualifier, field)| { + (name.clone(), qualifier.cloned(), field.data_type().clone()) + }) + .with_context(|| ColumnNotFoundSnafu { col: name.clone() }) + }) + .collect::>>()?; + let right_fields = right_context + .field_columns + .iter() + .map(|name| { + right + .schema() + .iter() + .find(|(_, field)| field.name() == name) + .map(|(qualifier, field)| { + (name.clone(), qualifier.cloned(), field.data_type().clone()) + }) + .with_context(|| ColumnNotFoundSnafu { col: name.clone() }) + }) + .collect::>>()?; + let left_field = &left_fields[0]; + let right_field = &right_fields[0]; + let left_field_col = &left_field.0; + let right_field_col = &right_field.0; + let fields_are_samples = |fields: &[(String, Option, ArrowDataType)]| { + fields.iter().all(|(_, _, data_type)| { + is_numeric(data_type) || data_type == &native_histogram_type + }) + }; + let mixed_sample_types = if left_has_alternative_samples || right_has_alternative_samples { + if !fields_are_samples(&left_fields) || !fields_are_samples(&right_fields) { + return UnexpectedPlanExprSnafu { + desc: format!( + "OR value fields have incompatible types: {:?} and {:?}", + left_fields + .iter() + .map(|(_, _, data_type)| data_type) + .collect::>(), + right_fields + .iter() + .map(|(_, _, data_type)| data_type) + .collect::>() + ), + } + .fail(); + } + true + } else { + (left_field.2 == native_histogram_type && is_numeric(&right_field.2)) + || (right_field.2 == native_histogram_type && is_numeric(&left_field.2)) + }; + let target_field_type = if mixed_sample_types { + // Mixed vectors use the existing response representation: one nullable float column + // and one nullable native-histogram column. + ArrowDataType::Float64 + } else if left_field.2 == right_field.2 { + left_field.2.clone() + } else if is_numeric(&left_field.2) && is_numeric(&right_field.2) { + ArrowDataType::Float64 + } else { + return UnexpectedPlanExprSnafu { + desc: format!( + "OR value fields have incompatible types: {:?} and {:?}", + left_field.2, right_field.2 + ), + } + .fail(); + }; + let (mixed_float_field_col, mixed_histogram_field_col) = if mixed_sample_types { + let mut reserved_names = left + .schema() + .fields() + .iter() + .chain(right.schema().fields().iter()) + .map(|field| field.name().clone()) + .collect::>(); + for (name, _, _) in left_fields.iter().chain(&right_fields) { + reserved_names.remove(name); + } + reserved_names.extend(all_tags.iter().cloned()); + let unique_name = |prefix: &str, reserved_names: &mut HashSet| { + let mut index = 0; + loop { + let name = format!("{prefix}{index}"); + index += 1; + if reserved_names.insert(name.clone()) { + break name; + } + } + }; + let float_field = unique_name(OR_FLOAT_FIELD_PREFIX, &mut reserved_names); + let histogram_field = unique_name(OR_HISTOGRAM_FIELD_PREFIX, &mut reserved_names); + (float_field, histogram_field) + } else { + (left_field_col.clone(), String::new()) + }; + let left_tag_types = left_tag_cols_set + .iter() + .map(|label| { + left.schema() + .fields() + .iter() + .find(|field| field.name() == label) + .map(|field| (label.clone(), field.data_type().clone())) + .with_context(|| ColumnNotFoundSnafu { col: label.clone() }) + }) + .collect::>>()?; + let right_tag_types = right_tag_cols_set + .iter() + .map(|label| { + right + .schema() + .fields() + .iter() + .find(|field| field.name() == label) + .map(|field| (label.clone(), field.data_type().clone())) + .with_context(|| ColumnNotFoundSnafu { col: label.clone() }) + }) + .collect::>>()?; + let mut target_tag_types = HashMap::with_capacity(all_tags.len()); + for label in &all_tags { + let Some(data_type) = + Self::common_label_data_type(left_tag_types.get(label), right_tag_types.get(label)) + else { + return UnexpectedPlanExprSnafu { + desc: format!( + "OR label {label} has incompatible types: {:?} and {:?}", + left_tag_types.get(label), + right_tag_types.get(label) + ), + } + .fail(); + }; + target_tag_types.insert(label.clone(), data_type); + } + let left_has_tsid = left + .schema() + .fields() + .iter() + .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME); + let right_has_tsid = right + .schema() + .fields() + .iter() + .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME); + + // step 0: fill all columns in output schema + let mut all_columns_set = left + .schema() + .fields() + .iter() + .chain(right.schema().fields().iter()) + .map(|field| field.name().clone()) + .collect::>(); + // Keep `__tsid` only when both sides contain it, otherwise it may break schema alignment + // (e.g. `unknown_metric or some_metric`). + if !(left_has_tsid && right_has_tsid) { + all_columns_set.remove(DATA_SCHEMA_TSID_COLUMN_NAME); + } + // remove time index column + all_columns_set.remove(&left_time_index_column); + all_columns_set.remove(&right_time_index_column); + if mixed_sample_types { + for (name, _, _) in left_fields.iter().chain(&right_fields) { + all_columns_set.remove(name); + } + all_columns_set.extend(all_tags.iter().cloned()); + all_columns_set.insert(mixed_float_field_col.clone()); + all_columns_set.insert(mixed_histogram_field_col.clone()); + } else if left_field_col != right_field_col { + // remove field column in the right + all_columns_set.remove(right_field_col); + } + let mut all_columns = all_columns_set.into_iter().collect::>(); + // sort to ensure the generated schema is not volatile + all_columns.sort_unstable(); + // use left time index column name as the result time index column name + all_columns.insert(0, left_time_index_column.clone()); + let mut occupied_column_names = left + .schema() + .fields() + .iter() + .chain(right.schema().fields().iter()) + .map(|field| field.name().clone()) + .collect::>(); + + // step 1: align schema using project, fill non-exist columns with null + let aligned_label_expr = |col: &String, source_types: &HashMap| { + let target_type = &target_tag_types[col]; + if let Some(source_type) = source_types.get(col) { + let expr = DfExpr::Column(Column::new(None::, col)); + if source_type == target_type { + expr + } else { + DfExpr::Cast(Cast::new(Box::new(expr), target_type.clone())).alias(col.clone()) + } + } else { + DfExpr::Literal( + Self::string_scalar_value(target_type, None) + .expect("target label type is a string"), + None, + ) + .alias(col.clone()) + } + }; + let null_histogram = + ScalarValue::try_new_null(&native_histogram_type).context(DataFusionPlanningSnafu)?; + let mixed_value_expr = |fields: &[(String, Option, ArrowDataType)], + output_col: &String| { + if output_col == &mixed_float_field_col { + if let Some((name, qualifier, data_type)) = fields + .iter() + .find(|(_, _, data_type)| is_numeric(data_type)) + { + let expr = DfExpr::Column(Column::new(qualifier.clone(), name)); + if data_type == &ArrowDataType::Float64 { + expr.alias(output_col) + } else { + DfExpr::Cast(Cast::new(Box::new(expr), ArrowDataType::Float64)) + .alias(output_col) + } + } else { + DfExpr::Literal(ScalarValue::Float64(None), None).alias(output_col) + } + } else { + fields + .iter() + .find(|(_, _, data_type)| data_type == &native_histogram_type) + .map(|(name, qualifier, _)| { + DfExpr::Column(Column::new(qualifier.clone(), name)).alias(output_col) + }) + .unwrap_or_else(|| { + DfExpr::Literal(null_histogram.clone(), None).alias(output_col) + }) + } + }; + let left_proj_exprs = all_columns.iter().map(|col| { + if mixed_sample_types + && (col == &mixed_float_field_col || col == &mixed_histogram_field_col) + { + mixed_value_expr(&left_fields, col) + } else if !mixed_sample_types + && col == left_field_col + && left_field.2 != target_field_type + { + DfExpr::Cast(Cast::new( + Box::new(DfExpr::Column(Column::new( + left_field.1.clone(), + left_field_col, + ))), + target_field_type.clone(), + )) + .alias(left_field_col.clone()) + } else if target_tag_types.contains_key(col) { + aligned_label_expr(col, &left_tag_types) + } else { + DfExpr::Column(Column::new(None::, col)) + } + }); + let right_time_index_expr = DfExpr::Column(Column::new( + right_qualifier.clone(), + right_time_index_column, + )) + .alias(left_time_index_column.clone()); + // The field column in right side may not have qualifier (it may be removed by join operation), + // so we need to find it from the schema. + // `skip（1)` to skip the time index column + let right_proj_exprs_without_time_index = all_columns.iter().skip(1).map(|col| { + // expr + if mixed_sample_types + && (col == &mixed_float_field_col || col == &mixed_histogram_field_col) + { + mixed_value_expr(&right_fields, col) + } else if !mixed_sample_types && col == left_field_col { + let expr = DfExpr::Column(Column::new(right_field.1.clone(), right_field_col)); + if right_field.2 != target_field_type { + DfExpr::Cast(Cast::new(Box::new(expr), target_field_type.clone())) + .alias(left_field_col.clone()) + } else if left_field_col != right_field_col { + expr.alias(left_field_col.clone()) + } else { + expr + } + } else if target_tag_types.contains_key(col) { + aligned_label_expr(col, &right_tag_types) + } else { + DfExpr::Column(Column::new(None::, col)) + } + }); + let right_proj_exprs = [right_time_index_expr] + .into_iter() + .chain(right_proj_exprs_without_time_index); + + let left_projected = LogicalPlanBuilder::from(left) + .project(left_proj_exprs) + .context(DataFusionPlanningSnafu)? + .alias(left_qualifier_string.clone()) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + let right_projected = LogicalPlanBuilder::from(right) + .project(right_proj_exprs) + .context(DataFusionPlanningSnafu)? + .alias(right_qualifier_string.clone()) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + + // step 2: compute match columns + let mut match_columns = if let Some(modifier) = modifier + && let Some(matching) = &modifier.matching + { + match matching { + // keeps columns mentioned in `on` + LabelModifier::Include(on) => on.labels.clone(), + // removes columns memtioned in `ignoring` + LabelModifier::Exclude(ignoring) => { + let ignoring = ignoring.labels.iter().cloned().collect::>(); + all_tags.difference(&ignoring).cloned().collect() + } + } + } else { + all_tags.iter().cloned().collect() + }; + // sort to ensure the generated plan is not volatile + match_columns.sort_unstable(); + match_columns.dedup(); + occupied_column_names.extend( + left_projected + .schema() + .fields() + .iter() + .chain(right_projected.schema().fields().iter()) + .map(|field| field.name().clone()), + ); + + let visible_schema = left_projected.schema().clone(); + let visible_left_exprs = left_projected + .schema() + .iter() + .map(|(qualifier, field)| { + DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) + }) + .collect::>(); + let visible_right_exprs = right_projected + .schema() + .iter() + .map(|(qualifier, field)| { + DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) + }) + .collect::>(); + let mut left_match_exprs = Vec::with_capacity(match_columns.len()); + let mut right_match_exprs = Vec::with_capacity(match_columns.len()); + let mut next_internal_column = 0; + + for label in &match_columns { + let left_field = if left_tag_cols_set.contains(label) { + Some( + left_projected + .schema() + .iter() + .find(|(_, field)| field.name() == label) + .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone())) + .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?, + ) + } else { + None + }; + let right_field = if right_tag_cols_set.contains(label) { + Some( + right_projected + .schema() + .iter() + .find(|(_, field)| field.name() == label) + .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone())) + .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?, + ) + } else { + None + }; + let data_type = match (left_field.as_ref(), right_field.as_ref()) { + (Some((_, left_type)), Some((_, right_type))) if left_type == right_type => { + left_type.clone() + } + (Some((_, left_type)), Some((_, right_type))) => { + return UnexpectedPlanExprSnafu { + desc: format!( + "OR match label {label} has incompatible types: {left_type:?} and {right_type:?}" + ), + } + .fail(); + } + (Some((_, data_type)), None) | (None, Some((_, data_type))) => data_type.clone(), + (None, None) => ArrowDataType::Utf8, + }; + let Some(value_type) = Self::string_value_data_type(&data_type).cloned() else { + return UnexpectedPlanExprSnafu { + desc: format!("OR match label {label} must be a string"), + } + .fail(); + }; + let internal_name = loop { + let name = format!("__promql_or_match_{next_internal_column}"); + next_internal_column += 1; + if occupied_column_names.insert(name.clone()) { + break name; + } + }; + left_match_exprs.push(Self::normalized_match_key_expr( + label, + left_field, + &value_type, + &internal_name, + )); + right_match_exprs.push(Self::normalized_match_key_expr( + label, + right_field, + &value_type, + &internal_name, + )); + } + + let left_augmented = LogicalPlanBuilder::from(left_projected) + .project(visible_left_exprs.into_iter().chain(left_match_exprs)) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + let right_augmented = LogicalPlanBuilder::from(right_projected) + .project(visible_right_exprs.into_iter().chain(right_match_exprs)) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + + // step 3: build `UnionDistinctOn` with normalized internal match keys. + let visible_field_count = visible_schema.fields().len(); + let compare_key_indices = + (visible_field_count..visible_field_count + match_columns.len()).collect::>(); + let (time_qualifier, _) = visible_schema + .iter() + .find(|(_, field)| field.name() == &left_time_index_column) + .with_context(|| TimeIndexNotFoundSnafu { + table: left_qualifier_string.clone(), + })?; + let ts_col_idx = left_augmented + .schema() + .iter() + .position(|(qualifier, field)| { + qualifier == time_qualifier && field.name() == &left_time_index_column + }) + .with_context(|| TimeIndexNotFoundSnafu { + table: left_qualifier_string.clone(), + })?; + let union_distinct_on = UnionDistinctOn::try_new( + left_augmented, + right_augmented, + compare_key_indices, + ts_col_idx, + ) + .context(DataFusionPlanningSnafu)?; + let augmented_result = LogicalPlan::Extension(Extension { + node: Arc::new(union_distinct_on), + }); + let result = LogicalPlanBuilder::from(augmented_result) + .project(visible_schema.iter().map(|(qualifier, field)| { + DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) + })) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + + // step 4: update context + let output_field_col = left_field_col.clone(); + let mut output_context = left_context; + let mut visible_tags = all_tags.into_iter().collect::>(); + visible_tags.sort_unstable(); + output_context.time_index_column = Some(left_time_index_column); + output_context.tag_columns = visible_tags; + output_context.field_columns = if mixed_sample_types { + vec![mixed_float_field_col, mixed_histogram_field_col] + } else { + vec![output_field_col] + }; + output_context.use_tsid = left_has_tsid && right_has_tsid; + self.ctx = output_context; + + Ok(result) + } + + /// Build a set operator (AND/OR/UNLESS) + pub(super) fn set_op_on_non_field_columns( + &mut self, + mut left: LogicalPlan, + mut right: LogicalPlan, + left_context: PromPlannerContext, + right_context: PromPlannerContext, + op: TokenType, + modifier: &Option, + ) -> Result { + let left_tag_col_set = left_context + .tag_columns + .iter() + .cloned() + .collect::>(); + let right_tag_col_set = right_context + .tag_columns + .iter() + .cloned() + .collect::>(); + + if matches!(op.id(), token::T_LOR) { + return self.or_operator( + left, + right, + left_tag_col_set, + right_tag_col_set, + left_context, + right_context, + modifier, + ); + } + + if let Some(modifier) = modifier { + ensure!( + matches!( + modifier.card, + VectorMatchCardinality::OneToOne | VectorMatchCardinality::ManyToMany + ), + UnsupportedVectorMatchSnafu { + name: modifier.card.clone(), + }, + ); + } + + let output_context = left_context.clone(); + let visible_left_schema = left.schema().clone(); + let mut left_context = left_context; + let mut right_context = right_context; + let added_marker_to_left = if Self::only_temporality_match_label_mismatches( + &left_context, + &right_context, + modifier, + ) { + let aligned = Self::align_temporality_match_column( + left, + right, + &mut left_context, + &mut right_context, + )?; + left = aligned.0; + right = aligned.1; + aligned.2 + } else { + false + }; + + let mut left_tag_col_set = left_context + .tag_columns + .iter() + .cloned() + .collect::>(); + let mut right_tag_col_set = right_context + .tag_columns + .iter() + .cloned() + .collect::>(); + if let Some(matching) = modifier + .as_ref() + .and_then(|modifier| modifier.matching.as_ref()) + { + match matching { + LabelModifier::Include(on) => { + let mask = on.labels.iter().cloned().collect::>(); + left_tag_col_set = left_tag_col_set.intersection(&mask).cloned().collect(); + right_tag_col_set = right_tag_col_set.intersection(&mask).cloned().collect(); + } + LabelModifier::Exclude(ignoring) => { + for label in &ignoring.labels { + let _ = left_tag_col_set.remove(label); + let _ = right_tag_col_set.remove(label); + } + } + } + } + ensure!( + left_tag_col_set == right_tag_col_set, + CombineTableColumnMismatchSnafu { + left: left_tag_col_set.iter().cloned().collect::>(), + right: right_tag_col_set.iter().cloned().collect::>(), + } + ); + + let left_time_index = left_context.time_index_column.clone().unwrap(); + let right_time_index = right_context.time_index_column.clone().unwrap(); + + // alias right time index column if necessary + if left_context.time_index_column != right_context.time_index_column { + let right_project_exprs = right + .schema() + .fields() + .iter() + .map(|field| { + if field.name() == &right_time_index { + DfExpr::Column(Column::from_name(&right_time_index)).alias(&left_time_index) + } else { + DfExpr::Column(Column::from_name(field.name())) + } + }) + .collect::>(); + + right = LogicalPlanBuilder::from(right) + .project(right_project_exprs) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + } + + let join_keys = left_tag_col_set + .into_iter() + .chain([left_time_index]) + .collect::>(); + + ensure!( + left_context.field_columns.len() == 1 + || Self::field_columns_are_alternative_samples( + left.schema(), + &left_context.field_columns, + ), + MultiFieldsNotSupportedSnafu { + operator: "AND/UNLESS operator" + } + ); + // Generate join plan. + // All set operations in PromQL are "distinct" + let result = match op.id() { + token::T_LAND => LogicalPlanBuilder::from(left) + .distinct() + .context(DataFusionPlanningSnafu)? + .join_detailed( + right, + JoinType::LeftSemi, + (join_keys.clone(), join_keys), + None, + NullEquality::NullEqualsNull, + ) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu), + token::T_LUNLESS => LogicalPlanBuilder::from(left) + .distinct() + .context(DataFusionPlanningSnafu)? + .join_detailed( + right, + JoinType::LeftAnti, + (join_keys.clone(), join_keys), + None, + NullEquality::NullEqualsNull, + ) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu), + token::T_LOR => { + // OR is handled at the beginning of this function, as it cannot + // be expressed using JOIN like AND and UNLESS. + unreachable!() + } + _ => UnexpectedTokenSnafu { token: op }.fail(), + }?; + let result = if added_marker_to_left { + LogicalPlanBuilder::from(result) + .project(visible_left_schema.iter().map(|(qualifier, field)| { + DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) + })) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)? + } else { + result + }; + + // AND/UNLESS preserve the complete left operand's visible columns and values; encoded + // markers are decoded. + self.ctx = output_context; + Ok(result) + } +} diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs new file mode 100644 index 00000000000..1634e6d9e1d --- /dev/null +++ b/src/query/src/promql/planner/test.rs @@ -0,0 +1,7585 @@ +// 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::collections::HashMap; +use std::time::{Duration, UNIX_EPOCH}; + +use catalog::RegisterTableRequest; +use catalog::memory::{MemoryCatalogManager, new_memory_catalog_manager}; +use common_base::Plugins; +use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; +use common_query::native_histogram::{ + CUSTOM_BUCKETS_SCHEMA, CounterResetHint, NativeHistogram, build_histogram_array, +}; +use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value}; +use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; +use common_query::test_util::DummyDecoder; +use common_recordbatch::RecordBatch as GreptimeRecordBatch; +use datafusion::arrow::array::{ + Array, ArrayRef, Float64Array, Int64Array, StringArray, TimestampMillisecondArray, +}; +use datafusion::arrow::datatypes::{Field, Schema as ArrowSchema}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::catalog::{CatalogProvider, MemoryCatalogProvider, MemorySchemaProvider}; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; +use datafusion::datasource::{MemTable, provider_as_source}; +use datafusion::execution::context::SessionContext; +use datafusion::logical_expr::Extension; +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::{ColumnSchema, Schema}; +use promql::extension_plan::HistogramFold; +use promql_parser::label::Labels; +use promql_parser::parser; +use session::context::QueryContext; +use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; +use table::Table; +use table::metadata::{FilterPushDownType, TableInfoBuilder, TableMetaBuilder}; +use table::test_util::{EmptyTable, MemTable as GreptimeMemTable}; + +use super::*; +use crate::QueryEngineContext; +use crate::options::QueryOptions; +use crate::parser::QueryLanguageParser; +use crate::query_engine::DefaultSerializer; + +mod delta; + +fn find_instant_manipulate(plan: &LogicalPlan) -> Option<&InstantManipulate> { + if let LogicalPlan::Extension(Extension { node }) = plan + && let Some(instant_manipulate) = node.as_any().downcast_ref::() + { + return Some(instant_manipulate); + } + + plan.inputs().into_iter().find_map(find_instant_manipulate) +} + +fn build_query_engine_state() -> QueryEngineState { + QueryEngineState::new( + new_memory_catalog_manager().unwrap(), + None, + None, + None, + None, + None, + false, + Plugins::default(), + QueryOptions::default(), + ) +} + +#[test] +fn common_label_type_preserves_only_shared_dictionary_encoding() { + let dictionary = ArrowDataType::Dictionary( + Box::new(ArrowDataType::UInt32), + Box::new(ArrowDataType::Utf8), + ); + let other_dictionary = ArrowDataType::Dictionary( + Box::new(ArrowDataType::Int32), + Box::new(ArrowDataType::Utf8), + ); + + assert_eq!( + Some(dictionary.clone()), + PromPlanner::common_label_data_type(Some(&dictionary), Some(&dictionary)) + ); + assert_eq!( + Some(ArrowDataType::Utf8), + PromPlanner::common_label_data_type(Some(&dictionary), Some(&ArrowDataType::Utf8)) + ); + assert_eq!( + Some(ArrowDataType::Utf8), + PromPlanner::common_label_data_type(Some(&dictionary), Some(&other_dictionary)) + ); + assert_eq!( + Some(ArrowDataType::Utf8), + PromPlanner::common_label_data_type(Some(&dictionary), None) + ); +} + +async fn build_optimized_promql_plan( + table_provider: DfTableSourceProvider, + eval_stmt: &EvalStmt, +) -> LogicalPlan { + let state = build_query_engine_state(); + let raw_plan = PromPlanner::stmt_to_plan(table_provider, eval_stmt, &state) + .await + .unwrap(); + let context = QueryEngineContext::new(state.session_state(), QueryContext::arc()); + state + .optimize_by_extension_rules(raw_plan, &context) + .unwrap() +} + +async fn build_optimized_tsid_plan( + query: &str, + num_tag: usize, + num_field: usize, + end_secs: u64, + lookback_secs: u64, +) -> String { + let eval_stmt = EvalStmt { + expr: parser::parse(query).unwrap(), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(end_secs)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(lookback_secs), + }; + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + num_tag, + num_field, + ) + .await; + + build_optimized_promql_plan(table_provider, &eval_stmt) + .await + .display_indent_schema() + .to_string() +} + +async fn assert_nested_count_rewrite_applies(query: &str, expected_outer_agg: &str) { + let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await; + + assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); + assert!(plan_str.contains("Projection: some_metric.timestamp, some_metric.tag_0")); + assert!(plan_str.contains("Distinct:")); + assert!(plan_str.contains(expected_outer_agg), "{plan_str}"); + assert!(!plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]")); +} + +async fn assert_nested_count_rewrite_missing(query: &str, num_tag: usize, lookback_secs: u64) { + let plan_str = build_optimized_tsid_plan(query, num_tag, 1, 100_000, lookback_secs).await; + assert!(!plan_str.contains("Distinct:"), "{plan_str}"); +} + +fn build_eval_stmt(expr: &str) -> EvalStmt { + EvalStmt { + expr: parser::parse(expr).unwrap(), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + } +} + +enum DirectOrValue { + Float64(f64), + Int64(i64), + NativeHistogram(NativeHistogram), + Utf8(&'static str), +} + +impl DirectOrValue { + fn data_type(&self) -> ArrowDataType { + match self { + Self::Float64(_) => ArrowDataType::Float64, + Self::Int64(_) => ArrowDataType::Int64, + Self::NativeHistogram(_) => native_histogram_value_type().as_arrow_type(), + Self::Utf8(_) => ArrowDataType::Utf8, + } + } + fn array(&self) -> Arc { + match self { + Self::Float64(v) => Arc::new(Float64Array::from(vec![*v])), + Self::Int64(v) => Arc::new(Int64Array::from(vec![*v])), + Self::NativeHistogram(v) => build_histogram_array(&[Some(v.clone())]), + Self::Utf8(v) => Arc::new(StringArray::from(vec![*v])), + } + } +} + +fn direct_or_histogram() -> NativeHistogram { + NativeHistogram { + schema: 0, + zero_threshold: 0.0, + sum: 1.0, + reset_hint: CounterResetHint::Unknown, + start_timestamp: None, + custom_values: vec![], + positive_spans: vec![], + negative_spans: vec![], + count: 1.0, + zero_count: 1.0, + positive_buckets: vec![], + negative_buckets: vec![], + } +} + +fn operator_metric_table( + name: &str, + table_id: u32, + tag: &str, + le: Option<&str>, + value: DirectOrValue, +) -> table::TableRef { + let value_type = match &value { + DirectOrValue::Float64(_) => ConcreteDataType::float64_datatype(), + DirectOrValue::Int64(_) => ConcreteDataType::int64_datatype(), + DirectOrValue::NativeHistogram(_) => native_histogram_value_type().clone(), + DirectOrValue::Utf8(_) => ConcreteDataType::string_datatype(), + }; + let tag_count = 1 + usize::from(le.is_some()); + let mut columns = vec![ColumnSchema::new( + "tag".to_string(), + ConcreteDataType::string_datatype(), + false, + )]; + if le.is_some() { + columns.push(ColumnSchema::new( + LE_COLUMN_NAME.to_string(), + ConcreteDataType::string_datatype(), + false, + )); + } + columns.extend([ + ColumnSchema::new( + "ts".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new("v".to_string(), value_type, true), + ]); + let schema = Arc::new(Schema::new(columns)); + let mut arrays = vec![Arc::new(StringArray::from(vec![tag])) as Arc]; + if let Some(le) = le { + arrays.push(Arc::new(StringArray::from(vec![le]))); + } + arrays.extend([ + Arc::new(TimestampMillisecondArray::from(vec![1_000])) as Arc, + value.array(), + ]); + let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap(); + let backing = GreptimeMemTable::new_with_catalog( + name, + GreptimeRecordBatch::from_df_record_batch(schema.clone(), batch), + table_id, + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + ); + let value_index = tag_count + 1; + let meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices((0..tag_count).collect()) + .value_indices(vec![value_index]) + .next_column_id((value_index + 1) as u32) + .build() + .unwrap(); + let info = Arc::new( + TableInfoBuilder::default() + .table_id(table_id) + .name(name) + .meta(meta) + .build() + .unwrap(), + ); + Arc::new(Table::new( + info, + FilterPushDownType::Unsupported, + backing.data_source(), + )) +} + +fn operator_table_provider() -> DfTableSourceProvider { + let catalog = MemoryCatalogManager::with_default_setup(); + let tables = [ + operator_metric_table("lf", 2_001, "a", None, DirectOrValue::Float64(2.0)), + operator_metric_table( + "lh", + 2_002, + "b", + None, + DirectOrValue::NativeHistogram(direct_or_histogram()), + ), + operator_metric_table("rf", 2_003, "b", None, DirectOrValue::Float64(3.0)), + operator_metric_table( + "rh", + 2_004, + "a", + None, + DirectOrValue::NativeHistogram(direct_or_histogram()), + ), + operator_metric_table("fallback", 2_005, "c", None, DirectOrValue::Float64(7.0)), + operator_metric_table( + "bad_classic", + 2_006, + "d", + Some("broken"), + DirectOrValue::Float64(1.0), + ), + operator_metric_table( + "bad_native", + 2_007, + "d", + None, + DirectOrValue::NativeHistogram(direct_or_histogram()), + ), + ]; + for table in tables { + let info = table.table_info(); + catalog + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: info.name.clone(), + table_id: info.ident.table_id, + table, + }) + .unwrap(); + } + DfTableSourceProvider::new( + catalog, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +fn operator_eval_stmt(expr: &str) -> EvalStmt { + let time = UNIX_EPOCH.checked_add(Duration::from_secs(1)).unwrap(); + EvalStmt { + expr: parser::parse(expr).unwrap(), + start: time, + end: time, + interval: Duration::from_secs(1), + lookback_delta: Duration::from_secs(5), + } +} + +struct DirectOrSource { + name: &'static str, + empty: bool, + timestamp: i64, + tags: Vec<(&'static str, Option<&'static str>)>, + value: DirectOrValue, +} + +fn source( + name: &'static str, + empty: bool, + timestamp: i64, + tags: Vec<(&'static str, Option<&'static str>)>, + value: DirectOrValue, +) -> DirectOrSource { + DirectOrSource { + name, + empty, + timestamp, + tags, + value, + } +} + +fn tagged_source( + name: &'static str, + empty: bool, + tag: (&'static str, Option<&'static str>), + value: DirectOrValue, +) -> DirectOrSource { + source(name, empty, 1, vec![("job", Some("job")), tag], value) +} + +fn job_source(name: &'static str, value: DirectOrValue) -> DirectOrSource { + source(name, true, 1, vec![("job", Some("job"))], value) +} + +fn table(source: &DirectOrSource) -> Arc { + let mut fields = vec![Field::new( + "ts", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + )]; + fields.extend( + source + .tags + .iter() + .map(|(name, _)| Field::new(*name, ArrowDataType::Utf8, true)), + ); + fields.push(Field::new("v", source.value.data_type(), true)); + let schema = Arc::new(ArrowSchema::new(fields)); + let partitions = if source.empty { + vec![vec![]] + } else { + let mut columns: Vec> = + vec![Arc::new(TimestampMillisecondArray::from(vec![ + source.timestamp, + ]))]; + columns.extend( + source + .tags + .iter() + .map(|(_, value)| Arc::new(StringArray::from(vec![*value])) as Arc), + ); + columns.push(source.value.array()); + vec![vec![RecordBatch::try_new(schema.clone(), columns).unwrap()]] + }; + Arc::new(MemTable::try_new(schema, partitions).unwrap()) +} + +fn scan(source: &DirectOrSource) -> LogicalPlan { + LogicalPlanBuilder::scan(source.name, provider_as_source(table(source)), None) + .unwrap() + .build() + .unwrap() +} + +fn direct_or_context(qualifier: &str, tags: &[&str], field: &str) -> PromPlannerContext { + PromPlannerContext { + table_name: Some(qualifier.to_string()), + time_index_column: Some("ts".to_string()), + field_columns: vec![field.to_string()], + tag_columns: tags.iter().map(|tag| (*tag).to_string()).collect(), + ..Default::default() + } +} + +fn or_modifier(expr: &str) -> Option { + let PromExpr::Binary(expr) = parser::parse(expr).unwrap() else { + unreachable!() + }; + expr.modifier +} + +async fn plan_direct_or( + left: LogicalPlan, + right: LogicalPlan, + left_context: PromPlannerContext, + right_context: PromPlannerContext, + modifier: &Option, +) -> LogicalPlan { + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let mut planner = PromPlanner { + table_provider, + ctx: PromPlannerContext::default(), + promql_annotations: None, + }; + planner + .or_operator( + left, + right, + left_context.tag_columns.iter().cloned().collect(), + right_context.tag_columns.iter().cloned().collect(), + left_context, + right_context, + modifier, + ) + .unwrap() +} + +async fn execute(plan: LogicalPlan, state: &QueryEngineState) -> (LogicalPlan, Vec) { + let context = QueryEngineContext::new(state.session_state(), QueryContext::arc()); + let optimized = state.optimize_by_extension_rules(plan, &context).unwrap(); + let physical = state + .session_state() + .create_physical_plan(&optimized) + .await + .unwrap(); + let batches = datafusion::physical_plan::collect(physical, state.session_state().task_ctx()) + .await + .unwrap(); + (optimized, batches) +} + +async fn run( + left: &DirectOrSource, + right: &DirectOrSource, + left_context: PromPlannerContext, + right_context: PromPlannerContext, + modifier: &Option, +) -> (LogicalPlan, Vec) { + let plan = plan_direct_or( + scan(left), + scan(right), + left_context, + right_context, + modifier, + ) + .await; + execute(plan, &build_query_engine_state()).await +} + +async fn mixed_direct_or(histogram_on_left: bool) -> (PromPlanner, LogicalPlan) { + let sample = |histogram: bool| { + if histogram { + DirectOrValue::NativeHistogram(direct_or_histogram()) + } else { + DirectOrValue::Float64(1.25) + } + }; + let left = tagged_source( + "lhs", + false, + ( + "k", + Some(if histogram_on_left { + "histogram" + } else { + "float" + }), + ), + sample(histogram_on_left), + ); + let right = tagged_source( + "rhs", + false, + ( + "k", + Some(if histogram_on_left { + "float" + } else { + "histogram" + }), + ), + sample(!histogram_on_left), + ); + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let mut planner = PromPlanner { + table_provider, + ctx: PromPlannerContext::default(), + promql_annotations: None, + }; + let left_context = direct_or_context("lhs", &["job", "k"], "v"); + let right_context = direct_or_context("rhs", &["job", "k"], "v"); + let plan = planner + .or_operator( + scan(&left), + scan(&right), + left_context.tag_columns.iter().cloned().collect(), + right_context.tag_columns.iter().cloned().collect(), + left_context, + right_context, + &or_modifier("lhs or on(k) rhs"), + ) + .unwrap(); + (planner, plan) +} + +async fn mixed_aggregate_input(histograms: Vec) -> (PromPlanner, LogicalPlan) { + let float_field = format!("{OR_FLOAT_FIELD_PREFIX}0"); + let histogram_field = format!("{OR_HISTOGRAM_FIELD_PREFIX}0"); + let row_count = histograms.len() + 1; + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "ts", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + ), + Field::new("k", ArrowDataType::Utf8, false), + Field::new(&float_field, ArrowDataType::Float64, true), + Field::new( + &histogram_field, + native_histogram_value_type().as_arrow_type(), + true, + ), + ])); + let mut histogram_values = Vec::with_capacity(row_count); + histogram_values.push(None); + histogram_values.extend(histograms.into_iter().map(Some)); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1; row_count])), + Arc::new(StringArray::from_iter_values( + (0..row_count).map(|row| format!("kind_{row}")), + )), + Arc::new(Float64Array::from_iter( + (0..row_count).map(|row| (row == 0).then_some(1.25)), + )), + build_histogram_array(&histogram_values), + ], + ) + .unwrap(); + let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); + let plan = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None) + .unwrap() + .build() + .unwrap(); + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let planner = PromPlanner { + table_provider, + ctx: PromPlannerContext { + table_name: Some("mixed".to_string()), + time_index_column: Some("ts".to_string()), + field_columns: vec![float_field, histogram_field], + tag_columns: vec!["k".to_string()], + ..Default::default() + }, + promql_annotations: None, + }; + (planner, plan) +} + +fn assert_no_internal_or_keys(schema: &DFSchema) { + assert!( + schema + .fields() + .iter() + .all(|field| !field.name().starts_with("__promql_or_match_")), + "{schema:?}" + ); +} + +fn values(batches: &[RecordBatch], column: &str) -> Vec { + batches + .iter() + .flat_map(|batch| { + batch + .column_by_name(column) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + }) + .collect() +} + +fn numeric_values(batches: &[RecordBatch], column: &str) -> Vec { + batches + .iter() + .flat_map(|batch| { + let values = datafusion::arrow::compute::cast( + batch.column_by_name(column).unwrap(), + &ArrowDataType::Float64, + ) + .unwrap(); + values + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + .collect::>() + }) + .collect() +} + +fn histograms(batches: &[RecordBatch], column: &str) -> Vec { + batches + .iter() + .flat_map(|batch| { + let values = batch + .column_by_name(column) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..values.len()).filter_map(|row| { + common_query::native_histogram::read_histogram(values, row).unwrap() + }) + }) + .collect() +} + +fn rows(batches: &[RecordBatch]) -> Vec<(f64, Option)> { + let mut rows = batches + .iter() + .flat_map(|batch| { + let values = batch + .column_by_name("v") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let labels = batch + .column_by_name("k") + .map(|column| column.as_any().downcast_ref::().unwrap()); + (0..batch.num_rows()).map(move |i| { + ( + values.value(i), + labels.and_then(|labels| { + (!labels.is_null(i)).then(|| labels.value(i).to_string()) + }), + ) + }) + }) + .collect::>(); + rows.sort_by(|left, right| left.0.total_cmp(&right.0)); + rows +} + +fn matrix_source( + name: &'static str, + k: Option>, + timestamp: i64, + value: f64, +) -> DirectOrSource { + let mut tags = vec![("job", Some("job"))]; + if let Some(k) = k { + tags.push(("k", k)); + } + source(name, false, timestamp, tags, DirectOrValue::Float64(value)) +} + +fn matrix_context(name: &str, k: Option>) -> PromPlannerContext { + direct_or_context( + name, + if k.is_some() { &["job", "k"] } else { &["job"] }, + "v", + ) +} + +async fn build_missing_le_or_normal_metric_table_provider() -> DfTableSourceProvider { + build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "non_existent_histogram_bucket".to_string(), + ), + (DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string()), + ], + &["pod", "instance"], + ) + .await +} + +fn assert_normal_metric_schema(plan: &LogicalPlan) { + let fields = plan.schema().fields(); + assert_eq!(fields.len(), 4, "{fields:?}"); + assert!( + fields.iter().any(|field| field.name() == "pod"), + "{fields:?}" + ); + assert!( + fields.iter().any(|field| field.name() == "instance"), + "{fields:?}" + ); + assert!( + fields + .iter() + .any(|field| field.name() == greptime_timestamp()), + "{fields:?}" + ); + assert!( + fields.iter().any(|field| { + field.name() == greptime_value() && field.data_type() == &ArrowDataType::Float64 + }), + "{fields:?}" + ); +} + +async fn build_test_table_provider_with_distinct_tags( + table_tags: &[(&str, &[&str])], +) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + for (table_name, tags) in table_tags { + let mut columns = tags + .iter() + .map(|tag| { + ColumnSchema::new( + (*tag).to_string(), + ConcreteDataType::string_datatype(), + false, + ) + }) + .collect::>(); + columns.push( + ColumnSchema::new( + greptime_timestamp().to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ); + columns.push(ColumnSchema::new( + greptime_value().to_string(), + ConcreteDataType::float64_datatype(), + true, + )); + let table_meta = TableMetaBuilder::empty() + .schema(Arc::new(Schema::new(columns))) + .primary_key_indices((0..tags.len()).collect()) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name((*table_name).to_string()) + .meta(table_meta) + .build() + .unwrap(); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: (*table_name).to_string(), + table_id: 1024, + table: EmptyTable::from_table_info(&table_info), + }) + .is_ok() + ); + } + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +fn contains_histogram_fold(plan: &LogicalPlan) -> bool { + matches!(plan, LogicalPlan::Extension(Extension { node }) if node.as_any().is::()) + || plan.inputs().into_iter().any(contains_histogram_fold) +} + +async fn build_set_op_context_table_provider() -> DfTableSourceProvider { + build_test_table_provider_with_distinct_tags(&[ + ("bucket_metric", &["job", "le"]), + ("normal_metric", &["job"]), + ("fallback_metric", &["instance"]), + ]) + .await +} + +async fn build_or_context_table_provider() -> DfTableSourceProvider { + build_test_table_provider_with_distinct_tags(&[ + ("normal_metric", &["job"]), + ("other_metric", &["instance"]), + ("non_hist_metric", &["instance"]), + ]) + .await +} + +async fn optimize_and_create_physical_plan( + state: &QueryEngineState, + plan: LogicalPlan, +) -> ( + LogicalPlan, + Arc, +) { + let context = QueryEngineContext::new(state.session_state(), QueryContext::arc()); + let optimized = state.optimize_by_extension_rules(plan, &context).unwrap(); + let physical = state + .session_state() + .create_physical_plan(&optimized) + .await + .unwrap(); + (optimized, physical) +} + +async fn build_test_table_provider( + table_name_tuples: &[(String, String)], + num_tag: usize, + num_field: usize, +) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + for (schema_name, table_name) in table_name_tuples { + let mut columns = vec![]; + for i in 0..num_tag { + columns.push(ColumnSchema::new( + format!("tag_{i}"), + ConcreteDataType::string_datatype(), + false, + )); + } + columns.push( + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ); + for i in 0..num_field { + columns.push(ColumnSchema::new( + format!("field_{i}"), + ConcreteDataType::float64_datatype(), + true, + )); + } + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices((0..num_tag).collect()) + .value_indices((num_tag + 1..num_tag + 1 + num_field).collect()) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name(table_name.clone()) + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: schema_name.clone(), + table_name: table_name.clone(), + table_id: 1024, + table, + }) + .is_ok() + ); + } + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +async fn build_test_native_histogram_table_provider(table_name: &str) -> DfTableSourceProvider { + build_test_native_histogram_table_provider_with_marker(table_name, false).await +} + +async fn build_test_native_histogram_table_provider_with_marker( + table_name: &str, + temporality_marker: bool, +) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + let mut columns = vec![ + ColumnSchema::new( + "tag_0".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + ColumnSchema::new( + LE_COLUMN_NAME.to_string(), + ConcreteDataType::string_datatype(), + true, + ), + ]; + if temporality_marker { + columns.push(ColumnSchema::new( + OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(), + ConcreteDataType::string_datatype(), + true, + )); + } + let tag_count = columns.len(); + columns.extend([ + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + greptime_native_histogram().to_string(), + native_histogram_value_type().clone(), + true, + ), + ]); + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices((0..tag_count).collect()) + .value_indices(vec![tag_count + 1]) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name(table_name) + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + table_id: 1024, + table, + }) + .is_ok() + ); + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +async fn build_test_multi_histogram_table_provider(table_name: &str) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + let columns = vec![ + ColumnSchema::new( + "tag_0".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + greptime_native_histogram().to_string(), + native_histogram_value_type().clone(), + true, + ), + ColumnSchema::new( + "native_histogram_2".to_string(), + native_histogram_value_type().clone(), + true, + ), + ]; + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![0]) + .value_indices(vec![2, 3]) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name(table_name) + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + table_id: 1024, + table, + }) + .is_ok() + ); + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +async fn build_test_mixed_native_histogram_table_provider( + table_name: &str, +) -> DfTableSourceProvider { + build_test_mixed_native_histogram_table_provider_with_marker(table_name, false).await +} + +async fn build_test_mixed_native_histogram_table_provider_with_marker( + table_name: &str, + temporality_marker: bool, +) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + let mut columns = vec![ColumnSchema::new( + "tag_0".to_string(), + ConcreteDataType::string_datatype(), + false, + )]; + if temporality_marker { + columns.push(ColumnSchema::new( + OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(), + ConcreteDataType::string_datatype(), + true, + )); + } + let tag_count = columns.len(); + columns.extend([ + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + greptime_native_histogram().to_string(), + native_histogram_value_type().clone(), + true, + ), + ColumnSchema::new( + greptime_value().to_string(), + ConcreteDataType::float64_datatype(), + true, + ), + ]); + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema.clone()) + .primary_key_indices((0..tag_count).collect()) + .value_indices(vec![tag_count + 1, tag_count + 2]) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = Arc::new( + TableInfoBuilder::default() + .name(table_name) + .meta(table_meta) + .build() + .unwrap(), + ); + let mut arrays: Vec> = + vec![Arc::new(StringArray::from(vec!["float", "histogram"]))]; + if temporality_marker { + arrays.push(Arc::new(StringArray::from(vec![ + Some(GREPTIME_TEMPORALITY_DELTA), + Some(GREPTIME_TEMPORALITY_DELTA), + ]))); + } + arrays.extend([ + Arc::new(TimestampMillisecondArray::from(vec![1_000, 1_000])) as Arc, + build_histogram_array(&[None, Some(direct_or_histogram())]), + Arc::new(Float64Array::from(vec![Some(2.0), None])), + ]); + let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap(); + let backing = GreptimeMemTable::new_with_catalog( + table_name, + GreptimeRecordBatch::from_df_record_batch(schema, batch), + 1024, + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + ); + let table = Arc::new(Table::new( + table_info, + FilterPushDownType::Unsupported, + backing.data_source(), + )); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + table_id: 1024, + table, + }) + .is_ok() + ); + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +fn classic_and_native_histogram_table_provider( + native_tag: &str, + native_le: Option<&str>, + native_histogram: NativeHistogram, +) -> DfTableSourceProvider { + let table_name = "mixed_histogram"; + let catalog = MemoryCatalogManager::with_default_setup(); + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "tag".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + ColumnSchema::new( + LE_COLUMN_NAME.to_string(), + ConcreteDataType::string_datatype(), + true, + ), + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + greptime_native_histogram().to_string(), + native_histogram_value_type().clone(), + true, + ), + ColumnSchema::new( + greptime_value().to_string(), + ConcreteDataType::float64_datatype(), + true, + ), + ])); + let table_meta = TableMetaBuilder::empty() + .schema(schema.clone()) + .primary_key_indices(vec![0, 1]) + .value_indices(vec![3, 4]) + .next_column_id(5) + .build() + .unwrap(); + let table_info = Arc::new( + TableInfoBuilder::default() + .name(table_name) + .meta(table_meta) + .build() + .unwrap(), + ); + let batch = RecordBatch::try_new( + schema.arrow_schema().clone(), + vec![ + Arc::new(StringArray::from(vec![ + "classic", "classic", native_tag, "classic", "classic", native_tag, + ])), + Arc::new(StringArray::from(vec![ + Some("1"), + Some("+Inf"), + native_le, + Some("1"), + Some("+Inf"), + native_le, + ])), + Arc::new(TimestampMillisecondArray::from(vec![ + 1_000, 1_000, 1_000, 2_000, 2_000, 2_000, + ])), + build_histogram_array(&[ + None, + None, + Some(native_histogram.clone()), + None, + None, + Some(native_histogram), + ]), + Arc::new(Float64Array::from(vec![ + Some(2.0), + Some(4.0), + None, + Some(2.0), + Some(4.0), + None, + ])), + ], + ) + .unwrap(); + let backing = GreptimeMemTable::new_with_catalog( + table_name, + GreptimeRecordBatch::from_df_record_batch(schema, batch), + 2_200, + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + ); + let table = Arc::new(Table::new( + table_info, + FilterPushDownType::Unsupported, + backing.data_source(), + )); + catalog + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + table_id: 2_200, + table, + }) + .unwrap(); + + DfTableSourceProvider::new( + catalog, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +async fn build_test_table_provider_with_tsid( + table_name_tuples: &[(String, String)], + num_tag: usize, + num_field: usize, +) -> DfTableSourceProvider { + let table_specs = table_name_tuples + .iter() + .map(|(schema_name, table_name)| ((schema_name.clone(), table_name.clone()), num_field)) + .collect::>(); + build_test_table_provider_with_tsid_fields(&table_specs, num_tag).await +} + +async fn build_test_table_provider_with_tsid_fields( + table_specs: &[((String, String), usize)], + num_tag: usize, +) -> DfTableSourceProvider { + let table_specs = table_specs + .iter() + .map(|(table_name_tuple, num_field)| (table_name_tuple.clone(), num_tag, *num_field)) + .collect::>(); + build_test_table_provider_with_tsid_tag_fields(&table_specs).await +} + +async fn build_test_table_provider_with_tsid_tag_fields( + table_specs: &[((String, String), usize, usize)], +) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + + let physical_table_name = "phy"; + let physical_table_id = 999u32; + let physical_num_tag = table_specs + .iter() + .map(|(_, num_tag, _)| *num_tag) + .max() + .unwrap_or(0); + let physical_num_field = table_specs + .iter() + .map(|(_, _, num_field)| *num_field) + .max() + .unwrap_or(0); + + // Register a metric engine physical table with internal columns. + { + let mut columns = vec![ + ColumnSchema::new( + DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(), + ConcreteDataType::uint32_datatype(), + false, + ), + ColumnSchema::new( + DATA_SCHEMA_TSID_COLUMN_NAME.to_string(), + ConcreteDataType::uint64_datatype(), + false, + ), + ]; + for i in 0..physical_num_tag { + columns.push(ColumnSchema::new( + format!("tag_{i}"), + ConcreteDataType::string_datatype(), + false, + )); + } + columns.push( + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ); + for i in 0..physical_num_field { + columns.push(ColumnSchema::new( + format!("field_{i}"), + ConcreteDataType::float64_datatype(), + true, + )); + } + + let schema = Arc::new(Schema::new(columns)); + let primary_key_indices = (0..(2 + physical_num_tag)).collect::>(); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(primary_key_indices) + .value_indices( + (2 + physical_num_tag..2 + physical_num_tag + 1 + physical_num_field).collect(), + ) + .engine(METRIC_ENGINE_NAME.to_string()) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .table_id(physical_table_id) + .name(physical_table_name) + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: physical_table_name.to_string(), + table_id: physical_table_id, + table, + }) + .is_ok() + ); + } + + // Register metric engine logical tables without `__tsid`, referencing the physical table. + for (idx, ((schema_name, table_name), num_tag, num_field)) in table_specs.iter().enumerate() { + let mut columns = vec![]; + for i in 0..*num_tag { + columns.push(ColumnSchema::new( + format!("tag_{i}"), + ConcreteDataType::string_datatype(), + false, + )); + } + columns.push( + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ); + for i in 0..*num_field { + columns.push(ColumnSchema::new( + format!("field_{i}"), + ConcreteDataType::float64_datatype(), + true, + )); + } + + let schema = Arc::new(Schema::new(columns)); + let mut options = table::requests::TableOptions::default(); + options.extra_options.insert( + LOGICAL_TABLE_METADATA_KEY.to_string(), + physical_table_name.to_string(), + ); + let table_id = 1024u32 + idx as u32; + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices((0..*num_tag).collect()) + .value_indices((*num_tag + 1..*num_tag + 1 + *num_field).collect()) + .engine(METRIC_ENGINE_NAME.to_string()) + .options(options) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .table_id(table_id) + .name(table_name.clone()) + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: schema_name.clone(), + table_name: table_name.clone(), + table_id, + table, + }) + .is_ok() + ); + } + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +async fn build_test_table_provider_with_fields( + table_name_tuples: &[(String, String)], + tags: &[&str], +) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + for (schema_name, table_name) in table_name_tuples { + let mut columns = vec![]; + let num_tag = tags.len(); + for tag in tags { + columns.push(ColumnSchema::new( + tag.to_string(), + ConcreteDataType::string_datatype(), + false, + )); + } + columns.push( + ColumnSchema::new( + greptime_timestamp().to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ); + columns.push(ColumnSchema::new( + greptime_value().to_string(), + ConcreteDataType::float64_datatype(), + true, + )); + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices((0..num_tag).collect()) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name(table_name.clone()) + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: schema_name.clone(), + table_name: table_name.clone(), + table_id: 1024, + table, + }) + .is_ok() + ); + } + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +// { +// input: `abs(some_metric{foo!="bar"})`, +// expected: &Call{ +// Func: MustGetFunction("abs"), +// Args: Expressions{ +// &VectorSelector{ +// Name: "some_metric", +// LabelMatchers: []*labels.Matcher{ +// MustLabelMatcher(labels.MatchNotEqual, "foo", "bar"), +// MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"), +// }, +// }, +// }, +// }, +// }, +async fn do_single_instant_function_call(fn_name: &'static str, plan_name: &str) { + let prom_expr = parser::parse(&format!("{fn_name}(some_metric{{tag_0!=\"bar\"}})")).unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let expected = String::from( + "Filter: TEMPLATE(field_0) IS NOT NULL [timestamp:Timestamp(ms), TEMPLATE(field_0):Float64;N, tag_0:Utf8]\ + \n Projection: some_metric.timestamp, TEMPLATE(some_metric.field_0) AS TEMPLATE(field_0), some_metric.tag_0 [timestamp:Timestamp(ms), TEMPLATE(field_0):Float64;N, tag_0:Utf8]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]" + ).replace("TEMPLATE", plan_name); + + assert_eq!(plan.display_indent_schema().to_string(), expected); +} + +#[tokio::test] +async fn single_abs() { + do_single_instant_function_call("abs", "abs").await; +} + +#[tokio::test] +#[should_panic] +async fn single_absent() { + do_single_instant_function_call("absent", "").await; +} + +#[tokio::test] +async fn single_ceil() { + do_single_instant_function_call("ceil", "ceil").await; +} + +#[tokio::test] +async fn single_exp() { + do_single_instant_function_call("exp", "exp").await; +} + +#[tokio::test] +async fn single_ln() { + do_single_instant_function_call("ln", "ln").await; +} + +#[tokio::test] +async fn single_log2() { + do_single_instant_function_call("log2", "log2").await; +} + +#[tokio::test] +async fn single_log10() { + do_single_instant_function_call("log10", "log10").await; +} + +#[tokio::test] +#[should_panic] +async fn single_scalar() { + do_single_instant_function_call("scalar", "").await; +} + +#[tokio::test] +#[should_panic] +async fn single_sgn() { + do_single_instant_function_call("sgn", "").await; +} + +#[tokio::test] +#[should_panic] +async fn single_sort() { + do_single_instant_function_call("sort", "").await; +} + +#[tokio::test] +#[should_panic] +async fn single_sort_desc() { + do_single_instant_function_call("sort_desc", "").await; +} + +#[tokio::test] +async fn single_sqrt() { + do_single_instant_function_call("sqrt", "sqrt").await; +} + +#[tokio::test] +async fn single_timestamp_plan_preserves_source_value() { + let eval_stmt = build_eval_stmt(r#"timestamp(some_metric{tag_0!="bar"})"#); + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let expected = String::from( + "Filter: value IS NOT NULL [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\ + \n Projection: some_metric.timestamp, value AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\ + \n Projection: some_metric.timestamp, __promql_timestamp_value_ AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\ + \n Projection: some_metric.tag_0, some_metric.timestamp, some_metric.field_0, CAST(CAST(CAST(CAST(some_metric.timestamp AS Int64) AS Decimal128(19, 0)) * Decimal128(1,1,0) + Decimal128(0,19,0) AS Int64) AS Float64) / Float64(1000) AS __promql_timestamp_value_ [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + assert_eq!(plan.display_indent_schema().to_string(), expected); +} + +#[tokio::test] +async fn single_acos() { + do_single_instant_function_call("acos", "acos").await; +} + +#[tokio::test] +#[should_panic] +async fn single_acosh() { + do_single_instant_function_call("acosh", "").await; +} + +#[tokio::test] +async fn single_asin() { + do_single_instant_function_call("asin", "asin").await; +} + +#[tokio::test] +#[should_panic] +async fn single_asinh() { + do_single_instant_function_call("asinh", "").await; +} + +#[tokio::test] +async fn single_atan() { + do_single_instant_function_call("atan", "atan").await; +} + +#[tokio::test] +#[should_panic] +async fn single_atanh() { + do_single_instant_function_call("atanh", "").await; +} + +#[tokio::test] +async fn single_cos() { + do_single_instant_function_call("cos", "cos").await; +} + +#[tokio::test] +#[should_panic] +async fn single_cosh() { + do_single_instant_function_call("cosh", "").await; +} + +#[tokio::test] +async fn single_sin() { + do_single_instant_function_call("sin", "sin").await; +} + +#[tokio::test] +#[should_panic] +async fn single_sinh() { + do_single_instant_function_call("sinh", "").await; +} + +#[tokio::test] +async fn single_tan() { + do_single_instant_function_call("tan", "tan").await; +} + +#[tokio::test] +#[should_panic] +async fn single_tanh() { + do_single_instant_function_call("tanh", "").await; +} + +#[tokio::test] +#[should_panic] +async fn single_deg() { + do_single_instant_function_call("deg", "").await; +} + +#[tokio::test] +#[should_panic] +async fn single_rad() { + do_single_instant_function_call("rad", "").await; +} + +// { +// input: "avg by (foo)(some_metric)", +// expected: &AggregateExpr{ +// Op: AVG, +// Expr: &VectorSelector{ +// Name: "some_metric", +// LabelMatchers: []*labels.Matcher{ +// MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"), +// }, +// PosRange: PositionRange{ +// Start: 13, +// End: 24, +// }, +// }, +// Grouping: []string{"foo"}, +// PosRange: PositionRange{ +// Start: 0, +// End: 25, +// }, +// }, +// }, +async fn do_aggregate_expr_plan(fn_name: &str, plan_name: &str) { + let prom_expr = parser::parse(&format!( + "{fn_name} by (tag_1)(some_metric{{tag_0!=\"bar\"}})", + )) + .unwrap(); + let mut eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + // test group by + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 2, + 2, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let expected_no_without = String::from( + "Sort: some_metric.tag_1 ASC NULLS LAST, some_metric.timestamp ASC NULLS LAST [tag_1:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ + \n Aggregate: groupBy=[[some_metric.tag_1, some_metric.timestamp]], aggr=[[TEMPLATE(some_metric.field_0), TEMPLATE(some_metric.field_1)]] [tag_1:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.tag_1 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ + \n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]" + ).replace("TEMPLATE", plan_name); + assert_eq!( + plan.display_indent_schema().to_string(), + expected_no_without + ); + + // test group without + if let PromExpr::Aggregate(AggregateExpr { modifier, .. }) = &mut eval_stmt.expr { + *modifier = Some(LabelModifier::Exclude(Labels { + labels: vec![String::from("tag_1")].into_iter().collect(), + })); + } + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 2, + 2, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let expected_without = String::from( + "Sort: some_metric.tag_0 ASC NULLS LAST, some_metric.timestamp ASC NULLS LAST [tag_0:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ + \n Aggregate: groupBy=[[some_metric.tag_0, some_metric.timestamp]], aggr=[[TEMPLATE(some_metric.field_0), TEMPLATE(some_metric.field_1)]] [tag_0:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.tag_1 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ + \n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]" + ).replace("TEMPLATE", plan_name); + assert_eq!(plan.display_indent_schema().to_string(), expected_without); +} + +#[tokio::test] +async fn aggregate_sum() { + do_aggregate_expr_plan("sum", "sum").await; +} + +#[tokio::test] +async fn tsid_is_used_for_series_divide_when_available() { + let prom_expr = parser::parse("some_metric").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); + assert!(plan_str.contains("__tsid ASC NULLS FIRST")); + assert!( + !plan + .schema() + .fields() + .iter() + .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME) + ); + + let manipulate = find_instant_manipulate(&plan).unwrap(); + let exec = manipulate.to_execution_plan(Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new( + &[], + Arc::new( + datafusion_expr::UserDefinedLogicalNodeCore::inputs(manipulate)[0] + .schema() + .as_arrow() + .clone(), + ), + None, + ) + .unwrap(), + )))); + assert!(format!("{exec:?}").contains("reuse_tsid_column: true")); +} + +#[tokio::test] +async fn default_binary_join_uses_tsid_when_available() { + let eval_stmt = build_eval_stmt("some_metric / some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), + "{plan_str}" + ); + assert!( + !plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"), + "{plan_str}" + ); +} + +#[tokio::test] +async fn reject_binary_fill_modifiers() { + let state = build_query_engine_state(); + + for query in [ + "some_metric + fill(0) some_alt_metric", + "some_metric + fill_left(0) some_alt_metric", + "some_metric + fill_right(0) some_alt_metric", + "(some_metric + fill(0) some_alt_metric) + some_metric", + ] { + let eval_stmt = build_eval_stmt(query); + let table_provider = build_test_table_provider(&[], 0, 0).await; + let err = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &state) + .await + .unwrap_err(); + + assert!( + matches!( + &err, + crate::promql::error::Error::UnsupportedExpr { name, .. } + if name == "PromQL fill modifiers" + ), + "{err}" + ); + } +} + +#[tokio::test] +async fn timestamp_binary_join_falls_back_when_tsid_is_projected_out() { + for query in [ + "timestamp(some_metric) / some_metric", + "some_metric / timestamp(some_metric)", + ] { + let eval_stmt = build_eval_stmt(query); + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(!plan_str.contains("__tsid ="), "{query}: {plan_str}"); + assert!( + plan_str.contains("lhs.tag_0 = rhs.tag_0"), + "{query}: {plan_str}" + ); + assert!( + !plan + .schema() + .fields() + .iter() + .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME), + "{query}: {plan_str}" + ); + } +} + +#[tokio::test] +async fn timestamp_binary_join_rejects_default_matching_on_mismatched_labels() { + let eval_stmt = build_eval_stmt("timestamp(left_host_job) / right_by_job"); + + let table_provider = build_test_table_provider_with_tsid_tag_fields(&[ + ( + (DEFAULT_SCHEMA_NAME.to_string(), "left_host_job".to_string()), + 2, + 1, + ), + ( + (DEFAULT_SCHEMA_NAME.to_string(), "right_by_job".to_string()), + 1, + 1, + ), + ]) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let plan_str = plan.display_indent_schema().to_string(); + + assert!( + plan_str.contains("Boolean(false)") || plan_str.contains("false"), + "{plan_str}" + ); +} + +#[tokio::test] +async fn tsid_is_preserved_for_nested_default_binary_joins() { + let eval_stmt = build_eval_stmt("(some_metric - some_alt_metric) / some_third_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_third_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); + assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); +} + +#[tokio::test] +async fn repeated_tsid_binary_operand_reuses_leaf_plan() { + let eval_stmt = build_eval_stmt("((some_metric - some_alt_metric) / some_metric) * 100"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}"); + assert_eq!( + plan_str + .matches("Filter: phy.__table_id = UInt32(1024)") + .count(), + 1, + "{plan_str}" + ); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 2, + "{plan_str}" + ); + assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); +} + +#[tokio::test] +async fn repeated_tsid_binary_operand_reuses_shorter_field_side() { + let eval_stmt = + build_eval_stmt("((two_field_metric - one_field_metric) / one_field_metric) * 100"); + + let table_provider = build_test_table_provider_with_tsid_fields( + &[ + ( + ( + DEFAULT_SCHEMA_NAME.to_string(), + "two_field_metric".to_string(), + ), + 2, + ), + ( + ( + DEFAULT_SCHEMA_NAME.to_string(), + "one_field_metric".to_string(), + ), + 1, + ), + ], + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let field_names = plan + .schema() + .fields() + .iter() + .map(|field| field.name().clone()) + .collect::>(); + let value_columns = field_names + .iter() + .filter(|name| { + *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME + }) + .count(); + assert_eq!(value_columns, 1, "{field_names:?}"); + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}"); + assert_eq!( + plan_str + .matches("Filter: phy.__table_id = UInt32(1025)") + .count(), + 1, + "{plan_str}" + ); + assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); +} + +#[tokio::test] +async fn binary_island_reuses_self_operand_without_join() { + let eval_stmt = build_eval_stmt("some_metric / some_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 0, "{plan_str}"); + assert_eq!( + plan_str + .matches("Filter: phy.__table_id = UInt32(1024)") + .count(), + 1, + "{plan_str}" + ); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 1, + "{plan_str}" + ); +} + +#[tokio::test] +async fn binary_island_reuses_leaf_across_two_branches() { + let eval_stmt = + build_eval_stmt("(some_metric + some_alt_metric) / (some_metric + third_metric)"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + (DEFAULT_SCHEMA_NAME.to_string(), "third_metric".to_string()), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); + assert_eq!( + plan_str + .matches("Filter: phy.__table_id = UInt32(1024)") + .count(), + 1, + "{plan_str}" + ); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 3, + "{plan_str}" + ); +} + +#[tokio::test] +async fn binary_island_generated_alias_avoids_user_column_names() { + let eval_stmt = build_eval_stmt("(some_metric + some_alt_metric) / some_metric"); + + let table_provider = build_test_table_provider_with_fields( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + &["prom_v0", "__prom_v0"], + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let field_names = plan.schema().field_names(); + assert!(field_names.iter().any(|name| name.ends_with(".prom_v0"))); + assert!(field_names.iter().any(|name| name.ends_with(".__prom_v0"))); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("SubqueryAlias: __prom_v0"), "{plan_str}"); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 2, + "{plan_str}" + ); +} + +#[tokio::test] +async fn binary_island_clears_qualifier_for_nested_unary_projection() { + let eval_stmt = build_eval_stmt("-((some_metric + some_alt_metric) / some_metric)"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}"); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 2, + "{plan_str}" + ); +} + +#[tokio::test] +async fn binary_island_keeps_distinct_matcher_leaves() { + let eval_stmt = build_eval_stmt( + "(some_metric{tag_0=\"foo\"} + some_alt_metric) / some_metric{tag_0=\"bar\"}", + ); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 3, + "{plan_str}" + ); +} + +#[tokio::test] +async fn binary_island_keeps_offset_leaves_distinct() { + let eval_stmt = build_eval_stmt("(some_metric offset 5m + some_alt_metric) / some_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 3, + "{plan_str}" + ); +} + +#[tokio::test] +async fn binary_island_falls_back_for_group_modifier() { + let eval_stmt = + build_eval_stmt("(some_metric + ignoring(tag_0) group_left some_alt_metric) / some_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 3, + "{plan_str}" + ); +} + +#[tokio::test] +async fn binary_island_falls_back_for_comparison_filter() { + let eval_stmt = build_eval_stmt("(some_metric > some_alt_metric) / some_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 3, + "{plan_str}" + ); +} + +#[tokio::test] +async fn tsid_binary_join_uses_shorter_field_side() { + let eval_stmt = build_eval_stmt("one_field_metric / two_field_metric"); + + let table_provider = build_test_table_provider_with_tsid_fields( + &[ + ( + ( + DEFAULT_SCHEMA_NAME.to_string(), + "one_field_metric".to_string(), + ), + 1, + ), + ( + ( + DEFAULT_SCHEMA_NAME.to_string(), + "two_field_metric".to_string(), + ), + 2, + ), + ], + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let field_names = plan + .schema() + .fields() + .iter() + .map(|field| field.name().clone()) + .collect::>(); + let value_columns = field_names + .iter() + .filter(|name| { + *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME + }) + .count(); + assert_eq!(value_columns, 1, "{field_names:?}"); +} + +#[tokio::test] +async fn comparison_binary_join_uses_shorter_field_side() { + let eval_stmt = build_eval_stmt("two_field_metric > one_field_metric"); + + let table_provider = build_test_table_provider_with_tsid_fields( + &[ + ( + ( + DEFAULT_SCHEMA_NAME.to_string(), + "two_field_metric".to_string(), + ), + 2, + ), + ( + ( + DEFAULT_SCHEMA_NAME.to_string(), + "one_field_metric".to_string(), + ), + 1, + ), + ], + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let field_names = plan + .schema() + .fields() + .iter() + .map(|field| field.name().clone()) + .collect::>(); + assert!( + field_names.iter().any(|name| name == "field_0"), + "{field_names:?}" + ); + assert!( + !field_names.iter().any(|name| name == "field_1"), + "{field_names:?}" + ); +} + +#[tokio::test] +async fn label_matching_modifier_disables_tsid_binary_join() { + let eval_stmt = build_eval_stmt("some_metric / ignoring(tag_0) some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 2, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(!plan_str.contains("__tsid ="), "{plan_str}"); + assert!( + plan_str.contains("some_metric.tag_1 = some_alt_metric.tag_1"), + "{plan_str}" + ); +} + +#[tokio::test] +async fn ignoring_absent_label_keeps_tsid_binary_join() { + let eval_stmt = build_eval_stmt("some_metric / ignoring(missing) some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 2, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), + "{plan_str}" + ); + assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); + assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); +} + +#[tokio::test] +async fn range_function_keeps_tsid_for_absent_ignoring_binary_join() { + let eval_stmt = build_eval_stmt("rate(some_metric[5m]) / ignoring(missing) some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 2, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), + "{plan_str}" + ); + assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); + assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); +} + +#[tokio::test] +async fn on_full_label_set_keeps_tsid_binary_join() { + let eval_stmt = build_eval_stmt("some_metric / on(tag_0, tag_1) some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 2, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), + "{plan_str}" + ); + assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); + assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); +} + +#[tokio::test] +async fn on_partial_label_set_disables_tsid_binary_join() { + let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 2, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(!plan_str.contains("__tsid ="), "{plan_str}"); + assert!( + plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"), + "{plan_str}" + ); + assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); +} + +#[tokio::test] +async fn on_label_set_must_cover_both_sides_to_use_tsid_binary_join() { + let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid_tag_fields(&[ + ( + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + 2, + 1, + ), + ( + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + 1, + 1, + ), + ]) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(!plan_str.contains("__tsid ="), "{plan_str}"); + assert!( + plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"), + "{plan_str}" + ); + assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); +} + +#[tokio::test] +async fn comparison_binary_join_uses_tsid_and_keeps_it_in_filtered_result() { + let eval_stmt = build_eval_stmt("some_metric > some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 2, + 1, + ) + .await; + let mut planner = PromPlanner { + table_provider, + ctx: PromPlannerContext::from_eval_stmt(&eval_stmt), + promql_annotations: None, + }; + let plan = planner + .prom_expr_to_plan(&eval_stmt.expr, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), + "{plan_str}" + ); + assert!( + plan.schema() + .fields() + .iter() + .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME), + "{plan_str}" + ); + assert!(planner.ctx.use_tsid, "{plan_str}"); +} + +#[tokio::test] +async fn comparison_bool_binary_join_uses_tsid_when_available() { + let eval_stmt = build_eval_stmt("some_metric > bool some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 2, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), + "{plan_str}" + ); + assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); + assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); +} + +#[tokio::test] +async fn scalar_count_count_range_keeps_full_window() { + let plan_str = build_optimized_tsid_plan( + "scalar(count(count(some_metric) by (tag_0)))", + 1, + 1, + 100_000, + 1, + ) + .await; + assert!(plan_str.contains("ScalarCalculate: tags=[]")); + assert!(plan_str.contains("PromInstantManipulate: range=[0..100000000]")); + assert!(!plan_str.contains("PromInstantManipulate: range=[99999000..99999000]")); +} + +#[tokio::test] +async fn scalar_count_count_rewrite_applies_inside_binary_expr_for_tsid_input() { + let plan_str = build_optimized_tsid_plan( + "sum(irate(some_metric[1h])) / scalar(count(count(some_metric) by (tag_0)))", + 2, + 1, + 10, + 300, + ) + .await; + assert!(plan_str.contains("Distinct:"), "{plan_str}"); +} + +#[tokio::test] +async fn nested_count_rewrite_keeps_full_series_key_with_tsid_input() { + assert_nested_count_rewrite_applies( + "count(count(some_metric) by (tag_0))", + "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(count(some_metric.field_0))]]" + ) + .await; +} + +#[tokio::test] +async fn nested_sum_count_rewrite_keeps_full_series_key_with_tsid_input() { + assert_nested_count_rewrite_applies( + "count(sum(some_metric) by (tag_0))", + "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(sum(some_metric.field_0))]]" + ) + .await; +} + +#[tokio::test] +async fn nested_supported_inner_aggs_rewrite_apply_for_tsid_input() { + for (query, expected_outer_agg) in [ + ( + "count(avg(some_metric) by (tag_0))", + "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(avg(some_metric.field_0))]]", + ), + ( + "count(min(some_metric) by (tag_0))", + "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(min(some_metric.field_0))]]", + ), + ( + "count(max(some_metric) by (tag_0))", + "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(max(some_metric.field_0))]]", + ), + ( + "count(stddev(some_metric) by (tag_0))", + "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(stddev_pop(some_metric.field_0))]]", + ), + ( + "count(stdvar(some_metric) by (tag_0))", + "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(var_pop(some_metric.field_0))]]", + ), + ] { + assert_nested_count_rewrite_applies(query, expected_outer_agg).await; + } +} + +#[tokio::test] +async fn nested_non_count_inner_aggs_rewrite_filter_null_values_for_tsid_input() { + let count_plan = + build_optimized_tsid_plan("count(count(some_metric) by (tag_0))", 2, 1, 100_000, 1).await; + assert!( + !count_plan.contains("some_metric.field_0 IS NOT NULL"), + "{count_plan}" + ); + + for query in [ + "count(sum(some_metric) by (tag_0))", + "count(avg(some_metric) by (tag_0))", + "count(min(some_metric) by (tag_0))", + "count(max(some_metric) by (tag_0))", + "count(stddev(some_metric) by (tag_0))", + "count(stdvar(some_metric) by (tag_0))", + ] { + let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await; + assert!( + plan_str.contains("Filter: some_metric.field_0 IS NOT NULL"), + "{query}: {plan_str}" + ); + } +} + +#[tokio::test] +async fn nested_unsupported_or_non_direct_inner_aggs_do_not_rewrite() { + assert_nested_count_rewrite_missing("count(group(some_metric) by (tag_0))", 2, 1).await; + assert_nested_count_rewrite_missing("count(sum(irate(some_metric[1h])) by (tag_0))", 2, 300) + .await; +} + +#[tokio::test] +async fn physical_table_name_is_not_leaked_in_plan() { + let prom_expr = parser::parse("some_metric").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("TableScan: phy"), "{plan}"); + assert!(plan_str.contains("SubqueryAlias: some_metric")); + assert!(plan_str.contains("Filter: phy.__table_id = UInt32(1024)")); + assert!(!plan_str.contains("TableScan: some_metric")); +} + +#[tokio::test] +async fn sum_without_does_not_group_by_tsid() { + let prom_expr = parser::parse("sum without (tag_0) (some_metric)").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); + + let aggr_line = plan_str + .lines() + .find(|line| line.contains("Aggregate: groupBy=")) + .unwrap(); + assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME)); +} + +#[tokio::test] +async fn topk_without_does_not_partition_by_tsid() { + let prom_expr = parser::parse("topk without (tag_0) (1, some_metric)").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); + + let window_line = plan_str + .lines() + .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()")) + .unwrap(); + let partition_by = window_line + .split("PARTITION BY [") + .nth(1) + .and_then(|s| s.split("] ORDER BY").next()) + .unwrap(); + assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME)); +} + +#[tokio::test] +async fn sum_by_does_not_group_by_tsid() { + let prom_expr = parser::parse("sum by (__tsid) (some_metric)").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); + + let aggr_line = plan_str + .lines() + .find(|line| line.contains("Aggregate: groupBy=")) + .unwrap(); + assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME)); +} + +#[tokio::test] +async fn aggregate_over_binary_time_function_expr() { + for op in ["sum", "min", "max", "avg"] { + let prom_expr = parser::parse(&format!( + "{op} by (tag_0, tag_1, tag_2) (time() - some_metric)" + )) + .unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 3, + 1, + ) + .await; + let plan = + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + let aggr_line = plan_str + .lines() + .find(|line| line.contains("Aggregate: groupBy=")) + .unwrap(); + assert!(aggr_line.contains(op), "{plan_str}"); + assert!(aggr_line.contains("first_value"), "{plan_str}"); + assert!( + !plan + .schema() + .fields() + .iter() + .any(|field| { field.name() == DATA_SCHEMA_TSID_COLUMN_NAME }) + ); + } +} + +#[tokio::test] +async fn topk_by_does_not_partition_by_tsid() { + let prom_expr = parser::parse("topk by (__tsid) (1, some_metric)").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); + + let window_line = plan_str + .lines() + .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()")) + .unwrap(); + let partition_by = window_line + .split("PARTITION BY [") + .nth(1) + .and_then(|s| s.split("] ORDER BY").next()) + .unwrap(); + assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME)); +} + +#[tokio::test] +async fn selector_matcher_on_tsid_does_not_use_internal_column() { + let prom_expr = parser::parse(r#"some_metric{__tsid="123"}"#).unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + fn collect_filter_cols(plan: &LogicalPlan, out: &mut HashSet) { + if let LogicalPlan::Filter(filter) = plan { + datafusion_expr::utils::expr_to_columns(&filter.predicate, out).unwrap(); + } + for input in plan.inputs() { + collect_filter_cols(input, out); + } + } + + let mut filter_cols = HashSet::new(); + collect_filter_cols(&plan, &mut filter_cols); + assert!( + !filter_cols + .iter() + .any(|c| c.name == DATA_SCHEMA_TSID_COLUMN_NAME) + ); +} + +#[tokio::test] +async fn tsid_is_not_used_when_physical_table_is_missing() { + let prom_expr = parser::parse("some_metric").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let catalog_list = MemoryCatalogManager::with_default_setup(); + + // Register a metric engine logical table referencing a missing physical table. + let mut columns = vec![ColumnSchema::new( + "tag_0".to_string(), + ConcreteDataType::string_datatype(), + false, + )]; + columns.push( + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ); + columns.push(ColumnSchema::new( + "field_0".to_string(), + ConcreteDataType::float64_datatype(), + true, + )); + let schema = Arc::new(Schema::new(columns)); + let mut options = table::requests::TableOptions::default(); + options + .extra_options + .insert(LOGICAL_TABLE_METADATA_KEY.to_string(), "phy".to_string()); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![0]) + .value_indices(vec![2]) + .engine(METRIC_ENGINE_NAME.to_string()) + .options(options) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .table_id(1024) + .name("some_metric") + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "some_metric".to_string(), + table_id: 1024, + table, + }) + .unwrap(); + + let table_provider = DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ); + + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]")); + assert!(!plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); +} + +#[tokio::test] +async fn tsid_is_carried_only_when_aggregate_preserves_label_set() { + let prom_expr = parser::parse("sum by (tag_0) (some_metric)").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("first_value") && plan_str.contains("__tsid")); + assert!( + !plan + .schema() + .fields() + .iter() + .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME) + ); + + // Merging aggregate: label set is reduced, tsid should not be carried. + let prom_expr = parser::parse("sum(some_metric)").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let plan_str = plan.display_indent_schema().to_string(); + assert!(!plan_str.contains("first_value")); +} + +#[tokio::test] +async fn or_operator_with_unknown_metric_does_not_require_tsid() { + let prom_expr = parser::parse("unknown_metric or some_metric").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + assert!( + !plan + .schema() + .fields() + .iter() + .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME) + ); +} + +#[tokio::test] +async fn aggregate_avg() { + do_aggregate_expr_plan("avg", "avg").await; +} + +#[tokio::test] +#[should_panic] // output type doesn't match +async fn aggregate_count() { + do_aggregate_expr_plan("count", "count").await; +} + +#[tokio::test] +async fn aggregate_min() { + do_aggregate_expr_plan("min", "min").await; +} + +#[tokio::test] +async fn aggregate_max() { + do_aggregate_expr_plan("max", "max").await; +} + +#[tokio::test] +async fn aggregate_group() { + // Regression test for `group()` aggregator. + // PromQL: sum(group by (cluster)(kubernetes_build_info{service="kubernetes",job="apiserver"})) + // should be plannable, and `group()` should produce constant 1 for each group. + let prom_expr = parser::parse( + "sum(group by (cluster)(kubernetes_build_info{service=\"kubernetes\",job=\"apiserver\"}))", + ) + .unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_fields( + &[( + DEFAULT_SCHEMA_NAME.to_string(), + "kubernetes_build_info".to_string(), + )], + &["cluster", "service", "job"], + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("max(Float64(1")); +} + +#[tokio::test] +async fn aggregate_stddev() { + do_aggregate_expr_plan("stddev", "stddev_pop").await; +} + +#[tokio::test] +async fn aggregate_stdvar() { + do_aggregate_expr_plan("stdvar", "var_pop").await; +} + +// TODO(ruihang): add range fn tests once exprs are ready. + +// { +// input: "some_metric{tag_0="foo"} + some_metric{tag_0="bar"}", +// expected: &BinaryExpr{ +// Op: ADD, +// LHS: &VectorSelector{ +// Name: "a", +// LabelMatchers: []*labels.Matcher{ +// MustLabelMatcher(labels.MatchEqual, "tag_0", "foo"), +// MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"), +// }, +// }, +// RHS: &VectorSelector{ +// Name: "sum", +// LabelMatchers: []*labels.Matcher{ +// MustLabelMatcher(labels.MatchxEqual, "tag_0", "bar"), +// MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"), +// }, +// }, +// VectorMatching: &VectorMatching{}, +// }, +// }, +#[tokio::test] +async fn binary_op_column_column() { + let prom_expr = + parser::parse(r#"some_metric{tag_0="foo"} + some_metric{tag_0="bar"}"#).unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let expected = String::from( + "Projection: rhs.tag_0, rhs.timestamp, CAST(lhs.field_0 AS Float64) + CAST(rhs.field_0 AS Float64) AS lhs.field_0 + rhs.field_0 [tag_0:Utf8, timestamp:Timestamp(ms), lhs.field_0 + rhs.field_0:Float64;N]\ + \n Inner Join: lhs.tag_0 = rhs.tag_0, lhs.timestamp = rhs.timestamp [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n SubqueryAlias: lhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.tag_0 = Utf8(\"foo\") AND some_metric.tag_0 = Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n SubqueryAlias: rhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.tag_0 = Utf8(\"bar\") AND some_metric.tag_0 = Utf8(\"foo\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + assert_eq!(plan.display_indent_schema().to_string(), expected); +} + +async fn indie_query_plan_compare>(query: &str, expected: T) { + let prom_expr = parser::parse(query).unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + "greptime_private".to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + assert_eq!(plan.display_indent_schema().to_string(), expected.as_ref()); +} + +#[tokio::test] +async fn binary_op_literal_column() { + let query = r#"1 + some_metric{tag_0="bar"}"#; + let expected = String::from( + "Projection: some_metric.tag_0, some_metric.timestamp, Float64(1) + CAST(some_metric.field_0 AS Float64) AS Float64(1) + field_0 [tag_0:Utf8, timestamp:Timestamp(ms), Float64(1) + field_0:Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.tag_0 = Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +#[tokio::test] +async fn binary_op_literal_literal() { + let query = r#"1 + 1"#; + let expected = r#"EmptyMetric: range=[0..100000000], interval=[5000] [time:Timestamp(ms), value:Float64;N] + TableScan: dummy [time:Timestamp(ms), value:Float64;N]"#; + indie_query_plan_compare(query, expected).await; +} + +#[tokio::test] +async fn simple_bool_grammar() { + let query = "some_metric != bool 1.2345"; + let expected = String::from( + "Projection: some_metric.tag_0, some_metric.timestamp, CAST(some_metric.field_0 != Float64(1.2345) AS Float64) AS field_0 != Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0 != Float64(1.2345):Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +#[tokio::test] +async fn bool_with_additional_arithmetic() { + let query = "some_metric + (1 == bool 2)"; + let expected = String::from( + "Projection: some_metric.tag_0, some_metric.timestamp, CAST(some_metric.field_0 AS Float64) + CAST(Float64(1) = Float64(2) AS Float64) AS field_0 + Float64(1) = Float64(2) [tag_0:Utf8, timestamp:Timestamp(ms), field_0 + Float64(1) = Float64(2):Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +#[tokio::test] +async fn simple_unary() { + let query = "-some_metric"; + let expected = String::from( + "Projection: some_metric.tag_0, some_metric.timestamp, (- some_metric.field_0) AS (- field_0) [tag_0:Utf8, timestamp:Timestamp(ms), (- field_0):Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +#[tokio::test] +async fn increase_aggr() { + let query = "increase(some_metric[5m])"; + let expected = String::from( + "Filter: prom_increase(timestamp_range,field_0,timestamp,Int64(300000)) IS NOT NULL [timestamp:Timestamp(ms), prom_increase(timestamp_range,field_0,timestamp,Int64(300000)):Float64;N, tag_0:Utf8]\ + \n Projection: some_metric.timestamp, prom_increase(timestamp_range, field_0, some_metric.timestamp, Int64(300000)) AS prom_increase(timestamp_range,field_0,timestamp,Int64(300000)), some_metric.tag_0 [timestamp:Timestamp(ms), prom_increase(timestamp_range,field_0,timestamp,Int64(300000)):Float64;N, tag_0:Utf8]\ + \n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[300000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\ + \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.timestamp >= TimestampMillisecond(-299999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +async fn native_histogram_plan(query: &str) -> String { + let table_provider = build_test_native_histogram_table_provider("some_metric").await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt(query), + &build_query_engine_state(), + ) + .await + .unwrap(); + plan.display_indent_schema().to_string() +} + +#[tokio::test] +async fn native_histogram_count_uses_native_udf() { + let plan = native_histogram_plan("histogram_count(some_metric)").await; + + assert!(plan.contains("prom_native_histogram_count"), "{plan}"); + assert!(!plan.contains("HistogramFold:"), "{plan}"); +} + +#[tokio::test] +async fn timestamp_filters_native_histogram_stale_marker_before_projection() { + let mut stale = direct_or_histogram(); + stale.sum = f64::from_bits(PROMETHEUS_STALE_NAN_BITS); + let table = operator_metric_table( + "stale_histogram", + 2_100, + "a", + None, + DirectOrValue::NativeHistogram(stale), + ); + let catalog = MemoryCatalogManager::with_default_setup(); + catalog + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "stale_histogram".to_string(), + table_id: 2_100, + table, + }) + .unwrap(); + let provider = DfTableSourceProvider::new( + catalog, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ); + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + provider, + &operator_eval_stmt("timestamp(stale_histogram)"), + &state, + ) + .await + .unwrap(); + let plan_text = plan.display_indent_schema().to_string(); + assert!(plan_text.contains(TIMESTAMP_VALUE_PREFIX), "{plan_text}"); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); +} + +#[tokio::test] +async fn timestamp_filters_stale_marker_from_mixed_sample_companion() { + let histograms = build_histogram_array(&[None]); + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "timestamp", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + ), + Field::new( + greptime_native_histogram(), + histograms.data_type().clone(), + true, + ), + Field::new(greptime_value(), ArrowDataType::Float64, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1_000])), + histograms, + Arc::new(Float64Array::from(vec![f64::from_bits( + PROMETHEUS_STALE_NAN_BITS, + )])), + ], + ) + .unwrap(); + let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); + let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None) + .unwrap() + .build() + .unwrap(); + let input = LogicalPlan::Extension(Extension { + node: Arc::new(SeriesDivide::new( + Vec::new(), + "timestamp".to_string(), + input, + )), + }); + let input = LogicalPlan::Extension(Extension { + node: Arc::new(InstantManipulate::new( + 1_000, + 1_000, + 5_000, + 1_000, + 0, + "timestamp".to_string(), + Vec::new(), + Some(greptime_native_histogram().to_string()), + input, + )), + }); + // Match timestamp()'s parent projection, which otherwise prunes the companion lane. + let plan = LogicalPlanBuilder::from(input) + .project([col("timestamp")]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); +} + +#[tokio::test] +async fn native_histogram_rate_can_feed_count() { + let plan = native_histogram_plan("histogram_count(rate(some_metric[5m]))").await; + + assert!(plan.contains("prom_native_histogram_rate"), "{plan}"); + assert!(plan.contains("prom_native_histogram_count"), "{plan}"); +} + +#[tokio::test] +async fn native_histogram_quantile_skips_classic_fold() { + let plan = native_histogram_plan("histogram_quantile(0.9, some_metric)").await; + + assert!(plan.contains("prom_native_histogram_quantile"), "{plan}"); + assert!(!plan.contains("HistogramFold:"), "{plan}"); + assert!(plan.contains("some_metric.le"), "{plan}"); + // The phi literal is threaded into the native quantile UDF as its second argument. + assert!(plan.contains("Float64(0.9)"), "{plan}"); + // The empty-values filter drops NULL quantile results so the output is empty + // when all native histogram samples are dropped. + assert!(plan.contains("IS NOT NULL"), "{plan}"); +} + +#[tokio::test] +async fn mixed_native_histogram_quantile_uses_histogram_field() { + let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt("histogram_quantile(0.9, some_metric)"), + &build_query_engine_state(), + ) + .await + .unwrap() + .display_indent_schema() + .to_string(); + + assert!( + plan.contains("prom_native_histogram_quantile(greptime_native_histogram"), + "{plan}" + ); + assert!(!plan.contains("EmptyRelation"), "{plan}"); +} + +#[tokio::test] +async fn mixed_histogram_helpers_execute_classic_and_native_samples() { + let state = build_query_engine_state(); + for (query, expected) in [ + ( + "histogram_quantile(0.5, mixed_histogram)", + vec![("classic", 1.0), ("native", 0.0)], + ), + ( + "histogram_fraction(-Inf, +Inf, mixed_histogram)", + vec![("classic", 1.0), ("native", 1.0)], + ), + ] { + let plan = PromPlanner::stmt_to_plan( + classic_and_native_histogram_table_provider("native", None, direct_or_histogram()), + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap(); + let plan_text = plan.display_indent_schema().to_string(); + assert!(plan_text.contains("HistogramFold:"), "{plan_text}"); + assert!(plan_text.contains("prom_native_histogram_"), "{plan_text}"); + let value_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + let mut actual = batches + .iter() + .flat_map(|batch| { + let tags = batch + .column_by_name("tag") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column_by_name(&value_field) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()).map(|row| (tags.value(row), values.value(row))) + }) + .collect::>(); + actual.sort_by_key(|(tag, _)| *tag); + assert_eq!(actual, expected, "{query}"); + } +} + +#[tokio::test] +async fn mixed_histogram_helpers_report_annotations() { + let state = build_query_engine_state(); + let mut native_histogram = direct_or_histogram(); + native_histogram.count = 2.0; + native_histogram.sum = f64::NAN; + for (native_tag, expected_rows, expected_warnings, expected_infos) in [ + ( + "classic", + 0, + vec!["vector contains a mix of classic and native histograms"], + vec![], + ), + ( + "native", + 2, + vec![], + vec!["input to histogram_quantile has NaN observations, result is skewed higher"], + ), + ] { + let collector = PromqlAnnotationCollector::default(); + let plan = PromPlanner::stmt_to_plan_with_annotations( + classic_and_native_histogram_table_provider(native_tag, None, native_histogram.clone()), + &operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"), + &state, + Some(collector.clone()), + ) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + expected_rows + ); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert_eq!(warnings, expected_warnings); + assert_eq!(infos, expected_infos); + } +} + +#[tokio::test] +async fn mixed_histogram_helper_preserves_native_le_and_scans_once() { + let state = build_query_engine_state(); + let mut stmt = operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"); + stmt.end = UNIX_EPOCH.checked_add(Duration::from_secs(2)).unwrap(); + let plan = PromPlanner::stmt_to_plan( + classic_and_native_histogram_table_provider( + "classic", + Some("native"), + direct_or_histogram(), + ), + &stmt, + &state, + ) + .await + .unwrap(); + let plan_text = plan.display_indent_schema().to_string(); + assert_eq!( + plan_text.matches("TableScan: mixed_histogram").count(), + 1, + "{plan_text}" + ); + + let value_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + let (_, batches) = execute(plan, &state).await; + let mut actual = batches + .iter() + .flat_map(|batch| { + let le = batch + .column_by_name(LE_COLUMN_NAME) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let timestamps = batch + .column_by_name("timestamp") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column_by_name(&value_field) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()).map(|row| { + ( + timestamps.value(row), + (!le.is_null(row)).then(|| le.value(row).to_string()), + values.value(row), + ) + }) + }) + .collect::>(); + actual.sort_by(|lhs, rhs| (lhs.0, &lhs.1).cmp(&(rhs.0, &rhs.1))); + assert_eq!( + actual, + vec![ + (1_000, None, 1.0), + (1_000, Some("native".to_string()), 0.0), + (2_000, None, 1.0), + (2_000, Some("native".to_string()), 0.0), + ] + ); +} + +#[tokio::test] +async fn nested_histogram_helpers_ignore_unparsable_bucket_labels() { + let state = build_query_engine_state(); + for native_le in [None, Some("native")] { + for query in [ + "histogram_quantile(0.5, histogram_quantile(0.5, mixed_histogram))", + "histogram_fraction(-Inf, +Inf, histogram_fraction(-Inf, +Inf, mixed_histogram))", + ] { + let plan = PromPlanner::stmt_to_plan( + classic_and_native_histogram_table_provider( + "native", + native_le, + direct_or_histogram(), + ), + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + 0, + "native_le={native_le:?}, query={query}" + ); + } + } +} + +#[tokio::test] +async fn native_histogram_quantile_rejects_multi_field_input() { + let table_provider = build_test_multi_histogram_table_provider("some_metric").await; + let result = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt("histogram_quantile(0.9, some_metric)"), + &build_query_engine_state(), + ) + .await; + + let err = result.expect_err("histogram_quantile on two native histogram fields must fail"); + assert!( + err.to_string() + .contains("Multi fields calculation is not supported in histogram_quantile"), + "{err}" + ); +} + +#[tokio::test] +async fn native_histogram_topk_uses_drop_udf() { + let plan = native_histogram_plan("topk(1, some_metric)").await; + + assert!(plan.contains("prom_native_histogram_drop_float"), "{plan}"); + assert!( + plan.contains("Filter: prom_native_histogram_drop_float") && plan.contains("IS NOT NULL"), + "{plan}" + ); +} + +#[tokio::test] +async fn mixed_or_topk_bottomk_ignore_native_histograms() { + for op in ["topk", "bottomk"] { + let collector = PromqlAnnotationCollector::default(); + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan_with_annotations( + operator_table_provider(), + &operator_eval_stmt(&format!("{op}(1, lf or on(tag) lh)")), + &state, + Some(collector.clone()), + ) + .await + .unwrap(); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + assert!( + plan.schema() + .fields() + .iter() + .all(|field| field.data_type() != &PromPlanner::native_histogram_arrow_type()), + "{plan:?}" + ); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(values(&batches, &float_field), vec![2.0], "{op}"); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty()); + assert_eq!( + infos, + vec![format!( + "{op}: dropped native histogram samples because this aggregation is not supported for native histograms" + )] + ); + } +} + +#[tokio::test] +async fn native_histogram_scalar_is_ignored_before_scalar_calculate() { + let plan = native_histogram_plan("scalar(some_metric)").await; + + assert!(plan.contains("ScalarCalculate"), "{plan}"); + assert!(plan.contains("Filter: Boolean(false)"), "{plan}"); + assert!(!plan.contains("prom_native_histogram_drop"), "{plan}"); +} + +#[tokio::test] +async fn native_histogram_value_sort_is_empty_but_label_sort_preserves_samples() { + for function in ["sort", "sort_desc"] { + let plan = native_histogram_plan(&format!("{function}(some_metric)")).await; + + assert!(plan.contains("Float64(NULL) IS NOT NULL"), "{plan}"); + assert!( + !plan.contains(&format!("Sort: {}", greptime_native_histogram())), + "{plan}" + ); + assert!(!plan.contains("prom_native_histogram_drop"), "{plan}"); + } + + for (function, direction) in [("sort_by_label", "ASC"), ("sort_by_label_desc", "DESC")] { + let plan = native_histogram_plan(&format!("{function}(some_metric, \"tag_0\")")).await; + + assert!(plan.contains(&format!("tag_0 {direction}")), "{plan}"); + assert!(plan.contains(greptime_native_histogram()), "{plan}"); + assert!(!plan.contains("Float64(NULL) IS NOT NULL"), "{plan}"); + } +} + +#[tokio::test] +async fn unsupported_native_histogram_functions_use_drop_udf() { + for query in [ + "deriv(some_metric[5m])", + "min_over_time(some_metric[5m])", + "quantile_over_time(0.9, some_metric[5m])", + "predict_linear(some_metric[5m], 60)", + "round(some_metric)", + "abs(some_metric)", + ] { + let plan = native_histogram_plan(query).await; + + assert!( + plan.contains("prom_native_histogram_drop_float"), + "{query}\n{plan}" + ); + } +} + +#[tokio::test] +async fn native_histogram_absent_over_time_uses_native_udf() { + let plan = native_histogram_plan("absent_over_time(some_metric[5m])").await; + + assert!( + plan.contains("prom_native_histogram_absent_over_time"), + "{plan}" + ); +} + +#[tokio::test] +async fn native_histogram_all_function_arms_route_correctly() { + // Every native-histogram match arm in `create_function_expr` must route to the + // expected UDF when all field columns are native histograms. `holt_winters` shares + // the `double_exponential_smoothing` arm but is not registered in the promql + // parser (0.10), so it cannot be exercised through a query string. + let cases = [ + // Range functions routed to native histogram UDFs. + ( + "increase(some_metric[5m])", + "prom_native_histogram_increase", + ), + ("rate(some_metric[5m])", "prom_native_histogram_rate"), + ("delta(some_metric[5m])", "prom_native_histogram_delta"), + ("idelta(some_metric[5m])", "prom_native_histogram_idelta"), + ("irate(some_metric[5m])", "prom_native_histogram_irate"), + ("resets(some_metric[5m])", "prom_native_histogram_resets"), + ("changes(some_metric[5m])", "prom_native_histogram_changes"), + ( + "avg_over_time(some_metric[5m])", + "prom_native_histogram_avg_over_time", + ), + ( + "sum_over_time(some_metric[5m])", + "prom_native_histogram_sum_over_time", + ), + ( + "count_over_time(some_metric[5m])", + "prom_native_histogram_count_over_time", + ), + ( + "last_over_time(some_metric[5m])", + "prom_native_histogram_last_over_time", + ), + ( + "present_over_time(some_metric[5m])", + "prom_native_histogram_present_over_time", + ), + // Unsupported functions dropped with the float-null UDF. + ("deriv(some_metric[5m])", "prom_native_histogram_drop_float"), + ( + "min_over_time(some_metric[5m])", + "prom_native_histogram_drop_float", + ), + ( + "max_over_time(some_metric[5m])", + "prom_native_histogram_drop_float", + ), + ( + "stddev_over_time(some_metric[5m])", + "prom_native_histogram_drop_float", + ), + ( + "stdvar_over_time(some_metric[5m])", + "prom_native_histogram_drop_float", + ), + ( + "quantile_over_time(0.9, some_metric[5m])", + "prom_native_histogram_drop_float", + ), + ( + "predict_linear(some_metric[5m], 60)", + "prom_native_histogram_drop_float", + ), + ( + "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)", + "prom_native_histogram_drop_float", + ), + ("round(some_metric)", "prom_native_histogram_drop_float"), + ("rad(some_metric)", "prom_native_histogram_drop_float"), + ("deg(some_metric)", "prom_native_histogram_drop_float"), + ("sgn(some_metric)", "prom_native_histogram_drop_float"), + // Instant helper functions routed to native histogram UDFs. + ( + "histogram_count(some_metric)", + "prom_native_histogram_count", + ), + ("histogram_sum(some_metric)", "prom_native_histogram_sum"), + ("histogram_avg(some_metric)", "prom_native_histogram_avg"), + ( + "histogram_stddev(some_metric)", + "prom_native_histogram_stddev", + ), + ( + "histogram_stdvar(some_metric)", + "prom_native_histogram_stdvar", + ), + ( + "histogram_fraction(-2 + 1, 2 / 2, some_metric)", + "prom_native_histogram_fraction", + ), + ]; + + for (query, expected_udf) in cases { + let plan = native_histogram_plan(query).await; + assert!(plan.contains(expected_udf), "{query}\n{plan}"); + if query.starts_with("histogram_fraction") { + assert!(plan.contains("Float64(-1)"), "{query}\n{plan}"); + } + } +} + +#[tokio::test] +async fn mixed_native_histogram_ranges_use_coordinated_udfs() { + let dual_output = [ + "increase(some_metric[5m])", + "rate(some_metric[5m])", + "delta(some_metric[5m])", + "idelta(some_metric[5m])", + "irate(some_metric[5m])", + "avg_over_time(some_metric[5m])", + "sum_over_time(some_metric[5m])", + "last_over_time(some_metric[5m])", + ]; + let float_output = [ + "resets(some_metric[5m])", + "changes(some_metric[5m])", + "deriv(some_metric[5m])", + "min_over_time(some_metric[5m])", + "max_over_time(some_metric[5m])", + "count_over_time(some_metric[5m])", + "absent_over_time(some_metric[5m])", + "present_over_time(some_metric[5m])", + "stddev_over_time(some_metric[5m])", + "stdvar_over_time(some_metric[5m])", + "quantile_over_time(0.9, some_metric[5m])", + "predict_linear(some_metric[5m], 60)", + "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)", + ]; + + for query in dual_output.iter().chain(float_output.iter()) { + let plan = PromPlanner::stmt_to_plan( + build_test_mixed_native_histogram_table_provider("some_metric").await, + &build_eval_stmt(query), + &build_query_engine_state(), + ) + .await + .unwrap() + .display_indent_schema() + .to_string(); + assert!(plan.contains("prom_mixed_range_float"), "{query}\n{plan}"); + assert_eq!( + plan.contains("prom_mixed_range_histogram"), + dual_output.contains(query), + "{query}\n{plan}" + ); + } + + let plan = PromPlanner::stmt_to_plan( + build_test_mixed_native_histogram_table_provider("some_metric").await, + &build_eval_stmt("sum_over_time(rate(some_metric[5m])[10m:1m])"), + &build_query_engine_state(), + ) + .await + .unwrap() + .display_indent_schema() + .to_string(); + let expected = r#"Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_value, prom_mixed_range_histogram(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=["greptime_value", "greptime_native_histogram"] [timestamp:Timestamp(ms), greptime_value:Dictionary(Int64, Float64);N, greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, tag_0:Utf8, timestamp_range:Dictionary(Int64, Timestamp(ms))] + PromSeriesDivide: tags=["tag_0"] [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_value, prom_mixed_range_histogram(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + PromRangeManipulate: req range=[-540000..100000000], interval=[60000], eval range=[300000], time index=[timestamp], values=["greptime_native_histogram", "greptime_value"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, greptime_value:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))] + PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + Filter: some_metric.timestamp >= TimestampMillisecond(-839999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]"#; + assert_eq!(plan, expected); +} + +#[tokio::test] +async fn mixed_native_histogram_rate_executes_real_ranges() { + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "timestamp", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + ), + Field::new(greptime_value(), ArrowDataType::Float64, true), + Field::new( + greptime_native_histogram(), + native_histogram_value_type().as_arrow_type(), + true, + ), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1000, 2000, 3000])), + Arc::new(Float64Array::from(vec![Some(1.0), None, Some(3.0)])), + build_histogram_array(&[None, Some(direct_or_histogram()), None]), + ], + ) + .unwrap(); + let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); + let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None) + .unwrap() + .build() + .unwrap(); + let collector = PromqlAnnotationCollector::default(); + let mut planner = PromPlanner { + table_provider: build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await, + ctx: PromPlannerContext { + start: 3000, + end: 3000, + interval: 1000, + range: Some(3000), + time_index_column: Some("timestamp".to_string()), + field_columns: vec![ + greptime_native_histogram().to_string(), + greptime_value().to_string(), + ], + ..Default::default() + }, + promql_annotations: Some(collector.clone()), + }; + let input = LogicalPlan::Extension(Extension { + node: Arc::new( + RangeManipulate::new( + 3000, + 3000, + 1000, + 0, + 3000, + "timestamp".to_string(), + planner.ctx.field_columns.clone(), + input, + ) + .unwrap(), + ), + }); + let PromExpr::Call(call) = parser::parse("rate(mixed[3s])").unwrap() else { + unreachable!() + }; + let preserve_any_value = PromPlanner::field_columns_are_alternative_samples( + input.schema(), + &planner.ctx.field_columns, + ); + let state = build_query_engine_state(); + let (mut exprs, _) = planner + .create_function_expr(&call.func, vec![], input.schema(), &state) + .unwrap(); + exprs.insert(0, planner.create_time_index_column_expr().unwrap()); + let plan = LogicalPlanBuilder::from(input) + .project(exprs) + .unwrap() + .filter( + planner + .create_empty_values_filter_expr(preserve_any_value) + .unwrap(), + ) + .unwrap() + .build() + .unwrap(); + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + let mut warnings = Vec::new(); + collector.append_to(&mut warnings, &mut Vec::new()); + assert!( + warnings + .iter() + .any(|warning| warning.contains("mix of float and native histogram")) + ); +} + +#[tokio::test] +async fn native_histogram_mixed_field_table_behaves() { + // Exercise function planning after float and histogram samples have already been + // represented as alternative nullable fields. Histogram functions must select the + // histogram field without adding a NULL float field that would reject every row. + let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt("histogram_count(some_metric)"), + &build_query_engine_state(), + ) + .await + .unwrap(); + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("prom_native_histogram_count"), + "{plan_str}" + ); + assert!(!plan_str.contains("Float64(NULL)"), "{plan_str}"); + assert!( + plan_str.contains("prom_native_histogram_count(greptime_native_histogram) IS NOT NULL"), + "{plan_str}" + ); + + // Value sorting keeps the float column and never sorts by the histogram column. + let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt("sort(some_metric)"), + &build_query_engine_state(), + ) + .await + .unwrap(); + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("greptime_value ASC NULLS FIRST"), + "{plan_str}" + ); + assert!( + !plan_str.contains("greptime_native_histogram ASC"), + "{plan_str}" + ); + + // scalar() ignores histogram samples and evaluates only the float field. + let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt("scalar(some_metric)"), + &build_query_engine_state(), + ) + .await + .unwrap(); + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("ScalarCalculate"), "{plan_str}"); + assert!( + plan_str.contains("greptime_value IS NOT NULL"), + "{plan_str}" + ); + + // Functions that preserve both alternative fields keep rows with either sample type. + let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt(r#"label_replace(some_metric, "copied", "$1", "tag_0", "(.*)")"#), + &build_query_engine_state(), + ) + .await + .unwrap(); + let plan_str = plan.display_indent_schema().to_string(); + let filter = plan_str.lines().next().unwrap(); + assert!( + filter.starts_with("Filter: ") + && filter.contains("greptime_native_histogram IS NOT NULL") + && filter.contains(" OR ") + && filter.contains("greptime_value IS NOT NULL"), + "{plan_str}" + ); +} + +#[tokio::test] +async fn less_filter_on_value() { + let query = "some_metric < 1.2345"; + let expected = String::from( + "Filter: some_metric.field_0 < Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +#[tokio::test] +async fn count_over_time() { + let query = "count_over_time(some_metric[5m])"; + let expected = String::from( + "Filter: prom_count_over_time(timestamp_range,field_0) IS NOT NULL [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ + \n Projection: some_metric.timestamp, prom_count_over_time(timestamp_range, field_0) AS prom_count_over_time(timestamp_range,field_0), some_metric.tag_0 [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ + \n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[300000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\ + \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.timestamp >= TimestampMillisecond(-299999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +/// The outer `PromRangeManipulate` from a subquery must be preceded by +/// `Sort` + `PromSeriesDivide`. +#[tokio::test] +async fn count_over_time_subquery() { + let query = "count_over_time(some_metric[10m:1m])"; + let expected = String::from( + "Filter: prom_count_over_time(timestamp_range,field_0) IS NOT NULL [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ + \n Projection: some_metric.timestamp, prom_count_over_time(timestamp_range, field_0) AS prom_count_over_time(timestamp_range,field_0), some_metric.tag_0 [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ + \n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromInstantManipulate: range=[-540000..100000000], lookback=[1000], interval=[60000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.timestamp >= TimestampMillisecond(-540999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + indie_query_plan_compare(query, expected).await; +} + +#[tokio::test] +async fn test_hash_join() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let case = r#"http_server_requests_seconds_sum{uri="/accounts/login"} / ignoring(kubernetes_pod_name,kubernetes_namespace) http_server_requests_seconds_count{uri="/accounts/login"}"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "http_server_requests_seconds_sum".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "http_server_requests_seconds_count".to_string(), + ), + ], + &["uri", "kubernetes_namespace", "kubernetes_pod_name"], + ) + .await; + // Should be ok + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let expected = "Projection: http_server_requests_seconds_sum.uri, http_server_requests_seconds_count.greptime_timestamp, CAST(http_server_requests_seconds_sum.greptime_value AS Float64) / CAST(http_server_requests_seconds_count.greptime_value AS Float64) AS http_server_requests_seconds_sum.greptime_value / http_server_requests_seconds_count.greptime_value\ + \n Projection: http_server_requests_seconds_sum.uri, http_server_requests_seconds_sum.kubernetes_namespace, http_server_requests_seconds_sum.kubernetes_pod_name, http_server_requests_seconds_sum.greptime_timestamp, http_server_requests_seconds_sum.greptime_value, http_server_requests_seconds_count.uri, http_server_requests_seconds_count.kubernetes_namespace, http_server_requests_seconds_count.kubernetes_pod_name, http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_count.greptime_value\ + \n Filter: prom_assert_unique_match_group(__promql_match_group_count, http_server_requests_seconds_sum.uri)\ + \n WindowAggr: windowExpr=[[count(Int64(1)) PARTITION BY [http_server_requests_seconds_sum.uri, http_server_requests_seconds_sum.greptime_timestamp] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS __promql_match_group_count]]\ + \n Inner Join: http_server_requests_seconds_sum.greptime_timestamp = http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_sum.uri = http_server_requests_seconds_count.uri\ + \n SubqueryAlias: http_server_requests_seconds_sum\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\ + \n PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\ + \n Sort: http_server_requests_seconds_sum.uri ASC NULLS FIRST, http_server_requests_seconds_sum.kubernetes_namespace ASC NULLS FIRST, http_server_requests_seconds_sum.kubernetes_pod_name ASC NULLS FIRST, http_server_requests_seconds_sum.greptime_timestamp ASC NULLS FIRST\ + \n Filter: http_server_requests_seconds_sum.uri = Utf8(\"/accounts/login\") AND http_server_requests_seconds_sum.greptime_timestamp >= TimestampMillisecond(-999, None) AND http_server_requests_seconds_sum.greptime_timestamp <= TimestampMillisecond(100000000, None)\ + \n TableScan: http_server_requests_seconds_sum\ + \n SubqueryAlias: http_server_requests_seconds_count\ + \n Projection: http_server_requests_seconds_count.uri, http_server_requests_seconds_count.kubernetes_namespace, http_server_requests_seconds_count.kubernetes_pod_name, http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_count.greptime_value\ + \n Filter: prom_assert_unique_match_group(__promql_match_group_count, http_server_requests_seconds_count.uri)\ + \n WindowAggr: windowExpr=[[count(Int64(1)) PARTITION BY [http_server_requests_seconds_count.uri, http_server_requests_seconds_count.greptime_timestamp] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS __promql_match_group_count]]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\ + \n PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\ + \n Sort: http_server_requests_seconds_count.uri ASC NULLS FIRST, http_server_requests_seconds_count.kubernetes_namespace ASC NULLS FIRST, http_server_requests_seconds_count.kubernetes_pod_name ASC NULLS FIRST, http_server_requests_seconds_count.greptime_timestamp ASC NULLS FIRST\ + \n Filter: http_server_requests_seconds_count.uri = Utf8(\"/accounts/login\") AND http_server_requests_seconds_count.greptime_timestamp >= TimestampMillisecond(-999, None) AND http_server_requests_seconds_count.greptime_timestamp <= TimestampMillisecond(100000000, None)\ + \n TableScan: http_server_requests_seconds_count"; + assert_eq!(plan.to_string(), expected); +} + +#[tokio::test] +async fn test_nested_histogram_quantile() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let case = r#"label_replace(histogram_quantile(0.99, sum by(pod, le, path, code) (rate(greptime_servers_grpc_requests_elapsed_bucket{container="frontend"}[1m0s]))), "pod_new", "$1", "pod", "greptimedb-frontend-[0-9a-z]*-(.*)")"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider_with_fields( + &[( + DEFAULT_SCHEMA_NAME.to_string(), + "greptime_servers_grpc_requests_elapsed_bucket".to_string(), + )], + &["pod", "le", "path", "code", "container"], + ) + .await; + // Should be ok + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_histogram_quantile_binary_op() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + // Arithmetic applied to a histogram_quantile() result. Regression for #8144: + // HistogramFold used to drop the input column qualifiers, so the binary-op + // projection failed to resolve the qualified tag column. + let case = r#"histogram_quantile(0.5, sum by (le, pod) (rate(http_request_duration_seconds_bucket[5m]))) + 0"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider_with_fields( + &[( + DEFAULT_SCHEMA_NAME.to_string(), + "http_request_duration_seconds_bucket".to_string(), + )], + &["pod", "le"], + ) + .await; + // Should plan without a "No field named ..." error. + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_parse_and_operator() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let cases = [ + r#"count (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} ) and (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} )) / (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_capacity_bytes{namespace=~".+"} )) >= (80 / 100)) or vector (0)"#, + r#"count (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} ) unless (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} )) / (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_capacity_bytes{namespace=~".+"} )) >= (80 / 100)) or vector (0)"#, + ]; + + for case in cases { + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "kubelet_volume_stats_used_bytes".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "kubelet_volume_stats_capacity_bytes".to_string(), + ), + ], + &["namespace", "persistentvolumeclaim"], + ) + .await; + // Should be ok + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + } +} + +#[tokio::test] +async fn test_nested_binary_op() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let case = r#"sum(rate(nginx_ingress_controller_requests{job=~".*"}[2m])) - + ( + sum(rate(nginx_ingress_controller_requests{namespace=~".*"}[2m])) + or + vector(0) + )"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider_with_fields( + &[( + DEFAULT_SCHEMA_NAME.to_string(), + "nginx_ingress_controller_requests".to_string(), + )], + &["namespace", "job"], + ) + .await; + // Should be ok + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_parse_or_operator() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let case = r#" + sum(rate(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}[120s])) by (cluster_name,tenant_name) / + (sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) * 100) + or + 200 * sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) / + sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)"#; + + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())], + &["tenant_name", "cluster_name"], + ) + .await; + eval_stmt.expr = parser::parse(case).unwrap(); + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let case = r#"sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) / + (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) + + sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) / + (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0 + or + sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) / + (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0 + or + sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) / + (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0"#; + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())], + &["tenant_name", "cluster_name"], + ) + .await; + eval_stmt.expr = parser::parse(case).unwrap(); + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let case = r#"(sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) + + sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or + (sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or + (sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name))"#; + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "background_waitevent_cnt".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "foreground_waitevent_cnt".to_string(), + ), + ], + &["tenant_name", "cluster_name"], + ) + .await; + eval_stmt.expr = parser::parse(case).unwrap(); + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let case = r#"avg(node_load1{cluster_name=~"cluster1"}) by (cluster_name,host_name) or max(container_cpu_load_average_10s{cluster_name=~"cluster1"}) by (cluster_name,host_name) * 100 / max(container_spec_cpu_quota{cluster_name=~"cluster1"}) by (cluster_name,host_name)"#; + let table_provider = build_test_table_provider_with_fields( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "node_load1".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "container_cpu_load_average_10s".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "container_spec_cpu_quota".to_string(), + ), + ], + &["cluster_name", "host_name"], + ) + .await; + eval_stmt.expr = parser::parse(case).unwrap(); + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); +} + +#[tokio::test] +async fn value_matcher() { + // template + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let cases = [ + // single equal matcher + ( + r#"some_metric{__field__="field_1"}"#, + vec![ + "some_metric.field_1", + "some_metric.tag_0", + "some_metric.tag_1", + "some_metric.tag_2", + "some_metric.timestamp", + ], + ), + // two equal matchers + ( + r#"some_metric{__field__="field_1", __field__="field_0"}"#, + vec![ + "some_metric.field_0", + "some_metric.field_1", + "some_metric.tag_0", + "some_metric.tag_1", + "some_metric.tag_2", + "some_metric.timestamp", + ], + ), + // single not_eq matcher + ( + r#"some_metric{__field__!="field_1"}"#, + vec![ + "some_metric.field_0", + "some_metric.field_2", + "some_metric.tag_0", + "some_metric.tag_1", + "some_metric.tag_2", + "some_metric.timestamp", + ], + ), + // two not_eq matchers + ( + r#"some_metric{__field__!="field_1", __field__!="field_2"}"#, + vec![ + "some_metric.field_0", + "some_metric.tag_0", + "some_metric.tag_1", + "some_metric.tag_2", + "some_metric.timestamp", + ], + ), + // equal and not_eq matchers (no conflict) + ( + r#"some_metric{__field__="field_1", __field__!="field_0"}"#, + vec![ + "some_metric.field_1", + "some_metric.tag_0", + "some_metric.tag_1", + "some_metric.tag_2", + "some_metric.timestamp", + ], + ), + // equal and not_eq matchers (conflict) + ( + r#"some_metric{__field__="field_2", __field__!="field_2"}"#, + vec![ + "some_metric.tag_0", + "some_metric.tag_1", + "some_metric.tag_2", + "some_metric.timestamp", + ], + ), + // single regex eq matcher + ( + r#"some_metric{__field__=~"field_1|field_2"}"#, + vec![ + "some_metric.field_1", + "some_metric.field_2", + "some_metric.tag_0", + "some_metric.tag_1", + "some_metric.tag_2", + "some_metric.timestamp", + ], + ), + // single regex not_eq matcher + ( + r#"some_metric{__field__!~"field_1|field_2"}"#, + vec![ + "some_metric.field_0", + "some_metric.tag_0", + "some_metric.tag_1", + "some_metric.tag_2", + "some_metric.timestamp", + ], + ), + ]; + + for case in cases { + let prom_expr = parser::parse(case.0).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 3, + 3, + ) + .await; + let plan = + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let mut fields = plan.schema().field_names(); + let mut expected = case.1.into_iter().map(String::from).collect::>(); + fields.sort(); + expected.sort(); + assert_eq!(fields, expected, "case: {:?}", case.0); + } + + let bad_cases = [ + r#"some_metric{__field__="nonexistent"}"#, + r#"some_metric{__field__!="nonexistent"}"#, + ]; + + for case in bad_cases { + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 3, + 3, + ) + .await; + let plan = + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await; + assert!(plan.is_err(), "case: {:?}", case); + } +} + +#[tokio::test] +async fn custom_schema() { + let query = "some_alt_metric{__schema__=\"greptime_private\"}"; + let expected = String::from( + "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: greptime_private.some_alt_metric.tag_0 ASC NULLS FIRST, greptime_private.some_alt_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: greptime_private.some_alt_metric.timestamp >= TimestampMillisecond(-999, None) AND greptime_private.some_alt_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; + + let query = "some_alt_metric{__database__=\"greptime_private\"}"; + let expected = String::from( + "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: greptime_private.some_alt_metric.tag_0 ASC NULLS FIRST, greptime_private.some_alt_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: greptime_private.some_alt_metric.timestamp >= TimestampMillisecond(-999, None) AND greptime_private.some_alt_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; + + let query = "some_alt_metric{__schema__=\"greptime_private\"} / some_metric"; + let expected = String::from( + "Projection: some_metric.tag_0, some_metric.timestamp, CAST(greptime_private.some_alt_metric.field_0 AS Float64) / CAST(some_metric.field_0 AS Float64) AS greptime_private.some_alt_metric.field_0 / some_metric.field_0 [tag_0:Utf8, timestamp:Timestamp(ms), greptime_private.some_alt_metric.field_0 / some_metric.field_0:Float64;N]\ + \n Inner Join: greptime_private.some_alt_metric.tag_0 = some_metric.tag_0, greptime_private.some_alt_metric.timestamp = some_metric.timestamp [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n SubqueryAlias: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: greptime_private.some_alt_metric.tag_0 ASC NULLS FIRST, greptime_private.some_alt_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: greptime_private.some_alt_metric.timestamp >= TimestampMillisecond(-999, None) AND greptime_private.some_alt_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n SubqueryAlias: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +#[tokio::test] +async fn only_equals_is_supported_for_special_matcher() { + let queries = &[ + "some_alt_metric{__schema__!=\"greptime_private\"}", + "some_alt_metric{__schema__=~\"lalala\"}", + "some_alt_metric{__database__!=\"greptime_private\"}", + "some_alt_metric{__database__=~\"lalala\"}", + ]; + + for query in queries { + let prom_expr = parser::parse(query).unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + "greptime_private".to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + + let plan = + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await; + assert!(plan.is_err(), "query: {:?}", query); + } +} + +#[tokio::test] +async fn native_scan_bounds_preserve_zero_lookback_and_overflow() { + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let mut planner = PromPlanner { + table_provider, + ctx: PromPlannerContext::from_eval_stmt(&build_eval_stmt("some_metric")), + promql_annotations: None, + }; + planner.ctx.time_index_column = Some("timestamp".to_string()); + planner.ctx.start = 1_000; + planner.ctx.lookback_delta = 0; + let schema = Arc::new( + DFSchema::try_from(ArrowSchema::new(vec![Field::new( + "timestamp", + ArrowDataType::Timestamp(ArrowTimeUnit::Nanosecond, None), + false, + )])) + .unwrap(), + ); + for (end, interval, windows) in [ + (1_000, 1_000, 1), + (2_000, 1_000, 1), + (7_201_000, 7_200_000, 2), + ] { + planner.ctx.end = end; + planner.ctx.interval = interval; + let filter = planner + .build_time_index_filter(0, &schema) + .unwrap() + .unwrap() + .to_string(); + assert_eq!(filter.matches(">=").count(), windows, "{filter}"); + assert!( + filter.contains("TimestampNanosecond(1000000000, None)"), + "{filter}" + ); + } + planner.ctx.end = i64::MAX; + let filter = planner + .build_time_index_filter(0, &schema) + .unwrap() + .unwrap() + .to_string(); + assert!( + filter.contains("timestamp >= TimestampNanosecond(1000000000, None)"), + "{filter}" + ); + + // A lookback subtraction can underflow milliseconds while the upper bound remains + // representable. Keep that upper bound so LastRow cannot select a future sample. + let ms_schema = Arc::new( + DFSchema::try_from(ArrowSchema::new(vec![Field::new( + "timestamp", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + )])) + .unwrap(), + ); + planner.ctx.start = i64::MIN + 100; + planner.ctx.end = planner.ctx.start; + planner.ctx.lookback_delta = 200; + let filter = planner + .build_time_index_filter(0, &ms_schema) + .unwrap() + .unwrap() + .to_string(); + assert_eq!( + filter, + format!( + "timestamp <= TimestampMillisecond({}, None)", + i64::MIN + 100 + ) + ); + + // The lower bound can also overflow while converting milliseconds to native nanoseconds. + // Its representable upper bound still has to reach the scan. + planner.ctx.start = 0; + planner.ctx.end = 0; + planner.ctx.lookback_delta = 300_000; + let filter = planner + .build_time_index_filter(9_223_372_036_854, &schema) + .unwrap() + .unwrap() + .to_string(); + assert_eq!( + filter, + "timestamp <= TimestampNanosecond(-9223372036854000000, None)" + ); +} + +#[tokio::test] +async fn test_non_ms_precision() { + let catalog_list = MemoryCatalogManager::with_default_setup(); + let columns = vec![ + ColumnSchema::new( + "tag".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_nanosecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + "field".to_string(), + ConcreteDataType::float64_datatype(), + true, + ), + ]; + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![0]) + .value_indices(vec![2]) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name("metrics".to_string()) + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "metrics".to_string(), + table_id: 1024, + table, + }) + .is_ok() + ); + + let plan = PromPlanner::stmt_to_plan( + DfTableSourceProvider::new( + catalog_list.clone(), + false, + QueryContext::arc(), + DummyDecoder::arc(), + true, + ), + &EvalStmt { + expr: parser::parse("metrics{tag = \"1\"}").unwrap(), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }, + &build_query_engine_state(), + ) + .await + .unwrap(); + assert_eq!( + plan.display_indent_schema().to_string(), + "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\n PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp > TimestampNanosecond(-1000000000, None) AND metrics.timestamp <= TimestampNanosecond(100000000000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Projection: metrics.field, metrics.tag, metrics.timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]" + ); + let plan = PromPlanner::stmt_to_plan( + DfTableSourceProvider::new( + catalog_list.clone(), + false, + QueryContext::arc(), + DummyDecoder::arc(), + true, + ), + &EvalStmt { + expr: parser::parse("avg_over_time(metrics{tag = \"1\"}[5s])").unwrap(), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }, + &build_query_engine_state(), + ) + .await + .unwrap(); + assert_eq!( + plan.display_indent_schema().to_string(), + "Filter: prom_avg_over_time(timestamp_range,field) IS NOT NULL [timestamp:Timestamp(ms), prom_avg_over_time(timestamp_range,field):Float64;N, tag:Utf8]\n Projection: metrics.timestamp, prom_avg_over_time(timestamp_range, field) AS prom_avg_over_time(timestamp_range,field), metrics.tag [timestamp:Timestamp(ms), prom_avg_over_time(timestamp_range,field):Float64;N, tag:Utf8]\n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[5000], time index=[timestamp], values=[\"field\"] [field:Dictionary(Int64, Float64);N, tag:Utf8, timestamp:Timestamp(ms), timestamp_range:Dictionary(Int64, Timestamp(ms))]\n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp > TimestampNanosecond(-5000000000, None) AND metrics.timestamp <= TimestampNanosecond(100000000000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Projection: metrics.field, metrics.tag, metrics.timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]" + ); +} + +#[tokio::test] +async fn test_nonexistent_label() { + // template + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let case = r#"some_metric{nonexistent="hi"}"#; + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 3, + 3, + ) + .await; + // Should be ok + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_label_join() { + let prom_expr = + parser::parse("label_join(up{tag_0='api-server'}, 'foo', ',', 'tag_1', 'tag_2', 'tag_3')") + .unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = + build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 4, 1) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let expected = r#" +Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8] + Projection: up.timestamp, up.field_0, concat_ws(Utf8(","), up.tag_1, up.tag_2, up.tag_3) AS foo, up.tag_0, up.tag_1, up.tag_2, up.tag_3 [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8] + PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] + PromSeriesDivide: tags=["tag_0", "tag_1", "tag_2", "tag_3"] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] + Sort: up.tag_0 ASC NULLS FIRST, up.tag_1 ASC NULLS FIRST, up.tag_2 ASC NULLS FIRST, up.tag_3 ASC NULLS FIRST, up.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] + Filter: up.tag_0 = Utf8("api-server") AND up.timestamp >= TimestampMillisecond(-999, None) AND up.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] + TableScan: up [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#; + + let ret = plan.display_indent_schema().to_string(); + assert_eq!(format!("\n{ret}"), expected, "\n{}", ret); +} + +#[tokio::test] +async fn test_label_replace() { + let prom_expr = + parser::parse("label_replace(up{tag_0=\"a:c\"}, \"foo\", \"$1\", \"tag_0\", \"(.*):.*\")") + .unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = + build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 1, 1) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let expected = r#" +Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8] + Projection: up.timestamp, up.field_0, regexp_replace(up.tag_0, Utf8("^(?s:(.*):.*)$"), Utf8("$1")) AS foo, up.tag_0 [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8] + PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] + PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] + Sort: up.tag_0 ASC NULLS FIRST, up.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] + Filter: up.tag_0 = Utf8("a:c") AND up.timestamp >= TimestampMillisecond(-999, None) AND up.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] + TableScan: up [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#; + + let ret = plan.display_indent_schema().to_string(); + assert_eq!(format!("\n{ret}"), expected, "\n{}", ret); +} + +#[tokio::test] +async fn label_replace_aggregation_queries_plan_successfully() { + let aggregate = r#"sum by (foo) (label_replace(some_metric, "foo", "$1", "tag_0", "(.*)"))"#; + let queries = [ + aggregate.to_string(), + format!("{aggregate} <= 10"), + format!("{aggregate} * 0.8"), + format!("0.8 * {aggregate}"), + format!("{aggregate} <= {aggregate} * 0.8"), + ]; + let state = build_query_engine_state(); + let mut failures = Vec::new(); + + for query in queries { + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + if let Err(error) = + PromPlanner::stmt_to_plan(table_provider, &build_eval_stmt(&query), &state).await + { + failures.push(format!("{query}: {error:?}")); + } + } + + assert!(failures.is_empty(), "{}", failures.join("\n")); +} + +#[tokio::test] +async fn test_matchers_to_expr() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + let case = + r#"sum(prometheus_tsdb_head_series{tag_1=~"(10.0.160.237:8080|10.0.160.237:9090)"})"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider( + &[( + DEFAULT_SCHEMA_NAME.to_string(), + "prometheus_tsdb_head_series".to_string(), + )], + 3, + 3, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let expected = "Sort: prometheus_tsdb_head_series.timestamp ASC NULLS LAST [timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.field_0):Float64;N, sum(prometheus_tsdb_head_series.field_1):Float64;N, sum(prometheus_tsdb_head_series.field_2):Float64;N]\ + \n Aggregate: groupBy=[[prometheus_tsdb_head_series.timestamp]], aggr=[[sum(prometheus_tsdb_head_series.field_0), sum(prometheus_tsdb_head_series.field_1), sum(prometheus_tsdb_head_series.field_2)]] [timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.field_0):Float64;N, sum(prometheus_tsdb_head_series.field_1):Float64;N, sum(prometheus_tsdb_head_series.field_2):Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\", \"tag_2\"] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\ + \n Sort: prometheus_tsdb_head_series.tag_0 ASC NULLS FIRST, prometheus_tsdb_head_series.tag_1 ASC NULLS FIRST, prometheus_tsdb_head_series.tag_2 ASC NULLS FIRST, prometheus_tsdb_head_series.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\ + \n Filter: prometheus_tsdb_head_series.tag_1 ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\ + \n TableScan: prometheus_tsdb_head_series [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]"; + assert_eq!(plan.display_indent_schema().to_string(), expected); +} + +#[tokio::test] +async fn test_topk_expr() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + let case = r#"topk(10, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "prometheus_tsdb_head_series".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "http_server_requests_seconds_count".to_string(), + ), + ], + &["ip"], + ) + .await; + + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let expected = "Projection: sum(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp [sum(prometheus_tsdb_head_series.greptime_value):Float64;N, ip:Utf8, greptime_timestamp:Timestamp(ms)]\ + \n Sort: prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\ + \n Filter: row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= Float64(10) [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\ + \n WindowAggr: windowExpr=[[row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\ + \n Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ + \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[sum(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"; + + assert_eq!(plan.display_indent_schema().to_string(), expected); +} + +#[tokio::test] +async fn test_count_values_expr() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "prometheus_tsdb_head_series".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "http_server_requests_seconds_count".to_string(), + ), + ], + &["ip"], + ) + .await; + + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let expected = "Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N]\ + \n Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value) AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N]\ + \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value)]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), prom_float_to_string(prometheus_tsdb_head_series.greptime_value):Utf8;N, count(prometheus_tsdb_head_series.greptime_value):Int64]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"; + + assert_eq!(plan.display_indent_schema().to_string(), expected); +} + +#[tokio::test] +async fn test_value_alias() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + eval_stmt = QueryLanguageParser::apply_alias_extension(eval_stmt, "my_series"); + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "prometheus_tsdb_head_series".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "http_server_requests_seconds_count".to_string(), + ), + ], + &["ip"], + ) + .await; + + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let expected = r#" +Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.series, prometheus_tsdb_head_series.greptime_timestamp [my_series:Int64, ip:Utf8, series:Utf8;N, greptime_timestamp:Timestamp(ms)] + Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N] + Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value) AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N] + Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value)]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), prom_float_to_string(prometheus_tsdb_head_series.greptime_value):Utf8;N, count(prometheus_tsdb_head_series.greptime_value):Int64] + PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] + PromSeriesDivide: tags=["ip"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] + Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] + Filter: prometheus_tsdb_head_series.ip ~ Utf8("^(?:(10.0.160.237:8080|10.0.160.237:9090))$") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] + TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"#; + assert_eq!(format!("\n{}", plan.display_indent_schema()), expected); +} + +#[tokio::test] +async fn test_quantile_expr() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + let case = r#"quantile(0.3, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "prometheus_tsdb_head_series".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "http_server_requests_seconds_count".to_string(), + ), + ], + &["ip"], + ) + .await; + + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let expected = "Sort: prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [greptime_timestamp:Timestamp(ms), quantile(Float64(0.3),sum(prometheus_tsdb_head_series.greptime_value)):Float64;N]\ + \n Aggregate: groupBy=[[prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[quantile(Float64(0.3), sum(prometheus_tsdb_head_series.greptime_value))]] [greptime_timestamp:Timestamp(ms), quantile(Float64(0.3),sum(prometheus_tsdb_head_series.greptime_value)):Float64;N]\ + \n Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ + \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[sum(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"; + + assert_eq!(plan.display_indent_schema().to_string(), expected); +} + +#[tokio::test] +async fn test_or_not_exists_table_label() { + let state = build_query_engine_state(); + let provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string())], + &["job"], + ) + .await; + let raw = PromPlanner::stmt_to_plan( + provider, + &build_eval_stmt(r#"missing_metric or on(absent_label) normal_metric"#), + &state, + ) + .await + .unwrap(); + assert!( + raw.display_indent_schema() + .to_string() + .contains("__promql_or_match_0@") + ); + let (optimized, batches) = execute(raw, &state).await; + assert_no_internal_or_keys(optimized.schema()); + assert!(batches.iter().all(|batch| { + batch + .schema() + .fields() + .iter() + .all(|field| !field.name().starts_with("__promql_or_match_")) + })); +} + +#[tokio::test] +async fn test_histogram_quantile_missing_le_column() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + // Test case: histogram_quantile with a table that doesn't have 'le' column + let case = r#"histogram_quantile(0.99, sum by(pod,instance,le) (rate(non_existent_histogram_bucket{instance=~"xxx"}[1m])))"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + + // Create a table provider with a table that doesn't have 'le' column + let table_provider = build_test_table_provider_with_fields( + &[( + DEFAULT_SCHEMA_NAME.to_string(), + "non_existent_histogram_bucket".to_string(), + )], + &["pod", "instance"], // Note: no 'le' column + ) + .await; + + // Should return empty result instead of error + let result = + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()).await; + + // This should succeed now (returning empty result) instead of failing with "Cannot find column le" + assert!( + result.is_ok(), + "Expected successful plan creation with empty result, but got error: {:?}", + result.err() + ); + + // Verify that the result is an EmptyRelation + let plan = result.unwrap(); + match plan { + LogicalPlan::EmptyRelation(_) => { + // This is what we expect + } + _ => panic!("Expected EmptyRelation, but got: {:?}", plan), + } +} + +#[tokio::test] +async fn test_direct_or_normalizes_missing_match_labels() { + type Case<'a> = ( + Option>, + Option>, + i64, + i64, + &'a [(f64, Option<&'a str>)], + ); + + let modifier = or_modifier("lhs or on(k) rhs"); + #[rustfmt::skip] + let cases: &[Case<'_>] = &[ + (None, None, 1, 1, &[(1.0, None)]), + (None, Some(Some("")), 1, 1, &[(1.0, None)]), + (Some(Some("")), None, 1, 1, &[(1.0, Some(""))]), + (None, Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]), + (Some(Some("l")), None, 1, 1, &[(1.0, Some("l")), (2.0, None)]), + (Some(None), Some(Some("")), 1, 1, &[(1.0, None)]), + (Some(None), Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]), + (Some(Some("same")), Some(Some("same")), 1, 2, &[(1.0, Some("same")), (2.0, Some("same"))]), + ]; + for &(left, right, left_ts, right_ts, expected) in cases { + let (optimized, batches) = run( + &matrix_source("lhs", left, left_ts, 1.0), + &matrix_source("rhs", right, right_ts, 2.0), + matrix_context("lhs", left), + matrix_context("rhs", right), + &modifier, + ) + .await; + assert_no_internal_or_keys(optimized.schema()); + assert_eq!( + rows(&batches), + expected + .iter() + .map(|(value, label)| (*value, label.map(str::to_string))) + .collect::>() + ); + } +} + +#[tokio::test] +async fn test_direct_or_match_modifiers() { + for (modifier, left, right, expected) in [ + (None, "left", "right", 2), + (or_modifier("lhs or on(k) rhs"), "same", "same", 1), + (or_modifier("lhs or on() rhs"), "left", "right", 1), + (or_modifier("lhs or ignoring(k) rhs"), "left", "right", 1), + ] { + let (_, batches) = run( + &matrix_source("lhs", Some(Some(left)), 1, 1.0), + &matrix_source("rhs", Some(Some(right)), 1, 2.0), + direct_or_context("lhs", &["job", "k"], "v"), + direct_or_context("rhs", &["job", "k"], "v"), + &modifier, + ) + .await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + expected + ); + } +} + +#[tokio::test] +async fn test_direct_or_nested_projection_uses_left_context() { + let left = matrix_source("lhs", Some(Some("k")), 1, 1.0); + let right = matrix_source("rhs", Some(Some("k")), 1, 2.0); + let raw = plan_direct_or( + scan(&left), + scan(&right), + direct_or_context("lhs", &["job", "k"], "v"), + direct_or_context("rhs", &["job", "k"], "v"), + &or_modifier("lhs or on(k) rhs"), + ) + .await; + assert!(raw.schema().iter().any(|(qualifier, field)| { + qualifier.as_ref().is_some_and(|q| q.to_string() == "lhs") && field.name() == "v" + })); + let nested = LogicalPlanBuilder::from(raw) + .project(vec![ + DfExpr::BinaryExpr(BinaryExpr { + left: Box::new(DfExpr::Column(Column::new( + Some(TableReference::bare("lhs")), + "v", + ))), + op: Operator::Plus, + right: Box::new(lit(1.0)), + }) + .alias("v_plus"), + ]) + .unwrap() + .build() + .unwrap(); + let (_, batches) = execute(nested, &build_query_engine_state()).await; + assert_eq!(values(&batches, "v_plus"), vec![2.0]); +} + +#[tokio::test] +async fn test_direct_or_skips_user_internal_key_name() { + const USER_TAG: &str = "__promql_or_match_0"; + let left = tagged_source( + "lhs", + false, + (USER_TAG, Some("left")), + DirectOrValue::Float64(1.0), + ); + let right = tagged_source( + "rhs", + false, + (USER_TAG, Some("right")), + DirectOrValue::Float64(2.0), + ); + let raw = plan_direct_or( + scan(&left), + scan(&right), + direct_or_context("lhs", &["job", USER_TAG], "v"), + direct_or_context("rhs", &["job", USER_TAG], "v"), + &or_modifier("lhs or on(missing_label) rhs"), + ) + .await; + assert!( + raw.display_indent_schema() + .to_string() + .contains("__promql_or_match_1@") + ); + let (_, batches) = execute(raw, &build_query_engine_state()).await; + assert!( + batches + .iter() + .all(|batch| batch.column_by_name(USER_TAG).is_some()) + ); +} + +#[tokio::test] +async fn test_direct_or_substrait_round_trip_with_normalized_key() { + let state = build_query_engine_state(); + let ctx = SessionContext::new_with_state(state.session_state()); + let catalog = Arc::new(MemoryCatalogProvider::new()); + catalog + .register_schema("public", Arc::new(MemorySchemaProvider::new())) + .unwrap(); + ctx.register_catalog("datafusion", catalog); + let left = matrix_source("lhs", Some(Some("")), 1, 1.0); + let right = matrix_source("rhs", None, 1, 2.0); + ctx.register_table( + TableReference::full("datafusion", "public", "lhs"), + table(&left), + ) + .unwrap(); + ctx.register_table( + TableReference::full("datafusion", "public", "rhs"), + table(&right), + ) + .unwrap(); + let raw = plan_direct_or( + ctx.table("datafusion.public.lhs") + .await + .unwrap() + .into_unoptimized_plan(), + ctx.table("datafusion.public.rhs") + .await + .unwrap() + .into_unoptimized_plan(), + direct_or_context("lhs", &["job", "k"], "v"), + direct_or_context("rhs", &["job"], "v"), + &or_modifier("lhs or on(k) rhs"), + ) + .await; + let decoded = DFLogicalSubstraitConvertor + .decode( + DFLogicalSubstraitConvertor + .encode(&raw, DefaultSerializer) + .unwrap(), + ctx.state(), + ) + .await + .unwrap(); + let (optimized, batches) = execute(decoded, &state).await; + assert_no_internal_or_keys(optimized.schema()); + assert!(batches.iter().all(|batch| { + batch + .schema() + .fields() + .iter() + .all(|field| !field.name().starts_with("__promql_or_match_")) + })); + assert_eq!(values(&batches, "v"), vec![1.0]); +} + +#[tokio::test] +async fn test_direct_or_numeric_value_types() { + let left = tagged_source("lhs", true, ("k", Some("lhs")), DirectOrValue::Int64(0)); + let right = tagged_source( + "rhs", + false, + ("k", Some("rhs")), + DirectOrValue::Float64(0.5), + ); + let (optimized, batches) = run( + &left, + &right, + direct_or_context("lhs", &["job", "k"], "v"), + direct_or_context("rhs", &["job", "k"], "v"), + &or_modifier("lhs or on(k) rhs"), + ) + .await; + assert_eq!( + optimized + .schema() + .field_with_name(None, "v") + .unwrap() + .data_type(), + &ArrowDataType::Float64 + ); + assert_eq!(values(&batches, "v"), vec![0.5]); + let provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let mut planner = PromPlanner { + table_provider: provider, + ctx: PromPlannerContext::default(), + promql_annotations: None, + }; + let left_context = direct_or_context("lhs", &["job"], "v"); + let right_context = direct_or_context("rhs", &["job"], "v"); + let error = planner + .or_operator( + scan(&job_source("lhs", DirectOrValue::Utf8("x"))), + scan(&job_source("rhs", DirectOrValue::Float64(1.0))), + left_context.tag_columns.iter().cloned().collect(), + right_context.tag_columns.iter().cloned().collect(), + left_context, + right_context, + &or_modifier("lhs or on() rhs"), + ) + .unwrap_err(); + assert!( + error + .to_string() + .contains("OR value fields have incompatible types") + ); +} + +#[tokio::test] +async fn test_or_with_histogram_quantile_missing_le_column() { + let case = r#"histogram_quantile(0.99, non_existent_histogram_bucket) or normal_metric"#; + let eval_stmt = build_eval_stmt(case); + let table_provider = build_missing_le_or_normal_metric_table_provider().await; + + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + assert_normal_metric_schema(&plan); +} + +#[tokio::test] +async fn test_or_with_right_empty_histogram_restores_left_context() { + let eval_stmt = build_eval_stmt( + r#"abs(sum by(instance) (normal_metric) or histogram_quantile(0.99, sum by(pod) (non_existent_histogram_bucket)))"#, + ); + let table_provider = build_missing_le_or_normal_metric_table_provider().await; + + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_or_with_both_empty_histograms() { + let eval_stmt = build_eval_stmt( + r#"histogram_quantile(0.99, sum by(pod) (left_histogram_bucket)) or histogram_quantile(0.99, sum by(instance) (right_histogram_bucket))"#, + ); + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "left_histogram_bucket".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "right_histogram_bucket".to_string(), + ), + ], + &["pod", "instance"], + ) + .await; + + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + match plan { + LogicalPlan::EmptyRelation(relation) => { + assert!(!relation.produce_one_row); + assert!(!relation.schema.fields().is_empty()); + assert!( + relation + .schema + .fields() + .iter() + .any(|field| field.data_type() == &ArrowDataType::Float64) + ); + assert!( + relation + .schema + .fields() + .iter() + .any(|field| field.name() == "pod") + ); + assert!( + !relation + .schema + .fields() + .iter() + .any(|field| field.name() == "instance") + ); + } + _ => panic!("Expected EmptyRelation, but got: {plan:?}"), + } +} + +#[tokio::test] +async fn test_nested_or_with_both_empty_histograms() { + for case in [ + r#"abs(histogram_quantile(0.99, left_histogram_bucket) or histogram_quantile(0.99, right_histogram_bucket))"#, + r#"(histogram_quantile(0.99, left_histogram_bucket) or histogram_quantile(0.99, right_histogram_bucket)) + 1"#, + ] { + let eval_stmt = build_eval_stmt(case); + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "left_histogram_bucket".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "right_histogram_bucket".to_string(), + ), + ], + &["pod", "instance"], + ) + .await; + + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + } +} + +#[tokio::test] +async fn test_or_with_empty_histogram_modifiers() { + for case in [ + r#"histogram_quantile(0.99, non_existent_histogram_bucket) or on(pod) normal_metric"#, + r#"normal_metric or ignoring(instance) histogram_quantile(0.99, non_existent_histogram_bucket)"#, + ] { + let eval_stmt = build_eval_stmt(case); + let table_provider = build_missing_le_or_normal_metric_table_provider().await; + + let plan = + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + assert_normal_metric_schema(&plan); + } +} + +#[tokio::test] +async fn test_unless_preserves_left_context_for_histogram() { + let eval_stmt = build_eval_stmt( + r#"histogram_quantile(0.99, bucket_metric unless on(job) normal_metric) or fallback_metric"#, + ); + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + build_set_op_context_table_provider().await, + &eval_stmt, + &state, + ) + .await + .unwrap(); + assert!(contains_histogram_fold(&plan), "{plan:?}"); + let (optimized, physical) = optimize_and_create_physical_plan(&state, plan).await; + assert!(contains_histogram_fold(&optimized), "{optimized:?}"); + let batches = datafusion::physical_plan::collect(physical, state.session_state().task_ctx()) + .await + .unwrap(); + assert!(batches.iter().all(|batch| batch.num_rows() == 0)); +} + +#[tokio::test] +async fn test_and_preserves_left_context_for_histogram() { + let eval_stmt = build_eval_stmt( + r#"histogram_quantile(0.99, bucket_metric and on(job) normal_metric) or fallback_metric"#, + ); + let plan = PromPlanner::stmt_to_plan( + build_set_op_context_table_provider().await, + &eval_stmt, + &build_query_engine_state(), + ) + .await + .unwrap(); + assert!(contains_histogram_fold(&plan), "{plan:?}"); +} + +#[tokio::test] +async fn test_and_preserves_left_context_when_le_is_missing() { + let eval_stmt = + build_eval_stmt(r#"histogram_quantile(0.99, normal_metric and on(job) bucket_metric)"#); + let plan = PromPlanner::stmt_to_plan( + build_set_op_context_table_provider().await, + &eval_stmt, + &build_query_engine_state(), + ) + .await + .unwrap(); + assert!(matches!(&plan, LogicalPlan::EmptyRelation(_)), "{plan:?}"); + assert!(!plan.schema().fields().is_empty()); + assert!(!contains_histogram_fold(&plan), "{plan:?}"); +} + +async fn build_matching_filter_plan(query: &str) -> String { + let table_provider = build_test_table_provider_with_distinct_tags(&[ + ("metric_a", &["host", "device"]), + ("metric_b", &["host", "device"]), + ]) + .await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt(query), + &build_query_engine_state(), + ) + .await + .unwrap(); + plan.display_indent().to_string() +} + +/// [`build_test_table_provider_with_distinct_tags`] plus a `status` string column: a value +/// field that is neither a primary key nor a value column of the metric, so +/// `count by(status) (...)` still reports it among the aggregation's tag columns. +async fn build_test_table_provider_with_string_field( + table_tags: &[(&str, &[&str])], +) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + for (table_name, tags) in table_tags { + let mut columns = tags + .iter() + .map(|tag| { + ColumnSchema::new( + (*tag).to_string(), + ConcreteDataType::string_datatype(), + false, + ) + }) + .collect::>(); + columns.push( + ColumnSchema::new( + greptime_timestamp().to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ); + columns.push(ColumnSchema::new( + greptime_value().to_string(), + ConcreteDataType::float64_datatype(), + true, + )); + columns.push(ColumnSchema::new( + "status".to_string(), + ConcreteDataType::string_datatype(), + true, + )); + let table_meta = TableMetaBuilder::empty() + .schema(Arc::new(Schema::new(columns))) + .primary_key_indices((0..tags.len()).collect()) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name((*table_name).to_string()) + .meta(table_meta) + .build() + .unwrap(); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: (*table_name).to_string(), + table_id: 1024, + table: EmptyTable::from_table_info(&table_info), + }) + .is_ok() + ); + } + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +async fn build_matching_filter_plan_with_string_field(query: &str) -> String { + let table_provider = build_test_table_provider_with_string_field(&[ + ("metric_a", &["host", "device"]), + ("metric_b", &["host", "device"]), + ]) + .await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt(query), + &build_query_engine_state(), + ) + .await + .unwrap(); + plan.display_indent().to_string() +} + +#[tokio::test] +async fn binary_matching_label_filter_reaches_both_operands() { + for query in [ + r#"metric_a / metric_b{host="foo"}"#, + r#"metric_a / on(host, device) metric_b{host="foo"}"#, + r#"count_over_time(metric_a[1m]) / on(host) count_over_time(metric_b{host="foo"}[1m])"#, + r#"metric_a / ignoring(device) metric_b{host="foo"}"#, + r#"sum by(host) (metric_a) / on(host) sum by(host) (metric_b{host="foo"})"#, + r#"metric_a / on(host) avg without(device) (metric_b{host="foo"})"#, + ] { + let plan = build_matching_filter_plan(query).await; + assert_eq!( + plan.matches(r#"host = Utf8("foo")"#).count(), + 2, + "{query}\n{plan}" + ); + } +} + +#[tokio::test] +async fn binary_matching_label_filter_reaches_scalar_ranking_and_grouped_operands() { + for query in [ + r#"(8 * metric_a{host="foo"}) / on(host) metric_b"#, + r#"topk(1, metric_a{host="foo"}) / on(host, device) metric_b"#, + r#"(8 * metric_a{host="foo"}) / on(host) group_left topk by(host)(1, max by(host)(metric_b))"#, + ] { + let plan = build_matching_filter_plan(query).await; + assert_eq!( + plan.matches(r#"host = Utf8("foo")"#).count(), + 2, + "{query}\n{plan}" + ); + } + // A global ranking one-side must see every host, so the matcher stays put. + let query = r#"metric_a{host="foo"} / on(host) group_left topk(1, max by(host)(metric_b))"#; + let plan = build_matching_filter_plan(query).await; + assert_eq!( + plan.matches(r#"host = Utf8("foo")"#).count(), + 1, + "{query}\n{plan}" + ); +} + +#[tokio::test] +async fn binary_matching_label_filter_skips_selecting_aggregations() { + // `topk` ranks its input, so filtering before it changes the candidate set. + let query = r#"topk(1, metric_a) / on(host, device) metric_b{host="foo"}"#; + let plan = build_matching_filter_plan(query).await; + assert_eq!(plan.matches("foo").count(), 1, "{query}\n{plan}"); +} + +#[tokio::test] +async fn binary_value_field_matcher_stays_on_its_own_operand() { + let value = greptime_value(); + for query in [ + format!(r#"metric_a / metric_b{{{value}="2"}}"#), + format!(r#"metric_a / on(host, device, {value}) metric_b{{{value}="2"}}"#), + ] { + let plan = build_matching_filter_plan(&query).await; + assert_eq!(plan.matches(r#"Utf8("2")"#).count(), 1, "{query}\n{plan}"); + } +} + +#[tokio::test] +async fn binary_value_field_matcher_is_not_copied_across_aggregations() { + // `status` varies between the samples of one series, so filtering the other operand by it + // would drop the newest sample before sample selection (#9242). + for query in [ + r#"count by(status) (metric_a) / on(status) count by(status) (metric_b{status="ready"})"#, + r#"(8 * count by(status)(metric_a{__field__="status"})) / on(status) topk by(status)(1, count by(status)(metric_b{__field__="status",status="ready"}))"#, + r#"count by(status)(metric_a{__field__="status"}) / on(status) group_left topk by(status)(1, count by(status)(metric_b{__field__="status",status="ready"}))"#, + ] { + let plan = build_matching_filter_plan_with_string_field(query).await; + assert_eq!( + plan.matches(r#"Utf8("ready")"#).count(), + 1, + "{query}\n{plan}" + ); + } +} + +#[tokio::test] +async fn binary_matching_label_filter_reaches_aggregations_grouping_by_tags() { + // `count`, not `sum`: the string value field is not summable. + let query = r#"count by(host) (metric_a) / on(host) count by(host) (metric_b{host="foo"})"#; + let plan = build_matching_filter_plan_with_string_field(query).await; + assert_eq!(plan.matches(r#"Utf8("foo")"#).count(), 2, "{query}\n{plan}"); +} + +#[tokio::test] +async fn binary_matching_label_filter_skips_unproven_expressions() { + for query in [ + r#"metric_a > on(host, device) metric_b{host="foo"}"#, + r#"metric_a / on(host) group_left metric_b{host="foo"}"#, + r#"metric_a / on(host) label_replace(metric_b{host="foo"},"extra","e","host",".*")"#, + r#"metric_a / on(device) metric_b{host="foo"}"#, + ] { + let plan = build_matching_filter_plan(query).await; + assert_eq!(plan.matches("foo").count(), 1, "{query}\n{plan}"); + } +} + +#[tokio::test] +async fn test_or_context_uses_left_qualified_output() { + let case = r#"(normal_metric or other_metric) + 1"#; + let eval_stmt = build_eval_stmt(case); + let state = build_query_engine_state(); + let plan = + PromPlanner::stmt_to_plan(build_or_context_table_provider().await, &eval_stmt, &state) + .await + .unwrap(); + assert!( + plan.schema() + .fields() + .iter() + .any(|field| field.data_type() == &ArrowDataType::Float64), + "{plan:?}" + ); + let (_optimized, _physical) = optimize_and_create_physical_plan(&state, plan).await; +} + +#[tokio::test] +async fn test_or_context_uses_left_qualified_empty_histogram_output() { + let case = r#"(abs(histogram_quantile(0.99, non_hist_metric)) or normal_metric) + 1"#; + let eval_stmt = build_eval_stmt(case); + let plan = PromPlanner::stmt_to_plan( + build_or_context_table_provider().await, + &eval_stmt, + &build_query_engine_state(), + ) + .await + .unwrap(); + assert!( + plan.schema() + .fields() + .iter() + .any(|field| field.data_type() == &ArrowDataType::Float64), + "{plan:?}" + ); +} + +#[tokio::test] +async fn test_direct_or_preserves_float_and_native_histogram_samples() { + for histogram_on_left in [false, true] { + let (planner, plan) = mixed_direct_or(histogram_on_left).await; + + let float_field = &planner.ctx.field_columns[0]; + let histogram_field = &planner.ctx.field_columns[1]; + assert!(float_field.starts_with(OR_FLOAT_FIELD_PREFIX)); + assert!(histogram_field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)); + assert_eq!( + plan.schema() + .field_with_name(None, float_field) + .unwrap() + .data_type(), + &ArrowDataType::Float64 + ); + assert_eq!( + plan.schema() + .field_with_name(None, histogram_field) + .unwrap() + .data_type(), + &native_histogram_value_type().as_arrow_type() + ); + + let (optimized, batches) = execute(plan, &build_query_engine_state()).await; + assert_no_internal_or_keys(optimized.schema()); + let mut sample_kinds = batches + .iter() + .flat_map(|batch| { + let values = batch.column_by_name(float_field).unwrap(); + let histograms = batch.column_by_name(histogram_field).unwrap(); + (0..batch.num_rows()).map(|row| (values.is_valid(row), histograms.is_valid(row))) + }) + .collect::>(); + sample_kinds.sort_unstable(); + assert_eq!(sample_kinds, vec![(false, true), (true, false)]); + } +} + +#[tokio::test] +async fn malformed_classic_bucket_does_not_drop_native_histogram() { + let state = build_query_engine_state(); + let collector = PromqlAnnotationCollector::default(); + let plan = PromPlanner::stmt_to_plan_with_annotations( + operator_table_provider(), + &operator_eval_stmt("histogram_quantile(0.5, bad_classic or bad_native)"), + &state, + Some(collector.clone()), + ) + .await + .unwrap(); + let value_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + assert_eq!(values(&batches, &value_field), vec![0.0]); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty()); + assert!(infos.is_empty()); +} + +#[tokio::test] +async fn test_mixed_binary_operator_aligns_both_alternative_inputs() { + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt("(lf or on(tag) lh) * on(tag) (rf or on(tag) rh)"), + &state, + ) + .await + .unwrap(); + let plan_text = plan.display_indent_schema().to_string(); + assert!( + plan_text.contains("prom_native_histogram_mul_scalar"), + "{plan_text}" + ); + assert!( + plan_text.contains("prom_native_histogram_scalar_mul"), + "{plan_text}" + ); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX)) + .unwrap() + .name() + .clone(); + let histogram_field = plan + .schema() + .fields() + .iter() + .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX)) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + assert!(values(&batches, &float_field).is_empty()); + let mut sums = histograms(&batches, &histogram_field) + .into_iter() + .map(|histogram| histogram.sum) + .collect::>(); + sums.sort_by(f64::total_cmp); + assert_eq!(sums, vec![2.0, 3.0]); +} + +#[tokio::test] +async fn test_mixed_binary_operator_reports_only_dropped_samples() { + for (query, expected_rows, expected_infos) in [ + ("(lf or on(tag) lh) + on(tag) (rf or on(tag) rh)", 0, 1), + ("(lf or on(tag) lh) + on(tag) (lf or on(tag) lh)", 2, 0), + ("(lf or on(tag) lh) % on(tag) lh", 0, 1), + ] { + let state = build_query_engine_state(); + let annotations = PromqlAnnotationCollector::default(); + let plan = PromPlanner::stmt_to_plan_with_annotations( + operator_table_provider(), + &operator_eval_stmt(query), + &state, + Some(annotations.clone()), + ) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + expected_rows, + "{query}" + ); + let mut warnings = vec![]; + let mut infos = vec![]; + annotations.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty(), "{query}: {warnings:?}"); + assert_eq!(infos.len(), expected_infos, "{query}: {infos:?}"); + } +} + +#[tokio::test] +async fn test_histogram_only_min_drops_empty_aggregate_group() { + // `min` over native-histogram-only input drops every sample in the group, so the + // NULL-valued aggregate row must be filtered out. Otherwise an outer expression + // like `group()` resurrects the group Prometheus considers unseen. + let state = build_query_engine_state(); + for query in ["min(lh)", "group(min(lh))"] { + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap(); + let (_, batches) = execute(plan, &state).await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + 0, + "{query}" + ); + } +} + +#[tokio::test] +async fn test_mixed_min_drops_histogram_only_group() { + // With alternative float/histogram fields, `min by (tag)` keeps float-only groups + // (tag=a from `lf`) and drops histogram-only groups (tag=b from `lh`) instead of + // emitting a NULL-valued row for them. + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt("min by (tag) (lf or on(tag) lh)"), + &state, + ) + .await + .unwrap(); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + assert_eq!(values(&batches, &float_field), vec![2.0]); +} + +#[tokio::test] +async fn test_mixed_or_can_feed_another_or() { + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt("lf or on(tag) lh or on(tag) fallback"), + &state, + ) + .await + .unwrap(); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX)) + .unwrap() + .name() + .clone(); + let histogram_field = plan + .schema() + .fields() + .iter() + .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX)) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 3); + let mut float_values = values(&batches, &float_field); + float_values.sort_by(f64::total_cmp); + assert_eq!(float_values, vec![2.0, 7.0]); + assert_eq!(histograms(&batches, &histogram_field).len(), 1); +} + +#[tokio::test] +async fn test_mixed_fields_align_with_single_float_vector() { + let (planner, mixed) = mixed_direct_or(false).await; + let scale = tagged_source( + "scale", + false, + ("k", Some("float")), + DirectOrValue::Float64(2.0), + ); + let scale = scan(&scale); + let scale_fields = vec!["v".to_string()]; + let PromExpr::Binary(binary) = parser::parse("lhs * rhs").unwrap() else { + unreachable!() + }; + + let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns( + mixed.schema(), + scale.schema(), + &planner.ctx.field_columns, + &scale_fields, + binary.op, + false, + false, + ); + assert!(invalid_pairs.is_empty()); + assert_eq!( + groups + .iter() + .map(|(output, _)| output.clone()) + .collect::>(), + planner.ctx.field_columns + ); + assert_eq!(groups.len(), 2); + assert!( + groups + .iter() + .flat_map(|(_, pairs)| pairs) + .all(|(_, right)| *right == &scale_fields[0]) + ); + + let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns( + scale.schema(), + mixed.schema(), + &scale_fields, + &planner.ctx.field_columns, + binary.op, + false, + false, + ); + assert!(invalid_pairs.is_empty()); + assert_eq!( + groups + .iter() + .map(|(output, _)| output.clone()) + .collect::>(), + planner.ctx.field_columns + ); + assert_eq!(groups.len(), 2); + assert!( + groups + .iter() + .flat_map(|(_, pairs)| pairs) + .all(|(left, _)| *left == &scale_fields[0]) + ); +} + +#[tokio::test] +async fn test_non_bool_comparison_filters_mixed_sample_lanes() { + let (planner, input) = mixed_direct_or(false).await; + let input_schema = input.schema().clone(); + let plan = planner + .filter_on_field_column(input, |field| { + if PromPlanner::field_column_is_native_histogram(&input_schema, field) { + Ok(lit(false)) + } else { + Ok(col(field).gt(lit(0.0))) + } + }) + .unwrap(); + let float_field = planner.ctx.field_columns[0].clone(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(values(&batches, &float_field), vec![1.25]); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); +} + +#[tokio::test] +async fn test_mixed_left_and_unless_preserve_sample_lanes() { + for (expression, expected_sample_kind) in [ + ("lhs and on(k) mask", (false, true)), + ("lhs unless on(k) mask", (true, false)), + ] { + let (mut planner, left) = mixed_direct_or(false).await; + let left_context = planner.ctx.clone(); + let float_field = left_context.field_columns[0].clone(); + let histogram_field = left_context.field_columns[1].clone(); + let mask = tagged_source( + "mask", + false, + ("k", Some("histogram")), + DirectOrValue::Float64(1.0), + ); + let PromExpr::Binary(binary) = parser::parse(expression).unwrap() else { + unreachable!() + }; + let plan = planner + .set_op_on_non_field_columns( + left, + scan(&mask), + left_context, + direct_or_context("mask", &["job", "k"], "v"), + binary.op, + &binary.modifier, + ) + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + let sample_kinds = batches + .iter() + .flat_map(|batch| { + let floats = batch.column_by_name(&float_field).unwrap(); + let histograms = batch.column_by_name(&histogram_field).unwrap(); + (0..batch.num_rows()).map(|row| (floats.is_valid(row), histograms.is_valid(row))) + }) + .collect::>(); + assert_eq!(sample_kinds, vec![expected_sample_kind], "{expression}"); + } +} + +#[tokio::test] +async fn test_mixed_fields_arithmetic_broadcasts_computed_scalar() { + let plan = PromPlanner::stmt_to_plan( + build_test_mixed_native_histogram_table_provider("some_metric").await, + &build_eval_stmt("some_metric * scalar(vector(2))"), + &build_query_engine_state(), + ) + .await + .unwrap(); + let schema = plan.schema(); + assert_eq!( + schema + .field_with_unqualified_name(greptime_value()) + .unwrap() + .data_type(), + &ArrowDataType::Float64 + ); + assert_eq!( + schema + .field_with_unqualified_name(greptime_native_histogram()) + .unwrap() + .data_type(), + &native_histogram_value_type().as_arrow_type() + ); + assert!( + plan.display_indent_schema() + .to_string() + .contains("prom_native_histogram_mul_scalar"), + "{plan:?}" + ); +} + +#[tokio::test] +async fn test_unsupported_histogram_binary_does_not_block_or_fallback() { + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt("((lf or on(tag) lh) % 2) or on(tag) lh"), + &state, + ) + .await + .unwrap(); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + let histogram_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &native_histogram_value_type().as_arrow_type()) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(values(&batches, &float_field), vec![0.0]); + assert_eq!(histograms(&batches, &histogram_field).len(), 1); +} + +#[tokio::test] +async fn test_unary_negates_mixed_float_and_native_histogram_samples() { + for histogram_on_left in [false, true] { + let (mut planner, input) = mixed_direct_or(histogram_on_left).await; + let plan = planner.negate_field_columns(input).unwrap(); + assert!(PromPlanner::field_columns_are_alternative_samples( + plan.schema(), + &planner.ctx.field_columns + )); + let float_field = planner + .ctx + .field_columns + .iter() + .find(|field| field.starts_with(OR_FLOAT_FIELD_PREFIX)) + .unwrap(); + let histogram_field = planner + .ctx + .field_columns + .iter() + .find(|field| field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)) + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(values(&batches, float_field), vec![-1.25]); + let histogram = batches + .iter() + .find_map(|batch| { + let values = batch + .column_by_name(histogram_field) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..values.len()).find_map(|row| { + common_query::native_histogram::read_histogram(values, row).unwrap() + }) + }) + .unwrap(); + assert_eq!(histogram.count, -1.0); + assert_eq!(histogram.sum, -1.0); + assert_eq!(histogram.reset_hint, CounterResetHint::Gauge); + } +} + +#[tokio::test] +async fn test_native_histogram_sum_and_avg_execute_real_batches() { + for op_name in ["sum", "avg"] { + for incompatible in [false, true] { + let mut second = direct_or_histogram(); + if incompatible { + second.schema = CUSTOM_BUCKETS_SCHEMA; + second.custom_values = vec![1.0]; + } + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = + mixed_aggregate_input(vec![direct_or_histogram(), second]).await; + planner.promql_annotations = Some(collector.clone()); + let histogram_column = planner.ctx.field_columns[1].clone(); + planner.ctx.field_columns = vec![histogram_column.clone()]; + let input = LogicalPlanBuilder::from(input) + .project([col("ts"), col(&histogram_column)]) + .unwrap() + .build() + .unwrap(); + let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = + parser::parse(&format!("{op_name}(mixed)")).unwrap() + else { + unreachable!() + }; + let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts")], aggregate_exprs) + .unwrap() + .filter(planner.create_empty_values_filter_expr(false).unwrap()) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(infos.is_empty()); + if incompatible { + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + assert!(warnings.iter().any(|warning| { + warning + == &format!( + "prom_native_histogram_agg_{op_name}: dropped native histogram aggregate with incompatible schemas" + ) + })); + } else { + let histograms = histograms(&batches, &histogram_column); + assert_eq!(histograms.len(), 1); + let expected = if op_name == "sum" { 2.0 } else { 1.0 }; + assert_eq!(histograms[0].count, expected); + assert_eq!(histograms[0].sum, expected); + assert!(warnings.is_empty()); + } + } + } +} + +#[tokio::test] +async fn test_canonical_mixed_count_group_and_count_values_execute() { + let state = build_query_engine_state(); + for (query, expected) in [ + ("count(some_metric)", vec![2.0]), + ("group(some_metric)", vec![1.0]), + (r#"count_values("sample", some_metric)"#, vec![1.0, 1.0]), + ] { + let plan = PromPlanner::stmt_to_plan( + build_test_mixed_native_histogram_table_provider("some_metric").await, + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap(); + assert!( + plan.schema() + .fields() + .iter() + .all(|field| !field.name().starts_with("__promql_sample_count")), + "{query}: {plan:?}" + ); + let value_fields = plan + .schema() + .fields() + .iter() + .filter(|field| { + matches!( + field.data_type(), + ArrowDataType::Float64 | ArrowDataType::Int64 | ArrowDataType::UInt64 + ) || field.data_type() == &native_histogram_value_type().as_arrow_type() + }) + .collect::>(); + assert_eq!(value_fields.len(), 1, "{query}: {plan:?}"); + assert_ne!( + value_fields[0].data_type(), + &native_histogram_value_type().as_arrow_type(), + "{query}: {plan:?}" + ); + let value_column = value_fields[0].name().clone(); + + let (_, batches) = execute(plan, &state).await; + let mut actual = numeric_values(&batches, &value_column); + actual.sort_by(f64::total_cmp); + assert_eq!(actual, expected, "{query}"); + + if query.starts_with("count_values") { + let mut sample_labels = batches + .iter() + .flat_map(|batch| { + batch + .column_by_name("sample") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + .map(str::to_string) + }) + .collect::>(); + sample_labels.sort(); + let mut expected_labels = vec!["2".to_string(), direct_or_histogram().promql_string()]; + expected_labels.sort(); + assert_eq!(sample_labels, expected_labels); + } + } +} + +#[tokio::test] +async fn test_mixed_or_sum_aggregates_each_sample_type() { + let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = parser::parse("sum(lhs)").unwrap() + else { + unreachable!() + }; + + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = mixed_direct_or(false).await; + planner.promql_annotations = Some(collector.clone()); + let float_column = planner.ctx.field_columns[0].clone(); + let histogram_column = planner.ctx.field_columns[1].clone(); + let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts"), col("k")], aggregate_exprs) + .unwrap() + .filter( + planner + .mixed_aggregate_filter_expr(op, &float_column, &histogram_column) + .unwrap(), + ) + .unwrap() + .project([ + col(&float_column), + col(&histogram_column), + col("ts"), + col("k"), + ]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(values(&batches, &float_column), vec![1.25]); + let histogram = batches + .iter() + .find_map(|batch| { + let values = batch + .column_by_name(&histogram_column)? + .as_any() + .downcast_ref::()?; + (0..values.len()).find_map(|row| { + common_query::native_histogram::read_histogram(values, row).unwrap() + }) + }) + .unwrap(); + assert_eq!(histogram.count, 1.0); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty()); + + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = mixed_direct_or(false).await; + planner.promql_annotations = Some(collector.clone()); + let float_column = planner.ctx.field_columns[0].clone(); + let histogram_column = planner.ctx.field_columns[1].clone(); + let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts")], aggregate_exprs) + .unwrap() + .filter( + planner + .mixed_aggregate_filter_expr(op, &float_column, &histogram_column) + .unwrap(), + ) + .unwrap() + .project([col(&float_column), col(&histogram_column), col("ts")]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert_eq!( + warnings, + vec!["sum: dropped aggregation result containing both float and native histogram samples"] + ); +} + +#[tokio::test] +async fn test_mixed_or_sum_drops_incompatible_mixed_group() { + let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = parser::parse("sum(lhs)").unwrap() + else { + unreachable!() + }; + let mut custom = direct_or_histogram(); + custom.schema = CUSTOM_BUCKETS_SCHEMA; + custom.custom_values = vec![1.0]; + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = mixed_aggregate_input(vec![direct_or_histogram(), custom]).await; + planner.promql_annotations = Some(collector.clone()); + let float_column = planner.ctx.field_columns[0].clone(); + let histogram_column = planner.ctx.field_columns[1].clone(); + let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts")], aggregate_exprs) + .unwrap() + .filter( + planner + .mixed_aggregate_filter_expr(op, &float_column, &histogram_column) + .unwrap(), + ) + .unwrap() + .project([col(&float_column), col(&histogram_column), col("ts")]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.iter().any(|warning| { + warning + == "sum: dropped aggregation result containing both float and native histogram samples" + })); +} + +#[tokio::test] +async fn test_mixed_or_min_records_only_present_histograms() { + let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = parser::parse("min(lhs)").unwrap() + else { + unreachable!() + }; + let expected_info = "min: dropped native histogram samples because this aggregation is not supported for native histograms"; + + for (histograms, expected_infos) in [ + (vec![], vec![]), + (vec![direct_or_histogram()], vec![expected_info]), + ] { + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = mixed_aggregate_input(histograms).await; + planner.promql_annotations = Some(collector.clone()); + let float_column = planner.ctx.field_columns[0].clone(); + let histogram_column = planner.ctx.field_columns[1].clone(); + let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts")], aggregate_exprs) + .unwrap() + .filter( + planner + .mixed_ignored_histogram_filter_expr(op, &histogram_column) + .unwrap(), + ) + .unwrap() + .project([col(&float_column), col("ts")]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(values(&batches, &float_column), vec![1.25]); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty()); + assert_eq!(infos, expected_infos); + } +} + +#[tokio::test] +async fn test_mixed_or_value_aliases_do_not_replace_labels() { + let left = source( + "lhs", + false, + 1, + vec![("job", Some("job")), ("k", Some("float"))], + DirectOrValue::Float64(1.0), + ); + let right = source( + "rhs", + false, + 1, + vec![ + ("job", Some("job")), + ("k", Some("histogram")), + (greptime_value(), Some("value-label")), + ], + DirectOrValue::NativeHistogram(direct_or_histogram()), + ); + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let mut planner = PromPlanner { + table_provider, + ctx: PromPlannerContext::default(), + promql_annotations: None, + }; + let left = LogicalPlanBuilder::from(scan(&left)) + .project(vec![ + col("ts"), + col("job"), + col("k"), + col("v").alias(greptime_value()), + ]) + .unwrap() + .build() + .unwrap(); + let left_context = direct_or_context("lhs", &["job", "k"], greptime_value()); + let right_context = direct_or_context("rhs", &["job", "k", greptime_value()], "v"); + let plan = planner + .or_operator( + left, + scan(&right), + left_context.tag_columns.iter().cloned().collect(), + right_context.tag_columns.iter().cloned().collect(), + left_context, + right_context, + &or_modifier("lhs or on(k) rhs"), + ) + .unwrap(); + + assert_eq!( + plan.schema() + .field_with_name(None, greptime_value()) + .unwrap() + .data_type(), + &ArrowDataType::Utf8 + ); + assert!( + planner + .ctx + .field_columns + .iter() + .all(|field| { field != greptime_value() && field != greptime_native_histogram() }) + ); + assert!(PromPlanner::field_columns_are_alternative_samples( + plan.schema(), + &planner.ctx.field_columns + )); + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + let labels = batches + .iter() + .flat_map(|batch| { + batch + .column_by_name(greptime_value()) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + }) + .collect::>(); + assert_eq!(labels, vec!["value-label"]); +} + +#[tokio::test] +async fn test_mixed_or_routes_float_histogram_and_label_functions() { + for (function, expected) in [("abs", 1.25), ("round", 1.0), ("histogram_count", 1.0)] { + let (mut planner, input) = mixed_direct_or(false).await; + let preserve_any_value = PromPlanner::field_columns_are_alternative_samples( + input.schema(), + &planner.ctx.field_columns, + ); + let PromExpr::Call(call) = parser::parse(&format!("{function}(lhs)")).unwrap() else { + unreachable!() + }; + let state = build_query_engine_state(); + let (mut exprs, _) = planner + .create_function_expr(&call.func, vec![], input.schema(), &state) + .unwrap(); + exprs.insert(0, planner.create_time_index_column_expr().unwrap()); + exprs.extend(planner.create_tag_column_exprs().unwrap()); + let plan = LogicalPlanBuilder::from(input) + .project(exprs) + .unwrap() + .filter( + planner + .create_empty_values_filter_expr(preserve_any_value) + .unwrap(), + ) + .unwrap() + .build() + .unwrap(); + let (_, batches) = execute(plan, &state).await; + let values = batches + .iter() + .flat_map(|batch| { + batch + .schema() + .fields() + .iter() + .position(|field| field.data_type() == &ArrowDataType::Float64) + .map(|index| { + batch + .column(index) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + }) + .into_iter() + .flatten() + }) + .collect::>(); + assert_eq!(values, vec![expected], "{function}"); + } + + let (mut planner, input) = mixed_direct_or(false).await; + let preserve_any_value = PromPlanner::field_columns_are_alternative_samples( + input.schema(), + &planner.ctx.field_columns, + ); + let PromExpr::Call(call) = + parser::parse(r#"label_replace(lhs, "copy", "$1", "k", "(.*)")"#).unwrap() + else { + unreachable!() + }; + let args = planner.create_function_args(&call.args.args).unwrap(); + let state = build_query_engine_state(); + let (mut exprs, _) = planner + .create_function_expr(&call.func, args.literals, input.schema(), &state) + .unwrap(); + exprs.insert(0, planner.create_time_index_column_expr().unwrap()); + exprs.extend(planner.create_tag_column_exprs().unwrap()); + let plan = LogicalPlanBuilder::from(input) + .project(exprs) + .unwrap() + .filter( + planner + .create_empty_values_filter_expr(preserve_any_value) + .unwrap(), + ) + .unwrap() + .build() + .unwrap(); + let (_, batches) = execute(plan, &state).await; + let sample_count = batches.iter().map(RecordBatch::num_rows).sum::(); + assert_eq!(sample_count, 2); +} + +/// Table provider with a single metric `cv_metric` holding three series at `ts=1000`: +/// `k="a"` and `k="c"` carry the value 1.0, `k="b"` carries 2.0. +async fn build_count_values_table_provider() -> DfTableSourceProvider { + build_count_values_table_provider_with_values(&[1.0, 2.0, 1.0]).await +} + +/// Table provider with a single metric `cv_metric` holding one series per given value at +/// `ts=1000` (`k` names the series, the sample value is the given value). +async fn build_count_values_table_provider_with_values(values: &[f64]) -> DfTableSourceProvider { + build_count_values_table_provider_with_value_array(Arc::new(Float64Array::from( + values.to_vec(), + ))) + .await +} + +/// Like [`build_count_values_table_provider_with_values`], but with a caller provided value +/// column, so tests can cover value columns that are not `Float64` (e.g. `BIGINT`). +async fn build_count_values_table_provider_with_value_array( + values: ArrayRef, +) -> DfTableSourceProvider { + let value_data_type = ConcreteDataType::from_arrow_type(values.data_type()); + let catalog_list = MemoryCatalogManager::with_default_setup(); + let columns = vec![ + ColumnSchema::new("k".to_string(), ConcreteDataType::string_datatype(), false), + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new(greptime_value().to_string(), value_data_type, true), + ]; + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema.clone()) + .primary_key_indices(vec![0]) + .value_indices(vec![2]) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = Arc::new( + TableInfoBuilder::default() + .table_id(3_001) + .name("cv_metric") + .meta(table_meta) + .build() + .unwrap(), + ); + let batch = RecordBatch::try_new( + schema.arrow_schema().clone(), + vec![ + Arc::new(StringArray::from( + (0..values.len()) + .map(|index| format!("k{index}")) + .collect::>(), + )), + Arc::new(TimestampMillisecondArray::from(vec![1_000; values.len()])), + values.clone(), + ], + ) + .unwrap(); + let backing = GreptimeMemTable::new_with_catalog( + "cv_metric", + GreptimeRecordBatch::from_df_record_batch(schema, batch), + 3_001, + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + ); + let table = Arc::new(Table::new( + table_info, + FilterPushDownType::Unsupported, + backing.data_source(), + )); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "cv_metric".to_string(), + table_id: 3_001, + table, + }) + .is_ok() + ); + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +/// Collects `(label, value)` pairs of a `count_values` result, where `label` is the +/// PromQL label generated by `count_values` and `value` is the aggregated sample value. +/// +/// The generated label holds the original sample value in PromQL's textual form, so it +/// is asserted as a string: comparing it as a number would not catch formatting bugs +/// (`1.0` instead of `1`, scientific notation, ...). +fn count_values_rows<'a>(batches: &'a [RecordBatch], label: &str) -> Vec<(&'a str, f64)> { + let mut rows = batches + .iter() + .flat_map(|batch| { + // The aggregated value is the only numeric column that is not the generated label. + let value_index = batch + .schema() + .fields() + .iter() + .position(|field| { + field.name() != label + && matches!( + field.data_type(), + ArrowDataType::Float64 | ArrowDataType::Int64 | ArrowDataType::UInt64 + ) + }) + .expect("no aggregated value column"); + let labels = batch + .column_by_name(label) + .expect("no generated label column") + .as_any() + .downcast_ref::() + .expect("the generated label must be a string column"); + let values = datafusion::arrow::compute::cast( + batch.column(value_index), + &ArrowDataType::Float64, + ) + .unwrap(); + let values = values.as_any().downcast_ref::().unwrap(); + labels + .iter() + .zip(values.iter()) + .map(|(label, value)| (label.unwrap(), value.unwrap())) + .collect::>() + }) + .collect::>(); + rows.sort_by(|left, right| left.0.cmp(right.0).then(left.1.total_cmp(&right.1))); + rows +} + +/// Asserts that a `count_values` result holds one sample per label set and evaluation +/// timestamp: Prometheus groups by the generated label, so a timestamp must never repeat +/// a label set (that would mean the samples were still grouped by the overwritten input +/// label). +fn assert_unique_label_set_per_timestamp(batches: &[RecordBatch], label: &str) { + let mut seen = HashMap::>::new(); + for batch in batches { + let timestamp_index = batch + .schema() + .fields() + .iter() + .position(|field| matches!(field.data_type(), ArrowDataType::Timestamp(..))) + .expect("no timestamp column"); + let timestamps = batch + .column(timestamp_index) + .as_any() + .downcast_ref::() + .expect("timestamp column is not a millisecond timestamp"); + let labels = batch + .column_by_name(label) + .expect("no generated label column") + .as_any() + .downcast_ref::() + .expect("the generated label must be a string column"); + for (timestamp, label) in timestamps.iter().zip(labels.iter()) { + let timestamp = timestamp.unwrap(); + let label = label.unwrap(); + assert!( + seen.entry(timestamp).or_default().insert(label.to_string()), + "duplicated label set `{label}` at timestamp {timestamp}" + ); + } + } +} + +#[tokio::test] +async fn test_count_values_generated_label_survives_enclosing_expr() { + // https://github.com/GreptimeTeam/greptimedb/issues/9181 + for (case, label) in [ + (r#"count_values("v", prometheus_tsdb_head_series)"#, "v"), + ( + r#"abs(count_values("v", prometheus_tsdb_head_series))"#, + "v", + ), + ( + r#"round(count_values("v", prometheus_tsdb_head_series))"#, + "v", + ), + (r#"count_values("v", prometheus_tsdb_head_series) + 1"#, "v"), + ( + r#"topk(1, count_values("v", prometheus_tsdb_head_series))"#, + "v", + ), + ( + r#"sum by (v) (count_values("v", prometheus_tsdb_head_series))"#, + "v", + ), + ( + r#"label_replace(count_values("v", prometheus_tsdb_head_series), "vcopy", "$1", "v", "(.*)")"#, + "v", + ), + ( + r#"count_values("v", prometheus_tsdb_head_series) by (ip) + 1"#, + "v", + ), + // The generated label overwrites an input label with the same name. + ( + r#"count_values("ip", prometheus_tsdb_head_series) by (ip)"#, + "ip", + ), + ( + r#"count_values("ip", prometheus_tsdb_head_series) by (ip) + 1"#, + "ip", + ), + ] { + let plan = PromPlanner::stmt_to_plan( + build_test_table_provider_with_fields( + &[( + DEFAULT_SCHEMA_NAME.to_string(), + "prometheus_tsdb_head_series".to_string(), + )], + &["ip"], + ) + .await, + &build_eval_stmt(case), + &build_query_engine_state(), + ) + .await + .unwrap(); + + let label_columns = plan + .schema() + .fields() + .iter() + .filter(|field| field.name() == label) + .count(); + assert_eq!( + label_columns, + 1, + "{case}: the `{label}` label must survive: {}", + plan.display_indent() + ); + } +} + +#[tokio::test] +async fn test_count_values_generated_label_in_enclosing_expr_execute() { + // https://github.com/GreptimeTeam/greptimedb/issues/9181 + let state = build_query_engine_state(); + // (query, generated label, expected `(label value, aggregated value)` pairs) + for (query, label, expected) in [ + ( + r#"count_values("v", cv_metric)"#, + "v", + vec![("1", 2.0), ("2", 1.0)], + ), + ( + r#"abs(count_values("v", cv_metric))"#, + "v", + vec![("1", 2.0), ("2", 1.0)], + ), + ( + r#"round(count_values("v", cv_metric))"#, + "v", + vec![("1", 2.0), ("2", 1.0)], + ), + ( + r#"count_values("v", cv_metric) + 1"#, + "v", + vec![("1", 3.0), ("2", 2.0)], + ), + ( + r#"sum by (v) (count_values("v", cv_metric))"#, + "v", + vec![("1", 2.0), ("2", 1.0)], + ), + ( + r#"topk(10, count_values("v", cv_metric))"#, + "v", + vec![("1", 2.0), ("2", 1.0)], + ), + // The generated label overwrites the input label with the same name, and the + // samples are grouped by the generated label only: `{k="1"}` holds the two + // samples of value `1.0` instead of one row per (overwritten label, value). + ( + r#"count_values("k", cv_metric) by (k)"#, + "k", + vec![("1", 2.0), ("2", 1.0)], + ), + ] { + let plan = PromPlanner::stmt_to_plan( + build_count_values_table_provider().await, + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap_or_else(|err| panic!("{query}: {err}")); + + assert_eq!( + plan.schema() + .fields() + .iter() + .filter(|field| field.name() == label) + .count(), + 1, + "{query}: {}", + plan.display_indent() + ); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(count_values_rows(&batches, label), expected, "{query}"); + assert_unique_label_set_per_timestamp(&batches, label); + } +} + +#[tokio::test] +async fn test_count_values_label_is_prometheus_formatted_value() { + // PromQL materializes the generated label with `strconv.FormatFloat(value, 'f', -1, 64)`: + // the shortest decimal form of the sample value without an exponent. The label is a + // label, so it must be a string column holding exactly that text: arrow's + // `Float64 -> Utf8` cast would render `1`/`200`/`1e21` as `1.0`/`200.0`/`1e21`. + let state = build_query_engine_state(); + // Samples are grouped by that formatted text, exactly like Prometheus groups by the + // generated label: `-0.0` and `0.0` are two series (`-0` and `0`), while values that + // round to the same text share one group. `0.0` formats as "0". + let plan = PromPlanner::stmt_to_plan( + build_count_values_table_provider_with_values(&[ + -0.0, 0.0, 1.0, 0.5, 200.0, 1e21, 1e-7, 2.5, + ]) + .await, + &operator_eval_stmt(r#"count_values("v", cv_metric)"#), + &state, + ) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + assert_unique_label_set_per_timestamp(&batches, "v"); + let mut labels = count_values_rows(&batches, "v") + .into_iter() + .map(|(label, _)| label) + .collect::>(); + labels.sort(); + assert_eq!( + labels, + vec![ + "-0", + "0", + "0.0000001", + "0.5", + "1", + "1000000000000000000000", + "2.5", + "200", + ] + ); +} + +#[tokio::test] +async fn test_count_values_groups_by_formatted_value_for_bigint_input() { + // The grouping key of `count_values` is the formatted sample value, not the raw input + // value. Two `BIGINT` values that differ below the `Float64` precision (`2^53` and + // `2^53 + 1`) cast and format to the same label, so they must share one group and one + // count, exactly like Prometheus, which groups by the generated label text. + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + build_count_values_table_provider_with_value_array(Arc::new(Int64Array::from(vec![ + 9_007_199_254_740_992_i64, + 9_007_199_254_740_993_i64, + ]))) + .await, + &operator_eval_stmt(r#"count_values("v", cv_metric)"#), + &state, + ) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + // One timestamp must never carry the same label set twice. + assert_unique_label_set_per_timestamp(&batches, "v"); + assert_eq!( + count_values_rows(&batches, "v"), + vec![("9007199254740992", 2.0)] + ); +}