1mod matching_filters;
16
17use std::collections::{BTreeSet, HashMap, HashSet, VecDeque};
18use std::sync::Arc;
19use std::time::UNIX_EPOCH;
20
21use arrow::datatypes::IntervalDayTime;
22use async_recursion::async_recursion;
23use catalog::table_source::DfTableSourceProvider;
24use common_error::ext::ErrorExt;
25use common_error::status_code::StatusCode;
26use common_function::function::FunctionContext;
27use common_query::native_histogram::native_histogram_value_type;
28use common_query::prelude::{
29 GREPTIME_TEMPORALITY_DELTA, OTLP_AGGREGATION_TEMPORALITY_LABEL, greptime_native_histogram,
30 greptime_value,
31};
32use common_query::promql_annotations::PromqlAnnotationCollector;
33use datafusion::common::DFSchemaRef;
34use datafusion::datasource::DefaultTableSource;
35use datafusion::functions_aggregate::average::avg_udaf;
36use datafusion::functions_aggregate::count::count_udaf;
37use datafusion::functions_aggregate::expr_fn::first_value;
38use datafusion::functions_aggregate::min_max::{max_udaf, min_udaf};
39use datafusion::functions_aggregate::stddev::stddev_pop_udaf;
40use datafusion::functions_aggregate::sum::sum_udaf;
41use datafusion::functions_aggregate::variance::var_pop_udaf;
42use datafusion::functions_window::row_number::RowNumber;
43use datafusion::logical_expr::expr::{Alias, ScalarFunction, WindowFunction};
44use datafusion::logical_expr::expr_rewriter::normalize_cols;
45use datafusion::logical_expr::{
46 BinaryExpr, Cast, Extension, LogicalPlan, LogicalPlanBuilder, Operator,
47 ScalarUDF as ScalarUdfDef, WindowFrame, WindowFunctionDefinition,
48};
49use datafusion::optimizer::simplify_expressions::ExprSimplifier;
50use datafusion::prelude as df_prelude;
51use datafusion::prelude::{Column, Expr as DfExpr, JoinType};
52use datafusion::scalar::ScalarValue;
53use datafusion_common::{DFSchema, NullEquality, TableReference};
54use datafusion_expr::expr::WindowFunctionParams;
55use datafusion_expr::expr_fn::when;
56use datafusion_expr::simplify::SimplifyContext;
57use datafusion_expr::utils::{conjunction, disjunction};
58use datafusion_expr::{ExprSchemable, Literal, SortExpr, TableSource, col, lit};
59use datafusion_functions::core::coalesce;
60use datatypes::arrow::datatypes::{DataType as ArrowDataType, TimeUnit as ArrowTimeUnit};
61use datatypes::data_type::{ConcreteDataType, DataType as GreptimeDataType};
62use itertools::Itertools;
63use once_cell::sync::Lazy;
64use promql::extension_plan::{
65 Absent, EmptyMetric, HistogramFold, HistogramFoldOperation, InstantManipulate, Millisecond,
66 RangeManipulate, ScalarCalculate, SeriesDivide, SeriesNormalize, UnionDistinctOn,
67 build_special_time_expr,
68};
69use promql::functions::{
70 AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing,
71 IDelta, Increase, LastOverTime, MatchGroupViolation, MaxOverTime, MinOverTime, MixedRange,
72 NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAggAvg,
73 NativeHistogramAggSum, NativeHistogramAvg, NativeHistogramAvgOverTime, NativeHistogramChanges,
74 NativeHistogramCount, NativeHistogramCountOverTime, NativeHistogramDelta,
75 NativeHistogramDivScalar, NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction,
76 NativeHistogramIDelta, NativeHistogramIRate, NativeHistogramIncrease,
77 NativeHistogramLastOverTime, NativeHistogramMulScalar, NativeHistogramNeg,
78 NativeHistogramNotEq, NativeHistogramPresentOverTime, NativeHistogramQuantile,
79 NativeHistogramRate, NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev,
80 NativeHistogramStdvar, NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime,
81 NativeHistogramToString, PredictLinear, PresentOverTime, PromqlFloatToString, QuantileOverTime,
82 Rate, Resets, Round, StddevOverTime, StdvarOverTime, SumOverTime, UniqueMatchGroup,
83 quantile_udaf,
84};
85use promql_parser::label::{METRIC_NAME, MatchOp, Matcher, Matchers};
86use promql_parser::parser::token::TokenType;
87use promql_parser::parser::value::ValueType;
88use promql_parser::parser::{
89 AggregateExpr, BinModifier, BinaryExpr as PromBinaryExpr, Call, EvalStmt, Expr as PromExpr,
90 Function, FunctionArgs as PromFunctionArgs, LabelModifier, MatrixSelector, NumberLiteral,
91 Offset, ParenExpr, StringLiteral, SubqueryExpr, UnaryExpr, VectorMatchCardinality,
92 VectorSelector, token,
93};
94use regex::{self, Regex};
95use snafu::{OptionExt, ResultExt, ensure};
96use store_api::metric_engine_consts::{
97 DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, LOGICAL_TABLE_METADATA_KEY,
98 METRIC_ENGINE_NAME, is_metric_engine_internal_column,
99};
100use table::table::adapter::DfTableProviderAdapter;
101
102use crate::parser::{
103 ALIAS_NODE_NAME, ANALYZE_NODE_NAME, ANALYZE_VERBOSE_NODE_NAME, AliasExpr, EXPLAIN_NODE_NAME,
104 EXPLAIN_VERBOSE_NODE_NAME,
105};
106use crate::promql::error::{
107 CatalogSnafu, ColumnNotFoundSnafu, CombineTableColumnMismatchSnafu, DataFusionPlanningSnafu,
108 ExpectRangeSelectorSnafu, FunctionInvalidArgumentSnafu, InvalidDestinationLabelNameSnafu,
109 InvalidRegularExpressionSnafu, InvalidTimeRangeSnafu, MultiFieldsNotSupportedSnafu,
110 MultipleMetricMatchersSnafu, MultipleVectorSnafu, NoMetricMatcherSnafu, PromqlPlanNodeSnafu,
111 Result, SameLabelSetSnafu, TableNameNotFoundSnafu, TimeIndexNotFoundSnafu,
112 UnexpectedPlanExprSnafu, UnexpectedTokenSnafu, UnknownTableSnafu, UnsupportedExprSnafu,
113 UnsupportedMatcherOpSnafu, UnsupportedVectorMatchSnafu, ValueNotFoundSnafu,
114 ZeroRangeSelectorSnafu,
115};
116use crate::query_engine::QueryEngineState;
117
118const SPECIAL_TIME_FUNCTION: &str = "time";
120const SCALAR_FUNCTION: &str = "scalar";
122const SPECIAL_ABSENT_FUNCTION: &str = "absent";
124const SPECIAL_HISTOGRAM_QUANTILE: &str = "histogram_quantile";
126const SPECIAL_HISTOGRAM_FRACTION: &str = "histogram_fraction";
128const SPECIAL_VECTOR_FUNCTION: &str = "vector";
130const LE_COLUMN_NAME: &str = "le";
132
133static LABEL_NAME_REGEX: Lazy<Regex> =
136 Lazy::new(|| Regex::new(r"^[a-zA-Z_][a-zA-Z0-9_]*$").unwrap());
137
138const DEFAULT_TIME_INDEX_COLUMN: &str = "time";
139
140const DEFAULT_FIELD_COLUMN: &str = "value";
142
143const FIELD_COLUMN_MATCHER: &str = "__field__";
145
146const SCHEMA_COLUMN_MATCHER: &str = "__schema__";
148const DB_COLUMN_MATCHER: &str = "__database__";
149
150const BINARY_ISLAND_LEAF_ALIAS_PREFIX: &str = "__prom_v";
152const OR_FLOAT_FIELD_PREFIX: &str = "__promql_or_float_";
153const OR_HISTOGRAM_FIELD_PREFIX: &str = "__promql_or_histogram_";
154const TIMESTAMP_VALUE_PREFIX: &str = "__promql_timestamp_value_";
155const MATCH_GROUP_COUNT_COLUMN: &str = "__promql_match_group_count";
157
158const MAX_SCATTER_POINTS: i64 = 400;
160
161const INTERVAL_1H: i64 = 60 * 60 * 1000;
163
164#[derive(Default, Debug, Clone)]
165struct PromPlannerContext {
166 start: Millisecond,
168 end: Millisecond,
169 interval: Millisecond,
170 lookback_delta: Millisecond,
171
172 table_name: Option<String>,
174 time_index_column: Option<String>,
175 field_columns: Vec<String>,
176 tag_columns: Vec<String>,
177 aggregation_field_labels: Vec<String>,
186 use_tsid: bool,
192 field_column_matcher: Option<Vec<Matcher>>,
194 selector_matcher: Vec<Matcher>,
196 schema_name: Option<String>,
197 range: Option<Millisecond>,
199}
200
201#[derive(Debug, Clone, PartialEq, Eq, Hash)]
202struct VectorLeafKey {
203 metric_name: String,
204 matchers: Vec<(String, String, String)>,
205 or_matchers: Vec<Vec<(String, String, String)>>,
206 offset_ms: i128,
207 at: String,
208}
209
210#[derive(Debug, Clone)]
211struct IslandLeaf {
212 selector: VectorSelector,
213 display_table: String,
214}
215
216#[derive(Debug, Clone)]
217enum IslandExpr {
218 VectorLeaf(usize),
219 Scalar(DfExpr),
220 Unary {
221 input: Box<IslandExpr>,
222 },
223 Binary {
224 op: TokenType,
225 lhs: Box<IslandExpr>,
226 rhs: Box<IslandExpr>,
227 },
228}
229
230impl IslandExpr {
231 fn try_new(expr: &PromExpr, env: &mut IslandCollectEnv) -> Option<Self> {
232 if let Some(expr) = PromPlanner::try_build_literal_expr(expr) {
233 return Some(Self::Scalar(expr));
234 }
235
236 match expr {
237 PromExpr::Paren(ParenExpr { expr }) => Self::try_new(expr, env),
238 PromExpr::VectorSelector(selector) => {
239 let leaf = env.intern_leaf(selector)?;
240 Some(Self::VectorLeaf(leaf))
241 }
242 PromExpr::Unary(UnaryExpr { expr }) => {
243 let input = Self::try_new(expr, env)?;
244 Some(Self::Unary {
245 input: Box::new(input),
246 })
247 }
248 PromExpr::Binary(PromBinaryExpr {
249 lhs,
250 rhs,
251 op,
252 modifier,
253 }) if matches!(
254 op.id(),
255 token::T_ADD
256 | token::T_SUB
257 | token::T_MUL
258 | token::T_DIV
259 | token::T_MOD
260 | token::T_POW
261 | token::T_ATAN2
262 ) && modifier.as_ref().is_none_or(|modifier| {
263 !modifier.return_bool
264 && modifier.matching.is_none()
265 && matches!(modifier.card, VectorMatchCardinality::OneToOne)
266 && modifier.fill_values.lhs.is_none()
267 && modifier.fill_values.rhs.is_none()
268 }) =>
269 {
270 let lhs = Self::try_new(lhs, env)?;
271 let rhs = Self::try_new(rhs, env)?;
272 Some(Self::Binary {
273 op: *op,
274 lhs: Box::new(lhs),
275 rhs: Box::new(rhs),
276 })
277 }
278 _ => None,
279 }
280 }
281}
282
283#[derive(Debug, Default)]
284struct IslandCollectEnv {
285 leaf_by_key: HashMap<VectorLeafKey, usize>,
286 leaves: Vec<IslandLeaf>,
287 vector_occurrences: usize,
288}
289
290#[derive(Debug)]
291struct PlannedIslandLeaf {
292 plan: LogicalPlan,
293 ctx: PromPlannerContext,
294 alias: TableReference,
295 display_table: String,
296}
297
298#[derive(Debug)]
301struct BinaryResultLabels {
302 exprs: Vec<DfExpr>,
303 names: Vec<String>,
304 aggregation_field_labels: Vec<String>,
305 tsid: Option<DfExpr>,
308}
309
310impl BinaryResultLabels {
311 fn apply(&self, ctx: &mut PromPlannerContext) {
312 ctx.tag_columns = self.names.clone();
313 ctx.aggregation_field_labels = self.aggregation_field_labels.clone();
314 ctx.use_tsid = self.tsid.is_some();
315 }
316
317 fn tsid_projection(&self, table_ref: Option<TableReference>) -> Option<DfExpr> {
321 self.tsid
322 .clone()
323 .map(|tsid| tsid.alias_qualified(table_ref, DATA_SCHEMA_TSID_COLUMN_NAME))
324 }
325}
326
327#[derive(Debug)]
328struct IslandFieldExprs {
329 exprs: Vec<DfExpr>,
330 names: Vec<String>,
331 scalar: bool,
332}
333
334impl VectorLeafKey {
335 fn from_selector(selector: &VectorSelector) -> Option<Self> {
336 let mut metric_name = selector.name.clone();
337 let mut matchers = Vec::with_capacity(selector.matchers.matchers.len());
338 let matcher_key = |matcher: &Matcher| {
339 (
340 matcher.name.clone(),
341 matcher.op.to_string(),
342 matcher.value.clone(),
343 )
344 };
345
346 for matcher in &selector.matchers.matchers {
347 if matcher.name == METRIC_NAME {
348 if matcher.op != MatchOp::Equal || metric_name.is_some() {
349 return None;
350 }
351 metric_name = Some(matcher.value.clone());
352 } else {
353 matchers.push(matcher_key(matcher));
354 }
355 }
356 matchers.sort();
357
358 let mut or_matchers = selector
359 .matchers
360 .or_matchers
361 .iter()
362 .map(|group| {
363 let mut group = group.iter().map(matcher_key).collect::<Vec<_>>();
364 group.sort();
365 group
366 })
367 .collect::<Vec<_>>();
368 or_matchers.sort();
369
370 Some(Self {
371 metric_name: metric_name?,
372 matchers,
373 or_matchers,
374 offset_ms: match &selector.offset {
375 Some(Offset::Pos(duration)) => duration.as_millis() as i128,
376 Some(Offset::Neg(duration)) => -(duration.as_millis() as i128),
377 None => 0,
378 },
379 at: format!("{:?}", selector.at),
380 })
381 }
382}
383
384impl IslandCollectEnv {
385 fn intern_leaf(&mut self, selector: &VectorSelector) -> Option<usize> {
386 self.vector_occurrences += 1;
387 let key = VectorLeafKey::from_selector(selector)?;
388 if let Some(id) = self.leaf_by_key.get(&key) {
389 return Some(*id);
390 }
391
392 let id = self.leaves.len();
393 self.leaves.push(IslandLeaf {
394 selector: selector.clone(),
395 display_table: key.metric_name.clone(),
396 });
397 self.leaf_by_key.insert(key, id);
398 Some(id)
399 }
400}
401
402impl PromPlannerContext {
403 fn from_eval_stmt(stmt: &EvalStmt) -> Self {
404 Self {
405 start: stmt.start.duration_since(UNIX_EPOCH).unwrap().as_millis() as _,
406 end: stmt.end.duration_since(UNIX_EPOCH).unwrap().as_millis() as _,
407 interval: stmt.interval.as_millis() as _,
408 lookback_delta: stmt.lookback_delta.as_millis() as _,
409 ..Default::default()
410 }
411 }
412
413 fn reset(&mut self) {
415 self.table_name = None;
416 self.time_index_column = None;
417 self.field_columns = vec![];
418 self.tag_columns = vec![];
419 self.use_tsid = false;
420 self.field_column_matcher = None;
421 self.selector_matcher.clear();
422 self.schema_name = None;
423 self.range = None;
424 }
425
426 fn reset_table_name_and_schema(&mut self) {
428 self.table_name = Some(String::new());
429 self.schema_name = None;
430 self.use_tsid = false;
431 }
432
433 fn has_le_tag(&self) -> bool {
435 self.tag_columns.iter().any(|c| c.eq(&LE_COLUMN_NAME))
436 }
437}
438
439pub struct PromPlanner {
440 table_provider: DfTableSourceProvider,
441 ctx: PromPlannerContext,
442 promql_annotations: Option<PromqlAnnotationCollector>,
444}
445
446type BinaryFieldPair<'a> = (&'a String, &'a String);
447
448impl PromPlanner {
449 pub async fn stmt_to_plan(
450 table_provider: DfTableSourceProvider,
451 stmt: &EvalStmt,
452 query_engine_state: &QueryEngineState,
453 ) -> Result<LogicalPlan> {
454 Self::stmt_to_plan_with_annotations(table_provider, stmt, query_engine_state, None).await
455 }
456
457 pub async fn stmt_to_plan_with_annotations(
459 table_provider: DfTableSourceProvider,
460 stmt: &EvalStmt,
461 query_engine_state: &QueryEngineState,
462 promql_annotations: Option<PromqlAnnotationCollector>,
463 ) -> Result<LogicalPlan> {
464 let mut planner = Self {
465 table_provider,
466 ctx: PromPlannerContext::from_eval_stmt(stmt),
467 promql_annotations,
468 };
469
470 let plan = planner
471 .prom_expr_to_plan(&stmt.expr, query_engine_state)
472 .await?;
473
474 planner.strip_tsid_column(plan)
476 }
477
478 pub async fn prom_expr_to_plan(
479 &mut self,
480 prom_expr: &PromExpr,
481 query_engine_state: &QueryEngineState,
482 ) -> Result<LogicalPlan> {
483 self.prom_expr_to_plan_inner(prom_expr, false, query_engine_state)
484 .await
485 }
486
487 #[async_recursion]
497 async fn prom_expr_to_plan_inner(
498 &mut self,
499 prom_expr: &PromExpr,
500 timestamp_fn: bool,
501 query_engine_state: &QueryEngineState,
502 ) -> Result<LogicalPlan> {
503 let res = match prom_expr {
504 PromExpr::Aggregate(expr) => {
505 self.prom_aggr_expr_to_plan(query_engine_state, expr)
506 .await?
507 }
508 PromExpr::Unary(expr) => {
509 self.prom_unary_expr_to_plan(query_engine_state, expr)
510 .await?
511 }
512 PromExpr::Binary(expr) => {
513 self.prom_binary_expr_to_plan(query_engine_state, expr)
514 .await?
515 }
516 PromExpr::Paren(ParenExpr { expr }) => {
517 self.prom_expr_to_plan_inner(expr, timestamp_fn, query_engine_state)
518 .await?
519 }
520 PromExpr::Subquery(expr) => {
521 self.prom_subquery_expr_to_plan(query_engine_state, expr)
522 .await?
523 }
524 PromExpr::NumberLiteral(lit) => self.prom_number_lit_to_plan(lit)?,
525 PromExpr::StringLiteral(lit) => self.prom_string_lit_to_plan(lit)?,
526 PromExpr::VectorSelector(selector) => {
527 self.prom_vector_selector_to_plan(selector, timestamp_fn)
528 .await?
529 }
530 PromExpr::MatrixSelector(selector) => {
531 self.prom_matrix_selector_to_plan(selector).await?
532 }
533 PromExpr::Call(expr) => {
534 self.prom_call_expr_to_plan(query_engine_state, expr)
535 .await?
536 }
537 PromExpr::Extension(expr) => {
538 self.prom_ext_expr_to_plan(query_engine_state, expr).await?
539 }
540 };
541
542 Ok(res)
543 }
544
545 async fn prom_subquery_expr_to_plan(
546 &mut self,
547 query_engine_state: &QueryEngineState,
548 subquery_expr: &SubqueryExpr,
549 ) -> Result<LogicalPlan> {
550 let SubqueryExpr {
551 expr, range, step, ..
552 } = subquery_expr;
553
554 let current_interval = self.ctx.interval;
555 if let Some(step) = step {
556 self.ctx.interval = step.as_millis() as _;
557 }
558 let current_start = self.ctx.start;
559 self.ctx.start -= range.as_millis() as i64 - self.ctx.interval;
560 let input = self.prom_expr_to_plan(expr, query_engine_state).await?;
561 self.ctx.interval = current_interval;
562 self.ctx.start = current_start;
563
564 ensure!(!range.is_zero(), ZeroRangeSelectorSnafu);
565 let range_ms = range.as_millis() as _;
566 self.ctx.range = Some(range_ms);
567
568 let time_index_column =
569 self.ctx
570 .time_index_column
571 .clone()
572 .with_context(|| TimeIndexNotFoundSnafu {
573 table: self.ctx.table_name.clone().unwrap_or_default(),
574 })?;
575
576 let input_schema = input.schema();
582 let input_has_tsid = input_schema.fields().iter().any(|field| {
583 field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
584 && field.data_type() == &ArrowDataType::UInt64
585 });
586 let (series_key_columns, mut sort_exprs) = if input_has_tsid {
587 (
588 vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()],
589 vec![
590 DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME))
591 .sort(true, true),
592 ],
593 )
594 } else {
595 let key_columns: Vec<String> = self
598 .ctx
599 .tag_columns
600 .iter()
601 .filter(|name| input_schema.has_column_with_unqualified_name(name))
602 .cloned()
603 .collect();
604 let sort = key_columns
605 .iter()
606 .map(|name| DfExpr::Column(Column::from_name(name)).sort(true, true))
607 .collect::<Vec<_>>();
608 (key_columns, sort)
609 };
610 sort_exprs.push(DfExpr::Column(Column::from_name(&time_index_column)).sort(true, true));
611
612 let sort_plan = LogicalPlanBuilder::from(input)
613 .sort(sort_exprs)
614 .context(DataFusionPlanningSnafu)?
615 .build()
616 .context(DataFusionPlanningSnafu)?;
617 let divide_plan = LogicalPlan::Extension(Extension {
618 node: Arc::new(SeriesDivide::new(
619 series_key_columns,
620 time_index_column.clone(),
621 sort_plan,
622 )),
623 });
624
625 let manipulate = RangeManipulate::new(
626 self.ctx.start,
627 self.ctx.end,
628 self.ctx.interval,
629 0,
630 range_ms,
631 time_index_column,
632 self.ctx.field_columns.clone(),
633 divide_plan,
634 )
635 .context(DataFusionPlanningSnafu)?;
636
637 Ok(LogicalPlan::Extension(Extension {
638 node: Arc::new(manipulate),
639 }))
640 }
641
642 async fn prom_aggr_expr_to_plan(
643 &mut self,
644 query_engine_state: &QueryEngineState,
645 aggr_expr: &AggregateExpr,
646 ) -> Result<LogicalPlan> {
647 let AggregateExpr {
648 op,
649 expr,
650 modifier,
651 param,
652 } = aggr_expr;
653
654 let input = self.prom_expr_to_plan(expr, query_engine_state).await?;
655 let input_has_tsid = input.schema().fields().iter().any(|field| {
656 field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
657 && field.data_type() == &ArrowDataType::UInt64
658 });
659
660 match (*op).id() {
661 token::T_TOPK | token::T_BOTTOMK => {
662 self.prom_topk_bottomk_to_plan(aggr_expr, input).await
663 }
664 _ => {
665 let input_tag_columns = if input_has_tsid {
669 self.collect_row_key_tag_columns_from_plan(&input)?
670 .into_iter()
671 .collect::<Vec<_>>()
672 } else {
673 self.ctx.tag_columns.clone()
674 };
675 let mut group_exprs = self.agg_modifier_to_col(input.schema(), modifier, true)?;
678 let mixed_sample_columns =
679 Self::alternative_sample_columns(input.schema(), &self.ctx.field_columns)
680 .map(|(float, histogram)| (float.to_string(), histogram.to_string()));
681 let preserve_any_value = mixed_sample_columns.is_some();
686 let has_native_histogram = preserve_any_value
687 || self.all_field_columns_are_native_histograms(input.schema());
688 let (mut aggr_exprs, prev_field_exprs) =
690 self.create_aggregate_exprs(*op, param, &input)?;
691 let prev_field_exprs =
692 normalize_cols(prev_field_exprs, &input).context(DataFusionPlanningSnafu)?;
693
694 let keep_tsid = op.id() != token::T_COUNT_VALUES
695 && input_has_tsid
696 && input_tag_columns.iter().collect::<HashSet<_>>()
697 == self.ctx.tag_columns.iter().collect::<HashSet<_>>();
698
699 if keep_tsid {
700 aggr_exprs.push(
701 first_value(
702 DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME)),
703 vec![],
704 )
705 .alias(DATA_SCHEMA_TSID_COLUMN_NAME),
706 );
707 }
708 self.ctx.use_tsid = keep_tsid;
709
710 let builder = LogicalPlanBuilder::from(input);
712 let builder = if op.id() == token::T_COUNT_VALUES {
713 let label = Self::get_param_value_as_str(*op, param)?;
714 self.ctx.tag_columns.retain(|tag| tag != label);
728 group_exprs.retain(
729 |expr| !matches!(expr, DfExpr::Column(column) if column.name == label),
730 );
731 let label_qualifier = self
736 .ctx
737 .time_index_column
738 .as_deref()
739 .and_then(|time_index| {
740 builder
741 .schema()
742 .qualified_field_with_unqualified_name(time_index)
743 .ok()
744 })
745 .and_then(|(qualifier, _)| qualifier.cloned());
746 let count_value_exprs = prev_field_exprs
753 .iter()
754 .map(|expr| {
755 let value = match expr {
756 DfExpr::Column(column) => {
757 let value = DfExpr::Column(column.clone());
758 let value = if matches!(
761 builder.schema().field_with_unqualified_name(&column.name),
762 Ok(field) if field.data_type() == &ArrowDataType::Float64
763 ) {
764 value
765 } else {
766 DfExpr::Cast(Cast::new(
767 Box::new(value),
768 ArrowDataType::Float64,
769 ))
770 };
771 DfExpr::ScalarFunction(ScalarFunction {
772 func: Arc::new(PromqlFloatToString::scalar_udf()),
773 args: vec![value],
774 })
775 }
776 _ => DfExpr::Column(Column::from_name(
780 expr.schema_name().to_string(),
781 )),
782 };
783 DfExpr::Alias(Alias::new(value, label_qualifier.clone(), label))
784 })
785 .collect::<Vec<_>>();
786 let aggregate_group_exprs = group_exprs
787 .iter()
788 .cloned()
789 .chain(prev_field_exprs.clone())
790 .collect::<Vec<_>>();
791 group_exprs.push(col(label));
792 let project_fields = self
793 .create_field_column_exprs()?
794 .into_iter()
795 .chain(self.create_tag_column_exprs()?)
796 .chain(Some(self.create_time_index_column_expr()?))
797 .chain(count_value_exprs);
798
799 let builder = builder
800 .aggregate(aggregate_group_exprs, aggr_exprs)
801 .context(DataFusionPlanningSnafu)?
802 .project(project_fields)
803 .context(DataFusionPlanningSnafu)?;
804 self.ctx.tag_columns.push(label.to_string());
806 builder
807 } else {
808 builder
809 .aggregate(group_exprs.clone(), aggr_exprs)
810 .context(DataFusionPlanningSnafu)?
811 };
812
813 let builder = if let Some((float, histogram)) = mixed_sample_columns {
814 let builder = match op.id() {
815 token::T_SUM | token::T_AVG => builder
816 .filter(self.mixed_aggregate_filter_expr(*op, &float, &histogram)?)
817 .context(DataFusionPlanningSnafu)?,
818 token::T_MIN
819 | token::T_MAX
820 | token::T_STDDEV
821 | token::T_STDVAR
822 | token::T_QUANTILE => builder
823 .filter(self.mixed_ignored_histogram_filter_expr(*op, &histogram)?)
824 .context(DataFusionPlanningSnafu)?,
825 _ => builder,
826 };
827
828 match op.id() {
829 token::T_SUM
830 | token::T_AVG
831 | token::T_MIN
832 | token::T_MAX
833 | token::T_STDDEV
834 | token::T_STDVAR
835 | token::T_QUANTILE => {
836 let project_fields = self
837 .create_field_column_exprs()?
838 .into_iter()
839 .chain(self.create_tag_column_exprs()?)
840 .chain(self.ctx.use_tsid.then_some(DfExpr::Column(
841 Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME),
842 )))
843 .chain(Some(self.create_time_index_column_expr()?));
844 builder
845 .project(project_fields)
846 .context(DataFusionPlanningSnafu)?
847 }
848 _ => builder,
849 }
850 } else {
851 builder
852 };
853
854 let builder = if has_native_histogram {
859 builder
860 .filter(self.create_empty_values_filter_expr(preserve_any_value)?)
861 .context(DataFusionPlanningSnafu)?
862 } else {
863 builder
864 };
865
866 let sort_expr = group_exprs.into_iter().map(|expr| expr.sort(true, false));
867
868 builder
869 .sort(sort_expr)
870 .context(DataFusionPlanningSnafu)?
871 .build()
872 .context(DataFusionPlanningSnafu)
873 }
874 }
875 }
876
877 async fn prom_topk_bottomk_to_plan(
879 &mut self,
880 aggr_expr: &AggregateExpr,
881 input: LogicalPlan,
882 ) -> Result<LogicalPlan> {
883 let AggregateExpr {
884 op,
885 param,
886 modifier,
887 ..
888 } = aggr_expr;
889
890 let input_has_tsid = input.schema().fields().iter().any(|field| {
891 field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
892 && field.data_type() == &ArrowDataType::UInt64
893 });
894 self.ctx.use_tsid = input_has_tsid;
895
896 let group_exprs = self.agg_modifier_to_col(input.schema(), modifier, false)?;
897
898 let mut input = input;
899 if let Some((float_column, histogram_column)) =
900 Self::alternative_sample_columns(input.schema(), &self.ctx.field_columns)
901 .map(|(float, histogram)| (float.to_string(), histogram.to_string()))
902 {
903 let drop_histogram = DfExpr::ScalarFunction(ScalarFunction {
904 func: Arc::new(NativeHistogramDrop::bool_false_udf(
905 format!(
906 "{}: dropped native histogram samples because this aggregation is not supported for native histograms",
907 op
908 ),
909 self.promql_annotations.clone(),
910 )),
911 args: vec![col(&histogram_column)],
912 });
913 let keep_float = when(col(&histogram_column).is_not_null(), drop_histogram)
914 .otherwise(col(&float_column).is_not_null())
915 .context(DataFusionPlanningSnafu)?;
916 input = LogicalPlanBuilder::from(input)
917 .filter(keep_float)
918 .context(DataFusionPlanningSnafu)?
919 .build()
920 .context(DataFusionPlanningSnafu)?;
921 self.ctx.field_columns = vec![float_column];
922 }
923
924 if self.all_field_columns_are_native_histograms(input.schema()) {
925 let promql_annotations = self.promql_annotations.clone();
926 let input = self.projection_for_each_field_column(input, |col| {
927 Ok(DfExpr::ScalarFunction(ScalarFunction {
928 func: Arc::new(NativeHistogramDrop::float_null_udf(
929 format!(
930 "{}: dropped native histogram samples because this aggregation is not supported for native histograms",
931 op
932 ),
933 promql_annotations.clone(),
934 )),
935 args: vec![DfExpr::Column(Column::from_name(col))],
936 }))
937 })?;
938 return LogicalPlanBuilder::from(input)
939 .filter(self.create_empty_values_filter_expr(false)?)
940 .context(DataFusionPlanningSnafu)?
941 .build()
942 .context(DataFusionPlanningSnafu);
943 }
944
945 let val = Self::get_param_as_literal_expr(
946 param.as_deref(),
947 Some(*op),
948 Some(ArrowDataType::Float64),
949 )?;
950
951 let window_exprs = self.create_window_exprs(*op, group_exprs.clone(), &input)?;
953
954 let rank_columns: Vec<_> = window_exprs
955 .iter()
956 .map(|expr| expr.schema_name().to_string())
957 .collect();
958
959 let filter: DfExpr = rank_columns
962 .iter()
963 .fold(None, |expr, rank| {
964 let predicate = DfExpr::BinaryExpr(BinaryExpr {
965 left: Box::new(col(rank)),
966 op: Operator::LtEq,
967 right: Box::new(val.clone()),
968 });
969
970 match expr {
971 None => Some(predicate),
972 Some(expr) => Some(DfExpr::BinaryExpr(BinaryExpr {
973 left: Box::new(expr),
974 op: Operator::Or,
975 right: Box::new(predicate),
976 })),
977 }
978 })
979 .unwrap();
980
981 let rank_columns: Vec<_> = rank_columns.into_iter().map(col).collect();
982
983 let mut new_group_exprs = group_exprs.clone();
984 new_group_exprs.extend(rank_columns);
986
987 let group_sort_expr = new_group_exprs
988 .into_iter()
989 .map(|expr| expr.sort(true, false));
990
991 let project_fields = self
992 .create_field_column_exprs()?
993 .into_iter()
994 .chain(self.create_tag_column_exprs()?)
995 .chain(
996 self.ctx
997 .use_tsid
998 .then_some(DfExpr::Column(Column::from_name(
999 DATA_SCHEMA_TSID_COLUMN_NAME,
1000 ))),
1001 )
1002 .chain(Some(self.create_time_index_column_expr()?));
1003
1004 LogicalPlanBuilder::from(input)
1005 .window(window_exprs)
1006 .context(DataFusionPlanningSnafu)?
1007 .filter(filter)
1008 .context(DataFusionPlanningSnafu)?
1009 .sort(group_sort_expr)
1010 .context(DataFusionPlanningSnafu)?
1011 .project(project_fields)
1012 .context(DataFusionPlanningSnafu)?
1013 .build()
1014 .context(DataFusionPlanningSnafu)
1015 }
1016
1017 async fn prom_unary_expr_to_plan(
1018 &mut self,
1019 query_engine_state: &QueryEngineState,
1020 unary_expr: &UnaryExpr,
1021 ) -> Result<LogicalPlan> {
1022 let UnaryExpr { expr } = unary_expr;
1023 let input = self.prom_expr_to_plan(expr, query_engine_state).await?;
1025 self.negate_field_columns(input)
1026 }
1027
1028 fn negate_field_columns(&mut self, input: LogicalPlan) -> Result<LogicalPlan> {
1029 let input_schema = input.schema().clone();
1030 self.projection_for_each_field_column(input, |col| {
1031 if Self::field_column_is_native_histogram(&input_schema, col) {
1032 Ok(DfExpr::ScalarFunction(ScalarFunction {
1033 func: Arc::new(NativeHistogramNeg::scalar_udf()),
1034 args: vec![DfExpr::Column(col.into())],
1035 }))
1036 } else {
1037 Ok(DfExpr::Negative(Box::new(DfExpr::Column(col.into()))))
1038 }
1039 })
1040 }
1041
1042 async fn try_plan_binary_island(
1043 &mut self,
1044 binary_expr: &PromBinaryExpr,
1045 ) -> Result<Option<LogicalPlan>> {
1046 let original_ctx = self.ctx.clone();
1047 let mut collect_env = IslandCollectEnv::default();
1048 let Some(island_expr) =
1049 IslandExpr::try_new(&PromExpr::Binary(binary_expr.clone()), &mut collect_env)
1050 else {
1051 return Ok(None);
1052 };
1053
1054 if collect_env.leaves.is_empty()
1055 || collect_env.vector_occurrences <= collect_env.leaves.len()
1056 {
1057 return Ok(None);
1058 }
1059
1060 let mut planned_leaves = Vec::with_capacity(collect_env.leaves.len());
1061 for (idx, leaf) in collect_env.leaves.iter().enumerate() {
1062 let plan = self
1063 .prom_vector_selector_to_plan(&leaf.selector, false)
1064 .await?;
1065 let ctx = self.ctx.clone();
1066 let alias = TableReference::bare(format!("{BINARY_ISLAND_LEAF_ALIAS_PREFIX}{idx}"));
1067 let plan = LogicalPlanBuilder::from(plan)
1068 .alias(alias.clone())
1069 .context(DataFusionPlanningSnafu)?
1070 .build()
1071 .context(DataFusionPlanningSnafu)?;
1072 planned_leaves.push(PlannedIslandLeaf {
1073 plan,
1074 ctx,
1075 alias,
1076 display_table: leaf.display_table.clone(),
1077 });
1078 }
1079
1080 if planned_leaves.iter().any(|leaf| {
1081 Self::field_columns_contain_native_histogram(
1082 leaf.plan.schema(),
1083 &leaf.ctx.field_columns,
1084 )
1085 }) {
1086 self.ctx = original_ctx;
1087 return Ok(None);
1088 }
1089
1090 if !Self::binary_island_join_contexts_supported(&planned_leaves) {
1091 self.ctx = original_ctx;
1092 return Ok(None);
1093 }
1094
1095 let mut input = planned_leaves[0].plan.clone();
1096 for right_idx in 1..planned_leaves.len() {
1097 input = self.join_binary_island_leaf(
1098 input,
1099 &planned_leaves[0],
1100 &planned_leaves[right_idx],
1101 )?;
1102 }
1103
1104 let field_exprs =
1105 Self::build_binary_island_field_exprs(&island_expr, &planned_leaves, input.schema())?;
1106 if field_exprs.scalar || field_exprs.exprs.is_empty() {
1107 self.ctx = original_ctx;
1108 return Ok(None);
1109 }
1110
1111 let plan = self.project_binary_island(
1112 input,
1113 &planned_leaves[0].alias,
1114 &planned_leaves[0].ctx,
1115 field_exprs,
1116 )?;
1117 Ok(Some(plan))
1118 }
1119
1120 fn binary_island_join_contexts_supported(leaves: &[PlannedIslandLeaf]) -> bool {
1121 if leaves
1122 .iter()
1123 .any(|leaf| leaf.ctx.time_index_column.is_none())
1124 {
1125 return false;
1126 }
1127
1128 if leaves.len() <= 1 {
1129 return true;
1130 }
1131
1132 let first_tags = leaves[0].ctx.tag_columns.iter().collect::<BTreeSet<_>>();
1133
1134 leaves.iter().skip(1).all(|leaf| {
1135 (Self::plan_has_tsid_column(&leaves[0].plan) && Self::plan_has_tsid_column(&leaf.plan))
1136 || leaf.ctx.tag_columns.iter().collect::<BTreeSet<_>>() == first_tags
1137 })
1138 }
1139
1140 fn join_binary_island_leaf(
1141 &self,
1142 left: LogicalPlan,
1143 first_leaf: &PlannedIslandLeaf,
1144 right_leaf: &PlannedIslandLeaf,
1145 ) -> Result<LogicalPlan> {
1146 let only_join_time_index = (first_leaf.ctx.tag_columns.is_empty()
1147 || right_leaf.ctx.tag_columns.is_empty())
1148 && !first_leaf
1149 .ctx
1150 .tag_columns
1151 .iter()
1152 .chain(&right_leaf.ctx.tag_columns)
1153 .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL);
1154 let (mut left_keys, mut right_keys, force_empty_join) = self.binary_join_key_columns(
1155 left.schema(),
1156 right_leaf.plan.schema(),
1157 &first_leaf.ctx,
1158 &right_leaf.ctx,
1159 only_join_time_index,
1160 &None,
1161 )?;
1162
1163 if let (Some(left_time_index_column), Some(right_time_index_column)) = (
1164 first_leaf.ctx.time_index_column.clone(),
1165 right_leaf.ctx.time_index_column.clone(),
1166 ) {
1167 left_keys.insert(left_time_index_column);
1168 right_keys.insert(right_time_index_column);
1169 }
1170
1171 LogicalPlanBuilder::from(left)
1172 .join_detailed(
1173 right_leaf.plan.clone(),
1174 JoinType::Inner,
1175 (
1176 left_keys
1177 .into_iter()
1178 .map(|name| Column::new(Some(first_leaf.alias.clone()), name))
1179 .collect::<Vec<_>>(),
1180 right_keys
1181 .into_iter()
1182 .map(|name| Column::new(Some(right_leaf.alias.clone()), name))
1183 .collect::<Vec<_>>(),
1184 ),
1185 force_empty_join.then_some(lit(false)),
1186 NullEquality::NullEqualsNull,
1187 )
1188 .context(DataFusionPlanningSnafu)?
1189 .build()
1190 .context(DataFusionPlanningSnafu)
1191 }
1192
1193 fn build_binary_island_field_exprs(
1194 expr: &IslandExpr,
1195 leaves: &[PlannedIslandLeaf],
1196 schema: &DFSchemaRef,
1197 ) -> Result<IslandFieldExprs> {
1198 match expr {
1199 IslandExpr::VectorLeaf(id) => {
1200 let leaf = &leaves[*id];
1201 let exprs = leaf
1202 .ctx
1203 .field_columns
1204 .iter()
1205 .map(|field| {
1206 schema
1207 .qualified_field_with_name(Some(&leaf.alias), field)
1208 .context(DataFusionPlanningSnafu)
1209 .map(|field| DfExpr::Column(field.into()))
1210 })
1211 .collect::<Result<Vec<_>>>()?;
1212 let names = leaf
1213 .ctx
1214 .field_columns
1215 .iter()
1216 .map(|field| format!("{}.{}", leaf.display_table, field))
1217 .collect();
1218 Ok(IslandFieldExprs {
1219 exprs,
1220 names,
1221 scalar: false,
1222 })
1223 }
1224 IslandExpr::Scalar(expr) => Ok(IslandFieldExprs {
1225 exprs: vec![expr.clone()],
1226 names: vec![expr.schema_name().to_string()],
1227 scalar: true,
1228 }),
1229 IslandExpr::Unary { input } => {
1230 let input = Self::build_binary_island_field_exprs(input, leaves, schema)?;
1231 let mut exprs = Vec::with_capacity(input.exprs.len());
1232 let mut names = Vec::with_capacity(input.names.len());
1233 for (expr, name) in input.exprs.into_iter().zip(input.names) {
1234 exprs.push(DfExpr::Negative(Box::new(expr)));
1235 names.push(format!("-{name}"));
1236 }
1237 Ok(IslandFieldExprs {
1238 exprs,
1239 names,
1240 scalar: input.scalar,
1241 })
1242 }
1243 IslandExpr::Binary { op, lhs, rhs } => {
1244 let same_leaf = match (&**lhs, &**rhs) {
1245 (IslandExpr::VectorLeaf(left), IslandExpr::VectorLeaf(right))
1246 if left == right =>
1247 {
1248 Some(*left)
1249 }
1250 _ => None,
1251 };
1252 let lhs = Self::build_binary_island_field_exprs(lhs, leaves, schema)?;
1253 let rhs = Self::build_binary_island_field_exprs(rhs, leaves, schema)?;
1254 let expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1255 let scalar = lhs.scalar && rhs.scalar;
1256 let op = op.to_string();
1257
1258 let (exprs, names) = match (lhs.scalar, rhs.scalar) {
1259 (true, true) => {
1260 let expr = expr_builder(lhs.exprs[0].clone(), rhs.exprs[0].clone())?;
1261 let name = format!("{} {op} {}", lhs.names[0], rhs.names[0]);
1262 (vec![expr], vec![name])
1263 }
1264 (true, false) => {
1265 let mut exprs = Vec::with_capacity(rhs.exprs.len());
1266 let mut names = Vec::with_capacity(rhs.names.len());
1267 for (rhs_expr, rhs_name) in rhs.exprs.into_iter().zip(rhs.names) {
1268 exprs.push(expr_builder(lhs.exprs[0].clone(), rhs_expr)?);
1269 names.push(format!("{} {op} {rhs_name}", lhs.names[0]));
1270 }
1271 (exprs, names)
1272 }
1273 (false, true) => {
1274 let mut exprs = Vec::with_capacity(lhs.exprs.len());
1275 let mut names = Vec::with_capacity(lhs.names.len());
1276 for (lhs_expr, lhs_name) in lhs.exprs.into_iter().zip(lhs.names) {
1277 exprs.push(expr_builder(lhs_expr, rhs.exprs[0].clone())?);
1278 names.push(format!("{lhs_name} {op} {}", rhs.names[0]));
1279 }
1280 (exprs, names)
1281 }
1282 (false, false) => {
1283 let mut exprs = Vec::new();
1284 let mut names = Vec::new();
1285 for (idx, ((lhs_expr, rhs_expr), (mut lhs_name, mut rhs_name))) in lhs
1286 .exprs
1287 .into_iter()
1288 .zip(rhs.exprs)
1289 .zip(lhs.names.into_iter().zip(rhs.names))
1290 .enumerate()
1291 {
1292 if let Some(leaf) = same_leaf {
1293 let field = leaves[leaf]
1294 .ctx
1295 .field_columns
1296 .get(idx)
1297 .cloned()
1298 .unwrap_or_else(|| lhs_name.clone());
1299 lhs_name = format!("lhs.{field}");
1300 rhs_name = format!("rhs.{field}");
1301 }
1302 exprs.push(expr_builder(lhs_expr, rhs_expr)?);
1303 names.push(format!("{lhs_name} {op} {rhs_name}"));
1304 }
1305 (exprs, names)
1306 }
1307 };
1308
1309 Ok(IslandFieldExprs {
1310 exprs,
1311 names,
1312 scalar,
1313 })
1314 }
1315 }
1316 }
1317
1318 fn project_binary_island(
1319 &mut self,
1320 input: LogicalPlan,
1321 base_alias: &TableReference,
1322 base_ctx: &PromPlannerContext,
1323 field_exprs: IslandFieldExprs,
1324 ) -> Result<LogicalPlan> {
1325 self.ctx = base_ctx.clone();
1326
1327 let schema = input.schema();
1328 let non_field_exprs = base_ctx
1329 .tag_columns
1330 .iter()
1331 .chain(base_ctx.time_index_column.iter())
1332 .map(|column| {
1333 schema
1334 .qualified_field_with_name(Some(base_alias), column)
1335 .context(DataFusionPlanningSnafu)
1336 .map(|field| DfExpr::Column(field.into()))
1337 });
1338 let tsid_expr = Self::optional_tsid_projection(schema, Some(base_alias), base_ctx.use_tsid)
1339 .into_iter()
1340 .map(Ok);
1341
1342 self.ctx.field_columns = field_exprs.names;
1343 let field_exprs = field_exprs
1344 .exprs
1345 .into_iter()
1346 .zip(self.ctx.field_columns.iter())
1347 .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::<String>, name))));
1348
1349 let project_exprs = non_field_exprs
1350 .chain(tsid_expr)
1351 .chain(field_exprs)
1352 .collect::<Result<Vec<_>>>()?;
1353
1354 let plan = LogicalPlanBuilder::from(input)
1355 .project(project_exprs)
1356 .context(DataFusionPlanningSnafu)?
1357 .build()
1358 .context(DataFusionPlanningSnafu)?;
1359
1360 self.ctx.table_name = None;
1361 self.ctx.schema_name = None;
1362
1363 Ok(plan)
1364 }
1365
1366 async fn prom_binary_expr_to_plan(
1367 &mut self,
1368 query_engine_state: &QueryEngineState,
1369 binary_expr: &PromBinaryExpr,
1370 ) -> Result<LogicalPlan> {
1371 if let Some(modifier) = &binary_expr.modifier {
1375 ensure!(
1376 modifier.fill_values.lhs.is_none() && modifier.fill_values.rhs.is_none(),
1377 UnsupportedExprSnafu {
1378 name: "PromQL fill modifiers"
1379 }
1380 );
1381 }
1382
1383 if let Some(plan) = self.try_plan_binary_island(binary_expr).await? {
1384 return Ok(plan);
1385 }
1386
1387 let PromBinaryExpr {
1388 lhs,
1389 rhs,
1390 op,
1391 modifier,
1392 } = binary_expr;
1393
1394 let should_return_bool = if let Some(m) = modifier {
1397 m.return_bool
1398 } else {
1399 false
1400 };
1401 let is_comparison_op = Self::is_token_a_comparison_op(*op);
1402
1403 match (
1406 Self::try_build_literal_expr(lhs),
1407 Self::try_build_literal_expr(rhs),
1408 ) {
1409 (Some(lhs), Some(rhs)) => {
1410 self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
1411 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1412 self.ctx.reset_table_name_and_schema();
1413 let field_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1414 let mut field_expr = field_expr_builder(lhs, rhs)?;
1415
1416 if is_comparison_op && should_return_bool {
1417 field_expr =
1418 DfExpr::Cast(Cast::new(Box::new(field_expr), ArrowDataType::Float64));
1419 }
1420
1421 Ok(LogicalPlan::Extension(Extension {
1422 node: Arc::new(
1423 EmptyMetric::new(
1424 self.ctx.start,
1425 self.ctx.end,
1426 self.ctx.interval,
1427 SPECIAL_TIME_FUNCTION.to_string(),
1428 DEFAULT_FIELD_COLUMN.to_string(),
1429 Some(field_expr),
1430 )
1431 .context(DataFusionPlanningSnafu)?,
1432 ),
1433 }))
1434 }
1435 (Some(mut expr), None) => {
1437 let input = self.prom_expr_to_plan(rhs, query_engine_state).await?;
1438 if let Some(time_expr) = self.try_build_special_time_expr_with_context(lhs) {
1442 expr = time_expr
1443 }
1444 let input_schema = input.schema().clone();
1445 let preserve_any_value = Self::field_columns_are_alternative_samples(
1446 &input_schema,
1447 &self.ctx.field_columns,
1448 );
1449 let has_native_histogram = Self::field_columns_contain_native_histogram(
1450 &input_schema,
1451 &self.ctx.field_columns,
1452 );
1453 let retain_field_columns = self
1454 .ctx
1455 .field_columns
1456 .iter()
1457 .map(|col| {
1458 Self::binary_result_is_histogram(
1459 *op,
1460 false,
1461 Self::field_column_is_native_histogram(&input_schema, col),
1462 )
1463 .is_some()
1464 })
1465 .collect();
1466 let promql_annotations = self.promql_annotations.clone();
1467 let bin_expr_builder = |col: &String| {
1468 let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1469 let rhs_is_histogram =
1470 Self::field_column_is_native_histogram(&input_schema, col);
1471 let rhs = DfExpr::Column(col.into());
1472 let mut binary_expr = match Self::native_histogram_binary_expr(
1473 *op,
1474 expr.clone(),
1475 false,
1476 rhs.clone(),
1477 rhs_is_histogram,
1478 is_comparison_op && !should_return_bool,
1479 promql_annotations.clone(),
1480 )? {
1481 Some(expr) => expr,
1482 None => binary_expr_builder(expr.clone(), rhs)?,
1483 };
1484
1485 if is_comparison_op && should_return_bool {
1486 binary_expr =
1487 DfExpr::Cast(Cast::new(Box::new(binary_expr), ArrowDataType::Float64));
1488 }
1489 Ok(binary_expr)
1490 };
1491 if is_comparison_op && !should_return_bool {
1492 self.filter_on_field_column(input, bin_expr_builder)
1493 } else {
1494 let projected =
1495 self.projection_for_each_field_column(input, bin_expr_builder)?;
1496 self.filter_binary_projection(
1497 projected,
1498 has_native_histogram,
1499 preserve_any_value,
1500 retain_field_columns,
1501 )
1502 }
1503 }
1504 (None, Some(mut expr)) => {
1506 let input = self.prom_expr_to_plan(lhs, query_engine_state).await?;
1507 if let Some(time_expr) = self.try_build_special_time_expr_with_context(rhs) {
1509 expr = time_expr
1510 }
1511 let input_schema = input.schema().clone();
1512 let preserve_any_value = Self::field_columns_are_alternative_samples(
1513 &input_schema,
1514 &self.ctx.field_columns,
1515 );
1516 let has_native_histogram = Self::field_columns_contain_native_histogram(
1517 &input_schema,
1518 &self.ctx.field_columns,
1519 );
1520 let retain_field_columns = self
1521 .ctx
1522 .field_columns
1523 .iter()
1524 .map(|col| {
1525 Self::binary_result_is_histogram(
1526 *op,
1527 Self::field_column_is_native_histogram(&input_schema, col),
1528 false,
1529 )
1530 .is_some()
1531 })
1532 .collect();
1533 let promql_annotations = self.promql_annotations.clone();
1534 let bin_expr_builder = |col: &String| {
1535 let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1536 let lhs_is_histogram =
1537 Self::field_column_is_native_histogram(&input_schema, col);
1538 let lhs = DfExpr::Column(col.into());
1539 let mut binary_expr = match Self::native_histogram_binary_expr(
1540 *op,
1541 lhs.clone(),
1542 lhs_is_histogram,
1543 expr.clone(),
1544 false,
1545 is_comparison_op && !should_return_bool,
1546 promql_annotations.clone(),
1547 )? {
1548 Some(expr) => expr,
1549 None => binary_expr_builder(lhs, expr.clone())?,
1550 };
1551
1552 if is_comparison_op && should_return_bool {
1553 binary_expr =
1554 DfExpr::Cast(Cast::new(Box::new(binary_expr), ArrowDataType::Float64));
1555 }
1556 Ok(binary_expr)
1557 };
1558 if is_comparison_op && !should_return_bool {
1559 self.filter_on_field_column(input, bin_expr_builder)
1560 } else {
1561 let projected =
1562 self.projection_for_each_field_column(input, bin_expr_builder)?;
1563 self.filter_binary_projection(
1564 projected,
1565 has_native_histogram,
1566 preserve_any_value,
1567 retain_field_columns,
1568 )
1569 }
1570 }
1571 (None, None) => {
1573 let mut left_input = self.prom_expr_to_plan(lhs, query_engine_state).await?;
1574 let left_field_columns = self.ctx.field_columns.clone();
1575 let left_time_index_column = self.ctx.time_index_column.clone();
1576 let mut left_table_ref = self
1577 .table_ref()
1578 .unwrap_or_else(|_| TableReference::bare(""));
1579 let mut left_context = self.ctx.clone();
1580
1581 let mut right_input = self.prom_expr_to_plan(rhs, query_engine_state).await?;
1582 let right_field_columns = self.ctx.field_columns.clone();
1583 let right_time_index_column = self.ctx.time_index_column.clone();
1584 let mut right_table_ref = self
1585 .table_ref()
1586 .unwrap_or_else(|_| TableReference::bare(""));
1587 let mut right_context = self.ctx.clone();
1588
1589 if let Some(rewritten) = matching_filters::propagate(
1593 binary_expr,
1594 &left_context.tag_columns,
1595 &right_context.tag_columns,
1596 &left_context.aggregation_field_labels,
1597 &right_context.aggregation_field_labels,
1598 ) {
1599 if rewritten.lhs.as_ref() != lhs.as_ref() {
1603 let selectors = std::mem::take(&mut left_context.selector_matcher);
1604 left_input = self
1605 .prom_expr_to_plan(&rewritten.lhs, query_engine_state)
1606 .await?;
1607 left_context = self.ctx.clone();
1608 left_context.selector_matcher = selectors;
1609 }
1610 if rewritten.rhs.as_ref() != rhs.as_ref() {
1611 let selectors = std::mem::take(&mut right_context.selector_matcher);
1612 right_input = self
1613 .prom_expr_to_plan(&rewritten.rhs, query_engine_state)
1614 .await?;
1615 right_context = self.ctx.clone();
1616 right_context.selector_matcher = selectors;
1617 }
1618 self.ctx = right_context.clone();
1620 }
1621
1622 let left_is_empty_metric = Self::is_empty_metric(&left_input);
1623 let right_is_empty_metric = Self::is_empty_metric(&right_input);
1624
1625 if Self::is_token_a_set_op(*op) {
1629 return self.set_op_on_non_field_columns(
1630 left_input,
1631 right_input,
1632 left_context,
1633 right_context,
1634 *op,
1635 modifier,
1636 );
1637 }
1638
1639 let has_native_histogram = Self::field_columns_contain_native_histogram(
1640 left_input.schema(),
1641 &left_field_columns,
1642 ) || Self::field_columns_contain_native_histogram(
1643 right_input.schema(),
1644 &right_field_columns,
1645 );
1646
1647 if left_table_ref == right_table_ref {
1649 left_table_ref = TableReference::bare("lhs");
1651 right_table_ref = TableReference::bare("rhs");
1652 if self.ctx.tag_columns.is_empty() {
1658 self.ctx = left_context.clone();
1659 self.ctx.table_name = Some("lhs".to_string());
1660 } else {
1661 self.ctx.table_name = Some("rhs".to_string());
1662 }
1663 } else if right_is_empty_metric && !left_is_empty_metric {
1664 self.ctx = left_context.clone();
1665 }
1666 let broadcast_scalar = !is_comparison_op;
1669 let (field_groups, invalid_field_pairs) = Self::align_binary_field_columns(
1670 left_input.schema(),
1671 right_input.schema(),
1672 &left_field_columns,
1673 &right_field_columns,
1674 *op,
1675 broadcast_scalar && lhs.value_type() == ValueType::Scalar,
1676 broadcast_scalar && rhs.value_type() == ValueType::Scalar,
1677 );
1678 let left_aligned_field_columns = field_groups
1679 .iter()
1680 .flat_map(|(_, pairs)| {
1681 pairs
1682 .iter()
1683 .map(|(left_col_name, _)| (*left_col_name).clone())
1684 })
1685 .collect::<Vec<_>>();
1686 let right_aligned_field_columns = field_groups
1687 .iter()
1688 .flat_map(|(_, pairs)| {
1689 pairs
1690 .iter()
1691 .map(|(_, right_col_name)| (*right_col_name).clone())
1692 })
1693 .collect::<Vec<_>>();
1694 self.ctx.field_columns = field_groups
1697 .iter()
1698 .map(|(output, _)| output.clone())
1699 .collect();
1700 let mut field_groups = field_groups.into_iter();
1701 let has_empty_metric_operand = left_is_empty_metric || right_is_empty_metric;
1703
1704 let only_join_time_index = lhs.value_type() == ValueType::Scalar
1705 || rhs.value_type() == ValueType::Scalar
1706 || has_empty_metric_operand
1707 || ((left_context.tag_columns.is_empty()
1708 || right_context.tag_columns.is_empty())
1709 && !left_context
1710 .tag_columns
1711 .iter()
1712 .chain(&right_context.tag_columns)
1713 .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL));
1714 let join_plan = self.join_on_non_field_columns(
1715 left_input,
1716 right_input,
1717 left_table_ref.clone(),
1718 right_table_ref.clone(),
1719 left_time_index_column,
1720 right_time_index_column,
1721 only_join_time_index,
1722 modifier,
1723 &left_context,
1724 &right_context,
1725 )?;
1726 let join_plan_schema = join_plan.schema().clone();
1727 let result_labels = if only_join_time_index {
1731 None
1732 } else {
1733 Self::binary_result_labels(&left_context, &right_context, modifier)
1734 .map(|labels| {
1735 Self::binary_result_label_projection(
1736 &join_plan_schema,
1737 &left_table_ref,
1738 &right_table_ref,
1739 &left_context,
1740 &right_context,
1741 labels,
1742 )
1743 })
1744 .transpose()?
1745 };
1746 if let Some(labels) = &result_labels {
1747 labels.apply(&mut self.ctx);
1748 }
1749 let promql_annotations = self.promql_annotations.clone();
1750 let invalid_pair_predicates = invalid_field_pairs
1753 .into_iter()
1754 .filter(|_| promql_annotations.is_some())
1755 .map(|(left_col_name, right_col_name)| {
1756 let left_field = join_plan_schema
1757 .qualified_field_with_name(Some(&left_table_ref), left_col_name)
1758 .context(DataFusionPlanningSnafu)?;
1759 let right_field = join_plan_schema
1760 .qualified_field_with_name(Some(&right_table_ref), right_col_name)
1761 .context(DataFusionPlanningSnafu)?;
1762 let left_is_histogram =
1763 left_field.1.data_type() == &Self::native_histogram_arrow_type();
1764 let right_is_histogram =
1765 right_field.1.data_type() == &Self::native_histogram_arrow_type();
1766 let drop_expr = Self::native_histogram_binary_expr(
1767 *op,
1768 DfExpr::Column(left_field.into()),
1769 left_is_histogram,
1770 DfExpr::Column(right_field.into()),
1771 right_is_histogram,
1772 true,
1773 promql_annotations.clone(),
1774 )?
1775 .with_context(|| UnexpectedPlanExprSnafu {
1776 desc: "invalid native histogram pair produced no drop expression",
1777 })?;
1778 Ok(DfExpr::Not(Box::new(drop_expr)))
1779 })
1780 .collect::<Result<Vec<_>>>()?;
1781 let join_plan = if let Some(predicate) = conjunction(invalid_pair_predicates) {
1782 LogicalPlanBuilder::from(join_plan)
1783 .filter(predicate)
1784 .context(DataFusionPlanningSnafu)?
1785 .build()
1786 .context(DataFusionPlanningSnafu)?
1787 } else {
1788 join_plan
1789 };
1790
1791 let bin_expr_builder = |_: &String| {
1792 let (_, field_pairs) =
1793 field_groups
1794 .next()
1795 .with_context(|| UnexpectedPlanExprSnafu {
1796 desc: "missing binary field group",
1797 })?;
1798 let binary_exprs = field_pairs
1799 .into_iter()
1800 .map(|(left_col_name, right_col_name)| {
1801 let left_field = join_plan_schema
1802 .qualified_field_with_name(Some(&left_table_ref), left_col_name)
1803 .context(DataFusionPlanningSnafu)?;
1804 let right_field = join_plan_schema
1805 .qualified_field_with_name(Some(&right_table_ref), right_col_name)
1806 .context(DataFusionPlanningSnafu)?;
1807 let left_is_histogram =
1808 left_field.1.data_type() == &Self::native_histogram_arrow_type();
1809 let right_is_histogram =
1810 right_field.1.data_type() == &Self::native_histogram_arrow_type();
1811 let left_col = left_field.into();
1812 let right_col = right_field.into();
1813
1814 let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1815 let lhs = DfExpr::Column(left_col);
1816 let rhs = DfExpr::Column(right_col);
1817 let mut binary_expr = match Self::native_histogram_binary_expr(
1818 *op,
1819 lhs.clone(),
1820 left_is_histogram,
1821 rhs.clone(),
1822 right_is_histogram,
1823 is_comparison_op && !should_return_bool,
1824 promql_annotations.clone(),
1825 )? {
1826 Some(expr) => expr,
1827 None => binary_expr_builder(lhs, rhs)?,
1828 };
1829 if is_comparison_op && should_return_bool {
1830 binary_expr = DfExpr::Cast(Cast::new(
1831 Box::new(binary_expr),
1832 ArrowDataType::Float64,
1833 ));
1834 }
1835 Ok(binary_expr)
1836 })
1837 .collect::<Result<Vec<_>>>()?;
1838 if let [binary_expr] = binary_exprs.as_slice() {
1839 Ok(binary_expr.clone())
1840 } else {
1841 Ok(DfExpr::ScalarFunction(ScalarFunction {
1842 func: coalesce(),
1843 args: binary_exprs,
1844 }))
1845 }
1846 };
1847 if is_comparison_op && !should_return_bool {
1848 let filtered = self.filter_on_field_column(join_plan, bin_expr_builder)?;
1855 let (project_table_ref, mut project_context, project_field_columns) =
1856 match (lhs.value_type(), rhs.value_type()) {
1857 (ValueType::Scalar, ValueType::Vector) => (
1858 &right_table_ref,
1859 right_context.clone(),
1860 right_aligned_field_columns,
1861 ),
1862 _ => (
1863 &left_table_ref,
1864 left_context.clone(),
1865 left_aligned_field_columns,
1866 ),
1867 };
1868 project_context.field_columns = project_field_columns;
1869 self.project_binary_join_side(
1870 filtered,
1871 project_table_ref,
1872 &project_context,
1873 result_labels.as_ref(),
1874 )
1875 } else {
1876 let projected = self.projection_for_each_field_column_with_labels(
1877 join_plan,
1878 result_labels.as_ref(),
1879 bin_expr_builder,
1880 )?;
1881 let preserve_any_value = Self::field_columns_are_alternative_samples(
1882 projected.schema(),
1883 &self.ctx.field_columns,
1884 );
1885 let retain_field_columns = vec![true; self.ctx.field_columns.len()];
1886 self.filter_binary_projection(
1887 projected,
1888 has_native_histogram,
1889 preserve_any_value,
1890 retain_field_columns,
1891 )
1892 }
1893 }
1894 }
1895 }
1896
1897 fn filter_binary_projection(
1898 &mut self,
1899 input: LogicalPlan,
1900 has_native_histogram: bool,
1901 preserve_any_value: bool,
1902 retain_field_columns: Vec<bool>,
1903 ) -> Result<LogicalPlan> {
1904 if !has_native_histogram {
1905 return Ok(input);
1906 }
1907
1908 ensure!(
1909 retain_field_columns.len() == self.ctx.field_columns.len(),
1910 UnexpectedPlanExprSnafu {
1911 desc: "binary output field count changed unexpectedly",
1912 }
1913 );
1914
1915 let filtered = LogicalPlanBuilder::from(input)
1916 .filter(self.create_empty_values_filter_expr(preserve_any_value)?)
1917 .context(DataFusionPlanningSnafu)?
1918 .build()
1919 .context(DataFusionPlanningSnafu)?;
1920 if retain_field_columns.iter().all(|retain| *retain) {
1921 return Ok(filtered);
1922 }
1923
1924 let retained = self
1925 .ctx
1926 .field_columns
1927 .iter()
1928 .zip(retain_field_columns)
1929 .filter(|(_, retain)| *retain)
1930 .map(|(field, _)| field.clone())
1931 .collect::<Vec<_>>();
1932 if retained.is_empty() {
1933 return Ok(filtered);
1934 }
1935 self.ctx.field_columns = retained;
1936
1937 let mut output_columns = self
1938 .ctx
1939 .field_columns
1940 .iter()
1941 .chain(&self.ctx.tag_columns)
1942 .cloned()
1943 .collect::<HashSet<_>>();
1944 output_columns.extend(self.ctx.time_index_column.iter().cloned());
1945 if self.ctx.use_tsid {
1946 output_columns.insert(DATA_SCHEMA_TSID_COLUMN_NAME.to_string());
1947 }
1948 let project_exprs = filtered
1949 .schema()
1950 .iter()
1951 .filter(|(_, field)| output_columns.contains(field.name()))
1952 .map(|(qualifier, field)| {
1953 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
1954 })
1955 .collect::<Vec<_>>();
1956 LogicalPlanBuilder::from(filtered)
1957 .project(project_exprs)
1958 .context(DataFusionPlanningSnafu)?
1959 .build()
1960 .context(DataFusionPlanningSnafu)
1961 }
1962
1963 fn project_binary_join_side(
1964 &mut self,
1965 input: LogicalPlan,
1966 table_ref: &TableReference,
1967 context: &PromPlannerContext,
1968 result_labels: Option<&BinaryResultLabels>,
1969 ) -> Result<LogicalPlan> {
1970 let schema = input.schema();
1971
1972 let mut project_exprs =
1973 Vec::with_capacity(context.tag_columns.len() + context.field_columns.len() + 2);
1974
1975 if let Some(time_index_column) = &context.time_index_column {
1977 let time_index_col = schema
1978 .qualified_field_with_name(Some(table_ref), time_index_column)
1979 .context(DataFusionPlanningSnafu)?
1980 .into();
1981 project_exprs.push(DfExpr::Column(time_index_col));
1982 }
1983
1984 for field_column in &context.field_columns {
1986 let field_col = schema
1987 .qualified_field_with_name(Some(table_ref), field_column)
1988 .context(DataFusionPlanningSnafu)?
1989 .into();
1990 project_exprs.push(DfExpr::Column(field_col));
1991 }
1992
1993 match result_labels {
1995 Some(labels) => project_exprs.extend(labels.exprs.iter().cloned()),
1996 None => {
1997 for tag_column in &context.tag_columns {
1998 let tag_col = schema
1999 .qualified_field_with_name(Some(table_ref), tag_column)
2000 .context(DataFusionPlanningSnafu)?
2001 .into();
2002 project_exprs.push(DfExpr::Column(tag_col));
2003 }
2004 }
2005 }
2006
2007 let tsid_col = match result_labels {
2010 Some(labels) => labels.tsid_projection(Some(table_ref.clone())),
2011 None => Self::optional_tsid_projection(schema, Some(table_ref), context.use_tsid),
2012 };
2013 if let Some(tsid_col) = tsid_col {
2014 project_exprs.push(tsid_col);
2015 }
2016
2017 let plan = LogicalPlanBuilder::from(input)
2018 .project(project_exprs)
2019 .context(DataFusionPlanningSnafu)?
2020 .build()
2021 .context(DataFusionPlanningSnafu)?;
2022
2023 self.ctx = context.clone();
2026 self.ctx.table_name = None;
2027 self.ctx.schema_name = None;
2028 if let Some(labels) = result_labels {
2029 labels.apply(&mut self.ctx);
2030 }
2031
2032 Ok(plan)
2033 }
2034
2035 fn prom_number_lit_to_plan(&mut self, number_literal: &NumberLiteral) -> Result<LogicalPlan> {
2036 let NumberLiteral { val } = number_literal;
2037 self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
2038 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
2039 self.ctx.reset_table_name_and_schema();
2040 let literal_expr = df_prelude::lit(*val);
2041
2042 let plan = LogicalPlan::Extension(Extension {
2043 node: Arc::new(
2044 EmptyMetric::new(
2045 self.ctx.start,
2046 self.ctx.end,
2047 self.ctx.interval,
2048 SPECIAL_TIME_FUNCTION.to_string(),
2049 DEFAULT_FIELD_COLUMN.to_string(),
2050 Some(literal_expr),
2051 )
2052 .context(DataFusionPlanningSnafu)?,
2053 ),
2054 });
2055 Ok(plan)
2056 }
2057
2058 fn prom_string_lit_to_plan(&mut self, string_literal: &StringLiteral) -> Result<LogicalPlan> {
2059 let StringLiteral { val } = string_literal;
2060 self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
2061 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
2062 self.ctx.reset_table_name_and_schema();
2063 let literal_expr = df_prelude::lit(val.clone());
2064
2065 let plan = LogicalPlan::Extension(Extension {
2066 node: Arc::new(
2067 EmptyMetric::new(
2068 self.ctx.start,
2069 self.ctx.end,
2070 self.ctx.interval,
2071 SPECIAL_TIME_FUNCTION.to_string(),
2072 DEFAULT_FIELD_COLUMN.to_string(),
2073 Some(literal_expr),
2074 )
2075 .context(DataFusionPlanningSnafu)?,
2076 ),
2077 });
2078 Ok(plan)
2079 }
2080
2081 async fn prom_vector_selector_to_plan(
2082 &mut self,
2083 vector_selector: &VectorSelector,
2084 timestamp_fn: bool,
2085 ) -> Result<LogicalPlan> {
2086 let VectorSelector {
2087 name,
2088 offset,
2089 matchers,
2090 at: _,
2091 } = vector_selector;
2092 let matchers = self.preprocess_label_matchers(matchers, name)?;
2093 if let Some(empty_plan) = self.setup_context().await? {
2094 return Ok(empty_plan);
2095 }
2096 let offset_ms = match offset {
2097 Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond,
2098 Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond),
2099 None => 0,
2100 };
2101 let normalize = self
2102 .selector_to_series_normalize_plan(offset, matchers, false)
2103 .await?;
2104 let time_index_column =
2105 self.ctx
2106 .time_index_column
2107 .clone()
2108 .with_context(|| TimeIndexNotFoundSnafu {
2109 table: self.ctx.table_name.clone().unwrap_or_default(),
2110 })?;
2111
2112 let (normalize, timestamp_value_column) = if timestamp_fn {
2113 let occupied = normalize
2116 .schema()
2117 .fields()
2118 .iter()
2119 .map(|field| field.name().as_str())
2120 .collect::<HashSet<_>>();
2121 let mut timestamp_value_column = TIMESTAMP_VALUE_PREFIX.to_string();
2122 while occupied.contains(timestamp_value_column.as_str()) {
2123 timestamp_value_column.push('_');
2124 }
2125 let mut project_exprs = normalize
2126 .schema()
2127 .iter()
2128 .map(|(qualifier, field)| {
2129 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
2130 })
2131 .collect::<Vec<_>>();
2132 let unit_factor = match col(&time_index_column)
2136 .get_type(normalize.schema())
2137 .context(DataFusionPlanningSnafu)?
2138 {
2139 ArrowDataType::Timestamp(ArrowTimeUnit::Second, _) => (1_000_i128, 4, 0),
2140 ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, _) => (1, 1, 0),
2141 ArrowDataType::Timestamp(ArrowTimeUnit::Microsecond, _) => (1, 4, 3),
2142 ArrowDataType::Timestamp(ArrowTimeUnit::Nanosecond, _) => (1, 7, 6),
2143 _ => unreachable!("time index is a timestamp"),
2144 };
2145 let sample_time = col(&time_index_column)
2146 .cast_to(&ArrowDataType::Int64, normalize.schema())
2147 .context(DataFusionPlanningSnafu)?
2148 .cast_to(&ArrowDataType::Decimal128(19, 0), normalize.schema())
2149 .context(DataFusionPlanningSnafu)?;
2150 let sample_time = DfExpr::BinaryExpr(BinaryExpr {
2151 left: Box::new(sample_time),
2152 op: Operator::Multiply,
2153 right: Box::new(lit(ScalarValue::Decimal128(
2154 Some(unit_factor.0),
2155 unit_factor.1,
2156 unit_factor.2,
2157 ))),
2158 });
2159 let sample_time = DfExpr::BinaryExpr(BinaryExpr {
2160 left: Box::new(sample_time),
2161 op: Operator::Plus,
2162 right: Box::new(lit(ScalarValue::Decimal128(Some(offset_ms as i128), 19, 0))),
2163 })
2164 .cast_to(&ArrowDataType::Int64, normalize.schema())
2165 .context(DataFusionPlanningSnafu)?
2166 .cast_to(&ArrowDataType::Float64, normalize.schema())
2167 .context(DataFusionPlanningSnafu)?;
2168 let sample_time = DfExpr::BinaryExpr(BinaryExpr {
2169 left: Box::new(sample_time),
2170 op: Operator::Divide,
2171 right: Box::new(lit(1000.0)),
2172 });
2173 project_exprs.push(sample_time.alias(×tamp_value_column));
2174 let normalize = LogicalPlanBuilder::from(normalize)
2175 .project(project_exprs)
2176 .context(DataFusionPlanningSnafu)?
2177 .build()
2178 .context(DataFusionPlanningSnafu)?;
2179 (normalize, Some(timestamp_value_column))
2180 } else {
2181 (normalize, None)
2182 };
2183
2184 let field_column = self.ctx.field_columns.first().cloned();
2185 let manipulate = InstantManipulate::new(
2186 self.ctx.start,
2187 self.ctx.end,
2188 self.ctx.lookback_delta,
2189 self.ctx.interval,
2190 offset_ms,
2191 time_index_column,
2192 if self.ctx.use_tsid {
2193 vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]
2194 } else {
2195 self.ctx.tag_columns.clone()
2196 },
2197 field_column,
2198 normalize,
2199 );
2200 let manipulate = LogicalPlan::Extension(Extension {
2201 node: Arc::new(manipulate),
2202 });
2203 if let Some(timestamp_value_column) = timestamp_value_column {
2204 self.create_timestamp_func_plan(manipulate, ×tamp_value_column)
2205 } else {
2206 Ok(manipulate)
2207 }
2208 }
2209
2210 fn create_timestamp_func_plan(
2233 &mut self,
2234 input: LogicalPlan,
2235 timestamp_value_column: &str,
2236 ) -> Result<LogicalPlan> {
2237 let time_expr = col(timestamp_value_column).alias(DEFAULT_FIELD_COLUMN);
2238 self.ctx.field_columns = vec![time_expr.schema_name().to_string()];
2239 let mut project_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 2);
2240 project_exprs.push(self.create_time_index_column_expr()?);
2241 project_exprs.push(time_expr);
2242 project_exprs.extend(self.create_tag_column_exprs()?);
2243
2244 LogicalPlanBuilder::from(input)
2245 .project(project_exprs)
2246 .context(DataFusionPlanningSnafu)?
2247 .build()
2248 .context(DataFusionPlanningSnafu)
2249 }
2250
2251 async fn prom_matrix_selector_to_plan(
2252 &mut self,
2253 matrix_selector: &MatrixSelector,
2254 ) -> Result<LogicalPlan> {
2255 let MatrixSelector { vs, range } = matrix_selector;
2256 let VectorSelector {
2257 name,
2258 offset,
2259 matchers,
2260 ..
2261 } = vs;
2262 let matchers = self.preprocess_label_matchers(matchers, name)?;
2263 ensure!(!range.is_zero(), ZeroRangeSelectorSnafu);
2264 let range_ms = range.as_millis() as _;
2265 self.ctx.range = Some(range_ms);
2266 let offset_ms = match offset {
2267 Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond,
2268 Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond),
2269 None => 0,
2270 };
2271
2272 let normalize = match self.setup_context().await? {
2275 Some(empty_plan) => empty_plan,
2276 None => {
2277 self.selector_to_series_normalize_plan(offset, matchers, true)
2278 .await?
2279 }
2280 };
2281 let manipulate = RangeManipulate::new(
2282 self.ctx.start,
2283 self.ctx.end,
2284 self.ctx.interval,
2285 offset_ms,
2286 range_ms,
2288 self.ctx
2289 .time_index_column
2290 .clone()
2291 .expect("time index should be set in `setup_context`"),
2292 self.ctx.field_columns.clone(),
2293 normalize,
2294 )
2295 .context(DataFusionPlanningSnafu)?;
2296
2297 Ok(LogicalPlan::Extension(Extension {
2298 node: Arc::new(manipulate),
2299 }))
2300 }
2301
2302 async fn prom_call_expr_to_plan(
2303 &mut self,
2304 query_engine_state: &QueryEngineState,
2305 call_expr: &Call,
2306 ) -> Result<LogicalPlan> {
2307 let Call { func, args } = call_expr;
2308 match func.name {
2310 SPECIAL_HISTOGRAM_QUANTILE | SPECIAL_HISTOGRAM_FRACTION => {
2311 return self
2312 .create_histogram_plan(func.name, args, query_engine_state)
2313 .await;
2314 }
2315 SPECIAL_VECTOR_FUNCTION => return self.create_vector_plan(args).await,
2316 SCALAR_FUNCTION => return self.create_scalar_plan(args, query_engine_state).await,
2317 SPECIAL_ABSENT_FUNCTION => {
2318 return self.create_absent_plan(args, query_engine_state).await;
2319 }
2320 _ => {}
2321 }
2322
2323 let args = self.create_function_args(&args.args)?;
2325 let input = if let Some(prom_expr) = &args.input {
2326 self.prom_expr_to_plan_inner(prom_expr, func.name == "timestamp", query_engine_state)
2327 .await?
2328 } else {
2329 self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
2330 self.ctx.reset_table_name_and_schema();
2331 self.ctx.tag_columns = vec![];
2332 self.ctx.aggregation_field_labels.clear();
2333 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
2334 LogicalPlan::Extension(Extension {
2335 node: Arc::new(
2336 EmptyMetric::new(
2337 self.ctx.start,
2338 self.ctx.end,
2339 self.ctx.interval,
2340 SPECIAL_TIME_FUNCTION.to_string(),
2341 DEFAULT_FIELD_COLUMN.to_string(),
2342 None,
2343 )
2344 .context(DataFusionPlanningSnafu)?,
2345 ),
2346 })
2347 };
2348 let (mut func_exprs, new_tags) = self.create_function_expr(
2349 func,
2350 args.literals.clone(),
2351 input.schema(),
2352 query_engine_state,
2353 )?;
2354 func_exprs.insert(0, self.create_time_index_column_expr()?);
2355 func_exprs.extend_from_slice(&self.create_tag_column_exprs()?);
2356 if let Some(tsid_col) =
2357 Self::optional_tsid_projection(input.schema(), None, self.ctx.use_tsid)
2358 {
2359 func_exprs.push(tsid_col);
2360 }
2361
2362 let builder = LogicalPlanBuilder::from(input)
2368 .project(func_exprs)
2369 .context(DataFusionPlanningSnafu)?
2370 .filter(self.create_empty_values_filter_expr(true)?)
2371 .context(DataFusionPlanningSnafu)?;
2372
2373 let builder = match func.name {
2374 "sort" => builder
2375 .sort(self.create_field_columns_sort_exprs(true))
2376 .context(DataFusionPlanningSnafu)?,
2377 "sort_desc" => builder
2378 .sort(self.create_field_columns_sort_exprs(false))
2379 .context(DataFusionPlanningSnafu)?,
2380 "sort_by_label" => builder
2381 .sort(Self::create_sort_exprs_by_tags(
2382 func.name,
2383 args.literals,
2384 true,
2385 )?)
2386 .context(DataFusionPlanningSnafu)?,
2387 "sort_by_label_desc" => builder
2388 .sort(Self::create_sort_exprs_by_tags(
2389 func.name,
2390 args.literals,
2391 false,
2392 )?)
2393 .context(DataFusionPlanningSnafu)?,
2394
2395 _ => builder,
2396 };
2397
2398 for tag in new_tags {
2401 self.ctx.tag_columns.push(tag);
2402 }
2403
2404 let plan = builder.build().context(DataFusionPlanningSnafu)?;
2405 common_telemetry::debug!("Created PromQL function plan: {plan:?} for {call_expr:?}");
2406
2407 Ok(plan)
2408 }
2409
2410 async fn prom_ext_expr_to_plan(
2411 &mut self,
2412 query_engine_state: &QueryEngineState,
2413 ext_expr: &promql_parser::parser::ast::Extension,
2414 ) -> Result<LogicalPlan> {
2415 let expr = &ext_expr.expr;
2417 let children = expr.children();
2418 let plan = self
2419 .prom_expr_to_plan(&children[0], query_engine_state)
2420 .await?;
2421 match expr.name() {
2427 ANALYZE_NODE_NAME => LogicalPlanBuilder::from(plan)
2428 .explain(false, true)
2429 .unwrap()
2430 .build()
2431 .context(DataFusionPlanningSnafu),
2432 ANALYZE_VERBOSE_NODE_NAME => LogicalPlanBuilder::from(plan)
2433 .explain(true, true)
2434 .unwrap()
2435 .build()
2436 .context(DataFusionPlanningSnafu),
2437 EXPLAIN_NODE_NAME => LogicalPlanBuilder::from(plan)
2438 .explain(false, false)
2439 .unwrap()
2440 .build()
2441 .context(DataFusionPlanningSnafu),
2442 EXPLAIN_VERBOSE_NODE_NAME => LogicalPlanBuilder::from(plan)
2443 .explain(true, false)
2444 .unwrap()
2445 .build()
2446 .context(DataFusionPlanningSnafu),
2447 ALIAS_NODE_NAME => {
2448 let alias = expr
2449 .as_any()
2450 .downcast_ref::<AliasExpr>()
2451 .context(UnexpectedPlanExprSnafu {
2452 desc: "Expected AliasExpr",
2453 })?
2454 .alias
2455 .clone();
2456 self.apply_alias(plan, alias)
2457 }
2458 _ => LogicalPlanBuilder::empty(true)
2459 .build()
2460 .context(DataFusionPlanningSnafu),
2461 }
2462 }
2463
2464 #[allow(clippy::mutable_key_type)]
2474 fn preprocess_label_matchers(
2475 &mut self,
2476 label_matchers: &Matchers,
2477 name: &Option<String>,
2478 ) -> Result<Matchers> {
2479 self.ctx.reset();
2480
2481 let metric_name;
2482 if let Some(name) = name.clone() {
2483 metric_name = Some(name);
2484 ensure!(
2485 label_matchers.find_matchers(METRIC_NAME).is_empty(),
2486 MultipleMetricMatchersSnafu
2487 );
2488 } else {
2489 let mut matches = label_matchers.find_matchers(METRIC_NAME);
2490 ensure!(!matches.is_empty(), NoMetricMatcherSnafu);
2491 ensure!(matches.len() == 1, MultipleMetricMatchersSnafu);
2492 ensure!(
2493 matches[0].op == MatchOp::Equal,
2494 UnsupportedMatcherOpSnafu {
2495 matcher_op: matches[0].op.to_string(),
2496 matcher: METRIC_NAME
2497 }
2498 );
2499 metric_name = matches.pop().map(|m| m.value);
2500 }
2501
2502 self.ctx.table_name = metric_name;
2503
2504 let mut matchers: Vec<Matcher> = Vec::with_capacity(label_matchers.matchers.len());
2507 for matcher in &label_matchers.matchers {
2508 if matcher.name == FIELD_COLUMN_MATCHER {
2510 self.ctx
2511 .field_column_matcher
2512 .get_or_insert_default()
2513 .push(matcher.clone());
2514 } else if matcher.name == SCHEMA_COLUMN_MATCHER || matcher.name == DB_COLUMN_MATCHER {
2515 ensure!(
2516 matcher.op == MatchOp::Equal,
2517 UnsupportedMatcherOpSnafu {
2518 matcher: matcher.name.clone(),
2519 matcher_op: matcher.op.to_string(),
2520 }
2521 );
2522 self.ctx.schema_name = Some(matcher.value.clone());
2523 } else if matcher.name != METRIC_NAME {
2524 self.ctx.selector_matcher.push(matcher.clone());
2525 if !matchers.contains(matcher) {
2526 matchers.push(matcher.clone());
2527 }
2528 }
2529 }
2530
2531 Ok(Matchers::new(matchers))
2532 }
2533
2534 async fn selector_to_series_normalize_plan(
2535 &mut self,
2536 offset: &Option<Offset>,
2537 label_matchers: Matchers,
2538 is_range_selector: bool,
2539 ) -> Result<LogicalPlan> {
2540 let table_ref = self.table_ref()?;
2542 let mut table_scan = self.create_table_scan_plan(table_ref.clone()).await?;
2543 let table_schema = table_scan.schema();
2544
2545 let offset_duration = match offset {
2547 Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond,
2548 Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond),
2549 None => 0,
2550 };
2551 let mut scan_filters = Self::matchers_to_expr(label_matchers.clone(), table_schema)?;
2552 if let Some(time_index_filter) =
2553 self.build_time_index_filter(offset_duration, table_schema)?
2554 {
2555 scan_filters.push(time_index_filter);
2556 }
2557 if let Some(filter) = conjunction(scan_filters) {
2558 table_scan = LogicalPlanBuilder::from(table_scan)
2559 .filter(filter)
2560 .context(DataFusionPlanningSnafu)?
2561 .build()
2562 .context(DataFusionPlanningSnafu)?;
2563 }
2564
2565 if let Some(field_matchers) = &self.ctx.field_column_matcher {
2567 let col_set = self.ctx.field_columns.iter().collect::<HashSet<_>>();
2568 let mut result_set = HashSet::new();
2570 let mut reverse_set = HashSet::new();
2572 for matcher in field_matchers {
2573 match &matcher.op {
2574 MatchOp::Equal => {
2575 if col_set.contains(&matcher.value) {
2576 let _ = result_set.insert(matcher.value.clone());
2577 } else {
2578 return Err(ColumnNotFoundSnafu {
2579 col: matcher.value.clone(),
2580 }
2581 .build());
2582 }
2583 }
2584 MatchOp::NotEqual => {
2585 if col_set.contains(&matcher.value) {
2586 let _ = reverse_set.insert(matcher.value.clone());
2587 } else {
2588 return Err(ColumnNotFoundSnafu {
2589 col: matcher.value.clone(),
2590 }
2591 .build());
2592 }
2593 }
2594 MatchOp::Re(regex) => {
2595 for col in &self.ctx.field_columns {
2596 if regex.is_match(col) {
2597 let _ = result_set.insert(col.clone());
2598 }
2599 }
2600 }
2601 MatchOp::NotRe(regex) => {
2602 for col in &self.ctx.field_columns {
2603 if regex.is_match(col) {
2604 let _ = reverse_set.insert(col.clone());
2605 }
2606 }
2607 }
2608 }
2609 }
2610 if result_set.is_empty() {
2612 result_set = col_set.into_iter().cloned().collect();
2613 }
2614 for col in reverse_set {
2615 let _ = result_set.remove(&col);
2616 }
2617
2618 self.ctx.field_columns = self
2620 .ctx
2621 .field_columns
2622 .drain(..)
2623 .filter(|col| result_set.contains(col))
2624 .collect();
2625
2626 let exprs = result_set
2627 .into_iter()
2628 .map(|col| DfExpr::Column(Column::new_unqualified(col)))
2629 .chain(self.create_tag_column_exprs()?)
2630 .chain(
2631 self.ctx
2632 .use_tsid
2633 .then_some(DfExpr::Column(Column::new_unqualified(
2634 DATA_SCHEMA_TSID_COLUMN_NAME,
2635 ))),
2636 )
2637 .chain(Some(self.create_time_index_column_expr()?))
2638 .collect::<Vec<_>>();
2639
2640 table_scan = LogicalPlanBuilder::from(table_scan)
2642 .project(exprs)
2643 .context(DataFusionPlanningSnafu)?
2644 .build()
2645 .context(DataFusionPlanningSnafu)?;
2646 }
2647
2648 let series_key_columns = if self.ctx.use_tsid {
2650 vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]
2651 } else {
2652 self.ctx.tag_columns.clone()
2653 };
2654
2655 let sort_exprs = if self.ctx.use_tsid {
2656 vec![
2657 DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME)).sort(true, true),
2658 self.create_time_index_column_expr()?.sort(true, true),
2659 ]
2660 } else {
2661 self.create_tag_and_time_index_column_sort_exprs()?
2662 };
2663
2664 let sort_plan = LogicalPlanBuilder::from(table_scan)
2665 .sort(sort_exprs)
2666 .context(DataFusionPlanningSnafu)?
2667 .build()
2668 .context(DataFusionPlanningSnafu)?;
2669
2670 let time_index_column =
2672 self.ctx
2673 .time_index_column
2674 .clone()
2675 .with_context(|| TimeIndexNotFoundSnafu {
2676 table: table_ref.to_string(),
2677 })?;
2678 let divide_plan = LogicalPlan::Extension(Extension {
2679 node: Arc::new(SeriesDivide::new(
2680 series_key_columns.clone(),
2681 time_index_column,
2682 sort_plan,
2683 )),
2684 });
2685
2686 if !is_range_selector && offset_duration == 0 {
2688 return Ok(divide_plan);
2689 }
2690 let series_normalize = SeriesNormalize::new(
2691 offset_duration,
2692 self.ctx
2693 .time_index_column
2694 .clone()
2695 .with_context(|| TimeIndexNotFoundSnafu {
2696 table: table_ref.to_quoted_string(),
2697 })?,
2698 is_range_selector,
2699 series_key_columns,
2700 divide_plan,
2701 );
2702 let logical_plan = LogicalPlan::Extension(Extension {
2703 node: Arc::new(series_normalize),
2704 });
2705
2706 Ok(logical_plan)
2707 }
2708
2709 fn agg_modifier_to_col(
2716 &mut self,
2717 input_schema: &DFSchemaRef,
2718 modifier: &Option<LabelModifier>,
2719 update_ctx: bool,
2720 ) -> Result<Vec<DfExpr>> {
2721 match modifier {
2722 None => {
2723 if update_ctx {
2724 self.ctx.tag_columns.clear();
2725 self.ctx.aggregation_field_labels.clear();
2726 }
2727 Ok(vec![self.create_time_index_column_expr()?])
2728 }
2729 Some(LabelModifier::Include(labels)) => {
2730 if update_ctx {
2731 self.ctx.aggregation_field_labels = labels
2737 .labels
2738 .iter()
2739 .filter(|label| {
2740 self.ctx.field_columns.contains(label)
2741 || !self.ctx.tag_columns.contains(label)
2742 || self.ctx.aggregation_field_labels.contains(label)
2743 })
2744 .cloned()
2745 .collect();
2746 self.ctx.tag_columns.clear();
2747 }
2748 let mut exprs = Vec::with_capacity(labels.labels.len());
2749 for label in &labels.labels {
2750 if is_metric_engine_internal_column(label) {
2751 continue;
2752 }
2753 if let Some(column_name) = Self::find_case_sensitive_column(input_schema, label)
2755 {
2756 exprs.push(DfExpr::Column(Column::from_name(column_name.clone())));
2757
2758 if update_ctx {
2759 self.ctx.tag_columns.push(column_name);
2761 }
2762 }
2763 }
2764 exprs.push(self.create_time_index_column_expr()?);
2766
2767 Ok(exprs)
2768 }
2769 Some(LabelModifier::Exclude(labels)) => {
2770 let mut all_fields = input_schema
2771 .fields()
2772 .iter()
2773 .map(|f| f.name())
2774 .collect::<BTreeSet<_>>();
2775
2776 all_fields.retain(|col| !is_metric_engine_internal_column(col.as_str()));
2779
2780 for label in &labels.labels {
2783 let _ = all_fields.remove(label);
2784 }
2785
2786 if let Some(time_index) = &self.ctx.time_index_column {
2788 let _ = all_fields.remove(time_index);
2789 }
2790 for value in &self.ctx.field_columns {
2791 let _ = all_fields.remove(value);
2792 }
2793
2794 if update_ctx {
2795 self.ctx.tag_columns = all_fields.iter().map(|col| (*col).clone()).collect();
2797 self.ctx
2801 .aggregation_field_labels
2802 .retain(|label| all_fields.iter().any(|col| *col == label));
2803 }
2804
2805 let mut exprs = all_fields
2807 .into_iter()
2808 .map(|c| DfExpr::Column(Column::from(c)))
2809 .collect::<Vec<_>>();
2810
2811 exprs.push(self.create_time_index_column_expr()?);
2813
2814 Ok(exprs)
2815 }
2816 }
2817 }
2818
2819 pub fn matchers_to_expr(
2821 label_matchers: Matchers,
2822 table_schema: &DFSchemaRef,
2823 ) -> Result<Vec<DfExpr>> {
2824 let mut exprs = Vec::with_capacity(label_matchers.matchers.len());
2825 for matcher in label_matchers.matchers {
2826 if matcher.name == SCHEMA_COLUMN_MATCHER
2827 || matcher.name == DB_COLUMN_MATCHER
2828 || matcher.name == FIELD_COLUMN_MATCHER
2829 {
2830 continue;
2831 }
2832
2833 let accepts_empty = matcher.is_match("");
2834 let column_name = Self::find_case_sensitive_column(table_schema, matcher.name.as_str());
2835 let col = if let Some(column_name) = column_name {
2841 let column = DfExpr::Column(Column::from_name(&column_name));
2842 let field = table_schema
2843 .index_of_column_by_name(None, &column_name)
2844 .map(|index| table_schema.field(index));
2845 if accepts_empty
2846 && let Some(data_type) = field
2847 .filter(|field| {
2848 field.is_nullable()
2849 && Self::string_value_data_type(field.data_type()).is_some()
2850 })
2851 .map(|field| field.data_type())
2852 {
2853 let empty = Self::string_scalar_value(data_type, Some(String::new()))
2854 .expect("nullable label has a string type");
2855 DfExpr::ScalarFunction(ScalarFunction {
2856 func: coalesce(),
2857 args: vec![column, DfExpr::Literal(empty, None)],
2858 })
2859 } else {
2860 column
2861 }
2862 } else {
2863 DfExpr::Literal(ScalarValue::Utf8(Some(String::new())), None)
2864 };
2865 let lit = DfExpr::Literal(ScalarValue::Utf8(Some(matcher.value)), None);
2866 let expr = match matcher.op {
2867 MatchOp::Equal => col.eq(lit),
2868 MatchOp::NotEqual => col.not_eq(lit),
2869 MatchOp::Re(re) => {
2870 if re.as_str() == "^(?:.*)$" {
2876 continue;
2877 }
2878 if re.as_str() == "^(?:.+)$" {
2879 col.not_eq(DfExpr::Literal(
2880 ScalarValue::Utf8(Some(String::new())),
2881 None,
2882 ))
2883 } else {
2884 DfExpr::BinaryExpr(BinaryExpr {
2885 left: Box::new(col),
2886 op: Operator::RegexMatch,
2887 right: Box::new(DfExpr::Literal(
2888 ScalarValue::Utf8(Some(re.as_str().to_string())),
2889 None,
2890 )),
2891 })
2892 }
2893 }
2894 MatchOp::NotRe(re) => {
2895 if re.as_str() == "^(?:.*)$" {
2896 DfExpr::Literal(ScalarValue::Boolean(Some(false)), None)
2897 } else if re.as_str() == "^(?:.+)$" {
2898 col.eq(DfExpr::Literal(
2899 ScalarValue::Utf8(Some(String::new())),
2900 None,
2901 ))
2902 } else {
2903 DfExpr::BinaryExpr(BinaryExpr {
2904 left: Box::new(col),
2905 op: Operator::RegexNotMatch,
2906 right: Box::new(DfExpr::Literal(
2907 ScalarValue::Utf8(Some(re.as_str().to_string())),
2908 None,
2909 )),
2910 })
2911 }
2912 }
2913 };
2914 exprs.push(expr);
2915 }
2916
2917 Ok(exprs)
2918 }
2919
2920 fn find_case_sensitive_column(schema: &DFSchemaRef, column: &str) -> Option<String> {
2921 if is_metric_engine_internal_column(column) {
2922 return None;
2923 }
2924 schema
2925 .fields()
2926 .iter()
2927 .find(|field| field.name() == column)
2928 .map(|field| field.name().clone())
2929 }
2930
2931 fn table_from_source(&self, source: &Arc<dyn TableSource>) -> Result<table::TableRef> {
2932 Ok(source
2933 .downcast_ref::<DefaultTableSource>()
2934 .context(UnknownTableSnafu)?
2935 .table_provider
2936 .downcast_ref::<DfTableProviderAdapter>()
2937 .context(UnknownTableSnafu)?
2938 .table())
2939 }
2940
2941 fn table_ref(&self) -> Result<TableReference> {
2942 let table_name = self
2943 .ctx
2944 .table_name
2945 .clone()
2946 .context(TableNameNotFoundSnafu)?;
2947
2948 let table_ref = if let Some(schema_name) = &self.ctx.schema_name {
2950 TableReference::partial(schema_name.as_str(), table_name.as_str())
2951 } else {
2952 TableReference::bare(table_name.as_str())
2953 };
2954
2955 Ok(table_ref)
2956 }
2957
2958 fn build_time_index_filter(
2959 &self,
2960 offset_duration: i64,
2961 schema: &DFSchemaRef,
2962 ) -> Result<Option<DfExpr>> {
2963 let start = self.ctx.start;
2964 let end = self.ctx.end;
2965 if end < start {
2966 return InvalidTimeRangeSnafu { start, end }.fail();
2967 }
2968 let time_index_expr = self.create_time_index_column_expr()?;
2969 let time_index_name = self.ctx.time_index_column.as_ref().unwrap();
2970 let unit = schema
2971 .index_of_column_by_name(None, time_index_name)
2972 .and_then(|index| match schema.field(index).data_type() {
2973 ArrowDataType::Timestamp(unit, _) => Some(*unit),
2974 _ => None,
2975 })
2976 .unwrap_or(ArrowTimeUnit::Millisecond);
2977 let native_value = |milliseconds: i128| match unit {
2978 ArrowTimeUnit::Second => milliseconds.div_euclid(1_000),
2979 ArrowTimeUnit::Millisecond => milliseconds,
2980 ArrowTimeUnit::Microsecond => milliseconds * 1_000,
2981 ArrowTimeUnit::Nanosecond => milliseconds * 1_000_000,
2982 };
2983 let scalar = |milliseconds: i128| -> Option<ScalarValue> {
2984 let value = i64::try_from(native_value(milliseconds)).ok()?;
2985 Some(match unit {
2986 ArrowTimeUnit::Second => ScalarValue::TimestampSecond(Some(value), None),
2987 ArrowTimeUnit::Millisecond => ScalarValue::TimestampMillisecond(Some(value), None),
2988 ArrowTimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(Some(value), None),
2989 ArrowTimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(Some(value), None),
2990 })
2991 };
2992 let window = self.ctx.range.unwrap_or(self.ctx.lookback_delta);
2993 let filter = |lower_ms: i128, upper_ms: i128| {
2994 let lower_value = native_value(lower_ms);
2995 let upper_value = native_value(upper_ms);
2996 if lower_value > i128::from(i64::MAX) || upper_value < i128::from(i64::MIN) {
2997 return Some(lit(false));
2998 }
2999 let lower_filter = (lower_value >= i128::from(i64::MIN)).then(|| {
3000 let lower = DfExpr::Literal(scalar(lower_ms).unwrap(), None);
3001 if window == 0 {
3002 time_index_expr.clone().gt_eq(lower)
3003 } else if unit == ArrowTimeUnit::Millisecond
3004 && let Some(inclusive_lower) = lower_ms
3005 .checked_add(1)
3006 .and_then(|lower| i64::try_from(lower).ok())
3007 .and_then(|lower| scalar(i128::from(lower)))
3008 {
3009 time_index_expr
3010 .clone()
3011 .gt_eq(DfExpr::Literal(inclusive_lower, None))
3012 } else {
3013 time_index_expr.clone().gt(lower)
3014 }
3015 });
3016 let upper_filter = (upper_value <= i128::from(i64::MAX)).then(|| {
3017 time_index_expr
3018 .clone()
3019 .lt_eq(DfExpr::Literal(scalar(upper_ms).unwrap(), None))
3020 });
3021
3022 match (lower_filter, upper_filter) {
3026 (Some(lower), Some(upper)) => Some(lower.and(upper)),
3027 (Some(filter), None) | (None, Some(filter)) => Some(filter),
3028 (None, None) => None,
3029 }
3030 };
3031 let bounds = |timestamp: i64| {
3032 let upper = i128::from(timestamp) - i128::from(offset_duration);
3033 (upper - i128::from(window), upper)
3034 };
3035 let num_points = (end as i128 - start as i128) / self.ctx.interval as i128;
3036 if num_points > MAX_SCATTER_POINTS as i128 || self.ctx.interval <= INTERVAL_1H {
3037 let (lower, _) = bounds(start);
3038 let (_, upper) = bounds(end);
3039 return Ok(filter(lower, upper));
3040 }
3041 let mut filters = Vec::new();
3042 for timestamp in (start..=end).step_by(self.ctx.interval as usize) {
3043 let (lower, upper) = bounds(timestamp);
3044 let Some(filter) = filter(lower, upper) else {
3045 return Ok(None);
3048 };
3049 filters.push(filter);
3050 }
3051 Ok(filters.into_iter().reduce(DfExpr::or))
3052 }
3053
3054 async fn create_table_scan_plan(&mut self, table_ref: TableReference) -> Result<LogicalPlan> {
3059 let provider = self
3060 .table_provider
3061 .resolve_table(table_ref.clone())
3062 .await
3063 .context(CatalogSnafu)?;
3064
3065 let logical_table = self.table_from_source(&provider)?;
3066
3067 let mut maybe_phy_table_ref = table_ref.clone();
3069 let mut scan_provider = provider;
3070 let mut table_id_filter: Option<u32> = None;
3071
3072 if logical_table.table_info().meta.engine == METRIC_ENGINE_NAME
3075 && let Some(physical_table_name) = logical_table
3076 .table_info()
3077 .meta
3078 .options
3079 .extra_options
3080 .get(LOGICAL_TABLE_METADATA_KEY)
3081 {
3082 let physical_table_ref = if let Some(schema_name) = &self.ctx.schema_name {
3083 TableReference::partial(schema_name.as_str(), physical_table_name.as_str())
3084 } else {
3085 TableReference::bare(physical_table_name.as_str())
3086 };
3087
3088 let physical_provider = match self
3089 .table_provider
3090 .resolve_table(physical_table_ref.clone())
3091 .await
3092 {
3093 Ok(provider) => provider,
3094 Err(e) if e.status_code() == StatusCode::TableNotFound => {
3095 scan_provider.clone()
3098 }
3099 Err(e) => return Err(e).context(CatalogSnafu),
3100 };
3101
3102 if !Arc::ptr_eq(&physical_provider, &scan_provider) {
3103 let physical_table = self.table_from_source(&physical_provider)?;
3105
3106 let has_table_id = physical_table
3107 .schema()
3108 .column_schema_by_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME)
3109 .is_some();
3110 let has_tsid = physical_table
3111 .schema()
3112 .column_schema_by_name(DATA_SCHEMA_TSID_COLUMN_NAME)
3113 .is_some_and(|col| matches!(col.data_type, ConcreteDataType::UInt64(_)));
3114
3115 if has_table_id && has_tsid {
3116 scan_provider = physical_provider;
3117 maybe_phy_table_ref = physical_table_ref;
3118 table_id_filter = Some(logical_table.table_info().ident.table_id);
3119 }
3120 }
3121 }
3122
3123 let scan_table = self.table_from_source(&scan_provider)?;
3124
3125 let use_tsid = table_id_filter.is_some()
3126 && scan_table
3127 .schema()
3128 .column_schema_by_name(DATA_SCHEMA_TSID_COLUMN_NAME)
3129 .is_some_and(|col| matches!(col.data_type, ConcreteDataType::UInt64(_)));
3130 self.ctx.use_tsid = use_tsid;
3131
3132 let all_table_tags = self.ctx.tag_columns.clone();
3133
3134 let scan_tag_columns = if use_tsid {
3135 let mut scan_tags = self.ctx.tag_columns.clone();
3136 for matcher in &self.ctx.selector_matcher {
3137 if is_metric_engine_internal_column(&matcher.name) {
3138 continue;
3139 }
3140 if all_table_tags.iter().any(|tag| tag == &matcher.name) {
3141 scan_tags.push(matcher.name.clone());
3142 }
3143 }
3144 scan_tags.sort_unstable();
3145 scan_tags.dedup();
3146 scan_tags
3147 } else {
3148 self.ctx.tag_columns.clone()
3149 };
3150
3151 let time_index_data_type = scan_table
3152 .schema()
3153 .timestamp_column()
3154 .with_context(|| TimeIndexNotFoundSnafu {
3155 table: maybe_phy_table_ref.to_quoted_string(),
3156 })?
3157 .data_type
3158 .clone();
3159 let is_time_index_second =
3160 time_index_data_type == ConcreteDataType::timestamp_second_datatype();
3161
3162 let scan_projection = if table_id_filter.is_some() {
3163 let mut required_columns = HashSet::new();
3164 required_columns.insert(DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string());
3165 required_columns.insert(self.ctx.time_index_column.clone().with_context(|| {
3166 TimeIndexNotFoundSnafu {
3167 table: maybe_phy_table_ref.to_quoted_string(),
3168 }
3169 })?);
3170 for col in &scan_tag_columns {
3171 required_columns.insert(col.clone());
3172 }
3173 for col in &self.ctx.field_columns {
3174 required_columns.insert(col.clone());
3175 }
3176 if use_tsid {
3177 required_columns.insert(DATA_SCHEMA_TSID_COLUMN_NAME.to_string());
3178 }
3179
3180 let arrow_schema = scan_table.schema().arrow_schema().clone();
3181 Some(
3182 arrow_schema
3183 .fields()
3184 .iter()
3185 .enumerate()
3186 .filter(|(_, field)| required_columns.contains(field.name().as_str()))
3187 .map(|(idx, _)| idx)
3188 .collect::<Vec<_>>(),
3189 )
3190 } else {
3191 None
3192 };
3193
3194 let mut scan_plan =
3195 LogicalPlanBuilder::scan(maybe_phy_table_ref.clone(), scan_provider, scan_projection)
3196 .context(DataFusionPlanningSnafu)?
3197 .build()
3198 .context(DataFusionPlanningSnafu)?;
3199
3200 if let Some(table_id) = table_id_filter {
3201 scan_plan = LogicalPlanBuilder::from(scan_plan)
3202 .filter(
3203 DfExpr::Column(Column::from_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME))
3204 .eq(lit(table_id)),
3205 )
3206 .context(DataFusionPlanningSnafu)?
3207 .alias(table_ref.clone()) .context(DataFusionPlanningSnafu)?
3209 .build()
3210 .context(DataFusionPlanningSnafu)?;
3211 }
3212
3213 if is_time_index_second {
3214 let expr: Vec<_> = self
3219 .create_field_column_exprs()?
3220 .into_iter()
3221 .chain(
3222 scan_tag_columns
3223 .iter()
3224 .map(|tag| DfExpr::Column(Column::from_name(tag))),
3225 )
3226 .chain(self.ctx.use_tsid.then_some(DfExpr::Column(Column::new(
3227 Some(table_ref.clone()),
3228 DATA_SCHEMA_TSID_COLUMN_NAME.to_string(),
3229 ))))
3230 .chain(Some(DfExpr::Alias(Alias {
3231 expr: Box::new(DfExpr::Cast(Cast::new(
3232 Box::new(self.create_time_index_column_expr()?),
3233 ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
3234 ))),
3235 relation: Some(table_ref.clone()),
3236 name: self
3237 .ctx
3238 .time_index_column
3239 .as_ref()
3240 .with_context(|| TimeIndexNotFoundSnafu {
3241 table: table_ref.to_quoted_string(),
3242 })?
3243 .clone(),
3244 metadata: None,
3245 })))
3246 .collect::<Vec<_>>();
3247 scan_plan = LogicalPlanBuilder::from(scan_plan)
3248 .project(expr)
3249 .context(DataFusionPlanningSnafu)?
3250 .build()
3251 .context(DataFusionPlanningSnafu)?;
3252 } else if table_id_filter.is_some()
3253 || time_index_data_type == ConcreteDataType::timestamp_microsecond_datatype()
3254 || time_index_data_type == ConcreteDataType::timestamp_nanosecond_datatype()
3255 {
3256 let project_exprs = self
3262 .create_field_column_exprs()?
3263 .into_iter()
3264 .chain(
3265 scan_tag_columns
3266 .iter()
3267 .map(|tag| DfExpr::Column(Column::from_name(tag))),
3268 )
3269 .chain(
3270 self.ctx
3271 .use_tsid
3272 .then_some(DfExpr::Column(Column::from_name(
3273 DATA_SCHEMA_TSID_COLUMN_NAME,
3274 ))),
3275 )
3276 .chain(Some(self.create_time_index_column_expr()?))
3277 .collect::<Vec<_>>();
3278
3279 scan_plan = LogicalPlanBuilder::from(scan_plan)
3280 .project(project_exprs)
3281 .context(DataFusionPlanningSnafu)?
3282 .build()
3283 .context(DataFusionPlanningSnafu)?;
3284 }
3285
3286 let result = LogicalPlanBuilder::from(scan_plan)
3287 .build()
3288 .context(DataFusionPlanningSnafu)?;
3289 Ok(result)
3290 }
3291
3292 fn collect_row_key_tag_columns_from_plan(
3293 &self,
3294 plan: &LogicalPlan,
3295 ) -> Result<BTreeSet<String>> {
3296 fn walk(
3297 planner: &PromPlanner,
3298 plan: &LogicalPlan,
3299 out: &mut BTreeSet<String>,
3300 ) -> Result<()> {
3301 if let LogicalPlan::TableScan(scan) = plan
3303 && let Ok(table) = planner.table_from_source(&scan.source)
3304 {
3305 for col in table.table_info().meta.row_key_column_names() {
3306 if col != DATA_SCHEMA_TABLE_ID_COLUMN_NAME
3307 && col != DATA_SCHEMA_TSID_COLUMN_NAME
3308 && !is_metric_engine_internal_column(col)
3309 {
3310 out.insert(col.clone());
3311 }
3312 }
3313 }
3314
3315 for input in plan.inputs() {
3316 walk(planner, input, out)?;
3317 }
3318 Ok(())
3319 }
3320
3321 let mut out = BTreeSet::new();
3322 walk(self, plan, &mut out)?;
3323 Ok(out)
3324 }
3325
3326 async fn setup_context(&mut self) -> Result<Option<LogicalPlan>> {
3330 let table_ref = self.table_ref()?;
3331 let source = match self.table_provider.resolve_table(table_ref.clone()).await {
3332 Err(e) if e.status_code() == StatusCode::TableNotFound => {
3333 let plan = self.setup_context_for_empty_metric()?;
3334 return Ok(Some(plan));
3335 }
3336 res => res.context(CatalogSnafu)?,
3337 };
3338 let table = self.table_from_source(&source)?;
3339
3340 let time_index = table
3342 .schema()
3343 .timestamp_column()
3344 .with_context(|| TimeIndexNotFoundSnafu {
3345 table: table_ref.to_quoted_string(),
3346 })?
3347 .name
3348 .clone();
3349 self.ctx.time_index_column = Some(time_index);
3350
3351 let values = table
3353 .table_info()
3354 .meta
3355 .field_column_names()
3356 .cloned()
3357 .collect();
3358 self.ctx.field_columns = values;
3359
3360 let tags = table
3362 .table_info()
3363 .meta
3364 .row_key_column_names()
3365 .filter(|col| {
3366 col != &DATA_SCHEMA_TABLE_ID_COLUMN_NAME && col != &DATA_SCHEMA_TSID_COLUMN_NAME
3368 })
3369 .cloned()
3370 .collect();
3371 self.ctx.tag_columns = tags;
3372 self.ctx.aggregation_field_labels.clear();
3374
3375 self.ctx.use_tsid = false;
3376
3377 Ok(None)
3378 }
3379
3380 fn setup_context_for_empty_metric(&mut self) -> Result<LogicalPlan> {
3383 self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
3384 self.ctx.reset_table_name_and_schema();
3385 self.ctx.tag_columns = vec![];
3386 self.ctx.aggregation_field_labels.clear();
3387 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
3388 self.ctx.use_tsid = false;
3389
3390 let plan = LogicalPlan::Extension(Extension {
3392 node: Arc::new(
3393 EmptyMetric::new(
3394 0,
3395 -1,
3396 self.ctx.interval,
3397 SPECIAL_TIME_FUNCTION.to_string(),
3398 DEFAULT_FIELD_COLUMN.to_string(),
3399 Some(lit(0.0f64)),
3400 )
3401 .context(DataFusionPlanningSnafu)?,
3402 ),
3403 });
3404 Ok(plan)
3405 }
3406
3407 fn create_function_args(&self, args: &[Box<PromExpr>]) -> Result<FunctionArgs> {
3409 let mut result = FunctionArgs::default();
3410
3411 for arg in args {
3412 if let Some(expr) = Self::try_build_literal_expr(arg) {
3414 result.literals.push(expr);
3415 } else {
3416 match arg.as_ref() {
3418 PromExpr::Subquery(_)
3419 | PromExpr::VectorSelector(_)
3420 | PromExpr::MatrixSelector(_)
3421 | PromExpr::Extension(_)
3422 | PromExpr::Aggregate(_)
3423 | PromExpr::Paren(_)
3424 | PromExpr::Call(_)
3425 | PromExpr::Binary(_)
3426 | PromExpr::Unary(_) => {
3427 if result.input.replace(*arg.clone()).is_some() {
3428 MultipleVectorSnafu { expr: *arg.clone() }.fail()?;
3429 }
3430 }
3431
3432 _ => {
3433 let expr = Self::get_param_as_literal_expr(Some(arg.as_ref()), None, None)?;
3434 result.literals.push(expr);
3435 }
3436 }
3437 }
3438 }
3439
3440 Ok(result)
3441 }
3442
3443 fn create_mixed_range_function_exprs(
3444 &mut self,
3445 func: &Function,
3446 mut other_input_exprs: VecDeque<DfExpr>,
3447 float_field: &str,
3448 histogram_field: &str,
3449 input_schema: &DFSchemaRef,
3450 ) -> Result<Option<Vec<DfExpr>>> {
3451 let returns_histogram = matches!(
3452 func.name,
3453 "rate"
3454 | "increase"
3455 | "delta"
3456 | "idelta"
3457 | "irate"
3458 | "avg_over_time"
3459 | "sum_over_time"
3460 | "last_over_time"
3461 );
3462 if !returns_histogram
3463 && !matches!(
3464 func.name,
3465 "changes"
3466 | "resets"
3467 | "deriv"
3468 | "min_over_time"
3469 | "max_over_time"
3470 | "count_over_time"
3471 | "absent_over_time"
3472 | "present_over_time"
3473 | "stddev_over_time"
3474 | "stdvar_over_time"
3475 | "quantile_over_time"
3476 | "predict_linear"
3477 | "double_exponential_smoothing"
3478 | "holt_winters"
3479 )
3480 {
3481 return Ok(None);
3482 }
3483
3484 if func.name == "predict_linear" {
3485 other_input_exprs[0] = DfExpr::Cast(Cast::new(
3486 Box::new(other_input_exprs[0].clone()),
3487 ArrowDataType::Int64,
3488 ));
3489 }
3490
3491 let timestamp_range = DfExpr::Column(Column::from_name(
3492 RangeManipulate::build_timestamp_range_name(
3493 self.ctx.time_index_column.as_ref().unwrap(),
3494 ),
3495 ));
3496 let float_range = DfExpr::Column(Column::from_name(float_field));
3497 let histogram_range = DfExpr::Column(Column::from_name(histogram_field));
3498 let mut args = Vec::with_capacity(other_input_exprs.len() + 6);
3499 args.push(lit(func.name));
3500 args.push(timestamp_range.clone());
3501 args.push(float_range.clone());
3502 args.push(histogram_range.clone());
3503 args.extend(other_input_exprs);
3504 if matches!(func.name, "rate" | "increase" | "delta") {
3505 args.push(self.create_time_index_column_expr()?);
3506 args.push(lit(self.ctx.range.context(ExpectRangeSelectorSnafu)?));
3507 }
3508
3509 let mut float_expr = DfExpr::ScalarFunction(ScalarFunction {
3510 func: Arc::new(MixedRange::float_udf(self.promql_annotations.clone())),
3511 args: args.clone(),
3512 });
3513 if matches!(func.name, "rate" | "increase") {
3514 let raw_delta_function = if func.name == "rate" {
3515 "raw_delta_rate"
3516 } else {
3517 "raw_delta_increase"
3518 };
3519 let delta_sum = DfExpr::ScalarFunction(ScalarFunction {
3520 func: Arc::new(MixedRange::float_udf(self.promql_annotations.clone())),
3521 args: vec![
3522 lit(raw_delta_function),
3523 timestamp_range,
3524 float_range,
3525 histogram_range,
3526 ],
3527 });
3528 float_expr = self.select_delta_range_math(
3529 func.name,
3530 input_schema,
3531 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3532 delta_sum,
3533 float_expr,
3534 )?;
3535 }
3536 let exprs = if returns_histogram {
3537 self.ctx.field_columns = vec![float_field.to_string(), histogram_field.to_string()];
3538 vec![
3539 float_expr.alias(float_field),
3540 DfExpr::ScalarFunction(ScalarFunction {
3541 func: Arc::new(MixedRange::histogram_udf(self.promql_annotations.clone())),
3542 args,
3543 })
3544 .alias(histogram_field),
3545 ]
3546 } else {
3547 let display_name = float_expr.schema_name().to_string();
3548 self.ctx.field_columns = vec![display_name.clone()];
3549 vec![float_expr.alias(display_name)]
3550 };
3551 Ok(Some(exprs))
3552 }
3553
3554 fn create_function_expr(
3560 &mut self,
3561 func: &Function,
3562 other_input_exprs: Vec<DfExpr>,
3563 input_schema: &DFSchemaRef,
3564 query_engine_state: &QueryEngineState,
3565 ) -> Result<(Vec<DfExpr>, Vec<String>)> {
3566 let mut other_input_exprs: VecDeque<DfExpr> = other_input_exprs.into();
3568 if let Some((float_field, histogram_field)) =
3569 Self::alternative_sample_range_columns(input_schema, &self.ctx.field_columns)
3570 .map(|(float, histogram)| (float.to_string(), histogram.to_string()))
3571 && let Some(exprs) = self.create_mixed_range_function_exprs(
3572 func,
3573 other_input_exprs.clone(),
3574 &float_field,
3575 &histogram_field,
3576 input_schema,
3577 )?
3578 {
3579 return Ok((exprs, vec![]));
3580 }
3581 let alternative_samples =
3582 Self::field_columns_are_alternative_samples(input_schema, &self.ctx.field_columns);
3583 let all_field_columns_are_native_histogram_ranges =
3584 self.all_field_columns_are_native_histogram_ranges(input_schema);
3585
3586 let field_column_pos = 0;
3588 let mut exprs = Vec::with_capacity(self.ctx.field_columns.len());
3589 let mut new_tags = vec![];
3591 let promql_annotations = self.promql_annotations.clone();
3592 let native_histogram_drop_udf = |name: &str| {
3593 Arc::new(NativeHistogramDrop::float_null_udf(
3594 format!(
3595 "{name}: dropped native histogram samples because this function is not supported for native histograms"
3596 ),
3597 promql_annotations.clone(),
3598 ))
3599 };
3600 let scalar_func = match func.name {
3601 "increase" => {
3602 if all_field_columns_are_native_histogram_ranges {
3603 ScalarFunc::ExtrapolateUdf(
3604 Arc::new(NativeHistogramIncrease::scalar_udf_with_collector(
3605 self.promql_annotations.clone(),
3606 )),
3607 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3608 )
3609 } else {
3610 ScalarFunc::ExtrapolateUdf(
3611 Arc::new(Increase::scalar_udf()),
3612 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3613 )
3614 }
3615 }
3616 "rate" => {
3617 if all_field_columns_are_native_histogram_ranges {
3618 ScalarFunc::ExtrapolateUdf(
3619 Arc::new(NativeHistogramRate::scalar_udf_with_collector(
3620 self.promql_annotations.clone(),
3621 )),
3622 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3623 )
3624 } else {
3625 ScalarFunc::ExtrapolateUdf(
3626 Arc::new(Rate::scalar_udf()),
3627 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3628 )
3629 }
3630 }
3631 "delta" => {
3632 if all_field_columns_are_native_histogram_ranges {
3633 ScalarFunc::ExtrapolateUdf(
3634 Arc::new(NativeHistogramDelta::scalar_udf_with_collector(
3635 self.promql_annotations.clone(),
3636 )),
3637 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3638 )
3639 } else {
3640 ScalarFunc::ExtrapolateUdf(
3641 Arc::new(Delta::scalar_udf()),
3642 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3643 )
3644 }
3645 }
3646 "idelta" => {
3647 if all_field_columns_are_native_histogram_ranges {
3648 ScalarFunc::Udf(Arc::new(NativeHistogramIDelta::scalar_udf_with_collector(
3649 self.promql_annotations.clone(),
3650 )))
3651 } else {
3652 ScalarFunc::Udf(Arc::new(IDelta::<false>::scalar_udf()))
3653 }
3654 }
3655 "irate" => {
3656 if all_field_columns_are_native_histogram_ranges {
3657 ScalarFunc::Udf(Arc::new(NativeHistogramIRate::scalar_udf_with_collector(
3658 self.promql_annotations.clone(),
3659 )))
3660 } else {
3661 ScalarFunc::Udf(Arc::new(IDelta::<true>::scalar_udf()))
3662 }
3663 }
3664 "resets" => {
3665 if all_field_columns_are_native_histogram_ranges {
3666 ScalarFunc::Udf(Arc::new(NativeHistogramResets::scalar_udf()))
3667 } else {
3668 ScalarFunc::Udf(Arc::new(Resets::scalar_udf()))
3669 }
3670 }
3671 "changes" => {
3672 if all_field_columns_are_native_histogram_ranges {
3673 ScalarFunc::Udf(Arc::new(NativeHistogramChanges::scalar_udf()))
3674 } else {
3675 ScalarFunc::Udf(Arc::new(Changes::scalar_udf()))
3676 }
3677 }
3678 "deriv" => {
3679 if all_field_columns_are_native_histogram_ranges {
3680 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3681 } else {
3682 ScalarFunc::Udf(Arc::new(Deriv::scalar_udf()))
3683 }
3684 }
3685 "avg_over_time" => {
3686 if all_field_columns_are_native_histogram_ranges {
3687 ScalarFunc::Udf(Arc::new(
3688 NativeHistogramAvgOverTime::scalar_udf_with_collector(
3689 self.promql_annotations.clone(),
3690 ),
3691 ))
3692 } else {
3693 ScalarFunc::Udf(Arc::new(AvgOverTime::scalar_udf()))
3694 }
3695 }
3696 "min_over_time" => {
3697 if all_field_columns_are_native_histogram_ranges {
3698 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3699 } else {
3700 ScalarFunc::Udf(Arc::new(MinOverTime::scalar_udf()))
3701 }
3702 }
3703 "max_over_time" => {
3704 if all_field_columns_are_native_histogram_ranges {
3705 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3706 } else {
3707 ScalarFunc::Udf(Arc::new(MaxOverTime::scalar_udf()))
3708 }
3709 }
3710 "sum_over_time" => {
3711 if all_field_columns_are_native_histogram_ranges {
3712 ScalarFunc::Udf(Arc::new(
3713 NativeHistogramSumOverTime::scalar_udf_with_collector(
3714 self.promql_annotations.clone(),
3715 ),
3716 ))
3717 } else {
3718 ScalarFunc::Udf(Arc::new(SumOverTime::scalar_udf()))
3719 }
3720 }
3721 "count_over_time" => {
3722 if all_field_columns_are_native_histogram_ranges {
3723 ScalarFunc::Udf(Arc::new(NativeHistogramCountOverTime::scalar_udf()))
3724 } else {
3725 ScalarFunc::Udf(Arc::new(CountOverTime::scalar_udf()))
3726 }
3727 }
3728 "last_over_time" => {
3729 if all_field_columns_are_native_histogram_ranges {
3730 ScalarFunc::Udf(Arc::new(NativeHistogramLastOverTime::scalar_udf()))
3731 } else {
3732 ScalarFunc::Udf(Arc::new(LastOverTime::scalar_udf()))
3733 }
3734 }
3735 "absent_over_time" => {
3736 if all_field_columns_are_native_histogram_ranges {
3737 ScalarFunc::Udf(Arc::new(NativeHistogramAbsentOverTime::scalar_udf()))
3738 } else {
3739 ScalarFunc::Udf(Arc::new(AbsentOverTime::scalar_udf()))
3740 }
3741 }
3742 "present_over_time" => {
3743 if all_field_columns_are_native_histogram_ranges {
3744 ScalarFunc::Udf(Arc::new(NativeHistogramPresentOverTime::scalar_udf()))
3745 } else {
3746 ScalarFunc::Udf(Arc::new(PresentOverTime::scalar_udf()))
3747 }
3748 }
3749 "stddev_over_time" => {
3750 if all_field_columns_are_native_histogram_ranges {
3751 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3752 } else {
3753 ScalarFunc::Udf(Arc::new(StddevOverTime::scalar_udf()))
3754 }
3755 }
3756 "stdvar_over_time" => {
3757 if all_field_columns_are_native_histogram_ranges {
3758 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3759 } else {
3760 ScalarFunc::Udf(Arc::new(StdvarOverTime::scalar_udf()))
3761 }
3762 }
3763 "quantile_over_time" => {
3764 if all_field_columns_are_native_histogram_ranges {
3765 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3766 } else {
3767 ScalarFunc::Udf(Arc::new(QuantileOverTime::scalar_udf()))
3768 }
3769 }
3770 "predict_linear" => {
3771 if all_field_columns_are_native_histogram_ranges {
3772 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3773 } else {
3774 other_input_exprs[0] = DfExpr::Cast(Cast::new(
3775 Box::new(other_input_exprs[0].clone()),
3776 ArrowDataType::Int64,
3777 ));
3778 ScalarFunc::Udf(Arc::new(PredictLinear::scalar_udf()))
3779 }
3780 }
3781 "double_exponential_smoothing" | "holt_winters" => {
3782 if all_field_columns_are_native_histogram_ranges {
3783 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3784 } else {
3785 ScalarFunc::Udf(Arc::new(DoubleExponentialSmoothing::scalar_udf()))
3786 }
3787 }
3788 "histogram_count" => {
3789 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramCount::scalar_udf()))
3790 }
3791 "histogram_sum" => {
3792 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramSum::scalar_udf()))
3793 }
3794 "histogram_avg" => {
3795 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramAvg::scalar_udf()))
3796 }
3797 "histogram_stddev" => {
3798 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramStddev::scalar_udf()))
3799 }
3800 "histogram_stdvar" => {
3801 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramStdvar::scalar_udf()))
3802 }
3803 "time" => {
3804 exprs.push(build_special_time_expr(
3805 self.ctx.time_index_column.as_ref().unwrap(),
3806 ));
3807 ScalarFunc::GeneratedExpr
3808 }
3809 "minute" => {
3810 let expr = self.date_part_on_time_index("minute")?;
3812 exprs.push(expr);
3813 ScalarFunc::GeneratedExpr
3814 }
3815 "hour" => {
3816 let expr = self.date_part_on_time_index("hour")?;
3818 exprs.push(expr);
3819 ScalarFunc::GeneratedExpr
3820 }
3821 "month" => {
3822 let expr = self.date_part_on_time_index("month")?;
3824 exprs.push(expr);
3825 ScalarFunc::GeneratedExpr
3826 }
3827 "year" => {
3828 let expr = self.date_part_on_time_index("year")?;
3830 exprs.push(expr);
3831 ScalarFunc::GeneratedExpr
3832 }
3833 "day_of_month" => {
3834 let expr = self.date_part_on_time_index("day")?;
3836 exprs.push(expr);
3837 ScalarFunc::GeneratedExpr
3838 }
3839 "day_of_week" => {
3840 let expr = self.date_part_on_time_index("dow")?;
3842 exprs.push(expr);
3843 ScalarFunc::GeneratedExpr
3844 }
3845 "day_of_year" => {
3846 let expr = self.date_part_on_time_index("doy")?;
3848 exprs.push(expr);
3849 ScalarFunc::GeneratedExpr
3850 }
3851 "days_in_month" => {
3852 let day_lit_expr = "day".lit();
3857 let month_lit_expr = "month".lit();
3858 let interval_1month_lit_expr =
3859 DfExpr::Literal(ScalarValue::IntervalYearMonth(Some(1)), None);
3860 let interval_1day_lit_expr = DfExpr::Literal(
3861 ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(1, 0))),
3862 None,
3863 );
3864 let the_1month_minus_1day_expr = DfExpr::BinaryExpr(BinaryExpr {
3865 left: Box::new(interval_1month_lit_expr),
3866 op: Operator::Minus,
3867 right: Box::new(interval_1day_lit_expr),
3868 });
3869 let date_trunc_expr = DfExpr::ScalarFunction(ScalarFunction {
3870 func: datafusion_functions::datetime::date_trunc(),
3871 args: vec![month_lit_expr, self.create_time_index_column_expr()?],
3872 });
3873 let date_trunc_plus_interval_expr = DfExpr::BinaryExpr(BinaryExpr {
3874 left: Box::new(date_trunc_expr),
3875 op: Operator::Plus,
3876 right: Box::new(the_1month_minus_1day_expr),
3877 });
3878 let date_part_expr = DfExpr::ScalarFunction(ScalarFunction {
3879 func: datafusion_functions::datetime::date_part(),
3880 args: vec![day_lit_expr, date_trunc_plus_interval_expr],
3881 });
3882
3883 exprs.push(date_part_expr);
3884 ScalarFunc::GeneratedExpr
3885 }
3886
3887 "label_join" => {
3888 self.ctx.use_tsid = false;
3889 let (concat_expr, dst_label) = Self::build_concat_labels_expr(
3890 &mut other_input_exprs,
3891 &self.ctx,
3892 query_engine_state,
3893 )?;
3894
3895 for value in &self.ctx.field_columns {
3897 if *value != dst_label {
3898 let expr = DfExpr::Column(Column::from_name(value));
3899 exprs.push(expr);
3900 }
3901 }
3902
3903 self.ctx.tag_columns.retain(|tag| *tag != dst_label);
3905 new_tags.push(dst_label);
3906 exprs.push(concat_expr);
3908
3909 ScalarFunc::GeneratedExpr
3910 }
3911 "label_replace" => {
3912 self.ctx.use_tsid = false;
3913 if let Some((replace_expr, dst_label)) = self
3914 .build_regexp_replace_label_expr(&mut other_input_exprs, query_engine_state)?
3915 {
3916 for value in &self.ctx.field_columns {
3918 if *value != dst_label {
3919 let expr = DfExpr::Column(Column::from_name(value));
3920 exprs.push(expr);
3921 }
3922 }
3923
3924 ensure!(
3925 !self.ctx.tag_columns.contains(&dst_label),
3926 SameLabelSetSnafu
3927 );
3928 new_tags.push(dst_label);
3929 exprs.push(replace_expr);
3931 } else {
3932 for value in &self.ctx.field_columns {
3934 let expr = DfExpr::Column(Column::from_name(value));
3935 exprs.push(expr);
3936 }
3937 }
3938
3939 ScalarFunc::GeneratedExpr
3940 }
3941 "sort" | "sort_desc" => {
3942 for value in &self.ctx.field_columns {
3944 if !Self::field_column_is_native_histogram(input_schema, value) {
3945 exprs.push(DfExpr::Column(Column::from_name(value)));
3946 }
3947 }
3948 if exprs.is_empty() {
3951 exprs.push(DfExpr::Literal(ScalarValue::Float64(None), None));
3952 }
3953
3954 ScalarFunc::GeneratedExpr
3955 }
3956 "sort_by_label" | "sort_by_label_desc" | "timestamp" => {
3957 for value in &self.ctx.field_columns {
3960 let expr = DfExpr::Column(Column::from_name(value));
3961 exprs.push(expr);
3962 }
3963
3964 ScalarFunc::GeneratedExpr
3965 }
3966 "round" if self.all_field_columns_are_native_histograms(input_schema) => {
3967 if other_input_exprs.is_empty() {
3968 other_input_exprs.push_front(0.0f64.lit());
3969 }
3970 ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3971 }
3972 "round" => {
3973 if other_input_exprs.is_empty() {
3974 other_input_exprs.push_front(0.0f64.lit());
3975 }
3976 ScalarFunc::DataFusionUdf(Arc::new(Round::scalar_udf()))
3977 }
3978 "rad" | "deg" | "sgn" if self.all_field_columns_are_native_histograms(input_schema) => {
3979 ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3980 }
3981 "rad" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::radians()),
3982 "deg" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::degrees()),
3983 "sgn" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::signum()),
3984 "pi" => {
3985 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
3987 func: datafusion::functions::math::pi(),
3988 args: vec![],
3989 });
3990 exprs.push(fn_expr);
3991
3992 ScalarFunc::GeneratedExpr
3993 }
3994 _ => {
3995 if let Some(f) = query_engine_state
3996 .session_state()
3997 .scalar_functions()
3998 .get(func.name)
3999 {
4000 if self.all_field_columns_are_native_histograms(input_schema) {
4001 ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
4002 } else {
4003 ScalarFunc::DataFusionBuiltin(f.clone())
4004 }
4005 } else if let Some(factory) = query_engine_state.scalar_function(func.name) {
4006 if self.all_field_columns_are_native_histograms(input_schema) {
4007 ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
4008 } else {
4009 let func_state = query_engine_state.function_state();
4010 let query_ctx = self.table_provider.query_ctx();
4011
4012 ScalarFunc::DataFusionUdf(Arc::new(factory.provide(FunctionContext {
4013 state: func_state,
4014 query_ctx: query_ctx.clone(),
4015 })))
4016 }
4017 } else if let Some(f) = datafusion_functions::math::functions()
4018 .iter()
4019 .find(|f| f.name() == func.name)
4020 {
4021 if self.all_field_columns_are_native_histograms(input_schema) {
4022 ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
4023 } else {
4024 ScalarFunc::DataFusionUdf(f.clone())
4025 }
4026 } else {
4027 return UnsupportedExprSnafu {
4028 name: func.name.to_string(),
4029 }
4030 .fail();
4031 }
4032 }
4033 };
4034
4035 for value in &self.ctx.field_columns {
4036 let col_expr = DfExpr::Column(Column::from_name(value));
4037 let value_is_histogram = Self::field_column_is_native_histogram(input_schema, value);
4038
4039 match scalar_func.clone() {
4040 ScalarFunc::DataFusionBuiltin(func) => {
4041 if alternative_samples && value_is_histogram {
4042 continue;
4043 }
4044 other_input_exprs.insert(field_column_pos, col_expr);
4045 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
4046 func,
4047 args: other_input_exprs.clone().into(),
4048 });
4049 exprs.push(fn_expr);
4050 let _ = other_input_exprs.remove(field_column_pos);
4051 }
4052 ScalarFunc::DataFusionUdf(func) => {
4053 if alternative_samples && value_is_histogram {
4054 continue;
4055 }
4056 let args = itertools::chain!(
4057 other_input_exprs.iter().take(field_column_pos).cloned(),
4058 std::iter::once(col_expr),
4059 other_input_exprs.iter().skip(field_column_pos).cloned()
4060 )
4061 .collect_vec();
4062 exprs.push(DfExpr::ScalarFunction(ScalarFunction { func, args }))
4063 }
4064 ScalarFunc::NativeHistogramUdf(func) => {
4065 if value_is_histogram {
4066 let args = itertools::chain!(
4067 other_input_exprs.iter().take(field_column_pos).cloned(),
4068 std::iter::once(col_expr),
4069 other_input_exprs.iter().skip(field_column_pos).cloned()
4070 )
4071 .collect_vec();
4072 exprs.push(DfExpr::ScalarFunction(ScalarFunction { func, args }));
4073 } else if !alternative_samples {
4074 exprs.push(
4075 DfExpr::Literal(ScalarValue::Float64(None), None).alias(format!(
4076 "{}_{}",
4077 func.name(),
4078 value
4079 )),
4080 );
4081 }
4082 }
4083 ScalarFunc::Udf(func) => {
4084 let ts_range_expr = DfExpr::Column(Column::from_name(
4085 RangeManipulate::build_timestamp_range_name(
4086 self.ctx.time_index_column.as_ref().unwrap(),
4087 ),
4088 ));
4089 other_input_exprs.insert(field_column_pos, ts_range_expr);
4090 other_input_exprs.insert(field_column_pos + 1, col_expr);
4091 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
4092 func,
4093 args: other_input_exprs.clone().into(),
4094 });
4095 exprs.push(fn_expr);
4096 let _ = other_input_exprs.remove(field_column_pos + 1);
4097 let _ = other_input_exprs.remove(field_column_pos);
4098 }
4099 ScalarFunc::ExtrapolateUdf(udf, range_length) => {
4100 let ts_range_expr = DfExpr::Column(Column::from_name(
4101 RangeManipulate::build_timestamp_range_name(
4102 self.ctx.time_index_column.as_ref().unwrap(),
4103 ),
4104 ));
4105 other_input_exprs.insert(field_column_pos, ts_range_expr.clone());
4106 other_input_exprs.insert(field_column_pos + 1, col_expr.clone());
4107 other_input_exprs
4108 .insert(field_column_pos + 2, self.create_time_index_column_expr()?);
4109 other_input_exprs.push_back(lit(range_length));
4110 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
4111 func: udf,
4112 args: other_input_exprs.clone().into(),
4113 });
4114 let fn_expr = if matches!(func.name, "rate" | "increase")
4115 && !all_field_columns_are_native_histogram_ranges
4116 {
4117 let delta_sum = DfExpr::ScalarFunction(ScalarFunction {
4118 func: Arc::new(SumOverTime::scalar_udf()),
4119 args: vec![ts_range_expr, col_expr],
4120 });
4121 self.select_delta_range_math(
4122 func.name,
4123 input_schema,
4124 range_length,
4125 delta_sum,
4126 fn_expr,
4127 )?
4128 } else {
4129 fn_expr
4130 };
4131 exprs.push(fn_expr);
4132 let _ = other_input_exprs.pop_back();
4133 let _ = other_input_exprs.remove(field_column_pos + 2);
4134 let _ = other_input_exprs.remove(field_column_pos + 1);
4135 let _ = other_input_exprs.remove(field_column_pos);
4136 }
4137 ScalarFunc::GeneratedExpr => {}
4138 }
4139 }
4140
4141 if !matches!(func.name, "label_join" | "label_replace") {
4145 let mut new_field_columns = Vec::with_capacity(exprs.len());
4146
4147 exprs = exprs
4148 .into_iter()
4149 .map(|expr| {
4150 let display_name = expr.schema_name().to_string();
4151 new_field_columns.push(display_name.clone());
4152 Ok(expr.alias(display_name))
4153 })
4154 .collect::<std::result::Result<Vec<_>, _>>()
4155 .context(DataFusionPlanningSnafu)?;
4156
4157 self.ctx.field_columns = new_field_columns;
4158 }
4159
4160 Ok((exprs, new_tags))
4161 }
4162
4163 fn select_delta_range_math(
4164 &self,
4165 function: &str,
4166 input_schema: &DFSchemaRef,
4167 range_length: Millisecond,
4168 delta_sum: DfExpr,
4169 cumulative: DfExpr,
4170 ) -> Result<DfExpr> {
4171 let marker_is_delta = if self
4172 .ctx
4173 .tag_columns
4174 .iter()
4175 .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL)
4176 {
4177 Self::field_column_type(input_schema, OTLP_AGGREGATION_TEMPORALITY_LABEL)
4178 .filter(|data_type| Self::string_value_data_type(data_type).is_some())
4179 .map(|_| {
4180 DfExpr::Column(Column::from_name(OTLP_AGGREGATION_TEMPORALITY_LABEL))
4181 .eq(lit(GREPTIME_TEMPORALITY_DELTA))
4182 })
4183 } else {
4184 None
4185 };
4186 let Some(marker_is_delta) = marker_is_delta else {
4187 return Ok(cumulative);
4188 };
4189
4190 let delta = if function == "rate" {
4191 DfExpr::BinaryExpr(BinaryExpr {
4192 left: Box::new(delta_sum),
4193 op: Operator::Divide,
4194 right: Box::new(lit(range_length as f64 / 1000.0)),
4195 })
4196 } else {
4197 delta_sum
4198 };
4199 let display_name = cumulative.schema_name().to_string();
4200 when(marker_is_delta, delta)
4201 .otherwise(cumulative)
4202 .context(DataFusionPlanningSnafu)
4203 .map(|expr| expr.alias(display_name))
4204 }
4205
4206 fn validate_label_name(label_name: &str) -> Result<()> {
4210 if label_name.starts_with("__") {
4212 return InvalidDestinationLabelNameSnafu { label_name }.fail();
4213 }
4214 if !LABEL_NAME_REGEX.is_match(label_name) {
4216 return InvalidDestinationLabelNameSnafu { label_name }.fail();
4217 }
4218
4219 Ok(())
4220 }
4221
4222 fn build_regexp_replace_label_expr(
4224 &self,
4225 other_input_exprs: &mut VecDeque<DfExpr>,
4226 query_engine_state: &QueryEngineState,
4227 ) -> Result<Option<(DfExpr, String)>> {
4228 let dst_label = match other_input_exprs.pop_front() {
4230 Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4231 other => UnexpectedPlanExprSnafu {
4232 desc: format!("expected dst_label string literal, but found {:?}", other),
4233 }
4234 .fail()?,
4235 };
4236
4237 Self::validate_label_name(&dst_label)?;
4239 let replacement = match other_input_exprs.pop_front() {
4240 Some(DfExpr::Literal(ScalarValue::Utf8(Some(r)), _)) => r,
4241 other => UnexpectedPlanExprSnafu {
4242 desc: format!("expected replacement string literal, but found {:?}", other),
4243 }
4244 .fail()?,
4245 };
4246 let src_label = match other_input_exprs.pop_front() {
4247 Some(DfExpr::Literal(ScalarValue::Utf8(Some(s)), None)) => s,
4248 other => UnexpectedPlanExprSnafu {
4249 desc: format!("expected src_label string literal, but found {:?}", other),
4250 }
4251 .fail()?,
4252 };
4253
4254 let regex = match other_input_exprs.pop_front() {
4255 Some(DfExpr::Literal(ScalarValue::Utf8(Some(r)), None)) => r,
4256 other => UnexpectedPlanExprSnafu {
4257 desc: format!("expected regex string literal, but found {:?}", other),
4258 }
4259 .fail()?,
4260 };
4261
4262 regex::Regex::new(®ex).map_err(|_| {
4265 InvalidRegularExpressionSnafu {
4266 regex: regex.clone(),
4267 }
4268 .build()
4269 })?;
4270
4271 if self.ctx.tag_columns.contains(&src_label) && regex.is_empty() {
4273 return Ok(None);
4274 }
4275
4276 if !self.ctx.tag_columns.contains(&src_label) {
4278 if replacement.is_empty() {
4279 return Ok(None);
4281 } else {
4282 return Ok(Some((
4284 lit(replacement).alias(&dst_label),
4286 dst_label,
4287 )));
4288 }
4289 }
4290
4291 let regex = format!("^(?s:{regex})$");
4294
4295 let session_state = query_engine_state.session_state();
4296 let func = session_state
4297 .scalar_functions()
4298 .get("regexp_replace")
4299 .context(UnsupportedExprSnafu {
4300 name: "regexp_replace",
4301 })?;
4302
4303 let args = vec![
4305 if src_label.is_empty() {
4306 DfExpr::Literal(ScalarValue::Utf8(Some(String::new())), None)
4307 } else {
4308 DfExpr::Column(Column::from_name(src_label))
4309 },
4310 DfExpr::Literal(ScalarValue::Utf8(Some(regex)), None),
4311 DfExpr::Literal(ScalarValue::Utf8(Some(replacement)), None),
4312 ];
4313
4314 Ok(Some((
4315 DfExpr::ScalarFunction(ScalarFunction {
4316 func: func.clone(),
4317 args,
4318 })
4319 .alias(&dst_label),
4320 dst_label,
4321 )))
4322 }
4323
4324 fn build_concat_labels_expr(
4326 other_input_exprs: &mut VecDeque<DfExpr>,
4327 ctx: &PromPlannerContext,
4328 query_engine_state: &QueryEngineState,
4329 ) -> Result<(DfExpr, String)> {
4330 let dst_label = match other_input_exprs.pop_front() {
4333 Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4334 other => UnexpectedPlanExprSnafu {
4335 desc: format!("expected dst_label string literal, but found {:?}", other),
4336 }
4337 .fail()?,
4338 };
4339 let separator = match other_input_exprs.pop_front() {
4340 Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4341 other => UnexpectedPlanExprSnafu {
4342 desc: format!("expected separator string literal, but found {:?}", other),
4343 }
4344 .fail()?,
4345 };
4346
4347 let available_columns: HashSet<&str> = ctx
4349 .tag_columns
4350 .iter()
4351 .chain(ctx.field_columns.iter())
4352 .chain(ctx.time_index_column.as_ref())
4353 .map(|s| s.as_str())
4354 .collect();
4355
4356 let src_labels = other_input_exprs
4357 .iter()
4358 .map(|expr| {
4359 match expr {
4361 DfExpr::Literal(ScalarValue::Utf8(Some(label)), None) => {
4362 if label.is_empty() {
4363 Ok(DfExpr::Literal(ScalarValue::Null, None))
4364 } else if available_columns.contains(label.as_str()) {
4365 Ok(DfExpr::Column(Column::from_name(label)))
4367 } else {
4368 Ok(DfExpr::Literal(ScalarValue::Null, None))
4370 }
4371 }
4372 other => UnexpectedPlanExprSnafu {
4373 desc: format!(
4374 "expected source label string literal, but found {:?}",
4375 other
4376 ),
4377 }
4378 .fail(),
4379 }
4380 })
4381 .collect::<Result<Vec<_>>>()?;
4382 ensure!(
4383 !src_labels.is_empty(),
4384 FunctionInvalidArgumentSnafu {
4385 fn_name: "label_join"
4386 }
4387 );
4388
4389 let session_state = query_engine_state.session_state();
4390 let func = session_state
4391 .scalar_functions()
4392 .get("concat_ws")
4393 .context(UnsupportedExprSnafu { name: "concat_ws" })?;
4394
4395 let mut args = Vec::with_capacity(1 + src_labels.len());
4397 args.push(DfExpr::Literal(ScalarValue::Utf8(Some(separator)), None));
4398 args.extend(src_labels);
4399
4400 Ok((
4401 DfExpr::ScalarFunction(ScalarFunction {
4402 func: func.clone(),
4403 args,
4404 })
4405 .alias(&dst_label),
4406 dst_label,
4407 ))
4408 }
4409
4410 fn create_time_index_column_expr(&self) -> Result<DfExpr> {
4411 Ok(DfExpr::Column(Column::from_name(
4412 self.ctx
4413 .time_index_column
4414 .clone()
4415 .with_context(|| TimeIndexNotFoundSnafu { table: "unknown" })?,
4416 )))
4417 }
4418
4419 fn create_tag_column_exprs(&self) -> Result<Vec<DfExpr>> {
4420 let mut result = Vec::with_capacity(self.ctx.tag_columns.len());
4421 for tag in &self.ctx.tag_columns {
4422 let expr = DfExpr::Column(Column::from_name(tag));
4423 result.push(expr);
4424 }
4425 Ok(result)
4426 }
4427
4428 fn create_field_column_exprs(&self) -> Result<Vec<DfExpr>> {
4429 let mut result = Vec::with_capacity(self.ctx.field_columns.len());
4430 for field in &self.ctx.field_columns {
4431 let expr = DfExpr::Column(Column::from_name(field));
4432 result.push(expr);
4433 }
4434 Ok(result)
4435 }
4436
4437 fn create_tag_and_time_index_column_sort_exprs(&self) -> Result<Vec<SortExpr>> {
4438 let mut result = self
4439 .ctx
4440 .tag_columns
4441 .iter()
4442 .map(|col| DfExpr::Column(Column::from_name(col)).sort(true, true))
4443 .collect::<Vec<_>>();
4444 result.push(self.create_time_index_column_expr()?.sort(true, true));
4445 Ok(result)
4446 }
4447
4448 fn create_field_columns_sort_exprs(&self, asc: bool) -> Vec<SortExpr> {
4449 self.ctx
4450 .field_columns
4451 .iter()
4452 .map(|col| DfExpr::Column(Column::from_name(col)).sort(asc, true))
4453 .collect::<Vec<_>>()
4454 }
4455
4456 fn create_sort_exprs_by_tags(
4457 func: &str,
4458 tags: Vec<DfExpr>,
4459 asc: bool,
4460 ) -> Result<Vec<SortExpr>> {
4461 ensure!(
4462 !tags.is_empty(),
4463 FunctionInvalidArgumentSnafu { fn_name: func }
4464 );
4465
4466 tags.iter()
4467 .map(|col| match col {
4468 DfExpr::Literal(ScalarValue::Utf8(Some(label)), _) => {
4469 Ok(DfExpr::Column(Column::from_name(label)).sort(asc, false))
4470 }
4471 other => UnexpectedPlanExprSnafu {
4472 desc: format!("expected label string literal, but found {:?}", other),
4473 }
4474 .fail(),
4475 })
4476 .collect::<Result<Vec<_>>>()
4477 }
4478
4479 fn create_empty_values_filter_expr(&self, preserve_any_value: bool) -> Result<DfExpr> {
4480 let mut exprs = Vec::with_capacity(self.ctx.field_columns.len());
4481 for value in &self.ctx.field_columns {
4482 let expr = DfExpr::Column(Column::from_name(value)).is_not_null();
4483 exprs.push(expr);
4484 }
4485
4486 let predicate = if preserve_any_value {
4491 disjunction(exprs)
4492 } else {
4493 conjunction(exprs)
4494 };
4495 predicate.with_context(|| ValueNotFoundSnafu {
4496 table: self
4497 .table_ref()
4498 .map(|t| t.to_quoted_string())
4499 .unwrap_or_else(|_| "unknown".to_string()),
4500 })
4501 }
4502
4503 fn create_aggregate_exprs(
4520 &mut self,
4521 op: TokenType,
4522 param: &Option<Box<PromExpr>>,
4523 input_plan: &LogicalPlan,
4524 ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4525 let mixed_sample_columns =
4526 Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns)
4527 .map(|(float, histogram)| (float.to_string(), histogram.to_string()));
4528 let is_group_agg = op.id() == token::T_GROUP;
4529 if is_group_agg && mixed_sample_columns.is_none() {
4530 ensure!(
4531 self.ctx.field_columns.len() == 1,
4532 MultiFieldsNotSupportedSnafu {
4533 operator: "group()"
4534 }
4535 );
4536 }
4537
4538 if let Some((float, histogram)) = mixed_sample_columns {
4539 return self.create_mixed_aggregate_exprs(op, param, &float, &histogram);
4540 }
4541
4542 if self.all_field_columns_are_native_histograms(input_plan.schema()) {
4543 return self.create_native_histogram_aggregate_exprs(op, input_plan);
4544 }
4545
4546 let exprs = self
4548 .ctx
4549 .field_columns
4550 .iter()
4551 .map(|col| {
4552 Self::create_numeric_aggregate_expr(
4553 op,
4554 param,
4555 DfExpr::Column(Column::from_name(col)),
4556 )
4557 })
4558 .collect::<Result<Vec<_>>>()?;
4559
4560 let prev_field_exprs = if op.id() == token::T_COUNT_VALUES {
4568 ensure!(
4569 self.ctx.field_columns.len() == 1,
4570 UnsupportedExprSnafu {
4571 name: "count_values on multi-value input"
4572 }
4573 );
4574
4575 self.ctx
4576 .field_columns
4577 .iter()
4578 .map(|col| {
4579 let value = DfExpr::Column(Column::from_name(col));
4580 let value = if Self::field_column_type(input_plan.schema(), col)
4585 == Some(&ArrowDataType::Float64)
4586 {
4587 value
4588 } else {
4589 DfExpr::Cast(Cast::new(Box::new(value), ArrowDataType::Float64))
4590 };
4591 DfExpr::ScalarFunction(ScalarFunction {
4592 func: Arc::new(PromqlFloatToString::scalar_udf()),
4593 args: vec![value],
4594 })
4595 })
4596 .collect()
4597 } else {
4598 vec![]
4599 };
4600
4601 let mut new_field_columns = Vec::with_capacity(self.ctx.field_columns.len());
4603
4604 let normalized_exprs =
4605 normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
4606 for expr in normalized_exprs {
4607 new_field_columns.push(expr.schema_name().to_string());
4608 }
4609 self.ctx.field_columns = new_field_columns;
4610
4611 Ok((exprs, prev_field_exprs))
4612 }
4613
4614 fn create_numeric_aggregate_expr(
4615 op: TokenType,
4616 param: &Option<Box<PromExpr>>,
4617 input: DfExpr,
4618 ) -> Result<DfExpr> {
4619 let expr = match op.id() {
4620 token::T_SUM => sum_udaf().call(vec![input]),
4621 token::T_QUANTILE => {
4622 let q = Self::get_param_as_literal_expr(
4623 param.as_deref(),
4624 Some(op),
4625 Some(ArrowDataType::Float64),
4626 )?;
4627 quantile_udaf().call(vec![q, input])
4628 }
4629 token::T_AVG => avg_udaf().call(vec![input]),
4630 token::T_COUNT_VALUES | token::T_COUNT => count_udaf().call(vec![input]),
4631 token::T_MIN => min_udaf().call(vec![input]),
4632 token::T_MAX => max_udaf().call(vec![input]),
4633 token::T_GROUP => max_udaf().call(vec![lit(1_f64)]),
4636 token::T_STDDEV => stddev_pop_udaf().call(vec![input]),
4637 token::T_STDVAR => var_pop_udaf().call(vec![input]),
4638 token::T_TOPK | token::T_BOTTOMK => {
4639 return UnsupportedExprSnafu {
4640 name: format!("{op:?}"),
4641 }
4642 .fail();
4643 }
4644 _ => return UnexpectedTokenSnafu { token: op }.fail(),
4645 };
4646 Ok(expr)
4647 }
4648
4649 fn create_mixed_aggregate_exprs(
4650 &mut self,
4651 op: TokenType,
4652 param: &Option<Box<PromExpr>>,
4653 float_column: &str,
4654 histogram_column: &str,
4655 ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4656 let float_input = DfExpr::Column(Column::from_name(float_column));
4657 let histogram_input = DfExpr::Column(Column::from_name(histogram_column));
4658 let float_count = count_udaf().call(vec![float_input.clone()]);
4659 let histogram_count = count_udaf().call(vec![histogram_input.clone()]);
4660 let mixed_sample_value = || {
4661 DfExpr::ScalarFunction(ScalarFunction {
4662 func: coalesce(),
4663 args: vec![
4664 DfExpr::ScalarFunction(ScalarFunction {
4665 func: Arc::new(PromqlFloatToString::scalar_udf()),
4666 args: vec![float_input.clone()],
4667 }),
4668 DfExpr::ScalarFunction(ScalarFunction {
4669 func: Arc::new(NativeHistogramToString::scalar_udf()),
4670 args: vec![histogram_input.clone()],
4671 }),
4672 ],
4673 })
4674 };
4675
4676 let (exprs, prev_field_exprs, field_columns) = match op.id() {
4677 token::T_SUM | token::T_AVG => (
4678 vec![
4679 Self::create_numeric_aggregate_expr(op, param, float_input)?
4680 .alias(float_column),
4681 self.create_native_histogram_aggregate_expr(op, histogram_column)?,
4682 float_count.alias(Self::mixed_sample_count_name(float_column)),
4683 histogram_count.alias(Self::mixed_sample_count_name(histogram_column)),
4684 ],
4685 vec![],
4686 vec![float_column.to_string(), histogram_column.to_string()],
4687 ),
4688 token::T_COUNT => {
4689 let present = when(
4690 float_input
4691 .clone()
4692 .is_not_null()
4693 .or(histogram_input.clone().is_not_null()),
4694 lit(1_i64),
4695 )
4696 .otherwise(lit(ScalarValue::Int64(None)))
4697 .context(DataFusionPlanningSnafu)?;
4698 (
4699 vec![count_udaf().call(vec![present]).alias(float_column)],
4700 vec![],
4701 vec![float_column.to_string()],
4702 )
4703 }
4704 token::T_GROUP => (
4705 vec![max_udaf().call(vec![lit(1_f64)]).alias(float_column)],
4706 vec![],
4707 vec![float_column.to_string()],
4708 ),
4709 token::T_COUNT_VALUES => {
4710 let value = mixed_sample_value();
4711 (
4712 vec![count_udaf().call(vec![value.clone()]).alias(float_column)],
4713 vec![value],
4714 vec![float_column.to_string()],
4715 )
4716 }
4717 token::T_MIN | token::T_MAX | token::T_STDDEV | token::T_STDVAR | token::T_QUANTILE => {
4718 (
4719 vec![
4720 Self::create_numeric_aggregate_expr(op, param, float_input)?
4721 .alias(float_column),
4722 histogram_count.alias(Self::mixed_sample_count_name(histogram_column)),
4723 ],
4724 vec![],
4725 vec![float_column.to_string()],
4726 )
4727 }
4728 token::T_TOPK | token::T_BOTTOMK => {
4729 return UnsupportedExprSnafu {
4730 name: format!("{op:?}"),
4731 }
4732 .fail();
4733 }
4734 _ => return UnexpectedTokenSnafu { token: op }.fail(),
4735 };
4736
4737 self.ctx.field_columns = field_columns;
4738 Ok((exprs, prev_field_exprs))
4739 }
4740
4741 fn mixed_sample_count_column(column: &str) -> DfExpr {
4742 DfExpr::Column(Column::from_name(Self::mixed_sample_count_name(column)))
4743 }
4744
4745 fn mixed_sample_count_name(column: &str) -> String {
4746 format!("__promql_sample_count({column})")
4747 }
4748
4749 fn mixed_aggregate_filter_expr(
4750 &self,
4751 op: TokenType,
4752 float_column: &str,
4753 histogram_column: &str,
4754 ) -> Result<DfExpr> {
4755 let float_count = Self::mixed_sample_count_column(float_column);
4756 let histogram_count = Self::mixed_sample_count_column(histogram_column);
4757 let mixed = float_count
4758 .clone()
4759 .gt(lit(0_i64))
4760 .and(histogram_count.clone().gt(lit(0_i64)));
4761 let drop_mixed = DfExpr::ScalarFunction(ScalarFunction {
4762 func: Arc::new(NativeHistogramDrop::warning_bool_false_udf(
4763 format!(
4764 "{op}: dropped aggregation result containing both float and native histogram samples"
4765 ),
4766 self.promql_annotations.clone(),
4767 )),
4768 args: vec![float_count, histogram_count],
4769 });
4770
4771 when(mixed, drop_mixed)
4772 .otherwise(lit(true))
4773 .context(DataFusionPlanningSnafu)
4774 }
4775
4776 fn mixed_ignored_histogram_filter_expr(
4777 &self,
4778 op: TokenType,
4779 histogram_column: &str,
4780 ) -> Result<DfExpr> {
4781 let histogram_count = Self::mixed_sample_count_column(histogram_column);
4782 let has_histograms = histogram_count.clone().gt(lit(0_i64));
4783 let record_info = DfExpr::ScalarFunction(ScalarFunction {
4784 func: Arc::new(NativeHistogramDrop::bool_true_udf(
4785 format!(
4786 "{op}: dropped native histogram samples because this aggregation is not supported for native histograms"
4787 ),
4788 self.promql_annotations.clone(),
4789 )),
4790 args: vec![histogram_count],
4791 });
4792
4793 when(has_histograms, record_info)
4794 .otherwise(lit(true))
4795 .context(DataFusionPlanningSnafu)
4796 }
4797
4798 fn create_native_histogram_aggregate_expr(
4799 &self,
4800 op: TokenType,
4801 column: &str,
4802 ) -> Result<DfExpr> {
4803 let input = DfExpr::Column(Column::from_name(column));
4804 let expr = match op.id() {
4805 token::T_SUM => Arc::new(NativeHistogramAggSum::aggregate_udf_with_collector(
4806 self.promql_annotations.clone(),
4807 ))
4808 .call(vec![input])
4809 .alias(column),
4810 token::T_AVG => Arc::new(NativeHistogramAggAvg::aggregate_udf_with_collector(
4811 self.promql_annotations.clone(),
4812 ))
4813 .call(vec![input])
4814 .alias(column),
4815 token::T_COUNT_VALUES | token::T_COUNT => {
4816 count_udaf().call(vec![input]).alias(column)
4817 }
4818 token::T_GROUP => max_udaf().call(vec![lit(1_f64)]).alias(column),
4819 token::T_MIN
4820 | token::T_MAX
4821 | token::T_STDDEV
4822 | token::T_STDVAR
4823 | token::T_QUANTILE
4824 | token::T_TOPK
4825 | token::T_BOTTOMK => sum_udaf()
4826 .call(vec![DfExpr::ScalarFunction(ScalarFunction {
4827 func: Arc::new(NativeHistogramDrop::float_null_udf(
4828 format!(
4829 "{op}: dropped native histogram samples because this aggregation is not supported for native histograms"
4830 ),
4831 self.promql_annotations.clone(),
4832 )),
4833 args: vec![input],
4834 })])
4835 .alias(column),
4836 _ => return UnexpectedTokenSnafu { token: op }.fail(),
4837 };
4838 Ok(expr)
4839 }
4840
4841 fn create_native_histogram_aggregate_exprs(
4842 &mut self,
4843 op: TokenType,
4844 input_plan: &LogicalPlan,
4845 ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4846 let prev_field_exprs = if op.id() == token::T_COUNT_VALUES {
4847 ensure!(
4848 self.ctx.field_columns.len() == 1,
4849 UnsupportedExprSnafu {
4850 name: "count_values on multi-value input"
4851 }
4852 );
4853 self.ctx
4854 .field_columns
4855 .iter()
4856 .map(|col| {
4857 DfExpr::ScalarFunction(ScalarFunction {
4858 func: Arc::new(NativeHistogramToString::scalar_udf()),
4859 args: vec![DfExpr::Column(Column::from_name(col))],
4860 })
4861 })
4862 .collect::<Vec<_>>()
4863 } else {
4864 vec![]
4865 };
4866
4867 let exprs = self
4868 .ctx
4869 .field_columns
4870 .iter()
4871 .map(|col| self.create_native_histogram_aggregate_expr(op, col))
4872 .collect::<Result<Vec<_>>>()?;
4873
4874 let normalized_exprs =
4875 normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
4876 self.ctx.field_columns = normalized_exprs
4877 .into_iter()
4878 .map(|expr| expr.schema_name().to_string())
4879 .collect();
4880
4881 Ok((exprs, prev_field_exprs))
4882 }
4883
4884 fn get_param_value_as_str(op: TokenType, param: &Option<Box<PromExpr>>) -> Result<&str> {
4885 let param = param
4886 .as_deref()
4887 .with_context(|| FunctionInvalidArgumentSnafu {
4888 fn_name: op.to_string(),
4889 })?;
4890 let PromExpr::StringLiteral(StringLiteral { val }) = param else {
4891 return FunctionInvalidArgumentSnafu {
4892 fn_name: op.to_string(),
4893 }
4894 .fail();
4895 };
4896
4897 Ok(val)
4898 }
4899
4900 fn get_param_as_literal_expr(
4901 param: Option<&PromExpr>,
4902 op: Option<TokenType>,
4903 expected_type: Option<ArrowDataType>,
4904 ) -> Result<DfExpr> {
4905 let prom_param = param.with_context(|| {
4906 if let Some(op) = op {
4907 FunctionInvalidArgumentSnafu {
4908 fn_name: op.to_string(),
4909 }
4910 } else {
4911 FunctionInvalidArgumentSnafu {
4912 fn_name: "unknown".to_string(),
4913 }
4914 }
4915 })?;
4916
4917 let expr = Self::try_build_literal_expr(prom_param).with_context(|| {
4918 if let Some(op) = op {
4919 FunctionInvalidArgumentSnafu {
4920 fn_name: op.to_string(),
4921 }
4922 } else {
4923 FunctionInvalidArgumentSnafu {
4924 fn_name: "unknown".to_string(),
4925 }
4926 }
4927 })?;
4928
4929 if let Some(expected_type) = expected_type {
4931 let expr_type = expr
4933 .get_type(&DFSchema::empty())
4934 .context(DataFusionPlanningSnafu)?;
4935 if expected_type != expr_type {
4936 return FunctionInvalidArgumentSnafu {
4937 fn_name: format!("expected {expected_type:?}, but found {expr_type:?}"),
4938 }
4939 .fail();
4940 }
4941 }
4942
4943 Ok(expr)
4944 }
4945
4946 fn create_window_exprs(
4949 &mut self,
4950 op: TokenType,
4951 group_exprs: Vec<DfExpr>,
4952 input_plan: &LogicalPlan,
4953 ) -> Result<Vec<DfExpr>> {
4954 ensure!(
4955 self.ctx.field_columns.len() == 1,
4956 UnsupportedExprSnafu {
4957 name: "topk or bottomk on multi-value input"
4958 }
4959 );
4960
4961 assert!(matches!(op.id(), token::T_TOPK | token::T_BOTTOMK));
4962
4963 let asc = matches!(op.id(), token::T_BOTTOMK);
4964
4965 let tag_sort_exprs = self
4966 .create_tag_column_exprs()?
4967 .into_iter()
4968 .map(|expr| expr.sort(asc, true));
4969
4970 let exprs: Vec<DfExpr> = self
4972 .ctx
4973 .field_columns
4974 .iter()
4975 .map(|col| {
4976 let mut sort_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 1);
4977 sort_exprs.push(DfExpr::Column(Column::from(col)).sort(asc, true));
4979 sort_exprs.extend(tag_sort_exprs.clone());
4982
4983 DfExpr::WindowFunction(Box::new(WindowFunction {
4984 fun: WindowFunctionDefinition::WindowUDF(Arc::new(RowNumber::new().into())),
4985 params: WindowFunctionParams {
4986 args: vec![],
4987 partition_by: group_exprs.clone(),
4988 order_by: sort_exprs,
4989 window_frame: WindowFrame::new(Some(true)),
4990 null_treatment: None,
4991 distinct: false,
4992 filter: None,
4993 },
4994 }))
4995 })
4996 .collect();
4997
4998 let normalized_exprs =
4999 normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
5000 Ok(normalized_exprs)
5001 }
5002
5003 async fn create_histogram_plan(
5005 &mut self,
5006 function_name: &str,
5007 args: &PromFunctionArgs,
5008 query_engine_state: &QueryEngineState,
5009 ) -> Result<LogicalPlan> {
5010 let float_literal = |param: &PromExpr| -> Result<f64> {
5011 let value = (|| {
5012 let expr = Self::get_param_as_literal_expr(
5013 Some(param),
5014 None,
5015 Some(ArrowDataType::Float64),
5016 )
5017 .ok()?;
5018 let simplifier = ExprSimplifier::new(SimplifyContext::default());
5019 let expr = simplifier.coerce(expr, &DFSchema::empty()).ok()?;
5020 let DfExpr::Literal(value, _) = simplifier.simplify(expr).ok()? else {
5021 return None;
5022 };
5023 let ScalarValue::Float64(Some(value)) =
5024 value.cast_to(&ArrowDataType::Float64).ok()?
5025 else {
5026 return None;
5027 };
5028 Some(value)
5029 })()
5030 .with_context(|| FunctionInvalidArgumentSnafu {
5031 fn_name: function_name.to_string(),
5032 })?;
5033 Ok(value)
5034 };
5035 let (function, input) = match (function_name, args.args.as_slice()) {
5036 (SPECIAL_HISTOGRAM_QUANTILE, [quantile, input]) => (
5037 HistogramFoldOperation::Quantile(float_literal(quantile)?.into()),
5038 input.as_ref().clone(),
5039 ),
5040 (SPECIAL_HISTOGRAM_FRACTION, [lower, upper, input]) => (
5041 HistogramFoldOperation::Fraction {
5042 lower: float_literal(lower)?.into(),
5043 upper: float_literal(upper)?.into(),
5044 },
5045 input.as_ref().clone(),
5046 ),
5047 _ => {
5048 return FunctionInvalidArgumentSnafu {
5049 fn_name: function_name.to_string(),
5050 }
5051 .fail();
5052 }
5053 };
5054
5055 let input_plan = self.prom_expr_to_plan(&input, query_engine_state).await?;
5056 let input_plan = self.strip_tsid_column(input_plan)?;
5059 self.ctx.use_tsid = false;
5060
5061 if let Some((float_field, histogram_field)) =
5062 Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns)
5063 .map(|(float, histogram)| (float.to_string(), histogram.to_string()))
5064 {
5065 if self.ctx.has_le_tag() {
5066 return self.create_mixed_histogram_plan(
5067 function,
5068 input_plan,
5069 float_field,
5070 histogram_field,
5071 );
5072 }
5073 self.ctx.field_columns = vec![histogram_field];
5074 }
5075 if self.all_field_columns_are_native_histograms(input_plan.schema()) {
5076 return self.create_native_histogram_plan(function, input_plan);
5077 }
5078
5079 if !self.ctx.has_le_tag() {
5080 return Ok(LogicalPlan::EmptyRelation(
5083 datafusion::logical_expr::EmptyRelation {
5084 produce_one_row: false,
5085 schema: input_plan.schema().clone(),
5086 },
5087 ));
5088 }
5089 let time_index_column =
5090 self.ctx
5091 .time_index_column
5092 .clone()
5093 .with_context(|| TimeIndexNotFoundSnafu {
5094 table: self.ctx.table_name.clone().unwrap_or_default(),
5095 })?;
5096 let field_column = self
5098 .ctx
5099 .field_columns
5100 .first()
5101 .with_context(|| FunctionInvalidArgumentSnafu {
5102 fn_name: function.function_name().to_string(),
5103 })?
5104 .clone();
5105 self.ctx.tag_columns.retain(|col| col != LE_COLUMN_NAME);
5107
5108 let fold = HistogramFold::new_with_operation(
5109 LE_COLUMN_NAME.to_string(),
5110 field_column,
5111 time_index_column,
5112 function,
5113 None,
5114 input_plan,
5115 )
5116 .context(DataFusionPlanningSnafu)?;
5117 Ok(LogicalPlan::Extension(Extension {
5118 node: Arc::new(fold),
5119 }))
5120 }
5121
5122 fn create_native_histogram_expr(
5123 &self,
5124 function: HistogramFoldOperation,
5125 field_column: &str,
5126 ) -> DfExpr {
5127 let field = DfExpr::Column(Column::from_name(field_column));
5128 let (func, args) = match function {
5129 HistogramFoldOperation::Quantile(quantile) => (
5130 Arc::new(NativeHistogramQuantile::scalar_udf_with_collector(
5131 self.promql_annotations.clone(),
5132 )),
5133 vec![field, lit(f64::from(quantile))],
5134 ),
5135 HistogramFoldOperation::Fraction { lower, upper } => (
5136 Arc::new(NativeHistogramFraction::scalar_udf_with_collector(
5137 self.promql_annotations.clone(),
5138 )),
5139 vec![field, lit(f64::from(lower)), lit(f64::from(upper))],
5140 ),
5141 };
5142 DfExpr::ScalarFunction(ScalarFunction { func, args })
5143 }
5144
5145 fn create_native_histogram_plan(
5146 &mut self,
5147 function: HistogramFoldOperation,
5148 input_plan: LogicalPlan,
5149 ) -> Result<LogicalPlan> {
5150 ensure!(
5151 self.ctx.field_columns.len() == 1,
5152 MultiFieldsNotSupportedSnafu {
5153 operator: function.function_name()
5154 },
5155 );
5156
5157 let field_column = self.ctx.field_columns[0].clone();
5158 let function_expr = self.create_native_histogram_expr(function, &field_column);
5159 let display_name = function_expr.schema_name().to_string();
5160 self.ctx.field_columns = vec![display_name.clone()];
5161
5162 let project_exprs = std::iter::once(self.create_time_index_column_expr()?)
5163 .chain(std::iter::once(function_expr.alias(display_name)))
5164 .chain(self.create_tag_column_exprs()?)
5165 .collect::<Vec<_>>();
5166
5167 LogicalPlanBuilder::from(input_plan)
5168 .project(project_exprs)
5169 .context(DataFusionPlanningSnafu)?
5170 .filter(self.create_empty_values_filter_expr(false)?)
5171 .context(DataFusionPlanningSnafu)?
5172 .build()
5173 .context(DataFusionPlanningSnafu)
5174 }
5175
5176 fn create_mixed_histogram_plan(
5177 &mut self,
5178 function: HistogramFoldOperation,
5179 input_plan: LogicalPlan,
5180 float_field: String,
5181 histogram_field: String,
5182 ) -> Result<LogicalPlan> {
5183 let time_index_column =
5184 self.ctx
5185 .time_index_column
5186 .clone()
5187 .with_context(|| TimeIndexNotFoundSnafu {
5188 table: self.ctx.table_name.clone().unwrap_or_default(),
5189 })?;
5190 let tag_columns = self.ctx.tag_columns.clone();
5191 let folded = HistogramFold::new_with_operation(
5192 LE_COLUMN_NAME.to_string(),
5193 float_field.clone(),
5194 time_index_column.clone(),
5195 function,
5196 Some(histogram_field.clone()),
5197 input_plan,
5198 )
5199 .context(DataFusionPlanningSnafu)?;
5200 let record_collision = DfExpr::ScalarFunction(ScalarFunction {
5201 func: Arc::new(NativeHistogramDrop::warning_bool_false_udf(
5202 "vector contains a mix of classic and native histograms".to_string(),
5203 self.promql_annotations.clone(),
5204 )),
5205 args: vec![col(&float_field), col(&histogram_field)],
5206 });
5207 let keep = when(
5208 col(&float_field)
5209 .is_not_null()
5210 .and(col(&histogram_field).is_not_null()),
5211 record_collision,
5212 )
5213 .otherwise(lit(true))
5214 .context(DataFusionPlanningSnafu)?;
5215
5216 let native_expr = self.create_native_histogram_expr(function, &histogram_field);
5217 let output_field = native_expr.schema_name().to_string();
5218 let value = DfExpr::ScalarFunction(ScalarFunction {
5219 func: coalesce(),
5220 args: vec![col(&float_field), native_expr],
5221 });
5222 self.ctx.field_columns = vec![output_field.clone()];
5223 LogicalPlanBuilder::from(LogicalPlan::Extension(Extension {
5224 node: Arc::new(folded),
5225 }))
5226 .filter(keep)
5227 .context(DataFusionPlanningSnafu)?
5228 .project(
5229 std::iter::once(col(&time_index_column))
5230 .chain(std::iter::once(value.alias(output_field)))
5231 .chain(tag_columns.iter().map(col)),
5232 )
5233 .context(DataFusionPlanningSnafu)?
5234 .build()
5235 .context(DataFusionPlanningSnafu)
5236 }
5237
5238 async fn create_vector_plan(&mut self, args: &PromFunctionArgs) -> Result<LogicalPlan> {
5240 if args.args.len() != 1 {
5241 return FunctionInvalidArgumentSnafu {
5242 fn_name: SPECIAL_VECTOR_FUNCTION.to_string(),
5243 }
5244 .fail();
5245 }
5246 let lit = Self::get_param_as_literal_expr(Some(args.args[0].as_ref()), None, None)?;
5247
5248 self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
5250 self.ctx.reset_table_name_and_schema();
5251 self.ctx.tag_columns = vec![];
5252 self.ctx.aggregation_field_labels.clear();
5253 self.ctx.field_columns = vec![greptime_value().to_string()];
5254 Ok(LogicalPlan::Extension(Extension {
5255 node: Arc::new(
5256 EmptyMetric::new(
5257 self.ctx.start,
5258 self.ctx.end,
5259 self.ctx.interval,
5260 SPECIAL_TIME_FUNCTION.to_string(),
5261 greptime_value().to_string(),
5262 Some(lit),
5263 )
5264 .context(DataFusionPlanningSnafu)?,
5265 ),
5266 }))
5267 }
5268
5269 async fn create_scalar_plan(
5271 &mut self,
5272 args: &PromFunctionArgs,
5273 query_engine_state: &QueryEngineState,
5274 ) -> Result<LogicalPlan> {
5275 ensure!(
5276 args.len() == 1,
5277 FunctionInvalidArgumentSnafu {
5278 fn_name: SCALAR_FUNCTION
5279 }
5280 );
5281 let input = self
5282 .prom_expr_to_plan(&args.args[0], query_engine_state)
5283 .await?;
5284 let input_schema = input.schema().clone();
5285 let alternative_samples =
5286 Self::field_columns_are_alternative_samples(&input_schema, &self.ctx.field_columns);
5287 let histogram_fields = self
5288 .ctx
5289 .field_columns
5290 .iter()
5291 .filter(|field| Self::field_column_is_native_histogram(&input_schema, field))
5292 .count();
5293 ensure!(
5294 self.ctx.field_columns.len() == 1 || alternative_samples,
5295 MultiFieldsNotSupportedSnafu {
5296 operator: SCALAR_FUNCTION
5297 },
5298 );
5299 let scalar_field = self
5300 .ctx
5301 .field_columns
5302 .iter()
5303 .find(|field| !Self::field_column_is_native_histogram(&input_schema, field))
5304 .or_else(|| self.ctx.field_columns.first())
5305 .cloned()
5306 .with_context(|| FunctionInvalidArgumentSnafu {
5307 fn_name: SCALAR_FUNCTION,
5308 })?;
5309 let input = if histogram_fields == self.ctx.field_columns.len() {
5310 LogicalPlanBuilder::from(input)
5313 .filter(lit(false))
5314 .context(DataFusionPlanningSnafu)?
5315 .build()
5316 .context(DataFusionPlanningSnafu)?
5317 } else if histogram_fields > 0 {
5318 LogicalPlanBuilder::from(input)
5320 .filter(DfExpr::Column(Column::from_name(&scalar_field)).is_not_null())
5321 .context(DataFusionPlanningSnafu)?
5322 .build()
5323 .context(DataFusionPlanningSnafu)?
5324 } else {
5325 input
5326 };
5327 let scalar_plan = LogicalPlan::Extension(Extension {
5328 node: Arc::new(
5329 ScalarCalculate::new(
5330 self.ctx.start,
5331 self.ctx.end,
5332 self.ctx.interval,
5333 input,
5334 self.ctx.time_index_column.as_ref().unwrap(),
5335 &self.ctx.tag_columns,
5336 &scalar_field,
5337 self.ctx.table_name.as_deref(),
5338 )
5339 .context(PromqlPlanNodeSnafu)?,
5340 ),
5341 });
5342 self.ctx.tag_columns.clear();
5344 self.ctx.aggregation_field_labels.clear();
5345 self.ctx.field_columns.clear();
5346 self.ctx
5347 .field_columns
5348 .push(scalar_plan.schema().field(1).name().clone());
5349 Ok(scalar_plan)
5350 }
5351
5352 async fn create_absent_plan(
5354 &mut self,
5355 args: &PromFunctionArgs,
5356 query_engine_state: &QueryEngineState,
5357 ) -> Result<LogicalPlan> {
5358 if args.args.len() != 1 {
5359 return FunctionInvalidArgumentSnafu {
5360 fn_name: SPECIAL_ABSENT_FUNCTION.to_string(),
5361 }
5362 .fail();
5363 }
5364 let input = self
5365 .prom_expr_to_plan(&args.args[0], query_engine_state)
5366 .await?;
5367
5368 let time_index_expr = self.create_time_index_column_expr()?;
5369 let first_field_expr =
5370 self.create_field_column_exprs()?
5371 .pop()
5372 .with_context(|| ValueNotFoundSnafu {
5373 table: self.ctx.table_name.clone().unwrap_or_default(),
5374 })?;
5375 let first_value_expr = first_value(first_field_expr, vec![]);
5376
5377 let ordered_aggregated_input = LogicalPlanBuilder::from(input)
5378 .aggregate(
5379 vec![time_index_expr.clone()],
5380 vec![first_value_expr.clone()],
5381 )
5382 .context(DataFusionPlanningSnafu)?
5383 .sort(vec![time_index_expr.sort(true, false)])
5384 .context(DataFusionPlanningSnafu)?
5385 .build()
5386 .context(DataFusionPlanningSnafu)?;
5387
5388 let fake_labels = self
5389 .ctx
5390 .selector_matcher
5391 .iter()
5392 .filter_map(|matcher| match matcher.op {
5393 MatchOp::Equal => Some((matcher.name.clone(), matcher.value.clone())),
5394 _ => None,
5395 })
5396 .collect::<Vec<_>>();
5397
5398 let absent_plan = LogicalPlan::Extension(Extension {
5400 node: Arc::new(
5401 Absent::try_new(
5402 self.ctx.start,
5403 self.ctx.end,
5404 self.ctx.interval,
5405 self.ctx.time_index_column.as_ref().unwrap().clone(),
5406 self.ctx.field_columns[0].clone(),
5407 fake_labels,
5408 ordered_aggregated_input,
5409 )
5410 .context(DataFusionPlanningSnafu)?,
5411 ),
5412 });
5413
5414 self.ctx.aggregation_field_labels.clear();
5417 Ok(absent_plan)
5418 }
5419
5420 fn try_build_literal_expr(expr: &PromExpr) -> Option<DfExpr> {
5423 match expr {
5424 PromExpr::NumberLiteral(NumberLiteral { val }) => Some(val.lit()),
5425 PromExpr::StringLiteral(StringLiteral { val }) => Some(val.lit()),
5426 PromExpr::VectorSelector(_)
5427 | PromExpr::MatrixSelector(_)
5428 | PromExpr::Extension(_)
5429 | PromExpr::Aggregate(_)
5430 | PromExpr::Subquery(_) => None,
5431 PromExpr::Call(Call { func, .. }) => {
5432 if func.name == SPECIAL_TIME_FUNCTION {
5433 None
5436 } else {
5437 None
5438 }
5439 }
5440 PromExpr::Paren(ParenExpr { expr }) => Self::try_build_literal_expr(expr),
5441 PromExpr::Unary(UnaryExpr { expr, .. }) => Some(DfExpr::Negative(Box::new(
5442 Self::try_build_literal_expr(expr)?,
5443 ))),
5444 PromExpr::Binary(PromBinaryExpr {
5445 lhs,
5446 rhs,
5447 op,
5448 modifier,
5449 }) => {
5450 let lhs = Self::try_build_literal_expr(lhs)?;
5451 let rhs = Self::try_build_literal_expr(rhs)?;
5452 let is_comparison_op = Self::is_token_a_comparison_op(*op);
5453 let expr_builder = Self::prom_token_to_binary_expr_builder(*op).ok()?;
5454 let expr = expr_builder(lhs, rhs).ok()?;
5455
5456 let should_return_bool = if let Some(m) = modifier {
5457 m.return_bool
5458 } else {
5459 false
5460 };
5461 if is_comparison_op && should_return_bool {
5462 Some(DfExpr::Cast(Cast::new(
5463 Box::new(expr),
5464 ArrowDataType::Float64,
5465 )))
5466 } else {
5467 Some(expr)
5468 }
5469 }
5470 }
5471 }
5472
5473 fn try_build_special_time_expr_with_context(&self, expr: &PromExpr) -> Option<DfExpr> {
5474 match expr {
5475 PromExpr::Call(Call { func, .. }) => {
5476 if func.name == SPECIAL_TIME_FUNCTION
5477 && let Some(time_index_col) = self.ctx.time_index_column.as_ref()
5478 {
5479 Some(build_special_time_expr(time_index_col))
5480 } else {
5481 None
5482 }
5483 }
5484 _ => None,
5485 }
5486 }
5487
5488 fn native_histogram_binary_expr(
5489 token: TokenType,
5490 lhs: DfExpr,
5491 lhs_is_histogram: bool,
5492 rhs: DfExpr,
5493 rhs_is_histogram: bool,
5494 filter_context: bool,
5495 promql_annotations: Option<PromqlAnnotationCollector>,
5496 ) -> Result<Option<DfExpr>> {
5497 if !lhs_is_histogram && !rhs_is_histogram {
5498 return Ok(None);
5499 }
5500
5501 let scalar_fn = |func: ScalarUdfDef, args| {
5502 DfExpr::ScalarFunction(ScalarFunction {
5503 func: Arc::new(func),
5504 args,
5505 })
5506 };
5507 let invalid_expr = || {
5508 let message = format!(
5509 "{}: dropped native histogram samples because this binary operation is not supported for native histograms",
5510 token
5511 );
5512 let func = if filter_context {
5513 NativeHistogramDrop::bool_false_udf(message, promql_annotations.clone())
5514 } else {
5515 NativeHistogramDrop::float_null_udf(message, promql_annotations.clone())
5516 };
5517 let args = vec![lhs.clone(), rhs.clone()];
5518 scalar_fn(func, args)
5519 };
5520
5521 let expr = match (token.id(), lhs_is_histogram, rhs_is_histogram) {
5522 (token::T_ADD, true, true) => scalar_fn(
5523 NativeHistogramAdd::scalar_udf_with_collector(promql_annotations.clone()),
5524 vec![lhs, rhs],
5525 ),
5526 (token::T_SUB, true, true) => scalar_fn(
5527 NativeHistogramSub::scalar_udf_with_collector(promql_annotations.clone()),
5528 vec![lhs, rhs],
5529 ),
5530 (token::T_MUL, true, false) => {
5531 scalar_fn(NativeHistogramMulScalar::scalar_udf(), vec![lhs, rhs])
5532 }
5533 (token::T_MUL, false, true) => {
5534 scalar_fn(NativeHistogramScalarMul::scalar_udf(), vec![lhs, rhs])
5535 }
5536 (token::T_DIV, true, false) => {
5537 scalar_fn(NativeHistogramDivScalar::scalar_udf(), vec![lhs, rhs])
5538 }
5539 (token::T_EQLC, true, true) => {
5540 scalar_fn(NativeHistogramEq::scalar_udf(), vec![lhs, rhs])
5541 }
5542 (token::T_NEQ, true, true) => {
5543 scalar_fn(NativeHistogramNotEq::scalar_udf(), vec![lhs, rhs])
5544 }
5545 _ => invalid_expr(),
5546 };
5547
5548 Ok(Some(expr))
5549 }
5550
5551 #[allow(clippy::type_complexity)]
5554 fn prom_token_to_binary_expr_builder(
5555 token: TokenType,
5556 ) -> Result<Box<dyn Fn(DfExpr, DfExpr) -> Result<DfExpr>>> {
5557 let cast_float = |expr| {
5558 if matches!(
5559 &expr,
5560 DfExpr::Cast(Cast { field, .. }) if field.data_type() == &ArrowDataType::Float64
5561 ) || matches!(&expr, DfExpr::Literal(ScalarValue::Float64(_), _))
5562 {
5563 expr
5564 } else {
5565 DfExpr::Cast(Cast::new(Box::new(expr), ArrowDataType::Float64))
5566 }
5567 };
5568 match token.id() {
5569 token::T_ADD => Ok(Box::new(move |lhs, rhs| {
5570 Ok(cast_float(lhs) + cast_float(rhs))
5571 })),
5572 token::T_SUB => Ok(Box::new(move |lhs, rhs| {
5573 Ok(cast_float(lhs) - cast_float(rhs))
5574 })),
5575 token::T_MUL => Ok(Box::new(move |lhs, rhs| {
5576 Ok(cast_float(lhs) * cast_float(rhs))
5577 })),
5578 token::T_DIV => Ok(Box::new(move |lhs, rhs| {
5579 Ok(cast_float(lhs) / cast_float(rhs))
5580 })),
5581 token::T_MOD => Ok(Box::new(move |lhs: DfExpr, rhs| {
5582 Ok(cast_float(lhs) % cast_float(rhs))
5583 })),
5584 token::T_EQLC => Ok(Box::new(|lhs, rhs| Ok(lhs.eq(rhs)))),
5585 token::T_NEQ => Ok(Box::new(|lhs, rhs| Ok(lhs.not_eq(rhs)))),
5586 token::T_GTR => Ok(Box::new(|lhs, rhs| Ok(lhs.gt(rhs)))),
5587 token::T_LSS => Ok(Box::new(|lhs, rhs| Ok(lhs.lt(rhs)))),
5588 token::T_GTE => Ok(Box::new(|lhs, rhs| Ok(lhs.gt_eq(rhs)))),
5589 token::T_LTE => Ok(Box::new(|lhs, rhs| Ok(lhs.lt_eq(rhs)))),
5590 token::T_POW => Ok(Box::new(move |lhs, rhs| {
5591 Ok(DfExpr::ScalarFunction(ScalarFunction {
5592 func: datafusion_functions::math::power(),
5593 args: vec![cast_float(lhs), cast_float(rhs)],
5594 }))
5595 })),
5596 token::T_ATAN2 => Ok(Box::new(move |lhs, rhs| {
5597 Ok(DfExpr::ScalarFunction(ScalarFunction {
5598 func: datafusion_functions::math::atan2(),
5599 args: vec![cast_float(lhs), cast_float(rhs)],
5600 }))
5601 })),
5602 _ => UnexpectedTokenSnafu { token }.fail(),
5603 }
5604 }
5605
5606 fn is_token_a_comparison_op(token: TokenType) -> bool {
5608 matches!(
5609 token.id(),
5610 token::T_EQLC
5611 | token::T_NEQ
5612 | token::T_GTR
5613 | token::T_LSS
5614 | token::T_GTE
5615 | token::T_LTE
5616 )
5617 }
5618
5619 fn is_token_a_set_op(token: TokenType) -> bool {
5621 matches!(
5622 token.id(),
5623 token::T_LAND | token::T_LOR | token::T_LUNLESS )
5627 }
5628
5629 fn align_binary_field_columns<'a>(
5630 left_schema: &DFSchemaRef,
5631 right_schema: &DFSchemaRef,
5632 left_field_columns: &'a [String],
5633 right_field_columns: &'a [String],
5634 op: TokenType,
5635 left_is_scalar: bool,
5636 right_is_scalar: bool,
5637 ) -> (
5638 Vec<(String, Vec<BinaryFieldPair<'a>>)>,
5639 Vec<BinaryFieldPair<'a>>,
5640 ) {
5641 let left_alternative = Self::alternative_sample_columns(left_schema, left_field_columns);
5644 let right_alternative = Self::alternative_sample_columns(right_schema, right_field_columns);
5645 let alternative_alignment = match (left_alternative, right_alternative) {
5646 (Some(output_names), Some(_)) => Some((
5647 output_names,
5648 left_field_columns
5649 .iter()
5650 .flat_map(|left| right_field_columns.iter().map(move |right| (left, right)))
5651 .collect::<Vec<_>>(),
5652 )),
5653 (Some(output_names), None) if right_field_columns.len() == 1 => Some((
5654 output_names,
5655 left_field_columns
5656 .iter()
5657 .map(|left| (left, &right_field_columns[0]))
5658 .collect::<Vec<_>>(),
5659 )),
5660 (None, Some(output_names)) if left_field_columns.len() == 1 => Some((
5661 output_names,
5662 right_field_columns
5663 .iter()
5664 .map(|right| (&left_field_columns[0], right))
5665 .collect::<Vec<_>>(),
5666 )),
5667 _ => None,
5668 };
5669 let mut invalid_pairs = Vec::new();
5670 if let Some(((float_output, histogram_output), field_pairs)) = alternative_alignment {
5671 let mut float_pairs = Vec::new();
5672 let mut histogram_pairs = Vec::new();
5673 for (left, right) in field_pairs {
5674 let left_is_histogram = Self::field_column_is_native_histogram(left_schema, left);
5675 let right_is_histogram =
5676 Self::field_column_is_native_histogram(right_schema, right);
5677 match Self::binary_result_is_histogram(op, left_is_histogram, right_is_histogram) {
5678 Some(false) => float_pairs.push((left, right)),
5679 Some(true) => histogram_pairs.push((left, right)),
5680 None => invalid_pairs.push((left, right)),
5681 }
5682 }
5683 if !float_pairs.is_empty() || !histogram_pairs.is_empty() {
5684 return (
5685 [
5686 (!float_pairs.is_empty()).then(|| (float_output.to_string(), float_pairs)),
5687 (!histogram_pairs.is_empty())
5688 .then(|| (histogram_output.to_string(), histogram_pairs)),
5689 ]
5690 .into_iter()
5691 .flatten()
5692 .collect(),
5693 invalid_pairs,
5694 );
5695 }
5696 }
5697
5698 if left_is_scalar && !right_is_scalar && left_field_columns.len() == 1 {
5699 return (
5700 right_field_columns
5701 .iter()
5702 .map(|right| (right.clone(), vec![(&left_field_columns[0], right)]))
5703 .collect(),
5704 invalid_pairs,
5705 );
5706 }
5707 if right_is_scalar && !left_is_scalar && right_field_columns.len() == 1 {
5708 return (
5709 left_field_columns
5710 .iter()
5711 .map(|left| (left.clone(), vec![(left, &right_field_columns[0])]))
5712 .collect(),
5713 invalid_pairs,
5714 );
5715 }
5716
5717 (
5718 left_field_columns
5719 .iter()
5720 .zip(right_field_columns.iter())
5721 .map(|(left, right)| (left.clone(), vec![(left, right)]))
5722 .collect(),
5723 invalid_pairs,
5724 )
5725 }
5726
5727 fn binary_result_is_histogram(
5728 token: TokenType,
5729 lhs_is_histogram: bool,
5730 rhs_is_histogram: bool,
5731 ) -> Option<bool> {
5732 match (token.id(), lhs_is_histogram, rhs_is_histogram) {
5733 (_, false, false) => Some(false),
5734 (token::T_ADD | token::T_SUB, true, true)
5735 | (token::T_MUL, true, false)
5736 | (token::T_MUL, false, true)
5737 | (token::T_DIV, true, false) => Some(true),
5738 (token::T_EQLC | token::T_NEQ, true, true) => Some(false),
5739 _ => None,
5740 }
5741 }
5742
5743 fn plan_has_tsid_column(plan: &LogicalPlan) -> bool {
5744 plan.schema()
5745 .fields()
5746 .iter()
5747 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
5748 }
5749
5750 fn is_empty_metric(plan: &LogicalPlan) -> bool {
5751 matches!(plan, LogicalPlan::Extension(Extension { node }) if node.as_any().is::<EmptyMetric>())
5752 }
5753
5754 fn native_histogram_arrow_type() -> ArrowDataType {
5755 native_histogram_value_type().as_arrow_type()
5756 }
5757
5758 fn field_column_type<'a>(
5759 schema: &'a DFSchemaRef,
5760 field_column: &str,
5761 ) -> Option<&'a ArrowDataType> {
5762 schema
5763 .index_of_column_by_name(None, field_column)
5764 .map(|idx| schema.field(idx).data_type())
5765 }
5766
5767 fn field_column_is_native_histogram(schema: &DFSchemaRef, field_column: &str) -> bool {
5768 Self::field_column_type(schema, field_column)
5769 .is_some_and(|data_type| data_type == &Self::native_histogram_arrow_type())
5770 }
5771
5772 fn field_columns_contain_native_histogram(
5773 schema: &DFSchemaRef,
5774 field_columns: &[String],
5775 ) -> bool {
5776 field_columns
5777 .iter()
5778 .any(|field| Self::field_column_is_native_histogram(schema, field))
5779 }
5780
5781 fn field_column_is_float_range(schema: &DFSchemaRef, field_column: &str) -> bool {
5782 Self::field_column_type(schema, field_column).is_some_and(|data_type| {
5783 matches!(
5784 data_type,
5785 ArrowDataType::Dictionary(key_type, value_type)
5786 if key_type.as_ref() == &ArrowDataType::Int64
5787 && value_type.as_ref() == &ArrowDataType::Float64
5788 )
5789 })
5790 }
5791
5792 fn field_columns_are_alternative_samples(
5793 schema: &DFSchemaRef,
5794 field_columns: &[String],
5795 ) -> bool {
5796 Self::alternative_sample_columns(schema, field_columns).is_some()
5797 }
5798
5799 fn alternative_sample_columns<'a>(
5800 schema: &DFSchemaRef,
5801 field_columns: &'a [String],
5802 ) -> Option<(&'a str, &'a str)> {
5803 if field_columns.len() != 2 {
5804 return None;
5805 }
5806
5807 let canonical_float = field_columns.iter().find(|field| {
5808 field.as_str() == greptime_value()
5809 && (Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64)
5810 || Self::field_column_is_float_range(schema, field))
5811 });
5812 let canonical_histogram = field_columns.iter().find(|field| {
5813 field.as_str() == greptime_native_histogram()
5814 && (Self::field_column_is_native_histogram(schema, field)
5815 || Self::field_column_is_native_histogram_range(schema, field))
5816 });
5817 if let (Some(float), Some(histogram)) = (canonical_float, canonical_histogram) {
5818 return Some((float, histogram));
5819 }
5820
5821 let float = field_columns.iter().find(|field| {
5822 field.starts_with(OR_FLOAT_FIELD_PREFIX)
5823 && (Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64)
5824 || Self::field_column_is_float_range(schema, field))
5825 })?;
5826 let histogram = field_columns.iter().find(|field| {
5827 field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)
5828 && (Self::field_column_is_native_histogram(schema, field)
5829 || Self::field_column_is_native_histogram_range(schema, field))
5830 })?;
5831 Some((float, histogram))
5832 }
5833
5834 fn alternative_sample_range_columns<'a>(
5835 schema: &DFSchemaRef,
5836 field_columns: &'a [String],
5837 ) -> Option<(&'a str, &'a str)> {
5838 Self::alternative_sample_columns(schema, field_columns).filter(|(float, histogram)| {
5839 Self::field_column_is_float_range(schema, float)
5840 && Self::field_column_is_native_histogram_range(schema, histogram)
5841 })
5842 }
5843
5844 fn field_column_is_native_histogram_range(schema: &DFSchemaRef, field_column: &str) -> bool {
5845 Self::field_column_type(schema, field_column).is_some_and(|data_type| {
5846 matches!(
5847 data_type,
5848 ArrowDataType::Dictionary(key_type, value_type)
5849 if key_type.as_ref() == &ArrowDataType::Int64
5850 && value_type.as_ref() == &Self::native_histogram_arrow_type()
5851 )
5852 })
5853 }
5854
5855 fn all_field_columns_are_native_histograms(&self, schema: &DFSchemaRef) -> bool {
5856 !self.ctx.field_columns.is_empty()
5857 && self
5858 .ctx
5859 .field_columns
5860 .iter()
5861 .all(|field| Self::field_column_is_native_histogram(schema, field))
5862 }
5863
5864 fn all_field_columns_are_native_histogram_ranges(&self, schema: &DFSchemaRef) -> bool {
5865 !self.ctx.field_columns.is_empty()
5866 && self
5867 .ctx
5868 .field_columns
5869 .iter()
5870 .all(|field| Self::field_column_is_native_histogram_range(schema, field))
5871 }
5872
5873 fn optional_tsid_projection(
5874 schema: &DFSchemaRef,
5875 table_ref: Option<&TableReference>,
5876 keep_tsid: bool,
5877 ) -> Option<DfExpr> {
5878 keep_tsid.then_some(()).and_then(|_| {
5879 schema
5880 .qualified_field_with_name(table_ref, DATA_SCHEMA_TSID_COLUMN_NAME)
5881 .ok()
5882 .map(|field| DfExpr::Column(field.into()))
5883 })
5884 }
5885
5886 fn binary_join_key_columns(
5887 &self,
5888 left_schema: &DFSchemaRef,
5889 right_schema: &DFSchemaRef,
5890 left_context: &PromPlannerContext,
5891 right_context: &PromPlannerContext,
5892 only_join_time_index: bool,
5893 modifier: &Option<BinModifier>,
5894 ) -> Result<(BTreeSet<String>, BTreeSet<String>, bool)> {
5895 let has_tsid = |schema: &DFSchemaRef| {
5896 schema
5897 .fields()
5898 .iter()
5899 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
5900 };
5901 let use_tsid_join = !only_join_time_index
5902 && self.binary_modifier_preserves_tsid_join_key(left_context, right_context, modifier)
5903 && left_context.use_tsid
5904 && right_context.use_tsid
5905 && has_tsid(left_schema)
5906 && has_tsid(right_schema);
5907
5908 let (mut left_tag_columns, mut right_tag_columns) = if use_tsid_join {
5909 (
5910 BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]),
5911 BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]),
5912 )
5913 } else {
5914 if only_join_time_index {
5915 (BTreeSet::new(), BTreeSet::new())
5916 } else {
5917 (
5918 left_context
5919 .tag_columns
5920 .iter()
5921 .cloned()
5922 .collect::<BTreeSet<_>>(),
5923 right_context
5924 .tag_columns
5925 .iter()
5926 .cloned()
5927 .collect::<BTreeSet<_>>(),
5928 )
5929 }
5930 };
5931
5932 if !use_tsid_join
5933 && let Some(modifier) = modifier
5934 && let Some(matching) = &modifier.matching
5935 {
5936 match matching {
5937 LabelModifier::Include(on) => {
5938 let mask = on.labels.iter().cloned().collect::<BTreeSet<_>>();
5939 left_tag_columns = left_tag_columns.intersection(&mask).cloned().collect();
5940 right_tag_columns = right_tag_columns.intersection(&mask).cloned().collect();
5941 }
5942 LabelModifier::Exclude(ignoring) => {
5943 for label in &ignoring.labels {
5944 let _ = left_tag_columns.remove(label);
5945 let _ = right_tag_columns.remove(label);
5946 }
5947 }
5948 }
5949 }
5950
5951 let force_empty_join =
5952 !use_tsid_join && !only_join_time_index && left_tag_columns != right_tag_columns;
5953 if force_empty_join {
5954 let common_tag_columns = left_tag_columns
5955 .intersection(&right_tag_columns)
5956 .cloned()
5957 .collect::<BTreeSet<_>>();
5958 left_tag_columns = common_tag_columns.clone();
5959 right_tag_columns = common_tag_columns;
5960 }
5961
5962 Ok((left_tag_columns, right_tag_columns, force_empty_join))
5963 }
5964
5965 fn binary_result_labels(
5972 left_context: &PromPlannerContext,
5973 right_context: &PromPlannerContext,
5974 modifier: &Option<BinModifier>,
5975 ) -> Option<Vec<(bool, String)>> {
5976 let modifier = modifier.as_ref()?;
5977 let (many_is_left, include) = match &modifier.card {
5978 VectorMatchCardinality::OneToOne => (true, None),
5979 VectorMatchCardinality::ManyToOne(labels) => (true, Some(labels)),
5980 VectorMatchCardinality::OneToMany(labels) => (false, Some(labels)),
5981 VectorMatchCardinality::ManyToMany => return None,
5983 };
5984
5985 let Some(include) = include else {
5986 let matching = modifier.matching.as_ref()?;
5987 let reduced = |keep: bool, labels: &BTreeSet<&String>| {
5988 left_context
5989 .tag_columns
5990 .iter()
5991 .filter(|tag| labels.contains(tag) == keep)
5992 .map(|tag| (true, tag.clone()))
5993 .collect()
5994 };
5995 return Some(match matching {
5996 LabelModifier::Include(on) => reduced(true, &on.labels.iter().collect()),
5997 LabelModifier::Exclude(ignoring) => {
5998 reduced(false, &ignoring.labels.iter().collect())
5999 }
6000 });
6001 };
6002
6003 let (many_context, one_context) = if many_is_left {
6004 (left_context, right_context)
6005 } else {
6006 (right_context, left_context)
6007 };
6008 let include = include.labels.iter().collect::<BTreeSet<_>>();
6009 let mut labels = many_context
6012 .tag_columns
6013 .iter()
6014 .filter(|tag| !include.contains(tag))
6015 .map(|tag| (many_is_left, tag.clone()))
6016 .collect::<Vec<_>>();
6017 labels.extend(
6018 include
6019 .into_iter()
6020 .filter(|label| one_context.tag_columns.contains(label))
6021 .map(|label| (!many_is_left, label.clone())),
6022 );
6023 Some(labels)
6024 }
6025
6026 fn binary_result_label_projection(
6028 schema: &DFSchemaRef,
6029 left_table_ref: &TableReference,
6030 right_table_ref: &TableReference,
6031 left_context: &PromPlannerContext,
6032 right_context: &PromPlannerContext,
6033 labels: Vec<(bool, String)>,
6034 ) -> Result<BinaryResultLabels> {
6035 let mut exprs = Vec::with_capacity(labels.len());
6036 let mut names = Vec::with_capacity(labels.len());
6037 let mut aggregation_field_labels = Vec::new();
6038 let mut sources = HashSet::new();
6039 for (from_left, label) in labels {
6040 let (table_ref, context) = if from_left {
6041 (left_table_ref, left_context)
6042 } else {
6043 (right_table_ref, right_context)
6044 };
6045 let field = schema
6046 .qualified_field_with_name(Some(table_ref), &label)
6047 .context(DataFusionPlanningSnafu)?;
6048 exprs.push(DfExpr::Column(field.into()));
6049 if context.aggregation_field_labels.contains(&label) {
6050 aggregation_field_labels.push(label.clone());
6051 }
6052 sources.insert(from_left);
6053 names.push(label);
6054 }
6055
6056 let source_context = match sources.into_iter().collect::<Vec<_>>().as_slice() {
6060 [true] => Some((left_table_ref, left_context)),
6061 [false] => Some((right_table_ref, right_context)),
6062 _ => None,
6063 };
6064 let tsid = source_context
6065 .filter(|(_, context)| {
6066 context.use_tsid
6067 && context.tag_columns.len() == names.len()
6068 && context.tag_columns.iter().all(|tag| names.contains(tag))
6069 })
6070 .and_then(|(table_ref, _)| {
6071 Self::optional_tsid_projection(schema, Some(table_ref), true)
6072 });
6073
6074 Ok(BinaryResultLabels {
6075 exprs,
6076 names,
6077 aggregation_field_labels,
6078 tsid,
6079 })
6080 }
6081
6082 fn binary_result_labels_may_repeat(
6086 left_context: &PromPlannerContext,
6087 right_context: &PromPlannerContext,
6088 modifier: &Option<BinModifier>,
6089 ) -> bool {
6090 let Some(modifier) = modifier else {
6091 return false;
6092 };
6093 match &modifier.card {
6094 VectorMatchCardinality::OneToOne => match &modifier.matching {
6095 None => false,
6096 Some(LabelModifier::Include(on)) => {
6097 let on = on.labels.iter().collect::<BTreeSet<_>>();
6098 !left_context.tag_columns.iter().all(|tag| on.contains(tag))
6099 }
6100 Some(LabelModifier::Exclude(ignoring)) => ignoring
6101 .labels
6102 .iter()
6103 .any(|label| left_context.tag_columns.contains(label)),
6104 },
6105 VectorMatchCardinality::ManyToOne(include) => include
6106 .labels
6107 .iter()
6108 .any(|label| left_context.tag_columns.contains(label)),
6109 VectorMatchCardinality::OneToMany(include) => include
6110 .labels
6111 .iter()
6112 .any(|label| right_context.tag_columns.contains(label)),
6113 VectorMatchCardinality::ManyToMany => false,
6114 }
6115 }
6116
6117 fn assert_unique_match_group(
6120 plan: LogicalPlan,
6121 group_exprs: Vec<DfExpr>,
6122 group_labels: Vec<String>,
6123 time_index_expr: DfExpr,
6124 violation: MatchGroupViolation,
6125 ) -> Result<LogicalPlan> {
6126 let mut partition_by = group_exprs.clone();
6127 partition_by.push(time_index_expr);
6128 let occupied_column_names = plan
6131 .schema()
6132 .fields()
6133 .iter()
6134 .map(|field| field.name().as_str())
6135 .collect::<HashSet<_>>();
6136 let mut next_suffix = 0;
6137 let count_column = loop {
6138 let name = match next_suffix {
6139 0 => MATCH_GROUP_COUNT_COLUMN.to_string(),
6140 suffix => format!("{MATCH_GROUP_COUNT_COLUMN}_{suffix}"),
6141 };
6142 next_suffix += 1;
6143 if !occupied_column_names.contains(name.as_str()) {
6144 break name;
6145 }
6146 };
6147 let count = DfExpr::WindowFunction(Box::new(WindowFunction {
6148 fun: WindowFunctionDefinition::AggregateUDF(count_udaf()),
6149 params: WindowFunctionParams {
6150 args: vec![lit(1i64)],
6151 partition_by,
6152 order_by: vec![],
6153 window_frame: WindowFrame::new(None),
6154 null_treatment: None,
6155 distinct: false,
6156 filter: None,
6157 },
6158 }))
6159 .alias(count_column.as_str());
6160
6161 let output_exprs = plan
6162 .schema()
6163 .iter()
6164 .map(|(qualifier, field)| DfExpr::Column(Column::new(qualifier.cloned(), field.name())))
6165 .collect::<Vec<_>>();
6166 let assert_expr = DfExpr::ScalarFunction(ScalarFunction {
6167 func: Arc::new(UniqueMatchGroup::scalar_udf(group_labels, violation)),
6168 args: std::iter::once(col(count_column.as_str()))
6169 .chain(group_exprs)
6170 .collect(),
6171 });
6172
6173 LogicalPlanBuilder::from(plan)
6174 .window(vec![count])
6175 .context(DataFusionPlanningSnafu)?
6176 .filter(assert_expr)
6177 .context(DataFusionPlanningSnafu)?
6178 .project(output_exprs)
6179 .context(DataFusionPlanningSnafu)?
6180 .build()
6181 .context(DataFusionPlanningSnafu)
6182 }
6183
6184 fn binary_modifier_preserves_tsid_join_key(
6185 &self,
6186 left_context: &PromPlannerContext,
6187 right_context: &PromPlannerContext,
6188 modifier: &Option<BinModifier>,
6189 ) -> bool {
6190 let Some(modifier) = modifier else {
6191 return true;
6192 };
6193
6194 if !matches!(modifier.card, VectorMatchCardinality::OneToOne) {
6195 return false;
6196 }
6197
6198 match &modifier.matching {
6199 None => true,
6200 Some(LabelModifier::Exclude(ignoring)) => ignoring.labels.iter().all(|label| {
6201 !left_context.tag_columns.contains(label)
6202 && !right_context.tag_columns.contains(label)
6203 }),
6204 Some(LabelModifier::Include(on)) => {
6205 let on_labels = on.labels.iter().cloned().collect::<BTreeSet<_>>();
6206 let left_labels = left_context
6207 .tag_columns
6208 .iter()
6209 .cloned()
6210 .collect::<BTreeSet<_>>();
6211 let right_labels = right_context
6212 .tag_columns
6213 .iter()
6214 .cloned()
6215 .collect::<BTreeSet<_>>();
6216
6217 on_labels == left_labels && on_labels == right_labels
6218 }
6219 }
6220 }
6221
6222 #[allow(clippy::too_many_arguments)]
6225 fn join_on_non_field_columns(
6226 &self,
6227 left: LogicalPlan,
6228 right: LogicalPlan,
6229 left_table_ref: TableReference,
6230 right_table_ref: TableReference,
6231 left_time_index_column: Option<String>,
6232 right_time_index_column: Option<String>,
6233 only_join_time_index: bool,
6234 modifier: &Option<BinModifier>,
6235 left_context: &PromPlannerContext,
6236 right_context: &PromPlannerContext,
6237 ) -> Result<LogicalPlan> {
6238 let (mut left_tag_columns, mut right_tag_columns, mut force_empty_join) = self
6239 .binary_join_key_columns(
6240 left.schema(),
6241 right.schema(),
6242 left_context,
6243 right_context,
6244 only_join_time_index,
6245 modifier,
6246 )?;
6247 let use_tsid_join = !only_join_time_index
6248 && !force_empty_join
6249 && left_tag_columns == BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()])
6250 && right_tag_columns == BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]);
6251 let (left, right, left_matched_tags, right_matched_tags) = if !only_join_time_index
6252 && !use_tsid_join
6253 && Self::only_temporality_match_label_mismatches(left_context, right_context, modifier)
6254 {
6255 let mut aligned_left_context = left_context.clone();
6256 let mut aligned_right_context = right_context.clone();
6257 let (left, right, _) = Self::align_temporality_match_column(
6258 left,
6259 right,
6260 &mut aligned_left_context,
6261 &mut aligned_right_context,
6262 )?;
6263 (left_tag_columns, right_tag_columns, force_empty_join) = self
6264 .binary_join_key_columns(
6265 left.schema(),
6266 right.schema(),
6267 &aligned_left_context,
6268 &aligned_right_context,
6269 false,
6270 modifier,
6271 )?;
6272 (
6273 left,
6274 right,
6275 aligned_left_context.tag_columns,
6276 aligned_right_context.tag_columns,
6277 )
6278 } else {
6279 (
6280 left,
6281 right,
6282 left_context.tag_columns.clone(),
6283 right_context.tag_columns.clone(),
6284 )
6285 };
6286
6287 let checked_matching = !only_join_time_index
6290 && !force_empty_join
6291 && !use_tsid_join
6292 && !matches!(
6293 modifier.as_ref().map(|modifier| &modifier.card),
6294 Some(VectorMatchCardinality::ManyToMany)
6295 );
6296 let one_side_is_left = matches!(
6298 modifier.as_ref().map(|modifier| &modifier.card),
6299 Some(VectorMatchCardinality::OneToMany(_))
6300 );
6301 let (left, right) = if !checked_matching {
6302 (left, right)
6303 } else if one_side_is_left {
6304 (
6305 Self::assert_unique_one_side(
6306 left,
6307 &left_tag_columns,
6308 &left_matched_tags,
6309 left_time_index_column.as_deref(),
6310 true,
6311 )?,
6312 right,
6313 )
6314 } else {
6315 (
6316 left,
6317 Self::assert_unique_one_side(
6318 right,
6319 &right_tag_columns,
6320 &right_matched_tags,
6321 right_time_index_column.as_deref(),
6322 false,
6323 )?,
6324 )
6325 };
6326
6327 if let (Some(left_time_index_column), Some(right_time_index_column)) = (
6329 left_time_index_column.clone(),
6330 right_time_index_column.clone(),
6331 ) {
6332 left_tag_columns.insert(left_time_index_column);
6333 right_tag_columns.insert(right_time_index_column);
6334 }
6335
6336 let right = LogicalPlanBuilder::from(right)
6337 .alias(right_table_ref.clone())
6338 .context(DataFusionPlanningSnafu)?
6339 .build()
6340 .context(DataFusionPlanningSnafu)?;
6341
6342 let join_plan = LogicalPlanBuilder::from(left)
6344 .alias(left_table_ref.clone())
6345 .context(DataFusionPlanningSnafu)?
6346 .join_detailed(
6347 right,
6348 JoinType::Inner,
6349 (
6350 left_tag_columns
6351 .into_iter()
6352 .map(Column::from_name)
6353 .collect::<Vec<_>>(),
6354 right_tag_columns
6355 .into_iter()
6356 .map(Column::from_name)
6357 .collect::<Vec<_>>(),
6358 ),
6359 force_empty_join.then_some(lit(false)),
6360 NullEquality::NullEqualsNull,
6361 )
6362 .context(DataFusionPlanningSnafu)?
6363 .build()
6364 .context(DataFusionPlanningSnafu)?;
6365
6366 let labels = checked_matching
6369 .then(|| Self::binary_result_labels(left_context, right_context, modifier))
6370 .flatten()
6371 .filter(|_| {
6372 Self::binary_result_labels_may_repeat(left_context, right_context, modifier)
6373 });
6374 let Some(labels) = labels else {
6375 return Ok(join_plan);
6376 };
6377 let schema = join_plan.schema().clone();
6378 let group_exprs = labels
6379 .iter()
6380 .map(|(from_left, label)| {
6381 let table_ref = if *from_left {
6382 &left_table_ref
6383 } else {
6384 &right_table_ref
6385 };
6386 schema
6387 .qualified_field_with_name(Some(table_ref), label)
6388 .map(|field| DfExpr::Column(field.into()))
6389 .context(DataFusionPlanningSnafu)
6390 })
6391 .collect::<Result<Vec<_>>>()?;
6392 let time_index_expr = left_time_index_column
6393 .as_deref()
6394 .map(|column| {
6395 schema
6396 .qualified_field_with_name(Some(&left_table_ref), column)
6397 .map(|field| DfExpr::Column(field.into()))
6398 .context(DataFusionPlanningSnafu)
6399 })
6400 .transpose()?
6401 .with_context(|| UnexpectedPlanExprSnafu {
6402 desc: "vector matching on a plan without a time index",
6403 })?;
6404 let violation = if matches!(
6405 modifier.as_ref().map(|modifier| &modifier.card),
6406 Some(VectorMatchCardinality::OneToOne)
6407 ) {
6408 MatchGroupViolation::ImplicitManyToOne
6409 } else {
6410 MatchGroupViolation::AmbiguousGroupLabels
6411 };
6412 Self::assert_unique_match_group(
6413 join_plan,
6414 group_exprs,
6415 labels.into_iter().map(|(_, label)| label).collect(),
6416 time_index_expr,
6417 violation,
6418 )
6419 }
6420
6421 fn assert_unique_one_side(
6423 plan: LogicalPlan,
6424 join_keys: &BTreeSet<String>,
6425 tag_columns: &[String],
6426 time_index_column: Option<&str>,
6427 one_side_is_left: bool,
6428 ) -> Result<LogicalPlan> {
6429 let Some(time_index_column) = time_index_column else {
6430 return Ok(plan);
6431 };
6432 if tag_columns.iter().all(|tag| join_keys.contains(tag)) {
6433 return Ok(plan);
6434 }
6435
6436 let group_labels = join_keys.iter().cloned().collect::<Vec<_>>();
6437 let group_exprs = group_labels
6438 .iter()
6439 .map(|label| DfExpr::Column(Column::from_name(label)))
6440 .collect();
6441 Self::assert_unique_match_group(
6442 plan,
6443 group_exprs,
6444 group_labels,
6445 DfExpr::Column(Column::from_name(time_index_column)),
6446 MatchGroupViolation::DuplicateOnOneSide { one_side_is_left },
6447 )
6448 }
6449
6450 fn selected_binary_match_labels(
6451 left_context: &PromPlannerContext,
6452 right_context: &PromPlannerContext,
6453 modifier: &Option<BinModifier>,
6454 ) -> BTreeSet<String> {
6455 let mut labels = left_context
6456 .tag_columns
6457 .iter()
6458 .chain(&right_context.tag_columns)
6459 .cloned()
6460 .collect::<BTreeSet<_>>();
6461 if let Some(matching) = modifier
6462 .as_ref()
6463 .and_then(|modifier| modifier.matching.as_ref())
6464 {
6465 match matching {
6466 LabelModifier::Include(on) => {
6467 labels = on
6468 .labels
6469 .iter()
6470 .filter(|label| {
6471 left_context.tag_columns.contains(label)
6472 || right_context.tag_columns.contains(label)
6473 })
6474 .cloned()
6475 .collect();
6476 }
6477 LabelModifier::Exclude(ignoring) => {
6478 for label in &ignoring.labels {
6479 labels.remove(label);
6480 }
6481 }
6482 }
6483 }
6484 labels
6485 }
6486
6487 fn only_temporality_match_label_mismatches(
6488 left_context: &PromPlannerContext,
6489 right_context: &PromPlannerContext,
6490 modifier: &Option<BinModifier>,
6491 ) -> bool {
6492 let mut mismatches =
6493 Self::selected_binary_match_labels(left_context, right_context, modifier)
6494 .into_iter()
6495 .filter(|label| {
6496 left_context.tag_columns.contains(label)
6497 != right_context.tag_columns.contains(label)
6498 });
6499 matches!(
6500 (mismatches.next(), mismatches.next()),
6501 (Some(label), None) if label == OTLP_AGGREGATION_TEMPORALITY_LABEL
6502 )
6503 }
6504
6505 fn align_temporality_match_column(
6506 mut left: LogicalPlan,
6507 mut right: LogicalPlan,
6508 left_context: &mut PromPlannerContext,
6509 right_context: &mut PromPlannerContext,
6510 ) -> Result<(LogicalPlan, LogicalPlan, bool)> {
6511 let marker = OTLP_AGGREGATION_TEMPORALITY_LABEL;
6512 let left_has_marker = left_context.tag_columns.iter().any(|tag| tag == marker);
6513 let (data_type, value_type, add_to_left) = {
6514 let (present, add_to_left) = if left_has_marker {
6515 (&left, false)
6516 } else {
6517 (&right, true)
6518 };
6519 let data_type = present
6520 .schema()
6521 .fields()
6522 .iter()
6523 .find(|field| field.name() == marker)
6524 .map(|field| field.data_type().clone())
6525 .with_context(|| ColumnNotFoundSnafu {
6526 col: marker.to_string(),
6527 })?;
6528 let value_type = Self::string_value_data_type(&data_type)
6529 .cloned()
6530 .with_context(|| UnexpectedPlanExprSnafu {
6531 desc: format!("temporality match label {marker} must be a string"),
6532 })?;
6533 (data_type, value_type, add_to_left)
6534 };
6535 let null = Self::string_scalar_value(&value_type, None).with_context(|| {
6536 UnexpectedPlanExprSnafu {
6537 desc: format!("temporality match label {marker} must be a string"),
6538 }
6539 })?;
6540 let add_marker = |plan: LogicalPlan| {
6541 let visible = plan
6542 .schema()
6543 .iter()
6544 .map(|(qualifier, field)| {
6545 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6546 })
6547 .collect::<Vec<_>>();
6548 LogicalPlanBuilder::from(plan)
6549 .project(
6550 visible
6551 .into_iter()
6552 .chain([DfExpr::Literal(null, None).alias(marker)]),
6553 )
6554 .context(DataFusionPlanningSnafu)?
6555 .build()
6556 .context(DataFusionPlanningSnafu)
6557 };
6558 if data_type != value_type {
6559 let present = if add_to_left { &mut right } else { &mut left };
6560 let visible = present
6561 .schema()
6562 .iter()
6563 .map(|(qualifier, field)| {
6564 let column =
6565 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()));
6566 if field.name() == marker {
6567 DfExpr::Cast(Cast::new(Box::new(column), value_type.clone()))
6568 .alias_qualified(qualifier.cloned(), field.name().clone())
6569 } else {
6570 column
6571 }
6572 })
6573 .collect::<Vec<_>>();
6574 *present = LogicalPlanBuilder::from(present.clone())
6575 .project(visible)
6576 .context(DataFusionPlanningSnafu)?
6577 .build()
6578 .context(DataFusionPlanningSnafu)?;
6579 }
6580 if add_to_left {
6581 left = add_marker(left)?;
6582 left_context.tag_columns.push(marker.to_string());
6583 } else {
6584 right = add_marker(right)?;
6585 right_context.tag_columns.push(marker.to_string());
6586 }
6587 Ok((left, right, add_to_left))
6588 }
6589
6590 fn normalized_match_key_expr(
6591 label: &str,
6592 field: Option<(Option<TableReference>, ArrowDataType)>,
6593 value_type: &ArrowDataType,
6594 internal_name: &str,
6595 ) -> DfExpr {
6596 let empty = Self::string_scalar_value(value_type, Some(String::new()))
6597 .expect("match label value type is a string");
6598 let expr = if let Some((qualifier, data_type)) = field {
6599 let column = DfExpr::Column(Column::new(qualifier, label));
6600 let column = if &data_type == value_type {
6601 column
6602 } else {
6603 DfExpr::Cast(Cast::new(Box::new(column), value_type.clone()))
6604 };
6605 DfExpr::ScalarFunction(ScalarFunction {
6606 func: coalesce(),
6607 args: vec![column, DfExpr::Literal(empty, None)],
6608 })
6609 } else {
6610 DfExpr::Literal(empty, None)
6611 };
6612 expr.alias(internal_name)
6613 }
6614
6615 fn is_zero_row_empty_relation(plan: &LogicalPlan) -> bool {
6616 matches!(plan, LogicalPlan::EmptyRelation(relation) if !relation.produce_one_row)
6619 }
6620
6621 fn set_op_on_non_field_columns(
6623 &mut self,
6624 mut left: LogicalPlan,
6625 mut right: LogicalPlan,
6626 left_context: PromPlannerContext,
6627 right_context: PromPlannerContext,
6628 op: TokenType,
6629 modifier: &Option<BinModifier>,
6630 ) -> Result<LogicalPlan> {
6631 let left_tag_col_set = left_context
6632 .tag_columns
6633 .iter()
6634 .cloned()
6635 .collect::<HashSet<_>>();
6636 let right_tag_col_set = right_context
6637 .tag_columns
6638 .iter()
6639 .cloned()
6640 .collect::<HashSet<_>>();
6641
6642 if matches!(op.id(), token::T_LOR) {
6643 return self.or_operator(
6644 left,
6645 right,
6646 left_tag_col_set,
6647 right_tag_col_set,
6648 left_context,
6649 right_context,
6650 modifier,
6651 );
6652 }
6653
6654 if let Some(modifier) = modifier {
6655 ensure!(
6656 matches!(
6657 modifier.card,
6658 VectorMatchCardinality::OneToOne | VectorMatchCardinality::ManyToMany
6659 ),
6660 UnsupportedVectorMatchSnafu {
6661 name: modifier.card.clone(),
6662 },
6663 );
6664 }
6665
6666 let output_context = left_context.clone();
6667 let visible_left_schema = left.schema().clone();
6668 let mut left_context = left_context;
6669 let mut right_context = right_context;
6670 let added_marker_to_left = if Self::only_temporality_match_label_mismatches(
6671 &left_context,
6672 &right_context,
6673 modifier,
6674 ) {
6675 let aligned = Self::align_temporality_match_column(
6676 left,
6677 right,
6678 &mut left_context,
6679 &mut right_context,
6680 )?;
6681 left = aligned.0;
6682 right = aligned.1;
6683 aligned.2
6684 } else {
6685 false
6686 };
6687
6688 let mut left_tag_col_set = left_context
6689 .tag_columns
6690 .iter()
6691 .cloned()
6692 .collect::<BTreeSet<_>>();
6693 let mut right_tag_col_set = right_context
6694 .tag_columns
6695 .iter()
6696 .cloned()
6697 .collect::<BTreeSet<_>>();
6698 if let Some(matching) = modifier
6699 .as_ref()
6700 .and_then(|modifier| modifier.matching.as_ref())
6701 {
6702 match matching {
6703 LabelModifier::Include(on) => {
6704 let mask = on.labels.iter().cloned().collect::<BTreeSet<_>>();
6705 left_tag_col_set = left_tag_col_set.intersection(&mask).cloned().collect();
6706 right_tag_col_set = right_tag_col_set.intersection(&mask).cloned().collect();
6707 }
6708 LabelModifier::Exclude(ignoring) => {
6709 for label in &ignoring.labels {
6710 let _ = left_tag_col_set.remove(label);
6711 let _ = right_tag_col_set.remove(label);
6712 }
6713 }
6714 }
6715 }
6716 ensure!(
6717 left_tag_col_set == right_tag_col_set,
6718 CombineTableColumnMismatchSnafu {
6719 left: left_tag_col_set.iter().cloned().collect::<Vec<_>>(),
6720 right: right_tag_col_set.iter().cloned().collect::<Vec<_>>(),
6721 }
6722 );
6723
6724 let left_time_index = left_context.time_index_column.clone().unwrap();
6725 let right_time_index = right_context.time_index_column.clone().unwrap();
6726
6727 if left_context.time_index_column != right_context.time_index_column {
6729 let right_project_exprs = right
6730 .schema()
6731 .fields()
6732 .iter()
6733 .map(|field| {
6734 if field.name() == &right_time_index {
6735 DfExpr::Column(Column::from_name(&right_time_index)).alias(&left_time_index)
6736 } else {
6737 DfExpr::Column(Column::from_name(field.name()))
6738 }
6739 })
6740 .collect::<Vec<_>>();
6741
6742 right = LogicalPlanBuilder::from(right)
6743 .project(right_project_exprs)
6744 .context(DataFusionPlanningSnafu)?
6745 .build()
6746 .context(DataFusionPlanningSnafu)?;
6747 }
6748
6749 let join_keys = left_tag_col_set
6750 .into_iter()
6751 .chain([left_time_index])
6752 .collect::<Vec<_>>();
6753
6754 ensure!(
6755 left_context.field_columns.len() == 1
6756 || Self::field_columns_are_alternative_samples(
6757 left.schema(),
6758 &left_context.field_columns,
6759 ),
6760 MultiFieldsNotSupportedSnafu {
6761 operator: "AND/UNLESS operator"
6762 }
6763 );
6764 let result = match op.id() {
6767 token::T_LAND => LogicalPlanBuilder::from(left)
6768 .distinct()
6769 .context(DataFusionPlanningSnafu)?
6770 .join_detailed(
6771 right,
6772 JoinType::LeftSemi,
6773 (join_keys.clone(), join_keys),
6774 None,
6775 NullEquality::NullEqualsNull,
6776 )
6777 .context(DataFusionPlanningSnafu)?
6778 .build()
6779 .context(DataFusionPlanningSnafu),
6780 token::T_LUNLESS => LogicalPlanBuilder::from(left)
6781 .distinct()
6782 .context(DataFusionPlanningSnafu)?
6783 .join_detailed(
6784 right,
6785 JoinType::LeftAnti,
6786 (join_keys.clone(), join_keys),
6787 None,
6788 NullEquality::NullEqualsNull,
6789 )
6790 .context(DataFusionPlanningSnafu)?
6791 .build()
6792 .context(DataFusionPlanningSnafu),
6793 token::T_LOR => {
6794 unreachable!()
6797 }
6798 _ => UnexpectedTokenSnafu { token: op }.fail(),
6799 }?;
6800 let result = if added_marker_to_left {
6801 LogicalPlanBuilder::from(result)
6802 .project(visible_left_schema.iter().map(|(qualifier, field)| {
6803 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6804 }))
6805 .context(DataFusionPlanningSnafu)?
6806 .build()
6807 .context(DataFusionPlanningSnafu)?
6808 } else {
6809 result
6810 };
6811
6812 self.ctx = output_context;
6815 Ok(result)
6816 }
6817
6818 fn string_value_data_type(data_type: &ArrowDataType) -> Option<&ArrowDataType> {
6819 match data_type {
6820 data_type if data_type.is_string() => Some(data_type),
6821 ArrowDataType::Dictionary(_, value_type) if value_type.is_string() => Some(value_type),
6822 _ => None,
6823 }
6824 }
6825
6826 fn string_scalar_value(
6827 data_type: &ArrowDataType,
6828 value: Option<String>,
6829 ) -> Option<ScalarValue> {
6830 match data_type {
6831 ArrowDataType::Utf8 => Some(ScalarValue::Utf8(value)),
6832 ArrowDataType::LargeUtf8 => Some(ScalarValue::LargeUtf8(value)),
6833 ArrowDataType::Utf8View => Some(ScalarValue::Utf8View(value)),
6834 ArrowDataType::Dictionary(key_type, value_type) => Some(ScalarValue::Dictionary(
6835 key_type.clone(),
6836 Box::new(Self::string_scalar_value(value_type, value)?),
6837 )),
6838 _ => None,
6839 }
6840 }
6841
6842 fn common_label_data_type(
6843 left: Option<&ArrowDataType>,
6844 right: Option<&ArrowDataType>,
6845 ) -> Option<ArrowDataType> {
6846 match (left, right) {
6847 (Some(left), Some(right)) if left == right => {
6848 Self::string_value_data_type(left).map(|_| left.clone())
6849 }
6850 (Some(left), Some(right)) => {
6851 let left_value_type = Self::string_value_data_type(left)?;
6852 let right_value_type = Self::string_value_data_type(right)?;
6853 match (left_value_type, right_value_type) {
6856 (left, right) if left == right => Some(left.clone()),
6857 (ArrowDataType::LargeUtf8, _) | (_, ArrowDataType::LargeUtf8) => {
6858 Some(ArrowDataType::LargeUtf8)
6859 }
6860 (ArrowDataType::Utf8View, ArrowDataType::Utf8View) => {
6861 Some(ArrowDataType::Utf8View)
6862 }
6863 _ => Some(ArrowDataType::Utf8),
6864 }
6865 }
6866 (Some(data_type), None) | (None, Some(data_type)) => {
6867 Self::string_value_data_type(data_type).cloned()
6868 }
6869 (None, None) => Some(ArrowDataType::Utf8),
6870 }
6871 }
6872
6873 #[allow(clippy::too_many_arguments)]
6875 fn or_operator(
6876 &mut self,
6877 left: LogicalPlan,
6878 right: LogicalPlan,
6879 left_tag_cols_set: HashSet<String>,
6880 right_tag_cols_set: HashSet<String>,
6881 left_context: PromPlannerContext,
6882 right_context: PromPlannerContext,
6883 modifier: &Option<BinModifier>,
6884 ) -> Result<LogicalPlan> {
6885 let left_is_empty = Self::is_zero_row_empty_relation(&left);
6886 let right_is_empty = Self::is_zero_row_empty_relation(&right);
6887 match (left_is_empty, right_is_empty) {
6888 (true, false) => {
6889 self.ctx = right_context;
6890 return Ok(right);
6891 }
6892 (false, true) => {
6893 self.ctx = left_context;
6894 return Ok(left);
6895 }
6896 (true, true) => {
6897 self.ctx = left_context;
6898 return Ok(left);
6899 }
6900 (false, false) => {}
6901 }
6902
6903 ensure!(
6904 !left.schema().fields().is_empty() && !right.schema().fields().is_empty(),
6905 UnexpectedPlanExprSnafu {
6906 desc: "OR operator input has zero columns",
6907 }
6908 );
6909 let left_has_alternative_samples =
6910 Self::field_columns_are_alternative_samples(left.schema(), &left_context.field_columns);
6911 let right_has_alternative_samples = Self::field_columns_are_alternative_samples(
6912 right.schema(),
6913 &right_context.field_columns,
6914 );
6915 ensure!(
6916 left_context.field_columns.len() == 1 || left_has_alternative_samples,
6917 MultiFieldsNotSupportedSnafu {
6918 operator: "OR operator"
6919 }
6920 );
6921 ensure!(
6922 right_context.field_columns.len() == 1 || right_has_alternative_samples,
6923 MultiFieldsNotSupportedSnafu {
6924 operator: "OR operator"
6925 }
6926 );
6927
6928 let all_tags = left_tag_cols_set
6930 .union(&right_tag_cols_set)
6931 .cloned()
6932 .collect::<HashSet<_>>();
6933 let left_qualifier = left.schema().qualified_field(0).0.cloned();
6934 let right_qualifier = right.schema().qualified_field(0).0.cloned();
6935 let left_qualifier_string = left_qualifier
6936 .as_ref()
6937 .map(|l| l.to_string())
6938 .unwrap_or_default();
6939 let right_qualifier_string = right_qualifier
6940 .as_ref()
6941 .map(|r| r.to_string())
6942 .unwrap_or_default();
6943 let left_time_index_column =
6944 left_context
6945 .time_index_column
6946 .clone()
6947 .with_context(|| TimeIndexNotFoundSnafu {
6948 table: left_qualifier_string.clone(),
6949 })?;
6950 let right_time_index_column =
6951 right_context
6952 .time_index_column
6953 .clone()
6954 .with_context(|| TimeIndexNotFoundSnafu {
6955 table: right_qualifier_string.clone(),
6956 })?;
6957 let native_histogram_type = Self::native_histogram_arrow_type();
6958 let is_numeric = |data_type: &ArrowDataType| {
6959 matches!(
6960 data_type,
6961 ArrowDataType::Int8
6962 | ArrowDataType::Int16
6963 | ArrowDataType::Int32
6964 | ArrowDataType::Int64
6965 | ArrowDataType::UInt8
6966 | ArrowDataType::UInt16
6967 | ArrowDataType::UInt32
6968 | ArrowDataType::UInt64
6969 | ArrowDataType::Float32
6970 | ArrowDataType::Float64
6971 )
6972 };
6973 let left_fields = left_context
6974 .field_columns
6975 .iter()
6976 .map(|name| {
6977 left.schema()
6978 .iter()
6979 .find(|(_, field)| field.name() == name)
6980 .map(|(qualifier, field)| {
6981 (name.clone(), qualifier.cloned(), field.data_type().clone())
6982 })
6983 .with_context(|| ColumnNotFoundSnafu { col: name.clone() })
6984 })
6985 .collect::<Result<Vec<_>>>()?;
6986 let right_fields = right_context
6987 .field_columns
6988 .iter()
6989 .map(|name| {
6990 right
6991 .schema()
6992 .iter()
6993 .find(|(_, field)| field.name() == name)
6994 .map(|(qualifier, field)| {
6995 (name.clone(), qualifier.cloned(), field.data_type().clone())
6996 })
6997 .with_context(|| ColumnNotFoundSnafu { col: name.clone() })
6998 })
6999 .collect::<Result<Vec<_>>>()?;
7000 let left_field = &left_fields[0];
7001 let right_field = &right_fields[0];
7002 let left_field_col = &left_field.0;
7003 let right_field_col = &right_field.0;
7004 let fields_are_samples = |fields: &[(String, Option<TableReference>, ArrowDataType)]| {
7005 fields.iter().all(|(_, _, data_type)| {
7006 is_numeric(data_type) || data_type == &native_histogram_type
7007 })
7008 };
7009 let mixed_sample_types = if left_has_alternative_samples || right_has_alternative_samples {
7010 if !fields_are_samples(&left_fields) || !fields_are_samples(&right_fields) {
7011 return UnexpectedPlanExprSnafu {
7012 desc: format!(
7013 "OR value fields have incompatible types: {:?} and {:?}",
7014 left_fields
7015 .iter()
7016 .map(|(_, _, data_type)| data_type)
7017 .collect::<Vec<_>>(),
7018 right_fields
7019 .iter()
7020 .map(|(_, _, data_type)| data_type)
7021 .collect::<Vec<_>>()
7022 ),
7023 }
7024 .fail();
7025 }
7026 true
7027 } else {
7028 (left_field.2 == native_histogram_type && is_numeric(&right_field.2))
7029 || (right_field.2 == native_histogram_type && is_numeric(&left_field.2))
7030 };
7031 let target_field_type = if mixed_sample_types {
7032 ArrowDataType::Float64
7035 } else if left_field.2 == right_field.2 {
7036 left_field.2.clone()
7037 } else if is_numeric(&left_field.2) && is_numeric(&right_field.2) {
7038 ArrowDataType::Float64
7039 } else {
7040 return UnexpectedPlanExprSnafu {
7041 desc: format!(
7042 "OR value fields have incompatible types: {:?} and {:?}",
7043 left_field.2, right_field.2
7044 ),
7045 }
7046 .fail();
7047 };
7048 let (mixed_float_field_col, mixed_histogram_field_col) = if mixed_sample_types {
7049 let mut reserved_names = left
7050 .schema()
7051 .fields()
7052 .iter()
7053 .chain(right.schema().fields().iter())
7054 .map(|field| field.name().clone())
7055 .collect::<HashSet<_>>();
7056 for (name, _, _) in left_fields.iter().chain(&right_fields) {
7057 reserved_names.remove(name);
7058 }
7059 reserved_names.extend(all_tags.iter().cloned());
7060 let unique_name = |prefix: &str, reserved_names: &mut HashSet<String>| {
7061 let mut index = 0;
7062 loop {
7063 let name = format!("{prefix}{index}");
7064 index += 1;
7065 if reserved_names.insert(name.clone()) {
7066 break name;
7067 }
7068 }
7069 };
7070 let float_field = unique_name(OR_FLOAT_FIELD_PREFIX, &mut reserved_names);
7071 let histogram_field = unique_name(OR_HISTOGRAM_FIELD_PREFIX, &mut reserved_names);
7072 (float_field, histogram_field)
7073 } else {
7074 (left_field_col.clone(), String::new())
7075 };
7076 let left_tag_types = left_tag_cols_set
7077 .iter()
7078 .map(|label| {
7079 left.schema()
7080 .fields()
7081 .iter()
7082 .find(|field| field.name() == label)
7083 .map(|field| (label.clone(), field.data_type().clone()))
7084 .with_context(|| ColumnNotFoundSnafu { col: label.clone() })
7085 })
7086 .collect::<Result<HashMap<_, _>>>()?;
7087 let right_tag_types = right_tag_cols_set
7088 .iter()
7089 .map(|label| {
7090 right
7091 .schema()
7092 .fields()
7093 .iter()
7094 .find(|field| field.name() == label)
7095 .map(|field| (label.clone(), field.data_type().clone()))
7096 .with_context(|| ColumnNotFoundSnafu { col: label.clone() })
7097 })
7098 .collect::<Result<HashMap<_, _>>>()?;
7099 let mut target_tag_types = HashMap::with_capacity(all_tags.len());
7100 for label in &all_tags {
7101 let Some(data_type) =
7102 Self::common_label_data_type(left_tag_types.get(label), right_tag_types.get(label))
7103 else {
7104 return UnexpectedPlanExprSnafu {
7105 desc: format!(
7106 "OR label {label} has incompatible types: {:?} and {:?}",
7107 left_tag_types.get(label),
7108 right_tag_types.get(label)
7109 ),
7110 }
7111 .fail();
7112 };
7113 target_tag_types.insert(label.clone(), data_type);
7114 }
7115 let left_has_tsid = left
7116 .schema()
7117 .fields()
7118 .iter()
7119 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME);
7120 let right_has_tsid = right
7121 .schema()
7122 .fields()
7123 .iter()
7124 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME);
7125
7126 let mut all_columns_set = left
7128 .schema()
7129 .fields()
7130 .iter()
7131 .chain(right.schema().fields().iter())
7132 .map(|field| field.name().clone())
7133 .collect::<HashSet<_>>();
7134 if !(left_has_tsid && right_has_tsid) {
7137 all_columns_set.remove(DATA_SCHEMA_TSID_COLUMN_NAME);
7138 }
7139 all_columns_set.remove(&left_time_index_column);
7141 all_columns_set.remove(&right_time_index_column);
7142 if mixed_sample_types {
7143 for (name, _, _) in left_fields.iter().chain(&right_fields) {
7144 all_columns_set.remove(name);
7145 }
7146 all_columns_set.extend(all_tags.iter().cloned());
7147 all_columns_set.insert(mixed_float_field_col.clone());
7148 all_columns_set.insert(mixed_histogram_field_col.clone());
7149 } else if left_field_col != right_field_col {
7150 all_columns_set.remove(right_field_col);
7152 }
7153 let mut all_columns = all_columns_set.into_iter().collect::<Vec<_>>();
7154 all_columns.sort_unstable();
7156 all_columns.insert(0, left_time_index_column.clone());
7158 let mut occupied_column_names = left
7159 .schema()
7160 .fields()
7161 .iter()
7162 .chain(right.schema().fields().iter())
7163 .map(|field| field.name().clone())
7164 .collect::<HashSet<_>>();
7165
7166 let aligned_label_expr = |col: &String, source_types: &HashMap<String, ArrowDataType>| {
7168 let target_type = &target_tag_types[col];
7169 if let Some(source_type) = source_types.get(col) {
7170 let expr = DfExpr::Column(Column::new(None::<String>, col));
7171 if source_type == target_type {
7172 expr
7173 } else {
7174 DfExpr::Cast(Cast::new(Box::new(expr), target_type.clone())).alias(col.clone())
7175 }
7176 } else {
7177 DfExpr::Literal(
7178 Self::string_scalar_value(target_type, None)
7179 .expect("target label type is a string"),
7180 None,
7181 )
7182 .alias(col.clone())
7183 }
7184 };
7185 let null_histogram =
7186 ScalarValue::try_new_null(&native_histogram_type).context(DataFusionPlanningSnafu)?;
7187 let mixed_value_expr = |fields: &[(String, Option<TableReference>, ArrowDataType)],
7188 output_col: &String| {
7189 if output_col == &mixed_float_field_col {
7190 if let Some((name, qualifier, data_type)) = fields
7191 .iter()
7192 .find(|(_, _, data_type)| is_numeric(data_type))
7193 {
7194 let expr = DfExpr::Column(Column::new(qualifier.clone(), name));
7195 if data_type == &ArrowDataType::Float64 {
7196 expr.alias(output_col)
7197 } else {
7198 DfExpr::Cast(Cast::new(Box::new(expr), ArrowDataType::Float64))
7199 .alias(output_col)
7200 }
7201 } else {
7202 DfExpr::Literal(ScalarValue::Float64(None), None).alias(output_col)
7203 }
7204 } else {
7205 fields
7206 .iter()
7207 .find(|(_, _, data_type)| data_type == &native_histogram_type)
7208 .map(|(name, qualifier, _)| {
7209 DfExpr::Column(Column::new(qualifier.clone(), name)).alias(output_col)
7210 })
7211 .unwrap_or_else(|| {
7212 DfExpr::Literal(null_histogram.clone(), None).alias(output_col)
7213 })
7214 }
7215 };
7216 let left_proj_exprs = all_columns.iter().map(|col| {
7217 if mixed_sample_types
7218 && (col == &mixed_float_field_col || col == &mixed_histogram_field_col)
7219 {
7220 mixed_value_expr(&left_fields, col)
7221 } else if !mixed_sample_types
7222 && col == left_field_col
7223 && left_field.2 != target_field_type
7224 {
7225 DfExpr::Cast(Cast::new(
7226 Box::new(DfExpr::Column(Column::new(
7227 left_field.1.clone(),
7228 left_field_col,
7229 ))),
7230 target_field_type.clone(),
7231 ))
7232 .alias(left_field_col.clone())
7233 } else if target_tag_types.contains_key(col) {
7234 aligned_label_expr(col, &left_tag_types)
7235 } else {
7236 DfExpr::Column(Column::new(None::<String>, col))
7237 }
7238 });
7239 let right_time_index_expr = DfExpr::Column(Column::new(
7240 right_qualifier.clone(),
7241 right_time_index_column,
7242 ))
7243 .alias(left_time_index_column.clone());
7244 let right_proj_exprs_without_time_index = all_columns.iter().skip(1).map(|col| {
7248 if mixed_sample_types
7250 && (col == &mixed_float_field_col || col == &mixed_histogram_field_col)
7251 {
7252 mixed_value_expr(&right_fields, col)
7253 } else if !mixed_sample_types && col == left_field_col {
7254 let expr = DfExpr::Column(Column::new(right_field.1.clone(), right_field_col));
7255 if right_field.2 != target_field_type {
7256 DfExpr::Cast(Cast::new(Box::new(expr), target_field_type.clone()))
7257 .alias(left_field_col.clone())
7258 } else if left_field_col != right_field_col {
7259 expr.alias(left_field_col.clone())
7260 } else {
7261 expr
7262 }
7263 } else if target_tag_types.contains_key(col) {
7264 aligned_label_expr(col, &right_tag_types)
7265 } else {
7266 DfExpr::Column(Column::new(None::<String>, col))
7267 }
7268 });
7269 let right_proj_exprs = [right_time_index_expr]
7270 .into_iter()
7271 .chain(right_proj_exprs_without_time_index);
7272
7273 let left_projected = LogicalPlanBuilder::from(left)
7274 .project(left_proj_exprs)
7275 .context(DataFusionPlanningSnafu)?
7276 .alias(left_qualifier_string.clone())
7277 .context(DataFusionPlanningSnafu)?
7278 .build()
7279 .context(DataFusionPlanningSnafu)?;
7280 let right_projected = LogicalPlanBuilder::from(right)
7281 .project(right_proj_exprs)
7282 .context(DataFusionPlanningSnafu)?
7283 .alias(right_qualifier_string.clone())
7284 .context(DataFusionPlanningSnafu)?
7285 .build()
7286 .context(DataFusionPlanningSnafu)?;
7287
7288 let mut match_columns = if let Some(modifier) = modifier
7290 && let Some(matching) = &modifier.matching
7291 {
7292 match matching {
7293 LabelModifier::Include(on) => on.labels.clone(),
7295 LabelModifier::Exclude(ignoring) => {
7297 let ignoring = ignoring.labels.iter().cloned().collect::<HashSet<_>>();
7298 all_tags.difference(&ignoring).cloned().collect()
7299 }
7300 }
7301 } else {
7302 all_tags.iter().cloned().collect()
7303 };
7304 match_columns.sort_unstable();
7306 match_columns.dedup();
7307 occupied_column_names.extend(
7308 left_projected
7309 .schema()
7310 .fields()
7311 .iter()
7312 .chain(right_projected.schema().fields().iter())
7313 .map(|field| field.name().clone()),
7314 );
7315
7316 let visible_schema = left_projected.schema().clone();
7317 let visible_left_exprs = left_projected
7318 .schema()
7319 .iter()
7320 .map(|(qualifier, field)| {
7321 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
7322 })
7323 .collect::<Vec<_>>();
7324 let visible_right_exprs = right_projected
7325 .schema()
7326 .iter()
7327 .map(|(qualifier, field)| {
7328 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
7329 })
7330 .collect::<Vec<_>>();
7331 let mut left_match_exprs = Vec::with_capacity(match_columns.len());
7332 let mut right_match_exprs = Vec::with_capacity(match_columns.len());
7333 let mut next_internal_column = 0;
7334
7335 for label in &match_columns {
7336 let left_field = if left_tag_cols_set.contains(label) {
7337 Some(
7338 left_projected
7339 .schema()
7340 .iter()
7341 .find(|(_, field)| field.name() == label)
7342 .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone()))
7343 .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?,
7344 )
7345 } else {
7346 None
7347 };
7348 let right_field = if right_tag_cols_set.contains(label) {
7349 Some(
7350 right_projected
7351 .schema()
7352 .iter()
7353 .find(|(_, field)| field.name() == label)
7354 .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone()))
7355 .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?,
7356 )
7357 } else {
7358 None
7359 };
7360 let data_type = match (left_field.as_ref(), right_field.as_ref()) {
7361 (Some((_, left_type)), Some((_, right_type))) if left_type == right_type => {
7362 left_type.clone()
7363 }
7364 (Some((_, left_type)), Some((_, right_type))) => {
7365 return UnexpectedPlanExprSnafu {
7366 desc: format!(
7367 "OR match label {label} has incompatible types: {left_type:?} and {right_type:?}"
7368 ),
7369 }
7370 .fail();
7371 }
7372 (Some((_, data_type)), None) | (None, Some((_, data_type))) => data_type.clone(),
7373 (None, None) => ArrowDataType::Utf8,
7374 };
7375 let Some(value_type) = Self::string_value_data_type(&data_type).cloned() else {
7376 return UnexpectedPlanExprSnafu {
7377 desc: format!("OR match label {label} must be a string"),
7378 }
7379 .fail();
7380 };
7381 let internal_name = loop {
7382 let name = format!("__promql_or_match_{next_internal_column}");
7383 next_internal_column += 1;
7384 if occupied_column_names.insert(name.clone()) {
7385 break name;
7386 }
7387 };
7388 left_match_exprs.push(Self::normalized_match_key_expr(
7389 label,
7390 left_field,
7391 &value_type,
7392 &internal_name,
7393 ));
7394 right_match_exprs.push(Self::normalized_match_key_expr(
7395 label,
7396 right_field,
7397 &value_type,
7398 &internal_name,
7399 ));
7400 }
7401
7402 let left_augmented = LogicalPlanBuilder::from(left_projected)
7403 .project(visible_left_exprs.into_iter().chain(left_match_exprs))
7404 .context(DataFusionPlanningSnafu)?
7405 .build()
7406 .context(DataFusionPlanningSnafu)?;
7407 let right_augmented = LogicalPlanBuilder::from(right_projected)
7408 .project(visible_right_exprs.into_iter().chain(right_match_exprs))
7409 .context(DataFusionPlanningSnafu)?
7410 .build()
7411 .context(DataFusionPlanningSnafu)?;
7412
7413 let visible_field_count = visible_schema.fields().len();
7415 let compare_key_indices =
7416 (visible_field_count..visible_field_count + match_columns.len()).collect::<Vec<_>>();
7417 let (time_qualifier, _) = visible_schema
7418 .iter()
7419 .find(|(_, field)| field.name() == &left_time_index_column)
7420 .with_context(|| TimeIndexNotFoundSnafu {
7421 table: left_qualifier_string.clone(),
7422 })?;
7423 let ts_col_idx = left_augmented
7424 .schema()
7425 .iter()
7426 .position(|(qualifier, field)| {
7427 qualifier == time_qualifier && field.name() == &left_time_index_column
7428 })
7429 .with_context(|| TimeIndexNotFoundSnafu {
7430 table: left_qualifier_string.clone(),
7431 })?;
7432 let union_distinct_on = UnionDistinctOn::try_new(
7433 left_augmented,
7434 right_augmented,
7435 compare_key_indices,
7436 ts_col_idx,
7437 )
7438 .context(DataFusionPlanningSnafu)?;
7439 let augmented_result = LogicalPlan::Extension(Extension {
7440 node: Arc::new(union_distinct_on),
7441 });
7442 let result = LogicalPlanBuilder::from(augmented_result)
7443 .project(visible_schema.iter().map(|(qualifier, field)| {
7444 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
7445 }))
7446 .context(DataFusionPlanningSnafu)?
7447 .build()
7448 .context(DataFusionPlanningSnafu)?;
7449
7450 let output_field_col = left_field_col.clone();
7452 let mut output_context = left_context;
7453 let mut visible_tags = all_tags.into_iter().collect::<Vec<_>>();
7454 visible_tags.sort_unstable();
7455 output_context.time_index_column = Some(left_time_index_column);
7456 output_context.tag_columns = visible_tags;
7457 output_context.field_columns = if mixed_sample_types {
7458 vec![mixed_float_field_col, mixed_histogram_field_col]
7459 } else {
7460 vec![output_field_col]
7461 };
7462 output_context.use_tsid = left_has_tsid && right_has_tsid;
7463 self.ctx = output_context;
7464
7465 Ok(result)
7466 }
7467
7468 fn projection_for_each_field_column<F>(
7476 &mut self,
7477 input: LogicalPlan,
7478 name_to_expr: F,
7479 ) -> Result<LogicalPlan>
7480 where
7481 F: FnMut(&String) -> Result<DfExpr>,
7482 {
7483 self.projection_for_each_field_column_with_labels(input, None, name_to_expr)
7484 }
7485
7486 fn projection_for_each_field_column_with_labels<F>(
7489 &mut self,
7490 input: LogicalPlan,
7491 result_labels: Option<&BinaryResultLabels>,
7492 name_to_expr: F,
7493 ) -> Result<LogicalPlan>
7494 where
7495 F: FnMut(&String) -> Result<DfExpr>,
7496 {
7497 let preserve_field_names =
7500 Self::field_columns_are_alternative_samples(input.schema(), &self.ctx.field_columns);
7501 let table_ref = self.ctx.table_name.clone().map(TableReference::bare);
7502 let input_schema = input.schema().clone();
7504 let lookup = |col: &String| {
7505 input_schema
7506 .qualified_field_with_name(table_ref.as_ref(), col)
7507 .or_else(|_| input_schema.qualified_field_with_unqualified_name(col))
7508 .map(|field| DfExpr::Column(field.into()))
7509 .context(DataFusionPlanningSnafu)
7510 };
7511 let tag_columns_iter = match result_labels {
7512 Some(labels) => labels.exprs.iter().cloned().map(Ok).collect::<Vec<_>>(),
7513 None => self.ctx.tag_columns.iter().map(lookup).collect::<Vec<_>>(),
7514 };
7515 let non_field_columns_iter = tag_columns_iter
7516 .into_iter()
7517 .chain(self.ctx.time_index_column.iter().map(lookup));
7518 let tsid_iter = match result_labels {
7519 Some(labels) => labels.tsid_projection(table_ref.clone()),
7520 None => Self::optional_tsid_projection(
7521 input.schema(),
7522 table_ref.as_ref(),
7523 self.ctx.use_tsid,
7524 ),
7525 }
7526 .into_iter()
7527 .map(Ok);
7528
7529 let result_field_columns = self
7531 .ctx
7532 .field_columns
7533 .iter()
7534 .map(name_to_expr)
7535 .collect::<Result<Vec<_>>>()?;
7536
7537 if !preserve_field_names {
7539 self.ctx.field_columns = result_field_columns
7540 .iter()
7541 .map(|expr| expr.schema_name().to_string())
7542 .collect();
7543 }
7544 let field_columns_iter = result_field_columns
7545 .into_iter()
7546 .zip(self.ctx.field_columns.iter())
7547 .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::<String>, name))));
7548
7549 let project_fields = non_field_columns_iter
7551 .chain(tsid_iter)
7552 .chain(field_columns_iter)
7553 .collect::<Result<Vec<_>>>()?;
7554
7555 LogicalPlanBuilder::from(input)
7556 .project(project_fields)
7557 .context(DataFusionPlanningSnafu)?
7558 .build()
7559 .context(DataFusionPlanningSnafu)
7560 }
7561
7562 fn filter_on_field_column<F>(&self, input: LogicalPlan, name_to_expr: F) -> Result<LogicalPlan>
7564 where
7565 F: FnMut(&String) -> Result<DfExpr>,
7566 {
7567 ensure!(
7568 self.ctx.field_columns.len() == 1
7569 || Self::field_columns_are_alternative_samples(
7570 input.schema(),
7571 &self.ctx.field_columns,
7572 ),
7573 UnsupportedExprSnafu {
7574 name: "filter on multi-value input"
7575 }
7576 );
7577
7578 let field_column_filters = self
7579 .ctx
7580 .field_columns
7581 .iter()
7582 .map(name_to_expr)
7583 .collect::<Result<Vec<_>>>()?;
7584 let field_column_filter =
7585 disjunction(field_column_filters).context(UnsupportedExprSnafu {
7586 name: "filter on empty input",
7587 })?;
7588
7589 LogicalPlanBuilder::from(input)
7590 .filter(field_column_filter)
7591 .context(DataFusionPlanningSnafu)?
7592 .build()
7593 .context(DataFusionPlanningSnafu)
7594 }
7595
7596 fn date_part_on_time_index(&self, date_part: &str) -> Result<DfExpr> {
7599 let input_expr = datafusion::logical_expr::col(
7600 self.ctx
7601 .time_index_column
7602 .as_ref()
7603 .with_context(|| TimeIndexNotFoundSnafu {
7605 table: "<doesn't matter>",
7606 })?,
7607 );
7608 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
7609 func: datafusion_functions::datetime::date_part(),
7610 args: vec![date_part.lit(), input_expr],
7611 });
7612 Ok(fn_expr)
7613 }
7614
7615 fn strip_tsid_column(&self, plan: LogicalPlan) -> Result<LogicalPlan> {
7616 let schema = plan.schema();
7617 if !schema
7618 .fields()
7619 .iter()
7620 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
7621 {
7622 return Ok(plan);
7623 }
7624
7625 let project_exprs = schema
7628 .iter()
7629 .filter(|(_, field)| field.name() != DATA_SCHEMA_TSID_COLUMN_NAME)
7630 .map(|(qualifier, field)| {
7631 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
7632 })
7633 .collect::<Vec<_>>();
7634
7635 LogicalPlanBuilder::from(plan)
7636 .project(project_exprs)
7637 .context(DataFusionPlanningSnafu)?
7638 .build()
7639 .context(DataFusionPlanningSnafu)
7640 }
7641
7642 fn apply_alias(&mut self, plan: LogicalPlan, alias_name: String) -> Result<LogicalPlan> {
7644 let fields_expr = self.create_field_column_exprs()?;
7645
7646 ensure!(
7648 fields_expr.len() == 1,
7649 UnsupportedExprSnafu {
7650 name: "alias on multi-value result"
7651 }
7652 );
7653
7654 let project_fields = fields_expr
7655 .into_iter()
7656 .map(|expr| expr.alias(&alias_name))
7657 .chain(self.create_tag_column_exprs()?)
7658 .chain(Some(self.create_time_index_column_expr()?));
7659
7660 LogicalPlanBuilder::from(plan)
7661 .project(project_fields)
7662 .context(DataFusionPlanningSnafu)?
7663 .build()
7664 .context(DataFusionPlanningSnafu)
7665 }
7666}
7667
7668#[derive(Default, Debug)]
7669struct FunctionArgs {
7670 input: Option<PromExpr>,
7671 literals: Vec<DfExpr>,
7672}
7673
7674#[derive(Debug, Clone)]
7677enum ScalarFunc {
7678 DataFusionBuiltin(Arc<ScalarUdfDef>),
7682 DataFusionUdf(Arc<ScalarUdfDef>),
7686 NativeHistogramUdf(Arc<ScalarUdfDef>),
7689 Udf(Arc<ScalarUdfDef>),
7694 ExtrapolateUdf(Arc<ScalarUdfDef>, i64),
7701 GeneratedExpr,
7705}
7706
7707#[cfg(test)]
7708mod test {
7709 use std::time::{Duration, UNIX_EPOCH};
7710
7711 use catalog::RegisterTableRequest;
7712 use catalog::memory::{MemoryCatalogManager, new_memory_catalog_manager};
7713 use common_base::Plugins;
7714 use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
7715 use common_query::native_histogram::{
7716 CUSTOM_BUCKETS_SCHEMA, CounterResetHint, NativeHistogram, build_histogram_array,
7717 };
7718 use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value};
7719 use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS;
7720 use common_query::test_util::DummyDecoder;
7721 use common_recordbatch::RecordBatch as GreptimeRecordBatch;
7722 use datafusion::arrow::array::{
7723 Array, ArrayRef, Float64Array, Int64Array, StringArray, TimestampMillisecondArray,
7724 };
7725 use datafusion::arrow::datatypes::{Field, Schema as ArrowSchema};
7726 use datafusion::arrow::record_batch::RecordBatch;
7727 use datafusion::catalog::{CatalogProvider, MemoryCatalogProvider, MemorySchemaProvider};
7728 use datafusion::datasource::memory::MemorySourceConfig;
7729 use datafusion::datasource::source::DataSourceExec;
7730 use datafusion::datasource::{MemTable, provider_as_source};
7731 use datafusion::execution::context::SessionContext;
7732 use datafusion::logical_expr::Extension;
7733 use datatypes::prelude::ConcreteDataType;
7734 use datatypes::schema::{ColumnSchema, Schema};
7735 use promql_parser::label::Labels;
7736 use promql_parser::parser;
7737 use session::context::QueryContext;
7738 use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan};
7739 use table::Table;
7740 use table::metadata::{FilterPushDownType, TableInfoBuilder, TableMetaBuilder};
7741 use table::test_util::{EmptyTable, MemTable as GreptimeMemTable};
7742
7743 use super::*;
7744 use crate::QueryEngineContext;
7745 use crate::options::QueryOptions;
7746 use crate::parser::QueryLanguageParser;
7747 use crate::query_engine::DefaultSerializer;
7748
7749 mod delta;
7750
7751 fn find_instant_manipulate(plan: &LogicalPlan) -> Option<&InstantManipulate> {
7752 if let LogicalPlan::Extension(Extension { node }) = plan
7753 && let Some(instant_manipulate) = node.as_any().downcast_ref::<InstantManipulate>()
7754 {
7755 return Some(instant_manipulate);
7756 }
7757
7758 plan.inputs().into_iter().find_map(find_instant_manipulate)
7759 }
7760
7761 fn build_query_engine_state() -> QueryEngineState {
7762 QueryEngineState::new(
7763 new_memory_catalog_manager().unwrap(),
7764 None,
7765 None,
7766 None,
7767 None,
7768 None,
7769 false,
7770 Plugins::default(),
7771 QueryOptions::default(),
7772 )
7773 }
7774
7775 #[test]
7776 fn common_label_type_preserves_only_shared_dictionary_encoding() {
7777 let dictionary = ArrowDataType::Dictionary(
7778 Box::new(ArrowDataType::UInt32),
7779 Box::new(ArrowDataType::Utf8),
7780 );
7781 let other_dictionary = ArrowDataType::Dictionary(
7782 Box::new(ArrowDataType::Int32),
7783 Box::new(ArrowDataType::Utf8),
7784 );
7785
7786 assert_eq!(
7787 Some(dictionary.clone()),
7788 PromPlanner::common_label_data_type(Some(&dictionary), Some(&dictionary))
7789 );
7790 assert_eq!(
7791 Some(ArrowDataType::Utf8),
7792 PromPlanner::common_label_data_type(Some(&dictionary), Some(&ArrowDataType::Utf8))
7793 );
7794 assert_eq!(
7795 Some(ArrowDataType::Utf8),
7796 PromPlanner::common_label_data_type(Some(&dictionary), Some(&other_dictionary))
7797 );
7798 assert_eq!(
7799 Some(ArrowDataType::Utf8),
7800 PromPlanner::common_label_data_type(Some(&dictionary), None)
7801 );
7802 }
7803
7804 async fn build_optimized_promql_plan(
7805 table_provider: DfTableSourceProvider,
7806 eval_stmt: &EvalStmt,
7807 ) -> LogicalPlan {
7808 let state = build_query_engine_state();
7809 let raw_plan = PromPlanner::stmt_to_plan(table_provider, eval_stmt, &state)
7810 .await
7811 .unwrap();
7812 let context = QueryEngineContext::new(state.session_state(), QueryContext::arc());
7813 state
7814 .optimize_by_extension_rules(raw_plan, &context)
7815 .unwrap()
7816 }
7817
7818 async fn build_optimized_tsid_plan(
7819 query: &str,
7820 num_tag: usize,
7821 num_field: usize,
7822 end_secs: u64,
7823 lookback_secs: u64,
7824 ) -> String {
7825 let eval_stmt = EvalStmt {
7826 expr: parser::parse(query).unwrap(),
7827 start: UNIX_EPOCH,
7828 end: UNIX_EPOCH
7829 .checked_add(Duration::from_secs(end_secs))
7830 .unwrap(),
7831 interval: Duration::from_secs(5),
7832 lookback_delta: Duration::from_secs(lookback_secs),
7833 };
7834 let table_provider = build_test_table_provider_with_tsid(
7835 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
7836 num_tag,
7837 num_field,
7838 )
7839 .await;
7840
7841 build_optimized_promql_plan(table_provider, &eval_stmt)
7842 .await
7843 .display_indent_schema()
7844 .to_string()
7845 }
7846
7847 async fn assert_nested_count_rewrite_applies(query: &str, expected_outer_agg: &str) {
7848 let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await;
7849
7850 assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
7851 assert!(plan_str.contains("Projection: some_metric.timestamp, some_metric.tag_0"));
7852 assert!(plan_str.contains("Distinct:"));
7853 assert!(plan_str.contains(expected_outer_agg), "{plan_str}");
7854 assert!(!plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]"));
7855 }
7856
7857 async fn assert_nested_count_rewrite_missing(query: &str, num_tag: usize, lookback_secs: u64) {
7858 let plan_str = build_optimized_tsid_plan(query, num_tag, 1, 100_000, lookback_secs).await;
7859 assert!(!plan_str.contains("Distinct:"), "{plan_str}");
7860 }
7861
7862 fn build_eval_stmt(expr: &str) -> EvalStmt {
7863 EvalStmt {
7864 expr: parser::parse(expr).unwrap(),
7865 start: UNIX_EPOCH,
7866 end: UNIX_EPOCH
7867 .checked_add(Duration::from_secs(100_000))
7868 .unwrap(),
7869 interval: Duration::from_secs(5),
7870 lookback_delta: Duration::from_secs(1),
7871 }
7872 }
7873
7874 enum DirectOrValue {
7875 Float64(f64),
7876 Int64(i64),
7877 NativeHistogram(NativeHistogram),
7878 Utf8(&'static str),
7879 }
7880
7881 impl DirectOrValue {
7882 fn data_type(&self) -> ArrowDataType {
7883 match self {
7884 Self::Float64(_) => ArrowDataType::Float64,
7885 Self::Int64(_) => ArrowDataType::Int64,
7886 Self::NativeHistogram(_) => native_histogram_value_type().as_arrow_type(),
7887 Self::Utf8(_) => ArrowDataType::Utf8,
7888 }
7889 }
7890 fn array(&self) -> Arc<dyn Array> {
7891 match self {
7892 Self::Float64(v) => Arc::new(Float64Array::from(vec![*v])),
7893 Self::Int64(v) => Arc::new(Int64Array::from(vec![*v])),
7894 Self::NativeHistogram(v) => build_histogram_array(&[Some(v.clone())]),
7895 Self::Utf8(v) => Arc::new(StringArray::from(vec![*v])),
7896 }
7897 }
7898 }
7899
7900 fn direct_or_histogram() -> NativeHistogram {
7901 NativeHistogram {
7902 schema: 0,
7903 zero_threshold: 0.0,
7904 sum: 1.0,
7905 reset_hint: CounterResetHint::Unknown,
7906 start_timestamp: None,
7907 custom_values: vec![],
7908 positive_spans: vec![],
7909 negative_spans: vec![],
7910 count: 1.0,
7911 zero_count: 1.0,
7912 positive_buckets: vec![],
7913 negative_buckets: vec![],
7914 }
7915 }
7916
7917 fn operator_metric_table(
7918 name: &str,
7919 table_id: u32,
7920 tag: &str,
7921 le: Option<&str>,
7922 value: DirectOrValue,
7923 ) -> table::TableRef {
7924 let value_type = match &value {
7925 DirectOrValue::Float64(_) => ConcreteDataType::float64_datatype(),
7926 DirectOrValue::Int64(_) => ConcreteDataType::int64_datatype(),
7927 DirectOrValue::NativeHistogram(_) => native_histogram_value_type().clone(),
7928 DirectOrValue::Utf8(_) => ConcreteDataType::string_datatype(),
7929 };
7930 let tag_count = 1 + usize::from(le.is_some());
7931 let mut columns = vec![ColumnSchema::new(
7932 "tag".to_string(),
7933 ConcreteDataType::string_datatype(),
7934 false,
7935 )];
7936 if le.is_some() {
7937 columns.push(ColumnSchema::new(
7938 LE_COLUMN_NAME.to_string(),
7939 ConcreteDataType::string_datatype(),
7940 false,
7941 ));
7942 }
7943 columns.extend([
7944 ColumnSchema::new(
7945 "ts".to_string(),
7946 ConcreteDataType::timestamp_millisecond_datatype(),
7947 false,
7948 )
7949 .with_time_index(true),
7950 ColumnSchema::new("v".to_string(), value_type, true),
7951 ]);
7952 let schema = Arc::new(Schema::new(columns));
7953 let mut arrays = vec![Arc::new(StringArray::from(vec![tag])) as Arc<dyn Array>];
7954 if let Some(le) = le {
7955 arrays.push(Arc::new(StringArray::from(vec![le])));
7956 }
7957 arrays.extend([
7958 Arc::new(TimestampMillisecondArray::from(vec![1_000])) as Arc<dyn Array>,
7959 value.array(),
7960 ]);
7961 let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap();
7962 let backing = GreptimeMemTable::new_with_catalog(
7963 name,
7964 GreptimeRecordBatch::from_df_record_batch(schema.clone(), batch),
7965 table_id,
7966 DEFAULT_CATALOG_NAME.to_string(),
7967 DEFAULT_SCHEMA_NAME.to_string(),
7968 );
7969 let value_index = tag_count + 1;
7970 let meta = TableMetaBuilder::empty()
7971 .schema(schema)
7972 .primary_key_indices((0..tag_count).collect())
7973 .value_indices(vec![value_index])
7974 .next_column_id((value_index + 1) as u32)
7975 .build()
7976 .unwrap();
7977 let info = Arc::new(
7978 TableInfoBuilder::default()
7979 .table_id(table_id)
7980 .name(name)
7981 .meta(meta)
7982 .build()
7983 .unwrap(),
7984 );
7985 Arc::new(Table::new(
7986 info,
7987 FilterPushDownType::Unsupported,
7988 backing.data_source(),
7989 ))
7990 }
7991
7992 fn operator_table_provider() -> DfTableSourceProvider {
7993 let catalog = MemoryCatalogManager::with_default_setup();
7994 let tables = [
7995 operator_metric_table("lf", 2_001, "a", None, DirectOrValue::Float64(2.0)),
7996 operator_metric_table(
7997 "lh",
7998 2_002,
7999 "b",
8000 None,
8001 DirectOrValue::NativeHistogram(direct_or_histogram()),
8002 ),
8003 operator_metric_table("rf", 2_003, "b", None, DirectOrValue::Float64(3.0)),
8004 operator_metric_table(
8005 "rh",
8006 2_004,
8007 "a",
8008 None,
8009 DirectOrValue::NativeHistogram(direct_or_histogram()),
8010 ),
8011 operator_metric_table("fallback", 2_005, "c", None, DirectOrValue::Float64(7.0)),
8012 operator_metric_table(
8013 "bad_classic",
8014 2_006,
8015 "d",
8016 Some("broken"),
8017 DirectOrValue::Float64(1.0),
8018 ),
8019 operator_metric_table(
8020 "bad_native",
8021 2_007,
8022 "d",
8023 None,
8024 DirectOrValue::NativeHistogram(direct_or_histogram()),
8025 ),
8026 ];
8027 for table in tables {
8028 let info = table.table_info();
8029 catalog
8030 .register_table_sync(RegisterTableRequest {
8031 catalog: DEFAULT_CATALOG_NAME.to_string(),
8032 schema: DEFAULT_SCHEMA_NAME.to_string(),
8033 table_name: info.name.clone(),
8034 table_id: info.ident.table_id,
8035 table,
8036 })
8037 .unwrap();
8038 }
8039 DfTableSourceProvider::new(
8040 catalog,
8041 false,
8042 QueryContext::arc(),
8043 DummyDecoder::arc(),
8044 false,
8045 )
8046 }
8047
8048 fn operator_eval_stmt(expr: &str) -> EvalStmt {
8049 let time = UNIX_EPOCH.checked_add(Duration::from_secs(1)).unwrap();
8050 EvalStmt {
8051 expr: parser::parse(expr).unwrap(),
8052 start: time,
8053 end: time,
8054 interval: Duration::from_secs(1),
8055 lookback_delta: Duration::from_secs(5),
8056 }
8057 }
8058
8059 struct DirectOrSource {
8060 name: &'static str,
8061 empty: bool,
8062 timestamp: i64,
8063 tags: Vec<(&'static str, Option<&'static str>)>,
8064 value: DirectOrValue,
8065 }
8066
8067 fn source(
8068 name: &'static str,
8069 empty: bool,
8070 timestamp: i64,
8071 tags: Vec<(&'static str, Option<&'static str>)>,
8072 value: DirectOrValue,
8073 ) -> DirectOrSource {
8074 DirectOrSource {
8075 name,
8076 empty,
8077 timestamp,
8078 tags,
8079 value,
8080 }
8081 }
8082
8083 fn tagged_source(
8084 name: &'static str,
8085 empty: bool,
8086 tag: (&'static str, Option<&'static str>),
8087 value: DirectOrValue,
8088 ) -> DirectOrSource {
8089 source(name, empty, 1, vec![("job", Some("job")), tag], value)
8090 }
8091
8092 fn job_source(name: &'static str, value: DirectOrValue) -> DirectOrSource {
8093 source(name, true, 1, vec![("job", Some("job"))], value)
8094 }
8095
8096 fn table(source: &DirectOrSource) -> Arc<MemTable> {
8097 let mut fields = vec![Field::new(
8098 "ts",
8099 ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
8100 false,
8101 )];
8102 fields.extend(
8103 source
8104 .tags
8105 .iter()
8106 .map(|(name, _)| Field::new(*name, ArrowDataType::Utf8, true)),
8107 );
8108 fields.push(Field::new("v", source.value.data_type(), true));
8109 let schema = Arc::new(ArrowSchema::new(fields));
8110 let partitions = if source.empty {
8111 vec![vec![]]
8112 } else {
8113 let mut columns: Vec<Arc<dyn Array>> =
8114 vec![Arc::new(TimestampMillisecondArray::from(vec![
8115 source.timestamp,
8116 ]))];
8117 columns.extend(
8118 source
8119 .tags
8120 .iter()
8121 .map(|(_, value)| Arc::new(StringArray::from(vec![*value])) as Arc<dyn Array>),
8122 );
8123 columns.push(source.value.array());
8124 vec![vec![RecordBatch::try_new(schema.clone(), columns).unwrap()]]
8125 };
8126 Arc::new(MemTable::try_new(schema, partitions).unwrap())
8127 }
8128
8129 fn scan(source: &DirectOrSource) -> LogicalPlan {
8130 LogicalPlanBuilder::scan(source.name, provider_as_source(table(source)), None)
8131 .unwrap()
8132 .build()
8133 .unwrap()
8134 }
8135
8136 fn direct_or_context(qualifier: &str, tags: &[&str], field: &str) -> PromPlannerContext {
8137 PromPlannerContext {
8138 table_name: Some(qualifier.to_string()),
8139 time_index_column: Some("ts".to_string()),
8140 field_columns: vec![field.to_string()],
8141 tag_columns: tags.iter().map(|tag| (*tag).to_string()).collect(),
8142 ..Default::default()
8143 }
8144 }
8145
8146 fn or_modifier(expr: &str) -> Option<BinModifier> {
8147 let PromExpr::Binary(expr) = parser::parse(expr).unwrap() else {
8148 unreachable!()
8149 };
8150 expr.modifier
8151 }
8152
8153 async fn plan_direct_or(
8154 left: LogicalPlan,
8155 right: LogicalPlan,
8156 left_context: PromPlannerContext,
8157 right_context: PromPlannerContext,
8158 modifier: &Option<BinModifier>,
8159 ) -> LogicalPlan {
8160 let table_provider = build_test_table_provider_with_fields(
8161 &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
8162 &[],
8163 )
8164 .await;
8165 let mut planner = PromPlanner {
8166 table_provider,
8167 ctx: PromPlannerContext::default(),
8168 promql_annotations: None,
8169 };
8170 planner
8171 .or_operator(
8172 left,
8173 right,
8174 left_context.tag_columns.iter().cloned().collect(),
8175 right_context.tag_columns.iter().cloned().collect(),
8176 left_context,
8177 right_context,
8178 modifier,
8179 )
8180 .unwrap()
8181 }
8182
8183 async fn execute(
8184 plan: LogicalPlan,
8185 state: &QueryEngineState,
8186 ) -> (LogicalPlan, Vec<RecordBatch>) {
8187 let context = QueryEngineContext::new(state.session_state(), QueryContext::arc());
8188 let optimized = state.optimize_by_extension_rules(plan, &context).unwrap();
8189 let physical = state
8190 .session_state()
8191 .create_physical_plan(&optimized)
8192 .await
8193 .unwrap();
8194 let batches =
8195 datafusion::physical_plan::collect(physical, state.session_state().task_ctx())
8196 .await
8197 .unwrap();
8198 (optimized, batches)
8199 }
8200
8201 async fn run(
8202 left: &DirectOrSource,
8203 right: &DirectOrSource,
8204 left_context: PromPlannerContext,
8205 right_context: PromPlannerContext,
8206 modifier: &Option<BinModifier>,
8207 ) -> (LogicalPlan, Vec<RecordBatch>) {
8208 let plan = plan_direct_or(
8209 scan(left),
8210 scan(right),
8211 left_context,
8212 right_context,
8213 modifier,
8214 )
8215 .await;
8216 execute(plan, &build_query_engine_state()).await
8217 }
8218
8219 async fn mixed_direct_or(histogram_on_left: bool) -> (PromPlanner, LogicalPlan) {
8220 let sample = |histogram: bool| {
8221 if histogram {
8222 DirectOrValue::NativeHistogram(direct_or_histogram())
8223 } else {
8224 DirectOrValue::Float64(1.25)
8225 }
8226 };
8227 let left = tagged_source(
8228 "lhs",
8229 false,
8230 (
8231 "k",
8232 Some(if histogram_on_left {
8233 "histogram"
8234 } else {
8235 "float"
8236 }),
8237 ),
8238 sample(histogram_on_left),
8239 );
8240 let right = tagged_source(
8241 "rhs",
8242 false,
8243 (
8244 "k",
8245 Some(if histogram_on_left {
8246 "float"
8247 } else {
8248 "histogram"
8249 }),
8250 ),
8251 sample(!histogram_on_left),
8252 );
8253 let table_provider = build_test_table_provider_with_fields(
8254 &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
8255 &[],
8256 )
8257 .await;
8258 let mut planner = PromPlanner {
8259 table_provider,
8260 ctx: PromPlannerContext::default(),
8261 promql_annotations: None,
8262 };
8263 let left_context = direct_or_context("lhs", &["job", "k"], "v");
8264 let right_context = direct_or_context("rhs", &["job", "k"], "v");
8265 let plan = planner
8266 .or_operator(
8267 scan(&left),
8268 scan(&right),
8269 left_context.tag_columns.iter().cloned().collect(),
8270 right_context.tag_columns.iter().cloned().collect(),
8271 left_context,
8272 right_context,
8273 &or_modifier("lhs or on(k) rhs"),
8274 )
8275 .unwrap();
8276 (planner, plan)
8277 }
8278
8279 async fn mixed_aggregate_input(histograms: Vec<NativeHistogram>) -> (PromPlanner, LogicalPlan) {
8280 let float_field = format!("{OR_FLOAT_FIELD_PREFIX}0");
8281 let histogram_field = format!("{OR_HISTOGRAM_FIELD_PREFIX}0");
8282 let row_count = histograms.len() + 1;
8283 let schema = Arc::new(ArrowSchema::new(vec![
8284 Field::new(
8285 "ts",
8286 ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
8287 false,
8288 ),
8289 Field::new("k", ArrowDataType::Utf8, false),
8290 Field::new(&float_field, ArrowDataType::Float64, true),
8291 Field::new(
8292 &histogram_field,
8293 native_histogram_value_type().as_arrow_type(),
8294 true,
8295 ),
8296 ]));
8297 let mut histogram_values = Vec::with_capacity(row_count);
8298 histogram_values.push(None);
8299 histogram_values.extend(histograms.into_iter().map(Some));
8300 let batch = RecordBatch::try_new(
8301 schema.clone(),
8302 vec![
8303 Arc::new(TimestampMillisecondArray::from(vec![1; row_count])),
8304 Arc::new(StringArray::from_iter_values(
8305 (0..row_count).map(|row| format!("kind_{row}")),
8306 )),
8307 Arc::new(Float64Array::from_iter(
8308 (0..row_count).map(|row| (row == 0).then_some(1.25)),
8309 )),
8310 build_histogram_array(&histogram_values),
8311 ],
8312 )
8313 .unwrap();
8314 let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap());
8315 let plan = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None)
8316 .unwrap()
8317 .build()
8318 .unwrap();
8319 let table_provider = build_test_table_provider_with_fields(
8320 &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
8321 &[],
8322 )
8323 .await;
8324 let planner = PromPlanner {
8325 table_provider,
8326 ctx: PromPlannerContext {
8327 table_name: Some("mixed".to_string()),
8328 time_index_column: Some("ts".to_string()),
8329 field_columns: vec![float_field, histogram_field],
8330 tag_columns: vec!["k".to_string()],
8331 ..Default::default()
8332 },
8333 promql_annotations: None,
8334 };
8335 (planner, plan)
8336 }
8337
8338 fn assert_no_internal_or_keys(schema: &DFSchema) {
8339 assert!(
8340 schema
8341 .fields()
8342 .iter()
8343 .all(|field| !field.name().starts_with("__promql_or_match_")),
8344 "{schema:?}"
8345 );
8346 }
8347
8348 fn values(batches: &[RecordBatch], column: &str) -> Vec<f64> {
8349 batches
8350 .iter()
8351 .flat_map(|batch| {
8352 batch
8353 .column_by_name(column)
8354 .unwrap()
8355 .as_any()
8356 .downcast_ref::<Float64Array>()
8357 .unwrap()
8358 .iter()
8359 .flatten()
8360 })
8361 .collect()
8362 }
8363
8364 fn numeric_values(batches: &[RecordBatch], column: &str) -> Vec<f64> {
8365 batches
8366 .iter()
8367 .flat_map(|batch| {
8368 let values = datafusion::arrow::compute::cast(
8369 batch.column_by_name(column).unwrap(),
8370 &ArrowDataType::Float64,
8371 )
8372 .unwrap();
8373 values
8374 .as_any()
8375 .downcast_ref::<Float64Array>()
8376 .unwrap()
8377 .iter()
8378 .flatten()
8379 .collect::<Vec<_>>()
8380 })
8381 .collect()
8382 }
8383
8384 fn histograms(batches: &[RecordBatch], column: &str) -> Vec<NativeHistogram> {
8385 batches
8386 .iter()
8387 .flat_map(|batch| {
8388 let values = batch
8389 .column_by_name(column)
8390 .unwrap()
8391 .as_any()
8392 .downcast_ref::<datafusion::arrow::array::StructArray>()
8393 .unwrap();
8394 (0..values.len()).filter_map(|row| {
8395 common_query::native_histogram::read_histogram(values, row).unwrap()
8396 })
8397 })
8398 .collect()
8399 }
8400
8401 fn rows(batches: &[RecordBatch]) -> Vec<(f64, Option<String>)> {
8402 let mut rows = batches
8403 .iter()
8404 .flat_map(|batch| {
8405 let values = batch
8406 .column_by_name("v")
8407 .unwrap()
8408 .as_any()
8409 .downcast_ref::<Float64Array>()
8410 .unwrap();
8411 let labels = batch
8412 .column_by_name("k")
8413 .map(|column| column.as_any().downcast_ref::<StringArray>().unwrap());
8414 (0..batch.num_rows()).map(move |i| {
8415 (
8416 values.value(i),
8417 labels.and_then(|labels| {
8418 (!labels.is_null(i)).then(|| labels.value(i).to_string())
8419 }),
8420 )
8421 })
8422 })
8423 .collect::<Vec<_>>();
8424 rows.sort_by(|left, right| left.0.total_cmp(&right.0));
8425 rows
8426 }
8427
8428 fn matrix_source(
8429 name: &'static str,
8430 k: Option<Option<&'static str>>,
8431 timestamp: i64,
8432 value: f64,
8433 ) -> DirectOrSource {
8434 let mut tags = vec![("job", Some("job"))];
8435 if let Some(k) = k {
8436 tags.push(("k", k));
8437 }
8438 source(name, false, timestamp, tags, DirectOrValue::Float64(value))
8439 }
8440
8441 fn matrix_context(name: &str, k: Option<Option<&str>>) -> PromPlannerContext {
8442 direct_or_context(
8443 name,
8444 if k.is_some() { &["job", "k"] } else { &["job"] },
8445 "v",
8446 )
8447 }
8448
8449 async fn build_missing_le_or_normal_metric_table_provider() -> DfTableSourceProvider {
8450 build_test_table_provider_with_fields(
8451 &[
8452 (
8453 DEFAULT_SCHEMA_NAME.to_string(),
8454 "non_existent_histogram_bucket".to_string(),
8455 ),
8456 (DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string()),
8457 ],
8458 &["pod", "instance"],
8459 )
8460 .await
8461 }
8462
8463 fn assert_normal_metric_schema(plan: &LogicalPlan) {
8464 let fields = plan.schema().fields();
8465 assert_eq!(fields.len(), 4, "{fields:?}");
8466 assert!(
8467 fields.iter().any(|field| field.name() == "pod"),
8468 "{fields:?}"
8469 );
8470 assert!(
8471 fields.iter().any(|field| field.name() == "instance"),
8472 "{fields:?}"
8473 );
8474 assert!(
8475 fields
8476 .iter()
8477 .any(|field| field.name() == greptime_timestamp()),
8478 "{fields:?}"
8479 );
8480 assert!(
8481 fields.iter().any(|field| {
8482 field.name() == greptime_value() && field.data_type() == &ArrowDataType::Float64
8483 }),
8484 "{fields:?}"
8485 );
8486 }
8487
8488 async fn build_test_table_provider_with_distinct_tags(
8489 table_tags: &[(&str, &[&str])],
8490 ) -> DfTableSourceProvider {
8491 let catalog_list = MemoryCatalogManager::with_default_setup();
8492 for (table_name, tags) in table_tags {
8493 let mut columns = tags
8494 .iter()
8495 .map(|tag| {
8496 ColumnSchema::new(
8497 (*tag).to_string(),
8498 ConcreteDataType::string_datatype(),
8499 false,
8500 )
8501 })
8502 .collect::<Vec<_>>();
8503 columns.push(
8504 ColumnSchema::new(
8505 greptime_timestamp().to_string(),
8506 ConcreteDataType::timestamp_millisecond_datatype(),
8507 false,
8508 )
8509 .with_time_index(true),
8510 );
8511 columns.push(ColumnSchema::new(
8512 greptime_value().to_string(),
8513 ConcreteDataType::float64_datatype(),
8514 true,
8515 ));
8516 let table_meta = TableMetaBuilder::empty()
8517 .schema(Arc::new(Schema::new(columns)))
8518 .primary_key_indices((0..tags.len()).collect())
8519 .next_column_id(1024)
8520 .build()
8521 .unwrap();
8522 let table_info = TableInfoBuilder::default()
8523 .name((*table_name).to_string())
8524 .meta(table_meta)
8525 .build()
8526 .unwrap();
8527
8528 assert!(
8529 catalog_list
8530 .register_table_sync(RegisterTableRequest {
8531 catalog: DEFAULT_CATALOG_NAME.to_string(),
8532 schema: DEFAULT_SCHEMA_NAME.to_string(),
8533 table_name: (*table_name).to_string(),
8534 table_id: 1024,
8535 table: EmptyTable::from_table_info(&table_info),
8536 })
8537 .is_ok()
8538 );
8539 }
8540
8541 DfTableSourceProvider::new(
8542 catalog_list,
8543 false,
8544 QueryContext::arc(),
8545 DummyDecoder::arc(),
8546 false,
8547 )
8548 }
8549
8550 fn contains_histogram_fold(plan: &LogicalPlan) -> bool {
8551 matches!(plan, LogicalPlan::Extension(Extension { node }) if node.as_any().is::<HistogramFold>())
8552 || plan.inputs().into_iter().any(contains_histogram_fold)
8553 }
8554
8555 async fn build_set_op_context_table_provider() -> DfTableSourceProvider {
8556 build_test_table_provider_with_distinct_tags(&[
8557 ("bucket_metric", &["job", "le"]),
8558 ("normal_metric", &["job"]),
8559 ("fallback_metric", &["instance"]),
8560 ])
8561 .await
8562 }
8563
8564 async fn build_or_context_table_provider() -> DfTableSourceProvider {
8565 build_test_table_provider_with_distinct_tags(&[
8566 ("normal_metric", &["job"]),
8567 ("other_metric", &["instance"]),
8568 ("non_hist_metric", &["instance"]),
8569 ])
8570 .await
8571 }
8572
8573 async fn optimize_and_create_physical_plan(
8574 state: &QueryEngineState,
8575 plan: LogicalPlan,
8576 ) -> (
8577 LogicalPlan,
8578 Arc<dyn datafusion::physical_plan::ExecutionPlan>,
8579 ) {
8580 let context = QueryEngineContext::new(state.session_state(), QueryContext::arc());
8581 let optimized = state.optimize_by_extension_rules(plan, &context).unwrap();
8582 let physical = state
8583 .session_state()
8584 .create_physical_plan(&optimized)
8585 .await
8586 .unwrap();
8587 (optimized, physical)
8588 }
8589
8590 async fn build_test_table_provider(
8591 table_name_tuples: &[(String, String)],
8592 num_tag: usize,
8593 num_field: usize,
8594 ) -> DfTableSourceProvider {
8595 let catalog_list = MemoryCatalogManager::with_default_setup();
8596 for (schema_name, table_name) in table_name_tuples {
8597 let mut columns = vec![];
8598 for i in 0..num_tag {
8599 columns.push(ColumnSchema::new(
8600 format!("tag_{i}"),
8601 ConcreteDataType::string_datatype(),
8602 false,
8603 ));
8604 }
8605 columns.push(
8606 ColumnSchema::new(
8607 "timestamp".to_string(),
8608 ConcreteDataType::timestamp_millisecond_datatype(),
8609 false,
8610 )
8611 .with_time_index(true),
8612 );
8613 for i in 0..num_field {
8614 columns.push(ColumnSchema::new(
8615 format!("field_{i}"),
8616 ConcreteDataType::float64_datatype(),
8617 true,
8618 ));
8619 }
8620 let schema = Arc::new(Schema::new(columns));
8621 let table_meta = TableMetaBuilder::empty()
8622 .schema(schema)
8623 .primary_key_indices((0..num_tag).collect())
8624 .value_indices((num_tag + 1..num_tag + 1 + num_field).collect())
8625 .next_column_id(1024)
8626 .build()
8627 .unwrap();
8628 let table_info = TableInfoBuilder::default()
8629 .name(table_name.clone())
8630 .meta(table_meta)
8631 .build()
8632 .unwrap();
8633 let table = EmptyTable::from_table_info(&table_info);
8634
8635 assert!(
8636 catalog_list
8637 .register_table_sync(RegisterTableRequest {
8638 catalog: DEFAULT_CATALOG_NAME.to_string(),
8639 schema: schema_name.clone(),
8640 table_name: table_name.clone(),
8641 table_id: 1024,
8642 table,
8643 })
8644 .is_ok()
8645 );
8646 }
8647
8648 DfTableSourceProvider::new(
8649 catalog_list,
8650 false,
8651 QueryContext::arc(),
8652 DummyDecoder::arc(),
8653 false,
8654 )
8655 }
8656
8657 async fn build_test_native_histogram_table_provider(table_name: &str) -> DfTableSourceProvider {
8658 build_test_native_histogram_table_provider_with_marker(table_name, false).await
8659 }
8660
8661 async fn build_test_native_histogram_table_provider_with_marker(
8662 table_name: &str,
8663 temporality_marker: bool,
8664 ) -> DfTableSourceProvider {
8665 let catalog_list = MemoryCatalogManager::with_default_setup();
8666 let mut columns = vec![
8667 ColumnSchema::new(
8668 "tag_0".to_string(),
8669 ConcreteDataType::string_datatype(),
8670 false,
8671 ),
8672 ColumnSchema::new(
8673 LE_COLUMN_NAME.to_string(),
8674 ConcreteDataType::string_datatype(),
8675 true,
8676 ),
8677 ];
8678 if temporality_marker {
8679 columns.push(ColumnSchema::new(
8680 OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(),
8681 ConcreteDataType::string_datatype(),
8682 true,
8683 ));
8684 }
8685 let tag_count = columns.len();
8686 columns.extend([
8687 ColumnSchema::new(
8688 "timestamp".to_string(),
8689 ConcreteDataType::timestamp_millisecond_datatype(),
8690 false,
8691 )
8692 .with_time_index(true),
8693 ColumnSchema::new(
8694 greptime_native_histogram().to_string(),
8695 native_histogram_value_type().clone(),
8696 true,
8697 ),
8698 ]);
8699 let schema = Arc::new(Schema::new(columns));
8700 let table_meta = TableMetaBuilder::empty()
8701 .schema(schema)
8702 .primary_key_indices((0..tag_count).collect())
8703 .value_indices(vec![tag_count + 1])
8704 .next_column_id(1024)
8705 .build()
8706 .unwrap();
8707 let table_info = TableInfoBuilder::default()
8708 .name(table_name)
8709 .meta(table_meta)
8710 .build()
8711 .unwrap();
8712 let table = EmptyTable::from_table_info(&table_info);
8713
8714 assert!(
8715 catalog_list
8716 .register_table_sync(RegisterTableRequest {
8717 catalog: DEFAULT_CATALOG_NAME.to_string(),
8718 schema: DEFAULT_SCHEMA_NAME.to_string(),
8719 table_name: table_name.to_string(),
8720 table_id: 1024,
8721 table,
8722 })
8723 .is_ok()
8724 );
8725
8726 DfTableSourceProvider::new(
8727 catalog_list,
8728 false,
8729 QueryContext::arc(),
8730 DummyDecoder::arc(),
8731 false,
8732 )
8733 }
8734
8735 async fn build_test_multi_histogram_table_provider(table_name: &str) -> DfTableSourceProvider {
8736 let catalog_list = MemoryCatalogManager::with_default_setup();
8737 let columns = vec![
8738 ColumnSchema::new(
8739 "tag_0".to_string(),
8740 ConcreteDataType::string_datatype(),
8741 false,
8742 ),
8743 ColumnSchema::new(
8744 "timestamp".to_string(),
8745 ConcreteDataType::timestamp_millisecond_datatype(),
8746 false,
8747 )
8748 .with_time_index(true),
8749 ColumnSchema::new(
8750 greptime_native_histogram().to_string(),
8751 native_histogram_value_type().clone(),
8752 true,
8753 ),
8754 ColumnSchema::new(
8755 "native_histogram_2".to_string(),
8756 native_histogram_value_type().clone(),
8757 true,
8758 ),
8759 ];
8760 let schema = Arc::new(Schema::new(columns));
8761 let table_meta = TableMetaBuilder::empty()
8762 .schema(schema)
8763 .primary_key_indices(vec![0])
8764 .value_indices(vec![2, 3])
8765 .next_column_id(1024)
8766 .build()
8767 .unwrap();
8768 let table_info = TableInfoBuilder::default()
8769 .name(table_name)
8770 .meta(table_meta)
8771 .build()
8772 .unwrap();
8773 let table = EmptyTable::from_table_info(&table_info);
8774
8775 assert!(
8776 catalog_list
8777 .register_table_sync(RegisterTableRequest {
8778 catalog: DEFAULT_CATALOG_NAME.to_string(),
8779 schema: DEFAULT_SCHEMA_NAME.to_string(),
8780 table_name: table_name.to_string(),
8781 table_id: 1024,
8782 table,
8783 })
8784 .is_ok()
8785 );
8786
8787 DfTableSourceProvider::new(
8788 catalog_list,
8789 false,
8790 QueryContext::arc(),
8791 DummyDecoder::arc(),
8792 false,
8793 )
8794 }
8795
8796 async fn build_test_mixed_native_histogram_table_provider(
8797 table_name: &str,
8798 ) -> DfTableSourceProvider {
8799 build_test_mixed_native_histogram_table_provider_with_marker(table_name, false).await
8800 }
8801
8802 async fn build_test_mixed_native_histogram_table_provider_with_marker(
8803 table_name: &str,
8804 temporality_marker: bool,
8805 ) -> DfTableSourceProvider {
8806 let catalog_list = MemoryCatalogManager::with_default_setup();
8807 let mut columns = vec![ColumnSchema::new(
8808 "tag_0".to_string(),
8809 ConcreteDataType::string_datatype(),
8810 false,
8811 )];
8812 if temporality_marker {
8813 columns.push(ColumnSchema::new(
8814 OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(),
8815 ConcreteDataType::string_datatype(),
8816 true,
8817 ));
8818 }
8819 let tag_count = columns.len();
8820 columns.extend([
8821 ColumnSchema::new(
8822 "timestamp".to_string(),
8823 ConcreteDataType::timestamp_millisecond_datatype(),
8824 false,
8825 )
8826 .with_time_index(true),
8827 ColumnSchema::new(
8828 greptime_native_histogram().to_string(),
8829 native_histogram_value_type().clone(),
8830 true,
8831 ),
8832 ColumnSchema::new(
8833 greptime_value().to_string(),
8834 ConcreteDataType::float64_datatype(),
8835 true,
8836 ),
8837 ]);
8838 let schema = Arc::new(Schema::new(columns));
8839 let table_meta = TableMetaBuilder::empty()
8840 .schema(schema.clone())
8841 .primary_key_indices((0..tag_count).collect())
8842 .value_indices(vec![tag_count + 1, tag_count + 2])
8843 .next_column_id(1024)
8844 .build()
8845 .unwrap();
8846 let table_info = Arc::new(
8847 TableInfoBuilder::default()
8848 .name(table_name)
8849 .meta(table_meta)
8850 .build()
8851 .unwrap(),
8852 );
8853 let mut arrays: Vec<Arc<dyn Array>> =
8854 vec![Arc::new(StringArray::from(vec!["float", "histogram"]))];
8855 if temporality_marker {
8856 arrays.push(Arc::new(StringArray::from(vec![
8857 Some(GREPTIME_TEMPORALITY_DELTA),
8858 Some(GREPTIME_TEMPORALITY_DELTA),
8859 ])));
8860 }
8861 arrays.extend([
8862 Arc::new(TimestampMillisecondArray::from(vec![1_000, 1_000])) as Arc<dyn Array>,
8863 build_histogram_array(&[None, Some(direct_or_histogram())]),
8864 Arc::new(Float64Array::from(vec![Some(2.0), None])),
8865 ]);
8866 let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap();
8867 let backing = GreptimeMemTable::new_with_catalog(
8868 table_name,
8869 GreptimeRecordBatch::from_df_record_batch(schema, batch),
8870 1024,
8871 DEFAULT_CATALOG_NAME.to_string(),
8872 DEFAULT_SCHEMA_NAME.to_string(),
8873 );
8874 let table = Arc::new(Table::new(
8875 table_info,
8876 FilterPushDownType::Unsupported,
8877 backing.data_source(),
8878 ));
8879
8880 assert!(
8881 catalog_list
8882 .register_table_sync(RegisterTableRequest {
8883 catalog: DEFAULT_CATALOG_NAME.to_string(),
8884 schema: DEFAULT_SCHEMA_NAME.to_string(),
8885 table_name: table_name.to_string(),
8886 table_id: 1024,
8887 table,
8888 })
8889 .is_ok()
8890 );
8891
8892 DfTableSourceProvider::new(
8893 catalog_list,
8894 false,
8895 QueryContext::arc(),
8896 DummyDecoder::arc(),
8897 false,
8898 )
8899 }
8900
8901 fn classic_and_native_histogram_table_provider(
8902 native_tag: &str,
8903 native_le: Option<&str>,
8904 native_histogram: NativeHistogram,
8905 ) -> DfTableSourceProvider {
8906 let table_name = "mixed_histogram";
8907 let catalog = MemoryCatalogManager::with_default_setup();
8908 let schema = Arc::new(Schema::new(vec![
8909 ColumnSchema::new(
8910 "tag".to_string(),
8911 ConcreteDataType::string_datatype(),
8912 false,
8913 ),
8914 ColumnSchema::new(
8915 LE_COLUMN_NAME.to_string(),
8916 ConcreteDataType::string_datatype(),
8917 true,
8918 ),
8919 ColumnSchema::new(
8920 "timestamp".to_string(),
8921 ConcreteDataType::timestamp_millisecond_datatype(),
8922 false,
8923 )
8924 .with_time_index(true),
8925 ColumnSchema::new(
8926 greptime_native_histogram().to_string(),
8927 native_histogram_value_type().clone(),
8928 true,
8929 ),
8930 ColumnSchema::new(
8931 greptime_value().to_string(),
8932 ConcreteDataType::float64_datatype(),
8933 true,
8934 ),
8935 ]));
8936 let table_meta = TableMetaBuilder::empty()
8937 .schema(schema.clone())
8938 .primary_key_indices(vec![0, 1])
8939 .value_indices(vec![3, 4])
8940 .next_column_id(5)
8941 .build()
8942 .unwrap();
8943 let table_info = Arc::new(
8944 TableInfoBuilder::default()
8945 .name(table_name)
8946 .meta(table_meta)
8947 .build()
8948 .unwrap(),
8949 );
8950 let batch = RecordBatch::try_new(
8951 schema.arrow_schema().clone(),
8952 vec![
8953 Arc::new(StringArray::from(vec![
8954 "classic", "classic", native_tag, "classic", "classic", native_tag,
8955 ])),
8956 Arc::new(StringArray::from(vec![
8957 Some("1"),
8958 Some("+Inf"),
8959 native_le,
8960 Some("1"),
8961 Some("+Inf"),
8962 native_le,
8963 ])),
8964 Arc::new(TimestampMillisecondArray::from(vec![
8965 1_000, 1_000, 1_000, 2_000, 2_000, 2_000,
8966 ])),
8967 build_histogram_array(&[
8968 None,
8969 None,
8970 Some(native_histogram.clone()),
8971 None,
8972 None,
8973 Some(native_histogram),
8974 ]),
8975 Arc::new(Float64Array::from(vec![
8976 Some(2.0),
8977 Some(4.0),
8978 None,
8979 Some(2.0),
8980 Some(4.0),
8981 None,
8982 ])),
8983 ],
8984 )
8985 .unwrap();
8986 let backing = GreptimeMemTable::new_with_catalog(
8987 table_name,
8988 GreptimeRecordBatch::from_df_record_batch(schema, batch),
8989 2_200,
8990 DEFAULT_CATALOG_NAME.to_string(),
8991 DEFAULT_SCHEMA_NAME.to_string(),
8992 );
8993 let table = Arc::new(Table::new(
8994 table_info,
8995 FilterPushDownType::Unsupported,
8996 backing.data_source(),
8997 ));
8998 catalog
8999 .register_table_sync(RegisterTableRequest {
9000 catalog: DEFAULT_CATALOG_NAME.to_string(),
9001 schema: DEFAULT_SCHEMA_NAME.to_string(),
9002 table_name: table_name.to_string(),
9003 table_id: 2_200,
9004 table,
9005 })
9006 .unwrap();
9007
9008 DfTableSourceProvider::new(
9009 catalog,
9010 false,
9011 QueryContext::arc(),
9012 DummyDecoder::arc(),
9013 false,
9014 )
9015 }
9016
9017 async fn build_test_table_provider_with_tsid(
9018 table_name_tuples: &[(String, String)],
9019 num_tag: usize,
9020 num_field: usize,
9021 ) -> DfTableSourceProvider {
9022 let table_specs = table_name_tuples
9023 .iter()
9024 .map(|(schema_name, table_name)| ((schema_name.clone(), table_name.clone()), num_field))
9025 .collect::<Vec<_>>();
9026 build_test_table_provider_with_tsid_fields(&table_specs, num_tag).await
9027 }
9028
9029 async fn build_test_table_provider_with_tsid_fields(
9030 table_specs: &[((String, String), usize)],
9031 num_tag: usize,
9032 ) -> DfTableSourceProvider {
9033 let table_specs = table_specs
9034 .iter()
9035 .map(|(table_name_tuple, num_field)| (table_name_tuple.clone(), num_tag, *num_field))
9036 .collect::<Vec<_>>();
9037 build_test_table_provider_with_tsid_tag_fields(&table_specs).await
9038 }
9039
9040 async fn build_test_table_provider_with_tsid_tag_fields(
9041 table_specs: &[((String, String), usize, usize)],
9042 ) -> DfTableSourceProvider {
9043 let catalog_list = MemoryCatalogManager::with_default_setup();
9044
9045 let physical_table_name = "phy";
9046 let physical_table_id = 999u32;
9047 let physical_num_tag = table_specs
9048 .iter()
9049 .map(|(_, num_tag, _)| *num_tag)
9050 .max()
9051 .unwrap_or(0);
9052 let physical_num_field = table_specs
9053 .iter()
9054 .map(|(_, _, num_field)| *num_field)
9055 .max()
9056 .unwrap_or(0);
9057
9058 {
9060 let mut columns = vec![
9061 ColumnSchema::new(
9062 DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(),
9063 ConcreteDataType::uint32_datatype(),
9064 false,
9065 ),
9066 ColumnSchema::new(
9067 DATA_SCHEMA_TSID_COLUMN_NAME.to_string(),
9068 ConcreteDataType::uint64_datatype(),
9069 false,
9070 ),
9071 ];
9072 for i in 0..physical_num_tag {
9073 columns.push(ColumnSchema::new(
9074 format!("tag_{i}"),
9075 ConcreteDataType::string_datatype(),
9076 false,
9077 ));
9078 }
9079 columns.push(
9080 ColumnSchema::new(
9081 "timestamp".to_string(),
9082 ConcreteDataType::timestamp_millisecond_datatype(),
9083 false,
9084 )
9085 .with_time_index(true),
9086 );
9087 for i in 0..physical_num_field {
9088 columns.push(ColumnSchema::new(
9089 format!("field_{i}"),
9090 ConcreteDataType::float64_datatype(),
9091 true,
9092 ));
9093 }
9094
9095 let schema = Arc::new(Schema::new(columns));
9096 let primary_key_indices = (0..(2 + physical_num_tag)).collect::<Vec<_>>();
9097 let table_meta = TableMetaBuilder::empty()
9098 .schema(schema)
9099 .primary_key_indices(primary_key_indices)
9100 .value_indices(
9101 (2 + physical_num_tag..2 + physical_num_tag + 1 + physical_num_field).collect(),
9102 )
9103 .engine(METRIC_ENGINE_NAME.to_string())
9104 .next_column_id(1024)
9105 .build()
9106 .unwrap();
9107 let table_info = TableInfoBuilder::default()
9108 .table_id(physical_table_id)
9109 .name(physical_table_name)
9110 .meta(table_meta)
9111 .build()
9112 .unwrap();
9113 let table = EmptyTable::from_table_info(&table_info);
9114
9115 assert!(
9116 catalog_list
9117 .register_table_sync(RegisterTableRequest {
9118 catalog: DEFAULT_CATALOG_NAME.to_string(),
9119 schema: DEFAULT_SCHEMA_NAME.to_string(),
9120 table_name: physical_table_name.to_string(),
9121 table_id: physical_table_id,
9122 table,
9123 })
9124 .is_ok()
9125 );
9126 }
9127
9128 for (idx, ((schema_name, table_name), num_tag, num_field)) in table_specs.iter().enumerate()
9130 {
9131 let mut columns = vec![];
9132 for i in 0..*num_tag {
9133 columns.push(ColumnSchema::new(
9134 format!("tag_{i}"),
9135 ConcreteDataType::string_datatype(),
9136 false,
9137 ));
9138 }
9139 columns.push(
9140 ColumnSchema::new(
9141 "timestamp".to_string(),
9142 ConcreteDataType::timestamp_millisecond_datatype(),
9143 false,
9144 )
9145 .with_time_index(true),
9146 );
9147 for i in 0..*num_field {
9148 columns.push(ColumnSchema::new(
9149 format!("field_{i}"),
9150 ConcreteDataType::float64_datatype(),
9151 true,
9152 ));
9153 }
9154
9155 let schema = Arc::new(Schema::new(columns));
9156 let mut options = table::requests::TableOptions::default();
9157 options.extra_options.insert(
9158 LOGICAL_TABLE_METADATA_KEY.to_string(),
9159 physical_table_name.to_string(),
9160 );
9161 let table_id = 1024u32 + idx as u32;
9162 let table_meta = TableMetaBuilder::empty()
9163 .schema(schema)
9164 .primary_key_indices((0..*num_tag).collect())
9165 .value_indices((*num_tag + 1..*num_tag + 1 + *num_field).collect())
9166 .engine(METRIC_ENGINE_NAME.to_string())
9167 .options(options)
9168 .next_column_id(1024)
9169 .build()
9170 .unwrap();
9171 let table_info = TableInfoBuilder::default()
9172 .table_id(table_id)
9173 .name(table_name.clone())
9174 .meta(table_meta)
9175 .build()
9176 .unwrap();
9177 let table = EmptyTable::from_table_info(&table_info);
9178
9179 assert!(
9180 catalog_list
9181 .register_table_sync(RegisterTableRequest {
9182 catalog: DEFAULT_CATALOG_NAME.to_string(),
9183 schema: schema_name.clone(),
9184 table_name: table_name.clone(),
9185 table_id,
9186 table,
9187 })
9188 .is_ok()
9189 );
9190 }
9191
9192 DfTableSourceProvider::new(
9193 catalog_list,
9194 false,
9195 QueryContext::arc(),
9196 DummyDecoder::arc(),
9197 false,
9198 )
9199 }
9200
9201 async fn build_test_table_provider_with_fields(
9202 table_name_tuples: &[(String, String)],
9203 tags: &[&str],
9204 ) -> DfTableSourceProvider {
9205 let catalog_list = MemoryCatalogManager::with_default_setup();
9206 for (schema_name, table_name) in table_name_tuples {
9207 let mut columns = vec![];
9208 let num_tag = tags.len();
9209 for tag in tags {
9210 columns.push(ColumnSchema::new(
9211 tag.to_string(),
9212 ConcreteDataType::string_datatype(),
9213 false,
9214 ));
9215 }
9216 columns.push(
9217 ColumnSchema::new(
9218 greptime_timestamp().to_string(),
9219 ConcreteDataType::timestamp_millisecond_datatype(),
9220 false,
9221 )
9222 .with_time_index(true),
9223 );
9224 columns.push(ColumnSchema::new(
9225 greptime_value().to_string(),
9226 ConcreteDataType::float64_datatype(),
9227 true,
9228 ));
9229 let schema = Arc::new(Schema::new(columns));
9230 let table_meta = TableMetaBuilder::empty()
9231 .schema(schema)
9232 .primary_key_indices((0..num_tag).collect())
9233 .next_column_id(1024)
9234 .build()
9235 .unwrap();
9236 let table_info = TableInfoBuilder::default()
9237 .name(table_name.clone())
9238 .meta(table_meta)
9239 .build()
9240 .unwrap();
9241 let table = EmptyTable::from_table_info(&table_info);
9242
9243 assert!(
9244 catalog_list
9245 .register_table_sync(RegisterTableRequest {
9246 catalog: DEFAULT_CATALOG_NAME.to_string(),
9247 schema: schema_name.clone(),
9248 table_name: table_name.clone(),
9249 table_id: 1024,
9250 table,
9251 })
9252 .is_ok()
9253 );
9254 }
9255
9256 DfTableSourceProvider::new(
9257 catalog_list,
9258 false,
9259 QueryContext::arc(),
9260 DummyDecoder::arc(),
9261 false,
9262 )
9263 }
9264
9265 async fn do_single_instant_function_call(fn_name: &'static str, plan_name: &str) {
9281 let prom_expr =
9282 parser::parse(&format!("{fn_name}(some_metric{{tag_0!=\"bar\"}})")).unwrap();
9283 let eval_stmt = EvalStmt {
9284 expr: prom_expr,
9285 start: UNIX_EPOCH,
9286 end: UNIX_EPOCH
9287 .checked_add(Duration::from_secs(100_000))
9288 .unwrap(),
9289 interval: Duration::from_secs(5),
9290 lookback_delta: Duration::from_secs(1),
9291 };
9292
9293 let table_provider = build_test_table_provider(
9294 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9295 1,
9296 1,
9297 )
9298 .await;
9299 let plan =
9300 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9301 .await
9302 .unwrap();
9303
9304 let expected = String::from(
9305 "Filter: TEMPLATE(field_0) IS NOT NULL [timestamp:Timestamp(ms), TEMPLATE(field_0):Float64;N, tag_0:Utf8]\
9306 \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]\
9307 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
9308 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
9309 \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]\
9310 \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]\
9311 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"
9312 ).replace("TEMPLATE", plan_name);
9313
9314 assert_eq!(plan.display_indent_schema().to_string(), expected);
9315 }
9316
9317 #[tokio::test]
9318 async fn single_abs() {
9319 do_single_instant_function_call("abs", "abs").await;
9320 }
9321
9322 #[tokio::test]
9323 #[should_panic]
9324 async fn single_absent() {
9325 do_single_instant_function_call("absent", "").await;
9326 }
9327
9328 #[tokio::test]
9329 async fn single_ceil() {
9330 do_single_instant_function_call("ceil", "ceil").await;
9331 }
9332
9333 #[tokio::test]
9334 async fn single_exp() {
9335 do_single_instant_function_call("exp", "exp").await;
9336 }
9337
9338 #[tokio::test]
9339 async fn single_ln() {
9340 do_single_instant_function_call("ln", "ln").await;
9341 }
9342
9343 #[tokio::test]
9344 async fn single_log2() {
9345 do_single_instant_function_call("log2", "log2").await;
9346 }
9347
9348 #[tokio::test]
9349 async fn single_log10() {
9350 do_single_instant_function_call("log10", "log10").await;
9351 }
9352
9353 #[tokio::test]
9354 #[should_panic]
9355 async fn single_scalar() {
9356 do_single_instant_function_call("scalar", "").await;
9357 }
9358
9359 #[tokio::test]
9360 #[should_panic]
9361 async fn single_sgn() {
9362 do_single_instant_function_call("sgn", "").await;
9363 }
9364
9365 #[tokio::test]
9366 #[should_panic]
9367 async fn single_sort() {
9368 do_single_instant_function_call("sort", "").await;
9369 }
9370
9371 #[tokio::test]
9372 #[should_panic]
9373 async fn single_sort_desc() {
9374 do_single_instant_function_call("sort_desc", "").await;
9375 }
9376
9377 #[tokio::test]
9378 async fn single_sqrt() {
9379 do_single_instant_function_call("sqrt", "sqrt").await;
9380 }
9381
9382 #[tokio::test]
9383 async fn single_timestamp_plan_preserves_source_value() {
9384 let eval_stmt = build_eval_stmt(r#"timestamp(some_metric{tag_0!="bar"})"#);
9385 let table_provider = build_test_table_provider(
9386 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9387 1,
9388 1,
9389 )
9390 .await;
9391 let plan =
9392 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9393 .await
9394 .unwrap();
9395
9396 let expected = String::from(
9397 "Filter: value IS NOT NULL [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\
9398 \n Projection: some_metric.timestamp, value AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\
9399 \n Projection: some_metric.timestamp, __promql_timestamp_value_ AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\
9400 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\
9401 \n Projection: some_metric.tag_0, some_metric.timestamp, some_metric.field_0, CAST(CAST(CAST(CAST(some_metric.timestamp AS Int64) AS Decimal128(19, 0)) * Decimal128(1,1,0) + Decimal128(0,19,0) AS Int64) AS Float64) / Float64(1000) AS __promql_timestamp_value_ [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\
9402 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
9403 \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]\
9404 \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]\
9405 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
9406 );
9407
9408 assert_eq!(plan.display_indent_schema().to_string(), expected);
9409 }
9410
9411 #[tokio::test]
9412 async fn single_acos() {
9413 do_single_instant_function_call("acos", "acos").await;
9414 }
9415
9416 #[tokio::test]
9417 #[should_panic]
9418 async fn single_acosh() {
9419 do_single_instant_function_call("acosh", "").await;
9420 }
9421
9422 #[tokio::test]
9423 async fn single_asin() {
9424 do_single_instant_function_call("asin", "asin").await;
9425 }
9426
9427 #[tokio::test]
9428 #[should_panic]
9429 async fn single_asinh() {
9430 do_single_instant_function_call("asinh", "").await;
9431 }
9432
9433 #[tokio::test]
9434 async fn single_atan() {
9435 do_single_instant_function_call("atan", "atan").await;
9436 }
9437
9438 #[tokio::test]
9439 #[should_panic]
9440 async fn single_atanh() {
9441 do_single_instant_function_call("atanh", "").await;
9442 }
9443
9444 #[tokio::test]
9445 async fn single_cos() {
9446 do_single_instant_function_call("cos", "cos").await;
9447 }
9448
9449 #[tokio::test]
9450 #[should_panic]
9451 async fn single_cosh() {
9452 do_single_instant_function_call("cosh", "").await;
9453 }
9454
9455 #[tokio::test]
9456 async fn single_sin() {
9457 do_single_instant_function_call("sin", "sin").await;
9458 }
9459
9460 #[tokio::test]
9461 #[should_panic]
9462 async fn single_sinh() {
9463 do_single_instant_function_call("sinh", "").await;
9464 }
9465
9466 #[tokio::test]
9467 async fn single_tan() {
9468 do_single_instant_function_call("tan", "tan").await;
9469 }
9470
9471 #[tokio::test]
9472 #[should_panic]
9473 async fn single_tanh() {
9474 do_single_instant_function_call("tanh", "").await;
9475 }
9476
9477 #[tokio::test]
9478 #[should_panic]
9479 async fn single_deg() {
9480 do_single_instant_function_call("deg", "").await;
9481 }
9482
9483 #[tokio::test]
9484 #[should_panic]
9485 async fn single_rad() {
9486 do_single_instant_function_call("rad", "").await;
9487 }
9488
9489 async fn do_aggregate_expr_plan(fn_name: &str, plan_name: &str) {
9511 let prom_expr = parser::parse(&format!(
9512 "{fn_name} by (tag_1)(some_metric{{tag_0!=\"bar\"}})",
9513 ))
9514 .unwrap();
9515 let mut eval_stmt = EvalStmt {
9516 expr: prom_expr,
9517 start: UNIX_EPOCH,
9518 end: UNIX_EPOCH
9519 .checked_add(Duration::from_secs(100_000))
9520 .unwrap(),
9521 interval: Duration::from_secs(5),
9522 lookback_delta: Duration::from_secs(1),
9523 };
9524
9525 let table_provider = build_test_table_provider(
9527 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9528 2,
9529 2,
9530 )
9531 .await;
9532 let plan =
9533 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9534 .await
9535 .unwrap();
9536 let expected_no_without = String::from(
9537 "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]\
9538 \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]\
9539 \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]\
9540 \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
9541 \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]\
9542 \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]\
9543 \n TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]"
9544 ).replace("TEMPLATE", plan_name);
9545 assert_eq!(
9546 plan.display_indent_schema().to_string(),
9547 expected_no_without
9548 );
9549
9550 if let PromExpr::Aggregate(AggregateExpr { modifier, .. }) = &mut eval_stmt.expr {
9552 *modifier = Some(LabelModifier::Exclude(Labels {
9553 labels: vec![String::from("tag_1")].into_iter().collect(),
9554 }));
9555 }
9556 let table_provider = build_test_table_provider(
9557 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9558 2,
9559 2,
9560 )
9561 .await;
9562 let plan =
9563 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9564 .await
9565 .unwrap();
9566 let expected_without = String::from(
9567 "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]\
9568 \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]\
9569 \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]\
9570 \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
9571 \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]\
9572 \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]\
9573 \n TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]"
9574 ).replace("TEMPLATE", plan_name);
9575 assert_eq!(plan.display_indent_schema().to_string(), expected_without);
9576 }
9577
9578 #[tokio::test]
9579 async fn aggregate_sum() {
9580 do_aggregate_expr_plan("sum", "sum").await;
9581 }
9582
9583 #[tokio::test]
9584 async fn tsid_is_used_for_series_divide_when_available() {
9585 let prom_expr = parser::parse("some_metric").unwrap();
9586 let eval_stmt = EvalStmt {
9587 expr: prom_expr,
9588 start: UNIX_EPOCH,
9589 end: UNIX_EPOCH
9590 .checked_add(Duration::from_secs(100_000))
9591 .unwrap(),
9592 interval: Duration::from_secs(5),
9593 lookback_delta: Duration::from_secs(1),
9594 };
9595
9596 let table_provider = build_test_table_provider_with_tsid(
9597 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9598 1,
9599 1,
9600 )
9601 .await;
9602 let plan =
9603 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9604 .await
9605 .unwrap();
9606
9607 let plan_str = plan.display_indent_schema().to_string();
9608 assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
9609 assert!(plan_str.contains("__tsid ASC NULLS FIRST"));
9610 assert!(
9611 !plan
9612 .schema()
9613 .fields()
9614 .iter()
9615 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
9616 );
9617
9618 let manipulate = find_instant_manipulate(&plan).unwrap();
9619 let exec = manipulate.to_execution_plan(Arc::new(DataSourceExec::new(Arc::new(
9620 MemorySourceConfig::try_new(
9621 &[],
9622 Arc::new(
9623 datafusion_expr::UserDefinedLogicalNodeCore::inputs(manipulate)[0]
9624 .schema()
9625 .as_arrow()
9626 .clone(),
9627 ),
9628 None,
9629 )
9630 .unwrap(),
9631 ))));
9632 assert!(format!("{exec:?}").contains("reuse_tsid_column: true"));
9633 }
9634
9635 #[tokio::test]
9636 async fn default_binary_join_uses_tsid_when_available() {
9637 let eval_stmt = build_eval_stmt("some_metric / some_alt_metric");
9638
9639 let table_provider = build_test_table_provider_with_tsid(
9640 &[
9641 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9642 (
9643 DEFAULT_SCHEMA_NAME.to_string(),
9644 "some_alt_metric".to_string(),
9645 ),
9646 ],
9647 1,
9648 1,
9649 )
9650 .await;
9651 let plan =
9652 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9653 .await
9654 .unwrap();
9655
9656 let plan_str = plan.display_indent_schema().to_string();
9657 assert!(
9658 plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
9659 "{plan_str}"
9660 );
9661 assert!(
9662 !plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
9663 "{plan_str}"
9664 );
9665 }
9666
9667 #[tokio::test]
9668 async fn reject_binary_fill_modifiers() {
9669 let state = build_query_engine_state();
9670
9671 for query in [
9672 "some_metric + fill(0) some_alt_metric",
9673 "some_metric + fill_left(0) some_alt_metric",
9674 "some_metric + fill_right(0) some_alt_metric",
9675 "(some_metric + fill(0) some_alt_metric) + some_metric",
9676 ] {
9677 let eval_stmt = build_eval_stmt(query);
9678 let table_provider = build_test_table_provider(&[], 0, 0).await;
9679 let err = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &state)
9680 .await
9681 .unwrap_err();
9682
9683 assert!(
9684 matches!(
9685 &err,
9686 crate::promql::error::Error::UnsupportedExpr { name, .. }
9687 if name == "PromQL fill modifiers"
9688 ),
9689 "{err}"
9690 );
9691 }
9692 }
9693
9694 #[tokio::test]
9695 async fn timestamp_binary_join_falls_back_when_tsid_is_projected_out() {
9696 for query in [
9697 "timestamp(some_metric) / some_metric",
9698 "some_metric / timestamp(some_metric)",
9699 ] {
9700 let eval_stmt = build_eval_stmt(query);
9701
9702 let table_provider = build_test_table_provider_with_tsid(
9703 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9704 1,
9705 1,
9706 )
9707 .await;
9708 let plan =
9709 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9710 .await
9711 .unwrap();
9712
9713 let plan_str = plan.display_indent_schema().to_string();
9714 assert!(!plan_str.contains("__tsid ="), "{query}: {plan_str}");
9715 assert!(
9716 plan_str.contains("lhs.tag_0 = rhs.tag_0"),
9717 "{query}: {plan_str}"
9718 );
9719 assert!(
9720 !plan
9721 .schema()
9722 .fields()
9723 .iter()
9724 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME),
9725 "{query}: {plan_str}"
9726 );
9727 }
9728 }
9729
9730 #[tokio::test]
9731 async fn timestamp_binary_join_rejects_default_matching_on_mismatched_labels() {
9732 let eval_stmt = build_eval_stmt("timestamp(left_host_job) / right_by_job");
9733
9734 let table_provider = build_test_table_provider_with_tsid_tag_fields(&[
9735 (
9736 (DEFAULT_SCHEMA_NAME.to_string(), "left_host_job".to_string()),
9737 2,
9738 1,
9739 ),
9740 (
9741 (DEFAULT_SCHEMA_NAME.to_string(), "right_by_job".to_string()),
9742 1,
9743 1,
9744 ),
9745 ])
9746 .await;
9747 let plan =
9748 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9749 .await
9750 .unwrap();
9751 let plan_str = plan.display_indent_schema().to_string();
9752
9753 assert!(
9754 plan_str.contains("Boolean(false)") || plan_str.contains("false"),
9755 "{plan_str}"
9756 );
9757 }
9758
9759 #[tokio::test]
9760 async fn tsid_is_preserved_for_nested_default_binary_joins() {
9761 let eval_stmt = build_eval_stmt("(some_metric - some_alt_metric) / some_third_metric");
9762
9763 let table_provider = build_test_table_provider_with_tsid(
9764 &[
9765 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9766 (
9767 DEFAULT_SCHEMA_NAME.to_string(),
9768 "some_alt_metric".to_string(),
9769 ),
9770 (
9771 DEFAULT_SCHEMA_NAME.to_string(),
9772 "some_third_metric".to_string(),
9773 ),
9774 ],
9775 1,
9776 1,
9777 )
9778 .await;
9779 let plan =
9780 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9781 .await
9782 .unwrap();
9783
9784 let plan_str = plan.display_indent_schema().to_string();
9785 assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
9786 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9787 }
9788
9789 #[tokio::test]
9790 async fn repeated_tsid_binary_operand_reuses_leaf_plan() {
9791 let eval_stmt = build_eval_stmt("((some_metric - some_alt_metric) / some_metric) * 100");
9792
9793 let table_provider = build_test_table_provider_with_tsid(
9794 &[
9795 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9796 (
9797 DEFAULT_SCHEMA_NAME.to_string(),
9798 "some_alt_metric".to_string(),
9799 ),
9800 ],
9801 1,
9802 1,
9803 )
9804 .await;
9805 let plan =
9806 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9807 .await
9808 .unwrap();
9809
9810 let plan_str = plan.display_indent_schema().to_string();
9811 assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}");
9812 assert_eq!(
9813 plan_str
9814 .matches("Filter: phy.__table_id = UInt32(1024)")
9815 .count(),
9816 1,
9817 "{plan_str}"
9818 );
9819 assert_eq!(
9820 plan_str.matches("PromInstantManipulate").count(),
9821 2,
9822 "{plan_str}"
9823 );
9824 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9825 }
9826
9827 #[tokio::test]
9828 async fn repeated_tsid_binary_operand_reuses_shorter_field_side() {
9829 let eval_stmt =
9830 build_eval_stmt("((two_field_metric - one_field_metric) / one_field_metric) * 100");
9831
9832 let table_provider = build_test_table_provider_with_tsid_fields(
9833 &[
9834 (
9835 (
9836 DEFAULT_SCHEMA_NAME.to_string(),
9837 "two_field_metric".to_string(),
9838 ),
9839 2,
9840 ),
9841 (
9842 (
9843 DEFAULT_SCHEMA_NAME.to_string(),
9844 "one_field_metric".to_string(),
9845 ),
9846 1,
9847 ),
9848 ],
9849 1,
9850 )
9851 .await;
9852 let plan =
9853 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9854 .await
9855 .unwrap();
9856
9857 let field_names = plan
9858 .schema()
9859 .fields()
9860 .iter()
9861 .map(|field| field.name().clone())
9862 .collect::<Vec<_>>();
9863 let value_columns = field_names
9864 .iter()
9865 .filter(|name| {
9866 *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME
9867 })
9868 .count();
9869 assert_eq!(value_columns, 1, "{field_names:?}");
9870 let plan_str = plan.display_indent_schema().to_string();
9871 assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}");
9872 assert_eq!(
9873 plan_str
9874 .matches("Filter: phy.__table_id = UInt32(1025)")
9875 .count(),
9876 1,
9877 "{plan_str}"
9878 );
9879 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9880 }
9881
9882 #[tokio::test]
9883 async fn binary_island_reuses_self_operand_without_join() {
9884 let eval_stmt = build_eval_stmt("some_metric / some_metric");
9885
9886 let table_provider = build_test_table_provider_with_tsid(
9887 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9888 1,
9889 1,
9890 )
9891 .await;
9892 let plan =
9893 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9894 .await
9895 .unwrap();
9896
9897 let plan_str = plan.display_indent_schema().to_string();
9898 assert_eq!(plan_str.matches("__tsid =").count(), 0, "{plan_str}");
9899 assert_eq!(
9900 plan_str
9901 .matches("Filter: phy.__table_id = UInt32(1024)")
9902 .count(),
9903 1,
9904 "{plan_str}"
9905 );
9906 assert_eq!(
9907 plan_str.matches("PromInstantManipulate").count(),
9908 1,
9909 "{plan_str}"
9910 );
9911 }
9912
9913 #[tokio::test]
9914 async fn binary_island_reuses_leaf_across_two_branches() {
9915 let eval_stmt =
9916 build_eval_stmt("(some_metric + some_alt_metric) / (some_metric + third_metric)");
9917
9918 let table_provider = build_test_table_provider_with_tsid(
9919 &[
9920 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9921 (
9922 DEFAULT_SCHEMA_NAME.to_string(),
9923 "some_alt_metric".to_string(),
9924 ),
9925 (DEFAULT_SCHEMA_NAME.to_string(), "third_metric".to_string()),
9926 ],
9927 1,
9928 1,
9929 )
9930 .await;
9931 let plan =
9932 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9933 .await
9934 .unwrap();
9935
9936 let plan_str = plan.display_indent_schema().to_string();
9937 assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
9938 assert_eq!(
9939 plan_str
9940 .matches("Filter: phy.__table_id = UInt32(1024)")
9941 .count(),
9942 1,
9943 "{plan_str}"
9944 );
9945 assert_eq!(
9946 plan_str.matches("PromInstantManipulate").count(),
9947 3,
9948 "{plan_str}"
9949 );
9950 }
9951
9952 #[tokio::test]
9953 async fn binary_island_generated_alias_avoids_user_column_names() {
9954 let eval_stmt = build_eval_stmt("(some_metric + some_alt_metric) / some_metric");
9955
9956 let table_provider = build_test_table_provider_with_fields(
9957 &[
9958 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9959 (
9960 DEFAULT_SCHEMA_NAME.to_string(),
9961 "some_alt_metric".to_string(),
9962 ),
9963 ],
9964 &["prom_v0", "__prom_v0"],
9965 )
9966 .await;
9967 let plan =
9968 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9969 .await
9970 .unwrap();
9971
9972 let field_names = plan.schema().field_names();
9973 assert!(field_names.iter().any(|name| name.ends_with(".prom_v0")));
9974 assert!(field_names.iter().any(|name| name.ends_with(".__prom_v0")));
9975
9976 let plan_str = plan.display_indent_schema().to_string();
9977 assert!(plan_str.contains("SubqueryAlias: __prom_v0"), "{plan_str}");
9978 assert_eq!(
9979 plan_str.matches("PromInstantManipulate").count(),
9980 2,
9981 "{plan_str}"
9982 );
9983 }
9984
9985 #[tokio::test]
9986 async fn binary_island_clears_qualifier_for_nested_unary_projection() {
9987 let eval_stmt = build_eval_stmt("-((some_metric + some_alt_metric) / some_metric)");
9988
9989 let table_provider = build_test_table_provider_with_tsid(
9990 &[
9991 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9992 (
9993 DEFAULT_SCHEMA_NAME.to_string(),
9994 "some_alt_metric".to_string(),
9995 ),
9996 ],
9997 1,
9998 1,
9999 )
10000 .await;
10001 let plan =
10002 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10003 .await
10004 .unwrap();
10005
10006 let plan_str = plan.display_indent_schema().to_string();
10007 assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}");
10008 assert_eq!(
10009 plan_str.matches("PromInstantManipulate").count(),
10010 2,
10011 "{plan_str}"
10012 );
10013 }
10014
10015 #[tokio::test]
10016 async fn binary_island_keeps_distinct_matcher_leaves() {
10017 let eval_stmt = build_eval_stmt(
10018 "(some_metric{tag_0=\"foo\"} + some_alt_metric) / some_metric{tag_0=\"bar\"}",
10019 );
10020
10021 let table_provider = build_test_table_provider_with_tsid(
10022 &[
10023 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10024 (
10025 DEFAULT_SCHEMA_NAME.to_string(),
10026 "some_alt_metric".to_string(),
10027 ),
10028 ],
10029 1,
10030 1,
10031 )
10032 .await;
10033 let plan =
10034 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10035 .await
10036 .unwrap();
10037
10038 let plan_str = plan.display_indent_schema().to_string();
10039 assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
10040 assert_eq!(
10041 plan_str.matches("PromInstantManipulate").count(),
10042 3,
10043 "{plan_str}"
10044 );
10045 }
10046
10047 #[tokio::test]
10048 async fn binary_island_keeps_offset_leaves_distinct() {
10049 let eval_stmt = build_eval_stmt("(some_metric offset 5m + some_alt_metric) / some_metric");
10050
10051 let table_provider = build_test_table_provider_with_tsid(
10052 &[
10053 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10054 (
10055 DEFAULT_SCHEMA_NAME.to_string(),
10056 "some_alt_metric".to_string(),
10057 ),
10058 ],
10059 1,
10060 1,
10061 )
10062 .await;
10063 let plan =
10064 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10065 .await
10066 .unwrap();
10067
10068 let plan_str = plan.display_indent_schema().to_string();
10069 assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
10070 assert_eq!(
10071 plan_str.matches("PromInstantManipulate").count(),
10072 3,
10073 "{plan_str}"
10074 );
10075 }
10076
10077 #[tokio::test]
10078 async fn binary_island_falls_back_for_group_modifier() {
10079 let eval_stmt = build_eval_stmt(
10080 "(some_metric + ignoring(tag_0) group_left some_alt_metric) / some_metric",
10081 );
10082
10083 let table_provider = build_test_table_provider_with_tsid(
10084 &[
10085 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10086 (
10087 DEFAULT_SCHEMA_NAME.to_string(),
10088 "some_alt_metric".to_string(),
10089 ),
10090 ],
10091 1,
10092 1,
10093 )
10094 .await;
10095 let plan =
10096 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10097 .await
10098 .unwrap();
10099
10100 let plan_str = plan.display_indent_schema().to_string();
10101 assert_eq!(
10102 plan_str.matches("PromInstantManipulate").count(),
10103 3,
10104 "{plan_str}"
10105 );
10106 }
10107
10108 #[tokio::test]
10109 async fn binary_island_falls_back_for_comparison_filter() {
10110 let eval_stmt = build_eval_stmt("(some_metric > some_alt_metric) / some_metric");
10111
10112 let table_provider = build_test_table_provider_with_tsid(
10113 &[
10114 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10115 (
10116 DEFAULT_SCHEMA_NAME.to_string(),
10117 "some_alt_metric".to_string(),
10118 ),
10119 ],
10120 1,
10121 1,
10122 )
10123 .await;
10124 let plan =
10125 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10126 .await
10127 .unwrap();
10128
10129 let plan_str = plan.display_indent_schema().to_string();
10130 assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
10131 assert_eq!(
10132 plan_str.matches("PromInstantManipulate").count(),
10133 3,
10134 "{plan_str}"
10135 );
10136 }
10137
10138 #[tokio::test]
10139 async fn tsid_binary_join_uses_shorter_field_side() {
10140 let eval_stmt = build_eval_stmt("one_field_metric / two_field_metric");
10141
10142 let table_provider = build_test_table_provider_with_tsid_fields(
10143 &[
10144 (
10145 (
10146 DEFAULT_SCHEMA_NAME.to_string(),
10147 "one_field_metric".to_string(),
10148 ),
10149 1,
10150 ),
10151 (
10152 (
10153 DEFAULT_SCHEMA_NAME.to_string(),
10154 "two_field_metric".to_string(),
10155 ),
10156 2,
10157 ),
10158 ],
10159 1,
10160 )
10161 .await;
10162 let plan =
10163 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10164 .await
10165 .unwrap();
10166
10167 let field_names = plan
10168 .schema()
10169 .fields()
10170 .iter()
10171 .map(|field| field.name().clone())
10172 .collect::<Vec<_>>();
10173 let value_columns = field_names
10174 .iter()
10175 .filter(|name| {
10176 *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME
10177 })
10178 .count();
10179 assert_eq!(value_columns, 1, "{field_names:?}");
10180 }
10181
10182 #[tokio::test]
10183 async fn comparison_binary_join_uses_shorter_field_side() {
10184 let eval_stmt = build_eval_stmt("two_field_metric > one_field_metric");
10185
10186 let table_provider = build_test_table_provider_with_tsid_fields(
10187 &[
10188 (
10189 (
10190 DEFAULT_SCHEMA_NAME.to_string(),
10191 "two_field_metric".to_string(),
10192 ),
10193 2,
10194 ),
10195 (
10196 (
10197 DEFAULT_SCHEMA_NAME.to_string(),
10198 "one_field_metric".to_string(),
10199 ),
10200 1,
10201 ),
10202 ],
10203 1,
10204 )
10205 .await;
10206 let plan =
10207 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10208 .await
10209 .unwrap();
10210
10211 let field_names = plan
10212 .schema()
10213 .fields()
10214 .iter()
10215 .map(|field| field.name().clone())
10216 .collect::<Vec<_>>();
10217 assert!(
10218 field_names.iter().any(|name| name == "field_0"),
10219 "{field_names:?}"
10220 );
10221 assert!(
10222 !field_names.iter().any(|name| name == "field_1"),
10223 "{field_names:?}"
10224 );
10225 }
10226
10227 #[tokio::test]
10228 async fn label_matching_modifier_disables_tsid_binary_join() {
10229 let eval_stmt = build_eval_stmt("some_metric / ignoring(tag_0) some_alt_metric");
10230
10231 let table_provider = build_test_table_provider_with_tsid(
10232 &[
10233 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10234 (
10235 DEFAULT_SCHEMA_NAME.to_string(),
10236 "some_alt_metric".to_string(),
10237 ),
10238 ],
10239 2,
10240 1,
10241 )
10242 .await;
10243 let plan =
10244 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10245 .await
10246 .unwrap();
10247
10248 let plan_str = plan.display_indent_schema().to_string();
10249 assert!(!plan_str.contains("__tsid ="), "{plan_str}");
10250 assert!(
10251 plan_str.contains("some_metric.tag_1 = some_alt_metric.tag_1"),
10252 "{plan_str}"
10253 );
10254 }
10255
10256 #[tokio::test]
10257 async fn ignoring_absent_label_keeps_tsid_binary_join() {
10258 let eval_stmt = build_eval_stmt("some_metric / ignoring(missing) some_alt_metric");
10259
10260 let table_provider = build_test_table_provider_with_tsid(
10261 &[
10262 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10263 (
10264 DEFAULT_SCHEMA_NAME.to_string(),
10265 "some_alt_metric".to_string(),
10266 ),
10267 ],
10268 2,
10269 1,
10270 )
10271 .await;
10272 let plan =
10273 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10274 .await
10275 .unwrap();
10276
10277 let plan_str = plan.display_indent_schema().to_string();
10278 assert!(
10279 plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
10280 "{plan_str}"
10281 );
10282 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
10283 assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
10284 }
10285
10286 #[tokio::test]
10287 async fn range_function_keeps_tsid_for_absent_ignoring_binary_join() {
10288 let eval_stmt =
10289 build_eval_stmt("rate(some_metric[5m]) / ignoring(missing) some_alt_metric");
10290
10291 let table_provider = build_test_table_provider_with_tsid(
10292 &[
10293 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10294 (
10295 DEFAULT_SCHEMA_NAME.to_string(),
10296 "some_alt_metric".to_string(),
10297 ),
10298 ],
10299 2,
10300 1,
10301 )
10302 .await;
10303 let plan =
10304 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10305 .await
10306 .unwrap();
10307
10308 let plan_str = plan.display_indent_schema().to_string();
10309 assert!(
10310 plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
10311 "{plan_str}"
10312 );
10313 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
10314 assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
10315 }
10316
10317 #[tokio::test]
10318 async fn on_full_label_set_keeps_tsid_binary_join() {
10319 let eval_stmt = build_eval_stmt("some_metric / on(tag_0, tag_1) some_alt_metric");
10320
10321 let table_provider = build_test_table_provider_with_tsid(
10322 &[
10323 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10324 (
10325 DEFAULT_SCHEMA_NAME.to_string(),
10326 "some_alt_metric".to_string(),
10327 ),
10328 ],
10329 2,
10330 1,
10331 )
10332 .await;
10333 let plan =
10334 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10335 .await
10336 .unwrap();
10337
10338 let plan_str = plan.display_indent_schema().to_string();
10339 assert!(
10340 plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
10341 "{plan_str}"
10342 );
10343 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
10344 assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
10345 }
10346
10347 #[tokio::test]
10348 async fn on_partial_label_set_disables_tsid_binary_join() {
10349 let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric");
10350
10351 let table_provider = build_test_table_provider_with_tsid(
10352 &[
10353 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10354 (
10355 DEFAULT_SCHEMA_NAME.to_string(),
10356 "some_alt_metric".to_string(),
10357 ),
10358 ],
10359 2,
10360 1,
10361 )
10362 .await;
10363 let plan =
10364 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10365 .await
10366 .unwrap();
10367
10368 let plan_str = plan.display_indent_schema().to_string();
10369 assert!(!plan_str.contains("__tsid ="), "{plan_str}");
10370 assert!(
10371 plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
10372 "{plan_str}"
10373 );
10374 assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
10375 }
10376
10377 #[tokio::test]
10378 async fn on_label_set_must_cover_both_sides_to_use_tsid_binary_join() {
10379 let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric");
10380
10381 let table_provider = build_test_table_provider_with_tsid_tag_fields(&[
10382 (
10383 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10384 2,
10385 1,
10386 ),
10387 (
10388 (
10389 DEFAULT_SCHEMA_NAME.to_string(),
10390 "some_alt_metric".to_string(),
10391 ),
10392 1,
10393 1,
10394 ),
10395 ])
10396 .await;
10397 let plan =
10398 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10399 .await
10400 .unwrap();
10401
10402 let plan_str = plan.display_indent_schema().to_string();
10403 assert!(!plan_str.contains("__tsid ="), "{plan_str}");
10404 assert!(
10405 plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
10406 "{plan_str}"
10407 );
10408 assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
10409 }
10410
10411 #[tokio::test]
10412 async fn comparison_binary_join_uses_tsid_and_keeps_it_in_filtered_result() {
10413 let eval_stmt = build_eval_stmt("some_metric > some_alt_metric");
10414
10415 let table_provider = build_test_table_provider_with_tsid(
10416 &[
10417 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10418 (
10419 DEFAULT_SCHEMA_NAME.to_string(),
10420 "some_alt_metric".to_string(),
10421 ),
10422 ],
10423 2,
10424 1,
10425 )
10426 .await;
10427 let mut planner = PromPlanner {
10428 table_provider,
10429 ctx: PromPlannerContext::from_eval_stmt(&eval_stmt),
10430 promql_annotations: None,
10431 };
10432 let plan = planner
10433 .prom_expr_to_plan(&eval_stmt.expr, &build_query_engine_state())
10434 .await
10435 .unwrap();
10436
10437 let plan_str = plan.display_indent_schema().to_string();
10438 assert!(
10439 plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
10440 "{plan_str}"
10441 );
10442 assert!(
10443 plan.schema()
10444 .fields()
10445 .iter()
10446 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME),
10447 "{plan_str}"
10448 );
10449 assert!(planner.ctx.use_tsid, "{plan_str}");
10450 }
10451
10452 #[tokio::test]
10453 async fn comparison_bool_binary_join_uses_tsid_when_available() {
10454 let eval_stmt = build_eval_stmt("some_metric > bool some_alt_metric");
10455
10456 let table_provider = build_test_table_provider_with_tsid(
10457 &[
10458 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10459 (
10460 DEFAULT_SCHEMA_NAME.to_string(),
10461 "some_alt_metric".to_string(),
10462 ),
10463 ],
10464 2,
10465 1,
10466 )
10467 .await;
10468 let plan =
10469 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10470 .await
10471 .unwrap();
10472
10473 let plan_str = plan.display_indent_schema().to_string();
10474 assert!(
10475 plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
10476 "{plan_str}"
10477 );
10478 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
10479 assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
10480 }
10481
10482 #[tokio::test]
10483 async fn scalar_count_count_range_keeps_full_window() {
10484 let plan_str = build_optimized_tsid_plan(
10485 "scalar(count(count(some_metric) by (tag_0)))",
10486 1,
10487 1,
10488 100_000,
10489 1,
10490 )
10491 .await;
10492 assert!(plan_str.contains("ScalarCalculate: tags=[]"));
10493 assert!(plan_str.contains("PromInstantManipulate: range=[0..100000000]"));
10494 assert!(!plan_str.contains("PromInstantManipulate: range=[99999000..99999000]"));
10495 }
10496
10497 #[tokio::test]
10498 async fn scalar_count_count_rewrite_applies_inside_binary_expr_for_tsid_input() {
10499 let plan_str = build_optimized_tsid_plan(
10500 "sum(irate(some_metric[1h])) / scalar(count(count(some_metric) by (tag_0)))",
10501 2,
10502 1,
10503 10,
10504 300,
10505 )
10506 .await;
10507 assert!(plan_str.contains("Distinct:"), "{plan_str}");
10508 }
10509
10510 #[tokio::test]
10511 async fn nested_count_rewrite_keeps_full_series_key_with_tsid_input() {
10512 assert_nested_count_rewrite_applies(
10513 "count(count(some_metric) by (tag_0))",
10514 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(count(some_metric.field_0))]]"
10515 )
10516 .await;
10517 }
10518
10519 #[tokio::test]
10520 async fn nested_sum_count_rewrite_keeps_full_series_key_with_tsid_input() {
10521 assert_nested_count_rewrite_applies(
10522 "count(sum(some_metric) by (tag_0))",
10523 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(sum(some_metric.field_0))]]"
10524 )
10525 .await;
10526 }
10527
10528 #[tokio::test]
10529 async fn nested_supported_inner_aggs_rewrite_apply_for_tsid_input() {
10530 for (query, expected_outer_agg) in [
10531 (
10532 "count(avg(some_metric) by (tag_0))",
10533 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(avg(some_metric.field_0))]]",
10534 ),
10535 (
10536 "count(min(some_metric) by (tag_0))",
10537 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(min(some_metric.field_0))]]",
10538 ),
10539 (
10540 "count(max(some_metric) by (tag_0))",
10541 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(max(some_metric.field_0))]]",
10542 ),
10543 (
10544 "count(stddev(some_metric) by (tag_0))",
10545 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(stddev_pop(some_metric.field_0))]]",
10546 ),
10547 (
10548 "count(stdvar(some_metric) by (tag_0))",
10549 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(var_pop(some_metric.field_0))]]",
10550 ),
10551 ] {
10552 assert_nested_count_rewrite_applies(query, expected_outer_agg).await;
10553 }
10554 }
10555
10556 #[tokio::test]
10557 async fn nested_non_count_inner_aggs_rewrite_filter_null_values_for_tsid_input() {
10558 let count_plan =
10559 build_optimized_tsid_plan("count(count(some_metric) by (tag_0))", 2, 1, 100_000, 1)
10560 .await;
10561 assert!(
10562 !count_plan.contains("some_metric.field_0 IS NOT NULL"),
10563 "{count_plan}"
10564 );
10565
10566 for query in [
10567 "count(sum(some_metric) by (tag_0))",
10568 "count(avg(some_metric) by (tag_0))",
10569 "count(min(some_metric) by (tag_0))",
10570 "count(max(some_metric) by (tag_0))",
10571 "count(stddev(some_metric) by (tag_0))",
10572 "count(stdvar(some_metric) by (tag_0))",
10573 ] {
10574 let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await;
10575 assert!(
10576 plan_str.contains("Filter: some_metric.field_0 IS NOT NULL"),
10577 "{query}: {plan_str}"
10578 );
10579 }
10580 }
10581
10582 #[tokio::test]
10583 async fn nested_unsupported_or_non_direct_inner_aggs_do_not_rewrite() {
10584 assert_nested_count_rewrite_missing("count(group(some_metric) by (tag_0))", 2, 1).await;
10585 assert_nested_count_rewrite_missing(
10586 "count(sum(irate(some_metric[1h])) by (tag_0))",
10587 2,
10588 300,
10589 )
10590 .await;
10591 }
10592
10593 #[tokio::test]
10594 async fn physical_table_name_is_not_leaked_in_plan() {
10595 let prom_expr = parser::parse("some_metric").unwrap();
10596 let eval_stmt = EvalStmt {
10597 expr: prom_expr,
10598 start: UNIX_EPOCH,
10599 end: UNIX_EPOCH
10600 .checked_add(Duration::from_secs(100_000))
10601 .unwrap(),
10602 interval: Duration::from_secs(5),
10603 lookback_delta: Duration::from_secs(1),
10604 };
10605
10606 let table_provider = build_test_table_provider_with_tsid(
10607 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10608 1,
10609 1,
10610 )
10611 .await;
10612 let plan =
10613 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10614 .await
10615 .unwrap();
10616
10617 let plan_str = plan.display_indent_schema().to_string();
10618 assert!(plan_str.contains("TableScan: phy"), "{plan}");
10619 assert!(plan_str.contains("SubqueryAlias: some_metric"));
10620 assert!(plan_str.contains("Filter: phy.__table_id = UInt32(1024)"));
10621 assert!(!plan_str.contains("TableScan: some_metric"));
10622 }
10623
10624 #[tokio::test]
10625 async fn sum_without_does_not_group_by_tsid() {
10626 let prom_expr = parser::parse("sum without (tag_0) (some_metric)").unwrap();
10627 let eval_stmt = EvalStmt {
10628 expr: prom_expr,
10629 start: UNIX_EPOCH,
10630 end: UNIX_EPOCH
10631 .checked_add(Duration::from_secs(100_000))
10632 .unwrap(),
10633 interval: Duration::from_secs(5),
10634 lookback_delta: Duration::from_secs(1),
10635 };
10636
10637 let table_provider = build_test_table_provider_with_tsid(
10638 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10639 1,
10640 1,
10641 )
10642 .await;
10643 let plan =
10644 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10645 .await
10646 .unwrap();
10647
10648 let plan_str = plan.display_indent_schema().to_string();
10649 assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10650
10651 let aggr_line = plan_str
10652 .lines()
10653 .find(|line| line.contains("Aggregate: groupBy="))
10654 .unwrap();
10655 assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
10656 }
10657
10658 #[tokio::test]
10659 async fn topk_without_does_not_partition_by_tsid() {
10660 let prom_expr = parser::parse("topk without (tag_0) (1, some_metric)").unwrap();
10661 let eval_stmt = EvalStmt {
10662 expr: prom_expr,
10663 start: UNIX_EPOCH,
10664 end: UNIX_EPOCH
10665 .checked_add(Duration::from_secs(100_000))
10666 .unwrap(),
10667 interval: Duration::from_secs(5),
10668 lookback_delta: Duration::from_secs(1),
10669 };
10670
10671 let table_provider = build_test_table_provider_with_tsid(
10672 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10673 1,
10674 1,
10675 )
10676 .await;
10677 let plan =
10678 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10679 .await
10680 .unwrap();
10681
10682 let plan_str = plan.display_indent_schema().to_string();
10683 assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10684
10685 let window_line = plan_str
10686 .lines()
10687 .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()"))
10688 .unwrap();
10689 let partition_by = window_line
10690 .split("PARTITION BY [")
10691 .nth(1)
10692 .and_then(|s| s.split("] ORDER BY").next())
10693 .unwrap();
10694 assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
10695 }
10696
10697 #[tokio::test]
10698 async fn sum_by_does_not_group_by_tsid() {
10699 let prom_expr = parser::parse("sum by (__tsid) (some_metric)").unwrap();
10700 let eval_stmt = EvalStmt {
10701 expr: prom_expr,
10702 start: UNIX_EPOCH,
10703 end: UNIX_EPOCH
10704 .checked_add(Duration::from_secs(100_000))
10705 .unwrap(),
10706 interval: Duration::from_secs(5),
10707 lookback_delta: Duration::from_secs(1),
10708 };
10709
10710 let table_provider = build_test_table_provider_with_tsid(
10711 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10712 1,
10713 1,
10714 )
10715 .await;
10716 let plan =
10717 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10718 .await
10719 .unwrap();
10720
10721 let plan_str = plan.display_indent_schema().to_string();
10722 assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10723
10724 let aggr_line = plan_str
10725 .lines()
10726 .find(|line| line.contains("Aggregate: groupBy="))
10727 .unwrap();
10728 assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
10729 }
10730
10731 #[tokio::test]
10732 async fn aggregate_over_binary_time_function_expr() {
10733 for op in ["sum", "min", "max", "avg"] {
10734 let prom_expr = parser::parse(&format!(
10735 "{op} by (tag_0, tag_1, tag_2) (time() - some_metric)"
10736 ))
10737 .unwrap();
10738 let eval_stmt = EvalStmt {
10739 expr: prom_expr,
10740 start: UNIX_EPOCH,
10741 end: UNIX_EPOCH
10742 .checked_add(Duration::from_secs(100_000))
10743 .unwrap(),
10744 interval: Duration::from_secs(5),
10745 lookback_delta: Duration::from_secs(1),
10746 };
10747
10748 let table_provider = build_test_table_provider_with_tsid(
10749 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10750 3,
10751 1,
10752 )
10753 .await;
10754 let plan =
10755 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10756 .await
10757 .unwrap();
10758
10759 let plan_str = plan.display_indent_schema().to_string();
10760 let aggr_line = plan_str
10761 .lines()
10762 .find(|line| line.contains("Aggregate: groupBy="))
10763 .unwrap();
10764 assert!(aggr_line.contains(op), "{plan_str}");
10765 assert!(aggr_line.contains("first_value"), "{plan_str}");
10766 assert!(
10767 !plan
10768 .schema()
10769 .fields()
10770 .iter()
10771 .any(|field| { field.name() == DATA_SCHEMA_TSID_COLUMN_NAME })
10772 );
10773 }
10774 }
10775
10776 #[tokio::test]
10777 async fn topk_by_does_not_partition_by_tsid() {
10778 let prom_expr = parser::parse("topk by (__tsid) (1, some_metric)").unwrap();
10779 let eval_stmt = EvalStmt {
10780 expr: prom_expr,
10781 start: UNIX_EPOCH,
10782 end: UNIX_EPOCH
10783 .checked_add(Duration::from_secs(100_000))
10784 .unwrap(),
10785 interval: Duration::from_secs(5),
10786 lookback_delta: Duration::from_secs(1),
10787 };
10788
10789 let table_provider = build_test_table_provider_with_tsid(
10790 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10791 1,
10792 1,
10793 )
10794 .await;
10795 let plan =
10796 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10797 .await
10798 .unwrap();
10799
10800 let plan_str = plan.display_indent_schema().to_string();
10801 assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10802
10803 let window_line = plan_str
10804 .lines()
10805 .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()"))
10806 .unwrap();
10807 let partition_by = window_line
10808 .split("PARTITION BY [")
10809 .nth(1)
10810 .and_then(|s| s.split("] ORDER BY").next())
10811 .unwrap();
10812 assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
10813 }
10814
10815 #[tokio::test]
10816 async fn selector_matcher_on_tsid_does_not_use_internal_column() {
10817 let prom_expr = parser::parse(r#"some_metric{__tsid="123"}"#).unwrap();
10818 let eval_stmt = EvalStmt {
10819 expr: prom_expr,
10820 start: UNIX_EPOCH,
10821 end: UNIX_EPOCH
10822 .checked_add(Duration::from_secs(100_000))
10823 .unwrap(),
10824 interval: Duration::from_secs(5),
10825 lookback_delta: Duration::from_secs(1),
10826 };
10827
10828 let table_provider = build_test_table_provider_with_tsid(
10829 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10830 1,
10831 1,
10832 )
10833 .await;
10834 let plan =
10835 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10836 .await
10837 .unwrap();
10838
10839 fn collect_filter_cols(plan: &LogicalPlan, out: &mut HashSet<Column>) {
10840 if let LogicalPlan::Filter(filter) = plan {
10841 datafusion_expr::utils::expr_to_columns(&filter.predicate, out).unwrap();
10842 }
10843 for input in plan.inputs() {
10844 collect_filter_cols(input, out);
10845 }
10846 }
10847
10848 let mut filter_cols = HashSet::new();
10849 collect_filter_cols(&plan, &mut filter_cols);
10850 assert!(
10851 !filter_cols
10852 .iter()
10853 .any(|c| c.name == DATA_SCHEMA_TSID_COLUMN_NAME)
10854 );
10855 }
10856
10857 #[tokio::test]
10858 async fn tsid_is_not_used_when_physical_table_is_missing() {
10859 let prom_expr = parser::parse("some_metric").unwrap();
10860 let eval_stmt = EvalStmt {
10861 expr: prom_expr,
10862 start: UNIX_EPOCH,
10863 end: UNIX_EPOCH
10864 .checked_add(Duration::from_secs(100_000))
10865 .unwrap(),
10866 interval: Duration::from_secs(5),
10867 lookback_delta: Duration::from_secs(1),
10868 };
10869
10870 let catalog_list = MemoryCatalogManager::with_default_setup();
10871
10872 let mut columns = vec![ColumnSchema::new(
10874 "tag_0".to_string(),
10875 ConcreteDataType::string_datatype(),
10876 false,
10877 )];
10878 columns.push(
10879 ColumnSchema::new(
10880 "timestamp".to_string(),
10881 ConcreteDataType::timestamp_millisecond_datatype(),
10882 false,
10883 )
10884 .with_time_index(true),
10885 );
10886 columns.push(ColumnSchema::new(
10887 "field_0".to_string(),
10888 ConcreteDataType::float64_datatype(),
10889 true,
10890 ));
10891 let schema = Arc::new(Schema::new(columns));
10892 let mut options = table::requests::TableOptions::default();
10893 options
10894 .extra_options
10895 .insert(LOGICAL_TABLE_METADATA_KEY.to_string(), "phy".to_string());
10896 let table_meta = TableMetaBuilder::empty()
10897 .schema(schema)
10898 .primary_key_indices(vec![0])
10899 .value_indices(vec![2])
10900 .engine(METRIC_ENGINE_NAME.to_string())
10901 .options(options)
10902 .next_column_id(1024)
10903 .build()
10904 .unwrap();
10905 let table_info = TableInfoBuilder::default()
10906 .table_id(1024)
10907 .name("some_metric")
10908 .meta(table_meta)
10909 .build()
10910 .unwrap();
10911 let table = EmptyTable::from_table_info(&table_info);
10912 catalog_list
10913 .register_table_sync(RegisterTableRequest {
10914 catalog: DEFAULT_CATALOG_NAME.to_string(),
10915 schema: DEFAULT_SCHEMA_NAME.to_string(),
10916 table_name: "some_metric".to_string(),
10917 table_id: 1024,
10918 table,
10919 })
10920 .unwrap();
10921
10922 let table_provider = DfTableSourceProvider::new(
10923 catalog_list,
10924 false,
10925 QueryContext::arc(),
10926 DummyDecoder::arc(),
10927 false,
10928 );
10929
10930 let plan =
10931 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10932 .await
10933 .unwrap();
10934
10935 let plan_str = plan.display_indent_schema().to_string();
10936 assert!(plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]"));
10937 assert!(!plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10938 }
10939
10940 #[tokio::test]
10941 async fn tsid_is_carried_only_when_aggregate_preserves_label_set() {
10942 let prom_expr = parser::parse("sum by (tag_0) (some_metric)").unwrap();
10943 let eval_stmt = EvalStmt {
10944 expr: prom_expr,
10945 start: UNIX_EPOCH,
10946 end: UNIX_EPOCH
10947 .checked_add(Duration::from_secs(100_000))
10948 .unwrap(),
10949 interval: Duration::from_secs(5),
10950 lookback_delta: Duration::from_secs(1),
10951 };
10952
10953 let table_provider = build_test_table_provider_with_tsid(
10954 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10955 1,
10956 1,
10957 )
10958 .await;
10959 let plan =
10960 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10961 .await
10962 .unwrap();
10963
10964 let plan_str = plan.display_indent_schema().to_string();
10965 assert!(plan_str.contains("first_value") && plan_str.contains("__tsid"));
10966 assert!(
10967 !plan
10968 .schema()
10969 .fields()
10970 .iter()
10971 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
10972 );
10973
10974 let prom_expr = parser::parse("sum(some_metric)").unwrap();
10976 let eval_stmt = EvalStmt {
10977 expr: prom_expr,
10978 start: UNIX_EPOCH,
10979 end: UNIX_EPOCH
10980 .checked_add(Duration::from_secs(100_000))
10981 .unwrap(),
10982 interval: Duration::from_secs(5),
10983 lookback_delta: Duration::from_secs(1),
10984 };
10985 let table_provider = build_test_table_provider_with_tsid(
10986 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10987 1,
10988 1,
10989 )
10990 .await;
10991 let plan =
10992 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10993 .await
10994 .unwrap();
10995 let plan_str = plan.display_indent_schema().to_string();
10996 assert!(!plan_str.contains("first_value"));
10997 }
10998
10999 #[tokio::test]
11000 async fn or_operator_with_unknown_metric_does_not_require_tsid() {
11001 let prom_expr = parser::parse("unknown_metric or some_metric").unwrap();
11002 let eval_stmt = EvalStmt {
11003 expr: prom_expr,
11004 start: UNIX_EPOCH,
11005 end: UNIX_EPOCH
11006 .checked_add(Duration::from_secs(100_000))
11007 .unwrap(),
11008 interval: Duration::from_secs(5),
11009 lookback_delta: Duration::from_secs(1),
11010 };
11011
11012 let table_provider = build_test_table_provider_with_tsid(
11013 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
11014 1,
11015 1,
11016 )
11017 .await;
11018
11019 let plan =
11020 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11021 .await
11022 .unwrap();
11023
11024 assert!(
11025 !plan
11026 .schema()
11027 .fields()
11028 .iter()
11029 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
11030 );
11031 }
11032
11033 #[tokio::test]
11034 async fn aggregate_avg() {
11035 do_aggregate_expr_plan("avg", "avg").await;
11036 }
11037
11038 #[tokio::test]
11039 #[should_panic] async fn aggregate_count() {
11041 do_aggregate_expr_plan("count", "count").await;
11042 }
11043
11044 #[tokio::test]
11045 async fn aggregate_min() {
11046 do_aggregate_expr_plan("min", "min").await;
11047 }
11048
11049 #[tokio::test]
11050 async fn aggregate_max() {
11051 do_aggregate_expr_plan("max", "max").await;
11052 }
11053
11054 #[tokio::test]
11055 async fn aggregate_group() {
11056 let prom_expr = parser::parse(
11060 "sum(group by (cluster)(kubernetes_build_info{service=\"kubernetes\",job=\"apiserver\"}))",
11061 )
11062 .unwrap();
11063 let eval_stmt = EvalStmt {
11064 expr: prom_expr,
11065 start: UNIX_EPOCH,
11066 end: UNIX_EPOCH
11067 .checked_add(Duration::from_secs(100_000))
11068 .unwrap(),
11069 interval: Duration::from_secs(5),
11070 lookback_delta: Duration::from_secs(1),
11071 };
11072
11073 let table_provider = build_test_table_provider_with_fields(
11074 &[(
11075 DEFAULT_SCHEMA_NAME.to_string(),
11076 "kubernetes_build_info".to_string(),
11077 )],
11078 &["cluster", "service", "job"],
11079 )
11080 .await;
11081 let plan =
11082 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11083 .await
11084 .unwrap();
11085
11086 let plan_str = plan.display_indent_schema().to_string();
11087 assert!(plan_str.contains("max(Float64(1"));
11088 }
11089
11090 #[tokio::test]
11091 async fn aggregate_stddev() {
11092 do_aggregate_expr_plan("stddev", "stddev_pop").await;
11093 }
11094
11095 #[tokio::test]
11096 async fn aggregate_stdvar() {
11097 do_aggregate_expr_plan("stdvar", "var_pop").await;
11098 }
11099
11100 #[tokio::test]
11124 async fn binary_op_column_column() {
11125 let prom_expr =
11126 parser::parse(r#"some_metric{tag_0="foo"} + some_metric{tag_0="bar"}"#).unwrap();
11127 let eval_stmt = EvalStmt {
11128 expr: prom_expr,
11129 start: UNIX_EPOCH,
11130 end: UNIX_EPOCH
11131 .checked_add(Duration::from_secs(100_000))
11132 .unwrap(),
11133 interval: Duration::from_secs(5),
11134 lookback_delta: Duration::from_secs(1),
11135 };
11136
11137 let table_provider = build_test_table_provider(
11138 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
11139 1,
11140 1,
11141 )
11142 .await;
11143 let plan =
11144 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11145 .await
11146 .unwrap();
11147
11148 let expected = String::from(
11149 "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]\
11150 \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]\
11151 \n SubqueryAlias: lhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11152 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11153 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11154 \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]\
11155 \n Filter: some_metric.tag_0 = Utf8(\"foo\") AND some_metric.tag_0 = Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11156 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11157 \n SubqueryAlias: rhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11158 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11159 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11160 \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]\
11161 \n Filter: some_metric.tag_0 = Utf8(\"bar\") AND some_metric.tag_0 = Utf8(\"foo\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11162 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
11163 );
11164
11165 assert_eq!(plan.display_indent_schema().to_string(), expected);
11166 }
11167
11168 async fn indie_query_plan_compare<T: AsRef<str>>(query: &str, expected: T) {
11169 let prom_expr = parser::parse(query).unwrap();
11170 let eval_stmt = EvalStmt {
11171 expr: prom_expr,
11172 start: UNIX_EPOCH,
11173 end: UNIX_EPOCH
11174 .checked_add(Duration::from_secs(100_000))
11175 .unwrap(),
11176 interval: Duration::from_secs(5),
11177 lookback_delta: Duration::from_secs(1),
11178 };
11179
11180 let table_provider = build_test_table_provider(
11181 &[
11182 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
11183 (
11184 "greptime_private".to_string(),
11185 "some_alt_metric".to_string(),
11186 ),
11187 ],
11188 1,
11189 1,
11190 )
11191 .await;
11192 let plan =
11193 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11194 .await
11195 .unwrap();
11196
11197 assert_eq!(plan.display_indent_schema().to_string(), expected.as_ref());
11198 }
11199
11200 #[tokio::test]
11201 async fn binary_op_literal_column() {
11202 let query = r#"1 + some_metric{tag_0="bar"}"#;
11203 let expected = String::from(
11204 "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]\
11205 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11206 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11207 \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]\
11208 \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]\
11209 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
11210 );
11211
11212 indie_query_plan_compare(query, expected).await;
11213 }
11214
11215 #[tokio::test]
11216 async fn binary_op_literal_literal() {
11217 let query = r#"1 + 1"#;
11218 let expected = r#"EmptyMetric: range=[0..100000000], interval=[5000] [time:Timestamp(ms), value:Float64;N]
11219 TableScan: dummy [time:Timestamp(ms), value:Float64;N]"#;
11220 indie_query_plan_compare(query, expected).await;
11221 }
11222
11223 #[tokio::test]
11224 async fn simple_bool_grammar() {
11225 let query = "some_metric != bool 1.2345";
11226 let expected = String::from(
11227 "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]\
11228 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11229 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11230 \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]\
11231 \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]\
11232 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
11233 );
11234
11235 indie_query_plan_compare(query, expected).await;
11236 }
11237
11238 #[tokio::test]
11239 async fn bool_with_additional_arithmetic() {
11240 let query = "some_metric + (1 == bool 2)";
11241 let expected = String::from(
11242 "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]\
11243 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11244 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11245 \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]\
11246 \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]\
11247 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
11248 );
11249
11250 indie_query_plan_compare(query, expected).await;
11251 }
11252
11253 #[tokio::test]
11254 async fn simple_unary() {
11255 let query = "-some_metric";
11256 let expected = String::from(
11257 "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]\
11258 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11259 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11260 \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]\
11261 \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]\
11262 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
11263 );
11264
11265 indie_query_plan_compare(query, expected).await;
11266 }
11267
11268 #[tokio::test]
11269 async fn increase_aggr() {
11270 let query = "increase(some_metric[5m])";
11271 let expected = String::from(
11272 "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]\
11273 \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]\
11274 \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))]\
11275 \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11276 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11277 \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]\
11278 \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]\
11279 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
11280 );
11281
11282 indie_query_plan_compare(query, expected).await;
11283 }
11284
11285 async fn native_histogram_plan(query: &str) -> String {
11286 let table_provider = build_test_native_histogram_table_provider("some_metric").await;
11287 let plan = PromPlanner::stmt_to_plan(
11288 table_provider,
11289 &build_eval_stmt(query),
11290 &build_query_engine_state(),
11291 )
11292 .await
11293 .unwrap();
11294 plan.display_indent_schema().to_string()
11295 }
11296
11297 #[tokio::test]
11298 async fn native_histogram_count_uses_native_udf() {
11299 let plan = native_histogram_plan("histogram_count(some_metric)").await;
11300
11301 assert!(plan.contains("prom_native_histogram_count"), "{plan}");
11302 assert!(!plan.contains("HistogramFold:"), "{plan}");
11303 }
11304
11305 #[tokio::test]
11306 async fn timestamp_filters_native_histogram_stale_marker_before_projection() {
11307 let mut stale = direct_or_histogram();
11308 stale.sum = f64::from_bits(PROMETHEUS_STALE_NAN_BITS);
11309 let table = operator_metric_table(
11310 "stale_histogram",
11311 2_100,
11312 "a",
11313 None,
11314 DirectOrValue::NativeHistogram(stale),
11315 );
11316 let catalog = MemoryCatalogManager::with_default_setup();
11317 catalog
11318 .register_table_sync(RegisterTableRequest {
11319 catalog: DEFAULT_CATALOG_NAME.to_string(),
11320 schema: DEFAULT_SCHEMA_NAME.to_string(),
11321 table_name: "stale_histogram".to_string(),
11322 table_id: 2_100,
11323 table,
11324 })
11325 .unwrap();
11326 let provider = DfTableSourceProvider::new(
11327 catalog,
11328 false,
11329 QueryContext::arc(),
11330 DummyDecoder::arc(),
11331 false,
11332 );
11333 let state = build_query_engine_state();
11334 let plan = PromPlanner::stmt_to_plan(
11335 provider,
11336 &operator_eval_stmt("timestamp(stale_histogram)"),
11337 &state,
11338 )
11339 .await
11340 .unwrap();
11341 let plan_text = plan.display_indent_schema().to_string();
11342 assert!(plan_text.contains(TIMESTAMP_VALUE_PREFIX), "{plan_text}");
11343
11344 let (_, batches) = execute(plan, &state).await;
11345 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
11346 }
11347
11348 #[tokio::test]
11349 async fn timestamp_filters_stale_marker_from_mixed_sample_companion() {
11350 let histograms = build_histogram_array(&[None]);
11351 let schema = Arc::new(ArrowSchema::new(vec![
11352 Field::new(
11353 "timestamp",
11354 ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
11355 false,
11356 ),
11357 Field::new(
11358 greptime_native_histogram(),
11359 histograms.data_type().clone(),
11360 true,
11361 ),
11362 Field::new(greptime_value(), ArrowDataType::Float64, true),
11363 ]));
11364 let batch = RecordBatch::try_new(
11365 schema.clone(),
11366 vec![
11367 Arc::new(TimestampMillisecondArray::from(vec![1_000])),
11368 histograms,
11369 Arc::new(Float64Array::from(vec![f64::from_bits(
11370 PROMETHEUS_STALE_NAN_BITS,
11371 )])),
11372 ],
11373 )
11374 .unwrap();
11375 let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap());
11376 let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None)
11377 .unwrap()
11378 .build()
11379 .unwrap();
11380 let input = LogicalPlan::Extension(Extension {
11381 node: Arc::new(SeriesDivide::new(
11382 Vec::new(),
11383 "timestamp".to_string(),
11384 input,
11385 )),
11386 });
11387 let input = LogicalPlan::Extension(Extension {
11388 node: Arc::new(InstantManipulate::new(
11389 1_000,
11390 1_000,
11391 5_000,
11392 1_000,
11393 0,
11394 "timestamp".to_string(),
11395 Vec::new(),
11396 Some(greptime_native_histogram().to_string()),
11397 input,
11398 )),
11399 });
11400 let plan = LogicalPlanBuilder::from(input)
11402 .project([col("timestamp")])
11403 .unwrap()
11404 .build()
11405 .unwrap();
11406
11407 let (_, batches) = execute(plan, &build_query_engine_state()).await;
11408 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
11409 }
11410
11411 #[tokio::test]
11412 async fn native_histogram_rate_can_feed_count() {
11413 let plan = native_histogram_plan("histogram_count(rate(some_metric[5m]))").await;
11414
11415 assert!(plan.contains("prom_native_histogram_rate"), "{plan}");
11416 assert!(plan.contains("prom_native_histogram_count"), "{plan}");
11417 }
11418
11419 #[tokio::test]
11420 async fn native_histogram_quantile_skips_classic_fold() {
11421 let plan = native_histogram_plan("histogram_quantile(0.9, some_metric)").await;
11422
11423 assert!(plan.contains("prom_native_histogram_quantile"), "{plan}");
11424 assert!(!plan.contains("HistogramFold:"), "{plan}");
11425 assert!(plan.contains("some_metric.le"), "{plan}");
11426 assert!(plan.contains("Float64(0.9)"), "{plan}");
11428 assert!(plan.contains("IS NOT NULL"), "{plan}");
11431 }
11432
11433 #[tokio::test]
11434 async fn mixed_native_histogram_quantile_uses_histogram_field() {
11435 let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
11436 let plan = PromPlanner::stmt_to_plan(
11437 table_provider,
11438 &build_eval_stmt("histogram_quantile(0.9, some_metric)"),
11439 &build_query_engine_state(),
11440 )
11441 .await
11442 .unwrap()
11443 .display_indent_schema()
11444 .to_string();
11445
11446 assert!(
11447 plan.contains("prom_native_histogram_quantile(greptime_native_histogram"),
11448 "{plan}"
11449 );
11450 assert!(!plan.contains("EmptyRelation"), "{plan}");
11451 }
11452
11453 #[tokio::test]
11454 async fn mixed_histogram_helpers_execute_classic_and_native_samples() {
11455 let state = build_query_engine_state();
11456 for (query, expected) in [
11457 (
11458 "histogram_quantile(0.5, mixed_histogram)",
11459 vec![("classic", 1.0), ("native", 0.0)],
11460 ),
11461 (
11462 "histogram_fraction(-Inf, +Inf, mixed_histogram)",
11463 vec![("classic", 1.0), ("native", 1.0)],
11464 ),
11465 ] {
11466 let plan = PromPlanner::stmt_to_plan(
11467 classic_and_native_histogram_table_provider("native", None, direct_or_histogram()),
11468 &operator_eval_stmt(query),
11469 &state,
11470 )
11471 .await
11472 .unwrap();
11473 let plan_text = plan.display_indent_schema().to_string();
11474 assert!(plan_text.contains("HistogramFold:"), "{plan_text}");
11475 assert!(plan_text.contains("prom_native_histogram_"), "{plan_text}");
11476 let value_field = plan
11477 .schema()
11478 .fields()
11479 .iter()
11480 .find(|field| field.data_type() == &ArrowDataType::Float64)
11481 .unwrap()
11482 .name()
11483 .clone();
11484
11485 let (_, batches) = execute(plan, &state).await;
11486 let mut actual = batches
11487 .iter()
11488 .flat_map(|batch| {
11489 let tags = batch
11490 .column_by_name("tag")
11491 .unwrap()
11492 .as_any()
11493 .downcast_ref::<StringArray>()
11494 .unwrap();
11495 let values = batch
11496 .column_by_name(&value_field)
11497 .unwrap()
11498 .as_any()
11499 .downcast_ref::<Float64Array>()
11500 .unwrap();
11501 (0..batch.num_rows()).map(|row| (tags.value(row), values.value(row)))
11502 })
11503 .collect::<Vec<_>>();
11504 actual.sort_by_key(|(tag, _)| *tag);
11505 assert_eq!(actual, expected, "{query}");
11506 }
11507 }
11508
11509 #[tokio::test]
11510 async fn mixed_histogram_helpers_report_annotations() {
11511 let state = build_query_engine_state();
11512 let mut native_histogram = direct_or_histogram();
11513 native_histogram.count = 2.0;
11514 native_histogram.sum = f64::NAN;
11515 for (native_tag, expected_rows, expected_warnings, expected_infos) in [
11516 (
11517 "classic",
11518 0,
11519 vec!["vector contains a mix of classic and native histograms"],
11520 vec![],
11521 ),
11522 (
11523 "native",
11524 2,
11525 vec![],
11526 vec!["input to histogram_quantile has NaN observations, result is skewed higher"],
11527 ),
11528 ] {
11529 let collector = PromqlAnnotationCollector::default();
11530 let plan = PromPlanner::stmt_to_plan_with_annotations(
11531 classic_and_native_histogram_table_provider(
11532 native_tag,
11533 None,
11534 native_histogram.clone(),
11535 ),
11536 &operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"),
11537 &state,
11538 Some(collector.clone()),
11539 )
11540 .await
11541 .unwrap();
11542
11543 let (_, batches) = execute(plan, &state).await;
11544 assert_eq!(
11545 batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
11546 expected_rows
11547 );
11548 let mut warnings = vec![];
11549 let mut infos = vec![];
11550 collector.append_to(&mut warnings, &mut infos);
11551 assert_eq!(warnings, expected_warnings);
11552 assert_eq!(infos, expected_infos);
11553 }
11554 }
11555
11556 #[tokio::test]
11557 async fn mixed_histogram_helper_preserves_native_le_and_scans_once() {
11558 let state = build_query_engine_state();
11559 let mut stmt = operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)");
11560 stmt.end = UNIX_EPOCH.checked_add(Duration::from_secs(2)).unwrap();
11561 let plan = PromPlanner::stmt_to_plan(
11562 classic_and_native_histogram_table_provider(
11563 "classic",
11564 Some("native"),
11565 direct_or_histogram(),
11566 ),
11567 &stmt,
11568 &state,
11569 )
11570 .await
11571 .unwrap();
11572 let plan_text = plan.display_indent_schema().to_string();
11573 assert_eq!(
11574 plan_text.matches("TableScan: mixed_histogram").count(),
11575 1,
11576 "{plan_text}"
11577 );
11578
11579 let value_field = plan
11580 .schema()
11581 .fields()
11582 .iter()
11583 .find(|field| field.data_type() == &ArrowDataType::Float64)
11584 .unwrap()
11585 .name()
11586 .clone();
11587 let (_, batches) = execute(plan, &state).await;
11588 let mut actual = batches
11589 .iter()
11590 .flat_map(|batch| {
11591 let le = batch
11592 .column_by_name(LE_COLUMN_NAME)
11593 .unwrap()
11594 .as_any()
11595 .downcast_ref::<StringArray>()
11596 .unwrap();
11597 let timestamps = batch
11598 .column_by_name("timestamp")
11599 .unwrap()
11600 .as_any()
11601 .downcast_ref::<TimestampMillisecondArray>()
11602 .unwrap();
11603 let values = batch
11604 .column_by_name(&value_field)
11605 .unwrap()
11606 .as_any()
11607 .downcast_ref::<Float64Array>()
11608 .unwrap();
11609 (0..batch.num_rows()).map(|row| {
11610 (
11611 timestamps.value(row),
11612 (!le.is_null(row)).then(|| le.value(row).to_string()),
11613 values.value(row),
11614 )
11615 })
11616 })
11617 .collect::<Vec<_>>();
11618 actual.sort_by(|lhs, rhs| (lhs.0, &lhs.1).cmp(&(rhs.0, &rhs.1)));
11619 assert_eq!(
11620 actual,
11621 vec![
11622 (1_000, None, 1.0),
11623 (1_000, Some("native".to_string()), 0.0),
11624 (2_000, None, 1.0),
11625 (2_000, Some("native".to_string()), 0.0),
11626 ]
11627 );
11628 }
11629
11630 #[tokio::test]
11631 async fn nested_histogram_helpers_ignore_unparsable_bucket_labels() {
11632 let state = build_query_engine_state();
11633 for native_le in [None, Some("native")] {
11634 for query in [
11635 "histogram_quantile(0.5, histogram_quantile(0.5, mixed_histogram))",
11636 "histogram_fraction(-Inf, +Inf, histogram_fraction(-Inf, +Inf, mixed_histogram))",
11637 ] {
11638 let plan = PromPlanner::stmt_to_plan(
11639 classic_and_native_histogram_table_provider(
11640 "native",
11641 native_le,
11642 direct_or_histogram(),
11643 ),
11644 &operator_eval_stmt(query),
11645 &state,
11646 )
11647 .await
11648 .unwrap();
11649
11650 let (_, batches) = execute(plan, &state).await;
11651 assert_eq!(
11652 batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
11653 0,
11654 "native_le={native_le:?}, query={query}"
11655 );
11656 }
11657 }
11658 }
11659
11660 #[tokio::test]
11661 async fn native_histogram_quantile_rejects_multi_field_input() {
11662 let table_provider = build_test_multi_histogram_table_provider("some_metric").await;
11663 let result = PromPlanner::stmt_to_plan(
11664 table_provider,
11665 &build_eval_stmt("histogram_quantile(0.9, some_metric)"),
11666 &build_query_engine_state(),
11667 )
11668 .await;
11669
11670 let err = result.expect_err("histogram_quantile on two native histogram fields must fail");
11671 assert!(
11672 err.to_string()
11673 .contains("Multi fields calculation is not supported in histogram_quantile"),
11674 "{err}"
11675 );
11676 }
11677
11678 #[tokio::test]
11679 async fn native_histogram_topk_uses_drop_udf() {
11680 let plan = native_histogram_plan("topk(1, some_metric)").await;
11681
11682 assert!(plan.contains("prom_native_histogram_drop_float"), "{plan}");
11683 assert!(
11684 plan.contains("Filter: prom_native_histogram_drop_float")
11685 && plan.contains("IS NOT NULL"),
11686 "{plan}"
11687 );
11688 }
11689
11690 #[tokio::test]
11691 async fn mixed_or_topk_bottomk_ignore_native_histograms() {
11692 for op in ["topk", "bottomk"] {
11693 let collector = PromqlAnnotationCollector::default();
11694 let state = build_query_engine_state();
11695 let plan = PromPlanner::stmt_to_plan_with_annotations(
11696 operator_table_provider(),
11697 &operator_eval_stmt(&format!("{op}(1, lf or on(tag) lh)")),
11698 &state,
11699 Some(collector.clone()),
11700 )
11701 .await
11702 .unwrap();
11703 let float_field = plan
11704 .schema()
11705 .fields()
11706 .iter()
11707 .find(|field| field.data_type() == &ArrowDataType::Float64)
11708 .unwrap()
11709 .name()
11710 .clone();
11711 assert!(
11712 plan.schema()
11713 .fields()
11714 .iter()
11715 .all(|field| field.data_type() != &PromPlanner::native_histogram_arrow_type()),
11716 "{plan:?}"
11717 );
11718
11719 let (_, batches) = execute(plan, &state).await;
11720 assert_eq!(values(&batches, &float_field), vec![2.0], "{op}");
11721 let mut warnings = vec![];
11722 let mut infos = vec![];
11723 collector.append_to(&mut warnings, &mut infos);
11724 assert!(warnings.is_empty());
11725 assert_eq!(
11726 infos,
11727 vec![format!(
11728 "{op}: dropped native histogram samples because this aggregation is not supported for native histograms"
11729 )]
11730 );
11731 }
11732 }
11733
11734 #[tokio::test]
11735 async fn native_histogram_scalar_is_ignored_before_scalar_calculate() {
11736 let plan = native_histogram_plan("scalar(some_metric)").await;
11737
11738 assert!(plan.contains("ScalarCalculate"), "{plan}");
11739 assert!(plan.contains("Filter: Boolean(false)"), "{plan}");
11740 assert!(!plan.contains("prom_native_histogram_drop"), "{plan}");
11741 }
11742
11743 #[tokio::test]
11744 async fn native_histogram_value_sort_is_empty_but_label_sort_preserves_samples() {
11745 for function in ["sort", "sort_desc"] {
11746 let plan = native_histogram_plan(&format!("{function}(some_metric)")).await;
11747
11748 assert!(plan.contains("Float64(NULL) IS NOT NULL"), "{plan}");
11749 assert!(
11750 !plan.contains(&format!("Sort: {}", greptime_native_histogram())),
11751 "{plan}"
11752 );
11753 assert!(!plan.contains("prom_native_histogram_drop"), "{plan}");
11754 }
11755
11756 for (function, direction) in [("sort_by_label", "ASC"), ("sort_by_label_desc", "DESC")] {
11757 let plan = native_histogram_plan(&format!("{function}(some_metric, \"tag_0\")")).await;
11758
11759 assert!(plan.contains(&format!("tag_0 {direction}")), "{plan}");
11760 assert!(plan.contains(greptime_native_histogram()), "{plan}");
11761 assert!(!plan.contains("Float64(NULL) IS NOT NULL"), "{plan}");
11762 }
11763 }
11764
11765 #[tokio::test]
11766 async fn unsupported_native_histogram_functions_use_drop_udf() {
11767 for query in [
11768 "deriv(some_metric[5m])",
11769 "min_over_time(some_metric[5m])",
11770 "quantile_over_time(0.9, some_metric[5m])",
11771 "predict_linear(some_metric[5m], 60)",
11772 "round(some_metric)",
11773 "abs(some_metric)",
11774 ] {
11775 let plan = native_histogram_plan(query).await;
11776
11777 assert!(
11778 plan.contains("prom_native_histogram_drop_float"),
11779 "{query}\n{plan}"
11780 );
11781 }
11782 }
11783
11784 #[tokio::test]
11785 async fn native_histogram_absent_over_time_uses_native_udf() {
11786 let plan = native_histogram_plan("absent_over_time(some_metric[5m])").await;
11787
11788 assert!(
11789 plan.contains("prom_native_histogram_absent_over_time"),
11790 "{plan}"
11791 );
11792 }
11793
11794 #[tokio::test]
11795 async fn native_histogram_all_function_arms_route_correctly() {
11796 let cases = [
11801 (
11803 "increase(some_metric[5m])",
11804 "prom_native_histogram_increase",
11805 ),
11806 ("rate(some_metric[5m])", "prom_native_histogram_rate"),
11807 ("delta(some_metric[5m])", "prom_native_histogram_delta"),
11808 ("idelta(some_metric[5m])", "prom_native_histogram_idelta"),
11809 ("irate(some_metric[5m])", "prom_native_histogram_irate"),
11810 ("resets(some_metric[5m])", "prom_native_histogram_resets"),
11811 ("changes(some_metric[5m])", "prom_native_histogram_changes"),
11812 (
11813 "avg_over_time(some_metric[5m])",
11814 "prom_native_histogram_avg_over_time",
11815 ),
11816 (
11817 "sum_over_time(some_metric[5m])",
11818 "prom_native_histogram_sum_over_time",
11819 ),
11820 (
11821 "count_over_time(some_metric[5m])",
11822 "prom_native_histogram_count_over_time",
11823 ),
11824 (
11825 "last_over_time(some_metric[5m])",
11826 "prom_native_histogram_last_over_time",
11827 ),
11828 (
11829 "present_over_time(some_metric[5m])",
11830 "prom_native_histogram_present_over_time",
11831 ),
11832 ("deriv(some_metric[5m])", "prom_native_histogram_drop_float"),
11834 (
11835 "min_over_time(some_metric[5m])",
11836 "prom_native_histogram_drop_float",
11837 ),
11838 (
11839 "max_over_time(some_metric[5m])",
11840 "prom_native_histogram_drop_float",
11841 ),
11842 (
11843 "stddev_over_time(some_metric[5m])",
11844 "prom_native_histogram_drop_float",
11845 ),
11846 (
11847 "stdvar_over_time(some_metric[5m])",
11848 "prom_native_histogram_drop_float",
11849 ),
11850 (
11851 "quantile_over_time(0.9, some_metric[5m])",
11852 "prom_native_histogram_drop_float",
11853 ),
11854 (
11855 "predict_linear(some_metric[5m], 60)",
11856 "prom_native_histogram_drop_float",
11857 ),
11858 (
11859 "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)",
11860 "prom_native_histogram_drop_float",
11861 ),
11862 ("round(some_metric)", "prom_native_histogram_drop_float"),
11863 ("rad(some_metric)", "prom_native_histogram_drop_float"),
11864 ("deg(some_metric)", "prom_native_histogram_drop_float"),
11865 ("sgn(some_metric)", "prom_native_histogram_drop_float"),
11866 (
11868 "histogram_count(some_metric)",
11869 "prom_native_histogram_count",
11870 ),
11871 ("histogram_sum(some_metric)", "prom_native_histogram_sum"),
11872 ("histogram_avg(some_metric)", "prom_native_histogram_avg"),
11873 (
11874 "histogram_stddev(some_metric)",
11875 "prom_native_histogram_stddev",
11876 ),
11877 (
11878 "histogram_stdvar(some_metric)",
11879 "prom_native_histogram_stdvar",
11880 ),
11881 (
11882 "histogram_fraction(-2 + 1, 2 / 2, some_metric)",
11883 "prom_native_histogram_fraction",
11884 ),
11885 ];
11886
11887 for (query, expected_udf) in cases {
11888 let plan = native_histogram_plan(query).await;
11889 assert!(plan.contains(expected_udf), "{query}\n{plan}");
11890 if query.starts_with("histogram_fraction") {
11891 assert!(plan.contains("Float64(-1)"), "{query}\n{plan}");
11892 }
11893 }
11894 }
11895
11896 #[tokio::test]
11897 async fn mixed_native_histogram_ranges_use_coordinated_udfs() {
11898 let dual_output = [
11899 "increase(some_metric[5m])",
11900 "rate(some_metric[5m])",
11901 "delta(some_metric[5m])",
11902 "idelta(some_metric[5m])",
11903 "irate(some_metric[5m])",
11904 "avg_over_time(some_metric[5m])",
11905 "sum_over_time(some_metric[5m])",
11906 "last_over_time(some_metric[5m])",
11907 ];
11908 let float_output = [
11909 "resets(some_metric[5m])",
11910 "changes(some_metric[5m])",
11911 "deriv(some_metric[5m])",
11912 "min_over_time(some_metric[5m])",
11913 "max_over_time(some_metric[5m])",
11914 "count_over_time(some_metric[5m])",
11915 "absent_over_time(some_metric[5m])",
11916 "present_over_time(some_metric[5m])",
11917 "stddev_over_time(some_metric[5m])",
11918 "stdvar_over_time(some_metric[5m])",
11919 "quantile_over_time(0.9, some_metric[5m])",
11920 "predict_linear(some_metric[5m], 60)",
11921 "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)",
11922 ];
11923
11924 for query in dual_output.iter().chain(float_output.iter()) {
11925 let plan = PromPlanner::stmt_to_plan(
11926 build_test_mixed_native_histogram_table_provider("some_metric").await,
11927 &build_eval_stmt(query),
11928 &build_query_engine_state(),
11929 )
11930 .await
11931 .unwrap()
11932 .display_indent_schema()
11933 .to_string();
11934 assert!(plan.contains("prom_mixed_range_float"), "{query}\n{plan}");
11935 assert_eq!(
11936 plan.contains("prom_mixed_range_histogram"),
11937 dual_output.contains(query),
11938 "{query}\n{plan}"
11939 );
11940 }
11941
11942 let plan = PromPlanner::stmt_to_plan(
11943 build_test_mixed_native_histogram_table_provider("some_metric").await,
11944 &build_eval_stmt("sum_over_time(rate(some_metric[5m])[10m:1m])"),
11945 &build_query_engine_state(),
11946 )
11947 .await
11948 .unwrap()
11949 .display_indent_schema()
11950 .to_string();
11951 let expected = r#"Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8]
11952 Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_value, prom_mixed_range_histogram(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8]
11953 PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=["greptime_value", "greptime_native_histogram"] [timestamp:Timestamp(ms), greptime_value:Dictionary(Int64, Float64);N, greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, tag_0:Utf8, timestamp_range:Dictionary(Int64, Timestamp(ms))]
11954 PromSeriesDivide: tags=["tag_0"] [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8]
11955 Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8]
11956 Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8]
11957 Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_value, prom_mixed_range_histogram(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8]
11958 PromRangeManipulate: req range=[-540000..100000000], interval=[60000], eval range=[300000], time index=[timestamp], values=["greptime_native_histogram", "greptime_value"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, greptime_value:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]
11959 PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]
11960 PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]
11961 Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]
11962 Filter: some_metric.timestamp >= TimestampMillisecond(-839999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]
11963 TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]"#;
11964 assert_eq!(plan, expected);
11965 }
11966
11967 #[tokio::test]
11968 async fn mixed_native_histogram_rate_executes_real_ranges() {
11969 let schema = Arc::new(ArrowSchema::new(vec![
11970 Field::new(
11971 "timestamp",
11972 ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
11973 false,
11974 ),
11975 Field::new(greptime_value(), ArrowDataType::Float64, true),
11976 Field::new(
11977 greptime_native_histogram(),
11978 native_histogram_value_type().as_arrow_type(),
11979 true,
11980 ),
11981 ]));
11982 let batch = RecordBatch::try_new(
11983 schema.clone(),
11984 vec![
11985 Arc::new(TimestampMillisecondArray::from(vec![1000, 2000, 3000])),
11986 Arc::new(Float64Array::from(vec![Some(1.0), None, Some(3.0)])),
11987 build_histogram_array(&[None, Some(direct_or_histogram()), None]),
11988 ],
11989 )
11990 .unwrap();
11991 let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap());
11992 let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None)
11993 .unwrap()
11994 .build()
11995 .unwrap();
11996 let collector = PromqlAnnotationCollector::default();
11997 let mut planner = PromPlanner {
11998 table_provider: build_test_table_provider_with_fields(
11999 &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
12000 &[],
12001 )
12002 .await,
12003 ctx: PromPlannerContext {
12004 start: 3000,
12005 end: 3000,
12006 interval: 1000,
12007 range: Some(3000),
12008 time_index_column: Some("timestamp".to_string()),
12009 field_columns: vec![
12010 greptime_native_histogram().to_string(),
12011 greptime_value().to_string(),
12012 ],
12013 ..Default::default()
12014 },
12015 promql_annotations: Some(collector.clone()),
12016 };
12017 let input = LogicalPlan::Extension(Extension {
12018 node: Arc::new(
12019 RangeManipulate::new(
12020 3000,
12021 3000,
12022 1000,
12023 0,
12024 3000,
12025 "timestamp".to_string(),
12026 planner.ctx.field_columns.clone(),
12027 input,
12028 )
12029 .unwrap(),
12030 ),
12031 });
12032 let PromExpr::Call(call) = parser::parse("rate(mixed[3s])").unwrap() else {
12033 unreachable!()
12034 };
12035 let preserve_any_value = PromPlanner::field_columns_are_alternative_samples(
12036 input.schema(),
12037 &planner.ctx.field_columns,
12038 );
12039 let state = build_query_engine_state();
12040 let (mut exprs, _) = planner
12041 .create_function_expr(&call.func, vec![], input.schema(), &state)
12042 .unwrap();
12043 exprs.insert(0, planner.create_time_index_column_expr().unwrap());
12044 let plan = LogicalPlanBuilder::from(input)
12045 .project(exprs)
12046 .unwrap()
12047 .filter(
12048 planner
12049 .create_empty_values_filter_expr(preserve_any_value)
12050 .unwrap(),
12051 )
12052 .unwrap()
12053 .build()
12054 .unwrap();
12055 let (_, batches) = execute(plan, &state).await;
12056 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
12057 let mut warnings = Vec::new();
12058 collector.append_to(&mut warnings, &mut Vec::new());
12059 assert!(
12060 warnings
12061 .iter()
12062 .any(|warning| warning.contains("mix of float and native histogram"))
12063 );
12064 }
12065
12066 #[tokio::test]
12067 async fn native_histogram_mixed_field_table_behaves() {
12068 let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
12072 let plan = PromPlanner::stmt_to_plan(
12073 table_provider,
12074 &build_eval_stmt("histogram_count(some_metric)"),
12075 &build_query_engine_state(),
12076 )
12077 .await
12078 .unwrap();
12079 let plan_str = plan.display_indent_schema().to_string();
12080 assert!(
12081 plan_str.contains("prom_native_histogram_count"),
12082 "{plan_str}"
12083 );
12084 assert!(!plan_str.contains("Float64(NULL)"), "{plan_str}");
12085 assert!(
12086 plan_str.contains("prom_native_histogram_count(greptime_native_histogram) IS NOT NULL"),
12087 "{plan_str}"
12088 );
12089
12090 let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
12092 let plan = PromPlanner::stmt_to_plan(
12093 table_provider,
12094 &build_eval_stmt("sort(some_metric)"),
12095 &build_query_engine_state(),
12096 )
12097 .await
12098 .unwrap();
12099 let plan_str = plan.display_indent_schema().to_string();
12100 assert!(
12101 plan_str.contains("greptime_value ASC NULLS FIRST"),
12102 "{plan_str}"
12103 );
12104 assert!(
12105 !plan_str.contains("greptime_native_histogram ASC"),
12106 "{plan_str}"
12107 );
12108
12109 let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
12111 let plan = PromPlanner::stmt_to_plan(
12112 table_provider,
12113 &build_eval_stmt("scalar(some_metric)"),
12114 &build_query_engine_state(),
12115 )
12116 .await
12117 .unwrap();
12118 let plan_str = plan.display_indent_schema().to_string();
12119 assert!(plan_str.contains("ScalarCalculate"), "{plan_str}");
12120 assert!(
12121 plan_str.contains("greptime_value IS NOT NULL"),
12122 "{plan_str}"
12123 );
12124
12125 let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
12127 let plan = PromPlanner::stmt_to_plan(
12128 table_provider,
12129 &build_eval_stmt(r#"label_replace(some_metric, "copied", "$1", "tag_0", "(.*)")"#),
12130 &build_query_engine_state(),
12131 )
12132 .await
12133 .unwrap();
12134 let plan_str = plan.display_indent_schema().to_string();
12135 let filter = plan_str.lines().next().unwrap();
12136 assert!(
12137 filter.starts_with("Filter: ")
12138 && filter.contains("greptime_native_histogram IS NOT NULL")
12139 && filter.contains(" OR ")
12140 && filter.contains("greptime_value IS NOT NULL"),
12141 "{plan_str}"
12142 );
12143 }
12144
12145 #[tokio::test]
12146 async fn less_filter_on_value() {
12147 let query = "some_metric < 1.2345";
12148 let expected = String::from(
12149 "Filter: some_metric.field_0 < Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12150 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12151 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12152 \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]\
12153 \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]\
12154 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
12155 );
12156
12157 indie_query_plan_compare(query, expected).await;
12158 }
12159
12160 #[tokio::test]
12161 async fn count_over_time() {
12162 let query = "count_over_time(some_metric[5m])";
12163 let expected = String::from(
12164 "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]\
12165 \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]\
12166 \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))]\
12167 \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12168 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12169 \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]\
12170 \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]\
12171 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
12172 );
12173
12174 indie_query_plan_compare(query, expected).await;
12175 }
12176
12177 #[tokio::test]
12180 async fn count_over_time_subquery() {
12181 let query = "count_over_time(some_metric[10m:1m])";
12182 let expected = String::from(
12183 "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]\
12184 \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]\
12185 \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))]\
12186 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12187 \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]\
12188 \n PromInstantManipulate: range=[-540000..100000000], lookback=[1000], interval=[60000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12189 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12190 \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]\
12191 \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]\
12192 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
12193 );
12194 indie_query_plan_compare(query, expected).await;
12195 }
12196
12197 #[tokio::test]
12198 async fn test_hash_join() {
12199 let mut eval_stmt = EvalStmt {
12200 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12201 start: UNIX_EPOCH,
12202 end: UNIX_EPOCH
12203 .checked_add(Duration::from_secs(100_000))
12204 .unwrap(),
12205 interval: Duration::from_secs(5),
12206 lookback_delta: Duration::from_secs(1),
12207 };
12208
12209 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"}"#;
12210
12211 let prom_expr = parser::parse(case).unwrap();
12212 eval_stmt.expr = prom_expr;
12213 let table_provider = build_test_table_provider_with_fields(
12214 &[
12215 (
12216 DEFAULT_SCHEMA_NAME.to_string(),
12217 "http_server_requests_seconds_sum".to_string(),
12218 ),
12219 (
12220 DEFAULT_SCHEMA_NAME.to_string(),
12221 "http_server_requests_seconds_count".to_string(),
12222 ),
12223 ],
12224 &["uri", "kubernetes_namespace", "kubernetes_pod_name"],
12225 )
12226 .await;
12227 let plan =
12229 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12230 .await
12231 .unwrap();
12232 let expected = "Projection: http_server_requests_seconds_sum.uri, http_server_requests_seconds_count.greptime_timestamp, CAST(http_server_requests_seconds_sum.greptime_value AS Float64) / CAST(http_server_requests_seconds_count.greptime_value AS Float64) AS http_server_requests_seconds_sum.greptime_value / http_server_requests_seconds_count.greptime_value\
12233 \n Projection: http_server_requests_seconds_sum.uri, http_server_requests_seconds_sum.kubernetes_namespace, http_server_requests_seconds_sum.kubernetes_pod_name, http_server_requests_seconds_sum.greptime_timestamp, http_server_requests_seconds_sum.greptime_value, http_server_requests_seconds_count.uri, http_server_requests_seconds_count.kubernetes_namespace, http_server_requests_seconds_count.kubernetes_pod_name, http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_count.greptime_value\
12234 \n Filter: prom_assert_unique_match_group(__promql_match_group_count, http_server_requests_seconds_sum.uri)\
12235 \n WindowAggr: windowExpr=[[count(Int64(1)) PARTITION BY [http_server_requests_seconds_sum.uri, http_server_requests_seconds_sum.greptime_timestamp] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS __promql_match_group_count]]\
12236 \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\
12237 \n SubqueryAlias: http_server_requests_seconds_sum\
12238 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\
12239 \n PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\
12240 \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\
12241 \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)\
12242 \n TableScan: http_server_requests_seconds_sum\
12243 \n SubqueryAlias: http_server_requests_seconds_count\
12244 \n Projection: http_server_requests_seconds_count.uri, http_server_requests_seconds_count.kubernetes_namespace, http_server_requests_seconds_count.kubernetes_pod_name, http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_count.greptime_value\
12245 \n Filter: prom_assert_unique_match_group(__promql_match_group_count, http_server_requests_seconds_count.uri)\
12246 \n WindowAggr: windowExpr=[[count(Int64(1)) PARTITION BY [http_server_requests_seconds_count.uri, http_server_requests_seconds_count.greptime_timestamp] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS __promql_match_group_count]]\
12247 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\
12248 \n PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\
12249 \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\
12250 \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)\
12251 \n TableScan: http_server_requests_seconds_count";
12252 assert_eq!(plan.to_string(), expected);
12253 }
12254
12255 #[tokio::test]
12256 async fn test_nested_histogram_quantile() {
12257 let mut eval_stmt = EvalStmt {
12258 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12259 start: UNIX_EPOCH,
12260 end: UNIX_EPOCH
12261 .checked_add(Duration::from_secs(100_000))
12262 .unwrap(),
12263 interval: Duration::from_secs(5),
12264 lookback_delta: Duration::from_secs(1),
12265 };
12266
12267 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]*-(.*)")"#;
12268
12269 let prom_expr = parser::parse(case).unwrap();
12270 eval_stmt.expr = prom_expr;
12271 let table_provider = build_test_table_provider_with_fields(
12272 &[(
12273 DEFAULT_SCHEMA_NAME.to_string(),
12274 "greptime_servers_grpc_requests_elapsed_bucket".to_string(),
12275 )],
12276 &["pod", "le", "path", "code", "container"],
12277 )
12278 .await;
12279 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12281 .await
12282 .unwrap();
12283 }
12284
12285 #[tokio::test]
12286 async fn test_histogram_quantile_binary_op() {
12287 let mut eval_stmt = EvalStmt {
12288 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12289 start: UNIX_EPOCH,
12290 end: UNIX_EPOCH
12291 .checked_add(Duration::from_secs(100_000))
12292 .unwrap(),
12293 interval: Duration::from_secs(5),
12294 lookback_delta: Duration::from_secs(1),
12295 };
12296
12297 let case = r#"histogram_quantile(0.5, sum by (le, pod) (rate(http_request_duration_seconds_bucket[5m]))) + 0"#;
12301
12302 let prom_expr = parser::parse(case).unwrap();
12303 eval_stmt.expr = prom_expr;
12304 let table_provider = build_test_table_provider_with_fields(
12305 &[(
12306 DEFAULT_SCHEMA_NAME.to_string(),
12307 "http_request_duration_seconds_bucket".to_string(),
12308 )],
12309 &["pod", "le"],
12310 )
12311 .await;
12312 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12314 .await
12315 .unwrap();
12316 }
12317
12318 #[tokio::test]
12319 async fn test_parse_and_operator() {
12320 let mut eval_stmt = EvalStmt {
12321 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12322 start: UNIX_EPOCH,
12323 end: UNIX_EPOCH
12324 .checked_add(Duration::from_secs(100_000))
12325 .unwrap(),
12326 interval: Duration::from_secs(5),
12327 lookback_delta: Duration::from_secs(1),
12328 };
12329
12330 let cases = [
12331 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)"#,
12332 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)"#,
12333 ];
12334
12335 for case in cases {
12336 let prom_expr = parser::parse(case).unwrap();
12337 eval_stmt.expr = prom_expr;
12338 let table_provider = build_test_table_provider_with_fields(
12339 &[
12340 (
12341 DEFAULT_SCHEMA_NAME.to_string(),
12342 "kubelet_volume_stats_used_bytes".to_string(),
12343 ),
12344 (
12345 DEFAULT_SCHEMA_NAME.to_string(),
12346 "kubelet_volume_stats_capacity_bytes".to_string(),
12347 ),
12348 ],
12349 &["namespace", "persistentvolumeclaim"],
12350 )
12351 .await;
12352 let _ =
12354 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12355 .await
12356 .unwrap();
12357 }
12358 }
12359
12360 #[tokio::test]
12361 async fn test_nested_binary_op() {
12362 let mut eval_stmt = EvalStmt {
12363 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12364 start: UNIX_EPOCH,
12365 end: UNIX_EPOCH
12366 .checked_add(Duration::from_secs(100_000))
12367 .unwrap(),
12368 interval: Duration::from_secs(5),
12369 lookback_delta: Duration::from_secs(1),
12370 };
12371
12372 let case = r#"sum(rate(nginx_ingress_controller_requests{job=~".*"}[2m])) -
12373 (
12374 sum(rate(nginx_ingress_controller_requests{namespace=~".*"}[2m]))
12375 or
12376 vector(0)
12377 )"#;
12378
12379 let prom_expr = parser::parse(case).unwrap();
12380 eval_stmt.expr = prom_expr;
12381 let table_provider = build_test_table_provider_with_fields(
12382 &[(
12383 DEFAULT_SCHEMA_NAME.to_string(),
12384 "nginx_ingress_controller_requests".to_string(),
12385 )],
12386 &["namespace", "job"],
12387 )
12388 .await;
12389 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12391 .await
12392 .unwrap();
12393 }
12394
12395 #[tokio::test]
12396 async fn test_parse_or_operator() {
12397 let mut eval_stmt = EvalStmt {
12398 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12399 start: UNIX_EPOCH,
12400 end: UNIX_EPOCH
12401 .checked_add(Duration::from_secs(100_000))
12402 .unwrap(),
12403 interval: Duration::from_secs(5),
12404 lookback_delta: Duration::from_secs(1),
12405 };
12406
12407 let case = r#"
12408 sum(rate(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}[120s])) by (cluster_name,tenant_name) /
12409 (sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) * 100)
12410 or
12411 200 * sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) /
12412 sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)"#;
12413
12414 let table_provider = build_test_table_provider_with_fields(
12415 &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())],
12416 &["tenant_name", "cluster_name"],
12417 )
12418 .await;
12419 eval_stmt.expr = parser::parse(case).unwrap();
12420 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12421 .await
12422 .unwrap();
12423
12424 let case = r#"sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
12425 (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) +
12426 sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
12427 (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0
12428 or
12429 sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
12430 (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0
12431 or
12432 sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
12433 (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0"#;
12434 let table_provider = build_test_table_provider_with_fields(
12435 &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())],
12436 &["tenant_name", "cluster_name"],
12437 )
12438 .await;
12439 eval_stmt.expr = parser::parse(case).unwrap();
12440 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12441 .await
12442 .unwrap();
12443
12444 let case = r#"(sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) +
12445 sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or
12446 (sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or
12447 (sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name))"#;
12448 let table_provider = build_test_table_provider_with_fields(
12449 &[
12450 (
12451 DEFAULT_SCHEMA_NAME.to_string(),
12452 "background_waitevent_cnt".to_string(),
12453 ),
12454 (
12455 DEFAULT_SCHEMA_NAME.to_string(),
12456 "foreground_waitevent_cnt".to_string(),
12457 ),
12458 ],
12459 &["tenant_name", "cluster_name"],
12460 )
12461 .await;
12462 eval_stmt.expr = parser::parse(case).unwrap();
12463 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12464 .await
12465 .unwrap();
12466
12467 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)"#;
12468 let table_provider = build_test_table_provider_with_fields(
12469 &[
12470 (DEFAULT_SCHEMA_NAME.to_string(), "node_load1".to_string()),
12471 (
12472 DEFAULT_SCHEMA_NAME.to_string(),
12473 "container_cpu_load_average_10s".to_string(),
12474 ),
12475 (
12476 DEFAULT_SCHEMA_NAME.to_string(),
12477 "container_spec_cpu_quota".to_string(),
12478 ),
12479 ],
12480 &["cluster_name", "host_name"],
12481 )
12482 .await;
12483 eval_stmt.expr = parser::parse(case).unwrap();
12484 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12485 .await
12486 .unwrap();
12487 }
12488
12489 #[tokio::test]
12490 async fn value_matcher() {
12491 let mut eval_stmt = EvalStmt {
12493 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12494 start: UNIX_EPOCH,
12495 end: UNIX_EPOCH
12496 .checked_add(Duration::from_secs(100_000))
12497 .unwrap(),
12498 interval: Duration::from_secs(5),
12499 lookback_delta: Duration::from_secs(1),
12500 };
12501
12502 let cases = [
12503 (
12505 r#"some_metric{__field__="field_1"}"#,
12506 vec![
12507 "some_metric.field_1",
12508 "some_metric.tag_0",
12509 "some_metric.tag_1",
12510 "some_metric.tag_2",
12511 "some_metric.timestamp",
12512 ],
12513 ),
12514 (
12516 r#"some_metric{__field__="field_1", __field__="field_0"}"#,
12517 vec![
12518 "some_metric.field_0",
12519 "some_metric.field_1",
12520 "some_metric.tag_0",
12521 "some_metric.tag_1",
12522 "some_metric.tag_2",
12523 "some_metric.timestamp",
12524 ],
12525 ),
12526 (
12528 r#"some_metric{__field__!="field_1"}"#,
12529 vec![
12530 "some_metric.field_0",
12531 "some_metric.field_2",
12532 "some_metric.tag_0",
12533 "some_metric.tag_1",
12534 "some_metric.tag_2",
12535 "some_metric.timestamp",
12536 ],
12537 ),
12538 (
12540 r#"some_metric{__field__!="field_1", __field__!="field_2"}"#,
12541 vec![
12542 "some_metric.field_0",
12543 "some_metric.tag_0",
12544 "some_metric.tag_1",
12545 "some_metric.tag_2",
12546 "some_metric.timestamp",
12547 ],
12548 ),
12549 (
12551 r#"some_metric{__field__="field_1", __field__!="field_0"}"#,
12552 vec![
12553 "some_metric.field_1",
12554 "some_metric.tag_0",
12555 "some_metric.tag_1",
12556 "some_metric.tag_2",
12557 "some_metric.timestamp",
12558 ],
12559 ),
12560 (
12562 r#"some_metric{__field__="field_2", __field__!="field_2"}"#,
12563 vec![
12564 "some_metric.tag_0",
12565 "some_metric.tag_1",
12566 "some_metric.tag_2",
12567 "some_metric.timestamp",
12568 ],
12569 ),
12570 (
12572 r#"some_metric{__field__=~"field_1|field_2"}"#,
12573 vec![
12574 "some_metric.field_1",
12575 "some_metric.field_2",
12576 "some_metric.tag_0",
12577 "some_metric.tag_1",
12578 "some_metric.tag_2",
12579 "some_metric.timestamp",
12580 ],
12581 ),
12582 (
12584 r#"some_metric{__field__!~"field_1|field_2"}"#,
12585 vec![
12586 "some_metric.field_0",
12587 "some_metric.tag_0",
12588 "some_metric.tag_1",
12589 "some_metric.tag_2",
12590 "some_metric.timestamp",
12591 ],
12592 ),
12593 ];
12594
12595 for case in cases {
12596 let prom_expr = parser::parse(case.0).unwrap();
12597 eval_stmt.expr = prom_expr;
12598 let table_provider = build_test_table_provider(
12599 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12600 3,
12601 3,
12602 )
12603 .await;
12604 let plan =
12605 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12606 .await
12607 .unwrap();
12608 let mut fields = plan.schema().field_names();
12609 let mut expected = case.1.into_iter().map(String::from).collect::<Vec<_>>();
12610 fields.sort();
12611 expected.sort();
12612 assert_eq!(fields, expected, "case: {:?}", case.0);
12613 }
12614
12615 let bad_cases = [
12616 r#"some_metric{__field__="nonexistent"}"#,
12617 r#"some_metric{__field__!="nonexistent"}"#,
12618 ];
12619
12620 for case in bad_cases {
12621 let prom_expr = parser::parse(case).unwrap();
12622 eval_stmt.expr = prom_expr;
12623 let table_provider = build_test_table_provider(
12624 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12625 3,
12626 3,
12627 )
12628 .await;
12629 let plan =
12630 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12631 .await;
12632 assert!(plan.is_err(), "case: {:?}", case);
12633 }
12634 }
12635
12636 #[tokio::test]
12637 async fn custom_schema() {
12638 let query = "some_alt_metric{__schema__=\"greptime_private\"}";
12639 let expected = String::from(
12640 "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12641 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12642 \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]\
12643 \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]\
12644 \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
12645 );
12646
12647 indie_query_plan_compare(query, expected).await;
12648
12649 let query = "some_alt_metric{__database__=\"greptime_private\"}";
12650 let expected = String::from(
12651 "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12652 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12653 \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]\
12654 \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]\
12655 \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
12656 );
12657
12658 indie_query_plan_compare(query, expected).await;
12659
12660 let query = "some_alt_metric{__schema__=\"greptime_private\"} / some_metric";
12661 let expected = String::from(
12662 "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]\
12663 \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]\
12664 \n SubqueryAlias: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12665 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12666 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12667 \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]\
12668 \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]\
12669 \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12670 \n SubqueryAlias: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12671 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12672 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12673 \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]\
12674 \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]\
12675 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
12676 );
12677
12678 indie_query_plan_compare(query, expected).await;
12679 }
12680
12681 #[tokio::test]
12682 async fn only_equals_is_supported_for_special_matcher() {
12683 let queries = &[
12684 "some_alt_metric{__schema__!=\"greptime_private\"}",
12685 "some_alt_metric{__schema__=~\"lalala\"}",
12686 "some_alt_metric{__database__!=\"greptime_private\"}",
12687 "some_alt_metric{__database__=~\"lalala\"}",
12688 ];
12689
12690 for query in queries {
12691 let prom_expr = parser::parse(query).unwrap();
12692 let eval_stmt = EvalStmt {
12693 expr: prom_expr,
12694 start: UNIX_EPOCH,
12695 end: UNIX_EPOCH
12696 .checked_add(Duration::from_secs(100_000))
12697 .unwrap(),
12698 interval: Duration::from_secs(5),
12699 lookback_delta: Duration::from_secs(1),
12700 };
12701
12702 let table_provider = build_test_table_provider(
12703 &[
12704 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
12705 (
12706 "greptime_private".to_string(),
12707 "some_alt_metric".to_string(),
12708 ),
12709 ],
12710 1,
12711 1,
12712 )
12713 .await;
12714
12715 let plan =
12716 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12717 .await;
12718 assert!(plan.is_err(), "query: {:?}", query);
12719 }
12720 }
12721
12722 #[tokio::test]
12723 async fn native_scan_bounds_preserve_zero_lookback_and_overflow() {
12724 let table_provider = build_test_table_provider(
12725 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12726 1,
12727 1,
12728 )
12729 .await;
12730 let mut planner = PromPlanner {
12731 table_provider,
12732 ctx: PromPlannerContext::from_eval_stmt(&build_eval_stmt("some_metric")),
12733 promql_annotations: None,
12734 };
12735 planner.ctx.time_index_column = Some("timestamp".to_string());
12736 planner.ctx.start = 1_000;
12737 planner.ctx.lookback_delta = 0;
12738 let schema = Arc::new(
12739 DFSchema::try_from(ArrowSchema::new(vec![Field::new(
12740 "timestamp",
12741 ArrowDataType::Timestamp(ArrowTimeUnit::Nanosecond, None),
12742 false,
12743 )]))
12744 .unwrap(),
12745 );
12746 for (end, interval, windows) in [
12747 (1_000, 1_000, 1),
12748 (2_000, 1_000, 1),
12749 (7_201_000, 7_200_000, 2),
12750 ] {
12751 planner.ctx.end = end;
12752 planner.ctx.interval = interval;
12753 let filter = planner
12754 .build_time_index_filter(0, &schema)
12755 .unwrap()
12756 .unwrap()
12757 .to_string();
12758 assert_eq!(filter.matches(">=").count(), windows, "{filter}");
12759 assert!(
12760 filter.contains("TimestampNanosecond(1000000000, None)"),
12761 "{filter}"
12762 );
12763 }
12764 planner.ctx.end = i64::MAX;
12765 let filter = planner
12766 .build_time_index_filter(0, &schema)
12767 .unwrap()
12768 .unwrap()
12769 .to_string();
12770 assert!(
12771 filter.contains("timestamp >= TimestampNanosecond(1000000000, None)"),
12772 "{filter}"
12773 );
12774
12775 let ms_schema = Arc::new(
12778 DFSchema::try_from(ArrowSchema::new(vec![Field::new(
12779 "timestamp",
12780 ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
12781 false,
12782 )]))
12783 .unwrap(),
12784 );
12785 planner.ctx.start = i64::MIN + 100;
12786 planner.ctx.end = planner.ctx.start;
12787 planner.ctx.lookback_delta = 200;
12788 let filter = planner
12789 .build_time_index_filter(0, &ms_schema)
12790 .unwrap()
12791 .unwrap()
12792 .to_string();
12793 assert_eq!(
12794 filter,
12795 format!(
12796 "timestamp <= TimestampMillisecond({}, None)",
12797 i64::MIN + 100
12798 )
12799 );
12800
12801 planner.ctx.start = 0;
12804 planner.ctx.end = 0;
12805 planner.ctx.lookback_delta = 300_000;
12806 let filter = planner
12807 .build_time_index_filter(9_223_372_036_854, &schema)
12808 .unwrap()
12809 .unwrap()
12810 .to_string();
12811 assert_eq!(
12812 filter,
12813 "timestamp <= TimestampNanosecond(-9223372036854000000, None)"
12814 );
12815 }
12816
12817 #[tokio::test]
12818 async fn test_non_ms_precision() {
12819 let catalog_list = MemoryCatalogManager::with_default_setup();
12820 let columns = vec![
12821 ColumnSchema::new(
12822 "tag".to_string(),
12823 ConcreteDataType::string_datatype(),
12824 false,
12825 ),
12826 ColumnSchema::new(
12827 "timestamp".to_string(),
12828 ConcreteDataType::timestamp_nanosecond_datatype(),
12829 false,
12830 )
12831 .with_time_index(true),
12832 ColumnSchema::new(
12833 "field".to_string(),
12834 ConcreteDataType::float64_datatype(),
12835 true,
12836 ),
12837 ];
12838 let schema = Arc::new(Schema::new(columns));
12839 let table_meta = TableMetaBuilder::empty()
12840 .schema(schema)
12841 .primary_key_indices(vec![0])
12842 .value_indices(vec![2])
12843 .next_column_id(1024)
12844 .build()
12845 .unwrap();
12846 let table_info = TableInfoBuilder::default()
12847 .name("metrics".to_string())
12848 .meta(table_meta)
12849 .build()
12850 .unwrap();
12851 let table = EmptyTable::from_table_info(&table_info);
12852 assert!(
12853 catalog_list
12854 .register_table_sync(RegisterTableRequest {
12855 catalog: DEFAULT_CATALOG_NAME.to_string(),
12856 schema: DEFAULT_SCHEMA_NAME.to_string(),
12857 table_name: "metrics".to_string(),
12858 table_id: 1024,
12859 table,
12860 })
12861 .is_ok()
12862 );
12863
12864 let plan = PromPlanner::stmt_to_plan(
12865 DfTableSourceProvider::new(
12866 catalog_list.clone(),
12867 false,
12868 QueryContext::arc(),
12869 DummyDecoder::arc(),
12870 true,
12871 ),
12872 &EvalStmt {
12873 expr: parser::parse("metrics{tag = \"1\"}").unwrap(),
12874 start: UNIX_EPOCH,
12875 end: UNIX_EPOCH
12876 .checked_add(Duration::from_secs(100_000))
12877 .unwrap(),
12878 interval: Duration::from_secs(5),
12879 lookback_delta: Duration::from_secs(1),
12880 },
12881 &build_query_engine_state(),
12882 )
12883 .await
12884 .unwrap();
12885 assert_eq!(
12886 plan.display_indent_schema().to_string(),
12887 "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\n PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp > TimestampNanosecond(-1000000000, None) AND metrics.timestamp <= TimestampNanosecond(100000000000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Projection: metrics.field, metrics.tag, metrics.timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]"
12888 );
12889 let plan = PromPlanner::stmt_to_plan(
12890 DfTableSourceProvider::new(
12891 catalog_list.clone(),
12892 false,
12893 QueryContext::arc(),
12894 DummyDecoder::arc(),
12895 true,
12896 ),
12897 &EvalStmt {
12898 expr: parser::parse("avg_over_time(metrics{tag = \"1\"}[5s])").unwrap(),
12899 start: UNIX_EPOCH,
12900 end: UNIX_EPOCH
12901 .checked_add(Duration::from_secs(100_000))
12902 .unwrap(),
12903 interval: Duration::from_secs(5),
12904 lookback_delta: Duration::from_secs(1),
12905 },
12906 &build_query_engine_state(),
12907 )
12908 .await
12909 .unwrap();
12910 assert_eq!(
12911 plan.display_indent_schema().to_string(),
12912 "Filter: prom_avg_over_time(timestamp_range,field) IS NOT NULL [timestamp:Timestamp(ms), prom_avg_over_time(timestamp_range,field):Float64;N, tag:Utf8]\n Projection: metrics.timestamp, prom_avg_over_time(timestamp_range, field) AS prom_avg_over_time(timestamp_range,field), metrics.tag [timestamp:Timestamp(ms), prom_avg_over_time(timestamp_range,field):Float64;N, tag:Utf8]\n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[5000], time index=[timestamp], values=[\"field\"] [field:Dictionary(Int64, Float64);N, tag:Utf8, timestamp:Timestamp(ms), timestamp_range:Dictionary(Int64, Timestamp(ms))]\n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp > TimestampNanosecond(-5000000000, None) AND metrics.timestamp <= TimestampNanosecond(100000000000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Projection: metrics.field, metrics.tag, metrics.timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]"
12913 );
12914 }
12915
12916 #[tokio::test]
12917 async fn test_nonexistent_label() {
12918 let mut eval_stmt = EvalStmt {
12920 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12921 start: UNIX_EPOCH,
12922 end: UNIX_EPOCH
12923 .checked_add(Duration::from_secs(100_000))
12924 .unwrap(),
12925 interval: Duration::from_secs(5),
12926 lookback_delta: Duration::from_secs(1),
12927 };
12928
12929 let case = r#"some_metric{nonexistent="hi"}"#;
12930 let prom_expr = parser::parse(case).unwrap();
12931 eval_stmt.expr = prom_expr;
12932 let table_provider = build_test_table_provider(
12933 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12934 3,
12935 3,
12936 )
12937 .await;
12938 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12940 .await
12941 .unwrap();
12942 }
12943
12944 #[tokio::test]
12945 async fn test_label_join() {
12946 let prom_expr = parser::parse(
12947 "label_join(up{tag_0='api-server'}, 'foo', ',', 'tag_1', 'tag_2', 'tag_3')",
12948 )
12949 .unwrap();
12950 let eval_stmt = EvalStmt {
12951 expr: prom_expr,
12952 start: UNIX_EPOCH,
12953 end: UNIX_EPOCH
12954 .checked_add(Duration::from_secs(100_000))
12955 .unwrap(),
12956 interval: Duration::from_secs(5),
12957 lookback_delta: Duration::from_secs(1),
12958 };
12959
12960 let table_provider =
12961 build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 4, 1)
12962 .await;
12963 let plan =
12964 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12965 .await
12966 .unwrap();
12967
12968 let expected = r#"
12969Filter: 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]
12970 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]
12971 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]
12972 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]
12973 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]
12974 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]
12975 TableScan: up [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#;
12976
12977 let ret = plan.display_indent_schema().to_string();
12978 assert_eq!(format!("\n{ret}"), expected, "\n{}", ret);
12979 }
12980
12981 #[tokio::test]
12982 async fn test_label_replace() {
12983 let prom_expr = parser::parse(
12984 "label_replace(up{tag_0=\"a:c\"}, \"foo\", \"$1\", \"tag_0\", \"(.*):.*\")",
12985 )
12986 .unwrap();
12987 let eval_stmt = EvalStmt {
12988 expr: prom_expr,
12989 start: UNIX_EPOCH,
12990 end: UNIX_EPOCH
12991 .checked_add(Duration::from_secs(100_000))
12992 .unwrap(),
12993 interval: Duration::from_secs(5),
12994 lookback_delta: Duration::from_secs(1),
12995 };
12996
12997 let table_provider =
12998 build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 1, 1)
12999 .await;
13000 let plan =
13001 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13002 .await
13003 .unwrap();
13004
13005 let expected = r#"
13006Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8]
13007 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]
13008 PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
13009 PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
13010 Sort: up.tag_0 ASC NULLS FIRST, up.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
13011 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]
13012 TableScan: up [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#;
13013
13014 let ret = plan.display_indent_schema().to_string();
13015 assert_eq!(format!("\n{ret}"), expected, "\n{}", ret);
13016 }
13017
13018 #[tokio::test]
13019 async fn label_replace_aggregation_queries_plan_successfully() {
13020 let aggregate =
13021 r#"sum by (foo) (label_replace(some_metric, "foo", "$1", "tag_0", "(.*)"))"#;
13022 let queries = [
13023 aggregate.to_string(),
13024 format!("{aggregate} <= 10"),
13025 format!("{aggregate} * 0.8"),
13026 format!("0.8 * {aggregate}"),
13027 format!("{aggregate} <= {aggregate} * 0.8"),
13028 ];
13029 let state = build_query_engine_state();
13030 let mut failures = Vec::new();
13031
13032 for query in queries {
13033 let table_provider = build_test_table_provider(
13034 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
13035 1,
13036 1,
13037 )
13038 .await;
13039 if let Err(error) =
13040 PromPlanner::stmt_to_plan(table_provider, &build_eval_stmt(&query), &state).await
13041 {
13042 failures.push(format!("{query}: {error:?}"));
13043 }
13044 }
13045
13046 assert!(failures.is_empty(), "{}", failures.join("\n"));
13047 }
13048
13049 #[tokio::test]
13050 async fn test_matchers_to_expr() {
13051 let mut eval_stmt = EvalStmt {
13052 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
13053 start: UNIX_EPOCH,
13054 end: UNIX_EPOCH
13055 .checked_add(Duration::from_secs(100_000))
13056 .unwrap(),
13057 interval: Duration::from_secs(5),
13058 lookback_delta: Duration::from_secs(1),
13059 };
13060 let case =
13061 r#"sum(prometheus_tsdb_head_series{tag_1=~"(10.0.160.237:8080|10.0.160.237:9090)"})"#;
13062
13063 let prom_expr = parser::parse(case).unwrap();
13064 eval_stmt.expr = prom_expr;
13065 let table_provider = build_test_table_provider(
13066 &[(
13067 DEFAULT_SCHEMA_NAME.to_string(),
13068 "prometheus_tsdb_head_series".to_string(),
13069 )],
13070 3,
13071 3,
13072 )
13073 .await;
13074 let plan =
13075 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13076 .await
13077 .unwrap();
13078 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]\
13079 \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]\
13080 \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]\
13081 \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]\
13082 \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]\
13083 \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]\
13084 \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]";
13085 assert_eq!(plan.display_indent_schema().to_string(), expected);
13086 }
13087
13088 #[tokio::test]
13089 async fn test_topk_expr() {
13090 let mut eval_stmt = EvalStmt {
13091 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
13092 start: UNIX_EPOCH,
13093 end: UNIX_EPOCH
13094 .checked_add(Duration::from_secs(100_000))
13095 .unwrap(),
13096 interval: Duration::from_secs(5),
13097 lookback_delta: Duration::from_secs(1),
13098 };
13099 let case = r#"topk(10, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#;
13100
13101 let prom_expr = parser::parse(case).unwrap();
13102 eval_stmt.expr = prom_expr;
13103 let table_provider = build_test_table_provider_with_fields(
13104 &[
13105 (
13106 DEFAULT_SCHEMA_NAME.to_string(),
13107 "prometheus_tsdb_head_series".to_string(),
13108 ),
13109 (
13110 DEFAULT_SCHEMA_NAME.to_string(),
13111 "http_server_requests_seconds_count".to_string(),
13112 ),
13113 ],
13114 &["ip"],
13115 )
13116 .await;
13117
13118 let plan =
13119 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13120 .await
13121 .unwrap();
13122 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)]\
13123 \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]\
13124 \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]\
13125 \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]\
13126 \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]\
13127 \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]\
13128 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
13129 \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
13130 \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]\
13131 \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]\
13132 \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]";
13133
13134 assert_eq!(plan.display_indent_schema().to_string(), expected);
13135 }
13136
13137 #[tokio::test]
13138 async fn test_count_values_expr() {
13139 let mut eval_stmt = EvalStmt {
13140 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
13141 start: UNIX_EPOCH,
13142 end: UNIX_EPOCH
13143 .checked_add(Duration::from_secs(100_000))
13144 .unwrap(),
13145 interval: Duration::from_secs(5),
13146 lookback_delta: Duration::from_secs(1),
13147 };
13148 let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#;
13149
13150 let prom_expr = parser::parse(case).unwrap();
13151 eval_stmt.expr = prom_expr;
13152 let table_provider = build_test_table_provider_with_fields(
13153 &[
13154 (
13155 DEFAULT_SCHEMA_NAME.to_string(),
13156 "prometheus_tsdb_head_series".to_string(),
13157 ),
13158 (
13159 DEFAULT_SCHEMA_NAME.to_string(),
13160 "http_server_requests_seconds_count".to_string(),
13161 ),
13162 ],
13163 &["ip"],
13164 )
13165 .await;
13166
13167 let plan =
13168 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13169 .await
13170 .unwrap();
13171 let expected = "Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N]\
13172 \n Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value) AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N]\
13173 \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value)]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), prom_float_to_string(prometheus_tsdb_head_series.greptime_value):Utf8;N, count(prometheus_tsdb_head_series.greptime_value):Int64]\
13174 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
13175 \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
13176 \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]\
13177 \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]\
13178 \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]";
13179
13180 assert_eq!(plan.display_indent_schema().to_string(), expected);
13181 }
13182
13183 #[tokio::test]
13184 async fn test_value_alias() {
13185 let mut eval_stmt = EvalStmt {
13186 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
13187 start: UNIX_EPOCH,
13188 end: UNIX_EPOCH
13189 .checked_add(Duration::from_secs(100_000))
13190 .unwrap(),
13191 interval: Duration::from_secs(5),
13192 lookback_delta: Duration::from_secs(1),
13193 };
13194 let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#;
13195
13196 let prom_expr = parser::parse(case).unwrap();
13197 eval_stmt.expr = prom_expr;
13198 eval_stmt = QueryLanguageParser::apply_alias_extension(eval_stmt, "my_series");
13199 let table_provider = build_test_table_provider_with_fields(
13200 &[
13201 (
13202 DEFAULT_SCHEMA_NAME.to_string(),
13203 "prometheus_tsdb_head_series".to_string(),
13204 ),
13205 (
13206 DEFAULT_SCHEMA_NAME.to_string(),
13207 "http_server_requests_seconds_count".to_string(),
13208 ),
13209 ],
13210 &["ip"],
13211 )
13212 .await;
13213
13214 let plan =
13215 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13216 .await
13217 .unwrap();
13218 let expected = r#"
13219Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.series, prometheus_tsdb_head_series.greptime_timestamp [my_series:Int64, ip:Utf8, series:Utf8;N, greptime_timestamp:Timestamp(ms)]
13220 Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N]
13221 Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value) AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N]
13222 Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value)]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), prom_float_to_string(prometheus_tsdb_head_series.greptime_value):Utf8;N, count(prometheus_tsdb_head_series.greptime_value):Int64]
13223 PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
13224 PromSeriesDivide: tags=["ip"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
13225 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]
13226 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]
13227 TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"#;
13228 assert_eq!(format!("\n{}", plan.display_indent_schema()), expected);
13229 }
13230
13231 #[tokio::test]
13232 async fn test_quantile_expr() {
13233 let mut eval_stmt = EvalStmt {
13234 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
13235 start: UNIX_EPOCH,
13236 end: UNIX_EPOCH
13237 .checked_add(Duration::from_secs(100_000))
13238 .unwrap(),
13239 interval: Duration::from_secs(5),
13240 lookback_delta: Duration::from_secs(1),
13241 };
13242 let case = r#"quantile(0.3, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#;
13243
13244 let prom_expr = parser::parse(case).unwrap();
13245 eval_stmt.expr = prom_expr;
13246 let table_provider = build_test_table_provider_with_fields(
13247 &[
13248 (
13249 DEFAULT_SCHEMA_NAME.to_string(),
13250 "prometheus_tsdb_head_series".to_string(),
13251 ),
13252 (
13253 DEFAULT_SCHEMA_NAME.to_string(),
13254 "http_server_requests_seconds_count".to_string(),
13255 ),
13256 ],
13257 &["ip"],
13258 )
13259 .await;
13260
13261 let plan =
13262 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13263 .await
13264 .unwrap();
13265 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]\
13266 \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]\
13267 \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]\
13268 \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]\
13269 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
13270 \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
13271 \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]\
13272 \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]\
13273 \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]";
13274
13275 assert_eq!(plan.display_indent_schema().to_string(), expected);
13276 }
13277
13278 #[tokio::test]
13279 async fn test_or_not_exists_table_label() {
13280 let state = build_query_engine_state();
13281 let provider = build_test_table_provider_with_fields(
13282 &[(DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string())],
13283 &["job"],
13284 )
13285 .await;
13286 let raw = PromPlanner::stmt_to_plan(
13287 provider,
13288 &build_eval_stmt(r#"missing_metric or on(absent_label) normal_metric"#),
13289 &state,
13290 )
13291 .await
13292 .unwrap();
13293 assert!(
13294 raw.display_indent_schema()
13295 .to_string()
13296 .contains("__promql_or_match_0@")
13297 );
13298 let (optimized, batches) = execute(raw, &state).await;
13299 assert_no_internal_or_keys(optimized.schema());
13300 assert!(batches.iter().all(|batch| {
13301 batch
13302 .schema()
13303 .fields()
13304 .iter()
13305 .all(|field| !field.name().starts_with("__promql_or_match_"))
13306 }));
13307 }
13308
13309 #[tokio::test]
13310 async fn test_histogram_quantile_missing_le_column() {
13311 let mut eval_stmt = EvalStmt {
13312 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
13313 start: UNIX_EPOCH,
13314 end: UNIX_EPOCH
13315 .checked_add(Duration::from_secs(100_000))
13316 .unwrap(),
13317 interval: Duration::from_secs(5),
13318 lookback_delta: Duration::from_secs(1),
13319 };
13320
13321 let case = r#"histogram_quantile(0.99, sum by(pod,instance,le) (rate(non_existent_histogram_bucket{instance=~"xxx"}[1m])))"#;
13323
13324 let prom_expr = parser::parse(case).unwrap();
13325 eval_stmt.expr = prom_expr;
13326
13327 let table_provider = build_test_table_provider_with_fields(
13329 &[(
13330 DEFAULT_SCHEMA_NAME.to_string(),
13331 "non_existent_histogram_bucket".to_string(),
13332 )],
13333 &["pod", "instance"], )
13335 .await;
13336
13337 let result =
13339 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13340 .await;
13341
13342 assert!(
13344 result.is_ok(),
13345 "Expected successful plan creation with empty result, but got error: {:?}",
13346 result.err()
13347 );
13348
13349 let plan = result.unwrap();
13351 match plan {
13352 LogicalPlan::EmptyRelation(_) => {
13353 }
13355 _ => panic!("Expected EmptyRelation, but got: {:?}", plan),
13356 }
13357 }
13358
13359 #[tokio::test]
13360 async fn test_direct_or_normalizes_missing_match_labels() {
13361 type Case<'a> = (
13362 Option<Option<&'a str>>,
13363 Option<Option<&'a str>>,
13364 i64,
13365 i64,
13366 &'a [(f64, Option<&'a str>)],
13367 );
13368
13369 let modifier = or_modifier("lhs or on(k) rhs");
13370 #[rustfmt::skip]
13371 let cases: &[Case<'_>] = &[
13372 (None, None, 1, 1, &[(1.0, None)]),
13373 (None, Some(Some("")), 1, 1, &[(1.0, None)]),
13374 (Some(Some("")), None, 1, 1, &[(1.0, Some(""))]),
13375 (None, Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]),
13376 (Some(Some("l")), None, 1, 1, &[(1.0, Some("l")), (2.0, None)]),
13377 (Some(None), Some(Some("")), 1, 1, &[(1.0, None)]),
13378 (Some(None), Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]),
13379 (Some(Some("same")), Some(Some("same")), 1, 2, &[(1.0, Some("same")), (2.0, Some("same"))]),
13380 ];
13381 for &(left, right, left_ts, right_ts, expected) in cases {
13382 let (optimized, batches) = run(
13383 &matrix_source("lhs", left, left_ts, 1.0),
13384 &matrix_source("rhs", right, right_ts, 2.0),
13385 matrix_context("lhs", left),
13386 matrix_context("rhs", right),
13387 &modifier,
13388 )
13389 .await;
13390 assert_no_internal_or_keys(optimized.schema());
13391 assert_eq!(
13392 rows(&batches),
13393 expected
13394 .iter()
13395 .map(|(value, label)| (*value, label.map(str::to_string)))
13396 .collect::<Vec<_>>()
13397 );
13398 }
13399 }
13400
13401 #[tokio::test]
13402 async fn test_direct_or_match_modifiers() {
13403 for (modifier, left, right, expected) in [
13404 (None, "left", "right", 2),
13405 (or_modifier("lhs or on(k) rhs"), "same", "same", 1),
13406 (or_modifier("lhs or on() rhs"), "left", "right", 1),
13407 (or_modifier("lhs or ignoring(k) rhs"), "left", "right", 1),
13408 ] {
13409 let (_, batches) = run(
13410 &matrix_source("lhs", Some(Some(left)), 1, 1.0),
13411 &matrix_source("rhs", Some(Some(right)), 1, 2.0),
13412 direct_or_context("lhs", &["job", "k"], "v"),
13413 direct_or_context("rhs", &["job", "k"], "v"),
13414 &modifier,
13415 )
13416 .await;
13417 assert_eq!(
13418 batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
13419 expected
13420 );
13421 }
13422 }
13423
13424 #[tokio::test]
13425 async fn test_direct_or_nested_projection_uses_left_context() {
13426 let left = matrix_source("lhs", Some(Some("k")), 1, 1.0);
13427 let right = matrix_source("rhs", Some(Some("k")), 1, 2.0);
13428 let raw = plan_direct_or(
13429 scan(&left),
13430 scan(&right),
13431 direct_or_context("lhs", &["job", "k"], "v"),
13432 direct_or_context("rhs", &["job", "k"], "v"),
13433 &or_modifier("lhs or on(k) rhs"),
13434 )
13435 .await;
13436 assert!(raw.schema().iter().any(|(qualifier, field)| {
13437 qualifier.as_ref().is_some_and(|q| q.to_string() == "lhs") && field.name() == "v"
13438 }));
13439 let nested = LogicalPlanBuilder::from(raw)
13440 .project(vec![
13441 DfExpr::BinaryExpr(BinaryExpr {
13442 left: Box::new(DfExpr::Column(Column::new(
13443 Some(TableReference::bare("lhs")),
13444 "v",
13445 ))),
13446 op: Operator::Plus,
13447 right: Box::new(lit(1.0)),
13448 })
13449 .alias("v_plus"),
13450 ])
13451 .unwrap()
13452 .build()
13453 .unwrap();
13454 let (_, batches) = execute(nested, &build_query_engine_state()).await;
13455 assert_eq!(values(&batches, "v_plus"), vec![2.0]);
13456 }
13457
13458 #[tokio::test]
13459 async fn test_direct_or_skips_user_internal_key_name() {
13460 const USER_TAG: &str = "__promql_or_match_0";
13461 let left = tagged_source(
13462 "lhs",
13463 false,
13464 (USER_TAG, Some("left")),
13465 DirectOrValue::Float64(1.0),
13466 );
13467 let right = tagged_source(
13468 "rhs",
13469 false,
13470 (USER_TAG, Some("right")),
13471 DirectOrValue::Float64(2.0),
13472 );
13473 let raw = plan_direct_or(
13474 scan(&left),
13475 scan(&right),
13476 direct_or_context("lhs", &["job", USER_TAG], "v"),
13477 direct_or_context("rhs", &["job", USER_TAG], "v"),
13478 &or_modifier("lhs or on(missing_label) rhs"),
13479 )
13480 .await;
13481 assert!(
13482 raw.display_indent_schema()
13483 .to_string()
13484 .contains("__promql_or_match_1@")
13485 );
13486 let (_, batches) = execute(raw, &build_query_engine_state()).await;
13487 assert!(
13488 batches
13489 .iter()
13490 .all(|batch| batch.column_by_name(USER_TAG).is_some())
13491 );
13492 }
13493
13494 #[tokio::test]
13495 async fn test_direct_or_substrait_round_trip_with_normalized_key() {
13496 let state = build_query_engine_state();
13497 let ctx = SessionContext::new_with_state(state.session_state());
13498 let catalog = Arc::new(MemoryCatalogProvider::new());
13499 catalog
13500 .register_schema("public", Arc::new(MemorySchemaProvider::new()))
13501 .unwrap();
13502 ctx.register_catalog("datafusion", catalog);
13503 let left = matrix_source("lhs", Some(Some("")), 1, 1.0);
13504 let right = matrix_source("rhs", None, 1, 2.0);
13505 ctx.register_table(
13506 TableReference::full("datafusion", "public", "lhs"),
13507 table(&left),
13508 )
13509 .unwrap();
13510 ctx.register_table(
13511 TableReference::full("datafusion", "public", "rhs"),
13512 table(&right),
13513 )
13514 .unwrap();
13515 let raw = plan_direct_or(
13516 ctx.table("datafusion.public.lhs")
13517 .await
13518 .unwrap()
13519 .into_unoptimized_plan(),
13520 ctx.table("datafusion.public.rhs")
13521 .await
13522 .unwrap()
13523 .into_unoptimized_plan(),
13524 direct_or_context("lhs", &["job", "k"], "v"),
13525 direct_or_context("rhs", &["job"], "v"),
13526 &or_modifier("lhs or on(k) rhs"),
13527 )
13528 .await;
13529 let decoded = DFLogicalSubstraitConvertor
13530 .decode(
13531 DFLogicalSubstraitConvertor
13532 .encode(&raw, DefaultSerializer)
13533 .unwrap(),
13534 ctx.state(),
13535 )
13536 .await
13537 .unwrap();
13538 let (optimized, batches) = execute(decoded, &state).await;
13539 assert_no_internal_or_keys(optimized.schema());
13540 assert!(batches.iter().all(|batch| {
13541 batch
13542 .schema()
13543 .fields()
13544 .iter()
13545 .all(|field| !field.name().starts_with("__promql_or_match_"))
13546 }));
13547 assert_eq!(values(&batches, "v"), vec![1.0]);
13548 }
13549
13550 #[tokio::test]
13551 async fn test_direct_or_numeric_value_types() {
13552 let left = tagged_source("lhs", true, ("k", Some("lhs")), DirectOrValue::Int64(0));
13553 let right = tagged_source(
13554 "rhs",
13555 false,
13556 ("k", Some("rhs")),
13557 DirectOrValue::Float64(0.5),
13558 );
13559 let (optimized, batches) = run(
13560 &left,
13561 &right,
13562 direct_or_context("lhs", &["job", "k"], "v"),
13563 direct_or_context("rhs", &["job", "k"], "v"),
13564 &or_modifier("lhs or on(k) rhs"),
13565 )
13566 .await;
13567 assert_eq!(
13568 optimized
13569 .schema()
13570 .field_with_name(None, "v")
13571 .unwrap()
13572 .data_type(),
13573 &ArrowDataType::Float64
13574 );
13575 assert_eq!(values(&batches, "v"), vec![0.5]);
13576 let provider = build_test_table_provider_with_fields(
13577 &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
13578 &[],
13579 )
13580 .await;
13581 let mut planner = PromPlanner {
13582 table_provider: provider,
13583 ctx: PromPlannerContext::default(),
13584 promql_annotations: None,
13585 };
13586 let left_context = direct_or_context("lhs", &["job"], "v");
13587 let right_context = direct_or_context("rhs", &["job"], "v");
13588 let error = planner
13589 .or_operator(
13590 scan(&job_source("lhs", DirectOrValue::Utf8("x"))),
13591 scan(&job_source("rhs", DirectOrValue::Float64(1.0))),
13592 left_context.tag_columns.iter().cloned().collect(),
13593 right_context.tag_columns.iter().cloned().collect(),
13594 left_context,
13595 right_context,
13596 &or_modifier("lhs or on() rhs"),
13597 )
13598 .unwrap_err();
13599 assert!(
13600 error
13601 .to_string()
13602 .contains("OR value fields have incompatible types")
13603 );
13604 }
13605
13606 #[tokio::test]
13607 async fn test_or_with_histogram_quantile_missing_le_column() {
13608 let case = r#"histogram_quantile(0.99, non_existent_histogram_bucket) or normal_metric"#;
13609 let eval_stmt = build_eval_stmt(case);
13610 let table_provider = build_missing_le_or_normal_metric_table_provider().await;
13611
13612 let plan =
13613 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13614 .await
13615 .unwrap();
13616 assert_normal_metric_schema(&plan);
13617 }
13618
13619 #[tokio::test]
13620 async fn test_or_with_right_empty_histogram_restores_left_context() {
13621 let eval_stmt = build_eval_stmt(
13622 r#"abs(sum by(instance) (normal_metric) or histogram_quantile(0.99, sum by(pod) (non_existent_histogram_bucket)))"#,
13623 );
13624 let table_provider = build_missing_le_or_normal_metric_table_provider().await;
13625
13626 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13627 .await
13628 .unwrap();
13629 }
13630
13631 #[tokio::test]
13632 async fn test_or_with_both_empty_histograms() {
13633 let eval_stmt = build_eval_stmt(
13634 r#"histogram_quantile(0.99, sum by(pod) (left_histogram_bucket)) or histogram_quantile(0.99, sum by(instance) (right_histogram_bucket))"#,
13635 );
13636 let table_provider = build_test_table_provider_with_fields(
13637 &[
13638 (
13639 DEFAULT_SCHEMA_NAME.to_string(),
13640 "left_histogram_bucket".to_string(),
13641 ),
13642 (
13643 DEFAULT_SCHEMA_NAME.to_string(),
13644 "right_histogram_bucket".to_string(),
13645 ),
13646 ],
13647 &["pod", "instance"],
13648 )
13649 .await;
13650
13651 let plan =
13652 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13653 .await
13654 .unwrap();
13655 match plan {
13656 LogicalPlan::EmptyRelation(relation) => {
13657 assert!(!relation.produce_one_row);
13658 assert!(!relation.schema.fields().is_empty());
13659 assert!(
13660 relation
13661 .schema
13662 .fields()
13663 .iter()
13664 .any(|field| field.data_type() == &ArrowDataType::Float64)
13665 );
13666 assert!(
13667 relation
13668 .schema
13669 .fields()
13670 .iter()
13671 .any(|field| field.name() == "pod")
13672 );
13673 assert!(
13674 !relation
13675 .schema
13676 .fields()
13677 .iter()
13678 .any(|field| field.name() == "instance")
13679 );
13680 }
13681 _ => panic!("Expected EmptyRelation, but got: {plan:?}"),
13682 }
13683 }
13684
13685 #[tokio::test]
13686 async fn test_nested_or_with_both_empty_histograms() {
13687 for case in [
13688 r#"abs(histogram_quantile(0.99, left_histogram_bucket) or histogram_quantile(0.99, right_histogram_bucket))"#,
13689 r#"(histogram_quantile(0.99, left_histogram_bucket) or histogram_quantile(0.99, right_histogram_bucket)) + 1"#,
13690 ] {
13691 let eval_stmt = build_eval_stmt(case);
13692 let table_provider = build_test_table_provider_with_fields(
13693 &[
13694 (
13695 DEFAULT_SCHEMA_NAME.to_string(),
13696 "left_histogram_bucket".to_string(),
13697 ),
13698 (
13699 DEFAULT_SCHEMA_NAME.to_string(),
13700 "right_histogram_bucket".to_string(),
13701 ),
13702 ],
13703 &["pod", "instance"],
13704 )
13705 .await;
13706
13707 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13708 .await
13709 .unwrap();
13710 }
13711 }
13712
13713 #[tokio::test]
13714 async fn test_or_with_empty_histogram_modifiers() {
13715 for case in [
13716 r#"histogram_quantile(0.99, non_existent_histogram_bucket) or on(pod) normal_metric"#,
13717 r#"normal_metric or ignoring(instance) histogram_quantile(0.99, non_existent_histogram_bucket)"#,
13718 ] {
13719 let eval_stmt = build_eval_stmt(case);
13720 let table_provider = build_missing_le_or_normal_metric_table_provider().await;
13721
13722 let plan =
13723 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13724 .await
13725 .unwrap();
13726 assert_normal_metric_schema(&plan);
13727 }
13728 }
13729
13730 #[tokio::test]
13731 async fn test_unless_preserves_left_context_for_histogram() {
13732 let eval_stmt = build_eval_stmt(
13733 r#"histogram_quantile(0.99, bucket_metric unless on(job) normal_metric) or fallback_metric"#,
13734 );
13735 let state = build_query_engine_state();
13736 let plan = PromPlanner::stmt_to_plan(
13737 build_set_op_context_table_provider().await,
13738 &eval_stmt,
13739 &state,
13740 )
13741 .await
13742 .unwrap();
13743 assert!(contains_histogram_fold(&plan), "{plan:?}");
13744 let (optimized, physical) = optimize_and_create_physical_plan(&state, plan).await;
13745 assert!(contains_histogram_fold(&optimized), "{optimized:?}");
13746 let batches =
13747 datafusion::physical_plan::collect(physical, state.session_state().task_ctx())
13748 .await
13749 .unwrap();
13750 assert!(batches.iter().all(|batch| batch.num_rows() == 0));
13751 }
13752
13753 #[tokio::test]
13754 async fn test_and_preserves_left_context_for_histogram() {
13755 let eval_stmt = build_eval_stmt(
13756 r#"histogram_quantile(0.99, bucket_metric and on(job) normal_metric) or fallback_metric"#,
13757 );
13758 let plan = PromPlanner::stmt_to_plan(
13759 build_set_op_context_table_provider().await,
13760 &eval_stmt,
13761 &build_query_engine_state(),
13762 )
13763 .await
13764 .unwrap();
13765 assert!(contains_histogram_fold(&plan), "{plan:?}");
13766 }
13767
13768 #[tokio::test]
13769 async fn test_and_preserves_left_context_when_le_is_missing() {
13770 let eval_stmt =
13771 build_eval_stmt(r#"histogram_quantile(0.99, normal_metric and on(job) bucket_metric)"#);
13772 let plan = PromPlanner::stmt_to_plan(
13773 build_set_op_context_table_provider().await,
13774 &eval_stmt,
13775 &build_query_engine_state(),
13776 )
13777 .await
13778 .unwrap();
13779 assert!(matches!(&plan, LogicalPlan::EmptyRelation(_)), "{plan:?}");
13780 assert!(!plan.schema().fields().is_empty());
13781 assert!(!contains_histogram_fold(&plan), "{plan:?}");
13782 }
13783
13784 async fn build_matching_filter_plan(query: &str) -> String {
13785 let table_provider = build_test_table_provider_with_distinct_tags(&[
13786 ("metric_a", &["host", "device"]),
13787 ("metric_b", &["host", "device"]),
13788 ])
13789 .await;
13790 let plan = PromPlanner::stmt_to_plan(
13791 table_provider,
13792 &build_eval_stmt(query),
13793 &build_query_engine_state(),
13794 )
13795 .await
13796 .unwrap();
13797 plan.display_indent().to_string()
13798 }
13799
13800 async fn build_test_table_provider_with_string_field(
13804 table_tags: &[(&str, &[&str])],
13805 ) -> DfTableSourceProvider {
13806 let catalog_list = MemoryCatalogManager::with_default_setup();
13807 for (table_name, tags) in table_tags {
13808 let mut columns = tags
13809 .iter()
13810 .map(|tag| {
13811 ColumnSchema::new(
13812 (*tag).to_string(),
13813 ConcreteDataType::string_datatype(),
13814 false,
13815 )
13816 })
13817 .collect::<Vec<_>>();
13818 columns.push(
13819 ColumnSchema::new(
13820 greptime_timestamp().to_string(),
13821 ConcreteDataType::timestamp_millisecond_datatype(),
13822 false,
13823 )
13824 .with_time_index(true),
13825 );
13826 columns.push(ColumnSchema::new(
13827 greptime_value().to_string(),
13828 ConcreteDataType::float64_datatype(),
13829 true,
13830 ));
13831 columns.push(ColumnSchema::new(
13832 "status".to_string(),
13833 ConcreteDataType::string_datatype(),
13834 true,
13835 ));
13836 let table_meta = TableMetaBuilder::empty()
13837 .schema(Arc::new(Schema::new(columns)))
13838 .primary_key_indices((0..tags.len()).collect())
13839 .next_column_id(1024)
13840 .build()
13841 .unwrap();
13842 let table_info = TableInfoBuilder::default()
13843 .name((*table_name).to_string())
13844 .meta(table_meta)
13845 .build()
13846 .unwrap();
13847
13848 assert!(
13849 catalog_list
13850 .register_table_sync(RegisterTableRequest {
13851 catalog: DEFAULT_CATALOG_NAME.to_string(),
13852 schema: DEFAULT_SCHEMA_NAME.to_string(),
13853 table_name: (*table_name).to_string(),
13854 table_id: 1024,
13855 table: EmptyTable::from_table_info(&table_info),
13856 })
13857 .is_ok()
13858 );
13859 }
13860
13861 DfTableSourceProvider::new(
13862 catalog_list,
13863 false,
13864 QueryContext::arc(),
13865 DummyDecoder::arc(),
13866 false,
13867 )
13868 }
13869
13870 async fn build_matching_filter_plan_with_string_field(query: &str) -> String {
13871 let table_provider = build_test_table_provider_with_string_field(&[
13872 ("metric_a", &["host", "device"]),
13873 ("metric_b", &["host", "device"]),
13874 ])
13875 .await;
13876 let plan = PromPlanner::stmt_to_plan(
13877 table_provider,
13878 &build_eval_stmt(query),
13879 &build_query_engine_state(),
13880 )
13881 .await
13882 .unwrap();
13883 plan.display_indent().to_string()
13884 }
13885
13886 #[tokio::test]
13887 async fn binary_matching_label_filter_reaches_both_operands() {
13888 for query in [
13889 r#"metric_a / metric_b{host="foo"}"#,
13890 r#"metric_a / on(host, device) metric_b{host="foo"}"#,
13891 r#"count_over_time(metric_a[1m]) / on(host) count_over_time(metric_b{host="foo"}[1m])"#,
13892 r#"metric_a / ignoring(device) metric_b{host="foo"}"#,
13893 r#"sum by(host) (metric_a) / on(host) sum by(host) (metric_b{host="foo"})"#,
13894 r#"metric_a / on(host) avg without(device) (metric_b{host="foo"})"#,
13895 ] {
13896 let plan = build_matching_filter_plan(query).await;
13897 assert_eq!(
13898 plan.matches(r#"host = Utf8("foo")"#).count(),
13899 2,
13900 "{query}\n{plan}"
13901 );
13902 }
13903 }
13904
13905 #[tokio::test]
13906 async fn binary_matching_label_filter_reaches_scalar_ranking_and_grouped_operands() {
13907 for query in [
13908 r#"(8 * metric_a{host="foo"}) / on(host) metric_b"#,
13909 r#"topk(1, metric_a{host="foo"}) / on(host, device) metric_b"#,
13910 r#"(8 * metric_a{host="foo"}) / on(host) group_left topk by(host)(1, max by(host)(metric_b))"#,
13911 ] {
13912 let plan = build_matching_filter_plan(query).await;
13913 assert_eq!(
13914 plan.matches(r#"host = Utf8("foo")"#).count(),
13915 2,
13916 "{query}\n{plan}"
13917 );
13918 }
13919 let query = r#"metric_a{host="foo"} / on(host) group_left topk(1, max by(host)(metric_b))"#;
13921 let plan = build_matching_filter_plan(query).await;
13922 assert_eq!(
13923 plan.matches(r#"host = Utf8("foo")"#).count(),
13924 1,
13925 "{query}\n{plan}"
13926 );
13927 }
13928
13929 #[tokio::test]
13930 async fn binary_matching_label_filter_skips_selecting_aggregations() {
13931 let query = r#"topk(1, metric_a) / on(host, device) metric_b{host="foo"}"#;
13933 let plan = build_matching_filter_plan(query).await;
13934 assert_eq!(plan.matches("foo").count(), 1, "{query}\n{plan}");
13935 }
13936
13937 #[tokio::test]
13938 async fn binary_value_field_matcher_stays_on_its_own_operand() {
13939 let value = greptime_value();
13940 for query in [
13941 format!(r#"metric_a / metric_b{{{value}="2"}}"#),
13942 format!(r#"metric_a / on(host, device, {value}) metric_b{{{value}="2"}}"#),
13943 ] {
13944 let plan = build_matching_filter_plan(&query).await;
13945 assert_eq!(plan.matches(r#"Utf8("2")"#).count(), 1, "{query}\n{plan}");
13946 }
13947 }
13948
13949 #[tokio::test]
13950 async fn binary_value_field_matcher_is_not_copied_across_aggregations() {
13951 for query in [
13954 r#"count by(status) (metric_a) / on(status) count by(status) (metric_b{status="ready"})"#,
13955 r#"(8 * count by(status)(metric_a{__field__="status"})) / on(status) topk by(status)(1, count by(status)(metric_b{__field__="status",status="ready"}))"#,
13956 r#"count by(status)(metric_a{__field__="status"}) / on(status) group_left topk by(status)(1, count by(status)(metric_b{__field__="status",status="ready"}))"#,
13957 ] {
13958 let plan = build_matching_filter_plan_with_string_field(query).await;
13959 assert_eq!(
13960 plan.matches(r#"Utf8("ready")"#).count(),
13961 1,
13962 "{query}\n{plan}"
13963 );
13964 }
13965 }
13966
13967 #[tokio::test]
13968 async fn binary_matching_label_filter_reaches_aggregations_grouping_by_tags() {
13969 let query = r#"count by(host) (metric_a) / on(host) count by(host) (metric_b{host="foo"})"#;
13971 let plan = build_matching_filter_plan_with_string_field(query).await;
13972 assert_eq!(plan.matches(r#"Utf8("foo")"#).count(), 2, "{query}\n{plan}");
13973 }
13974
13975 #[tokio::test]
13976 async fn binary_matching_label_filter_skips_unproven_expressions() {
13977 for query in [
13978 r#"metric_a > on(host, device) metric_b{host="foo"}"#,
13979 r#"metric_a / on(host) group_left metric_b{host="foo"}"#,
13980 r#"metric_a / on(host) label_replace(metric_b{host="foo"},"extra","e","host",".*")"#,
13981 r#"metric_a / on(device) metric_b{host="foo"}"#,
13982 ] {
13983 let plan = build_matching_filter_plan(query).await;
13984 assert_eq!(plan.matches("foo").count(), 1, "{query}\n{plan}");
13985 }
13986 }
13987
13988 #[tokio::test]
13989 async fn test_or_context_uses_left_qualified_output() {
13990 let case = r#"(normal_metric or other_metric) + 1"#;
13991 let eval_stmt = build_eval_stmt(case);
13992 let state = build_query_engine_state();
13993 let plan =
13994 PromPlanner::stmt_to_plan(build_or_context_table_provider().await, &eval_stmt, &state)
13995 .await
13996 .unwrap();
13997 assert!(
13998 plan.schema()
13999 .fields()
14000 .iter()
14001 .any(|field| field.data_type() == &ArrowDataType::Float64),
14002 "{plan:?}"
14003 );
14004 let (_optimized, _physical) = optimize_and_create_physical_plan(&state, plan).await;
14005 }
14006
14007 #[tokio::test]
14008 async fn test_or_context_uses_left_qualified_empty_histogram_output() {
14009 let case = r#"(abs(histogram_quantile(0.99, non_hist_metric)) or normal_metric) + 1"#;
14010 let eval_stmt = build_eval_stmt(case);
14011 let plan = PromPlanner::stmt_to_plan(
14012 build_or_context_table_provider().await,
14013 &eval_stmt,
14014 &build_query_engine_state(),
14015 )
14016 .await
14017 .unwrap();
14018 assert!(
14019 plan.schema()
14020 .fields()
14021 .iter()
14022 .any(|field| field.data_type() == &ArrowDataType::Float64),
14023 "{plan:?}"
14024 );
14025 }
14026
14027 #[tokio::test]
14028 async fn test_direct_or_preserves_float_and_native_histogram_samples() {
14029 for histogram_on_left in [false, true] {
14030 let (planner, plan) = mixed_direct_or(histogram_on_left).await;
14031
14032 let float_field = &planner.ctx.field_columns[0];
14033 let histogram_field = &planner.ctx.field_columns[1];
14034 assert!(float_field.starts_with(OR_FLOAT_FIELD_PREFIX));
14035 assert!(histogram_field.starts_with(OR_HISTOGRAM_FIELD_PREFIX));
14036 assert_eq!(
14037 plan.schema()
14038 .field_with_name(None, float_field)
14039 .unwrap()
14040 .data_type(),
14041 &ArrowDataType::Float64
14042 );
14043 assert_eq!(
14044 plan.schema()
14045 .field_with_name(None, histogram_field)
14046 .unwrap()
14047 .data_type(),
14048 &native_histogram_value_type().as_arrow_type()
14049 );
14050
14051 let (optimized, batches) = execute(plan, &build_query_engine_state()).await;
14052 assert_no_internal_or_keys(optimized.schema());
14053 let mut sample_kinds = batches
14054 .iter()
14055 .flat_map(|batch| {
14056 let values = batch.column_by_name(float_field).unwrap();
14057 let histograms = batch.column_by_name(histogram_field).unwrap();
14058 (0..batch.num_rows())
14059 .map(|row| (values.is_valid(row), histograms.is_valid(row)))
14060 })
14061 .collect::<Vec<_>>();
14062 sample_kinds.sort_unstable();
14063 assert_eq!(sample_kinds, vec![(false, true), (true, false)]);
14064 }
14065 }
14066
14067 #[tokio::test]
14068 async fn malformed_classic_bucket_does_not_drop_native_histogram() {
14069 let state = build_query_engine_state();
14070 let collector = PromqlAnnotationCollector::default();
14071 let plan = PromPlanner::stmt_to_plan_with_annotations(
14072 operator_table_provider(),
14073 &operator_eval_stmt("histogram_quantile(0.5, bad_classic or bad_native)"),
14074 &state,
14075 Some(collector.clone()),
14076 )
14077 .await
14078 .unwrap();
14079 let value_field = plan
14080 .schema()
14081 .fields()
14082 .iter()
14083 .find(|field| field.data_type() == &ArrowDataType::Float64)
14084 .unwrap()
14085 .name()
14086 .clone();
14087
14088 let (_, batches) = execute(plan, &state).await;
14089 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
14090 assert_eq!(values(&batches, &value_field), vec![0.0]);
14091 let mut warnings = vec![];
14092 let mut infos = vec![];
14093 collector.append_to(&mut warnings, &mut infos);
14094 assert!(warnings.is_empty());
14095 assert!(infos.is_empty());
14096 }
14097
14098 #[tokio::test]
14099 async fn test_mixed_binary_operator_aligns_both_alternative_inputs() {
14100 let state = build_query_engine_state();
14101 let plan = PromPlanner::stmt_to_plan(
14102 operator_table_provider(),
14103 &operator_eval_stmt("(lf or on(tag) lh) * on(tag) (rf or on(tag) rh)"),
14104 &state,
14105 )
14106 .await
14107 .unwrap();
14108 let plan_text = plan.display_indent_schema().to_string();
14109 assert!(
14110 plan_text.contains("prom_native_histogram_mul_scalar"),
14111 "{plan_text}"
14112 );
14113 assert!(
14114 plan_text.contains("prom_native_histogram_scalar_mul"),
14115 "{plan_text}"
14116 );
14117 let float_field = plan
14118 .schema()
14119 .fields()
14120 .iter()
14121 .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX))
14122 .unwrap()
14123 .name()
14124 .clone();
14125 let histogram_field = plan
14126 .schema()
14127 .fields()
14128 .iter()
14129 .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX))
14130 .unwrap()
14131 .name()
14132 .clone();
14133
14134 let (_, batches) = execute(plan, &state).await;
14135 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
14136 assert!(values(&batches, &float_field).is_empty());
14137 let mut sums = histograms(&batches, &histogram_field)
14138 .into_iter()
14139 .map(|histogram| histogram.sum)
14140 .collect::<Vec<_>>();
14141 sums.sort_by(f64::total_cmp);
14142 assert_eq!(sums, vec![2.0, 3.0]);
14143 }
14144
14145 #[tokio::test]
14146 async fn test_mixed_binary_operator_reports_only_dropped_samples() {
14147 for (query, expected_rows, expected_infos) in [
14148 ("(lf or on(tag) lh) + on(tag) (rf or on(tag) rh)", 0, 1),
14149 ("(lf or on(tag) lh) + on(tag) (lf or on(tag) lh)", 2, 0),
14150 ("(lf or on(tag) lh) % on(tag) lh", 0, 1),
14151 ] {
14152 let state = build_query_engine_state();
14153 let annotations = PromqlAnnotationCollector::default();
14154 let plan = PromPlanner::stmt_to_plan_with_annotations(
14155 operator_table_provider(),
14156 &operator_eval_stmt(query),
14157 &state,
14158 Some(annotations.clone()),
14159 )
14160 .await
14161 .unwrap();
14162
14163 let (_, batches) = execute(plan, &state).await;
14164 assert_eq!(
14165 batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
14166 expected_rows,
14167 "{query}"
14168 );
14169 let mut warnings = vec![];
14170 let mut infos = vec![];
14171 annotations.append_to(&mut warnings, &mut infos);
14172 assert!(warnings.is_empty(), "{query}: {warnings:?}");
14173 assert_eq!(infos.len(), expected_infos, "{query}: {infos:?}");
14174 }
14175 }
14176
14177 #[tokio::test]
14178 async fn test_histogram_only_min_drops_empty_aggregate_group() {
14179 let state = build_query_engine_state();
14183 for query in ["min(lh)", "group(min(lh))"] {
14184 let plan = PromPlanner::stmt_to_plan(
14185 operator_table_provider(),
14186 &operator_eval_stmt(query),
14187 &state,
14188 )
14189 .await
14190 .unwrap();
14191 let (_, batches) = execute(plan, &state).await;
14192 assert_eq!(
14193 batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
14194 0,
14195 "{query}"
14196 );
14197 }
14198 }
14199
14200 #[tokio::test]
14201 async fn test_mixed_min_drops_histogram_only_group() {
14202 let state = build_query_engine_state();
14206 let plan = PromPlanner::stmt_to_plan(
14207 operator_table_provider(),
14208 &operator_eval_stmt("min by (tag) (lf or on(tag) lh)"),
14209 &state,
14210 )
14211 .await
14212 .unwrap();
14213 let float_field = plan
14214 .schema()
14215 .fields()
14216 .iter()
14217 .find(|field| field.data_type() == &ArrowDataType::Float64)
14218 .unwrap()
14219 .name()
14220 .clone();
14221 let (_, batches) = execute(plan, &state).await;
14222 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
14223 assert_eq!(values(&batches, &float_field), vec![2.0]);
14224 }
14225
14226 #[tokio::test]
14227 async fn test_mixed_or_can_feed_another_or() {
14228 let state = build_query_engine_state();
14229 let plan = PromPlanner::stmt_to_plan(
14230 operator_table_provider(),
14231 &operator_eval_stmt("lf or on(tag) lh or on(tag) fallback"),
14232 &state,
14233 )
14234 .await
14235 .unwrap();
14236 let float_field = plan
14237 .schema()
14238 .fields()
14239 .iter()
14240 .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX))
14241 .unwrap()
14242 .name()
14243 .clone();
14244 let histogram_field = plan
14245 .schema()
14246 .fields()
14247 .iter()
14248 .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX))
14249 .unwrap()
14250 .name()
14251 .clone();
14252
14253 let (_, batches) = execute(plan, &state).await;
14254 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 3);
14255 let mut float_values = values(&batches, &float_field);
14256 float_values.sort_by(f64::total_cmp);
14257 assert_eq!(float_values, vec![2.0, 7.0]);
14258 assert_eq!(histograms(&batches, &histogram_field).len(), 1);
14259 }
14260
14261 #[tokio::test]
14262 async fn test_mixed_fields_align_with_single_float_vector() {
14263 let (planner, mixed) = mixed_direct_or(false).await;
14264 let scale = tagged_source(
14265 "scale",
14266 false,
14267 ("k", Some("float")),
14268 DirectOrValue::Float64(2.0),
14269 );
14270 let scale = scan(&scale);
14271 let scale_fields = vec!["v".to_string()];
14272 let PromExpr::Binary(binary) = parser::parse("lhs * rhs").unwrap() else {
14273 unreachable!()
14274 };
14275
14276 let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns(
14277 mixed.schema(),
14278 scale.schema(),
14279 &planner.ctx.field_columns,
14280 &scale_fields,
14281 binary.op,
14282 false,
14283 false,
14284 );
14285 assert!(invalid_pairs.is_empty());
14286 assert_eq!(
14287 groups
14288 .iter()
14289 .map(|(output, _)| output.clone())
14290 .collect::<Vec<_>>(),
14291 planner.ctx.field_columns
14292 );
14293 assert_eq!(groups.len(), 2);
14294 assert!(
14295 groups
14296 .iter()
14297 .flat_map(|(_, pairs)| pairs)
14298 .all(|(_, right)| *right == &scale_fields[0])
14299 );
14300
14301 let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns(
14302 scale.schema(),
14303 mixed.schema(),
14304 &scale_fields,
14305 &planner.ctx.field_columns,
14306 binary.op,
14307 false,
14308 false,
14309 );
14310 assert!(invalid_pairs.is_empty());
14311 assert_eq!(
14312 groups
14313 .iter()
14314 .map(|(output, _)| output.clone())
14315 .collect::<Vec<_>>(),
14316 planner.ctx.field_columns
14317 );
14318 assert_eq!(groups.len(), 2);
14319 assert!(
14320 groups
14321 .iter()
14322 .flat_map(|(_, pairs)| pairs)
14323 .all(|(left, _)| *left == &scale_fields[0])
14324 );
14325 }
14326
14327 #[tokio::test]
14328 async fn test_non_bool_comparison_filters_mixed_sample_lanes() {
14329 let (planner, input) = mixed_direct_or(false).await;
14330 let input_schema = input.schema().clone();
14331 let plan = planner
14332 .filter_on_field_column(input, |field| {
14333 if PromPlanner::field_column_is_native_histogram(&input_schema, field) {
14334 Ok(lit(false))
14335 } else {
14336 Ok(col(field).gt(lit(0.0)))
14337 }
14338 })
14339 .unwrap();
14340 let float_field = planner.ctx.field_columns[0].clone();
14341
14342 let (_, batches) = execute(plan, &build_query_engine_state()).await;
14343 assert_eq!(values(&batches, &float_field), vec![1.25]);
14344 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
14345 }
14346
14347 #[tokio::test]
14348 async fn test_mixed_left_and_unless_preserve_sample_lanes() {
14349 for (expression, expected_sample_kind) in [
14350 ("lhs and on(k) mask", (false, true)),
14351 ("lhs unless on(k) mask", (true, false)),
14352 ] {
14353 let (mut planner, left) = mixed_direct_or(false).await;
14354 let left_context = planner.ctx.clone();
14355 let float_field = left_context.field_columns[0].clone();
14356 let histogram_field = left_context.field_columns[1].clone();
14357 let mask = tagged_source(
14358 "mask",
14359 false,
14360 ("k", Some("histogram")),
14361 DirectOrValue::Float64(1.0),
14362 );
14363 let PromExpr::Binary(binary) = parser::parse(expression).unwrap() else {
14364 unreachable!()
14365 };
14366 let plan = planner
14367 .set_op_on_non_field_columns(
14368 left,
14369 scan(&mask),
14370 left_context,
14371 direct_or_context("mask", &["job", "k"], "v"),
14372 binary.op,
14373 &binary.modifier,
14374 )
14375 .unwrap();
14376
14377 let (_, batches) = execute(plan, &build_query_engine_state()).await;
14378 let sample_kinds = batches
14379 .iter()
14380 .flat_map(|batch| {
14381 let floats = batch.column_by_name(&float_field).unwrap();
14382 let histograms = batch.column_by_name(&histogram_field).unwrap();
14383 (0..batch.num_rows())
14384 .map(|row| (floats.is_valid(row), histograms.is_valid(row)))
14385 })
14386 .collect::<Vec<_>>();
14387 assert_eq!(sample_kinds, vec![expected_sample_kind], "{expression}");
14388 }
14389 }
14390
14391 #[tokio::test]
14392 async fn test_mixed_fields_arithmetic_broadcasts_computed_scalar() {
14393 let plan = PromPlanner::stmt_to_plan(
14394 build_test_mixed_native_histogram_table_provider("some_metric").await,
14395 &build_eval_stmt("some_metric * scalar(vector(2))"),
14396 &build_query_engine_state(),
14397 )
14398 .await
14399 .unwrap();
14400 let schema = plan.schema();
14401 assert_eq!(
14402 schema
14403 .field_with_unqualified_name(greptime_value())
14404 .unwrap()
14405 .data_type(),
14406 &ArrowDataType::Float64
14407 );
14408 assert_eq!(
14409 schema
14410 .field_with_unqualified_name(greptime_native_histogram())
14411 .unwrap()
14412 .data_type(),
14413 &native_histogram_value_type().as_arrow_type()
14414 );
14415 assert!(
14416 plan.display_indent_schema()
14417 .to_string()
14418 .contains("prom_native_histogram_mul_scalar"),
14419 "{plan:?}"
14420 );
14421 }
14422
14423 #[tokio::test]
14424 async fn test_unsupported_histogram_binary_does_not_block_or_fallback() {
14425 let state = build_query_engine_state();
14426 let plan = PromPlanner::stmt_to_plan(
14427 operator_table_provider(),
14428 &operator_eval_stmt("((lf or on(tag) lh) % 2) or on(tag) lh"),
14429 &state,
14430 )
14431 .await
14432 .unwrap();
14433 let float_field = plan
14434 .schema()
14435 .fields()
14436 .iter()
14437 .find(|field| field.data_type() == &ArrowDataType::Float64)
14438 .unwrap()
14439 .name()
14440 .clone();
14441 let histogram_field = plan
14442 .schema()
14443 .fields()
14444 .iter()
14445 .find(|field| field.data_type() == &native_histogram_value_type().as_arrow_type())
14446 .unwrap()
14447 .name()
14448 .clone();
14449
14450 let (_, batches) = execute(plan, &state).await;
14451 assert_eq!(values(&batches, &float_field), vec![0.0]);
14452 assert_eq!(histograms(&batches, &histogram_field).len(), 1);
14453 }
14454
14455 #[tokio::test]
14456 async fn test_unary_negates_mixed_float_and_native_histogram_samples() {
14457 for histogram_on_left in [false, true] {
14458 let (mut planner, input) = mixed_direct_or(histogram_on_left).await;
14459 let plan = planner.negate_field_columns(input).unwrap();
14460 assert!(PromPlanner::field_columns_are_alternative_samples(
14461 plan.schema(),
14462 &planner.ctx.field_columns
14463 ));
14464 let float_field = planner
14465 .ctx
14466 .field_columns
14467 .iter()
14468 .find(|field| field.starts_with(OR_FLOAT_FIELD_PREFIX))
14469 .unwrap();
14470 let histogram_field = planner
14471 .ctx
14472 .field_columns
14473 .iter()
14474 .find(|field| field.starts_with(OR_HISTOGRAM_FIELD_PREFIX))
14475 .unwrap();
14476
14477 let (_, batches) = execute(plan, &build_query_engine_state()).await;
14478 assert_eq!(values(&batches, float_field), vec![-1.25]);
14479 let histogram = batches
14480 .iter()
14481 .find_map(|batch| {
14482 let values = batch
14483 .column_by_name(histogram_field)
14484 .unwrap()
14485 .as_any()
14486 .downcast_ref::<datafusion::arrow::array::StructArray>()
14487 .unwrap();
14488 (0..values.len()).find_map(|row| {
14489 common_query::native_histogram::read_histogram(values, row).unwrap()
14490 })
14491 })
14492 .unwrap();
14493 assert_eq!(histogram.count, -1.0);
14494 assert_eq!(histogram.sum, -1.0);
14495 assert_eq!(histogram.reset_hint, CounterResetHint::Gauge);
14496 }
14497 }
14498
14499 #[tokio::test]
14500 async fn test_native_histogram_sum_and_avg_execute_real_batches() {
14501 for op_name in ["sum", "avg"] {
14502 for incompatible in [false, true] {
14503 let mut second = direct_or_histogram();
14504 if incompatible {
14505 second.schema = CUSTOM_BUCKETS_SCHEMA;
14506 second.custom_values = vec![1.0];
14507 }
14508 let collector = PromqlAnnotationCollector::default();
14509 let (mut planner, input) =
14510 mixed_aggregate_input(vec![direct_or_histogram(), second]).await;
14511 planner.promql_annotations = Some(collector.clone());
14512 let histogram_column = planner.ctx.field_columns[1].clone();
14513 planner.ctx.field_columns = vec![histogram_column.clone()];
14514 let input = LogicalPlanBuilder::from(input)
14515 .project([col("ts"), col(&histogram_column)])
14516 .unwrap()
14517 .build()
14518 .unwrap();
14519 let PromExpr::Aggregate(AggregateExpr { op, param, .. }) =
14520 parser::parse(&format!("{op_name}(mixed)")).unwrap()
14521 else {
14522 unreachable!()
14523 };
14524 let (aggregate_exprs, _) =
14525 planner.create_aggregate_exprs(op, ¶m, &input).unwrap();
14526 let plan = LogicalPlanBuilder::from(input)
14527 .aggregate(vec![col("ts")], aggregate_exprs)
14528 .unwrap()
14529 .filter(planner.create_empty_values_filter_expr(false).unwrap())
14530 .unwrap()
14531 .build()
14532 .unwrap();
14533
14534 let (_, batches) = execute(plan, &build_query_engine_state()).await;
14535 let mut warnings = vec![];
14536 let mut infos = vec![];
14537 collector.append_to(&mut warnings, &mut infos);
14538 assert!(infos.is_empty());
14539 if incompatible {
14540 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
14541 assert!(warnings.iter().any(|warning| {
14542 warning
14543 == &format!(
14544 "prom_native_histogram_agg_{op_name}: dropped native histogram aggregate with incompatible schemas"
14545 )
14546 }));
14547 } else {
14548 let histograms = histograms(&batches, &histogram_column);
14549 assert_eq!(histograms.len(), 1);
14550 let expected = if op_name == "sum" { 2.0 } else { 1.0 };
14551 assert_eq!(histograms[0].count, expected);
14552 assert_eq!(histograms[0].sum, expected);
14553 assert!(warnings.is_empty());
14554 }
14555 }
14556 }
14557 }
14558
14559 #[tokio::test]
14560 async fn test_canonical_mixed_count_group_and_count_values_execute() {
14561 let state = build_query_engine_state();
14562 for (query, expected) in [
14563 ("count(some_metric)", vec![2.0]),
14564 ("group(some_metric)", vec![1.0]),
14565 (r#"count_values("sample", some_metric)"#, vec![1.0, 1.0]),
14566 ] {
14567 let plan = PromPlanner::stmt_to_plan(
14568 build_test_mixed_native_histogram_table_provider("some_metric").await,
14569 &operator_eval_stmt(query),
14570 &state,
14571 )
14572 .await
14573 .unwrap();
14574 assert!(
14575 plan.schema()
14576 .fields()
14577 .iter()
14578 .all(|field| !field.name().starts_with("__promql_sample_count")),
14579 "{query}: {plan:?}"
14580 );
14581 let value_fields = plan
14582 .schema()
14583 .fields()
14584 .iter()
14585 .filter(|field| {
14586 matches!(
14587 field.data_type(),
14588 ArrowDataType::Float64 | ArrowDataType::Int64 | ArrowDataType::UInt64
14589 ) || field.data_type() == &native_histogram_value_type().as_arrow_type()
14590 })
14591 .collect::<Vec<_>>();
14592 assert_eq!(value_fields.len(), 1, "{query}: {plan:?}");
14593 assert_ne!(
14594 value_fields[0].data_type(),
14595 &native_histogram_value_type().as_arrow_type(),
14596 "{query}: {plan:?}"
14597 );
14598 let value_column = value_fields[0].name().clone();
14599
14600 let (_, batches) = execute(plan, &state).await;
14601 let mut actual = numeric_values(&batches, &value_column);
14602 actual.sort_by(f64::total_cmp);
14603 assert_eq!(actual, expected, "{query}");
14604
14605 if query.starts_with("count_values") {
14606 let mut sample_labels = batches
14607 .iter()
14608 .flat_map(|batch| {
14609 batch
14610 .column_by_name("sample")
14611 .unwrap()
14612 .as_any()
14613 .downcast_ref::<StringArray>()
14614 .unwrap()
14615 .iter()
14616 .flatten()
14617 .map(str::to_string)
14618 })
14619 .collect::<Vec<_>>();
14620 sample_labels.sort();
14621 let mut expected_labels =
14622 vec!["2".to_string(), direct_or_histogram().promql_string()];
14623 expected_labels.sort();
14624 assert_eq!(sample_labels, expected_labels);
14625 }
14626 }
14627 }
14628
14629 #[tokio::test]
14630 async fn test_mixed_or_sum_aggregates_each_sample_type() {
14631 let PromExpr::Aggregate(AggregateExpr { op, param, .. }) =
14632 parser::parse("sum(lhs)").unwrap()
14633 else {
14634 unreachable!()
14635 };
14636
14637 let collector = PromqlAnnotationCollector::default();
14638 let (mut planner, input) = mixed_direct_or(false).await;
14639 planner.promql_annotations = Some(collector.clone());
14640 let float_column = planner.ctx.field_columns[0].clone();
14641 let histogram_column = planner.ctx.field_columns[1].clone();
14642 let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap();
14643 let plan = LogicalPlanBuilder::from(input)
14644 .aggregate(vec![col("ts"), col("k")], aggregate_exprs)
14645 .unwrap()
14646 .filter(
14647 planner
14648 .mixed_aggregate_filter_expr(op, &float_column, &histogram_column)
14649 .unwrap(),
14650 )
14651 .unwrap()
14652 .project([
14653 col(&float_column),
14654 col(&histogram_column),
14655 col("ts"),
14656 col("k"),
14657 ])
14658 .unwrap()
14659 .build()
14660 .unwrap();
14661
14662 let (_, batches) = execute(plan, &build_query_engine_state()).await;
14663 assert_eq!(values(&batches, &float_column), vec![1.25]);
14664 let histogram = batches
14665 .iter()
14666 .find_map(|batch| {
14667 let values = batch
14668 .column_by_name(&histogram_column)?
14669 .as_any()
14670 .downcast_ref::<datafusion::arrow::array::StructArray>()?;
14671 (0..values.len()).find_map(|row| {
14672 common_query::native_histogram::read_histogram(values, row).unwrap()
14673 })
14674 })
14675 .unwrap();
14676 assert_eq!(histogram.count, 1.0);
14677 let mut warnings = vec![];
14678 let mut infos = vec![];
14679 collector.append_to(&mut warnings, &mut infos);
14680 assert!(warnings.is_empty());
14681
14682 let collector = PromqlAnnotationCollector::default();
14683 let (mut planner, input) = mixed_direct_or(false).await;
14684 planner.promql_annotations = Some(collector.clone());
14685 let float_column = planner.ctx.field_columns[0].clone();
14686 let histogram_column = planner.ctx.field_columns[1].clone();
14687 let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap();
14688 let plan = LogicalPlanBuilder::from(input)
14689 .aggregate(vec![col("ts")], aggregate_exprs)
14690 .unwrap()
14691 .filter(
14692 planner
14693 .mixed_aggregate_filter_expr(op, &float_column, &histogram_column)
14694 .unwrap(),
14695 )
14696 .unwrap()
14697 .project([col(&float_column), col(&histogram_column), col("ts")])
14698 .unwrap()
14699 .build()
14700 .unwrap();
14701
14702 let (_, batches) = execute(plan, &build_query_engine_state()).await;
14703 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
14704 let mut warnings = vec![];
14705 let mut infos = vec![];
14706 collector.append_to(&mut warnings, &mut infos);
14707 assert_eq!(
14708 warnings,
14709 vec![
14710 "sum: dropped aggregation result containing both float and native histogram samples"
14711 ]
14712 );
14713 }
14714
14715 #[tokio::test]
14716 async fn test_mixed_or_sum_drops_incompatible_mixed_group() {
14717 let PromExpr::Aggregate(AggregateExpr { op, param, .. }) =
14718 parser::parse("sum(lhs)").unwrap()
14719 else {
14720 unreachable!()
14721 };
14722 let mut custom = direct_or_histogram();
14723 custom.schema = CUSTOM_BUCKETS_SCHEMA;
14724 custom.custom_values = vec![1.0];
14725 let collector = PromqlAnnotationCollector::default();
14726 let (mut planner, input) = mixed_aggregate_input(vec![direct_or_histogram(), custom]).await;
14727 planner.promql_annotations = Some(collector.clone());
14728 let float_column = planner.ctx.field_columns[0].clone();
14729 let histogram_column = planner.ctx.field_columns[1].clone();
14730 let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap();
14731 let plan = LogicalPlanBuilder::from(input)
14732 .aggregate(vec![col("ts")], aggregate_exprs)
14733 .unwrap()
14734 .filter(
14735 planner
14736 .mixed_aggregate_filter_expr(op, &float_column, &histogram_column)
14737 .unwrap(),
14738 )
14739 .unwrap()
14740 .project([col(&float_column), col(&histogram_column), col("ts")])
14741 .unwrap()
14742 .build()
14743 .unwrap();
14744
14745 let (_, batches) = execute(plan, &build_query_engine_state()).await;
14746 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
14747 let mut warnings = vec![];
14748 let mut infos = vec![];
14749 collector.append_to(&mut warnings, &mut infos);
14750 assert!(warnings.iter().any(|warning| {
14751 warning
14752 == "sum: dropped aggregation result containing both float and native histogram samples"
14753 }));
14754 }
14755
14756 #[tokio::test]
14757 async fn test_mixed_or_min_records_only_present_histograms() {
14758 let PromExpr::Aggregate(AggregateExpr { op, param, .. }) =
14759 parser::parse("min(lhs)").unwrap()
14760 else {
14761 unreachable!()
14762 };
14763 let expected_info = "min: dropped native histogram samples because this aggregation is not supported for native histograms";
14764
14765 for (histograms, expected_infos) in [
14766 (vec![], vec![]),
14767 (vec![direct_or_histogram()], vec![expected_info]),
14768 ] {
14769 let collector = PromqlAnnotationCollector::default();
14770 let (mut planner, input) = mixed_aggregate_input(histograms).await;
14771 planner.promql_annotations = Some(collector.clone());
14772 let float_column = planner.ctx.field_columns[0].clone();
14773 let histogram_column = planner.ctx.field_columns[1].clone();
14774 let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap();
14775 let plan = LogicalPlanBuilder::from(input)
14776 .aggregate(vec![col("ts")], aggregate_exprs)
14777 .unwrap()
14778 .filter(
14779 planner
14780 .mixed_ignored_histogram_filter_expr(op, &histogram_column)
14781 .unwrap(),
14782 )
14783 .unwrap()
14784 .project([col(&float_column), col("ts")])
14785 .unwrap()
14786 .build()
14787 .unwrap();
14788
14789 let (_, batches) = execute(plan, &build_query_engine_state()).await;
14790 assert_eq!(values(&batches, &float_column), vec![1.25]);
14791 let mut warnings = vec![];
14792 let mut infos = vec![];
14793 collector.append_to(&mut warnings, &mut infos);
14794 assert!(warnings.is_empty());
14795 assert_eq!(infos, expected_infos);
14796 }
14797 }
14798
14799 #[tokio::test]
14800 async fn test_mixed_or_value_aliases_do_not_replace_labels() {
14801 let left = source(
14802 "lhs",
14803 false,
14804 1,
14805 vec![("job", Some("job")), ("k", Some("float"))],
14806 DirectOrValue::Float64(1.0),
14807 );
14808 let right = source(
14809 "rhs",
14810 false,
14811 1,
14812 vec![
14813 ("job", Some("job")),
14814 ("k", Some("histogram")),
14815 (greptime_value(), Some("value-label")),
14816 ],
14817 DirectOrValue::NativeHistogram(direct_or_histogram()),
14818 );
14819 let table_provider = build_test_table_provider_with_fields(
14820 &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
14821 &[],
14822 )
14823 .await;
14824 let mut planner = PromPlanner {
14825 table_provider,
14826 ctx: PromPlannerContext::default(),
14827 promql_annotations: None,
14828 };
14829 let left = LogicalPlanBuilder::from(scan(&left))
14830 .project(vec![
14831 col("ts"),
14832 col("job"),
14833 col("k"),
14834 col("v").alias(greptime_value()),
14835 ])
14836 .unwrap()
14837 .build()
14838 .unwrap();
14839 let left_context = direct_or_context("lhs", &["job", "k"], greptime_value());
14840 let right_context = direct_or_context("rhs", &["job", "k", greptime_value()], "v");
14841 let plan = planner
14842 .or_operator(
14843 left,
14844 scan(&right),
14845 left_context.tag_columns.iter().cloned().collect(),
14846 right_context.tag_columns.iter().cloned().collect(),
14847 left_context,
14848 right_context,
14849 &or_modifier("lhs or on(k) rhs"),
14850 )
14851 .unwrap();
14852
14853 assert_eq!(
14854 plan.schema()
14855 .field_with_name(None, greptime_value())
14856 .unwrap()
14857 .data_type(),
14858 &ArrowDataType::Utf8
14859 );
14860 assert!(
14861 planner
14862 .ctx
14863 .field_columns
14864 .iter()
14865 .all(|field| { field != greptime_value() && field != greptime_native_histogram() })
14866 );
14867 assert!(PromPlanner::field_columns_are_alternative_samples(
14868 plan.schema(),
14869 &planner.ctx.field_columns
14870 ));
14871 let (_, batches) = execute(plan, &build_query_engine_state()).await;
14872 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
14873 let labels = batches
14874 .iter()
14875 .flat_map(|batch| {
14876 batch
14877 .column_by_name(greptime_value())
14878 .unwrap()
14879 .as_any()
14880 .downcast_ref::<StringArray>()
14881 .unwrap()
14882 .iter()
14883 .flatten()
14884 })
14885 .collect::<Vec<_>>();
14886 assert_eq!(labels, vec!["value-label"]);
14887 }
14888
14889 #[tokio::test]
14890 async fn test_mixed_or_routes_float_histogram_and_label_functions() {
14891 for (function, expected) in [("abs", 1.25), ("round", 1.0), ("histogram_count", 1.0)] {
14892 let (mut planner, input) = mixed_direct_or(false).await;
14893 let preserve_any_value = PromPlanner::field_columns_are_alternative_samples(
14894 input.schema(),
14895 &planner.ctx.field_columns,
14896 );
14897 let PromExpr::Call(call) = parser::parse(&format!("{function}(lhs)")).unwrap() else {
14898 unreachable!()
14899 };
14900 let state = build_query_engine_state();
14901 let (mut exprs, _) = planner
14902 .create_function_expr(&call.func, vec![], input.schema(), &state)
14903 .unwrap();
14904 exprs.insert(0, planner.create_time_index_column_expr().unwrap());
14905 exprs.extend(planner.create_tag_column_exprs().unwrap());
14906 let plan = LogicalPlanBuilder::from(input)
14907 .project(exprs)
14908 .unwrap()
14909 .filter(
14910 planner
14911 .create_empty_values_filter_expr(preserve_any_value)
14912 .unwrap(),
14913 )
14914 .unwrap()
14915 .build()
14916 .unwrap();
14917 let (_, batches) = execute(plan, &state).await;
14918 let values = batches
14919 .iter()
14920 .flat_map(|batch| {
14921 batch
14922 .schema()
14923 .fields()
14924 .iter()
14925 .position(|field| field.data_type() == &ArrowDataType::Float64)
14926 .map(|index| {
14927 batch
14928 .column(index)
14929 .as_any()
14930 .downcast_ref::<Float64Array>()
14931 .unwrap()
14932 .iter()
14933 .flatten()
14934 })
14935 .into_iter()
14936 .flatten()
14937 })
14938 .collect::<Vec<_>>();
14939 assert_eq!(values, vec![expected], "{function}");
14940 }
14941
14942 let (mut planner, input) = mixed_direct_or(false).await;
14943 let preserve_any_value = PromPlanner::field_columns_are_alternative_samples(
14944 input.schema(),
14945 &planner.ctx.field_columns,
14946 );
14947 let PromExpr::Call(call) =
14948 parser::parse(r#"label_replace(lhs, "copy", "$1", "k", "(.*)")"#).unwrap()
14949 else {
14950 unreachable!()
14951 };
14952 let args = planner.create_function_args(&call.args.args).unwrap();
14953 let state = build_query_engine_state();
14954 let (mut exprs, _) = planner
14955 .create_function_expr(&call.func, args.literals, input.schema(), &state)
14956 .unwrap();
14957 exprs.insert(0, planner.create_time_index_column_expr().unwrap());
14958 exprs.extend(planner.create_tag_column_exprs().unwrap());
14959 let plan = LogicalPlanBuilder::from(input)
14960 .project(exprs)
14961 .unwrap()
14962 .filter(
14963 planner
14964 .create_empty_values_filter_expr(preserve_any_value)
14965 .unwrap(),
14966 )
14967 .unwrap()
14968 .build()
14969 .unwrap();
14970 let (_, batches) = execute(plan, &state).await;
14971 let sample_count = batches.iter().map(RecordBatch::num_rows).sum::<usize>();
14972 assert_eq!(sample_count, 2);
14973 }
14974
14975 async fn build_count_values_table_provider() -> DfTableSourceProvider {
14978 build_count_values_table_provider_with_values(&[1.0, 2.0, 1.0]).await
14979 }
14980
14981 async fn build_count_values_table_provider_with_values(
14984 values: &[f64],
14985 ) -> DfTableSourceProvider {
14986 build_count_values_table_provider_with_value_array(Arc::new(Float64Array::from(
14987 values.to_vec(),
14988 )))
14989 .await
14990 }
14991
14992 async fn build_count_values_table_provider_with_value_array(
14995 values: ArrayRef,
14996 ) -> DfTableSourceProvider {
14997 let value_data_type = ConcreteDataType::from_arrow_type(values.data_type());
14998 let catalog_list = MemoryCatalogManager::with_default_setup();
14999 let columns = vec![
15000 ColumnSchema::new("k".to_string(), ConcreteDataType::string_datatype(), false),
15001 ColumnSchema::new(
15002 "timestamp".to_string(),
15003 ConcreteDataType::timestamp_millisecond_datatype(),
15004 false,
15005 )
15006 .with_time_index(true),
15007 ColumnSchema::new(greptime_value().to_string(), value_data_type, true),
15008 ];
15009 let schema = Arc::new(Schema::new(columns));
15010 let table_meta = TableMetaBuilder::empty()
15011 .schema(schema.clone())
15012 .primary_key_indices(vec![0])
15013 .value_indices(vec![2])
15014 .next_column_id(1024)
15015 .build()
15016 .unwrap();
15017 let table_info = Arc::new(
15018 TableInfoBuilder::default()
15019 .table_id(3_001)
15020 .name("cv_metric")
15021 .meta(table_meta)
15022 .build()
15023 .unwrap(),
15024 );
15025 let batch = RecordBatch::try_new(
15026 schema.arrow_schema().clone(),
15027 vec![
15028 Arc::new(StringArray::from(
15029 (0..values.len())
15030 .map(|index| format!("k{index}"))
15031 .collect::<Vec<_>>(),
15032 )),
15033 Arc::new(TimestampMillisecondArray::from(vec![1_000; values.len()])),
15034 values.clone(),
15035 ],
15036 )
15037 .unwrap();
15038 let backing = GreptimeMemTable::new_with_catalog(
15039 "cv_metric",
15040 GreptimeRecordBatch::from_df_record_batch(schema, batch),
15041 3_001,
15042 DEFAULT_CATALOG_NAME.to_string(),
15043 DEFAULT_SCHEMA_NAME.to_string(),
15044 );
15045 let table = Arc::new(Table::new(
15046 table_info,
15047 FilterPushDownType::Unsupported,
15048 backing.data_source(),
15049 ));
15050
15051 assert!(
15052 catalog_list
15053 .register_table_sync(RegisterTableRequest {
15054 catalog: DEFAULT_CATALOG_NAME.to_string(),
15055 schema: DEFAULT_SCHEMA_NAME.to_string(),
15056 table_name: "cv_metric".to_string(),
15057 table_id: 3_001,
15058 table,
15059 })
15060 .is_ok()
15061 );
15062
15063 DfTableSourceProvider::new(
15064 catalog_list,
15065 false,
15066 QueryContext::arc(),
15067 DummyDecoder::arc(),
15068 false,
15069 )
15070 }
15071
15072 fn count_values_rows<'a>(batches: &'a [RecordBatch], label: &str) -> Vec<(&'a str, f64)> {
15079 let mut rows = batches
15080 .iter()
15081 .flat_map(|batch| {
15082 let value_index = batch
15084 .schema()
15085 .fields()
15086 .iter()
15087 .position(|field| {
15088 field.name() != label
15089 && matches!(
15090 field.data_type(),
15091 ArrowDataType::Float64
15092 | ArrowDataType::Int64
15093 | ArrowDataType::UInt64
15094 )
15095 })
15096 .expect("no aggregated value column");
15097 let labels = batch
15098 .column_by_name(label)
15099 .expect("no generated label column")
15100 .as_any()
15101 .downcast_ref::<StringArray>()
15102 .expect("the generated label must be a string column");
15103 let values = datafusion::arrow::compute::cast(
15104 batch.column(value_index),
15105 &ArrowDataType::Float64,
15106 )
15107 .unwrap();
15108 let values = values.as_any().downcast_ref::<Float64Array>().unwrap();
15109 labels
15110 .iter()
15111 .zip(values.iter())
15112 .map(|(label, value)| (label.unwrap(), value.unwrap()))
15113 .collect::<Vec<_>>()
15114 })
15115 .collect::<Vec<_>>();
15116 rows.sort_by(|left, right| left.0.cmp(right.0).then(left.1.total_cmp(&right.1)));
15117 rows
15118 }
15119
15120 fn assert_unique_label_set_per_timestamp(batches: &[RecordBatch], label: &str) {
15125 let mut seen = HashMap::<i64, HashSet<String>>::new();
15126 for batch in batches {
15127 let timestamp_index = batch
15128 .schema()
15129 .fields()
15130 .iter()
15131 .position(|field| matches!(field.data_type(), ArrowDataType::Timestamp(..)))
15132 .expect("no timestamp column");
15133 let timestamps = batch
15134 .column(timestamp_index)
15135 .as_any()
15136 .downcast_ref::<TimestampMillisecondArray>()
15137 .expect("timestamp column is not a millisecond timestamp");
15138 let labels = batch
15139 .column_by_name(label)
15140 .expect("no generated label column")
15141 .as_any()
15142 .downcast_ref::<StringArray>()
15143 .expect("the generated label must be a string column");
15144 for (timestamp, label) in timestamps.iter().zip(labels.iter()) {
15145 let timestamp = timestamp.unwrap();
15146 let label = label.unwrap();
15147 assert!(
15148 seen.entry(timestamp).or_default().insert(label.to_string()),
15149 "duplicated label set `{label}` at timestamp {timestamp}"
15150 );
15151 }
15152 }
15153 }
15154
15155 #[tokio::test]
15156 async fn test_count_values_generated_label_survives_enclosing_expr() {
15157 for (case, label) in [
15159 (r#"count_values("v", prometheus_tsdb_head_series)"#, "v"),
15160 (
15161 r#"abs(count_values("v", prometheus_tsdb_head_series))"#,
15162 "v",
15163 ),
15164 (
15165 r#"round(count_values("v", prometheus_tsdb_head_series))"#,
15166 "v",
15167 ),
15168 (r#"count_values("v", prometheus_tsdb_head_series) + 1"#, "v"),
15169 (
15170 r#"topk(1, count_values("v", prometheus_tsdb_head_series))"#,
15171 "v",
15172 ),
15173 (
15174 r#"sum by (v) (count_values("v", prometheus_tsdb_head_series))"#,
15175 "v",
15176 ),
15177 (
15178 r#"label_replace(count_values("v", prometheus_tsdb_head_series), "vcopy", "$1", "v", "(.*)")"#,
15179 "v",
15180 ),
15181 (
15182 r#"count_values("v", prometheus_tsdb_head_series) by (ip) + 1"#,
15183 "v",
15184 ),
15185 (
15187 r#"count_values("ip", prometheus_tsdb_head_series) by (ip)"#,
15188 "ip",
15189 ),
15190 (
15191 r#"count_values("ip", prometheus_tsdb_head_series) by (ip) + 1"#,
15192 "ip",
15193 ),
15194 ] {
15195 let plan = PromPlanner::stmt_to_plan(
15196 build_test_table_provider_with_fields(
15197 &[(
15198 DEFAULT_SCHEMA_NAME.to_string(),
15199 "prometheus_tsdb_head_series".to_string(),
15200 )],
15201 &["ip"],
15202 )
15203 .await,
15204 &build_eval_stmt(case),
15205 &build_query_engine_state(),
15206 )
15207 .await
15208 .unwrap();
15209
15210 let label_columns = plan
15211 .schema()
15212 .fields()
15213 .iter()
15214 .filter(|field| field.name() == label)
15215 .count();
15216 assert_eq!(
15217 label_columns,
15218 1,
15219 "{case}: the `{label}` label must survive: {}",
15220 plan.display_indent()
15221 );
15222 }
15223 }
15224
15225 #[tokio::test]
15226 async fn test_count_values_generated_label_in_enclosing_expr_execute() {
15227 let state = build_query_engine_state();
15229 for (query, label, expected) in [
15231 (
15232 r#"count_values("v", cv_metric)"#,
15233 "v",
15234 vec![("1", 2.0), ("2", 1.0)],
15235 ),
15236 (
15237 r#"abs(count_values("v", cv_metric))"#,
15238 "v",
15239 vec![("1", 2.0), ("2", 1.0)],
15240 ),
15241 (
15242 r#"round(count_values("v", cv_metric))"#,
15243 "v",
15244 vec![("1", 2.0), ("2", 1.0)],
15245 ),
15246 (
15247 r#"count_values("v", cv_metric) + 1"#,
15248 "v",
15249 vec![("1", 3.0), ("2", 2.0)],
15250 ),
15251 (
15252 r#"sum by (v) (count_values("v", cv_metric))"#,
15253 "v",
15254 vec![("1", 2.0), ("2", 1.0)],
15255 ),
15256 (
15257 r#"topk(10, count_values("v", cv_metric))"#,
15258 "v",
15259 vec![("1", 2.0), ("2", 1.0)],
15260 ),
15261 (
15265 r#"count_values("k", cv_metric) by (k)"#,
15266 "k",
15267 vec![("1", 2.0), ("2", 1.0)],
15268 ),
15269 ] {
15270 let plan = PromPlanner::stmt_to_plan(
15271 build_count_values_table_provider().await,
15272 &operator_eval_stmt(query),
15273 &state,
15274 )
15275 .await
15276 .unwrap_or_else(|err| panic!("{query}: {err}"));
15277
15278 assert_eq!(
15279 plan.schema()
15280 .fields()
15281 .iter()
15282 .filter(|field| field.name() == label)
15283 .count(),
15284 1,
15285 "{query}: {}",
15286 plan.display_indent()
15287 );
15288
15289 let (_, batches) = execute(plan, &state).await;
15290 assert_eq!(count_values_rows(&batches, label), expected, "{query}");
15291 assert_unique_label_set_per_timestamp(&batches, label);
15292 }
15293 }
15294
15295 #[tokio::test]
15296 async fn test_count_values_label_is_prometheus_formatted_value() {
15297 let state = build_query_engine_state();
15302 let plan = PromPlanner::stmt_to_plan(
15306 build_count_values_table_provider_with_values(&[
15307 -0.0, 0.0, 1.0, 0.5, 200.0, 1e21, 1e-7, 2.5,
15308 ])
15309 .await,
15310 &operator_eval_stmt(r#"count_values("v", cv_metric)"#),
15311 &state,
15312 )
15313 .await
15314 .unwrap();
15315
15316 let (_, batches) = execute(plan, &state).await;
15317 assert_unique_label_set_per_timestamp(&batches, "v");
15318 let mut labels = count_values_rows(&batches, "v")
15319 .into_iter()
15320 .map(|(label, _)| label)
15321 .collect::<Vec<_>>();
15322 labels.sort();
15323 assert_eq!(
15324 labels,
15325 vec![
15326 "-0",
15327 "0",
15328 "0.0000001",
15329 "0.5",
15330 "1",
15331 "1000000000000000000000",
15332 "2.5",
15333 "200",
15334 ]
15335 );
15336 }
15337
15338 #[tokio::test]
15339 async fn test_count_values_groups_by_formatted_value_for_bigint_input() {
15340 let state = build_query_engine_state();
15345 let plan = PromPlanner::stmt_to_plan(
15346 build_count_values_table_provider_with_value_array(Arc::new(Int64Array::from(vec![
15347 9_007_199_254_740_992_i64,
15348 9_007_199_254_740_993_i64,
15349 ])))
15350 .await,
15351 &operator_eval_stmt(r#"count_values("v", cv_metric)"#),
15352 &state,
15353 )
15354 .await
15355 .unwrap();
15356
15357 let (_, batches) = execute(plan, &state).await;
15358 assert_unique_label_set_per_timestamp(&batches, "v");
15360 assert_eq!(
15361 count_values_rows(&batches, "v"),
15362 vec![("9007199254740992", 2.0)]
15363 );
15364 }
15365}