1use std::collections::{BTreeSet, HashMap, HashSet, VecDeque};
16use std::sync::Arc;
17use std::time::UNIX_EPOCH;
18
19use arrow::datatypes::IntervalDayTime;
20use async_recursion::async_recursion;
21use catalog::table_source::DfTableSourceProvider;
22use common_error::ext::ErrorExt;
23use common_error::status_code::StatusCode;
24use common_function::function::FunctionContext;
25use common_query::native_histogram::native_histogram_value_type;
26use common_query::prelude::{
27 GREPTIME_TEMPORALITY_DELTA, OTLP_AGGREGATION_TEMPORALITY_LABEL, greptime_native_histogram,
28 greptime_value,
29};
30use common_query::promql_annotations::PromqlAnnotationCollector;
31use datafusion::common::DFSchemaRef;
32use datafusion::datasource::DefaultTableSource;
33use datafusion::functions_aggregate::average::avg_udaf;
34use datafusion::functions_aggregate::count::count_udaf;
35use datafusion::functions_aggregate::expr_fn::first_value;
36use datafusion::functions_aggregate::min_max::{max_udaf, min_udaf};
37use datafusion::functions_aggregate::stddev::stddev_pop_udaf;
38use datafusion::functions_aggregate::sum::sum_udaf;
39use datafusion::functions_aggregate::variance::var_pop_udaf;
40use datafusion::functions_window::row_number::RowNumber;
41use datafusion::logical_expr::expr::{Alias, ScalarFunction, WindowFunction};
42use datafusion::logical_expr::expr_rewriter::normalize_cols;
43use datafusion::logical_expr::{
44 BinaryExpr, Cast, Extension, LogicalPlan, LogicalPlanBuilder, Operator,
45 ScalarUDF as ScalarUdfDef, WindowFrame, WindowFunctionDefinition,
46};
47use datafusion::optimizer::simplify_expressions::ExprSimplifier;
48use datafusion::prelude as df_prelude;
49use datafusion::prelude::{Column, Expr as DfExpr, JoinType};
50use datafusion::scalar::ScalarValue;
51use datafusion::sql::TableReference;
52use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter};
53use datafusion_common::{DFSchema, NullEquality};
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::{
59 ExprSchemable, Literal, Projection, SortExpr, TableScan, TableSource, col, lit,
60};
61use datafusion_functions::core::coalesce;
62use datatypes::arrow::datatypes::{DataType as ArrowDataType, TimeUnit as ArrowTimeUnit};
63use datatypes::data_type::{ConcreteDataType, DataType as GreptimeDataType};
64use itertools::Itertools;
65use once_cell::sync::Lazy;
66use promql::extension_plan::{
67 Absent, EmptyMetric, HistogramFold, HistogramFoldOperation, InstantManipulate, Millisecond,
68 RangeManipulate, ScalarCalculate, SeriesDivide, SeriesNormalize, UnionDistinctOn,
69 build_special_time_expr,
70};
71use promql::functions::{
72 AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing,
73 IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, MixedRange,
74 NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAggAvg,
75 NativeHistogramAggSum, NativeHistogramAvg, NativeHistogramAvgOverTime, NativeHistogramChanges,
76 NativeHistogramCount, NativeHistogramCountOverTime, NativeHistogramDelta,
77 NativeHistogramDivScalar, NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction,
78 NativeHistogramIDelta, NativeHistogramIRate, NativeHistogramIncrease,
79 NativeHistogramLastOverTime, NativeHistogramMulScalar, NativeHistogramNeg,
80 NativeHistogramNotEq, NativeHistogramPresentOverTime, NativeHistogramQuantile,
81 NativeHistogramRate, NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev,
82 NativeHistogramStdvar, NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime,
83 NativeHistogramToString, PredictLinear, PresentOverTime, PromqlFloatToString, QuantileOverTime,
84 Rate, Resets, Round, StddevOverTime, StdvarOverTime, SumOverTime, quantile_udaf,
85};
86use promql_parser::label::{METRIC_NAME, MatchOp, Matcher, Matchers};
87use promql_parser::parser::token::TokenType;
88use promql_parser::parser::value::ValueType;
89use promql_parser::parser::{
90 AggregateExpr, BinModifier, BinaryExpr as PromBinaryExpr, Call, EvalStmt, Expr as PromExpr,
91 Function, FunctionArgs as PromFunctionArgs, LabelModifier, MatrixSelector, NumberLiteral,
92 Offset, ParenExpr, StringLiteral, SubqueryExpr, UnaryExpr, VectorMatchCardinality,
93 VectorSelector, token,
94};
95use regex::{self, Regex};
96use snafu::{OptionExt, ResultExt, ensure};
97use store_api::metric_engine_consts::{
98 DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, LOGICAL_TABLE_METADATA_KEY,
99 METRIC_ENGINE_NAME, is_metric_engine_internal_column,
100};
101use table::table::adapter::DfTableProviderAdapter;
102
103use crate::parser::{
104 ALIAS_NODE_NAME, ANALYZE_NODE_NAME, ANALYZE_VERBOSE_NODE_NAME, AliasExpr, EXPLAIN_NODE_NAME,
105 EXPLAIN_VERBOSE_NODE_NAME,
106};
107use crate::promql::error::{
108 CatalogSnafu, ColumnNotFoundSnafu, CombineTableColumnMismatchSnafu, DataFusionPlanningSnafu,
109 ExpectRangeSelectorSnafu, FunctionInvalidArgumentSnafu, InvalidDestinationLabelNameSnafu,
110 InvalidRegularExpressionSnafu, InvalidTimeRangeSnafu, MultiFieldsNotSupportedSnafu,
111 MultipleMetricMatchersSnafu, MultipleVectorSnafu, NoMetricMatcherSnafu, PromqlPlanNodeSnafu,
112 Result, SameLabelSetSnafu, TableNameNotFoundSnafu, TimeIndexNotFoundSnafu,
113 UnexpectedPlanExprSnafu, UnexpectedTokenSnafu, UnknownTableSnafu, UnsupportedExprSnafu,
114 UnsupportedMatcherOpSnafu, UnsupportedVectorMatchSnafu, ValueNotFoundSnafu,
115 ZeroRangeSelectorSnafu,
116};
117use crate::query_engine::QueryEngineState;
118
119const SPECIAL_TIME_FUNCTION: &str = "time";
121const SCALAR_FUNCTION: &str = "scalar";
123const SPECIAL_ABSENT_FUNCTION: &str = "absent";
125const SPECIAL_HISTOGRAM_QUANTILE: &str = "histogram_quantile";
127const SPECIAL_HISTOGRAM_FRACTION: &str = "histogram_fraction";
129const SPECIAL_VECTOR_FUNCTION: &str = "vector";
131const LE_COLUMN_NAME: &str = "le";
133
134static LABEL_NAME_REGEX: Lazy<Regex> =
137 Lazy::new(|| Regex::new(r"^[a-zA-Z_][a-zA-Z0-9_]*$").unwrap());
138
139const DEFAULT_TIME_INDEX_COLUMN: &str = "time";
140
141const DEFAULT_FIELD_COLUMN: &str = "value";
143
144const FIELD_COLUMN_MATCHER: &str = "__field__";
146
147const SCHEMA_COLUMN_MATCHER: &str = "__schema__";
149const DB_COLUMN_MATCHER: &str = "__database__";
150
151const BINARY_ISLAND_LEAF_ALIAS_PREFIX: &str = "__prom_v";
153const OR_FLOAT_FIELD_PREFIX: &str = "__promql_or_float_";
154const OR_HISTOGRAM_FIELD_PREFIX: &str = "__promql_or_histogram_";
155const TIMESTAMP_VALUE_PREFIX: &str = "__promql_timestamp_value_";
156
157const MAX_SCATTER_POINTS: i64 = 400;
159
160const INTERVAL_1H: i64 = 60 * 60 * 1000;
162
163#[derive(Default, Debug, Clone)]
164struct PromPlannerContext {
165 start: Millisecond,
167 end: Millisecond,
168 interval: Millisecond,
169 lookback_delta: Millisecond,
170
171 table_name: Option<String>,
173 time_index_column: Option<String>,
174 field_columns: Vec<String>,
175 tag_columns: Vec<String>,
176 use_tsid: bool,
182 field_column_matcher: Option<Vec<Matcher>>,
184 selector_matcher: Vec<Matcher>,
186 schema_name: Option<String>,
187 range: Option<Millisecond>,
189}
190
191#[derive(Debug, Clone, PartialEq, Eq, Hash)]
192struct VectorLeafKey {
193 metric_name: String,
194 matchers: Vec<(String, String, String)>,
195 or_matchers: Vec<Vec<(String, String, String)>>,
196 offset_ms: i128,
197 at: String,
198}
199
200#[derive(Debug, Clone)]
201struct IslandLeaf {
202 selector: VectorSelector,
203 display_table: String,
204}
205
206#[derive(Debug, Clone)]
207enum IslandExpr {
208 VectorLeaf(usize),
209 Scalar(DfExpr),
210 Unary {
211 input: Box<IslandExpr>,
212 },
213 Binary {
214 op: TokenType,
215 lhs: Box<IslandExpr>,
216 rhs: Box<IslandExpr>,
217 },
218}
219
220impl IslandExpr {
221 fn try_new(expr: &PromExpr, env: &mut IslandCollectEnv) -> Option<Self> {
222 if let Some(expr) = PromPlanner::try_build_literal_expr(expr) {
223 return Some(Self::Scalar(expr));
224 }
225
226 match expr {
227 PromExpr::Paren(ParenExpr { expr }) => Self::try_new(expr, env),
228 PromExpr::VectorSelector(selector) => {
229 let leaf = env.intern_leaf(selector)?;
230 Some(Self::VectorLeaf(leaf))
231 }
232 PromExpr::Unary(UnaryExpr { expr }) => {
233 let input = Self::try_new(expr, env)?;
234 Some(Self::Unary {
235 input: Box::new(input),
236 })
237 }
238 PromExpr::Binary(PromBinaryExpr {
239 lhs,
240 rhs,
241 op,
242 modifier,
243 }) if matches!(
244 op.id(),
245 token::T_ADD
246 | token::T_SUB
247 | token::T_MUL
248 | token::T_DIV
249 | token::T_MOD
250 | token::T_POW
251 | token::T_ATAN2
252 ) && modifier.as_ref().is_none_or(|modifier| {
253 !modifier.return_bool
254 && modifier.matching.is_none()
255 && matches!(modifier.card, VectorMatchCardinality::OneToOne)
256 && modifier.fill_values.lhs.is_none()
257 && modifier.fill_values.rhs.is_none()
258 }) =>
259 {
260 let lhs = Self::try_new(lhs, env)?;
261 let rhs = Self::try_new(rhs, env)?;
262 Some(Self::Binary {
263 op: *op,
264 lhs: Box::new(lhs),
265 rhs: Box::new(rhs),
266 })
267 }
268 _ => None,
269 }
270 }
271}
272
273#[derive(Debug, Default)]
274struct IslandCollectEnv {
275 leaf_by_key: HashMap<VectorLeafKey, usize>,
276 leaves: Vec<IslandLeaf>,
277 vector_occurrences: usize,
278}
279
280#[derive(Debug)]
281struct PlannedIslandLeaf {
282 plan: LogicalPlan,
283 ctx: PromPlannerContext,
284 alias: TableReference,
285 display_table: String,
286}
287
288#[derive(Debug)]
289struct IslandFieldExprs {
290 exprs: Vec<DfExpr>,
291 names: Vec<String>,
292 scalar: bool,
293}
294
295impl VectorLeafKey {
296 fn from_selector(selector: &VectorSelector) -> Option<Self> {
297 let mut metric_name = selector.name.clone();
298 let mut matchers = Vec::with_capacity(selector.matchers.matchers.len());
299 let matcher_key = |matcher: &Matcher| {
300 (
301 matcher.name.clone(),
302 matcher.op.to_string(),
303 matcher.value.clone(),
304 )
305 };
306
307 for matcher in &selector.matchers.matchers {
308 if matcher.name == METRIC_NAME {
309 if matcher.op != MatchOp::Equal || metric_name.is_some() {
310 return None;
311 }
312 metric_name = Some(matcher.value.clone());
313 } else {
314 matchers.push(matcher_key(matcher));
315 }
316 }
317 matchers.sort();
318
319 let mut or_matchers = selector
320 .matchers
321 .or_matchers
322 .iter()
323 .map(|group| {
324 let mut group = group.iter().map(matcher_key).collect::<Vec<_>>();
325 group.sort();
326 group
327 })
328 .collect::<Vec<_>>();
329 or_matchers.sort();
330
331 Some(Self {
332 metric_name: metric_name?,
333 matchers,
334 or_matchers,
335 offset_ms: match &selector.offset {
336 Some(Offset::Pos(duration)) => duration.as_millis() as i128,
337 Some(Offset::Neg(duration)) => -(duration.as_millis() as i128),
338 None => 0,
339 },
340 at: format!("{:?}", selector.at),
341 })
342 }
343}
344
345impl IslandCollectEnv {
346 fn intern_leaf(&mut self, selector: &VectorSelector) -> Option<usize> {
347 self.vector_occurrences += 1;
348 let key = VectorLeafKey::from_selector(selector)?;
349 if let Some(id) = self.leaf_by_key.get(&key) {
350 return Some(*id);
351 }
352
353 let id = self.leaves.len();
354 self.leaves.push(IslandLeaf {
355 selector: selector.clone(),
356 display_table: key.metric_name.clone(),
357 });
358 self.leaf_by_key.insert(key, id);
359 Some(id)
360 }
361}
362
363impl PromPlannerContext {
364 fn from_eval_stmt(stmt: &EvalStmt) -> Self {
365 Self {
366 start: stmt.start.duration_since(UNIX_EPOCH).unwrap().as_millis() as _,
367 end: stmt.end.duration_since(UNIX_EPOCH).unwrap().as_millis() as _,
368 interval: stmt.interval.as_millis() as _,
369 lookback_delta: stmt.lookback_delta.as_millis() as _,
370 ..Default::default()
371 }
372 }
373
374 fn reset(&mut self) {
376 self.table_name = None;
377 self.time_index_column = None;
378 self.field_columns = vec![];
379 self.tag_columns = vec![];
380 self.use_tsid = false;
381 self.field_column_matcher = None;
382 self.selector_matcher.clear();
383 self.schema_name = None;
384 self.range = None;
385 }
386
387 fn reset_table_name_and_schema(&mut self) {
389 self.table_name = Some(String::new());
390 self.schema_name = None;
391 self.use_tsid = false;
392 }
393
394 fn has_le_tag(&self) -> bool {
396 self.tag_columns.iter().any(|c| c.eq(&LE_COLUMN_NAME))
397 }
398}
399
400pub struct PromPlanner {
401 table_provider: DfTableSourceProvider,
402 ctx: PromPlannerContext,
403 promql_annotations: Option<PromqlAnnotationCollector>,
405}
406
407type BinaryFieldPair<'a> = (&'a String, &'a String);
408
409impl PromPlanner {
410 pub async fn stmt_to_plan(
411 table_provider: DfTableSourceProvider,
412 stmt: &EvalStmt,
413 query_engine_state: &QueryEngineState,
414 ) -> Result<LogicalPlan> {
415 Self::stmt_to_plan_with_annotations(table_provider, stmt, query_engine_state, None).await
416 }
417
418 pub async fn stmt_to_plan_with_annotations(
420 table_provider: DfTableSourceProvider,
421 stmt: &EvalStmt,
422 query_engine_state: &QueryEngineState,
423 promql_annotations: Option<PromqlAnnotationCollector>,
424 ) -> Result<LogicalPlan> {
425 let mut planner = Self {
426 table_provider,
427 ctx: PromPlannerContext::from_eval_stmt(stmt),
428 promql_annotations,
429 };
430
431 let plan = planner
432 .prom_expr_to_plan(&stmt.expr, query_engine_state)
433 .await?;
434
435 planner.strip_tsid_column(plan)
437 }
438
439 pub async fn prom_expr_to_plan(
440 &mut self,
441 prom_expr: &PromExpr,
442 query_engine_state: &QueryEngineState,
443 ) -> Result<LogicalPlan> {
444 self.prom_expr_to_plan_inner(prom_expr, false, query_engine_state)
445 .await
446 }
447
448 #[async_recursion]
458 async fn prom_expr_to_plan_inner(
459 &mut self,
460 prom_expr: &PromExpr,
461 timestamp_fn: bool,
462 query_engine_state: &QueryEngineState,
463 ) -> Result<LogicalPlan> {
464 let res = match prom_expr {
465 PromExpr::Aggregate(expr) => {
466 self.prom_aggr_expr_to_plan(query_engine_state, expr)
467 .await?
468 }
469 PromExpr::Unary(expr) => {
470 self.prom_unary_expr_to_plan(query_engine_state, expr)
471 .await?
472 }
473 PromExpr::Binary(expr) => {
474 self.prom_binary_expr_to_plan(query_engine_state, expr)
475 .await?
476 }
477 PromExpr::Paren(ParenExpr { expr }) => {
478 self.prom_expr_to_plan_inner(expr, timestamp_fn, query_engine_state)
479 .await?
480 }
481 PromExpr::Subquery(expr) => {
482 self.prom_subquery_expr_to_plan(query_engine_state, expr)
483 .await?
484 }
485 PromExpr::NumberLiteral(lit) => self.prom_number_lit_to_plan(lit)?,
486 PromExpr::StringLiteral(lit) => self.prom_string_lit_to_plan(lit)?,
487 PromExpr::VectorSelector(selector) => {
488 self.prom_vector_selector_to_plan(selector, timestamp_fn)
489 .await?
490 }
491 PromExpr::MatrixSelector(selector) => {
492 self.prom_matrix_selector_to_plan(selector).await?
493 }
494 PromExpr::Call(expr) => {
495 self.prom_call_expr_to_plan(query_engine_state, expr)
496 .await?
497 }
498 PromExpr::Extension(expr) => {
499 self.prom_ext_expr_to_plan(query_engine_state, expr).await?
500 }
501 };
502
503 Ok(res)
504 }
505
506 async fn prom_subquery_expr_to_plan(
507 &mut self,
508 query_engine_state: &QueryEngineState,
509 subquery_expr: &SubqueryExpr,
510 ) -> Result<LogicalPlan> {
511 let SubqueryExpr {
512 expr, range, step, ..
513 } = subquery_expr;
514
515 let current_interval = self.ctx.interval;
516 if let Some(step) = step {
517 self.ctx.interval = step.as_millis() as _;
518 }
519 let current_start = self.ctx.start;
520 self.ctx.start -= range.as_millis() as i64 - self.ctx.interval;
521 let input = self.prom_expr_to_plan(expr, query_engine_state).await?;
522 self.ctx.interval = current_interval;
523 self.ctx.start = current_start;
524
525 ensure!(!range.is_zero(), ZeroRangeSelectorSnafu);
526 let range_ms = range.as_millis() as _;
527 self.ctx.range = Some(range_ms);
528
529 let time_index_column =
530 self.ctx
531 .time_index_column
532 .clone()
533 .with_context(|| TimeIndexNotFoundSnafu {
534 table: self.ctx.table_name.clone().unwrap_or_default(),
535 })?;
536
537 let input_schema = input.schema();
543 let input_has_tsid = input_schema.fields().iter().any(|field| {
544 field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
545 && field.data_type() == &ArrowDataType::UInt64
546 });
547 let (series_key_columns, mut sort_exprs) = if input_has_tsid {
548 (
549 vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()],
550 vec![
551 DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME))
552 .sort(true, true),
553 ],
554 )
555 } else {
556 let key_columns: Vec<String> = self
559 .ctx
560 .tag_columns
561 .iter()
562 .filter(|name| input_schema.has_column_with_unqualified_name(name))
563 .cloned()
564 .collect();
565 let sort = key_columns
566 .iter()
567 .map(|name| DfExpr::Column(Column::from_name(name)).sort(true, true))
568 .collect::<Vec<_>>();
569 (key_columns, sort)
570 };
571 sort_exprs.push(DfExpr::Column(Column::from_name(&time_index_column)).sort(true, true));
572
573 let sort_plan = LogicalPlanBuilder::from(input)
574 .sort(sort_exprs)
575 .context(DataFusionPlanningSnafu)?
576 .build()
577 .context(DataFusionPlanningSnafu)?;
578 let divide_plan = LogicalPlan::Extension(Extension {
579 node: Arc::new(SeriesDivide::new(
580 series_key_columns,
581 time_index_column.clone(),
582 sort_plan,
583 )),
584 });
585
586 let manipulate = RangeManipulate::new(
587 self.ctx.start,
588 self.ctx.end,
589 self.ctx.interval,
590 range_ms,
591 time_index_column,
592 self.ctx.field_columns.clone(),
593 divide_plan,
594 )
595 .context(DataFusionPlanningSnafu)?;
596
597 Ok(LogicalPlan::Extension(Extension {
598 node: Arc::new(manipulate),
599 }))
600 }
601
602 async fn prom_aggr_expr_to_plan(
603 &mut self,
604 query_engine_state: &QueryEngineState,
605 aggr_expr: &AggregateExpr,
606 ) -> Result<LogicalPlan> {
607 let AggregateExpr {
608 op,
609 expr,
610 modifier,
611 param,
612 } = aggr_expr;
613
614 let mut input = self.prom_expr_to_plan(expr, query_engine_state).await?;
615 let input_has_tsid = input.schema().fields().iter().any(|field| {
616 field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
617 && field.data_type() == &ArrowDataType::UInt64
618 });
619
620 let required_group_tags = match modifier {
623 None => BTreeSet::new(),
624 Some(LabelModifier::Include(labels)) => labels
625 .labels
626 .iter()
627 .filter(|label| !is_metric_engine_internal_column(label.as_str()))
628 .cloned()
629 .collect(),
630 Some(LabelModifier::Exclude(labels)) => {
631 let mut all_tags = self.collect_row_key_tag_columns_from_plan(&input)?;
632 for label in &labels.labels {
633 let _ = all_tags.remove(label);
634 }
635 all_tags
636 }
637 };
638
639 if !required_group_tags.is_empty()
640 && required_group_tags
641 .iter()
642 .any(|tag| Self::find_case_sensitive_column(input.schema(), tag.as_str()).is_none())
643 {
644 input = self.ensure_tag_columns_available(input, &required_group_tags)?;
645 self.refresh_tag_columns_from_schema(input.schema());
646 }
647
648 match (*op).id() {
649 token::T_TOPK | token::T_BOTTOMK => {
650 self.prom_topk_bottomk_to_plan(aggr_expr, input).await
651 }
652 _ => {
653 let input_tag_columns = if input_has_tsid {
657 self.collect_row_key_tag_columns_from_plan(&input)?
658 .into_iter()
659 .collect::<Vec<_>>()
660 } else {
661 self.ctx.tag_columns.clone()
662 };
663 let mut group_exprs = self.agg_modifier_to_col(input.schema(), modifier, true)?;
666 let mixed_sample_columns =
667 Self::alternative_sample_columns(input.schema(), &self.ctx.field_columns)
668 .map(|(float, histogram)| (float.to_string(), histogram.to_string()));
669 let preserve_any_value = mixed_sample_columns.is_some();
674 let has_native_histogram = preserve_any_value
675 || self.all_field_columns_are_native_histograms(input.schema());
676 let (mut aggr_exprs, prev_field_exprs) =
678 self.create_aggregate_exprs(*op, param, &input)?;
679 let prev_field_exprs =
680 normalize_cols(prev_field_exprs, &input).context(DataFusionPlanningSnafu)?;
681
682 let keep_tsid = op.id() != token::T_COUNT_VALUES
683 && input_has_tsid
684 && input_tag_columns.iter().collect::<HashSet<_>>()
685 == self.ctx.tag_columns.iter().collect::<HashSet<_>>();
686
687 if keep_tsid {
688 aggr_exprs.push(
689 first_value(
690 DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME)),
691 vec![],
692 )
693 .alias(DATA_SCHEMA_TSID_COLUMN_NAME),
694 );
695 }
696 self.ctx.use_tsid = keep_tsid;
697
698 let builder = LogicalPlanBuilder::from(input);
700 let builder = if op.id() == token::T_COUNT_VALUES {
701 let label = Self::get_param_value_as_str(*op, param)?;
702 let count_value_exprs = prev_field_exprs.iter().map(|expr| {
705 match expr {
706 DfExpr::Column(column) => DfExpr::Column(column.clone()),
707 _ => DfExpr::Column(Column::from_name(expr.schema_name().to_string())),
708 }
709 .alias(label)
710 });
711 let aggregate_group_exprs = group_exprs
712 .iter()
713 .cloned()
714 .chain(prev_field_exprs.clone())
715 .collect::<Vec<_>>();
716 group_exprs.push(col(label));
717 let project_fields = self
718 .create_field_column_exprs()?
719 .into_iter()
720 .chain(self.create_tag_column_exprs()?)
721 .chain(Some(self.create_time_index_column_expr()?))
722 .chain(count_value_exprs);
723
724 builder
725 .aggregate(aggregate_group_exprs, aggr_exprs)
726 .context(DataFusionPlanningSnafu)?
727 .project(project_fields)
728 .context(DataFusionPlanningSnafu)?
729 } else {
730 builder
731 .aggregate(group_exprs.clone(), aggr_exprs)
732 .context(DataFusionPlanningSnafu)?
733 };
734
735 let builder = if let Some((float, histogram)) = mixed_sample_columns {
736 let builder = match op.id() {
737 token::T_SUM | token::T_AVG => builder
738 .filter(self.mixed_aggregate_filter_expr(*op, &float, &histogram)?)
739 .context(DataFusionPlanningSnafu)?,
740 token::T_MIN
741 | token::T_MAX
742 | token::T_STDDEV
743 | token::T_STDVAR
744 | token::T_QUANTILE => builder
745 .filter(self.mixed_ignored_histogram_filter_expr(*op, &histogram)?)
746 .context(DataFusionPlanningSnafu)?,
747 _ => builder,
748 };
749
750 match op.id() {
751 token::T_SUM
752 | token::T_AVG
753 | token::T_MIN
754 | token::T_MAX
755 | token::T_STDDEV
756 | token::T_STDVAR
757 | token::T_QUANTILE => {
758 let project_fields = self
759 .create_field_column_exprs()?
760 .into_iter()
761 .chain(self.create_tag_column_exprs()?)
762 .chain(self.ctx.use_tsid.then_some(DfExpr::Column(
763 Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME),
764 )))
765 .chain(Some(self.create_time_index_column_expr()?));
766 builder
767 .project(project_fields)
768 .context(DataFusionPlanningSnafu)?
769 }
770 _ => builder,
771 }
772 } else {
773 builder
774 };
775
776 let builder = if has_native_histogram {
781 builder
782 .filter(self.create_empty_values_filter_expr(preserve_any_value)?)
783 .context(DataFusionPlanningSnafu)?
784 } else {
785 builder
786 };
787
788 let sort_expr = group_exprs.into_iter().map(|expr| expr.sort(true, false));
789
790 builder
791 .sort(sort_expr)
792 .context(DataFusionPlanningSnafu)?
793 .build()
794 .context(DataFusionPlanningSnafu)
795 }
796 }
797 }
798
799 async fn prom_topk_bottomk_to_plan(
801 &mut self,
802 aggr_expr: &AggregateExpr,
803 input: LogicalPlan,
804 ) -> Result<LogicalPlan> {
805 let AggregateExpr {
806 op,
807 param,
808 modifier,
809 ..
810 } = aggr_expr;
811
812 let input_has_tsid = input.schema().fields().iter().any(|field| {
813 field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
814 && field.data_type() == &ArrowDataType::UInt64
815 });
816 self.ctx.use_tsid = input_has_tsid;
817
818 let group_exprs = self.agg_modifier_to_col(input.schema(), modifier, false)?;
819
820 let mut input = input;
821 if let Some((float_column, histogram_column)) =
822 Self::alternative_sample_columns(input.schema(), &self.ctx.field_columns)
823 .map(|(float, histogram)| (float.to_string(), histogram.to_string()))
824 {
825 let drop_histogram = DfExpr::ScalarFunction(ScalarFunction {
826 func: Arc::new(NativeHistogramDrop::bool_false_udf(
827 format!(
828 "{}: dropped native histogram samples because this aggregation is not supported for native histograms",
829 op
830 ),
831 self.promql_annotations.clone(),
832 )),
833 args: vec![col(&histogram_column)],
834 });
835 let keep_float = when(col(&histogram_column).is_not_null(), drop_histogram)
836 .otherwise(col(&float_column).is_not_null())
837 .context(DataFusionPlanningSnafu)?;
838 input = LogicalPlanBuilder::from(input)
839 .filter(keep_float)
840 .context(DataFusionPlanningSnafu)?
841 .build()
842 .context(DataFusionPlanningSnafu)?;
843 self.ctx.field_columns = vec![float_column];
844 }
845
846 if self.all_field_columns_are_native_histograms(input.schema()) {
847 let promql_annotations = self.promql_annotations.clone();
848 let input = self.projection_for_each_field_column(input, |col| {
849 Ok(DfExpr::ScalarFunction(ScalarFunction {
850 func: Arc::new(NativeHistogramDrop::float_null_udf(
851 format!(
852 "{}: dropped native histogram samples because this aggregation is not supported for native histograms",
853 op
854 ),
855 promql_annotations.clone(),
856 )),
857 args: vec![DfExpr::Column(Column::from_name(col))],
858 }))
859 })?;
860 return LogicalPlanBuilder::from(input)
861 .filter(self.create_empty_values_filter_expr(false)?)
862 .context(DataFusionPlanningSnafu)?
863 .build()
864 .context(DataFusionPlanningSnafu);
865 }
866
867 let val = Self::get_param_as_literal_expr(
868 param.as_deref(),
869 Some(*op),
870 Some(ArrowDataType::Float64),
871 )?;
872
873 let window_exprs = self.create_window_exprs(*op, group_exprs.clone(), &input)?;
875
876 let rank_columns: Vec<_> = window_exprs
877 .iter()
878 .map(|expr| expr.schema_name().to_string())
879 .collect();
880
881 let filter: DfExpr = rank_columns
884 .iter()
885 .fold(None, |expr, rank| {
886 let predicate = DfExpr::BinaryExpr(BinaryExpr {
887 left: Box::new(col(rank)),
888 op: Operator::LtEq,
889 right: Box::new(val.clone()),
890 });
891
892 match expr {
893 None => Some(predicate),
894 Some(expr) => Some(DfExpr::BinaryExpr(BinaryExpr {
895 left: Box::new(expr),
896 op: Operator::Or,
897 right: Box::new(predicate),
898 })),
899 }
900 })
901 .unwrap();
902
903 let rank_columns: Vec<_> = rank_columns.into_iter().map(col).collect();
904
905 let mut new_group_exprs = group_exprs.clone();
906 new_group_exprs.extend(rank_columns);
908
909 let group_sort_expr = new_group_exprs
910 .into_iter()
911 .map(|expr| expr.sort(true, false));
912
913 let project_fields = self
914 .create_field_column_exprs()?
915 .into_iter()
916 .chain(self.create_tag_column_exprs()?)
917 .chain(
918 self.ctx
919 .use_tsid
920 .then_some(DfExpr::Column(Column::from_name(
921 DATA_SCHEMA_TSID_COLUMN_NAME,
922 ))),
923 )
924 .chain(Some(self.create_time_index_column_expr()?));
925
926 LogicalPlanBuilder::from(input)
927 .window(window_exprs)
928 .context(DataFusionPlanningSnafu)?
929 .filter(filter)
930 .context(DataFusionPlanningSnafu)?
931 .sort(group_sort_expr)
932 .context(DataFusionPlanningSnafu)?
933 .project(project_fields)
934 .context(DataFusionPlanningSnafu)?
935 .build()
936 .context(DataFusionPlanningSnafu)
937 }
938
939 async fn prom_unary_expr_to_plan(
940 &mut self,
941 query_engine_state: &QueryEngineState,
942 unary_expr: &UnaryExpr,
943 ) -> Result<LogicalPlan> {
944 let UnaryExpr { expr } = unary_expr;
945 let input = self.prom_expr_to_plan(expr, query_engine_state).await?;
947 self.negate_field_columns(input)
948 }
949
950 fn negate_field_columns(&mut self, input: LogicalPlan) -> Result<LogicalPlan> {
951 let input_schema = input.schema().clone();
952 self.projection_for_each_field_column(input, |col| {
953 if Self::field_column_is_native_histogram(&input_schema, col) {
954 Ok(DfExpr::ScalarFunction(ScalarFunction {
955 func: Arc::new(NativeHistogramNeg::scalar_udf()),
956 args: vec![DfExpr::Column(col.into())],
957 }))
958 } else {
959 Ok(DfExpr::Negative(Box::new(DfExpr::Column(col.into()))))
960 }
961 })
962 }
963
964 async fn try_plan_binary_island(
965 &mut self,
966 binary_expr: &PromBinaryExpr,
967 ) -> Result<Option<LogicalPlan>> {
968 let original_ctx = self.ctx.clone();
969 let mut collect_env = IslandCollectEnv::default();
970 let Some(island_expr) =
971 IslandExpr::try_new(&PromExpr::Binary(binary_expr.clone()), &mut collect_env)
972 else {
973 return Ok(None);
974 };
975
976 if collect_env.leaves.is_empty()
977 || collect_env.vector_occurrences <= collect_env.leaves.len()
978 {
979 return Ok(None);
980 }
981
982 let mut planned_leaves = Vec::with_capacity(collect_env.leaves.len());
983 for (idx, leaf) in collect_env.leaves.iter().enumerate() {
984 let plan = self
985 .prom_vector_selector_to_plan(&leaf.selector, false)
986 .await?;
987 let ctx = self.ctx.clone();
988 let alias = TableReference::bare(format!("{BINARY_ISLAND_LEAF_ALIAS_PREFIX}{idx}"));
989 let plan = LogicalPlanBuilder::from(plan)
990 .alias(alias.clone())
991 .context(DataFusionPlanningSnafu)?
992 .build()
993 .context(DataFusionPlanningSnafu)?;
994 planned_leaves.push(PlannedIslandLeaf {
995 plan,
996 ctx,
997 alias,
998 display_table: leaf.display_table.clone(),
999 });
1000 }
1001
1002 if planned_leaves.iter().any(|leaf| {
1003 Self::field_columns_contain_native_histogram(
1004 leaf.plan.schema(),
1005 &leaf.ctx.field_columns,
1006 )
1007 }) {
1008 self.ctx = original_ctx;
1009 return Ok(None);
1010 }
1011
1012 if !Self::binary_island_join_contexts_supported(&planned_leaves) {
1013 self.ctx = original_ctx;
1014 return Ok(None);
1015 }
1016
1017 let mut input = planned_leaves[0].plan.clone();
1018 for right_idx in 1..planned_leaves.len() {
1019 input = self.join_binary_island_leaf(
1020 input,
1021 &planned_leaves[0],
1022 &planned_leaves[right_idx],
1023 )?;
1024 }
1025
1026 let field_exprs =
1027 Self::build_binary_island_field_exprs(&island_expr, &planned_leaves, input.schema())?;
1028 if field_exprs.scalar || field_exprs.exprs.is_empty() {
1029 self.ctx = original_ctx;
1030 return Ok(None);
1031 }
1032
1033 let plan = self.project_binary_island(
1034 input,
1035 &planned_leaves[0].alias,
1036 &planned_leaves[0].ctx,
1037 field_exprs,
1038 )?;
1039 Ok(Some(plan))
1040 }
1041
1042 fn binary_island_join_contexts_supported(leaves: &[PlannedIslandLeaf]) -> bool {
1043 if leaves
1044 .iter()
1045 .any(|leaf| leaf.ctx.time_index_column.is_none())
1046 {
1047 return false;
1048 }
1049
1050 if leaves.len() <= 1 {
1051 return true;
1052 }
1053
1054 let first_tags = leaves[0].ctx.tag_columns.iter().collect::<BTreeSet<_>>();
1055
1056 leaves.iter().skip(1).all(|leaf| {
1057 (Self::plan_has_tsid_column(&leaves[0].plan) && Self::plan_has_tsid_column(&leaf.plan))
1058 || leaf.ctx.tag_columns.iter().collect::<BTreeSet<_>>() == first_tags
1059 })
1060 }
1061
1062 fn join_binary_island_leaf(
1063 &self,
1064 left: LogicalPlan,
1065 first_leaf: &PlannedIslandLeaf,
1066 right_leaf: &PlannedIslandLeaf,
1067 ) -> Result<LogicalPlan> {
1068 let only_join_time_index = (first_leaf.ctx.tag_columns.is_empty()
1069 || right_leaf.ctx.tag_columns.is_empty())
1070 && !first_leaf
1071 .ctx
1072 .tag_columns
1073 .iter()
1074 .chain(&right_leaf.ctx.tag_columns)
1075 .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL);
1076 let (mut left_keys, mut right_keys, force_empty_join) = self.binary_join_key_columns(
1077 left.schema(),
1078 right_leaf.plan.schema(),
1079 &first_leaf.ctx,
1080 &right_leaf.ctx,
1081 only_join_time_index,
1082 &None,
1083 )?;
1084
1085 if let (Some(left_time_index_column), Some(right_time_index_column)) = (
1086 first_leaf.ctx.time_index_column.clone(),
1087 right_leaf.ctx.time_index_column.clone(),
1088 ) {
1089 left_keys.insert(left_time_index_column);
1090 right_keys.insert(right_time_index_column);
1091 }
1092
1093 LogicalPlanBuilder::from(left)
1094 .join_detailed(
1095 right_leaf.plan.clone(),
1096 JoinType::Inner,
1097 (
1098 left_keys
1099 .into_iter()
1100 .map(|name| Column::new(Some(first_leaf.alias.clone()), name))
1101 .collect::<Vec<_>>(),
1102 right_keys
1103 .into_iter()
1104 .map(|name| Column::new(Some(right_leaf.alias.clone()), name))
1105 .collect::<Vec<_>>(),
1106 ),
1107 force_empty_join.then_some(lit(false)),
1108 NullEquality::NullEqualsNull,
1109 )
1110 .context(DataFusionPlanningSnafu)?
1111 .build()
1112 .context(DataFusionPlanningSnafu)
1113 }
1114
1115 fn build_binary_island_field_exprs(
1116 expr: &IslandExpr,
1117 leaves: &[PlannedIslandLeaf],
1118 schema: &DFSchemaRef,
1119 ) -> Result<IslandFieldExprs> {
1120 match expr {
1121 IslandExpr::VectorLeaf(id) => {
1122 let leaf = &leaves[*id];
1123 let exprs = leaf
1124 .ctx
1125 .field_columns
1126 .iter()
1127 .map(|field| {
1128 schema
1129 .qualified_field_with_name(Some(&leaf.alias), field)
1130 .context(DataFusionPlanningSnafu)
1131 .map(|field| DfExpr::Column(field.into()))
1132 })
1133 .collect::<Result<Vec<_>>>()?;
1134 let names = leaf
1135 .ctx
1136 .field_columns
1137 .iter()
1138 .map(|field| format!("{}.{}", leaf.display_table, field))
1139 .collect();
1140 Ok(IslandFieldExprs {
1141 exprs,
1142 names,
1143 scalar: false,
1144 })
1145 }
1146 IslandExpr::Scalar(expr) => Ok(IslandFieldExprs {
1147 exprs: vec![expr.clone()],
1148 names: vec![expr.schema_name().to_string()],
1149 scalar: true,
1150 }),
1151 IslandExpr::Unary { input } => {
1152 let input = Self::build_binary_island_field_exprs(input, leaves, schema)?;
1153 let mut exprs = Vec::with_capacity(input.exprs.len());
1154 let mut names = Vec::with_capacity(input.names.len());
1155 for (expr, name) in input.exprs.into_iter().zip(input.names) {
1156 exprs.push(DfExpr::Negative(Box::new(expr)));
1157 names.push(format!("-{name}"));
1158 }
1159 Ok(IslandFieldExprs {
1160 exprs,
1161 names,
1162 scalar: input.scalar,
1163 })
1164 }
1165 IslandExpr::Binary { op, lhs, rhs } => {
1166 let same_leaf = match (&**lhs, &**rhs) {
1167 (IslandExpr::VectorLeaf(left), IslandExpr::VectorLeaf(right))
1168 if left == right =>
1169 {
1170 Some(*left)
1171 }
1172 _ => None,
1173 };
1174 let lhs = Self::build_binary_island_field_exprs(lhs, leaves, schema)?;
1175 let rhs = Self::build_binary_island_field_exprs(rhs, leaves, schema)?;
1176 let expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1177 let scalar = lhs.scalar && rhs.scalar;
1178 let op = op.to_string();
1179
1180 let (exprs, names) = match (lhs.scalar, rhs.scalar) {
1181 (true, true) => {
1182 let expr = expr_builder(lhs.exprs[0].clone(), rhs.exprs[0].clone())?;
1183 let name = format!("{} {op} {}", lhs.names[0], rhs.names[0]);
1184 (vec![expr], vec![name])
1185 }
1186 (true, false) => {
1187 let mut exprs = Vec::with_capacity(rhs.exprs.len());
1188 let mut names = Vec::with_capacity(rhs.names.len());
1189 for (rhs_expr, rhs_name) in rhs.exprs.into_iter().zip(rhs.names) {
1190 exprs.push(expr_builder(lhs.exprs[0].clone(), rhs_expr)?);
1191 names.push(format!("{} {op} {rhs_name}", lhs.names[0]));
1192 }
1193 (exprs, names)
1194 }
1195 (false, true) => {
1196 let mut exprs = Vec::with_capacity(lhs.exprs.len());
1197 let mut names = Vec::with_capacity(lhs.names.len());
1198 for (lhs_expr, lhs_name) in lhs.exprs.into_iter().zip(lhs.names) {
1199 exprs.push(expr_builder(lhs_expr, rhs.exprs[0].clone())?);
1200 names.push(format!("{lhs_name} {op} {}", rhs.names[0]));
1201 }
1202 (exprs, names)
1203 }
1204 (false, false) => {
1205 let mut exprs = Vec::new();
1206 let mut names = Vec::new();
1207 for (idx, ((lhs_expr, rhs_expr), (mut lhs_name, mut rhs_name))) in lhs
1208 .exprs
1209 .into_iter()
1210 .zip(rhs.exprs)
1211 .zip(lhs.names.into_iter().zip(rhs.names))
1212 .enumerate()
1213 {
1214 if let Some(leaf) = same_leaf {
1215 let field = leaves[leaf]
1216 .ctx
1217 .field_columns
1218 .get(idx)
1219 .cloned()
1220 .unwrap_or_else(|| lhs_name.clone());
1221 lhs_name = format!("lhs.{field}");
1222 rhs_name = format!("rhs.{field}");
1223 }
1224 exprs.push(expr_builder(lhs_expr, rhs_expr)?);
1225 names.push(format!("{lhs_name} {op} {rhs_name}"));
1226 }
1227 (exprs, names)
1228 }
1229 };
1230
1231 Ok(IslandFieldExprs {
1232 exprs,
1233 names,
1234 scalar,
1235 })
1236 }
1237 }
1238 }
1239
1240 fn project_binary_island(
1241 &mut self,
1242 input: LogicalPlan,
1243 base_alias: &TableReference,
1244 base_ctx: &PromPlannerContext,
1245 field_exprs: IslandFieldExprs,
1246 ) -> Result<LogicalPlan> {
1247 self.ctx = base_ctx.clone();
1248
1249 let schema = input.schema();
1250 let non_field_exprs = base_ctx
1251 .tag_columns
1252 .iter()
1253 .chain(base_ctx.time_index_column.iter())
1254 .map(|column| {
1255 schema
1256 .qualified_field_with_name(Some(base_alias), column)
1257 .context(DataFusionPlanningSnafu)
1258 .map(|field| DfExpr::Column(field.into()))
1259 });
1260 let tsid_expr = Self::optional_tsid_projection(schema, Some(base_alias), base_ctx.use_tsid)
1261 .into_iter()
1262 .map(Ok);
1263
1264 self.ctx.field_columns = field_exprs.names;
1265 let field_exprs = field_exprs
1266 .exprs
1267 .into_iter()
1268 .zip(self.ctx.field_columns.iter())
1269 .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::<String>, name))));
1270
1271 let project_exprs = non_field_exprs
1272 .chain(tsid_expr)
1273 .chain(field_exprs)
1274 .collect::<Result<Vec<_>>>()?;
1275
1276 let plan = LogicalPlanBuilder::from(input)
1277 .project(project_exprs)
1278 .context(DataFusionPlanningSnafu)?
1279 .build()
1280 .context(DataFusionPlanningSnafu)?;
1281
1282 self.ctx.table_name = None;
1283 self.ctx.schema_name = None;
1284
1285 Ok(plan)
1286 }
1287
1288 async fn prom_binary_expr_to_plan(
1289 &mut self,
1290 query_engine_state: &QueryEngineState,
1291 binary_expr: &PromBinaryExpr,
1292 ) -> Result<LogicalPlan> {
1293 if let Some(modifier) = &binary_expr.modifier {
1297 ensure!(
1298 modifier.fill_values.lhs.is_none() && modifier.fill_values.rhs.is_none(),
1299 UnsupportedExprSnafu {
1300 name: "PromQL fill modifiers"
1301 }
1302 );
1303 }
1304
1305 if let Some(plan) = self.try_plan_binary_island(binary_expr).await? {
1306 return Ok(plan);
1307 }
1308
1309 let PromBinaryExpr {
1310 lhs,
1311 rhs,
1312 op,
1313 modifier,
1314 } = binary_expr;
1315
1316 let should_return_bool = if let Some(m) = modifier {
1319 m.return_bool
1320 } else {
1321 false
1322 };
1323 let is_comparison_op = Self::is_token_a_comparison_op(*op);
1324
1325 match (
1328 Self::try_build_literal_expr(lhs),
1329 Self::try_build_literal_expr(rhs),
1330 ) {
1331 (Some(lhs), Some(rhs)) => {
1332 self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
1333 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1334 self.ctx.reset_table_name_and_schema();
1335 let field_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1336 let mut field_expr = field_expr_builder(lhs, rhs)?;
1337
1338 if is_comparison_op && should_return_bool {
1339 field_expr = DfExpr::Cast(Cast {
1340 expr: Box::new(field_expr),
1341 data_type: ArrowDataType::Float64,
1342 });
1343 }
1344
1345 Ok(LogicalPlan::Extension(Extension {
1346 node: Arc::new(
1347 EmptyMetric::new(
1348 self.ctx.start,
1349 self.ctx.end,
1350 self.ctx.interval,
1351 SPECIAL_TIME_FUNCTION.to_string(),
1352 DEFAULT_FIELD_COLUMN.to_string(),
1353 Some(field_expr),
1354 )
1355 .context(DataFusionPlanningSnafu)?,
1356 ),
1357 }))
1358 }
1359 (Some(mut expr), None) => {
1361 let input = self.prom_expr_to_plan(rhs, query_engine_state).await?;
1362 if let Some(time_expr) = self.try_build_special_time_expr_with_context(lhs) {
1364 expr = time_expr
1365 }
1366 let input_schema = input.schema().clone();
1367 let preserve_any_value = Self::field_columns_are_alternative_samples(
1368 &input_schema,
1369 &self.ctx.field_columns,
1370 );
1371 let has_native_histogram = Self::field_columns_contain_native_histogram(
1372 &input_schema,
1373 &self.ctx.field_columns,
1374 );
1375 let retain_field_columns = self
1376 .ctx
1377 .field_columns
1378 .iter()
1379 .map(|col| {
1380 Self::binary_result_is_histogram(
1381 *op,
1382 false,
1383 Self::field_column_is_native_histogram(&input_schema, col),
1384 )
1385 .is_some()
1386 })
1387 .collect();
1388 let promql_annotations = self.promql_annotations.clone();
1389 let bin_expr_builder = |col: &String| {
1390 let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1391 let rhs_is_histogram =
1392 Self::field_column_is_native_histogram(&input_schema, col);
1393 let rhs = DfExpr::Column(col.into());
1394 let mut binary_expr = match Self::native_histogram_binary_expr(
1395 *op,
1396 expr.clone(),
1397 false,
1398 rhs.clone(),
1399 rhs_is_histogram,
1400 is_comparison_op && !should_return_bool,
1401 promql_annotations.clone(),
1402 )? {
1403 Some(expr) => expr,
1404 None => binary_expr_builder(expr.clone(), rhs)?,
1405 };
1406
1407 if is_comparison_op && should_return_bool {
1408 binary_expr = DfExpr::Cast(Cast {
1409 expr: Box::new(binary_expr),
1410 data_type: ArrowDataType::Float64,
1411 });
1412 }
1413 Ok(binary_expr)
1414 };
1415 if is_comparison_op && !should_return_bool {
1416 self.filter_on_field_column(input, bin_expr_builder)
1417 } else {
1418 let projected =
1419 self.projection_for_each_field_column(input, bin_expr_builder)?;
1420 self.filter_binary_projection(
1421 projected,
1422 has_native_histogram,
1423 preserve_any_value,
1424 retain_field_columns,
1425 )
1426 }
1427 }
1428 (None, Some(mut expr)) => {
1430 let input = self.prom_expr_to_plan(lhs, query_engine_state).await?;
1431 if let Some(time_expr) = self.try_build_special_time_expr_with_context(rhs) {
1433 expr = time_expr
1434 }
1435 let input_schema = input.schema().clone();
1436 let preserve_any_value = Self::field_columns_are_alternative_samples(
1437 &input_schema,
1438 &self.ctx.field_columns,
1439 );
1440 let has_native_histogram = Self::field_columns_contain_native_histogram(
1441 &input_schema,
1442 &self.ctx.field_columns,
1443 );
1444 let retain_field_columns = self
1445 .ctx
1446 .field_columns
1447 .iter()
1448 .map(|col| {
1449 Self::binary_result_is_histogram(
1450 *op,
1451 Self::field_column_is_native_histogram(&input_schema, col),
1452 false,
1453 )
1454 .is_some()
1455 })
1456 .collect();
1457 let promql_annotations = self.promql_annotations.clone();
1458 let bin_expr_builder = |col: &String| {
1459 let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1460 let lhs_is_histogram =
1461 Self::field_column_is_native_histogram(&input_schema, col);
1462 let lhs = DfExpr::Column(col.into());
1463 let mut binary_expr = match Self::native_histogram_binary_expr(
1464 *op,
1465 lhs.clone(),
1466 lhs_is_histogram,
1467 expr.clone(),
1468 false,
1469 is_comparison_op && !should_return_bool,
1470 promql_annotations.clone(),
1471 )? {
1472 Some(expr) => expr,
1473 None => binary_expr_builder(lhs, expr.clone())?,
1474 };
1475
1476 if is_comparison_op && should_return_bool {
1477 binary_expr = DfExpr::Cast(Cast {
1478 expr: Box::new(binary_expr),
1479 data_type: ArrowDataType::Float64,
1480 });
1481 }
1482 Ok(binary_expr)
1483 };
1484 if is_comparison_op && !should_return_bool {
1485 self.filter_on_field_column(input, bin_expr_builder)
1486 } else {
1487 let projected =
1488 self.projection_for_each_field_column(input, bin_expr_builder)?;
1489 self.filter_binary_projection(
1490 projected,
1491 has_native_histogram,
1492 preserve_any_value,
1493 retain_field_columns,
1494 )
1495 }
1496 }
1497 (None, None) => {
1499 let left_input = self.prom_expr_to_plan(lhs, query_engine_state).await?;
1500 let left_field_columns = self.ctx.field_columns.clone();
1501 let left_time_index_column = self.ctx.time_index_column.clone();
1502 let mut left_table_ref = self
1503 .table_ref()
1504 .unwrap_or_else(|_| TableReference::bare(""));
1505 let left_context = self.ctx.clone();
1506
1507 let right_input = self.prom_expr_to_plan(rhs, query_engine_state).await?;
1508 let right_field_columns = self.ctx.field_columns.clone();
1509 let right_time_index_column = self.ctx.time_index_column.clone();
1510 let mut right_table_ref = self
1511 .table_ref()
1512 .unwrap_or_else(|_| TableReference::bare(""));
1513 let right_context = self.ctx.clone();
1514 let left_is_empty_metric = Self::is_empty_metric(&left_input);
1515 let right_is_empty_metric = Self::is_empty_metric(&right_input);
1516
1517 if Self::is_token_a_set_op(*op) {
1521 return self.set_op_on_non_field_columns(
1522 left_input,
1523 right_input,
1524 left_context,
1525 right_context,
1526 *op,
1527 modifier,
1528 );
1529 }
1530
1531 let has_native_histogram = Self::field_columns_contain_native_histogram(
1532 left_input.schema(),
1533 &left_field_columns,
1534 ) || Self::field_columns_contain_native_histogram(
1535 right_input.schema(),
1536 &right_field_columns,
1537 );
1538
1539 if left_table_ref == right_table_ref {
1541 left_table_ref = TableReference::bare("lhs");
1543 right_table_ref = TableReference::bare("rhs");
1544 if self.ctx.tag_columns.is_empty() {
1550 self.ctx = left_context.clone();
1551 self.ctx.table_name = Some("lhs".to_string());
1552 } else {
1553 self.ctx.table_name = Some("rhs".to_string());
1554 }
1555 } else if right_is_empty_metric && !left_is_empty_metric {
1556 self.ctx = left_context.clone();
1557 }
1558 let broadcast_scalar = !is_comparison_op;
1561 let (field_groups, invalid_field_pairs) = Self::align_binary_field_columns(
1562 left_input.schema(),
1563 right_input.schema(),
1564 &left_field_columns,
1565 &right_field_columns,
1566 *op,
1567 broadcast_scalar && lhs.value_type() == ValueType::Scalar,
1568 broadcast_scalar && rhs.value_type() == ValueType::Scalar,
1569 );
1570 let left_aligned_field_columns = field_groups
1571 .iter()
1572 .flat_map(|(_, pairs)| {
1573 pairs
1574 .iter()
1575 .map(|(left_col_name, _)| (*left_col_name).clone())
1576 })
1577 .collect::<Vec<_>>();
1578 let right_aligned_field_columns = field_groups
1579 .iter()
1580 .flat_map(|(_, pairs)| {
1581 pairs
1582 .iter()
1583 .map(|(_, right_col_name)| (*right_col_name).clone())
1584 })
1585 .collect::<Vec<_>>();
1586 self.ctx.field_columns = field_groups
1589 .iter()
1590 .map(|(output, _)| output.clone())
1591 .collect();
1592 let mut field_groups = field_groups.into_iter();
1593 let has_empty_metric_operand = left_is_empty_metric || right_is_empty_metric;
1595
1596 let join_plan = self.join_on_non_field_columns(
1597 left_input,
1598 right_input,
1599 left_table_ref.clone(),
1600 right_table_ref.clone(),
1601 left_time_index_column,
1602 right_time_index_column,
1603 lhs.value_type() == ValueType::Scalar
1604 || rhs.value_type() == ValueType::Scalar
1605 || has_empty_metric_operand
1606 || ((left_context.tag_columns.is_empty()
1607 || right_context.tag_columns.is_empty())
1608 && !left_context
1609 .tag_columns
1610 .iter()
1611 .chain(&right_context.tag_columns)
1612 .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL)),
1613 modifier,
1614 &left_context,
1615 &right_context,
1616 )?;
1617 let join_plan_schema = join_plan.schema().clone();
1618 let promql_annotations = self.promql_annotations.clone();
1619 let invalid_pair_predicates = invalid_field_pairs
1622 .into_iter()
1623 .filter(|_| promql_annotations.is_some())
1624 .map(|(left_col_name, right_col_name)| {
1625 let left_field = join_plan_schema
1626 .qualified_field_with_name(Some(&left_table_ref), left_col_name)
1627 .context(DataFusionPlanningSnafu)?;
1628 let right_field = join_plan_schema
1629 .qualified_field_with_name(Some(&right_table_ref), right_col_name)
1630 .context(DataFusionPlanningSnafu)?;
1631 let left_is_histogram =
1632 left_field.1.data_type() == &Self::native_histogram_arrow_type();
1633 let right_is_histogram =
1634 right_field.1.data_type() == &Self::native_histogram_arrow_type();
1635 let drop_expr = Self::native_histogram_binary_expr(
1636 *op,
1637 DfExpr::Column(left_field.into()),
1638 left_is_histogram,
1639 DfExpr::Column(right_field.into()),
1640 right_is_histogram,
1641 true,
1642 promql_annotations.clone(),
1643 )?
1644 .with_context(|| UnexpectedPlanExprSnafu {
1645 desc: "invalid native histogram pair produced no drop expression",
1646 })?;
1647 Ok(DfExpr::Not(Box::new(drop_expr)))
1648 })
1649 .collect::<Result<Vec<_>>>()?;
1650 let join_plan = if let Some(predicate) = conjunction(invalid_pair_predicates) {
1651 LogicalPlanBuilder::from(join_plan)
1652 .filter(predicate)
1653 .context(DataFusionPlanningSnafu)?
1654 .build()
1655 .context(DataFusionPlanningSnafu)?
1656 } else {
1657 join_plan
1658 };
1659
1660 let bin_expr_builder = |_: &String| {
1661 let (_, field_pairs) =
1662 field_groups
1663 .next()
1664 .with_context(|| UnexpectedPlanExprSnafu {
1665 desc: "missing binary field group",
1666 })?;
1667 let binary_exprs = field_pairs
1668 .into_iter()
1669 .map(|(left_col_name, right_col_name)| {
1670 let left_field = join_plan_schema
1671 .qualified_field_with_name(Some(&left_table_ref), left_col_name)
1672 .context(DataFusionPlanningSnafu)?;
1673 let right_field = join_plan_schema
1674 .qualified_field_with_name(Some(&right_table_ref), right_col_name)
1675 .context(DataFusionPlanningSnafu)?;
1676 let left_is_histogram =
1677 left_field.1.data_type() == &Self::native_histogram_arrow_type();
1678 let right_is_histogram =
1679 right_field.1.data_type() == &Self::native_histogram_arrow_type();
1680 let left_col = left_field.into();
1681 let right_col = right_field.into();
1682
1683 let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1684 let lhs = DfExpr::Column(left_col);
1685 let rhs = DfExpr::Column(right_col);
1686 let mut binary_expr = match Self::native_histogram_binary_expr(
1687 *op,
1688 lhs.clone(),
1689 left_is_histogram,
1690 rhs.clone(),
1691 right_is_histogram,
1692 is_comparison_op && !should_return_bool,
1693 promql_annotations.clone(),
1694 )? {
1695 Some(expr) => expr,
1696 None => binary_expr_builder(lhs, rhs)?,
1697 };
1698 if is_comparison_op && should_return_bool {
1699 binary_expr = DfExpr::Cast(Cast {
1700 expr: Box::new(binary_expr),
1701 data_type: ArrowDataType::Float64,
1702 });
1703 }
1704 Ok(binary_expr)
1705 })
1706 .collect::<Result<Vec<_>>>()?;
1707 if let [binary_expr] = binary_exprs.as_slice() {
1708 Ok(binary_expr.clone())
1709 } else {
1710 Ok(DfExpr::ScalarFunction(ScalarFunction {
1711 func: coalesce(),
1712 args: binary_exprs,
1713 }))
1714 }
1715 };
1716 if is_comparison_op && !should_return_bool {
1717 let filtered = self.filter_on_field_column(join_plan, bin_expr_builder)?;
1724 let (project_table_ref, mut project_context, project_field_columns) =
1725 match (lhs.value_type(), rhs.value_type()) {
1726 (ValueType::Scalar, ValueType::Vector) => (
1727 &right_table_ref,
1728 right_context.clone(),
1729 right_aligned_field_columns,
1730 ),
1731 _ => (
1732 &left_table_ref,
1733 left_context.clone(),
1734 left_aligned_field_columns,
1735 ),
1736 };
1737 project_context.field_columns = project_field_columns;
1738 self.project_binary_join_side(filtered, project_table_ref, &project_context)
1739 } else {
1740 let projected =
1741 self.projection_for_each_field_column(join_plan, bin_expr_builder)?;
1742 let preserve_any_value = Self::field_columns_are_alternative_samples(
1743 projected.schema(),
1744 &self.ctx.field_columns,
1745 );
1746 let retain_field_columns = vec![true; self.ctx.field_columns.len()];
1747 self.filter_binary_projection(
1748 projected,
1749 has_native_histogram,
1750 preserve_any_value,
1751 retain_field_columns,
1752 )
1753 }
1754 }
1755 }
1756 }
1757
1758 fn filter_binary_projection(
1759 &mut self,
1760 input: LogicalPlan,
1761 has_native_histogram: bool,
1762 preserve_any_value: bool,
1763 retain_field_columns: Vec<bool>,
1764 ) -> Result<LogicalPlan> {
1765 if !has_native_histogram {
1766 return Ok(input);
1767 }
1768
1769 ensure!(
1770 retain_field_columns.len() == self.ctx.field_columns.len(),
1771 UnexpectedPlanExprSnafu {
1772 desc: "binary output field count changed unexpectedly",
1773 }
1774 );
1775
1776 let filtered = LogicalPlanBuilder::from(input)
1777 .filter(self.create_empty_values_filter_expr(preserve_any_value)?)
1778 .context(DataFusionPlanningSnafu)?
1779 .build()
1780 .context(DataFusionPlanningSnafu)?;
1781 if retain_field_columns.iter().all(|retain| *retain) {
1782 return Ok(filtered);
1783 }
1784
1785 let retained = self
1786 .ctx
1787 .field_columns
1788 .iter()
1789 .zip(retain_field_columns)
1790 .filter(|(_, retain)| *retain)
1791 .map(|(field, _)| field.clone())
1792 .collect::<Vec<_>>();
1793 if retained.is_empty() {
1794 return Ok(filtered);
1795 }
1796 self.ctx.field_columns = retained;
1797
1798 let mut output_columns = self
1799 .ctx
1800 .field_columns
1801 .iter()
1802 .chain(&self.ctx.tag_columns)
1803 .cloned()
1804 .collect::<HashSet<_>>();
1805 output_columns.extend(self.ctx.time_index_column.iter().cloned());
1806 if self.ctx.use_tsid {
1807 output_columns.insert(DATA_SCHEMA_TSID_COLUMN_NAME.to_string());
1808 }
1809 let project_exprs = filtered
1810 .schema()
1811 .iter()
1812 .filter(|(_, field)| output_columns.contains(field.name()))
1813 .map(|(qualifier, field)| {
1814 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
1815 })
1816 .collect::<Vec<_>>();
1817 LogicalPlanBuilder::from(filtered)
1818 .project(project_exprs)
1819 .context(DataFusionPlanningSnafu)?
1820 .build()
1821 .context(DataFusionPlanningSnafu)
1822 }
1823
1824 fn project_binary_join_side(
1825 &mut self,
1826 input: LogicalPlan,
1827 table_ref: &TableReference,
1828 context: &PromPlannerContext,
1829 ) -> Result<LogicalPlan> {
1830 let schema = input.schema();
1831
1832 let mut project_exprs =
1833 Vec::with_capacity(context.tag_columns.len() + context.field_columns.len() + 2);
1834
1835 if let Some(time_index_column) = &context.time_index_column {
1837 let time_index_col = schema
1838 .qualified_field_with_name(Some(table_ref), time_index_column)
1839 .context(DataFusionPlanningSnafu)?
1840 .into();
1841 project_exprs.push(DfExpr::Column(time_index_col));
1842 }
1843
1844 for field_column in &context.field_columns {
1846 let field_col = schema
1847 .qualified_field_with_name(Some(table_ref), field_column)
1848 .context(DataFusionPlanningSnafu)?
1849 .into();
1850 project_exprs.push(DfExpr::Column(field_col));
1851 }
1852
1853 for tag_column in &context.tag_columns {
1855 let tag_col = schema
1856 .qualified_field_with_name(Some(table_ref), tag_column)
1857 .context(DataFusionPlanningSnafu)?
1858 .into();
1859 project_exprs.push(DfExpr::Column(tag_col));
1860 }
1861
1862 if let Some(tsid_col) =
1865 Self::optional_tsid_projection(schema, Some(table_ref), context.use_tsid)
1866 {
1867 project_exprs.push(tsid_col);
1868 }
1869
1870 let plan = LogicalPlanBuilder::from(input)
1871 .project(project_exprs)
1872 .context(DataFusionPlanningSnafu)?
1873 .build()
1874 .context(DataFusionPlanningSnafu)?;
1875
1876 self.ctx = context.clone();
1879 self.ctx.table_name = None;
1880 self.ctx.schema_name = None;
1881
1882 Ok(plan)
1883 }
1884
1885 fn prom_number_lit_to_plan(&mut self, number_literal: &NumberLiteral) -> Result<LogicalPlan> {
1886 let NumberLiteral { val } = number_literal;
1887 self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
1888 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1889 self.ctx.reset_table_name_and_schema();
1890 let literal_expr = df_prelude::lit(*val);
1891
1892 let plan = LogicalPlan::Extension(Extension {
1893 node: Arc::new(
1894 EmptyMetric::new(
1895 self.ctx.start,
1896 self.ctx.end,
1897 self.ctx.interval,
1898 SPECIAL_TIME_FUNCTION.to_string(),
1899 DEFAULT_FIELD_COLUMN.to_string(),
1900 Some(literal_expr),
1901 )
1902 .context(DataFusionPlanningSnafu)?,
1903 ),
1904 });
1905 Ok(plan)
1906 }
1907
1908 fn prom_string_lit_to_plan(&mut self, string_literal: &StringLiteral) -> Result<LogicalPlan> {
1909 let StringLiteral { val } = string_literal;
1910 self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
1911 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1912 self.ctx.reset_table_name_and_schema();
1913 let literal_expr = df_prelude::lit(val.clone());
1914
1915 let plan = LogicalPlan::Extension(Extension {
1916 node: Arc::new(
1917 EmptyMetric::new(
1918 self.ctx.start,
1919 self.ctx.end,
1920 self.ctx.interval,
1921 SPECIAL_TIME_FUNCTION.to_string(),
1922 DEFAULT_FIELD_COLUMN.to_string(),
1923 Some(literal_expr),
1924 )
1925 .context(DataFusionPlanningSnafu)?,
1926 ),
1927 });
1928 Ok(plan)
1929 }
1930
1931 async fn prom_vector_selector_to_plan(
1932 &mut self,
1933 vector_selector: &VectorSelector,
1934 timestamp_fn: bool,
1935 ) -> Result<LogicalPlan> {
1936 let VectorSelector {
1937 name,
1938 offset,
1939 matchers,
1940 at: _,
1941 } = vector_selector;
1942 let matchers = self.preprocess_label_matchers(matchers, name)?;
1943 if let Some(empty_plan) = self.setup_context().await? {
1944 return Ok(empty_plan);
1945 }
1946 let normalize = self
1947 .selector_to_series_normalize_plan(offset, matchers, false)
1948 .await?;
1949 let time_index_column =
1950 self.ctx
1951 .time_index_column
1952 .clone()
1953 .with_context(|| TimeIndexNotFoundSnafu {
1954 table: self.ctx.table_name.clone().unwrap_or_default(),
1955 })?;
1956
1957 let (normalize, timestamp_value_column) = if timestamp_fn {
1958 let occupied = normalize
1961 .schema()
1962 .fields()
1963 .iter()
1964 .map(|field| field.name().as_str())
1965 .collect::<HashSet<_>>();
1966 let mut timestamp_value_column = TIMESTAMP_VALUE_PREFIX.to_string();
1967 while occupied.contains(timestamp_value_column.as_str()) {
1968 timestamp_value_column.push('_');
1969 }
1970 let mut project_exprs = normalize
1971 .schema()
1972 .iter()
1973 .map(|(qualifier, field)| {
1974 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
1975 })
1976 .collect::<Vec<_>>();
1977 project_exprs
1978 .push(build_special_time_expr(&time_index_column).alias(×tamp_value_column));
1979 let normalize = LogicalPlanBuilder::from(normalize)
1980 .project(project_exprs)
1981 .context(DataFusionPlanningSnafu)?
1982 .build()
1983 .context(DataFusionPlanningSnafu)?;
1984 (normalize, Some(timestamp_value_column))
1985 } else {
1986 (normalize, None)
1987 };
1988
1989 let field_column = self.ctx.field_columns.first().cloned();
1990 let manipulate = InstantManipulate::new(
1991 self.ctx.start,
1992 self.ctx.end,
1993 self.ctx.lookback_delta,
1994 self.ctx.interval,
1995 time_index_column,
1996 if self.ctx.use_tsid {
1997 vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]
1998 } else {
1999 self.ctx.tag_columns.clone()
2000 },
2001 field_column,
2002 normalize,
2003 );
2004 let manipulate = LogicalPlan::Extension(Extension {
2005 node: Arc::new(manipulate),
2006 });
2007 if let Some(timestamp_value_column) = timestamp_value_column {
2008 self.create_timestamp_func_plan(manipulate, ×tamp_value_column)
2009 } else {
2010 Ok(manipulate)
2011 }
2012 }
2013
2014 fn create_timestamp_func_plan(
2037 &mut self,
2038 input: LogicalPlan,
2039 timestamp_value_column: &str,
2040 ) -> Result<LogicalPlan> {
2041 let time_expr = col(timestamp_value_column).alias(DEFAULT_FIELD_COLUMN);
2042 self.ctx.field_columns = vec![time_expr.schema_name().to_string()];
2043 let mut project_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 2);
2044 project_exprs.push(self.create_time_index_column_expr()?);
2045 project_exprs.push(time_expr);
2046 project_exprs.extend(self.create_tag_column_exprs()?);
2047
2048 LogicalPlanBuilder::from(input)
2049 .project(project_exprs)
2050 .context(DataFusionPlanningSnafu)?
2051 .build()
2052 .context(DataFusionPlanningSnafu)
2053 }
2054
2055 async fn prom_matrix_selector_to_plan(
2056 &mut self,
2057 matrix_selector: &MatrixSelector,
2058 ) -> Result<LogicalPlan> {
2059 let MatrixSelector { vs, range } = matrix_selector;
2060 let VectorSelector {
2061 name,
2062 offset,
2063 matchers,
2064 ..
2065 } = vs;
2066 let matchers = self.preprocess_label_matchers(matchers, name)?;
2067 ensure!(!range.is_zero(), ZeroRangeSelectorSnafu);
2068 let range_ms = range.as_millis() as _;
2069 self.ctx.range = Some(range_ms);
2070
2071 let normalize = match self.setup_context().await? {
2074 Some(empty_plan) => empty_plan,
2075 None => {
2076 self.selector_to_series_normalize_plan(offset, matchers, true)
2077 .await?
2078 }
2079 };
2080 let manipulate = RangeManipulate::new(
2081 self.ctx.start,
2082 self.ctx.end,
2083 self.ctx.interval,
2084 range_ms,
2086 self.ctx
2087 .time_index_column
2088 .clone()
2089 .expect("time index should be set in `setup_context`"),
2090 self.ctx.field_columns.clone(),
2091 normalize,
2092 )
2093 .context(DataFusionPlanningSnafu)?;
2094
2095 Ok(LogicalPlan::Extension(Extension {
2096 node: Arc::new(manipulate),
2097 }))
2098 }
2099
2100 async fn prom_call_expr_to_plan(
2101 &mut self,
2102 query_engine_state: &QueryEngineState,
2103 call_expr: &Call,
2104 ) -> Result<LogicalPlan> {
2105 let Call { func, args } = call_expr;
2106 match func.name {
2108 SPECIAL_HISTOGRAM_QUANTILE | SPECIAL_HISTOGRAM_FRACTION => {
2109 return self
2110 .create_histogram_plan(func.name, args, query_engine_state)
2111 .await;
2112 }
2113 SPECIAL_VECTOR_FUNCTION => return self.create_vector_plan(args).await,
2114 SCALAR_FUNCTION => return self.create_scalar_plan(args, query_engine_state).await,
2115 SPECIAL_ABSENT_FUNCTION => {
2116 return self.create_absent_plan(args, query_engine_state).await;
2117 }
2118 _ => {}
2119 }
2120
2121 let args = self.create_function_args(&args.args)?;
2123 let input = if let Some(prom_expr) = &args.input {
2124 self.prom_expr_to_plan_inner(prom_expr, func.name == "timestamp", query_engine_state)
2125 .await?
2126 } else {
2127 self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
2128 self.ctx.reset_table_name_and_schema();
2129 self.ctx.tag_columns = vec![];
2130 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
2131 LogicalPlan::Extension(Extension {
2132 node: Arc::new(
2133 EmptyMetric::new(
2134 self.ctx.start,
2135 self.ctx.end,
2136 self.ctx.interval,
2137 SPECIAL_TIME_FUNCTION.to_string(),
2138 DEFAULT_FIELD_COLUMN.to_string(),
2139 None,
2140 )
2141 .context(DataFusionPlanningSnafu)?,
2142 ),
2143 })
2144 };
2145 let preserve_any_value =
2146 Self::field_columns_are_alternative_samples(input.schema(), &self.ctx.field_columns);
2147 let (mut func_exprs, new_tags) = self.create_function_expr(
2148 func,
2149 args.literals.clone(),
2150 input.schema(),
2151 query_engine_state,
2152 )?;
2153 func_exprs.insert(0, self.create_time_index_column_expr()?);
2154 func_exprs.extend_from_slice(&self.create_tag_column_exprs()?);
2155 if let Some(tsid_col) =
2156 Self::optional_tsid_projection(input.schema(), None, self.ctx.use_tsid)
2157 {
2158 func_exprs.push(tsid_col);
2159 }
2160
2161 let builder = LogicalPlanBuilder::from(input)
2162 .project(func_exprs)
2163 .context(DataFusionPlanningSnafu)?
2164 .filter(self.create_empty_values_filter_expr(preserve_any_value)?)
2165 .context(DataFusionPlanningSnafu)?;
2166
2167 let builder = match func.name {
2168 "sort" => builder
2169 .sort(self.create_field_columns_sort_exprs(true))
2170 .context(DataFusionPlanningSnafu)?,
2171 "sort_desc" => builder
2172 .sort(self.create_field_columns_sort_exprs(false))
2173 .context(DataFusionPlanningSnafu)?,
2174 "sort_by_label" => builder
2175 .sort(Self::create_sort_exprs_by_tags(
2176 func.name,
2177 args.literals,
2178 true,
2179 )?)
2180 .context(DataFusionPlanningSnafu)?,
2181 "sort_by_label_desc" => builder
2182 .sort(Self::create_sort_exprs_by_tags(
2183 func.name,
2184 args.literals,
2185 false,
2186 )?)
2187 .context(DataFusionPlanningSnafu)?,
2188
2189 _ => builder,
2190 };
2191
2192 for tag in new_tags {
2195 self.ctx.tag_columns.push(tag);
2196 }
2197
2198 let plan = builder.build().context(DataFusionPlanningSnafu)?;
2199 common_telemetry::debug!("Created PromQL function plan: {plan:?} for {call_expr:?}");
2200
2201 Ok(plan)
2202 }
2203
2204 async fn prom_ext_expr_to_plan(
2205 &mut self,
2206 query_engine_state: &QueryEngineState,
2207 ext_expr: &promql_parser::parser::ast::Extension,
2208 ) -> Result<LogicalPlan> {
2209 let expr = &ext_expr.expr;
2211 let children = expr.children();
2212 let plan = self
2213 .prom_expr_to_plan(&children[0], query_engine_state)
2214 .await?;
2215 match expr.name() {
2221 ANALYZE_NODE_NAME => LogicalPlanBuilder::from(plan)
2222 .explain(false, true)
2223 .unwrap()
2224 .build()
2225 .context(DataFusionPlanningSnafu),
2226 ANALYZE_VERBOSE_NODE_NAME => LogicalPlanBuilder::from(plan)
2227 .explain(true, true)
2228 .unwrap()
2229 .build()
2230 .context(DataFusionPlanningSnafu),
2231 EXPLAIN_NODE_NAME => LogicalPlanBuilder::from(plan)
2232 .explain(false, false)
2233 .unwrap()
2234 .build()
2235 .context(DataFusionPlanningSnafu),
2236 EXPLAIN_VERBOSE_NODE_NAME => LogicalPlanBuilder::from(plan)
2237 .explain(true, false)
2238 .unwrap()
2239 .build()
2240 .context(DataFusionPlanningSnafu),
2241 ALIAS_NODE_NAME => {
2242 let alias = expr
2243 .as_any()
2244 .downcast_ref::<AliasExpr>()
2245 .context(UnexpectedPlanExprSnafu {
2246 desc: "Expected AliasExpr",
2247 })?
2248 .alias
2249 .clone();
2250 self.apply_alias(plan, alias)
2251 }
2252 _ => LogicalPlanBuilder::empty(true)
2253 .build()
2254 .context(DataFusionPlanningSnafu),
2255 }
2256 }
2257
2258 #[allow(clippy::mutable_key_type)]
2268 fn preprocess_label_matchers(
2269 &mut self,
2270 label_matchers: &Matchers,
2271 name: &Option<String>,
2272 ) -> Result<Matchers> {
2273 self.ctx.reset();
2274
2275 let metric_name;
2276 if let Some(name) = name.clone() {
2277 metric_name = Some(name);
2278 ensure!(
2279 label_matchers.find_matchers(METRIC_NAME).is_empty(),
2280 MultipleMetricMatchersSnafu
2281 );
2282 } else {
2283 let mut matches = label_matchers.find_matchers(METRIC_NAME);
2284 ensure!(!matches.is_empty(), NoMetricMatcherSnafu);
2285 ensure!(matches.len() == 1, MultipleMetricMatchersSnafu);
2286 ensure!(
2287 matches[0].op == MatchOp::Equal,
2288 UnsupportedMatcherOpSnafu {
2289 matcher_op: matches[0].op.to_string(),
2290 matcher: METRIC_NAME
2291 }
2292 );
2293 metric_name = matches.pop().map(|m| m.value);
2294 }
2295
2296 self.ctx.table_name = metric_name;
2297
2298 let mut matchers = HashSet::new();
2299 for matcher in &label_matchers.matchers {
2300 if matcher.name == FIELD_COLUMN_MATCHER {
2302 self.ctx
2303 .field_column_matcher
2304 .get_or_insert_default()
2305 .push(matcher.clone());
2306 } else if matcher.name == SCHEMA_COLUMN_MATCHER || matcher.name == DB_COLUMN_MATCHER {
2307 ensure!(
2308 matcher.op == MatchOp::Equal,
2309 UnsupportedMatcherOpSnafu {
2310 matcher: matcher.name.clone(),
2311 matcher_op: matcher.op.to_string(),
2312 }
2313 );
2314 self.ctx.schema_name = Some(matcher.value.clone());
2315 } else if matcher.name != METRIC_NAME {
2316 self.ctx.selector_matcher.push(matcher.clone());
2317 let _ = matchers.insert(matcher.clone());
2318 }
2319 }
2320
2321 Ok(Matchers::new(matchers.into_iter().collect()))
2322 }
2323
2324 async fn selector_to_series_normalize_plan(
2325 &mut self,
2326 offset: &Option<Offset>,
2327 label_matchers: Matchers,
2328 is_range_selector: bool,
2329 ) -> Result<LogicalPlan> {
2330 let table_ref = self.table_ref()?;
2332 let mut table_scan = self.create_table_scan_plan(table_ref.clone()).await?;
2333 let table_schema = table_scan.schema();
2334
2335 let offset_duration = match offset {
2337 Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond,
2338 Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond),
2339 None => 0,
2340 };
2341 let mut scan_filters = Self::matchers_to_expr(label_matchers.clone(), table_schema)?;
2342 if let Some(time_index_filter) = self.build_time_index_filter(offset_duration)? {
2343 scan_filters.push(time_index_filter);
2344 }
2345 table_scan = LogicalPlanBuilder::from(table_scan)
2346 .filter(conjunction(scan_filters).unwrap()) .context(DataFusionPlanningSnafu)?
2348 .build()
2349 .context(DataFusionPlanningSnafu)?;
2350
2351 if let Some(field_matchers) = &self.ctx.field_column_matcher {
2353 let col_set = self.ctx.field_columns.iter().collect::<HashSet<_>>();
2354 let mut result_set = HashSet::new();
2356 let mut reverse_set = HashSet::new();
2358 for matcher in field_matchers {
2359 match &matcher.op {
2360 MatchOp::Equal => {
2361 if col_set.contains(&matcher.value) {
2362 let _ = result_set.insert(matcher.value.clone());
2363 } else {
2364 return Err(ColumnNotFoundSnafu {
2365 col: matcher.value.clone(),
2366 }
2367 .build());
2368 }
2369 }
2370 MatchOp::NotEqual => {
2371 if col_set.contains(&matcher.value) {
2372 let _ = reverse_set.insert(matcher.value.clone());
2373 } else {
2374 return Err(ColumnNotFoundSnafu {
2375 col: matcher.value.clone(),
2376 }
2377 .build());
2378 }
2379 }
2380 MatchOp::Re(regex) => {
2381 for col in &self.ctx.field_columns {
2382 if regex.is_match(col) {
2383 let _ = result_set.insert(col.clone());
2384 }
2385 }
2386 }
2387 MatchOp::NotRe(regex) => {
2388 for col in &self.ctx.field_columns {
2389 if regex.is_match(col) {
2390 let _ = reverse_set.insert(col.clone());
2391 }
2392 }
2393 }
2394 }
2395 }
2396 if result_set.is_empty() {
2398 result_set = col_set.into_iter().cloned().collect();
2399 }
2400 for col in reverse_set {
2401 let _ = result_set.remove(&col);
2402 }
2403
2404 self.ctx.field_columns = self
2406 .ctx
2407 .field_columns
2408 .drain(..)
2409 .filter(|col| result_set.contains(col))
2410 .collect();
2411
2412 let exprs = result_set
2413 .into_iter()
2414 .map(|col| DfExpr::Column(Column::new_unqualified(col)))
2415 .chain(self.create_tag_column_exprs()?)
2416 .chain(
2417 self.ctx
2418 .use_tsid
2419 .then_some(DfExpr::Column(Column::new_unqualified(
2420 DATA_SCHEMA_TSID_COLUMN_NAME,
2421 ))),
2422 )
2423 .chain(Some(self.create_time_index_column_expr()?))
2424 .collect::<Vec<_>>();
2425
2426 table_scan = LogicalPlanBuilder::from(table_scan)
2428 .project(exprs)
2429 .context(DataFusionPlanningSnafu)?
2430 .build()
2431 .context(DataFusionPlanningSnafu)?;
2432 }
2433
2434 let series_key_columns = if self.ctx.use_tsid {
2436 vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]
2437 } else {
2438 self.ctx.tag_columns.clone()
2439 };
2440
2441 let sort_exprs = if self.ctx.use_tsid {
2442 vec![
2443 DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME)).sort(true, true),
2444 self.create_time_index_column_expr()?.sort(true, true),
2445 ]
2446 } else {
2447 self.create_tag_and_time_index_column_sort_exprs()?
2448 };
2449
2450 let sort_plan = LogicalPlanBuilder::from(table_scan)
2451 .sort(sort_exprs)
2452 .context(DataFusionPlanningSnafu)?
2453 .build()
2454 .context(DataFusionPlanningSnafu)?;
2455
2456 let time_index_column =
2458 self.ctx
2459 .time_index_column
2460 .clone()
2461 .with_context(|| TimeIndexNotFoundSnafu {
2462 table: table_ref.to_string(),
2463 })?;
2464 let divide_plan = LogicalPlan::Extension(Extension {
2465 node: Arc::new(SeriesDivide::new(
2466 series_key_columns.clone(),
2467 time_index_column,
2468 sort_plan,
2469 )),
2470 });
2471
2472 if !is_range_selector && offset_duration == 0 {
2474 return Ok(divide_plan);
2475 }
2476 let series_normalize = SeriesNormalize::new(
2477 offset_duration,
2478 self.ctx
2479 .time_index_column
2480 .clone()
2481 .with_context(|| TimeIndexNotFoundSnafu {
2482 table: table_ref.to_quoted_string(),
2483 })?,
2484 is_range_selector,
2485 series_key_columns,
2486 divide_plan,
2487 );
2488 let logical_plan = LogicalPlan::Extension(Extension {
2489 node: Arc::new(series_normalize),
2490 });
2491
2492 Ok(logical_plan)
2493 }
2494
2495 fn agg_modifier_to_col(
2502 &mut self,
2503 input_schema: &DFSchemaRef,
2504 modifier: &Option<LabelModifier>,
2505 update_ctx: bool,
2506 ) -> Result<Vec<DfExpr>> {
2507 match modifier {
2508 None => {
2509 if update_ctx {
2510 self.ctx.tag_columns.clear();
2511 }
2512 Ok(vec![self.create_time_index_column_expr()?])
2513 }
2514 Some(LabelModifier::Include(labels)) => {
2515 if update_ctx {
2516 self.ctx.tag_columns.clear();
2517 }
2518 let mut exprs = Vec::with_capacity(labels.labels.len());
2519 for label in &labels.labels {
2520 if is_metric_engine_internal_column(label) {
2521 continue;
2522 }
2523 if let Some(column_name) = Self::find_case_sensitive_column(input_schema, label)
2525 {
2526 exprs.push(DfExpr::Column(Column::from_name(column_name.clone())));
2527
2528 if update_ctx {
2529 self.ctx.tag_columns.push(column_name);
2531 }
2532 }
2533 }
2534 exprs.push(self.create_time_index_column_expr()?);
2536
2537 Ok(exprs)
2538 }
2539 Some(LabelModifier::Exclude(labels)) => {
2540 let mut all_fields = input_schema
2541 .fields()
2542 .iter()
2543 .map(|f| f.name())
2544 .collect::<BTreeSet<_>>();
2545
2546 all_fields.retain(|col| !is_metric_engine_internal_column(col.as_str()));
2549
2550 for label in &labels.labels {
2553 let _ = all_fields.remove(label);
2554 }
2555
2556 if let Some(time_index) = &self.ctx.time_index_column {
2558 let _ = all_fields.remove(time_index);
2559 }
2560 for value in &self.ctx.field_columns {
2561 let _ = all_fields.remove(value);
2562 }
2563
2564 if update_ctx {
2565 self.ctx.tag_columns = all_fields.iter().map(|col| (*col).clone()).collect();
2567 }
2568
2569 let mut exprs = all_fields
2571 .into_iter()
2572 .map(|c| DfExpr::Column(Column::from(c)))
2573 .collect::<Vec<_>>();
2574
2575 exprs.push(self.create_time_index_column_expr()?);
2577
2578 Ok(exprs)
2579 }
2580 }
2581 }
2582
2583 pub fn matchers_to_expr(
2585 label_matchers: Matchers,
2586 table_schema: &DFSchemaRef,
2587 ) -> Result<Vec<DfExpr>> {
2588 let mut exprs = Vec::with_capacity(label_matchers.matchers.len());
2589 for matcher in label_matchers.matchers {
2590 if matcher.name == SCHEMA_COLUMN_MATCHER
2591 || matcher.name == DB_COLUMN_MATCHER
2592 || matcher.name == FIELD_COLUMN_MATCHER
2593 {
2594 continue;
2595 }
2596
2597 let accepts_empty = matcher.is_match("");
2598 let column_name = Self::find_case_sensitive_column(table_schema, matcher.name.as_str());
2599 let col = if let Some(column_name) = column_name {
2600 let column = DfExpr::Column(Column::from_name(&column_name));
2601 let field = table_schema
2602 .index_of_column_by_name(None, &column_name)
2603 .map(|index| table_schema.field(index));
2604 if accepts_empty
2605 && column_name == OTLP_AGGREGATION_TEMPORALITY_LABEL
2606 && let Some(data_type) = field
2607 .filter(|field| {
2608 field.is_nullable()
2609 && Self::string_value_data_type(field.data_type()).is_some()
2610 })
2611 .map(|field| field.data_type())
2612 {
2613 let empty = Self::string_scalar_value(data_type, Some(String::new()))
2614 .expect("nullable label has a string type");
2615 DfExpr::ScalarFunction(ScalarFunction {
2616 func: coalesce(),
2617 args: vec![column, DfExpr::Literal(empty, None)],
2618 })
2619 } else {
2620 column
2621 }
2622 } else {
2623 DfExpr::Literal(ScalarValue::Utf8(Some(String::new())), None)
2624 .alias(matcher.name.clone())
2625 };
2626 let lit = DfExpr::Literal(ScalarValue::Utf8(Some(matcher.value)), None);
2627 let expr = match matcher.op {
2628 MatchOp::Equal => col.eq(lit),
2629 MatchOp::NotEqual => col.not_eq(lit),
2630 MatchOp::Re(re) => {
2631 if re.as_str() == "^(?:.*)$" {
2637 continue;
2638 }
2639 if re.as_str() == "^(?:.+)$" {
2640 col.not_eq(DfExpr::Literal(
2641 ScalarValue::Utf8(Some(String::new())),
2642 None,
2643 ))
2644 } else {
2645 DfExpr::BinaryExpr(BinaryExpr {
2646 left: Box::new(col),
2647 op: Operator::RegexMatch,
2648 right: Box::new(DfExpr::Literal(
2649 ScalarValue::Utf8(Some(re.as_str().to_string())),
2650 None,
2651 )),
2652 })
2653 }
2654 }
2655 MatchOp::NotRe(re) => {
2656 if re.as_str() == "^(?:.*)$" {
2657 DfExpr::Literal(ScalarValue::Boolean(Some(false)), None)
2658 } else if re.as_str() == "^(?:.+)$" {
2659 col.eq(DfExpr::Literal(
2660 ScalarValue::Utf8(Some(String::new())),
2661 None,
2662 ))
2663 } else {
2664 DfExpr::BinaryExpr(BinaryExpr {
2665 left: Box::new(col),
2666 op: Operator::RegexNotMatch,
2667 right: Box::new(DfExpr::Literal(
2668 ScalarValue::Utf8(Some(re.as_str().to_string())),
2669 None,
2670 )),
2671 })
2672 }
2673 }
2674 };
2675 exprs.push(expr);
2676 }
2677
2678 Ok(exprs)
2679 }
2680
2681 fn find_case_sensitive_column(schema: &DFSchemaRef, column: &str) -> Option<String> {
2682 if is_metric_engine_internal_column(column) {
2683 return None;
2684 }
2685 schema
2686 .fields()
2687 .iter()
2688 .find(|field| field.name() == column)
2689 .map(|field| field.name().clone())
2690 }
2691
2692 fn table_from_source(&self, source: &Arc<dyn TableSource>) -> Result<table::TableRef> {
2693 Ok(source
2694 .as_any()
2695 .downcast_ref::<DefaultTableSource>()
2696 .context(UnknownTableSnafu)?
2697 .table_provider
2698 .as_any()
2699 .downcast_ref::<DfTableProviderAdapter>()
2700 .context(UnknownTableSnafu)?
2701 .table())
2702 }
2703
2704 fn table_ref(&self) -> Result<TableReference> {
2705 let table_name = self
2706 .ctx
2707 .table_name
2708 .clone()
2709 .context(TableNameNotFoundSnafu)?;
2710
2711 let table_ref = if let Some(schema_name) = &self.ctx.schema_name {
2713 TableReference::partial(schema_name.as_str(), table_name.as_str())
2714 } else {
2715 TableReference::bare(table_name.as_str())
2716 };
2717
2718 Ok(table_ref)
2719 }
2720
2721 fn build_time_index_filter(&self, offset_duration: i64) -> Result<Option<DfExpr>> {
2722 let start = self.ctx.start;
2723 let end = self.ctx.end;
2724 if end < start {
2725 return InvalidTimeRangeSnafu { start, end }.fail();
2726 }
2727 let lookback_delta = self.ctx.lookback_delta;
2728 let range = self.ctx.range.unwrap_or_default();
2729 let interval = self.ctx.interval;
2730 let time_index_expr = self.create_time_index_column_expr()?;
2731 let num_points = (end - start) / interval;
2732
2733 let selector_window = if range == 0 { lookback_delta } else { range };
2741 let lower_exclusive_adjustment = if selector_window > 0 { 1 } else { 0 };
2742
2743 if (end - start) / interval > MAX_SCATTER_POINTS || interval <= INTERVAL_1H {
2745 let single_time_range = time_index_expr
2746 .clone()
2747 .gt_eq(DfExpr::Literal(
2748 ScalarValue::TimestampMillisecond(
2749 Some(
2750 self.ctx.start - offset_duration - selector_window
2751 + lower_exclusive_adjustment,
2752 ),
2753 None,
2754 ),
2755 None,
2756 ))
2757 .and(time_index_expr.lt_eq(DfExpr::Literal(
2758 ScalarValue::TimestampMillisecond(Some(self.ctx.end - offset_duration), None),
2759 None,
2760 )));
2761 return Ok(Some(single_time_range));
2762 }
2763
2764 let mut filters = Vec::with_capacity(num_points as usize + 1);
2766 for timestamp in (start..=end).step_by(interval as usize) {
2767 filters.push(
2768 time_index_expr
2769 .clone()
2770 .gt_eq(DfExpr::Literal(
2771 ScalarValue::TimestampMillisecond(
2772 Some(
2773 timestamp - offset_duration - selector_window
2774 + lower_exclusive_adjustment,
2775 ),
2776 None,
2777 ),
2778 None,
2779 ))
2780 .and(time_index_expr.clone().lt_eq(DfExpr::Literal(
2781 ScalarValue::TimestampMillisecond(Some(timestamp - offset_duration), None),
2782 None,
2783 ))),
2784 )
2785 }
2786
2787 Ok(filters.into_iter().reduce(DfExpr::or))
2788 }
2789
2790 async fn create_table_scan_plan(&mut self, table_ref: TableReference) -> Result<LogicalPlan> {
2795 let provider = self
2796 .table_provider
2797 .resolve_table(table_ref.clone())
2798 .await
2799 .context(CatalogSnafu)?;
2800
2801 let logical_table = self.table_from_source(&provider)?;
2802
2803 let mut maybe_phy_table_ref = table_ref.clone();
2805 let mut scan_provider = provider;
2806 let mut table_id_filter: Option<u32> = None;
2807
2808 if logical_table.table_info().meta.engine == METRIC_ENGINE_NAME
2811 && let Some(physical_table_name) = logical_table
2812 .table_info()
2813 .meta
2814 .options
2815 .extra_options
2816 .get(LOGICAL_TABLE_METADATA_KEY)
2817 {
2818 let physical_table_ref = if let Some(schema_name) = &self.ctx.schema_name {
2819 TableReference::partial(schema_name.as_str(), physical_table_name.as_str())
2820 } else {
2821 TableReference::bare(physical_table_name.as_str())
2822 };
2823
2824 let physical_provider = match self
2825 .table_provider
2826 .resolve_table(physical_table_ref.clone())
2827 .await
2828 {
2829 Ok(provider) => provider,
2830 Err(e) if e.status_code() == StatusCode::TableNotFound => {
2831 scan_provider.clone()
2834 }
2835 Err(e) => return Err(e).context(CatalogSnafu),
2836 };
2837
2838 if !Arc::ptr_eq(&physical_provider, &scan_provider) {
2839 let physical_table = self.table_from_source(&physical_provider)?;
2841
2842 let has_table_id = physical_table
2843 .schema()
2844 .column_schema_by_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME)
2845 .is_some();
2846 let has_tsid = physical_table
2847 .schema()
2848 .column_schema_by_name(DATA_SCHEMA_TSID_COLUMN_NAME)
2849 .is_some_and(|col| matches!(col.data_type, ConcreteDataType::UInt64(_)));
2850
2851 if has_table_id && has_tsid {
2852 scan_provider = physical_provider;
2853 maybe_phy_table_ref = physical_table_ref;
2854 table_id_filter = Some(logical_table.table_info().ident.table_id);
2855 }
2856 }
2857 }
2858
2859 let scan_table = self.table_from_source(&scan_provider)?;
2860
2861 let use_tsid = table_id_filter.is_some()
2862 && scan_table
2863 .schema()
2864 .column_schema_by_name(DATA_SCHEMA_TSID_COLUMN_NAME)
2865 .is_some_and(|col| matches!(col.data_type, ConcreteDataType::UInt64(_)));
2866 self.ctx.use_tsid = use_tsid;
2867
2868 let all_table_tags = self.ctx.tag_columns.clone();
2869
2870 let scan_tag_columns = if use_tsid {
2871 let mut scan_tags = self.ctx.tag_columns.clone();
2872 for matcher in &self.ctx.selector_matcher {
2873 if is_metric_engine_internal_column(&matcher.name) {
2874 continue;
2875 }
2876 if all_table_tags.iter().any(|tag| tag == &matcher.name) {
2877 scan_tags.push(matcher.name.clone());
2878 }
2879 }
2880 scan_tags.sort_unstable();
2881 scan_tags.dedup();
2882 scan_tags
2883 } else {
2884 self.ctx.tag_columns.clone()
2885 };
2886
2887 let is_time_index_ms = scan_table
2888 .schema()
2889 .timestamp_column()
2890 .with_context(|| TimeIndexNotFoundSnafu {
2891 table: maybe_phy_table_ref.to_quoted_string(),
2892 })?
2893 .data_type
2894 == ConcreteDataType::timestamp_millisecond_datatype();
2895
2896 let scan_projection = if table_id_filter.is_some() {
2897 let mut required_columns = HashSet::new();
2898 required_columns.insert(DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string());
2899 required_columns.insert(self.ctx.time_index_column.clone().with_context(|| {
2900 TimeIndexNotFoundSnafu {
2901 table: maybe_phy_table_ref.to_quoted_string(),
2902 }
2903 })?);
2904 for col in &scan_tag_columns {
2905 required_columns.insert(col.clone());
2906 }
2907 for col in &self.ctx.field_columns {
2908 required_columns.insert(col.clone());
2909 }
2910 if use_tsid {
2911 required_columns.insert(DATA_SCHEMA_TSID_COLUMN_NAME.to_string());
2912 }
2913
2914 let arrow_schema = scan_table.schema().arrow_schema().clone();
2915 Some(
2916 arrow_schema
2917 .fields()
2918 .iter()
2919 .enumerate()
2920 .filter(|(_, field)| required_columns.contains(field.name().as_str()))
2921 .map(|(idx, _)| idx)
2922 .collect::<Vec<_>>(),
2923 )
2924 } else {
2925 None
2926 };
2927
2928 let mut scan_plan =
2929 LogicalPlanBuilder::scan(maybe_phy_table_ref.clone(), scan_provider, scan_projection)
2930 .context(DataFusionPlanningSnafu)?
2931 .build()
2932 .context(DataFusionPlanningSnafu)?;
2933
2934 if let Some(table_id) = table_id_filter {
2935 scan_plan = LogicalPlanBuilder::from(scan_plan)
2936 .filter(
2937 DfExpr::Column(Column::from_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME))
2938 .eq(lit(table_id)),
2939 )
2940 .context(DataFusionPlanningSnafu)?
2941 .alias(table_ref.clone()) .context(DataFusionPlanningSnafu)?
2943 .build()
2944 .context(DataFusionPlanningSnafu)?;
2945 }
2946
2947 if !is_time_index_ms {
2948 let expr: Vec<_> = self
2950 .create_field_column_exprs()?
2951 .into_iter()
2952 .chain(
2953 scan_tag_columns
2954 .iter()
2955 .map(|tag| DfExpr::Column(Column::from_name(tag))),
2956 )
2957 .chain(self.ctx.use_tsid.then_some(DfExpr::Column(Column::new(
2958 Some(table_ref.clone()),
2959 DATA_SCHEMA_TSID_COLUMN_NAME.to_string(),
2960 ))))
2961 .chain(Some(DfExpr::Alias(Alias {
2962 expr: Box::new(DfExpr::Cast(Cast {
2963 expr: Box::new(self.create_time_index_column_expr()?),
2964 data_type: ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
2965 })),
2966 relation: Some(table_ref.clone()),
2967 name: self
2968 .ctx
2969 .time_index_column
2970 .as_ref()
2971 .with_context(|| TimeIndexNotFoundSnafu {
2972 table: table_ref.to_quoted_string(),
2973 })?
2974 .clone(),
2975 metadata: None,
2976 })))
2977 .collect::<Vec<_>>();
2978 scan_plan = LogicalPlanBuilder::from(scan_plan)
2979 .project(expr)
2980 .context(DataFusionPlanningSnafu)?
2981 .build()
2982 .context(DataFusionPlanningSnafu)?;
2983 } else if table_id_filter.is_some() {
2984 let project_exprs = self
2986 .create_field_column_exprs()?
2987 .into_iter()
2988 .chain(
2989 scan_tag_columns
2990 .iter()
2991 .map(|tag| DfExpr::Column(Column::from_name(tag))),
2992 )
2993 .chain(
2994 self.ctx
2995 .use_tsid
2996 .then_some(DfExpr::Column(Column::from_name(
2997 DATA_SCHEMA_TSID_COLUMN_NAME,
2998 ))),
2999 )
3000 .chain(Some(self.create_time_index_column_expr()?))
3001 .collect::<Vec<_>>();
3002
3003 scan_plan = LogicalPlanBuilder::from(scan_plan)
3004 .project(project_exprs)
3005 .context(DataFusionPlanningSnafu)?
3006 .build()
3007 .context(DataFusionPlanningSnafu)?;
3008 }
3009
3010 let result = LogicalPlanBuilder::from(scan_plan)
3011 .build()
3012 .context(DataFusionPlanningSnafu)?;
3013 Ok(result)
3014 }
3015
3016 fn collect_row_key_tag_columns_from_plan(
3017 &self,
3018 plan: &LogicalPlan,
3019 ) -> Result<BTreeSet<String>> {
3020 fn walk(
3021 planner: &PromPlanner,
3022 plan: &LogicalPlan,
3023 out: &mut BTreeSet<String>,
3024 ) -> Result<()> {
3025 if let LogicalPlan::TableScan(scan) = plan
3027 && let Ok(table) = planner.table_from_source(&scan.source)
3028 {
3029 for col in table.table_info().meta.row_key_column_names() {
3030 if col != DATA_SCHEMA_TABLE_ID_COLUMN_NAME
3031 && col != DATA_SCHEMA_TSID_COLUMN_NAME
3032 && !is_metric_engine_internal_column(col)
3033 {
3034 out.insert(col.clone());
3035 }
3036 }
3037 }
3038
3039 for input in plan.inputs() {
3040 walk(planner, input, out)?;
3041 }
3042 Ok(())
3043 }
3044
3045 let mut out = BTreeSet::new();
3046 walk(self, plan, &mut out)?;
3047 Ok(out)
3048 }
3049
3050 fn ensure_tag_columns_available(
3051 &self,
3052 plan: LogicalPlan,
3053 required_tags: &BTreeSet<String>,
3054 ) -> Result<LogicalPlan> {
3055 if required_tags.is_empty() {
3056 return Ok(plan);
3057 }
3058
3059 struct Rewriter {
3060 required_tags: BTreeSet<String>,
3061 }
3062
3063 impl TreeNodeRewriter for Rewriter {
3064 type Node = LogicalPlan;
3065
3066 fn f_up(
3067 &mut self,
3068 node: Self::Node,
3069 ) -> datafusion_common::Result<Transformed<Self::Node>> {
3070 match node {
3071 LogicalPlan::TableScan(scan) => {
3072 let schema = scan.source.schema();
3073 let mut projection = match scan.projection.clone() {
3074 Some(p) => p,
3075 None => {
3076 return Ok(Transformed::no(LogicalPlan::TableScan(scan)));
3078 }
3079 };
3080
3081 let mut changed = false;
3082 for tag in &self.required_tags {
3083 if let Some((idx, _)) = schema
3084 .fields()
3085 .iter()
3086 .enumerate()
3087 .find(|(_, field)| field.name() == tag)
3088 && !projection.contains(&idx)
3089 {
3090 projection.push(idx);
3091 changed = true;
3092 }
3093 }
3094
3095 if !changed {
3096 return Ok(Transformed::no(LogicalPlan::TableScan(scan)));
3097 }
3098
3099 projection.sort_unstable();
3100 projection.dedup();
3101
3102 let new_scan = TableScan::try_new(
3103 scan.table_name.clone(),
3104 scan.source.clone(),
3105 Some(projection),
3106 scan.filters,
3107 scan.fetch,
3108 )?;
3109 Ok(Transformed::yes(LogicalPlan::TableScan(new_scan)))
3110 }
3111 LogicalPlan::Projection(proj) => {
3112 let input_schema = proj.input.schema();
3113
3114 let existing = proj
3115 .schema
3116 .fields()
3117 .iter()
3118 .map(|f| f.name().as_str())
3119 .collect::<HashSet<_>>();
3120
3121 let mut expr = proj.expr.clone();
3122 let mut has_changed = false;
3123 for tag in &self.required_tags {
3124 if existing.contains(tag.as_str()) {
3125 continue;
3126 }
3127
3128 if let Some(idx) = input_schema.index_of_column_by_name(None, tag) {
3129 expr.push(DfExpr::Column(Column::from(
3130 input_schema.qualified_field(idx),
3131 )));
3132 has_changed = true;
3133 }
3134 }
3135
3136 if !has_changed {
3137 return Ok(Transformed::no(LogicalPlan::Projection(proj)));
3138 }
3139
3140 let new_proj = Projection::try_new(expr, proj.input)?;
3141 Ok(Transformed::yes(LogicalPlan::Projection(new_proj)))
3142 }
3143 other => Ok(Transformed::no(other)),
3144 }
3145 }
3146 }
3147
3148 let mut rewriter = Rewriter {
3149 required_tags: required_tags.clone(),
3150 };
3151 let rewritten = plan
3152 .rewrite(&mut rewriter)
3153 .context(DataFusionPlanningSnafu)?;
3154 Ok(rewritten.data)
3155 }
3156
3157 fn refresh_tag_columns_from_schema(&mut self, schema: &DFSchemaRef) {
3158 let time_index = self.ctx.time_index_column.as_deref();
3159 let field_columns = self.ctx.field_columns.iter().collect::<HashSet<_>>();
3160
3161 let mut tags = schema
3162 .fields()
3163 .iter()
3164 .map(|f| f.name())
3165 .filter(|name| Some(name.as_str()) != time_index)
3166 .filter(|name| !field_columns.contains(name))
3167 .filter(|name| !is_metric_engine_internal_column(name))
3168 .cloned()
3169 .collect::<Vec<_>>();
3170 tags.sort_unstable();
3171 tags.dedup();
3172 self.ctx.tag_columns = tags;
3173 }
3174
3175 async fn setup_context(&mut self) -> Result<Option<LogicalPlan>> {
3179 let table_ref = self.table_ref()?;
3180 let source = match self.table_provider.resolve_table(table_ref.clone()).await {
3181 Err(e) if e.status_code() == StatusCode::TableNotFound => {
3182 let plan = self.setup_context_for_empty_metric()?;
3183 return Ok(Some(plan));
3184 }
3185 res => res.context(CatalogSnafu)?,
3186 };
3187 let table = self.table_from_source(&source)?;
3188
3189 let time_index = table
3191 .schema()
3192 .timestamp_column()
3193 .with_context(|| TimeIndexNotFoundSnafu {
3194 table: table_ref.to_quoted_string(),
3195 })?
3196 .name
3197 .clone();
3198 self.ctx.time_index_column = Some(time_index);
3199
3200 let values = table
3202 .table_info()
3203 .meta
3204 .field_column_names()
3205 .cloned()
3206 .collect();
3207 self.ctx.field_columns = values;
3208
3209 let tags = table
3211 .table_info()
3212 .meta
3213 .row_key_column_names()
3214 .filter(|col| {
3215 col != &DATA_SCHEMA_TABLE_ID_COLUMN_NAME && col != &DATA_SCHEMA_TSID_COLUMN_NAME
3217 })
3218 .cloned()
3219 .collect();
3220 self.ctx.tag_columns = tags;
3221
3222 self.ctx.use_tsid = false;
3223
3224 Ok(None)
3225 }
3226
3227 fn setup_context_for_empty_metric(&mut self) -> Result<LogicalPlan> {
3230 self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
3231 self.ctx.reset_table_name_and_schema();
3232 self.ctx.tag_columns = vec![];
3233 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
3234 self.ctx.use_tsid = false;
3235
3236 let plan = LogicalPlan::Extension(Extension {
3238 node: Arc::new(
3239 EmptyMetric::new(
3240 0,
3241 -1,
3242 self.ctx.interval,
3243 SPECIAL_TIME_FUNCTION.to_string(),
3244 DEFAULT_FIELD_COLUMN.to_string(),
3245 Some(lit(0.0f64)),
3246 )
3247 .context(DataFusionPlanningSnafu)?,
3248 ),
3249 });
3250 Ok(plan)
3251 }
3252
3253 fn create_function_args(&self, args: &[Box<PromExpr>]) -> Result<FunctionArgs> {
3255 let mut result = FunctionArgs::default();
3256
3257 for arg in args {
3258 if let Some(expr) = Self::try_build_literal_expr(arg) {
3260 result.literals.push(expr);
3261 } else {
3262 match arg.as_ref() {
3264 PromExpr::Subquery(_)
3265 | PromExpr::VectorSelector(_)
3266 | PromExpr::MatrixSelector(_)
3267 | PromExpr::Extension(_)
3268 | PromExpr::Aggregate(_)
3269 | PromExpr::Paren(_)
3270 | PromExpr::Call(_)
3271 | PromExpr::Binary(_)
3272 | PromExpr::Unary(_) => {
3273 if result.input.replace(*arg.clone()).is_some() {
3274 MultipleVectorSnafu { expr: *arg.clone() }.fail()?;
3275 }
3276 }
3277
3278 _ => {
3279 let expr = Self::get_param_as_literal_expr(Some(arg.as_ref()), None, None)?;
3280 result.literals.push(expr);
3281 }
3282 }
3283 }
3284 }
3285
3286 Ok(result)
3287 }
3288
3289 fn create_mixed_range_function_exprs(
3290 &mut self,
3291 func: &Function,
3292 mut other_input_exprs: VecDeque<DfExpr>,
3293 float_field: &str,
3294 histogram_field: &str,
3295 input_schema: &DFSchemaRef,
3296 ) -> Result<Option<Vec<DfExpr>>> {
3297 let returns_histogram = matches!(
3298 func.name,
3299 "rate"
3300 | "increase"
3301 | "delta"
3302 | "idelta"
3303 | "irate"
3304 | "avg_over_time"
3305 | "sum_over_time"
3306 | "last_over_time"
3307 );
3308 if !returns_histogram
3309 && !matches!(
3310 func.name,
3311 "changes"
3312 | "resets"
3313 | "deriv"
3314 | "min_over_time"
3315 | "max_over_time"
3316 | "count_over_time"
3317 | "absent_over_time"
3318 | "present_over_time"
3319 | "stddev_over_time"
3320 | "stdvar_over_time"
3321 | "quantile_over_time"
3322 | "predict_linear"
3323 | "double_exponential_smoothing"
3324 | "holt_winters"
3325 )
3326 {
3327 return Ok(None);
3328 }
3329
3330 if func.name == "predict_linear" {
3331 other_input_exprs[0] = DfExpr::Cast(Cast {
3332 expr: Box::new(other_input_exprs[0].clone()),
3333 data_type: ArrowDataType::Int64,
3334 });
3335 }
3336
3337 let timestamp_range = DfExpr::Column(Column::from_name(
3338 RangeManipulate::build_timestamp_range_name(
3339 self.ctx.time_index_column.as_ref().unwrap(),
3340 ),
3341 ));
3342 let float_range = DfExpr::Column(Column::from_name(float_field));
3343 let histogram_range = DfExpr::Column(Column::from_name(histogram_field));
3344 let mut args = Vec::with_capacity(other_input_exprs.len() + 6);
3345 args.push(lit(func.name));
3346 args.push(timestamp_range.clone());
3347 args.push(float_range.clone());
3348 args.push(histogram_range.clone());
3349 args.extend(other_input_exprs);
3350 if matches!(func.name, "rate" | "increase" | "delta") {
3351 args.push(self.create_time_index_column_expr()?);
3352 args.push(lit(self.ctx.range.context(ExpectRangeSelectorSnafu)?));
3353 }
3354
3355 let mut float_expr = DfExpr::ScalarFunction(ScalarFunction {
3356 func: Arc::new(MixedRange::float_udf(self.promql_annotations.clone())),
3357 args: args.clone(),
3358 });
3359 if matches!(func.name, "rate" | "increase") {
3360 let raw_delta_function = if func.name == "rate" {
3361 "raw_delta_rate"
3362 } else {
3363 "raw_delta_increase"
3364 };
3365 let delta_sum = DfExpr::ScalarFunction(ScalarFunction {
3366 func: Arc::new(MixedRange::float_udf(self.promql_annotations.clone())),
3367 args: vec![
3368 lit(raw_delta_function),
3369 timestamp_range,
3370 float_range,
3371 histogram_range,
3372 ],
3373 });
3374 float_expr = self.select_delta_range_math(
3375 func.name,
3376 input_schema,
3377 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3378 delta_sum,
3379 float_expr,
3380 )?;
3381 }
3382 let exprs = if returns_histogram {
3383 self.ctx.field_columns = vec![float_field.to_string(), histogram_field.to_string()];
3384 vec![
3385 float_expr.alias(float_field),
3386 DfExpr::ScalarFunction(ScalarFunction {
3387 func: Arc::new(MixedRange::histogram_udf(self.promql_annotations.clone())),
3388 args,
3389 })
3390 .alias(histogram_field),
3391 ]
3392 } else {
3393 let display_name = float_expr.schema_name().to_string();
3394 self.ctx.field_columns = vec![display_name.clone()];
3395 vec![float_expr.alias(display_name)]
3396 };
3397 Ok(Some(exprs))
3398 }
3399
3400 fn create_function_expr(
3406 &mut self,
3407 func: &Function,
3408 other_input_exprs: Vec<DfExpr>,
3409 input_schema: &DFSchemaRef,
3410 query_engine_state: &QueryEngineState,
3411 ) -> Result<(Vec<DfExpr>, Vec<String>)> {
3412 let mut other_input_exprs: VecDeque<DfExpr> = other_input_exprs.into();
3414 if let Some((float_field, histogram_field)) =
3415 Self::alternative_sample_range_columns(input_schema, &self.ctx.field_columns)
3416 .map(|(float, histogram)| (float.to_string(), histogram.to_string()))
3417 && let Some(exprs) = self.create_mixed_range_function_exprs(
3418 func,
3419 other_input_exprs.clone(),
3420 &float_field,
3421 &histogram_field,
3422 input_schema,
3423 )?
3424 {
3425 return Ok((exprs, vec![]));
3426 }
3427 let alternative_samples =
3428 Self::field_columns_are_alternative_samples(input_schema, &self.ctx.field_columns);
3429 let all_field_columns_are_native_histogram_ranges =
3430 self.all_field_columns_are_native_histogram_ranges(input_schema);
3431
3432 let field_column_pos = 0;
3434 let mut exprs = Vec::with_capacity(self.ctx.field_columns.len());
3435 let mut new_tags = vec![];
3437 let promql_annotations = self.promql_annotations.clone();
3438 let native_histogram_drop_udf = |name: &str| {
3439 Arc::new(NativeHistogramDrop::float_null_udf(
3440 format!(
3441 "{name}: dropped native histogram samples because this function is not supported for native histograms"
3442 ),
3443 promql_annotations.clone(),
3444 ))
3445 };
3446 let scalar_func = match func.name {
3447 "increase" => {
3448 if all_field_columns_are_native_histogram_ranges {
3449 ScalarFunc::ExtrapolateUdf(
3450 Arc::new(NativeHistogramIncrease::scalar_udf_with_collector(
3451 self.promql_annotations.clone(),
3452 )),
3453 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3454 )
3455 } else {
3456 ScalarFunc::ExtrapolateUdf(
3457 Arc::new(Increase::scalar_udf()),
3458 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3459 )
3460 }
3461 }
3462 "rate" => {
3463 if all_field_columns_are_native_histogram_ranges {
3464 ScalarFunc::ExtrapolateUdf(
3465 Arc::new(NativeHistogramRate::scalar_udf_with_collector(
3466 self.promql_annotations.clone(),
3467 )),
3468 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3469 )
3470 } else {
3471 ScalarFunc::ExtrapolateUdf(
3472 Arc::new(Rate::scalar_udf()),
3473 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3474 )
3475 }
3476 }
3477 "delta" => {
3478 if all_field_columns_are_native_histogram_ranges {
3479 ScalarFunc::ExtrapolateUdf(
3480 Arc::new(NativeHistogramDelta::scalar_udf_with_collector(
3481 self.promql_annotations.clone(),
3482 )),
3483 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3484 )
3485 } else {
3486 ScalarFunc::ExtrapolateUdf(
3487 Arc::new(Delta::scalar_udf()),
3488 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3489 )
3490 }
3491 }
3492 "idelta" => {
3493 if all_field_columns_are_native_histogram_ranges {
3494 ScalarFunc::Udf(Arc::new(NativeHistogramIDelta::scalar_udf_with_collector(
3495 self.promql_annotations.clone(),
3496 )))
3497 } else {
3498 ScalarFunc::Udf(Arc::new(IDelta::<false>::scalar_udf()))
3499 }
3500 }
3501 "irate" => {
3502 if all_field_columns_are_native_histogram_ranges {
3503 ScalarFunc::Udf(Arc::new(NativeHistogramIRate::scalar_udf_with_collector(
3504 self.promql_annotations.clone(),
3505 )))
3506 } else {
3507 ScalarFunc::Udf(Arc::new(IDelta::<true>::scalar_udf()))
3508 }
3509 }
3510 "resets" => {
3511 if all_field_columns_are_native_histogram_ranges {
3512 ScalarFunc::Udf(Arc::new(NativeHistogramResets::scalar_udf()))
3513 } else {
3514 ScalarFunc::Udf(Arc::new(Resets::scalar_udf()))
3515 }
3516 }
3517 "changes" => {
3518 if all_field_columns_are_native_histogram_ranges {
3519 ScalarFunc::Udf(Arc::new(NativeHistogramChanges::scalar_udf()))
3520 } else {
3521 ScalarFunc::Udf(Arc::new(Changes::scalar_udf()))
3522 }
3523 }
3524 "deriv" => {
3525 if all_field_columns_are_native_histogram_ranges {
3526 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3527 } else {
3528 ScalarFunc::Udf(Arc::new(Deriv::scalar_udf()))
3529 }
3530 }
3531 "avg_over_time" => {
3532 if all_field_columns_are_native_histogram_ranges {
3533 ScalarFunc::Udf(Arc::new(
3534 NativeHistogramAvgOverTime::scalar_udf_with_collector(
3535 self.promql_annotations.clone(),
3536 ),
3537 ))
3538 } else {
3539 ScalarFunc::Udf(Arc::new(AvgOverTime::scalar_udf()))
3540 }
3541 }
3542 "min_over_time" => {
3543 if all_field_columns_are_native_histogram_ranges {
3544 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3545 } else {
3546 ScalarFunc::Udf(Arc::new(MinOverTime::scalar_udf()))
3547 }
3548 }
3549 "max_over_time" => {
3550 if all_field_columns_are_native_histogram_ranges {
3551 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3552 } else {
3553 ScalarFunc::Udf(Arc::new(MaxOverTime::scalar_udf()))
3554 }
3555 }
3556 "sum_over_time" => {
3557 if all_field_columns_are_native_histogram_ranges {
3558 ScalarFunc::Udf(Arc::new(
3559 NativeHistogramSumOverTime::scalar_udf_with_collector(
3560 self.promql_annotations.clone(),
3561 ),
3562 ))
3563 } else {
3564 ScalarFunc::Udf(Arc::new(SumOverTime::scalar_udf()))
3565 }
3566 }
3567 "count_over_time" => {
3568 if all_field_columns_are_native_histogram_ranges {
3569 ScalarFunc::Udf(Arc::new(NativeHistogramCountOverTime::scalar_udf()))
3570 } else {
3571 ScalarFunc::Udf(Arc::new(CountOverTime::scalar_udf()))
3572 }
3573 }
3574 "last_over_time" => {
3575 if all_field_columns_are_native_histogram_ranges {
3576 ScalarFunc::Udf(Arc::new(NativeHistogramLastOverTime::scalar_udf()))
3577 } else {
3578 ScalarFunc::Udf(Arc::new(LastOverTime::scalar_udf()))
3579 }
3580 }
3581 "absent_over_time" => {
3582 if all_field_columns_are_native_histogram_ranges {
3583 ScalarFunc::Udf(Arc::new(NativeHistogramAbsentOverTime::scalar_udf()))
3584 } else {
3585 ScalarFunc::Udf(Arc::new(AbsentOverTime::scalar_udf()))
3586 }
3587 }
3588 "present_over_time" => {
3589 if all_field_columns_are_native_histogram_ranges {
3590 ScalarFunc::Udf(Arc::new(NativeHistogramPresentOverTime::scalar_udf()))
3591 } else {
3592 ScalarFunc::Udf(Arc::new(PresentOverTime::scalar_udf()))
3593 }
3594 }
3595 "stddev_over_time" => {
3596 if all_field_columns_are_native_histogram_ranges {
3597 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3598 } else {
3599 ScalarFunc::Udf(Arc::new(StddevOverTime::scalar_udf()))
3600 }
3601 }
3602 "stdvar_over_time" => {
3603 if all_field_columns_are_native_histogram_ranges {
3604 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3605 } else {
3606 ScalarFunc::Udf(Arc::new(StdvarOverTime::scalar_udf()))
3607 }
3608 }
3609 "quantile_over_time" => {
3610 if all_field_columns_are_native_histogram_ranges {
3611 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3612 } else {
3613 ScalarFunc::Udf(Arc::new(QuantileOverTime::scalar_udf()))
3614 }
3615 }
3616 "predict_linear" => {
3617 if all_field_columns_are_native_histogram_ranges {
3618 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3619 } else {
3620 other_input_exprs[0] = DfExpr::Cast(Cast {
3621 expr: Box::new(other_input_exprs[0].clone()),
3622 data_type: ArrowDataType::Int64,
3623 });
3624 ScalarFunc::Udf(Arc::new(PredictLinear::scalar_udf()))
3625 }
3626 }
3627 "double_exponential_smoothing" | "holt_winters" => {
3628 if all_field_columns_are_native_histogram_ranges {
3629 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3630 } else {
3631 ScalarFunc::Udf(Arc::new(DoubleExponentialSmoothing::scalar_udf()))
3632 }
3633 }
3634 "histogram_count" => {
3635 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramCount::scalar_udf()))
3636 }
3637 "histogram_sum" => {
3638 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramSum::scalar_udf()))
3639 }
3640 "histogram_avg" => {
3641 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramAvg::scalar_udf()))
3642 }
3643 "histogram_stddev" => {
3644 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramStddev::scalar_udf()))
3645 }
3646 "histogram_stdvar" => {
3647 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramStdvar::scalar_udf()))
3648 }
3649 "time" => {
3650 exprs.push(build_special_time_expr(
3651 self.ctx.time_index_column.as_ref().unwrap(),
3652 ));
3653 ScalarFunc::GeneratedExpr
3654 }
3655 "minute" => {
3656 let expr = self.date_part_on_time_index("minute")?;
3658 exprs.push(expr);
3659 ScalarFunc::GeneratedExpr
3660 }
3661 "hour" => {
3662 let expr = self.date_part_on_time_index("hour")?;
3664 exprs.push(expr);
3665 ScalarFunc::GeneratedExpr
3666 }
3667 "month" => {
3668 let expr = self.date_part_on_time_index("month")?;
3670 exprs.push(expr);
3671 ScalarFunc::GeneratedExpr
3672 }
3673 "year" => {
3674 let expr = self.date_part_on_time_index("year")?;
3676 exprs.push(expr);
3677 ScalarFunc::GeneratedExpr
3678 }
3679 "day_of_month" => {
3680 let expr = self.date_part_on_time_index("day")?;
3682 exprs.push(expr);
3683 ScalarFunc::GeneratedExpr
3684 }
3685 "day_of_week" => {
3686 let expr = self.date_part_on_time_index("dow")?;
3688 exprs.push(expr);
3689 ScalarFunc::GeneratedExpr
3690 }
3691 "day_of_year" => {
3692 let expr = self.date_part_on_time_index("doy")?;
3694 exprs.push(expr);
3695 ScalarFunc::GeneratedExpr
3696 }
3697 "days_in_month" => {
3698 let day_lit_expr = "day".lit();
3703 let month_lit_expr = "month".lit();
3704 let interval_1month_lit_expr =
3705 DfExpr::Literal(ScalarValue::IntervalYearMonth(Some(1)), None);
3706 let interval_1day_lit_expr = DfExpr::Literal(
3707 ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(1, 0))),
3708 None,
3709 );
3710 let the_1month_minus_1day_expr = DfExpr::BinaryExpr(BinaryExpr {
3711 left: Box::new(interval_1month_lit_expr),
3712 op: Operator::Minus,
3713 right: Box::new(interval_1day_lit_expr),
3714 });
3715 let date_trunc_expr = DfExpr::ScalarFunction(ScalarFunction {
3716 func: datafusion_functions::datetime::date_trunc(),
3717 args: vec![month_lit_expr, self.create_time_index_column_expr()?],
3718 });
3719 let date_trunc_plus_interval_expr = DfExpr::BinaryExpr(BinaryExpr {
3720 left: Box::new(date_trunc_expr),
3721 op: Operator::Plus,
3722 right: Box::new(the_1month_minus_1day_expr),
3723 });
3724 let date_part_expr = DfExpr::ScalarFunction(ScalarFunction {
3725 func: datafusion_functions::datetime::date_part(),
3726 args: vec![day_lit_expr, date_trunc_plus_interval_expr],
3727 });
3728
3729 exprs.push(date_part_expr);
3730 ScalarFunc::GeneratedExpr
3731 }
3732
3733 "label_join" => {
3734 self.ctx.use_tsid = false;
3735 let (concat_expr, dst_label) = Self::build_concat_labels_expr(
3736 &mut other_input_exprs,
3737 &self.ctx,
3738 query_engine_state,
3739 )?;
3740
3741 for value in &self.ctx.field_columns {
3743 if *value != dst_label {
3744 let expr = DfExpr::Column(Column::from_name(value));
3745 exprs.push(expr);
3746 }
3747 }
3748
3749 self.ctx.tag_columns.retain(|tag| *tag != dst_label);
3751 new_tags.push(dst_label);
3752 exprs.push(concat_expr);
3754
3755 ScalarFunc::GeneratedExpr
3756 }
3757 "label_replace" => {
3758 self.ctx.use_tsid = false;
3759 if let Some((replace_expr, dst_label)) = self
3760 .build_regexp_replace_label_expr(&mut other_input_exprs, query_engine_state)?
3761 {
3762 for value in &self.ctx.field_columns {
3764 if *value != dst_label {
3765 let expr = DfExpr::Column(Column::from_name(value));
3766 exprs.push(expr);
3767 }
3768 }
3769
3770 ensure!(
3771 !self.ctx.tag_columns.contains(&dst_label),
3772 SameLabelSetSnafu
3773 );
3774 new_tags.push(dst_label);
3775 exprs.push(replace_expr);
3777 } else {
3778 for value in &self.ctx.field_columns {
3780 let expr = DfExpr::Column(Column::from_name(value));
3781 exprs.push(expr);
3782 }
3783 }
3784
3785 ScalarFunc::GeneratedExpr
3786 }
3787 "sort" | "sort_desc" => {
3788 for value in &self.ctx.field_columns {
3790 if !Self::field_column_is_native_histogram(input_schema, value) {
3791 exprs.push(DfExpr::Column(Column::from_name(value)));
3792 }
3793 }
3794 if exprs.is_empty() {
3797 exprs.push(DfExpr::Literal(ScalarValue::Float64(None), None));
3798 }
3799
3800 ScalarFunc::GeneratedExpr
3801 }
3802 "sort_by_label" | "sort_by_label_desc" | "timestamp" => {
3803 for value in &self.ctx.field_columns {
3806 let expr = DfExpr::Column(Column::from_name(value));
3807 exprs.push(expr);
3808 }
3809
3810 ScalarFunc::GeneratedExpr
3811 }
3812 "round" if self.all_field_columns_are_native_histograms(input_schema) => {
3813 if other_input_exprs.is_empty() {
3814 other_input_exprs.push_front(0.0f64.lit());
3815 }
3816 ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3817 }
3818 "round" => {
3819 if other_input_exprs.is_empty() {
3820 other_input_exprs.push_front(0.0f64.lit());
3821 }
3822 ScalarFunc::DataFusionUdf(Arc::new(Round::scalar_udf()))
3823 }
3824 "rad" | "deg" | "sgn" if self.all_field_columns_are_native_histograms(input_schema) => {
3825 ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3826 }
3827 "rad" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::radians()),
3828 "deg" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::degrees()),
3829 "sgn" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::signum()),
3830 "pi" => {
3831 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
3833 func: datafusion::functions::math::pi(),
3834 args: vec![],
3835 });
3836 exprs.push(fn_expr);
3837
3838 ScalarFunc::GeneratedExpr
3839 }
3840 _ => {
3841 if let Some(f) = query_engine_state
3842 .session_state()
3843 .scalar_functions()
3844 .get(func.name)
3845 {
3846 if self.all_field_columns_are_native_histograms(input_schema) {
3847 ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3848 } else {
3849 ScalarFunc::DataFusionBuiltin(f.clone())
3850 }
3851 } else if let Some(factory) = query_engine_state.scalar_function(func.name) {
3852 if self.all_field_columns_are_native_histograms(input_schema) {
3853 ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3854 } else {
3855 let func_state = query_engine_state.function_state();
3856 let query_ctx = self.table_provider.query_ctx();
3857
3858 ScalarFunc::DataFusionUdf(Arc::new(factory.provide(FunctionContext {
3859 state: func_state,
3860 query_ctx: query_ctx.clone(),
3861 })))
3862 }
3863 } else if let Some(f) = datafusion_functions::math::functions()
3864 .iter()
3865 .find(|f| f.name() == func.name)
3866 {
3867 if self.all_field_columns_are_native_histograms(input_schema) {
3868 ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3869 } else {
3870 ScalarFunc::DataFusionUdf(f.clone())
3871 }
3872 } else {
3873 return UnsupportedExprSnafu {
3874 name: func.name.to_string(),
3875 }
3876 .fail();
3877 }
3878 }
3879 };
3880
3881 for value in &self.ctx.field_columns {
3882 let col_expr = DfExpr::Column(Column::from_name(value));
3883 let value_is_histogram = Self::field_column_is_native_histogram(input_schema, value);
3884
3885 match scalar_func.clone() {
3886 ScalarFunc::DataFusionBuiltin(func) => {
3887 if alternative_samples && value_is_histogram {
3888 continue;
3889 }
3890 other_input_exprs.insert(field_column_pos, col_expr);
3891 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
3892 func,
3893 args: other_input_exprs.clone().into(),
3894 });
3895 exprs.push(fn_expr);
3896 let _ = other_input_exprs.remove(field_column_pos);
3897 }
3898 ScalarFunc::DataFusionUdf(func) => {
3899 if alternative_samples && value_is_histogram {
3900 continue;
3901 }
3902 let args = itertools::chain!(
3903 other_input_exprs.iter().take(field_column_pos).cloned(),
3904 std::iter::once(col_expr),
3905 other_input_exprs.iter().skip(field_column_pos).cloned()
3906 )
3907 .collect_vec();
3908 exprs.push(DfExpr::ScalarFunction(ScalarFunction { func, args }))
3909 }
3910 ScalarFunc::NativeHistogramUdf(func) => {
3911 if value_is_histogram {
3912 let args = itertools::chain!(
3913 other_input_exprs.iter().take(field_column_pos).cloned(),
3914 std::iter::once(col_expr),
3915 other_input_exprs.iter().skip(field_column_pos).cloned()
3916 )
3917 .collect_vec();
3918 exprs.push(DfExpr::ScalarFunction(ScalarFunction { func, args }));
3919 } else if !alternative_samples {
3920 exprs.push(
3921 DfExpr::Literal(ScalarValue::Float64(None), None).alias(format!(
3922 "{}_{}",
3923 func.name(),
3924 value
3925 )),
3926 );
3927 }
3928 }
3929 ScalarFunc::Udf(func) => {
3930 let ts_range_expr = DfExpr::Column(Column::from_name(
3931 RangeManipulate::build_timestamp_range_name(
3932 self.ctx.time_index_column.as_ref().unwrap(),
3933 ),
3934 ));
3935 other_input_exprs.insert(field_column_pos, ts_range_expr);
3936 other_input_exprs.insert(field_column_pos + 1, col_expr);
3937 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
3938 func,
3939 args: other_input_exprs.clone().into(),
3940 });
3941 exprs.push(fn_expr);
3942 let _ = other_input_exprs.remove(field_column_pos + 1);
3943 let _ = other_input_exprs.remove(field_column_pos);
3944 }
3945 ScalarFunc::ExtrapolateUdf(udf, range_length) => {
3946 let ts_range_expr = DfExpr::Column(Column::from_name(
3947 RangeManipulate::build_timestamp_range_name(
3948 self.ctx.time_index_column.as_ref().unwrap(),
3949 ),
3950 ));
3951 other_input_exprs.insert(field_column_pos, ts_range_expr.clone());
3952 other_input_exprs.insert(field_column_pos + 1, col_expr.clone());
3953 other_input_exprs
3954 .insert(field_column_pos + 2, self.create_time_index_column_expr()?);
3955 other_input_exprs.push_back(lit(range_length));
3956 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
3957 func: udf,
3958 args: other_input_exprs.clone().into(),
3959 });
3960 let fn_expr = if matches!(func.name, "rate" | "increase")
3961 && !all_field_columns_are_native_histogram_ranges
3962 {
3963 let delta_sum = DfExpr::ScalarFunction(ScalarFunction {
3964 func: Arc::new(SumOverTime::scalar_udf()),
3965 args: vec![ts_range_expr, col_expr],
3966 });
3967 self.select_delta_range_math(
3968 func.name,
3969 input_schema,
3970 range_length,
3971 delta_sum,
3972 fn_expr,
3973 )?
3974 } else {
3975 fn_expr
3976 };
3977 exprs.push(fn_expr);
3978 let _ = other_input_exprs.pop_back();
3979 let _ = other_input_exprs.remove(field_column_pos + 2);
3980 let _ = other_input_exprs.remove(field_column_pos + 1);
3981 let _ = other_input_exprs.remove(field_column_pos);
3982 }
3983 ScalarFunc::GeneratedExpr => {}
3984 }
3985 }
3986
3987 if !matches!(func.name, "label_join" | "label_replace") {
3991 let mut new_field_columns = Vec::with_capacity(exprs.len());
3992
3993 exprs = exprs
3994 .into_iter()
3995 .map(|expr| {
3996 let display_name = expr.schema_name().to_string();
3997 new_field_columns.push(display_name.clone());
3998 Ok(expr.alias(display_name))
3999 })
4000 .collect::<std::result::Result<Vec<_>, _>>()
4001 .context(DataFusionPlanningSnafu)?;
4002
4003 self.ctx.field_columns = new_field_columns;
4004 }
4005
4006 Ok((exprs, new_tags))
4007 }
4008
4009 fn select_delta_range_math(
4010 &self,
4011 function: &str,
4012 input_schema: &DFSchemaRef,
4013 range_length: Millisecond,
4014 delta_sum: DfExpr,
4015 cumulative: DfExpr,
4016 ) -> Result<DfExpr> {
4017 let marker_is_delta = if self
4018 .ctx
4019 .tag_columns
4020 .iter()
4021 .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL)
4022 {
4023 Self::field_column_type(input_schema, OTLP_AGGREGATION_TEMPORALITY_LABEL)
4024 .filter(|data_type| Self::string_value_data_type(data_type).is_some())
4025 .map(|_| {
4026 DfExpr::Column(Column::from_name(OTLP_AGGREGATION_TEMPORALITY_LABEL))
4027 .eq(lit(GREPTIME_TEMPORALITY_DELTA))
4028 })
4029 } else {
4030 None
4031 };
4032 let Some(marker_is_delta) = marker_is_delta else {
4033 return Ok(cumulative);
4034 };
4035
4036 let delta = if function == "rate" {
4037 DfExpr::BinaryExpr(BinaryExpr {
4038 left: Box::new(delta_sum),
4039 op: Operator::Divide,
4040 right: Box::new(lit(range_length as f64 / 1000.0)),
4041 })
4042 } else {
4043 delta_sum
4044 };
4045 let display_name = cumulative.schema_name().to_string();
4046 when(marker_is_delta, delta)
4047 .otherwise(cumulative)
4048 .context(DataFusionPlanningSnafu)
4049 .map(|expr| expr.alias(display_name))
4050 }
4051
4052 fn validate_label_name(label_name: &str) -> Result<()> {
4056 if label_name.starts_with("__") {
4058 return InvalidDestinationLabelNameSnafu { label_name }.fail();
4059 }
4060 if !LABEL_NAME_REGEX.is_match(label_name) {
4062 return InvalidDestinationLabelNameSnafu { label_name }.fail();
4063 }
4064
4065 Ok(())
4066 }
4067
4068 fn build_regexp_replace_label_expr(
4070 &self,
4071 other_input_exprs: &mut VecDeque<DfExpr>,
4072 query_engine_state: &QueryEngineState,
4073 ) -> Result<Option<(DfExpr, String)>> {
4074 let dst_label = match other_input_exprs.pop_front() {
4076 Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4077 other => UnexpectedPlanExprSnafu {
4078 desc: format!("expected dst_label string literal, but found {:?}", other),
4079 }
4080 .fail()?,
4081 };
4082
4083 Self::validate_label_name(&dst_label)?;
4085 let replacement = match other_input_exprs.pop_front() {
4086 Some(DfExpr::Literal(ScalarValue::Utf8(Some(r)), _)) => r,
4087 other => UnexpectedPlanExprSnafu {
4088 desc: format!("expected replacement string literal, but found {:?}", other),
4089 }
4090 .fail()?,
4091 };
4092 let src_label = match other_input_exprs.pop_front() {
4093 Some(DfExpr::Literal(ScalarValue::Utf8(Some(s)), None)) => s,
4094 other => UnexpectedPlanExprSnafu {
4095 desc: format!("expected src_label string literal, but found {:?}", other),
4096 }
4097 .fail()?,
4098 };
4099
4100 let regex = match other_input_exprs.pop_front() {
4101 Some(DfExpr::Literal(ScalarValue::Utf8(Some(r)), None)) => r,
4102 other => UnexpectedPlanExprSnafu {
4103 desc: format!("expected regex string literal, but found {:?}", other),
4104 }
4105 .fail()?,
4106 };
4107
4108 regex::Regex::new(®ex).map_err(|_| {
4111 InvalidRegularExpressionSnafu {
4112 regex: regex.clone(),
4113 }
4114 .build()
4115 })?;
4116
4117 if self.ctx.tag_columns.contains(&src_label) && regex.is_empty() {
4119 return Ok(None);
4120 }
4121
4122 if !self.ctx.tag_columns.contains(&src_label) {
4124 if replacement.is_empty() {
4125 return Ok(None);
4127 } else {
4128 return Ok(Some((
4130 lit(replacement).alias(&dst_label),
4132 dst_label,
4133 )));
4134 }
4135 }
4136
4137 let regex = format!("^(?s:{regex})$");
4140
4141 let session_state = query_engine_state.session_state();
4142 let func = session_state
4143 .scalar_functions()
4144 .get("regexp_replace")
4145 .context(UnsupportedExprSnafu {
4146 name: "regexp_replace",
4147 })?;
4148
4149 let args = vec![
4151 if src_label.is_empty() {
4152 DfExpr::Literal(ScalarValue::Utf8(Some(String::new())), None)
4153 } else {
4154 DfExpr::Column(Column::from_name(src_label))
4155 },
4156 DfExpr::Literal(ScalarValue::Utf8(Some(regex)), None),
4157 DfExpr::Literal(ScalarValue::Utf8(Some(replacement)), None),
4158 ];
4159
4160 Ok(Some((
4161 DfExpr::ScalarFunction(ScalarFunction {
4162 func: func.clone(),
4163 args,
4164 })
4165 .alias(&dst_label),
4166 dst_label,
4167 )))
4168 }
4169
4170 fn build_concat_labels_expr(
4172 other_input_exprs: &mut VecDeque<DfExpr>,
4173 ctx: &PromPlannerContext,
4174 query_engine_state: &QueryEngineState,
4175 ) -> Result<(DfExpr, String)> {
4176 let dst_label = match other_input_exprs.pop_front() {
4179 Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4180 other => UnexpectedPlanExprSnafu {
4181 desc: format!("expected dst_label string literal, but found {:?}", other),
4182 }
4183 .fail()?,
4184 };
4185 let separator = match other_input_exprs.pop_front() {
4186 Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4187 other => UnexpectedPlanExprSnafu {
4188 desc: format!("expected separator string literal, but found {:?}", other),
4189 }
4190 .fail()?,
4191 };
4192
4193 let available_columns: HashSet<&str> = ctx
4195 .tag_columns
4196 .iter()
4197 .chain(ctx.field_columns.iter())
4198 .chain(ctx.time_index_column.as_ref())
4199 .map(|s| s.as_str())
4200 .collect();
4201
4202 let src_labels = other_input_exprs
4203 .iter()
4204 .map(|expr| {
4205 match expr {
4207 DfExpr::Literal(ScalarValue::Utf8(Some(label)), None) => {
4208 if label.is_empty() {
4209 Ok(DfExpr::Literal(ScalarValue::Null, None))
4210 } else if available_columns.contains(label.as_str()) {
4211 Ok(DfExpr::Column(Column::from_name(label)))
4213 } else {
4214 Ok(DfExpr::Literal(ScalarValue::Null, None))
4216 }
4217 }
4218 other => UnexpectedPlanExprSnafu {
4219 desc: format!(
4220 "expected source label string literal, but found {:?}",
4221 other
4222 ),
4223 }
4224 .fail(),
4225 }
4226 })
4227 .collect::<Result<Vec<_>>>()?;
4228 ensure!(
4229 !src_labels.is_empty(),
4230 FunctionInvalidArgumentSnafu {
4231 fn_name: "label_join"
4232 }
4233 );
4234
4235 let session_state = query_engine_state.session_state();
4236 let func = session_state
4237 .scalar_functions()
4238 .get("concat_ws")
4239 .context(UnsupportedExprSnafu { name: "concat_ws" })?;
4240
4241 let mut args = Vec::with_capacity(1 + src_labels.len());
4243 args.push(DfExpr::Literal(ScalarValue::Utf8(Some(separator)), None));
4244 args.extend(src_labels);
4245
4246 Ok((
4247 DfExpr::ScalarFunction(ScalarFunction {
4248 func: func.clone(),
4249 args,
4250 })
4251 .alias(&dst_label),
4252 dst_label,
4253 ))
4254 }
4255
4256 fn create_time_index_column_expr(&self) -> Result<DfExpr> {
4257 Ok(DfExpr::Column(Column::from_name(
4258 self.ctx
4259 .time_index_column
4260 .clone()
4261 .with_context(|| TimeIndexNotFoundSnafu { table: "unknown" })?,
4262 )))
4263 }
4264
4265 fn create_tag_column_exprs(&self) -> Result<Vec<DfExpr>> {
4266 let mut result = Vec::with_capacity(self.ctx.tag_columns.len());
4267 for tag in &self.ctx.tag_columns {
4268 let expr = DfExpr::Column(Column::from_name(tag));
4269 result.push(expr);
4270 }
4271 Ok(result)
4272 }
4273
4274 fn create_field_column_exprs(&self) -> Result<Vec<DfExpr>> {
4275 let mut result = Vec::with_capacity(self.ctx.field_columns.len());
4276 for field in &self.ctx.field_columns {
4277 let expr = DfExpr::Column(Column::from_name(field));
4278 result.push(expr);
4279 }
4280 Ok(result)
4281 }
4282
4283 fn create_tag_and_time_index_column_sort_exprs(&self) -> Result<Vec<SortExpr>> {
4284 let mut result = self
4285 .ctx
4286 .tag_columns
4287 .iter()
4288 .map(|col| DfExpr::Column(Column::from_name(col)).sort(true, true))
4289 .collect::<Vec<_>>();
4290 result.push(self.create_time_index_column_expr()?.sort(true, true));
4291 Ok(result)
4292 }
4293
4294 fn create_field_columns_sort_exprs(&self, asc: bool) -> Vec<SortExpr> {
4295 self.ctx
4296 .field_columns
4297 .iter()
4298 .map(|col| DfExpr::Column(Column::from_name(col)).sort(asc, true))
4299 .collect::<Vec<_>>()
4300 }
4301
4302 fn create_sort_exprs_by_tags(
4303 func: &str,
4304 tags: Vec<DfExpr>,
4305 asc: bool,
4306 ) -> Result<Vec<SortExpr>> {
4307 ensure!(
4308 !tags.is_empty(),
4309 FunctionInvalidArgumentSnafu { fn_name: func }
4310 );
4311
4312 tags.iter()
4313 .map(|col| match col {
4314 DfExpr::Literal(ScalarValue::Utf8(Some(label)), _) => {
4315 Ok(DfExpr::Column(Column::from_name(label)).sort(asc, false))
4316 }
4317 other => UnexpectedPlanExprSnafu {
4318 desc: format!("expected label string literal, but found {:?}", other),
4319 }
4320 .fail(),
4321 })
4322 .collect::<Result<Vec<_>>>()
4323 }
4324
4325 fn create_empty_values_filter_expr(&self, preserve_any_value: bool) -> Result<DfExpr> {
4326 let mut exprs = Vec::with_capacity(self.ctx.field_columns.len());
4327 for value in &self.ctx.field_columns {
4328 let expr = DfExpr::Column(Column::from_name(value)).is_not_null();
4329 exprs.push(expr);
4330 }
4331
4332 let predicate = if preserve_any_value {
4337 disjunction(exprs)
4338 } else {
4339 conjunction(exprs)
4340 };
4341 predicate.with_context(|| ValueNotFoundSnafu {
4342 table: self
4343 .table_ref()
4344 .map(|t| t.to_quoted_string())
4345 .unwrap_or_else(|_| "unknown".to_string()),
4346 })
4347 }
4348
4349 fn create_aggregate_exprs(
4365 &mut self,
4366 op: TokenType,
4367 param: &Option<Box<PromExpr>>,
4368 input_plan: &LogicalPlan,
4369 ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4370 let mixed_sample_columns =
4371 Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns)
4372 .map(|(float, histogram)| (float.to_string(), histogram.to_string()));
4373 let is_group_agg = op.id() == token::T_GROUP;
4374 if is_group_agg && mixed_sample_columns.is_none() {
4375 ensure!(
4376 self.ctx.field_columns.len() == 1,
4377 MultiFieldsNotSupportedSnafu {
4378 operator: "group()"
4379 }
4380 );
4381 }
4382
4383 if let Some((float, histogram)) = mixed_sample_columns {
4384 return self.create_mixed_aggregate_exprs(op, param, &float, &histogram);
4385 }
4386
4387 if self.all_field_columns_are_native_histograms(input_plan.schema()) {
4388 return self.create_native_histogram_aggregate_exprs(op, input_plan);
4389 }
4390
4391 let exprs = self
4393 .ctx
4394 .field_columns
4395 .iter()
4396 .map(|col| {
4397 Self::create_numeric_aggregate_expr(
4398 op,
4399 param,
4400 DfExpr::Column(Column::from_name(col)),
4401 )
4402 })
4403 .collect::<Result<Vec<_>>>()?;
4404
4405 let prev_field_exprs = if op.id() == token::T_COUNT_VALUES {
4407 let prev_field_exprs: Vec<_> = self
4408 .ctx
4409 .field_columns
4410 .iter()
4411 .map(|col| DfExpr::Column(Column::from_name(col)))
4412 .collect();
4413
4414 ensure!(
4415 self.ctx.field_columns.len() == 1,
4416 UnsupportedExprSnafu {
4417 name: "count_values on multi-value input"
4418 }
4419 );
4420
4421 prev_field_exprs
4422 } else {
4423 vec![]
4424 };
4425
4426 let mut new_field_columns = Vec::with_capacity(self.ctx.field_columns.len());
4428
4429 let normalized_exprs =
4430 normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
4431 for expr in normalized_exprs {
4432 new_field_columns.push(expr.schema_name().to_string());
4433 }
4434 self.ctx.field_columns = new_field_columns;
4435
4436 Ok((exprs, prev_field_exprs))
4437 }
4438
4439 fn create_numeric_aggregate_expr(
4440 op: TokenType,
4441 param: &Option<Box<PromExpr>>,
4442 input: DfExpr,
4443 ) -> Result<DfExpr> {
4444 let expr = match op.id() {
4445 token::T_SUM => sum_udaf().call(vec![input]),
4446 token::T_QUANTILE => {
4447 let q = Self::get_param_as_literal_expr(
4448 param.as_deref(),
4449 Some(op),
4450 Some(ArrowDataType::Float64),
4451 )?;
4452 quantile_udaf().call(vec![q, input])
4453 }
4454 token::T_AVG => avg_udaf().call(vec![input]),
4455 token::T_COUNT_VALUES | token::T_COUNT => count_udaf().call(vec![input]),
4456 token::T_MIN => min_udaf().call(vec![input]),
4457 token::T_MAX => max_udaf().call(vec![input]),
4458 token::T_GROUP => max_udaf().call(vec![lit(1_f64)]),
4461 token::T_STDDEV => stddev_pop_udaf().call(vec![input]),
4462 token::T_STDVAR => var_pop_udaf().call(vec![input]),
4463 token::T_TOPK | token::T_BOTTOMK => {
4464 return UnsupportedExprSnafu {
4465 name: format!("{op:?}"),
4466 }
4467 .fail();
4468 }
4469 _ => return UnexpectedTokenSnafu { token: op }.fail(),
4470 };
4471 Ok(expr)
4472 }
4473
4474 fn create_mixed_aggregate_exprs(
4475 &mut self,
4476 op: TokenType,
4477 param: &Option<Box<PromExpr>>,
4478 float_column: &str,
4479 histogram_column: &str,
4480 ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4481 let float_input = DfExpr::Column(Column::from_name(float_column));
4482 let histogram_input = DfExpr::Column(Column::from_name(histogram_column));
4483 let float_count = count_udaf().call(vec![float_input.clone()]);
4484 let histogram_count = count_udaf().call(vec![histogram_input.clone()]);
4485 let mixed_sample_value = || {
4486 DfExpr::ScalarFunction(ScalarFunction {
4487 func: coalesce(),
4488 args: vec![
4489 DfExpr::ScalarFunction(ScalarFunction {
4490 func: Arc::new(PromqlFloatToString::scalar_udf()),
4491 args: vec![float_input.clone()],
4492 }),
4493 DfExpr::ScalarFunction(ScalarFunction {
4494 func: Arc::new(NativeHistogramToString::scalar_udf()),
4495 args: vec![histogram_input.clone()],
4496 }),
4497 ],
4498 })
4499 };
4500
4501 let (exprs, prev_field_exprs, field_columns) = match op.id() {
4502 token::T_SUM | token::T_AVG => (
4503 vec![
4504 Self::create_numeric_aggregate_expr(op, param, float_input)?
4505 .alias(float_column),
4506 self.create_native_histogram_aggregate_expr(op, histogram_column)?,
4507 float_count.alias(Self::mixed_sample_count_name(float_column)),
4508 histogram_count.alias(Self::mixed_sample_count_name(histogram_column)),
4509 ],
4510 vec![],
4511 vec![float_column.to_string(), histogram_column.to_string()],
4512 ),
4513 token::T_COUNT => {
4514 let present = when(
4515 float_input
4516 .clone()
4517 .is_not_null()
4518 .or(histogram_input.clone().is_not_null()),
4519 lit(1_i64),
4520 )
4521 .otherwise(lit(ScalarValue::Int64(None)))
4522 .context(DataFusionPlanningSnafu)?;
4523 (
4524 vec![count_udaf().call(vec![present]).alias(float_column)],
4525 vec![],
4526 vec![float_column.to_string()],
4527 )
4528 }
4529 token::T_GROUP => (
4530 vec![max_udaf().call(vec![lit(1_f64)]).alias(float_column)],
4531 vec![],
4532 vec![float_column.to_string()],
4533 ),
4534 token::T_COUNT_VALUES => {
4535 let value = mixed_sample_value();
4536 (
4537 vec![count_udaf().call(vec![value.clone()]).alias(float_column)],
4538 vec![value],
4539 vec![float_column.to_string()],
4540 )
4541 }
4542 token::T_MIN | token::T_MAX | token::T_STDDEV | token::T_STDVAR | token::T_QUANTILE => {
4543 (
4544 vec![
4545 Self::create_numeric_aggregate_expr(op, param, float_input)?
4546 .alias(float_column),
4547 histogram_count.alias(Self::mixed_sample_count_name(histogram_column)),
4548 ],
4549 vec![],
4550 vec![float_column.to_string()],
4551 )
4552 }
4553 token::T_TOPK | token::T_BOTTOMK => {
4554 return UnsupportedExprSnafu {
4555 name: format!("{op:?}"),
4556 }
4557 .fail();
4558 }
4559 _ => return UnexpectedTokenSnafu { token: op }.fail(),
4560 };
4561
4562 self.ctx.field_columns = field_columns;
4563 Ok((exprs, prev_field_exprs))
4564 }
4565
4566 fn mixed_sample_count_column(column: &str) -> DfExpr {
4567 DfExpr::Column(Column::from_name(Self::mixed_sample_count_name(column)))
4568 }
4569
4570 fn mixed_sample_count_name(column: &str) -> String {
4571 format!("__promql_sample_count({column})")
4572 }
4573
4574 fn mixed_aggregate_filter_expr(
4575 &self,
4576 op: TokenType,
4577 float_column: &str,
4578 histogram_column: &str,
4579 ) -> Result<DfExpr> {
4580 let float_count = Self::mixed_sample_count_column(float_column);
4581 let histogram_count = Self::mixed_sample_count_column(histogram_column);
4582 let mixed = float_count
4583 .clone()
4584 .gt(lit(0_i64))
4585 .and(histogram_count.clone().gt(lit(0_i64)));
4586 let drop_mixed = DfExpr::ScalarFunction(ScalarFunction {
4587 func: Arc::new(NativeHistogramDrop::warning_bool_false_udf(
4588 format!(
4589 "{op}: dropped aggregation result containing both float and native histogram samples"
4590 ),
4591 self.promql_annotations.clone(),
4592 )),
4593 args: vec![float_count, histogram_count],
4594 });
4595
4596 when(mixed, drop_mixed)
4597 .otherwise(lit(true))
4598 .context(DataFusionPlanningSnafu)
4599 }
4600
4601 fn mixed_ignored_histogram_filter_expr(
4602 &self,
4603 op: TokenType,
4604 histogram_column: &str,
4605 ) -> Result<DfExpr> {
4606 let histogram_count = Self::mixed_sample_count_column(histogram_column);
4607 let has_histograms = histogram_count.clone().gt(lit(0_i64));
4608 let record_info = DfExpr::ScalarFunction(ScalarFunction {
4609 func: Arc::new(NativeHistogramDrop::bool_true_udf(
4610 format!(
4611 "{op}: dropped native histogram samples because this aggregation is not supported for native histograms"
4612 ),
4613 self.promql_annotations.clone(),
4614 )),
4615 args: vec![histogram_count],
4616 });
4617
4618 when(has_histograms, record_info)
4619 .otherwise(lit(true))
4620 .context(DataFusionPlanningSnafu)
4621 }
4622
4623 fn create_native_histogram_aggregate_expr(
4624 &self,
4625 op: TokenType,
4626 column: &str,
4627 ) -> Result<DfExpr> {
4628 let input = DfExpr::Column(Column::from_name(column));
4629 let expr = match op.id() {
4630 token::T_SUM => Arc::new(NativeHistogramAggSum::aggregate_udf_with_collector(
4631 self.promql_annotations.clone(),
4632 ))
4633 .call(vec![input])
4634 .alias(column),
4635 token::T_AVG => Arc::new(NativeHistogramAggAvg::aggregate_udf_with_collector(
4636 self.promql_annotations.clone(),
4637 ))
4638 .call(vec![input])
4639 .alias(column),
4640 token::T_COUNT_VALUES | token::T_COUNT => {
4641 count_udaf().call(vec![input]).alias(column)
4642 }
4643 token::T_GROUP => max_udaf().call(vec![lit(1_f64)]).alias(column),
4644 token::T_MIN
4645 | token::T_MAX
4646 | token::T_STDDEV
4647 | token::T_STDVAR
4648 | token::T_QUANTILE
4649 | token::T_TOPK
4650 | token::T_BOTTOMK => sum_udaf()
4651 .call(vec![DfExpr::ScalarFunction(ScalarFunction {
4652 func: Arc::new(NativeHistogramDrop::float_null_udf(
4653 format!(
4654 "{op}: dropped native histogram samples because this aggregation is not supported for native histograms"
4655 ),
4656 self.promql_annotations.clone(),
4657 )),
4658 args: vec![input],
4659 })])
4660 .alias(column),
4661 _ => return UnexpectedTokenSnafu { token: op }.fail(),
4662 };
4663 Ok(expr)
4664 }
4665
4666 fn create_native_histogram_aggregate_exprs(
4667 &mut self,
4668 op: TokenType,
4669 input_plan: &LogicalPlan,
4670 ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4671 let prev_field_exprs = if op.id() == token::T_COUNT_VALUES {
4672 ensure!(
4673 self.ctx.field_columns.len() == 1,
4674 UnsupportedExprSnafu {
4675 name: "count_values on multi-value input"
4676 }
4677 );
4678 self.ctx
4679 .field_columns
4680 .iter()
4681 .map(|col| {
4682 DfExpr::ScalarFunction(ScalarFunction {
4683 func: Arc::new(NativeHistogramToString::scalar_udf()),
4684 args: vec![DfExpr::Column(Column::from_name(col))],
4685 })
4686 })
4687 .collect::<Vec<_>>()
4688 } else {
4689 vec![]
4690 };
4691
4692 let exprs = self
4693 .ctx
4694 .field_columns
4695 .iter()
4696 .map(|col| self.create_native_histogram_aggregate_expr(op, col))
4697 .collect::<Result<Vec<_>>>()?;
4698
4699 let normalized_exprs =
4700 normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
4701 self.ctx.field_columns = normalized_exprs
4702 .into_iter()
4703 .map(|expr| expr.schema_name().to_string())
4704 .collect();
4705
4706 Ok((exprs, prev_field_exprs))
4707 }
4708
4709 fn get_param_value_as_str(op: TokenType, param: &Option<Box<PromExpr>>) -> Result<&str> {
4710 let param = param
4711 .as_deref()
4712 .with_context(|| FunctionInvalidArgumentSnafu {
4713 fn_name: op.to_string(),
4714 })?;
4715 let PromExpr::StringLiteral(StringLiteral { val }) = param else {
4716 return FunctionInvalidArgumentSnafu {
4717 fn_name: op.to_string(),
4718 }
4719 .fail();
4720 };
4721
4722 Ok(val)
4723 }
4724
4725 fn get_param_as_literal_expr(
4726 param: Option<&PromExpr>,
4727 op: Option<TokenType>,
4728 expected_type: Option<ArrowDataType>,
4729 ) -> Result<DfExpr> {
4730 let prom_param = param.with_context(|| {
4731 if let Some(op) = op {
4732 FunctionInvalidArgumentSnafu {
4733 fn_name: op.to_string(),
4734 }
4735 } else {
4736 FunctionInvalidArgumentSnafu {
4737 fn_name: "unknown".to_string(),
4738 }
4739 }
4740 })?;
4741
4742 let expr = Self::try_build_literal_expr(prom_param).with_context(|| {
4743 if let Some(op) = op {
4744 FunctionInvalidArgumentSnafu {
4745 fn_name: op.to_string(),
4746 }
4747 } else {
4748 FunctionInvalidArgumentSnafu {
4749 fn_name: "unknown".to_string(),
4750 }
4751 }
4752 })?;
4753
4754 if let Some(expected_type) = expected_type {
4756 let expr_type = expr
4758 .get_type(&DFSchema::empty())
4759 .context(DataFusionPlanningSnafu)?;
4760 if expected_type != expr_type {
4761 return FunctionInvalidArgumentSnafu {
4762 fn_name: format!("expected {expected_type:?}, but found {expr_type:?}"),
4763 }
4764 .fail();
4765 }
4766 }
4767
4768 Ok(expr)
4769 }
4770
4771 fn create_window_exprs(
4774 &mut self,
4775 op: TokenType,
4776 group_exprs: Vec<DfExpr>,
4777 input_plan: &LogicalPlan,
4778 ) -> Result<Vec<DfExpr>> {
4779 ensure!(
4780 self.ctx.field_columns.len() == 1,
4781 UnsupportedExprSnafu {
4782 name: "topk or bottomk on multi-value input"
4783 }
4784 );
4785
4786 assert!(matches!(op.id(), token::T_TOPK | token::T_BOTTOMK));
4787
4788 let asc = matches!(op.id(), token::T_BOTTOMK);
4789
4790 let tag_sort_exprs = self
4791 .create_tag_column_exprs()?
4792 .into_iter()
4793 .map(|expr| expr.sort(asc, true));
4794
4795 let exprs: Vec<DfExpr> = self
4797 .ctx
4798 .field_columns
4799 .iter()
4800 .map(|col| {
4801 let mut sort_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 1);
4802 sort_exprs.push(DfExpr::Column(Column::from(col)).sort(asc, true));
4804 sort_exprs.extend(tag_sort_exprs.clone());
4807
4808 DfExpr::WindowFunction(Box::new(WindowFunction {
4809 fun: WindowFunctionDefinition::WindowUDF(Arc::new(RowNumber::new().into())),
4810 params: WindowFunctionParams {
4811 args: vec![],
4812 partition_by: group_exprs.clone(),
4813 order_by: sort_exprs,
4814 window_frame: WindowFrame::new(Some(true)),
4815 null_treatment: None,
4816 distinct: false,
4817 filter: None,
4818 },
4819 }))
4820 })
4821 .collect();
4822
4823 let normalized_exprs =
4824 normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
4825 Ok(normalized_exprs)
4826 }
4827
4828 async fn create_histogram_plan(
4830 &mut self,
4831 function_name: &str,
4832 args: &PromFunctionArgs,
4833 query_engine_state: &QueryEngineState,
4834 ) -> Result<LogicalPlan> {
4835 let float_literal = |param: &PromExpr| -> Result<f64> {
4836 let value = (|| {
4837 let expr = Self::get_param_as_literal_expr(
4838 Some(param),
4839 None,
4840 Some(ArrowDataType::Float64),
4841 )
4842 .ok()?;
4843 let simplifier = ExprSimplifier::new(SimplifyContext::default());
4844 let expr = simplifier.coerce(expr, &DFSchema::empty()).ok()?;
4845 let DfExpr::Literal(value, _) = simplifier.simplify(expr).ok()? else {
4846 return None;
4847 };
4848 let ScalarValue::Float64(Some(value)) =
4849 value.cast_to(&ArrowDataType::Float64).ok()?
4850 else {
4851 return None;
4852 };
4853 Some(value)
4854 })()
4855 .with_context(|| FunctionInvalidArgumentSnafu {
4856 fn_name: function_name.to_string(),
4857 })?;
4858 Ok(value)
4859 };
4860 let (function, input) = match (function_name, args.args.as_slice()) {
4861 (SPECIAL_HISTOGRAM_QUANTILE, [quantile, input]) => (
4862 HistogramFoldOperation::Quantile(float_literal(quantile)?.into()),
4863 input.as_ref().clone(),
4864 ),
4865 (SPECIAL_HISTOGRAM_FRACTION, [lower, upper, input]) => (
4866 HistogramFoldOperation::Fraction {
4867 lower: float_literal(lower)?.into(),
4868 upper: float_literal(upper)?.into(),
4869 },
4870 input.as_ref().clone(),
4871 ),
4872 _ => {
4873 return FunctionInvalidArgumentSnafu {
4874 fn_name: function_name.to_string(),
4875 }
4876 .fail();
4877 }
4878 };
4879
4880 let input_plan = self.prom_expr_to_plan(&input, query_engine_state).await?;
4881 let input_plan = self.strip_tsid_column(input_plan)?;
4884 self.ctx.use_tsid = false;
4885
4886 if let Some((float_field, histogram_field)) =
4887 Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns)
4888 .map(|(float, histogram)| (float.to_string(), histogram.to_string()))
4889 {
4890 if self.ctx.has_le_tag() {
4891 return self.create_mixed_histogram_plan(
4892 function,
4893 input_plan,
4894 float_field,
4895 histogram_field,
4896 );
4897 }
4898 self.ctx.field_columns = vec![histogram_field];
4899 }
4900 if self.all_field_columns_are_native_histograms(input_plan.schema()) {
4901 return self.create_native_histogram_plan(function, input_plan);
4902 }
4903
4904 if !self.ctx.has_le_tag() {
4905 return Ok(LogicalPlan::EmptyRelation(
4908 datafusion::logical_expr::EmptyRelation {
4909 produce_one_row: false,
4910 schema: input_plan.schema().clone(),
4911 },
4912 ));
4913 }
4914 let time_index_column =
4915 self.ctx
4916 .time_index_column
4917 .clone()
4918 .with_context(|| TimeIndexNotFoundSnafu {
4919 table: self.ctx.table_name.clone().unwrap_or_default(),
4920 })?;
4921 let field_column = self
4923 .ctx
4924 .field_columns
4925 .first()
4926 .with_context(|| FunctionInvalidArgumentSnafu {
4927 fn_name: function.function_name().to_string(),
4928 })?
4929 .clone();
4930 self.ctx.tag_columns.retain(|col| col != LE_COLUMN_NAME);
4932
4933 let fold = HistogramFold::new_with_operation(
4934 LE_COLUMN_NAME.to_string(),
4935 field_column,
4936 time_index_column,
4937 function,
4938 None,
4939 input_plan,
4940 )
4941 .context(DataFusionPlanningSnafu)?;
4942 Ok(LogicalPlan::Extension(Extension {
4943 node: Arc::new(fold),
4944 }))
4945 }
4946
4947 fn create_native_histogram_expr(
4948 &self,
4949 function: HistogramFoldOperation,
4950 field_column: &str,
4951 ) -> DfExpr {
4952 let field = DfExpr::Column(Column::from_name(field_column));
4953 let (func, args) = match function {
4954 HistogramFoldOperation::Quantile(quantile) => (
4955 Arc::new(NativeHistogramQuantile::scalar_udf_with_collector(
4956 self.promql_annotations.clone(),
4957 )),
4958 vec![field, lit(f64::from(quantile))],
4959 ),
4960 HistogramFoldOperation::Fraction { lower, upper } => (
4961 Arc::new(NativeHistogramFraction::scalar_udf_with_collector(
4962 self.promql_annotations.clone(),
4963 )),
4964 vec![field, lit(f64::from(lower)), lit(f64::from(upper))],
4965 ),
4966 };
4967 DfExpr::ScalarFunction(ScalarFunction { func, args })
4968 }
4969
4970 fn create_native_histogram_plan(
4971 &mut self,
4972 function: HistogramFoldOperation,
4973 input_plan: LogicalPlan,
4974 ) -> Result<LogicalPlan> {
4975 ensure!(
4976 self.ctx.field_columns.len() == 1,
4977 MultiFieldsNotSupportedSnafu {
4978 operator: function.function_name()
4979 },
4980 );
4981
4982 let field_column = self.ctx.field_columns[0].clone();
4983 let function_expr = self.create_native_histogram_expr(function, &field_column);
4984 let display_name = function_expr.schema_name().to_string();
4985 self.ctx.field_columns = vec![display_name.clone()];
4986
4987 let project_exprs = std::iter::once(self.create_time_index_column_expr()?)
4988 .chain(std::iter::once(function_expr.alias(display_name)))
4989 .chain(self.create_tag_column_exprs()?)
4990 .collect::<Vec<_>>();
4991
4992 LogicalPlanBuilder::from(input_plan)
4993 .project(project_exprs)
4994 .context(DataFusionPlanningSnafu)?
4995 .filter(self.create_empty_values_filter_expr(false)?)
4996 .context(DataFusionPlanningSnafu)?
4997 .build()
4998 .context(DataFusionPlanningSnafu)
4999 }
5000
5001 fn create_mixed_histogram_plan(
5002 &mut self,
5003 function: HistogramFoldOperation,
5004 input_plan: LogicalPlan,
5005 float_field: String,
5006 histogram_field: String,
5007 ) -> Result<LogicalPlan> {
5008 let time_index_column =
5009 self.ctx
5010 .time_index_column
5011 .clone()
5012 .with_context(|| TimeIndexNotFoundSnafu {
5013 table: self.ctx.table_name.clone().unwrap_or_default(),
5014 })?;
5015 let tag_columns = self.ctx.tag_columns.clone();
5016 let folded = HistogramFold::new_with_operation(
5017 LE_COLUMN_NAME.to_string(),
5018 float_field.clone(),
5019 time_index_column.clone(),
5020 function,
5021 Some(histogram_field.clone()),
5022 input_plan,
5023 )
5024 .context(DataFusionPlanningSnafu)?;
5025 let record_collision = DfExpr::ScalarFunction(ScalarFunction {
5026 func: Arc::new(NativeHistogramDrop::warning_bool_false_udf(
5027 "vector contains a mix of classic and native histograms".to_string(),
5028 self.promql_annotations.clone(),
5029 )),
5030 args: vec![col(&float_field), col(&histogram_field)],
5031 });
5032 let keep = when(
5033 col(&float_field)
5034 .is_not_null()
5035 .and(col(&histogram_field).is_not_null()),
5036 record_collision,
5037 )
5038 .otherwise(lit(true))
5039 .context(DataFusionPlanningSnafu)?;
5040
5041 let native_expr = self.create_native_histogram_expr(function, &histogram_field);
5042 let output_field = native_expr.schema_name().to_string();
5043 let value = DfExpr::ScalarFunction(ScalarFunction {
5044 func: coalesce(),
5045 args: vec![col(&float_field), native_expr],
5046 });
5047 self.ctx.field_columns = vec![output_field.clone()];
5048 LogicalPlanBuilder::from(LogicalPlan::Extension(Extension {
5049 node: Arc::new(folded),
5050 }))
5051 .filter(keep)
5052 .context(DataFusionPlanningSnafu)?
5053 .project(
5054 std::iter::once(col(&time_index_column))
5055 .chain(std::iter::once(value.alias(output_field)))
5056 .chain(tag_columns.iter().map(col)),
5057 )
5058 .context(DataFusionPlanningSnafu)?
5059 .build()
5060 .context(DataFusionPlanningSnafu)
5061 }
5062
5063 async fn create_vector_plan(&mut self, args: &PromFunctionArgs) -> Result<LogicalPlan> {
5065 if args.args.len() != 1 {
5066 return FunctionInvalidArgumentSnafu {
5067 fn_name: SPECIAL_VECTOR_FUNCTION.to_string(),
5068 }
5069 .fail();
5070 }
5071 let lit = Self::get_param_as_literal_expr(Some(args.args[0].as_ref()), None, None)?;
5072
5073 self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
5075 self.ctx.reset_table_name_and_schema();
5076 self.ctx.tag_columns = vec![];
5077 self.ctx.field_columns = vec![greptime_value().to_string()];
5078 Ok(LogicalPlan::Extension(Extension {
5079 node: Arc::new(
5080 EmptyMetric::new(
5081 self.ctx.start,
5082 self.ctx.end,
5083 self.ctx.interval,
5084 SPECIAL_TIME_FUNCTION.to_string(),
5085 greptime_value().to_string(),
5086 Some(lit),
5087 )
5088 .context(DataFusionPlanningSnafu)?,
5089 ),
5090 }))
5091 }
5092
5093 async fn create_scalar_plan(
5095 &mut self,
5096 args: &PromFunctionArgs,
5097 query_engine_state: &QueryEngineState,
5098 ) -> Result<LogicalPlan> {
5099 ensure!(
5100 args.len() == 1,
5101 FunctionInvalidArgumentSnafu {
5102 fn_name: SCALAR_FUNCTION
5103 }
5104 );
5105 let input = self
5106 .prom_expr_to_plan(&args.args[0], query_engine_state)
5107 .await?;
5108 let input_schema = input.schema().clone();
5109 let alternative_samples =
5110 Self::field_columns_are_alternative_samples(&input_schema, &self.ctx.field_columns);
5111 let histogram_fields = self
5112 .ctx
5113 .field_columns
5114 .iter()
5115 .filter(|field| Self::field_column_is_native_histogram(&input_schema, field))
5116 .count();
5117 ensure!(
5118 self.ctx.field_columns.len() == 1 || alternative_samples,
5119 MultiFieldsNotSupportedSnafu {
5120 operator: SCALAR_FUNCTION
5121 },
5122 );
5123 let scalar_field = self
5124 .ctx
5125 .field_columns
5126 .iter()
5127 .find(|field| !Self::field_column_is_native_histogram(&input_schema, field))
5128 .or_else(|| self.ctx.field_columns.first())
5129 .cloned()
5130 .with_context(|| FunctionInvalidArgumentSnafu {
5131 fn_name: SCALAR_FUNCTION,
5132 })?;
5133 let input = if histogram_fields == self.ctx.field_columns.len() {
5134 LogicalPlanBuilder::from(input)
5137 .filter(lit(false))
5138 .context(DataFusionPlanningSnafu)?
5139 .build()
5140 .context(DataFusionPlanningSnafu)?
5141 } else if histogram_fields > 0 {
5142 LogicalPlanBuilder::from(input)
5144 .filter(DfExpr::Column(Column::from_name(&scalar_field)).is_not_null())
5145 .context(DataFusionPlanningSnafu)?
5146 .build()
5147 .context(DataFusionPlanningSnafu)?
5148 } else {
5149 input
5150 };
5151 let scalar_plan = LogicalPlan::Extension(Extension {
5152 node: Arc::new(
5153 ScalarCalculate::new(
5154 self.ctx.start,
5155 self.ctx.end,
5156 self.ctx.interval,
5157 input,
5158 self.ctx.time_index_column.as_ref().unwrap(),
5159 &self.ctx.tag_columns,
5160 &scalar_field,
5161 self.ctx.table_name.as_deref(),
5162 )
5163 .context(PromqlPlanNodeSnafu)?,
5164 ),
5165 });
5166 self.ctx.tag_columns.clear();
5168 self.ctx.field_columns.clear();
5169 self.ctx
5170 .field_columns
5171 .push(scalar_plan.schema().field(1).name().clone());
5172 Ok(scalar_plan)
5173 }
5174
5175 async fn create_absent_plan(
5177 &mut self,
5178 args: &PromFunctionArgs,
5179 query_engine_state: &QueryEngineState,
5180 ) -> Result<LogicalPlan> {
5181 if args.args.len() != 1 {
5182 return FunctionInvalidArgumentSnafu {
5183 fn_name: SPECIAL_ABSENT_FUNCTION.to_string(),
5184 }
5185 .fail();
5186 }
5187 let input = self
5188 .prom_expr_to_plan(&args.args[0], query_engine_state)
5189 .await?;
5190
5191 let time_index_expr = self.create_time_index_column_expr()?;
5192 let first_field_expr =
5193 self.create_field_column_exprs()?
5194 .pop()
5195 .with_context(|| ValueNotFoundSnafu {
5196 table: self.ctx.table_name.clone().unwrap_or_default(),
5197 })?;
5198 let first_value_expr = first_value(first_field_expr, vec![]);
5199
5200 let ordered_aggregated_input = LogicalPlanBuilder::from(input)
5201 .aggregate(
5202 vec![time_index_expr.clone()],
5203 vec![first_value_expr.clone()],
5204 )
5205 .context(DataFusionPlanningSnafu)?
5206 .sort(vec![time_index_expr.sort(true, false)])
5207 .context(DataFusionPlanningSnafu)?
5208 .build()
5209 .context(DataFusionPlanningSnafu)?;
5210
5211 let fake_labels = self
5212 .ctx
5213 .selector_matcher
5214 .iter()
5215 .filter_map(|matcher| match matcher.op {
5216 MatchOp::Equal => Some((matcher.name.clone(), matcher.value.clone())),
5217 _ => None,
5218 })
5219 .collect::<Vec<_>>();
5220
5221 let absent_plan = LogicalPlan::Extension(Extension {
5223 node: Arc::new(
5224 Absent::try_new(
5225 self.ctx.start,
5226 self.ctx.end,
5227 self.ctx.interval,
5228 self.ctx.time_index_column.as_ref().unwrap().clone(),
5229 self.ctx.field_columns[0].clone(),
5230 fake_labels,
5231 ordered_aggregated_input,
5232 )
5233 .context(DataFusionPlanningSnafu)?,
5234 ),
5235 });
5236
5237 Ok(absent_plan)
5238 }
5239
5240 fn try_build_literal_expr(expr: &PromExpr) -> Option<DfExpr> {
5243 match expr {
5244 PromExpr::NumberLiteral(NumberLiteral { val }) => Some(val.lit()),
5245 PromExpr::StringLiteral(StringLiteral { val }) => Some(val.lit()),
5246 PromExpr::VectorSelector(_)
5247 | PromExpr::MatrixSelector(_)
5248 | PromExpr::Extension(_)
5249 | PromExpr::Aggregate(_)
5250 | PromExpr::Subquery(_) => None,
5251 PromExpr::Call(Call { func, .. }) => {
5252 if func.name == SPECIAL_TIME_FUNCTION {
5253 None
5256 } else {
5257 None
5258 }
5259 }
5260 PromExpr::Paren(ParenExpr { expr }) => Self::try_build_literal_expr(expr),
5261 PromExpr::Unary(UnaryExpr { expr, .. }) => Some(DfExpr::Negative(Box::new(
5262 Self::try_build_literal_expr(expr)?,
5263 ))),
5264 PromExpr::Binary(PromBinaryExpr {
5265 lhs,
5266 rhs,
5267 op,
5268 modifier,
5269 }) => {
5270 let lhs = Self::try_build_literal_expr(lhs)?;
5271 let rhs = Self::try_build_literal_expr(rhs)?;
5272 let is_comparison_op = Self::is_token_a_comparison_op(*op);
5273 let expr_builder = Self::prom_token_to_binary_expr_builder(*op).ok()?;
5274 let expr = expr_builder(lhs, rhs).ok()?;
5275
5276 let should_return_bool = if let Some(m) = modifier {
5277 m.return_bool
5278 } else {
5279 false
5280 };
5281 if is_comparison_op && should_return_bool {
5282 Some(DfExpr::Cast(Cast {
5283 expr: Box::new(expr),
5284 data_type: ArrowDataType::Float64,
5285 }))
5286 } else {
5287 Some(expr)
5288 }
5289 }
5290 }
5291 }
5292
5293 fn try_build_special_time_expr_with_context(&self, expr: &PromExpr) -> Option<DfExpr> {
5294 match expr {
5295 PromExpr::Call(Call { func, .. }) => {
5296 if func.name == SPECIAL_TIME_FUNCTION
5297 && let Some(time_index_col) = self.ctx.time_index_column.as_ref()
5298 {
5299 Some(build_special_time_expr(time_index_col))
5300 } else {
5301 None
5302 }
5303 }
5304 _ => None,
5305 }
5306 }
5307
5308 fn native_histogram_binary_expr(
5309 token: TokenType,
5310 lhs: DfExpr,
5311 lhs_is_histogram: bool,
5312 rhs: DfExpr,
5313 rhs_is_histogram: bool,
5314 filter_context: bool,
5315 promql_annotations: Option<PromqlAnnotationCollector>,
5316 ) -> Result<Option<DfExpr>> {
5317 if !lhs_is_histogram && !rhs_is_histogram {
5318 return Ok(None);
5319 }
5320
5321 let scalar_fn = |func: ScalarUdfDef, args| {
5322 DfExpr::ScalarFunction(ScalarFunction {
5323 func: Arc::new(func),
5324 args,
5325 })
5326 };
5327 let invalid_expr = || {
5328 let message = format!(
5329 "{}: dropped native histogram samples because this binary operation is not supported for native histograms",
5330 token
5331 );
5332 let func = if filter_context {
5333 NativeHistogramDrop::bool_false_udf(message, promql_annotations.clone())
5334 } else {
5335 NativeHistogramDrop::float_null_udf(message, promql_annotations.clone())
5336 };
5337 let args = vec![lhs.clone(), rhs.clone()];
5338 scalar_fn(func, args)
5339 };
5340
5341 let expr = match (token.id(), lhs_is_histogram, rhs_is_histogram) {
5342 (token::T_ADD, true, true) => scalar_fn(
5343 NativeHistogramAdd::scalar_udf_with_collector(promql_annotations.clone()),
5344 vec![lhs, rhs],
5345 ),
5346 (token::T_SUB, true, true) => scalar_fn(
5347 NativeHistogramSub::scalar_udf_with_collector(promql_annotations.clone()),
5348 vec![lhs, rhs],
5349 ),
5350 (token::T_MUL, true, false) => {
5351 scalar_fn(NativeHistogramMulScalar::scalar_udf(), vec![lhs, rhs])
5352 }
5353 (token::T_MUL, false, true) => {
5354 scalar_fn(NativeHistogramScalarMul::scalar_udf(), vec![lhs, rhs])
5355 }
5356 (token::T_DIV, true, false) => {
5357 scalar_fn(NativeHistogramDivScalar::scalar_udf(), vec![lhs, rhs])
5358 }
5359 (token::T_EQLC, true, true) => {
5360 scalar_fn(NativeHistogramEq::scalar_udf(), vec![lhs, rhs])
5361 }
5362 (token::T_NEQ, true, true) => {
5363 scalar_fn(NativeHistogramNotEq::scalar_udf(), vec![lhs, rhs])
5364 }
5365 _ => invalid_expr(),
5366 };
5367
5368 Ok(Some(expr))
5369 }
5370
5371 #[allow(clippy::type_complexity)]
5374 fn prom_token_to_binary_expr_builder(
5375 token: TokenType,
5376 ) -> Result<Box<dyn Fn(DfExpr, DfExpr) -> Result<DfExpr>>> {
5377 let cast_float = |expr| {
5378 if matches!(
5379 &expr,
5380 DfExpr::Cast(Cast {
5381 data_type: ArrowDataType::Float64,
5382 ..
5383 })
5384 ) || matches!(&expr, DfExpr::Literal(ScalarValue::Float64(_), _))
5385 {
5386 expr
5387 } else {
5388 DfExpr::Cast(Cast {
5389 expr: Box::new(expr),
5390 data_type: ArrowDataType::Float64,
5391 })
5392 }
5393 };
5394 match token.id() {
5395 token::T_ADD => Ok(Box::new(move |lhs, rhs| {
5396 Ok(cast_float(lhs) + cast_float(rhs))
5397 })),
5398 token::T_SUB => Ok(Box::new(move |lhs, rhs| {
5399 Ok(cast_float(lhs) - cast_float(rhs))
5400 })),
5401 token::T_MUL => Ok(Box::new(move |lhs, rhs| {
5402 Ok(cast_float(lhs) * cast_float(rhs))
5403 })),
5404 token::T_DIV => Ok(Box::new(move |lhs, rhs| {
5405 Ok(cast_float(lhs) / cast_float(rhs))
5406 })),
5407 token::T_MOD => Ok(Box::new(move |lhs: DfExpr, rhs| {
5408 Ok(cast_float(lhs) % cast_float(rhs))
5409 })),
5410 token::T_EQLC => Ok(Box::new(|lhs, rhs| Ok(lhs.eq(rhs)))),
5411 token::T_NEQ => Ok(Box::new(|lhs, rhs| Ok(lhs.not_eq(rhs)))),
5412 token::T_GTR => Ok(Box::new(|lhs, rhs| Ok(lhs.gt(rhs)))),
5413 token::T_LSS => Ok(Box::new(|lhs, rhs| Ok(lhs.lt(rhs)))),
5414 token::T_GTE => Ok(Box::new(|lhs, rhs| Ok(lhs.gt_eq(rhs)))),
5415 token::T_LTE => Ok(Box::new(|lhs, rhs| Ok(lhs.lt_eq(rhs)))),
5416 token::T_POW => Ok(Box::new(move |lhs, rhs| {
5417 Ok(DfExpr::ScalarFunction(ScalarFunction {
5418 func: datafusion_functions::math::power(),
5419 args: vec![cast_float(lhs), cast_float(rhs)],
5420 }))
5421 })),
5422 token::T_ATAN2 => Ok(Box::new(move |lhs, rhs| {
5423 Ok(DfExpr::ScalarFunction(ScalarFunction {
5424 func: datafusion_functions::math::atan2(),
5425 args: vec![cast_float(lhs), cast_float(rhs)],
5426 }))
5427 })),
5428 _ => UnexpectedTokenSnafu { token }.fail(),
5429 }
5430 }
5431
5432 fn is_token_a_comparison_op(token: TokenType) -> bool {
5434 matches!(
5435 token.id(),
5436 token::T_EQLC
5437 | token::T_NEQ
5438 | token::T_GTR
5439 | token::T_LSS
5440 | token::T_GTE
5441 | token::T_LTE
5442 )
5443 }
5444
5445 fn is_token_a_set_op(token: TokenType) -> bool {
5447 matches!(
5448 token.id(),
5449 token::T_LAND | token::T_LOR | token::T_LUNLESS )
5453 }
5454
5455 fn align_binary_field_columns<'a>(
5456 left_schema: &DFSchemaRef,
5457 right_schema: &DFSchemaRef,
5458 left_field_columns: &'a [String],
5459 right_field_columns: &'a [String],
5460 op: TokenType,
5461 left_is_scalar: bool,
5462 right_is_scalar: bool,
5463 ) -> (
5464 Vec<(String, Vec<BinaryFieldPair<'a>>)>,
5465 Vec<BinaryFieldPair<'a>>,
5466 ) {
5467 let left_alternative = Self::alternative_sample_columns(left_schema, left_field_columns);
5470 let right_alternative = Self::alternative_sample_columns(right_schema, right_field_columns);
5471 let alternative_alignment = match (left_alternative, right_alternative) {
5472 (Some(output_names), Some(_)) => Some((
5473 output_names,
5474 left_field_columns
5475 .iter()
5476 .flat_map(|left| right_field_columns.iter().map(move |right| (left, right)))
5477 .collect::<Vec<_>>(),
5478 )),
5479 (Some(output_names), None) if right_field_columns.len() == 1 => Some((
5480 output_names,
5481 left_field_columns
5482 .iter()
5483 .map(|left| (left, &right_field_columns[0]))
5484 .collect::<Vec<_>>(),
5485 )),
5486 (None, Some(output_names)) if left_field_columns.len() == 1 => Some((
5487 output_names,
5488 right_field_columns
5489 .iter()
5490 .map(|right| (&left_field_columns[0], right))
5491 .collect::<Vec<_>>(),
5492 )),
5493 _ => None,
5494 };
5495 let mut invalid_pairs = Vec::new();
5496 if let Some(((float_output, histogram_output), field_pairs)) = alternative_alignment {
5497 let mut float_pairs = Vec::new();
5498 let mut histogram_pairs = Vec::new();
5499 for (left, right) in field_pairs {
5500 let left_is_histogram = Self::field_column_is_native_histogram(left_schema, left);
5501 let right_is_histogram =
5502 Self::field_column_is_native_histogram(right_schema, right);
5503 match Self::binary_result_is_histogram(op, left_is_histogram, right_is_histogram) {
5504 Some(false) => float_pairs.push((left, right)),
5505 Some(true) => histogram_pairs.push((left, right)),
5506 None => invalid_pairs.push((left, right)),
5507 }
5508 }
5509 if !float_pairs.is_empty() || !histogram_pairs.is_empty() {
5510 return (
5511 [
5512 (!float_pairs.is_empty()).then(|| (float_output.to_string(), float_pairs)),
5513 (!histogram_pairs.is_empty())
5514 .then(|| (histogram_output.to_string(), histogram_pairs)),
5515 ]
5516 .into_iter()
5517 .flatten()
5518 .collect(),
5519 invalid_pairs,
5520 );
5521 }
5522 }
5523
5524 if left_is_scalar && !right_is_scalar && left_field_columns.len() == 1 {
5525 return (
5526 right_field_columns
5527 .iter()
5528 .map(|right| (right.clone(), vec![(&left_field_columns[0], right)]))
5529 .collect(),
5530 invalid_pairs,
5531 );
5532 }
5533 if right_is_scalar && !left_is_scalar && right_field_columns.len() == 1 {
5534 return (
5535 left_field_columns
5536 .iter()
5537 .map(|left| (left.clone(), vec![(left, &right_field_columns[0])]))
5538 .collect(),
5539 invalid_pairs,
5540 );
5541 }
5542
5543 (
5544 left_field_columns
5545 .iter()
5546 .zip(right_field_columns.iter())
5547 .map(|(left, right)| (left.clone(), vec![(left, right)]))
5548 .collect(),
5549 invalid_pairs,
5550 )
5551 }
5552
5553 fn binary_result_is_histogram(
5554 token: TokenType,
5555 lhs_is_histogram: bool,
5556 rhs_is_histogram: bool,
5557 ) -> Option<bool> {
5558 match (token.id(), lhs_is_histogram, rhs_is_histogram) {
5559 (_, false, false) => Some(false),
5560 (token::T_ADD | token::T_SUB, true, true)
5561 | (token::T_MUL, true, false)
5562 | (token::T_MUL, false, true)
5563 | (token::T_DIV, true, false) => Some(true),
5564 (token::T_EQLC | token::T_NEQ, true, true) => Some(false),
5565 _ => None,
5566 }
5567 }
5568
5569 fn plan_has_tsid_column(plan: &LogicalPlan) -> bool {
5570 plan.schema()
5571 .fields()
5572 .iter()
5573 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
5574 }
5575
5576 fn is_empty_metric(plan: &LogicalPlan) -> bool {
5577 matches!(plan, LogicalPlan::Extension(Extension { node }) if node.as_any().is::<EmptyMetric>())
5578 }
5579
5580 fn native_histogram_arrow_type() -> ArrowDataType {
5581 native_histogram_value_type().as_arrow_type()
5582 }
5583
5584 fn field_column_type<'a>(
5585 schema: &'a DFSchemaRef,
5586 field_column: &str,
5587 ) -> Option<&'a ArrowDataType> {
5588 schema
5589 .index_of_column_by_name(None, field_column)
5590 .map(|idx| schema.field(idx).data_type())
5591 }
5592
5593 fn field_column_is_native_histogram(schema: &DFSchemaRef, field_column: &str) -> bool {
5594 Self::field_column_type(schema, field_column)
5595 .is_some_and(|data_type| data_type == &Self::native_histogram_arrow_type())
5596 }
5597
5598 fn field_columns_contain_native_histogram(
5599 schema: &DFSchemaRef,
5600 field_columns: &[String],
5601 ) -> bool {
5602 field_columns
5603 .iter()
5604 .any(|field| Self::field_column_is_native_histogram(schema, field))
5605 }
5606
5607 fn field_column_is_float_range(schema: &DFSchemaRef, field_column: &str) -> bool {
5608 Self::field_column_type(schema, field_column).is_some_and(|data_type| {
5609 matches!(
5610 data_type,
5611 ArrowDataType::Dictionary(key_type, value_type)
5612 if key_type.as_ref() == &ArrowDataType::Int64
5613 && value_type.as_ref() == &ArrowDataType::Float64
5614 )
5615 })
5616 }
5617
5618 fn field_columns_are_alternative_samples(
5619 schema: &DFSchemaRef,
5620 field_columns: &[String],
5621 ) -> bool {
5622 Self::alternative_sample_columns(schema, field_columns).is_some()
5623 }
5624
5625 fn alternative_sample_columns<'a>(
5626 schema: &DFSchemaRef,
5627 field_columns: &'a [String],
5628 ) -> Option<(&'a str, &'a str)> {
5629 if field_columns.len() != 2 {
5630 return None;
5631 }
5632
5633 let canonical_float = field_columns.iter().find(|field| {
5634 field.as_str() == greptime_value()
5635 && (Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64)
5636 || Self::field_column_is_float_range(schema, field))
5637 });
5638 let canonical_histogram = field_columns.iter().find(|field| {
5639 field.as_str() == greptime_native_histogram()
5640 && (Self::field_column_is_native_histogram(schema, field)
5641 || Self::field_column_is_native_histogram_range(schema, field))
5642 });
5643 if let (Some(float), Some(histogram)) = (canonical_float, canonical_histogram) {
5644 return Some((float, histogram));
5645 }
5646
5647 let float = field_columns.iter().find(|field| {
5648 field.starts_with(OR_FLOAT_FIELD_PREFIX)
5649 && (Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64)
5650 || Self::field_column_is_float_range(schema, field))
5651 })?;
5652 let histogram = field_columns.iter().find(|field| {
5653 field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)
5654 && (Self::field_column_is_native_histogram(schema, field)
5655 || Self::field_column_is_native_histogram_range(schema, field))
5656 })?;
5657 Some((float, histogram))
5658 }
5659
5660 fn alternative_sample_range_columns<'a>(
5661 schema: &DFSchemaRef,
5662 field_columns: &'a [String],
5663 ) -> Option<(&'a str, &'a str)> {
5664 Self::alternative_sample_columns(schema, field_columns).filter(|(float, histogram)| {
5665 Self::field_column_is_float_range(schema, float)
5666 && Self::field_column_is_native_histogram_range(schema, histogram)
5667 })
5668 }
5669
5670 fn field_column_is_native_histogram_range(schema: &DFSchemaRef, field_column: &str) -> bool {
5671 Self::field_column_type(schema, field_column).is_some_and(|data_type| {
5672 matches!(
5673 data_type,
5674 ArrowDataType::Dictionary(key_type, value_type)
5675 if key_type.as_ref() == &ArrowDataType::Int64
5676 && value_type.as_ref() == &Self::native_histogram_arrow_type()
5677 )
5678 })
5679 }
5680
5681 fn all_field_columns_are_native_histograms(&self, schema: &DFSchemaRef) -> bool {
5682 !self.ctx.field_columns.is_empty()
5683 && self
5684 .ctx
5685 .field_columns
5686 .iter()
5687 .all(|field| Self::field_column_is_native_histogram(schema, field))
5688 }
5689
5690 fn all_field_columns_are_native_histogram_ranges(&self, schema: &DFSchemaRef) -> bool {
5691 !self.ctx.field_columns.is_empty()
5692 && self
5693 .ctx
5694 .field_columns
5695 .iter()
5696 .all(|field| Self::field_column_is_native_histogram_range(schema, field))
5697 }
5698
5699 fn optional_tsid_projection(
5700 schema: &DFSchemaRef,
5701 table_ref: Option<&TableReference>,
5702 keep_tsid: bool,
5703 ) -> Option<DfExpr> {
5704 keep_tsid.then_some(()).and_then(|_| {
5705 schema
5706 .qualified_field_with_name(table_ref, DATA_SCHEMA_TSID_COLUMN_NAME)
5707 .ok()
5708 .map(|field| DfExpr::Column(field.into()))
5709 })
5710 }
5711
5712 fn binary_join_key_columns(
5713 &self,
5714 left_schema: &DFSchemaRef,
5715 right_schema: &DFSchemaRef,
5716 left_context: &PromPlannerContext,
5717 right_context: &PromPlannerContext,
5718 only_join_time_index: bool,
5719 modifier: &Option<BinModifier>,
5720 ) -> Result<(BTreeSet<String>, BTreeSet<String>, bool)> {
5721 let has_tsid = |schema: &DFSchemaRef| {
5722 schema
5723 .fields()
5724 .iter()
5725 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
5726 };
5727 let use_tsid_join = !only_join_time_index
5728 && self.binary_modifier_preserves_tsid_join_key(left_context, right_context, modifier)
5729 && left_context.use_tsid
5730 && right_context.use_tsid
5731 && has_tsid(left_schema)
5732 && has_tsid(right_schema);
5733
5734 let (mut left_tag_columns, mut right_tag_columns) = if use_tsid_join {
5735 (
5736 BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]),
5737 BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]),
5738 )
5739 } else {
5740 if only_join_time_index {
5741 (BTreeSet::new(), BTreeSet::new())
5742 } else {
5743 (
5744 left_context
5745 .tag_columns
5746 .iter()
5747 .cloned()
5748 .collect::<BTreeSet<_>>(),
5749 right_context
5750 .tag_columns
5751 .iter()
5752 .cloned()
5753 .collect::<BTreeSet<_>>(),
5754 )
5755 }
5756 };
5757
5758 if !use_tsid_join
5759 && let Some(modifier) = modifier
5760 && let Some(matching) = &modifier.matching
5761 {
5762 match matching {
5763 LabelModifier::Include(on) => {
5764 let mask = on.labels.iter().cloned().collect::<BTreeSet<_>>();
5765 left_tag_columns = left_tag_columns.intersection(&mask).cloned().collect();
5766 right_tag_columns = right_tag_columns.intersection(&mask).cloned().collect();
5767 }
5768 LabelModifier::Exclude(ignoring) => {
5769 for label in &ignoring.labels {
5770 let _ = left_tag_columns.remove(label);
5771 let _ = right_tag_columns.remove(label);
5772 }
5773 }
5774 }
5775 }
5776
5777 let force_empty_join =
5778 !use_tsid_join && !only_join_time_index && left_tag_columns != right_tag_columns;
5779 if force_empty_join {
5780 let common_tag_columns = left_tag_columns
5781 .intersection(&right_tag_columns)
5782 .cloned()
5783 .collect::<BTreeSet<_>>();
5784 left_tag_columns = common_tag_columns.clone();
5785 right_tag_columns = common_tag_columns;
5786 }
5787
5788 Ok((left_tag_columns, right_tag_columns, force_empty_join))
5789 }
5790
5791 fn binary_modifier_preserves_tsid_join_key(
5792 &self,
5793 left_context: &PromPlannerContext,
5794 right_context: &PromPlannerContext,
5795 modifier: &Option<BinModifier>,
5796 ) -> bool {
5797 let Some(modifier) = modifier else {
5798 return true;
5799 };
5800
5801 if !matches!(modifier.card, VectorMatchCardinality::OneToOne) {
5802 return false;
5803 }
5804
5805 match &modifier.matching {
5806 None => true,
5807 Some(LabelModifier::Exclude(ignoring)) => ignoring.labels.iter().all(|label| {
5808 !left_context.tag_columns.contains(label)
5809 && !right_context.tag_columns.contains(label)
5810 }),
5811 Some(LabelModifier::Include(on)) => {
5812 let on_labels = on.labels.iter().cloned().collect::<BTreeSet<_>>();
5813 let left_labels = left_context
5814 .tag_columns
5815 .iter()
5816 .cloned()
5817 .collect::<BTreeSet<_>>();
5818 let right_labels = right_context
5819 .tag_columns
5820 .iter()
5821 .cloned()
5822 .collect::<BTreeSet<_>>();
5823
5824 on_labels == left_labels && on_labels == right_labels
5825 }
5826 }
5827 }
5828
5829 #[allow(clippy::too_many_arguments)]
5832 fn join_on_non_field_columns(
5833 &self,
5834 left: LogicalPlan,
5835 right: LogicalPlan,
5836 left_table_ref: TableReference,
5837 right_table_ref: TableReference,
5838 left_time_index_column: Option<String>,
5839 right_time_index_column: Option<String>,
5840 only_join_time_index: bool,
5841 modifier: &Option<BinModifier>,
5842 left_context: &PromPlannerContext,
5843 right_context: &PromPlannerContext,
5844 ) -> Result<LogicalPlan> {
5845 let (mut left_tag_columns, mut right_tag_columns, mut force_empty_join) = self
5846 .binary_join_key_columns(
5847 left.schema(),
5848 right.schema(),
5849 left_context,
5850 right_context,
5851 only_join_time_index,
5852 modifier,
5853 )?;
5854 let use_tsid_join = !only_join_time_index
5855 && !force_empty_join
5856 && left_tag_columns == BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()])
5857 && right_tag_columns == BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]);
5858 let (left, right) = if !only_join_time_index
5859 && !use_tsid_join
5860 && Self::only_temporality_match_label_mismatches(left_context, right_context, modifier)
5861 {
5862 let mut aligned_left_context = left_context.clone();
5863 let mut aligned_right_context = right_context.clone();
5864 let (left, right, _) = Self::align_temporality_match_column(
5865 left,
5866 right,
5867 &mut aligned_left_context,
5868 &mut aligned_right_context,
5869 )?;
5870 (left_tag_columns, right_tag_columns, force_empty_join) = self
5871 .binary_join_key_columns(
5872 left.schema(),
5873 right.schema(),
5874 &aligned_left_context,
5875 &aligned_right_context,
5876 false,
5877 modifier,
5878 )?;
5879 (left, right)
5880 } else {
5881 (left, right)
5882 };
5883
5884 if let (Some(left_time_index_column), Some(right_time_index_column)) =
5886 (left_time_index_column, right_time_index_column)
5887 {
5888 left_tag_columns.insert(left_time_index_column);
5889 right_tag_columns.insert(right_time_index_column);
5890 }
5891
5892 let right = LogicalPlanBuilder::from(right)
5893 .alias(right_table_ref)
5894 .context(DataFusionPlanningSnafu)?
5895 .build()
5896 .context(DataFusionPlanningSnafu)?;
5897
5898 LogicalPlanBuilder::from(left)
5900 .alias(left_table_ref)
5901 .context(DataFusionPlanningSnafu)?
5902 .join_detailed(
5903 right,
5904 JoinType::Inner,
5905 (
5906 left_tag_columns
5907 .into_iter()
5908 .map(Column::from_name)
5909 .collect::<Vec<_>>(),
5910 right_tag_columns
5911 .into_iter()
5912 .map(Column::from_name)
5913 .collect::<Vec<_>>(),
5914 ),
5915 force_empty_join.then_some(lit(false)),
5916 NullEquality::NullEqualsNull,
5917 )
5918 .context(DataFusionPlanningSnafu)?
5919 .build()
5920 .context(DataFusionPlanningSnafu)
5921 }
5922
5923 fn selected_binary_match_labels(
5924 left_context: &PromPlannerContext,
5925 right_context: &PromPlannerContext,
5926 modifier: &Option<BinModifier>,
5927 ) -> BTreeSet<String> {
5928 let mut labels = left_context
5929 .tag_columns
5930 .iter()
5931 .chain(&right_context.tag_columns)
5932 .cloned()
5933 .collect::<BTreeSet<_>>();
5934 if let Some(matching) = modifier
5935 .as_ref()
5936 .and_then(|modifier| modifier.matching.as_ref())
5937 {
5938 match matching {
5939 LabelModifier::Include(on) => {
5940 labels = on
5941 .labels
5942 .iter()
5943 .filter(|label| {
5944 left_context.tag_columns.contains(label)
5945 || right_context.tag_columns.contains(label)
5946 })
5947 .cloned()
5948 .collect();
5949 }
5950 LabelModifier::Exclude(ignoring) => {
5951 for label in &ignoring.labels {
5952 labels.remove(label);
5953 }
5954 }
5955 }
5956 }
5957 labels
5958 }
5959
5960 fn only_temporality_match_label_mismatches(
5961 left_context: &PromPlannerContext,
5962 right_context: &PromPlannerContext,
5963 modifier: &Option<BinModifier>,
5964 ) -> bool {
5965 let mut mismatches =
5966 Self::selected_binary_match_labels(left_context, right_context, modifier)
5967 .into_iter()
5968 .filter(|label| {
5969 left_context.tag_columns.contains(label)
5970 != right_context.tag_columns.contains(label)
5971 });
5972 matches!(
5973 (mismatches.next(), mismatches.next()),
5974 (Some(label), None) if label == OTLP_AGGREGATION_TEMPORALITY_LABEL
5975 )
5976 }
5977
5978 fn align_temporality_match_column(
5979 mut left: LogicalPlan,
5980 mut right: LogicalPlan,
5981 left_context: &mut PromPlannerContext,
5982 right_context: &mut PromPlannerContext,
5983 ) -> Result<(LogicalPlan, LogicalPlan, bool)> {
5984 let marker = OTLP_AGGREGATION_TEMPORALITY_LABEL;
5985 let left_has_marker = left_context.tag_columns.iter().any(|tag| tag == marker);
5986 let (present, add_to_left) = if left_has_marker {
5987 (&left, false)
5988 } else {
5989 (&right, true)
5990 };
5991 let data_type = present
5992 .schema()
5993 .fields()
5994 .iter()
5995 .find(|field| field.name() == marker)
5996 .map(|field| field.data_type().clone())
5997 .with_context(|| ColumnNotFoundSnafu {
5998 col: marker.to_string(),
5999 })?;
6000 let null = Self::string_scalar_value(&data_type, None).with_context(|| {
6001 UnexpectedPlanExprSnafu {
6002 desc: format!("temporality match label {marker} must be a string"),
6003 }
6004 })?;
6005 let add_marker = |plan: LogicalPlan| {
6006 let visible = plan
6007 .schema()
6008 .iter()
6009 .map(|(qualifier, field)| {
6010 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6011 })
6012 .collect::<Vec<_>>();
6013 LogicalPlanBuilder::from(plan)
6014 .project(
6015 visible
6016 .into_iter()
6017 .chain([DfExpr::Literal(null, None).alias(marker)]),
6018 )
6019 .context(DataFusionPlanningSnafu)?
6020 .build()
6021 .context(DataFusionPlanningSnafu)
6022 };
6023
6024 if add_to_left {
6025 left = add_marker(left)?;
6026 left_context.tag_columns.push(marker.to_string());
6027 } else {
6028 right = add_marker(right)?;
6029 right_context.tag_columns.push(marker.to_string());
6030 }
6031 Ok((left, right, add_to_left))
6032 }
6033
6034 fn normalized_match_key_expr(
6035 label: &str,
6036 field: Option<(Option<TableReference>, ArrowDataType)>,
6037 value_type: &ArrowDataType,
6038 internal_name: &str,
6039 ) -> DfExpr {
6040 let empty = Self::string_scalar_value(value_type, Some(String::new()))
6041 .expect("match label value type is a string");
6042 let expr = if let Some((qualifier, data_type)) = field {
6043 let column = DfExpr::Column(Column::new(qualifier, label));
6044 let column = if &data_type == value_type {
6045 column
6046 } else {
6047 DfExpr::Cast(Cast {
6048 expr: Box::new(column),
6049 data_type: value_type.clone(),
6050 })
6051 };
6052 DfExpr::ScalarFunction(ScalarFunction {
6053 func: coalesce(),
6054 args: vec![column, DfExpr::Literal(empty, None)],
6055 })
6056 } else {
6057 DfExpr::Literal(empty, None)
6058 };
6059 expr.alias(internal_name)
6060 }
6061
6062 fn is_zero_row_empty_relation(plan: &LogicalPlan) -> bool {
6063 matches!(plan, LogicalPlan::EmptyRelation(relation) if !relation.produce_one_row)
6066 }
6067
6068 fn set_op_on_non_field_columns(
6070 &mut self,
6071 mut left: LogicalPlan,
6072 mut right: LogicalPlan,
6073 left_context: PromPlannerContext,
6074 right_context: PromPlannerContext,
6075 op: TokenType,
6076 modifier: &Option<BinModifier>,
6077 ) -> Result<LogicalPlan> {
6078 let left_tag_col_set = left_context
6079 .tag_columns
6080 .iter()
6081 .cloned()
6082 .collect::<HashSet<_>>();
6083 let right_tag_col_set = right_context
6084 .tag_columns
6085 .iter()
6086 .cloned()
6087 .collect::<HashSet<_>>();
6088
6089 if matches!(op.id(), token::T_LOR) {
6090 return self.or_operator(
6091 left,
6092 right,
6093 left_tag_col_set,
6094 right_tag_col_set,
6095 left_context,
6096 right_context,
6097 modifier,
6098 );
6099 }
6100
6101 if let Some(modifier) = modifier {
6102 ensure!(
6103 matches!(
6104 modifier.card,
6105 VectorMatchCardinality::OneToOne | VectorMatchCardinality::ManyToMany
6106 ),
6107 UnsupportedVectorMatchSnafu {
6108 name: modifier.card.clone(),
6109 },
6110 );
6111 }
6112
6113 let output_context = left_context.clone();
6114 let visible_left_schema = left.schema().clone();
6115 let mut left_context = left_context;
6116 let mut right_context = right_context;
6117 let added_marker_to_left = if Self::only_temporality_match_label_mismatches(
6118 &left_context,
6119 &right_context,
6120 modifier,
6121 ) {
6122 let aligned = Self::align_temporality_match_column(
6123 left,
6124 right,
6125 &mut left_context,
6126 &mut right_context,
6127 )?;
6128 left = aligned.0;
6129 right = aligned.1;
6130 aligned.2
6131 } else {
6132 false
6133 };
6134
6135 let mut left_tag_col_set = left_context
6136 .tag_columns
6137 .iter()
6138 .cloned()
6139 .collect::<BTreeSet<_>>();
6140 let mut right_tag_col_set = right_context
6141 .tag_columns
6142 .iter()
6143 .cloned()
6144 .collect::<BTreeSet<_>>();
6145 if let Some(matching) = modifier
6146 .as_ref()
6147 .and_then(|modifier| modifier.matching.as_ref())
6148 {
6149 match matching {
6150 LabelModifier::Include(on) => {
6151 let mask = on.labels.iter().cloned().collect::<BTreeSet<_>>();
6152 left_tag_col_set = left_tag_col_set.intersection(&mask).cloned().collect();
6153 right_tag_col_set = right_tag_col_set.intersection(&mask).cloned().collect();
6154 }
6155 LabelModifier::Exclude(ignoring) => {
6156 for label in &ignoring.labels {
6157 let _ = left_tag_col_set.remove(label);
6158 let _ = right_tag_col_set.remove(label);
6159 }
6160 }
6161 }
6162 }
6163 ensure!(
6164 left_tag_col_set == right_tag_col_set,
6165 CombineTableColumnMismatchSnafu {
6166 left: left_tag_col_set.iter().cloned().collect::<Vec<_>>(),
6167 right: right_tag_col_set.iter().cloned().collect::<Vec<_>>(),
6168 }
6169 );
6170
6171 let left_time_index = left_context.time_index_column.clone().unwrap();
6172 let right_time_index = right_context.time_index_column.clone().unwrap();
6173
6174 if left_context.time_index_column != right_context.time_index_column {
6176 let right_project_exprs = right
6177 .schema()
6178 .fields()
6179 .iter()
6180 .map(|field| {
6181 if field.name() == &right_time_index {
6182 DfExpr::Column(Column::from_name(&right_time_index)).alias(&left_time_index)
6183 } else {
6184 DfExpr::Column(Column::from_name(field.name()))
6185 }
6186 })
6187 .collect::<Vec<_>>();
6188
6189 right = LogicalPlanBuilder::from(right)
6190 .project(right_project_exprs)
6191 .context(DataFusionPlanningSnafu)?
6192 .build()
6193 .context(DataFusionPlanningSnafu)?;
6194 }
6195
6196 let join_keys = left_tag_col_set
6197 .into_iter()
6198 .chain([left_time_index])
6199 .collect::<Vec<_>>();
6200
6201 ensure!(
6202 left_context.field_columns.len() == 1
6203 || Self::field_columns_are_alternative_samples(
6204 left.schema(),
6205 &left_context.field_columns,
6206 ),
6207 MultiFieldsNotSupportedSnafu {
6208 operator: "AND/UNLESS operator"
6209 }
6210 );
6211 let result = match op.id() {
6214 token::T_LAND => LogicalPlanBuilder::from(left)
6215 .distinct()
6216 .context(DataFusionPlanningSnafu)?
6217 .join_detailed(
6218 right,
6219 JoinType::LeftSemi,
6220 (join_keys.clone(), join_keys),
6221 None,
6222 NullEquality::NullEqualsNull,
6223 )
6224 .context(DataFusionPlanningSnafu)?
6225 .build()
6226 .context(DataFusionPlanningSnafu),
6227 token::T_LUNLESS => LogicalPlanBuilder::from(left)
6228 .distinct()
6229 .context(DataFusionPlanningSnafu)?
6230 .join_detailed(
6231 right,
6232 JoinType::LeftAnti,
6233 (join_keys.clone(), join_keys),
6234 None,
6235 NullEquality::NullEqualsNull,
6236 )
6237 .context(DataFusionPlanningSnafu)?
6238 .build()
6239 .context(DataFusionPlanningSnafu),
6240 token::T_LOR => {
6241 unreachable!()
6244 }
6245 _ => UnexpectedTokenSnafu { token: op }.fail(),
6246 }?;
6247 let result = if added_marker_to_left {
6248 LogicalPlanBuilder::from(result)
6249 .project(visible_left_schema.iter().map(|(qualifier, field)| {
6250 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6251 }))
6252 .context(DataFusionPlanningSnafu)?
6253 .build()
6254 .context(DataFusionPlanningSnafu)?
6255 } else {
6256 result
6257 };
6258
6259 self.ctx = output_context;
6261 Ok(result)
6262 }
6263
6264 fn string_value_data_type(data_type: &ArrowDataType) -> Option<&ArrowDataType> {
6265 match data_type {
6266 data_type if data_type.is_string() => Some(data_type),
6267 ArrowDataType::Dictionary(_, value_type) if value_type.is_string() => Some(value_type),
6268 _ => None,
6269 }
6270 }
6271
6272 fn string_scalar_value(
6273 data_type: &ArrowDataType,
6274 value: Option<String>,
6275 ) -> Option<ScalarValue> {
6276 match data_type {
6277 ArrowDataType::Utf8 => Some(ScalarValue::Utf8(value)),
6278 ArrowDataType::LargeUtf8 => Some(ScalarValue::LargeUtf8(value)),
6279 ArrowDataType::Utf8View => Some(ScalarValue::Utf8View(value)),
6280 ArrowDataType::Dictionary(key_type, value_type) => Some(ScalarValue::Dictionary(
6281 key_type.clone(),
6282 Box::new(Self::string_scalar_value(value_type, value)?),
6283 )),
6284 _ => None,
6285 }
6286 }
6287
6288 fn common_label_data_type(
6289 left: Option<&ArrowDataType>,
6290 right: Option<&ArrowDataType>,
6291 ) -> Option<ArrowDataType> {
6292 match (left, right) {
6293 (Some(left), Some(right)) if left == right => {
6294 Self::string_value_data_type(left).map(|_| left.clone())
6295 }
6296 (Some(left), Some(right)) => {
6297 let left_value_type = Self::string_value_data_type(left)?;
6298 let right_value_type = Self::string_value_data_type(right)?;
6299 match (left_value_type, right_value_type) {
6302 (left, right) if left == right => Some(left.clone()),
6303 (ArrowDataType::LargeUtf8, _) | (_, ArrowDataType::LargeUtf8) => {
6304 Some(ArrowDataType::LargeUtf8)
6305 }
6306 (ArrowDataType::Utf8View, ArrowDataType::Utf8View) => {
6307 Some(ArrowDataType::Utf8View)
6308 }
6309 _ => Some(ArrowDataType::Utf8),
6310 }
6311 }
6312 (Some(data_type), None) | (None, Some(data_type)) => {
6313 Self::string_value_data_type(data_type).cloned()
6314 }
6315 (None, None) => Some(ArrowDataType::Utf8),
6316 }
6317 }
6318
6319 #[allow(clippy::too_many_arguments)]
6321 fn or_operator(
6322 &mut self,
6323 left: LogicalPlan,
6324 right: LogicalPlan,
6325 left_tag_cols_set: HashSet<String>,
6326 right_tag_cols_set: HashSet<String>,
6327 left_context: PromPlannerContext,
6328 right_context: PromPlannerContext,
6329 modifier: &Option<BinModifier>,
6330 ) -> Result<LogicalPlan> {
6331 let left_is_empty = Self::is_zero_row_empty_relation(&left);
6332 let right_is_empty = Self::is_zero_row_empty_relation(&right);
6333 match (left_is_empty, right_is_empty) {
6334 (true, false) => {
6335 self.ctx = right_context;
6336 return Ok(right);
6337 }
6338 (false, true) => {
6339 self.ctx = left_context;
6340 return Ok(left);
6341 }
6342 (true, true) => {
6343 self.ctx = left_context;
6344 return Ok(left);
6345 }
6346 (false, false) => {}
6347 }
6348
6349 ensure!(
6350 !left.schema().fields().is_empty() && !right.schema().fields().is_empty(),
6351 UnexpectedPlanExprSnafu {
6352 desc: "OR operator input has zero columns",
6353 }
6354 );
6355 let left_has_alternative_samples =
6356 Self::field_columns_are_alternative_samples(left.schema(), &left_context.field_columns);
6357 let right_has_alternative_samples = Self::field_columns_are_alternative_samples(
6358 right.schema(),
6359 &right_context.field_columns,
6360 );
6361 ensure!(
6362 left_context.field_columns.len() == 1 || left_has_alternative_samples,
6363 MultiFieldsNotSupportedSnafu {
6364 operator: "OR operator"
6365 }
6366 );
6367 ensure!(
6368 right_context.field_columns.len() == 1 || right_has_alternative_samples,
6369 MultiFieldsNotSupportedSnafu {
6370 operator: "OR operator"
6371 }
6372 );
6373
6374 let all_tags = left_tag_cols_set
6376 .union(&right_tag_cols_set)
6377 .cloned()
6378 .collect::<HashSet<_>>();
6379 let left_qualifier = left.schema().qualified_field(0).0.cloned();
6380 let right_qualifier = right.schema().qualified_field(0).0.cloned();
6381 let left_qualifier_string = left_qualifier
6382 .as_ref()
6383 .map(|l| l.to_string())
6384 .unwrap_or_default();
6385 let right_qualifier_string = right_qualifier
6386 .as_ref()
6387 .map(|r| r.to_string())
6388 .unwrap_or_default();
6389 let left_time_index_column =
6390 left_context
6391 .time_index_column
6392 .clone()
6393 .with_context(|| TimeIndexNotFoundSnafu {
6394 table: left_qualifier_string.clone(),
6395 })?;
6396 let right_time_index_column =
6397 right_context
6398 .time_index_column
6399 .clone()
6400 .with_context(|| TimeIndexNotFoundSnafu {
6401 table: right_qualifier_string.clone(),
6402 })?;
6403 let native_histogram_type = Self::native_histogram_arrow_type();
6404 let is_numeric = |data_type: &ArrowDataType| {
6405 matches!(
6406 data_type,
6407 ArrowDataType::Int8
6408 | ArrowDataType::Int16
6409 | ArrowDataType::Int32
6410 | ArrowDataType::Int64
6411 | ArrowDataType::UInt8
6412 | ArrowDataType::UInt16
6413 | ArrowDataType::UInt32
6414 | ArrowDataType::UInt64
6415 | ArrowDataType::Float32
6416 | ArrowDataType::Float64
6417 )
6418 };
6419 let left_fields = left_context
6420 .field_columns
6421 .iter()
6422 .map(|name| {
6423 left.schema()
6424 .iter()
6425 .find(|(_, field)| field.name() == name)
6426 .map(|(qualifier, field)| {
6427 (name.clone(), qualifier.cloned(), field.data_type().clone())
6428 })
6429 .with_context(|| ColumnNotFoundSnafu { col: name.clone() })
6430 })
6431 .collect::<Result<Vec<_>>>()?;
6432 let right_fields = right_context
6433 .field_columns
6434 .iter()
6435 .map(|name| {
6436 right
6437 .schema()
6438 .iter()
6439 .find(|(_, field)| field.name() == name)
6440 .map(|(qualifier, field)| {
6441 (name.clone(), qualifier.cloned(), field.data_type().clone())
6442 })
6443 .with_context(|| ColumnNotFoundSnafu { col: name.clone() })
6444 })
6445 .collect::<Result<Vec<_>>>()?;
6446 let left_field = &left_fields[0];
6447 let right_field = &right_fields[0];
6448 let left_field_col = &left_field.0;
6449 let right_field_col = &right_field.0;
6450 let fields_are_samples = |fields: &[(String, Option<TableReference>, ArrowDataType)]| {
6451 fields.iter().all(|(_, _, data_type)| {
6452 is_numeric(data_type) || data_type == &native_histogram_type
6453 })
6454 };
6455 let mixed_sample_types = if left_has_alternative_samples || right_has_alternative_samples {
6456 if !fields_are_samples(&left_fields) || !fields_are_samples(&right_fields) {
6457 return UnexpectedPlanExprSnafu {
6458 desc: format!(
6459 "OR value fields have incompatible types: {:?} and {:?}",
6460 left_fields
6461 .iter()
6462 .map(|(_, _, data_type)| data_type)
6463 .collect::<Vec<_>>(),
6464 right_fields
6465 .iter()
6466 .map(|(_, _, data_type)| data_type)
6467 .collect::<Vec<_>>()
6468 ),
6469 }
6470 .fail();
6471 }
6472 true
6473 } else {
6474 (left_field.2 == native_histogram_type && is_numeric(&right_field.2))
6475 || (right_field.2 == native_histogram_type && is_numeric(&left_field.2))
6476 };
6477 let target_field_type = if mixed_sample_types {
6478 ArrowDataType::Float64
6481 } else if left_field.2 == right_field.2 {
6482 left_field.2.clone()
6483 } else if is_numeric(&left_field.2) && is_numeric(&right_field.2) {
6484 ArrowDataType::Float64
6485 } else {
6486 return UnexpectedPlanExprSnafu {
6487 desc: format!(
6488 "OR value fields have incompatible types: {:?} and {:?}",
6489 left_field.2, right_field.2
6490 ),
6491 }
6492 .fail();
6493 };
6494 let (mixed_float_field_col, mixed_histogram_field_col) = if mixed_sample_types {
6495 let mut reserved_names = left
6496 .schema()
6497 .fields()
6498 .iter()
6499 .chain(right.schema().fields().iter())
6500 .map(|field| field.name().clone())
6501 .collect::<HashSet<_>>();
6502 for (name, _, _) in left_fields.iter().chain(&right_fields) {
6503 reserved_names.remove(name);
6504 }
6505 reserved_names.extend(all_tags.iter().cloned());
6506 let unique_name = |prefix: &str, reserved_names: &mut HashSet<String>| {
6507 let mut index = 0;
6508 loop {
6509 let name = format!("{prefix}{index}");
6510 index += 1;
6511 if reserved_names.insert(name.clone()) {
6512 break name;
6513 }
6514 }
6515 };
6516 let float_field = unique_name(OR_FLOAT_FIELD_PREFIX, &mut reserved_names);
6517 let histogram_field = unique_name(OR_HISTOGRAM_FIELD_PREFIX, &mut reserved_names);
6518 (float_field, histogram_field)
6519 } else {
6520 (left_field_col.clone(), String::new())
6521 };
6522 let left_tag_types = left_tag_cols_set
6523 .iter()
6524 .map(|label| {
6525 left.schema()
6526 .fields()
6527 .iter()
6528 .find(|field| field.name() == label)
6529 .map(|field| (label.clone(), field.data_type().clone()))
6530 .with_context(|| ColumnNotFoundSnafu { col: label.clone() })
6531 })
6532 .collect::<Result<HashMap<_, _>>>()?;
6533 let right_tag_types = right_tag_cols_set
6534 .iter()
6535 .map(|label| {
6536 right
6537 .schema()
6538 .fields()
6539 .iter()
6540 .find(|field| field.name() == label)
6541 .map(|field| (label.clone(), field.data_type().clone()))
6542 .with_context(|| ColumnNotFoundSnafu { col: label.clone() })
6543 })
6544 .collect::<Result<HashMap<_, _>>>()?;
6545 let mut target_tag_types = HashMap::with_capacity(all_tags.len());
6546 for label in &all_tags {
6547 let Some(data_type) =
6548 Self::common_label_data_type(left_tag_types.get(label), right_tag_types.get(label))
6549 else {
6550 return UnexpectedPlanExprSnafu {
6551 desc: format!(
6552 "OR label {label} has incompatible types: {:?} and {:?}",
6553 left_tag_types.get(label),
6554 right_tag_types.get(label)
6555 ),
6556 }
6557 .fail();
6558 };
6559 target_tag_types.insert(label.clone(), data_type);
6560 }
6561 let left_has_tsid = left
6562 .schema()
6563 .fields()
6564 .iter()
6565 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME);
6566 let right_has_tsid = right
6567 .schema()
6568 .fields()
6569 .iter()
6570 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME);
6571
6572 let mut all_columns_set = left
6574 .schema()
6575 .fields()
6576 .iter()
6577 .chain(right.schema().fields().iter())
6578 .map(|field| field.name().clone())
6579 .collect::<HashSet<_>>();
6580 if !(left_has_tsid && right_has_tsid) {
6583 all_columns_set.remove(DATA_SCHEMA_TSID_COLUMN_NAME);
6584 }
6585 all_columns_set.remove(&left_time_index_column);
6587 all_columns_set.remove(&right_time_index_column);
6588 if mixed_sample_types {
6589 for (name, _, _) in left_fields.iter().chain(&right_fields) {
6590 all_columns_set.remove(name);
6591 }
6592 all_columns_set.extend(all_tags.iter().cloned());
6593 all_columns_set.insert(mixed_float_field_col.clone());
6594 all_columns_set.insert(mixed_histogram_field_col.clone());
6595 } else if left_field_col != right_field_col {
6596 all_columns_set.remove(right_field_col);
6598 }
6599 let mut all_columns = all_columns_set.into_iter().collect::<Vec<_>>();
6600 all_columns.sort_unstable();
6602 all_columns.insert(0, left_time_index_column.clone());
6604 let mut occupied_column_names = left
6605 .schema()
6606 .fields()
6607 .iter()
6608 .chain(right.schema().fields().iter())
6609 .map(|field| field.name().clone())
6610 .collect::<HashSet<_>>();
6611
6612 let aligned_label_expr = |col: &String, source_types: &HashMap<String, ArrowDataType>| {
6614 let target_type = &target_tag_types[col];
6615 if let Some(source_type) = source_types.get(col) {
6616 let expr = DfExpr::Column(Column::new(None::<String>, col));
6617 if source_type == target_type {
6618 expr
6619 } else {
6620 DfExpr::Cast(Cast {
6621 expr: Box::new(expr),
6622 data_type: target_type.clone(),
6623 })
6624 .alias(col.clone())
6625 }
6626 } else {
6627 DfExpr::Literal(
6628 Self::string_scalar_value(target_type, None)
6629 .expect("target label type is a string"),
6630 None,
6631 )
6632 .alias(col.clone())
6633 }
6634 };
6635 let null_histogram =
6636 ScalarValue::try_new_null(&native_histogram_type).context(DataFusionPlanningSnafu)?;
6637 let mixed_value_expr = |fields: &[(String, Option<TableReference>, ArrowDataType)],
6638 output_col: &String| {
6639 if output_col == &mixed_float_field_col {
6640 if let Some((name, qualifier, data_type)) = fields
6641 .iter()
6642 .find(|(_, _, data_type)| is_numeric(data_type))
6643 {
6644 let expr = DfExpr::Column(Column::new(qualifier.clone(), name));
6645 if data_type == &ArrowDataType::Float64 {
6646 expr.alias(output_col)
6647 } else {
6648 DfExpr::Cast(Cast {
6649 expr: Box::new(expr),
6650 data_type: ArrowDataType::Float64,
6651 })
6652 .alias(output_col)
6653 }
6654 } else {
6655 DfExpr::Literal(ScalarValue::Float64(None), None).alias(output_col)
6656 }
6657 } else {
6658 fields
6659 .iter()
6660 .find(|(_, _, data_type)| data_type == &native_histogram_type)
6661 .map(|(name, qualifier, _)| {
6662 DfExpr::Column(Column::new(qualifier.clone(), name)).alias(output_col)
6663 })
6664 .unwrap_or_else(|| {
6665 DfExpr::Literal(null_histogram.clone(), None).alias(output_col)
6666 })
6667 }
6668 };
6669 let left_proj_exprs = all_columns.iter().map(|col| {
6670 if mixed_sample_types
6671 && (col == &mixed_float_field_col || col == &mixed_histogram_field_col)
6672 {
6673 mixed_value_expr(&left_fields, col)
6674 } else if !mixed_sample_types
6675 && col == left_field_col
6676 && left_field.2 != target_field_type
6677 {
6678 DfExpr::Cast(Cast {
6679 expr: Box::new(DfExpr::Column(Column::new(
6680 left_field.1.clone(),
6681 left_field_col,
6682 ))),
6683 data_type: target_field_type.clone(),
6684 })
6685 .alias(left_field_col.clone())
6686 } else if target_tag_types.contains_key(col) {
6687 aligned_label_expr(col, &left_tag_types)
6688 } else {
6689 DfExpr::Column(Column::new(None::<String>, col))
6690 }
6691 });
6692 let right_time_index_expr = DfExpr::Column(Column::new(
6693 right_qualifier.clone(),
6694 right_time_index_column,
6695 ))
6696 .alias(left_time_index_column.clone());
6697 let right_proj_exprs_without_time_index = all_columns.iter().skip(1).map(|col| {
6701 if mixed_sample_types
6703 && (col == &mixed_float_field_col || col == &mixed_histogram_field_col)
6704 {
6705 mixed_value_expr(&right_fields, col)
6706 } else if !mixed_sample_types && col == left_field_col {
6707 let expr = DfExpr::Column(Column::new(right_field.1.clone(), right_field_col));
6708 if right_field.2 != target_field_type {
6709 DfExpr::Cast(Cast {
6710 expr: Box::new(expr),
6711 data_type: target_field_type.clone(),
6712 })
6713 .alias(left_field_col.clone())
6714 } else if left_field_col != right_field_col {
6715 expr.alias(left_field_col.clone())
6716 } else {
6717 expr
6718 }
6719 } else if target_tag_types.contains_key(col) {
6720 aligned_label_expr(col, &right_tag_types)
6721 } else {
6722 DfExpr::Column(Column::new(None::<String>, col))
6723 }
6724 });
6725 let right_proj_exprs = [right_time_index_expr]
6726 .into_iter()
6727 .chain(right_proj_exprs_without_time_index);
6728
6729 let left_projected = LogicalPlanBuilder::from(left)
6730 .project(left_proj_exprs)
6731 .context(DataFusionPlanningSnafu)?
6732 .alias(left_qualifier_string.clone())
6733 .context(DataFusionPlanningSnafu)?
6734 .build()
6735 .context(DataFusionPlanningSnafu)?;
6736 let right_projected = LogicalPlanBuilder::from(right)
6737 .project(right_proj_exprs)
6738 .context(DataFusionPlanningSnafu)?
6739 .alias(right_qualifier_string.clone())
6740 .context(DataFusionPlanningSnafu)?
6741 .build()
6742 .context(DataFusionPlanningSnafu)?;
6743
6744 let mut match_columns = if let Some(modifier) = modifier
6746 && let Some(matching) = &modifier.matching
6747 {
6748 match matching {
6749 LabelModifier::Include(on) => on.labels.clone(),
6751 LabelModifier::Exclude(ignoring) => {
6753 let ignoring = ignoring.labels.iter().cloned().collect::<HashSet<_>>();
6754 all_tags.difference(&ignoring).cloned().collect()
6755 }
6756 }
6757 } else {
6758 all_tags.iter().cloned().collect()
6759 };
6760 match_columns.sort_unstable();
6762 match_columns.dedup();
6763 occupied_column_names.extend(
6764 left_projected
6765 .schema()
6766 .fields()
6767 .iter()
6768 .chain(right_projected.schema().fields().iter())
6769 .map(|field| field.name().clone()),
6770 );
6771
6772 let visible_schema = left_projected.schema().clone();
6773 let visible_left_exprs = left_projected
6774 .schema()
6775 .iter()
6776 .map(|(qualifier, field)| {
6777 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6778 })
6779 .collect::<Vec<_>>();
6780 let visible_right_exprs = right_projected
6781 .schema()
6782 .iter()
6783 .map(|(qualifier, field)| {
6784 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6785 })
6786 .collect::<Vec<_>>();
6787 let mut left_match_exprs = Vec::with_capacity(match_columns.len());
6788 let mut right_match_exprs = Vec::with_capacity(match_columns.len());
6789 let mut next_internal_column = 0;
6790
6791 for label in &match_columns {
6792 let left_field = if left_tag_cols_set.contains(label) {
6793 Some(
6794 left_projected
6795 .schema()
6796 .iter()
6797 .find(|(_, field)| field.name() == label)
6798 .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone()))
6799 .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?,
6800 )
6801 } else {
6802 None
6803 };
6804 let right_field = if right_tag_cols_set.contains(label) {
6805 Some(
6806 right_projected
6807 .schema()
6808 .iter()
6809 .find(|(_, field)| field.name() == label)
6810 .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone()))
6811 .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?,
6812 )
6813 } else {
6814 None
6815 };
6816 let data_type = match (left_field.as_ref(), right_field.as_ref()) {
6817 (Some((_, left_type)), Some((_, right_type))) if left_type == right_type => {
6818 left_type.clone()
6819 }
6820 (Some((_, left_type)), Some((_, right_type))) => {
6821 return UnexpectedPlanExprSnafu {
6822 desc: format!(
6823 "OR match label {label} has incompatible types: {left_type:?} and {right_type:?}"
6824 ),
6825 }
6826 .fail();
6827 }
6828 (Some((_, data_type)), None) | (None, Some((_, data_type))) => data_type.clone(),
6829 (None, None) => ArrowDataType::Utf8,
6830 };
6831 let Some(value_type) = Self::string_value_data_type(&data_type).cloned() else {
6832 return UnexpectedPlanExprSnafu {
6833 desc: format!("OR match label {label} must be a string"),
6834 }
6835 .fail();
6836 };
6837 let internal_name = loop {
6838 let name = format!("__promql_or_match_{next_internal_column}");
6839 next_internal_column += 1;
6840 if occupied_column_names.insert(name.clone()) {
6841 break name;
6842 }
6843 };
6844 left_match_exprs.push(Self::normalized_match_key_expr(
6845 label,
6846 left_field,
6847 &value_type,
6848 &internal_name,
6849 ));
6850 right_match_exprs.push(Self::normalized_match_key_expr(
6851 label,
6852 right_field,
6853 &value_type,
6854 &internal_name,
6855 ));
6856 }
6857
6858 let left_augmented = LogicalPlanBuilder::from(left_projected)
6859 .project(visible_left_exprs.into_iter().chain(left_match_exprs))
6860 .context(DataFusionPlanningSnafu)?
6861 .build()
6862 .context(DataFusionPlanningSnafu)?;
6863 let right_augmented = LogicalPlanBuilder::from(right_projected)
6864 .project(visible_right_exprs.into_iter().chain(right_match_exprs))
6865 .context(DataFusionPlanningSnafu)?
6866 .build()
6867 .context(DataFusionPlanningSnafu)?;
6868
6869 let visible_field_count = visible_schema.fields().len();
6871 let compare_key_indices =
6872 (visible_field_count..visible_field_count + match_columns.len()).collect::<Vec<_>>();
6873 let (time_qualifier, _) = visible_schema
6874 .iter()
6875 .find(|(_, field)| field.name() == &left_time_index_column)
6876 .with_context(|| TimeIndexNotFoundSnafu {
6877 table: left_qualifier_string.clone(),
6878 })?;
6879 let ts_col_idx = left_augmented
6880 .schema()
6881 .iter()
6882 .position(|(qualifier, field)| {
6883 qualifier == time_qualifier && field.name() == &left_time_index_column
6884 })
6885 .with_context(|| TimeIndexNotFoundSnafu {
6886 table: left_qualifier_string.clone(),
6887 })?;
6888 let union_distinct_on = UnionDistinctOn::try_new(
6889 left_augmented,
6890 right_augmented,
6891 compare_key_indices,
6892 ts_col_idx,
6893 )
6894 .context(DataFusionPlanningSnafu)?;
6895 let augmented_result = LogicalPlan::Extension(Extension {
6896 node: Arc::new(union_distinct_on),
6897 });
6898 let result = LogicalPlanBuilder::from(augmented_result)
6899 .project(visible_schema.iter().map(|(qualifier, field)| {
6900 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6901 }))
6902 .context(DataFusionPlanningSnafu)?
6903 .build()
6904 .context(DataFusionPlanningSnafu)?;
6905
6906 let output_field_col = left_field_col.clone();
6908 let mut output_context = left_context;
6909 let mut visible_tags = all_tags.into_iter().collect::<Vec<_>>();
6910 visible_tags.sort_unstable();
6911 output_context.time_index_column = Some(left_time_index_column);
6912 output_context.tag_columns = visible_tags;
6913 output_context.field_columns = if mixed_sample_types {
6914 vec![mixed_float_field_col, mixed_histogram_field_col]
6915 } else {
6916 vec![output_field_col]
6917 };
6918 output_context.use_tsid = left_has_tsid && right_has_tsid;
6919 self.ctx = output_context;
6920
6921 Ok(result)
6922 }
6923
6924 fn projection_for_each_field_column<F>(
6932 &mut self,
6933 input: LogicalPlan,
6934 name_to_expr: F,
6935 ) -> Result<LogicalPlan>
6936 where
6937 F: FnMut(&String) -> Result<DfExpr>,
6938 {
6939 let preserve_field_names =
6942 Self::field_columns_are_alternative_samples(input.schema(), &self.ctx.field_columns);
6943 let table_ref = self.ctx.table_name.clone().map(TableReference::bare);
6944 let input_schema = input.schema().clone();
6946 let non_field_columns_iter = self
6947 .ctx
6948 .tag_columns
6949 .iter()
6950 .chain(self.ctx.time_index_column.iter())
6951 .map(|col| {
6952 input_schema
6953 .qualified_field_with_name(table_ref.as_ref(), col)
6954 .or_else(|_| input_schema.qualified_field_with_unqualified_name(col))
6955 .map(|field| DfExpr::Column(field.into()))
6956 .context(DataFusionPlanningSnafu)
6957 });
6958 let tsid_iter =
6959 Self::optional_tsid_projection(input.schema(), table_ref.as_ref(), self.ctx.use_tsid)
6960 .into_iter()
6961 .map(Ok);
6962
6963 let result_field_columns = self
6965 .ctx
6966 .field_columns
6967 .iter()
6968 .map(name_to_expr)
6969 .collect::<Result<Vec<_>>>()?;
6970
6971 if !preserve_field_names {
6973 self.ctx.field_columns = result_field_columns
6974 .iter()
6975 .map(|expr| expr.schema_name().to_string())
6976 .collect();
6977 }
6978 let field_columns_iter = result_field_columns
6979 .into_iter()
6980 .zip(self.ctx.field_columns.iter())
6981 .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::<String>, name))));
6982
6983 let project_fields = non_field_columns_iter
6985 .chain(tsid_iter)
6986 .chain(field_columns_iter)
6987 .collect::<Result<Vec<_>>>()?;
6988
6989 LogicalPlanBuilder::from(input)
6990 .project(project_fields)
6991 .context(DataFusionPlanningSnafu)?
6992 .build()
6993 .context(DataFusionPlanningSnafu)
6994 }
6995
6996 fn filter_on_field_column<F>(&self, input: LogicalPlan, name_to_expr: F) -> Result<LogicalPlan>
6998 where
6999 F: FnMut(&String) -> Result<DfExpr>,
7000 {
7001 ensure!(
7002 self.ctx.field_columns.len() == 1
7003 || Self::field_columns_are_alternative_samples(
7004 input.schema(),
7005 &self.ctx.field_columns,
7006 ),
7007 UnsupportedExprSnafu {
7008 name: "filter on multi-value input"
7009 }
7010 );
7011
7012 let field_column_filters = self
7013 .ctx
7014 .field_columns
7015 .iter()
7016 .map(name_to_expr)
7017 .collect::<Result<Vec<_>>>()?;
7018 let field_column_filter =
7019 disjunction(field_column_filters).context(UnsupportedExprSnafu {
7020 name: "filter on empty input",
7021 })?;
7022
7023 LogicalPlanBuilder::from(input)
7024 .filter(field_column_filter)
7025 .context(DataFusionPlanningSnafu)?
7026 .build()
7027 .context(DataFusionPlanningSnafu)
7028 }
7029
7030 fn date_part_on_time_index(&self, date_part: &str) -> Result<DfExpr> {
7033 let input_expr = datafusion::logical_expr::col(
7034 self.ctx
7035 .time_index_column
7036 .as_ref()
7037 .with_context(|| TimeIndexNotFoundSnafu {
7039 table: "<doesn't matter>",
7040 })?,
7041 );
7042 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
7043 func: datafusion_functions::datetime::date_part(),
7044 args: vec![date_part.lit(), input_expr],
7045 });
7046 Ok(fn_expr)
7047 }
7048
7049 fn strip_tsid_column(&self, plan: LogicalPlan) -> Result<LogicalPlan> {
7050 let schema = plan.schema();
7051 if !schema
7052 .fields()
7053 .iter()
7054 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
7055 {
7056 return Ok(plan);
7057 }
7058
7059 let project_exprs = schema
7062 .iter()
7063 .filter(|(_, field)| field.name() != DATA_SCHEMA_TSID_COLUMN_NAME)
7064 .map(|(qualifier, field)| {
7065 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
7066 })
7067 .collect::<Vec<_>>();
7068
7069 LogicalPlanBuilder::from(plan)
7070 .project(project_exprs)
7071 .context(DataFusionPlanningSnafu)?
7072 .build()
7073 .context(DataFusionPlanningSnafu)
7074 }
7075
7076 fn apply_alias(&mut self, plan: LogicalPlan, alias_name: String) -> Result<LogicalPlan> {
7078 let fields_expr = self.create_field_column_exprs()?;
7079
7080 ensure!(
7082 fields_expr.len() == 1,
7083 UnsupportedExprSnafu {
7084 name: "alias on multi-value result"
7085 }
7086 );
7087
7088 let project_fields = fields_expr
7089 .into_iter()
7090 .map(|expr| expr.alias(&alias_name))
7091 .chain(self.create_tag_column_exprs()?)
7092 .chain(Some(self.create_time_index_column_expr()?));
7093
7094 LogicalPlanBuilder::from(plan)
7095 .project(project_fields)
7096 .context(DataFusionPlanningSnafu)?
7097 .build()
7098 .context(DataFusionPlanningSnafu)
7099 }
7100}
7101
7102#[derive(Default, Debug)]
7103struct FunctionArgs {
7104 input: Option<PromExpr>,
7105 literals: Vec<DfExpr>,
7106}
7107
7108#[derive(Debug, Clone)]
7111enum ScalarFunc {
7112 DataFusionBuiltin(Arc<ScalarUdfDef>),
7116 DataFusionUdf(Arc<ScalarUdfDef>),
7120 NativeHistogramUdf(Arc<ScalarUdfDef>),
7123 Udf(Arc<ScalarUdfDef>),
7128 ExtrapolateUdf(Arc<ScalarUdfDef>, i64),
7135 GeneratedExpr,
7139}
7140
7141#[cfg(test)]
7142mod test {
7143 use std::time::{Duration, UNIX_EPOCH};
7144
7145 use catalog::RegisterTableRequest;
7146 use catalog::memory::{MemoryCatalogManager, new_memory_catalog_manager};
7147 use common_base::Plugins;
7148 use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
7149 use common_query::native_histogram::{
7150 CUSTOM_BUCKETS_SCHEMA, CounterResetHint, NativeHistogram, build_histogram_array,
7151 };
7152 use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value};
7153 use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS;
7154 use common_query::test_util::DummyDecoder;
7155 use common_recordbatch::RecordBatch as GreptimeRecordBatch;
7156 use datafusion::arrow::array::{
7157 Array, Float64Array, Int64Array, StringArray, TimestampMillisecondArray,
7158 };
7159 use datafusion::arrow::datatypes::{Field, Schema as ArrowSchema};
7160 use datafusion::arrow::record_batch::RecordBatch;
7161 use datafusion::catalog::{CatalogProvider, MemoryCatalogProvider, MemorySchemaProvider};
7162 use datafusion::datasource::memory::MemorySourceConfig;
7163 use datafusion::datasource::source::DataSourceExec;
7164 use datafusion::datasource::{MemTable, provider_as_source};
7165 use datafusion::execution::context::SessionContext;
7166 use datafusion::logical_expr::Extension;
7167 use datatypes::prelude::ConcreteDataType;
7168 use datatypes::schema::{ColumnSchema, Schema};
7169 use promql_parser::label::Labels;
7170 use promql_parser::parser;
7171 use session::context::QueryContext;
7172 use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan};
7173 use table::Table;
7174 use table::metadata::{FilterPushDownType, TableInfoBuilder, TableMetaBuilder};
7175 use table::test_util::{EmptyTable, MemTable as GreptimeMemTable};
7176
7177 use super::*;
7178 use crate::QueryEngineContext;
7179 use crate::options::QueryOptions;
7180 use crate::parser::QueryLanguageParser;
7181 use crate::query_engine::DefaultSerializer;
7182
7183 mod delta;
7184
7185 fn find_instant_manipulate(plan: &LogicalPlan) -> Option<&InstantManipulate> {
7186 if let LogicalPlan::Extension(Extension { node }) = plan
7187 && let Some(instant_manipulate) = node.as_any().downcast_ref::<InstantManipulate>()
7188 {
7189 return Some(instant_manipulate);
7190 }
7191
7192 plan.inputs().into_iter().find_map(find_instant_manipulate)
7193 }
7194
7195 fn build_query_engine_state() -> QueryEngineState {
7196 QueryEngineState::new(
7197 new_memory_catalog_manager().unwrap(),
7198 None,
7199 None,
7200 None,
7201 None,
7202 None,
7203 false,
7204 Plugins::default(),
7205 QueryOptions::default(),
7206 )
7207 }
7208
7209 #[test]
7210 fn common_label_type_preserves_only_shared_dictionary_encoding() {
7211 let dictionary = ArrowDataType::Dictionary(
7212 Box::new(ArrowDataType::UInt32),
7213 Box::new(ArrowDataType::Utf8),
7214 );
7215 let other_dictionary = ArrowDataType::Dictionary(
7216 Box::new(ArrowDataType::Int32),
7217 Box::new(ArrowDataType::Utf8),
7218 );
7219
7220 assert_eq!(
7221 Some(dictionary.clone()),
7222 PromPlanner::common_label_data_type(Some(&dictionary), Some(&dictionary))
7223 );
7224 assert_eq!(
7225 Some(ArrowDataType::Utf8),
7226 PromPlanner::common_label_data_type(Some(&dictionary), Some(&ArrowDataType::Utf8))
7227 );
7228 assert_eq!(
7229 Some(ArrowDataType::Utf8),
7230 PromPlanner::common_label_data_type(Some(&dictionary), Some(&other_dictionary))
7231 );
7232 assert_eq!(
7233 Some(ArrowDataType::Utf8),
7234 PromPlanner::common_label_data_type(Some(&dictionary), None)
7235 );
7236 }
7237
7238 async fn build_optimized_promql_plan(
7239 table_provider: DfTableSourceProvider,
7240 eval_stmt: &EvalStmt,
7241 ) -> LogicalPlan {
7242 let state = build_query_engine_state();
7243 let raw_plan = PromPlanner::stmt_to_plan(table_provider, eval_stmt, &state)
7244 .await
7245 .unwrap();
7246 let context = QueryEngineContext::new(state.session_state(), QueryContext::arc());
7247 state
7248 .optimize_by_extension_rules(raw_plan, &context)
7249 .unwrap()
7250 }
7251
7252 async fn build_optimized_tsid_plan(
7253 query: &str,
7254 num_tag: usize,
7255 num_field: usize,
7256 end_secs: u64,
7257 lookback_secs: u64,
7258 ) -> String {
7259 let eval_stmt = EvalStmt {
7260 expr: parser::parse(query).unwrap(),
7261 start: UNIX_EPOCH,
7262 end: UNIX_EPOCH
7263 .checked_add(Duration::from_secs(end_secs))
7264 .unwrap(),
7265 interval: Duration::from_secs(5),
7266 lookback_delta: Duration::from_secs(lookback_secs),
7267 };
7268 let table_provider = build_test_table_provider_with_tsid(
7269 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
7270 num_tag,
7271 num_field,
7272 )
7273 .await;
7274
7275 build_optimized_promql_plan(table_provider, &eval_stmt)
7276 .await
7277 .display_indent_schema()
7278 .to_string()
7279 }
7280
7281 async fn assert_nested_count_rewrite_applies(query: &str, expected_outer_agg: &str) {
7282 let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await;
7283
7284 assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
7285 assert!(plan_str.contains("Projection: some_metric.timestamp, some_metric.tag_0"));
7286 assert!(plan_str.contains("Distinct:"));
7287 assert!(plan_str.contains(expected_outer_agg), "{plan_str}");
7288 assert!(!plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]"));
7289 }
7290
7291 async fn assert_nested_count_rewrite_missing(query: &str, num_tag: usize, lookback_secs: u64) {
7292 let plan_str = build_optimized_tsid_plan(query, num_tag, 1, 100_000, lookback_secs).await;
7293 assert!(!plan_str.contains("Distinct:"), "{plan_str}");
7294 }
7295
7296 fn build_eval_stmt(expr: &str) -> EvalStmt {
7297 EvalStmt {
7298 expr: parser::parse(expr).unwrap(),
7299 start: UNIX_EPOCH,
7300 end: UNIX_EPOCH
7301 .checked_add(Duration::from_secs(100_000))
7302 .unwrap(),
7303 interval: Duration::from_secs(5),
7304 lookback_delta: Duration::from_secs(1),
7305 }
7306 }
7307
7308 enum DirectOrValue {
7309 Float64(f64),
7310 Int64(i64),
7311 NativeHistogram(NativeHistogram),
7312 Utf8(&'static str),
7313 }
7314
7315 impl DirectOrValue {
7316 fn data_type(&self) -> ArrowDataType {
7317 match self {
7318 Self::Float64(_) => ArrowDataType::Float64,
7319 Self::Int64(_) => ArrowDataType::Int64,
7320 Self::NativeHistogram(_) => native_histogram_value_type().as_arrow_type(),
7321 Self::Utf8(_) => ArrowDataType::Utf8,
7322 }
7323 }
7324 fn array(&self) -> Arc<dyn Array> {
7325 match self {
7326 Self::Float64(v) => Arc::new(Float64Array::from(vec![*v])),
7327 Self::Int64(v) => Arc::new(Int64Array::from(vec![*v])),
7328 Self::NativeHistogram(v) => build_histogram_array(&[Some(v.clone())]),
7329 Self::Utf8(v) => Arc::new(StringArray::from(vec![*v])),
7330 }
7331 }
7332 }
7333
7334 fn direct_or_histogram() -> NativeHistogram {
7335 NativeHistogram {
7336 schema: 0,
7337 zero_threshold: 0.0,
7338 sum: 1.0,
7339 reset_hint: CounterResetHint::Unknown,
7340 start_timestamp: None,
7341 custom_values: vec![],
7342 positive_spans: vec![],
7343 negative_spans: vec![],
7344 count: 1.0,
7345 zero_count: 1.0,
7346 positive_buckets: vec![],
7347 negative_buckets: vec![],
7348 }
7349 }
7350
7351 fn operator_metric_table(
7352 name: &str,
7353 table_id: u32,
7354 tag: &str,
7355 le: Option<&str>,
7356 value: DirectOrValue,
7357 ) -> table::TableRef {
7358 let value_type = match &value {
7359 DirectOrValue::Float64(_) => ConcreteDataType::float64_datatype(),
7360 DirectOrValue::Int64(_) => ConcreteDataType::int64_datatype(),
7361 DirectOrValue::NativeHistogram(_) => native_histogram_value_type().clone(),
7362 DirectOrValue::Utf8(_) => ConcreteDataType::string_datatype(),
7363 };
7364 let tag_count = 1 + usize::from(le.is_some());
7365 let mut columns = vec![ColumnSchema::new(
7366 "tag".to_string(),
7367 ConcreteDataType::string_datatype(),
7368 false,
7369 )];
7370 if le.is_some() {
7371 columns.push(ColumnSchema::new(
7372 LE_COLUMN_NAME.to_string(),
7373 ConcreteDataType::string_datatype(),
7374 false,
7375 ));
7376 }
7377 columns.extend([
7378 ColumnSchema::new(
7379 "ts".to_string(),
7380 ConcreteDataType::timestamp_millisecond_datatype(),
7381 false,
7382 )
7383 .with_time_index(true),
7384 ColumnSchema::new("v".to_string(), value_type, true),
7385 ]);
7386 let schema = Arc::new(Schema::new(columns));
7387 let mut arrays = vec![Arc::new(StringArray::from(vec![tag])) as Arc<dyn Array>];
7388 if let Some(le) = le {
7389 arrays.push(Arc::new(StringArray::from(vec![le])));
7390 }
7391 arrays.extend([
7392 Arc::new(TimestampMillisecondArray::from(vec![1_000])) as Arc<dyn Array>,
7393 value.array(),
7394 ]);
7395 let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap();
7396 let backing = GreptimeMemTable::new_with_catalog(
7397 name,
7398 GreptimeRecordBatch::from_df_record_batch(schema.clone(), batch),
7399 table_id,
7400 DEFAULT_CATALOG_NAME.to_string(),
7401 DEFAULT_SCHEMA_NAME.to_string(),
7402 );
7403 let value_index = tag_count + 1;
7404 let meta = TableMetaBuilder::empty()
7405 .schema(schema)
7406 .primary_key_indices((0..tag_count).collect())
7407 .value_indices(vec![value_index])
7408 .next_column_id((value_index + 1) as u32)
7409 .build()
7410 .unwrap();
7411 let info = Arc::new(
7412 TableInfoBuilder::default()
7413 .table_id(table_id)
7414 .name(name)
7415 .meta(meta)
7416 .build()
7417 .unwrap(),
7418 );
7419 Arc::new(Table::new(
7420 info,
7421 FilterPushDownType::Unsupported,
7422 backing.data_source(),
7423 ))
7424 }
7425
7426 fn operator_table_provider() -> DfTableSourceProvider {
7427 let catalog = MemoryCatalogManager::with_default_setup();
7428 let tables = [
7429 operator_metric_table("lf", 2_001, "a", None, DirectOrValue::Float64(2.0)),
7430 operator_metric_table(
7431 "lh",
7432 2_002,
7433 "b",
7434 None,
7435 DirectOrValue::NativeHistogram(direct_or_histogram()),
7436 ),
7437 operator_metric_table("rf", 2_003, "b", None, DirectOrValue::Float64(3.0)),
7438 operator_metric_table(
7439 "rh",
7440 2_004,
7441 "a",
7442 None,
7443 DirectOrValue::NativeHistogram(direct_or_histogram()),
7444 ),
7445 operator_metric_table("fallback", 2_005, "c", None, DirectOrValue::Float64(7.0)),
7446 operator_metric_table(
7447 "bad_classic",
7448 2_006,
7449 "d",
7450 Some("broken"),
7451 DirectOrValue::Float64(1.0),
7452 ),
7453 operator_metric_table(
7454 "bad_native",
7455 2_007,
7456 "d",
7457 None,
7458 DirectOrValue::NativeHistogram(direct_or_histogram()),
7459 ),
7460 ];
7461 for table in tables {
7462 let info = table.table_info();
7463 catalog
7464 .register_table_sync(RegisterTableRequest {
7465 catalog: DEFAULT_CATALOG_NAME.to_string(),
7466 schema: DEFAULT_SCHEMA_NAME.to_string(),
7467 table_name: info.name.clone(),
7468 table_id: info.ident.table_id,
7469 table,
7470 })
7471 .unwrap();
7472 }
7473 DfTableSourceProvider::new(
7474 catalog,
7475 false,
7476 QueryContext::arc(),
7477 DummyDecoder::arc(),
7478 false,
7479 )
7480 }
7481
7482 fn operator_eval_stmt(expr: &str) -> EvalStmt {
7483 let time = UNIX_EPOCH.checked_add(Duration::from_secs(1)).unwrap();
7484 EvalStmt {
7485 expr: parser::parse(expr).unwrap(),
7486 start: time,
7487 end: time,
7488 interval: Duration::from_secs(1),
7489 lookback_delta: Duration::from_secs(5),
7490 }
7491 }
7492
7493 struct DirectOrSource {
7494 name: &'static str,
7495 empty: bool,
7496 timestamp: i64,
7497 tags: Vec<(&'static str, Option<&'static str>)>,
7498 value: DirectOrValue,
7499 }
7500
7501 fn source(
7502 name: &'static str,
7503 empty: bool,
7504 timestamp: i64,
7505 tags: Vec<(&'static str, Option<&'static str>)>,
7506 value: DirectOrValue,
7507 ) -> DirectOrSource {
7508 DirectOrSource {
7509 name,
7510 empty,
7511 timestamp,
7512 tags,
7513 value,
7514 }
7515 }
7516
7517 fn tagged_source(
7518 name: &'static str,
7519 empty: bool,
7520 tag: (&'static str, Option<&'static str>),
7521 value: DirectOrValue,
7522 ) -> DirectOrSource {
7523 source(name, empty, 1, vec![("job", Some("job")), tag], value)
7524 }
7525
7526 fn job_source(name: &'static str, value: DirectOrValue) -> DirectOrSource {
7527 source(name, true, 1, vec![("job", Some("job"))], value)
7528 }
7529
7530 fn table(source: &DirectOrSource) -> Arc<MemTable> {
7531 let mut fields = vec![Field::new(
7532 "ts",
7533 ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
7534 false,
7535 )];
7536 fields.extend(
7537 source
7538 .tags
7539 .iter()
7540 .map(|(name, _)| Field::new(*name, ArrowDataType::Utf8, true)),
7541 );
7542 fields.push(Field::new("v", source.value.data_type(), true));
7543 let schema = Arc::new(ArrowSchema::new(fields));
7544 let partitions = if source.empty {
7545 vec![vec![]]
7546 } else {
7547 let mut columns: Vec<Arc<dyn Array>> =
7548 vec![Arc::new(TimestampMillisecondArray::from(vec![
7549 source.timestamp,
7550 ]))];
7551 columns.extend(
7552 source
7553 .tags
7554 .iter()
7555 .map(|(_, value)| Arc::new(StringArray::from(vec![*value])) as Arc<dyn Array>),
7556 );
7557 columns.push(source.value.array());
7558 vec![vec![RecordBatch::try_new(schema.clone(), columns).unwrap()]]
7559 };
7560 Arc::new(MemTable::try_new(schema, partitions).unwrap())
7561 }
7562
7563 fn scan(source: &DirectOrSource) -> LogicalPlan {
7564 LogicalPlanBuilder::scan(source.name, provider_as_source(table(source)), None)
7565 .unwrap()
7566 .build()
7567 .unwrap()
7568 }
7569
7570 fn direct_or_context(qualifier: &str, tags: &[&str], field: &str) -> PromPlannerContext {
7571 PromPlannerContext {
7572 table_name: Some(qualifier.to_string()),
7573 time_index_column: Some("ts".to_string()),
7574 field_columns: vec![field.to_string()],
7575 tag_columns: tags.iter().map(|tag| (*tag).to_string()).collect(),
7576 ..Default::default()
7577 }
7578 }
7579
7580 fn or_modifier(expr: &str) -> Option<BinModifier> {
7581 let PromExpr::Binary(expr) = parser::parse(expr).unwrap() else {
7582 unreachable!()
7583 };
7584 expr.modifier
7585 }
7586
7587 async fn plan_direct_or(
7588 left: LogicalPlan,
7589 right: LogicalPlan,
7590 left_context: PromPlannerContext,
7591 right_context: PromPlannerContext,
7592 modifier: &Option<BinModifier>,
7593 ) -> LogicalPlan {
7594 let table_provider = build_test_table_provider_with_fields(
7595 &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
7596 &[],
7597 )
7598 .await;
7599 let mut planner = PromPlanner {
7600 table_provider,
7601 ctx: PromPlannerContext::default(),
7602 promql_annotations: None,
7603 };
7604 planner
7605 .or_operator(
7606 left,
7607 right,
7608 left_context.tag_columns.iter().cloned().collect(),
7609 right_context.tag_columns.iter().cloned().collect(),
7610 left_context,
7611 right_context,
7612 modifier,
7613 )
7614 .unwrap()
7615 }
7616
7617 async fn execute(
7618 plan: LogicalPlan,
7619 state: &QueryEngineState,
7620 ) -> (LogicalPlan, Vec<RecordBatch>) {
7621 let context = QueryEngineContext::new(state.session_state(), QueryContext::arc());
7622 let optimized = state.optimize_by_extension_rules(plan, &context).unwrap();
7623 let physical = state
7624 .session_state()
7625 .create_physical_plan(&optimized)
7626 .await
7627 .unwrap();
7628 let batches =
7629 datafusion::physical_plan::collect(physical, state.session_state().task_ctx())
7630 .await
7631 .unwrap();
7632 (optimized, batches)
7633 }
7634
7635 async fn run(
7636 left: &DirectOrSource,
7637 right: &DirectOrSource,
7638 left_context: PromPlannerContext,
7639 right_context: PromPlannerContext,
7640 modifier: &Option<BinModifier>,
7641 ) -> (LogicalPlan, Vec<RecordBatch>) {
7642 let plan = plan_direct_or(
7643 scan(left),
7644 scan(right),
7645 left_context,
7646 right_context,
7647 modifier,
7648 )
7649 .await;
7650 execute(plan, &build_query_engine_state()).await
7651 }
7652
7653 async fn mixed_direct_or(histogram_on_left: bool) -> (PromPlanner, LogicalPlan) {
7654 let sample = |histogram: bool| {
7655 if histogram {
7656 DirectOrValue::NativeHistogram(direct_or_histogram())
7657 } else {
7658 DirectOrValue::Float64(1.25)
7659 }
7660 };
7661 let left = tagged_source(
7662 "lhs",
7663 false,
7664 (
7665 "k",
7666 Some(if histogram_on_left {
7667 "histogram"
7668 } else {
7669 "float"
7670 }),
7671 ),
7672 sample(histogram_on_left),
7673 );
7674 let right = tagged_source(
7675 "rhs",
7676 false,
7677 (
7678 "k",
7679 Some(if histogram_on_left {
7680 "float"
7681 } else {
7682 "histogram"
7683 }),
7684 ),
7685 sample(!histogram_on_left),
7686 );
7687 let table_provider = build_test_table_provider_with_fields(
7688 &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
7689 &[],
7690 )
7691 .await;
7692 let mut planner = PromPlanner {
7693 table_provider,
7694 ctx: PromPlannerContext::default(),
7695 promql_annotations: None,
7696 };
7697 let left_context = direct_or_context("lhs", &["job", "k"], "v");
7698 let right_context = direct_or_context("rhs", &["job", "k"], "v");
7699 let plan = planner
7700 .or_operator(
7701 scan(&left),
7702 scan(&right),
7703 left_context.tag_columns.iter().cloned().collect(),
7704 right_context.tag_columns.iter().cloned().collect(),
7705 left_context,
7706 right_context,
7707 &or_modifier("lhs or on(k) rhs"),
7708 )
7709 .unwrap();
7710 (planner, plan)
7711 }
7712
7713 async fn mixed_aggregate_input(histograms: Vec<NativeHistogram>) -> (PromPlanner, LogicalPlan) {
7714 let float_field = format!("{OR_FLOAT_FIELD_PREFIX}0");
7715 let histogram_field = format!("{OR_HISTOGRAM_FIELD_PREFIX}0");
7716 let row_count = histograms.len() + 1;
7717 let schema = Arc::new(ArrowSchema::new(vec![
7718 Field::new(
7719 "ts",
7720 ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
7721 false,
7722 ),
7723 Field::new("k", ArrowDataType::Utf8, false),
7724 Field::new(&float_field, ArrowDataType::Float64, true),
7725 Field::new(
7726 &histogram_field,
7727 native_histogram_value_type().as_arrow_type(),
7728 true,
7729 ),
7730 ]));
7731 let mut histogram_values = Vec::with_capacity(row_count);
7732 histogram_values.push(None);
7733 histogram_values.extend(histograms.into_iter().map(Some));
7734 let batch = RecordBatch::try_new(
7735 schema.clone(),
7736 vec![
7737 Arc::new(TimestampMillisecondArray::from(vec![1; row_count])),
7738 Arc::new(StringArray::from_iter_values(
7739 (0..row_count).map(|row| format!("kind_{row}")),
7740 )),
7741 Arc::new(Float64Array::from_iter(
7742 (0..row_count).map(|row| (row == 0).then_some(1.25)),
7743 )),
7744 build_histogram_array(&histogram_values),
7745 ],
7746 )
7747 .unwrap();
7748 let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap());
7749 let plan = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None)
7750 .unwrap()
7751 .build()
7752 .unwrap();
7753 let table_provider = build_test_table_provider_with_fields(
7754 &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
7755 &[],
7756 )
7757 .await;
7758 let planner = PromPlanner {
7759 table_provider,
7760 ctx: PromPlannerContext {
7761 table_name: Some("mixed".to_string()),
7762 time_index_column: Some("ts".to_string()),
7763 field_columns: vec![float_field, histogram_field],
7764 tag_columns: vec!["k".to_string()],
7765 ..Default::default()
7766 },
7767 promql_annotations: None,
7768 };
7769 (planner, plan)
7770 }
7771
7772 fn assert_no_internal_or_keys(schema: &DFSchema) {
7773 assert!(
7774 schema
7775 .fields()
7776 .iter()
7777 .all(|field| !field.name().starts_with("__promql_or_match_")),
7778 "{schema:?}"
7779 );
7780 }
7781
7782 fn values(batches: &[RecordBatch], column: &str) -> Vec<f64> {
7783 batches
7784 .iter()
7785 .flat_map(|batch| {
7786 batch
7787 .column_by_name(column)
7788 .unwrap()
7789 .as_any()
7790 .downcast_ref::<Float64Array>()
7791 .unwrap()
7792 .iter()
7793 .flatten()
7794 })
7795 .collect()
7796 }
7797
7798 fn numeric_values(batches: &[RecordBatch], column: &str) -> Vec<f64> {
7799 batches
7800 .iter()
7801 .flat_map(|batch| {
7802 let values = datafusion::arrow::compute::cast(
7803 batch.column_by_name(column).unwrap(),
7804 &ArrowDataType::Float64,
7805 )
7806 .unwrap();
7807 values
7808 .as_any()
7809 .downcast_ref::<Float64Array>()
7810 .unwrap()
7811 .iter()
7812 .flatten()
7813 .collect::<Vec<_>>()
7814 })
7815 .collect()
7816 }
7817
7818 fn histograms(batches: &[RecordBatch], column: &str) -> Vec<NativeHistogram> {
7819 batches
7820 .iter()
7821 .flat_map(|batch| {
7822 let values = batch
7823 .column_by_name(column)
7824 .unwrap()
7825 .as_any()
7826 .downcast_ref::<datafusion::arrow::array::StructArray>()
7827 .unwrap();
7828 (0..values.len()).filter_map(|row| {
7829 common_query::native_histogram::read_histogram(values, row).unwrap()
7830 })
7831 })
7832 .collect()
7833 }
7834
7835 fn rows(batches: &[RecordBatch]) -> Vec<(f64, Option<String>)> {
7836 let mut rows = batches
7837 .iter()
7838 .flat_map(|batch| {
7839 let values = batch
7840 .column_by_name("v")
7841 .unwrap()
7842 .as_any()
7843 .downcast_ref::<Float64Array>()
7844 .unwrap();
7845 let labels = batch
7846 .column_by_name("k")
7847 .map(|column| column.as_any().downcast_ref::<StringArray>().unwrap());
7848 (0..batch.num_rows()).map(move |i| {
7849 (
7850 values.value(i),
7851 labels.and_then(|labels| {
7852 (!labels.is_null(i)).then(|| labels.value(i).to_string())
7853 }),
7854 )
7855 })
7856 })
7857 .collect::<Vec<_>>();
7858 rows.sort_by(|left, right| left.0.total_cmp(&right.0));
7859 rows
7860 }
7861
7862 fn matrix_source(
7863 name: &'static str,
7864 k: Option<Option<&'static str>>,
7865 timestamp: i64,
7866 value: f64,
7867 ) -> DirectOrSource {
7868 let mut tags = vec![("job", Some("job"))];
7869 if let Some(k) = k {
7870 tags.push(("k", k));
7871 }
7872 source(name, false, timestamp, tags, DirectOrValue::Float64(value))
7873 }
7874
7875 fn matrix_context(name: &str, k: Option<Option<&str>>) -> PromPlannerContext {
7876 direct_or_context(
7877 name,
7878 if k.is_some() { &["job", "k"] } else { &["job"] },
7879 "v",
7880 )
7881 }
7882
7883 async fn build_missing_le_or_normal_metric_table_provider() -> DfTableSourceProvider {
7884 build_test_table_provider_with_fields(
7885 &[
7886 (
7887 DEFAULT_SCHEMA_NAME.to_string(),
7888 "non_existent_histogram_bucket".to_string(),
7889 ),
7890 (DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string()),
7891 ],
7892 &["pod", "instance"],
7893 )
7894 .await
7895 }
7896
7897 fn assert_normal_metric_schema(plan: &LogicalPlan) {
7898 let fields = plan.schema().fields();
7899 assert_eq!(fields.len(), 4, "{fields:?}");
7900 assert!(
7901 fields.iter().any(|field| field.name() == "pod"),
7902 "{fields:?}"
7903 );
7904 assert!(
7905 fields.iter().any(|field| field.name() == "instance"),
7906 "{fields:?}"
7907 );
7908 assert!(
7909 fields
7910 .iter()
7911 .any(|field| field.name() == greptime_timestamp()),
7912 "{fields:?}"
7913 );
7914 assert!(
7915 fields.iter().any(|field| {
7916 field.name() == greptime_value() && field.data_type() == &ArrowDataType::Float64
7917 }),
7918 "{fields:?}"
7919 );
7920 }
7921
7922 async fn build_test_table_provider_with_distinct_tags(
7923 table_tags: &[(&str, &[&str])],
7924 ) -> DfTableSourceProvider {
7925 let catalog_list = MemoryCatalogManager::with_default_setup();
7926 for (table_name, tags) in table_tags {
7927 let mut columns = tags
7928 .iter()
7929 .map(|tag| {
7930 ColumnSchema::new(
7931 (*tag).to_string(),
7932 ConcreteDataType::string_datatype(),
7933 false,
7934 )
7935 })
7936 .collect::<Vec<_>>();
7937 columns.push(
7938 ColumnSchema::new(
7939 greptime_timestamp().to_string(),
7940 ConcreteDataType::timestamp_millisecond_datatype(),
7941 false,
7942 )
7943 .with_time_index(true),
7944 );
7945 columns.push(ColumnSchema::new(
7946 greptime_value().to_string(),
7947 ConcreteDataType::float64_datatype(),
7948 true,
7949 ));
7950 let table_meta = TableMetaBuilder::empty()
7951 .schema(Arc::new(Schema::new(columns)))
7952 .primary_key_indices((0..tags.len()).collect())
7953 .next_column_id(1024)
7954 .build()
7955 .unwrap();
7956 let table_info = TableInfoBuilder::default()
7957 .name((*table_name).to_string())
7958 .meta(table_meta)
7959 .build()
7960 .unwrap();
7961
7962 assert!(
7963 catalog_list
7964 .register_table_sync(RegisterTableRequest {
7965 catalog: DEFAULT_CATALOG_NAME.to_string(),
7966 schema: DEFAULT_SCHEMA_NAME.to_string(),
7967 table_name: (*table_name).to_string(),
7968 table_id: 1024,
7969 table: EmptyTable::from_table_info(&table_info),
7970 })
7971 .is_ok()
7972 );
7973 }
7974
7975 DfTableSourceProvider::new(
7976 catalog_list,
7977 false,
7978 QueryContext::arc(),
7979 DummyDecoder::arc(),
7980 false,
7981 )
7982 }
7983
7984 fn contains_histogram_fold(plan: &LogicalPlan) -> bool {
7985 matches!(plan, LogicalPlan::Extension(Extension { node }) if node.as_any().is::<HistogramFold>())
7986 || plan.inputs().into_iter().any(contains_histogram_fold)
7987 }
7988
7989 async fn build_set_op_context_table_provider() -> DfTableSourceProvider {
7990 build_test_table_provider_with_distinct_tags(&[
7991 ("bucket_metric", &["job", "le"]),
7992 ("normal_metric", &["job"]),
7993 ("fallback_metric", &["instance"]),
7994 ])
7995 .await
7996 }
7997
7998 async fn build_or_context_table_provider() -> DfTableSourceProvider {
7999 build_test_table_provider_with_distinct_tags(&[
8000 ("normal_metric", &["job"]),
8001 ("other_metric", &["instance"]),
8002 ("non_hist_metric", &["instance"]),
8003 ])
8004 .await
8005 }
8006
8007 async fn optimize_and_create_physical_plan(
8008 state: &QueryEngineState,
8009 plan: LogicalPlan,
8010 ) -> (
8011 LogicalPlan,
8012 Arc<dyn datafusion::physical_plan::ExecutionPlan>,
8013 ) {
8014 let context = QueryEngineContext::new(state.session_state(), QueryContext::arc());
8015 let optimized = state.optimize_by_extension_rules(plan, &context).unwrap();
8016 let physical = state
8017 .session_state()
8018 .create_physical_plan(&optimized)
8019 .await
8020 .unwrap();
8021 (optimized, physical)
8022 }
8023
8024 async fn build_test_table_provider(
8025 table_name_tuples: &[(String, String)],
8026 num_tag: usize,
8027 num_field: usize,
8028 ) -> DfTableSourceProvider {
8029 let catalog_list = MemoryCatalogManager::with_default_setup();
8030 for (schema_name, table_name) in table_name_tuples {
8031 let mut columns = vec![];
8032 for i in 0..num_tag {
8033 columns.push(ColumnSchema::new(
8034 format!("tag_{i}"),
8035 ConcreteDataType::string_datatype(),
8036 false,
8037 ));
8038 }
8039 columns.push(
8040 ColumnSchema::new(
8041 "timestamp".to_string(),
8042 ConcreteDataType::timestamp_millisecond_datatype(),
8043 false,
8044 )
8045 .with_time_index(true),
8046 );
8047 for i in 0..num_field {
8048 columns.push(ColumnSchema::new(
8049 format!("field_{i}"),
8050 ConcreteDataType::float64_datatype(),
8051 true,
8052 ));
8053 }
8054 let schema = Arc::new(Schema::new(columns));
8055 let table_meta = TableMetaBuilder::empty()
8056 .schema(schema)
8057 .primary_key_indices((0..num_tag).collect())
8058 .value_indices((num_tag + 1..num_tag + 1 + num_field).collect())
8059 .next_column_id(1024)
8060 .build()
8061 .unwrap();
8062 let table_info = TableInfoBuilder::default()
8063 .name(table_name.clone())
8064 .meta(table_meta)
8065 .build()
8066 .unwrap();
8067 let table = EmptyTable::from_table_info(&table_info);
8068
8069 assert!(
8070 catalog_list
8071 .register_table_sync(RegisterTableRequest {
8072 catalog: DEFAULT_CATALOG_NAME.to_string(),
8073 schema: schema_name.clone(),
8074 table_name: table_name.clone(),
8075 table_id: 1024,
8076 table,
8077 })
8078 .is_ok()
8079 );
8080 }
8081
8082 DfTableSourceProvider::new(
8083 catalog_list,
8084 false,
8085 QueryContext::arc(),
8086 DummyDecoder::arc(),
8087 false,
8088 )
8089 }
8090
8091 async fn build_test_native_histogram_table_provider(table_name: &str) -> DfTableSourceProvider {
8092 build_test_native_histogram_table_provider_with_marker(table_name, false).await
8093 }
8094
8095 async fn build_test_native_histogram_table_provider_with_marker(
8096 table_name: &str,
8097 temporality_marker: bool,
8098 ) -> DfTableSourceProvider {
8099 let catalog_list = MemoryCatalogManager::with_default_setup();
8100 let mut columns = vec![
8101 ColumnSchema::new(
8102 "tag_0".to_string(),
8103 ConcreteDataType::string_datatype(),
8104 false,
8105 ),
8106 ColumnSchema::new(
8107 LE_COLUMN_NAME.to_string(),
8108 ConcreteDataType::string_datatype(),
8109 true,
8110 ),
8111 ];
8112 if temporality_marker {
8113 columns.push(ColumnSchema::new(
8114 OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(),
8115 ConcreteDataType::string_datatype(),
8116 true,
8117 ));
8118 }
8119 let tag_count = columns.len();
8120 columns.extend([
8121 ColumnSchema::new(
8122 "timestamp".to_string(),
8123 ConcreteDataType::timestamp_millisecond_datatype(),
8124 false,
8125 )
8126 .with_time_index(true),
8127 ColumnSchema::new(
8128 greptime_native_histogram().to_string(),
8129 native_histogram_value_type().clone(),
8130 true,
8131 ),
8132 ]);
8133 let schema = Arc::new(Schema::new(columns));
8134 let table_meta = TableMetaBuilder::empty()
8135 .schema(schema)
8136 .primary_key_indices((0..tag_count).collect())
8137 .value_indices(vec![tag_count + 1])
8138 .next_column_id(1024)
8139 .build()
8140 .unwrap();
8141 let table_info = TableInfoBuilder::default()
8142 .name(table_name)
8143 .meta(table_meta)
8144 .build()
8145 .unwrap();
8146 let table = EmptyTable::from_table_info(&table_info);
8147
8148 assert!(
8149 catalog_list
8150 .register_table_sync(RegisterTableRequest {
8151 catalog: DEFAULT_CATALOG_NAME.to_string(),
8152 schema: DEFAULT_SCHEMA_NAME.to_string(),
8153 table_name: table_name.to_string(),
8154 table_id: 1024,
8155 table,
8156 })
8157 .is_ok()
8158 );
8159
8160 DfTableSourceProvider::new(
8161 catalog_list,
8162 false,
8163 QueryContext::arc(),
8164 DummyDecoder::arc(),
8165 false,
8166 )
8167 }
8168
8169 async fn build_test_multi_histogram_table_provider(table_name: &str) -> DfTableSourceProvider {
8170 let catalog_list = MemoryCatalogManager::with_default_setup();
8171 let columns = vec![
8172 ColumnSchema::new(
8173 "tag_0".to_string(),
8174 ConcreteDataType::string_datatype(),
8175 false,
8176 ),
8177 ColumnSchema::new(
8178 "timestamp".to_string(),
8179 ConcreteDataType::timestamp_millisecond_datatype(),
8180 false,
8181 )
8182 .with_time_index(true),
8183 ColumnSchema::new(
8184 greptime_native_histogram().to_string(),
8185 native_histogram_value_type().clone(),
8186 true,
8187 ),
8188 ColumnSchema::new(
8189 "native_histogram_2".to_string(),
8190 native_histogram_value_type().clone(),
8191 true,
8192 ),
8193 ];
8194 let schema = Arc::new(Schema::new(columns));
8195 let table_meta = TableMetaBuilder::empty()
8196 .schema(schema)
8197 .primary_key_indices(vec![0])
8198 .value_indices(vec![2, 3])
8199 .next_column_id(1024)
8200 .build()
8201 .unwrap();
8202 let table_info = TableInfoBuilder::default()
8203 .name(table_name)
8204 .meta(table_meta)
8205 .build()
8206 .unwrap();
8207 let table = EmptyTable::from_table_info(&table_info);
8208
8209 assert!(
8210 catalog_list
8211 .register_table_sync(RegisterTableRequest {
8212 catalog: DEFAULT_CATALOG_NAME.to_string(),
8213 schema: DEFAULT_SCHEMA_NAME.to_string(),
8214 table_name: table_name.to_string(),
8215 table_id: 1024,
8216 table,
8217 })
8218 .is_ok()
8219 );
8220
8221 DfTableSourceProvider::new(
8222 catalog_list,
8223 false,
8224 QueryContext::arc(),
8225 DummyDecoder::arc(),
8226 false,
8227 )
8228 }
8229
8230 async fn build_test_mixed_native_histogram_table_provider(
8231 table_name: &str,
8232 ) -> DfTableSourceProvider {
8233 build_test_mixed_native_histogram_table_provider_with_marker(table_name, false).await
8234 }
8235
8236 async fn build_test_mixed_native_histogram_table_provider_with_marker(
8237 table_name: &str,
8238 temporality_marker: bool,
8239 ) -> DfTableSourceProvider {
8240 let catalog_list = MemoryCatalogManager::with_default_setup();
8241 let mut columns = vec![ColumnSchema::new(
8242 "tag_0".to_string(),
8243 ConcreteDataType::string_datatype(),
8244 false,
8245 )];
8246 if temporality_marker {
8247 columns.push(ColumnSchema::new(
8248 OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(),
8249 ConcreteDataType::string_datatype(),
8250 true,
8251 ));
8252 }
8253 let tag_count = columns.len();
8254 columns.extend([
8255 ColumnSchema::new(
8256 "timestamp".to_string(),
8257 ConcreteDataType::timestamp_millisecond_datatype(),
8258 false,
8259 )
8260 .with_time_index(true),
8261 ColumnSchema::new(
8262 greptime_native_histogram().to_string(),
8263 native_histogram_value_type().clone(),
8264 true,
8265 ),
8266 ColumnSchema::new(
8267 greptime_value().to_string(),
8268 ConcreteDataType::float64_datatype(),
8269 true,
8270 ),
8271 ]);
8272 let schema = Arc::new(Schema::new(columns));
8273 let table_meta = TableMetaBuilder::empty()
8274 .schema(schema.clone())
8275 .primary_key_indices((0..tag_count).collect())
8276 .value_indices(vec![tag_count + 1, tag_count + 2])
8277 .next_column_id(1024)
8278 .build()
8279 .unwrap();
8280 let table_info = Arc::new(
8281 TableInfoBuilder::default()
8282 .name(table_name)
8283 .meta(table_meta)
8284 .build()
8285 .unwrap(),
8286 );
8287 let mut arrays: Vec<Arc<dyn Array>> =
8288 vec![Arc::new(StringArray::from(vec!["float", "histogram"]))];
8289 if temporality_marker {
8290 arrays.push(Arc::new(StringArray::from(vec![
8291 Some(GREPTIME_TEMPORALITY_DELTA),
8292 Some(GREPTIME_TEMPORALITY_DELTA),
8293 ])));
8294 }
8295 arrays.extend([
8296 Arc::new(TimestampMillisecondArray::from(vec![1_000, 1_000])) as Arc<dyn Array>,
8297 build_histogram_array(&[None, Some(direct_or_histogram())]),
8298 Arc::new(Float64Array::from(vec![Some(2.0), None])),
8299 ]);
8300 let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap();
8301 let backing = GreptimeMemTable::new_with_catalog(
8302 table_name,
8303 GreptimeRecordBatch::from_df_record_batch(schema, batch),
8304 1024,
8305 DEFAULT_CATALOG_NAME.to_string(),
8306 DEFAULT_SCHEMA_NAME.to_string(),
8307 );
8308 let table = Arc::new(Table::new(
8309 table_info,
8310 FilterPushDownType::Unsupported,
8311 backing.data_source(),
8312 ));
8313
8314 assert!(
8315 catalog_list
8316 .register_table_sync(RegisterTableRequest {
8317 catalog: DEFAULT_CATALOG_NAME.to_string(),
8318 schema: DEFAULT_SCHEMA_NAME.to_string(),
8319 table_name: table_name.to_string(),
8320 table_id: 1024,
8321 table,
8322 })
8323 .is_ok()
8324 );
8325
8326 DfTableSourceProvider::new(
8327 catalog_list,
8328 false,
8329 QueryContext::arc(),
8330 DummyDecoder::arc(),
8331 false,
8332 )
8333 }
8334
8335 fn classic_and_native_histogram_table_provider(
8336 native_tag: &str,
8337 native_le: Option<&str>,
8338 native_histogram: NativeHistogram,
8339 ) -> DfTableSourceProvider {
8340 let table_name = "mixed_histogram";
8341 let catalog = MemoryCatalogManager::with_default_setup();
8342 let schema = Arc::new(Schema::new(vec![
8343 ColumnSchema::new(
8344 "tag".to_string(),
8345 ConcreteDataType::string_datatype(),
8346 false,
8347 ),
8348 ColumnSchema::new(
8349 LE_COLUMN_NAME.to_string(),
8350 ConcreteDataType::string_datatype(),
8351 true,
8352 ),
8353 ColumnSchema::new(
8354 "timestamp".to_string(),
8355 ConcreteDataType::timestamp_millisecond_datatype(),
8356 false,
8357 )
8358 .with_time_index(true),
8359 ColumnSchema::new(
8360 greptime_native_histogram().to_string(),
8361 native_histogram_value_type().clone(),
8362 true,
8363 ),
8364 ColumnSchema::new(
8365 greptime_value().to_string(),
8366 ConcreteDataType::float64_datatype(),
8367 true,
8368 ),
8369 ]));
8370 let table_meta = TableMetaBuilder::empty()
8371 .schema(schema.clone())
8372 .primary_key_indices(vec![0, 1])
8373 .value_indices(vec![3, 4])
8374 .next_column_id(5)
8375 .build()
8376 .unwrap();
8377 let table_info = Arc::new(
8378 TableInfoBuilder::default()
8379 .name(table_name)
8380 .meta(table_meta)
8381 .build()
8382 .unwrap(),
8383 );
8384 let batch = RecordBatch::try_new(
8385 schema.arrow_schema().clone(),
8386 vec![
8387 Arc::new(StringArray::from(vec![
8388 "classic", "classic", native_tag, "classic", "classic", native_tag,
8389 ])),
8390 Arc::new(StringArray::from(vec![
8391 Some("1"),
8392 Some("+Inf"),
8393 native_le,
8394 Some("1"),
8395 Some("+Inf"),
8396 native_le,
8397 ])),
8398 Arc::new(TimestampMillisecondArray::from(vec![
8399 1_000, 1_000, 1_000, 2_000, 2_000, 2_000,
8400 ])),
8401 build_histogram_array(&[
8402 None,
8403 None,
8404 Some(native_histogram.clone()),
8405 None,
8406 None,
8407 Some(native_histogram),
8408 ]),
8409 Arc::new(Float64Array::from(vec![
8410 Some(2.0),
8411 Some(4.0),
8412 None,
8413 Some(2.0),
8414 Some(4.0),
8415 None,
8416 ])),
8417 ],
8418 )
8419 .unwrap();
8420 let backing = GreptimeMemTable::new_with_catalog(
8421 table_name,
8422 GreptimeRecordBatch::from_df_record_batch(schema, batch),
8423 2_200,
8424 DEFAULT_CATALOG_NAME.to_string(),
8425 DEFAULT_SCHEMA_NAME.to_string(),
8426 );
8427 let table = Arc::new(Table::new(
8428 table_info,
8429 FilterPushDownType::Unsupported,
8430 backing.data_source(),
8431 ));
8432 catalog
8433 .register_table_sync(RegisterTableRequest {
8434 catalog: DEFAULT_CATALOG_NAME.to_string(),
8435 schema: DEFAULT_SCHEMA_NAME.to_string(),
8436 table_name: table_name.to_string(),
8437 table_id: 2_200,
8438 table,
8439 })
8440 .unwrap();
8441
8442 DfTableSourceProvider::new(
8443 catalog,
8444 false,
8445 QueryContext::arc(),
8446 DummyDecoder::arc(),
8447 false,
8448 )
8449 }
8450
8451 async fn build_test_table_provider_with_tsid(
8452 table_name_tuples: &[(String, String)],
8453 num_tag: usize,
8454 num_field: usize,
8455 ) -> DfTableSourceProvider {
8456 let table_specs = table_name_tuples
8457 .iter()
8458 .map(|(schema_name, table_name)| ((schema_name.clone(), table_name.clone()), num_field))
8459 .collect::<Vec<_>>();
8460 build_test_table_provider_with_tsid_fields(&table_specs, num_tag).await
8461 }
8462
8463 async fn build_test_table_provider_with_tsid_fields(
8464 table_specs: &[((String, String), usize)],
8465 num_tag: usize,
8466 ) -> DfTableSourceProvider {
8467 let table_specs = table_specs
8468 .iter()
8469 .map(|(table_name_tuple, num_field)| (table_name_tuple.clone(), num_tag, *num_field))
8470 .collect::<Vec<_>>();
8471 build_test_table_provider_with_tsid_tag_fields(&table_specs).await
8472 }
8473
8474 async fn build_test_table_provider_with_tsid_tag_fields(
8475 table_specs: &[((String, String), usize, usize)],
8476 ) -> DfTableSourceProvider {
8477 let catalog_list = MemoryCatalogManager::with_default_setup();
8478
8479 let physical_table_name = "phy";
8480 let physical_table_id = 999u32;
8481 let physical_num_tag = table_specs
8482 .iter()
8483 .map(|(_, num_tag, _)| *num_tag)
8484 .max()
8485 .unwrap_or(0);
8486 let physical_num_field = table_specs
8487 .iter()
8488 .map(|(_, _, num_field)| *num_field)
8489 .max()
8490 .unwrap_or(0);
8491
8492 {
8494 let mut columns = vec![
8495 ColumnSchema::new(
8496 DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(),
8497 ConcreteDataType::uint32_datatype(),
8498 false,
8499 ),
8500 ColumnSchema::new(
8501 DATA_SCHEMA_TSID_COLUMN_NAME.to_string(),
8502 ConcreteDataType::uint64_datatype(),
8503 false,
8504 ),
8505 ];
8506 for i in 0..physical_num_tag {
8507 columns.push(ColumnSchema::new(
8508 format!("tag_{i}"),
8509 ConcreteDataType::string_datatype(),
8510 false,
8511 ));
8512 }
8513 columns.push(
8514 ColumnSchema::new(
8515 "timestamp".to_string(),
8516 ConcreteDataType::timestamp_millisecond_datatype(),
8517 false,
8518 )
8519 .with_time_index(true),
8520 );
8521 for i in 0..physical_num_field {
8522 columns.push(ColumnSchema::new(
8523 format!("field_{i}"),
8524 ConcreteDataType::float64_datatype(),
8525 true,
8526 ));
8527 }
8528
8529 let schema = Arc::new(Schema::new(columns));
8530 let primary_key_indices = (0..(2 + physical_num_tag)).collect::<Vec<_>>();
8531 let table_meta = TableMetaBuilder::empty()
8532 .schema(schema)
8533 .primary_key_indices(primary_key_indices)
8534 .value_indices(
8535 (2 + physical_num_tag..2 + physical_num_tag + 1 + physical_num_field).collect(),
8536 )
8537 .engine(METRIC_ENGINE_NAME.to_string())
8538 .next_column_id(1024)
8539 .build()
8540 .unwrap();
8541 let table_info = TableInfoBuilder::default()
8542 .table_id(physical_table_id)
8543 .name(physical_table_name)
8544 .meta(table_meta)
8545 .build()
8546 .unwrap();
8547 let table = EmptyTable::from_table_info(&table_info);
8548
8549 assert!(
8550 catalog_list
8551 .register_table_sync(RegisterTableRequest {
8552 catalog: DEFAULT_CATALOG_NAME.to_string(),
8553 schema: DEFAULT_SCHEMA_NAME.to_string(),
8554 table_name: physical_table_name.to_string(),
8555 table_id: physical_table_id,
8556 table,
8557 })
8558 .is_ok()
8559 );
8560 }
8561
8562 for (idx, ((schema_name, table_name), num_tag, num_field)) in table_specs.iter().enumerate()
8564 {
8565 let mut columns = vec![];
8566 for i in 0..*num_tag {
8567 columns.push(ColumnSchema::new(
8568 format!("tag_{i}"),
8569 ConcreteDataType::string_datatype(),
8570 false,
8571 ));
8572 }
8573 columns.push(
8574 ColumnSchema::new(
8575 "timestamp".to_string(),
8576 ConcreteDataType::timestamp_millisecond_datatype(),
8577 false,
8578 )
8579 .with_time_index(true),
8580 );
8581 for i in 0..*num_field {
8582 columns.push(ColumnSchema::new(
8583 format!("field_{i}"),
8584 ConcreteDataType::float64_datatype(),
8585 true,
8586 ));
8587 }
8588
8589 let schema = Arc::new(Schema::new(columns));
8590 let mut options = table::requests::TableOptions::default();
8591 options.extra_options.insert(
8592 LOGICAL_TABLE_METADATA_KEY.to_string(),
8593 physical_table_name.to_string(),
8594 );
8595 let table_id = 1024u32 + idx as u32;
8596 let table_meta = TableMetaBuilder::empty()
8597 .schema(schema)
8598 .primary_key_indices((0..*num_tag).collect())
8599 .value_indices((*num_tag + 1..*num_tag + 1 + *num_field).collect())
8600 .engine(METRIC_ENGINE_NAME.to_string())
8601 .options(options)
8602 .next_column_id(1024)
8603 .build()
8604 .unwrap();
8605 let table_info = TableInfoBuilder::default()
8606 .table_id(table_id)
8607 .name(table_name.clone())
8608 .meta(table_meta)
8609 .build()
8610 .unwrap();
8611 let table = EmptyTable::from_table_info(&table_info);
8612
8613 assert!(
8614 catalog_list
8615 .register_table_sync(RegisterTableRequest {
8616 catalog: DEFAULT_CATALOG_NAME.to_string(),
8617 schema: schema_name.clone(),
8618 table_name: table_name.clone(),
8619 table_id,
8620 table,
8621 })
8622 .is_ok()
8623 );
8624 }
8625
8626 DfTableSourceProvider::new(
8627 catalog_list,
8628 false,
8629 QueryContext::arc(),
8630 DummyDecoder::arc(),
8631 false,
8632 )
8633 }
8634
8635 async fn build_test_table_provider_with_fields(
8636 table_name_tuples: &[(String, String)],
8637 tags: &[&str],
8638 ) -> DfTableSourceProvider {
8639 let catalog_list = MemoryCatalogManager::with_default_setup();
8640 for (schema_name, table_name) in table_name_tuples {
8641 let mut columns = vec![];
8642 let num_tag = tags.len();
8643 for tag in tags {
8644 columns.push(ColumnSchema::new(
8645 tag.to_string(),
8646 ConcreteDataType::string_datatype(),
8647 false,
8648 ));
8649 }
8650 columns.push(
8651 ColumnSchema::new(
8652 greptime_timestamp().to_string(),
8653 ConcreteDataType::timestamp_millisecond_datatype(),
8654 false,
8655 )
8656 .with_time_index(true),
8657 );
8658 columns.push(ColumnSchema::new(
8659 greptime_value().to_string(),
8660 ConcreteDataType::float64_datatype(),
8661 true,
8662 ));
8663 let schema = Arc::new(Schema::new(columns));
8664 let table_meta = TableMetaBuilder::empty()
8665 .schema(schema)
8666 .primary_key_indices((0..num_tag).collect())
8667 .next_column_id(1024)
8668 .build()
8669 .unwrap();
8670 let table_info = TableInfoBuilder::default()
8671 .name(table_name.clone())
8672 .meta(table_meta)
8673 .build()
8674 .unwrap();
8675 let table = EmptyTable::from_table_info(&table_info);
8676
8677 assert!(
8678 catalog_list
8679 .register_table_sync(RegisterTableRequest {
8680 catalog: DEFAULT_CATALOG_NAME.to_string(),
8681 schema: schema_name.clone(),
8682 table_name: table_name.clone(),
8683 table_id: 1024,
8684 table,
8685 })
8686 .is_ok()
8687 );
8688 }
8689
8690 DfTableSourceProvider::new(
8691 catalog_list,
8692 false,
8693 QueryContext::arc(),
8694 DummyDecoder::arc(),
8695 false,
8696 )
8697 }
8698
8699 async fn do_single_instant_function_call(fn_name: &'static str, plan_name: &str) {
8715 let prom_expr =
8716 parser::parse(&format!("{fn_name}(some_metric{{tag_0!=\"bar\"}})")).unwrap();
8717 let eval_stmt = EvalStmt {
8718 expr: prom_expr,
8719 start: UNIX_EPOCH,
8720 end: UNIX_EPOCH
8721 .checked_add(Duration::from_secs(100_000))
8722 .unwrap(),
8723 interval: Duration::from_secs(5),
8724 lookback_delta: Duration::from_secs(1),
8725 };
8726
8727 let table_provider = build_test_table_provider(
8728 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
8729 1,
8730 1,
8731 )
8732 .await;
8733 let plan =
8734 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
8735 .await
8736 .unwrap();
8737
8738 let expected = String::from(
8739 "Filter: TEMPLATE(field_0) IS NOT NULL [timestamp:Timestamp(ms), TEMPLATE(field_0):Float64;N, tag_0:Utf8]\
8740 \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]\
8741 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
8742 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
8743 \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]\
8744 \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]\
8745 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"
8746 ).replace("TEMPLATE", plan_name);
8747
8748 assert_eq!(plan.display_indent_schema().to_string(), expected);
8749 }
8750
8751 #[tokio::test]
8752 async fn single_abs() {
8753 do_single_instant_function_call("abs", "abs").await;
8754 }
8755
8756 #[tokio::test]
8757 #[should_panic]
8758 async fn single_absent() {
8759 do_single_instant_function_call("absent", "").await;
8760 }
8761
8762 #[tokio::test]
8763 async fn single_ceil() {
8764 do_single_instant_function_call("ceil", "ceil").await;
8765 }
8766
8767 #[tokio::test]
8768 async fn single_exp() {
8769 do_single_instant_function_call("exp", "exp").await;
8770 }
8771
8772 #[tokio::test]
8773 async fn single_ln() {
8774 do_single_instant_function_call("ln", "ln").await;
8775 }
8776
8777 #[tokio::test]
8778 async fn single_log2() {
8779 do_single_instant_function_call("log2", "log2").await;
8780 }
8781
8782 #[tokio::test]
8783 async fn single_log10() {
8784 do_single_instant_function_call("log10", "log10").await;
8785 }
8786
8787 #[tokio::test]
8788 #[should_panic]
8789 async fn single_scalar() {
8790 do_single_instant_function_call("scalar", "").await;
8791 }
8792
8793 #[tokio::test]
8794 #[should_panic]
8795 async fn single_sgn() {
8796 do_single_instant_function_call("sgn", "").await;
8797 }
8798
8799 #[tokio::test]
8800 #[should_panic]
8801 async fn single_sort() {
8802 do_single_instant_function_call("sort", "").await;
8803 }
8804
8805 #[tokio::test]
8806 #[should_panic]
8807 async fn single_sort_desc() {
8808 do_single_instant_function_call("sort_desc", "").await;
8809 }
8810
8811 #[tokio::test]
8812 async fn single_sqrt() {
8813 do_single_instant_function_call("sqrt", "sqrt").await;
8814 }
8815
8816 #[tokio::test]
8817 async fn single_timestamp_plan_preserves_source_value() {
8818 let eval_stmt = build_eval_stmt(r#"timestamp(some_metric{tag_0!="bar"})"#);
8819 let table_provider = build_test_table_provider(
8820 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
8821 1,
8822 1,
8823 )
8824 .await;
8825 let plan =
8826 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
8827 .await
8828 .unwrap();
8829
8830 let expected = String::from(
8831 "Filter: value IS NOT NULL [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\
8832 \n Projection: some_metric.timestamp, value AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\
8833 \n Projection: some_metric.timestamp, __promql_timestamp_value_ AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\
8834 \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]\
8835 \n Projection: some_metric.tag_0, some_metric.timestamp, some_metric.field_0, CAST(CAST(some_metric.timestamp AS Int64) AS Float64) / Float64(1000) AS __promql_timestamp_value_ [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\
8836 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
8837 \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]\
8838 \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]\
8839 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
8840 );
8841
8842 assert_eq!(plan.display_indent_schema().to_string(), expected);
8843 }
8844
8845 #[tokio::test]
8846 async fn single_acos() {
8847 do_single_instant_function_call("acos", "acos").await;
8848 }
8849
8850 #[tokio::test]
8851 #[should_panic]
8852 async fn single_acosh() {
8853 do_single_instant_function_call("acosh", "").await;
8854 }
8855
8856 #[tokio::test]
8857 async fn single_asin() {
8858 do_single_instant_function_call("asin", "asin").await;
8859 }
8860
8861 #[tokio::test]
8862 #[should_panic]
8863 async fn single_asinh() {
8864 do_single_instant_function_call("asinh", "").await;
8865 }
8866
8867 #[tokio::test]
8868 async fn single_atan() {
8869 do_single_instant_function_call("atan", "atan").await;
8870 }
8871
8872 #[tokio::test]
8873 #[should_panic]
8874 async fn single_atanh() {
8875 do_single_instant_function_call("atanh", "").await;
8876 }
8877
8878 #[tokio::test]
8879 async fn single_cos() {
8880 do_single_instant_function_call("cos", "cos").await;
8881 }
8882
8883 #[tokio::test]
8884 #[should_panic]
8885 async fn single_cosh() {
8886 do_single_instant_function_call("cosh", "").await;
8887 }
8888
8889 #[tokio::test]
8890 async fn single_sin() {
8891 do_single_instant_function_call("sin", "sin").await;
8892 }
8893
8894 #[tokio::test]
8895 #[should_panic]
8896 async fn single_sinh() {
8897 do_single_instant_function_call("sinh", "").await;
8898 }
8899
8900 #[tokio::test]
8901 async fn single_tan() {
8902 do_single_instant_function_call("tan", "tan").await;
8903 }
8904
8905 #[tokio::test]
8906 #[should_panic]
8907 async fn single_tanh() {
8908 do_single_instant_function_call("tanh", "").await;
8909 }
8910
8911 #[tokio::test]
8912 #[should_panic]
8913 async fn single_deg() {
8914 do_single_instant_function_call("deg", "").await;
8915 }
8916
8917 #[tokio::test]
8918 #[should_panic]
8919 async fn single_rad() {
8920 do_single_instant_function_call("rad", "").await;
8921 }
8922
8923 async fn do_aggregate_expr_plan(fn_name: &str, plan_name: &str) {
8945 let prom_expr = parser::parse(&format!(
8946 "{fn_name} by (tag_1)(some_metric{{tag_0!=\"bar\"}})",
8947 ))
8948 .unwrap();
8949 let mut eval_stmt = EvalStmt {
8950 expr: prom_expr,
8951 start: UNIX_EPOCH,
8952 end: UNIX_EPOCH
8953 .checked_add(Duration::from_secs(100_000))
8954 .unwrap(),
8955 interval: Duration::from_secs(5),
8956 lookback_delta: Duration::from_secs(1),
8957 };
8958
8959 let table_provider = build_test_table_provider(
8961 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
8962 2,
8963 2,
8964 )
8965 .await;
8966 let plan =
8967 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
8968 .await
8969 .unwrap();
8970 let expected_no_without = String::from(
8971 "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]\
8972 \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]\
8973 \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]\
8974 \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
8975 \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]\
8976 \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]\
8977 \n TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]"
8978 ).replace("TEMPLATE", plan_name);
8979 assert_eq!(
8980 plan.display_indent_schema().to_string(),
8981 expected_no_without
8982 );
8983
8984 if let PromExpr::Aggregate(AggregateExpr { modifier, .. }) = &mut eval_stmt.expr {
8986 *modifier = Some(LabelModifier::Exclude(Labels {
8987 labels: vec![String::from("tag_1")].into_iter().collect(),
8988 }));
8989 }
8990 let table_provider = build_test_table_provider(
8991 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
8992 2,
8993 2,
8994 )
8995 .await;
8996 let plan =
8997 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
8998 .await
8999 .unwrap();
9000 let expected_without = String::from(
9001 "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]\
9002 \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]\
9003 \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]\
9004 \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
9005 \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]\
9006 \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]\
9007 \n TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]"
9008 ).replace("TEMPLATE", plan_name);
9009 assert_eq!(plan.display_indent_schema().to_string(), expected_without);
9010 }
9011
9012 #[tokio::test]
9013 async fn aggregate_sum() {
9014 do_aggregate_expr_plan("sum", "sum").await;
9015 }
9016
9017 #[tokio::test]
9018 async fn tsid_is_used_for_series_divide_when_available() {
9019 let prom_expr = parser::parse("some_metric").unwrap();
9020 let eval_stmt = EvalStmt {
9021 expr: prom_expr,
9022 start: UNIX_EPOCH,
9023 end: UNIX_EPOCH
9024 .checked_add(Duration::from_secs(100_000))
9025 .unwrap(),
9026 interval: Duration::from_secs(5),
9027 lookback_delta: Duration::from_secs(1),
9028 };
9029
9030 let table_provider = build_test_table_provider_with_tsid(
9031 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9032 1,
9033 1,
9034 )
9035 .await;
9036 let plan =
9037 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9038 .await
9039 .unwrap();
9040
9041 let plan_str = plan.display_indent_schema().to_string();
9042 assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
9043 assert!(plan_str.contains("__tsid ASC NULLS FIRST"));
9044 assert!(
9045 !plan
9046 .schema()
9047 .fields()
9048 .iter()
9049 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
9050 );
9051
9052 let manipulate = find_instant_manipulate(&plan).unwrap();
9053 let exec = manipulate.to_execution_plan(Arc::new(DataSourceExec::new(Arc::new(
9054 MemorySourceConfig::try_new(&[], Arc::new(ArrowSchema::empty()), None).unwrap(),
9055 ))));
9056 assert!(format!("{exec:?}").contains("reuse_tsid_column: true"));
9057 }
9058
9059 #[tokio::test]
9060 async fn default_binary_join_uses_tsid_when_available() {
9061 let eval_stmt = build_eval_stmt("some_metric / some_alt_metric");
9062
9063 let table_provider = build_test_table_provider_with_tsid(
9064 &[
9065 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9066 (
9067 DEFAULT_SCHEMA_NAME.to_string(),
9068 "some_alt_metric".to_string(),
9069 ),
9070 ],
9071 1,
9072 1,
9073 )
9074 .await;
9075 let plan =
9076 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9077 .await
9078 .unwrap();
9079
9080 let plan_str = plan.display_indent_schema().to_string();
9081 assert!(
9082 plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
9083 "{plan_str}"
9084 );
9085 assert!(
9086 !plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
9087 "{plan_str}"
9088 );
9089 }
9090
9091 #[tokio::test]
9092 async fn reject_binary_fill_modifiers() {
9093 let state = build_query_engine_state();
9094
9095 for query in [
9096 "some_metric + fill(0) some_alt_metric",
9097 "some_metric + fill_left(0) some_alt_metric",
9098 "some_metric + fill_right(0) some_alt_metric",
9099 "(some_metric + fill(0) some_alt_metric) + some_metric",
9100 ] {
9101 let eval_stmt = build_eval_stmt(query);
9102 let table_provider = build_test_table_provider(&[], 0, 0).await;
9103 let err = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &state)
9104 .await
9105 .unwrap_err();
9106
9107 assert!(
9108 matches!(
9109 &err,
9110 crate::promql::error::Error::UnsupportedExpr { name, .. }
9111 if name == "PromQL fill modifiers"
9112 ),
9113 "{err}"
9114 );
9115 }
9116 }
9117
9118 #[tokio::test]
9119 async fn timestamp_binary_join_falls_back_when_tsid_is_projected_out() {
9120 for query in [
9121 "timestamp(some_metric) / some_metric",
9122 "some_metric / timestamp(some_metric)",
9123 ] {
9124 let eval_stmt = build_eval_stmt(query);
9125
9126 let table_provider = build_test_table_provider_with_tsid(
9127 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9128 1,
9129 1,
9130 )
9131 .await;
9132 let plan =
9133 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9134 .await
9135 .unwrap();
9136
9137 let plan_str = plan.display_indent_schema().to_string();
9138 assert!(!plan_str.contains("__tsid ="), "{query}: {plan_str}");
9139 assert!(
9140 plan_str.contains("lhs.tag_0 = rhs.tag_0"),
9141 "{query}: {plan_str}"
9142 );
9143 assert!(
9144 !plan
9145 .schema()
9146 .fields()
9147 .iter()
9148 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME),
9149 "{query}: {plan_str}"
9150 );
9151 }
9152 }
9153
9154 #[tokio::test]
9155 async fn timestamp_binary_join_rejects_default_matching_on_mismatched_labels() {
9156 let eval_stmt = build_eval_stmt("timestamp(left_host_job) / right_by_job");
9157
9158 let table_provider = build_test_table_provider_with_tsid_tag_fields(&[
9159 (
9160 (DEFAULT_SCHEMA_NAME.to_string(), "left_host_job".to_string()),
9161 2,
9162 1,
9163 ),
9164 (
9165 (DEFAULT_SCHEMA_NAME.to_string(), "right_by_job".to_string()),
9166 1,
9167 1,
9168 ),
9169 ])
9170 .await;
9171 let plan =
9172 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9173 .await
9174 .unwrap();
9175 let plan_str = plan.display_indent_schema().to_string();
9176
9177 assert!(
9178 plan_str.contains("Boolean(false)") || plan_str.contains("false"),
9179 "{plan_str}"
9180 );
9181 }
9182
9183 #[tokio::test]
9184 async fn tsid_is_preserved_for_nested_default_binary_joins() {
9185 let eval_stmt = build_eval_stmt("(some_metric - some_alt_metric) / some_third_metric");
9186
9187 let table_provider = build_test_table_provider_with_tsid(
9188 &[
9189 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9190 (
9191 DEFAULT_SCHEMA_NAME.to_string(),
9192 "some_alt_metric".to_string(),
9193 ),
9194 (
9195 DEFAULT_SCHEMA_NAME.to_string(),
9196 "some_third_metric".to_string(),
9197 ),
9198 ],
9199 1,
9200 1,
9201 )
9202 .await;
9203 let plan =
9204 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9205 .await
9206 .unwrap();
9207
9208 let plan_str = plan.display_indent_schema().to_string();
9209 assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
9210 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9211 }
9212
9213 #[tokio::test]
9214 async fn repeated_tsid_binary_operand_reuses_leaf_plan() {
9215 let eval_stmt = build_eval_stmt("((some_metric - some_alt_metric) / some_metric) * 100");
9216
9217 let table_provider = build_test_table_provider_with_tsid(
9218 &[
9219 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9220 (
9221 DEFAULT_SCHEMA_NAME.to_string(),
9222 "some_alt_metric".to_string(),
9223 ),
9224 ],
9225 1,
9226 1,
9227 )
9228 .await;
9229 let plan =
9230 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9231 .await
9232 .unwrap();
9233
9234 let plan_str = plan.display_indent_schema().to_string();
9235 assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}");
9236 assert_eq!(
9237 plan_str
9238 .matches("Filter: phy.__table_id = UInt32(1024)")
9239 .count(),
9240 1,
9241 "{plan_str}"
9242 );
9243 assert_eq!(
9244 plan_str.matches("PromInstantManipulate").count(),
9245 2,
9246 "{plan_str}"
9247 );
9248 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9249 }
9250
9251 #[tokio::test]
9252 async fn repeated_tsid_binary_operand_reuses_shorter_field_side() {
9253 let eval_stmt =
9254 build_eval_stmt("((two_field_metric - one_field_metric) / one_field_metric) * 100");
9255
9256 let table_provider = build_test_table_provider_with_tsid_fields(
9257 &[
9258 (
9259 (
9260 DEFAULT_SCHEMA_NAME.to_string(),
9261 "two_field_metric".to_string(),
9262 ),
9263 2,
9264 ),
9265 (
9266 (
9267 DEFAULT_SCHEMA_NAME.to_string(),
9268 "one_field_metric".to_string(),
9269 ),
9270 1,
9271 ),
9272 ],
9273 1,
9274 )
9275 .await;
9276 let plan =
9277 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9278 .await
9279 .unwrap();
9280
9281 let field_names = plan
9282 .schema()
9283 .fields()
9284 .iter()
9285 .map(|field| field.name().clone())
9286 .collect::<Vec<_>>();
9287 let value_columns = field_names
9288 .iter()
9289 .filter(|name| {
9290 *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME
9291 })
9292 .count();
9293 assert_eq!(value_columns, 1, "{field_names:?}");
9294 let plan_str = plan.display_indent_schema().to_string();
9295 assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}");
9296 assert_eq!(
9297 plan_str
9298 .matches("Filter: phy.__table_id = UInt32(1025)")
9299 .count(),
9300 1,
9301 "{plan_str}"
9302 );
9303 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9304 }
9305
9306 #[tokio::test]
9307 async fn binary_island_reuses_self_operand_without_join() {
9308 let eval_stmt = build_eval_stmt("some_metric / some_metric");
9309
9310 let table_provider = build_test_table_provider_with_tsid(
9311 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9312 1,
9313 1,
9314 )
9315 .await;
9316 let plan =
9317 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9318 .await
9319 .unwrap();
9320
9321 let plan_str = plan.display_indent_schema().to_string();
9322 assert_eq!(plan_str.matches("__tsid =").count(), 0, "{plan_str}");
9323 assert_eq!(
9324 plan_str
9325 .matches("Filter: phy.__table_id = UInt32(1024)")
9326 .count(),
9327 1,
9328 "{plan_str}"
9329 );
9330 assert_eq!(
9331 plan_str.matches("PromInstantManipulate").count(),
9332 1,
9333 "{plan_str}"
9334 );
9335 }
9336
9337 #[tokio::test]
9338 async fn binary_island_reuses_leaf_across_two_branches() {
9339 let eval_stmt =
9340 build_eval_stmt("(some_metric + some_alt_metric) / (some_metric + third_metric)");
9341
9342 let table_provider = build_test_table_provider_with_tsid(
9343 &[
9344 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9345 (
9346 DEFAULT_SCHEMA_NAME.to_string(),
9347 "some_alt_metric".to_string(),
9348 ),
9349 (DEFAULT_SCHEMA_NAME.to_string(), "third_metric".to_string()),
9350 ],
9351 1,
9352 1,
9353 )
9354 .await;
9355 let plan =
9356 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9357 .await
9358 .unwrap();
9359
9360 let plan_str = plan.display_indent_schema().to_string();
9361 assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
9362 assert_eq!(
9363 plan_str
9364 .matches("Filter: phy.__table_id = UInt32(1024)")
9365 .count(),
9366 1,
9367 "{plan_str}"
9368 );
9369 assert_eq!(
9370 plan_str.matches("PromInstantManipulate").count(),
9371 3,
9372 "{plan_str}"
9373 );
9374 }
9375
9376 #[tokio::test]
9377 async fn binary_island_generated_alias_avoids_user_column_names() {
9378 let eval_stmt = build_eval_stmt("(some_metric + some_alt_metric) / some_metric");
9379
9380 let table_provider = build_test_table_provider_with_fields(
9381 &[
9382 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9383 (
9384 DEFAULT_SCHEMA_NAME.to_string(),
9385 "some_alt_metric".to_string(),
9386 ),
9387 ],
9388 &["prom_v0", "__prom_v0"],
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 field_names = plan.schema().field_names();
9397 assert!(field_names.iter().any(|name| name.ends_with(".prom_v0")));
9398 assert!(field_names.iter().any(|name| name.ends_with(".__prom_v0")));
9399
9400 let plan_str = plan.display_indent_schema().to_string();
9401 assert!(plan_str.contains("SubqueryAlias: __prom_v0"), "{plan_str}");
9402 assert_eq!(
9403 plan_str.matches("PromInstantManipulate").count(),
9404 2,
9405 "{plan_str}"
9406 );
9407 }
9408
9409 #[tokio::test]
9410 async fn binary_island_clears_qualifier_for_nested_unary_projection() {
9411 let eval_stmt = build_eval_stmt("-((some_metric + some_alt_metric) / some_metric)");
9412
9413 let table_provider = build_test_table_provider_with_tsid(
9414 &[
9415 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9416 (
9417 DEFAULT_SCHEMA_NAME.to_string(),
9418 "some_alt_metric".to_string(),
9419 ),
9420 ],
9421 1,
9422 1,
9423 )
9424 .await;
9425 let plan =
9426 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9427 .await
9428 .unwrap();
9429
9430 let plan_str = plan.display_indent_schema().to_string();
9431 assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}");
9432 assert_eq!(
9433 plan_str.matches("PromInstantManipulate").count(),
9434 2,
9435 "{plan_str}"
9436 );
9437 }
9438
9439 #[tokio::test]
9440 async fn binary_island_keeps_distinct_matcher_leaves() {
9441 let eval_stmt = build_eval_stmt(
9442 "(some_metric{tag_0=\"foo\"} + some_alt_metric) / some_metric{tag_0=\"bar\"}",
9443 );
9444
9445 let table_provider = build_test_table_provider_with_tsid(
9446 &[
9447 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9448 (
9449 DEFAULT_SCHEMA_NAME.to_string(),
9450 "some_alt_metric".to_string(),
9451 ),
9452 ],
9453 1,
9454 1,
9455 )
9456 .await;
9457 let plan =
9458 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9459 .await
9460 .unwrap();
9461
9462 let plan_str = plan.display_indent_schema().to_string();
9463 assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
9464 assert_eq!(
9465 plan_str.matches("PromInstantManipulate").count(),
9466 3,
9467 "{plan_str}"
9468 );
9469 }
9470
9471 #[tokio::test]
9472 async fn binary_island_keeps_offset_leaves_distinct() {
9473 let eval_stmt = build_eval_stmt("(some_metric offset 5m + some_alt_metric) / some_metric");
9474
9475 let table_provider = build_test_table_provider_with_tsid(
9476 &[
9477 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9478 (
9479 DEFAULT_SCHEMA_NAME.to_string(),
9480 "some_alt_metric".to_string(),
9481 ),
9482 ],
9483 1,
9484 1,
9485 )
9486 .await;
9487 let plan =
9488 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9489 .await
9490 .unwrap();
9491
9492 let plan_str = plan.display_indent_schema().to_string();
9493 assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
9494 assert_eq!(
9495 plan_str.matches("PromInstantManipulate").count(),
9496 3,
9497 "{plan_str}"
9498 );
9499 }
9500
9501 #[tokio::test]
9502 async fn binary_island_falls_back_for_group_modifier() {
9503 let eval_stmt = build_eval_stmt(
9504 "(some_metric + ignoring(tag_0) group_left some_alt_metric) / some_metric",
9505 );
9506
9507 let table_provider = build_test_table_provider_with_tsid(
9508 &[
9509 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9510 (
9511 DEFAULT_SCHEMA_NAME.to_string(),
9512 "some_alt_metric".to_string(),
9513 ),
9514 ],
9515 1,
9516 1,
9517 )
9518 .await;
9519 let plan =
9520 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9521 .await
9522 .unwrap();
9523
9524 let plan_str = plan.display_indent_schema().to_string();
9525 assert_eq!(
9526 plan_str.matches("PromInstantManipulate").count(),
9527 3,
9528 "{plan_str}"
9529 );
9530 }
9531
9532 #[tokio::test]
9533 async fn binary_island_falls_back_for_comparison_filter() {
9534 let eval_stmt = build_eval_stmt("(some_metric > some_alt_metric) / some_metric");
9535
9536 let table_provider = build_test_table_provider_with_tsid(
9537 &[
9538 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9539 (
9540 DEFAULT_SCHEMA_NAME.to_string(),
9541 "some_alt_metric".to_string(),
9542 ),
9543 ],
9544 1,
9545 1,
9546 )
9547 .await;
9548 let plan =
9549 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9550 .await
9551 .unwrap();
9552
9553 let plan_str = plan.display_indent_schema().to_string();
9554 assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
9555 assert_eq!(
9556 plan_str.matches("PromInstantManipulate").count(),
9557 3,
9558 "{plan_str}"
9559 );
9560 }
9561
9562 #[tokio::test]
9563 async fn tsid_binary_join_uses_shorter_field_side() {
9564 let eval_stmt = build_eval_stmt("one_field_metric / two_field_metric");
9565
9566 let table_provider = build_test_table_provider_with_tsid_fields(
9567 &[
9568 (
9569 (
9570 DEFAULT_SCHEMA_NAME.to_string(),
9571 "one_field_metric".to_string(),
9572 ),
9573 1,
9574 ),
9575 (
9576 (
9577 DEFAULT_SCHEMA_NAME.to_string(),
9578 "two_field_metric".to_string(),
9579 ),
9580 2,
9581 ),
9582 ],
9583 1,
9584 )
9585 .await;
9586 let plan =
9587 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9588 .await
9589 .unwrap();
9590
9591 let field_names = plan
9592 .schema()
9593 .fields()
9594 .iter()
9595 .map(|field| field.name().clone())
9596 .collect::<Vec<_>>();
9597 let value_columns = field_names
9598 .iter()
9599 .filter(|name| {
9600 *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME
9601 })
9602 .count();
9603 assert_eq!(value_columns, 1, "{field_names:?}");
9604 }
9605
9606 #[tokio::test]
9607 async fn comparison_binary_join_uses_shorter_field_side() {
9608 let eval_stmt = build_eval_stmt("two_field_metric > one_field_metric");
9609
9610 let table_provider = build_test_table_provider_with_tsid_fields(
9611 &[
9612 (
9613 (
9614 DEFAULT_SCHEMA_NAME.to_string(),
9615 "two_field_metric".to_string(),
9616 ),
9617 2,
9618 ),
9619 (
9620 (
9621 DEFAULT_SCHEMA_NAME.to_string(),
9622 "one_field_metric".to_string(),
9623 ),
9624 1,
9625 ),
9626 ],
9627 1,
9628 )
9629 .await;
9630 let plan =
9631 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9632 .await
9633 .unwrap();
9634
9635 let field_names = plan
9636 .schema()
9637 .fields()
9638 .iter()
9639 .map(|field| field.name().clone())
9640 .collect::<Vec<_>>();
9641 assert!(
9642 field_names.iter().any(|name| name == "field_0"),
9643 "{field_names:?}"
9644 );
9645 assert!(
9646 !field_names.iter().any(|name| name == "field_1"),
9647 "{field_names:?}"
9648 );
9649 }
9650
9651 #[tokio::test]
9652 async fn label_matching_modifier_disables_tsid_binary_join() {
9653 let eval_stmt = build_eval_stmt("some_metric / ignoring(tag_0) some_alt_metric");
9654
9655 let table_provider = build_test_table_provider_with_tsid(
9656 &[
9657 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9658 (
9659 DEFAULT_SCHEMA_NAME.to_string(),
9660 "some_alt_metric".to_string(),
9661 ),
9662 ],
9663 2,
9664 1,
9665 )
9666 .await;
9667 let plan =
9668 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9669 .await
9670 .unwrap();
9671
9672 let plan_str = plan.display_indent_schema().to_string();
9673 assert!(!plan_str.contains("__tsid ="), "{plan_str}");
9674 assert!(
9675 plan_str.contains("some_metric.tag_1 = some_alt_metric.tag_1"),
9676 "{plan_str}"
9677 );
9678 }
9679
9680 #[tokio::test]
9681 async fn ignoring_absent_label_keeps_tsid_binary_join() {
9682 let eval_stmt = build_eval_stmt("some_metric / ignoring(missing) some_alt_metric");
9683
9684 let table_provider = build_test_table_provider_with_tsid(
9685 &[
9686 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9687 (
9688 DEFAULT_SCHEMA_NAME.to_string(),
9689 "some_alt_metric".to_string(),
9690 ),
9691 ],
9692 2,
9693 1,
9694 )
9695 .await;
9696 let plan =
9697 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9698 .await
9699 .unwrap();
9700
9701 let plan_str = plan.display_indent_schema().to_string();
9702 assert!(
9703 plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
9704 "{plan_str}"
9705 );
9706 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9707 assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
9708 }
9709
9710 #[tokio::test]
9711 async fn range_function_keeps_tsid_for_absent_ignoring_binary_join() {
9712 let eval_stmt =
9713 build_eval_stmt("rate(some_metric[5m]) / ignoring(missing) some_alt_metric");
9714
9715 let table_provider = build_test_table_provider_with_tsid(
9716 &[
9717 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9718 (
9719 DEFAULT_SCHEMA_NAME.to_string(),
9720 "some_alt_metric".to_string(),
9721 ),
9722 ],
9723 2,
9724 1,
9725 )
9726 .await;
9727 let plan =
9728 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9729 .await
9730 .unwrap();
9731
9732 let plan_str = plan.display_indent_schema().to_string();
9733 assert!(
9734 plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
9735 "{plan_str}"
9736 );
9737 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9738 assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
9739 }
9740
9741 #[tokio::test]
9742 async fn on_full_label_set_keeps_tsid_binary_join() {
9743 let eval_stmt = build_eval_stmt("some_metric / on(tag_0, tag_1) some_alt_metric");
9744
9745 let table_provider = build_test_table_provider_with_tsid(
9746 &[
9747 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9748 (
9749 DEFAULT_SCHEMA_NAME.to_string(),
9750 "some_alt_metric".to_string(),
9751 ),
9752 ],
9753 2,
9754 1,
9755 )
9756 .await;
9757 let plan =
9758 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9759 .await
9760 .unwrap();
9761
9762 let plan_str = plan.display_indent_schema().to_string();
9763 assert!(
9764 plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
9765 "{plan_str}"
9766 );
9767 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9768 assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
9769 }
9770
9771 #[tokio::test]
9772 async fn on_partial_label_set_disables_tsid_binary_join() {
9773 let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric");
9774
9775 let table_provider = build_test_table_provider_with_tsid(
9776 &[
9777 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9778 (
9779 DEFAULT_SCHEMA_NAME.to_string(),
9780 "some_alt_metric".to_string(),
9781 ),
9782 ],
9783 2,
9784 1,
9785 )
9786 .await;
9787 let plan =
9788 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9789 .await
9790 .unwrap();
9791
9792 let plan_str = plan.display_indent_schema().to_string();
9793 assert!(!plan_str.contains("__tsid ="), "{plan_str}");
9794 assert!(
9795 plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
9796 "{plan_str}"
9797 );
9798 assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
9799 }
9800
9801 #[tokio::test]
9802 async fn on_label_set_must_cover_both_sides_to_use_tsid_binary_join() {
9803 let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric");
9804
9805 let table_provider = build_test_table_provider_with_tsid_tag_fields(&[
9806 (
9807 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9808 2,
9809 1,
9810 ),
9811 (
9812 (
9813 DEFAULT_SCHEMA_NAME.to_string(),
9814 "some_alt_metric".to_string(),
9815 ),
9816 1,
9817 1,
9818 ),
9819 ])
9820 .await;
9821 let plan =
9822 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9823 .await
9824 .unwrap();
9825
9826 let plan_str = plan.display_indent_schema().to_string();
9827 assert!(!plan_str.contains("__tsid ="), "{plan_str}");
9828 assert!(
9829 plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
9830 "{plan_str}"
9831 );
9832 assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
9833 }
9834
9835 #[tokio::test]
9836 async fn comparison_binary_join_uses_tsid_and_keeps_it_in_filtered_result() {
9837 let eval_stmt = build_eval_stmt("some_metric > some_alt_metric");
9838
9839 let table_provider = build_test_table_provider_with_tsid(
9840 &[
9841 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9842 (
9843 DEFAULT_SCHEMA_NAME.to_string(),
9844 "some_alt_metric".to_string(),
9845 ),
9846 ],
9847 2,
9848 1,
9849 )
9850 .await;
9851 let mut planner = PromPlanner {
9852 table_provider,
9853 ctx: PromPlannerContext::from_eval_stmt(&eval_stmt),
9854 promql_annotations: None,
9855 };
9856 let plan = planner
9857 .prom_expr_to_plan(&eval_stmt.expr, &build_query_engine_state())
9858 .await
9859 .unwrap();
9860
9861 let plan_str = plan.display_indent_schema().to_string();
9862 assert!(
9863 plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
9864 "{plan_str}"
9865 );
9866 assert!(
9867 plan.schema()
9868 .fields()
9869 .iter()
9870 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME),
9871 "{plan_str}"
9872 );
9873 assert!(planner.ctx.use_tsid, "{plan_str}");
9874 }
9875
9876 #[tokio::test]
9877 async fn comparison_bool_binary_join_uses_tsid_when_available() {
9878 let eval_stmt = build_eval_stmt("some_metric > bool some_alt_metric");
9879
9880 let table_provider = build_test_table_provider_with_tsid(
9881 &[
9882 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9883 (
9884 DEFAULT_SCHEMA_NAME.to_string(),
9885 "some_alt_metric".to_string(),
9886 ),
9887 ],
9888 2,
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!(
9899 plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
9900 "{plan_str}"
9901 );
9902 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9903 assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
9904 }
9905
9906 #[tokio::test]
9907 async fn scalar_count_count_range_keeps_full_window() {
9908 let plan_str = build_optimized_tsid_plan(
9909 "scalar(count(count(some_metric) by (tag_0)))",
9910 1,
9911 1,
9912 100_000,
9913 1,
9914 )
9915 .await;
9916 assert!(plan_str.contains("ScalarCalculate: tags=[]"));
9917 assert!(plan_str.contains("PromInstantManipulate: range=[0..100000000]"));
9918 assert!(!plan_str.contains("PromInstantManipulate: range=[99999000..99999000]"));
9919 }
9920
9921 #[tokio::test]
9922 async fn scalar_count_count_rewrite_applies_inside_binary_expr_for_tsid_input() {
9923 let plan_str = build_optimized_tsid_plan(
9924 "sum(irate(some_metric[1h])) / scalar(count(count(some_metric) by (tag_0)))",
9925 2,
9926 1,
9927 10,
9928 300,
9929 )
9930 .await;
9931 assert!(plan_str.contains("Distinct:"), "{plan_str}");
9932 }
9933
9934 #[tokio::test]
9935 async fn nested_count_rewrite_keeps_full_series_key_with_tsid_input() {
9936 assert_nested_count_rewrite_applies(
9937 "count(count(some_metric) by (tag_0))",
9938 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(count(some_metric.field_0))]]"
9939 )
9940 .await;
9941 }
9942
9943 #[tokio::test]
9944 async fn nested_sum_count_rewrite_keeps_full_series_key_with_tsid_input() {
9945 assert_nested_count_rewrite_applies(
9946 "count(sum(some_metric) by (tag_0))",
9947 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(sum(some_metric.field_0))]]"
9948 )
9949 .await;
9950 }
9951
9952 #[tokio::test]
9953 async fn nested_supported_inner_aggs_rewrite_apply_for_tsid_input() {
9954 for (query, expected_outer_agg) in [
9955 (
9956 "count(avg(some_metric) by (tag_0))",
9957 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(avg(some_metric.field_0))]]",
9958 ),
9959 (
9960 "count(min(some_metric) by (tag_0))",
9961 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(min(some_metric.field_0))]]",
9962 ),
9963 (
9964 "count(max(some_metric) by (tag_0))",
9965 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(max(some_metric.field_0))]]",
9966 ),
9967 (
9968 "count(stddev(some_metric) by (tag_0))",
9969 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(stddev_pop(some_metric.field_0))]]",
9970 ),
9971 (
9972 "count(stdvar(some_metric) by (tag_0))",
9973 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(var_pop(some_metric.field_0))]]",
9974 ),
9975 ] {
9976 assert_nested_count_rewrite_applies(query, expected_outer_agg).await;
9977 }
9978 }
9979
9980 #[tokio::test]
9981 async fn nested_non_count_inner_aggs_rewrite_filter_null_values_for_tsid_input() {
9982 let count_plan =
9983 build_optimized_tsid_plan("count(count(some_metric) by (tag_0))", 2, 1, 100_000, 1)
9984 .await;
9985 assert!(
9986 !count_plan.contains("some_metric.field_0 IS NOT NULL"),
9987 "{count_plan}"
9988 );
9989
9990 for query in [
9991 "count(sum(some_metric) by (tag_0))",
9992 "count(avg(some_metric) by (tag_0))",
9993 "count(min(some_metric) by (tag_0))",
9994 "count(max(some_metric) by (tag_0))",
9995 "count(stddev(some_metric) by (tag_0))",
9996 "count(stdvar(some_metric) by (tag_0))",
9997 ] {
9998 let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await;
9999 assert!(
10000 plan_str.contains("Filter: some_metric.field_0 IS NOT NULL"),
10001 "{query}: {plan_str}"
10002 );
10003 }
10004 }
10005
10006 #[tokio::test]
10007 async fn nested_unsupported_or_non_direct_inner_aggs_do_not_rewrite() {
10008 assert_nested_count_rewrite_missing("count(group(some_metric) by (tag_0))", 2, 1).await;
10009 assert_nested_count_rewrite_missing(
10010 "count(sum(irate(some_metric[1h])) by (tag_0))",
10011 2,
10012 300,
10013 )
10014 .await;
10015 }
10016
10017 #[tokio::test]
10018 async fn physical_table_name_is_not_leaked_in_plan() {
10019 let prom_expr = parser::parse("some_metric").unwrap();
10020 let eval_stmt = EvalStmt {
10021 expr: prom_expr,
10022 start: UNIX_EPOCH,
10023 end: UNIX_EPOCH
10024 .checked_add(Duration::from_secs(100_000))
10025 .unwrap(),
10026 interval: Duration::from_secs(5),
10027 lookback_delta: Duration::from_secs(1),
10028 };
10029
10030 let table_provider = build_test_table_provider_with_tsid(
10031 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10032 1,
10033 1,
10034 )
10035 .await;
10036 let plan =
10037 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10038 .await
10039 .unwrap();
10040
10041 let plan_str = plan.display_indent_schema().to_string();
10042 assert!(plan_str.contains("TableScan: phy"), "{plan}");
10043 assert!(plan_str.contains("SubqueryAlias: some_metric"));
10044 assert!(plan_str.contains("Filter: phy.__table_id = UInt32(1024)"));
10045 assert!(!plan_str.contains("TableScan: some_metric"));
10046 }
10047
10048 #[tokio::test]
10049 async fn sum_without_does_not_group_by_tsid() {
10050 let prom_expr = parser::parse("sum without (tag_0) (some_metric)").unwrap();
10051 let eval_stmt = EvalStmt {
10052 expr: prom_expr,
10053 start: UNIX_EPOCH,
10054 end: UNIX_EPOCH
10055 .checked_add(Duration::from_secs(100_000))
10056 .unwrap(),
10057 interval: Duration::from_secs(5),
10058 lookback_delta: Duration::from_secs(1),
10059 };
10060
10061 let table_provider = build_test_table_provider_with_tsid(
10062 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10063 1,
10064 1,
10065 )
10066 .await;
10067 let plan =
10068 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10069 .await
10070 .unwrap();
10071
10072 let plan_str = plan.display_indent_schema().to_string();
10073 assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10074
10075 let aggr_line = plan_str
10076 .lines()
10077 .find(|line| line.contains("Aggregate: groupBy="))
10078 .unwrap();
10079 assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
10080 }
10081
10082 #[tokio::test]
10083 async fn topk_without_does_not_partition_by_tsid() {
10084 let prom_expr = parser::parse("topk without (tag_0) (1, some_metric)").unwrap();
10085 let eval_stmt = EvalStmt {
10086 expr: prom_expr,
10087 start: UNIX_EPOCH,
10088 end: UNIX_EPOCH
10089 .checked_add(Duration::from_secs(100_000))
10090 .unwrap(),
10091 interval: Duration::from_secs(5),
10092 lookback_delta: Duration::from_secs(1),
10093 };
10094
10095 let table_provider = build_test_table_provider_with_tsid(
10096 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10097 1,
10098 1,
10099 )
10100 .await;
10101 let plan =
10102 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10103 .await
10104 .unwrap();
10105
10106 let plan_str = plan.display_indent_schema().to_string();
10107 assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10108
10109 let window_line = plan_str
10110 .lines()
10111 .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()"))
10112 .unwrap();
10113 let partition_by = window_line
10114 .split("PARTITION BY [")
10115 .nth(1)
10116 .and_then(|s| s.split("] ORDER BY").next())
10117 .unwrap();
10118 assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
10119 }
10120
10121 #[tokio::test]
10122 async fn sum_by_does_not_group_by_tsid() {
10123 let prom_expr = parser::parse("sum by (__tsid) (some_metric)").unwrap();
10124 let eval_stmt = EvalStmt {
10125 expr: prom_expr,
10126 start: UNIX_EPOCH,
10127 end: UNIX_EPOCH
10128 .checked_add(Duration::from_secs(100_000))
10129 .unwrap(),
10130 interval: Duration::from_secs(5),
10131 lookback_delta: Duration::from_secs(1),
10132 };
10133
10134 let table_provider = build_test_table_provider_with_tsid(
10135 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10136 1,
10137 1,
10138 )
10139 .await;
10140 let plan =
10141 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10142 .await
10143 .unwrap();
10144
10145 let plan_str = plan.display_indent_schema().to_string();
10146 assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10147
10148 let aggr_line = plan_str
10149 .lines()
10150 .find(|line| line.contains("Aggregate: groupBy="))
10151 .unwrap();
10152 assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
10153 }
10154
10155 #[tokio::test]
10156 async fn aggregate_over_binary_time_function_expr() {
10157 for op in ["sum", "min", "max", "avg"] {
10158 let prom_expr = parser::parse(&format!(
10159 "{op} by (tag_0, tag_1, tag_2) (time() - some_metric)"
10160 ))
10161 .unwrap();
10162 let eval_stmt = EvalStmt {
10163 expr: prom_expr,
10164 start: UNIX_EPOCH,
10165 end: UNIX_EPOCH
10166 .checked_add(Duration::from_secs(100_000))
10167 .unwrap(),
10168 interval: Duration::from_secs(5),
10169 lookback_delta: Duration::from_secs(1),
10170 };
10171
10172 let table_provider = build_test_table_provider_with_tsid(
10173 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10174 3,
10175 1,
10176 )
10177 .await;
10178 let plan =
10179 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10180 .await
10181 .unwrap();
10182
10183 let plan_str = plan.display_indent_schema().to_string();
10184 let aggr_line = plan_str
10185 .lines()
10186 .find(|line| line.contains("Aggregate: groupBy="))
10187 .unwrap();
10188 assert!(aggr_line.contains(op), "{plan_str}");
10189 assert!(aggr_line.contains("first_value"), "{plan_str}");
10190 assert!(
10191 !plan
10192 .schema()
10193 .fields()
10194 .iter()
10195 .any(|field| { field.name() == DATA_SCHEMA_TSID_COLUMN_NAME })
10196 );
10197 }
10198 }
10199
10200 #[tokio::test]
10201 async fn topk_by_does_not_partition_by_tsid() {
10202 let prom_expr = parser::parse("topk by (__tsid) (1, some_metric)").unwrap();
10203 let eval_stmt = EvalStmt {
10204 expr: prom_expr,
10205 start: UNIX_EPOCH,
10206 end: UNIX_EPOCH
10207 .checked_add(Duration::from_secs(100_000))
10208 .unwrap(),
10209 interval: Duration::from_secs(5),
10210 lookback_delta: Duration::from_secs(1),
10211 };
10212
10213 let table_provider = build_test_table_provider_with_tsid(
10214 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10215 1,
10216 1,
10217 )
10218 .await;
10219 let plan =
10220 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10221 .await
10222 .unwrap();
10223
10224 let plan_str = plan.display_indent_schema().to_string();
10225 assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10226
10227 let window_line = plan_str
10228 .lines()
10229 .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()"))
10230 .unwrap();
10231 let partition_by = window_line
10232 .split("PARTITION BY [")
10233 .nth(1)
10234 .and_then(|s| s.split("] ORDER BY").next())
10235 .unwrap();
10236 assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
10237 }
10238
10239 #[tokio::test]
10240 async fn selector_matcher_on_tsid_does_not_use_internal_column() {
10241 let prom_expr = parser::parse(r#"some_metric{__tsid="123"}"#).unwrap();
10242 let eval_stmt = EvalStmt {
10243 expr: prom_expr,
10244 start: UNIX_EPOCH,
10245 end: UNIX_EPOCH
10246 .checked_add(Duration::from_secs(100_000))
10247 .unwrap(),
10248 interval: Duration::from_secs(5),
10249 lookback_delta: Duration::from_secs(1),
10250 };
10251
10252 let table_provider = build_test_table_provider_with_tsid(
10253 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10254 1,
10255 1,
10256 )
10257 .await;
10258 let plan =
10259 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10260 .await
10261 .unwrap();
10262
10263 fn collect_filter_cols(plan: &LogicalPlan, out: &mut HashSet<Column>) {
10264 if let LogicalPlan::Filter(filter) = plan {
10265 datafusion_expr::utils::expr_to_columns(&filter.predicate, out).unwrap();
10266 }
10267 for input in plan.inputs() {
10268 collect_filter_cols(input, out);
10269 }
10270 }
10271
10272 let mut filter_cols = HashSet::new();
10273 collect_filter_cols(&plan, &mut filter_cols);
10274 assert!(
10275 !filter_cols
10276 .iter()
10277 .any(|c| c.name == DATA_SCHEMA_TSID_COLUMN_NAME)
10278 );
10279 }
10280
10281 #[tokio::test]
10282 async fn tsid_is_not_used_when_physical_table_is_missing() {
10283 let prom_expr = parser::parse("some_metric").unwrap();
10284 let eval_stmt = EvalStmt {
10285 expr: prom_expr,
10286 start: UNIX_EPOCH,
10287 end: UNIX_EPOCH
10288 .checked_add(Duration::from_secs(100_000))
10289 .unwrap(),
10290 interval: Duration::from_secs(5),
10291 lookback_delta: Duration::from_secs(1),
10292 };
10293
10294 let catalog_list = MemoryCatalogManager::with_default_setup();
10295
10296 let mut columns = vec![ColumnSchema::new(
10298 "tag_0".to_string(),
10299 ConcreteDataType::string_datatype(),
10300 false,
10301 )];
10302 columns.push(
10303 ColumnSchema::new(
10304 "timestamp".to_string(),
10305 ConcreteDataType::timestamp_millisecond_datatype(),
10306 false,
10307 )
10308 .with_time_index(true),
10309 );
10310 columns.push(ColumnSchema::new(
10311 "field_0".to_string(),
10312 ConcreteDataType::float64_datatype(),
10313 true,
10314 ));
10315 let schema = Arc::new(Schema::new(columns));
10316 let mut options = table::requests::TableOptions::default();
10317 options
10318 .extra_options
10319 .insert(LOGICAL_TABLE_METADATA_KEY.to_string(), "phy".to_string());
10320 let table_meta = TableMetaBuilder::empty()
10321 .schema(schema)
10322 .primary_key_indices(vec![0])
10323 .value_indices(vec![2])
10324 .engine(METRIC_ENGINE_NAME.to_string())
10325 .options(options)
10326 .next_column_id(1024)
10327 .build()
10328 .unwrap();
10329 let table_info = TableInfoBuilder::default()
10330 .table_id(1024)
10331 .name("some_metric")
10332 .meta(table_meta)
10333 .build()
10334 .unwrap();
10335 let table = EmptyTable::from_table_info(&table_info);
10336 catalog_list
10337 .register_table_sync(RegisterTableRequest {
10338 catalog: DEFAULT_CATALOG_NAME.to_string(),
10339 schema: DEFAULT_SCHEMA_NAME.to_string(),
10340 table_name: "some_metric".to_string(),
10341 table_id: 1024,
10342 table,
10343 })
10344 .unwrap();
10345
10346 let table_provider = DfTableSourceProvider::new(
10347 catalog_list,
10348 false,
10349 QueryContext::arc(),
10350 DummyDecoder::arc(),
10351 false,
10352 );
10353
10354 let plan =
10355 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10356 .await
10357 .unwrap();
10358
10359 let plan_str = plan.display_indent_schema().to_string();
10360 assert!(plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]"));
10361 assert!(!plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10362 }
10363
10364 #[tokio::test]
10365 async fn tsid_is_carried_only_when_aggregate_preserves_label_set() {
10366 let prom_expr = parser::parse("sum by (tag_0) (some_metric)").unwrap();
10367 let eval_stmt = EvalStmt {
10368 expr: prom_expr,
10369 start: UNIX_EPOCH,
10370 end: UNIX_EPOCH
10371 .checked_add(Duration::from_secs(100_000))
10372 .unwrap(),
10373 interval: Duration::from_secs(5),
10374 lookback_delta: Duration::from_secs(1),
10375 };
10376
10377 let table_provider = build_test_table_provider_with_tsid(
10378 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10379 1,
10380 1,
10381 )
10382 .await;
10383 let plan =
10384 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10385 .await
10386 .unwrap();
10387
10388 let plan_str = plan.display_indent_schema().to_string();
10389 assert!(plan_str.contains("first_value") && plan_str.contains("__tsid"));
10390 assert!(
10391 !plan
10392 .schema()
10393 .fields()
10394 .iter()
10395 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
10396 );
10397
10398 let prom_expr = parser::parse("sum(some_metric)").unwrap();
10400 let eval_stmt = EvalStmt {
10401 expr: prom_expr,
10402 start: UNIX_EPOCH,
10403 end: UNIX_EPOCH
10404 .checked_add(Duration::from_secs(100_000))
10405 .unwrap(),
10406 interval: Duration::from_secs(5),
10407 lookback_delta: Duration::from_secs(1),
10408 };
10409 let table_provider = build_test_table_provider_with_tsid(
10410 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10411 1,
10412 1,
10413 )
10414 .await;
10415 let plan =
10416 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10417 .await
10418 .unwrap();
10419 let plan_str = plan.display_indent_schema().to_string();
10420 assert!(!plan_str.contains("first_value"));
10421 }
10422
10423 #[tokio::test]
10424 async fn or_operator_with_unknown_metric_does_not_require_tsid() {
10425 let prom_expr = parser::parse("unknown_metric or some_metric").unwrap();
10426 let eval_stmt = EvalStmt {
10427 expr: prom_expr,
10428 start: UNIX_EPOCH,
10429 end: UNIX_EPOCH
10430 .checked_add(Duration::from_secs(100_000))
10431 .unwrap(),
10432 interval: Duration::from_secs(5),
10433 lookback_delta: Duration::from_secs(1),
10434 };
10435
10436 let table_provider = build_test_table_provider_with_tsid(
10437 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10438 1,
10439 1,
10440 )
10441 .await;
10442
10443 let plan =
10444 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10445 .await
10446 .unwrap();
10447
10448 assert!(
10449 !plan
10450 .schema()
10451 .fields()
10452 .iter()
10453 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
10454 );
10455 }
10456
10457 #[tokio::test]
10458 async fn aggregate_avg() {
10459 do_aggregate_expr_plan("avg", "avg").await;
10460 }
10461
10462 #[tokio::test]
10463 #[should_panic] async fn aggregate_count() {
10465 do_aggregate_expr_plan("count", "count").await;
10466 }
10467
10468 #[tokio::test]
10469 async fn aggregate_min() {
10470 do_aggregate_expr_plan("min", "min").await;
10471 }
10472
10473 #[tokio::test]
10474 async fn aggregate_max() {
10475 do_aggregate_expr_plan("max", "max").await;
10476 }
10477
10478 #[tokio::test]
10479 async fn aggregate_group() {
10480 let prom_expr = parser::parse(
10484 "sum(group by (cluster)(kubernetes_build_info{service=\"kubernetes\",job=\"apiserver\"}))",
10485 )
10486 .unwrap();
10487 let eval_stmt = EvalStmt {
10488 expr: prom_expr,
10489 start: UNIX_EPOCH,
10490 end: UNIX_EPOCH
10491 .checked_add(Duration::from_secs(100_000))
10492 .unwrap(),
10493 interval: Duration::from_secs(5),
10494 lookback_delta: Duration::from_secs(1),
10495 };
10496
10497 let table_provider = build_test_table_provider_with_fields(
10498 &[(
10499 DEFAULT_SCHEMA_NAME.to_string(),
10500 "kubernetes_build_info".to_string(),
10501 )],
10502 &["cluster", "service", "job"],
10503 )
10504 .await;
10505 let plan =
10506 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10507 .await
10508 .unwrap();
10509
10510 let plan_str = plan.display_indent_schema().to_string();
10511 assert!(plan_str.contains("max(Float64(1"));
10512 }
10513
10514 #[tokio::test]
10515 async fn aggregate_stddev() {
10516 do_aggregate_expr_plan("stddev", "stddev_pop").await;
10517 }
10518
10519 #[tokio::test]
10520 async fn aggregate_stdvar() {
10521 do_aggregate_expr_plan("stdvar", "var_pop").await;
10522 }
10523
10524 #[tokio::test]
10548 async fn binary_op_column_column() {
10549 let prom_expr =
10550 parser::parse(r#"some_metric{tag_0="foo"} + some_metric{tag_0="bar"}"#).unwrap();
10551 let eval_stmt = EvalStmt {
10552 expr: prom_expr,
10553 start: UNIX_EPOCH,
10554 end: UNIX_EPOCH
10555 .checked_add(Duration::from_secs(100_000))
10556 .unwrap(),
10557 interval: Duration::from_secs(5),
10558 lookback_delta: Duration::from_secs(1),
10559 };
10560
10561 let table_provider = build_test_table_provider(
10562 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10563 1,
10564 1,
10565 )
10566 .await;
10567 let plan =
10568 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10569 .await
10570 .unwrap();
10571
10572 let expected = String::from(
10573 "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]\
10574 \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]\
10575 \n SubqueryAlias: lhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10576 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10577 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10578 \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]\
10579 \n Filter: some_metric.tag_0 = Utf8(\"foo\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10580 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10581 \n SubqueryAlias: rhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10582 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10583 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10584 \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]\
10585 \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]\
10586 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
10587 );
10588
10589 assert_eq!(plan.display_indent_schema().to_string(), expected);
10590 }
10591
10592 async fn indie_query_plan_compare<T: AsRef<str>>(query: &str, expected: T) {
10593 let prom_expr = parser::parse(query).unwrap();
10594 let eval_stmt = EvalStmt {
10595 expr: prom_expr,
10596 start: UNIX_EPOCH,
10597 end: UNIX_EPOCH
10598 .checked_add(Duration::from_secs(100_000))
10599 .unwrap(),
10600 interval: Duration::from_secs(5),
10601 lookback_delta: Duration::from_secs(1),
10602 };
10603
10604 let table_provider = build_test_table_provider(
10605 &[
10606 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10607 (
10608 "greptime_private".to_string(),
10609 "some_alt_metric".to_string(),
10610 ),
10611 ],
10612 1,
10613 1,
10614 )
10615 .await;
10616 let plan =
10617 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10618 .await
10619 .unwrap();
10620
10621 assert_eq!(plan.display_indent_schema().to_string(), expected.as_ref());
10622 }
10623
10624 #[tokio::test]
10625 async fn binary_op_literal_column() {
10626 let query = r#"1 + some_metric{tag_0="bar"}"#;
10627 let expected = String::from(
10628 "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]\
10629 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10630 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10631 \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]\
10632 \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]\
10633 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
10634 );
10635
10636 indie_query_plan_compare(query, expected).await;
10637 }
10638
10639 #[tokio::test]
10640 async fn binary_op_literal_literal() {
10641 let query = r#"1 + 1"#;
10642 let expected = r#"EmptyMetric: range=[0..100000000], interval=[5000] [time:Timestamp(ms), value:Float64;N]
10643 TableScan: dummy [time:Timestamp(ms), value:Float64;N]"#;
10644 indie_query_plan_compare(query, expected).await;
10645 }
10646
10647 #[tokio::test]
10648 async fn simple_bool_grammar() {
10649 let query = "some_metric != bool 1.2345";
10650 let expected = String::from(
10651 "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]\
10652 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10653 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10654 \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]\
10655 \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]\
10656 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
10657 );
10658
10659 indie_query_plan_compare(query, expected).await;
10660 }
10661
10662 #[tokio::test]
10663 async fn bool_with_additional_arithmetic() {
10664 let query = "some_metric + (1 == bool 2)";
10665 let expected = String::from(
10666 "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]\
10667 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10668 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10669 \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]\
10670 \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]\
10671 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
10672 );
10673
10674 indie_query_plan_compare(query, expected).await;
10675 }
10676
10677 #[tokio::test]
10678 async fn simple_unary() {
10679 let query = "-some_metric";
10680 let expected = String::from(
10681 "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]\
10682 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10683 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10684 \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]\
10685 \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]\
10686 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
10687 );
10688
10689 indie_query_plan_compare(query, expected).await;
10690 }
10691
10692 #[tokio::test]
10693 async fn increase_aggr() {
10694 let query = "increase(some_metric[5m])";
10695 let expected = String::from(
10696 "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]\
10697 \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]\
10698 \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))]\
10699 \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10700 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10701 \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]\
10702 \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]\
10703 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
10704 );
10705
10706 indie_query_plan_compare(query, expected).await;
10707 }
10708
10709 async fn native_histogram_plan(query: &str) -> String {
10710 let table_provider = build_test_native_histogram_table_provider("some_metric").await;
10711 let plan = PromPlanner::stmt_to_plan(
10712 table_provider,
10713 &build_eval_stmt(query),
10714 &build_query_engine_state(),
10715 )
10716 .await
10717 .unwrap();
10718 plan.display_indent_schema().to_string()
10719 }
10720
10721 #[tokio::test]
10722 async fn native_histogram_count_uses_native_udf() {
10723 let plan = native_histogram_plan("histogram_count(some_metric)").await;
10724
10725 assert!(plan.contains("prom_native_histogram_count"), "{plan}");
10726 assert!(!plan.contains("HistogramFold:"), "{plan}");
10727 }
10728
10729 #[tokio::test]
10730 async fn timestamp_filters_native_histogram_stale_marker_before_projection() {
10731 let mut stale = direct_or_histogram();
10732 stale.sum = f64::from_bits(PROMETHEUS_STALE_NAN_BITS);
10733 let table = operator_metric_table(
10734 "stale_histogram",
10735 2_100,
10736 "a",
10737 None,
10738 DirectOrValue::NativeHistogram(stale),
10739 );
10740 let catalog = MemoryCatalogManager::with_default_setup();
10741 catalog
10742 .register_table_sync(RegisterTableRequest {
10743 catalog: DEFAULT_CATALOG_NAME.to_string(),
10744 schema: DEFAULT_SCHEMA_NAME.to_string(),
10745 table_name: "stale_histogram".to_string(),
10746 table_id: 2_100,
10747 table,
10748 })
10749 .unwrap();
10750 let provider = DfTableSourceProvider::new(
10751 catalog,
10752 false,
10753 QueryContext::arc(),
10754 DummyDecoder::arc(),
10755 false,
10756 );
10757 let state = build_query_engine_state();
10758 let plan = PromPlanner::stmt_to_plan(
10759 provider,
10760 &operator_eval_stmt("timestamp(stale_histogram)"),
10761 &state,
10762 )
10763 .await
10764 .unwrap();
10765 let plan_text = plan.display_indent_schema().to_string();
10766 assert!(plan_text.contains(TIMESTAMP_VALUE_PREFIX), "{plan_text}");
10767
10768 let (_, batches) = execute(plan, &state).await;
10769 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
10770 }
10771
10772 #[tokio::test]
10773 async fn timestamp_filters_stale_marker_from_mixed_sample_companion() {
10774 let histograms = build_histogram_array(&[None]);
10775 let schema = Arc::new(ArrowSchema::new(vec![
10776 Field::new(
10777 "timestamp",
10778 ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
10779 false,
10780 ),
10781 Field::new(
10782 greptime_native_histogram(),
10783 histograms.data_type().clone(),
10784 true,
10785 ),
10786 Field::new(greptime_value(), ArrowDataType::Float64, true),
10787 ]));
10788 let batch = RecordBatch::try_new(
10789 schema.clone(),
10790 vec![
10791 Arc::new(TimestampMillisecondArray::from(vec![1_000])),
10792 histograms,
10793 Arc::new(Float64Array::from(vec![f64::from_bits(
10794 PROMETHEUS_STALE_NAN_BITS,
10795 )])),
10796 ],
10797 )
10798 .unwrap();
10799 let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap());
10800 let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None)
10801 .unwrap()
10802 .build()
10803 .unwrap();
10804 let input = LogicalPlan::Extension(Extension {
10805 node: Arc::new(SeriesDivide::new(
10806 Vec::new(),
10807 "timestamp".to_string(),
10808 input,
10809 )),
10810 });
10811 let input = LogicalPlan::Extension(Extension {
10812 node: Arc::new(InstantManipulate::new(
10813 1_000,
10814 1_000,
10815 5_000,
10816 1_000,
10817 "timestamp".to_string(),
10818 Vec::new(),
10819 Some(greptime_native_histogram().to_string()),
10820 input,
10821 )),
10822 });
10823 let plan = LogicalPlanBuilder::from(input)
10825 .project([col("timestamp")])
10826 .unwrap()
10827 .build()
10828 .unwrap();
10829
10830 let (_, batches) = execute(plan, &build_query_engine_state()).await;
10831 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
10832 }
10833
10834 #[tokio::test]
10835 async fn native_histogram_rate_can_feed_count() {
10836 let plan = native_histogram_plan("histogram_count(rate(some_metric[5m]))").await;
10837
10838 assert!(plan.contains("prom_native_histogram_rate"), "{plan}");
10839 assert!(plan.contains("prom_native_histogram_count"), "{plan}");
10840 }
10841
10842 #[tokio::test]
10843 async fn native_histogram_quantile_skips_classic_fold() {
10844 let plan = native_histogram_plan("histogram_quantile(0.9, some_metric)").await;
10845
10846 assert!(plan.contains("prom_native_histogram_quantile"), "{plan}");
10847 assert!(!plan.contains("HistogramFold:"), "{plan}");
10848 assert!(plan.contains("some_metric.le"), "{plan}");
10849 assert!(plan.contains("Float64(0.9)"), "{plan}");
10851 assert!(plan.contains("IS NOT NULL"), "{plan}");
10854 }
10855
10856 #[tokio::test]
10857 async fn mixed_native_histogram_quantile_uses_histogram_field() {
10858 let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
10859 let plan = PromPlanner::stmt_to_plan(
10860 table_provider,
10861 &build_eval_stmt("histogram_quantile(0.9, some_metric)"),
10862 &build_query_engine_state(),
10863 )
10864 .await
10865 .unwrap()
10866 .display_indent_schema()
10867 .to_string();
10868
10869 assert!(
10870 plan.contains("prom_native_histogram_quantile(greptime_native_histogram"),
10871 "{plan}"
10872 );
10873 assert!(!plan.contains("EmptyRelation"), "{plan}");
10874 }
10875
10876 #[tokio::test]
10877 async fn mixed_histogram_helpers_execute_classic_and_native_samples() {
10878 let state = build_query_engine_state();
10879 for (query, expected) in [
10880 (
10881 "histogram_quantile(0.5, mixed_histogram)",
10882 vec![("classic", 1.0), ("native", 0.0)],
10883 ),
10884 (
10885 "histogram_fraction(-Inf, +Inf, mixed_histogram)",
10886 vec![("classic", 1.0), ("native", 1.0)],
10887 ),
10888 ] {
10889 let plan = PromPlanner::stmt_to_plan(
10890 classic_and_native_histogram_table_provider("native", None, direct_or_histogram()),
10891 &operator_eval_stmt(query),
10892 &state,
10893 )
10894 .await
10895 .unwrap();
10896 let plan_text = plan.display_indent_schema().to_string();
10897 assert!(plan_text.contains("HistogramFold:"), "{plan_text}");
10898 assert!(plan_text.contains("prom_native_histogram_"), "{plan_text}");
10899 let value_field = plan
10900 .schema()
10901 .fields()
10902 .iter()
10903 .find(|field| field.data_type() == &ArrowDataType::Float64)
10904 .unwrap()
10905 .name()
10906 .clone();
10907
10908 let (_, batches) = execute(plan, &state).await;
10909 let mut actual = batches
10910 .iter()
10911 .flat_map(|batch| {
10912 let tags = batch
10913 .column_by_name("tag")
10914 .unwrap()
10915 .as_any()
10916 .downcast_ref::<StringArray>()
10917 .unwrap();
10918 let values = batch
10919 .column_by_name(&value_field)
10920 .unwrap()
10921 .as_any()
10922 .downcast_ref::<Float64Array>()
10923 .unwrap();
10924 (0..batch.num_rows()).map(|row| (tags.value(row), values.value(row)))
10925 })
10926 .collect::<Vec<_>>();
10927 actual.sort_by_key(|(tag, _)| *tag);
10928 assert_eq!(actual, expected, "{query}");
10929 }
10930 }
10931
10932 #[tokio::test]
10933 async fn mixed_histogram_helpers_report_annotations() {
10934 let state = build_query_engine_state();
10935 let mut native_histogram = direct_or_histogram();
10936 native_histogram.count = 2.0;
10937 native_histogram.sum = f64::NAN;
10938 for (native_tag, expected_rows, expected_warnings, expected_infos) in [
10939 (
10940 "classic",
10941 0,
10942 vec!["vector contains a mix of classic and native histograms"],
10943 vec![],
10944 ),
10945 (
10946 "native",
10947 2,
10948 vec![],
10949 vec!["input to histogram_quantile has NaN observations, result is skewed higher"],
10950 ),
10951 ] {
10952 let collector = PromqlAnnotationCollector::default();
10953 let plan = PromPlanner::stmt_to_plan_with_annotations(
10954 classic_and_native_histogram_table_provider(
10955 native_tag,
10956 None,
10957 native_histogram.clone(),
10958 ),
10959 &operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"),
10960 &state,
10961 Some(collector.clone()),
10962 )
10963 .await
10964 .unwrap();
10965
10966 let (_, batches) = execute(plan, &state).await;
10967 assert_eq!(
10968 batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
10969 expected_rows
10970 );
10971 let mut warnings = vec![];
10972 let mut infos = vec![];
10973 collector.append_to(&mut warnings, &mut infos);
10974 assert_eq!(warnings, expected_warnings);
10975 assert_eq!(infos, expected_infos);
10976 }
10977 }
10978
10979 #[tokio::test]
10980 async fn mixed_histogram_helper_preserves_native_le_and_scans_once() {
10981 let state = build_query_engine_state();
10982 let mut stmt = operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)");
10983 stmt.end = UNIX_EPOCH.checked_add(Duration::from_secs(2)).unwrap();
10984 let plan = PromPlanner::stmt_to_plan(
10985 classic_and_native_histogram_table_provider(
10986 "classic",
10987 Some("native"),
10988 direct_or_histogram(),
10989 ),
10990 &stmt,
10991 &state,
10992 )
10993 .await
10994 .unwrap();
10995 let plan_text = plan.display_indent_schema().to_string();
10996 assert_eq!(
10997 plan_text.matches("TableScan: mixed_histogram").count(),
10998 1,
10999 "{plan_text}"
11000 );
11001
11002 let value_field = plan
11003 .schema()
11004 .fields()
11005 .iter()
11006 .find(|field| field.data_type() == &ArrowDataType::Float64)
11007 .unwrap()
11008 .name()
11009 .clone();
11010 let (_, batches) = execute(plan, &state).await;
11011 let mut actual = batches
11012 .iter()
11013 .flat_map(|batch| {
11014 let le = batch
11015 .column_by_name(LE_COLUMN_NAME)
11016 .unwrap()
11017 .as_any()
11018 .downcast_ref::<StringArray>()
11019 .unwrap();
11020 let timestamps = batch
11021 .column_by_name("timestamp")
11022 .unwrap()
11023 .as_any()
11024 .downcast_ref::<TimestampMillisecondArray>()
11025 .unwrap();
11026 let values = batch
11027 .column_by_name(&value_field)
11028 .unwrap()
11029 .as_any()
11030 .downcast_ref::<Float64Array>()
11031 .unwrap();
11032 (0..batch.num_rows()).map(|row| {
11033 (
11034 timestamps.value(row),
11035 (!le.is_null(row)).then(|| le.value(row).to_string()),
11036 values.value(row),
11037 )
11038 })
11039 })
11040 .collect::<Vec<_>>();
11041 actual.sort_by(|lhs, rhs| (lhs.0, &lhs.1).cmp(&(rhs.0, &rhs.1)));
11042 assert_eq!(
11043 actual,
11044 vec![
11045 (1_000, None, 1.0),
11046 (1_000, Some("native".to_string()), 0.0),
11047 (2_000, None, 1.0),
11048 (2_000, Some("native".to_string()), 0.0),
11049 ]
11050 );
11051 }
11052
11053 #[tokio::test]
11054 async fn nested_histogram_helpers_ignore_unparsable_bucket_labels() {
11055 let state = build_query_engine_state();
11056 for native_le in [None, Some("native")] {
11057 for query in [
11058 "histogram_quantile(0.5, histogram_quantile(0.5, mixed_histogram))",
11059 "histogram_fraction(-Inf, +Inf, histogram_fraction(-Inf, +Inf, mixed_histogram))",
11060 ] {
11061 let plan = PromPlanner::stmt_to_plan(
11062 classic_and_native_histogram_table_provider(
11063 "native",
11064 native_le,
11065 direct_or_histogram(),
11066 ),
11067 &operator_eval_stmt(query),
11068 &state,
11069 )
11070 .await
11071 .unwrap();
11072
11073 let (_, batches) = execute(plan, &state).await;
11074 assert_eq!(
11075 batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
11076 0,
11077 "native_le={native_le:?}, query={query}"
11078 );
11079 }
11080 }
11081 }
11082
11083 #[tokio::test]
11084 async fn native_histogram_quantile_rejects_multi_field_input() {
11085 let table_provider = build_test_multi_histogram_table_provider("some_metric").await;
11086 let result = PromPlanner::stmt_to_plan(
11087 table_provider,
11088 &build_eval_stmt("histogram_quantile(0.9, some_metric)"),
11089 &build_query_engine_state(),
11090 )
11091 .await;
11092
11093 let err = result.expect_err("histogram_quantile on two native histogram fields must fail");
11094 assert!(
11095 err.to_string()
11096 .contains("Multi fields calculation is not supported in histogram_quantile"),
11097 "{err}"
11098 );
11099 }
11100
11101 #[tokio::test]
11102 async fn native_histogram_topk_uses_drop_udf() {
11103 let plan = native_histogram_plan("topk(1, some_metric)").await;
11104
11105 assert!(plan.contains("prom_native_histogram_drop_float"), "{plan}");
11106 assert!(
11107 plan.contains("Filter: prom_native_histogram_drop_float")
11108 && plan.contains("IS NOT NULL"),
11109 "{plan}"
11110 );
11111 }
11112
11113 #[tokio::test]
11114 async fn mixed_or_topk_bottomk_ignore_native_histograms() {
11115 for op in ["topk", "bottomk"] {
11116 let collector = PromqlAnnotationCollector::default();
11117 let state = build_query_engine_state();
11118 let plan = PromPlanner::stmt_to_plan_with_annotations(
11119 operator_table_provider(),
11120 &operator_eval_stmt(&format!("{op}(1, lf or on(tag) lh)")),
11121 &state,
11122 Some(collector.clone()),
11123 )
11124 .await
11125 .unwrap();
11126 let float_field = plan
11127 .schema()
11128 .fields()
11129 .iter()
11130 .find(|field| field.data_type() == &ArrowDataType::Float64)
11131 .unwrap()
11132 .name()
11133 .clone();
11134 assert!(
11135 plan.schema()
11136 .fields()
11137 .iter()
11138 .all(|field| field.data_type() != &PromPlanner::native_histogram_arrow_type()),
11139 "{plan:?}"
11140 );
11141
11142 let (_, batches) = execute(plan, &state).await;
11143 assert_eq!(values(&batches, &float_field), vec![2.0], "{op}");
11144 let mut warnings = vec![];
11145 let mut infos = vec![];
11146 collector.append_to(&mut warnings, &mut infos);
11147 assert!(warnings.is_empty());
11148 assert_eq!(
11149 infos,
11150 vec![format!(
11151 "{op}: dropped native histogram samples because this aggregation is not supported for native histograms"
11152 )]
11153 );
11154 }
11155 }
11156
11157 #[tokio::test]
11158 async fn native_histogram_scalar_is_ignored_before_scalar_calculate() {
11159 let plan = native_histogram_plan("scalar(some_metric)").await;
11160
11161 assert!(plan.contains("ScalarCalculate"), "{plan}");
11162 assert!(plan.contains("Filter: Boolean(false)"), "{plan}");
11163 assert!(!plan.contains("prom_native_histogram_drop"), "{plan}");
11164 }
11165
11166 #[tokio::test]
11167 async fn native_histogram_value_sort_is_empty_but_label_sort_preserves_samples() {
11168 for function in ["sort", "sort_desc"] {
11169 let plan = native_histogram_plan(&format!("{function}(some_metric)")).await;
11170
11171 assert!(plan.contains("Float64(NULL) IS NOT NULL"), "{plan}");
11172 assert!(
11173 !plan.contains(&format!("Sort: {}", greptime_native_histogram())),
11174 "{plan}"
11175 );
11176 assert!(!plan.contains("prom_native_histogram_drop"), "{plan}");
11177 }
11178
11179 for (function, direction) in [("sort_by_label", "ASC"), ("sort_by_label_desc", "DESC")] {
11180 let plan = native_histogram_plan(&format!("{function}(some_metric, \"tag_0\")")).await;
11181
11182 assert!(plan.contains(&format!("tag_0 {direction}")), "{plan}");
11183 assert!(plan.contains(greptime_native_histogram()), "{plan}");
11184 assert!(!plan.contains("Float64(NULL) IS NOT NULL"), "{plan}");
11185 }
11186 }
11187
11188 #[tokio::test]
11189 async fn unsupported_native_histogram_functions_use_drop_udf() {
11190 for query in [
11191 "deriv(some_metric[5m])",
11192 "min_over_time(some_metric[5m])",
11193 "quantile_over_time(0.9, some_metric[5m])",
11194 "predict_linear(some_metric[5m], 60)",
11195 "round(some_metric)",
11196 "abs(some_metric)",
11197 ] {
11198 let plan = native_histogram_plan(query).await;
11199
11200 assert!(
11201 plan.contains("prom_native_histogram_drop_float"),
11202 "{query}\n{plan}"
11203 );
11204 }
11205 }
11206
11207 #[tokio::test]
11208 async fn native_histogram_absent_over_time_uses_native_udf() {
11209 let plan = native_histogram_plan("absent_over_time(some_metric[5m])").await;
11210
11211 assert!(
11212 plan.contains("prom_native_histogram_absent_over_time"),
11213 "{plan}"
11214 );
11215 }
11216
11217 #[tokio::test]
11218 async fn native_histogram_all_function_arms_route_correctly() {
11219 let cases = [
11224 (
11226 "increase(some_metric[5m])",
11227 "prom_native_histogram_increase",
11228 ),
11229 ("rate(some_metric[5m])", "prom_native_histogram_rate"),
11230 ("delta(some_metric[5m])", "prom_native_histogram_delta"),
11231 ("idelta(some_metric[5m])", "prom_native_histogram_idelta"),
11232 ("irate(some_metric[5m])", "prom_native_histogram_irate"),
11233 ("resets(some_metric[5m])", "prom_native_histogram_resets"),
11234 ("changes(some_metric[5m])", "prom_native_histogram_changes"),
11235 (
11236 "avg_over_time(some_metric[5m])",
11237 "prom_native_histogram_avg_over_time",
11238 ),
11239 (
11240 "sum_over_time(some_metric[5m])",
11241 "prom_native_histogram_sum_over_time",
11242 ),
11243 (
11244 "count_over_time(some_metric[5m])",
11245 "prom_native_histogram_count_over_time",
11246 ),
11247 (
11248 "last_over_time(some_metric[5m])",
11249 "prom_native_histogram_last_over_time",
11250 ),
11251 (
11252 "present_over_time(some_metric[5m])",
11253 "prom_native_histogram_present_over_time",
11254 ),
11255 ("deriv(some_metric[5m])", "prom_native_histogram_drop_float"),
11257 (
11258 "min_over_time(some_metric[5m])",
11259 "prom_native_histogram_drop_float",
11260 ),
11261 (
11262 "max_over_time(some_metric[5m])",
11263 "prom_native_histogram_drop_float",
11264 ),
11265 (
11266 "stddev_over_time(some_metric[5m])",
11267 "prom_native_histogram_drop_float",
11268 ),
11269 (
11270 "stdvar_over_time(some_metric[5m])",
11271 "prom_native_histogram_drop_float",
11272 ),
11273 (
11274 "quantile_over_time(0.9, some_metric[5m])",
11275 "prom_native_histogram_drop_float",
11276 ),
11277 (
11278 "predict_linear(some_metric[5m], 60)",
11279 "prom_native_histogram_drop_float",
11280 ),
11281 (
11282 "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)",
11283 "prom_native_histogram_drop_float",
11284 ),
11285 ("round(some_metric)", "prom_native_histogram_drop_float"),
11286 ("rad(some_metric)", "prom_native_histogram_drop_float"),
11287 ("deg(some_metric)", "prom_native_histogram_drop_float"),
11288 ("sgn(some_metric)", "prom_native_histogram_drop_float"),
11289 (
11291 "histogram_count(some_metric)",
11292 "prom_native_histogram_count",
11293 ),
11294 ("histogram_sum(some_metric)", "prom_native_histogram_sum"),
11295 ("histogram_avg(some_metric)", "prom_native_histogram_avg"),
11296 (
11297 "histogram_stddev(some_metric)",
11298 "prom_native_histogram_stddev",
11299 ),
11300 (
11301 "histogram_stdvar(some_metric)",
11302 "prom_native_histogram_stdvar",
11303 ),
11304 (
11305 "histogram_fraction(-2 + 1, 2 / 2, some_metric)",
11306 "prom_native_histogram_fraction",
11307 ),
11308 ];
11309
11310 for (query, expected_udf) in cases {
11311 let plan = native_histogram_plan(query).await;
11312 assert!(plan.contains(expected_udf), "{query}\n{plan}");
11313 if query.starts_with("histogram_fraction") {
11314 assert!(plan.contains("Float64(-1)"), "{query}\n{plan}");
11315 }
11316 }
11317 }
11318
11319 #[tokio::test]
11320 async fn mixed_native_histogram_ranges_use_coordinated_udfs() {
11321 let dual_output = [
11322 "increase(some_metric[5m])",
11323 "rate(some_metric[5m])",
11324 "delta(some_metric[5m])",
11325 "idelta(some_metric[5m])",
11326 "irate(some_metric[5m])",
11327 "avg_over_time(some_metric[5m])",
11328 "sum_over_time(some_metric[5m])",
11329 "last_over_time(some_metric[5m])",
11330 ];
11331 let float_output = [
11332 "resets(some_metric[5m])",
11333 "changes(some_metric[5m])",
11334 "deriv(some_metric[5m])",
11335 "min_over_time(some_metric[5m])",
11336 "max_over_time(some_metric[5m])",
11337 "count_over_time(some_metric[5m])",
11338 "absent_over_time(some_metric[5m])",
11339 "present_over_time(some_metric[5m])",
11340 "stddev_over_time(some_metric[5m])",
11341 "stdvar_over_time(some_metric[5m])",
11342 "quantile_over_time(0.9, some_metric[5m])",
11343 "predict_linear(some_metric[5m], 60)",
11344 "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)",
11345 ];
11346
11347 for query in dual_output.iter().chain(float_output.iter()) {
11348 let plan = PromPlanner::stmt_to_plan(
11349 build_test_mixed_native_histogram_table_provider("some_metric").await,
11350 &build_eval_stmt(query),
11351 &build_query_engine_state(),
11352 )
11353 .await
11354 .unwrap()
11355 .display_indent_schema()
11356 .to_string();
11357 assert!(plan.contains("prom_mixed_range_float"), "{query}\n{plan}");
11358 assert_eq!(
11359 plan.contains("prom_mixed_range_histogram"),
11360 dual_output.contains(query),
11361 "{query}\n{plan}"
11362 );
11363 }
11364
11365 let plan = PromPlanner::stmt_to_plan(
11366 build_test_mixed_native_histogram_table_provider("some_metric").await,
11367 &build_eval_stmt("sum_over_time(rate(some_metric[5m])[10m:1m])"),
11368 &build_query_engine_state(),
11369 )
11370 .await
11371 .unwrap()
11372 .display_indent_schema()
11373 .to_string();
11374 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]
11375 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]
11376 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))]
11377 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]
11378 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]
11379 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]
11380 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]
11381 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))]
11382 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]
11383 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]
11384 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]
11385 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]
11386 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]"#;
11387 assert_eq!(plan, expected);
11388 }
11389
11390 #[tokio::test]
11391 async fn mixed_native_histogram_rate_executes_real_ranges() {
11392 let schema = Arc::new(ArrowSchema::new(vec![
11393 Field::new(
11394 "timestamp",
11395 ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
11396 false,
11397 ),
11398 Field::new(greptime_value(), ArrowDataType::Float64, true),
11399 Field::new(
11400 greptime_native_histogram(),
11401 native_histogram_value_type().as_arrow_type(),
11402 true,
11403 ),
11404 ]));
11405 let batch = RecordBatch::try_new(
11406 schema.clone(),
11407 vec![
11408 Arc::new(TimestampMillisecondArray::from(vec![1000, 2000, 3000])),
11409 Arc::new(Float64Array::from(vec![Some(1.0), None, Some(3.0)])),
11410 build_histogram_array(&[None, Some(direct_or_histogram()), None]),
11411 ],
11412 )
11413 .unwrap();
11414 let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap());
11415 let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None)
11416 .unwrap()
11417 .build()
11418 .unwrap();
11419 let collector = PromqlAnnotationCollector::default();
11420 let mut planner = PromPlanner {
11421 table_provider: build_test_table_provider_with_fields(
11422 &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
11423 &[],
11424 )
11425 .await,
11426 ctx: PromPlannerContext {
11427 start: 3000,
11428 end: 3000,
11429 interval: 1000,
11430 range: Some(3000),
11431 time_index_column: Some("timestamp".to_string()),
11432 field_columns: vec![
11433 greptime_native_histogram().to_string(),
11434 greptime_value().to_string(),
11435 ],
11436 ..Default::default()
11437 },
11438 promql_annotations: Some(collector.clone()),
11439 };
11440 let input = LogicalPlan::Extension(Extension {
11441 node: Arc::new(
11442 RangeManipulate::new(
11443 3000,
11444 3000,
11445 1000,
11446 3000,
11447 "timestamp".to_string(),
11448 planner.ctx.field_columns.clone(),
11449 input,
11450 )
11451 .unwrap(),
11452 ),
11453 });
11454 let PromExpr::Call(call) = parser::parse("rate(mixed[3s])").unwrap() else {
11455 unreachable!()
11456 };
11457 let preserve_any_value = PromPlanner::field_columns_are_alternative_samples(
11458 input.schema(),
11459 &planner.ctx.field_columns,
11460 );
11461 let state = build_query_engine_state();
11462 let (mut exprs, _) = planner
11463 .create_function_expr(&call.func, vec![], input.schema(), &state)
11464 .unwrap();
11465 exprs.insert(0, planner.create_time_index_column_expr().unwrap());
11466 let plan = LogicalPlanBuilder::from(input)
11467 .project(exprs)
11468 .unwrap()
11469 .filter(
11470 planner
11471 .create_empty_values_filter_expr(preserve_any_value)
11472 .unwrap(),
11473 )
11474 .unwrap()
11475 .build()
11476 .unwrap();
11477 let (_, batches) = execute(plan, &state).await;
11478 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
11479 let mut warnings = Vec::new();
11480 collector.append_to(&mut warnings, &mut Vec::new());
11481 assert!(
11482 warnings
11483 .iter()
11484 .any(|warning| warning.contains("mix of float and native histogram"))
11485 );
11486 }
11487
11488 #[tokio::test]
11489 async fn native_histogram_mixed_field_table_behaves() {
11490 let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
11494 let plan = PromPlanner::stmt_to_plan(
11495 table_provider,
11496 &build_eval_stmt("histogram_count(some_metric)"),
11497 &build_query_engine_state(),
11498 )
11499 .await
11500 .unwrap();
11501 let plan_str = plan.display_indent_schema().to_string();
11502 assert!(
11503 plan_str.contains("prom_native_histogram_count"),
11504 "{plan_str}"
11505 );
11506 assert!(!plan_str.contains("Float64(NULL)"), "{plan_str}");
11507 assert!(
11508 plan_str.contains("prom_native_histogram_count(greptime_native_histogram) IS NOT NULL"),
11509 "{plan_str}"
11510 );
11511
11512 let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
11514 let plan = PromPlanner::stmt_to_plan(
11515 table_provider,
11516 &build_eval_stmt("sort(some_metric)"),
11517 &build_query_engine_state(),
11518 )
11519 .await
11520 .unwrap();
11521 let plan_str = plan.display_indent_schema().to_string();
11522 assert!(
11523 plan_str.contains("greptime_value ASC NULLS FIRST"),
11524 "{plan_str}"
11525 );
11526 assert!(
11527 !plan_str.contains("greptime_native_histogram ASC"),
11528 "{plan_str}"
11529 );
11530
11531 let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
11533 let plan = PromPlanner::stmt_to_plan(
11534 table_provider,
11535 &build_eval_stmt("scalar(some_metric)"),
11536 &build_query_engine_state(),
11537 )
11538 .await
11539 .unwrap();
11540 let plan_str = plan.display_indent_schema().to_string();
11541 assert!(plan_str.contains("ScalarCalculate"), "{plan_str}");
11542 assert!(
11543 plan_str.contains("greptime_value IS NOT NULL"),
11544 "{plan_str}"
11545 );
11546
11547 let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
11549 let plan = PromPlanner::stmt_to_plan(
11550 table_provider,
11551 &build_eval_stmt(r#"label_replace(some_metric, "copied", "$1", "tag_0", "(.*)")"#),
11552 &build_query_engine_state(),
11553 )
11554 .await
11555 .unwrap();
11556 let plan_str = plan.display_indent_schema().to_string();
11557 let filter = plan_str.lines().next().unwrap();
11558 assert!(
11559 filter.starts_with("Filter: ")
11560 && filter.contains("greptime_native_histogram IS NOT NULL")
11561 && filter.contains(" OR ")
11562 && filter.contains("greptime_value IS NOT NULL"),
11563 "{plan_str}"
11564 );
11565 }
11566
11567 #[tokio::test]
11568 async fn less_filter_on_value() {
11569 let query = "some_metric < 1.2345";
11570 let expected = String::from(
11571 "Filter: some_metric.field_0 < Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11572 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11573 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11574 \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]\
11575 \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]\
11576 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
11577 );
11578
11579 indie_query_plan_compare(query, expected).await;
11580 }
11581
11582 #[tokio::test]
11583 async fn count_over_time() {
11584 let query = "count_over_time(some_metric[5m])";
11585 let expected = String::from(
11586 "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]\
11587 \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]\
11588 \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))]\
11589 \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11590 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11591 \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]\
11592 \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]\
11593 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
11594 );
11595
11596 indie_query_plan_compare(query, expected).await;
11597 }
11598
11599 #[tokio::test]
11602 async fn count_over_time_subquery() {
11603 let query = "count_over_time(some_metric[10m:1m])";
11604 let expected = String::from(
11605 "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]\
11606 \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]\
11607 \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))]\
11608 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11609 \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]\
11610 \n PromInstantManipulate: range=[-540000..100000000], lookback=[1000], interval=[60000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11611 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11612 \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]\
11613 \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]\
11614 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
11615 );
11616 indie_query_plan_compare(query, expected).await;
11617 }
11618
11619 #[tokio::test]
11620 async fn test_hash_join() {
11621 let mut eval_stmt = EvalStmt {
11622 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11623 start: UNIX_EPOCH,
11624 end: UNIX_EPOCH
11625 .checked_add(Duration::from_secs(100_000))
11626 .unwrap(),
11627 interval: Duration::from_secs(5),
11628 lookback_delta: Duration::from_secs(1),
11629 };
11630
11631 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"}"#;
11632
11633 let prom_expr = parser::parse(case).unwrap();
11634 eval_stmt.expr = prom_expr;
11635 let table_provider = build_test_table_provider_with_fields(
11636 &[
11637 (
11638 DEFAULT_SCHEMA_NAME.to_string(),
11639 "http_server_requests_seconds_sum".to_string(),
11640 ),
11641 (
11642 DEFAULT_SCHEMA_NAME.to_string(),
11643 "http_server_requests_seconds_count".to_string(),
11644 ),
11645 ],
11646 &["uri", "kubernetes_namespace", "kubernetes_pod_name"],
11647 )
11648 .await;
11649 let plan =
11651 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11652 .await
11653 .unwrap();
11654 let expected = "Projection: http_server_requests_seconds_count.uri, http_server_requests_seconds_count.kubernetes_namespace, http_server_requests_seconds_count.kubernetes_pod_name, http_server_requests_seconds_count.greptime_timestamp, CAST(http_server_requests_seconds_sum.greptime_value AS Float64) / CAST(http_server_requests_seconds_count.greptime_value AS Float64) AS http_server_requests_seconds_sum.greptime_value / http_server_requests_seconds_count.greptime_value\
11655 \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\
11656 \n SubqueryAlias: http_server_requests_seconds_sum\
11657 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\
11658 \n PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\
11659 \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\
11660 \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)\
11661 \n TableScan: http_server_requests_seconds_sum\
11662 \n SubqueryAlias: http_server_requests_seconds_count\
11663 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\
11664 \n PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\
11665 \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\
11666 \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)\
11667 \n TableScan: http_server_requests_seconds_count";
11668 assert_eq!(plan.to_string(), expected);
11669 }
11670
11671 #[tokio::test]
11672 async fn test_nested_histogram_quantile() {
11673 let mut eval_stmt = EvalStmt {
11674 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11675 start: UNIX_EPOCH,
11676 end: UNIX_EPOCH
11677 .checked_add(Duration::from_secs(100_000))
11678 .unwrap(),
11679 interval: Duration::from_secs(5),
11680 lookback_delta: Duration::from_secs(1),
11681 };
11682
11683 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]*-(.*)")"#;
11684
11685 let prom_expr = parser::parse(case).unwrap();
11686 eval_stmt.expr = prom_expr;
11687 let table_provider = build_test_table_provider_with_fields(
11688 &[(
11689 DEFAULT_SCHEMA_NAME.to_string(),
11690 "greptime_servers_grpc_requests_elapsed_bucket".to_string(),
11691 )],
11692 &["pod", "le", "path", "code", "container"],
11693 )
11694 .await;
11695 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11697 .await
11698 .unwrap();
11699 }
11700
11701 #[tokio::test]
11702 async fn test_histogram_quantile_binary_op() {
11703 let mut eval_stmt = EvalStmt {
11704 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11705 start: UNIX_EPOCH,
11706 end: UNIX_EPOCH
11707 .checked_add(Duration::from_secs(100_000))
11708 .unwrap(),
11709 interval: Duration::from_secs(5),
11710 lookback_delta: Duration::from_secs(1),
11711 };
11712
11713 let case = r#"histogram_quantile(0.5, sum by (le, pod) (rate(http_request_duration_seconds_bucket[5m]))) + 0"#;
11717
11718 let prom_expr = parser::parse(case).unwrap();
11719 eval_stmt.expr = prom_expr;
11720 let table_provider = build_test_table_provider_with_fields(
11721 &[(
11722 DEFAULT_SCHEMA_NAME.to_string(),
11723 "http_request_duration_seconds_bucket".to_string(),
11724 )],
11725 &["pod", "le"],
11726 )
11727 .await;
11728 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11730 .await
11731 .unwrap();
11732 }
11733
11734 #[tokio::test]
11735 async fn test_parse_and_operator() {
11736 let mut eval_stmt = EvalStmt {
11737 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11738 start: UNIX_EPOCH,
11739 end: UNIX_EPOCH
11740 .checked_add(Duration::from_secs(100_000))
11741 .unwrap(),
11742 interval: Duration::from_secs(5),
11743 lookback_delta: Duration::from_secs(1),
11744 };
11745
11746 let cases = [
11747 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)"#,
11748 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)"#,
11749 ];
11750
11751 for case in cases {
11752 let prom_expr = parser::parse(case).unwrap();
11753 eval_stmt.expr = prom_expr;
11754 let table_provider = build_test_table_provider_with_fields(
11755 &[
11756 (
11757 DEFAULT_SCHEMA_NAME.to_string(),
11758 "kubelet_volume_stats_used_bytes".to_string(),
11759 ),
11760 (
11761 DEFAULT_SCHEMA_NAME.to_string(),
11762 "kubelet_volume_stats_capacity_bytes".to_string(),
11763 ),
11764 ],
11765 &["namespace", "persistentvolumeclaim"],
11766 )
11767 .await;
11768 let _ =
11770 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11771 .await
11772 .unwrap();
11773 }
11774 }
11775
11776 #[tokio::test]
11777 async fn test_nested_binary_op() {
11778 let mut eval_stmt = EvalStmt {
11779 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11780 start: UNIX_EPOCH,
11781 end: UNIX_EPOCH
11782 .checked_add(Duration::from_secs(100_000))
11783 .unwrap(),
11784 interval: Duration::from_secs(5),
11785 lookback_delta: Duration::from_secs(1),
11786 };
11787
11788 let case = r#"sum(rate(nginx_ingress_controller_requests{job=~".*"}[2m])) -
11789 (
11790 sum(rate(nginx_ingress_controller_requests{namespace=~".*"}[2m]))
11791 or
11792 vector(0)
11793 )"#;
11794
11795 let prom_expr = parser::parse(case).unwrap();
11796 eval_stmt.expr = prom_expr;
11797 let table_provider = build_test_table_provider_with_fields(
11798 &[(
11799 DEFAULT_SCHEMA_NAME.to_string(),
11800 "nginx_ingress_controller_requests".to_string(),
11801 )],
11802 &["namespace", "job"],
11803 )
11804 .await;
11805 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11807 .await
11808 .unwrap();
11809 }
11810
11811 #[tokio::test]
11812 async fn test_parse_or_operator() {
11813 let mut eval_stmt = EvalStmt {
11814 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11815 start: UNIX_EPOCH,
11816 end: UNIX_EPOCH
11817 .checked_add(Duration::from_secs(100_000))
11818 .unwrap(),
11819 interval: Duration::from_secs(5),
11820 lookback_delta: Duration::from_secs(1),
11821 };
11822
11823 let case = r#"
11824 sum(rate(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}[120s])) by (cluster_name,tenant_name) /
11825 (sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) * 100)
11826 or
11827 200 * sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) /
11828 sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)"#;
11829
11830 let table_provider = build_test_table_provider_with_fields(
11831 &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())],
11832 &["tenant_name", "cluster_name"],
11833 )
11834 .await;
11835 eval_stmt.expr = parser::parse(case).unwrap();
11836 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11837 .await
11838 .unwrap();
11839
11840 let case = r#"sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
11841 (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) +
11842 sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
11843 (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0
11844 or
11845 sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
11846 (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0
11847 or
11848 sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
11849 (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0"#;
11850 let table_provider = build_test_table_provider_with_fields(
11851 &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())],
11852 &["tenant_name", "cluster_name"],
11853 )
11854 .await;
11855 eval_stmt.expr = parser::parse(case).unwrap();
11856 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11857 .await
11858 .unwrap();
11859
11860 let case = r#"(sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) +
11861 sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or
11862 (sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or
11863 (sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name))"#;
11864 let table_provider = build_test_table_provider_with_fields(
11865 &[
11866 (
11867 DEFAULT_SCHEMA_NAME.to_string(),
11868 "background_waitevent_cnt".to_string(),
11869 ),
11870 (
11871 DEFAULT_SCHEMA_NAME.to_string(),
11872 "foreground_waitevent_cnt".to_string(),
11873 ),
11874 ],
11875 &["tenant_name", "cluster_name"],
11876 )
11877 .await;
11878 eval_stmt.expr = parser::parse(case).unwrap();
11879 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11880 .await
11881 .unwrap();
11882
11883 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)"#;
11884 let table_provider = build_test_table_provider_with_fields(
11885 &[
11886 (DEFAULT_SCHEMA_NAME.to_string(), "node_load1".to_string()),
11887 (
11888 DEFAULT_SCHEMA_NAME.to_string(),
11889 "container_cpu_load_average_10s".to_string(),
11890 ),
11891 (
11892 DEFAULT_SCHEMA_NAME.to_string(),
11893 "container_spec_cpu_quota".to_string(),
11894 ),
11895 ],
11896 &["cluster_name", "host_name"],
11897 )
11898 .await;
11899 eval_stmt.expr = parser::parse(case).unwrap();
11900 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11901 .await
11902 .unwrap();
11903 }
11904
11905 #[tokio::test]
11906 async fn value_matcher() {
11907 let mut eval_stmt = EvalStmt {
11909 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11910 start: UNIX_EPOCH,
11911 end: UNIX_EPOCH
11912 .checked_add(Duration::from_secs(100_000))
11913 .unwrap(),
11914 interval: Duration::from_secs(5),
11915 lookback_delta: Duration::from_secs(1),
11916 };
11917
11918 let cases = [
11919 (
11921 r#"some_metric{__field__="field_1"}"#,
11922 vec![
11923 "some_metric.field_1",
11924 "some_metric.tag_0",
11925 "some_metric.tag_1",
11926 "some_metric.tag_2",
11927 "some_metric.timestamp",
11928 ],
11929 ),
11930 (
11932 r#"some_metric{__field__="field_1", __field__="field_0"}"#,
11933 vec![
11934 "some_metric.field_0",
11935 "some_metric.field_1",
11936 "some_metric.tag_0",
11937 "some_metric.tag_1",
11938 "some_metric.tag_2",
11939 "some_metric.timestamp",
11940 ],
11941 ),
11942 (
11944 r#"some_metric{__field__!="field_1"}"#,
11945 vec![
11946 "some_metric.field_0",
11947 "some_metric.field_2",
11948 "some_metric.tag_0",
11949 "some_metric.tag_1",
11950 "some_metric.tag_2",
11951 "some_metric.timestamp",
11952 ],
11953 ),
11954 (
11956 r#"some_metric{__field__!="field_1", __field__!="field_2"}"#,
11957 vec![
11958 "some_metric.field_0",
11959 "some_metric.tag_0",
11960 "some_metric.tag_1",
11961 "some_metric.tag_2",
11962 "some_metric.timestamp",
11963 ],
11964 ),
11965 (
11967 r#"some_metric{__field__="field_1", __field__!="field_0"}"#,
11968 vec![
11969 "some_metric.field_1",
11970 "some_metric.tag_0",
11971 "some_metric.tag_1",
11972 "some_metric.tag_2",
11973 "some_metric.timestamp",
11974 ],
11975 ),
11976 (
11978 r#"some_metric{__field__="field_2", __field__!="field_2"}"#,
11979 vec![
11980 "some_metric.tag_0",
11981 "some_metric.tag_1",
11982 "some_metric.tag_2",
11983 "some_metric.timestamp",
11984 ],
11985 ),
11986 (
11988 r#"some_metric{__field__=~"field_1|field_2"}"#,
11989 vec![
11990 "some_metric.field_1",
11991 "some_metric.field_2",
11992 "some_metric.tag_0",
11993 "some_metric.tag_1",
11994 "some_metric.tag_2",
11995 "some_metric.timestamp",
11996 ],
11997 ),
11998 (
12000 r#"some_metric{__field__!~"field_1|field_2"}"#,
12001 vec![
12002 "some_metric.field_0",
12003 "some_metric.tag_0",
12004 "some_metric.tag_1",
12005 "some_metric.tag_2",
12006 "some_metric.timestamp",
12007 ],
12008 ),
12009 ];
12010
12011 for case in cases {
12012 let prom_expr = parser::parse(case.0).unwrap();
12013 eval_stmt.expr = prom_expr;
12014 let table_provider = build_test_table_provider(
12015 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12016 3,
12017 3,
12018 )
12019 .await;
12020 let plan =
12021 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12022 .await
12023 .unwrap();
12024 let mut fields = plan.schema().field_names();
12025 let mut expected = case.1.into_iter().map(String::from).collect::<Vec<_>>();
12026 fields.sort();
12027 expected.sort();
12028 assert_eq!(fields, expected, "case: {:?}", case.0);
12029 }
12030
12031 let bad_cases = [
12032 r#"some_metric{__field__="nonexistent"}"#,
12033 r#"some_metric{__field__!="nonexistent"}"#,
12034 ];
12035
12036 for case in bad_cases {
12037 let prom_expr = parser::parse(case).unwrap();
12038 eval_stmt.expr = prom_expr;
12039 let table_provider = build_test_table_provider(
12040 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12041 3,
12042 3,
12043 )
12044 .await;
12045 let plan =
12046 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12047 .await;
12048 assert!(plan.is_err(), "case: {:?}", case);
12049 }
12050 }
12051
12052 #[tokio::test]
12053 async fn custom_schema() {
12054 let query = "some_alt_metric{__schema__=\"greptime_private\"}";
12055 let expected = String::from(
12056 "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12057 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12058 \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]\
12059 \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]\
12060 \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
12061 );
12062
12063 indie_query_plan_compare(query, expected).await;
12064
12065 let query = "some_alt_metric{__database__=\"greptime_private\"}";
12066 let expected = String::from(
12067 "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12068 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12069 \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]\
12070 \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]\
12071 \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
12072 );
12073
12074 indie_query_plan_compare(query, expected).await;
12075
12076 let query = "some_alt_metric{__schema__=\"greptime_private\"} / some_metric";
12077 let expected = String::from(
12078 "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]\
12079 \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]\
12080 \n SubqueryAlias: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12081 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12082 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12083 \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]\
12084 \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]\
12085 \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12086 \n SubqueryAlias: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12087 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12088 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12089 \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]\
12090 \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]\
12091 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
12092 );
12093
12094 indie_query_plan_compare(query, expected).await;
12095 }
12096
12097 #[tokio::test]
12098 async fn only_equals_is_supported_for_special_matcher() {
12099 let queries = &[
12100 "some_alt_metric{__schema__!=\"greptime_private\"}",
12101 "some_alt_metric{__schema__=~\"lalala\"}",
12102 "some_alt_metric{__database__!=\"greptime_private\"}",
12103 "some_alt_metric{__database__=~\"lalala\"}",
12104 ];
12105
12106 for query in queries {
12107 let prom_expr = parser::parse(query).unwrap();
12108 let eval_stmt = EvalStmt {
12109 expr: prom_expr,
12110 start: UNIX_EPOCH,
12111 end: UNIX_EPOCH
12112 .checked_add(Duration::from_secs(100_000))
12113 .unwrap(),
12114 interval: Duration::from_secs(5),
12115 lookback_delta: Duration::from_secs(1),
12116 };
12117
12118 let table_provider = build_test_table_provider(
12119 &[
12120 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
12121 (
12122 "greptime_private".to_string(),
12123 "some_alt_metric".to_string(),
12124 ),
12125 ],
12126 1,
12127 1,
12128 )
12129 .await;
12130
12131 let plan =
12132 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12133 .await;
12134 assert!(plan.is_err(), "query: {:?}", query);
12135 }
12136 }
12137
12138 #[tokio::test]
12139 async fn test_non_ms_precision() {
12140 let catalog_list = MemoryCatalogManager::with_default_setup();
12141 let columns = vec![
12142 ColumnSchema::new(
12143 "tag".to_string(),
12144 ConcreteDataType::string_datatype(),
12145 false,
12146 ),
12147 ColumnSchema::new(
12148 "timestamp".to_string(),
12149 ConcreteDataType::timestamp_nanosecond_datatype(),
12150 false,
12151 )
12152 .with_time_index(true),
12153 ColumnSchema::new(
12154 "field".to_string(),
12155 ConcreteDataType::float64_datatype(),
12156 true,
12157 ),
12158 ];
12159 let schema = Arc::new(Schema::new(columns));
12160 let table_meta = TableMetaBuilder::empty()
12161 .schema(schema)
12162 .primary_key_indices(vec![0])
12163 .value_indices(vec![2])
12164 .next_column_id(1024)
12165 .build()
12166 .unwrap();
12167 let table_info = TableInfoBuilder::default()
12168 .name("metrics".to_string())
12169 .meta(table_meta)
12170 .build()
12171 .unwrap();
12172 let table = EmptyTable::from_table_info(&table_info);
12173 assert!(
12174 catalog_list
12175 .register_table_sync(RegisterTableRequest {
12176 catalog: DEFAULT_CATALOG_NAME.to_string(),
12177 schema: DEFAULT_SCHEMA_NAME.to_string(),
12178 table_name: "metrics".to_string(),
12179 table_id: 1024,
12180 table,
12181 })
12182 .is_ok()
12183 );
12184
12185 let plan = PromPlanner::stmt_to_plan(
12186 DfTableSourceProvider::new(
12187 catalog_list.clone(),
12188 false,
12189 QueryContext::arc(),
12190 DummyDecoder::arc(),
12191 true,
12192 ),
12193 &EvalStmt {
12194 expr: parser::parse("metrics{tag = \"1\"}").unwrap(),
12195 start: UNIX_EPOCH,
12196 end: UNIX_EPOCH
12197 .checked_add(Duration::from_secs(100_000))
12198 .unwrap(),
12199 interval: Duration::from_secs(5),
12200 lookback_delta: Duration::from_secs(1),
12201 },
12202 &build_query_engine_state(),
12203 )
12204 .await
12205 .unwrap();
12206 assert_eq!(
12207 plan.display_indent_schema().to_string(),
12208 "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
12209 \n PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
12210 \n Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
12211 \n Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp >= TimestampMillisecond(-999, None) AND metrics.timestamp <= TimestampMillisecond(100000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
12212 \n Projection: metrics.field, metrics.tag, CAST(metrics.timestamp AS Timestamp(ms)) AS timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
12213 \n TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]"
12214 );
12215 let plan = PromPlanner::stmt_to_plan(
12216 DfTableSourceProvider::new(
12217 catalog_list.clone(),
12218 false,
12219 QueryContext::arc(),
12220 DummyDecoder::arc(),
12221 true,
12222 ),
12223 &EvalStmt {
12224 expr: parser::parse("avg_over_time(metrics{tag = \"1\"}[5s])").unwrap(),
12225 start: UNIX_EPOCH,
12226 end: UNIX_EPOCH
12227 .checked_add(Duration::from_secs(100_000))
12228 .unwrap(),
12229 interval: Duration::from_secs(5),
12230 lookback_delta: Duration::from_secs(1),
12231 },
12232 &build_query_engine_state(),
12233 )
12234 .await
12235 .unwrap();
12236 assert_eq!(
12237 plan.display_indent_schema().to_string(),
12238 "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]\
12239 \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]\
12240 \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))]\
12241 \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
12242 \n PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
12243 \n Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
12244 \n Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp >= TimestampMillisecond(-4999, None) AND metrics.timestamp <= TimestampMillisecond(100000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
12245 \n Projection: metrics.field, metrics.tag, CAST(metrics.timestamp AS Timestamp(ms)) AS timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
12246 \n TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]"
12247 );
12248 }
12249
12250 #[tokio::test]
12251 async fn test_nonexistent_label() {
12252 let mut eval_stmt = EvalStmt {
12254 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12255 start: UNIX_EPOCH,
12256 end: UNIX_EPOCH
12257 .checked_add(Duration::from_secs(100_000))
12258 .unwrap(),
12259 interval: Duration::from_secs(5),
12260 lookback_delta: Duration::from_secs(1),
12261 };
12262
12263 let case = r#"some_metric{nonexistent="hi"}"#;
12264 let prom_expr = parser::parse(case).unwrap();
12265 eval_stmt.expr = prom_expr;
12266 let table_provider = build_test_table_provider(
12267 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12268 3,
12269 3,
12270 )
12271 .await;
12272 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12274 .await
12275 .unwrap();
12276 }
12277
12278 #[tokio::test]
12279 async fn test_label_join() {
12280 let prom_expr = parser::parse(
12281 "label_join(up{tag_0='api-server'}, 'foo', ',', 'tag_1', 'tag_2', 'tag_3')",
12282 )
12283 .unwrap();
12284 let eval_stmt = EvalStmt {
12285 expr: prom_expr,
12286 start: UNIX_EPOCH,
12287 end: UNIX_EPOCH
12288 .checked_add(Duration::from_secs(100_000))
12289 .unwrap(),
12290 interval: Duration::from_secs(5),
12291 lookback_delta: Duration::from_secs(1),
12292 };
12293
12294 let table_provider =
12295 build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 4, 1)
12296 .await;
12297 let plan =
12298 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12299 .await
12300 .unwrap();
12301
12302 let expected = r#"
12303Filter: 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]
12304 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]
12305 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]
12306 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]
12307 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]
12308 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]
12309 TableScan: up [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#;
12310
12311 let ret = plan.display_indent_schema().to_string();
12312 assert_eq!(format!("\n{ret}"), expected, "\n{}", ret);
12313 }
12314
12315 #[tokio::test]
12316 async fn test_label_replace() {
12317 let prom_expr = parser::parse(
12318 "label_replace(up{tag_0=\"a:c\"}, \"foo\", \"$1\", \"tag_0\", \"(.*):.*\")",
12319 )
12320 .unwrap();
12321 let eval_stmt = EvalStmt {
12322 expr: prom_expr,
12323 start: UNIX_EPOCH,
12324 end: UNIX_EPOCH
12325 .checked_add(Duration::from_secs(100_000))
12326 .unwrap(),
12327 interval: Duration::from_secs(5),
12328 lookback_delta: Duration::from_secs(1),
12329 };
12330
12331 let table_provider =
12332 build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 1, 1)
12333 .await;
12334 let plan =
12335 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12336 .await
12337 .unwrap();
12338
12339 let expected = r#"
12340Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8]
12341 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]
12342 PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
12343 PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
12344 Sort: up.tag_0 ASC NULLS FIRST, up.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
12345 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]
12346 TableScan: up [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#;
12347
12348 let ret = plan.display_indent_schema().to_string();
12349 assert_eq!(format!("\n{ret}"), expected, "\n{}", ret);
12350 }
12351
12352 #[tokio::test]
12353 async fn label_replace_aggregation_queries_plan_successfully() {
12354 let aggregate =
12355 r#"sum by (foo) (label_replace(some_metric, "foo", "$1", "tag_0", "(.*)"))"#;
12356 let queries = [
12357 aggregate.to_string(),
12358 format!("{aggregate} <= 10"),
12359 format!("{aggregate} * 0.8"),
12360 format!("0.8 * {aggregate}"),
12361 format!("{aggregate} <= {aggregate} * 0.8"),
12362 ];
12363 let state = build_query_engine_state();
12364 let mut failures = Vec::new();
12365
12366 for query in queries {
12367 let table_provider = build_test_table_provider(
12368 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12369 1,
12370 1,
12371 )
12372 .await;
12373 if let Err(error) =
12374 PromPlanner::stmt_to_plan(table_provider, &build_eval_stmt(&query), &state).await
12375 {
12376 failures.push(format!("{query}: {error:?}"));
12377 }
12378 }
12379
12380 assert!(failures.is_empty(), "{}", failures.join("\n"));
12381 }
12382
12383 #[tokio::test]
12384 async fn test_matchers_to_expr() {
12385 let mut eval_stmt = EvalStmt {
12386 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12387 start: UNIX_EPOCH,
12388 end: UNIX_EPOCH
12389 .checked_add(Duration::from_secs(100_000))
12390 .unwrap(),
12391 interval: Duration::from_secs(5),
12392 lookback_delta: Duration::from_secs(1),
12393 };
12394 let case =
12395 r#"sum(prometheus_tsdb_head_series{tag_1=~"(10.0.160.237:8080|10.0.160.237:9090)"})"#;
12396
12397 let prom_expr = parser::parse(case).unwrap();
12398 eval_stmt.expr = prom_expr;
12399 let table_provider = build_test_table_provider(
12400 &[(
12401 DEFAULT_SCHEMA_NAME.to_string(),
12402 "prometheus_tsdb_head_series".to_string(),
12403 )],
12404 3,
12405 3,
12406 )
12407 .await;
12408 let plan =
12409 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12410 .await
12411 .unwrap();
12412 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]\
12413 \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]\
12414 \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]\
12415 \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]\
12416 \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]\
12417 \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]\
12418 \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]";
12419 assert_eq!(plan.display_indent_schema().to_string(), expected);
12420 }
12421
12422 #[tokio::test]
12423 async fn test_topk_expr() {
12424 let mut eval_stmt = EvalStmt {
12425 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12426 start: UNIX_EPOCH,
12427 end: UNIX_EPOCH
12428 .checked_add(Duration::from_secs(100_000))
12429 .unwrap(),
12430 interval: Duration::from_secs(5),
12431 lookback_delta: Duration::from_secs(1),
12432 };
12433 let case = r#"topk(10, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#;
12434
12435 let prom_expr = parser::parse(case).unwrap();
12436 eval_stmt.expr = prom_expr;
12437 let table_provider = build_test_table_provider_with_fields(
12438 &[
12439 (
12440 DEFAULT_SCHEMA_NAME.to_string(),
12441 "prometheus_tsdb_head_series".to_string(),
12442 ),
12443 (
12444 DEFAULT_SCHEMA_NAME.to_string(),
12445 "http_server_requests_seconds_count".to_string(),
12446 ),
12447 ],
12448 &["ip"],
12449 )
12450 .await;
12451
12452 let plan =
12453 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12454 .await
12455 .unwrap();
12456 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)]\
12457 \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]\
12458 \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]\
12459 \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]\
12460 \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]\
12461 \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]\
12462 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
12463 \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
12464 \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]\
12465 \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]\
12466 \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]";
12467
12468 assert_eq!(plan.display_indent_schema().to_string(), expected);
12469 }
12470
12471 #[tokio::test]
12472 async fn test_count_values_expr() {
12473 let mut eval_stmt = EvalStmt {
12474 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12475 start: UNIX_EPOCH,
12476 end: UNIX_EPOCH
12477 .checked_add(Duration::from_secs(100_000))
12478 .unwrap(),
12479 interval: Duration::from_secs(5),
12480 lookback_delta: Duration::from_secs(1),
12481 };
12482 let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#;
12483
12484 let prom_expr = parser::parse(case).unwrap();
12485 eval_stmt.expr = prom_expr;
12486 let table_provider = build_test_table_provider_with_fields(
12487 &[
12488 (
12489 DEFAULT_SCHEMA_NAME.to_string(),
12490 "prometheus_tsdb_head_series".to_string(),
12491 ),
12492 (
12493 DEFAULT_SCHEMA_NAME.to_string(),
12494 "http_server_requests_seconds_count".to_string(),
12495 ),
12496 ],
12497 &["ip"],
12498 )
12499 .await;
12500
12501 let plan =
12502 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12503 .await
12504 .unwrap();
12505 let expected = "Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N]\
12506 \n Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N]\
12507 \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N, count(prometheus_tsdb_head_series.greptime_value):Int64]\
12508 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
12509 \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
12510 \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]\
12511 \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]\
12512 \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]";
12513
12514 assert_eq!(plan.display_indent_schema().to_string(), expected);
12515 }
12516
12517 #[tokio::test]
12518 async fn test_value_alias() {
12519 let mut eval_stmt = EvalStmt {
12520 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12521 start: UNIX_EPOCH,
12522 end: UNIX_EPOCH
12523 .checked_add(Duration::from_secs(100_000))
12524 .unwrap(),
12525 interval: Duration::from_secs(5),
12526 lookback_delta: Duration::from_secs(1),
12527 };
12528 let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#;
12529
12530 let prom_expr = parser::parse(case).unwrap();
12531 eval_stmt.expr = prom_expr;
12532 eval_stmt = QueryLanguageParser::apply_alias_extension(eval_stmt, "my_series");
12533 let table_provider = build_test_table_provider_with_fields(
12534 &[
12535 (
12536 DEFAULT_SCHEMA_NAME.to_string(),
12537 "prometheus_tsdb_head_series".to_string(),
12538 ),
12539 (
12540 DEFAULT_SCHEMA_NAME.to_string(),
12541 "http_server_requests_seconds_count".to_string(),
12542 ),
12543 ],
12544 &["ip"],
12545 )
12546 .await;
12547
12548 let plan =
12549 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12550 .await
12551 .unwrap();
12552 let expected = r#"
12553Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp [my_series:Int64, ip:Utf8, greptime_timestamp:Timestamp(ms)]
12554 Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N]
12555 Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N]
12556 Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N, count(prometheus_tsdb_head_series.greptime_value):Int64]
12557 PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
12558 PromSeriesDivide: tags=["ip"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
12559 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]
12560 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]
12561 TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"#;
12562 assert_eq!(format!("\n{}", plan.display_indent_schema()), expected);
12563 }
12564
12565 #[tokio::test]
12566 async fn test_quantile_expr() {
12567 let mut eval_stmt = EvalStmt {
12568 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12569 start: UNIX_EPOCH,
12570 end: UNIX_EPOCH
12571 .checked_add(Duration::from_secs(100_000))
12572 .unwrap(),
12573 interval: Duration::from_secs(5),
12574 lookback_delta: Duration::from_secs(1),
12575 };
12576 let case = r#"quantile(0.3, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#;
12577
12578 let prom_expr = parser::parse(case).unwrap();
12579 eval_stmt.expr = prom_expr;
12580 let table_provider = build_test_table_provider_with_fields(
12581 &[
12582 (
12583 DEFAULT_SCHEMA_NAME.to_string(),
12584 "prometheus_tsdb_head_series".to_string(),
12585 ),
12586 (
12587 DEFAULT_SCHEMA_NAME.to_string(),
12588 "http_server_requests_seconds_count".to_string(),
12589 ),
12590 ],
12591 &["ip"],
12592 )
12593 .await;
12594
12595 let plan =
12596 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12597 .await
12598 .unwrap();
12599 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]\
12600 \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]\
12601 \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]\
12602 \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]\
12603 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
12604 \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
12605 \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]\
12606 \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]\
12607 \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]";
12608
12609 assert_eq!(plan.display_indent_schema().to_string(), expected);
12610 }
12611
12612 #[tokio::test]
12613 async fn test_or_not_exists_table_label() {
12614 let state = build_query_engine_state();
12615 let provider = build_test_table_provider_with_fields(
12616 &[(DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string())],
12617 &["job"],
12618 )
12619 .await;
12620 let raw = PromPlanner::stmt_to_plan(
12621 provider,
12622 &build_eval_stmt(r#"missing_metric or on(absent_label) normal_metric"#),
12623 &state,
12624 )
12625 .await
12626 .unwrap();
12627 assert!(
12628 raw.display_indent_schema()
12629 .to_string()
12630 .contains("__promql_or_match_0@")
12631 );
12632 let (optimized, batches) = execute(raw, &state).await;
12633 assert_no_internal_or_keys(optimized.schema());
12634 assert!(batches.iter().all(|batch| {
12635 batch
12636 .schema()
12637 .fields()
12638 .iter()
12639 .all(|field| !field.name().starts_with("__promql_or_match_"))
12640 }));
12641 }
12642
12643 #[tokio::test]
12644 async fn test_histogram_quantile_missing_le_column() {
12645 let mut eval_stmt = EvalStmt {
12646 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12647 start: UNIX_EPOCH,
12648 end: UNIX_EPOCH
12649 .checked_add(Duration::from_secs(100_000))
12650 .unwrap(),
12651 interval: Duration::from_secs(5),
12652 lookback_delta: Duration::from_secs(1),
12653 };
12654
12655 let case = r#"histogram_quantile(0.99, sum by(pod,instance,le) (rate(non_existent_histogram_bucket{instance=~"xxx"}[1m])))"#;
12657
12658 let prom_expr = parser::parse(case).unwrap();
12659 eval_stmt.expr = prom_expr;
12660
12661 let table_provider = build_test_table_provider_with_fields(
12663 &[(
12664 DEFAULT_SCHEMA_NAME.to_string(),
12665 "non_existent_histogram_bucket".to_string(),
12666 )],
12667 &["pod", "instance"], )
12669 .await;
12670
12671 let result =
12673 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12674 .await;
12675
12676 assert!(
12678 result.is_ok(),
12679 "Expected successful plan creation with empty result, but got error: {:?}",
12680 result.err()
12681 );
12682
12683 let plan = result.unwrap();
12685 match plan {
12686 LogicalPlan::EmptyRelation(_) => {
12687 }
12689 _ => panic!("Expected EmptyRelation, but got: {:?}", plan),
12690 }
12691 }
12692
12693 #[tokio::test]
12694 async fn test_direct_or_normalizes_missing_match_labels() {
12695 type Case<'a> = (
12696 Option<Option<&'a str>>,
12697 Option<Option<&'a str>>,
12698 i64,
12699 i64,
12700 &'a [(f64, Option<&'a str>)],
12701 );
12702
12703 let modifier = or_modifier("lhs or on(k) rhs");
12704 #[rustfmt::skip]
12705 let cases: &[Case<'_>] = &[
12706 (None, None, 1, 1, &[(1.0, None)]),
12707 (None, Some(Some("")), 1, 1, &[(1.0, None)]),
12708 (Some(Some("")), None, 1, 1, &[(1.0, Some(""))]),
12709 (None, Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]),
12710 (Some(Some("l")), None, 1, 1, &[(1.0, Some("l")), (2.0, None)]),
12711 (Some(None), Some(Some("")), 1, 1, &[(1.0, None)]),
12712 (Some(None), Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]),
12713 (Some(Some("same")), Some(Some("same")), 1, 2, &[(1.0, Some("same")), (2.0, Some("same"))]),
12714 ];
12715 for &(left, right, left_ts, right_ts, expected) in cases {
12716 let (optimized, batches) = run(
12717 &matrix_source("lhs", left, left_ts, 1.0),
12718 &matrix_source("rhs", right, right_ts, 2.0),
12719 matrix_context("lhs", left),
12720 matrix_context("rhs", right),
12721 &modifier,
12722 )
12723 .await;
12724 assert_no_internal_or_keys(optimized.schema());
12725 assert_eq!(
12726 rows(&batches),
12727 expected
12728 .iter()
12729 .map(|(value, label)| (*value, label.map(str::to_string)))
12730 .collect::<Vec<_>>()
12731 );
12732 }
12733 }
12734
12735 #[tokio::test]
12736 async fn test_direct_or_match_modifiers() {
12737 for (modifier, left, right, expected) in [
12738 (None, "left", "right", 2),
12739 (or_modifier("lhs or on(k) rhs"), "same", "same", 1),
12740 (or_modifier("lhs or on() rhs"), "left", "right", 1),
12741 (or_modifier("lhs or ignoring(k) rhs"), "left", "right", 1),
12742 ] {
12743 let (_, batches) = run(
12744 &matrix_source("lhs", Some(Some(left)), 1, 1.0),
12745 &matrix_source("rhs", Some(Some(right)), 1, 2.0),
12746 direct_or_context("lhs", &["job", "k"], "v"),
12747 direct_or_context("rhs", &["job", "k"], "v"),
12748 &modifier,
12749 )
12750 .await;
12751 assert_eq!(
12752 batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
12753 expected
12754 );
12755 }
12756 }
12757
12758 #[tokio::test]
12759 async fn test_direct_or_nested_projection_uses_left_context() {
12760 let left = matrix_source("lhs", Some(Some("k")), 1, 1.0);
12761 let right = matrix_source("rhs", Some(Some("k")), 1, 2.0);
12762 let raw = plan_direct_or(
12763 scan(&left),
12764 scan(&right),
12765 direct_or_context("lhs", &["job", "k"], "v"),
12766 direct_or_context("rhs", &["job", "k"], "v"),
12767 &or_modifier("lhs or on(k) rhs"),
12768 )
12769 .await;
12770 assert!(raw.schema().iter().any(|(qualifier, field)| {
12771 qualifier.as_ref().is_some_and(|q| q.to_string() == "lhs") && field.name() == "v"
12772 }));
12773 let nested = LogicalPlanBuilder::from(raw)
12774 .project(vec![
12775 DfExpr::BinaryExpr(BinaryExpr {
12776 left: Box::new(DfExpr::Column(Column::new(
12777 Some(TableReference::bare("lhs")),
12778 "v",
12779 ))),
12780 op: Operator::Plus,
12781 right: Box::new(lit(1.0)),
12782 })
12783 .alias("v_plus"),
12784 ])
12785 .unwrap()
12786 .build()
12787 .unwrap();
12788 let (_, batches) = execute(nested, &build_query_engine_state()).await;
12789 assert_eq!(values(&batches, "v_plus"), vec![2.0]);
12790 }
12791
12792 #[tokio::test]
12793 async fn test_direct_or_skips_user_internal_key_name() {
12794 const USER_TAG: &str = "__promql_or_match_0";
12795 let left = tagged_source(
12796 "lhs",
12797 false,
12798 (USER_TAG, Some("left")),
12799 DirectOrValue::Float64(1.0),
12800 );
12801 let right = tagged_source(
12802 "rhs",
12803 false,
12804 (USER_TAG, Some("right")),
12805 DirectOrValue::Float64(2.0),
12806 );
12807 let raw = plan_direct_or(
12808 scan(&left),
12809 scan(&right),
12810 direct_or_context("lhs", &["job", USER_TAG], "v"),
12811 direct_or_context("rhs", &["job", USER_TAG], "v"),
12812 &or_modifier("lhs or on(missing_label) rhs"),
12813 )
12814 .await;
12815 assert!(
12816 raw.display_indent_schema()
12817 .to_string()
12818 .contains("__promql_or_match_1@")
12819 );
12820 let (_, batches) = execute(raw, &build_query_engine_state()).await;
12821 assert!(
12822 batches
12823 .iter()
12824 .all(|batch| batch.column_by_name(USER_TAG).is_some())
12825 );
12826 }
12827
12828 #[tokio::test]
12829 async fn test_direct_or_substrait_round_trip_with_normalized_key() {
12830 let state = build_query_engine_state();
12831 let ctx = SessionContext::new_with_state(state.session_state());
12832 let catalog = Arc::new(MemoryCatalogProvider::new());
12833 catalog
12834 .register_schema("public", Arc::new(MemorySchemaProvider::new()))
12835 .unwrap();
12836 ctx.register_catalog("datafusion", catalog);
12837 let left = matrix_source("lhs", Some(Some("")), 1, 1.0);
12838 let right = matrix_source("rhs", None, 1, 2.0);
12839 ctx.register_table(
12840 TableReference::full("datafusion", "public", "lhs"),
12841 table(&left),
12842 )
12843 .unwrap();
12844 ctx.register_table(
12845 TableReference::full("datafusion", "public", "rhs"),
12846 table(&right),
12847 )
12848 .unwrap();
12849 let raw = plan_direct_or(
12850 ctx.table("datafusion.public.lhs")
12851 .await
12852 .unwrap()
12853 .into_unoptimized_plan(),
12854 ctx.table("datafusion.public.rhs")
12855 .await
12856 .unwrap()
12857 .into_unoptimized_plan(),
12858 direct_or_context("lhs", &["job", "k"], "v"),
12859 direct_or_context("rhs", &["job"], "v"),
12860 &or_modifier("lhs or on(k) rhs"),
12861 )
12862 .await;
12863 let decoded = DFLogicalSubstraitConvertor
12864 .decode(
12865 DFLogicalSubstraitConvertor
12866 .encode(&raw, DefaultSerializer)
12867 .unwrap(),
12868 ctx.state(),
12869 )
12870 .await
12871 .unwrap();
12872 let (optimized, batches) = execute(decoded, &state).await;
12873 assert_no_internal_or_keys(optimized.schema());
12874 assert!(batches.iter().all(|batch| {
12875 batch
12876 .schema()
12877 .fields()
12878 .iter()
12879 .all(|field| !field.name().starts_with("__promql_or_match_"))
12880 }));
12881 assert_eq!(values(&batches, "v"), vec![1.0]);
12882 }
12883
12884 #[tokio::test]
12885 async fn test_direct_or_numeric_value_types() {
12886 let left = tagged_source("lhs", true, ("k", Some("lhs")), DirectOrValue::Int64(0));
12887 let right = tagged_source(
12888 "rhs",
12889 false,
12890 ("k", Some("rhs")),
12891 DirectOrValue::Float64(0.5),
12892 );
12893 let (optimized, batches) = run(
12894 &left,
12895 &right,
12896 direct_or_context("lhs", &["job", "k"], "v"),
12897 direct_or_context("rhs", &["job", "k"], "v"),
12898 &or_modifier("lhs or on(k) rhs"),
12899 )
12900 .await;
12901 assert_eq!(
12902 optimized
12903 .schema()
12904 .field_with_name(None, "v")
12905 .unwrap()
12906 .data_type(),
12907 &ArrowDataType::Float64
12908 );
12909 assert_eq!(values(&batches, "v"), vec![0.5]);
12910 let provider = build_test_table_provider_with_fields(
12911 &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
12912 &[],
12913 )
12914 .await;
12915 let mut planner = PromPlanner {
12916 table_provider: provider,
12917 ctx: PromPlannerContext::default(),
12918 promql_annotations: None,
12919 };
12920 let left_context = direct_or_context("lhs", &["job"], "v");
12921 let right_context = direct_or_context("rhs", &["job"], "v");
12922 let error = planner
12923 .or_operator(
12924 scan(&job_source("lhs", DirectOrValue::Utf8("x"))),
12925 scan(&job_source("rhs", DirectOrValue::Float64(1.0))),
12926 left_context.tag_columns.iter().cloned().collect(),
12927 right_context.tag_columns.iter().cloned().collect(),
12928 left_context,
12929 right_context,
12930 &or_modifier("lhs or on() rhs"),
12931 )
12932 .unwrap_err();
12933 assert!(
12934 error
12935 .to_string()
12936 .contains("OR value fields have incompatible types")
12937 );
12938 }
12939
12940 #[tokio::test]
12941 async fn test_or_with_histogram_quantile_missing_le_column() {
12942 let case = r#"histogram_quantile(0.99, non_existent_histogram_bucket) or normal_metric"#;
12943 let eval_stmt = build_eval_stmt(case);
12944 let table_provider = build_missing_le_or_normal_metric_table_provider().await;
12945
12946 let plan =
12947 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12948 .await
12949 .unwrap();
12950 assert_normal_metric_schema(&plan);
12951 }
12952
12953 #[tokio::test]
12954 async fn test_or_with_right_empty_histogram_restores_left_context() {
12955 let eval_stmt = build_eval_stmt(
12956 r#"abs(sum by(instance) (normal_metric) or histogram_quantile(0.99, sum by(pod) (non_existent_histogram_bucket)))"#,
12957 );
12958 let table_provider = build_missing_le_or_normal_metric_table_provider().await;
12959
12960 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12961 .await
12962 .unwrap();
12963 }
12964
12965 #[tokio::test]
12966 async fn test_or_with_both_empty_histograms() {
12967 let eval_stmt = build_eval_stmt(
12968 r#"histogram_quantile(0.99, sum by(pod) (left_histogram_bucket)) or histogram_quantile(0.99, sum by(instance) (right_histogram_bucket))"#,
12969 );
12970 let table_provider = build_test_table_provider_with_fields(
12971 &[
12972 (
12973 DEFAULT_SCHEMA_NAME.to_string(),
12974 "left_histogram_bucket".to_string(),
12975 ),
12976 (
12977 DEFAULT_SCHEMA_NAME.to_string(),
12978 "right_histogram_bucket".to_string(),
12979 ),
12980 ],
12981 &["pod", "instance"],
12982 )
12983 .await;
12984
12985 let plan =
12986 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12987 .await
12988 .unwrap();
12989 match plan {
12990 LogicalPlan::EmptyRelation(relation) => {
12991 assert!(!relation.produce_one_row);
12992 assert!(!relation.schema.fields().is_empty());
12993 assert!(
12994 relation
12995 .schema
12996 .fields()
12997 .iter()
12998 .any(|field| field.data_type() == &ArrowDataType::Float64)
12999 );
13000 assert!(
13001 relation
13002 .schema
13003 .fields()
13004 .iter()
13005 .any(|field| field.name() == "pod")
13006 );
13007 assert!(
13008 !relation
13009 .schema
13010 .fields()
13011 .iter()
13012 .any(|field| field.name() == "instance")
13013 );
13014 }
13015 _ => panic!("Expected EmptyRelation, but got: {plan:?}"),
13016 }
13017 }
13018
13019 #[tokio::test]
13020 async fn test_nested_or_with_both_empty_histograms() {
13021 for case in [
13022 r#"abs(histogram_quantile(0.99, left_histogram_bucket) or histogram_quantile(0.99, right_histogram_bucket))"#,
13023 r#"(histogram_quantile(0.99, left_histogram_bucket) or histogram_quantile(0.99, right_histogram_bucket)) + 1"#,
13024 ] {
13025 let eval_stmt = build_eval_stmt(case);
13026 let table_provider = build_test_table_provider_with_fields(
13027 &[
13028 (
13029 DEFAULT_SCHEMA_NAME.to_string(),
13030 "left_histogram_bucket".to_string(),
13031 ),
13032 (
13033 DEFAULT_SCHEMA_NAME.to_string(),
13034 "right_histogram_bucket".to_string(),
13035 ),
13036 ],
13037 &["pod", "instance"],
13038 )
13039 .await;
13040
13041 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13042 .await
13043 .unwrap();
13044 }
13045 }
13046
13047 #[tokio::test]
13048 async fn test_or_with_empty_histogram_modifiers() {
13049 for case in [
13050 r#"histogram_quantile(0.99, non_existent_histogram_bucket) or on(pod) normal_metric"#,
13051 r#"normal_metric or ignoring(instance) histogram_quantile(0.99, non_existent_histogram_bucket)"#,
13052 ] {
13053 let eval_stmt = build_eval_stmt(case);
13054 let table_provider = build_missing_le_or_normal_metric_table_provider().await;
13055
13056 let plan =
13057 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13058 .await
13059 .unwrap();
13060 assert_normal_metric_schema(&plan);
13061 }
13062 }
13063
13064 #[tokio::test]
13065 async fn test_unless_preserves_left_context_for_histogram() {
13066 let eval_stmt = build_eval_stmt(
13067 r#"histogram_quantile(0.99, bucket_metric unless on(job) normal_metric) or fallback_metric"#,
13068 );
13069 let state = build_query_engine_state();
13070 let plan = PromPlanner::stmt_to_plan(
13071 build_set_op_context_table_provider().await,
13072 &eval_stmt,
13073 &state,
13074 )
13075 .await
13076 .unwrap();
13077 assert!(contains_histogram_fold(&plan), "{plan:?}");
13078 let (optimized, physical) = optimize_and_create_physical_plan(&state, plan).await;
13079 assert!(contains_histogram_fold(&optimized), "{optimized:?}");
13080 let batches =
13081 datafusion::physical_plan::collect(physical, state.session_state().task_ctx())
13082 .await
13083 .unwrap();
13084 assert!(batches.iter().all(|batch| batch.num_rows() == 0));
13085 }
13086
13087 #[tokio::test]
13088 async fn test_and_preserves_left_context_for_histogram() {
13089 let eval_stmt = build_eval_stmt(
13090 r#"histogram_quantile(0.99, bucket_metric and on(job) normal_metric) or fallback_metric"#,
13091 );
13092 let plan = PromPlanner::stmt_to_plan(
13093 build_set_op_context_table_provider().await,
13094 &eval_stmt,
13095 &build_query_engine_state(),
13096 )
13097 .await
13098 .unwrap();
13099 assert!(contains_histogram_fold(&plan), "{plan:?}");
13100 }
13101
13102 #[tokio::test]
13103 async fn test_and_preserves_left_context_when_le_is_missing() {
13104 let eval_stmt =
13105 build_eval_stmt(r#"histogram_quantile(0.99, normal_metric and on(job) bucket_metric)"#);
13106 let plan = PromPlanner::stmt_to_plan(
13107 build_set_op_context_table_provider().await,
13108 &eval_stmt,
13109 &build_query_engine_state(),
13110 )
13111 .await
13112 .unwrap();
13113 assert!(matches!(&plan, LogicalPlan::EmptyRelation(_)), "{plan:?}");
13114 assert!(!plan.schema().fields().is_empty());
13115 assert!(!contains_histogram_fold(&plan), "{plan:?}");
13116 }
13117
13118 #[tokio::test]
13119 async fn test_or_context_uses_left_qualified_output() {
13120 let case = r#"(normal_metric or other_metric) + 1"#;
13121 let eval_stmt = build_eval_stmt(case);
13122 let state = build_query_engine_state();
13123 let plan =
13124 PromPlanner::stmt_to_plan(build_or_context_table_provider().await, &eval_stmt, &state)
13125 .await
13126 .unwrap();
13127 assert!(
13128 plan.schema()
13129 .fields()
13130 .iter()
13131 .any(|field| field.data_type() == &ArrowDataType::Float64),
13132 "{plan:?}"
13133 );
13134 let (_optimized, _physical) = optimize_and_create_physical_plan(&state, plan).await;
13135 }
13136
13137 #[tokio::test]
13138 async fn test_or_context_uses_left_qualified_empty_histogram_output() {
13139 let case = r#"(abs(histogram_quantile(0.99, non_hist_metric)) or normal_metric) + 1"#;
13140 let eval_stmt = build_eval_stmt(case);
13141 let plan = PromPlanner::stmt_to_plan(
13142 build_or_context_table_provider().await,
13143 &eval_stmt,
13144 &build_query_engine_state(),
13145 )
13146 .await
13147 .unwrap();
13148 assert!(
13149 plan.schema()
13150 .fields()
13151 .iter()
13152 .any(|field| field.data_type() == &ArrowDataType::Float64),
13153 "{plan:?}"
13154 );
13155 }
13156
13157 #[tokio::test]
13158 async fn test_direct_or_preserves_float_and_native_histogram_samples() {
13159 for histogram_on_left in [false, true] {
13160 let (planner, plan) = mixed_direct_or(histogram_on_left).await;
13161
13162 let float_field = &planner.ctx.field_columns[0];
13163 let histogram_field = &planner.ctx.field_columns[1];
13164 assert!(float_field.starts_with(OR_FLOAT_FIELD_PREFIX));
13165 assert!(histogram_field.starts_with(OR_HISTOGRAM_FIELD_PREFIX));
13166 assert_eq!(
13167 plan.schema()
13168 .field_with_name(None, float_field)
13169 .unwrap()
13170 .data_type(),
13171 &ArrowDataType::Float64
13172 );
13173 assert_eq!(
13174 plan.schema()
13175 .field_with_name(None, histogram_field)
13176 .unwrap()
13177 .data_type(),
13178 &native_histogram_value_type().as_arrow_type()
13179 );
13180
13181 let (optimized, batches) = execute(plan, &build_query_engine_state()).await;
13182 assert_no_internal_or_keys(optimized.schema());
13183 let mut sample_kinds = batches
13184 .iter()
13185 .flat_map(|batch| {
13186 let values = batch.column_by_name(float_field).unwrap();
13187 let histograms = batch.column_by_name(histogram_field).unwrap();
13188 (0..batch.num_rows())
13189 .map(|row| (values.is_valid(row), histograms.is_valid(row)))
13190 })
13191 .collect::<Vec<_>>();
13192 sample_kinds.sort_unstable();
13193 assert_eq!(sample_kinds, vec![(false, true), (true, false)]);
13194 }
13195 }
13196
13197 #[tokio::test]
13198 async fn malformed_classic_bucket_does_not_drop_native_histogram() {
13199 let state = build_query_engine_state();
13200 let collector = PromqlAnnotationCollector::default();
13201 let plan = PromPlanner::stmt_to_plan_with_annotations(
13202 operator_table_provider(),
13203 &operator_eval_stmt("histogram_quantile(0.5, bad_classic or bad_native)"),
13204 &state,
13205 Some(collector.clone()),
13206 )
13207 .await
13208 .unwrap();
13209 let value_field = plan
13210 .schema()
13211 .fields()
13212 .iter()
13213 .find(|field| field.data_type() == &ArrowDataType::Float64)
13214 .unwrap()
13215 .name()
13216 .clone();
13217
13218 let (_, batches) = execute(plan, &state).await;
13219 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
13220 assert_eq!(values(&batches, &value_field), vec![0.0]);
13221 let mut warnings = vec![];
13222 let mut infos = vec![];
13223 collector.append_to(&mut warnings, &mut infos);
13224 assert!(warnings.is_empty());
13225 assert!(infos.is_empty());
13226 }
13227
13228 #[tokio::test]
13229 async fn test_mixed_binary_operator_aligns_both_alternative_inputs() {
13230 let state = build_query_engine_state();
13231 let plan = PromPlanner::stmt_to_plan(
13232 operator_table_provider(),
13233 &operator_eval_stmt("(lf or on(tag) lh) * on(tag) (rf or on(tag) rh)"),
13234 &state,
13235 )
13236 .await
13237 .unwrap();
13238 let plan_text = plan.display_indent_schema().to_string();
13239 assert!(
13240 plan_text.contains("prom_native_histogram_mul_scalar"),
13241 "{plan_text}"
13242 );
13243 assert!(
13244 plan_text.contains("prom_native_histogram_scalar_mul"),
13245 "{plan_text}"
13246 );
13247 let float_field = plan
13248 .schema()
13249 .fields()
13250 .iter()
13251 .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX))
13252 .unwrap()
13253 .name()
13254 .clone();
13255 let histogram_field = plan
13256 .schema()
13257 .fields()
13258 .iter()
13259 .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX))
13260 .unwrap()
13261 .name()
13262 .clone();
13263
13264 let (_, batches) = execute(plan, &state).await;
13265 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
13266 assert!(values(&batches, &float_field).is_empty());
13267 let mut sums = histograms(&batches, &histogram_field)
13268 .into_iter()
13269 .map(|histogram| histogram.sum)
13270 .collect::<Vec<_>>();
13271 sums.sort_by(f64::total_cmp);
13272 assert_eq!(sums, vec![2.0, 3.0]);
13273 }
13274
13275 #[tokio::test]
13276 async fn test_mixed_binary_operator_reports_only_dropped_samples() {
13277 for (query, expected_rows, expected_infos) in [
13278 ("(lf or on(tag) lh) + on(tag) (rf or on(tag) rh)", 0, 1),
13279 ("(lf or on(tag) lh) + on(tag) (lf or on(tag) lh)", 2, 0),
13280 ("(lf or on(tag) lh) % on(tag) lh", 0, 1),
13281 ] {
13282 let state = build_query_engine_state();
13283 let annotations = PromqlAnnotationCollector::default();
13284 let plan = PromPlanner::stmt_to_plan_with_annotations(
13285 operator_table_provider(),
13286 &operator_eval_stmt(query),
13287 &state,
13288 Some(annotations.clone()),
13289 )
13290 .await
13291 .unwrap();
13292
13293 let (_, batches) = execute(plan, &state).await;
13294 assert_eq!(
13295 batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
13296 expected_rows,
13297 "{query}"
13298 );
13299 let mut warnings = vec![];
13300 let mut infos = vec![];
13301 annotations.append_to(&mut warnings, &mut infos);
13302 assert!(warnings.is_empty(), "{query}: {warnings:?}");
13303 assert_eq!(infos.len(), expected_infos, "{query}: {infos:?}");
13304 }
13305 }
13306
13307 #[tokio::test]
13308 async fn test_histogram_only_min_drops_empty_aggregate_group() {
13309 let state = build_query_engine_state();
13313 for query in ["min(lh)", "group(min(lh))"] {
13314 let plan = PromPlanner::stmt_to_plan(
13315 operator_table_provider(),
13316 &operator_eval_stmt(query),
13317 &state,
13318 )
13319 .await
13320 .unwrap();
13321 let (_, batches) = execute(plan, &state).await;
13322 assert_eq!(
13323 batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
13324 0,
13325 "{query}"
13326 );
13327 }
13328 }
13329
13330 #[tokio::test]
13331 async fn test_mixed_min_drops_histogram_only_group() {
13332 let state = build_query_engine_state();
13336 let plan = PromPlanner::stmt_to_plan(
13337 operator_table_provider(),
13338 &operator_eval_stmt("min by (tag) (lf or on(tag) lh)"),
13339 &state,
13340 )
13341 .await
13342 .unwrap();
13343 let float_field = plan
13344 .schema()
13345 .fields()
13346 .iter()
13347 .find(|field| field.data_type() == &ArrowDataType::Float64)
13348 .unwrap()
13349 .name()
13350 .clone();
13351 let (_, batches) = execute(plan, &state).await;
13352 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
13353 assert_eq!(values(&batches, &float_field), vec![2.0]);
13354 }
13355
13356 #[tokio::test]
13357 async fn test_mixed_or_can_feed_another_or() {
13358 let state = build_query_engine_state();
13359 let plan = PromPlanner::stmt_to_plan(
13360 operator_table_provider(),
13361 &operator_eval_stmt("lf or on(tag) lh or on(tag) fallback"),
13362 &state,
13363 )
13364 .await
13365 .unwrap();
13366 let float_field = plan
13367 .schema()
13368 .fields()
13369 .iter()
13370 .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX))
13371 .unwrap()
13372 .name()
13373 .clone();
13374 let histogram_field = plan
13375 .schema()
13376 .fields()
13377 .iter()
13378 .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX))
13379 .unwrap()
13380 .name()
13381 .clone();
13382
13383 let (_, batches) = execute(plan, &state).await;
13384 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 3);
13385 let mut float_values = values(&batches, &float_field);
13386 float_values.sort_by(f64::total_cmp);
13387 assert_eq!(float_values, vec![2.0, 7.0]);
13388 assert_eq!(histograms(&batches, &histogram_field).len(), 1);
13389 }
13390
13391 #[tokio::test]
13392 async fn test_mixed_fields_align_with_single_float_vector() {
13393 let (planner, mixed) = mixed_direct_or(false).await;
13394 let scale = tagged_source(
13395 "scale",
13396 false,
13397 ("k", Some("float")),
13398 DirectOrValue::Float64(2.0),
13399 );
13400 let scale = scan(&scale);
13401 let scale_fields = vec!["v".to_string()];
13402 let PromExpr::Binary(binary) = parser::parse("lhs * rhs").unwrap() else {
13403 unreachable!()
13404 };
13405
13406 let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns(
13407 mixed.schema(),
13408 scale.schema(),
13409 &planner.ctx.field_columns,
13410 &scale_fields,
13411 binary.op,
13412 false,
13413 false,
13414 );
13415 assert!(invalid_pairs.is_empty());
13416 assert_eq!(
13417 groups
13418 .iter()
13419 .map(|(output, _)| output.clone())
13420 .collect::<Vec<_>>(),
13421 planner.ctx.field_columns
13422 );
13423 assert_eq!(groups.len(), 2);
13424 assert!(
13425 groups
13426 .iter()
13427 .flat_map(|(_, pairs)| pairs)
13428 .all(|(_, right)| *right == &scale_fields[0])
13429 );
13430
13431 let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns(
13432 scale.schema(),
13433 mixed.schema(),
13434 &scale_fields,
13435 &planner.ctx.field_columns,
13436 binary.op,
13437 false,
13438 false,
13439 );
13440 assert!(invalid_pairs.is_empty());
13441 assert_eq!(
13442 groups
13443 .iter()
13444 .map(|(output, _)| output.clone())
13445 .collect::<Vec<_>>(),
13446 planner.ctx.field_columns
13447 );
13448 assert_eq!(groups.len(), 2);
13449 assert!(
13450 groups
13451 .iter()
13452 .flat_map(|(_, pairs)| pairs)
13453 .all(|(left, _)| *left == &scale_fields[0])
13454 );
13455 }
13456
13457 #[tokio::test]
13458 async fn test_non_bool_comparison_filters_mixed_sample_lanes() {
13459 let (planner, input) = mixed_direct_or(false).await;
13460 let input_schema = input.schema().clone();
13461 let plan = planner
13462 .filter_on_field_column(input, |field| {
13463 if PromPlanner::field_column_is_native_histogram(&input_schema, field) {
13464 Ok(lit(false))
13465 } else {
13466 Ok(col(field).gt(lit(0.0)))
13467 }
13468 })
13469 .unwrap();
13470 let float_field = planner.ctx.field_columns[0].clone();
13471
13472 let (_, batches) = execute(plan, &build_query_engine_state()).await;
13473 assert_eq!(values(&batches, &float_field), vec![1.25]);
13474 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
13475 }
13476
13477 #[tokio::test]
13478 async fn test_mixed_left_and_unless_preserve_sample_lanes() {
13479 for (expression, expected_sample_kind) in [
13480 ("lhs and on(k) mask", (false, true)),
13481 ("lhs unless on(k) mask", (true, false)),
13482 ] {
13483 let (mut planner, left) = mixed_direct_or(false).await;
13484 let left_context = planner.ctx.clone();
13485 let float_field = left_context.field_columns[0].clone();
13486 let histogram_field = left_context.field_columns[1].clone();
13487 let mask = tagged_source(
13488 "mask",
13489 false,
13490 ("k", Some("histogram")),
13491 DirectOrValue::Float64(1.0),
13492 );
13493 let PromExpr::Binary(binary) = parser::parse(expression).unwrap() else {
13494 unreachable!()
13495 };
13496 let plan = planner
13497 .set_op_on_non_field_columns(
13498 left,
13499 scan(&mask),
13500 left_context,
13501 direct_or_context("mask", &["job", "k"], "v"),
13502 binary.op,
13503 &binary.modifier,
13504 )
13505 .unwrap();
13506
13507 let (_, batches) = execute(plan, &build_query_engine_state()).await;
13508 let sample_kinds = batches
13509 .iter()
13510 .flat_map(|batch| {
13511 let floats = batch.column_by_name(&float_field).unwrap();
13512 let histograms = batch.column_by_name(&histogram_field).unwrap();
13513 (0..batch.num_rows())
13514 .map(|row| (floats.is_valid(row), histograms.is_valid(row)))
13515 })
13516 .collect::<Vec<_>>();
13517 assert_eq!(sample_kinds, vec![expected_sample_kind], "{expression}");
13518 }
13519 }
13520
13521 #[tokio::test]
13522 async fn test_mixed_fields_arithmetic_broadcasts_computed_scalar() {
13523 let plan = PromPlanner::stmt_to_plan(
13524 build_test_mixed_native_histogram_table_provider("some_metric").await,
13525 &build_eval_stmt("some_metric * scalar(vector(2))"),
13526 &build_query_engine_state(),
13527 )
13528 .await
13529 .unwrap();
13530 let schema = plan.schema();
13531 assert_eq!(
13532 schema
13533 .field_with_unqualified_name(greptime_value())
13534 .unwrap()
13535 .data_type(),
13536 &ArrowDataType::Float64
13537 );
13538 assert_eq!(
13539 schema
13540 .field_with_unqualified_name(greptime_native_histogram())
13541 .unwrap()
13542 .data_type(),
13543 &native_histogram_value_type().as_arrow_type()
13544 );
13545 assert!(
13546 plan.display_indent_schema()
13547 .to_string()
13548 .contains("prom_native_histogram_mul_scalar"),
13549 "{plan:?}"
13550 );
13551 }
13552
13553 #[tokio::test]
13554 async fn test_unsupported_histogram_binary_does_not_block_or_fallback() {
13555 let state = build_query_engine_state();
13556 let plan = PromPlanner::stmt_to_plan(
13557 operator_table_provider(),
13558 &operator_eval_stmt("((lf or on(tag) lh) % 2) or on(tag) lh"),
13559 &state,
13560 )
13561 .await
13562 .unwrap();
13563 let float_field = plan
13564 .schema()
13565 .fields()
13566 .iter()
13567 .find(|field| field.data_type() == &ArrowDataType::Float64)
13568 .unwrap()
13569 .name()
13570 .clone();
13571 let histogram_field = plan
13572 .schema()
13573 .fields()
13574 .iter()
13575 .find(|field| field.data_type() == &native_histogram_value_type().as_arrow_type())
13576 .unwrap()
13577 .name()
13578 .clone();
13579
13580 let (_, batches) = execute(plan, &state).await;
13581 assert_eq!(values(&batches, &float_field), vec![0.0]);
13582 assert_eq!(histograms(&batches, &histogram_field).len(), 1);
13583 }
13584
13585 #[tokio::test]
13586 async fn test_unary_negates_mixed_float_and_native_histogram_samples() {
13587 for histogram_on_left in [false, true] {
13588 let (mut planner, input) = mixed_direct_or(histogram_on_left).await;
13589 let plan = planner.negate_field_columns(input).unwrap();
13590 assert!(PromPlanner::field_columns_are_alternative_samples(
13591 plan.schema(),
13592 &planner.ctx.field_columns
13593 ));
13594 let float_field = planner
13595 .ctx
13596 .field_columns
13597 .iter()
13598 .find(|field| field.starts_with(OR_FLOAT_FIELD_PREFIX))
13599 .unwrap();
13600 let histogram_field = planner
13601 .ctx
13602 .field_columns
13603 .iter()
13604 .find(|field| field.starts_with(OR_HISTOGRAM_FIELD_PREFIX))
13605 .unwrap();
13606
13607 let (_, batches) = execute(plan, &build_query_engine_state()).await;
13608 assert_eq!(values(&batches, float_field), vec![-1.25]);
13609 let histogram = batches
13610 .iter()
13611 .find_map(|batch| {
13612 let values = batch
13613 .column_by_name(histogram_field)
13614 .unwrap()
13615 .as_any()
13616 .downcast_ref::<datafusion::arrow::array::StructArray>()
13617 .unwrap();
13618 (0..values.len()).find_map(|row| {
13619 common_query::native_histogram::read_histogram(values, row).unwrap()
13620 })
13621 })
13622 .unwrap();
13623 assert_eq!(histogram.count, -1.0);
13624 assert_eq!(histogram.sum, -1.0);
13625 assert_eq!(histogram.reset_hint, CounterResetHint::Gauge);
13626 }
13627 }
13628
13629 #[tokio::test]
13630 async fn test_native_histogram_sum_and_avg_execute_real_batches() {
13631 for op_name in ["sum", "avg"] {
13632 for incompatible in [false, true] {
13633 let mut second = direct_or_histogram();
13634 if incompatible {
13635 second.schema = CUSTOM_BUCKETS_SCHEMA;
13636 second.custom_values = vec![1.0];
13637 }
13638 let collector = PromqlAnnotationCollector::default();
13639 let (mut planner, input) =
13640 mixed_aggregate_input(vec![direct_or_histogram(), second]).await;
13641 planner.promql_annotations = Some(collector.clone());
13642 let histogram_column = planner.ctx.field_columns[1].clone();
13643 planner.ctx.field_columns = vec![histogram_column.clone()];
13644 let input = LogicalPlanBuilder::from(input)
13645 .project([col("ts"), col(&histogram_column)])
13646 .unwrap()
13647 .build()
13648 .unwrap();
13649 let PromExpr::Aggregate(AggregateExpr { op, param, .. }) =
13650 parser::parse(&format!("{op_name}(mixed)")).unwrap()
13651 else {
13652 unreachable!()
13653 };
13654 let (aggregate_exprs, _) =
13655 planner.create_aggregate_exprs(op, ¶m, &input).unwrap();
13656 let plan = LogicalPlanBuilder::from(input)
13657 .aggregate(vec![col("ts")], aggregate_exprs)
13658 .unwrap()
13659 .filter(planner.create_empty_values_filter_expr(false).unwrap())
13660 .unwrap()
13661 .build()
13662 .unwrap();
13663
13664 let (_, batches) = execute(plan, &build_query_engine_state()).await;
13665 let mut warnings = vec![];
13666 let mut infos = vec![];
13667 collector.append_to(&mut warnings, &mut infos);
13668 assert!(infos.is_empty());
13669 if incompatible {
13670 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
13671 assert!(warnings.iter().any(|warning| {
13672 warning
13673 == &format!(
13674 "prom_native_histogram_agg_{op_name}: dropped native histogram aggregate with incompatible schemas"
13675 )
13676 }));
13677 } else {
13678 let histograms = histograms(&batches, &histogram_column);
13679 assert_eq!(histograms.len(), 1);
13680 let expected = if op_name == "sum" { 2.0 } else { 1.0 };
13681 assert_eq!(histograms[0].count, expected);
13682 assert_eq!(histograms[0].sum, expected);
13683 assert!(warnings.is_empty());
13684 }
13685 }
13686 }
13687 }
13688
13689 #[tokio::test]
13690 async fn test_canonical_mixed_count_group_and_count_values_execute() {
13691 let state = build_query_engine_state();
13692 for (query, expected) in [
13693 ("count(some_metric)", vec![2.0]),
13694 ("group(some_metric)", vec![1.0]),
13695 (r#"count_values("sample", some_metric)"#, vec![1.0, 1.0]),
13696 ] {
13697 let plan = PromPlanner::stmt_to_plan(
13698 build_test_mixed_native_histogram_table_provider("some_metric").await,
13699 &operator_eval_stmt(query),
13700 &state,
13701 )
13702 .await
13703 .unwrap();
13704 assert!(
13705 plan.schema()
13706 .fields()
13707 .iter()
13708 .all(|field| !field.name().starts_with("__promql_sample_count")),
13709 "{query}: {plan:?}"
13710 );
13711 let value_fields = plan
13712 .schema()
13713 .fields()
13714 .iter()
13715 .filter(|field| {
13716 matches!(
13717 field.data_type(),
13718 ArrowDataType::Float64 | ArrowDataType::Int64 | ArrowDataType::UInt64
13719 ) || field.data_type() == &native_histogram_value_type().as_arrow_type()
13720 })
13721 .collect::<Vec<_>>();
13722 assert_eq!(value_fields.len(), 1, "{query}: {plan:?}");
13723 assert_ne!(
13724 value_fields[0].data_type(),
13725 &native_histogram_value_type().as_arrow_type(),
13726 "{query}: {plan:?}"
13727 );
13728 let value_column = value_fields[0].name().clone();
13729
13730 let (_, batches) = execute(plan, &state).await;
13731 let mut actual = numeric_values(&batches, &value_column);
13732 actual.sort_by(f64::total_cmp);
13733 assert_eq!(actual, expected, "{query}");
13734
13735 if query.starts_with("count_values") {
13736 let mut sample_labels = batches
13737 .iter()
13738 .flat_map(|batch| {
13739 batch
13740 .column_by_name("sample")
13741 .unwrap()
13742 .as_any()
13743 .downcast_ref::<StringArray>()
13744 .unwrap()
13745 .iter()
13746 .flatten()
13747 .map(str::to_string)
13748 })
13749 .collect::<Vec<_>>();
13750 sample_labels.sort();
13751 let mut expected_labels =
13752 vec!["2".to_string(), direct_or_histogram().promql_string()];
13753 expected_labels.sort();
13754 assert_eq!(sample_labels, expected_labels);
13755 }
13756 }
13757 }
13758
13759 #[tokio::test]
13760 async fn test_mixed_or_sum_aggregates_each_sample_type() {
13761 let PromExpr::Aggregate(AggregateExpr { op, param, .. }) =
13762 parser::parse("sum(lhs)").unwrap()
13763 else {
13764 unreachable!()
13765 };
13766
13767 let collector = PromqlAnnotationCollector::default();
13768 let (mut planner, input) = mixed_direct_or(false).await;
13769 planner.promql_annotations = Some(collector.clone());
13770 let float_column = planner.ctx.field_columns[0].clone();
13771 let histogram_column = planner.ctx.field_columns[1].clone();
13772 let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap();
13773 let plan = LogicalPlanBuilder::from(input)
13774 .aggregate(vec![col("ts"), col("k")], aggregate_exprs)
13775 .unwrap()
13776 .filter(
13777 planner
13778 .mixed_aggregate_filter_expr(op, &float_column, &histogram_column)
13779 .unwrap(),
13780 )
13781 .unwrap()
13782 .project([
13783 col(&float_column),
13784 col(&histogram_column),
13785 col("ts"),
13786 col("k"),
13787 ])
13788 .unwrap()
13789 .build()
13790 .unwrap();
13791
13792 let (_, batches) = execute(plan, &build_query_engine_state()).await;
13793 assert_eq!(values(&batches, &float_column), vec![1.25]);
13794 let histogram = batches
13795 .iter()
13796 .find_map(|batch| {
13797 let values = batch
13798 .column_by_name(&histogram_column)?
13799 .as_any()
13800 .downcast_ref::<datafusion::arrow::array::StructArray>()?;
13801 (0..values.len()).find_map(|row| {
13802 common_query::native_histogram::read_histogram(values, row).unwrap()
13803 })
13804 })
13805 .unwrap();
13806 assert_eq!(histogram.count, 1.0);
13807 let mut warnings = vec![];
13808 let mut infos = vec![];
13809 collector.append_to(&mut warnings, &mut infos);
13810 assert!(warnings.is_empty());
13811
13812 let collector = PromqlAnnotationCollector::default();
13813 let (mut planner, input) = mixed_direct_or(false).await;
13814 planner.promql_annotations = Some(collector.clone());
13815 let float_column = planner.ctx.field_columns[0].clone();
13816 let histogram_column = planner.ctx.field_columns[1].clone();
13817 let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap();
13818 let plan = LogicalPlanBuilder::from(input)
13819 .aggregate(vec![col("ts")], aggregate_exprs)
13820 .unwrap()
13821 .filter(
13822 planner
13823 .mixed_aggregate_filter_expr(op, &float_column, &histogram_column)
13824 .unwrap(),
13825 )
13826 .unwrap()
13827 .project([col(&float_column), col(&histogram_column), col("ts")])
13828 .unwrap()
13829 .build()
13830 .unwrap();
13831
13832 let (_, batches) = execute(plan, &build_query_engine_state()).await;
13833 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
13834 let mut warnings = vec![];
13835 let mut infos = vec![];
13836 collector.append_to(&mut warnings, &mut infos);
13837 assert_eq!(
13838 warnings,
13839 vec![
13840 "sum: dropped aggregation result containing both float and native histogram samples"
13841 ]
13842 );
13843 }
13844
13845 #[tokio::test]
13846 async fn test_mixed_or_sum_drops_incompatible_mixed_group() {
13847 let PromExpr::Aggregate(AggregateExpr { op, param, .. }) =
13848 parser::parse("sum(lhs)").unwrap()
13849 else {
13850 unreachable!()
13851 };
13852 let mut custom = direct_or_histogram();
13853 custom.schema = CUSTOM_BUCKETS_SCHEMA;
13854 custom.custom_values = vec![1.0];
13855 let collector = PromqlAnnotationCollector::default();
13856 let (mut planner, input) = mixed_aggregate_input(vec![direct_or_histogram(), custom]).await;
13857 planner.promql_annotations = Some(collector.clone());
13858 let float_column = planner.ctx.field_columns[0].clone();
13859 let histogram_column = planner.ctx.field_columns[1].clone();
13860 let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap();
13861 let plan = LogicalPlanBuilder::from(input)
13862 .aggregate(vec![col("ts")], aggregate_exprs)
13863 .unwrap()
13864 .filter(
13865 planner
13866 .mixed_aggregate_filter_expr(op, &float_column, &histogram_column)
13867 .unwrap(),
13868 )
13869 .unwrap()
13870 .project([col(&float_column), col(&histogram_column), col("ts")])
13871 .unwrap()
13872 .build()
13873 .unwrap();
13874
13875 let (_, batches) = execute(plan, &build_query_engine_state()).await;
13876 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
13877 let mut warnings = vec![];
13878 let mut infos = vec![];
13879 collector.append_to(&mut warnings, &mut infos);
13880 assert!(warnings.iter().any(|warning| {
13881 warning
13882 == "sum: dropped aggregation result containing both float and native histogram samples"
13883 }));
13884 }
13885
13886 #[tokio::test]
13887 async fn test_mixed_or_min_records_only_present_histograms() {
13888 let PromExpr::Aggregate(AggregateExpr { op, param, .. }) =
13889 parser::parse("min(lhs)").unwrap()
13890 else {
13891 unreachable!()
13892 };
13893 let expected_info = "min: dropped native histogram samples because this aggregation is not supported for native histograms";
13894
13895 for (histograms, expected_infos) in [
13896 (vec![], vec![]),
13897 (vec![direct_or_histogram()], vec![expected_info]),
13898 ] {
13899 let collector = PromqlAnnotationCollector::default();
13900 let (mut planner, input) = mixed_aggregate_input(histograms).await;
13901 planner.promql_annotations = Some(collector.clone());
13902 let float_column = planner.ctx.field_columns[0].clone();
13903 let histogram_column = planner.ctx.field_columns[1].clone();
13904 let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap();
13905 let plan = LogicalPlanBuilder::from(input)
13906 .aggregate(vec![col("ts")], aggregate_exprs)
13907 .unwrap()
13908 .filter(
13909 planner
13910 .mixed_ignored_histogram_filter_expr(op, &histogram_column)
13911 .unwrap(),
13912 )
13913 .unwrap()
13914 .project([col(&float_column), col("ts")])
13915 .unwrap()
13916 .build()
13917 .unwrap();
13918
13919 let (_, batches) = execute(plan, &build_query_engine_state()).await;
13920 assert_eq!(values(&batches, &float_column), vec![1.25]);
13921 let mut warnings = vec![];
13922 let mut infos = vec![];
13923 collector.append_to(&mut warnings, &mut infos);
13924 assert!(warnings.is_empty());
13925 assert_eq!(infos, expected_infos);
13926 }
13927 }
13928
13929 #[tokio::test]
13930 async fn test_mixed_or_value_aliases_do_not_replace_labels() {
13931 let left = source(
13932 "lhs",
13933 false,
13934 1,
13935 vec![("job", Some("job")), ("k", Some("float"))],
13936 DirectOrValue::Float64(1.0),
13937 );
13938 let right = source(
13939 "rhs",
13940 false,
13941 1,
13942 vec![
13943 ("job", Some("job")),
13944 ("k", Some("histogram")),
13945 (greptime_value(), Some("value-label")),
13946 ],
13947 DirectOrValue::NativeHistogram(direct_or_histogram()),
13948 );
13949 let table_provider = build_test_table_provider_with_fields(
13950 &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
13951 &[],
13952 )
13953 .await;
13954 let mut planner = PromPlanner {
13955 table_provider,
13956 ctx: PromPlannerContext::default(),
13957 promql_annotations: None,
13958 };
13959 let left = LogicalPlanBuilder::from(scan(&left))
13960 .project(vec![
13961 col("ts"),
13962 col("job"),
13963 col("k"),
13964 col("v").alias(greptime_value()),
13965 ])
13966 .unwrap()
13967 .build()
13968 .unwrap();
13969 let left_context = direct_or_context("lhs", &["job", "k"], greptime_value());
13970 let right_context = direct_or_context("rhs", &["job", "k", greptime_value()], "v");
13971 let plan = planner
13972 .or_operator(
13973 left,
13974 scan(&right),
13975 left_context.tag_columns.iter().cloned().collect(),
13976 right_context.tag_columns.iter().cloned().collect(),
13977 left_context,
13978 right_context,
13979 &or_modifier("lhs or on(k) rhs"),
13980 )
13981 .unwrap();
13982
13983 assert_eq!(
13984 plan.schema()
13985 .field_with_name(None, greptime_value())
13986 .unwrap()
13987 .data_type(),
13988 &ArrowDataType::Utf8
13989 );
13990 assert!(
13991 planner
13992 .ctx
13993 .field_columns
13994 .iter()
13995 .all(|field| { field != greptime_value() && field != greptime_native_histogram() })
13996 );
13997 assert!(PromPlanner::field_columns_are_alternative_samples(
13998 plan.schema(),
13999 &planner.ctx.field_columns
14000 ));
14001 let (_, batches) = execute(plan, &build_query_engine_state()).await;
14002 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
14003 let labels = batches
14004 .iter()
14005 .flat_map(|batch| {
14006 batch
14007 .column_by_name(greptime_value())
14008 .unwrap()
14009 .as_any()
14010 .downcast_ref::<StringArray>()
14011 .unwrap()
14012 .iter()
14013 .flatten()
14014 })
14015 .collect::<Vec<_>>();
14016 assert_eq!(labels, vec!["value-label"]);
14017 }
14018
14019 #[tokio::test]
14020 async fn test_mixed_or_routes_float_histogram_and_label_functions() {
14021 for (function, expected) in [("abs", 1.25), ("round", 1.0), ("histogram_count", 1.0)] {
14022 let (mut planner, input) = mixed_direct_or(false).await;
14023 let preserve_any_value = PromPlanner::field_columns_are_alternative_samples(
14024 input.schema(),
14025 &planner.ctx.field_columns,
14026 );
14027 let PromExpr::Call(call) = parser::parse(&format!("{function}(lhs)")).unwrap() else {
14028 unreachable!()
14029 };
14030 let state = build_query_engine_state();
14031 let (mut exprs, _) = planner
14032 .create_function_expr(&call.func, vec![], input.schema(), &state)
14033 .unwrap();
14034 exprs.insert(0, planner.create_time_index_column_expr().unwrap());
14035 exprs.extend(planner.create_tag_column_exprs().unwrap());
14036 let plan = LogicalPlanBuilder::from(input)
14037 .project(exprs)
14038 .unwrap()
14039 .filter(
14040 planner
14041 .create_empty_values_filter_expr(preserve_any_value)
14042 .unwrap(),
14043 )
14044 .unwrap()
14045 .build()
14046 .unwrap();
14047 let (_, batches) = execute(plan, &state).await;
14048 let values = batches
14049 .iter()
14050 .flat_map(|batch| {
14051 batch
14052 .schema()
14053 .fields()
14054 .iter()
14055 .position(|field| field.data_type() == &ArrowDataType::Float64)
14056 .map(|index| {
14057 batch
14058 .column(index)
14059 .as_any()
14060 .downcast_ref::<Float64Array>()
14061 .unwrap()
14062 .iter()
14063 .flatten()
14064 })
14065 .into_iter()
14066 .flatten()
14067 })
14068 .collect::<Vec<_>>();
14069 assert_eq!(values, vec![expected], "{function}");
14070 }
14071
14072 let (mut planner, input) = mixed_direct_or(false).await;
14073 let preserve_any_value = PromPlanner::field_columns_are_alternative_samples(
14074 input.schema(),
14075 &planner.ctx.field_columns,
14076 );
14077 let PromExpr::Call(call) =
14078 parser::parse(r#"label_replace(lhs, "copy", "$1", "k", "(.*)")"#).unwrap()
14079 else {
14080 unreachable!()
14081 };
14082 let args = planner.create_function_args(&call.args.args).unwrap();
14083 let state = build_query_engine_state();
14084 let (mut exprs, _) = planner
14085 .create_function_expr(&call.func, args.literals, input.schema(), &state)
14086 .unwrap();
14087 exprs.insert(0, planner.create_time_index_column_expr().unwrap());
14088 exprs.extend(planner.create_tag_column_exprs().unwrap());
14089 let plan = LogicalPlanBuilder::from(input)
14090 .project(exprs)
14091 .unwrap()
14092 .filter(
14093 planner
14094 .create_empty_values_filter_expr(preserve_any_value)
14095 .unwrap(),
14096 )
14097 .unwrap()
14098 .build()
14099 .unwrap();
14100 let (_, batches) = execute(plan, &state).await;
14101 let sample_count = batches.iter().map(RecordBatch::num_rows).sum::<usize>();
14102 assert_eq!(sample_count, 2);
14103 }
14104}