1use std::collections::{BTreeSet, HashMap, HashSet, VecDeque};
16use std::sync::Arc;
17use std::time::UNIX_EPOCH;
18
19use arrow::datatypes::IntervalDayTime;
20use async_recursion::async_recursion;
21use catalog::table_source::DfTableSourceProvider;
22use common_error::ext::ErrorExt;
23use common_error::status_code::StatusCode;
24use common_function::function::FunctionContext;
25use common_query::prelude::greptime_value;
26use datafusion::common::DFSchemaRef;
27use datafusion::datasource::DefaultTableSource;
28use datafusion::functions_aggregate::average::avg_udaf;
29use datafusion::functions_aggregate::count::count_udaf;
30use datafusion::functions_aggregate::expr_fn::first_value;
31use datafusion::functions_aggregate::min_max::{max_udaf, min_udaf};
32use datafusion::functions_aggregate::stddev::stddev_pop_udaf;
33use datafusion::functions_aggregate::sum::sum_udaf;
34use datafusion::functions_aggregate::variance::var_pop_udaf;
35use datafusion::functions_window::row_number::RowNumber;
36use datafusion::logical_expr::expr::{Alias, ScalarFunction, WindowFunction};
37use datafusion::logical_expr::expr_rewriter::normalize_cols;
38use datafusion::logical_expr::{
39 BinaryExpr, Cast, Extension, LogicalPlan, LogicalPlanBuilder, Operator,
40 ScalarUDF as ScalarUdfDef, WindowFrame, WindowFunctionDefinition,
41};
42use datafusion::prelude as df_prelude;
43use datafusion::prelude::{Column, Expr as DfExpr, JoinType};
44use datafusion::scalar::ScalarValue;
45use datafusion::sql::TableReference;
46use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter};
47use datafusion_common::{DFSchema, NullEquality};
48use datafusion_expr::expr::WindowFunctionParams;
49use datafusion_expr::utils::conjunction;
50use datafusion_expr::{
51 ExprSchemable, Literal, Projection, SortExpr, TableScan, TableSource, col, lit,
52};
53use datatypes::arrow::datatypes::{DataType as ArrowDataType, TimeUnit as ArrowTimeUnit};
54use datatypes::data_type::ConcreteDataType;
55use itertools::Itertools;
56use once_cell::sync::Lazy;
57use promql::extension_plan::{
58 Absent, EmptyMetric, HistogramFold, InstantManipulate, Millisecond, RangeManipulate,
59 ScalarCalculate, SeriesDivide, SeriesNormalize, UnionDistinctOn, build_special_time_expr,
60};
61use promql::functions::{
62 AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing,
63 IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, PredictLinear, PresentOverTime,
64 QuantileOverTime, Rate, Resets, Round, StddevOverTime, StdvarOverTime, SumOverTime,
65 quantile_udaf,
66};
67use promql_parser::label::{METRIC_NAME, MatchOp, Matcher, Matchers};
68use promql_parser::parser::token::TokenType;
69use promql_parser::parser::value::ValueType;
70use promql_parser::parser::{
71 AggregateExpr, BinModifier, BinaryExpr as PromBinaryExpr, Call, EvalStmt, Expr as PromExpr,
72 Function, FunctionArgs as PromFunctionArgs, LabelModifier, MatrixSelector, NumberLiteral,
73 Offset, ParenExpr, StringLiteral, SubqueryExpr, UnaryExpr, VectorMatchCardinality,
74 VectorSelector, token,
75};
76use regex::{self, Regex};
77use snafu::{OptionExt, ResultExt, ensure};
78use store_api::metric_engine_consts::{
79 DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, LOGICAL_TABLE_METADATA_KEY,
80 METRIC_ENGINE_NAME, is_metric_engine_internal_column,
81};
82use table::table::adapter::DfTableProviderAdapter;
83
84use crate::parser::{
85 ALIAS_NODE_NAME, ANALYZE_NODE_NAME, ANALYZE_VERBOSE_NODE_NAME, AliasExpr, EXPLAIN_NODE_NAME,
86 EXPLAIN_VERBOSE_NODE_NAME,
87};
88use crate::promql::error::{
89 CatalogSnafu, ColumnNotFoundSnafu, CombineTableColumnMismatchSnafu, DataFusionPlanningSnafu,
90 ExpectRangeSelectorSnafu, FunctionInvalidArgumentSnafu, InvalidDestinationLabelNameSnafu,
91 InvalidRegularExpressionSnafu, InvalidTimeRangeSnafu, MultiFieldsNotSupportedSnafu,
92 MultipleMetricMatchersSnafu, MultipleVectorSnafu, NoMetricMatcherSnafu, PromqlPlanNodeSnafu,
93 Result, SameLabelSetSnafu, TableNameNotFoundSnafu, TimeIndexNotFoundSnafu,
94 UnexpectedPlanExprSnafu, UnexpectedTokenSnafu, UnknownTableSnafu, UnsupportedExprSnafu,
95 UnsupportedMatcherOpSnafu, UnsupportedVectorMatchSnafu, ValueNotFoundSnafu,
96 ZeroRangeSelectorSnafu,
97};
98use crate::query_engine::QueryEngineState;
99
100const SPECIAL_TIME_FUNCTION: &str = "time";
102const SCALAR_FUNCTION: &str = "scalar";
104const SPECIAL_ABSENT_FUNCTION: &str = "absent";
106const SPECIAL_HISTOGRAM_QUANTILE: &str = "histogram_quantile";
108const SPECIAL_VECTOR_FUNCTION: &str = "vector";
110const LE_COLUMN_NAME: &str = "le";
112
113static LABEL_NAME_REGEX: Lazy<Regex> =
116 Lazy::new(|| Regex::new(r"^[a-zA-Z_][a-zA-Z0-9_]*$").unwrap());
117
118const DEFAULT_TIME_INDEX_COLUMN: &str = "time";
119
120const DEFAULT_FIELD_COLUMN: &str = "value";
122
123const FIELD_COLUMN_MATCHER: &str = "__field__";
125
126const SCHEMA_COLUMN_MATCHER: &str = "__schema__";
128const DB_COLUMN_MATCHER: &str = "__database__";
129
130const BINARY_ISLAND_LEAF_ALIAS_PREFIX: &str = "__prom_v";
132
133const MAX_SCATTER_POINTS: i64 = 400;
135
136const INTERVAL_1H: i64 = 60 * 60 * 1000;
138
139#[derive(Default, Debug, Clone)]
140struct PromPlannerContext {
141 start: Millisecond,
143 end: Millisecond,
144 interval: Millisecond,
145 lookback_delta: Millisecond,
146
147 table_name: Option<String>,
149 time_index_column: Option<String>,
150 field_columns: Vec<String>,
151 tag_columns: Vec<String>,
152 use_tsid: bool,
158 field_column_matcher: Option<Vec<Matcher>>,
160 selector_matcher: Vec<Matcher>,
162 schema_name: Option<String>,
163 range: Option<Millisecond>,
165}
166
167#[derive(Debug, Clone, PartialEq, Eq, Hash)]
168struct VectorLeafKey {
169 metric_name: String,
170 matchers: Vec<(String, String, String)>,
171 or_matchers: Vec<Vec<(String, String, String)>>,
172 offset_ms: i128,
173 at: String,
174}
175
176#[derive(Debug, Clone)]
177struct IslandLeaf {
178 selector: VectorSelector,
179 display_table: String,
180}
181
182#[derive(Debug, Clone)]
183enum IslandExpr {
184 VectorLeaf(usize),
185 Scalar(DfExpr),
186 Unary {
187 input: Box<IslandExpr>,
188 },
189 Binary {
190 op: TokenType,
191 lhs: Box<IslandExpr>,
192 rhs: Box<IslandExpr>,
193 },
194}
195
196impl IslandExpr {
197 fn try_new(expr: &PromExpr, env: &mut IslandCollectEnv) -> Option<Self> {
198 if let Some(expr) = PromPlanner::try_build_literal_expr(expr) {
199 return Some(Self::Scalar(expr));
200 }
201
202 match expr {
203 PromExpr::Paren(ParenExpr { expr }) => Self::try_new(expr, env),
204 PromExpr::VectorSelector(selector) => {
205 let leaf = env.intern_leaf(selector)?;
206 Some(Self::VectorLeaf(leaf))
207 }
208 PromExpr::Unary(UnaryExpr { expr }) => {
209 let input = Self::try_new(expr, env)?;
210 Some(Self::Unary {
211 input: Box::new(input),
212 })
213 }
214 PromExpr::Binary(PromBinaryExpr {
215 lhs,
216 rhs,
217 op,
218 modifier,
219 }) if matches!(
220 op.id(),
221 token::T_ADD
222 | token::T_SUB
223 | token::T_MUL
224 | token::T_DIV
225 | token::T_MOD
226 | token::T_POW
227 | token::T_ATAN2
228 ) && modifier.as_ref().is_none_or(|modifier| {
229 !modifier.return_bool
230 && modifier.matching.is_none()
231 && matches!(modifier.card, VectorMatchCardinality::OneToOne)
232 }) =>
233 {
234 let lhs = Self::try_new(lhs, env)?;
235 let rhs = Self::try_new(rhs, env)?;
236 Some(Self::Binary {
237 op: *op,
238 lhs: Box::new(lhs),
239 rhs: Box::new(rhs),
240 })
241 }
242 _ => None,
243 }
244 }
245}
246
247#[derive(Debug, Default)]
248struct IslandCollectEnv {
249 leaf_by_key: HashMap<VectorLeafKey, usize>,
250 leaves: Vec<IslandLeaf>,
251 vector_occurrences: usize,
252}
253
254#[derive(Debug)]
255struct PlannedIslandLeaf {
256 plan: LogicalPlan,
257 ctx: PromPlannerContext,
258 alias: TableReference,
259 display_table: String,
260}
261
262#[derive(Debug)]
263struct IslandFieldExprs {
264 exprs: Vec<DfExpr>,
265 names: Vec<String>,
266 scalar: bool,
267}
268
269impl VectorLeafKey {
270 fn from_selector(selector: &VectorSelector) -> Option<Self> {
271 let mut metric_name = selector.name.clone();
272 let mut matchers = Vec::with_capacity(selector.matchers.matchers.len());
273 let matcher_key = |matcher: &Matcher| {
274 (
275 matcher.name.clone(),
276 matcher.op.to_string(),
277 matcher.value.clone(),
278 )
279 };
280
281 for matcher in &selector.matchers.matchers {
282 if matcher.name == METRIC_NAME {
283 if matcher.op != MatchOp::Equal || metric_name.is_some() {
284 return None;
285 }
286 metric_name = Some(matcher.value.clone());
287 } else {
288 matchers.push(matcher_key(matcher));
289 }
290 }
291 matchers.sort();
292
293 let mut or_matchers = selector
294 .matchers
295 .or_matchers
296 .iter()
297 .map(|group| {
298 let mut group = group.iter().map(matcher_key).collect::<Vec<_>>();
299 group.sort();
300 group
301 })
302 .collect::<Vec<_>>();
303 or_matchers.sort();
304
305 Some(Self {
306 metric_name: metric_name?,
307 matchers,
308 or_matchers,
309 offset_ms: match &selector.offset {
310 Some(Offset::Pos(duration)) => duration.as_millis() as i128,
311 Some(Offset::Neg(duration)) => -(duration.as_millis() as i128),
312 None => 0,
313 },
314 at: format!("{:?}", selector.at),
315 })
316 }
317}
318
319impl IslandCollectEnv {
320 fn intern_leaf(&mut self, selector: &VectorSelector) -> Option<usize> {
321 self.vector_occurrences += 1;
322 let key = VectorLeafKey::from_selector(selector)?;
323 if let Some(id) = self.leaf_by_key.get(&key) {
324 return Some(*id);
325 }
326
327 let id = self.leaves.len();
328 self.leaves.push(IslandLeaf {
329 selector: selector.clone(),
330 display_table: key.metric_name.clone(),
331 });
332 self.leaf_by_key.insert(key, id);
333 Some(id)
334 }
335}
336
337impl PromPlannerContext {
338 fn from_eval_stmt(stmt: &EvalStmt) -> Self {
339 Self {
340 start: stmt.start.duration_since(UNIX_EPOCH).unwrap().as_millis() as _,
341 end: stmt.end.duration_since(UNIX_EPOCH).unwrap().as_millis() as _,
342 interval: stmt.interval.as_millis() as _,
343 lookback_delta: stmt.lookback_delta.as_millis() as _,
344 ..Default::default()
345 }
346 }
347
348 fn reset(&mut self) {
350 self.table_name = None;
351 self.time_index_column = None;
352 self.field_columns = vec![];
353 self.tag_columns = vec![];
354 self.use_tsid = false;
355 self.field_column_matcher = None;
356 self.selector_matcher.clear();
357 self.schema_name = None;
358 self.range = None;
359 }
360
361 fn reset_table_name_and_schema(&mut self) {
363 self.table_name = Some(String::new());
364 self.schema_name = None;
365 self.use_tsid = false;
366 }
367
368 fn has_le_tag(&self) -> bool {
370 self.tag_columns.iter().any(|c| c.eq(&LE_COLUMN_NAME))
371 }
372}
373
374pub struct PromPlanner {
375 table_provider: DfTableSourceProvider,
376 ctx: PromPlannerContext,
377}
378
379impl PromPlanner {
380 pub async fn stmt_to_plan(
381 table_provider: DfTableSourceProvider,
382 stmt: &EvalStmt,
383 query_engine_state: &QueryEngineState,
384 ) -> Result<LogicalPlan> {
385 let mut planner = Self {
386 table_provider,
387 ctx: PromPlannerContext::from_eval_stmt(stmt),
388 };
389
390 let plan = planner
391 .prom_expr_to_plan(&stmt.expr, query_engine_state)
392 .await?;
393
394 planner.strip_tsid_column(plan)
396 }
397
398 pub async fn prom_expr_to_plan(
399 &mut self,
400 prom_expr: &PromExpr,
401 query_engine_state: &QueryEngineState,
402 ) -> Result<LogicalPlan> {
403 self.prom_expr_to_plan_inner(prom_expr, false, query_engine_state)
404 .await
405 }
406
407 #[async_recursion]
417 async fn prom_expr_to_plan_inner(
418 &mut self,
419 prom_expr: &PromExpr,
420 timestamp_fn: bool,
421 query_engine_state: &QueryEngineState,
422 ) -> Result<LogicalPlan> {
423 let res = match prom_expr {
424 PromExpr::Aggregate(expr) => {
425 self.prom_aggr_expr_to_plan(query_engine_state, expr)
426 .await?
427 }
428 PromExpr::Unary(expr) => {
429 self.prom_unary_expr_to_plan(query_engine_state, expr)
430 .await?
431 }
432 PromExpr::Binary(expr) => {
433 self.prom_binary_expr_to_plan(query_engine_state, expr)
434 .await?
435 }
436 PromExpr::Paren(ParenExpr { expr }) => {
437 self.prom_expr_to_plan_inner(expr, timestamp_fn, query_engine_state)
438 .await?
439 }
440 PromExpr::Subquery(expr) => {
441 self.prom_subquery_expr_to_plan(query_engine_state, expr)
442 .await?
443 }
444 PromExpr::NumberLiteral(lit) => self.prom_number_lit_to_plan(lit)?,
445 PromExpr::StringLiteral(lit) => self.prom_string_lit_to_plan(lit)?,
446 PromExpr::VectorSelector(selector) => {
447 self.prom_vector_selector_to_plan(selector, timestamp_fn)
448 .await?
449 }
450 PromExpr::MatrixSelector(selector) => {
451 self.prom_matrix_selector_to_plan(selector).await?
452 }
453 PromExpr::Call(expr) => {
454 self.prom_call_expr_to_plan(query_engine_state, expr)
455 .await?
456 }
457 PromExpr::Extension(expr) => {
458 self.prom_ext_expr_to_plan(query_engine_state, expr).await?
459 }
460 };
461
462 Ok(res)
463 }
464
465 async fn prom_subquery_expr_to_plan(
466 &mut self,
467 query_engine_state: &QueryEngineState,
468 subquery_expr: &SubqueryExpr,
469 ) -> Result<LogicalPlan> {
470 let SubqueryExpr {
471 expr, range, step, ..
472 } = subquery_expr;
473
474 let current_interval = self.ctx.interval;
475 if let Some(step) = step {
476 self.ctx.interval = step.as_millis() as _;
477 }
478 let current_start = self.ctx.start;
479 self.ctx.start -= range.as_millis() as i64 - self.ctx.interval;
480 let input = self.prom_expr_to_plan(expr, query_engine_state).await?;
481 self.ctx.interval = current_interval;
482 self.ctx.start = current_start;
483
484 ensure!(!range.is_zero(), ZeroRangeSelectorSnafu);
485 let range_ms = range.as_millis() as _;
486 self.ctx.range = Some(range_ms);
487
488 let time_index_column =
489 self.ctx
490 .time_index_column
491 .clone()
492 .with_context(|| TimeIndexNotFoundSnafu {
493 table: self.ctx.table_name.clone().unwrap_or_default(),
494 })?;
495
496 let input_schema = input.schema();
502 let input_has_tsid = input_schema.fields().iter().any(|field| {
503 field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
504 && field.data_type() == &ArrowDataType::UInt64
505 });
506 let (series_key_columns, mut sort_exprs) = if input_has_tsid {
507 (
508 vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()],
509 vec![
510 DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME))
511 .sort(true, true),
512 ],
513 )
514 } else {
515 let key_columns: Vec<String> = self
518 .ctx
519 .tag_columns
520 .iter()
521 .filter(|name| input_schema.has_column_with_unqualified_name(name))
522 .cloned()
523 .collect();
524 let sort = key_columns
525 .iter()
526 .map(|name| DfExpr::Column(Column::from_name(name)).sort(true, true))
527 .collect::<Vec<_>>();
528 (key_columns, sort)
529 };
530 sort_exprs.push(DfExpr::Column(Column::from_name(&time_index_column)).sort(true, true));
531
532 let sort_plan = LogicalPlanBuilder::from(input)
533 .sort(sort_exprs)
534 .context(DataFusionPlanningSnafu)?
535 .build()
536 .context(DataFusionPlanningSnafu)?;
537 let divide_plan = LogicalPlan::Extension(Extension {
538 node: Arc::new(SeriesDivide::new(
539 series_key_columns,
540 time_index_column.clone(),
541 sort_plan,
542 )),
543 });
544
545 let manipulate = RangeManipulate::new(
546 self.ctx.start,
547 self.ctx.end,
548 self.ctx.interval,
549 range_ms,
550 time_index_column,
551 self.ctx.field_columns.clone(),
552 divide_plan,
553 )
554 .context(DataFusionPlanningSnafu)?;
555
556 Ok(LogicalPlan::Extension(Extension {
557 node: Arc::new(manipulate),
558 }))
559 }
560
561 async fn prom_aggr_expr_to_plan(
562 &mut self,
563 query_engine_state: &QueryEngineState,
564 aggr_expr: &AggregateExpr,
565 ) -> Result<LogicalPlan> {
566 let AggregateExpr {
567 op,
568 expr,
569 modifier,
570 param,
571 } = aggr_expr;
572
573 let mut input = self.prom_expr_to_plan(expr, query_engine_state).await?;
574 let input_has_tsid = input.schema().fields().iter().any(|field| {
575 field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
576 && field.data_type() == &ArrowDataType::UInt64
577 });
578
579 let required_group_tags = match modifier {
582 None => BTreeSet::new(),
583 Some(LabelModifier::Include(labels)) => labels
584 .labels
585 .iter()
586 .filter(|label| !is_metric_engine_internal_column(label.as_str()))
587 .cloned()
588 .collect(),
589 Some(LabelModifier::Exclude(labels)) => {
590 let mut all_tags = self.collect_row_key_tag_columns_from_plan(&input)?;
591 for label in &labels.labels {
592 let _ = all_tags.remove(label);
593 }
594 all_tags
595 }
596 };
597
598 if !required_group_tags.is_empty()
599 && required_group_tags
600 .iter()
601 .any(|tag| Self::find_case_sensitive_column(input.schema(), tag.as_str()).is_none())
602 {
603 input = self.ensure_tag_columns_available(input, &required_group_tags)?;
604 self.refresh_tag_columns_from_schema(input.schema());
605 }
606
607 match (*op).id() {
608 token::T_TOPK | token::T_BOTTOMK => {
609 self.prom_topk_bottomk_to_plan(aggr_expr, input).await
610 }
611 _ => {
612 let input_tag_columns = if input_has_tsid {
616 self.collect_row_key_tag_columns_from_plan(&input)?
617 .into_iter()
618 .collect::<Vec<_>>()
619 } else {
620 self.ctx.tag_columns.clone()
621 };
622 let mut group_exprs = self.agg_modifier_to_col(input.schema(), modifier, true)?;
625 let (mut aggr_exprs, prev_field_exprs) =
627 self.create_aggregate_exprs(*op, param, &input)?;
628
629 let keep_tsid = op.id() != token::T_COUNT_VALUES
630 && input_has_tsid
631 && input_tag_columns.iter().collect::<HashSet<_>>()
632 == self.ctx.tag_columns.iter().collect::<HashSet<_>>();
633
634 if keep_tsid {
635 aggr_exprs.push(
636 first_value(
637 DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME)),
638 vec![],
639 )
640 .alias(DATA_SCHEMA_TSID_COLUMN_NAME),
641 );
642 }
643 self.ctx.use_tsid = keep_tsid;
644
645 let builder = LogicalPlanBuilder::from(input);
647 let builder = if op.id() == token::T_COUNT_VALUES {
648 let label = Self::get_param_value_as_str(*op, param)?;
649 group_exprs.extend(prev_field_exprs.clone());
652 let project_fields = self
653 .create_field_column_exprs()?
654 .into_iter()
655 .chain(self.create_tag_column_exprs()?)
656 .chain(Some(self.create_time_index_column_expr()?))
657 .chain(prev_field_exprs.into_iter().map(|expr| expr.alias(label)));
658
659 builder
660 .aggregate(group_exprs.clone(), aggr_exprs)
661 .context(DataFusionPlanningSnafu)?
662 .project(project_fields)
663 .context(DataFusionPlanningSnafu)?
664 } else {
665 builder
666 .aggregate(group_exprs.clone(), aggr_exprs)
667 .context(DataFusionPlanningSnafu)?
668 };
669
670 let sort_expr = group_exprs.into_iter().map(|expr| expr.sort(true, false));
671
672 builder
673 .sort(sort_expr)
674 .context(DataFusionPlanningSnafu)?
675 .build()
676 .context(DataFusionPlanningSnafu)
677 }
678 }
679 }
680
681 async fn prom_topk_bottomk_to_plan(
683 &mut self,
684 aggr_expr: &AggregateExpr,
685 input: LogicalPlan,
686 ) -> Result<LogicalPlan> {
687 let AggregateExpr {
688 op,
689 param,
690 modifier,
691 ..
692 } = aggr_expr;
693
694 let input_has_tsid = input.schema().fields().iter().any(|field| {
695 field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
696 && field.data_type() == &ArrowDataType::UInt64
697 });
698 self.ctx.use_tsid = input_has_tsid;
699
700 let group_exprs = self.agg_modifier_to_col(input.schema(), modifier, false)?;
701
702 let val = Self::get_param_as_literal_expr(param, Some(*op), Some(ArrowDataType::Float64))?;
703
704 let window_exprs = self.create_window_exprs(*op, group_exprs.clone(), &input)?;
706
707 let rank_columns: Vec<_> = window_exprs
708 .iter()
709 .map(|expr| expr.schema_name().to_string())
710 .collect();
711
712 let filter: DfExpr = rank_columns
715 .iter()
716 .fold(None, |expr, rank| {
717 let predicate = DfExpr::BinaryExpr(BinaryExpr {
718 left: Box::new(col(rank)),
719 op: Operator::LtEq,
720 right: Box::new(val.clone()),
721 });
722
723 match expr {
724 None => Some(predicate),
725 Some(expr) => Some(DfExpr::BinaryExpr(BinaryExpr {
726 left: Box::new(expr),
727 op: Operator::Or,
728 right: Box::new(predicate),
729 })),
730 }
731 })
732 .unwrap();
733
734 let rank_columns: Vec<_> = rank_columns.into_iter().map(col).collect();
735
736 let mut new_group_exprs = group_exprs.clone();
737 new_group_exprs.extend(rank_columns);
739
740 let group_sort_expr = new_group_exprs
741 .into_iter()
742 .map(|expr| expr.sort(true, false));
743
744 let project_fields = self
745 .create_field_column_exprs()?
746 .into_iter()
747 .chain(self.create_tag_column_exprs()?)
748 .chain(
749 self.ctx
750 .use_tsid
751 .then_some(DfExpr::Column(Column::from_name(
752 DATA_SCHEMA_TSID_COLUMN_NAME,
753 ))),
754 )
755 .chain(Some(self.create_time_index_column_expr()?));
756
757 LogicalPlanBuilder::from(input)
758 .window(window_exprs)
759 .context(DataFusionPlanningSnafu)?
760 .filter(filter)
761 .context(DataFusionPlanningSnafu)?
762 .sort(group_sort_expr)
763 .context(DataFusionPlanningSnafu)?
764 .project(project_fields)
765 .context(DataFusionPlanningSnafu)?
766 .build()
767 .context(DataFusionPlanningSnafu)
768 }
769
770 async fn prom_unary_expr_to_plan(
771 &mut self,
772 query_engine_state: &QueryEngineState,
773 unary_expr: &UnaryExpr,
774 ) -> Result<LogicalPlan> {
775 let UnaryExpr { expr } = unary_expr;
776 let input = self.prom_expr_to_plan(expr, query_engine_state).await?;
778 self.projection_for_each_field_column(input, |col| {
779 Ok(DfExpr::Negative(Box::new(DfExpr::Column(col.into()))))
780 })
781 }
782
783 async fn try_plan_binary_island(
784 &mut self,
785 binary_expr: &PromBinaryExpr,
786 ) -> Result<Option<LogicalPlan>> {
787 let original_ctx = self.ctx.clone();
788 let mut collect_env = IslandCollectEnv::default();
789 let Some(island_expr) =
790 IslandExpr::try_new(&PromExpr::Binary(binary_expr.clone()), &mut collect_env)
791 else {
792 return Ok(None);
793 };
794
795 if collect_env.leaves.is_empty()
796 || collect_env.vector_occurrences <= collect_env.leaves.len()
797 {
798 return Ok(None);
799 }
800
801 let mut planned_leaves = Vec::with_capacity(collect_env.leaves.len());
802 for (idx, leaf) in collect_env.leaves.iter().enumerate() {
803 let plan = self
804 .prom_vector_selector_to_plan(&leaf.selector, false)
805 .await?;
806 let ctx = self.ctx.clone();
807 let alias = TableReference::bare(format!("{BINARY_ISLAND_LEAF_ALIAS_PREFIX}{idx}"));
808 let plan = LogicalPlanBuilder::from(plan)
809 .alias(alias.clone())
810 .context(DataFusionPlanningSnafu)?
811 .build()
812 .context(DataFusionPlanningSnafu)?;
813 planned_leaves.push(PlannedIslandLeaf {
814 plan,
815 ctx,
816 alias,
817 display_table: leaf.display_table.clone(),
818 });
819 }
820
821 if !Self::binary_island_join_contexts_supported(&planned_leaves) {
822 self.ctx = original_ctx;
823 return Ok(None);
824 }
825
826 let mut input = planned_leaves[0].plan.clone();
827 for right_idx in 1..planned_leaves.len() {
828 input = self.join_binary_island_leaf(
829 input,
830 &planned_leaves[0],
831 &planned_leaves[right_idx],
832 )?;
833 }
834
835 let field_exprs =
836 Self::build_binary_island_field_exprs(&island_expr, &planned_leaves, input.schema())?;
837 if field_exprs.scalar || field_exprs.exprs.is_empty() {
838 self.ctx = original_ctx;
839 return Ok(None);
840 }
841
842 let plan = self.project_binary_island(
843 input,
844 &planned_leaves[0].alias,
845 &planned_leaves[0].ctx,
846 field_exprs,
847 )?;
848 Ok(Some(plan))
849 }
850
851 fn binary_island_join_contexts_supported(leaves: &[PlannedIslandLeaf]) -> bool {
852 if leaves
853 .iter()
854 .any(|leaf| leaf.ctx.time_index_column.is_none())
855 {
856 return false;
857 }
858
859 if leaves.len() <= 1 {
860 return true;
861 }
862
863 let first_tags = leaves[0].ctx.tag_columns.iter().collect::<BTreeSet<_>>();
864
865 leaves.iter().skip(1).all(|leaf| {
866 (Self::plan_has_tsid_column(&leaves[0].plan) && Self::plan_has_tsid_column(&leaf.plan))
867 || leaf.ctx.tag_columns.iter().collect::<BTreeSet<_>>() == first_tags
868 })
869 }
870
871 fn join_binary_island_leaf(
872 &self,
873 left: LogicalPlan,
874 first_leaf: &PlannedIslandLeaf,
875 right_leaf: &PlannedIslandLeaf,
876 ) -> Result<LogicalPlan> {
877 let only_join_time_index =
878 first_leaf.ctx.tag_columns.is_empty() || right_leaf.ctx.tag_columns.is_empty();
879 let (mut left_keys, mut right_keys, force_empty_join) = self.binary_join_key_columns(
880 left.schema(),
881 right_leaf.plan.schema(),
882 &first_leaf.ctx,
883 &right_leaf.ctx,
884 only_join_time_index,
885 &None,
886 )?;
887
888 if let (Some(left_time_index_column), Some(right_time_index_column)) = (
889 first_leaf.ctx.time_index_column.clone(),
890 right_leaf.ctx.time_index_column.clone(),
891 ) {
892 left_keys.insert(left_time_index_column);
893 right_keys.insert(right_time_index_column);
894 }
895
896 LogicalPlanBuilder::from(left)
897 .join_detailed(
898 right_leaf.plan.clone(),
899 JoinType::Inner,
900 (
901 left_keys
902 .into_iter()
903 .map(|name| Column::new(Some(first_leaf.alias.clone()), name))
904 .collect::<Vec<_>>(),
905 right_keys
906 .into_iter()
907 .map(|name| Column::new(Some(right_leaf.alias.clone()), name))
908 .collect::<Vec<_>>(),
909 ),
910 force_empty_join.then_some(lit(false)),
911 NullEquality::NullEqualsNull,
912 )
913 .context(DataFusionPlanningSnafu)?
914 .build()
915 .context(DataFusionPlanningSnafu)
916 }
917
918 fn build_binary_island_field_exprs(
919 expr: &IslandExpr,
920 leaves: &[PlannedIslandLeaf],
921 schema: &DFSchemaRef,
922 ) -> Result<IslandFieldExprs> {
923 match expr {
924 IslandExpr::VectorLeaf(id) => {
925 let leaf = &leaves[*id];
926 let exprs = leaf
927 .ctx
928 .field_columns
929 .iter()
930 .map(|field| {
931 schema
932 .qualified_field_with_name(Some(&leaf.alias), field)
933 .context(DataFusionPlanningSnafu)
934 .map(|field| DfExpr::Column(field.into()))
935 })
936 .collect::<Result<Vec<_>>>()?;
937 let names = leaf
938 .ctx
939 .field_columns
940 .iter()
941 .map(|field| format!("{}.{}", leaf.display_table, field))
942 .collect();
943 Ok(IslandFieldExprs {
944 exprs,
945 names,
946 scalar: false,
947 })
948 }
949 IslandExpr::Scalar(expr) => Ok(IslandFieldExprs {
950 exprs: vec![expr.clone()],
951 names: vec![expr.schema_name().to_string()],
952 scalar: true,
953 }),
954 IslandExpr::Unary { input } => {
955 let input = Self::build_binary_island_field_exprs(input, leaves, schema)?;
956 let mut exprs = Vec::with_capacity(input.exprs.len());
957 let mut names = Vec::with_capacity(input.names.len());
958 for (expr, name) in input.exprs.into_iter().zip(input.names) {
959 exprs.push(DfExpr::Negative(Box::new(expr)));
960 names.push(format!("-{name}"));
961 }
962 Ok(IslandFieldExprs {
963 exprs,
964 names,
965 scalar: input.scalar,
966 })
967 }
968 IslandExpr::Binary { op, lhs, rhs } => {
969 let same_leaf = match (&**lhs, &**rhs) {
970 (IslandExpr::VectorLeaf(left), IslandExpr::VectorLeaf(right))
971 if left == right =>
972 {
973 Some(*left)
974 }
975 _ => None,
976 };
977 let lhs = Self::build_binary_island_field_exprs(lhs, leaves, schema)?;
978 let rhs = Self::build_binary_island_field_exprs(rhs, leaves, schema)?;
979 let expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
980 let scalar = lhs.scalar && rhs.scalar;
981 let op = op.to_string();
982
983 let (exprs, names) = match (lhs.scalar, rhs.scalar) {
984 (true, true) => {
985 let expr = expr_builder(lhs.exprs[0].clone(), rhs.exprs[0].clone())?;
986 let name = format!("{} {op} {}", lhs.names[0], rhs.names[0]);
987 (vec![expr], vec![name])
988 }
989 (true, false) => {
990 let mut exprs = Vec::with_capacity(rhs.exprs.len());
991 let mut names = Vec::with_capacity(rhs.names.len());
992 for (rhs_expr, rhs_name) in rhs.exprs.into_iter().zip(rhs.names) {
993 exprs.push(expr_builder(lhs.exprs[0].clone(), rhs_expr)?);
994 names.push(format!("{} {op} {rhs_name}", lhs.names[0]));
995 }
996 (exprs, names)
997 }
998 (false, true) => {
999 let mut exprs = Vec::with_capacity(lhs.exprs.len());
1000 let mut names = Vec::with_capacity(lhs.names.len());
1001 for (lhs_expr, lhs_name) in lhs.exprs.into_iter().zip(lhs.names) {
1002 exprs.push(expr_builder(lhs_expr, rhs.exprs[0].clone())?);
1003 names.push(format!("{lhs_name} {op} {}", rhs.names[0]));
1004 }
1005 (exprs, names)
1006 }
1007 (false, false) => {
1008 let mut exprs = Vec::new();
1009 let mut names = Vec::new();
1010 for (idx, ((lhs_expr, rhs_expr), (mut lhs_name, mut rhs_name))) in lhs
1011 .exprs
1012 .into_iter()
1013 .zip(rhs.exprs)
1014 .zip(lhs.names.into_iter().zip(rhs.names))
1015 .enumerate()
1016 {
1017 if let Some(leaf) = same_leaf {
1018 let field = leaves[leaf]
1019 .ctx
1020 .field_columns
1021 .get(idx)
1022 .cloned()
1023 .unwrap_or_else(|| lhs_name.clone());
1024 lhs_name = format!("lhs.{field}");
1025 rhs_name = format!("rhs.{field}");
1026 }
1027 exprs.push(expr_builder(lhs_expr, rhs_expr)?);
1028 names.push(format!("{lhs_name} {op} {rhs_name}"));
1029 }
1030 (exprs, names)
1031 }
1032 };
1033
1034 Ok(IslandFieldExprs {
1035 exprs,
1036 names,
1037 scalar,
1038 })
1039 }
1040 }
1041 }
1042
1043 fn project_binary_island(
1044 &mut self,
1045 input: LogicalPlan,
1046 base_alias: &TableReference,
1047 base_ctx: &PromPlannerContext,
1048 field_exprs: IslandFieldExprs,
1049 ) -> Result<LogicalPlan> {
1050 self.ctx = base_ctx.clone();
1051
1052 let schema = input.schema();
1053 let non_field_exprs = base_ctx
1054 .tag_columns
1055 .iter()
1056 .chain(base_ctx.time_index_column.iter())
1057 .map(|column| {
1058 schema
1059 .qualified_field_with_name(Some(base_alias), column)
1060 .context(DataFusionPlanningSnafu)
1061 .map(|field| DfExpr::Column(field.into()))
1062 });
1063 let tsid_expr = Self::optional_tsid_projection(schema, Some(base_alias), base_ctx.use_tsid)
1064 .into_iter()
1065 .map(Ok);
1066
1067 self.ctx.field_columns = field_exprs.names;
1068 let field_exprs = field_exprs
1069 .exprs
1070 .into_iter()
1071 .zip(self.ctx.field_columns.iter())
1072 .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::<String>, name))));
1073
1074 let project_exprs = non_field_exprs
1075 .chain(tsid_expr)
1076 .chain(field_exprs)
1077 .collect::<Result<Vec<_>>>()?;
1078
1079 let plan = LogicalPlanBuilder::from(input)
1080 .project(project_exprs)
1081 .context(DataFusionPlanningSnafu)?
1082 .build()
1083 .context(DataFusionPlanningSnafu)?;
1084
1085 self.ctx.table_name = None;
1086 self.ctx.schema_name = None;
1087
1088 Ok(plan)
1089 }
1090
1091 async fn prom_binary_expr_to_plan(
1092 &mut self,
1093 query_engine_state: &QueryEngineState,
1094 binary_expr: &PromBinaryExpr,
1095 ) -> Result<LogicalPlan> {
1096 if let Some(plan) = self.try_plan_binary_island(binary_expr).await? {
1097 return Ok(plan);
1098 }
1099
1100 let PromBinaryExpr {
1101 lhs,
1102 rhs,
1103 op,
1104 modifier,
1105 } = binary_expr;
1106
1107 let should_return_bool = if let Some(m) = modifier {
1110 m.return_bool
1111 } else {
1112 false
1113 };
1114 let is_comparison_op = Self::is_token_a_comparison_op(*op);
1115
1116 match (
1119 Self::try_build_literal_expr(lhs),
1120 Self::try_build_literal_expr(rhs),
1121 ) {
1122 (Some(lhs), Some(rhs)) => {
1123 self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
1124 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1125 self.ctx.reset_table_name_and_schema();
1126 let field_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1127 let mut field_expr = field_expr_builder(lhs, rhs)?;
1128
1129 if is_comparison_op && should_return_bool {
1130 field_expr = DfExpr::Cast(Cast {
1131 expr: Box::new(field_expr),
1132 data_type: ArrowDataType::Float64,
1133 });
1134 }
1135
1136 Ok(LogicalPlan::Extension(Extension {
1137 node: Arc::new(
1138 EmptyMetric::new(
1139 self.ctx.start,
1140 self.ctx.end,
1141 self.ctx.interval,
1142 SPECIAL_TIME_FUNCTION.to_string(),
1143 DEFAULT_FIELD_COLUMN.to_string(),
1144 Some(field_expr),
1145 )
1146 .context(DataFusionPlanningSnafu)?,
1147 ),
1148 }))
1149 }
1150 (Some(mut expr), None) => {
1152 let input = self.prom_expr_to_plan(rhs, query_engine_state).await?;
1153 if let Some(time_expr) = self.try_build_special_time_expr_with_context(lhs) {
1155 expr = time_expr
1156 }
1157 let bin_expr_builder = |col: &String| {
1158 let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1159 let mut binary_expr =
1160 binary_expr_builder(expr.clone(), DfExpr::Column(col.into()))?;
1161
1162 if is_comparison_op && should_return_bool {
1163 binary_expr = DfExpr::Cast(Cast {
1164 expr: Box::new(binary_expr),
1165 data_type: ArrowDataType::Float64,
1166 });
1167 }
1168 Ok(binary_expr)
1169 };
1170 if is_comparison_op && !should_return_bool {
1171 self.filter_on_field_column(input, bin_expr_builder)
1172 } else {
1173 self.projection_for_each_field_column(input, bin_expr_builder)
1174 }
1175 }
1176 (None, Some(mut expr)) => {
1178 let input = self.prom_expr_to_plan(lhs, query_engine_state).await?;
1179 if let Some(time_expr) = self.try_build_special_time_expr_with_context(rhs) {
1181 expr = time_expr
1182 }
1183 let bin_expr_builder = |col: &String| {
1184 let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1185 let mut binary_expr =
1186 binary_expr_builder(DfExpr::Column(col.into()), expr.clone())?;
1187
1188 if is_comparison_op && should_return_bool {
1189 binary_expr = DfExpr::Cast(Cast {
1190 expr: Box::new(binary_expr),
1191 data_type: ArrowDataType::Float64,
1192 });
1193 }
1194 Ok(binary_expr)
1195 };
1196 if is_comparison_op && !should_return_bool {
1197 self.filter_on_field_column(input, bin_expr_builder)
1198 } else {
1199 self.projection_for_each_field_column(input, bin_expr_builder)
1200 }
1201 }
1202 (None, None) => {
1204 let left_input = self.prom_expr_to_plan(lhs, query_engine_state).await?;
1205 let left_field_columns = self.ctx.field_columns.clone();
1206 let left_time_index_column = self.ctx.time_index_column.clone();
1207 let mut left_table_ref = self
1208 .table_ref()
1209 .unwrap_or_else(|_| TableReference::bare(""));
1210 let left_context = self.ctx.clone();
1211
1212 let right_input = self.prom_expr_to_plan(rhs, query_engine_state).await?;
1213 let right_field_columns = self.ctx.field_columns.clone();
1214 let right_time_index_column = self.ctx.time_index_column.clone();
1215 let mut right_table_ref = self
1216 .table_ref()
1217 .unwrap_or_else(|_| TableReference::bare(""));
1218 let right_context = self.ctx.clone();
1219
1220 if Self::is_token_a_set_op(*op) {
1224 return self.set_op_on_non_field_columns(
1225 left_input,
1226 right_input,
1227 left_context,
1228 right_context,
1229 *op,
1230 modifier,
1231 );
1232 }
1233
1234 if left_table_ref == right_table_ref {
1236 left_table_ref = TableReference::bare("lhs");
1238 right_table_ref = TableReference::bare("rhs");
1239 if self.ctx.tag_columns.is_empty() {
1245 self.ctx = left_context.clone();
1246 self.ctx.table_name = Some("lhs".to_string());
1247 } else {
1248 self.ctx.table_name = Some("rhs".to_string());
1249 }
1250 }
1251 let (output_field_columns, field_columns) =
1252 Self::align_binary_field_columns(&left_field_columns, &right_field_columns);
1253 let left_aligned_field_columns = field_columns
1254 .iter()
1255 .map(|(left_col_name, _)| (*left_col_name).clone())
1256 .collect::<Vec<_>>();
1257 let right_aligned_field_columns = field_columns
1258 .iter()
1259 .map(|(_, right_col_name)| (*right_col_name).clone())
1260 .collect::<Vec<_>>();
1261 self.ctx.field_columns = output_field_columns;
1265 let mut field_columns = field_columns.into_iter();
1266
1267 let join_plan = self.join_on_non_field_columns(
1268 left_input,
1269 right_input,
1270 left_table_ref.clone(),
1271 right_table_ref.clone(),
1272 left_time_index_column,
1273 right_time_index_column,
1274 left_context.tag_columns.is_empty() || right_context.tag_columns.is_empty(),
1277 modifier,
1278 &left_context,
1279 &right_context,
1280 )?;
1281 let join_plan_schema = join_plan.schema().clone();
1282
1283 let bin_expr_builder = |_: &String| {
1284 let (left_col_name, right_col_name) = field_columns.next().unwrap();
1285 let left_col = join_plan_schema
1286 .qualified_field_with_name(Some(&left_table_ref), left_col_name)
1287 .context(DataFusionPlanningSnafu)?
1288 .into();
1289 let right_col = join_plan_schema
1290 .qualified_field_with_name(Some(&right_table_ref), right_col_name)
1291 .context(DataFusionPlanningSnafu)?
1292 .into();
1293
1294 let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1295 let mut binary_expr =
1296 binary_expr_builder(DfExpr::Column(left_col), DfExpr::Column(right_col))?;
1297 if is_comparison_op && should_return_bool {
1298 binary_expr = DfExpr::Cast(Cast {
1299 expr: Box::new(binary_expr),
1300 data_type: ArrowDataType::Float64,
1301 });
1302 }
1303 Ok(binary_expr)
1304 };
1305 if is_comparison_op && !should_return_bool {
1306 let filtered = self.filter_on_field_column(join_plan, bin_expr_builder)?;
1313 let (project_table_ref, mut project_context, project_field_columns) =
1314 match (lhs.value_type(), rhs.value_type()) {
1315 (ValueType::Scalar, ValueType::Vector) => (
1316 &right_table_ref,
1317 right_context.clone(),
1318 right_aligned_field_columns,
1319 ),
1320 _ => (
1321 &left_table_ref,
1322 left_context.clone(),
1323 left_aligned_field_columns,
1324 ),
1325 };
1326 project_context.field_columns = project_field_columns;
1327 self.project_binary_join_side(filtered, project_table_ref, &project_context)
1328 } else {
1329 self.projection_for_each_field_column(join_plan, bin_expr_builder)
1330 }
1331 }
1332 }
1333 }
1334
1335 fn project_binary_join_side(
1336 &mut self,
1337 input: LogicalPlan,
1338 table_ref: &TableReference,
1339 context: &PromPlannerContext,
1340 ) -> Result<LogicalPlan> {
1341 let schema = input.schema();
1342
1343 let mut project_exprs =
1344 Vec::with_capacity(context.tag_columns.len() + context.field_columns.len() + 2);
1345
1346 if let Some(time_index_column) = &context.time_index_column {
1348 let time_index_col = schema
1349 .qualified_field_with_name(Some(table_ref), time_index_column)
1350 .context(DataFusionPlanningSnafu)?
1351 .into();
1352 project_exprs.push(DfExpr::Column(time_index_col));
1353 }
1354
1355 for field_column in &context.field_columns {
1357 let field_col = schema
1358 .qualified_field_with_name(Some(table_ref), field_column)
1359 .context(DataFusionPlanningSnafu)?
1360 .into();
1361 project_exprs.push(DfExpr::Column(field_col));
1362 }
1363
1364 for tag_column in &context.tag_columns {
1366 let tag_col = schema
1367 .qualified_field_with_name(Some(table_ref), tag_column)
1368 .context(DataFusionPlanningSnafu)?
1369 .into();
1370 project_exprs.push(DfExpr::Column(tag_col));
1371 }
1372
1373 if let Some(tsid_col) =
1376 Self::optional_tsid_projection(schema, Some(table_ref), context.use_tsid)
1377 {
1378 project_exprs.push(tsid_col);
1379 }
1380
1381 let plan = LogicalPlanBuilder::from(input)
1382 .project(project_exprs)
1383 .context(DataFusionPlanningSnafu)?
1384 .build()
1385 .context(DataFusionPlanningSnafu)?;
1386
1387 self.ctx = context.clone();
1390 self.ctx.table_name = None;
1391 self.ctx.schema_name = None;
1392
1393 Ok(plan)
1394 }
1395
1396 fn prom_number_lit_to_plan(&mut self, number_literal: &NumberLiteral) -> Result<LogicalPlan> {
1397 let NumberLiteral { val } = number_literal;
1398 self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
1399 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1400 self.ctx.reset_table_name_and_schema();
1401 let literal_expr = df_prelude::lit(*val);
1402
1403 let plan = LogicalPlan::Extension(Extension {
1404 node: Arc::new(
1405 EmptyMetric::new(
1406 self.ctx.start,
1407 self.ctx.end,
1408 self.ctx.interval,
1409 SPECIAL_TIME_FUNCTION.to_string(),
1410 DEFAULT_FIELD_COLUMN.to_string(),
1411 Some(literal_expr),
1412 )
1413 .context(DataFusionPlanningSnafu)?,
1414 ),
1415 });
1416 Ok(plan)
1417 }
1418
1419 fn prom_string_lit_to_plan(&mut self, string_literal: &StringLiteral) -> Result<LogicalPlan> {
1420 let StringLiteral { val } = string_literal;
1421 self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
1422 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1423 self.ctx.reset_table_name_and_schema();
1424 let literal_expr = df_prelude::lit(val.clone());
1425
1426 let plan = LogicalPlan::Extension(Extension {
1427 node: Arc::new(
1428 EmptyMetric::new(
1429 self.ctx.start,
1430 self.ctx.end,
1431 self.ctx.interval,
1432 SPECIAL_TIME_FUNCTION.to_string(),
1433 DEFAULT_FIELD_COLUMN.to_string(),
1434 Some(literal_expr),
1435 )
1436 .context(DataFusionPlanningSnafu)?,
1437 ),
1438 });
1439 Ok(plan)
1440 }
1441
1442 async fn prom_vector_selector_to_plan(
1443 &mut self,
1444 vector_selector: &VectorSelector,
1445 timestamp_fn: bool,
1446 ) -> Result<LogicalPlan> {
1447 let VectorSelector {
1448 name,
1449 offset,
1450 matchers,
1451 at: _,
1452 } = vector_selector;
1453 let matchers = self.preprocess_label_matchers(matchers, name)?;
1454 if let Some(empty_plan) = self.setup_context().await? {
1455 return Ok(empty_plan);
1456 }
1457 let normalize = self
1458 .selector_to_series_normalize_plan(offset, matchers, false)
1459 .await?;
1460
1461 let normalize = if timestamp_fn {
1462 self.create_timestamp_func_plan(normalize)?
1465 } else {
1466 normalize
1467 };
1468
1469 let manipulate = InstantManipulate::new(
1470 self.ctx.start,
1471 self.ctx.end,
1472 self.ctx.lookback_delta,
1473 self.ctx.interval,
1474 self.ctx
1475 .time_index_column
1476 .clone()
1477 .expect("time index should be set in `setup_context`"),
1478 if self.ctx.use_tsid {
1479 vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]
1480 } else {
1481 self.ctx.tag_columns.clone()
1482 },
1483 self.ctx.field_columns.first().cloned(),
1484 normalize,
1485 );
1486 Ok(LogicalPlan::Extension(Extension {
1487 node: Arc::new(manipulate),
1488 }))
1489 }
1490
1491 fn create_timestamp_func_plan(&mut self, normalize: LogicalPlan) -> Result<LogicalPlan> {
1513 let time_expr = build_special_time_expr(self.ctx.time_index_column.as_ref().unwrap())
1514 .alias(DEFAULT_FIELD_COLUMN);
1515 self.ctx.field_columns = vec![time_expr.schema_name().to_string()];
1516 let mut project_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 2);
1517 project_exprs.push(self.create_time_index_column_expr()?);
1518 project_exprs.push(time_expr);
1519 project_exprs.extend(self.create_tag_column_exprs()?);
1520
1521 LogicalPlanBuilder::from(normalize)
1522 .project(project_exprs)
1523 .context(DataFusionPlanningSnafu)?
1524 .build()
1525 .context(DataFusionPlanningSnafu)
1526 }
1527
1528 async fn prom_matrix_selector_to_plan(
1529 &mut self,
1530 matrix_selector: &MatrixSelector,
1531 ) -> Result<LogicalPlan> {
1532 let MatrixSelector { vs, range } = matrix_selector;
1533 let VectorSelector {
1534 name,
1535 offset,
1536 matchers,
1537 ..
1538 } = vs;
1539 let matchers = self.preprocess_label_matchers(matchers, name)?;
1540 ensure!(!range.is_zero(), ZeroRangeSelectorSnafu);
1541 let range_ms = range.as_millis() as _;
1542 self.ctx.range = Some(range_ms);
1543
1544 let normalize = match self.setup_context().await? {
1547 Some(empty_plan) => empty_plan,
1548 None => {
1549 self.selector_to_series_normalize_plan(offset, matchers, true)
1550 .await?
1551 }
1552 };
1553 let manipulate = RangeManipulate::new(
1554 self.ctx.start,
1555 self.ctx.end,
1556 self.ctx.interval,
1557 range_ms,
1559 self.ctx
1560 .time_index_column
1561 .clone()
1562 .expect("time index should be set in `setup_context`"),
1563 self.ctx.field_columns.clone(),
1564 normalize,
1565 )
1566 .context(DataFusionPlanningSnafu)?;
1567
1568 Ok(LogicalPlan::Extension(Extension {
1569 node: Arc::new(manipulate),
1570 }))
1571 }
1572
1573 async fn prom_call_expr_to_plan(
1574 &mut self,
1575 query_engine_state: &QueryEngineState,
1576 call_expr: &Call,
1577 ) -> Result<LogicalPlan> {
1578 let Call { func, args } = call_expr;
1579 match func.name {
1581 SPECIAL_HISTOGRAM_QUANTILE => {
1582 return self.create_histogram_plan(args, query_engine_state).await;
1583 }
1584 SPECIAL_VECTOR_FUNCTION => return self.create_vector_plan(args).await,
1585 SCALAR_FUNCTION => return self.create_scalar_plan(args, query_engine_state).await,
1586 SPECIAL_ABSENT_FUNCTION => {
1587 return self.create_absent_plan(args, query_engine_state).await;
1588 }
1589 _ => {}
1590 }
1591
1592 let args = self.create_function_args(&args.args)?;
1594 let input = if let Some(prom_expr) = &args.input {
1595 self.prom_expr_to_plan_inner(prom_expr, func.name == "timestamp", query_engine_state)
1596 .await?
1597 } else {
1598 self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
1599 self.ctx.reset_table_name_and_schema();
1600 self.ctx.tag_columns = vec![];
1601 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1602 LogicalPlan::Extension(Extension {
1603 node: Arc::new(
1604 EmptyMetric::new(
1605 self.ctx.start,
1606 self.ctx.end,
1607 self.ctx.interval,
1608 SPECIAL_TIME_FUNCTION.to_string(),
1609 DEFAULT_FIELD_COLUMN.to_string(),
1610 None,
1611 )
1612 .context(DataFusionPlanningSnafu)?,
1613 ),
1614 })
1615 };
1616 let (mut func_exprs, new_tags) =
1617 self.create_function_expr(func, args.literals.clone(), query_engine_state)?;
1618 func_exprs.insert(0, self.create_time_index_column_expr()?);
1619 func_exprs.extend_from_slice(&self.create_tag_column_exprs()?);
1620 if let Some(tsid_col) =
1621 Self::optional_tsid_projection(input.schema(), None, self.ctx.use_tsid)
1622 {
1623 func_exprs.push(tsid_col);
1624 }
1625
1626 let builder = LogicalPlanBuilder::from(input)
1627 .project(func_exprs)
1628 .context(DataFusionPlanningSnafu)?
1629 .filter(self.create_empty_values_filter_expr()?)
1630 .context(DataFusionPlanningSnafu)?;
1631
1632 let builder = match func.name {
1633 "sort" => builder
1634 .sort(self.create_field_columns_sort_exprs(true))
1635 .context(DataFusionPlanningSnafu)?,
1636 "sort_desc" => builder
1637 .sort(self.create_field_columns_sort_exprs(false))
1638 .context(DataFusionPlanningSnafu)?,
1639 "sort_by_label" => builder
1640 .sort(Self::create_sort_exprs_by_tags(
1641 func.name,
1642 args.literals,
1643 true,
1644 )?)
1645 .context(DataFusionPlanningSnafu)?,
1646 "sort_by_label_desc" => builder
1647 .sort(Self::create_sort_exprs_by_tags(
1648 func.name,
1649 args.literals,
1650 false,
1651 )?)
1652 .context(DataFusionPlanningSnafu)?,
1653
1654 _ => builder,
1655 };
1656
1657 for tag in new_tags {
1660 self.ctx.tag_columns.push(tag);
1661 }
1662
1663 let plan = builder.build().context(DataFusionPlanningSnafu)?;
1664 common_telemetry::debug!("Created PromQL function plan: {plan:?} for {call_expr:?}");
1665
1666 Ok(plan)
1667 }
1668
1669 async fn prom_ext_expr_to_plan(
1670 &mut self,
1671 query_engine_state: &QueryEngineState,
1672 ext_expr: &promql_parser::parser::ast::Extension,
1673 ) -> Result<LogicalPlan> {
1674 let expr = &ext_expr.expr;
1676 let children = expr.children();
1677 let plan = self
1678 .prom_expr_to_plan(&children[0], query_engine_state)
1679 .await?;
1680 match expr.name() {
1686 ANALYZE_NODE_NAME => LogicalPlanBuilder::from(plan)
1687 .explain(false, true)
1688 .unwrap()
1689 .build()
1690 .context(DataFusionPlanningSnafu),
1691 ANALYZE_VERBOSE_NODE_NAME => LogicalPlanBuilder::from(plan)
1692 .explain(true, true)
1693 .unwrap()
1694 .build()
1695 .context(DataFusionPlanningSnafu),
1696 EXPLAIN_NODE_NAME => LogicalPlanBuilder::from(plan)
1697 .explain(false, false)
1698 .unwrap()
1699 .build()
1700 .context(DataFusionPlanningSnafu),
1701 EXPLAIN_VERBOSE_NODE_NAME => LogicalPlanBuilder::from(plan)
1702 .explain(true, false)
1703 .unwrap()
1704 .build()
1705 .context(DataFusionPlanningSnafu),
1706 ALIAS_NODE_NAME => {
1707 let alias = expr
1708 .as_any()
1709 .downcast_ref::<AliasExpr>()
1710 .context(UnexpectedPlanExprSnafu {
1711 desc: "Expected AliasExpr",
1712 })?
1713 .alias
1714 .clone();
1715 self.apply_alias(plan, alias)
1716 }
1717 _ => LogicalPlanBuilder::empty(true)
1718 .build()
1719 .context(DataFusionPlanningSnafu),
1720 }
1721 }
1722
1723 #[allow(clippy::mutable_key_type)]
1733 fn preprocess_label_matchers(
1734 &mut self,
1735 label_matchers: &Matchers,
1736 name: &Option<String>,
1737 ) -> Result<Matchers> {
1738 self.ctx.reset();
1739
1740 let metric_name;
1741 if let Some(name) = name.clone() {
1742 metric_name = Some(name);
1743 ensure!(
1744 label_matchers.find_matchers(METRIC_NAME).is_empty(),
1745 MultipleMetricMatchersSnafu
1746 );
1747 } else {
1748 let mut matches = label_matchers.find_matchers(METRIC_NAME);
1749 ensure!(!matches.is_empty(), NoMetricMatcherSnafu);
1750 ensure!(matches.len() == 1, MultipleMetricMatchersSnafu);
1751 ensure!(
1752 matches[0].op == MatchOp::Equal,
1753 UnsupportedMatcherOpSnafu {
1754 matcher_op: matches[0].op.to_string(),
1755 matcher: METRIC_NAME
1756 }
1757 );
1758 metric_name = matches.pop().map(|m| m.value);
1759 }
1760
1761 self.ctx.table_name = metric_name;
1762
1763 let mut matchers = HashSet::new();
1764 for matcher in &label_matchers.matchers {
1765 if matcher.name == FIELD_COLUMN_MATCHER {
1767 self.ctx
1768 .field_column_matcher
1769 .get_or_insert_default()
1770 .push(matcher.clone());
1771 } else if matcher.name == SCHEMA_COLUMN_MATCHER || matcher.name == DB_COLUMN_MATCHER {
1772 ensure!(
1773 matcher.op == MatchOp::Equal,
1774 UnsupportedMatcherOpSnafu {
1775 matcher: matcher.name.clone(),
1776 matcher_op: matcher.op.to_string(),
1777 }
1778 );
1779 self.ctx.schema_name = Some(matcher.value.clone());
1780 } else if matcher.name != METRIC_NAME {
1781 self.ctx.selector_matcher.push(matcher.clone());
1782 let _ = matchers.insert(matcher.clone());
1783 }
1784 }
1785
1786 Ok(Matchers::new(matchers.into_iter().collect()))
1787 }
1788
1789 async fn selector_to_series_normalize_plan(
1790 &mut self,
1791 offset: &Option<Offset>,
1792 label_matchers: Matchers,
1793 is_range_selector: bool,
1794 ) -> Result<LogicalPlan> {
1795 let table_ref = self.table_ref()?;
1797 let mut table_scan = self.create_table_scan_plan(table_ref.clone()).await?;
1798 let table_schema = table_scan.schema();
1799
1800 let offset_duration = match offset {
1802 Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond,
1803 Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond),
1804 None => 0,
1805 };
1806 let mut scan_filters = Self::matchers_to_expr(label_matchers.clone(), table_schema)?;
1807 if let Some(time_index_filter) = self.build_time_index_filter(offset_duration)? {
1808 scan_filters.push(time_index_filter);
1809 }
1810 table_scan = LogicalPlanBuilder::from(table_scan)
1811 .filter(conjunction(scan_filters).unwrap()) .context(DataFusionPlanningSnafu)?
1813 .build()
1814 .context(DataFusionPlanningSnafu)?;
1815
1816 if let Some(field_matchers) = &self.ctx.field_column_matcher {
1818 let col_set = self.ctx.field_columns.iter().collect::<HashSet<_>>();
1819 let mut result_set = HashSet::new();
1821 let mut reverse_set = HashSet::new();
1823 for matcher in field_matchers {
1824 match &matcher.op {
1825 MatchOp::Equal => {
1826 if col_set.contains(&matcher.value) {
1827 let _ = result_set.insert(matcher.value.clone());
1828 } else {
1829 return Err(ColumnNotFoundSnafu {
1830 col: matcher.value.clone(),
1831 }
1832 .build());
1833 }
1834 }
1835 MatchOp::NotEqual => {
1836 if col_set.contains(&matcher.value) {
1837 let _ = reverse_set.insert(matcher.value.clone());
1838 } else {
1839 return Err(ColumnNotFoundSnafu {
1840 col: matcher.value.clone(),
1841 }
1842 .build());
1843 }
1844 }
1845 MatchOp::Re(regex) => {
1846 for col in &self.ctx.field_columns {
1847 if regex.is_match(col) {
1848 let _ = result_set.insert(col.clone());
1849 }
1850 }
1851 }
1852 MatchOp::NotRe(regex) => {
1853 for col in &self.ctx.field_columns {
1854 if regex.is_match(col) {
1855 let _ = reverse_set.insert(col.clone());
1856 }
1857 }
1858 }
1859 }
1860 }
1861 if result_set.is_empty() {
1863 result_set = col_set.into_iter().cloned().collect();
1864 }
1865 for col in reverse_set {
1866 let _ = result_set.remove(&col);
1867 }
1868
1869 self.ctx.field_columns = self
1871 .ctx
1872 .field_columns
1873 .drain(..)
1874 .filter(|col| result_set.contains(col))
1875 .collect();
1876
1877 let exprs = result_set
1878 .into_iter()
1879 .map(|col| DfExpr::Column(Column::new_unqualified(col)))
1880 .chain(self.create_tag_column_exprs()?)
1881 .chain(
1882 self.ctx
1883 .use_tsid
1884 .then_some(DfExpr::Column(Column::new_unqualified(
1885 DATA_SCHEMA_TSID_COLUMN_NAME,
1886 ))),
1887 )
1888 .chain(Some(self.create_time_index_column_expr()?))
1889 .collect::<Vec<_>>();
1890
1891 table_scan = LogicalPlanBuilder::from(table_scan)
1893 .project(exprs)
1894 .context(DataFusionPlanningSnafu)?
1895 .build()
1896 .context(DataFusionPlanningSnafu)?;
1897 }
1898
1899 let series_key_columns = if self.ctx.use_tsid {
1901 vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]
1902 } else {
1903 self.ctx.tag_columns.clone()
1904 };
1905
1906 let sort_exprs = if self.ctx.use_tsid {
1907 vec![
1908 DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME)).sort(true, true),
1909 self.create_time_index_column_expr()?.sort(true, true),
1910 ]
1911 } else {
1912 self.create_tag_and_time_index_column_sort_exprs()?
1913 };
1914
1915 let sort_plan = LogicalPlanBuilder::from(table_scan)
1916 .sort(sort_exprs)
1917 .context(DataFusionPlanningSnafu)?
1918 .build()
1919 .context(DataFusionPlanningSnafu)?;
1920
1921 let time_index_column =
1923 self.ctx
1924 .time_index_column
1925 .clone()
1926 .with_context(|| TimeIndexNotFoundSnafu {
1927 table: table_ref.to_string(),
1928 })?;
1929 let divide_plan = LogicalPlan::Extension(Extension {
1930 node: Arc::new(SeriesDivide::new(
1931 series_key_columns.clone(),
1932 time_index_column,
1933 sort_plan,
1934 )),
1935 });
1936
1937 if !is_range_selector && offset_duration == 0 {
1939 return Ok(divide_plan);
1940 }
1941 let series_normalize = SeriesNormalize::new(
1942 offset_duration,
1943 self.ctx
1944 .time_index_column
1945 .clone()
1946 .with_context(|| TimeIndexNotFoundSnafu {
1947 table: table_ref.to_quoted_string(),
1948 })?,
1949 is_range_selector,
1950 series_key_columns,
1951 divide_plan,
1952 );
1953 let logical_plan = LogicalPlan::Extension(Extension {
1954 node: Arc::new(series_normalize),
1955 });
1956
1957 Ok(logical_plan)
1958 }
1959
1960 fn agg_modifier_to_col(
1967 &mut self,
1968 input_schema: &DFSchemaRef,
1969 modifier: &Option<LabelModifier>,
1970 update_ctx: bool,
1971 ) -> Result<Vec<DfExpr>> {
1972 match modifier {
1973 None => {
1974 if update_ctx {
1975 self.ctx.tag_columns.clear();
1976 }
1977 Ok(vec![self.create_time_index_column_expr()?])
1978 }
1979 Some(LabelModifier::Include(labels)) => {
1980 if update_ctx {
1981 self.ctx.tag_columns.clear();
1982 }
1983 let mut exprs = Vec::with_capacity(labels.labels.len());
1984 for label in &labels.labels {
1985 if is_metric_engine_internal_column(label) {
1986 continue;
1987 }
1988 if let Some(column_name) = Self::find_case_sensitive_column(input_schema, label)
1990 {
1991 exprs.push(DfExpr::Column(Column::from_name(column_name.clone())));
1992
1993 if update_ctx {
1994 self.ctx.tag_columns.push(column_name);
1996 }
1997 }
1998 }
1999 exprs.push(self.create_time_index_column_expr()?);
2001
2002 Ok(exprs)
2003 }
2004 Some(LabelModifier::Exclude(labels)) => {
2005 let mut all_fields = input_schema
2006 .fields()
2007 .iter()
2008 .map(|f| f.name())
2009 .collect::<BTreeSet<_>>();
2010
2011 all_fields.retain(|col| !is_metric_engine_internal_column(col.as_str()));
2014
2015 for label in &labels.labels {
2018 let _ = all_fields.remove(label);
2019 }
2020
2021 if let Some(time_index) = &self.ctx.time_index_column {
2023 let _ = all_fields.remove(time_index);
2024 }
2025 for value in &self.ctx.field_columns {
2026 let _ = all_fields.remove(value);
2027 }
2028
2029 if update_ctx {
2030 self.ctx.tag_columns = all_fields.iter().map(|col| (*col).clone()).collect();
2032 }
2033
2034 let mut exprs = all_fields
2036 .into_iter()
2037 .map(|c| DfExpr::Column(Column::from(c)))
2038 .collect::<Vec<_>>();
2039
2040 exprs.push(self.create_time_index_column_expr()?);
2042
2043 Ok(exprs)
2044 }
2045 }
2046 }
2047
2048 pub fn matchers_to_expr(
2050 label_matchers: Matchers,
2051 table_schema: &DFSchemaRef,
2052 ) -> Result<Vec<DfExpr>> {
2053 let mut exprs = Vec::with_capacity(label_matchers.matchers.len());
2054 for matcher in label_matchers.matchers {
2055 if matcher.name == SCHEMA_COLUMN_MATCHER
2056 || matcher.name == DB_COLUMN_MATCHER
2057 || matcher.name == FIELD_COLUMN_MATCHER
2058 {
2059 continue;
2060 }
2061
2062 let column_name = Self::find_case_sensitive_column(table_schema, matcher.name.as_str());
2063 let col = if let Some(column_name) = column_name {
2064 DfExpr::Column(Column::from_name(column_name))
2065 } else {
2066 DfExpr::Literal(ScalarValue::Utf8(Some(String::new())), None)
2067 .alias(matcher.name.clone())
2068 };
2069 let lit = DfExpr::Literal(ScalarValue::Utf8(Some(matcher.value)), None);
2070 let expr = match matcher.op {
2071 MatchOp::Equal => col.eq(lit),
2072 MatchOp::NotEqual => col.not_eq(lit),
2073 MatchOp::Re(re) => {
2074 if re.as_str() == "^(?:.*)$" {
2080 continue;
2081 }
2082 if re.as_str() == "^(?:.+)$" {
2083 col.not_eq(DfExpr::Literal(
2084 ScalarValue::Utf8(Some(String::new())),
2085 None,
2086 ))
2087 } else {
2088 DfExpr::BinaryExpr(BinaryExpr {
2089 left: Box::new(col),
2090 op: Operator::RegexMatch,
2091 right: Box::new(DfExpr::Literal(
2092 ScalarValue::Utf8(Some(re.as_str().to_string())),
2093 None,
2094 )),
2095 })
2096 }
2097 }
2098 MatchOp::NotRe(re) => {
2099 if re.as_str() == "^(?:.*)$" {
2100 DfExpr::Literal(ScalarValue::Boolean(Some(false)), None)
2101 } else if re.as_str() == "^(?:.+)$" {
2102 col.eq(DfExpr::Literal(
2103 ScalarValue::Utf8(Some(String::new())),
2104 None,
2105 ))
2106 } else {
2107 DfExpr::BinaryExpr(BinaryExpr {
2108 left: Box::new(col),
2109 op: Operator::RegexNotMatch,
2110 right: Box::new(DfExpr::Literal(
2111 ScalarValue::Utf8(Some(re.as_str().to_string())),
2112 None,
2113 )),
2114 })
2115 }
2116 }
2117 };
2118 exprs.push(expr);
2119 }
2120
2121 Ok(exprs)
2122 }
2123
2124 fn find_case_sensitive_column(schema: &DFSchemaRef, column: &str) -> Option<String> {
2125 if is_metric_engine_internal_column(column) {
2126 return None;
2127 }
2128 schema
2129 .fields()
2130 .iter()
2131 .find(|field| field.name() == column)
2132 .map(|field| field.name().clone())
2133 }
2134
2135 fn table_from_source(&self, source: &Arc<dyn TableSource>) -> Result<table::TableRef> {
2136 Ok(source
2137 .as_any()
2138 .downcast_ref::<DefaultTableSource>()
2139 .context(UnknownTableSnafu)?
2140 .table_provider
2141 .as_any()
2142 .downcast_ref::<DfTableProviderAdapter>()
2143 .context(UnknownTableSnafu)?
2144 .table())
2145 }
2146
2147 fn table_ref(&self) -> Result<TableReference> {
2148 let table_name = self
2149 .ctx
2150 .table_name
2151 .clone()
2152 .context(TableNameNotFoundSnafu)?;
2153
2154 let table_ref = if let Some(schema_name) = &self.ctx.schema_name {
2156 TableReference::partial(schema_name.as_str(), table_name.as_str())
2157 } else {
2158 TableReference::bare(table_name.as_str())
2159 };
2160
2161 Ok(table_ref)
2162 }
2163
2164 fn build_time_index_filter(&self, offset_duration: i64) -> Result<Option<DfExpr>> {
2165 let start = self.ctx.start;
2166 let end = self.ctx.end;
2167 if end < start {
2168 return InvalidTimeRangeSnafu { start, end }.fail();
2169 }
2170 let lookback_delta = self.ctx.lookback_delta;
2171 let range = self.ctx.range.unwrap_or_default();
2172 let interval = self.ctx.interval;
2173 let time_index_expr = self.create_time_index_column_expr()?;
2174 let num_points = (end - start) / interval;
2175
2176 let selector_window = if range == 0 { lookback_delta } else { range };
2184 let lower_exclusive_adjustment = if selector_window > 0 { 1 } else { 0 };
2185
2186 if (end - start) / interval > MAX_SCATTER_POINTS || interval <= INTERVAL_1H {
2188 let single_time_range = time_index_expr
2189 .clone()
2190 .gt_eq(DfExpr::Literal(
2191 ScalarValue::TimestampMillisecond(
2192 Some(
2193 self.ctx.start - offset_duration - selector_window
2194 + lower_exclusive_adjustment,
2195 ),
2196 None,
2197 ),
2198 None,
2199 ))
2200 .and(time_index_expr.lt_eq(DfExpr::Literal(
2201 ScalarValue::TimestampMillisecond(Some(self.ctx.end - offset_duration), None),
2202 None,
2203 )));
2204 return Ok(Some(single_time_range));
2205 }
2206
2207 let mut filters = Vec::with_capacity(num_points as usize + 1);
2209 for timestamp in (start..=end).step_by(interval as usize) {
2210 filters.push(
2211 time_index_expr
2212 .clone()
2213 .gt_eq(DfExpr::Literal(
2214 ScalarValue::TimestampMillisecond(
2215 Some(
2216 timestamp - offset_duration - selector_window
2217 + lower_exclusive_adjustment,
2218 ),
2219 None,
2220 ),
2221 None,
2222 ))
2223 .and(time_index_expr.clone().lt_eq(DfExpr::Literal(
2224 ScalarValue::TimestampMillisecond(Some(timestamp - offset_duration), None),
2225 None,
2226 ))),
2227 )
2228 }
2229
2230 Ok(filters.into_iter().reduce(DfExpr::or))
2231 }
2232
2233 async fn create_table_scan_plan(&mut self, table_ref: TableReference) -> Result<LogicalPlan> {
2238 let provider = self
2239 .table_provider
2240 .resolve_table(table_ref.clone())
2241 .await
2242 .context(CatalogSnafu)?;
2243
2244 let logical_table = self.table_from_source(&provider)?;
2245
2246 let mut maybe_phy_table_ref = table_ref.clone();
2248 let mut scan_provider = provider;
2249 let mut table_id_filter: Option<u32> = None;
2250
2251 if logical_table.table_info().meta.engine == METRIC_ENGINE_NAME
2254 && let Some(physical_table_name) = logical_table
2255 .table_info()
2256 .meta
2257 .options
2258 .extra_options
2259 .get(LOGICAL_TABLE_METADATA_KEY)
2260 {
2261 let physical_table_ref = if let Some(schema_name) = &self.ctx.schema_name {
2262 TableReference::partial(schema_name.as_str(), physical_table_name.as_str())
2263 } else {
2264 TableReference::bare(physical_table_name.as_str())
2265 };
2266
2267 let physical_provider = match self
2268 .table_provider
2269 .resolve_table(physical_table_ref.clone())
2270 .await
2271 {
2272 Ok(provider) => provider,
2273 Err(e) if e.status_code() == StatusCode::TableNotFound => {
2274 scan_provider.clone()
2277 }
2278 Err(e) => return Err(e).context(CatalogSnafu),
2279 };
2280
2281 if !Arc::ptr_eq(&physical_provider, &scan_provider) {
2282 let physical_table = self.table_from_source(&physical_provider)?;
2284
2285 let has_table_id = physical_table
2286 .schema()
2287 .column_schema_by_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME)
2288 .is_some();
2289 let has_tsid = physical_table
2290 .schema()
2291 .column_schema_by_name(DATA_SCHEMA_TSID_COLUMN_NAME)
2292 .is_some_and(|col| matches!(col.data_type, ConcreteDataType::UInt64(_)));
2293
2294 if has_table_id && has_tsid {
2295 scan_provider = physical_provider;
2296 maybe_phy_table_ref = physical_table_ref;
2297 table_id_filter = Some(logical_table.table_info().ident.table_id);
2298 }
2299 }
2300 }
2301
2302 let scan_table = self.table_from_source(&scan_provider)?;
2303
2304 let use_tsid = table_id_filter.is_some()
2305 && scan_table
2306 .schema()
2307 .column_schema_by_name(DATA_SCHEMA_TSID_COLUMN_NAME)
2308 .is_some_and(|col| matches!(col.data_type, ConcreteDataType::UInt64(_)));
2309 self.ctx.use_tsid = use_tsid;
2310
2311 let all_table_tags = self.ctx.tag_columns.clone();
2312
2313 let scan_tag_columns = if use_tsid {
2314 let mut scan_tags = self.ctx.tag_columns.clone();
2315 for matcher in &self.ctx.selector_matcher {
2316 if is_metric_engine_internal_column(&matcher.name) {
2317 continue;
2318 }
2319 if all_table_tags.iter().any(|tag| tag == &matcher.name) {
2320 scan_tags.push(matcher.name.clone());
2321 }
2322 }
2323 scan_tags.sort_unstable();
2324 scan_tags.dedup();
2325 scan_tags
2326 } else {
2327 self.ctx.tag_columns.clone()
2328 };
2329
2330 let is_time_index_ms = scan_table
2331 .schema()
2332 .timestamp_column()
2333 .with_context(|| TimeIndexNotFoundSnafu {
2334 table: maybe_phy_table_ref.to_quoted_string(),
2335 })?
2336 .data_type
2337 == ConcreteDataType::timestamp_millisecond_datatype();
2338
2339 let scan_projection = if table_id_filter.is_some() {
2340 let mut required_columns = HashSet::new();
2341 required_columns.insert(DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string());
2342 required_columns.insert(self.ctx.time_index_column.clone().with_context(|| {
2343 TimeIndexNotFoundSnafu {
2344 table: maybe_phy_table_ref.to_quoted_string(),
2345 }
2346 })?);
2347 for col in &scan_tag_columns {
2348 required_columns.insert(col.clone());
2349 }
2350 for col in &self.ctx.field_columns {
2351 required_columns.insert(col.clone());
2352 }
2353 if use_tsid {
2354 required_columns.insert(DATA_SCHEMA_TSID_COLUMN_NAME.to_string());
2355 }
2356
2357 let arrow_schema = scan_table.schema().arrow_schema().clone();
2358 Some(
2359 arrow_schema
2360 .fields()
2361 .iter()
2362 .enumerate()
2363 .filter(|(_, field)| required_columns.contains(field.name().as_str()))
2364 .map(|(idx, _)| idx)
2365 .collect::<Vec<_>>(),
2366 )
2367 } else {
2368 None
2369 };
2370
2371 let mut scan_plan =
2372 LogicalPlanBuilder::scan(maybe_phy_table_ref.clone(), scan_provider, scan_projection)
2373 .context(DataFusionPlanningSnafu)?
2374 .build()
2375 .context(DataFusionPlanningSnafu)?;
2376
2377 if let Some(table_id) = table_id_filter {
2378 scan_plan = LogicalPlanBuilder::from(scan_plan)
2379 .filter(
2380 DfExpr::Column(Column::from_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME))
2381 .eq(lit(table_id)),
2382 )
2383 .context(DataFusionPlanningSnafu)?
2384 .alias(table_ref.clone()) .context(DataFusionPlanningSnafu)?
2386 .build()
2387 .context(DataFusionPlanningSnafu)?;
2388 }
2389
2390 if !is_time_index_ms {
2391 let expr: Vec<_> = self
2393 .create_field_column_exprs()?
2394 .into_iter()
2395 .chain(
2396 scan_tag_columns
2397 .iter()
2398 .map(|tag| DfExpr::Column(Column::from_name(tag))),
2399 )
2400 .chain(self.ctx.use_tsid.then_some(DfExpr::Column(Column::new(
2401 Some(table_ref.clone()),
2402 DATA_SCHEMA_TSID_COLUMN_NAME.to_string(),
2403 ))))
2404 .chain(Some(DfExpr::Alias(Alias {
2405 expr: Box::new(DfExpr::Cast(Cast {
2406 expr: Box::new(self.create_time_index_column_expr()?),
2407 data_type: ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
2408 })),
2409 relation: Some(table_ref.clone()),
2410 name: self
2411 .ctx
2412 .time_index_column
2413 .as_ref()
2414 .with_context(|| TimeIndexNotFoundSnafu {
2415 table: table_ref.to_quoted_string(),
2416 })?
2417 .clone(),
2418 metadata: None,
2419 })))
2420 .collect::<Vec<_>>();
2421 scan_plan = LogicalPlanBuilder::from(scan_plan)
2422 .project(expr)
2423 .context(DataFusionPlanningSnafu)?
2424 .build()
2425 .context(DataFusionPlanningSnafu)?;
2426 } else if table_id_filter.is_some() {
2427 let project_exprs = self
2429 .create_field_column_exprs()?
2430 .into_iter()
2431 .chain(
2432 scan_tag_columns
2433 .iter()
2434 .map(|tag| DfExpr::Column(Column::from_name(tag))),
2435 )
2436 .chain(
2437 self.ctx
2438 .use_tsid
2439 .then_some(DfExpr::Column(Column::from_name(
2440 DATA_SCHEMA_TSID_COLUMN_NAME,
2441 ))),
2442 )
2443 .chain(Some(self.create_time_index_column_expr()?))
2444 .collect::<Vec<_>>();
2445
2446 scan_plan = LogicalPlanBuilder::from(scan_plan)
2447 .project(project_exprs)
2448 .context(DataFusionPlanningSnafu)?
2449 .build()
2450 .context(DataFusionPlanningSnafu)?;
2451 }
2452
2453 let result = LogicalPlanBuilder::from(scan_plan)
2454 .build()
2455 .context(DataFusionPlanningSnafu)?;
2456 Ok(result)
2457 }
2458
2459 fn collect_row_key_tag_columns_from_plan(
2460 &self,
2461 plan: &LogicalPlan,
2462 ) -> Result<BTreeSet<String>> {
2463 fn walk(
2464 planner: &PromPlanner,
2465 plan: &LogicalPlan,
2466 out: &mut BTreeSet<String>,
2467 ) -> Result<()> {
2468 if let LogicalPlan::TableScan(scan) = plan
2470 && let Ok(table) = planner.table_from_source(&scan.source)
2471 {
2472 for col in table.table_info().meta.row_key_column_names() {
2473 if col != DATA_SCHEMA_TABLE_ID_COLUMN_NAME
2474 && col != DATA_SCHEMA_TSID_COLUMN_NAME
2475 && !is_metric_engine_internal_column(col)
2476 {
2477 out.insert(col.clone());
2478 }
2479 }
2480 }
2481
2482 for input in plan.inputs() {
2483 walk(planner, input, out)?;
2484 }
2485 Ok(())
2486 }
2487
2488 let mut out = BTreeSet::new();
2489 walk(self, plan, &mut out)?;
2490 Ok(out)
2491 }
2492
2493 fn ensure_tag_columns_available(
2494 &self,
2495 plan: LogicalPlan,
2496 required_tags: &BTreeSet<String>,
2497 ) -> Result<LogicalPlan> {
2498 if required_tags.is_empty() {
2499 return Ok(plan);
2500 }
2501
2502 struct Rewriter {
2503 required_tags: BTreeSet<String>,
2504 }
2505
2506 impl TreeNodeRewriter for Rewriter {
2507 type Node = LogicalPlan;
2508
2509 fn f_up(
2510 &mut self,
2511 node: Self::Node,
2512 ) -> datafusion_common::Result<Transformed<Self::Node>> {
2513 match node {
2514 LogicalPlan::TableScan(scan) => {
2515 let schema = scan.source.schema();
2516 let mut projection = match scan.projection.clone() {
2517 Some(p) => p,
2518 None => {
2519 return Ok(Transformed::no(LogicalPlan::TableScan(scan)));
2521 }
2522 };
2523
2524 let mut changed = false;
2525 for tag in &self.required_tags {
2526 if let Some((idx, _)) = schema
2527 .fields()
2528 .iter()
2529 .enumerate()
2530 .find(|(_, field)| field.name() == tag)
2531 && !projection.contains(&idx)
2532 {
2533 projection.push(idx);
2534 changed = true;
2535 }
2536 }
2537
2538 if !changed {
2539 return Ok(Transformed::no(LogicalPlan::TableScan(scan)));
2540 }
2541
2542 projection.sort_unstable();
2543 projection.dedup();
2544
2545 let new_scan = TableScan::try_new(
2546 scan.table_name.clone(),
2547 scan.source.clone(),
2548 Some(projection),
2549 scan.filters,
2550 scan.fetch,
2551 )?;
2552 Ok(Transformed::yes(LogicalPlan::TableScan(new_scan)))
2553 }
2554 LogicalPlan::Projection(proj) => {
2555 let input_schema = proj.input.schema();
2556
2557 let existing = proj
2558 .schema
2559 .fields()
2560 .iter()
2561 .map(|f| f.name().as_str())
2562 .collect::<HashSet<_>>();
2563
2564 let mut expr = proj.expr.clone();
2565 let mut has_changed = false;
2566 for tag in &self.required_tags {
2567 if existing.contains(tag.as_str()) {
2568 continue;
2569 }
2570
2571 if let Some(idx) = input_schema.index_of_column_by_name(None, tag) {
2572 expr.push(DfExpr::Column(Column::from(
2573 input_schema.qualified_field(idx),
2574 )));
2575 has_changed = true;
2576 }
2577 }
2578
2579 if !has_changed {
2580 return Ok(Transformed::no(LogicalPlan::Projection(proj)));
2581 }
2582
2583 let new_proj = Projection::try_new(expr, proj.input)?;
2584 Ok(Transformed::yes(LogicalPlan::Projection(new_proj)))
2585 }
2586 other => Ok(Transformed::no(other)),
2587 }
2588 }
2589 }
2590
2591 let mut rewriter = Rewriter {
2592 required_tags: required_tags.clone(),
2593 };
2594 let rewritten = plan
2595 .rewrite(&mut rewriter)
2596 .context(DataFusionPlanningSnafu)?;
2597 Ok(rewritten.data)
2598 }
2599
2600 fn refresh_tag_columns_from_schema(&mut self, schema: &DFSchemaRef) {
2601 let time_index = self.ctx.time_index_column.as_deref();
2602 let field_columns = self.ctx.field_columns.iter().collect::<HashSet<_>>();
2603
2604 let mut tags = schema
2605 .fields()
2606 .iter()
2607 .map(|f| f.name())
2608 .filter(|name| Some(name.as_str()) != time_index)
2609 .filter(|name| !field_columns.contains(name))
2610 .filter(|name| !is_metric_engine_internal_column(name))
2611 .cloned()
2612 .collect::<Vec<_>>();
2613 tags.sort_unstable();
2614 tags.dedup();
2615 self.ctx.tag_columns = tags;
2616 }
2617
2618 async fn setup_context(&mut self) -> Result<Option<LogicalPlan>> {
2622 let table_ref = self.table_ref()?;
2623 let source = match self.table_provider.resolve_table(table_ref.clone()).await {
2624 Err(e) if e.status_code() == StatusCode::TableNotFound => {
2625 let plan = self.setup_context_for_empty_metric()?;
2626 return Ok(Some(plan));
2627 }
2628 res => res.context(CatalogSnafu)?,
2629 };
2630 let table = self.table_from_source(&source)?;
2631
2632 let time_index = table
2634 .schema()
2635 .timestamp_column()
2636 .with_context(|| TimeIndexNotFoundSnafu {
2637 table: table_ref.to_quoted_string(),
2638 })?
2639 .name
2640 .clone();
2641 self.ctx.time_index_column = Some(time_index);
2642
2643 let values = table
2645 .table_info()
2646 .meta
2647 .field_column_names()
2648 .cloned()
2649 .collect();
2650 self.ctx.field_columns = values;
2651
2652 let tags = table
2654 .table_info()
2655 .meta
2656 .row_key_column_names()
2657 .filter(|col| {
2658 col != &DATA_SCHEMA_TABLE_ID_COLUMN_NAME && col != &DATA_SCHEMA_TSID_COLUMN_NAME
2660 })
2661 .cloned()
2662 .collect();
2663 self.ctx.tag_columns = tags;
2664
2665 self.ctx.use_tsid = false;
2666
2667 Ok(None)
2668 }
2669
2670 fn setup_context_for_empty_metric(&mut self) -> Result<LogicalPlan> {
2673 self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
2674 self.ctx.reset_table_name_and_schema();
2675 self.ctx.tag_columns = vec![];
2676 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
2677 self.ctx.use_tsid = false;
2678
2679 let plan = LogicalPlan::Extension(Extension {
2681 node: Arc::new(
2682 EmptyMetric::new(
2683 0,
2684 -1,
2685 self.ctx.interval,
2686 SPECIAL_TIME_FUNCTION.to_string(),
2687 DEFAULT_FIELD_COLUMN.to_string(),
2688 Some(lit(0.0f64)),
2689 )
2690 .context(DataFusionPlanningSnafu)?,
2691 ),
2692 });
2693 Ok(plan)
2694 }
2695
2696 fn create_function_args(&self, args: &[Box<PromExpr>]) -> Result<FunctionArgs> {
2698 let mut result = FunctionArgs::default();
2699
2700 for arg in args {
2701 if let Some(expr) = Self::try_build_literal_expr(arg) {
2703 result.literals.push(expr);
2704 } else {
2705 match arg.as_ref() {
2707 PromExpr::Subquery(_)
2708 | PromExpr::VectorSelector(_)
2709 | PromExpr::MatrixSelector(_)
2710 | PromExpr::Extension(_)
2711 | PromExpr::Aggregate(_)
2712 | PromExpr::Paren(_)
2713 | PromExpr::Call(_)
2714 | PromExpr::Binary(_)
2715 | PromExpr::Unary(_) => {
2716 if result.input.replace(*arg.clone()).is_some() {
2717 MultipleVectorSnafu { expr: *arg.clone() }.fail()?;
2718 }
2719 }
2720
2721 _ => {
2722 let expr = Self::get_param_as_literal_expr(&Some(arg.clone()), None, None)?;
2723 result.literals.push(expr);
2724 }
2725 }
2726 }
2727 }
2728
2729 Ok(result)
2730 }
2731
2732 fn create_function_expr(
2738 &mut self,
2739 func: &Function,
2740 other_input_exprs: Vec<DfExpr>,
2741 query_engine_state: &QueryEngineState,
2742 ) -> Result<(Vec<DfExpr>, Vec<String>)> {
2743 let mut other_input_exprs: VecDeque<DfExpr> = other_input_exprs.into();
2745
2746 let field_column_pos = 0;
2748 let mut exprs = Vec::with_capacity(self.ctx.field_columns.len());
2749 let mut new_tags = vec![];
2751 let scalar_func = match func.name {
2752 "increase" => ScalarFunc::ExtrapolateUdf(
2753 Arc::new(Increase::scalar_udf()),
2754 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
2755 ),
2756 "rate" => ScalarFunc::ExtrapolateUdf(
2757 Arc::new(Rate::scalar_udf()),
2758 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
2759 ),
2760 "delta" => ScalarFunc::ExtrapolateUdf(
2761 Arc::new(Delta::scalar_udf()),
2762 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
2763 ),
2764 "idelta" => ScalarFunc::Udf(Arc::new(IDelta::<false>::scalar_udf())),
2765 "irate" => ScalarFunc::Udf(Arc::new(IDelta::<true>::scalar_udf())),
2766 "resets" => ScalarFunc::Udf(Arc::new(Resets::scalar_udf())),
2767 "changes" => ScalarFunc::Udf(Arc::new(Changes::scalar_udf())),
2768 "deriv" => ScalarFunc::Udf(Arc::new(Deriv::scalar_udf())),
2769 "avg_over_time" => ScalarFunc::Udf(Arc::new(AvgOverTime::scalar_udf())),
2770 "min_over_time" => ScalarFunc::Udf(Arc::new(MinOverTime::scalar_udf())),
2771 "max_over_time" => ScalarFunc::Udf(Arc::new(MaxOverTime::scalar_udf())),
2772 "sum_over_time" => ScalarFunc::Udf(Arc::new(SumOverTime::scalar_udf())),
2773 "count_over_time" => ScalarFunc::Udf(Arc::new(CountOverTime::scalar_udf())),
2774 "last_over_time" => ScalarFunc::Udf(Arc::new(LastOverTime::scalar_udf())),
2775 "absent_over_time" => ScalarFunc::Udf(Arc::new(AbsentOverTime::scalar_udf())),
2776 "present_over_time" => ScalarFunc::Udf(Arc::new(PresentOverTime::scalar_udf())),
2777 "stddev_over_time" => ScalarFunc::Udf(Arc::new(StddevOverTime::scalar_udf())),
2778 "stdvar_over_time" => ScalarFunc::Udf(Arc::new(StdvarOverTime::scalar_udf())),
2779 "quantile_over_time" => ScalarFunc::Udf(Arc::new(QuantileOverTime::scalar_udf())),
2780 "predict_linear" => {
2781 other_input_exprs[0] = DfExpr::Cast(Cast {
2782 expr: Box::new(other_input_exprs[0].clone()),
2783 data_type: ArrowDataType::Int64,
2784 });
2785 ScalarFunc::Udf(Arc::new(PredictLinear::scalar_udf()))
2786 }
2787 "double_exponential_smoothing" | "holt_winters" => {
2788 ScalarFunc::Udf(Arc::new(DoubleExponentialSmoothing::scalar_udf()))
2789 }
2790 "time" => {
2791 exprs.push(build_special_time_expr(
2792 self.ctx.time_index_column.as_ref().unwrap(),
2793 ));
2794 ScalarFunc::GeneratedExpr
2795 }
2796 "minute" => {
2797 let expr = self.date_part_on_time_index("minute")?;
2799 exprs.push(expr);
2800 ScalarFunc::GeneratedExpr
2801 }
2802 "hour" => {
2803 let expr = self.date_part_on_time_index("hour")?;
2805 exprs.push(expr);
2806 ScalarFunc::GeneratedExpr
2807 }
2808 "month" => {
2809 let expr = self.date_part_on_time_index("month")?;
2811 exprs.push(expr);
2812 ScalarFunc::GeneratedExpr
2813 }
2814 "year" => {
2815 let expr = self.date_part_on_time_index("year")?;
2817 exprs.push(expr);
2818 ScalarFunc::GeneratedExpr
2819 }
2820 "day_of_month" => {
2821 let expr = self.date_part_on_time_index("day")?;
2823 exprs.push(expr);
2824 ScalarFunc::GeneratedExpr
2825 }
2826 "day_of_week" => {
2827 let expr = self.date_part_on_time_index("dow")?;
2829 exprs.push(expr);
2830 ScalarFunc::GeneratedExpr
2831 }
2832 "day_of_year" => {
2833 let expr = self.date_part_on_time_index("doy")?;
2835 exprs.push(expr);
2836 ScalarFunc::GeneratedExpr
2837 }
2838 "days_in_month" => {
2839 let day_lit_expr = "day".lit();
2844 let month_lit_expr = "month".lit();
2845 let interval_1month_lit_expr =
2846 DfExpr::Literal(ScalarValue::IntervalYearMonth(Some(1)), None);
2847 let interval_1day_lit_expr = DfExpr::Literal(
2848 ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(1, 0))),
2849 None,
2850 );
2851 let the_1month_minus_1day_expr = DfExpr::BinaryExpr(BinaryExpr {
2852 left: Box::new(interval_1month_lit_expr),
2853 op: Operator::Minus,
2854 right: Box::new(interval_1day_lit_expr),
2855 });
2856 let date_trunc_expr = DfExpr::ScalarFunction(ScalarFunction {
2857 func: datafusion_functions::datetime::date_trunc(),
2858 args: vec![month_lit_expr, self.create_time_index_column_expr()?],
2859 });
2860 let date_trunc_plus_interval_expr = DfExpr::BinaryExpr(BinaryExpr {
2861 left: Box::new(date_trunc_expr),
2862 op: Operator::Plus,
2863 right: Box::new(the_1month_minus_1day_expr),
2864 });
2865 let date_part_expr = DfExpr::ScalarFunction(ScalarFunction {
2866 func: datafusion_functions::datetime::date_part(),
2867 args: vec![day_lit_expr, date_trunc_plus_interval_expr],
2868 });
2869
2870 exprs.push(date_part_expr);
2871 ScalarFunc::GeneratedExpr
2872 }
2873
2874 "label_join" => {
2875 self.ctx.use_tsid = false;
2876 let (concat_expr, dst_label) = Self::build_concat_labels_expr(
2877 &mut other_input_exprs,
2878 &self.ctx,
2879 query_engine_state,
2880 )?;
2881
2882 for value in &self.ctx.field_columns {
2884 if *value != dst_label {
2885 let expr = DfExpr::Column(Column::from_name(value));
2886 exprs.push(expr);
2887 }
2888 }
2889
2890 self.ctx.tag_columns.retain(|tag| *tag != dst_label);
2892 new_tags.push(dst_label);
2893 exprs.push(concat_expr);
2895
2896 ScalarFunc::GeneratedExpr
2897 }
2898 "label_replace" => {
2899 self.ctx.use_tsid = false;
2900 if let Some((replace_expr, dst_label)) = self
2901 .build_regexp_replace_label_expr(&mut other_input_exprs, query_engine_state)?
2902 {
2903 for value in &self.ctx.field_columns {
2905 if *value != dst_label {
2906 let expr = DfExpr::Column(Column::from_name(value));
2907 exprs.push(expr);
2908 }
2909 }
2910
2911 ensure!(
2912 !self.ctx.tag_columns.contains(&dst_label),
2913 SameLabelSetSnafu
2914 );
2915 new_tags.push(dst_label);
2916 exprs.push(replace_expr);
2918 } else {
2919 for value in &self.ctx.field_columns {
2921 let expr = DfExpr::Column(Column::from_name(value));
2922 exprs.push(expr);
2923 }
2924 }
2925
2926 ScalarFunc::GeneratedExpr
2927 }
2928 "sort" | "sort_desc" | "sort_by_label" | "sort_by_label_desc" | "timestamp" => {
2929 for value in &self.ctx.field_columns {
2932 let expr = DfExpr::Column(Column::from_name(value));
2933 exprs.push(expr);
2934 }
2935
2936 ScalarFunc::GeneratedExpr
2937 }
2938 "round" => {
2939 if other_input_exprs.is_empty() {
2940 other_input_exprs.push_front(0.0f64.lit());
2941 }
2942 ScalarFunc::DataFusionUdf(Arc::new(Round::scalar_udf()))
2943 }
2944 "rad" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::radians()),
2945 "deg" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::degrees()),
2946 "sgn" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::signum()),
2947 "pi" => {
2948 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
2950 func: datafusion::functions::math::pi(),
2951 args: vec![],
2952 });
2953 exprs.push(fn_expr);
2954
2955 ScalarFunc::GeneratedExpr
2956 }
2957 _ => {
2958 if let Some(f) = query_engine_state
2959 .session_state()
2960 .scalar_functions()
2961 .get(func.name)
2962 {
2963 ScalarFunc::DataFusionBuiltin(f.clone())
2964 } else if let Some(factory) = query_engine_state.scalar_function(func.name) {
2965 let func_state = query_engine_state.function_state();
2966 let query_ctx = self.table_provider.query_ctx();
2967
2968 ScalarFunc::DataFusionUdf(Arc::new(factory.provide(FunctionContext {
2969 state: func_state,
2970 query_ctx: query_ctx.clone(),
2971 })))
2972 } else if let Some(f) = datafusion_functions::math::functions()
2973 .iter()
2974 .find(|f| f.name() == func.name)
2975 {
2976 ScalarFunc::DataFusionUdf(f.clone())
2977 } else {
2978 return UnsupportedExprSnafu {
2979 name: func.name.to_string(),
2980 }
2981 .fail();
2982 }
2983 }
2984 };
2985
2986 for value in &self.ctx.field_columns {
2987 let col_expr = DfExpr::Column(Column::from_name(value));
2988
2989 match scalar_func.clone() {
2990 ScalarFunc::DataFusionBuiltin(func) => {
2991 other_input_exprs.insert(field_column_pos, col_expr);
2992 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
2993 func,
2994 args: other_input_exprs.clone().into(),
2995 });
2996 exprs.push(fn_expr);
2997 let _ = other_input_exprs.remove(field_column_pos);
2998 }
2999 ScalarFunc::DataFusionUdf(func) => {
3000 let args = itertools::chain!(
3001 other_input_exprs.iter().take(field_column_pos).cloned(),
3002 std::iter::once(col_expr),
3003 other_input_exprs.iter().skip(field_column_pos).cloned()
3004 )
3005 .collect_vec();
3006 exprs.push(DfExpr::ScalarFunction(ScalarFunction { func, args }))
3007 }
3008 ScalarFunc::Udf(func) => {
3009 let ts_range_expr = DfExpr::Column(Column::from_name(
3010 RangeManipulate::build_timestamp_range_name(
3011 self.ctx.time_index_column.as_ref().unwrap(),
3012 ),
3013 ));
3014 other_input_exprs.insert(field_column_pos, ts_range_expr);
3015 other_input_exprs.insert(field_column_pos + 1, col_expr);
3016 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
3017 func,
3018 args: other_input_exprs.clone().into(),
3019 });
3020 exprs.push(fn_expr);
3021 let _ = other_input_exprs.remove(field_column_pos + 1);
3022 let _ = other_input_exprs.remove(field_column_pos);
3023 }
3024 ScalarFunc::ExtrapolateUdf(func, range_length) => {
3025 let ts_range_expr = DfExpr::Column(Column::from_name(
3026 RangeManipulate::build_timestamp_range_name(
3027 self.ctx.time_index_column.as_ref().unwrap(),
3028 ),
3029 ));
3030 other_input_exprs.insert(field_column_pos, ts_range_expr);
3031 other_input_exprs.insert(field_column_pos + 1, col_expr);
3032 other_input_exprs
3033 .insert(field_column_pos + 2, self.create_time_index_column_expr()?);
3034 other_input_exprs.push_back(lit(range_length));
3035 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
3036 func,
3037 args: other_input_exprs.clone().into(),
3038 });
3039 exprs.push(fn_expr);
3040 let _ = other_input_exprs.pop_back();
3041 let _ = other_input_exprs.remove(field_column_pos + 2);
3042 let _ = other_input_exprs.remove(field_column_pos + 1);
3043 let _ = other_input_exprs.remove(field_column_pos);
3044 }
3045 ScalarFunc::GeneratedExpr => {}
3046 }
3047 }
3048
3049 if !matches!(func.name, "label_join" | "label_replace") {
3053 let mut new_field_columns = Vec::with_capacity(exprs.len());
3054
3055 exprs = exprs
3056 .into_iter()
3057 .map(|expr| {
3058 let display_name = expr.schema_name().to_string();
3059 new_field_columns.push(display_name.clone());
3060 Ok(expr.alias(display_name))
3061 })
3062 .collect::<std::result::Result<Vec<_>, _>>()
3063 .context(DataFusionPlanningSnafu)?;
3064
3065 self.ctx.field_columns = new_field_columns;
3066 }
3067
3068 Ok((exprs, new_tags))
3069 }
3070
3071 fn validate_label_name(label_name: &str) -> Result<()> {
3075 if label_name.starts_with("__") {
3077 return InvalidDestinationLabelNameSnafu { label_name }.fail();
3078 }
3079 if !LABEL_NAME_REGEX.is_match(label_name) {
3081 return InvalidDestinationLabelNameSnafu { label_name }.fail();
3082 }
3083
3084 Ok(())
3085 }
3086
3087 fn build_regexp_replace_label_expr(
3089 &self,
3090 other_input_exprs: &mut VecDeque<DfExpr>,
3091 query_engine_state: &QueryEngineState,
3092 ) -> Result<Option<(DfExpr, String)>> {
3093 let dst_label = match other_input_exprs.pop_front() {
3095 Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
3096 other => UnexpectedPlanExprSnafu {
3097 desc: format!("expected dst_label string literal, but found {:?}", other),
3098 }
3099 .fail()?,
3100 };
3101
3102 Self::validate_label_name(&dst_label)?;
3104 let replacement = match other_input_exprs.pop_front() {
3105 Some(DfExpr::Literal(ScalarValue::Utf8(Some(r)), _)) => r,
3106 other => UnexpectedPlanExprSnafu {
3107 desc: format!("expected replacement string literal, but found {:?}", other),
3108 }
3109 .fail()?,
3110 };
3111 let src_label = match other_input_exprs.pop_front() {
3112 Some(DfExpr::Literal(ScalarValue::Utf8(Some(s)), None)) => s,
3113 other => UnexpectedPlanExprSnafu {
3114 desc: format!("expected src_label string literal, but found {:?}", other),
3115 }
3116 .fail()?,
3117 };
3118
3119 let regex = match other_input_exprs.pop_front() {
3120 Some(DfExpr::Literal(ScalarValue::Utf8(Some(r)), None)) => r,
3121 other => UnexpectedPlanExprSnafu {
3122 desc: format!("expected regex string literal, but found {:?}", other),
3123 }
3124 .fail()?,
3125 };
3126
3127 regex::Regex::new(®ex).map_err(|_| {
3130 InvalidRegularExpressionSnafu {
3131 regex: regex.clone(),
3132 }
3133 .build()
3134 })?;
3135
3136 if self.ctx.tag_columns.contains(&src_label) && regex.is_empty() {
3138 return Ok(None);
3139 }
3140
3141 if !self.ctx.tag_columns.contains(&src_label) {
3143 if replacement.is_empty() {
3144 return Ok(None);
3146 } else {
3147 return Ok(Some((
3149 lit(replacement).alias(&dst_label),
3151 dst_label,
3152 )));
3153 }
3154 }
3155
3156 let regex = format!("^(?s:{regex})$");
3159
3160 let session_state = query_engine_state.session_state();
3161 let func = session_state
3162 .scalar_functions()
3163 .get("regexp_replace")
3164 .context(UnsupportedExprSnafu {
3165 name: "regexp_replace",
3166 })?;
3167
3168 let args = vec![
3170 if src_label.is_empty() {
3171 DfExpr::Literal(ScalarValue::Utf8(Some(String::new())), None)
3172 } else {
3173 DfExpr::Column(Column::from_name(src_label))
3174 },
3175 DfExpr::Literal(ScalarValue::Utf8(Some(regex)), None),
3176 DfExpr::Literal(ScalarValue::Utf8(Some(replacement)), None),
3177 ];
3178
3179 Ok(Some((
3180 DfExpr::ScalarFunction(ScalarFunction {
3181 func: func.clone(),
3182 args,
3183 })
3184 .alias(&dst_label),
3185 dst_label,
3186 )))
3187 }
3188
3189 fn build_concat_labels_expr(
3191 other_input_exprs: &mut VecDeque<DfExpr>,
3192 ctx: &PromPlannerContext,
3193 query_engine_state: &QueryEngineState,
3194 ) -> Result<(DfExpr, String)> {
3195 let dst_label = match other_input_exprs.pop_front() {
3198 Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
3199 other => UnexpectedPlanExprSnafu {
3200 desc: format!("expected dst_label string literal, but found {:?}", other),
3201 }
3202 .fail()?,
3203 };
3204 let separator = match other_input_exprs.pop_front() {
3205 Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
3206 other => UnexpectedPlanExprSnafu {
3207 desc: format!("expected separator string literal, but found {:?}", other),
3208 }
3209 .fail()?,
3210 };
3211
3212 let available_columns: HashSet<&str> = ctx
3214 .tag_columns
3215 .iter()
3216 .chain(ctx.field_columns.iter())
3217 .chain(ctx.time_index_column.as_ref())
3218 .map(|s| s.as_str())
3219 .collect();
3220
3221 let src_labels = other_input_exprs
3222 .iter()
3223 .map(|expr| {
3224 match expr {
3226 DfExpr::Literal(ScalarValue::Utf8(Some(label)), None) => {
3227 if label.is_empty() {
3228 Ok(DfExpr::Literal(ScalarValue::Null, None))
3229 } else if available_columns.contains(label.as_str()) {
3230 Ok(DfExpr::Column(Column::from_name(label)))
3232 } else {
3233 Ok(DfExpr::Literal(ScalarValue::Null, None))
3235 }
3236 }
3237 other => UnexpectedPlanExprSnafu {
3238 desc: format!(
3239 "expected source label string literal, but found {:?}",
3240 other
3241 ),
3242 }
3243 .fail(),
3244 }
3245 })
3246 .collect::<Result<Vec<_>>>()?;
3247 ensure!(
3248 !src_labels.is_empty(),
3249 FunctionInvalidArgumentSnafu {
3250 fn_name: "label_join"
3251 }
3252 );
3253
3254 let session_state = query_engine_state.session_state();
3255 let func = session_state
3256 .scalar_functions()
3257 .get("concat_ws")
3258 .context(UnsupportedExprSnafu { name: "concat_ws" })?;
3259
3260 let mut args = Vec::with_capacity(1 + src_labels.len());
3262 args.push(DfExpr::Literal(ScalarValue::Utf8(Some(separator)), None));
3263 args.extend(src_labels);
3264
3265 Ok((
3266 DfExpr::ScalarFunction(ScalarFunction {
3267 func: func.clone(),
3268 args,
3269 })
3270 .alias(&dst_label),
3271 dst_label,
3272 ))
3273 }
3274
3275 fn create_time_index_column_expr(&self) -> Result<DfExpr> {
3276 Ok(DfExpr::Column(Column::from_name(
3277 self.ctx
3278 .time_index_column
3279 .clone()
3280 .with_context(|| TimeIndexNotFoundSnafu { table: "unknown" })?,
3281 )))
3282 }
3283
3284 fn create_tag_column_exprs(&self) -> Result<Vec<DfExpr>> {
3285 let mut result = Vec::with_capacity(self.ctx.tag_columns.len());
3286 for tag in &self.ctx.tag_columns {
3287 let expr = DfExpr::Column(Column::from_name(tag));
3288 result.push(expr);
3289 }
3290 Ok(result)
3291 }
3292
3293 fn create_field_column_exprs(&self) -> Result<Vec<DfExpr>> {
3294 let mut result = Vec::with_capacity(self.ctx.field_columns.len());
3295 for field in &self.ctx.field_columns {
3296 let expr = DfExpr::Column(Column::from_name(field));
3297 result.push(expr);
3298 }
3299 Ok(result)
3300 }
3301
3302 fn create_tag_and_time_index_column_sort_exprs(&self) -> Result<Vec<SortExpr>> {
3303 let mut result = self
3304 .ctx
3305 .tag_columns
3306 .iter()
3307 .map(|col| DfExpr::Column(Column::from_name(col)).sort(true, true))
3308 .collect::<Vec<_>>();
3309 result.push(self.create_time_index_column_expr()?.sort(true, true));
3310 Ok(result)
3311 }
3312
3313 fn create_field_columns_sort_exprs(&self, asc: bool) -> Vec<SortExpr> {
3314 self.ctx
3315 .field_columns
3316 .iter()
3317 .map(|col| DfExpr::Column(Column::from_name(col)).sort(asc, true))
3318 .collect::<Vec<_>>()
3319 }
3320
3321 fn create_sort_exprs_by_tags(
3322 func: &str,
3323 tags: Vec<DfExpr>,
3324 asc: bool,
3325 ) -> Result<Vec<SortExpr>> {
3326 ensure!(
3327 !tags.is_empty(),
3328 FunctionInvalidArgumentSnafu { fn_name: func }
3329 );
3330
3331 tags.iter()
3332 .map(|col| match col {
3333 DfExpr::Literal(ScalarValue::Utf8(Some(label)), _) => {
3334 Ok(DfExpr::Column(Column::from_name(label)).sort(asc, false))
3335 }
3336 other => UnexpectedPlanExprSnafu {
3337 desc: format!("expected label string literal, but found {:?}", other),
3338 }
3339 .fail(),
3340 })
3341 .collect::<Result<Vec<_>>>()
3342 }
3343
3344 fn create_empty_values_filter_expr(&self) -> Result<DfExpr> {
3345 let mut exprs = Vec::with_capacity(self.ctx.field_columns.len());
3346 for value in &self.ctx.field_columns {
3347 let expr = DfExpr::Column(Column::from_name(value)).is_not_null();
3348 exprs.push(expr);
3349 }
3350
3351 conjunction(exprs).with_context(|| ValueNotFoundSnafu {
3356 table: self
3357 .table_ref()
3358 .map(|t| t.to_quoted_string())
3359 .unwrap_or_else(|_| "unknown".to_string()),
3360 })
3361 }
3362
3363 fn create_aggregate_exprs(
3379 &mut self,
3380 op: TokenType,
3381 param: &Option<Box<PromExpr>>,
3382 input_plan: &LogicalPlan,
3383 ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
3384 let mut non_col_args = Vec::new();
3385 let is_group_agg = op.id() == token::T_GROUP;
3386 if is_group_agg {
3387 ensure!(
3388 self.ctx.field_columns.len() == 1,
3389 MultiFieldsNotSupportedSnafu {
3390 operator: "group()"
3391 }
3392 );
3393 }
3394 let aggr = match op.id() {
3395 token::T_SUM => sum_udaf(),
3396 token::T_QUANTILE => {
3397 let q =
3398 Self::get_param_as_literal_expr(param, Some(op), Some(ArrowDataType::Float64))?;
3399 non_col_args.push(q);
3400 quantile_udaf()
3401 }
3402 token::T_AVG => avg_udaf(),
3403 token::T_COUNT_VALUES | token::T_COUNT => count_udaf(),
3404 token::T_MIN => min_udaf(),
3405 token::T_MAX => max_udaf(),
3406 token::T_GROUP => max_udaf(),
3409 token::T_STDDEV => stddev_pop_udaf(),
3410 token::T_STDVAR => var_pop_udaf(),
3411 token::T_TOPK | token::T_BOTTOMK => UnsupportedExprSnafu {
3412 name: format!("{op:?}"),
3413 }
3414 .fail()?,
3415 _ => UnexpectedTokenSnafu { token: op }.fail()?,
3416 };
3417
3418 let exprs: Vec<DfExpr> = self
3420 .ctx
3421 .field_columns
3422 .iter()
3423 .map(|col| {
3424 if is_group_agg {
3425 aggr.call(vec![lit(1_f64)])
3426 } else {
3427 non_col_args.push(DfExpr::Column(Column::from_name(col)));
3428 let expr = aggr.call(non_col_args.clone());
3429 non_col_args.pop();
3430 expr
3431 }
3432 })
3433 .collect::<Vec<_>>();
3434
3435 let prev_field_exprs = if op.id() == token::T_COUNT_VALUES {
3437 let prev_field_exprs: Vec<_> = self
3438 .ctx
3439 .field_columns
3440 .iter()
3441 .map(|col| DfExpr::Column(Column::from_name(col)))
3442 .collect();
3443
3444 ensure!(
3445 self.ctx.field_columns.len() == 1,
3446 UnsupportedExprSnafu {
3447 name: "count_values on multi-value input"
3448 }
3449 );
3450
3451 prev_field_exprs
3452 } else {
3453 vec![]
3454 };
3455
3456 let mut new_field_columns = Vec::with_capacity(self.ctx.field_columns.len());
3458
3459 let normalized_exprs =
3460 normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
3461 for expr in normalized_exprs {
3462 new_field_columns.push(expr.schema_name().to_string());
3463 }
3464 self.ctx.field_columns = new_field_columns;
3465
3466 Ok((exprs, prev_field_exprs))
3467 }
3468
3469 fn get_param_value_as_str(op: TokenType, param: &Option<Box<PromExpr>>) -> Result<&str> {
3470 let param = param
3471 .as_deref()
3472 .with_context(|| FunctionInvalidArgumentSnafu {
3473 fn_name: op.to_string(),
3474 })?;
3475 let PromExpr::StringLiteral(StringLiteral { val }) = param else {
3476 return FunctionInvalidArgumentSnafu {
3477 fn_name: op.to_string(),
3478 }
3479 .fail();
3480 };
3481
3482 Ok(val)
3483 }
3484
3485 fn get_param_as_literal_expr(
3486 param: &Option<Box<PromExpr>>,
3487 op: Option<TokenType>,
3488 expected_type: Option<ArrowDataType>,
3489 ) -> Result<DfExpr> {
3490 let prom_param = param.as_deref().with_context(|| {
3491 if let Some(op) = op {
3492 FunctionInvalidArgumentSnafu {
3493 fn_name: op.to_string(),
3494 }
3495 } else {
3496 FunctionInvalidArgumentSnafu {
3497 fn_name: "unknown".to_string(),
3498 }
3499 }
3500 })?;
3501
3502 let expr = Self::try_build_literal_expr(prom_param).with_context(|| {
3503 if let Some(op) = op {
3504 FunctionInvalidArgumentSnafu {
3505 fn_name: op.to_string(),
3506 }
3507 } else {
3508 FunctionInvalidArgumentSnafu {
3509 fn_name: "unknown".to_string(),
3510 }
3511 }
3512 })?;
3513
3514 if let Some(expected_type) = expected_type {
3516 let expr_type = expr
3518 .get_type(&DFSchema::empty())
3519 .context(DataFusionPlanningSnafu)?;
3520 if expected_type != expr_type {
3521 return FunctionInvalidArgumentSnafu {
3522 fn_name: format!("expected {expected_type:?}, but found {expr_type:?}"),
3523 }
3524 .fail();
3525 }
3526 }
3527
3528 Ok(expr)
3529 }
3530
3531 fn create_window_exprs(
3534 &mut self,
3535 op: TokenType,
3536 group_exprs: Vec<DfExpr>,
3537 input_plan: &LogicalPlan,
3538 ) -> Result<Vec<DfExpr>> {
3539 ensure!(
3540 self.ctx.field_columns.len() == 1,
3541 UnsupportedExprSnafu {
3542 name: "topk or bottomk on multi-value input"
3543 }
3544 );
3545
3546 assert!(matches!(op.id(), token::T_TOPK | token::T_BOTTOMK));
3547
3548 let asc = matches!(op.id(), token::T_BOTTOMK);
3549
3550 let tag_sort_exprs = self
3551 .create_tag_column_exprs()?
3552 .into_iter()
3553 .map(|expr| expr.sort(asc, true));
3554
3555 let exprs: Vec<DfExpr> = self
3557 .ctx
3558 .field_columns
3559 .iter()
3560 .map(|col| {
3561 let mut sort_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 1);
3562 sort_exprs.push(DfExpr::Column(Column::from(col)).sort(asc, true));
3564 sort_exprs.extend(tag_sort_exprs.clone());
3567
3568 DfExpr::WindowFunction(Box::new(WindowFunction {
3569 fun: WindowFunctionDefinition::WindowUDF(Arc::new(RowNumber::new().into())),
3570 params: WindowFunctionParams {
3571 args: vec![],
3572 partition_by: group_exprs.clone(),
3573 order_by: sort_exprs,
3574 window_frame: WindowFrame::new(Some(true)),
3575 null_treatment: None,
3576 distinct: false,
3577 filter: None,
3578 },
3579 }))
3580 })
3581 .collect();
3582
3583 let normalized_exprs =
3584 normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
3585 Ok(normalized_exprs)
3586 }
3587
3588 #[deprecated(
3590 note = "use `Self::get_param_as_literal_expr` instead. This is only for `create_histogram_plan`"
3591 )]
3592 fn try_build_float_literal(expr: &PromExpr) -> Option<f64> {
3593 match expr {
3594 PromExpr::NumberLiteral(NumberLiteral { val }) => Some(*val),
3595 PromExpr::Paren(ParenExpr { expr }) => Self::try_build_float_literal(expr),
3596 PromExpr::Unary(UnaryExpr { expr, .. }) => {
3597 Self::try_build_float_literal(expr).map(|f| -f)
3598 }
3599 PromExpr::StringLiteral(_)
3600 | PromExpr::Binary(_)
3601 | PromExpr::VectorSelector(_)
3602 | PromExpr::MatrixSelector(_)
3603 | PromExpr::Call(_)
3604 | PromExpr::Extension(_)
3605 | PromExpr::Aggregate(_)
3606 | PromExpr::Subquery(_) => None,
3607 }
3608 }
3609
3610 async fn create_histogram_plan(
3612 &mut self,
3613 args: &PromFunctionArgs,
3614 query_engine_state: &QueryEngineState,
3615 ) -> Result<LogicalPlan> {
3616 if args.args.len() != 2 {
3617 return FunctionInvalidArgumentSnafu {
3618 fn_name: SPECIAL_HISTOGRAM_QUANTILE.to_string(),
3619 }
3620 .fail();
3621 }
3622 #[allow(deprecated)]
3623 let phi = Self::try_build_float_literal(&args.args[0]).with_context(|| {
3624 FunctionInvalidArgumentSnafu {
3625 fn_name: SPECIAL_HISTOGRAM_QUANTILE.to_string(),
3626 }
3627 })?;
3628
3629 let input = args.args[1].as_ref().clone();
3630 let input_plan = self.prom_expr_to_plan(&input, query_engine_state).await?;
3631 let input_plan = self.strip_tsid_column(input_plan)?;
3635 self.ctx.use_tsid = false;
3636
3637 if !self.ctx.has_le_tag() {
3638 return Ok(LogicalPlan::EmptyRelation(
3641 datafusion::logical_expr::EmptyRelation {
3642 produce_one_row: false,
3643 schema: Arc::new(DFSchema::empty()),
3644 },
3645 ));
3646 }
3647 let time_index_column =
3648 self.ctx
3649 .time_index_column
3650 .clone()
3651 .with_context(|| TimeIndexNotFoundSnafu {
3652 table: self.ctx.table_name.clone().unwrap_or_default(),
3653 })?;
3654 let field_column = self
3656 .ctx
3657 .field_columns
3658 .first()
3659 .with_context(|| FunctionInvalidArgumentSnafu {
3660 fn_name: SPECIAL_HISTOGRAM_QUANTILE.to_string(),
3661 })?
3662 .clone();
3663 self.ctx.tag_columns.retain(|col| col != LE_COLUMN_NAME);
3665
3666 Ok(LogicalPlan::Extension(Extension {
3667 node: Arc::new(
3668 HistogramFold::new(
3669 LE_COLUMN_NAME.to_string(),
3670 field_column,
3671 time_index_column,
3672 phi,
3673 input_plan,
3674 )
3675 .context(DataFusionPlanningSnafu)?,
3676 ),
3677 }))
3678 }
3679
3680 async fn create_vector_plan(&mut self, args: &PromFunctionArgs) -> Result<LogicalPlan> {
3682 if args.args.len() != 1 {
3683 return FunctionInvalidArgumentSnafu {
3684 fn_name: SPECIAL_VECTOR_FUNCTION.to_string(),
3685 }
3686 .fail();
3687 }
3688 let lit = Self::get_param_as_literal_expr(&Some(args.args[0].clone()), None, None)?;
3689
3690 self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
3692 self.ctx.reset_table_name_and_schema();
3693 self.ctx.tag_columns = vec![];
3694 self.ctx.field_columns = vec![greptime_value().to_string()];
3695 Ok(LogicalPlan::Extension(Extension {
3696 node: Arc::new(
3697 EmptyMetric::new(
3698 self.ctx.start,
3699 self.ctx.end,
3700 self.ctx.interval,
3701 SPECIAL_TIME_FUNCTION.to_string(),
3702 greptime_value().to_string(),
3703 Some(lit),
3704 )
3705 .context(DataFusionPlanningSnafu)?,
3706 ),
3707 }))
3708 }
3709
3710 async fn create_scalar_plan(
3712 &mut self,
3713 args: &PromFunctionArgs,
3714 query_engine_state: &QueryEngineState,
3715 ) -> Result<LogicalPlan> {
3716 ensure!(
3717 args.len() == 1,
3718 FunctionInvalidArgumentSnafu {
3719 fn_name: SCALAR_FUNCTION
3720 }
3721 );
3722 let input = self
3723 .prom_expr_to_plan(&args.args[0], query_engine_state)
3724 .await?;
3725 ensure!(
3726 self.ctx.field_columns.len() == 1,
3727 MultiFieldsNotSupportedSnafu {
3728 operator: SCALAR_FUNCTION
3729 },
3730 );
3731 let scalar_plan = LogicalPlan::Extension(Extension {
3732 node: Arc::new(
3733 ScalarCalculate::new(
3734 self.ctx.start,
3735 self.ctx.end,
3736 self.ctx.interval,
3737 input,
3738 self.ctx.time_index_column.as_ref().unwrap(),
3739 &self.ctx.tag_columns,
3740 &self.ctx.field_columns[0],
3741 self.ctx.table_name.as_deref(),
3742 )
3743 .context(PromqlPlanNodeSnafu)?,
3744 ),
3745 });
3746 self.ctx.tag_columns.clear();
3748 self.ctx.field_columns.clear();
3749 self.ctx
3750 .field_columns
3751 .push(scalar_plan.schema().field(1).name().clone());
3752 Ok(scalar_plan)
3753 }
3754
3755 async fn create_absent_plan(
3757 &mut self,
3758 args: &PromFunctionArgs,
3759 query_engine_state: &QueryEngineState,
3760 ) -> Result<LogicalPlan> {
3761 if args.args.len() != 1 {
3762 return FunctionInvalidArgumentSnafu {
3763 fn_name: SPECIAL_ABSENT_FUNCTION.to_string(),
3764 }
3765 .fail();
3766 }
3767 let input = self
3768 .prom_expr_to_plan(&args.args[0], query_engine_state)
3769 .await?;
3770
3771 let time_index_expr = self.create_time_index_column_expr()?;
3772 let first_field_expr =
3773 self.create_field_column_exprs()?
3774 .pop()
3775 .with_context(|| ValueNotFoundSnafu {
3776 table: self.ctx.table_name.clone().unwrap_or_default(),
3777 })?;
3778 let first_value_expr = first_value(first_field_expr, vec![]);
3779
3780 let ordered_aggregated_input = LogicalPlanBuilder::from(input)
3781 .aggregate(
3782 vec![time_index_expr.clone()],
3783 vec![first_value_expr.clone()],
3784 )
3785 .context(DataFusionPlanningSnafu)?
3786 .sort(vec![time_index_expr.sort(true, false)])
3787 .context(DataFusionPlanningSnafu)?
3788 .build()
3789 .context(DataFusionPlanningSnafu)?;
3790
3791 let fake_labels = self
3792 .ctx
3793 .selector_matcher
3794 .iter()
3795 .filter_map(|matcher| match matcher.op {
3796 MatchOp::Equal => Some((matcher.name.clone(), matcher.value.clone())),
3797 _ => None,
3798 })
3799 .collect::<Vec<_>>();
3800
3801 let absent_plan = LogicalPlan::Extension(Extension {
3803 node: Arc::new(
3804 Absent::try_new(
3805 self.ctx.start,
3806 self.ctx.end,
3807 self.ctx.interval,
3808 self.ctx.time_index_column.as_ref().unwrap().clone(),
3809 self.ctx.field_columns[0].clone(),
3810 fake_labels,
3811 ordered_aggregated_input,
3812 )
3813 .context(DataFusionPlanningSnafu)?,
3814 ),
3815 });
3816
3817 Ok(absent_plan)
3818 }
3819
3820 fn try_build_literal_expr(expr: &PromExpr) -> Option<DfExpr> {
3823 match expr {
3824 PromExpr::NumberLiteral(NumberLiteral { val }) => Some(val.lit()),
3825 PromExpr::StringLiteral(StringLiteral { val }) => Some(val.lit()),
3826 PromExpr::VectorSelector(_)
3827 | PromExpr::MatrixSelector(_)
3828 | PromExpr::Extension(_)
3829 | PromExpr::Aggregate(_)
3830 | PromExpr::Subquery(_) => None,
3831 PromExpr::Call(Call { func, .. }) => {
3832 if func.name == SPECIAL_TIME_FUNCTION {
3833 None
3836 } else {
3837 None
3838 }
3839 }
3840 PromExpr::Paren(ParenExpr { expr }) => Self::try_build_literal_expr(expr),
3841 PromExpr::Unary(UnaryExpr { expr, .. }) => Self::try_build_literal_expr(expr),
3843 PromExpr::Binary(PromBinaryExpr {
3844 lhs,
3845 rhs,
3846 op,
3847 modifier,
3848 }) => {
3849 let lhs = Self::try_build_literal_expr(lhs)?;
3850 let rhs = Self::try_build_literal_expr(rhs)?;
3851 let is_comparison_op = Self::is_token_a_comparison_op(*op);
3852 let expr_builder = Self::prom_token_to_binary_expr_builder(*op).ok()?;
3853 let expr = expr_builder(lhs, rhs).ok()?;
3854
3855 let should_return_bool = if let Some(m) = modifier {
3856 m.return_bool
3857 } else {
3858 false
3859 };
3860 if is_comparison_op && should_return_bool {
3861 Some(DfExpr::Cast(Cast {
3862 expr: Box::new(expr),
3863 data_type: ArrowDataType::Float64,
3864 }))
3865 } else {
3866 Some(expr)
3867 }
3868 }
3869 }
3870 }
3871
3872 fn try_build_special_time_expr_with_context(&self, expr: &PromExpr) -> Option<DfExpr> {
3873 match expr {
3874 PromExpr::Call(Call { func, .. }) => {
3875 if func.name == SPECIAL_TIME_FUNCTION
3876 && let Some(time_index_col) = self.ctx.time_index_column.as_ref()
3877 {
3878 Some(build_special_time_expr(time_index_col))
3879 } else {
3880 None
3881 }
3882 }
3883 _ => None,
3884 }
3885 }
3886
3887 #[allow(clippy::type_complexity)]
3890 fn prom_token_to_binary_expr_builder(
3891 token: TokenType,
3892 ) -> Result<Box<dyn Fn(DfExpr, DfExpr) -> Result<DfExpr>>> {
3893 let cast_float = |expr| {
3894 if matches!(
3895 &expr,
3896 DfExpr::Cast(Cast {
3897 data_type: ArrowDataType::Float64,
3898 ..
3899 })
3900 ) || matches!(&expr, DfExpr::Literal(ScalarValue::Float64(_), _))
3901 {
3902 expr
3903 } else {
3904 DfExpr::Cast(Cast {
3905 expr: Box::new(expr),
3906 data_type: ArrowDataType::Float64,
3907 })
3908 }
3909 };
3910 match token.id() {
3911 token::T_ADD => Ok(Box::new(move |lhs, rhs| {
3912 Ok(cast_float(lhs) + cast_float(rhs))
3913 })),
3914 token::T_SUB => Ok(Box::new(move |lhs, rhs| {
3915 Ok(cast_float(lhs) - cast_float(rhs))
3916 })),
3917 token::T_MUL => Ok(Box::new(move |lhs, rhs| {
3918 Ok(cast_float(lhs) * cast_float(rhs))
3919 })),
3920 token::T_DIV => Ok(Box::new(move |lhs, rhs| {
3921 Ok(cast_float(lhs) / cast_float(rhs))
3922 })),
3923 token::T_MOD => Ok(Box::new(move |lhs: DfExpr, rhs| {
3924 Ok(cast_float(lhs) % cast_float(rhs))
3925 })),
3926 token::T_EQLC => Ok(Box::new(|lhs, rhs| Ok(lhs.eq(rhs)))),
3927 token::T_NEQ => Ok(Box::new(|lhs, rhs| Ok(lhs.not_eq(rhs)))),
3928 token::T_GTR => Ok(Box::new(|lhs, rhs| Ok(lhs.gt(rhs)))),
3929 token::T_LSS => Ok(Box::new(|lhs, rhs| Ok(lhs.lt(rhs)))),
3930 token::T_GTE => Ok(Box::new(|lhs, rhs| Ok(lhs.gt_eq(rhs)))),
3931 token::T_LTE => Ok(Box::new(|lhs, rhs| Ok(lhs.lt_eq(rhs)))),
3932 token::T_POW => Ok(Box::new(move |lhs, rhs| {
3933 Ok(DfExpr::ScalarFunction(ScalarFunction {
3934 func: datafusion_functions::math::power(),
3935 args: vec![cast_float(lhs), cast_float(rhs)],
3936 }))
3937 })),
3938 token::T_ATAN2 => Ok(Box::new(move |lhs, rhs| {
3939 Ok(DfExpr::ScalarFunction(ScalarFunction {
3940 func: datafusion_functions::math::atan2(),
3941 args: vec![cast_float(lhs), cast_float(rhs)],
3942 }))
3943 })),
3944 _ => UnexpectedTokenSnafu { token }.fail(),
3945 }
3946 }
3947
3948 fn is_token_a_comparison_op(token: TokenType) -> bool {
3950 matches!(
3951 token.id(),
3952 token::T_EQLC
3953 | token::T_NEQ
3954 | token::T_GTR
3955 | token::T_LSS
3956 | token::T_GTE
3957 | token::T_LTE
3958 )
3959 }
3960
3961 fn is_token_a_set_op(token: TokenType) -> bool {
3963 matches!(
3964 token.id(),
3965 token::T_LAND | token::T_LOR | token::T_LUNLESS )
3969 }
3970
3971 fn align_binary_field_columns<'a>(
3972 left_field_columns: &'a [String],
3973 right_field_columns: &'a [String],
3974 ) -> (Vec<String>, Vec<(&'a String, &'a String)>) {
3975 let field_pairs = left_field_columns
3976 .iter()
3977 .zip(right_field_columns.iter())
3978 .collect::<Vec<_>>();
3979 let output_field_columns = field_pairs
3980 .iter()
3981 .map(|(left_col_name, _)| (*left_col_name).clone())
3982 .collect();
3983 (output_field_columns, field_pairs)
3984 }
3985
3986 fn plan_has_tsid_column(plan: &LogicalPlan) -> bool {
3987 plan.schema()
3988 .fields()
3989 .iter()
3990 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
3991 }
3992
3993 fn optional_tsid_projection(
3994 schema: &DFSchemaRef,
3995 table_ref: Option<&TableReference>,
3996 keep_tsid: bool,
3997 ) -> Option<DfExpr> {
3998 keep_tsid.then_some(()).and_then(|_| {
3999 schema
4000 .qualified_field_with_name(table_ref, DATA_SCHEMA_TSID_COLUMN_NAME)
4001 .ok()
4002 .map(|field| DfExpr::Column(field.into()))
4003 })
4004 }
4005
4006 fn binary_join_key_columns(
4007 &self,
4008 left_schema: &DFSchemaRef,
4009 right_schema: &DFSchemaRef,
4010 left_context: &PromPlannerContext,
4011 right_context: &PromPlannerContext,
4012 only_join_time_index: bool,
4013 modifier: &Option<BinModifier>,
4014 ) -> Result<(BTreeSet<String>, BTreeSet<String>, bool)> {
4015 let has_tsid = |schema: &DFSchemaRef| {
4016 schema
4017 .fields()
4018 .iter()
4019 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
4020 };
4021 let use_tsid_join = !only_join_time_index
4022 && self.binary_modifier_preserves_tsid_join_key(left_context, right_context, modifier)
4023 && left_context.use_tsid
4024 && right_context.use_tsid
4025 && has_tsid(left_schema)
4026 && has_tsid(right_schema);
4027
4028 let (mut left_tag_columns, mut right_tag_columns) = if use_tsid_join {
4029 (
4030 BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]),
4031 BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]),
4032 )
4033 } else {
4034 if only_join_time_index {
4035 (BTreeSet::new(), BTreeSet::new())
4036 } else {
4037 (
4038 left_context
4039 .tag_columns
4040 .iter()
4041 .cloned()
4042 .collect::<BTreeSet<_>>(),
4043 right_context
4044 .tag_columns
4045 .iter()
4046 .cloned()
4047 .collect::<BTreeSet<_>>(),
4048 )
4049 }
4050 };
4051
4052 if !use_tsid_join
4053 && let Some(modifier) = modifier
4054 && let Some(matching) = &modifier.matching
4055 {
4056 match matching {
4057 LabelModifier::Include(on) => {
4058 let mask = on.labels.iter().cloned().collect::<BTreeSet<_>>();
4059 left_tag_columns = left_tag_columns.intersection(&mask).cloned().collect();
4060 right_tag_columns = right_tag_columns.intersection(&mask).cloned().collect();
4061 }
4062 LabelModifier::Exclude(ignoring) => {
4063 for label in &ignoring.labels {
4064 let _ = left_tag_columns.remove(label);
4065 let _ = right_tag_columns.remove(label);
4066 }
4067 }
4068 }
4069 }
4070
4071 let force_empty_join =
4072 !use_tsid_join && !only_join_time_index && left_tag_columns != right_tag_columns;
4073 if force_empty_join {
4074 let common_tag_columns = left_tag_columns
4075 .intersection(&right_tag_columns)
4076 .cloned()
4077 .collect::<BTreeSet<_>>();
4078 left_tag_columns = common_tag_columns.clone();
4079 right_tag_columns = common_tag_columns;
4080 }
4081
4082 Ok((left_tag_columns, right_tag_columns, force_empty_join))
4083 }
4084
4085 fn binary_modifier_preserves_tsid_join_key(
4086 &self,
4087 left_context: &PromPlannerContext,
4088 right_context: &PromPlannerContext,
4089 modifier: &Option<BinModifier>,
4090 ) -> bool {
4091 let Some(modifier) = modifier else {
4092 return true;
4093 };
4094
4095 if !matches!(modifier.card, VectorMatchCardinality::OneToOne) {
4096 return false;
4097 }
4098
4099 match &modifier.matching {
4100 None => true,
4101 Some(LabelModifier::Exclude(ignoring)) => ignoring.labels.iter().all(|label| {
4102 !left_context.tag_columns.contains(label)
4103 && !right_context.tag_columns.contains(label)
4104 }),
4105 Some(LabelModifier::Include(on)) => {
4106 let on_labels = on.labels.iter().cloned().collect::<BTreeSet<_>>();
4107 let left_labels = left_context
4108 .tag_columns
4109 .iter()
4110 .cloned()
4111 .collect::<BTreeSet<_>>();
4112 let right_labels = right_context
4113 .tag_columns
4114 .iter()
4115 .cloned()
4116 .collect::<BTreeSet<_>>();
4117
4118 on_labels == left_labels && on_labels == right_labels
4119 }
4120 }
4121 }
4122
4123 #[allow(clippy::too_many_arguments)]
4126 fn join_on_non_field_columns(
4127 &self,
4128 left: LogicalPlan,
4129 right: LogicalPlan,
4130 left_table_ref: TableReference,
4131 right_table_ref: TableReference,
4132 left_time_index_column: Option<String>,
4133 right_time_index_column: Option<String>,
4134 only_join_time_index: bool,
4135 modifier: &Option<BinModifier>,
4136 left_context: &PromPlannerContext,
4137 right_context: &PromPlannerContext,
4138 ) -> Result<LogicalPlan> {
4139 let (mut left_tag_columns, mut right_tag_columns, force_empty_join) = self
4140 .binary_join_key_columns(
4141 left.schema(),
4142 right.schema(),
4143 left_context,
4144 right_context,
4145 only_join_time_index,
4146 modifier,
4147 )?;
4148
4149 if let (Some(left_time_index_column), Some(right_time_index_column)) =
4151 (left_time_index_column, right_time_index_column)
4152 {
4153 left_tag_columns.insert(left_time_index_column);
4154 right_tag_columns.insert(right_time_index_column);
4155 }
4156
4157 let right = LogicalPlanBuilder::from(right)
4158 .alias(right_table_ref)
4159 .context(DataFusionPlanningSnafu)?
4160 .build()
4161 .context(DataFusionPlanningSnafu)?;
4162
4163 LogicalPlanBuilder::from(left)
4165 .alias(left_table_ref)
4166 .context(DataFusionPlanningSnafu)?
4167 .join_detailed(
4168 right,
4169 JoinType::Inner,
4170 (
4171 left_tag_columns
4172 .into_iter()
4173 .map(Column::from_name)
4174 .collect::<Vec<_>>(),
4175 right_tag_columns
4176 .into_iter()
4177 .map(Column::from_name)
4178 .collect::<Vec<_>>(),
4179 ),
4180 force_empty_join.then_some(lit(false)),
4181 NullEquality::NullEqualsNull,
4182 )
4183 .context(DataFusionPlanningSnafu)?
4184 .build()
4185 .context(DataFusionPlanningSnafu)
4186 }
4187
4188 fn set_op_on_non_field_columns(
4190 &mut self,
4191 left: LogicalPlan,
4192 mut right: LogicalPlan,
4193 left_context: PromPlannerContext,
4194 right_context: PromPlannerContext,
4195 op: TokenType,
4196 modifier: &Option<BinModifier>,
4197 ) -> Result<LogicalPlan> {
4198 let mut left_tag_col_set = left_context
4199 .tag_columns
4200 .iter()
4201 .cloned()
4202 .collect::<HashSet<_>>();
4203 let mut right_tag_col_set = right_context
4204 .tag_columns
4205 .iter()
4206 .cloned()
4207 .collect::<HashSet<_>>();
4208
4209 if matches!(op.id(), token::T_LOR) {
4210 return self.or_operator(
4211 left,
4212 right,
4213 left_tag_col_set,
4214 right_tag_col_set,
4215 left_context,
4216 right_context,
4217 modifier,
4218 );
4219 }
4220
4221 if let Some(modifier) = modifier {
4223 ensure!(
4225 matches!(
4226 modifier.card,
4227 VectorMatchCardinality::OneToOne | VectorMatchCardinality::ManyToMany
4228 ),
4229 UnsupportedVectorMatchSnafu {
4230 name: modifier.card.clone(),
4231 },
4232 );
4233 if let Some(matching) = &modifier.matching {
4235 match matching {
4236 LabelModifier::Include(on) => {
4238 let mask = on.labels.iter().cloned().collect::<HashSet<_>>();
4239 left_tag_col_set = left_tag_col_set.intersection(&mask).cloned().collect();
4240 right_tag_col_set =
4241 right_tag_col_set.intersection(&mask).cloned().collect();
4242 }
4243 LabelModifier::Exclude(ignoring) => {
4245 for label in &ignoring.labels {
4247 let _ = left_tag_col_set.remove(label);
4248 let _ = right_tag_col_set.remove(label);
4249 }
4250 }
4251 }
4252 }
4253 }
4254 if !matches!(op.id(), token::T_LOR) {
4256 ensure!(
4257 left_tag_col_set == right_tag_col_set,
4258 CombineTableColumnMismatchSnafu {
4259 left: left_tag_col_set.into_iter().collect::<Vec<_>>(),
4260 right: right_tag_col_set.into_iter().collect::<Vec<_>>(),
4261 }
4262 )
4263 };
4264 let left_time_index = left_context.time_index_column.clone().unwrap();
4265 let right_time_index = right_context.time_index_column.clone().unwrap();
4266 let join_keys = left_tag_col_set
4267 .iter()
4268 .cloned()
4269 .chain([left_time_index.clone()])
4270 .collect::<Vec<_>>();
4271 self.ctx.time_index_column = Some(left_time_index.clone());
4272 self.ctx.use_tsid = left_context.use_tsid;
4273
4274 if left_context.time_index_column != right_context.time_index_column {
4276 let right_project_exprs = right
4277 .schema()
4278 .fields()
4279 .iter()
4280 .map(|field| {
4281 if field.name() == &right_time_index {
4282 DfExpr::Column(Column::from_name(&right_time_index)).alias(&left_time_index)
4283 } else {
4284 DfExpr::Column(Column::from_name(field.name()))
4285 }
4286 })
4287 .collect::<Vec<_>>();
4288
4289 right = LogicalPlanBuilder::from(right)
4290 .project(right_project_exprs)
4291 .context(DataFusionPlanningSnafu)?
4292 .build()
4293 .context(DataFusionPlanningSnafu)?;
4294 }
4295
4296 ensure!(
4297 left_context.field_columns.len() == 1,
4298 MultiFieldsNotSupportedSnafu {
4299 operator: "AND operator"
4300 }
4301 );
4302 let left_field_col = left_context.field_columns.first().unwrap();
4305 self.ctx.field_columns = vec![left_field_col.clone()];
4306
4307 match op.id() {
4310 token::T_LAND => LogicalPlanBuilder::from(left)
4311 .distinct()
4312 .context(DataFusionPlanningSnafu)?
4313 .join_detailed(
4314 right,
4315 JoinType::LeftSemi,
4316 (join_keys.clone(), join_keys),
4317 None,
4318 NullEquality::NullEqualsNull,
4319 )
4320 .context(DataFusionPlanningSnafu)?
4321 .build()
4322 .context(DataFusionPlanningSnafu),
4323 token::T_LUNLESS => LogicalPlanBuilder::from(left)
4324 .distinct()
4325 .context(DataFusionPlanningSnafu)?
4326 .join_detailed(
4327 right,
4328 JoinType::LeftAnti,
4329 (join_keys.clone(), join_keys),
4330 None,
4331 NullEquality::NullEqualsNull,
4332 )
4333 .context(DataFusionPlanningSnafu)?
4334 .build()
4335 .context(DataFusionPlanningSnafu),
4336 token::T_LOR => {
4337 unreachable!()
4340 }
4341 _ => UnexpectedTokenSnafu { token: op }.fail(),
4342 }
4343 }
4344
4345 #[allow(clippy::too_many_arguments)]
4347 fn or_operator(
4348 &mut self,
4349 left: LogicalPlan,
4350 right: LogicalPlan,
4351 left_tag_cols_set: HashSet<String>,
4352 right_tag_cols_set: HashSet<String>,
4353 left_context: PromPlannerContext,
4354 right_context: PromPlannerContext,
4355 modifier: &Option<BinModifier>,
4356 ) -> Result<LogicalPlan> {
4357 ensure!(
4359 left_context.field_columns.len() == right_context.field_columns.len(),
4360 CombineTableColumnMismatchSnafu {
4361 left: left_context.field_columns.clone(),
4362 right: right_context.field_columns.clone()
4363 }
4364 );
4365 ensure!(
4366 left_context.field_columns.len() == 1,
4367 MultiFieldsNotSupportedSnafu {
4368 operator: "OR operator"
4369 }
4370 );
4371
4372 let all_tags = left_tag_cols_set
4374 .union(&right_tag_cols_set)
4375 .cloned()
4376 .collect::<HashSet<_>>();
4377 let tags_not_in_left = all_tags
4378 .difference(&left_tag_cols_set)
4379 .cloned()
4380 .collect::<Vec<_>>();
4381 let tags_not_in_right = all_tags
4382 .difference(&right_tag_cols_set)
4383 .cloned()
4384 .collect::<Vec<_>>();
4385 let left_qualifier = left.schema().qualified_field(0).0.cloned();
4386 let right_qualifier = right.schema().qualified_field(0).0.cloned();
4387 let left_qualifier_string = left_qualifier
4388 .as_ref()
4389 .map(|l| l.to_string())
4390 .unwrap_or_default();
4391 let right_qualifier_string = right_qualifier
4392 .as_ref()
4393 .map(|r| r.to_string())
4394 .unwrap_or_default();
4395 let left_time_index_column =
4396 left_context
4397 .time_index_column
4398 .clone()
4399 .with_context(|| TimeIndexNotFoundSnafu {
4400 table: left_qualifier_string.clone(),
4401 })?;
4402 let right_time_index_column =
4403 right_context
4404 .time_index_column
4405 .clone()
4406 .with_context(|| TimeIndexNotFoundSnafu {
4407 table: right_qualifier_string.clone(),
4408 })?;
4409 let left_field_col = left_context.field_columns.first().unwrap();
4411 let right_field_col = right_context.field_columns.first().unwrap();
4412 let left_has_tsid = left
4413 .schema()
4414 .fields()
4415 .iter()
4416 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME);
4417 let right_has_tsid = right
4418 .schema()
4419 .fields()
4420 .iter()
4421 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME);
4422
4423 let mut all_columns_set = left
4425 .schema()
4426 .fields()
4427 .iter()
4428 .chain(right.schema().fields().iter())
4429 .map(|field| field.name().clone())
4430 .collect::<HashSet<_>>();
4431 if !(left_has_tsid && right_has_tsid) {
4434 all_columns_set.remove(DATA_SCHEMA_TSID_COLUMN_NAME);
4435 }
4436 all_columns_set.remove(&left_time_index_column);
4438 all_columns_set.remove(&right_time_index_column);
4439 if left_field_col != right_field_col {
4441 all_columns_set.remove(right_field_col);
4442 }
4443 let mut all_columns = all_columns_set.into_iter().collect::<Vec<_>>();
4444 all_columns.sort_unstable();
4446 all_columns.insert(0, left_time_index_column.clone());
4448
4449 let left_proj_exprs = all_columns.iter().map(|col| {
4451 if tags_not_in_left.contains(col) {
4452 DfExpr::Literal(ScalarValue::Utf8(None), None).alias(col.clone())
4453 } else {
4454 DfExpr::Column(Column::new(None::<String>, col))
4455 }
4456 });
4457 let right_time_index_expr = DfExpr::Column(Column::new(
4458 right_qualifier.clone(),
4459 right_time_index_column,
4460 ))
4461 .alias(left_time_index_column.clone());
4462 let right_qualifier_for_field = right
4465 .schema()
4466 .iter()
4467 .find(|(_, f)| f.name() == right_field_col)
4468 .map(|(q, _)| q)
4469 .with_context(|| ColumnNotFoundSnafu {
4470 col: right_field_col.clone(),
4471 })?
4472 .cloned();
4473
4474 let right_proj_exprs_without_time_index = all_columns.iter().skip(1).map(|col| {
4476 if col == left_field_col && left_field_col != right_field_col {
4478 DfExpr::Column(Column::new(
4480 right_qualifier_for_field.clone(),
4481 right_field_col,
4482 ))
4483 } else if tags_not_in_right.contains(col) {
4484 DfExpr::Literal(ScalarValue::Utf8(None), None).alias(col.clone())
4485 } else {
4486 DfExpr::Column(Column::new(None::<String>, col))
4487 }
4488 });
4489 let right_proj_exprs = [right_time_index_expr]
4490 .into_iter()
4491 .chain(right_proj_exprs_without_time_index);
4492
4493 let left_projected = LogicalPlanBuilder::from(left)
4494 .project(left_proj_exprs)
4495 .context(DataFusionPlanningSnafu)?
4496 .alias(left_qualifier_string.clone())
4497 .context(DataFusionPlanningSnafu)?
4498 .build()
4499 .context(DataFusionPlanningSnafu)?;
4500 let right_projected = LogicalPlanBuilder::from(right)
4501 .project(right_proj_exprs)
4502 .context(DataFusionPlanningSnafu)?
4503 .alias(right_qualifier_string.clone())
4504 .context(DataFusionPlanningSnafu)?
4505 .build()
4506 .context(DataFusionPlanningSnafu)?;
4507
4508 let mut match_columns = if let Some(modifier) = modifier
4510 && let Some(matching) = &modifier.matching
4511 {
4512 match matching {
4513 LabelModifier::Include(on) => on.labels.clone(),
4515 LabelModifier::Exclude(ignoring) => {
4517 let ignoring = ignoring.labels.iter().cloned().collect::<HashSet<_>>();
4518 all_tags.difference(&ignoring).cloned().collect()
4519 }
4520 }
4521 } else {
4522 all_tags.iter().cloned().collect()
4523 };
4524 match_columns.sort_unstable();
4526 let schema = left_projected.schema().clone();
4528 let union_distinct_on = UnionDistinctOn::new(
4529 left_projected,
4530 right_projected,
4531 match_columns,
4532 left_time_index_column.clone(),
4533 schema,
4534 );
4535 let result = LogicalPlan::Extension(Extension {
4536 node: Arc::new(union_distinct_on),
4537 });
4538
4539 self.ctx.time_index_column = Some(left_time_index_column);
4541 self.ctx.tag_columns = all_tags.into_iter().collect();
4542 self.ctx.field_columns = vec![left_field_col.clone()];
4543 self.ctx.use_tsid = left_has_tsid && right_has_tsid;
4544
4545 Ok(result)
4546 }
4547
4548 fn projection_for_each_field_column<F>(
4556 &mut self,
4557 input: LogicalPlan,
4558 name_to_expr: F,
4559 ) -> Result<LogicalPlan>
4560 where
4561 F: FnMut(&String) -> Result<DfExpr>,
4562 {
4563 let table_ref = self.ctx.table_name.clone().map(TableReference::bare);
4564 let non_field_columns_iter = self
4565 .ctx
4566 .tag_columns
4567 .iter()
4568 .chain(self.ctx.time_index_column.iter())
4569 .map(|col| Ok(DfExpr::Column(Column::new(table_ref.clone(), col))));
4570 let tsid_iter =
4571 Self::optional_tsid_projection(input.schema(), table_ref.as_ref(), self.ctx.use_tsid)
4572 .into_iter()
4573 .map(Ok);
4574
4575 let result_field_columns = self
4577 .ctx
4578 .field_columns
4579 .iter()
4580 .map(name_to_expr)
4581 .collect::<Result<Vec<_>>>()?;
4582
4583 self.ctx.field_columns = result_field_columns
4585 .iter()
4586 .map(|expr| expr.schema_name().to_string())
4587 .collect();
4588 let field_columns_iter = result_field_columns
4589 .into_iter()
4590 .zip(self.ctx.field_columns.iter())
4591 .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::<String>, name))));
4592
4593 let project_fields = non_field_columns_iter
4595 .chain(tsid_iter)
4596 .chain(field_columns_iter)
4597 .collect::<Result<Vec<_>>>()?;
4598
4599 LogicalPlanBuilder::from(input)
4600 .project(project_fields)
4601 .context(DataFusionPlanningSnafu)?
4602 .build()
4603 .context(DataFusionPlanningSnafu)
4604 }
4605
4606 fn filter_on_field_column<F>(
4609 &self,
4610 input: LogicalPlan,
4611 mut name_to_expr: F,
4612 ) -> Result<LogicalPlan>
4613 where
4614 F: FnMut(&String) -> Result<DfExpr>,
4615 {
4616 ensure!(
4617 self.ctx.field_columns.len() == 1,
4618 UnsupportedExprSnafu {
4619 name: "filter on multi-value input"
4620 }
4621 );
4622
4623 let field_column_filter = name_to_expr(&self.ctx.field_columns[0])?;
4624
4625 LogicalPlanBuilder::from(input)
4626 .filter(field_column_filter)
4627 .context(DataFusionPlanningSnafu)?
4628 .build()
4629 .context(DataFusionPlanningSnafu)
4630 }
4631
4632 fn date_part_on_time_index(&self, date_part: &str) -> Result<DfExpr> {
4635 let input_expr = datafusion::logical_expr::col(
4636 self.ctx
4637 .time_index_column
4638 .as_ref()
4639 .with_context(|| TimeIndexNotFoundSnafu {
4641 table: "<doesn't matter>",
4642 })?,
4643 );
4644 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
4645 func: datafusion_functions::datetime::date_part(),
4646 args: vec![date_part.lit(), input_expr],
4647 });
4648 Ok(fn_expr)
4649 }
4650
4651 fn strip_tsid_column(&self, plan: LogicalPlan) -> Result<LogicalPlan> {
4652 let schema = plan.schema();
4653 if !schema
4654 .fields()
4655 .iter()
4656 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
4657 {
4658 return Ok(plan);
4659 }
4660
4661 let project_exprs = schema
4664 .iter()
4665 .filter(|(_, field)| field.name() != DATA_SCHEMA_TSID_COLUMN_NAME)
4666 .map(|(qualifier, field)| {
4667 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
4668 })
4669 .collect::<Vec<_>>();
4670
4671 LogicalPlanBuilder::from(plan)
4672 .project(project_exprs)
4673 .context(DataFusionPlanningSnafu)?
4674 .build()
4675 .context(DataFusionPlanningSnafu)
4676 }
4677
4678 fn apply_alias(&mut self, plan: LogicalPlan, alias_name: String) -> Result<LogicalPlan> {
4680 let fields_expr = self.create_field_column_exprs()?;
4681
4682 ensure!(
4684 fields_expr.len() == 1,
4685 UnsupportedExprSnafu {
4686 name: "alias on multi-value result"
4687 }
4688 );
4689
4690 let project_fields = fields_expr
4691 .into_iter()
4692 .map(|expr| expr.alias(&alias_name))
4693 .chain(self.create_tag_column_exprs()?)
4694 .chain(Some(self.create_time_index_column_expr()?));
4695
4696 LogicalPlanBuilder::from(plan)
4697 .project(project_fields)
4698 .context(DataFusionPlanningSnafu)?
4699 .build()
4700 .context(DataFusionPlanningSnafu)
4701 }
4702}
4703
4704#[derive(Default, Debug)]
4705struct FunctionArgs {
4706 input: Option<PromExpr>,
4707 literals: Vec<DfExpr>,
4708}
4709
4710#[derive(Debug, Clone)]
4713enum ScalarFunc {
4714 DataFusionBuiltin(Arc<ScalarUdfDef>),
4718 DataFusionUdf(Arc<ScalarUdfDef>),
4722 Udf(Arc<ScalarUdfDef>),
4727 ExtrapolateUdf(Arc<ScalarUdfDef>, i64),
4734 GeneratedExpr,
4738}
4739
4740#[cfg(test)]
4741mod test {
4742 use std::time::{Duration, UNIX_EPOCH};
4743
4744 use catalog::RegisterTableRequest;
4745 use catalog::memory::{MemoryCatalogManager, new_memory_catalog_manager};
4746 use common_base::Plugins;
4747 use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
4748 use common_query::prelude::greptime_timestamp;
4749 use common_query::test_util::DummyDecoder;
4750 use datafusion::arrow::datatypes::Schema as ArrowSchema;
4751 use datafusion::datasource::memory::MemorySourceConfig;
4752 use datafusion::datasource::source::DataSourceExec;
4753 use datafusion::logical_expr::Extension;
4754 use datatypes::prelude::ConcreteDataType;
4755 use datatypes::schema::{ColumnSchema, Schema};
4756 use promql_parser::label::Labels;
4757 use promql_parser::parser;
4758 use session::context::QueryContext;
4759 use table::metadata::{TableInfoBuilder, TableMetaBuilder};
4760 use table::test_util::EmptyTable;
4761
4762 use super::*;
4763 use crate::QueryEngineContext;
4764 use crate::options::QueryOptions;
4765 use crate::parser::QueryLanguageParser;
4766
4767 fn find_instant_manipulate(plan: &LogicalPlan) -> Option<&InstantManipulate> {
4768 if let LogicalPlan::Extension(Extension { node }) = plan
4769 && let Some(instant_manipulate) = node.as_any().downcast_ref::<InstantManipulate>()
4770 {
4771 return Some(instant_manipulate);
4772 }
4773
4774 plan.inputs().into_iter().find_map(find_instant_manipulate)
4775 }
4776
4777 fn build_query_engine_state() -> QueryEngineState {
4778 QueryEngineState::new(
4779 new_memory_catalog_manager().unwrap(),
4780 None,
4781 None,
4782 None,
4783 None,
4784 None,
4785 false,
4786 Plugins::default(),
4787 QueryOptions::default(),
4788 )
4789 }
4790
4791 async fn build_optimized_promql_plan(
4792 table_provider: DfTableSourceProvider,
4793 eval_stmt: &EvalStmt,
4794 ) -> LogicalPlan {
4795 let state = build_query_engine_state();
4796 let raw_plan = PromPlanner::stmt_to_plan(table_provider, eval_stmt, &state)
4797 .await
4798 .unwrap();
4799 let context = QueryEngineContext::new(state.session_state(), QueryContext::arc());
4800 state
4801 .optimize_by_extension_rules(raw_plan, &context)
4802 .unwrap()
4803 }
4804
4805 async fn build_optimized_tsid_plan(
4806 query: &str,
4807 num_tag: usize,
4808 num_field: usize,
4809 end_secs: u64,
4810 lookback_secs: u64,
4811 ) -> String {
4812 let eval_stmt = EvalStmt {
4813 expr: parser::parse(query).unwrap(),
4814 start: UNIX_EPOCH,
4815 end: UNIX_EPOCH
4816 .checked_add(Duration::from_secs(end_secs))
4817 .unwrap(),
4818 interval: Duration::from_secs(5),
4819 lookback_delta: Duration::from_secs(lookback_secs),
4820 };
4821 let table_provider = build_test_table_provider_with_tsid(
4822 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
4823 num_tag,
4824 num_field,
4825 )
4826 .await;
4827
4828 build_optimized_promql_plan(table_provider, &eval_stmt)
4829 .await
4830 .display_indent_schema()
4831 .to_string()
4832 }
4833
4834 async fn assert_nested_count_rewrite_applies(query: &str, expected_outer_agg: &str) {
4835 let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await;
4836
4837 assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
4838 assert!(plan_str.contains("Projection: some_metric.timestamp, some_metric.tag_0"));
4839 assert!(plan_str.contains("Distinct:"));
4840 assert!(plan_str.contains(expected_outer_agg), "{plan_str}");
4841 assert!(!plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]"));
4842 }
4843
4844 async fn assert_nested_count_rewrite_missing(query: &str, num_tag: usize, lookback_secs: u64) {
4845 let plan_str = build_optimized_tsid_plan(query, num_tag, 1, 100_000, lookback_secs).await;
4846 assert!(!plan_str.contains("Distinct:"), "{plan_str}");
4847 }
4848
4849 fn build_eval_stmt(expr: &str) -> EvalStmt {
4850 EvalStmt {
4851 expr: parser::parse(expr).unwrap(),
4852 start: UNIX_EPOCH,
4853 end: UNIX_EPOCH
4854 .checked_add(Duration::from_secs(100_000))
4855 .unwrap(),
4856 interval: Duration::from_secs(5),
4857 lookback_delta: Duration::from_secs(1),
4858 }
4859 }
4860
4861 async fn build_test_table_provider(
4862 table_name_tuples: &[(String, String)],
4863 num_tag: usize,
4864 num_field: usize,
4865 ) -> DfTableSourceProvider {
4866 let catalog_list = MemoryCatalogManager::with_default_setup();
4867 for (schema_name, table_name) in table_name_tuples {
4868 let mut columns = vec![];
4869 for i in 0..num_tag {
4870 columns.push(ColumnSchema::new(
4871 format!("tag_{i}"),
4872 ConcreteDataType::string_datatype(),
4873 false,
4874 ));
4875 }
4876 columns.push(
4877 ColumnSchema::new(
4878 "timestamp".to_string(),
4879 ConcreteDataType::timestamp_millisecond_datatype(),
4880 false,
4881 )
4882 .with_time_index(true),
4883 );
4884 for i in 0..num_field {
4885 columns.push(ColumnSchema::new(
4886 format!("field_{i}"),
4887 ConcreteDataType::float64_datatype(),
4888 true,
4889 ));
4890 }
4891 let schema = Arc::new(Schema::new(columns));
4892 let table_meta = TableMetaBuilder::empty()
4893 .schema(schema)
4894 .primary_key_indices((0..num_tag).collect())
4895 .value_indices((num_tag + 1..num_tag + 1 + num_field).collect())
4896 .next_column_id(1024)
4897 .build()
4898 .unwrap();
4899 let table_info = TableInfoBuilder::default()
4900 .name(table_name.clone())
4901 .meta(table_meta)
4902 .build()
4903 .unwrap();
4904 let table = EmptyTable::from_table_info(&table_info);
4905
4906 assert!(
4907 catalog_list
4908 .register_table_sync(RegisterTableRequest {
4909 catalog: DEFAULT_CATALOG_NAME.to_string(),
4910 schema: schema_name.clone(),
4911 table_name: table_name.clone(),
4912 table_id: 1024,
4913 table,
4914 })
4915 .is_ok()
4916 );
4917 }
4918
4919 DfTableSourceProvider::new(
4920 catalog_list,
4921 false,
4922 QueryContext::arc(),
4923 DummyDecoder::arc(),
4924 false,
4925 )
4926 }
4927
4928 async fn build_test_table_provider_with_tsid(
4929 table_name_tuples: &[(String, String)],
4930 num_tag: usize,
4931 num_field: usize,
4932 ) -> DfTableSourceProvider {
4933 let table_specs = table_name_tuples
4934 .iter()
4935 .map(|(schema_name, table_name)| ((schema_name.clone(), table_name.clone()), num_field))
4936 .collect::<Vec<_>>();
4937 build_test_table_provider_with_tsid_fields(&table_specs, num_tag).await
4938 }
4939
4940 async fn build_test_table_provider_with_tsid_fields(
4941 table_specs: &[((String, String), usize)],
4942 num_tag: usize,
4943 ) -> DfTableSourceProvider {
4944 let table_specs = table_specs
4945 .iter()
4946 .map(|(table_name_tuple, num_field)| (table_name_tuple.clone(), num_tag, *num_field))
4947 .collect::<Vec<_>>();
4948 build_test_table_provider_with_tsid_tag_fields(&table_specs).await
4949 }
4950
4951 async fn build_test_table_provider_with_tsid_tag_fields(
4952 table_specs: &[((String, String), usize, usize)],
4953 ) -> DfTableSourceProvider {
4954 let catalog_list = MemoryCatalogManager::with_default_setup();
4955
4956 let physical_table_name = "phy";
4957 let physical_table_id = 999u32;
4958 let physical_num_tag = table_specs
4959 .iter()
4960 .map(|(_, num_tag, _)| *num_tag)
4961 .max()
4962 .unwrap_or(0);
4963 let physical_num_field = table_specs
4964 .iter()
4965 .map(|(_, _, num_field)| *num_field)
4966 .max()
4967 .unwrap_or(0);
4968
4969 {
4971 let mut columns = vec![
4972 ColumnSchema::new(
4973 DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(),
4974 ConcreteDataType::uint32_datatype(),
4975 false,
4976 ),
4977 ColumnSchema::new(
4978 DATA_SCHEMA_TSID_COLUMN_NAME.to_string(),
4979 ConcreteDataType::uint64_datatype(),
4980 false,
4981 ),
4982 ];
4983 for i in 0..physical_num_tag {
4984 columns.push(ColumnSchema::new(
4985 format!("tag_{i}"),
4986 ConcreteDataType::string_datatype(),
4987 false,
4988 ));
4989 }
4990 columns.push(
4991 ColumnSchema::new(
4992 "timestamp".to_string(),
4993 ConcreteDataType::timestamp_millisecond_datatype(),
4994 false,
4995 )
4996 .with_time_index(true),
4997 );
4998 for i in 0..physical_num_field {
4999 columns.push(ColumnSchema::new(
5000 format!("field_{i}"),
5001 ConcreteDataType::float64_datatype(),
5002 true,
5003 ));
5004 }
5005
5006 let schema = Arc::new(Schema::new(columns));
5007 let primary_key_indices = (0..(2 + physical_num_tag)).collect::<Vec<_>>();
5008 let table_meta = TableMetaBuilder::empty()
5009 .schema(schema)
5010 .primary_key_indices(primary_key_indices)
5011 .value_indices(
5012 (2 + physical_num_tag..2 + physical_num_tag + 1 + physical_num_field).collect(),
5013 )
5014 .engine(METRIC_ENGINE_NAME.to_string())
5015 .next_column_id(1024)
5016 .build()
5017 .unwrap();
5018 let table_info = TableInfoBuilder::default()
5019 .table_id(physical_table_id)
5020 .name(physical_table_name)
5021 .meta(table_meta)
5022 .build()
5023 .unwrap();
5024 let table = EmptyTable::from_table_info(&table_info);
5025
5026 assert!(
5027 catalog_list
5028 .register_table_sync(RegisterTableRequest {
5029 catalog: DEFAULT_CATALOG_NAME.to_string(),
5030 schema: DEFAULT_SCHEMA_NAME.to_string(),
5031 table_name: physical_table_name.to_string(),
5032 table_id: physical_table_id,
5033 table,
5034 })
5035 .is_ok()
5036 );
5037 }
5038
5039 for (idx, ((schema_name, table_name), num_tag, num_field)) in table_specs.iter().enumerate()
5041 {
5042 let mut columns = vec![];
5043 for i in 0..*num_tag {
5044 columns.push(ColumnSchema::new(
5045 format!("tag_{i}"),
5046 ConcreteDataType::string_datatype(),
5047 false,
5048 ));
5049 }
5050 columns.push(
5051 ColumnSchema::new(
5052 "timestamp".to_string(),
5053 ConcreteDataType::timestamp_millisecond_datatype(),
5054 false,
5055 )
5056 .with_time_index(true),
5057 );
5058 for i in 0..*num_field {
5059 columns.push(ColumnSchema::new(
5060 format!("field_{i}"),
5061 ConcreteDataType::float64_datatype(),
5062 true,
5063 ));
5064 }
5065
5066 let schema = Arc::new(Schema::new(columns));
5067 let mut options = table::requests::TableOptions::default();
5068 options.extra_options.insert(
5069 LOGICAL_TABLE_METADATA_KEY.to_string(),
5070 physical_table_name.to_string(),
5071 );
5072 let table_id = 1024u32 + idx as u32;
5073 let table_meta = TableMetaBuilder::empty()
5074 .schema(schema)
5075 .primary_key_indices((0..*num_tag).collect())
5076 .value_indices((*num_tag + 1..*num_tag + 1 + *num_field).collect())
5077 .engine(METRIC_ENGINE_NAME.to_string())
5078 .options(options)
5079 .next_column_id(1024)
5080 .build()
5081 .unwrap();
5082 let table_info = TableInfoBuilder::default()
5083 .table_id(table_id)
5084 .name(table_name.clone())
5085 .meta(table_meta)
5086 .build()
5087 .unwrap();
5088 let table = EmptyTable::from_table_info(&table_info);
5089
5090 assert!(
5091 catalog_list
5092 .register_table_sync(RegisterTableRequest {
5093 catalog: DEFAULT_CATALOG_NAME.to_string(),
5094 schema: schema_name.clone(),
5095 table_name: table_name.clone(),
5096 table_id,
5097 table,
5098 })
5099 .is_ok()
5100 );
5101 }
5102
5103 DfTableSourceProvider::new(
5104 catalog_list,
5105 false,
5106 QueryContext::arc(),
5107 DummyDecoder::arc(),
5108 false,
5109 )
5110 }
5111
5112 async fn build_test_table_provider_with_fields(
5113 table_name_tuples: &[(String, String)],
5114 tags: &[&str],
5115 ) -> DfTableSourceProvider {
5116 let catalog_list = MemoryCatalogManager::with_default_setup();
5117 for (schema_name, table_name) in table_name_tuples {
5118 let mut columns = vec![];
5119 let num_tag = tags.len();
5120 for tag in tags {
5121 columns.push(ColumnSchema::new(
5122 tag.to_string(),
5123 ConcreteDataType::string_datatype(),
5124 false,
5125 ));
5126 }
5127 columns.push(
5128 ColumnSchema::new(
5129 greptime_timestamp().to_string(),
5130 ConcreteDataType::timestamp_millisecond_datatype(),
5131 false,
5132 )
5133 .with_time_index(true),
5134 );
5135 columns.push(ColumnSchema::new(
5136 greptime_value().to_string(),
5137 ConcreteDataType::float64_datatype(),
5138 true,
5139 ));
5140 let schema = Arc::new(Schema::new(columns));
5141 let table_meta = TableMetaBuilder::empty()
5142 .schema(schema)
5143 .primary_key_indices((0..num_tag).collect())
5144 .next_column_id(1024)
5145 .build()
5146 .unwrap();
5147 let table_info = TableInfoBuilder::default()
5148 .name(table_name.clone())
5149 .meta(table_meta)
5150 .build()
5151 .unwrap();
5152 let table = EmptyTable::from_table_info(&table_info);
5153
5154 assert!(
5155 catalog_list
5156 .register_table_sync(RegisterTableRequest {
5157 catalog: DEFAULT_CATALOG_NAME.to_string(),
5158 schema: schema_name.clone(),
5159 table_name: table_name.clone(),
5160 table_id: 1024,
5161 table,
5162 })
5163 .is_ok()
5164 );
5165 }
5166
5167 DfTableSourceProvider::new(
5168 catalog_list,
5169 false,
5170 QueryContext::arc(),
5171 DummyDecoder::arc(),
5172 false,
5173 )
5174 }
5175
5176 async fn do_single_instant_function_call(fn_name: &'static str, plan_name: &str) {
5192 let prom_expr =
5193 parser::parse(&format!("{fn_name}(some_metric{{tag_0!=\"bar\"}})")).unwrap();
5194 let eval_stmt = EvalStmt {
5195 expr: prom_expr,
5196 start: UNIX_EPOCH,
5197 end: UNIX_EPOCH
5198 .checked_add(Duration::from_secs(100_000))
5199 .unwrap(),
5200 interval: Duration::from_secs(5),
5201 lookback_delta: Duration::from_secs(1),
5202 };
5203
5204 let table_provider = build_test_table_provider(
5205 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
5206 1,
5207 1,
5208 )
5209 .await;
5210 let plan =
5211 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5212 .await
5213 .unwrap();
5214
5215 let expected = String::from(
5216 "Filter: TEMPLATE(field_0) IS NOT NULL [timestamp:Timestamp(ms), TEMPLATE(field_0):Float64;N, tag_0:Utf8]\
5217 \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]\
5218 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
5219 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
5220 \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]\
5221 \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]\
5222 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"
5223 ).replace("TEMPLATE", plan_name);
5224
5225 assert_eq!(plan.display_indent_schema().to_string(), expected);
5226 }
5227
5228 #[tokio::test]
5229 async fn single_abs() {
5230 do_single_instant_function_call("abs", "abs").await;
5231 }
5232
5233 #[tokio::test]
5234 #[should_panic]
5235 async fn single_absent() {
5236 do_single_instant_function_call("absent", "").await;
5237 }
5238
5239 #[tokio::test]
5240 async fn single_ceil() {
5241 do_single_instant_function_call("ceil", "ceil").await;
5242 }
5243
5244 #[tokio::test]
5245 async fn single_exp() {
5246 do_single_instant_function_call("exp", "exp").await;
5247 }
5248
5249 #[tokio::test]
5250 async fn single_ln() {
5251 do_single_instant_function_call("ln", "ln").await;
5252 }
5253
5254 #[tokio::test]
5255 async fn single_log2() {
5256 do_single_instant_function_call("log2", "log2").await;
5257 }
5258
5259 #[tokio::test]
5260 async fn single_log10() {
5261 do_single_instant_function_call("log10", "log10").await;
5262 }
5263
5264 #[tokio::test]
5265 #[should_panic]
5266 async fn single_scalar() {
5267 do_single_instant_function_call("scalar", "").await;
5268 }
5269
5270 #[tokio::test]
5271 #[should_panic]
5272 async fn single_sgn() {
5273 do_single_instant_function_call("sgn", "").await;
5274 }
5275
5276 #[tokio::test]
5277 #[should_panic]
5278 async fn single_sort() {
5279 do_single_instant_function_call("sort", "").await;
5280 }
5281
5282 #[tokio::test]
5283 #[should_panic]
5284 async fn single_sort_desc() {
5285 do_single_instant_function_call("sort_desc", "").await;
5286 }
5287
5288 #[tokio::test]
5289 async fn single_sqrt() {
5290 do_single_instant_function_call("sqrt", "sqrt").await;
5291 }
5292
5293 #[tokio::test]
5294 #[should_panic]
5295 async fn single_timestamp() {
5296 do_single_instant_function_call("timestamp", "").await;
5297 }
5298
5299 #[tokio::test]
5300 async fn single_acos() {
5301 do_single_instant_function_call("acos", "acos").await;
5302 }
5303
5304 #[tokio::test]
5305 #[should_panic]
5306 async fn single_acosh() {
5307 do_single_instant_function_call("acosh", "").await;
5308 }
5309
5310 #[tokio::test]
5311 async fn single_asin() {
5312 do_single_instant_function_call("asin", "asin").await;
5313 }
5314
5315 #[tokio::test]
5316 #[should_panic]
5317 async fn single_asinh() {
5318 do_single_instant_function_call("asinh", "").await;
5319 }
5320
5321 #[tokio::test]
5322 async fn single_atan() {
5323 do_single_instant_function_call("atan", "atan").await;
5324 }
5325
5326 #[tokio::test]
5327 #[should_panic]
5328 async fn single_atanh() {
5329 do_single_instant_function_call("atanh", "").await;
5330 }
5331
5332 #[tokio::test]
5333 async fn single_cos() {
5334 do_single_instant_function_call("cos", "cos").await;
5335 }
5336
5337 #[tokio::test]
5338 #[should_panic]
5339 async fn single_cosh() {
5340 do_single_instant_function_call("cosh", "").await;
5341 }
5342
5343 #[tokio::test]
5344 async fn single_sin() {
5345 do_single_instant_function_call("sin", "sin").await;
5346 }
5347
5348 #[tokio::test]
5349 #[should_panic]
5350 async fn single_sinh() {
5351 do_single_instant_function_call("sinh", "").await;
5352 }
5353
5354 #[tokio::test]
5355 async fn single_tan() {
5356 do_single_instant_function_call("tan", "tan").await;
5357 }
5358
5359 #[tokio::test]
5360 #[should_panic]
5361 async fn single_tanh() {
5362 do_single_instant_function_call("tanh", "").await;
5363 }
5364
5365 #[tokio::test]
5366 #[should_panic]
5367 async fn single_deg() {
5368 do_single_instant_function_call("deg", "").await;
5369 }
5370
5371 #[tokio::test]
5372 #[should_panic]
5373 async fn single_rad() {
5374 do_single_instant_function_call("rad", "").await;
5375 }
5376
5377 async fn do_aggregate_expr_plan(fn_name: &str, plan_name: &str) {
5399 let prom_expr = parser::parse(&format!(
5400 "{fn_name} by (tag_1)(some_metric{{tag_0!=\"bar\"}})",
5401 ))
5402 .unwrap();
5403 let mut eval_stmt = EvalStmt {
5404 expr: prom_expr,
5405 start: UNIX_EPOCH,
5406 end: UNIX_EPOCH
5407 .checked_add(Duration::from_secs(100_000))
5408 .unwrap(),
5409 interval: Duration::from_secs(5),
5410 lookback_delta: Duration::from_secs(1),
5411 };
5412
5413 let table_provider = build_test_table_provider(
5415 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
5416 2,
5417 2,
5418 )
5419 .await;
5420 let plan =
5421 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5422 .await
5423 .unwrap();
5424 let expected_no_without = String::from(
5425 "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]\
5426 \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]\
5427 \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]\
5428 \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
5429 \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]\
5430 \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]\
5431 \n TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]"
5432 ).replace("TEMPLATE", plan_name);
5433 assert_eq!(
5434 plan.display_indent_schema().to_string(),
5435 expected_no_without
5436 );
5437
5438 if let PromExpr::Aggregate(AggregateExpr { modifier, .. }) = &mut eval_stmt.expr {
5440 *modifier = Some(LabelModifier::Exclude(Labels {
5441 labels: vec![String::from("tag_1")].into_iter().collect(),
5442 }));
5443 }
5444 let table_provider = build_test_table_provider(
5445 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
5446 2,
5447 2,
5448 )
5449 .await;
5450 let plan =
5451 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5452 .await
5453 .unwrap();
5454 let expected_without = String::from(
5455 "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]\
5456 \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]\
5457 \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]\
5458 \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
5459 \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]\
5460 \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]\
5461 \n TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]"
5462 ).replace("TEMPLATE", plan_name);
5463 assert_eq!(plan.display_indent_schema().to_string(), expected_without);
5464 }
5465
5466 #[tokio::test]
5467 async fn aggregate_sum() {
5468 do_aggregate_expr_plan("sum", "sum").await;
5469 }
5470
5471 #[tokio::test]
5472 async fn tsid_is_used_for_series_divide_when_available() {
5473 let prom_expr = parser::parse("some_metric").unwrap();
5474 let eval_stmt = EvalStmt {
5475 expr: prom_expr,
5476 start: UNIX_EPOCH,
5477 end: UNIX_EPOCH
5478 .checked_add(Duration::from_secs(100_000))
5479 .unwrap(),
5480 interval: Duration::from_secs(5),
5481 lookback_delta: Duration::from_secs(1),
5482 };
5483
5484 let table_provider = build_test_table_provider_with_tsid(
5485 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
5486 1,
5487 1,
5488 )
5489 .await;
5490 let plan =
5491 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5492 .await
5493 .unwrap();
5494
5495 let plan_str = plan.display_indent_schema().to_string();
5496 assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
5497 assert!(plan_str.contains("__tsid ASC NULLS FIRST"));
5498 assert!(
5499 !plan
5500 .schema()
5501 .fields()
5502 .iter()
5503 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
5504 );
5505
5506 let manipulate = find_instant_manipulate(&plan).unwrap();
5507 let exec = manipulate.to_execution_plan(Arc::new(DataSourceExec::new(Arc::new(
5508 MemorySourceConfig::try_new(&[], Arc::new(ArrowSchema::empty()), None).unwrap(),
5509 ))));
5510 assert!(format!("{exec:?}").contains("reuse_tsid_column: true"));
5511 }
5512
5513 #[tokio::test]
5514 async fn default_binary_join_uses_tsid_when_available() {
5515 let eval_stmt = build_eval_stmt("some_metric / some_alt_metric");
5516
5517 let table_provider = build_test_table_provider_with_tsid(
5518 &[
5519 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
5520 (
5521 DEFAULT_SCHEMA_NAME.to_string(),
5522 "some_alt_metric".to_string(),
5523 ),
5524 ],
5525 1,
5526 1,
5527 )
5528 .await;
5529 let plan =
5530 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5531 .await
5532 .unwrap();
5533
5534 let plan_str = plan.display_indent_schema().to_string();
5535 assert!(
5536 plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
5537 "{plan_str}"
5538 );
5539 assert!(
5540 !plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
5541 "{plan_str}"
5542 );
5543 }
5544
5545 #[tokio::test]
5546 async fn timestamp_binary_join_falls_back_when_tsid_is_projected_out() {
5547 for query in [
5548 "timestamp(some_metric) / some_metric",
5549 "some_metric / timestamp(some_metric)",
5550 ] {
5551 let eval_stmt = build_eval_stmt(query);
5552
5553 let table_provider = build_test_table_provider_with_tsid(
5554 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
5555 1,
5556 1,
5557 )
5558 .await;
5559 let plan =
5560 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5561 .await
5562 .unwrap();
5563
5564 let plan_str = plan.display_indent_schema().to_string();
5565 assert!(!plan_str.contains("__tsid ="), "{query}: {plan_str}");
5566 assert!(
5567 plan_str.contains("lhs.tag_0 = rhs.tag_0"),
5568 "{query}: {plan_str}"
5569 );
5570 assert!(
5571 !plan
5572 .schema()
5573 .fields()
5574 .iter()
5575 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME),
5576 "{query}: {plan_str}"
5577 );
5578 }
5579 }
5580
5581 #[tokio::test]
5582 async fn timestamp_binary_join_rejects_default_matching_on_mismatched_labels() {
5583 let eval_stmt = build_eval_stmt("timestamp(left_host_job) / right_by_job");
5584
5585 let table_provider = build_test_table_provider_with_tsid_tag_fields(&[
5586 (
5587 (DEFAULT_SCHEMA_NAME.to_string(), "left_host_job".to_string()),
5588 2,
5589 1,
5590 ),
5591 (
5592 (DEFAULT_SCHEMA_NAME.to_string(), "right_by_job".to_string()),
5593 1,
5594 1,
5595 ),
5596 ])
5597 .await;
5598 let plan =
5599 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5600 .await
5601 .unwrap();
5602 let plan_str = plan.display_indent_schema().to_string();
5603
5604 assert!(
5605 plan_str.contains("Boolean(false)") || plan_str.contains("false"),
5606 "{plan_str}"
5607 );
5608 }
5609
5610 #[tokio::test]
5611 async fn tsid_is_preserved_for_nested_default_binary_joins() {
5612 let eval_stmt = build_eval_stmt("(some_metric - some_alt_metric) / some_third_metric");
5613
5614 let table_provider = build_test_table_provider_with_tsid(
5615 &[
5616 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
5617 (
5618 DEFAULT_SCHEMA_NAME.to_string(),
5619 "some_alt_metric".to_string(),
5620 ),
5621 (
5622 DEFAULT_SCHEMA_NAME.to_string(),
5623 "some_third_metric".to_string(),
5624 ),
5625 ],
5626 1,
5627 1,
5628 )
5629 .await;
5630 let plan =
5631 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5632 .await
5633 .unwrap();
5634
5635 let plan_str = plan.display_indent_schema().to_string();
5636 assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
5637 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
5638 }
5639
5640 #[tokio::test]
5641 async fn repeated_tsid_binary_operand_reuses_leaf_plan() {
5642 let eval_stmt = build_eval_stmt("((some_metric - some_alt_metric) / some_metric) * 100");
5643
5644 let table_provider = build_test_table_provider_with_tsid(
5645 &[
5646 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
5647 (
5648 DEFAULT_SCHEMA_NAME.to_string(),
5649 "some_alt_metric".to_string(),
5650 ),
5651 ],
5652 1,
5653 1,
5654 )
5655 .await;
5656 let plan =
5657 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5658 .await
5659 .unwrap();
5660
5661 let plan_str = plan.display_indent_schema().to_string();
5662 assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}");
5663 assert_eq!(
5664 plan_str
5665 .matches("Filter: phy.__table_id = UInt32(1024)")
5666 .count(),
5667 1,
5668 "{plan_str}"
5669 );
5670 assert_eq!(
5671 plan_str.matches("PromInstantManipulate").count(),
5672 2,
5673 "{plan_str}"
5674 );
5675 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
5676 }
5677
5678 #[tokio::test]
5679 async fn repeated_tsid_binary_operand_reuses_shorter_field_side() {
5680 let eval_stmt =
5681 build_eval_stmt("((two_field_metric - one_field_metric) / one_field_metric) * 100");
5682
5683 let table_provider = build_test_table_provider_with_tsid_fields(
5684 &[
5685 (
5686 (
5687 DEFAULT_SCHEMA_NAME.to_string(),
5688 "two_field_metric".to_string(),
5689 ),
5690 2,
5691 ),
5692 (
5693 (
5694 DEFAULT_SCHEMA_NAME.to_string(),
5695 "one_field_metric".to_string(),
5696 ),
5697 1,
5698 ),
5699 ],
5700 1,
5701 )
5702 .await;
5703 let plan =
5704 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5705 .await
5706 .unwrap();
5707
5708 let field_names = plan
5709 .schema()
5710 .fields()
5711 .iter()
5712 .map(|field| field.name().clone())
5713 .collect::<Vec<_>>();
5714 let value_columns = field_names
5715 .iter()
5716 .filter(|name| {
5717 *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME
5718 })
5719 .count();
5720 assert_eq!(value_columns, 1, "{field_names:?}");
5721 let plan_str = plan.display_indent_schema().to_string();
5722 assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}");
5723 assert_eq!(
5724 plan_str
5725 .matches("Filter: phy.__table_id = UInt32(1025)")
5726 .count(),
5727 1,
5728 "{plan_str}"
5729 );
5730 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
5731 }
5732
5733 #[tokio::test]
5734 async fn binary_island_reuses_self_operand_without_join() {
5735 let eval_stmt = build_eval_stmt("some_metric / some_metric");
5736
5737 let table_provider = build_test_table_provider_with_tsid(
5738 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
5739 1,
5740 1,
5741 )
5742 .await;
5743 let plan =
5744 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5745 .await
5746 .unwrap();
5747
5748 let plan_str = plan.display_indent_schema().to_string();
5749 assert_eq!(plan_str.matches("__tsid =").count(), 0, "{plan_str}");
5750 assert_eq!(
5751 plan_str
5752 .matches("Filter: phy.__table_id = UInt32(1024)")
5753 .count(),
5754 1,
5755 "{plan_str}"
5756 );
5757 assert_eq!(
5758 plan_str.matches("PromInstantManipulate").count(),
5759 1,
5760 "{plan_str}"
5761 );
5762 }
5763
5764 #[tokio::test]
5765 async fn binary_island_reuses_leaf_across_two_branches() {
5766 let eval_stmt =
5767 build_eval_stmt("(some_metric + some_alt_metric) / (some_metric + third_metric)");
5768
5769 let table_provider = build_test_table_provider_with_tsid(
5770 &[
5771 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
5772 (
5773 DEFAULT_SCHEMA_NAME.to_string(),
5774 "some_alt_metric".to_string(),
5775 ),
5776 (DEFAULT_SCHEMA_NAME.to_string(), "third_metric".to_string()),
5777 ],
5778 1,
5779 1,
5780 )
5781 .await;
5782 let plan =
5783 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5784 .await
5785 .unwrap();
5786
5787 let plan_str = plan.display_indent_schema().to_string();
5788 assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
5789 assert_eq!(
5790 plan_str
5791 .matches("Filter: phy.__table_id = UInt32(1024)")
5792 .count(),
5793 1,
5794 "{plan_str}"
5795 );
5796 assert_eq!(
5797 plan_str.matches("PromInstantManipulate").count(),
5798 3,
5799 "{plan_str}"
5800 );
5801 }
5802
5803 #[tokio::test]
5804 async fn binary_island_generated_alias_avoids_user_column_names() {
5805 let eval_stmt = build_eval_stmt("(some_metric + some_alt_metric) / some_metric");
5806
5807 let table_provider = build_test_table_provider_with_fields(
5808 &[
5809 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
5810 (
5811 DEFAULT_SCHEMA_NAME.to_string(),
5812 "some_alt_metric".to_string(),
5813 ),
5814 ],
5815 &["prom_v0", "__prom_v0"],
5816 )
5817 .await;
5818 let plan =
5819 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5820 .await
5821 .unwrap();
5822
5823 let field_names = plan.schema().field_names();
5824 assert!(field_names.iter().any(|name| name.ends_with(".prom_v0")));
5825 assert!(field_names.iter().any(|name| name.ends_with(".__prom_v0")));
5826
5827 let plan_str = plan.display_indent_schema().to_string();
5828 assert!(plan_str.contains("SubqueryAlias: __prom_v0"), "{plan_str}");
5829 assert_eq!(
5830 plan_str.matches("PromInstantManipulate").count(),
5831 2,
5832 "{plan_str}"
5833 );
5834 }
5835
5836 #[tokio::test]
5837 async fn binary_island_clears_qualifier_for_nested_unary_projection() {
5838 let eval_stmt = build_eval_stmt("-((some_metric + some_alt_metric) / some_metric)");
5839
5840 let table_provider = build_test_table_provider_with_tsid(
5841 &[
5842 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
5843 (
5844 DEFAULT_SCHEMA_NAME.to_string(),
5845 "some_alt_metric".to_string(),
5846 ),
5847 ],
5848 1,
5849 1,
5850 )
5851 .await;
5852 let plan =
5853 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5854 .await
5855 .unwrap();
5856
5857 let plan_str = plan.display_indent_schema().to_string();
5858 assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}");
5859 assert_eq!(
5860 plan_str.matches("PromInstantManipulate").count(),
5861 2,
5862 "{plan_str}"
5863 );
5864 }
5865
5866 #[tokio::test]
5867 async fn binary_island_keeps_distinct_matcher_leaves() {
5868 let eval_stmt = build_eval_stmt(
5869 "(some_metric{tag_0=\"foo\"} + some_alt_metric) / some_metric{tag_0=\"bar\"}",
5870 );
5871
5872 let table_provider = build_test_table_provider_with_tsid(
5873 &[
5874 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
5875 (
5876 DEFAULT_SCHEMA_NAME.to_string(),
5877 "some_alt_metric".to_string(),
5878 ),
5879 ],
5880 1,
5881 1,
5882 )
5883 .await;
5884 let plan =
5885 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5886 .await
5887 .unwrap();
5888
5889 let plan_str = plan.display_indent_schema().to_string();
5890 assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
5891 assert_eq!(
5892 plan_str.matches("PromInstantManipulate").count(),
5893 3,
5894 "{plan_str}"
5895 );
5896 }
5897
5898 #[tokio::test]
5899 async fn binary_island_keeps_offset_leaves_distinct() {
5900 let eval_stmt = build_eval_stmt("(some_metric offset 5m + some_alt_metric) / some_metric");
5901
5902 let table_provider = build_test_table_provider_with_tsid(
5903 &[
5904 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
5905 (
5906 DEFAULT_SCHEMA_NAME.to_string(),
5907 "some_alt_metric".to_string(),
5908 ),
5909 ],
5910 1,
5911 1,
5912 )
5913 .await;
5914 let plan =
5915 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5916 .await
5917 .unwrap();
5918
5919 let plan_str = plan.display_indent_schema().to_string();
5920 assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
5921 assert_eq!(
5922 plan_str.matches("PromInstantManipulate").count(),
5923 3,
5924 "{plan_str}"
5925 );
5926 }
5927
5928 #[tokio::test]
5929 async fn binary_island_falls_back_for_group_modifier() {
5930 let eval_stmt = build_eval_stmt(
5931 "(some_metric + ignoring(tag_0) group_left some_alt_metric) / some_metric",
5932 );
5933
5934 let table_provider = build_test_table_provider_with_tsid(
5935 &[
5936 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
5937 (
5938 DEFAULT_SCHEMA_NAME.to_string(),
5939 "some_alt_metric".to_string(),
5940 ),
5941 ],
5942 1,
5943 1,
5944 )
5945 .await;
5946 let plan =
5947 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5948 .await
5949 .unwrap();
5950
5951 let plan_str = plan.display_indent_schema().to_string();
5952 assert_eq!(
5953 plan_str.matches("PromInstantManipulate").count(),
5954 3,
5955 "{plan_str}"
5956 );
5957 }
5958
5959 #[tokio::test]
5960 async fn binary_island_falls_back_for_comparison_filter() {
5961 let eval_stmt = build_eval_stmt("(some_metric > some_alt_metric) / some_metric");
5962
5963 let table_provider = build_test_table_provider_with_tsid(
5964 &[
5965 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
5966 (
5967 DEFAULT_SCHEMA_NAME.to_string(),
5968 "some_alt_metric".to_string(),
5969 ),
5970 ],
5971 1,
5972 1,
5973 )
5974 .await;
5975 let plan =
5976 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5977 .await
5978 .unwrap();
5979
5980 let plan_str = plan.display_indent_schema().to_string();
5981 assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
5982 assert_eq!(
5983 plan_str.matches("PromInstantManipulate").count(),
5984 3,
5985 "{plan_str}"
5986 );
5987 }
5988
5989 #[tokio::test]
5990 async fn tsid_binary_join_uses_shorter_field_side() {
5991 let eval_stmt = build_eval_stmt("one_field_metric / two_field_metric");
5992
5993 let table_provider = build_test_table_provider_with_tsid_fields(
5994 &[
5995 (
5996 (
5997 DEFAULT_SCHEMA_NAME.to_string(),
5998 "one_field_metric".to_string(),
5999 ),
6000 1,
6001 ),
6002 (
6003 (
6004 DEFAULT_SCHEMA_NAME.to_string(),
6005 "two_field_metric".to_string(),
6006 ),
6007 2,
6008 ),
6009 ],
6010 1,
6011 )
6012 .await;
6013 let plan =
6014 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6015 .await
6016 .unwrap();
6017
6018 let field_names = plan
6019 .schema()
6020 .fields()
6021 .iter()
6022 .map(|field| field.name().clone())
6023 .collect::<Vec<_>>();
6024 let value_columns = field_names
6025 .iter()
6026 .filter(|name| {
6027 *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME
6028 })
6029 .count();
6030 assert_eq!(value_columns, 1, "{field_names:?}");
6031 }
6032
6033 #[tokio::test]
6034 async fn comparison_binary_join_uses_shorter_field_side() {
6035 let eval_stmt = build_eval_stmt("two_field_metric > one_field_metric");
6036
6037 let table_provider = build_test_table_provider_with_tsid_fields(
6038 &[
6039 (
6040 (
6041 DEFAULT_SCHEMA_NAME.to_string(),
6042 "two_field_metric".to_string(),
6043 ),
6044 2,
6045 ),
6046 (
6047 (
6048 DEFAULT_SCHEMA_NAME.to_string(),
6049 "one_field_metric".to_string(),
6050 ),
6051 1,
6052 ),
6053 ],
6054 1,
6055 )
6056 .await;
6057 let plan =
6058 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6059 .await
6060 .unwrap();
6061
6062 let field_names = plan
6063 .schema()
6064 .fields()
6065 .iter()
6066 .map(|field| field.name().clone())
6067 .collect::<Vec<_>>();
6068 assert!(
6069 field_names.iter().any(|name| name == "field_0"),
6070 "{field_names:?}"
6071 );
6072 assert!(
6073 !field_names.iter().any(|name| name == "field_1"),
6074 "{field_names:?}"
6075 );
6076 }
6077
6078 #[tokio::test]
6079 async fn label_matching_modifier_disables_tsid_binary_join() {
6080 let eval_stmt = build_eval_stmt("some_metric / ignoring(tag_0) some_alt_metric");
6081
6082 let table_provider = build_test_table_provider_with_tsid(
6083 &[
6084 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
6085 (
6086 DEFAULT_SCHEMA_NAME.to_string(),
6087 "some_alt_metric".to_string(),
6088 ),
6089 ],
6090 2,
6091 1,
6092 )
6093 .await;
6094 let plan =
6095 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6096 .await
6097 .unwrap();
6098
6099 let plan_str = plan.display_indent_schema().to_string();
6100 assert!(!plan_str.contains("__tsid ="), "{plan_str}");
6101 assert!(
6102 plan_str.contains("some_metric.tag_1 = some_alt_metric.tag_1"),
6103 "{plan_str}"
6104 );
6105 }
6106
6107 #[tokio::test]
6108 async fn ignoring_absent_label_keeps_tsid_binary_join() {
6109 let eval_stmt = build_eval_stmt("some_metric / ignoring(missing) some_alt_metric");
6110
6111 let table_provider = build_test_table_provider_with_tsid(
6112 &[
6113 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
6114 (
6115 DEFAULT_SCHEMA_NAME.to_string(),
6116 "some_alt_metric".to_string(),
6117 ),
6118 ],
6119 2,
6120 1,
6121 )
6122 .await;
6123 let plan =
6124 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6125 .await
6126 .unwrap();
6127
6128 let plan_str = plan.display_indent_schema().to_string();
6129 assert!(
6130 plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
6131 "{plan_str}"
6132 );
6133 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
6134 assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
6135 }
6136
6137 #[tokio::test]
6138 async fn range_function_keeps_tsid_for_absent_ignoring_binary_join() {
6139 let eval_stmt =
6140 build_eval_stmt("rate(some_metric[5m]) / ignoring(missing) some_alt_metric");
6141
6142 let table_provider = build_test_table_provider_with_tsid(
6143 &[
6144 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
6145 (
6146 DEFAULT_SCHEMA_NAME.to_string(),
6147 "some_alt_metric".to_string(),
6148 ),
6149 ],
6150 2,
6151 1,
6152 )
6153 .await;
6154 let plan =
6155 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6156 .await
6157 .unwrap();
6158
6159 let plan_str = plan.display_indent_schema().to_string();
6160 assert!(
6161 plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
6162 "{plan_str}"
6163 );
6164 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
6165 assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
6166 }
6167
6168 #[tokio::test]
6169 async fn on_full_label_set_keeps_tsid_binary_join() {
6170 let eval_stmt = build_eval_stmt("some_metric / on(tag_0, tag_1) some_alt_metric");
6171
6172 let table_provider = build_test_table_provider_with_tsid(
6173 &[
6174 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
6175 (
6176 DEFAULT_SCHEMA_NAME.to_string(),
6177 "some_alt_metric".to_string(),
6178 ),
6179 ],
6180 2,
6181 1,
6182 )
6183 .await;
6184 let plan =
6185 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6186 .await
6187 .unwrap();
6188
6189 let plan_str = plan.display_indent_schema().to_string();
6190 assert!(
6191 plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
6192 "{plan_str}"
6193 );
6194 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
6195 assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
6196 }
6197
6198 #[tokio::test]
6199 async fn on_partial_label_set_disables_tsid_binary_join() {
6200 let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric");
6201
6202 let table_provider = build_test_table_provider_with_tsid(
6203 &[
6204 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
6205 (
6206 DEFAULT_SCHEMA_NAME.to_string(),
6207 "some_alt_metric".to_string(),
6208 ),
6209 ],
6210 2,
6211 1,
6212 )
6213 .await;
6214 let plan =
6215 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6216 .await
6217 .unwrap();
6218
6219 let plan_str = plan.display_indent_schema().to_string();
6220 assert!(!plan_str.contains("__tsid ="), "{plan_str}");
6221 assert!(
6222 plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
6223 "{plan_str}"
6224 );
6225 assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
6226 }
6227
6228 #[tokio::test]
6229 async fn on_label_set_must_cover_both_sides_to_use_tsid_binary_join() {
6230 let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric");
6231
6232 let table_provider = build_test_table_provider_with_tsid_tag_fields(&[
6233 (
6234 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
6235 2,
6236 1,
6237 ),
6238 (
6239 (
6240 DEFAULT_SCHEMA_NAME.to_string(),
6241 "some_alt_metric".to_string(),
6242 ),
6243 1,
6244 1,
6245 ),
6246 ])
6247 .await;
6248 let plan =
6249 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6250 .await
6251 .unwrap();
6252
6253 let plan_str = plan.display_indent_schema().to_string();
6254 assert!(!plan_str.contains("__tsid ="), "{plan_str}");
6255 assert!(
6256 plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
6257 "{plan_str}"
6258 );
6259 assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
6260 }
6261
6262 #[tokio::test]
6263 async fn comparison_binary_join_uses_tsid_and_keeps_it_in_filtered_result() {
6264 let eval_stmt = build_eval_stmt("some_metric > some_alt_metric");
6265
6266 let table_provider = build_test_table_provider_with_tsid(
6267 &[
6268 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
6269 (
6270 DEFAULT_SCHEMA_NAME.to_string(),
6271 "some_alt_metric".to_string(),
6272 ),
6273 ],
6274 2,
6275 1,
6276 )
6277 .await;
6278 let mut planner = PromPlanner {
6279 table_provider,
6280 ctx: PromPlannerContext::from_eval_stmt(&eval_stmt),
6281 };
6282 let plan = planner
6283 .prom_expr_to_plan(&eval_stmt.expr, &build_query_engine_state())
6284 .await
6285 .unwrap();
6286
6287 let plan_str = plan.display_indent_schema().to_string();
6288 assert!(
6289 plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
6290 "{plan_str}"
6291 );
6292 assert!(
6293 plan.schema()
6294 .fields()
6295 .iter()
6296 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME),
6297 "{plan_str}"
6298 );
6299 assert!(planner.ctx.use_tsid, "{plan_str}");
6300 }
6301
6302 #[tokio::test]
6303 async fn comparison_bool_binary_join_uses_tsid_when_available() {
6304 let eval_stmt = build_eval_stmt("some_metric > bool some_alt_metric");
6305
6306 let table_provider = build_test_table_provider_with_tsid(
6307 &[
6308 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
6309 (
6310 DEFAULT_SCHEMA_NAME.to_string(),
6311 "some_alt_metric".to_string(),
6312 ),
6313 ],
6314 2,
6315 1,
6316 )
6317 .await;
6318 let plan =
6319 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6320 .await
6321 .unwrap();
6322
6323 let plan_str = plan.display_indent_schema().to_string();
6324 assert!(
6325 plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
6326 "{plan_str}"
6327 );
6328 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
6329 assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
6330 }
6331
6332 #[tokio::test]
6333 async fn scalar_count_count_range_keeps_full_window() {
6334 let plan_str = build_optimized_tsid_plan(
6335 "scalar(count(count(some_metric) by (tag_0)))",
6336 1,
6337 1,
6338 100_000,
6339 1,
6340 )
6341 .await;
6342 assert!(plan_str.contains("ScalarCalculate: tags=[]"));
6343 assert!(plan_str.contains("PromInstantManipulate: range=[0..100000000]"));
6344 assert!(!plan_str.contains("PromInstantManipulate: range=[99999000..99999000]"));
6345 }
6346
6347 #[tokio::test]
6348 async fn scalar_count_count_rewrite_applies_inside_binary_expr_for_tsid_input() {
6349 let plan_str = build_optimized_tsid_plan(
6350 "sum(irate(some_metric[1h])) / scalar(count(count(some_metric) by (tag_0)))",
6351 2,
6352 1,
6353 10,
6354 300,
6355 )
6356 .await;
6357 assert!(plan_str.contains("Distinct:"), "{plan_str}");
6358 }
6359
6360 #[tokio::test]
6361 async fn nested_count_rewrite_keeps_full_series_key_with_tsid_input() {
6362 assert_nested_count_rewrite_applies(
6363 "count(count(some_metric) by (tag_0))",
6364 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(count(some_metric.field_0))]]"
6365 )
6366 .await;
6367 }
6368
6369 #[tokio::test]
6370 async fn nested_sum_count_rewrite_keeps_full_series_key_with_tsid_input() {
6371 assert_nested_count_rewrite_applies(
6372 "count(sum(some_metric) by (tag_0))",
6373 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(sum(some_metric.field_0))]]"
6374 )
6375 .await;
6376 }
6377
6378 #[tokio::test]
6379 async fn nested_supported_inner_aggs_rewrite_apply_for_tsid_input() {
6380 for (query, expected_outer_agg) in [
6381 (
6382 "count(avg(some_metric) by (tag_0))",
6383 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(avg(some_metric.field_0))]]",
6384 ),
6385 (
6386 "count(min(some_metric) by (tag_0))",
6387 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(min(some_metric.field_0))]]",
6388 ),
6389 (
6390 "count(max(some_metric) by (tag_0))",
6391 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(max(some_metric.field_0))]]",
6392 ),
6393 (
6394 "count(stddev(some_metric) by (tag_0))",
6395 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(stddev_pop(some_metric.field_0))]]",
6396 ),
6397 (
6398 "count(stdvar(some_metric) by (tag_0))",
6399 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(var_pop(some_metric.field_0))]]",
6400 ),
6401 ] {
6402 assert_nested_count_rewrite_applies(query, expected_outer_agg).await;
6403 }
6404 }
6405
6406 #[tokio::test]
6407 async fn nested_non_count_inner_aggs_rewrite_filter_null_values_for_tsid_input() {
6408 let count_plan =
6409 build_optimized_tsid_plan("count(count(some_metric) by (tag_0))", 2, 1, 100_000, 1)
6410 .await;
6411 assert!(
6412 !count_plan.contains("some_metric.field_0 IS NOT NULL"),
6413 "{count_plan}"
6414 );
6415
6416 for query in [
6417 "count(sum(some_metric) by (tag_0))",
6418 "count(avg(some_metric) by (tag_0))",
6419 "count(min(some_metric) by (tag_0))",
6420 "count(max(some_metric) by (tag_0))",
6421 "count(stddev(some_metric) by (tag_0))",
6422 "count(stdvar(some_metric) by (tag_0))",
6423 ] {
6424 let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await;
6425 assert!(
6426 plan_str.contains("Filter: some_metric.field_0 IS NOT NULL"),
6427 "{query}: {plan_str}"
6428 );
6429 }
6430 }
6431
6432 #[tokio::test]
6433 async fn nested_unsupported_or_non_direct_inner_aggs_do_not_rewrite() {
6434 assert_nested_count_rewrite_missing("count(group(some_metric) by (tag_0))", 2, 1).await;
6435 assert_nested_count_rewrite_missing(
6436 "count(sum(irate(some_metric[1h])) by (tag_0))",
6437 2,
6438 300,
6439 )
6440 .await;
6441 }
6442
6443 #[tokio::test]
6444 async fn physical_table_name_is_not_leaked_in_plan() {
6445 let prom_expr = parser::parse("some_metric").unwrap();
6446 let eval_stmt = EvalStmt {
6447 expr: prom_expr,
6448 start: UNIX_EPOCH,
6449 end: UNIX_EPOCH
6450 .checked_add(Duration::from_secs(100_000))
6451 .unwrap(),
6452 interval: Duration::from_secs(5),
6453 lookback_delta: Duration::from_secs(1),
6454 };
6455
6456 let table_provider = build_test_table_provider_with_tsid(
6457 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
6458 1,
6459 1,
6460 )
6461 .await;
6462 let plan =
6463 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6464 .await
6465 .unwrap();
6466
6467 let plan_str = plan.display_indent_schema().to_string();
6468 assert!(plan_str.contains("TableScan: phy"), "{plan}");
6469 assert!(plan_str.contains("SubqueryAlias: some_metric"));
6470 assert!(plan_str.contains("Filter: phy.__table_id = UInt32(1024)"));
6471 assert!(!plan_str.contains("TableScan: some_metric"));
6472 }
6473
6474 #[tokio::test]
6475 async fn sum_without_does_not_group_by_tsid() {
6476 let prom_expr = parser::parse("sum without (tag_0) (some_metric)").unwrap();
6477 let eval_stmt = EvalStmt {
6478 expr: prom_expr,
6479 start: UNIX_EPOCH,
6480 end: UNIX_EPOCH
6481 .checked_add(Duration::from_secs(100_000))
6482 .unwrap(),
6483 interval: Duration::from_secs(5),
6484 lookback_delta: Duration::from_secs(1),
6485 };
6486
6487 let table_provider = build_test_table_provider_with_tsid(
6488 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
6489 1,
6490 1,
6491 )
6492 .await;
6493 let plan =
6494 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6495 .await
6496 .unwrap();
6497
6498 let plan_str = plan.display_indent_schema().to_string();
6499 assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
6500
6501 let aggr_line = plan_str
6502 .lines()
6503 .find(|line| line.contains("Aggregate: groupBy="))
6504 .unwrap();
6505 assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
6506 }
6507
6508 #[tokio::test]
6509 async fn topk_without_does_not_partition_by_tsid() {
6510 let prom_expr = parser::parse("topk without (tag_0) (1, some_metric)").unwrap();
6511 let eval_stmt = EvalStmt {
6512 expr: prom_expr,
6513 start: UNIX_EPOCH,
6514 end: UNIX_EPOCH
6515 .checked_add(Duration::from_secs(100_000))
6516 .unwrap(),
6517 interval: Duration::from_secs(5),
6518 lookback_delta: Duration::from_secs(1),
6519 };
6520
6521 let table_provider = build_test_table_provider_with_tsid(
6522 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
6523 1,
6524 1,
6525 )
6526 .await;
6527 let plan =
6528 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6529 .await
6530 .unwrap();
6531
6532 let plan_str = plan.display_indent_schema().to_string();
6533 assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
6534
6535 let window_line = plan_str
6536 .lines()
6537 .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()"))
6538 .unwrap();
6539 let partition_by = window_line
6540 .split("PARTITION BY [")
6541 .nth(1)
6542 .and_then(|s| s.split("] ORDER BY").next())
6543 .unwrap();
6544 assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
6545 }
6546
6547 #[tokio::test]
6548 async fn sum_by_does_not_group_by_tsid() {
6549 let prom_expr = parser::parse("sum by (__tsid) (some_metric)").unwrap();
6550 let eval_stmt = EvalStmt {
6551 expr: prom_expr,
6552 start: UNIX_EPOCH,
6553 end: UNIX_EPOCH
6554 .checked_add(Duration::from_secs(100_000))
6555 .unwrap(),
6556 interval: Duration::from_secs(5),
6557 lookback_delta: Duration::from_secs(1),
6558 };
6559
6560 let table_provider = build_test_table_provider_with_tsid(
6561 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
6562 1,
6563 1,
6564 )
6565 .await;
6566 let plan =
6567 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6568 .await
6569 .unwrap();
6570
6571 let plan_str = plan.display_indent_schema().to_string();
6572 assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
6573
6574 let aggr_line = plan_str
6575 .lines()
6576 .find(|line| line.contains("Aggregate: groupBy="))
6577 .unwrap();
6578 assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
6579 }
6580
6581 #[tokio::test]
6582 async fn aggregate_over_binary_time_function_expr() {
6583 for op in ["sum", "min", "max", "avg"] {
6584 let prom_expr = parser::parse(&format!(
6585 "{op} by (tag_0, tag_1, tag_2) (time() - some_metric)"
6586 ))
6587 .unwrap();
6588 let eval_stmt = EvalStmt {
6589 expr: prom_expr,
6590 start: UNIX_EPOCH,
6591 end: UNIX_EPOCH
6592 .checked_add(Duration::from_secs(100_000))
6593 .unwrap(),
6594 interval: Duration::from_secs(5),
6595 lookback_delta: Duration::from_secs(1),
6596 };
6597
6598 let table_provider = build_test_table_provider_with_tsid(
6599 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
6600 3,
6601 1,
6602 )
6603 .await;
6604 let plan =
6605 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6606 .await
6607 .unwrap();
6608
6609 let plan_str = plan.display_indent_schema().to_string();
6610 let aggr_line = plan_str
6611 .lines()
6612 .find(|line| line.contains("Aggregate: groupBy="))
6613 .unwrap();
6614 assert!(aggr_line.contains(op), "{plan_str}");
6615 assert!(aggr_line.contains("first_value"), "{plan_str}");
6616 assert!(
6617 !plan
6618 .schema()
6619 .fields()
6620 .iter()
6621 .any(|field| { field.name() == DATA_SCHEMA_TSID_COLUMN_NAME })
6622 );
6623 }
6624 }
6625
6626 #[tokio::test]
6627 async fn topk_by_does_not_partition_by_tsid() {
6628 let prom_expr = parser::parse("topk by (__tsid) (1, some_metric)").unwrap();
6629 let eval_stmt = EvalStmt {
6630 expr: prom_expr,
6631 start: UNIX_EPOCH,
6632 end: UNIX_EPOCH
6633 .checked_add(Duration::from_secs(100_000))
6634 .unwrap(),
6635 interval: Duration::from_secs(5),
6636 lookback_delta: Duration::from_secs(1),
6637 };
6638
6639 let table_provider = build_test_table_provider_with_tsid(
6640 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
6641 1,
6642 1,
6643 )
6644 .await;
6645 let plan =
6646 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6647 .await
6648 .unwrap();
6649
6650 let plan_str = plan.display_indent_schema().to_string();
6651 assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
6652
6653 let window_line = plan_str
6654 .lines()
6655 .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()"))
6656 .unwrap();
6657 let partition_by = window_line
6658 .split("PARTITION BY [")
6659 .nth(1)
6660 .and_then(|s| s.split("] ORDER BY").next())
6661 .unwrap();
6662 assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
6663 }
6664
6665 #[tokio::test]
6666 async fn selector_matcher_on_tsid_does_not_use_internal_column() {
6667 let prom_expr = parser::parse(r#"some_metric{__tsid="123"}"#).unwrap();
6668 let eval_stmt = EvalStmt {
6669 expr: prom_expr,
6670 start: UNIX_EPOCH,
6671 end: UNIX_EPOCH
6672 .checked_add(Duration::from_secs(100_000))
6673 .unwrap(),
6674 interval: Duration::from_secs(5),
6675 lookback_delta: Duration::from_secs(1),
6676 };
6677
6678 let table_provider = build_test_table_provider_with_tsid(
6679 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
6680 1,
6681 1,
6682 )
6683 .await;
6684 let plan =
6685 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6686 .await
6687 .unwrap();
6688
6689 fn collect_filter_cols(plan: &LogicalPlan, out: &mut HashSet<Column>) {
6690 if let LogicalPlan::Filter(filter) = plan {
6691 datafusion_expr::utils::expr_to_columns(&filter.predicate, out).unwrap();
6692 }
6693 for input in plan.inputs() {
6694 collect_filter_cols(input, out);
6695 }
6696 }
6697
6698 let mut filter_cols = HashSet::new();
6699 collect_filter_cols(&plan, &mut filter_cols);
6700 assert!(
6701 !filter_cols
6702 .iter()
6703 .any(|c| c.name == DATA_SCHEMA_TSID_COLUMN_NAME)
6704 );
6705 }
6706
6707 #[tokio::test]
6708 async fn tsid_is_not_used_when_physical_table_is_missing() {
6709 let prom_expr = parser::parse("some_metric").unwrap();
6710 let eval_stmt = EvalStmt {
6711 expr: prom_expr,
6712 start: UNIX_EPOCH,
6713 end: UNIX_EPOCH
6714 .checked_add(Duration::from_secs(100_000))
6715 .unwrap(),
6716 interval: Duration::from_secs(5),
6717 lookback_delta: Duration::from_secs(1),
6718 };
6719
6720 let catalog_list = MemoryCatalogManager::with_default_setup();
6721
6722 let mut columns = vec![ColumnSchema::new(
6724 "tag_0".to_string(),
6725 ConcreteDataType::string_datatype(),
6726 false,
6727 )];
6728 columns.push(
6729 ColumnSchema::new(
6730 "timestamp".to_string(),
6731 ConcreteDataType::timestamp_millisecond_datatype(),
6732 false,
6733 )
6734 .with_time_index(true),
6735 );
6736 columns.push(ColumnSchema::new(
6737 "field_0".to_string(),
6738 ConcreteDataType::float64_datatype(),
6739 true,
6740 ));
6741 let schema = Arc::new(Schema::new(columns));
6742 let mut options = table::requests::TableOptions::default();
6743 options
6744 .extra_options
6745 .insert(LOGICAL_TABLE_METADATA_KEY.to_string(), "phy".to_string());
6746 let table_meta = TableMetaBuilder::empty()
6747 .schema(schema)
6748 .primary_key_indices(vec![0])
6749 .value_indices(vec![2])
6750 .engine(METRIC_ENGINE_NAME.to_string())
6751 .options(options)
6752 .next_column_id(1024)
6753 .build()
6754 .unwrap();
6755 let table_info = TableInfoBuilder::default()
6756 .table_id(1024)
6757 .name("some_metric")
6758 .meta(table_meta)
6759 .build()
6760 .unwrap();
6761 let table = EmptyTable::from_table_info(&table_info);
6762 catalog_list
6763 .register_table_sync(RegisterTableRequest {
6764 catalog: DEFAULT_CATALOG_NAME.to_string(),
6765 schema: DEFAULT_SCHEMA_NAME.to_string(),
6766 table_name: "some_metric".to_string(),
6767 table_id: 1024,
6768 table,
6769 })
6770 .unwrap();
6771
6772 let table_provider = DfTableSourceProvider::new(
6773 catalog_list,
6774 false,
6775 QueryContext::arc(),
6776 DummyDecoder::arc(),
6777 false,
6778 );
6779
6780 let plan =
6781 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6782 .await
6783 .unwrap();
6784
6785 let plan_str = plan.display_indent_schema().to_string();
6786 assert!(plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]"));
6787 assert!(!plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
6788 }
6789
6790 #[tokio::test]
6791 async fn tsid_is_carried_only_when_aggregate_preserves_label_set() {
6792 let prom_expr = parser::parse("sum by (tag_0) (some_metric)").unwrap();
6793 let eval_stmt = EvalStmt {
6794 expr: prom_expr,
6795 start: UNIX_EPOCH,
6796 end: UNIX_EPOCH
6797 .checked_add(Duration::from_secs(100_000))
6798 .unwrap(),
6799 interval: Duration::from_secs(5),
6800 lookback_delta: Duration::from_secs(1),
6801 };
6802
6803 let table_provider = build_test_table_provider_with_tsid(
6804 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
6805 1,
6806 1,
6807 )
6808 .await;
6809 let plan =
6810 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6811 .await
6812 .unwrap();
6813
6814 let plan_str = plan.display_indent_schema().to_string();
6815 assert!(plan_str.contains("first_value") && plan_str.contains("__tsid"));
6816 assert!(
6817 !plan
6818 .schema()
6819 .fields()
6820 .iter()
6821 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
6822 );
6823
6824 let prom_expr = parser::parse("sum(some_metric)").unwrap();
6826 let eval_stmt = EvalStmt {
6827 expr: prom_expr,
6828 start: UNIX_EPOCH,
6829 end: UNIX_EPOCH
6830 .checked_add(Duration::from_secs(100_000))
6831 .unwrap(),
6832 interval: Duration::from_secs(5),
6833 lookback_delta: Duration::from_secs(1),
6834 };
6835 let table_provider = build_test_table_provider_with_tsid(
6836 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
6837 1,
6838 1,
6839 )
6840 .await;
6841 let plan =
6842 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6843 .await
6844 .unwrap();
6845 let plan_str = plan.display_indent_schema().to_string();
6846 assert!(!plan_str.contains("first_value"));
6847 }
6848
6849 #[tokio::test]
6850 async fn or_operator_with_unknown_metric_does_not_require_tsid() {
6851 let prom_expr = parser::parse("unknown_metric or some_metric").unwrap();
6852 let eval_stmt = EvalStmt {
6853 expr: prom_expr,
6854 start: UNIX_EPOCH,
6855 end: UNIX_EPOCH
6856 .checked_add(Duration::from_secs(100_000))
6857 .unwrap(),
6858 interval: Duration::from_secs(5),
6859 lookback_delta: Duration::from_secs(1),
6860 };
6861
6862 let table_provider = build_test_table_provider_with_tsid(
6863 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
6864 1,
6865 1,
6866 )
6867 .await;
6868
6869 let plan =
6870 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6871 .await
6872 .unwrap();
6873
6874 assert!(
6875 !plan
6876 .schema()
6877 .fields()
6878 .iter()
6879 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
6880 );
6881 }
6882
6883 #[tokio::test]
6884 async fn aggregate_avg() {
6885 do_aggregate_expr_plan("avg", "avg").await;
6886 }
6887
6888 #[tokio::test]
6889 #[should_panic] async fn aggregate_count() {
6891 do_aggregate_expr_plan("count", "count").await;
6892 }
6893
6894 #[tokio::test]
6895 async fn aggregate_min() {
6896 do_aggregate_expr_plan("min", "min").await;
6897 }
6898
6899 #[tokio::test]
6900 async fn aggregate_max() {
6901 do_aggregate_expr_plan("max", "max").await;
6902 }
6903
6904 #[tokio::test]
6905 async fn aggregate_group() {
6906 let prom_expr = parser::parse(
6910 "sum(group by (cluster)(kubernetes_build_info{service=\"kubernetes\",job=\"apiserver\"}))",
6911 )
6912 .unwrap();
6913 let eval_stmt = EvalStmt {
6914 expr: prom_expr,
6915 start: UNIX_EPOCH,
6916 end: UNIX_EPOCH
6917 .checked_add(Duration::from_secs(100_000))
6918 .unwrap(),
6919 interval: Duration::from_secs(5),
6920 lookback_delta: Duration::from_secs(1),
6921 };
6922
6923 let table_provider = build_test_table_provider_with_fields(
6924 &[(
6925 DEFAULT_SCHEMA_NAME.to_string(),
6926 "kubernetes_build_info".to_string(),
6927 )],
6928 &["cluster", "service", "job"],
6929 )
6930 .await;
6931 let plan =
6932 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6933 .await
6934 .unwrap();
6935
6936 let plan_str = plan.display_indent_schema().to_string();
6937 assert!(plan_str.contains("max(Float64(1"));
6938 }
6939
6940 #[tokio::test]
6941 async fn aggregate_stddev() {
6942 do_aggregate_expr_plan("stddev", "stddev_pop").await;
6943 }
6944
6945 #[tokio::test]
6946 async fn aggregate_stdvar() {
6947 do_aggregate_expr_plan("stdvar", "var_pop").await;
6948 }
6949
6950 #[tokio::test]
6974 async fn binary_op_column_column() {
6975 let prom_expr =
6976 parser::parse(r#"some_metric{tag_0="foo"} + some_metric{tag_0="bar"}"#).unwrap();
6977 let eval_stmt = EvalStmt {
6978 expr: prom_expr,
6979 start: UNIX_EPOCH,
6980 end: UNIX_EPOCH
6981 .checked_add(Duration::from_secs(100_000))
6982 .unwrap(),
6983 interval: Duration::from_secs(5),
6984 lookback_delta: Duration::from_secs(1),
6985 };
6986
6987 let table_provider = build_test_table_provider(
6988 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
6989 1,
6990 1,
6991 )
6992 .await;
6993 let plan =
6994 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6995 .await
6996 .unwrap();
6997
6998 let expected = String::from(
6999 "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]\
7000 \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]\
7001 \n SubqueryAlias: lhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7002 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7003 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7004 \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]\
7005 \n Filter: 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]\
7006 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7007 \n SubqueryAlias: rhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7008 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7009 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7010 \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]\
7011 \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]\
7012 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
7013 );
7014
7015 assert_eq!(plan.display_indent_schema().to_string(), expected);
7016 }
7017
7018 async fn indie_query_plan_compare<T: AsRef<str>>(query: &str, expected: T) {
7019 let prom_expr = parser::parse(query).unwrap();
7020 let eval_stmt = EvalStmt {
7021 expr: prom_expr,
7022 start: UNIX_EPOCH,
7023 end: UNIX_EPOCH
7024 .checked_add(Duration::from_secs(100_000))
7025 .unwrap(),
7026 interval: Duration::from_secs(5),
7027 lookback_delta: Duration::from_secs(1),
7028 };
7029
7030 let table_provider = build_test_table_provider(
7031 &[
7032 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
7033 (
7034 "greptime_private".to_string(),
7035 "some_alt_metric".to_string(),
7036 ),
7037 ],
7038 1,
7039 1,
7040 )
7041 .await;
7042 let plan =
7043 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7044 .await
7045 .unwrap();
7046
7047 assert_eq!(plan.display_indent_schema().to_string(), expected.as_ref());
7048 }
7049
7050 #[tokio::test]
7051 async fn binary_op_literal_column() {
7052 let query = r#"1 + some_metric{tag_0="bar"}"#;
7053 let expected = String::from(
7054 "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]\
7055 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7056 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7057 \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]\
7058 \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]\
7059 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
7060 );
7061
7062 indie_query_plan_compare(query, expected).await;
7063 }
7064
7065 #[tokio::test]
7066 async fn binary_op_literal_literal() {
7067 let query = r#"1 + 1"#;
7068 let expected = r#"EmptyMetric: range=[0..100000000], interval=[5000] [time:Timestamp(ms), value:Float64;N]
7069 TableScan: dummy [time:Timestamp(ms), value:Float64;N]"#;
7070 indie_query_plan_compare(query, expected).await;
7071 }
7072
7073 #[tokio::test]
7074 async fn simple_bool_grammar() {
7075 let query = "some_metric != bool 1.2345";
7076 let expected = String::from(
7077 "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]\
7078 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7079 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7080 \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]\
7081 \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]\
7082 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
7083 );
7084
7085 indie_query_plan_compare(query, expected).await;
7086 }
7087
7088 #[tokio::test]
7089 async fn bool_with_additional_arithmetic() {
7090 let query = "some_metric + (1 == bool 2)";
7091 let expected = String::from(
7092 "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]\
7093 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7094 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7095 \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]\
7096 \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]\
7097 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
7098 );
7099
7100 indie_query_plan_compare(query, expected).await;
7101 }
7102
7103 #[tokio::test]
7104 async fn simple_unary() {
7105 let query = "-some_metric";
7106 let expected = String::from(
7107 "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]\
7108 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7109 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7110 \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]\
7111 \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]\
7112 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
7113 );
7114
7115 indie_query_plan_compare(query, expected).await;
7116 }
7117
7118 #[tokio::test]
7119 async fn increase_aggr() {
7120 let query = "increase(some_metric[5m])";
7121 let expected = String::from(
7122 "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]\
7123 \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]\
7124 \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))]\
7125 \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7126 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7127 \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]\
7128 \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]\
7129 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
7130 );
7131
7132 indie_query_plan_compare(query, expected).await;
7133 }
7134
7135 #[tokio::test]
7136 async fn less_filter_on_value() {
7137 let query = "some_metric < 1.2345";
7138 let expected = String::from(
7139 "Filter: some_metric.field_0 < Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7140 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7141 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7142 \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]\
7143 \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]\
7144 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
7145 );
7146
7147 indie_query_plan_compare(query, expected).await;
7148 }
7149
7150 #[tokio::test]
7151 async fn count_over_time() {
7152 let query = "count_over_time(some_metric[5m])";
7153 let expected = String::from(
7154 "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]\
7155 \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]\
7156 \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))]\
7157 \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7158 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7159 \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]\
7160 \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]\
7161 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
7162 );
7163
7164 indie_query_plan_compare(query, expected).await;
7165 }
7166
7167 #[tokio::test]
7170 async fn count_over_time_subquery() {
7171 let query = "count_over_time(some_metric[10m:1m])";
7172 let expected = String::from(
7173 "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]\
7174 \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]\
7175 \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))]\
7176 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7177 \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]\
7178 \n PromInstantManipulate: range=[-540000..100000000], lookback=[1000], interval=[60000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7179 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7180 \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]\
7181 \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]\
7182 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
7183 );
7184 indie_query_plan_compare(query, expected).await;
7185 }
7186
7187 #[tokio::test]
7188 async fn test_hash_join() {
7189 let mut eval_stmt = EvalStmt {
7190 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
7191 start: UNIX_EPOCH,
7192 end: UNIX_EPOCH
7193 .checked_add(Duration::from_secs(100_000))
7194 .unwrap(),
7195 interval: Duration::from_secs(5),
7196 lookback_delta: Duration::from_secs(1),
7197 };
7198
7199 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"}"#;
7200
7201 let prom_expr = parser::parse(case).unwrap();
7202 eval_stmt.expr = prom_expr;
7203 let table_provider = build_test_table_provider_with_fields(
7204 &[
7205 (
7206 DEFAULT_SCHEMA_NAME.to_string(),
7207 "http_server_requests_seconds_sum".to_string(),
7208 ),
7209 (
7210 DEFAULT_SCHEMA_NAME.to_string(),
7211 "http_server_requests_seconds_count".to_string(),
7212 ),
7213 ],
7214 &["uri", "kubernetes_namespace", "kubernetes_pod_name"],
7215 )
7216 .await;
7217 let plan =
7219 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7220 .await
7221 .unwrap();
7222 let expected = "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, 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\
7223 \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\
7224 \n SubqueryAlias: http_server_requests_seconds_sum\
7225 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\
7226 \n PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\
7227 \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\
7228 \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)\
7229 \n TableScan: http_server_requests_seconds_sum\
7230 \n SubqueryAlias: http_server_requests_seconds_count\
7231 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\
7232 \n PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\
7233 \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\
7234 \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)\
7235 \n TableScan: http_server_requests_seconds_count";
7236 assert_eq!(plan.to_string(), expected);
7237 }
7238
7239 #[tokio::test]
7240 async fn test_nested_histogram_quantile() {
7241 let mut eval_stmt = EvalStmt {
7242 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
7243 start: UNIX_EPOCH,
7244 end: UNIX_EPOCH
7245 .checked_add(Duration::from_secs(100_000))
7246 .unwrap(),
7247 interval: Duration::from_secs(5),
7248 lookback_delta: Duration::from_secs(1),
7249 };
7250
7251 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]*-(.*)")"#;
7252
7253 let prom_expr = parser::parse(case).unwrap();
7254 eval_stmt.expr = prom_expr;
7255 let table_provider = build_test_table_provider_with_fields(
7256 &[(
7257 DEFAULT_SCHEMA_NAME.to_string(),
7258 "greptime_servers_grpc_requests_elapsed_bucket".to_string(),
7259 )],
7260 &["pod", "le", "path", "code", "container"],
7261 )
7262 .await;
7263 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7265 .await
7266 .unwrap();
7267 }
7268
7269 #[tokio::test]
7270 async fn test_histogram_quantile_binary_op() {
7271 let mut eval_stmt = EvalStmt {
7272 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
7273 start: UNIX_EPOCH,
7274 end: UNIX_EPOCH
7275 .checked_add(Duration::from_secs(100_000))
7276 .unwrap(),
7277 interval: Duration::from_secs(5),
7278 lookback_delta: Duration::from_secs(1),
7279 };
7280
7281 let case = r#"histogram_quantile(0.5, sum by (le, pod) (rate(http_request_duration_seconds_bucket[5m]))) + 0"#;
7285
7286 let prom_expr = parser::parse(case).unwrap();
7287 eval_stmt.expr = prom_expr;
7288 let table_provider = build_test_table_provider_with_fields(
7289 &[(
7290 DEFAULT_SCHEMA_NAME.to_string(),
7291 "http_request_duration_seconds_bucket".to_string(),
7292 )],
7293 &["pod", "le"],
7294 )
7295 .await;
7296 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7298 .await
7299 .unwrap();
7300 }
7301
7302 #[tokio::test]
7303 async fn test_parse_and_operator() {
7304 let mut eval_stmt = EvalStmt {
7305 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
7306 start: UNIX_EPOCH,
7307 end: UNIX_EPOCH
7308 .checked_add(Duration::from_secs(100_000))
7309 .unwrap(),
7310 interval: Duration::from_secs(5),
7311 lookback_delta: Duration::from_secs(1),
7312 };
7313
7314 let cases = [
7315 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)"#,
7316 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)"#,
7317 ];
7318
7319 for case in cases {
7320 let prom_expr = parser::parse(case).unwrap();
7321 eval_stmt.expr = prom_expr;
7322 let table_provider = build_test_table_provider_with_fields(
7323 &[
7324 (
7325 DEFAULT_SCHEMA_NAME.to_string(),
7326 "kubelet_volume_stats_used_bytes".to_string(),
7327 ),
7328 (
7329 DEFAULT_SCHEMA_NAME.to_string(),
7330 "kubelet_volume_stats_capacity_bytes".to_string(),
7331 ),
7332 ],
7333 &["namespace", "persistentvolumeclaim"],
7334 )
7335 .await;
7336 let _ =
7338 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7339 .await
7340 .unwrap();
7341 }
7342 }
7343
7344 #[tokio::test]
7345 async fn test_nested_binary_op() {
7346 let mut eval_stmt = EvalStmt {
7347 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
7348 start: UNIX_EPOCH,
7349 end: UNIX_EPOCH
7350 .checked_add(Duration::from_secs(100_000))
7351 .unwrap(),
7352 interval: Duration::from_secs(5),
7353 lookback_delta: Duration::from_secs(1),
7354 };
7355
7356 let case = r#"sum(rate(nginx_ingress_controller_requests{job=~".*"}[2m])) -
7357 (
7358 sum(rate(nginx_ingress_controller_requests{namespace=~".*"}[2m]))
7359 or
7360 vector(0)
7361 )"#;
7362
7363 let prom_expr = parser::parse(case).unwrap();
7364 eval_stmt.expr = prom_expr;
7365 let table_provider = build_test_table_provider_with_fields(
7366 &[(
7367 DEFAULT_SCHEMA_NAME.to_string(),
7368 "nginx_ingress_controller_requests".to_string(),
7369 )],
7370 &["namespace", "job"],
7371 )
7372 .await;
7373 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7375 .await
7376 .unwrap();
7377 }
7378
7379 #[tokio::test]
7380 async fn test_parse_or_operator() {
7381 let mut eval_stmt = EvalStmt {
7382 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
7383 start: UNIX_EPOCH,
7384 end: UNIX_EPOCH
7385 .checked_add(Duration::from_secs(100_000))
7386 .unwrap(),
7387 interval: Duration::from_secs(5),
7388 lookback_delta: Duration::from_secs(1),
7389 };
7390
7391 let case = r#"
7392 sum(rate(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}[120s])) by (cluster_name,tenant_name) /
7393 (sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) * 100)
7394 or
7395 200 * sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) /
7396 sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)"#;
7397
7398 let table_provider = build_test_table_provider_with_fields(
7399 &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())],
7400 &["tenant_name", "cluster_name"],
7401 )
7402 .await;
7403 eval_stmt.expr = parser::parse(case).unwrap();
7404 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7405 .await
7406 .unwrap();
7407
7408 let case = r#"sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
7409 (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) +
7410 sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
7411 (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0
7412 or
7413 sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
7414 (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0
7415 or
7416 sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
7417 (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0"#;
7418 let table_provider = build_test_table_provider_with_fields(
7419 &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())],
7420 &["tenant_name", "cluster_name"],
7421 )
7422 .await;
7423 eval_stmt.expr = parser::parse(case).unwrap();
7424 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7425 .await
7426 .unwrap();
7427
7428 let case = r#"(sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) +
7429 sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or
7430 (sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or
7431 (sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name))"#;
7432 let table_provider = build_test_table_provider_with_fields(
7433 &[
7434 (
7435 DEFAULT_SCHEMA_NAME.to_string(),
7436 "background_waitevent_cnt".to_string(),
7437 ),
7438 (
7439 DEFAULT_SCHEMA_NAME.to_string(),
7440 "foreground_waitevent_cnt".to_string(),
7441 ),
7442 ],
7443 &["tenant_name", "cluster_name"],
7444 )
7445 .await;
7446 eval_stmt.expr = parser::parse(case).unwrap();
7447 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7448 .await
7449 .unwrap();
7450
7451 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)"#;
7452 let table_provider = build_test_table_provider_with_fields(
7453 &[
7454 (DEFAULT_SCHEMA_NAME.to_string(), "node_load1".to_string()),
7455 (
7456 DEFAULT_SCHEMA_NAME.to_string(),
7457 "container_cpu_load_average_10s".to_string(),
7458 ),
7459 (
7460 DEFAULT_SCHEMA_NAME.to_string(),
7461 "container_spec_cpu_quota".to_string(),
7462 ),
7463 ],
7464 &["cluster_name", "host_name"],
7465 )
7466 .await;
7467 eval_stmt.expr = parser::parse(case).unwrap();
7468 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7469 .await
7470 .unwrap();
7471 }
7472
7473 #[tokio::test]
7474 async fn value_matcher() {
7475 let mut eval_stmt = EvalStmt {
7477 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
7478 start: UNIX_EPOCH,
7479 end: UNIX_EPOCH
7480 .checked_add(Duration::from_secs(100_000))
7481 .unwrap(),
7482 interval: Duration::from_secs(5),
7483 lookback_delta: Duration::from_secs(1),
7484 };
7485
7486 let cases = [
7487 (
7489 r#"some_metric{__field__="field_1"}"#,
7490 vec![
7491 "some_metric.field_1",
7492 "some_metric.tag_0",
7493 "some_metric.tag_1",
7494 "some_metric.tag_2",
7495 "some_metric.timestamp",
7496 ],
7497 ),
7498 (
7500 r#"some_metric{__field__="field_1", __field__="field_0"}"#,
7501 vec![
7502 "some_metric.field_0",
7503 "some_metric.field_1",
7504 "some_metric.tag_0",
7505 "some_metric.tag_1",
7506 "some_metric.tag_2",
7507 "some_metric.timestamp",
7508 ],
7509 ),
7510 (
7512 r#"some_metric{__field__!="field_1"}"#,
7513 vec![
7514 "some_metric.field_0",
7515 "some_metric.field_2",
7516 "some_metric.tag_0",
7517 "some_metric.tag_1",
7518 "some_metric.tag_2",
7519 "some_metric.timestamp",
7520 ],
7521 ),
7522 (
7524 r#"some_metric{__field__!="field_1", __field__!="field_2"}"#,
7525 vec![
7526 "some_metric.field_0",
7527 "some_metric.tag_0",
7528 "some_metric.tag_1",
7529 "some_metric.tag_2",
7530 "some_metric.timestamp",
7531 ],
7532 ),
7533 (
7535 r#"some_metric{__field__="field_1", __field__!="field_0"}"#,
7536 vec![
7537 "some_metric.field_1",
7538 "some_metric.tag_0",
7539 "some_metric.tag_1",
7540 "some_metric.tag_2",
7541 "some_metric.timestamp",
7542 ],
7543 ),
7544 (
7546 r#"some_metric{__field__="field_2", __field__!="field_2"}"#,
7547 vec![
7548 "some_metric.tag_0",
7549 "some_metric.tag_1",
7550 "some_metric.tag_2",
7551 "some_metric.timestamp",
7552 ],
7553 ),
7554 (
7556 r#"some_metric{__field__=~"field_1|field_2"}"#,
7557 vec![
7558 "some_metric.field_1",
7559 "some_metric.field_2",
7560 "some_metric.tag_0",
7561 "some_metric.tag_1",
7562 "some_metric.tag_2",
7563 "some_metric.timestamp",
7564 ],
7565 ),
7566 (
7568 r#"some_metric{__field__!~"field_1|field_2"}"#,
7569 vec![
7570 "some_metric.field_0",
7571 "some_metric.tag_0",
7572 "some_metric.tag_1",
7573 "some_metric.tag_2",
7574 "some_metric.timestamp",
7575 ],
7576 ),
7577 ];
7578
7579 for case in cases {
7580 let prom_expr = parser::parse(case.0).unwrap();
7581 eval_stmt.expr = prom_expr;
7582 let table_provider = build_test_table_provider(
7583 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
7584 3,
7585 3,
7586 )
7587 .await;
7588 let plan =
7589 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7590 .await
7591 .unwrap();
7592 let mut fields = plan.schema().field_names();
7593 let mut expected = case.1.into_iter().map(String::from).collect::<Vec<_>>();
7594 fields.sort();
7595 expected.sort();
7596 assert_eq!(fields, expected, "case: {:?}", case.0);
7597 }
7598
7599 let bad_cases = [
7600 r#"some_metric{__field__="nonexistent"}"#,
7601 r#"some_metric{__field__!="nonexistent"}"#,
7602 ];
7603
7604 for case in bad_cases {
7605 let prom_expr = parser::parse(case).unwrap();
7606 eval_stmt.expr = prom_expr;
7607 let table_provider = build_test_table_provider(
7608 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
7609 3,
7610 3,
7611 )
7612 .await;
7613 let plan =
7614 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7615 .await;
7616 assert!(plan.is_err(), "case: {:?}", case);
7617 }
7618 }
7619
7620 #[tokio::test]
7621 async fn custom_schema() {
7622 let query = "some_alt_metric{__schema__=\"greptime_private\"}";
7623 let expected = String::from(
7624 "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7625 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7626 \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]\
7627 \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]\
7628 \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
7629 );
7630
7631 indie_query_plan_compare(query, expected).await;
7632
7633 let query = "some_alt_metric{__database__=\"greptime_private\"}";
7634 let expected = String::from(
7635 "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7636 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7637 \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]\
7638 \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]\
7639 \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
7640 );
7641
7642 indie_query_plan_compare(query, expected).await;
7643
7644 let query = "some_alt_metric{__schema__=\"greptime_private\"} / some_metric";
7645 let expected = String::from(
7646 "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]\
7647 \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]\
7648 \n SubqueryAlias: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7649 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7650 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7651 \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]\
7652 \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]\
7653 \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7654 \n SubqueryAlias: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7655 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7656 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7657 \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]\
7658 \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]\
7659 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
7660 );
7661
7662 indie_query_plan_compare(query, expected).await;
7663 }
7664
7665 #[tokio::test]
7666 async fn only_equals_is_supported_for_special_matcher() {
7667 let queries = &[
7668 "some_alt_metric{__schema__!=\"greptime_private\"}",
7669 "some_alt_metric{__schema__=~\"lalala\"}",
7670 "some_alt_metric{__database__!=\"greptime_private\"}",
7671 "some_alt_metric{__database__=~\"lalala\"}",
7672 ];
7673
7674 for query in queries {
7675 let prom_expr = parser::parse(query).unwrap();
7676 let eval_stmt = EvalStmt {
7677 expr: prom_expr,
7678 start: UNIX_EPOCH,
7679 end: UNIX_EPOCH
7680 .checked_add(Duration::from_secs(100_000))
7681 .unwrap(),
7682 interval: Duration::from_secs(5),
7683 lookback_delta: Duration::from_secs(1),
7684 };
7685
7686 let table_provider = build_test_table_provider(
7687 &[
7688 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
7689 (
7690 "greptime_private".to_string(),
7691 "some_alt_metric".to_string(),
7692 ),
7693 ],
7694 1,
7695 1,
7696 )
7697 .await;
7698
7699 let plan =
7700 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7701 .await;
7702 assert!(plan.is_err(), "query: {:?}", query);
7703 }
7704 }
7705
7706 #[tokio::test]
7707 async fn test_non_ms_precision() {
7708 let catalog_list = MemoryCatalogManager::with_default_setup();
7709 let columns = vec![
7710 ColumnSchema::new(
7711 "tag".to_string(),
7712 ConcreteDataType::string_datatype(),
7713 false,
7714 ),
7715 ColumnSchema::new(
7716 "timestamp".to_string(),
7717 ConcreteDataType::timestamp_nanosecond_datatype(),
7718 false,
7719 )
7720 .with_time_index(true),
7721 ColumnSchema::new(
7722 "field".to_string(),
7723 ConcreteDataType::float64_datatype(),
7724 true,
7725 ),
7726 ];
7727 let schema = Arc::new(Schema::new(columns));
7728 let table_meta = TableMetaBuilder::empty()
7729 .schema(schema)
7730 .primary_key_indices(vec![0])
7731 .value_indices(vec![2])
7732 .next_column_id(1024)
7733 .build()
7734 .unwrap();
7735 let table_info = TableInfoBuilder::default()
7736 .name("metrics".to_string())
7737 .meta(table_meta)
7738 .build()
7739 .unwrap();
7740 let table = EmptyTable::from_table_info(&table_info);
7741 assert!(
7742 catalog_list
7743 .register_table_sync(RegisterTableRequest {
7744 catalog: DEFAULT_CATALOG_NAME.to_string(),
7745 schema: DEFAULT_SCHEMA_NAME.to_string(),
7746 table_name: "metrics".to_string(),
7747 table_id: 1024,
7748 table,
7749 })
7750 .is_ok()
7751 );
7752
7753 let plan = PromPlanner::stmt_to_plan(
7754 DfTableSourceProvider::new(
7755 catalog_list.clone(),
7756 false,
7757 QueryContext::arc(),
7758 DummyDecoder::arc(),
7759 true,
7760 ),
7761 &EvalStmt {
7762 expr: parser::parse("metrics{tag = \"1\"}").unwrap(),
7763 start: UNIX_EPOCH,
7764 end: UNIX_EPOCH
7765 .checked_add(Duration::from_secs(100_000))
7766 .unwrap(),
7767 interval: Duration::from_secs(5),
7768 lookback_delta: Duration::from_secs(1),
7769 },
7770 &build_query_engine_state(),
7771 )
7772 .await
7773 .unwrap();
7774 assert_eq!(
7775 plan.display_indent_schema().to_string(),
7776 "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
7777 \n PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
7778 \n Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
7779 \n Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp >= TimestampMillisecond(-999, None) AND metrics.timestamp <= TimestampMillisecond(100000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
7780 \n Projection: metrics.field, metrics.tag, CAST(metrics.timestamp AS Timestamp(ms)) AS timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
7781 \n TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]"
7782 );
7783 let plan = PromPlanner::stmt_to_plan(
7784 DfTableSourceProvider::new(
7785 catalog_list.clone(),
7786 false,
7787 QueryContext::arc(),
7788 DummyDecoder::arc(),
7789 true,
7790 ),
7791 &EvalStmt {
7792 expr: parser::parse("avg_over_time(metrics{tag = \"1\"}[5s])").unwrap(),
7793 start: UNIX_EPOCH,
7794 end: UNIX_EPOCH
7795 .checked_add(Duration::from_secs(100_000))
7796 .unwrap(),
7797 interval: Duration::from_secs(5),
7798 lookback_delta: Duration::from_secs(1),
7799 },
7800 &build_query_engine_state(),
7801 )
7802 .await
7803 .unwrap();
7804 assert_eq!(
7805 plan.display_indent_schema().to_string(),
7806 "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]\
7807 \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]\
7808 \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))]\
7809 \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
7810 \n PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
7811 \n Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
7812 \n Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp >= TimestampMillisecond(-4999, None) AND metrics.timestamp <= TimestampMillisecond(100000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
7813 \n Projection: metrics.field, metrics.tag, CAST(metrics.timestamp AS Timestamp(ms)) AS timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
7814 \n TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]"
7815 );
7816 }
7817
7818 #[tokio::test]
7819 async fn test_nonexistent_label() {
7820 let mut eval_stmt = EvalStmt {
7822 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
7823 start: UNIX_EPOCH,
7824 end: UNIX_EPOCH
7825 .checked_add(Duration::from_secs(100_000))
7826 .unwrap(),
7827 interval: Duration::from_secs(5),
7828 lookback_delta: Duration::from_secs(1),
7829 };
7830
7831 let case = r#"some_metric{nonexistent="hi"}"#;
7832 let prom_expr = parser::parse(case).unwrap();
7833 eval_stmt.expr = prom_expr;
7834 let table_provider = build_test_table_provider(
7835 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
7836 3,
7837 3,
7838 )
7839 .await;
7840 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7842 .await
7843 .unwrap();
7844 }
7845
7846 #[tokio::test]
7847 async fn test_label_join() {
7848 let prom_expr = parser::parse(
7849 "label_join(up{tag_0='api-server'}, 'foo', ',', 'tag_1', 'tag_2', 'tag_3')",
7850 )
7851 .unwrap();
7852 let eval_stmt = EvalStmt {
7853 expr: prom_expr,
7854 start: UNIX_EPOCH,
7855 end: UNIX_EPOCH
7856 .checked_add(Duration::from_secs(100_000))
7857 .unwrap(),
7858 interval: Duration::from_secs(5),
7859 lookback_delta: Duration::from_secs(1),
7860 };
7861
7862 let table_provider =
7863 build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 4, 1)
7864 .await;
7865 let plan =
7866 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7867 .await
7868 .unwrap();
7869
7870 let expected = r#"
7871Filter: 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]
7872 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]
7873 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]
7874 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]
7875 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]
7876 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]
7877 TableScan: up [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#;
7878
7879 let ret = plan.display_indent_schema().to_string();
7880 assert_eq!(format!("\n{ret}"), expected, "\n{}", ret);
7881 }
7882
7883 #[tokio::test]
7884 async fn test_label_replace() {
7885 let prom_expr = parser::parse(
7886 "label_replace(up{tag_0=\"a:c\"}, \"foo\", \"$1\", \"tag_0\", \"(.*):.*\")",
7887 )
7888 .unwrap();
7889 let eval_stmt = EvalStmt {
7890 expr: prom_expr,
7891 start: UNIX_EPOCH,
7892 end: UNIX_EPOCH
7893 .checked_add(Duration::from_secs(100_000))
7894 .unwrap(),
7895 interval: Duration::from_secs(5),
7896 lookback_delta: Duration::from_secs(1),
7897 };
7898
7899 let table_provider =
7900 build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 1, 1)
7901 .await;
7902 let plan =
7903 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7904 .await
7905 .unwrap();
7906
7907 let expected = r#"
7908Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8]
7909 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]
7910 PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
7911 PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
7912 Sort: up.tag_0 ASC NULLS FIRST, up.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
7913 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]
7914 TableScan: up [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#;
7915
7916 let ret = plan.display_indent_schema().to_string();
7917 assert_eq!(format!("\n{ret}"), expected, "\n{}", ret);
7918 }
7919
7920 #[tokio::test]
7921 async fn test_matchers_to_expr() {
7922 let mut eval_stmt = EvalStmt {
7923 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
7924 start: UNIX_EPOCH,
7925 end: UNIX_EPOCH
7926 .checked_add(Duration::from_secs(100_000))
7927 .unwrap(),
7928 interval: Duration::from_secs(5),
7929 lookback_delta: Duration::from_secs(1),
7930 };
7931 let case =
7932 r#"sum(prometheus_tsdb_head_series{tag_1=~"(10.0.160.237:8080|10.0.160.237:9090)"})"#;
7933
7934 let prom_expr = parser::parse(case).unwrap();
7935 eval_stmt.expr = prom_expr;
7936 let table_provider = build_test_table_provider(
7937 &[(
7938 DEFAULT_SCHEMA_NAME.to_string(),
7939 "prometheus_tsdb_head_series".to_string(),
7940 )],
7941 3,
7942 3,
7943 )
7944 .await;
7945 let plan =
7946 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7947 .await
7948 .unwrap();
7949 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]\
7950 \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]\
7951 \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]\
7952 \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]\
7953 \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]\
7954 \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]\
7955 \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]";
7956 assert_eq!(plan.display_indent_schema().to_string(), expected);
7957 }
7958
7959 #[tokio::test]
7960 async fn test_topk_expr() {
7961 let mut eval_stmt = EvalStmt {
7962 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
7963 start: UNIX_EPOCH,
7964 end: UNIX_EPOCH
7965 .checked_add(Duration::from_secs(100_000))
7966 .unwrap(),
7967 interval: Duration::from_secs(5),
7968 lookback_delta: Duration::from_secs(1),
7969 };
7970 let case = r#"topk(10, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#;
7971
7972 let prom_expr = parser::parse(case).unwrap();
7973 eval_stmt.expr = prom_expr;
7974 let table_provider = build_test_table_provider_with_fields(
7975 &[
7976 (
7977 DEFAULT_SCHEMA_NAME.to_string(),
7978 "prometheus_tsdb_head_series".to_string(),
7979 ),
7980 (
7981 DEFAULT_SCHEMA_NAME.to_string(),
7982 "http_server_requests_seconds_count".to_string(),
7983 ),
7984 ],
7985 &["ip"],
7986 )
7987 .await;
7988
7989 let plan =
7990 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7991 .await
7992 .unwrap();
7993 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)]\
7994 \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]\
7995 \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]\
7996 \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]\
7997 \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]\
7998 \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]\
7999 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
8000 \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
8001 \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]\
8002 \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]\
8003 \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]";
8004
8005 assert_eq!(plan.display_indent_schema().to_string(), expected);
8006 }
8007
8008 #[tokio::test]
8009 async fn test_count_values_expr() {
8010 let mut eval_stmt = EvalStmt {
8011 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
8012 start: UNIX_EPOCH,
8013 end: UNIX_EPOCH
8014 .checked_add(Duration::from_secs(100_000))
8015 .unwrap(),
8016 interval: Duration::from_secs(5),
8017 lookback_delta: Duration::from_secs(1),
8018 };
8019 let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#;
8020
8021 let prom_expr = parser::parse(case).unwrap();
8022 eval_stmt.expr = prom_expr;
8023 let table_provider = build_test_table_provider_with_fields(
8024 &[
8025 (
8026 DEFAULT_SCHEMA_NAME.to_string(),
8027 "prometheus_tsdb_head_series".to_string(),
8028 ),
8029 (
8030 DEFAULT_SCHEMA_NAME.to_string(),
8031 "http_server_requests_seconds_count".to_string(),
8032 ),
8033 ],
8034 &["ip"],
8035 )
8036 .await;
8037
8038 let plan =
8039 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
8040 .await
8041 .unwrap();
8042 let expected = "Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N]\
8043 \n Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.greptime_value ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N, greptime_value:Float64;N]\
8044 \n Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value AS series, prometheus_tsdb_head_series.greptime_value [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N, greptime_value:Float64;N]\
8045 \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N, count(prometheus_tsdb_head_series.greptime_value):Int64]\
8046 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
8047 \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
8048 \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]\
8049 \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]\
8050 \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]";
8051
8052 assert_eq!(plan.display_indent_schema().to_string(), expected);
8053 }
8054
8055 #[tokio::test]
8056 async fn test_value_alias() {
8057 let mut eval_stmt = EvalStmt {
8058 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
8059 start: UNIX_EPOCH,
8060 end: UNIX_EPOCH
8061 .checked_add(Duration::from_secs(100_000))
8062 .unwrap(),
8063 interval: Duration::from_secs(5),
8064 lookback_delta: Duration::from_secs(1),
8065 };
8066 let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#;
8067
8068 let prom_expr = parser::parse(case).unwrap();
8069 eval_stmt.expr = prom_expr;
8070 eval_stmt = QueryLanguageParser::apply_alias_extension(eval_stmt, "my_series");
8071 let table_provider = build_test_table_provider_with_fields(
8072 &[
8073 (
8074 DEFAULT_SCHEMA_NAME.to_string(),
8075 "prometheus_tsdb_head_series".to_string(),
8076 ),
8077 (
8078 DEFAULT_SCHEMA_NAME.to_string(),
8079 "http_server_requests_seconds_count".to_string(),
8080 ),
8081 ],
8082 &["ip"],
8083 )
8084 .await;
8085
8086 let plan =
8087 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
8088 .await
8089 .unwrap();
8090 let expected = r#"
8091Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp [my_series:Int64, ip:Utf8, greptime_timestamp:Timestamp(ms)]
8092 Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N]
8093 Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.greptime_value ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N, greptime_value:Float64;N]
8094 Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value AS series, prometheus_tsdb_head_series.greptime_value [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N, greptime_value:Float64;N]
8095 Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N, count(prometheus_tsdb_head_series.greptime_value):Int64]
8096 PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
8097 PromSeriesDivide: tags=["ip"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
8098 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]
8099 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]
8100 TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"#;
8101 assert_eq!(format!("\n{}", plan.display_indent_schema()), expected);
8102 }
8103
8104 #[tokio::test]
8105 async fn test_quantile_expr() {
8106 let mut eval_stmt = EvalStmt {
8107 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
8108 start: UNIX_EPOCH,
8109 end: UNIX_EPOCH
8110 .checked_add(Duration::from_secs(100_000))
8111 .unwrap(),
8112 interval: Duration::from_secs(5),
8113 lookback_delta: Duration::from_secs(1),
8114 };
8115 let case = r#"quantile(0.3, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#;
8116
8117 let prom_expr = parser::parse(case).unwrap();
8118 eval_stmt.expr = prom_expr;
8119 let table_provider = build_test_table_provider_with_fields(
8120 &[
8121 (
8122 DEFAULT_SCHEMA_NAME.to_string(),
8123 "prometheus_tsdb_head_series".to_string(),
8124 ),
8125 (
8126 DEFAULT_SCHEMA_NAME.to_string(),
8127 "http_server_requests_seconds_count".to_string(),
8128 ),
8129 ],
8130 &["ip"],
8131 )
8132 .await;
8133
8134 let plan =
8135 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
8136 .await
8137 .unwrap();
8138 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]\
8139 \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]\
8140 \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]\
8141 \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]\
8142 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
8143 \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
8144 \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]\
8145 \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]\
8146 \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]";
8147
8148 assert_eq!(plan.display_indent_schema().to_string(), expected);
8149 }
8150
8151 #[tokio::test]
8152 async fn test_or_not_exists_table_label() {
8153 let mut eval_stmt = EvalStmt {
8154 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
8155 start: UNIX_EPOCH,
8156 end: UNIX_EPOCH
8157 .checked_add(Duration::from_secs(100_000))
8158 .unwrap(),
8159 interval: Duration::from_secs(5),
8160 lookback_delta: Duration::from_secs(1),
8161 };
8162 let case = r#"sum by (job, tag0, tag2) (metric_exists) or sum by (job, tag0, tag2) (metric_not_exists)"#;
8163
8164 let prom_expr = parser::parse(case).unwrap();
8165 eval_stmt.expr = prom_expr;
8166 let table_provider = build_test_table_provider_with_fields(
8167 &[(DEFAULT_SCHEMA_NAME.to_string(), "metric_exists".to_string())],
8168 &["job"],
8169 )
8170 .await;
8171
8172 let plan =
8173 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
8174 .await
8175 .unwrap();
8176 let expected = r#"UnionDistinctOn: on col=[["job"]], ts_col=[greptime_timestamp] [greptime_timestamp:Timestamp(ms), job:Utf8, sum(metric_exists.greptime_value):Float64;N]
8177 SubqueryAlias: metric_exists [greptime_timestamp:Timestamp(ms), job:Utf8, sum(metric_exists.greptime_value):Float64;N]
8178 Projection: metric_exists.greptime_timestamp, metric_exists.job, sum(metric_exists.greptime_value) [greptime_timestamp:Timestamp(ms), job:Utf8, sum(metric_exists.greptime_value):Float64;N]
8179 Sort: metric_exists.job ASC NULLS LAST, metric_exists.greptime_timestamp ASC NULLS LAST [job:Utf8, greptime_timestamp:Timestamp(ms), sum(metric_exists.greptime_value):Float64;N]
8180 Aggregate: groupBy=[[metric_exists.job, metric_exists.greptime_timestamp]], aggr=[[sum(metric_exists.greptime_value)]] [job:Utf8, greptime_timestamp:Timestamp(ms), sum(metric_exists.greptime_value):Float64;N]
8181 PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [job:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
8182 PromSeriesDivide: tags=["job"] [job:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
8183 Sort: metric_exists.job ASC NULLS FIRST, metric_exists.greptime_timestamp ASC NULLS FIRST [job:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
8184 Filter: metric_exists.greptime_timestamp >= TimestampMillisecond(-999, None) AND metric_exists.greptime_timestamp <= TimestampMillisecond(100000000, None) [job:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
8185 TableScan: metric_exists [job:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
8186 SubqueryAlias: [greptime_timestamp:Timestamp(ms), job:Utf8;N, sum(.value):Float64;N]
8187 Projection: .time AS greptime_timestamp, Utf8(NULL) AS job, sum(.value) [greptime_timestamp:Timestamp(ms), job:Utf8;N, sum(.value):Float64;N]
8188 Sort: .time ASC NULLS LAST [time:Timestamp(ms), sum(.value):Float64;N]
8189 Aggregate: groupBy=[[.time]], aggr=[[sum(.value)]] [time:Timestamp(ms), sum(.value):Float64;N]
8190 EmptyMetric: range=[0..-1], interval=[5000] [time:Timestamp(ms), value:Float64;N]
8191 TableScan: dummy [time:Timestamp(ms), value:Float64;N]"#;
8192
8193 assert_eq!(plan.display_indent_schema().to_string(), expected);
8194 }
8195
8196 #[tokio::test]
8197 async fn test_histogram_quantile_missing_le_column() {
8198 let mut eval_stmt = EvalStmt {
8199 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
8200 start: UNIX_EPOCH,
8201 end: UNIX_EPOCH
8202 .checked_add(Duration::from_secs(100_000))
8203 .unwrap(),
8204 interval: Duration::from_secs(5),
8205 lookback_delta: Duration::from_secs(1),
8206 };
8207
8208 let case = r#"histogram_quantile(0.99, sum by(pod,instance,le) (rate(non_existent_histogram_bucket{instance=~"xxx"}[1m])))"#;
8210
8211 let prom_expr = parser::parse(case).unwrap();
8212 eval_stmt.expr = prom_expr;
8213
8214 let table_provider = build_test_table_provider_with_fields(
8216 &[(
8217 DEFAULT_SCHEMA_NAME.to_string(),
8218 "non_existent_histogram_bucket".to_string(),
8219 )],
8220 &["pod", "instance"], )
8222 .await;
8223
8224 let result =
8226 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
8227 .await;
8228
8229 assert!(
8231 result.is_ok(),
8232 "Expected successful plan creation with empty result, but got error: {:?}",
8233 result.err()
8234 );
8235
8236 let plan = result.unwrap();
8238 match plan {
8239 LogicalPlan::EmptyRelation(_) => {
8240 }
8242 _ => panic!("Expected EmptyRelation, but got: {:?}", plan),
8243 }
8244 }
8245}