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 0,
591 range_ms,
592 time_index_column,
593 self.ctx.field_columns.clone(),
594 divide_plan,
595 )
596 .context(DataFusionPlanningSnafu)?;
597
598 Ok(LogicalPlan::Extension(Extension {
599 node: Arc::new(manipulate),
600 }))
601 }
602
603 async fn prom_aggr_expr_to_plan(
604 &mut self,
605 query_engine_state: &QueryEngineState,
606 aggr_expr: &AggregateExpr,
607 ) -> Result<LogicalPlan> {
608 let AggregateExpr {
609 op,
610 expr,
611 modifier,
612 param,
613 } = aggr_expr;
614
615 let mut input = self.prom_expr_to_plan(expr, query_engine_state).await?;
616 let input_has_tsid = input.schema().fields().iter().any(|field| {
617 field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
618 && field.data_type() == &ArrowDataType::UInt64
619 });
620
621 let required_group_tags = match modifier {
624 None => BTreeSet::new(),
625 Some(LabelModifier::Include(labels)) => labels
626 .labels
627 .iter()
628 .filter(|label| !is_metric_engine_internal_column(label.as_str()))
629 .cloned()
630 .collect(),
631 Some(LabelModifier::Exclude(labels)) => {
632 let mut all_tags = self.collect_row_key_tag_columns_from_plan(&input)?;
633 for label in &labels.labels {
634 let _ = all_tags.remove(label);
635 }
636 all_tags
637 }
638 };
639
640 if !required_group_tags.is_empty()
641 && required_group_tags
642 .iter()
643 .any(|tag| Self::find_case_sensitive_column(input.schema(), tag.as_str()).is_none())
644 {
645 input = self.ensure_tag_columns_available(input, &required_group_tags)?;
646 self.refresh_tag_columns_from_schema(input.schema());
647 }
648
649 match (*op).id() {
650 token::T_TOPK | token::T_BOTTOMK => {
651 self.prom_topk_bottomk_to_plan(aggr_expr, input).await
652 }
653 _ => {
654 let input_tag_columns = if input_has_tsid {
658 self.collect_row_key_tag_columns_from_plan(&input)?
659 .into_iter()
660 .collect::<Vec<_>>()
661 } else {
662 self.ctx.tag_columns.clone()
663 };
664 let mut group_exprs = self.agg_modifier_to_col(input.schema(), modifier, true)?;
667 let mixed_sample_columns =
668 Self::alternative_sample_columns(input.schema(), &self.ctx.field_columns)
669 .map(|(float, histogram)| (float.to_string(), histogram.to_string()));
670 let preserve_any_value = mixed_sample_columns.is_some();
675 let has_native_histogram = preserve_any_value
676 || self.all_field_columns_are_native_histograms(input.schema());
677 let (mut aggr_exprs, prev_field_exprs) =
679 self.create_aggregate_exprs(*op, param, &input)?;
680 let prev_field_exprs =
681 normalize_cols(prev_field_exprs, &input).context(DataFusionPlanningSnafu)?;
682
683 let keep_tsid = op.id() != token::T_COUNT_VALUES
684 && input_has_tsid
685 && input_tag_columns.iter().collect::<HashSet<_>>()
686 == self.ctx.tag_columns.iter().collect::<HashSet<_>>();
687
688 if keep_tsid {
689 aggr_exprs.push(
690 first_value(
691 DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME)),
692 vec![],
693 )
694 .alias(DATA_SCHEMA_TSID_COLUMN_NAME),
695 );
696 }
697 self.ctx.use_tsid = keep_tsid;
698
699 let builder = LogicalPlanBuilder::from(input);
701 let builder = if op.id() == token::T_COUNT_VALUES {
702 let label = Self::get_param_value_as_str(*op, param)?;
703 let count_value_exprs = prev_field_exprs.iter().map(|expr| {
706 match expr {
707 DfExpr::Column(column) => DfExpr::Column(column.clone()),
708 _ => DfExpr::Column(Column::from_name(expr.schema_name().to_string())),
709 }
710 .alias(label)
711 });
712 let aggregate_group_exprs = group_exprs
713 .iter()
714 .cloned()
715 .chain(prev_field_exprs.clone())
716 .collect::<Vec<_>>();
717 group_exprs.push(col(label));
718 let project_fields = self
719 .create_field_column_exprs()?
720 .into_iter()
721 .chain(self.create_tag_column_exprs()?)
722 .chain(Some(self.create_time_index_column_expr()?))
723 .chain(count_value_exprs);
724
725 builder
726 .aggregate(aggregate_group_exprs, aggr_exprs)
727 .context(DataFusionPlanningSnafu)?
728 .project(project_fields)
729 .context(DataFusionPlanningSnafu)?
730 } else {
731 builder
732 .aggregate(group_exprs.clone(), aggr_exprs)
733 .context(DataFusionPlanningSnafu)?
734 };
735
736 let builder = if let Some((float, histogram)) = mixed_sample_columns {
737 let builder = match op.id() {
738 token::T_SUM | token::T_AVG => builder
739 .filter(self.mixed_aggregate_filter_expr(*op, &float, &histogram)?)
740 .context(DataFusionPlanningSnafu)?,
741 token::T_MIN
742 | token::T_MAX
743 | token::T_STDDEV
744 | token::T_STDVAR
745 | token::T_QUANTILE => builder
746 .filter(self.mixed_ignored_histogram_filter_expr(*op, &histogram)?)
747 .context(DataFusionPlanningSnafu)?,
748 _ => builder,
749 };
750
751 match op.id() {
752 token::T_SUM
753 | token::T_AVG
754 | token::T_MIN
755 | token::T_MAX
756 | token::T_STDDEV
757 | token::T_STDVAR
758 | token::T_QUANTILE => {
759 let project_fields = self
760 .create_field_column_exprs()?
761 .into_iter()
762 .chain(self.create_tag_column_exprs()?)
763 .chain(self.ctx.use_tsid.then_some(DfExpr::Column(
764 Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME),
765 )))
766 .chain(Some(self.create_time_index_column_expr()?));
767 builder
768 .project(project_fields)
769 .context(DataFusionPlanningSnafu)?
770 }
771 _ => builder,
772 }
773 } else {
774 builder
775 };
776
777 let builder = if has_native_histogram {
782 builder
783 .filter(self.create_empty_values_filter_expr(preserve_any_value)?)
784 .context(DataFusionPlanningSnafu)?
785 } else {
786 builder
787 };
788
789 let sort_expr = group_exprs.into_iter().map(|expr| expr.sort(true, false));
790
791 builder
792 .sort(sort_expr)
793 .context(DataFusionPlanningSnafu)?
794 .build()
795 .context(DataFusionPlanningSnafu)
796 }
797 }
798 }
799
800 async fn prom_topk_bottomk_to_plan(
802 &mut self,
803 aggr_expr: &AggregateExpr,
804 input: LogicalPlan,
805 ) -> Result<LogicalPlan> {
806 let AggregateExpr {
807 op,
808 param,
809 modifier,
810 ..
811 } = aggr_expr;
812
813 let input_has_tsid = input.schema().fields().iter().any(|field| {
814 field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
815 && field.data_type() == &ArrowDataType::UInt64
816 });
817 self.ctx.use_tsid = input_has_tsid;
818
819 let group_exprs = self.agg_modifier_to_col(input.schema(), modifier, false)?;
820
821 let mut input = input;
822 if let Some((float_column, histogram_column)) =
823 Self::alternative_sample_columns(input.schema(), &self.ctx.field_columns)
824 .map(|(float, histogram)| (float.to_string(), histogram.to_string()))
825 {
826 let drop_histogram = DfExpr::ScalarFunction(ScalarFunction {
827 func: Arc::new(NativeHistogramDrop::bool_false_udf(
828 format!(
829 "{}: dropped native histogram samples because this aggregation is not supported for native histograms",
830 op
831 ),
832 self.promql_annotations.clone(),
833 )),
834 args: vec![col(&histogram_column)],
835 });
836 let keep_float = when(col(&histogram_column).is_not_null(), drop_histogram)
837 .otherwise(col(&float_column).is_not_null())
838 .context(DataFusionPlanningSnafu)?;
839 input = LogicalPlanBuilder::from(input)
840 .filter(keep_float)
841 .context(DataFusionPlanningSnafu)?
842 .build()
843 .context(DataFusionPlanningSnafu)?;
844 self.ctx.field_columns = vec![float_column];
845 }
846
847 if self.all_field_columns_are_native_histograms(input.schema()) {
848 let promql_annotations = self.promql_annotations.clone();
849 let input = self.projection_for_each_field_column(input, |col| {
850 Ok(DfExpr::ScalarFunction(ScalarFunction {
851 func: Arc::new(NativeHistogramDrop::float_null_udf(
852 format!(
853 "{}: dropped native histogram samples because this aggregation is not supported for native histograms",
854 op
855 ),
856 promql_annotations.clone(),
857 )),
858 args: vec![DfExpr::Column(Column::from_name(col))],
859 }))
860 })?;
861 return LogicalPlanBuilder::from(input)
862 .filter(self.create_empty_values_filter_expr(false)?)
863 .context(DataFusionPlanningSnafu)?
864 .build()
865 .context(DataFusionPlanningSnafu);
866 }
867
868 let val = Self::get_param_as_literal_expr(
869 param.as_deref(),
870 Some(*op),
871 Some(ArrowDataType::Float64),
872 )?;
873
874 let window_exprs = self.create_window_exprs(*op, group_exprs.clone(), &input)?;
876
877 let rank_columns: Vec<_> = window_exprs
878 .iter()
879 .map(|expr| expr.schema_name().to_string())
880 .collect();
881
882 let filter: DfExpr = rank_columns
885 .iter()
886 .fold(None, |expr, rank| {
887 let predicate = DfExpr::BinaryExpr(BinaryExpr {
888 left: Box::new(col(rank)),
889 op: Operator::LtEq,
890 right: Box::new(val.clone()),
891 });
892
893 match expr {
894 None => Some(predicate),
895 Some(expr) => Some(DfExpr::BinaryExpr(BinaryExpr {
896 left: Box::new(expr),
897 op: Operator::Or,
898 right: Box::new(predicate),
899 })),
900 }
901 })
902 .unwrap();
903
904 let rank_columns: Vec<_> = rank_columns.into_iter().map(col).collect();
905
906 let mut new_group_exprs = group_exprs.clone();
907 new_group_exprs.extend(rank_columns);
909
910 let group_sort_expr = new_group_exprs
911 .into_iter()
912 .map(|expr| expr.sort(true, false));
913
914 let project_fields = self
915 .create_field_column_exprs()?
916 .into_iter()
917 .chain(self.create_tag_column_exprs()?)
918 .chain(
919 self.ctx
920 .use_tsid
921 .then_some(DfExpr::Column(Column::from_name(
922 DATA_SCHEMA_TSID_COLUMN_NAME,
923 ))),
924 )
925 .chain(Some(self.create_time_index_column_expr()?));
926
927 LogicalPlanBuilder::from(input)
928 .window(window_exprs)
929 .context(DataFusionPlanningSnafu)?
930 .filter(filter)
931 .context(DataFusionPlanningSnafu)?
932 .sort(group_sort_expr)
933 .context(DataFusionPlanningSnafu)?
934 .project(project_fields)
935 .context(DataFusionPlanningSnafu)?
936 .build()
937 .context(DataFusionPlanningSnafu)
938 }
939
940 async fn prom_unary_expr_to_plan(
941 &mut self,
942 query_engine_state: &QueryEngineState,
943 unary_expr: &UnaryExpr,
944 ) -> Result<LogicalPlan> {
945 let UnaryExpr { expr } = unary_expr;
946 let input = self.prom_expr_to_plan(expr, query_engine_state).await?;
948 self.negate_field_columns(input)
949 }
950
951 fn negate_field_columns(&mut self, input: LogicalPlan) -> Result<LogicalPlan> {
952 let input_schema = input.schema().clone();
953 self.projection_for_each_field_column(input, |col| {
954 if Self::field_column_is_native_histogram(&input_schema, col) {
955 Ok(DfExpr::ScalarFunction(ScalarFunction {
956 func: Arc::new(NativeHistogramNeg::scalar_udf()),
957 args: vec![DfExpr::Column(col.into())],
958 }))
959 } else {
960 Ok(DfExpr::Negative(Box::new(DfExpr::Column(col.into()))))
961 }
962 })
963 }
964
965 async fn try_plan_binary_island(
966 &mut self,
967 binary_expr: &PromBinaryExpr,
968 ) -> Result<Option<LogicalPlan>> {
969 let original_ctx = self.ctx.clone();
970 let mut collect_env = IslandCollectEnv::default();
971 let Some(island_expr) =
972 IslandExpr::try_new(&PromExpr::Binary(binary_expr.clone()), &mut collect_env)
973 else {
974 return Ok(None);
975 };
976
977 if collect_env.leaves.is_empty()
978 || collect_env.vector_occurrences <= collect_env.leaves.len()
979 {
980 return Ok(None);
981 }
982
983 let mut planned_leaves = Vec::with_capacity(collect_env.leaves.len());
984 for (idx, leaf) in collect_env.leaves.iter().enumerate() {
985 let plan = self
986 .prom_vector_selector_to_plan(&leaf.selector, false)
987 .await?;
988 let ctx = self.ctx.clone();
989 let alias = TableReference::bare(format!("{BINARY_ISLAND_LEAF_ALIAS_PREFIX}{idx}"));
990 let plan = LogicalPlanBuilder::from(plan)
991 .alias(alias.clone())
992 .context(DataFusionPlanningSnafu)?
993 .build()
994 .context(DataFusionPlanningSnafu)?;
995 planned_leaves.push(PlannedIslandLeaf {
996 plan,
997 ctx,
998 alias,
999 display_table: leaf.display_table.clone(),
1000 });
1001 }
1002
1003 if planned_leaves.iter().any(|leaf| {
1004 Self::field_columns_contain_native_histogram(
1005 leaf.plan.schema(),
1006 &leaf.ctx.field_columns,
1007 )
1008 }) {
1009 self.ctx = original_ctx;
1010 return Ok(None);
1011 }
1012
1013 if !Self::binary_island_join_contexts_supported(&planned_leaves) {
1014 self.ctx = original_ctx;
1015 return Ok(None);
1016 }
1017
1018 let mut input = planned_leaves[0].plan.clone();
1019 for right_idx in 1..planned_leaves.len() {
1020 input = self.join_binary_island_leaf(
1021 input,
1022 &planned_leaves[0],
1023 &planned_leaves[right_idx],
1024 )?;
1025 }
1026
1027 let field_exprs =
1028 Self::build_binary_island_field_exprs(&island_expr, &planned_leaves, input.schema())?;
1029 if field_exprs.scalar || field_exprs.exprs.is_empty() {
1030 self.ctx = original_ctx;
1031 return Ok(None);
1032 }
1033
1034 let plan = self.project_binary_island(
1035 input,
1036 &planned_leaves[0].alias,
1037 &planned_leaves[0].ctx,
1038 field_exprs,
1039 )?;
1040 Ok(Some(plan))
1041 }
1042
1043 fn binary_island_join_contexts_supported(leaves: &[PlannedIslandLeaf]) -> bool {
1044 if leaves
1045 .iter()
1046 .any(|leaf| leaf.ctx.time_index_column.is_none())
1047 {
1048 return false;
1049 }
1050
1051 if leaves.len() <= 1 {
1052 return true;
1053 }
1054
1055 let first_tags = leaves[0].ctx.tag_columns.iter().collect::<BTreeSet<_>>();
1056
1057 leaves.iter().skip(1).all(|leaf| {
1058 (Self::plan_has_tsid_column(&leaves[0].plan) && Self::plan_has_tsid_column(&leaf.plan))
1059 || leaf.ctx.tag_columns.iter().collect::<BTreeSet<_>>() == first_tags
1060 })
1061 }
1062
1063 fn join_binary_island_leaf(
1064 &self,
1065 left: LogicalPlan,
1066 first_leaf: &PlannedIslandLeaf,
1067 right_leaf: &PlannedIslandLeaf,
1068 ) -> Result<LogicalPlan> {
1069 let only_join_time_index = (first_leaf.ctx.tag_columns.is_empty()
1070 || right_leaf.ctx.tag_columns.is_empty())
1071 && !first_leaf
1072 .ctx
1073 .tag_columns
1074 .iter()
1075 .chain(&right_leaf.ctx.tag_columns)
1076 .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL);
1077 let (mut left_keys, mut right_keys, force_empty_join) = self.binary_join_key_columns(
1078 left.schema(),
1079 right_leaf.plan.schema(),
1080 &first_leaf.ctx,
1081 &right_leaf.ctx,
1082 only_join_time_index,
1083 &None,
1084 )?;
1085
1086 if let (Some(left_time_index_column), Some(right_time_index_column)) = (
1087 first_leaf.ctx.time_index_column.clone(),
1088 right_leaf.ctx.time_index_column.clone(),
1089 ) {
1090 left_keys.insert(left_time_index_column);
1091 right_keys.insert(right_time_index_column);
1092 }
1093
1094 LogicalPlanBuilder::from(left)
1095 .join_detailed(
1096 right_leaf.plan.clone(),
1097 JoinType::Inner,
1098 (
1099 left_keys
1100 .into_iter()
1101 .map(|name| Column::new(Some(first_leaf.alias.clone()), name))
1102 .collect::<Vec<_>>(),
1103 right_keys
1104 .into_iter()
1105 .map(|name| Column::new(Some(right_leaf.alias.clone()), name))
1106 .collect::<Vec<_>>(),
1107 ),
1108 force_empty_join.then_some(lit(false)),
1109 NullEquality::NullEqualsNull,
1110 )
1111 .context(DataFusionPlanningSnafu)?
1112 .build()
1113 .context(DataFusionPlanningSnafu)
1114 }
1115
1116 fn build_binary_island_field_exprs(
1117 expr: &IslandExpr,
1118 leaves: &[PlannedIslandLeaf],
1119 schema: &DFSchemaRef,
1120 ) -> Result<IslandFieldExprs> {
1121 match expr {
1122 IslandExpr::VectorLeaf(id) => {
1123 let leaf = &leaves[*id];
1124 let exprs = leaf
1125 .ctx
1126 .field_columns
1127 .iter()
1128 .map(|field| {
1129 schema
1130 .qualified_field_with_name(Some(&leaf.alias), field)
1131 .context(DataFusionPlanningSnafu)
1132 .map(|field| DfExpr::Column(field.into()))
1133 })
1134 .collect::<Result<Vec<_>>>()?;
1135 let names = leaf
1136 .ctx
1137 .field_columns
1138 .iter()
1139 .map(|field| format!("{}.{}", leaf.display_table, field))
1140 .collect();
1141 Ok(IslandFieldExprs {
1142 exprs,
1143 names,
1144 scalar: false,
1145 })
1146 }
1147 IslandExpr::Scalar(expr) => Ok(IslandFieldExprs {
1148 exprs: vec![expr.clone()],
1149 names: vec![expr.schema_name().to_string()],
1150 scalar: true,
1151 }),
1152 IslandExpr::Unary { input } => {
1153 let input = Self::build_binary_island_field_exprs(input, leaves, schema)?;
1154 let mut exprs = Vec::with_capacity(input.exprs.len());
1155 let mut names = Vec::with_capacity(input.names.len());
1156 for (expr, name) in input.exprs.into_iter().zip(input.names) {
1157 exprs.push(DfExpr::Negative(Box::new(expr)));
1158 names.push(format!("-{name}"));
1159 }
1160 Ok(IslandFieldExprs {
1161 exprs,
1162 names,
1163 scalar: input.scalar,
1164 })
1165 }
1166 IslandExpr::Binary { op, lhs, rhs } => {
1167 let same_leaf = match (&**lhs, &**rhs) {
1168 (IslandExpr::VectorLeaf(left), IslandExpr::VectorLeaf(right))
1169 if left == right =>
1170 {
1171 Some(*left)
1172 }
1173 _ => None,
1174 };
1175 let lhs = Self::build_binary_island_field_exprs(lhs, leaves, schema)?;
1176 let rhs = Self::build_binary_island_field_exprs(rhs, leaves, schema)?;
1177 let expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1178 let scalar = lhs.scalar && rhs.scalar;
1179 let op = op.to_string();
1180
1181 let (exprs, names) = match (lhs.scalar, rhs.scalar) {
1182 (true, true) => {
1183 let expr = expr_builder(lhs.exprs[0].clone(), rhs.exprs[0].clone())?;
1184 let name = format!("{} {op} {}", lhs.names[0], rhs.names[0]);
1185 (vec![expr], vec![name])
1186 }
1187 (true, false) => {
1188 let mut exprs = Vec::with_capacity(rhs.exprs.len());
1189 let mut names = Vec::with_capacity(rhs.names.len());
1190 for (rhs_expr, rhs_name) in rhs.exprs.into_iter().zip(rhs.names) {
1191 exprs.push(expr_builder(lhs.exprs[0].clone(), rhs_expr)?);
1192 names.push(format!("{} {op} {rhs_name}", lhs.names[0]));
1193 }
1194 (exprs, names)
1195 }
1196 (false, true) => {
1197 let mut exprs = Vec::with_capacity(lhs.exprs.len());
1198 let mut names = Vec::with_capacity(lhs.names.len());
1199 for (lhs_expr, lhs_name) in lhs.exprs.into_iter().zip(lhs.names) {
1200 exprs.push(expr_builder(lhs_expr, rhs.exprs[0].clone())?);
1201 names.push(format!("{lhs_name} {op} {}", rhs.names[0]));
1202 }
1203 (exprs, names)
1204 }
1205 (false, false) => {
1206 let mut exprs = Vec::new();
1207 let mut names = Vec::new();
1208 for (idx, ((lhs_expr, rhs_expr), (mut lhs_name, mut rhs_name))) in lhs
1209 .exprs
1210 .into_iter()
1211 .zip(rhs.exprs)
1212 .zip(lhs.names.into_iter().zip(rhs.names))
1213 .enumerate()
1214 {
1215 if let Some(leaf) = same_leaf {
1216 let field = leaves[leaf]
1217 .ctx
1218 .field_columns
1219 .get(idx)
1220 .cloned()
1221 .unwrap_or_else(|| lhs_name.clone());
1222 lhs_name = format!("lhs.{field}");
1223 rhs_name = format!("rhs.{field}");
1224 }
1225 exprs.push(expr_builder(lhs_expr, rhs_expr)?);
1226 names.push(format!("{lhs_name} {op} {rhs_name}"));
1227 }
1228 (exprs, names)
1229 }
1230 };
1231
1232 Ok(IslandFieldExprs {
1233 exprs,
1234 names,
1235 scalar,
1236 })
1237 }
1238 }
1239 }
1240
1241 fn project_binary_island(
1242 &mut self,
1243 input: LogicalPlan,
1244 base_alias: &TableReference,
1245 base_ctx: &PromPlannerContext,
1246 field_exprs: IslandFieldExprs,
1247 ) -> Result<LogicalPlan> {
1248 self.ctx = base_ctx.clone();
1249
1250 let schema = input.schema();
1251 let non_field_exprs = base_ctx
1252 .tag_columns
1253 .iter()
1254 .chain(base_ctx.time_index_column.iter())
1255 .map(|column| {
1256 schema
1257 .qualified_field_with_name(Some(base_alias), column)
1258 .context(DataFusionPlanningSnafu)
1259 .map(|field| DfExpr::Column(field.into()))
1260 });
1261 let tsid_expr = Self::optional_tsid_projection(schema, Some(base_alias), base_ctx.use_tsid)
1262 .into_iter()
1263 .map(Ok);
1264
1265 self.ctx.field_columns = field_exprs.names;
1266 let field_exprs = field_exprs
1267 .exprs
1268 .into_iter()
1269 .zip(self.ctx.field_columns.iter())
1270 .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::<String>, name))));
1271
1272 let project_exprs = non_field_exprs
1273 .chain(tsid_expr)
1274 .chain(field_exprs)
1275 .collect::<Result<Vec<_>>>()?;
1276
1277 let plan = LogicalPlanBuilder::from(input)
1278 .project(project_exprs)
1279 .context(DataFusionPlanningSnafu)?
1280 .build()
1281 .context(DataFusionPlanningSnafu)?;
1282
1283 self.ctx.table_name = None;
1284 self.ctx.schema_name = None;
1285
1286 Ok(plan)
1287 }
1288
1289 async fn prom_binary_expr_to_plan(
1290 &mut self,
1291 query_engine_state: &QueryEngineState,
1292 binary_expr: &PromBinaryExpr,
1293 ) -> Result<LogicalPlan> {
1294 if let Some(modifier) = &binary_expr.modifier {
1298 ensure!(
1299 modifier.fill_values.lhs.is_none() && modifier.fill_values.rhs.is_none(),
1300 UnsupportedExprSnafu {
1301 name: "PromQL fill modifiers"
1302 }
1303 );
1304 }
1305
1306 if let Some(plan) = self.try_plan_binary_island(binary_expr).await? {
1307 return Ok(plan);
1308 }
1309
1310 let PromBinaryExpr {
1311 lhs,
1312 rhs,
1313 op,
1314 modifier,
1315 } = binary_expr;
1316
1317 let should_return_bool = if let Some(m) = modifier {
1320 m.return_bool
1321 } else {
1322 false
1323 };
1324 let is_comparison_op = Self::is_token_a_comparison_op(*op);
1325
1326 match (
1329 Self::try_build_literal_expr(lhs),
1330 Self::try_build_literal_expr(rhs),
1331 ) {
1332 (Some(lhs), Some(rhs)) => {
1333 self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
1334 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1335 self.ctx.reset_table_name_and_schema();
1336 let field_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1337 let mut field_expr = field_expr_builder(lhs, rhs)?;
1338
1339 if is_comparison_op && should_return_bool {
1340 field_expr = DfExpr::Cast(Cast {
1341 expr: Box::new(field_expr),
1342 data_type: ArrowDataType::Float64,
1343 });
1344 }
1345
1346 Ok(LogicalPlan::Extension(Extension {
1347 node: Arc::new(
1348 EmptyMetric::new(
1349 self.ctx.start,
1350 self.ctx.end,
1351 self.ctx.interval,
1352 SPECIAL_TIME_FUNCTION.to_string(),
1353 DEFAULT_FIELD_COLUMN.to_string(),
1354 Some(field_expr),
1355 )
1356 .context(DataFusionPlanningSnafu)?,
1357 ),
1358 }))
1359 }
1360 (Some(mut expr), None) => {
1362 let input = self.prom_expr_to_plan(rhs, query_engine_state).await?;
1363 if let Some(time_expr) = self.try_build_special_time_expr_with_context(lhs) {
1365 expr = time_expr
1366 }
1367 let input_schema = input.schema().clone();
1368 let preserve_any_value = Self::field_columns_are_alternative_samples(
1369 &input_schema,
1370 &self.ctx.field_columns,
1371 );
1372 let has_native_histogram = Self::field_columns_contain_native_histogram(
1373 &input_schema,
1374 &self.ctx.field_columns,
1375 );
1376 let retain_field_columns = self
1377 .ctx
1378 .field_columns
1379 .iter()
1380 .map(|col| {
1381 Self::binary_result_is_histogram(
1382 *op,
1383 false,
1384 Self::field_column_is_native_histogram(&input_schema, col),
1385 )
1386 .is_some()
1387 })
1388 .collect();
1389 let promql_annotations = self.promql_annotations.clone();
1390 let bin_expr_builder = |col: &String| {
1391 let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1392 let rhs_is_histogram =
1393 Self::field_column_is_native_histogram(&input_schema, col);
1394 let rhs = DfExpr::Column(col.into());
1395 let mut binary_expr = match Self::native_histogram_binary_expr(
1396 *op,
1397 expr.clone(),
1398 false,
1399 rhs.clone(),
1400 rhs_is_histogram,
1401 is_comparison_op && !should_return_bool,
1402 promql_annotations.clone(),
1403 )? {
1404 Some(expr) => expr,
1405 None => binary_expr_builder(expr.clone(), rhs)?,
1406 };
1407
1408 if is_comparison_op && should_return_bool {
1409 binary_expr = DfExpr::Cast(Cast {
1410 expr: Box::new(binary_expr),
1411 data_type: ArrowDataType::Float64,
1412 });
1413 }
1414 Ok(binary_expr)
1415 };
1416 if is_comparison_op && !should_return_bool {
1417 self.filter_on_field_column(input, bin_expr_builder)
1418 } else {
1419 let projected =
1420 self.projection_for_each_field_column(input, bin_expr_builder)?;
1421 self.filter_binary_projection(
1422 projected,
1423 has_native_histogram,
1424 preserve_any_value,
1425 retain_field_columns,
1426 )
1427 }
1428 }
1429 (None, Some(mut expr)) => {
1431 let input = self.prom_expr_to_plan(lhs, query_engine_state).await?;
1432 if let Some(time_expr) = self.try_build_special_time_expr_with_context(rhs) {
1434 expr = time_expr
1435 }
1436 let input_schema = input.schema().clone();
1437 let preserve_any_value = Self::field_columns_are_alternative_samples(
1438 &input_schema,
1439 &self.ctx.field_columns,
1440 );
1441 let has_native_histogram = Self::field_columns_contain_native_histogram(
1442 &input_schema,
1443 &self.ctx.field_columns,
1444 );
1445 let retain_field_columns = self
1446 .ctx
1447 .field_columns
1448 .iter()
1449 .map(|col| {
1450 Self::binary_result_is_histogram(
1451 *op,
1452 Self::field_column_is_native_histogram(&input_schema, col),
1453 false,
1454 )
1455 .is_some()
1456 })
1457 .collect();
1458 let promql_annotations = self.promql_annotations.clone();
1459 let bin_expr_builder = |col: &String| {
1460 let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1461 let lhs_is_histogram =
1462 Self::field_column_is_native_histogram(&input_schema, col);
1463 let lhs = DfExpr::Column(col.into());
1464 let mut binary_expr = match Self::native_histogram_binary_expr(
1465 *op,
1466 lhs.clone(),
1467 lhs_is_histogram,
1468 expr.clone(),
1469 false,
1470 is_comparison_op && !should_return_bool,
1471 promql_annotations.clone(),
1472 )? {
1473 Some(expr) => expr,
1474 None => binary_expr_builder(lhs, expr.clone())?,
1475 };
1476
1477 if is_comparison_op && should_return_bool {
1478 binary_expr = DfExpr::Cast(Cast {
1479 expr: Box::new(binary_expr),
1480 data_type: ArrowDataType::Float64,
1481 });
1482 }
1483 Ok(binary_expr)
1484 };
1485 if is_comparison_op && !should_return_bool {
1486 self.filter_on_field_column(input, bin_expr_builder)
1487 } else {
1488 let projected =
1489 self.projection_for_each_field_column(input, bin_expr_builder)?;
1490 self.filter_binary_projection(
1491 projected,
1492 has_native_histogram,
1493 preserve_any_value,
1494 retain_field_columns,
1495 )
1496 }
1497 }
1498 (None, None) => {
1500 let left_input = self.prom_expr_to_plan(lhs, query_engine_state).await?;
1501 let left_field_columns = self.ctx.field_columns.clone();
1502 let left_time_index_column = self.ctx.time_index_column.clone();
1503 let mut left_table_ref = self
1504 .table_ref()
1505 .unwrap_or_else(|_| TableReference::bare(""));
1506 let left_context = self.ctx.clone();
1507
1508 let right_input = self.prom_expr_to_plan(rhs, query_engine_state).await?;
1509 let right_field_columns = self.ctx.field_columns.clone();
1510 let right_time_index_column = self.ctx.time_index_column.clone();
1511 let mut right_table_ref = self
1512 .table_ref()
1513 .unwrap_or_else(|_| TableReference::bare(""));
1514 let right_context = self.ctx.clone();
1515 let left_is_empty_metric = Self::is_empty_metric(&left_input);
1516 let right_is_empty_metric = Self::is_empty_metric(&right_input);
1517
1518 if Self::is_token_a_set_op(*op) {
1522 return self.set_op_on_non_field_columns(
1523 left_input,
1524 right_input,
1525 left_context,
1526 right_context,
1527 *op,
1528 modifier,
1529 );
1530 }
1531
1532 let has_native_histogram = Self::field_columns_contain_native_histogram(
1533 left_input.schema(),
1534 &left_field_columns,
1535 ) || Self::field_columns_contain_native_histogram(
1536 right_input.schema(),
1537 &right_field_columns,
1538 );
1539
1540 if left_table_ref == right_table_ref {
1542 left_table_ref = TableReference::bare("lhs");
1544 right_table_ref = TableReference::bare("rhs");
1545 if self.ctx.tag_columns.is_empty() {
1551 self.ctx = left_context.clone();
1552 self.ctx.table_name = Some("lhs".to_string());
1553 } else {
1554 self.ctx.table_name = Some("rhs".to_string());
1555 }
1556 } else if right_is_empty_metric && !left_is_empty_metric {
1557 self.ctx = left_context.clone();
1558 }
1559 let broadcast_scalar = !is_comparison_op;
1562 let (field_groups, invalid_field_pairs) = Self::align_binary_field_columns(
1563 left_input.schema(),
1564 right_input.schema(),
1565 &left_field_columns,
1566 &right_field_columns,
1567 *op,
1568 broadcast_scalar && lhs.value_type() == ValueType::Scalar,
1569 broadcast_scalar && rhs.value_type() == ValueType::Scalar,
1570 );
1571 let left_aligned_field_columns = field_groups
1572 .iter()
1573 .flat_map(|(_, pairs)| {
1574 pairs
1575 .iter()
1576 .map(|(left_col_name, _)| (*left_col_name).clone())
1577 })
1578 .collect::<Vec<_>>();
1579 let right_aligned_field_columns = field_groups
1580 .iter()
1581 .flat_map(|(_, pairs)| {
1582 pairs
1583 .iter()
1584 .map(|(_, right_col_name)| (*right_col_name).clone())
1585 })
1586 .collect::<Vec<_>>();
1587 self.ctx.field_columns = field_groups
1590 .iter()
1591 .map(|(output, _)| output.clone())
1592 .collect();
1593 let mut field_groups = field_groups.into_iter();
1594 let has_empty_metric_operand = left_is_empty_metric || right_is_empty_metric;
1596
1597 let join_plan = self.join_on_non_field_columns(
1598 left_input,
1599 right_input,
1600 left_table_ref.clone(),
1601 right_table_ref.clone(),
1602 left_time_index_column,
1603 right_time_index_column,
1604 lhs.value_type() == ValueType::Scalar
1605 || rhs.value_type() == ValueType::Scalar
1606 || has_empty_metric_operand
1607 || ((left_context.tag_columns.is_empty()
1608 || right_context.tag_columns.is_empty())
1609 && !left_context
1610 .tag_columns
1611 .iter()
1612 .chain(&right_context.tag_columns)
1613 .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL)),
1614 modifier,
1615 &left_context,
1616 &right_context,
1617 )?;
1618 let join_plan_schema = join_plan.schema().clone();
1619 let promql_annotations = self.promql_annotations.clone();
1620 let invalid_pair_predicates = invalid_field_pairs
1623 .into_iter()
1624 .filter(|_| promql_annotations.is_some())
1625 .map(|(left_col_name, right_col_name)| {
1626 let left_field = join_plan_schema
1627 .qualified_field_with_name(Some(&left_table_ref), left_col_name)
1628 .context(DataFusionPlanningSnafu)?;
1629 let right_field = join_plan_schema
1630 .qualified_field_with_name(Some(&right_table_ref), right_col_name)
1631 .context(DataFusionPlanningSnafu)?;
1632 let left_is_histogram =
1633 left_field.1.data_type() == &Self::native_histogram_arrow_type();
1634 let right_is_histogram =
1635 right_field.1.data_type() == &Self::native_histogram_arrow_type();
1636 let drop_expr = Self::native_histogram_binary_expr(
1637 *op,
1638 DfExpr::Column(left_field.into()),
1639 left_is_histogram,
1640 DfExpr::Column(right_field.into()),
1641 right_is_histogram,
1642 true,
1643 promql_annotations.clone(),
1644 )?
1645 .with_context(|| UnexpectedPlanExprSnafu {
1646 desc: "invalid native histogram pair produced no drop expression",
1647 })?;
1648 Ok(DfExpr::Not(Box::new(drop_expr)))
1649 })
1650 .collect::<Result<Vec<_>>>()?;
1651 let join_plan = if let Some(predicate) = conjunction(invalid_pair_predicates) {
1652 LogicalPlanBuilder::from(join_plan)
1653 .filter(predicate)
1654 .context(DataFusionPlanningSnafu)?
1655 .build()
1656 .context(DataFusionPlanningSnafu)?
1657 } else {
1658 join_plan
1659 };
1660
1661 let bin_expr_builder = |_: &String| {
1662 let (_, field_pairs) =
1663 field_groups
1664 .next()
1665 .with_context(|| UnexpectedPlanExprSnafu {
1666 desc: "missing binary field group",
1667 })?;
1668 let binary_exprs = field_pairs
1669 .into_iter()
1670 .map(|(left_col_name, right_col_name)| {
1671 let left_field = join_plan_schema
1672 .qualified_field_with_name(Some(&left_table_ref), left_col_name)
1673 .context(DataFusionPlanningSnafu)?;
1674 let right_field = join_plan_schema
1675 .qualified_field_with_name(Some(&right_table_ref), right_col_name)
1676 .context(DataFusionPlanningSnafu)?;
1677 let left_is_histogram =
1678 left_field.1.data_type() == &Self::native_histogram_arrow_type();
1679 let right_is_histogram =
1680 right_field.1.data_type() == &Self::native_histogram_arrow_type();
1681 let left_col = left_field.into();
1682 let right_col = right_field.into();
1683
1684 let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1685 let lhs = DfExpr::Column(left_col);
1686 let rhs = DfExpr::Column(right_col);
1687 let mut binary_expr = match Self::native_histogram_binary_expr(
1688 *op,
1689 lhs.clone(),
1690 left_is_histogram,
1691 rhs.clone(),
1692 right_is_histogram,
1693 is_comparison_op && !should_return_bool,
1694 promql_annotations.clone(),
1695 )? {
1696 Some(expr) => expr,
1697 None => binary_expr_builder(lhs, rhs)?,
1698 };
1699 if is_comparison_op && should_return_bool {
1700 binary_expr = DfExpr::Cast(Cast {
1701 expr: Box::new(binary_expr),
1702 data_type: ArrowDataType::Float64,
1703 });
1704 }
1705 Ok(binary_expr)
1706 })
1707 .collect::<Result<Vec<_>>>()?;
1708 if let [binary_expr] = binary_exprs.as_slice() {
1709 Ok(binary_expr.clone())
1710 } else {
1711 Ok(DfExpr::ScalarFunction(ScalarFunction {
1712 func: coalesce(),
1713 args: binary_exprs,
1714 }))
1715 }
1716 };
1717 if is_comparison_op && !should_return_bool {
1718 let filtered = self.filter_on_field_column(join_plan, bin_expr_builder)?;
1725 let (project_table_ref, mut project_context, project_field_columns) =
1726 match (lhs.value_type(), rhs.value_type()) {
1727 (ValueType::Scalar, ValueType::Vector) => (
1728 &right_table_ref,
1729 right_context.clone(),
1730 right_aligned_field_columns,
1731 ),
1732 _ => (
1733 &left_table_ref,
1734 left_context.clone(),
1735 left_aligned_field_columns,
1736 ),
1737 };
1738 project_context.field_columns = project_field_columns;
1739 self.project_binary_join_side(filtered, project_table_ref, &project_context)
1740 } else {
1741 let projected =
1742 self.projection_for_each_field_column(join_plan, bin_expr_builder)?;
1743 let preserve_any_value = Self::field_columns_are_alternative_samples(
1744 projected.schema(),
1745 &self.ctx.field_columns,
1746 );
1747 let retain_field_columns = vec![true; self.ctx.field_columns.len()];
1748 self.filter_binary_projection(
1749 projected,
1750 has_native_histogram,
1751 preserve_any_value,
1752 retain_field_columns,
1753 )
1754 }
1755 }
1756 }
1757 }
1758
1759 fn filter_binary_projection(
1760 &mut self,
1761 input: LogicalPlan,
1762 has_native_histogram: bool,
1763 preserve_any_value: bool,
1764 retain_field_columns: Vec<bool>,
1765 ) -> Result<LogicalPlan> {
1766 if !has_native_histogram {
1767 return Ok(input);
1768 }
1769
1770 ensure!(
1771 retain_field_columns.len() == self.ctx.field_columns.len(),
1772 UnexpectedPlanExprSnafu {
1773 desc: "binary output field count changed unexpectedly",
1774 }
1775 );
1776
1777 let filtered = LogicalPlanBuilder::from(input)
1778 .filter(self.create_empty_values_filter_expr(preserve_any_value)?)
1779 .context(DataFusionPlanningSnafu)?
1780 .build()
1781 .context(DataFusionPlanningSnafu)?;
1782 if retain_field_columns.iter().all(|retain| *retain) {
1783 return Ok(filtered);
1784 }
1785
1786 let retained = self
1787 .ctx
1788 .field_columns
1789 .iter()
1790 .zip(retain_field_columns)
1791 .filter(|(_, retain)| *retain)
1792 .map(|(field, _)| field.clone())
1793 .collect::<Vec<_>>();
1794 if retained.is_empty() {
1795 return Ok(filtered);
1796 }
1797 self.ctx.field_columns = retained;
1798
1799 let mut output_columns = self
1800 .ctx
1801 .field_columns
1802 .iter()
1803 .chain(&self.ctx.tag_columns)
1804 .cloned()
1805 .collect::<HashSet<_>>();
1806 output_columns.extend(self.ctx.time_index_column.iter().cloned());
1807 if self.ctx.use_tsid {
1808 output_columns.insert(DATA_SCHEMA_TSID_COLUMN_NAME.to_string());
1809 }
1810 let project_exprs = filtered
1811 .schema()
1812 .iter()
1813 .filter(|(_, field)| output_columns.contains(field.name()))
1814 .map(|(qualifier, field)| {
1815 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
1816 })
1817 .collect::<Vec<_>>();
1818 LogicalPlanBuilder::from(filtered)
1819 .project(project_exprs)
1820 .context(DataFusionPlanningSnafu)?
1821 .build()
1822 .context(DataFusionPlanningSnafu)
1823 }
1824
1825 fn project_binary_join_side(
1826 &mut self,
1827 input: LogicalPlan,
1828 table_ref: &TableReference,
1829 context: &PromPlannerContext,
1830 ) -> Result<LogicalPlan> {
1831 let schema = input.schema();
1832
1833 let mut project_exprs =
1834 Vec::with_capacity(context.tag_columns.len() + context.field_columns.len() + 2);
1835
1836 if let Some(time_index_column) = &context.time_index_column {
1838 let time_index_col = schema
1839 .qualified_field_with_name(Some(table_ref), time_index_column)
1840 .context(DataFusionPlanningSnafu)?
1841 .into();
1842 project_exprs.push(DfExpr::Column(time_index_col));
1843 }
1844
1845 for field_column in &context.field_columns {
1847 let field_col = schema
1848 .qualified_field_with_name(Some(table_ref), field_column)
1849 .context(DataFusionPlanningSnafu)?
1850 .into();
1851 project_exprs.push(DfExpr::Column(field_col));
1852 }
1853
1854 for tag_column in &context.tag_columns {
1856 let tag_col = schema
1857 .qualified_field_with_name(Some(table_ref), tag_column)
1858 .context(DataFusionPlanningSnafu)?
1859 .into();
1860 project_exprs.push(DfExpr::Column(tag_col));
1861 }
1862
1863 if let Some(tsid_col) =
1866 Self::optional_tsid_projection(schema, Some(table_ref), context.use_tsid)
1867 {
1868 project_exprs.push(tsid_col);
1869 }
1870
1871 let plan = LogicalPlanBuilder::from(input)
1872 .project(project_exprs)
1873 .context(DataFusionPlanningSnafu)?
1874 .build()
1875 .context(DataFusionPlanningSnafu)?;
1876
1877 self.ctx = context.clone();
1880 self.ctx.table_name = None;
1881 self.ctx.schema_name = None;
1882
1883 Ok(plan)
1884 }
1885
1886 fn prom_number_lit_to_plan(&mut self, number_literal: &NumberLiteral) -> Result<LogicalPlan> {
1887 let NumberLiteral { val } = number_literal;
1888 self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
1889 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1890 self.ctx.reset_table_name_and_schema();
1891 let literal_expr = df_prelude::lit(*val);
1892
1893 let plan = LogicalPlan::Extension(Extension {
1894 node: Arc::new(
1895 EmptyMetric::new(
1896 self.ctx.start,
1897 self.ctx.end,
1898 self.ctx.interval,
1899 SPECIAL_TIME_FUNCTION.to_string(),
1900 DEFAULT_FIELD_COLUMN.to_string(),
1901 Some(literal_expr),
1902 )
1903 .context(DataFusionPlanningSnafu)?,
1904 ),
1905 });
1906 Ok(plan)
1907 }
1908
1909 fn prom_string_lit_to_plan(&mut self, string_literal: &StringLiteral) -> Result<LogicalPlan> {
1910 let StringLiteral { val } = string_literal;
1911 self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
1912 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1913 self.ctx.reset_table_name_and_schema();
1914 let literal_expr = df_prelude::lit(val.clone());
1915
1916 let plan = LogicalPlan::Extension(Extension {
1917 node: Arc::new(
1918 EmptyMetric::new(
1919 self.ctx.start,
1920 self.ctx.end,
1921 self.ctx.interval,
1922 SPECIAL_TIME_FUNCTION.to_string(),
1923 DEFAULT_FIELD_COLUMN.to_string(),
1924 Some(literal_expr),
1925 )
1926 .context(DataFusionPlanningSnafu)?,
1927 ),
1928 });
1929 Ok(plan)
1930 }
1931
1932 async fn prom_vector_selector_to_plan(
1933 &mut self,
1934 vector_selector: &VectorSelector,
1935 timestamp_fn: bool,
1936 ) -> Result<LogicalPlan> {
1937 let VectorSelector {
1938 name,
1939 offset,
1940 matchers,
1941 at: _,
1942 } = vector_selector;
1943 let matchers = self.preprocess_label_matchers(matchers, name)?;
1944 if let Some(empty_plan) = self.setup_context().await? {
1945 return Ok(empty_plan);
1946 }
1947 let offset_ms = match offset {
1948 Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond,
1949 Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond),
1950 None => 0,
1951 };
1952 let normalize = self
1953 .selector_to_series_normalize_plan(offset, matchers, false)
1954 .await?;
1955 let time_index_column =
1956 self.ctx
1957 .time_index_column
1958 .clone()
1959 .with_context(|| TimeIndexNotFoundSnafu {
1960 table: self.ctx.table_name.clone().unwrap_or_default(),
1961 })?;
1962
1963 let (normalize, timestamp_value_column) = if timestamp_fn {
1964 let occupied = normalize
1967 .schema()
1968 .fields()
1969 .iter()
1970 .map(|field| field.name().as_str())
1971 .collect::<HashSet<_>>();
1972 let mut timestamp_value_column = TIMESTAMP_VALUE_PREFIX.to_string();
1973 while occupied.contains(timestamp_value_column.as_str()) {
1974 timestamp_value_column.push('_');
1975 }
1976 let mut project_exprs = normalize
1977 .schema()
1978 .iter()
1979 .map(|(qualifier, field)| {
1980 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
1981 })
1982 .collect::<Vec<_>>();
1983 let unit_factor = match col(&time_index_column)
1987 .get_type(normalize.schema())
1988 .context(DataFusionPlanningSnafu)?
1989 {
1990 ArrowDataType::Timestamp(ArrowTimeUnit::Second, _) => (1_000_i128, 4, 0),
1991 ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, _) => (1, 1, 0),
1992 ArrowDataType::Timestamp(ArrowTimeUnit::Microsecond, _) => (1, 4, 3),
1993 ArrowDataType::Timestamp(ArrowTimeUnit::Nanosecond, _) => (1, 7, 6),
1994 _ => unreachable!("time index is a timestamp"),
1995 };
1996 let sample_time = col(&time_index_column)
1997 .cast_to(&ArrowDataType::Int64, normalize.schema())
1998 .context(DataFusionPlanningSnafu)?
1999 .cast_to(&ArrowDataType::Decimal128(19, 0), normalize.schema())
2000 .context(DataFusionPlanningSnafu)?;
2001 let sample_time = DfExpr::BinaryExpr(BinaryExpr {
2002 left: Box::new(sample_time),
2003 op: Operator::Multiply,
2004 right: Box::new(lit(ScalarValue::Decimal128(
2005 Some(unit_factor.0),
2006 unit_factor.1,
2007 unit_factor.2,
2008 ))),
2009 });
2010 let sample_time = DfExpr::BinaryExpr(BinaryExpr {
2011 left: Box::new(sample_time),
2012 op: Operator::Plus,
2013 right: Box::new(lit(ScalarValue::Decimal128(Some(offset_ms as i128), 19, 0))),
2014 })
2015 .cast_to(&ArrowDataType::Int64, normalize.schema())
2016 .context(DataFusionPlanningSnafu)?
2017 .cast_to(&ArrowDataType::Float64, normalize.schema())
2018 .context(DataFusionPlanningSnafu)?;
2019 let sample_time = DfExpr::BinaryExpr(BinaryExpr {
2020 left: Box::new(sample_time),
2021 op: Operator::Divide,
2022 right: Box::new(lit(1000.0)),
2023 });
2024 project_exprs.push(sample_time.alias(×tamp_value_column));
2025 let normalize = LogicalPlanBuilder::from(normalize)
2026 .project(project_exprs)
2027 .context(DataFusionPlanningSnafu)?
2028 .build()
2029 .context(DataFusionPlanningSnafu)?;
2030 (normalize, Some(timestamp_value_column))
2031 } else {
2032 (normalize, None)
2033 };
2034
2035 let field_column = self.ctx.field_columns.first().cloned();
2036 let manipulate = InstantManipulate::new(
2037 self.ctx.start,
2038 self.ctx.end,
2039 self.ctx.lookback_delta,
2040 self.ctx.interval,
2041 offset_ms,
2042 time_index_column,
2043 if self.ctx.use_tsid {
2044 vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]
2045 } else {
2046 self.ctx.tag_columns.clone()
2047 },
2048 field_column,
2049 normalize,
2050 );
2051 let manipulate = LogicalPlan::Extension(Extension {
2052 node: Arc::new(manipulate),
2053 });
2054 if let Some(timestamp_value_column) = timestamp_value_column {
2055 self.create_timestamp_func_plan(manipulate, ×tamp_value_column)
2056 } else {
2057 Ok(manipulate)
2058 }
2059 }
2060
2061 fn create_timestamp_func_plan(
2084 &mut self,
2085 input: LogicalPlan,
2086 timestamp_value_column: &str,
2087 ) -> Result<LogicalPlan> {
2088 let time_expr = col(timestamp_value_column).alias(DEFAULT_FIELD_COLUMN);
2089 self.ctx.field_columns = vec![time_expr.schema_name().to_string()];
2090 let mut project_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 2);
2091 project_exprs.push(self.create_time_index_column_expr()?);
2092 project_exprs.push(time_expr);
2093 project_exprs.extend(self.create_tag_column_exprs()?);
2094
2095 LogicalPlanBuilder::from(input)
2096 .project(project_exprs)
2097 .context(DataFusionPlanningSnafu)?
2098 .build()
2099 .context(DataFusionPlanningSnafu)
2100 }
2101
2102 async fn prom_matrix_selector_to_plan(
2103 &mut self,
2104 matrix_selector: &MatrixSelector,
2105 ) -> Result<LogicalPlan> {
2106 let MatrixSelector { vs, range } = matrix_selector;
2107 let VectorSelector {
2108 name,
2109 offset,
2110 matchers,
2111 ..
2112 } = vs;
2113 let matchers = self.preprocess_label_matchers(matchers, name)?;
2114 ensure!(!range.is_zero(), ZeroRangeSelectorSnafu);
2115 let range_ms = range.as_millis() as _;
2116 self.ctx.range = Some(range_ms);
2117 let offset_ms = match offset {
2118 Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond,
2119 Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond),
2120 None => 0,
2121 };
2122
2123 let normalize = match self.setup_context().await? {
2126 Some(empty_plan) => empty_plan,
2127 None => {
2128 self.selector_to_series_normalize_plan(offset, matchers, true)
2129 .await?
2130 }
2131 };
2132 let manipulate = RangeManipulate::new(
2133 self.ctx.start,
2134 self.ctx.end,
2135 self.ctx.interval,
2136 offset_ms,
2137 range_ms,
2139 self.ctx
2140 .time_index_column
2141 .clone()
2142 .expect("time index should be set in `setup_context`"),
2143 self.ctx.field_columns.clone(),
2144 normalize,
2145 )
2146 .context(DataFusionPlanningSnafu)?;
2147
2148 Ok(LogicalPlan::Extension(Extension {
2149 node: Arc::new(manipulate),
2150 }))
2151 }
2152
2153 async fn prom_call_expr_to_plan(
2154 &mut self,
2155 query_engine_state: &QueryEngineState,
2156 call_expr: &Call,
2157 ) -> Result<LogicalPlan> {
2158 let Call { func, args } = call_expr;
2159 match func.name {
2161 SPECIAL_HISTOGRAM_QUANTILE | SPECIAL_HISTOGRAM_FRACTION => {
2162 return self
2163 .create_histogram_plan(func.name, args, query_engine_state)
2164 .await;
2165 }
2166 SPECIAL_VECTOR_FUNCTION => return self.create_vector_plan(args).await,
2167 SCALAR_FUNCTION => return self.create_scalar_plan(args, query_engine_state).await,
2168 SPECIAL_ABSENT_FUNCTION => {
2169 return self.create_absent_plan(args, query_engine_state).await;
2170 }
2171 _ => {}
2172 }
2173
2174 let args = self.create_function_args(&args.args)?;
2176 let input = if let Some(prom_expr) = &args.input {
2177 self.prom_expr_to_plan_inner(prom_expr, func.name == "timestamp", query_engine_state)
2178 .await?
2179 } else {
2180 self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
2181 self.ctx.reset_table_name_and_schema();
2182 self.ctx.tag_columns = vec![];
2183 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
2184 LogicalPlan::Extension(Extension {
2185 node: Arc::new(
2186 EmptyMetric::new(
2187 self.ctx.start,
2188 self.ctx.end,
2189 self.ctx.interval,
2190 SPECIAL_TIME_FUNCTION.to_string(),
2191 DEFAULT_FIELD_COLUMN.to_string(),
2192 None,
2193 )
2194 .context(DataFusionPlanningSnafu)?,
2195 ),
2196 })
2197 };
2198 let (mut func_exprs, new_tags) = self.create_function_expr(
2199 func,
2200 args.literals.clone(),
2201 input.schema(),
2202 query_engine_state,
2203 )?;
2204 func_exprs.insert(0, self.create_time_index_column_expr()?);
2205 func_exprs.extend_from_slice(&self.create_tag_column_exprs()?);
2206 if let Some(tsid_col) =
2207 Self::optional_tsid_projection(input.schema(), None, self.ctx.use_tsid)
2208 {
2209 func_exprs.push(tsid_col);
2210 }
2211
2212 let builder = LogicalPlanBuilder::from(input)
2218 .project(func_exprs)
2219 .context(DataFusionPlanningSnafu)?
2220 .filter(self.create_empty_values_filter_expr(true)?)
2221 .context(DataFusionPlanningSnafu)?;
2222
2223 let builder = match func.name {
2224 "sort" => builder
2225 .sort(self.create_field_columns_sort_exprs(true))
2226 .context(DataFusionPlanningSnafu)?,
2227 "sort_desc" => builder
2228 .sort(self.create_field_columns_sort_exprs(false))
2229 .context(DataFusionPlanningSnafu)?,
2230 "sort_by_label" => builder
2231 .sort(Self::create_sort_exprs_by_tags(
2232 func.name,
2233 args.literals,
2234 true,
2235 )?)
2236 .context(DataFusionPlanningSnafu)?,
2237 "sort_by_label_desc" => builder
2238 .sort(Self::create_sort_exprs_by_tags(
2239 func.name,
2240 args.literals,
2241 false,
2242 )?)
2243 .context(DataFusionPlanningSnafu)?,
2244
2245 _ => builder,
2246 };
2247
2248 for tag in new_tags {
2251 self.ctx.tag_columns.push(tag);
2252 }
2253
2254 let plan = builder.build().context(DataFusionPlanningSnafu)?;
2255 common_telemetry::debug!("Created PromQL function plan: {plan:?} for {call_expr:?}");
2256
2257 Ok(plan)
2258 }
2259
2260 async fn prom_ext_expr_to_plan(
2261 &mut self,
2262 query_engine_state: &QueryEngineState,
2263 ext_expr: &promql_parser::parser::ast::Extension,
2264 ) -> Result<LogicalPlan> {
2265 let expr = &ext_expr.expr;
2267 let children = expr.children();
2268 let plan = self
2269 .prom_expr_to_plan(&children[0], query_engine_state)
2270 .await?;
2271 match expr.name() {
2277 ANALYZE_NODE_NAME => LogicalPlanBuilder::from(plan)
2278 .explain(false, true)
2279 .unwrap()
2280 .build()
2281 .context(DataFusionPlanningSnafu),
2282 ANALYZE_VERBOSE_NODE_NAME => LogicalPlanBuilder::from(plan)
2283 .explain(true, true)
2284 .unwrap()
2285 .build()
2286 .context(DataFusionPlanningSnafu),
2287 EXPLAIN_NODE_NAME => LogicalPlanBuilder::from(plan)
2288 .explain(false, false)
2289 .unwrap()
2290 .build()
2291 .context(DataFusionPlanningSnafu),
2292 EXPLAIN_VERBOSE_NODE_NAME => LogicalPlanBuilder::from(plan)
2293 .explain(true, false)
2294 .unwrap()
2295 .build()
2296 .context(DataFusionPlanningSnafu),
2297 ALIAS_NODE_NAME => {
2298 let alias = expr
2299 .as_any()
2300 .downcast_ref::<AliasExpr>()
2301 .context(UnexpectedPlanExprSnafu {
2302 desc: "Expected AliasExpr",
2303 })?
2304 .alias
2305 .clone();
2306 self.apply_alias(plan, alias)
2307 }
2308 _ => LogicalPlanBuilder::empty(true)
2309 .build()
2310 .context(DataFusionPlanningSnafu),
2311 }
2312 }
2313
2314 #[allow(clippy::mutable_key_type)]
2324 fn preprocess_label_matchers(
2325 &mut self,
2326 label_matchers: &Matchers,
2327 name: &Option<String>,
2328 ) -> Result<Matchers> {
2329 self.ctx.reset();
2330
2331 let metric_name;
2332 if let Some(name) = name.clone() {
2333 metric_name = Some(name);
2334 ensure!(
2335 label_matchers.find_matchers(METRIC_NAME).is_empty(),
2336 MultipleMetricMatchersSnafu
2337 );
2338 } else {
2339 let mut matches = label_matchers.find_matchers(METRIC_NAME);
2340 ensure!(!matches.is_empty(), NoMetricMatcherSnafu);
2341 ensure!(matches.len() == 1, MultipleMetricMatchersSnafu);
2342 ensure!(
2343 matches[0].op == MatchOp::Equal,
2344 UnsupportedMatcherOpSnafu {
2345 matcher_op: matches[0].op.to_string(),
2346 matcher: METRIC_NAME
2347 }
2348 );
2349 metric_name = matches.pop().map(|m| m.value);
2350 }
2351
2352 self.ctx.table_name = metric_name;
2353
2354 let mut matchers = HashSet::new();
2355 for matcher in &label_matchers.matchers {
2356 if matcher.name == FIELD_COLUMN_MATCHER {
2358 self.ctx
2359 .field_column_matcher
2360 .get_or_insert_default()
2361 .push(matcher.clone());
2362 } else if matcher.name == SCHEMA_COLUMN_MATCHER || matcher.name == DB_COLUMN_MATCHER {
2363 ensure!(
2364 matcher.op == MatchOp::Equal,
2365 UnsupportedMatcherOpSnafu {
2366 matcher: matcher.name.clone(),
2367 matcher_op: matcher.op.to_string(),
2368 }
2369 );
2370 self.ctx.schema_name = Some(matcher.value.clone());
2371 } else if matcher.name != METRIC_NAME {
2372 self.ctx.selector_matcher.push(matcher.clone());
2373 let _ = matchers.insert(matcher.clone());
2374 }
2375 }
2376
2377 Ok(Matchers::new(matchers.into_iter().collect()))
2378 }
2379
2380 async fn selector_to_series_normalize_plan(
2381 &mut self,
2382 offset: &Option<Offset>,
2383 label_matchers: Matchers,
2384 is_range_selector: bool,
2385 ) -> Result<LogicalPlan> {
2386 let table_ref = self.table_ref()?;
2388 let mut table_scan = self.create_table_scan_plan(table_ref.clone()).await?;
2389 let table_schema = table_scan.schema();
2390
2391 let offset_duration = match offset {
2393 Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond,
2394 Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond),
2395 None => 0,
2396 };
2397 let mut scan_filters = Self::matchers_to_expr(label_matchers.clone(), table_schema)?;
2398 if let Some(time_index_filter) =
2399 self.build_time_index_filter(offset_duration, table_schema)?
2400 {
2401 scan_filters.push(time_index_filter);
2402 }
2403 if let Some(filter) = conjunction(scan_filters) {
2404 table_scan = LogicalPlanBuilder::from(table_scan)
2405 .filter(filter)
2406 .context(DataFusionPlanningSnafu)?
2407 .build()
2408 .context(DataFusionPlanningSnafu)?;
2409 }
2410
2411 if let Some(field_matchers) = &self.ctx.field_column_matcher {
2413 let col_set = self.ctx.field_columns.iter().collect::<HashSet<_>>();
2414 let mut result_set = HashSet::new();
2416 let mut reverse_set = HashSet::new();
2418 for matcher in field_matchers {
2419 match &matcher.op {
2420 MatchOp::Equal => {
2421 if col_set.contains(&matcher.value) {
2422 let _ = result_set.insert(matcher.value.clone());
2423 } else {
2424 return Err(ColumnNotFoundSnafu {
2425 col: matcher.value.clone(),
2426 }
2427 .build());
2428 }
2429 }
2430 MatchOp::NotEqual => {
2431 if col_set.contains(&matcher.value) {
2432 let _ = reverse_set.insert(matcher.value.clone());
2433 } else {
2434 return Err(ColumnNotFoundSnafu {
2435 col: matcher.value.clone(),
2436 }
2437 .build());
2438 }
2439 }
2440 MatchOp::Re(regex) => {
2441 for col in &self.ctx.field_columns {
2442 if regex.is_match(col) {
2443 let _ = result_set.insert(col.clone());
2444 }
2445 }
2446 }
2447 MatchOp::NotRe(regex) => {
2448 for col in &self.ctx.field_columns {
2449 if regex.is_match(col) {
2450 let _ = reverse_set.insert(col.clone());
2451 }
2452 }
2453 }
2454 }
2455 }
2456 if result_set.is_empty() {
2458 result_set = col_set.into_iter().cloned().collect();
2459 }
2460 for col in reverse_set {
2461 let _ = result_set.remove(&col);
2462 }
2463
2464 self.ctx.field_columns = self
2466 .ctx
2467 .field_columns
2468 .drain(..)
2469 .filter(|col| result_set.contains(col))
2470 .collect();
2471
2472 let exprs = result_set
2473 .into_iter()
2474 .map(|col| DfExpr::Column(Column::new_unqualified(col)))
2475 .chain(self.create_tag_column_exprs()?)
2476 .chain(
2477 self.ctx
2478 .use_tsid
2479 .then_some(DfExpr::Column(Column::new_unqualified(
2480 DATA_SCHEMA_TSID_COLUMN_NAME,
2481 ))),
2482 )
2483 .chain(Some(self.create_time_index_column_expr()?))
2484 .collect::<Vec<_>>();
2485
2486 table_scan = LogicalPlanBuilder::from(table_scan)
2488 .project(exprs)
2489 .context(DataFusionPlanningSnafu)?
2490 .build()
2491 .context(DataFusionPlanningSnafu)?;
2492 }
2493
2494 let series_key_columns = if self.ctx.use_tsid {
2496 vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]
2497 } else {
2498 self.ctx.tag_columns.clone()
2499 };
2500
2501 let sort_exprs = if self.ctx.use_tsid {
2502 vec![
2503 DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME)).sort(true, true),
2504 self.create_time_index_column_expr()?.sort(true, true),
2505 ]
2506 } else {
2507 self.create_tag_and_time_index_column_sort_exprs()?
2508 };
2509
2510 let sort_plan = LogicalPlanBuilder::from(table_scan)
2511 .sort(sort_exprs)
2512 .context(DataFusionPlanningSnafu)?
2513 .build()
2514 .context(DataFusionPlanningSnafu)?;
2515
2516 let time_index_column =
2518 self.ctx
2519 .time_index_column
2520 .clone()
2521 .with_context(|| TimeIndexNotFoundSnafu {
2522 table: table_ref.to_string(),
2523 })?;
2524 let divide_plan = LogicalPlan::Extension(Extension {
2525 node: Arc::new(SeriesDivide::new(
2526 series_key_columns.clone(),
2527 time_index_column,
2528 sort_plan,
2529 )),
2530 });
2531
2532 if !is_range_selector && offset_duration == 0 {
2534 return Ok(divide_plan);
2535 }
2536 let series_normalize = SeriesNormalize::new(
2537 offset_duration,
2538 self.ctx
2539 .time_index_column
2540 .clone()
2541 .with_context(|| TimeIndexNotFoundSnafu {
2542 table: table_ref.to_quoted_string(),
2543 })?,
2544 is_range_selector,
2545 series_key_columns,
2546 divide_plan,
2547 );
2548 let logical_plan = LogicalPlan::Extension(Extension {
2549 node: Arc::new(series_normalize),
2550 });
2551
2552 Ok(logical_plan)
2553 }
2554
2555 fn agg_modifier_to_col(
2562 &mut self,
2563 input_schema: &DFSchemaRef,
2564 modifier: &Option<LabelModifier>,
2565 update_ctx: bool,
2566 ) -> Result<Vec<DfExpr>> {
2567 match modifier {
2568 None => {
2569 if update_ctx {
2570 self.ctx.tag_columns.clear();
2571 }
2572 Ok(vec![self.create_time_index_column_expr()?])
2573 }
2574 Some(LabelModifier::Include(labels)) => {
2575 if update_ctx {
2576 self.ctx.tag_columns.clear();
2577 }
2578 let mut exprs = Vec::with_capacity(labels.labels.len());
2579 for label in &labels.labels {
2580 if is_metric_engine_internal_column(label) {
2581 continue;
2582 }
2583 if let Some(column_name) = Self::find_case_sensitive_column(input_schema, label)
2585 {
2586 exprs.push(DfExpr::Column(Column::from_name(column_name.clone())));
2587
2588 if update_ctx {
2589 self.ctx.tag_columns.push(column_name);
2591 }
2592 }
2593 }
2594 exprs.push(self.create_time_index_column_expr()?);
2596
2597 Ok(exprs)
2598 }
2599 Some(LabelModifier::Exclude(labels)) => {
2600 let mut all_fields = input_schema
2601 .fields()
2602 .iter()
2603 .map(|f| f.name())
2604 .collect::<BTreeSet<_>>();
2605
2606 all_fields.retain(|col| !is_metric_engine_internal_column(col.as_str()));
2609
2610 for label in &labels.labels {
2613 let _ = all_fields.remove(label);
2614 }
2615
2616 if let Some(time_index) = &self.ctx.time_index_column {
2618 let _ = all_fields.remove(time_index);
2619 }
2620 for value in &self.ctx.field_columns {
2621 let _ = all_fields.remove(value);
2622 }
2623
2624 if update_ctx {
2625 self.ctx.tag_columns = all_fields.iter().map(|col| (*col).clone()).collect();
2627 }
2628
2629 let mut exprs = all_fields
2631 .into_iter()
2632 .map(|c| DfExpr::Column(Column::from(c)))
2633 .collect::<Vec<_>>();
2634
2635 exprs.push(self.create_time_index_column_expr()?);
2637
2638 Ok(exprs)
2639 }
2640 }
2641 }
2642
2643 pub fn matchers_to_expr(
2645 label_matchers: Matchers,
2646 table_schema: &DFSchemaRef,
2647 ) -> Result<Vec<DfExpr>> {
2648 let mut exprs = Vec::with_capacity(label_matchers.matchers.len());
2649 for matcher in label_matchers.matchers {
2650 if matcher.name == SCHEMA_COLUMN_MATCHER
2651 || matcher.name == DB_COLUMN_MATCHER
2652 || matcher.name == FIELD_COLUMN_MATCHER
2653 {
2654 continue;
2655 }
2656
2657 let accepts_empty = matcher.is_match("");
2658 let column_name = Self::find_case_sensitive_column(table_schema, matcher.name.as_str());
2659 let col = if let Some(column_name) = column_name {
2660 let column = DfExpr::Column(Column::from_name(&column_name));
2661 let field = table_schema
2662 .index_of_column_by_name(None, &column_name)
2663 .map(|index| table_schema.field(index));
2664 if accepts_empty
2665 && column_name == OTLP_AGGREGATION_TEMPORALITY_LABEL
2666 && let Some(data_type) = field
2667 .filter(|field| {
2668 field.is_nullable()
2669 && Self::string_value_data_type(field.data_type()).is_some()
2670 })
2671 .map(|field| field.data_type())
2672 {
2673 let empty = Self::string_scalar_value(data_type, Some(String::new()))
2674 .expect("nullable label has a string type");
2675 DfExpr::ScalarFunction(ScalarFunction {
2676 func: coalesce(),
2677 args: vec![column, DfExpr::Literal(empty, None)],
2678 })
2679 } else {
2680 column
2681 }
2682 } else {
2683 DfExpr::Literal(ScalarValue::Utf8(Some(String::new())), None)
2684 .alias(matcher.name.clone())
2685 };
2686 let lit = DfExpr::Literal(ScalarValue::Utf8(Some(matcher.value)), None);
2687 let expr = match matcher.op {
2688 MatchOp::Equal => col.eq(lit),
2689 MatchOp::NotEqual => col.not_eq(lit),
2690 MatchOp::Re(re) => {
2691 if re.as_str() == "^(?:.*)$" {
2697 continue;
2698 }
2699 if re.as_str() == "^(?:.+)$" {
2700 col.not_eq(DfExpr::Literal(
2701 ScalarValue::Utf8(Some(String::new())),
2702 None,
2703 ))
2704 } else {
2705 DfExpr::BinaryExpr(BinaryExpr {
2706 left: Box::new(col),
2707 op: Operator::RegexMatch,
2708 right: Box::new(DfExpr::Literal(
2709 ScalarValue::Utf8(Some(re.as_str().to_string())),
2710 None,
2711 )),
2712 })
2713 }
2714 }
2715 MatchOp::NotRe(re) => {
2716 if re.as_str() == "^(?:.*)$" {
2717 DfExpr::Literal(ScalarValue::Boolean(Some(false)), None)
2718 } else if re.as_str() == "^(?:.+)$" {
2719 col.eq(DfExpr::Literal(
2720 ScalarValue::Utf8(Some(String::new())),
2721 None,
2722 ))
2723 } else {
2724 DfExpr::BinaryExpr(BinaryExpr {
2725 left: Box::new(col),
2726 op: Operator::RegexNotMatch,
2727 right: Box::new(DfExpr::Literal(
2728 ScalarValue::Utf8(Some(re.as_str().to_string())),
2729 None,
2730 )),
2731 })
2732 }
2733 }
2734 };
2735 exprs.push(expr);
2736 }
2737
2738 Ok(exprs)
2739 }
2740
2741 fn find_case_sensitive_column(schema: &DFSchemaRef, column: &str) -> Option<String> {
2742 if is_metric_engine_internal_column(column) {
2743 return None;
2744 }
2745 schema
2746 .fields()
2747 .iter()
2748 .find(|field| field.name() == column)
2749 .map(|field| field.name().clone())
2750 }
2751
2752 fn table_from_source(&self, source: &Arc<dyn TableSource>) -> Result<table::TableRef> {
2753 Ok(source
2754 .as_any()
2755 .downcast_ref::<DefaultTableSource>()
2756 .context(UnknownTableSnafu)?
2757 .table_provider
2758 .as_any()
2759 .downcast_ref::<DfTableProviderAdapter>()
2760 .context(UnknownTableSnafu)?
2761 .table())
2762 }
2763
2764 fn table_ref(&self) -> Result<TableReference> {
2765 let table_name = self
2766 .ctx
2767 .table_name
2768 .clone()
2769 .context(TableNameNotFoundSnafu)?;
2770
2771 let table_ref = if let Some(schema_name) = &self.ctx.schema_name {
2773 TableReference::partial(schema_name.as_str(), table_name.as_str())
2774 } else {
2775 TableReference::bare(table_name.as_str())
2776 };
2777
2778 Ok(table_ref)
2779 }
2780
2781 fn build_time_index_filter(
2782 &self,
2783 offset_duration: i64,
2784 schema: &DFSchemaRef,
2785 ) -> Result<Option<DfExpr>> {
2786 let start = self.ctx.start;
2787 let end = self.ctx.end;
2788 if end < start {
2789 return InvalidTimeRangeSnafu { start, end }.fail();
2790 }
2791 let time_index_expr = self.create_time_index_column_expr()?;
2792 let time_index_name = self.ctx.time_index_column.as_ref().unwrap();
2793 let unit = schema
2794 .index_of_column_by_name(None, time_index_name)
2795 .and_then(|index| match schema.field(index).data_type() {
2796 ArrowDataType::Timestamp(unit, _) => Some(*unit),
2797 _ => None,
2798 })
2799 .unwrap_or(ArrowTimeUnit::Millisecond);
2800 let native_value = |milliseconds: i128| match unit {
2801 ArrowTimeUnit::Second => milliseconds.div_euclid(1_000),
2802 ArrowTimeUnit::Millisecond => milliseconds,
2803 ArrowTimeUnit::Microsecond => milliseconds * 1_000,
2804 ArrowTimeUnit::Nanosecond => milliseconds * 1_000_000,
2805 };
2806 let scalar = |milliseconds: i128| -> Option<ScalarValue> {
2807 let value = i64::try_from(native_value(milliseconds)).ok()?;
2808 Some(match unit {
2809 ArrowTimeUnit::Second => ScalarValue::TimestampSecond(Some(value), None),
2810 ArrowTimeUnit::Millisecond => ScalarValue::TimestampMillisecond(Some(value), None),
2811 ArrowTimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(Some(value), None),
2812 ArrowTimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(Some(value), None),
2813 })
2814 };
2815 let window = self.ctx.range.unwrap_or(self.ctx.lookback_delta);
2816 let filter = |lower_ms: i128, upper_ms: i128| {
2817 let lower_value = native_value(lower_ms);
2818 let upper_value = native_value(upper_ms);
2819 if lower_value > i128::from(i64::MAX) || upper_value < i128::from(i64::MIN) {
2820 return Some(lit(false));
2821 }
2822 let lower_filter = (lower_value >= i128::from(i64::MIN)).then(|| {
2823 let lower = DfExpr::Literal(scalar(lower_ms).unwrap(), None);
2824 if window == 0 {
2825 time_index_expr.clone().gt_eq(lower)
2826 } else if unit == ArrowTimeUnit::Millisecond
2827 && let Some(inclusive_lower) = lower_ms
2828 .checked_add(1)
2829 .and_then(|lower| i64::try_from(lower).ok())
2830 .and_then(|lower| scalar(i128::from(lower)))
2831 {
2832 time_index_expr
2833 .clone()
2834 .gt_eq(DfExpr::Literal(inclusive_lower, None))
2835 } else {
2836 time_index_expr.clone().gt(lower)
2837 }
2838 });
2839 let upper_filter = (upper_value <= i128::from(i64::MAX)).then(|| {
2840 time_index_expr
2841 .clone()
2842 .lt_eq(DfExpr::Literal(scalar(upper_ms).unwrap(), None))
2843 });
2844
2845 match (lower_filter, upper_filter) {
2849 (Some(lower), Some(upper)) => Some(lower.and(upper)),
2850 (Some(filter), None) | (None, Some(filter)) => Some(filter),
2851 (None, None) => None,
2852 }
2853 };
2854 let bounds = |timestamp: i64| {
2855 let upper = i128::from(timestamp) - i128::from(offset_duration);
2856 (upper - i128::from(window), upper)
2857 };
2858 let num_points = (end as i128 - start as i128) / self.ctx.interval as i128;
2859 if num_points > MAX_SCATTER_POINTS as i128 || self.ctx.interval <= INTERVAL_1H {
2860 let (lower, _) = bounds(start);
2861 let (_, upper) = bounds(end);
2862 return Ok(filter(lower, upper));
2863 }
2864 let mut filters = Vec::new();
2865 for timestamp in (start..=end).step_by(self.ctx.interval as usize) {
2866 let (lower, upper) = bounds(timestamp);
2867 let Some(filter) = filter(lower, upper) else {
2868 return Ok(None);
2871 };
2872 filters.push(filter);
2873 }
2874 Ok(filters.into_iter().reduce(DfExpr::or))
2875 }
2876
2877 async fn create_table_scan_plan(&mut self, table_ref: TableReference) -> Result<LogicalPlan> {
2882 let provider = self
2883 .table_provider
2884 .resolve_table(table_ref.clone())
2885 .await
2886 .context(CatalogSnafu)?;
2887
2888 let logical_table = self.table_from_source(&provider)?;
2889
2890 let mut maybe_phy_table_ref = table_ref.clone();
2892 let mut scan_provider = provider;
2893 let mut table_id_filter: Option<u32> = None;
2894
2895 if logical_table.table_info().meta.engine == METRIC_ENGINE_NAME
2898 && let Some(physical_table_name) = logical_table
2899 .table_info()
2900 .meta
2901 .options
2902 .extra_options
2903 .get(LOGICAL_TABLE_METADATA_KEY)
2904 {
2905 let physical_table_ref = if let Some(schema_name) = &self.ctx.schema_name {
2906 TableReference::partial(schema_name.as_str(), physical_table_name.as_str())
2907 } else {
2908 TableReference::bare(physical_table_name.as_str())
2909 };
2910
2911 let physical_provider = match self
2912 .table_provider
2913 .resolve_table(physical_table_ref.clone())
2914 .await
2915 {
2916 Ok(provider) => provider,
2917 Err(e) if e.status_code() == StatusCode::TableNotFound => {
2918 scan_provider.clone()
2921 }
2922 Err(e) => return Err(e).context(CatalogSnafu),
2923 };
2924
2925 if !Arc::ptr_eq(&physical_provider, &scan_provider) {
2926 let physical_table = self.table_from_source(&physical_provider)?;
2928
2929 let has_table_id = physical_table
2930 .schema()
2931 .column_schema_by_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME)
2932 .is_some();
2933 let has_tsid = physical_table
2934 .schema()
2935 .column_schema_by_name(DATA_SCHEMA_TSID_COLUMN_NAME)
2936 .is_some_and(|col| matches!(col.data_type, ConcreteDataType::UInt64(_)));
2937
2938 if has_table_id && has_tsid {
2939 scan_provider = physical_provider;
2940 maybe_phy_table_ref = physical_table_ref;
2941 table_id_filter = Some(logical_table.table_info().ident.table_id);
2942 }
2943 }
2944 }
2945
2946 let scan_table = self.table_from_source(&scan_provider)?;
2947
2948 let use_tsid = table_id_filter.is_some()
2949 && scan_table
2950 .schema()
2951 .column_schema_by_name(DATA_SCHEMA_TSID_COLUMN_NAME)
2952 .is_some_and(|col| matches!(col.data_type, ConcreteDataType::UInt64(_)));
2953 self.ctx.use_tsid = use_tsid;
2954
2955 let all_table_tags = self.ctx.tag_columns.clone();
2956
2957 let scan_tag_columns = if use_tsid {
2958 let mut scan_tags = self.ctx.tag_columns.clone();
2959 for matcher in &self.ctx.selector_matcher {
2960 if is_metric_engine_internal_column(&matcher.name) {
2961 continue;
2962 }
2963 if all_table_tags.iter().any(|tag| tag == &matcher.name) {
2964 scan_tags.push(matcher.name.clone());
2965 }
2966 }
2967 scan_tags.sort_unstable();
2968 scan_tags.dedup();
2969 scan_tags
2970 } else {
2971 self.ctx.tag_columns.clone()
2972 };
2973
2974 let time_index_data_type = scan_table
2975 .schema()
2976 .timestamp_column()
2977 .with_context(|| TimeIndexNotFoundSnafu {
2978 table: maybe_phy_table_ref.to_quoted_string(),
2979 })?
2980 .data_type
2981 .clone();
2982 let is_time_index_second =
2983 time_index_data_type == ConcreteDataType::timestamp_second_datatype();
2984
2985 let scan_projection = if table_id_filter.is_some() {
2986 let mut required_columns = HashSet::new();
2987 required_columns.insert(DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string());
2988 required_columns.insert(self.ctx.time_index_column.clone().with_context(|| {
2989 TimeIndexNotFoundSnafu {
2990 table: maybe_phy_table_ref.to_quoted_string(),
2991 }
2992 })?);
2993 for col in &scan_tag_columns {
2994 required_columns.insert(col.clone());
2995 }
2996 for col in &self.ctx.field_columns {
2997 required_columns.insert(col.clone());
2998 }
2999 if use_tsid {
3000 required_columns.insert(DATA_SCHEMA_TSID_COLUMN_NAME.to_string());
3001 }
3002
3003 let arrow_schema = scan_table.schema().arrow_schema().clone();
3004 Some(
3005 arrow_schema
3006 .fields()
3007 .iter()
3008 .enumerate()
3009 .filter(|(_, field)| required_columns.contains(field.name().as_str()))
3010 .map(|(idx, _)| idx)
3011 .collect::<Vec<_>>(),
3012 )
3013 } else {
3014 None
3015 };
3016
3017 let mut scan_plan =
3018 LogicalPlanBuilder::scan(maybe_phy_table_ref.clone(), scan_provider, scan_projection)
3019 .context(DataFusionPlanningSnafu)?
3020 .build()
3021 .context(DataFusionPlanningSnafu)?;
3022
3023 if let Some(table_id) = table_id_filter {
3024 scan_plan = LogicalPlanBuilder::from(scan_plan)
3025 .filter(
3026 DfExpr::Column(Column::from_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME))
3027 .eq(lit(table_id)),
3028 )
3029 .context(DataFusionPlanningSnafu)?
3030 .alias(table_ref.clone()) .context(DataFusionPlanningSnafu)?
3032 .build()
3033 .context(DataFusionPlanningSnafu)?;
3034 }
3035
3036 if is_time_index_second {
3037 let expr: Vec<_> = self
3042 .create_field_column_exprs()?
3043 .into_iter()
3044 .chain(
3045 scan_tag_columns
3046 .iter()
3047 .map(|tag| DfExpr::Column(Column::from_name(tag))),
3048 )
3049 .chain(self.ctx.use_tsid.then_some(DfExpr::Column(Column::new(
3050 Some(table_ref.clone()),
3051 DATA_SCHEMA_TSID_COLUMN_NAME.to_string(),
3052 ))))
3053 .chain(Some(DfExpr::Alias(Alias {
3054 expr: Box::new(DfExpr::Cast(Cast {
3055 expr: Box::new(self.create_time_index_column_expr()?),
3056 data_type: ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
3057 })),
3058 relation: Some(table_ref.clone()),
3059 name: self
3060 .ctx
3061 .time_index_column
3062 .as_ref()
3063 .with_context(|| TimeIndexNotFoundSnafu {
3064 table: table_ref.to_quoted_string(),
3065 })?
3066 .clone(),
3067 metadata: None,
3068 })))
3069 .collect::<Vec<_>>();
3070 scan_plan = LogicalPlanBuilder::from(scan_plan)
3071 .project(expr)
3072 .context(DataFusionPlanningSnafu)?
3073 .build()
3074 .context(DataFusionPlanningSnafu)?;
3075 } else if table_id_filter.is_some()
3076 || time_index_data_type == ConcreteDataType::timestamp_microsecond_datatype()
3077 || time_index_data_type == ConcreteDataType::timestamp_nanosecond_datatype()
3078 {
3079 let project_exprs = self
3085 .create_field_column_exprs()?
3086 .into_iter()
3087 .chain(
3088 scan_tag_columns
3089 .iter()
3090 .map(|tag| DfExpr::Column(Column::from_name(tag))),
3091 )
3092 .chain(
3093 self.ctx
3094 .use_tsid
3095 .then_some(DfExpr::Column(Column::from_name(
3096 DATA_SCHEMA_TSID_COLUMN_NAME,
3097 ))),
3098 )
3099 .chain(Some(self.create_time_index_column_expr()?))
3100 .collect::<Vec<_>>();
3101
3102 scan_plan = LogicalPlanBuilder::from(scan_plan)
3103 .project(project_exprs)
3104 .context(DataFusionPlanningSnafu)?
3105 .build()
3106 .context(DataFusionPlanningSnafu)?;
3107 }
3108
3109 let result = LogicalPlanBuilder::from(scan_plan)
3110 .build()
3111 .context(DataFusionPlanningSnafu)?;
3112 Ok(result)
3113 }
3114
3115 fn collect_row_key_tag_columns_from_plan(
3116 &self,
3117 plan: &LogicalPlan,
3118 ) -> Result<BTreeSet<String>> {
3119 fn walk(
3120 planner: &PromPlanner,
3121 plan: &LogicalPlan,
3122 out: &mut BTreeSet<String>,
3123 ) -> Result<()> {
3124 if let LogicalPlan::TableScan(scan) = plan
3126 && let Ok(table) = planner.table_from_source(&scan.source)
3127 {
3128 for col in table.table_info().meta.row_key_column_names() {
3129 if col != DATA_SCHEMA_TABLE_ID_COLUMN_NAME
3130 && col != DATA_SCHEMA_TSID_COLUMN_NAME
3131 && !is_metric_engine_internal_column(col)
3132 {
3133 out.insert(col.clone());
3134 }
3135 }
3136 }
3137
3138 for input in plan.inputs() {
3139 walk(planner, input, out)?;
3140 }
3141 Ok(())
3142 }
3143
3144 let mut out = BTreeSet::new();
3145 walk(self, plan, &mut out)?;
3146 Ok(out)
3147 }
3148
3149 fn ensure_tag_columns_available(
3150 &self,
3151 plan: LogicalPlan,
3152 required_tags: &BTreeSet<String>,
3153 ) -> Result<LogicalPlan> {
3154 if required_tags.is_empty() {
3155 return Ok(plan);
3156 }
3157
3158 struct Rewriter {
3159 required_tags: BTreeSet<String>,
3160 }
3161
3162 impl TreeNodeRewriter for Rewriter {
3163 type Node = LogicalPlan;
3164
3165 fn f_up(
3166 &mut self,
3167 node: Self::Node,
3168 ) -> datafusion_common::Result<Transformed<Self::Node>> {
3169 match node {
3170 LogicalPlan::TableScan(scan) => {
3171 let schema = scan.source.schema();
3172 let mut projection = match scan.projection.clone() {
3173 Some(p) => p,
3174 None => {
3175 return Ok(Transformed::no(LogicalPlan::TableScan(scan)));
3177 }
3178 };
3179
3180 let mut changed = false;
3181 for tag in &self.required_tags {
3182 if let Some((idx, _)) = schema
3183 .fields()
3184 .iter()
3185 .enumerate()
3186 .find(|(_, field)| field.name() == tag)
3187 && !projection.contains(&idx)
3188 {
3189 projection.push(idx);
3190 changed = true;
3191 }
3192 }
3193
3194 if !changed {
3195 return Ok(Transformed::no(LogicalPlan::TableScan(scan)));
3196 }
3197
3198 projection.sort_unstable();
3199 projection.dedup();
3200
3201 let new_scan = TableScan::try_new(
3202 scan.table_name.clone(),
3203 scan.source.clone(),
3204 Some(projection),
3205 scan.filters,
3206 scan.fetch,
3207 )?;
3208 Ok(Transformed::yes(LogicalPlan::TableScan(new_scan)))
3209 }
3210 LogicalPlan::Projection(proj) => {
3211 let input_schema = proj.input.schema();
3212
3213 let existing = proj
3214 .schema
3215 .fields()
3216 .iter()
3217 .map(|f| f.name().as_str())
3218 .collect::<HashSet<_>>();
3219
3220 let mut expr = proj.expr.clone();
3221 let mut has_changed = false;
3222 for tag in &self.required_tags {
3223 if existing.contains(tag.as_str()) {
3224 continue;
3225 }
3226
3227 if let Some(idx) = input_schema.index_of_column_by_name(None, tag) {
3228 expr.push(DfExpr::Column(Column::from(
3229 input_schema.qualified_field(idx),
3230 )));
3231 has_changed = true;
3232 }
3233 }
3234
3235 if !has_changed {
3236 return Ok(Transformed::no(LogicalPlan::Projection(proj)));
3237 }
3238
3239 let new_proj = Projection::try_new(expr, proj.input)?;
3240 Ok(Transformed::yes(LogicalPlan::Projection(new_proj)))
3241 }
3242 other => Ok(Transformed::no(other)),
3243 }
3244 }
3245 }
3246
3247 let mut rewriter = Rewriter {
3248 required_tags: required_tags.clone(),
3249 };
3250 let rewritten = plan
3251 .rewrite(&mut rewriter)
3252 .context(DataFusionPlanningSnafu)?;
3253 Ok(rewritten.data)
3254 }
3255
3256 fn refresh_tag_columns_from_schema(&mut self, schema: &DFSchemaRef) {
3257 let time_index = self.ctx.time_index_column.as_deref();
3258 let field_columns = self.ctx.field_columns.iter().collect::<HashSet<_>>();
3259
3260 let mut tags = schema
3261 .fields()
3262 .iter()
3263 .map(|f| f.name())
3264 .filter(|name| Some(name.as_str()) != time_index)
3265 .filter(|name| !field_columns.contains(name))
3266 .filter(|name| !is_metric_engine_internal_column(name))
3267 .cloned()
3268 .collect::<Vec<_>>();
3269 tags.sort_unstable();
3270 tags.dedup();
3271 self.ctx.tag_columns = tags;
3272 }
3273
3274 async fn setup_context(&mut self) -> Result<Option<LogicalPlan>> {
3278 let table_ref = self.table_ref()?;
3279 let source = match self.table_provider.resolve_table(table_ref.clone()).await {
3280 Err(e) if e.status_code() == StatusCode::TableNotFound => {
3281 let plan = self.setup_context_for_empty_metric()?;
3282 return Ok(Some(plan));
3283 }
3284 res => res.context(CatalogSnafu)?,
3285 };
3286 let table = self.table_from_source(&source)?;
3287
3288 let time_index = table
3290 .schema()
3291 .timestamp_column()
3292 .with_context(|| TimeIndexNotFoundSnafu {
3293 table: table_ref.to_quoted_string(),
3294 })?
3295 .name
3296 .clone();
3297 self.ctx.time_index_column = Some(time_index);
3298
3299 let values = table
3301 .table_info()
3302 .meta
3303 .field_column_names()
3304 .cloned()
3305 .collect();
3306 self.ctx.field_columns = values;
3307
3308 let tags = table
3310 .table_info()
3311 .meta
3312 .row_key_column_names()
3313 .filter(|col| {
3314 col != &DATA_SCHEMA_TABLE_ID_COLUMN_NAME && col != &DATA_SCHEMA_TSID_COLUMN_NAME
3316 })
3317 .cloned()
3318 .collect();
3319 self.ctx.tag_columns = tags;
3320
3321 self.ctx.use_tsid = false;
3322
3323 Ok(None)
3324 }
3325
3326 fn setup_context_for_empty_metric(&mut self) -> Result<LogicalPlan> {
3329 self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
3330 self.ctx.reset_table_name_and_schema();
3331 self.ctx.tag_columns = vec![];
3332 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
3333 self.ctx.use_tsid = false;
3334
3335 let plan = LogicalPlan::Extension(Extension {
3337 node: Arc::new(
3338 EmptyMetric::new(
3339 0,
3340 -1,
3341 self.ctx.interval,
3342 SPECIAL_TIME_FUNCTION.to_string(),
3343 DEFAULT_FIELD_COLUMN.to_string(),
3344 Some(lit(0.0f64)),
3345 )
3346 .context(DataFusionPlanningSnafu)?,
3347 ),
3348 });
3349 Ok(plan)
3350 }
3351
3352 fn create_function_args(&self, args: &[Box<PromExpr>]) -> Result<FunctionArgs> {
3354 let mut result = FunctionArgs::default();
3355
3356 for arg in args {
3357 if let Some(expr) = Self::try_build_literal_expr(arg) {
3359 result.literals.push(expr);
3360 } else {
3361 match arg.as_ref() {
3363 PromExpr::Subquery(_)
3364 | PromExpr::VectorSelector(_)
3365 | PromExpr::MatrixSelector(_)
3366 | PromExpr::Extension(_)
3367 | PromExpr::Aggregate(_)
3368 | PromExpr::Paren(_)
3369 | PromExpr::Call(_)
3370 | PromExpr::Binary(_)
3371 | PromExpr::Unary(_) => {
3372 if result.input.replace(*arg.clone()).is_some() {
3373 MultipleVectorSnafu { expr: *arg.clone() }.fail()?;
3374 }
3375 }
3376
3377 _ => {
3378 let expr = Self::get_param_as_literal_expr(Some(arg.as_ref()), None, None)?;
3379 result.literals.push(expr);
3380 }
3381 }
3382 }
3383 }
3384
3385 Ok(result)
3386 }
3387
3388 fn create_mixed_range_function_exprs(
3389 &mut self,
3390 func: &Function,
3391 mut other_input_exprs: VecDeque<DfExpr>,
3392 float_field: &str,
3393 histogram_field: &str,
3394 input_schema: &DFSchemaRef,
3395 ) -> Result<Option<Vec<DfExpr>>> {
3396 let returns_histogram = matches!(
3397 func.name,
3398 "rate"
3399 | "increase"
3400 | "delta"
3401 | "idelta"
3402 | "irate"
3403 | "avg_over_time"
3404 | "sum_over_time"
3405 | "last_over_time"
3406 );
3407 if !returns_histogram
3408 && !matches!(
3409 func.name,
3410 "changes"
3411 | "resets"
3412 | "deriv"
3413 | "min_over_time"
3414 | "max_over_time"
3415 | "count_over_time"
3416 | "absent_over_time"
3417 | "present_over_time"
3418 | "stddev_over_time"
3419 | "stdvar_over_time"
3420 | "quantile_over_time"
3421 | "predict_linear"
3422 | "double_exponential_smoothing"
3423 | "holt_winters"
3424 )
3425 {
3426 return Ok(None);
3427 }
3428
3429 if func.name == "predict_linear" {
3430 other_input_exprs[0] = DfExpr::Cast(Cast {
3431 expr: Box::new(other_input_exprs[0].clone()),
3432 data_type: ArrowDataType::Int64,
3433 });
3434 }
3435
3436 let timestamp_range = DfExpr::Column(Column::from_name(
3437 RangeManipulate::build_timestamp_range_name(
3438 self.ctx.time_index_column.as_ref().unwrap(),
3439 ),
3440 ));
3441 let float_range = DfExpr::Column(Column::from_name(float_field));
3442 let histogram_range = DfExpr::Column(Column::from_name(histogram_field));
3443 let mut args = Vec::with_capacity(other_input_exprs.len() + 6);
3444 args.push(lit(func.name));
3445 args.push(timestamp_range.clone());
3446 args.push(float_range.clone());
3447 args.push(histogram_range.clone());
3448 args.extend(other_input_exprs);
3449 if matches!(func.name, "rate" | "increase" | "delta") {
3450 args.push(self.create_time_index_column_expr()?);
3451 args.push(lit(self.ctx.range.context(ExpectRangeSelectorSnafu)?));
3452 }
3453
3454 let mut float_expr = DfExpr::ScalarFunction(ScalarFunction {
3455 func: Arc::new(MixedRange::float_udf(self.promql_annotations.clone())),
3456 args: args.clone(),
3457 });
3458 if matches!(func.name, "rate" | "increase") {
3459 let raw_delta_function = if func.name == "rate" {
3460 "raw_delta_rate"
3461 } else {
3462 "raw_delta_increase"
3463 };
3464 let delta_sum = DfExpr::ScalarFunction(ScalarFunction {
3465 func: Arc::new(MixedRange::float_udf(self.promql_annotations.clone())),
3466 args: vec![
3467 lit(raw_delta_function),
3468 timestamp_range,
3469 float_range,
3470 histogram_range,
3471 ],
3472 });
3473 float_expr = self.select_delta_range_math(
3474 func.name,
3475 input_schema,
3476 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3477 delta_sum,
3478 float_expr,
3479 )?;
3480 }
3481 let exprs = if returns_histogram {
3482 self.ctx.field_columns = vec![float_field.to_string(), histogram_field.to_string()];
3483 vec![
3484 float_expr.alias(float_field),
3485 DfExpr::ScalarFunction(ScalarFunction {
3486 func: Arc::new(MixedRange::histogram_udf(self.promql_annotations.clone())),
3487 args,
3488 })
3489 .alias(histogram_field),
3490 ]
3491 } else {
3492 let display_name = float_expr.schema_name().to_string();
3493 self.ctx.field_columns = vec![display_name.clone()];
3494 vec![float_expr.alias(display_name)]
3495 };
3496 Ok(Some(exprs))
3497 }
3498
3499 fn create_function_expr(
3505 &mut self,
3506 func: &Function,
3507 other_input_exprs: Vec<DfExpr>,
3508 input_schema: &DFSchemaRef,
3509 query_engine_state: &QueryEngineState,
3510 ) -> Result<(Vec<DfExpr>, Vec<String>)> {
3511 let mut other_input_exprs: VecDeque<DfExpr> = other_input_exprs.into();
3513 if let Some((float_field, histogram_field)) =
3514 Self::alternative_sample_range_columns(input_schema, &self.ctx.field_columns)
3515 .map(|(float, histogram)| (float.to_string(), histogram.to_string()))
3516 && let Some(exprs) = self.create_mixed_range_function_exprs(
3517 func,
3518 other_input_exprs.clone(),
3519 &float_field,
3520 &histogram_field,
3521 input_schema,
3522 )?
3523 {
3524 return Ok((exprs, vec![]));
3525 }
3526 let alternative_samples =
3527 Self::field_columns_are_alternative_samples(input_schema, &self.ctx.field_columns);
3528 let all_field_columns_are_native_histogram_ranges =
3529 self.all_field_columns_are_native_histogram_ranges(input_schema);
3530
3531 let field_column_pos = 0;
3533 let mut exprs = Vec::with_capacity(self.ctx.field_columns.len());
3534 let mut new_tags = vec![];
3536 let promql_annotations = self.promql_annotations.clone();
3537 let native_histogram_drop_udf = |name: &str| {
3538 Arc::new(NativeHistogramDrop::float_null_udf(
3539 format!(
3540 "{name}: dropped native histogram samples because this function is not supported for native histograms"
3541 ),
3542 promql_annotations.clone(),
3543 ))
3544 };
3545 let scalar_func = match func.name {
3546 "increase" => {
3547 if all_field_columns_are_native_histogram_ranges {
3548 ScalarFunc::ExtrapolateUdf(
3549 Arc::new(NativeHistogramIncrease::scalar_udf_with_collector(
3550 self.promql_annotations.clone(),
3551 )),
3552 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3553 )
3554 } else {
3555 ScalarFunc::ExtrapolateUdf(
3556 Arc::new(Increase::scalar_udf()),
3557 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3558 )
3559 }
3560 }
3561 "rate" => {
3562 if all_field_columns_are_native_histogram_ranges {
3563 ScalarFunc::ExtrapolateUdf(
3564 Arc::new(NativeHistogramRate::scalar_udf_with_collector(
3565 self.promql_annotations.clone(),
3566 )),
3567 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3568 )
3569 } else {
3570 ScalarFunc::ExtrapolateUdf(
3571 Arc::new(Rate::scalar_udf()),
3572 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3573 )
3574 }
3575 }
3576 "delta" => {
3577 if all_field_columns_are_native_histogram_ranges {
3578 ScalarFunc::ExtrapolateUdf(
3579 Arc::new(NativeHistogramDelta::scalar_udf_with_collector(
3580 self.promql_annotations.clone(),
3581 )),
3582 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3583 )
3584 } else {
3585 ScalarFunc::ExtrapolateUdf(
3586 Arc::new(Delta::scalar_udf()),
3587 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3588 )
3589 }
3590 }
3591 "idelta" => {
3592 if all_field_columns_are_native_histogram_ranges {
3593 ScalarFunc::Udf(Arc::new(NativeHistogramIDelta::scalar_udf_with_collector(
3594 self.promql_annotations.clone(),
3595 )))
3596 } else {
3597 ScalarFunc::Udf(Arc::new(IDelta::<false>::scalar_udf()))
3598 }
3599 }
3600 "irate" => {
3601 if all_field_columns_are_native_histogram_ranges {
3602 ScalarFunc::Udf(Arc::new(NativeHistogramIRate::scalar_udf_with_collector(
3603 self.promql_annotations.clone(),
3604 )))
3605 } else {
3606 ScalarFunc::Udf(Arc::new(IDelta::<true>::scalar_udf()))
3607 }
3608 }
3609 "resets" => {
3610 if all_field_columns_are_native_histogram_ranges {
3611 ScalarFunc::Udf(Arc::new(NativeHistogramResets::scalar_udf()))
3612 } else {
3613 ScalarFunc::Udf(Arc::new(Resets::scalar_udf()))
3614 }
3615 }
3616 "changes" => {
3617 if all_field_columns_are_native_histogram_ranges {
3618 ScalarFunc::Udf(Arc::new(NativeHistogramChanges::scalar_udf()))
3619 } else {
3620 ScalarFunc::Udf(Arc::new(Changes::scalar_udf()))
3621 }
3622 }
3623 "deriv" => {
3624 if all_field_columns_are_native_histogram_ranges {
3625 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3626 } else {
3627 ScalarFunc::Udf(Arc::new(Deriv::scalar_udf()))
3628 }
3629 }
3630 "avg_over_time" => {
3631 if all_field_columns_are_native_histogram_ranges {
3632 ScalarFunc::Udf(Arc::new(
3633 NativeHistogramAvgOverTime::scalar_udf_with_collector(
3634 self.promql_annotations.clone(),
3635 ),
3636 ))
3637 } else {
3638 ScalarFunc::Udf(Arc::new(AvgOverTime::scalar_udf()))
3639 }
3640 }
3641 "min_over_time" => {
3642 if all_field_columns_are_native_histogram_ranges {
3643 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3644 } else {
3645 ScalarFunc::Udf(Arc::new(MinOverTime::scalar_udf()))
3646 }
3647 }
3648 "max_over_time" => {
3649 if all_field_columns_are_native_histogram_ranges {
3650 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3651 } else {
3652 ScalarFunc::Udf(Arc::new(MaxOverTime::scalar_udf()))
3653 }
3654 }
3655 "sum_over_time" => {
3656 if all_field_columns_are_native_histogram_ranges {
3657 ScalarFunc::Udf(Arc::new(
3658 NativeHistogramSumOverTime::scalar_udf_with_collector(
3659 self.promql_annotations.clone(),
3660 ),
3661 ))
3662 } else {
3663 ScalarFunc::Udf(Arc::new(SumOverTime::scalar_udf()))
3664 }
3665 }
3666 "count_over_time" => {
3667 if all_field_columns_are_native_histogram_ranges {
3668 ScalarFunc::Udf(Arc::new(NativeHistogramCountOverTime::scalar_udf()))
3669 } else {
3670 ScalarFunc::Udf(Arc::new(CountOverTime::scalar_udf()))
3671 }
3672 }
3673 "last_over_time" => {
3674 if all_field_columns_are_native_histogram_ranges {
3675 ScalarFunc::Udf(Arc::new(NativeHistogramLastOverTime::scalar_udf()))
3676 } else {
3677 ScalarFunc::Udf(Arc::new(LastOverTime::scalar_udf()))
3678 }
3679 }
3680 "absent_over_time" => {
3681 if all_field_columns_are_native_histogram_ranges {
3682 ScalarFunc::Udf(Arc::new(NativeHistogramAbsentOverTime::scalar_udf()))
3683 } else {
3684 ScalarFunc::Udf(Arc::new(AbsentOverTime::scalar_udf()))
3685 }
3686 }
3687 "present_over_time" => {
3688 if all_field_columns_are_native_histogram_ranges {
3689 ScalarFunc::Udf(Arc::new(NativeHistogramPresentOverTime::scalar_udf()))
3690 } else {
3691 ScalarFunc::Udf(Arc::new(PresentOverTime::scalar_udf()))
3692 }
3693 }
3694 "stddev_over_time" => {
3695 if all_field_columns_are_native_histogram_ranges {
3696 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3697 } else {
3698 ScalarFunc::Udf(Arc::new(StddevOverTime::scalar_udf()))
3699 }
3700 }
3701 "stdvar_over_time" => {
3702 if all_field_columns_are_native_histogram_ranges {
3703 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3704 } else {
3705 ScalarFunc::Udf(Arc::new(StdvarOverTime::scalar_udf()))
3706 }
3707 }
3708 "quantile_over_time" => {
3709 if all_field_columns_are_native_histogram_ranges {
3710 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3711 } else {
3712 ScalarFunc::Udf(Arc::new(QuantileOverTime::scalar_udf()))
3713 }
3714 }
3715 "predict_linear" => {
3716 if all_field_columns_are_native_histogram_ranges {
3717 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3718 } else {
3719 other_input_exprs[0] = DfExpr::Cast(Cast {
3720 expr: Box::new(other_input_exprs[0].clone()),
3721 data_type: ArrowDataType::Int64,
3722 });
3723 ScalarFunc::Udf(Arc::new(PredictLinear::scalar_udf()))
3724 }
3725 }
3726 "double_exponential_smoothing" | "holt_winters" => {
3727 if all_field_columns_are_native_histogram_ranges {
3728 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3729 } else {
3730 ScalarFunc::Udf(Arc::new(DoubleExponentialSmoothing::scalar_udf()))
3731 }
3732 }
3733 "histogram_count" => {
3734 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramCount::scalar_udf()))
3735 }
3736 "histogram_sum" => {
3737 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramSum::scalar_udf()))
3738 }
3739 "histogram_avg" => {
3740 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramAvg::scalar_udf()))
3741 }
3742 "histogram_stddev" => {
3743 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramStddev::scalar_udf()))
3744 }
3745 "histogram_stdvar" => {
3746 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramStdvar::scalar_udf()))
3747 }
3748 "time" => {
3749 exprs.push(build_special_time_expr(
3750 self.ctx.time_index_column.as_ref().unwrap(),
3751 ));
3752 ScalarFunc::GeneratedExpr
3753 }
3754 "minute" => {
3755 let expr = self.date_part_on_time_index("minute")?;
3757 exprs.push(expr);
3758 ScalarFunc::GeneratedExpr
3759 }
3760 "hour" => {
3761 let expr = self.date_part_on_time_index("hour")?;
3763 exprs.push(expr);
3764 ScalarFunc::GeneratedExpr
3765 }
3766 "month" => {
3767 let expr = self.date_part_on_time_index("month")?;
3769 exprs.push(expr);
3770 ScalarFunc::GeneratedExpr
3771 }
3772 "year" => {
3773 let expr = self.date_part_on_time_index("year")?;
3775 exprs.push(expr);
3776 ScalarFunc::GeneratedExpr
3777 }
3778 "day_of_month" => {
3779 let expr = self.date_part_on_time_index("day")?;
3781 exprs.push(expr);
3782 ScalarFunc::GeneratedExpr
3783 }
3784 "day_of_week" => {
3785 let expr = self.date_part_on_time_index("dow")?;
3787 exprs.push(expr);
3788 ScalarFunc::GeneratedExpr
3789 }
3790 "day_of_year" => {
3791 let expr = self.date_part_on_time_index("doy")?;
3793 exprs.push(expr);
3794 ScalarFunc::GeneratedExpr
3795 }
3796 "days_in_month" => {
3797 let day_lit_expr = "day".lit();
3802 let month_lit_expr = "month".lit();
3803 let interval_1month_lit_expr =
3804 DfExpr::Literal(ScalarValue::IntervalYearMonth(Some(1)), None);
3805 let interval_1day_lit_expr = DfExpr::Literal(
3806 ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(1, 0))),
3807 None,
3808 );
3809 let the_1month_minus_1day_expr = DfExpr::BinaryExpr(BinaryExpr {
3810 left: Box::new(interval_1month_lit_expr),
3811 op: Operator::Minus,
3812 right: Box::new(interval_1day_lit_expr),
3813 });
3814 let date_trunc_expr = DfExpr::ScalarFunction(ScalarFunction {
3815 func: datafusion_functions::datetime::date_trunc(),
3816 args: vec![month_lit_expr, self.create_time_index_column_expr()?],
3817 });
3818 let date_trunc_plus_interval_expr = DfExpr::BinaryExpr(BinaryExpr {
3819 left: Box::new(date_trunc_expr),
3820 op: Operator::Plus,
3821 right: Box::new(the_1month_minus_1day_expr),
3822 });
3823 let date_part_expr = DfExpr::ScalarFunction(ScalarFunction {
3824 func: datafusion_functions::datetime::date_part(),
3825 args: vec![day_lit_expr, date_trunc_plus_interval_expr],
3826 });
3827
3828 exprs.push(date_part_expr);
3829 ScalarFunc::GeneratedExpr
3830 }
3831
3832 "label_join" => {
3833 self.ctx.use_tsid = false;
3834 let (concat_expr, dst_label) = Self::build_concat_labels_expr(
3835 &mut other_input_exprs,
3836 &self.ctx,
3837 query_engine_state,
3838 )?;
3839
3840 for value in &self.ctx.field_columns {
3842 if *value != dst_label {
3843 let expr = DfExpr::Column(Column::from_name(value));
3844 exprs.push(expr);
3845 }
3846 }
3847
3848 self.ctx.tag_columns.retain(|tag| *tag != dst_label);
3850 new_tags.push(dst_label);
3851 exprs.push(concat_expr);
3853
3854 ScalarFunc::GeneratedExpr
3855 }
3856 "label_replace" => {
3857 self.ctx.use_tsid = false;
3858 if let Some((replace_expr, dst_label)) = self
3859 .build_regexp_replace_label_expr(&mut other_input_exprs, query_engine_state)?
3860 {
3861 for value in &self.ctx.field_columns {
3863 if *value != dst_label {
3864 let expr = DfExpr::Column(Column::from_name(value));
3865 exprs.push(expr);
3866 }
3867 }
3868
3869 ensure!(
3870 !self.ctx.tag_columns.contains(&dst_label),
3871 SameLabelSetSnafu
3872 );
3873 new_tags.push(dst_label);
3874 exprs.push(replace_expr);
3876 } else {
3877 for value in &self.ctx.field_columns {
3879 let expr = DfExpr::Column(Column::from_name(value));
3880 exprs.push(expr);
3881 }
3882 }
3883
3884 ScalarFunc::GeneratedExpr
3885 }
3886 "sort" | "sort_desc" => {
3887 for value in &self.ctx.field_columns {
3889 if !Self::field_column_is_native_histogram(input_schema, value) {
3890 exprs.push(DfExpr::Column(Column::from_name(value)));
3891 }
3892 }
3893 if exprs.is_empty() {
3896 exprs.push(DfExpr::Literal(ScalarValue::Float64(None), None));
3897 }
3898
3899 ScalarFunc::GeneratedExpr
3900 }
3901 "sort_by_label" | "sort_by_label_desc" | "timestamp" => {
3902 for value in &self.ctx.field_columns {
3905 let expr = DfExpr::Column(Column::from_name(value));
3906 exprs.push(expr);
3907 }
3908
3909 ScalarFunc::GeneratedExpr
3910 }
3911 "round" if self.all_field_columns_are_native_histograms(input_schema) => {
3912 if other_input_exprs.is_empty() {
3913 other_input_exprs.push_front(0.0f64.lit());
3914 }
3915 ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3916 }
3917 "round" => {
3918 if other_input_exprs.is_empty() {
3919 other_input_exprs.push_front(0.0f64.lit());
3920 }
3921 ScalarFunc::DataFusionUdf(Arc::new(Round::scalar_udf()))
3922 }
3923 "rad" | "deg" | "sgn" if self.all_field_columns_are_native_histograms(input_schema) => {
3924 ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3925 }
3926 "rad" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::radians()),
3927 "deg" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::degrees()),
3928 "sgn" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::signum()),
3929 "pi" => {
3930 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
3932 func: datafusion::functions::math::pi(),
3933 args: vec![],
3934 });
3935 exprs.push(fn_expr);
3936
3937 ScalarFunc::GeneratedExpr
3938 }
3939 _ => {
3940 if let Some(f) = query_engine_state
3941 .session_state()
3942 .scalar_functions()
3943 .get(func.name)
3944 {
3945 if self.all_field_columns_are_native_histograms(input_schema) {
3946 ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3947 } else {
3948 ScalarFunc::DataFusionBuiltin(f.clone())
3949 }
3950 } else if let Some(factory) = query_engine_state.scalar_function(func.name) {
3951 if self.all_field_columns_are_native_histograms(input_schema) {
3952 ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3953 } else {
3954 let func_state = query_engine_state.function_state();
3955 let query_ctx = self.table_provider.query_ctx();
3956
3957 ScalarFunc::DataFusionUdf(Arc::new(factory.provide(FunctionContext {
3958 state: func_state,
3959 query_ctx: query_ctx.clone(),
3960 })))
3961 }
3962 } else if let Some(f) = datafusion_functions::math::functions()
3963 .iter()
3964 .find(|f| f.name() == func.name)
3965 {
3966 if self.all_field_columns_are_native_histograms(input_schema) {
3967 ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3968 } else {
3969 ScalarFunc::DataFusionUdf(f.clone())
3970 }
3971 } else {
3972 return UnsupportedExprSnafu {
3973 name: func.name.to_string(),
3974 }
3975 .fail();
3976 }
3977 }
3978 };
3979
3980 for value in &self.ctx.field_columns {
3981 let col_expr = DfExpr::Column(Column::from_name(value));
3982 let value_is_histogram = Self::field_column_is_native_histogram(input_schema, value);
3983
3984 match scalar_func.clone() {
3985 ScalarFunc::DataFusionBuiltin(func) => {
3986 if alternative_samples && value_is_histogram {
3987 continue;
3988 }
3989 other_input_exprs.insert(field_column_pos, col_expr);
3990 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
3991 func,
3992 args: other_input_exprs.clone().into(),
3993 });
3994 exprs.push(fn_expr);
3995 let _ = other_input_exprs.remove(field_column_pos);
3996 }
3997 ScalarFunc::DataFusionUdf(func) => {
3998 if alternative_samples && value_is_histogram {
3999 continue;
4000 }
4001 let args = itertools::chain!(
4002 other_input_exprs.iter().take(field_column_pos).cloned(),
4003 std::iter::once(col_expr),
4004 other_input_exprs.iter().skip(field_column_pos).cloned()
4005 )
4006 .collect_vec();
4007 exprs.push(DfExpr::ScalarFunction(ScalarFunction { func, args }))
4008 }
4009 ScalarFunc::NativeHistogramUdf(func) => {
4010 if value_is_histogram {
4011 let args = itertools::chain!(
4012 other_input_exprs.iter().take(field_column_pos).cloned(),
4013 std::iter::once(col_expr),
4014 other_input_exprs.iter().skip(field_column_pos).cloned()
4015 )
4016 .collect_vec();
4017 exprs.push(DfExpr::ScalarFunction(ScalarFunction { func, args }));
4018 } else if !alternative_samples {
4019 exprs.push(
4020 DfExpr::Literal(ScalarValue::Float64(None), None).alias(format!(
4021 "{}_{}",
4022 func.name(),
4023 value
4024 )),
4025 );
4026 }
4027 }
4028 ScalarFunc::Udf(func) => {
4029 let ts_range_expr = DfExpr::Column(Column::from_name(
4030 RangeManipulate::build_timestamp_range_name(
4031 self.ctx.time_index_column.as_ref().unwrap(),
4032 ),
4033 ));
4034 other_input_exprs.insert(field_column_pos, ts_range_expr);
4035 other_input_exprs.insert(field_column_pos + 1, col_expr);
4036 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
4037 func,
4038 args: other_input_exprs.clone().into(),
4039 });
4040 exprs.push(fn_expr);
4041 let _ = other_input_exprs.remove(field_column_pos + 1);
4042 let _ = other_input_exprs.remove(field_column_pos);
4043 }
4044 ScalarFunc::ExtrapolateUdf(udf, range_length) => {
4045 let ts_range_expr = DfExpr::Column(Column::from_name(
4046 RangeManipulate::build_timestamp_range_name(
4047 self.ctx.time_index_column.as_ref().unwrap(),
4048 ),
4049 ));
4050 other_input_exprs.insert(field_column_pos, ts_range_expr.clone());
4051 other_input_exprs.insert(field_column_pos + 1, col_expr.clone());
4052 other_input_exprs
4053 .insert(field_column_pos + 2, self.create_time_index_column_expr()?);
4054 other_input_exprs.push_back(lit(range_length));
4055 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
4056 func: udf,
4057 args: other_input_exprs.clone().into(),
4058 });
4059 let fn_expr = if matches!(func.name, "rate" | "increase")
4060 && !all_field_columns_are_native_histogram_ranges
4061 {
4062 let delta_sum = DfExpr::ScalarFunction(ScalarFunction {
4063 func: Arc::new(SumOverTime::scalar_udf()),
4064 args: vec![ts_range_expr, col_expr],
4065 });
4066 self.select_delta_range_math(
4067 func.name,
4068 input_schema,
4069 range_length,
4070 delta_sum,
4071 fn_expr,
4072 )?
4073 } else {
4074 fn_expr
4075 };
4076 exprs.push(fn_expr);
4077 let _ = other_input_exprs.pop_back();
4078 let _ = other_input_exprs.remove(field_column_pos + 2);
4079 let _ = other_input_exprs.remove(field_column_pos + 1);
4080 let _ = other_input_exprs.remove(field_column_pos);
4081 }
4082 ScalarFunc::GeneratedExpr => {}
4083 }
4084 }
4085
4086 if !matches!(func.name, "label_join" | "label_replace") {
4090 let mut new_field_columns = Vec::with_capacity(exprs.len());
4091
4092 exprs = exprs
4093 .into_iter()
4094 .map(|expr| {
4095 let display_name = expr.schema_name().to_string();
4096 new_field_columns.push(display_name.clone());
4097 Ok(expr.alias(display_name))
4098 })
4099 .collect::<std::result::Result<Vec<_>, _>>()
4100 .context(DataFusionPlanningSnafu)?;
4101
4102 self.ctx.field_columns = new_field_columns;
4103 }
4104
4105 Ok((exprs, new_tags))
4106 }
4107
4108 fn select_delta_range_math(
4109 &self,
4110 function: &str,
4111 input_schema: &DFSchemaRef,
4112 range_length: Millisecond,
4113 delta_sum: DfExpr,
4114 cumulative: DfExpr,
4115 ) -> Result<DfExpr> {
4116 let marker_is_delta = if self
4117 .ctx
4118 .tag_columns
4119 .iter()
4120 .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL)
4121 {
4122 Self::field_column_type(input_schema, OTLP_AGGREGATION_TEMPORALITY_LABEL)
4123 .filter(|data_type| Self::string_value_data_type(data_type).is_some())
4124 .map(|_| {
4125 DfExpr::Column(Column::from_name(OTLP_AGGREGATION_TEMPORALITY_LABEL))
4126 .eq(lit(GREPTIME_TEMPORALITY_DELTA))
4127 })
4128 } else {
4129 None
4130 };
4131 let Some(marker_is_delta) = marker_is_delta else {
4132 return Ok(cumulative);
4133 };
4134
4135 let delta = if function == "rate" {
4136 DfExpr::BinaryExpr(BinaryExpr {
4137 left: Box::new(delta_sum),
4138 op: Operator::Divide,
4139 right: Box::new(lit(range_length as f64 / 1000.0)),
4140 })
4141 } else {
4142 delta_sum
4143 };
4144 let display_name = cumulative.schema_name().to_string();
4145 when(marker_is_delta, delta)
4146 .otherwise(cumulative)
4147 .context(DataFusionPlanningSnafu)
4148 .map(|expr| expr.alias(display_name))
4149 }
4150
4151 fn validate_label_name(label_name: &str) -> Result<()> {
4155 if label_name.starts_with("__") {
4157 return InvalidDestinationLabelNameSnafu { label_name }.fail();
4158 }
4159 if !LABEL_NAME_REGEX.is_match(label_name) {
4161 return InvalidDestinationLabelNameSnafu { label_name }.fail();
4162 }
4163
4164 Ok(())
4165 }
4166
4167 fn build_regexp_replace_label_expr(
4169 &self,
4170 other_input_exprs: &mut VecDeque<DfExpr>,
4171 query_engine_state: &QueryEngineState,
4172 ) -> Result<Option<(DfExpr, String)>> {
4173 let dst_label = match other_input_exprs.pop_front() {
4175 Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4176 other => UnexpectedPlanExprSnafu {
4177 desc: format!("expected dst_label string literal, but found {:?}", other),
4178 }
4179 .fail()?,
4180 };
4181
4182 Self::validate_label_name(&dst_label)?;
4184 let replacement = match other_input_exprs.pop_front() {
4185 Some(DfExpr::Literal(ScalarValue::Utf8(Some(r)), _)) => r,
4186 other => UnexpectedPlanExprSnafu {
4187 desc: format!("expected replacement string literal, but found {:?}", other),
4188 }
4189 .fail()?,
4190 };
4191 let src_label = match other_input_exprs.pop_front() {
4192 Some(DfExpr::Literal(ScalarValue::Utf8(Some(s)), None)) => s,
4193 other => UnexpectedPlanExprSnafu {
4194 desc: format!("expected src_label string literal, but found {:?}", other),
4195 }
4196 .fail()?,
4197 };
4198
4199 let regex = match other_input_exprs.pop_front() {
4200 Some(DfExpr::Literal(ScalarValue::Utf8(Some(r)), None)) => r,
4201 other => UnexpectedPlanExprSnafu {
4202 desc: format!("expected regex string literal, but found {:?}", other),
4203 }
4204 .fail()?,
4205 };
4206
4207 regex::Regex::new(®ex).map_err(|_| {
4210 InvalidRegularExpressionSnafu {
4211 regex: regex.clone(),
4212 }
4213 .build()
4214 })?;
4215
4216 if self.ctx.tag_columns.contains(&src_label) && regex.is_empty() {
4218 return Ok(None);
4219 }
4220
4221 if !self.ctx.tag_columns.contains(&src_label) {
4223 if replacement.is_empty() {
4224 return Ok(None);
4226 } else {
4227 return Ok(Some((
4229 lit(replacement).alias(&dst_label),
4231 dst_label,
4232 )));
4233 }
4234 }
4235
4236 let regex = format!("^(?s:{regex})$");
4239
4240 let session_state = query_engine_state.session_state();
4241 let func = session_state
4242 .scalar_functions()
4243 .get("regexp_replace")
4244 .context(UnsupportedExprSnafu {
4245 name: "regexp_replace",
4246 })?;
4247
4248 let args = vec![
4250 if src_label.is_empty() {
4251 DfExpr::Literal(ScalarValue::Utf8(Some(String::new())), None)
4252 } else {
4253 DfExpr::Column(Column::from_name(src_label))
4254 },
4255 DfExpr::Literal(ScalarValue::Utf8(Some(regex)), None),
4256 DfExpr::Literal(ScalarValue::Utf8(Some(replacement)), None),
4257 ];
4258
4259 Ok(Some((
4260 DfExpr::ScalarFunction(ScalarFunction {
4261 func: func.clone(),
4262 args,
4263 })
4264 .alias(&dst_label),
4265 dst_label,
4266 )))
4267 }
4268
4269 fn build_concat_labels_expr(
4271 other_input_exprs: &mut VecDeque<DfExpr>,
4272 ctx: &PromPlannerContext,
4273 query_engine_state: &QueryEngineState,
4274 ) -> Result<(DfExpr, String)> {
4275 let dst_label = match other_input_exprs.pop_front() {
4278 Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4279 other => UnexpectedPlanExprSnafu {
4280 desc: format!("expected dst_label string literal, but found {:?}", other),
4281 }
4282 .fail()?,
4283 };
4284 let separator = match other_input_exprs.pop_front() {
4285 Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4286 other => UnexpectedPlanExprSnafu {
4287 desc: format!("expected separator string literal, but found {:?}", other),
4288 }
4289 .fail()?,
4290 };
4291
4292 let available_columns: HashSet<&str> = ctx
4294 .tag_columns
4295 .iter()
4296 .chain(ctx.field_columns.iter())
4297 .chain(ctx.time_index_column.as_ref())
4298 .map(|s| s.as_str())
4299 .collect();
4300
4301 let src_labels = other_input_exprs
4302 .iter()
4303 .map(|expr| {
4304 match expr {
4306 DfExpr::Literal(ScalarValue::Utf8(Some(label)), None) => {
4307 if label.is_empty() {
4308 Ok(DfExpr::Literal(ScalarValue::Null, None))
4309 } else if available_columns.contains(label.as_str()) {
4310 Ok(DfExpr::Column(Column::from_name(label)))
4312 } else {
4313 Ok(DfExpr::Literal(ScalarValue::Null, None))
4315 }
4316 }
4317 other => UnexpectedPlanExprSnafu {
4318 desc: format!(
4319 "expected source label string literal, but found {:?}",
4320 other
4321 ),
4322 }
4323 .fail(),
4324 }
4325 })
4326 .collect::<Result<Vec<_>>>()?;
4327 ensure!(
4328 !src_labels.is_empty(),
4329 FunctionInvalidArgumentSnafu {
4330 fn_name: "label_join"
4331 }
4332 );
4333
4334 let session_state = query_engine_state.session_state();
4335 let func = session_state
4336 .scalar_functions()
4337 .get("concat_ws")
4338 .context(UnsupportedExprSnafu { name: "concat_ws" })?;
4339
4340 let mut args = Vec::with_capacity(1 + src_labels.len());
4342 args.push(DfExpr::Literal(ScalarValue::Utf8(Some(separator)), None));
4343 args.extend(src_labels);
4344
4345 Ok((
4346 DfExpr::ScalarFunction(ScalarFunction {
4347 func: func.clone(),
4348 args,
4349 })
4350 .alias(&dst_label),
4351 dst_label,
4352 ))
4353 }
4354
4355 fn create_time_index_column_expr(&self) -> Result<DfExpr> {
4356 Ok(DfExpr::Column(Column::from_name(
4357 self.ctx
4358 .time_index_column
4359 .clone()
4360 .with_context(|| TimeIndexNotFoundSnafu { table: "unknown" })?,
4361 )))
4362 }
4363
4364 fn create_tag_column_exprs(&self) -> Result<Vec<DfExpr>> {
4365 let mut result = Vec::with_capacity(self.ctx.tag_columns.len());
4366 for tag in &self.ctx.tag_columns {
4367 let expr = DfExpr::Column(Column::from_name(tag));
4368 result.push(expr);
4369 }
4370 Ok(result)
4371 }
4372
4373 fn create_field_column_exprs(&self) -> Result<Vec<DfExpr>> {
4374 let mut result = Vec::with_capacity(self.ctx.field_columns.len());
4375 for field in &self.ctx.field_columns {
4376 let expr = DfExpr::Column(Column::from_name(field));
4377 result.push(expr);
4378 }
4379 Ok(result)
4380 }
4381
4382 fn create_tag_and_time_index_column_sort_exprs(&self) -> Result<Vec<SortExpr>> {
4383 let mut result = self
4384 .ctx
4385 .tag_columns
4386 .iter()
4387 .map(|col| DfExpr::Column(Column::from_name(col)).sort(true, true))
4388 .collect::<Vec<_>>();
4389 result.push(self.create_time_index_column_expr()?.sort(true, true));
4390 Ok(result)
4391 }
4392
4393 fn create_field_columns_sort_exprs(&self, asc: bool) -> Vec<SortExpr> {
4394 self.ctx
4395 .field_columns
4396 .iter()
4397 .map(|col| DfExpr::Column(Column::from_name(col)).sort(asc, true))
4398 .collect::<Vec<_>>()
4399 }
4400
4401 fn create_sort_exprs_by_tags(
4402 func: &str,
4403 tags: Vec<DfExpr>,
4404 asc: bool,
4405 ) -> Result<Vec<SortExpr>> {
4406 ensure!(
4407 !tags.is_empty(),
4408 FunctionInvalidArgumentSnafu { fn_name: func }
4409 );
4410
4411 tags.iter()
4412 .map(|col| match col {
4413 DfExpr::Literal(ScalarValue::Utf8(Some(label)), _) => {
4414 Ok(DfExpr::Column(Column::from_name(label)).sort(asc, false))
4415 }
4416 other => UnexpectedPlanExprSnafu {
4417 desc: format!("expected label string literal, but found {:?}", other),
4418 }
4419 .fail(),
4420 })
4421 .collect::<Result<Vec<_>>>()
4422 }
4423
4424 fn create_empty_values_filter_expr(&self, preserve_any_value: bool) -> Result<DfExpr> {
4425 let mut exprs = Vec::with_capacity(self.ctx.field_columns.len());
4426 for value in &self.ctx.field_columns {
4427 let expr = DfExpr::Column(Column::from_name(value)).is_not_null();
4428 exprs.push(expr);
4429 }
4430
4431 let predicate = if preserve_any_value {
4436 disjunction(exprs)
4437 } else {
4438 conjunction(exprs)
4439 };
4440 predicate.with_context(|| ValueNotFoundSnafu {
4441 table: self
4442 .table_ref()
4443 .map(|t| t.to_quoted_string())
4444 .unwrap_or_else(|_| "unknown".to_string()),
4445 })
4446 }
4447
4448 fn create_aggregate_exprs(
4464 &mut self,
4465 op: TokenType,
4466 param: &Option<Box<PromExpr>>,
4467 input_plan: &LogicalPlan,
4468 ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4469 let mixed_sample_columns =
4470 Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns)
4471 .map(|(float, histogram)| (float.to_string(), histogram.to_string()));
4472 let is_group_agg = op.id() == token::T_GROUP;
4473 if is_group_agg && mixed_sample_columns.is_none() {
4474 ensure!(
4475 self.ctx.field_columns.len() == 1,
4476 MultiFieldsNotSupportedSnafu {
4477 operator: "group()"
4478 }
4479 );
4480 }
4481
4482 if let Some((float, histogram)) = mixed_sample_columns {
4483 return self.create_mixed_aggregate_exprs(op, param, &float, &histogram);
4484 }
4485
4486 if self.all_field_columns_are_native_histograms(input_plan.schema()) {
4487 return self.create_native_histogram_aggregate_exprs(op, input_plan);
4488 }
4489
4490 let exprs = self
4492 .ctx
4493 .field_columns
4494 .iter()
4495 .map(|col| {
4496 Self::create_numeric_aggregate_expr(
4497 op,
4498 param,
4499 DfExpr::Column(Column::from_name(col)),
4500 )
4501 })
4502 .collect::<Result<Vec<_>>>()?;
4503
4504 let prev_field_exprs = if op.id() == token::T_COUNT_VALUES {
4506 let prev_field_exprs: Vec<_> = self
4507 .ctx
4508 .field_columns
4509 .iter()
4510 .map(|col| DfExpr::Column(Column::from_name(col)))
4511 .collect();
4512
4513 ensure!(
4514 self.ctx.field_columns.len() == 1,
4515 UnsupportedExprSnafu {
4516 name: "count_values on multi-value input"
4517 }
4518 );
4519
4520 prev_field_exprs
4521 } else {
4522 vec![]
4523 };
4524
4525 let mut new_field_columns = Vec::with_capacity(self.ctx.field_columns.len());
4527
4528 let normalized_exprs =
4529 normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
4530 for expr in normalized_exprs {
4531 new_field_columns.push(expr.schema_name().to_string());
4532 }
4533 self.ctx.field_columns = new_field_columns;
4534
4535 Ok((exprs, prev_field_exprs))
4536 }
4537
4538 fn create_numeric_aggregate_expr(
4539 op: TokenType,
4540 param: &Option<Box<PromExpr>>,
4541 input: DfExpr,
4542 ) -> Result<DfExpr> {
4543 let expr = match op.id() {
4544 token::T_SUM => sum_udaf().call(vec![input]),
4545 token::T_QUANTILE => {
4546 let q = Self::get_param_as_literal_expr(
4547 param.as_deref(),
4548 Some(op),
4549 Some(ArrowDataType::Float64),
4550 )?;
4551 quantile_udaf().call(vec![q, input])
4552 }
4553 token::T_AVG => avg_udaf().call(vec![input]),
4554 token::T_COUNT_VALUES | token::T_COUNT => count_udaf().call(vec![input]),
4555 token::T_MIN => min_udaf().call(vec![input]),
4556 token::T_MAX => max_udaf().call(vec![input]),
4557 token::T_GROUP => max_udaf().call(vec![lit(1_f64)]),
4560 token::T_STDDEV => stddev_pop_udaf().call(vec![input]),
4561 token::T_STDVAR => var_pop_udaf().call(vec![input]),
4562 token::T_TOPK | token::T_BOTTOMK => {
4563 return UnsupportedExprSnafu {
4564 name: format!("{op:?}"),
4565 }
4566 .fail();
4567 }
4568 _ => return UnexpectedTokenSnafu { token: op }.fail(),
4569 };
4570 Ok(expr)
4571 }
4572
4573 fn create_mixed_aggregate_exprs(
4574 &mut self,
4575 op: TokenType,
4576 param: &Option<Box<PromExpr>>,
4577 float_column: &str,
4578 histogram_column: &str,
4579 ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4580 let float_input = DfExpr::Column(Column::from_name(float_column));
4581 let histogram_input = DfExpr::Column(Column::from_name(histogram_column));
4582 let float_count = count_udaf().call(vec![float_input.clone()]);
4583 let histogram_count = count_udaf().call(vec![histogram_input.clone()]);
4584 let mixed_sample_value = || {
4585 DfExpr::ScalarFunction(ScalarFunction {
4586 func: coalesce(),
4587 args: vec![
4588 DfExpr::ScalarFunction(ScalarFunction {
4589 func: Arc::new(PromqlFloatToString::scalar_udf()),
4590 args: vec![float_input.clone()],
4591 }),
4592 DfExpr::ScalarFunction(ScalarFunction {
4593 func: Arc::new(NativeHistogramToString::scalar_udf()),
4594 args: vec![histogram_input.clone()],
4595 }),
4596 ],
4597 })
4598 };
4599
4600 let (exprs, prev_field_exprs, field_columns) = match op.id() {
4601 token::T_SUM | token::T_AVG => (
4602 vec![
4603 Self::create_numeric_aggregate_expr(op, param, float_input)?
4604 .alias(float_column),
4605 self.create_native_histogram_aggregate_expr(op, histogram_column)?,
4606 float_count.alias(Self::mixed_sample_count_name(float_column)),
4607 histogram_count.alias(Self::mixed_sample_count_name(histogram_column)),
4608 ],
4609 vec![],
4610 vec![float_column.to_string(), histogram_column.to_string()],
4611 ),
4612 token::T_COUNT => {
4613 let present = when(
4614 float_input
4615 .clone()
4616 .is_not_null()
4617 .or(histogram_input.clone().is_not_null()),
4618 lit(1_i64),
4619 )
4620 .otherwise(lit(ScalarValue::Int64(None)))
4621 .context(DataFusionPlanningSnafu)?;
4622 (
4623 vec![count_udaf().call(vec![present]).alias(float_column)],
4624 vec![],
4625 vec![float_column.to_string()],
4626 )
4627 }
4628 token::T_GROUP => (
4629 vec![max_udaf().call(vec![lit(1_f64)]).alias(float_column)],
4630 vec![],
4631 vec![float_column.to_string()],
4632 ),
4633 token::T_COUNT_VALUES => {
4634 let value = mixed_sample_value();
4635 (
4636 vec![count_udaf().call(vec![value.clone()]).alias(float_column)],
4637 vec![value],
4638 vec![float_column.to_string()],
4639 )
4640 }
4641 token::T_MIN | token::T_MAX | token::T_STDDEV | token::T_STDVAR | token::T_QUANTILE => {
4642 (
4643 vec![
4644 Self::create_numeric_aggregate_expr(op, param, float_input)?
4645 .alias(float_column),
4646 histogram_count.alias(Self::mixed_sample_count_name(histogram_column)),
4647 ],
4648 vec![],
4649 vec![float_column.to_string()],
4650 )
4651 }
4652 token::T_TOPK | token::T_BOTTOMK => {
4653 return UnsupportedExprSnafu {
4654 name: format!("{op:?}"),
4655 }
4656 .fail();
4657 }
4658 _ => return UnexpectedTokenSnafu { token: op }.fail(),
4659 };
4660
4661 self.ctx.field_columns = field_columns;
4662 Ok((exprs, prev_field_exprs))
4663 }
4664
4665 fn mixed_sample_count_column(column: &str) -> DfExpr {
4666 DfExpr::Column(Column::from_name(Self::mixed_sample_count_name(column)))
4667 }
4668
4669 fn mixed_sample_count_name(column: &str) -> String {
4670 format!("__promql_sample_count({column})")
4671 }
4672
4673 fn mixed_aggregate_filter_expr(
4674 &self,
4675 op: TokenType,
4676 float_column: &str,
4677 histogram_column: &str,
4678 ) -> Result<DfExpr> {
4679 let float_count = Self::mixed_sample_count_column(float_column);
4680 let histogram_count = Self::mixed_sample_count_column(histogram_column);
4681 let mixed = float_count
4682 .clone()
4683 .gt(lit(0_i64))
4684 .and(histogram_count.clone().gt(lit(0_i64)));
4685 let drop_mixed = DfExpr::ScalarFunction(ScalarFunction {
4686 func: Arc::new(NativeHistogramDrop::warning_bool_false_udf(
4687 format!(
4688 "{op}: dropped aggregation result containing both float and native histogram samples"
4689 ),
4690 self.promql_annotations.clone(),
4691 )),
4692 args: vec![float_count, histogram_count],
4693 });
4694
4695 when(mixed, drop_mixed)
4696 .otherwise(lit(true))
4697 .context(DataFusionPlanningSnafu)
4698 }
4699
4700 fn mixed_ignored_histogram_filter_expr(
4701 &self,
4702 op: TokenType,
4703 histogram_column: &str,
4704 ) -> Result<DfExpr> {
4705 let histogram_count = Self::mixed_sample_count_column(histogram_column);
4706 let has_histograms = histogram_count.clone().gt(lit(0_i64));
4707 let record_info = DfExpr::ScalarFunction(ScalarFunction {
4708 func: Arc::new(NativeHistogramDrop::bool_true_udf(
4709 format!(
4710 "{op}: dropped native histogram samples because this aggregation is not supported for native histograms"
4711 ),
4712 self.promql_annotations.clone(),
4713 )),
4714 args: vec![histogram_count],
4715 });
4716
4717 when(has_histograms, record_info)
4718 .otherwise(lit(true))
4719 .context(DataFusionPlanningSnafu)
4720 }
4721
4722 fn create_native_histogram_aggregate_expr(
4723 &self,
4724 op: TokenType,
4725 column: &str,
4726 ) -> Result<DfExpr> {
4727 let input = DfExpr::Column(Column::from_name(column));
4728 let expr = match op.id() {
4729 token::T_SUM => Arc::new(NativeHistogramAggSum::aggregate_udf_with_collector(
4730 self.promql_annotations.clone(),
4731 ))
4732 .call(vec![input])
4733 .alias(column),
4734 token::T_AVG => Arc::new(NativeHistogramAggAvg::aggregate_udf_with_collector(
4735 self.promql_annotations.clone(),
4736 ))
4737 .call(vec![input])
4738 .alias(column),
4739 token::T_COUNT_VALUES | token::T_COUNT => {
4740 count_udaf().call(vec![input]).alias(column)
4741 }
4742 token::T_GROUP => max_udaf().call(vec![lit(1_f64)]).alias(column),
4743 token::T_MIN
4744 | token::T_MAX
4745 | token::T_STDDEV
4746 | token::T_STDVAR
4747 | token::T_QUANTILE
4748 | token::T_TOPK
4749 | token::T_BOTTOMK => sum_udaf()
4750 .call(vec![DfExpr::ScalarFunction(ScalarFunction {
4751 func: Arc::new(NativeHistogramDrop::float_null_udf(
4752 format!(
4753 "{op}: dropped native histogram samples because this aggregation is not supported for native histograms"
4754 ),
4755 self.promql_annotations.clone(),
4756 )),
4757 args: vec![input],
4758 })])
4759 .alias(column),
4760 _ => return UnexpectedTokenSnafu { token: op }.fail(),
4761 };
4762 Ok(expr)
4763 }
4764
4765 fn create_native_histogram_aggregate_exprs(
4766 &mut self,
4767 op: TokenType,
4768 input_plan: &LogicalPlan,
4769 ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4770 let prev_field_exprs = if op.id() == token::T_COUNT_VALUES {
4771 ensure!(
4772 self.ctx.field_columns.len() == 1,
4773 UnsupportedExprSnafu {
4774 name: "count_values on multi-value input"
4775 }
4776 );
4777 self.ctx
4778 .field_columns
4779 .iter()
4780 .map(|col| {
4781 DfExpr::ScalarFunction(ScalarFunction {
4782 func: Arc::new(NativeHistogramToString::scalar_udf()),
4783 args: vec![DfExpr::Column(Column::from_name(col))],
4784 })
4785 })
4786 .collect::<Vec<_>>()
4787 } else {
4788 vec![]
4789 };
4790
4791 let exprs = self
4792 .ctx
4793 .field_columns
4794 .iter()
4795 .map(|col| self.create_native_histogram_aggregate_expr(op, col))
4796 .collect::<Result<Vec<_>>>()?;
4797
4798 let normalized_exprs =
4799 normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
4800 self.ctx.field_columns = normalized_exprs
4801 .into_iter()
4802 .map(|expr| expr.schema_name().to_string())
4803 .collect();
4804
4805 Ok((exprs, prev_field_exprs))
4806 }
4807
4808 fn get_param_value_as_str(op: TokenType, param: &Option<Box<PromExpr>>) -> Result<&str> {
4809 let param = param
4810 .as_deref()
4811 .with_context(|| FunctionInvalidArgumentSnafu {
4812 fn_name: op.to_string(),
4813 })?;
4814 let PromExpr::StringLiteral(StringLiteral { val }) = param else {
4815 return FunctionInvalidArgumentSnafu {
4816 fn_name: op.to_string(),
4817 }
4818 .fail();
4819 };
4820
4821 Ok(val)
4822 }
4823
4824 fn get_param_as_literal_expr(
4825 param: Option<&PromExpr>,
4826 op: Option<TokenType>,
4827 expected_type: Option<ArrowDataType>,
4828 ) -> Result<DfExpr> {
4829 let prom_param = param.with_context(|| {
4830 if let Some(op) = op {
4831 FunctionInvalidArgumentSnafu {
4832 fn_name: op.to_string(),
4833 }
4834 } else {
4835 FunctionInvalidArgumentSnafu {
4836 fn_name: "unknown".to_string(),
4837 }
4838 }
4839 })?;
4840
4841 let expr = Self::try_build_literal_expr(prom_param).with_context(|| {
4842 if let Some(op) = op {
4843 FunctionInvalidArgumentSnafu {
4844 fn_name: op.to_string(),
4845 }
4846 } else {
4847 FunctionInvalidArgumentSnafu {
4848 fn_name: "unknown".to_string(),
4849 }
4850 }
4851 })?;
4852
4853 if let Some(expected_type) = expected_type {
4855 let expr_type = expr
4857 .get_type(&DFSchema::empty())
4858 .context(DataFusionPlanningSnafu)?;
4859 if expected_type != expr_type {
4860 return FunctionInvalidArgumentSnafu {
4861 fn_name: format!("expected {expected_type:?}, but found {expr_type:?}"),
4862 }
4863 .fail();
4864 }
4865 }
4866
4867 Ok(expr)
4868 }
4869
4870 fn create_window_exprs(
4873 &mut self,
4874 op: TokenType,
4875 group_exprs: Vec<DfExpr>,
4876 input_plan: &LogicalPlan,
4877 ) -> Result<Vec<DfExpr>> {
4878 ensure!(
4879 self.ctx.field_columns.len() == 1,
4880 UnsupportedExprSnafu {
4881 name: "topk or bottomk on multi-value input"
4882 }
4883 );
4884
4885 assert!(matches!(op.id(), token::T_TOPK | token::T_BOTTOMK));
4886
4887 let asc = matches!(op.id(), token::T_BOTTOMK);
4888
4889 let tag_sort_exprs = self
4890 .create_tag_column_exprs()?
4891 .into_iter()
4892 .map(|expr| expr.sort(asc, true));
4893
4894 let exprs: Vec<DfExpr> = self
4896 .ctx
4897 .field_columns
4898 .iter()
4899 .map(|col| {
4900 let mut sort_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 1);
4901 sort_exprs.push(DfExpr::Column(Column::from(col)).sort(asc, true));
4903 sort_exprs.extend(tag_sort_exprs.clone());
4906
4907 DfExpr::WindowFunction(Box::new(WindowFunction {
4908 fun: WindowFunctionDefinition::WindowUDF(Arc::new(RowNumber::new().into())),
4909 params: WindowFunctionParams {
4910 args: vec![],
4911 partition_by: group_exprs.clone(),
4912 order_by: sort_exprs,
4913 window_frame: WindowFrame::new(Some(true)),
4914 null_treatment: None,
4915 distinct: false,
4916 filter: None,
4917 },
4918 }))
4919 })
4920 .collect();
4921
4922 let normalized_exprs =
4923 normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
4924 Ok(normalized_exprs)
4925 }
4926
4927 async fn create_histogram_plan(
4929 &mut self,
4930 function_name: &str,
4931 args: &PromFunctionArgs,
4932 query_engine_state: &QueryEngineState,
4933 ) -> Result<LogicalPlan> {
4934 let float_literal = |param: &PromExpr| -> Result<f64> {
4935 let value = (|| {
4936 let expr = Self::get_param_as_literal_expr(
4937 Some(param),
4938 None,
4939 Some(ArrowDataType::Float64),
4940 )
4941 .ok()?;
4942 let simplifier = ExprSimplifier::new(SimplifyContext::default());
4943 let expr = simplifier.coerce(expr, &DFSchema::empty()).ok()?;
4944 let DfExpr::Literal(value, _) = simplifier.simplify(expr).ok()? else {
4945 return None;
4946 };
4947 let ScalarValue::Float64(Some(value)) =
4948 value.cast_to(&ArrowDataType::Float64).ok()?
4949 else {
4950 return None;
4951 };
4952 Some(value)
4953 })()
4954 .with_context(|| FunctionInvalidArgumentSnafu {
4955 fn_name: function_name.to_string(),
4956 })?;
4957 Ok(value)
4958 };
4959 let (function, input) = match (function_name, args.args.as_slice()) {
4960 (SPECIAL_HISTOGRAM_QUANTILE, [quantile, input]) => (
4961 HistogramFoldOperation::Quantile(float_literal(quantile)?.into()),
4962 input.as_ref().clone(),
4963 ),
4964 (SPECIAL_HISTOGRAM_FRACTION, [lower, upper, input]) => (
4965 HistogramFoldOperation::Fraction {
4966 lower: float_literal(lower)?.into(),
4967 upper: float_literal(upper)?.into(),
4968 },
4969 input.as_ref().clone(),
4970 ),
4971 _ => {
4972 return FunctionInvalidArgumentSnafu {
4973 fn_name: function_name.to_string(),
4974 }
4975 .fail();
4976 }
4977 };
4978
4979 let input_plan = self.prom_expr_to_plan(&input, query_engine_state).await?;
4980 let input_plan = self.strip_tsid_column(input_plan)?;
4983 self.ctx.use_tsid = false;
4984
4985 if let Some((float_field, histogram_field)) =
4986 Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns)
4987 .map(|(float, histogram)| (float.to_string(), histogram.to_string()))
4988 {
4989 if self.ctx.has_le_tag() {
4990 return self.create_mixed_histogram_plan(
4991 function,
4992 input_plan,
4993 float_field,
4994 histogram_field,
4995 );
4996 }
4997 self.ctx.field_columns = vec![histogram_field];
4998 }
4999 if self.all_field_columns_are_native_histograms(input_plan.schema()) {
5000 return self.create_native_histogram_plan(function, input_plan);
5001 }
5002
5003 if !self.ctx.has_le_tag() {
5004 return Ok(LogicalPlan::EmptyRelation(
5007 datafusion::logical_expr::EmptyRelation {
5008 produce_one_row: false,
5009 schema: input_plan.schema().clone(),
5010 },
5011 ));
5012 }
5013 let time_index_column =
5014 self.ctx
5015 .time_index_column
5016 .clone()
5017 .with_context(|| TimeIndexNotFoundSnafu {
5018 table: self.ctx.table_name.clone().unwrap_or_default(),
5019 })?;
5020 let field_column = self
5022 .ctx
5023 .field_columns
5024 .first()
5025 .with_context(|| FunctionInvalidArgumentSnafu {
5026 fn_name: function.function_name().to_string(),
5027 })?
5028 .clone();
5029 self.ctx.tag_columns.retain(|col| col != LE_COLUMN_NAME);
5031
5032 let fold = HistogramFold::new_with_operation(
5033 LE_COLUMN_NAME.to_string(),
5034 field_column,
5035 time_index_column,
5036 function,
5037 None,
5038 input_plan,
5039 )
5040 .context(DataFusionPlanningSnafu)?;
5041 Ok(LogicalPlan::Extension(Extension {
5042 node: Arc::new(fold),
5043 }))
5044 }
5045
5046 fn create_native_histogram_expr(
5047 &self,
5048 function: HistogramFoldOperation,
5049 field_column: &str,
5050 ) -> DfExpr {
5051 let field = DfExpr::Column(Column::from_name(field_column));
5052 let (func, args) = match function {
5053 HistogramFoldOperation::Quantile(quantile) => (
5054 Arc::new(NativeHistogramQuantile::scalar_udf_with_collector(
5055 self.promql_annotations.clone(),
5056 )),
5057 vec![field, lit(f64::from(quantile))],
5058 ),
5059 HistogramFoldOperation::Fraction { lower, upper } => (
5060 Arc::new(NativeHistogramFraction::scalar_udf_with_collector(
5061 self.promql_annotations.clone(),
5062 )),
5063 vec![field, lit(f64::from(lower)), lit(f64::from(upper))],
5064 ),
5065 };
5066 DfExpr::ScalarFunction(ScalarFunction { func, args })
5067 }
5068
5069 fn create_native_histogram_plan(
5070 &mut self,
5071 function: HistogramFoldOperation,
5072 input_plan: LogicalPlan,
5073 ) -> Result<LogicalPlan> {
5074 ensure!(
5075 self.ctx.field_columns.len() == 1,
5076 MultiFieldsNotSupportedSnafu {
5077 operator: function.function_name()
5078 },
5079 );
5080
5081 let field_column = self.ctx.field_columns[0].clone();
5082 let function_expr = self.create_native_histogram_expr(function, &field_column);
5083 let display_name = function_expr.schema_name().to_string();
5084 self.ctx.field_columns = vec![display_name.clone()];
5085
5086 let project_exprs = std::iter::once(self.create_time_index_column_expr()?)
5087 .chain(std::iter::once(function_expr.alias(display_name)))
5088 .chain(self.create_tag_column_exprs()?)
5089 .collect::<Vec<_>>();
5090
5091 LogicalPlanBuilder::from(input_plan)
5092 .project(project_exprs)
5093 .context(DataFusionPlanningSnafu)?
5094 .filter(self.create_empty_values_filter_expr(false)?)
5095 .context(DataFusionPlanningSnafu)?
5096 .build()
5097 .context(DataFusionPlanningSnafu)
5098 }
5099
5100 fn create_mixed_histogram_plan(
5101 &mut self,
5102 function: HistogramFoldOperation,
5103 input_plan: LogicalPlan,
5104 float_field: String,
5105 histogram_field: String,
5106 ) -> Result<LogicalPlan> {
5107 let time_index_column =
5108 self.ctx
5109 .time_index_column
5110 .clone()
5111 .with_context(|| TimeIndexNotFoundSnafu {
5112 table: self.ctx.table_name.clone().unwrap_or_default(),
5113 })?;
5114 let tag_columns = self.ctx.tag_columns.clone();
5115 let folded = HistogramFold::new_with_operation(
5116 LE_COLUMN_NAME.to_string(),
5117 float_field.clone(),
5118 time_index_column.clone(),
5119 function,
5120 Some(histogram_field.clone()),
5121 input_plan,
5122 )
5123 .context(DataFusionPlanningSnafu)?;
5124 let record_collision = DfExpr::ScalarFunction(ScalarFunction {
5125 func: Arc::new(NativeHistogramDrop::warning_bool_false_udf(
5126 "vector contains a mix of classic and native histograms".to_string(),
5127 self.promql_annotations.clone(),
5128 )),
5129 args: vec![col(&float_field), col(&histogram_field)],
5130 });
5131 let keep = when(
5132 col(&float_field)
5133 .is_not_null()
5134 .and(col(&histogram_field).is_not_null()),
5135 record_collision,
5136 )
5137 .otherwise(lit(true))
5138 .context(DataFusionPlanningSnafu)?;
5139
5140 let native_expr = self.create_native_histogram_expr(function, &histogram_field);
5141 let output_field = native_expr.schema_name().to_string();
5142 let value = DfExpr::ScalarFunction(ScalarFunction {
5143 func: coalesce(),
5144 args: vec![col(&float_field), native_expr],
5145 });
5146 self.ctx.field_columns = vec![output_field.clone()];
5147 LogicalPlanBuilder::from(LogicalPlan::Extension(Extension {
5148 node: Arc::new(folded),
5149 }))
5150 .filter(keep)
5151 .context(DataFusionPlanningSnafu)?
5152 .project(
5153 std::iter::once(col(&time_index_column))
5154 .chain(std::iter::once(value.alias(output_field)))
5155 .chain(tag_columns.iter().map(col)),
5156 )
5157 .context(DataFusionPlanningSnafu)?
5158 .build()
5159 .context(DataFusionPlanningSnafu)
5160 }
5161
5162 async fn create_vector_plan(&mut self, args: &PromFunctionArgs) -> Result<LogicalPlan> {
5164 if args.args.len() != 1 {
5165 return FunctionInvalidArgumentSnafu {
5166 fn_name: SPECIAL_VECTOR_FUNCTION.to_string(),
5167 }
5168 .fail();
5169 }
5170 let lit = Self::get_param_as_literal_expr(Some(args.args[0].as_ref()), None, None)?;
5171
5172 self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
5174 self.ctx.reset_table_name_and_schema();
5175 self.ctx.tag_columns = vec![];
5176 self.ctx.field_columns = vec![greptime_value().to_string()];
5177 Ok(LogicalPlan::Extension(Extension {
5178 node: Arc::new(
5179 EmptyMetric::new(
5180 self.ctx.start,
5181 self.ctx.end,
5182 self.ctx.interval,
5183 SPECIAL_TIME_FUNCTION.to_string(),
5184 greptime_value().to_string(),
5185 Some(lit),
5186 )
5187 .context(DataFusionPlanningSnafu)?,
5188 ),
5189 }))
5190 }
5191
5192 async fn create_scalar_plan(
5194 &mut self,
5195 args: &PromFunctionArgs,
5196 query_engine_state: &QueryEngineState,
5197 ) -> Result<LogicalPlan> {
5198 ensure!(
5199 args.len() == 1,
5200 FunctionInvalidArgumentSnafu {
5201 fn_name: SCALAR_FUNCTION
5202 }
5203 );
5204 let input = self
5205 .prom_expr_to_plan(&args.args[0], query_engine_state)
5206 .await?;
5207 let input_schema = input.schema().clone();
5208 let alternative_samples =
5209 Self::field_columns_are_alternative_samples(&input_schema, &self.ctx.field_columns);
5210 let histogram_fields = self
5211 .ctx
5212 .field_columns
5213 .iter()
5214 .filter(|field| Self::field_column_is_native_histogram(&input_schema, field))
5215 .count();
5216 ensure!(
5217 self.ctx.field_columns.len() == 1 || alternative_samples,
5218 MultiFieldsNotSupportedSnafu {
5219 operator: SCALAR_FUNCTION
5220 },
5221 );
5222 let scalar_field = self
5223 .ctx
5224 .field_columns
5225 .iter()
5226 .find(|field| !Self::field_column_is_native_histogram(&input_schema, field))
5227 .or_else(|| self.ctx.field_columns.first())
5228 .cloned()
5229 .with_context(|| FunctionInvalidArgumentSnafu {
5230 fn_name: SCALAR_FUNCTION,
5231 })?;
5232 let input = if histogram_fields == self.ctx.field_columns.len() {
5233 LogicalPlanBuilder::from(input)
5236 .filter(lit(false))
5237 .context(DataFusionPlanningSnafu)?
5238 .build()
5239 .context(DataFusionPlanningSnafu)?
5240 } else if histogram_fields > 0 {
5241 LogicalPlanBuilder::from(input)
5243 .filter(DfExpr::Column(Column::from_name(&scalar_field)).is_not_null())
5244 .context(DataFusionPlanningSnafu)?
5245 .build()
5246 .context(DataFusionPlanningSnafu)?
5247 } else {
5248 input
5249 };
5250 let scalar_plan = LogicalPlan::Extension(Extension {
5251 node: Arc::new(
5252 ScalarCalculate::new(
5253 self.ctx.start,
5254 self.ctx.end,
5255 self.ctx.interval,
5256 input,
5257 self.ctx.time_index_column.as_ref().unwrap(),
5258 &self.ctx.tag_columns,
5259 &scalar_field,
5260 self.ctx.table_name.as_deref(),
5261 )
5262 .context(PromqlPlanNodeSnafu)?,
5263 ),
5264 });
5265 self.ctx.tag_columns.clear();
5267 self.ctx.field_columns.clear();
5268 self.ctx
5269 .field_columns
5270 .push(scalar_plan.schema().field(1).name().clone());
5271 Ok(scalar_plan)
5272 }
5273
5274 async fn create_absent_plan(
5276 &mut self,
5277 args: &PromFunctionArgs,
5278 query_engine_state: &QueryEngineState,
5279 ) -> Result<LogicalPlan> {
5280 if args.args.len() != 1 {
5281 return FunctionInvalidArgumentSnafu {
5282 fn_name: SPECIAL_ABSENT_FUNCTION.to_string(),
5283 }
5284 .fail();
5285 }
5286 let input = self
5287 .prom_expr_to_plan(&args.args[0], query_engine_state)
5288 .await?;
5289
5290 let time_index_expr = self.create_time_index_column_expr()?;
5291 let first_field_expr =
5292 self.create_field_column_exprs()?
5293 .pop()
5294 .with_context(|| ValueNotFoundSnafu {
5295 table: self.ctx.table_name.clone().unwrap_or_default(),
5296 })?;
5297 let first_value_expr = first_value(first_field_expr, vec![]);
5298
5299 let ordered_aggregated_input = LogicalPlanBuilder::from(input)
5300 .aggregate(
5301 vec![time_index_expr.clone()],
5302 vec![first_value_expr.clone()],
5303 )
5304 .context(DataFusionPlanningSnafu)?
5305 .sort(vec![time_index_expr.sort(true, false)])
5306 .context(DataFusionPlanningSnafu)?
5307 .build()
5308 .context(DataFusionPlanningSnafu)?;
5309
5310 let fake_labels = self
5311 .ctx
5312 .selector_matcher
5313 .iter()
5314 .filter_map(|matcher| match matcher.op {
5315 MatchOp::Equal => Some((matcher.name.clone(), matcher.value.clone())),
5316 _ => None,
5317 })
5318 .collect::<Vec<_>>();
5319
5320 let absent_plan = LogicalPlan::Extension(Extension {
5322 node: Arc::new(
5323 Absent::try_new(
5324 self.ctx.start,
5325 self.ctx.end,
5326 self.ctx.interval,
5327 self.ctx.time_index_column.as_ref().unwrap().clone(),
5328 self.ctx.field_columns[0].clone(),
5329 fake_labels,
5330 ordered_aggregated_input,
5331 )
5332 .context(DataFusionPlanningSnafu)?,
5333 ),
5334 });
5335
5336 Ok(absent_plan)
5337 }
5338
5339 fn try_build_literal_expr(expr: &PromExpr) -> Option<DfExpr> {
5342 match expr {
5343 PromExpr::NumberLiteral(NumberLiteral { val }) => Some(val.lit()),
5344 PromExpr::StringLiteral(StringLiteral { val }) => Some(val.lit()),
5345 PromExpr::VectorSelector(_)
5346 | PromExpr::MatrixSelector(_)
5347 | PromExpr::Extension(_)
5348 | PromExpr::Aggregate(_)
5349 | PromExpr::Subquery(_) => None,
5350 PromExpr::Call(Call { func, .. }) => {
5351 if func.name == SPECIAL_TIME_FUNCTION {
5352 None
5355 } else {
5356 None
5357 }
5358 }
5359 PromExpr::Paren(ParenExpr { expr }) => Self::try_build_literal_expr(expr),
5360 PromExpr::Unary(UnaryExpr { expr, .. }) => Some(DfExpr::Negative(Box::new(
5361 Self::try_build_literal_expr(expr)?,
5362 ))),
5363 PromExpr::Binary(PromBinaryExpr {
5364 lhs,
5365 rhs,
5366 op,
5367 modifier,
5368 }) => {
5369 let lhs = Self::try_build_literal_expr(lhs)?;
5370 let rhs = Self::try_build_literal_expr(rhs)?;
5371 let is_comparison_op = Self::is_token_a_comparison_op(*op);
5372 let expr_builder = Self::prom_token_to_binary_expr_builder(*op).ok()?;
5373 let expr = expr_builder(lhs, rhs).ok()?;
5374
5375 let should_return_bool = if let Some(m) = modifier {
5376 m.return_bool
5377 } else {
5378 false
5379 };
5380 if is_comparison_op && should_return_bool {
5381 Some(DfExpr::Cast(Cast {
5382 expr: Box::new(expr),
5383 data_type: ArrowDataType::Float64,
5384 }))
5385 } else {
5386 Some(expr)
5387 }
5388 }
5389 }
5390 }
5391
5392 fn try_build_special_time_expr_with_context(&self, expr: &PromExpr) -> Option<DfExpr> {
5393 match expr {
5394 PromExpr::Call(Call { func, .. }) => {
5395 if func.name == SPECIAL_TIME_FUNCTION
5396 && let Some(time_index_col) = self.ctx.time_index_column.as_ref()
5397 {
5398 Some(build_special_time_expr(time_index_col))
5399 } else {
5400 None
5401 }
5402 }
5403 _ => None,
5404 }
5405 }
5406
5407 fn native_histogram_binary_expr(
5408 token: TokenType,
5409 lhs: DfExpr,
5410 lhs_is_histogram: bool,
5411 rhs: DfExpr,
5412 rhs_is_histogram: bool,
5413 filter_context: bool,
5414 promql_annotations: Option<PromqlAnnotationCollector>,
5415 ) -> Result<Option<DfExpr>> {
5416 if !lhs_is_histogram && !rhs_is_histogram {
5417 return Ok(None);
5418 }
5419
5420 let scalar_fn = |func: ScalarUdfDef, args| {
5421 DfExpr::ScalarFunction(ScalarFunction {
5422 func: Arc::new(func),
5423 args,
5424 })
5425 };
5426 let invalid_expr = || {
5427 let message = format!(
5428 "{}: dropped native histogram samples because this binary operation is not supported for native histograms",
5429 token
5430 );
5431 let func = if filter_context {
5432 NativeHistogramDrop::bool_false_udf(message, promql_annotations.clone())
5433 } else {
5434 NativeHistogramDrop::float_null_udf(message, promql_annotations.clone())
5435 };
5436 let args = vec![lhs.clone(), rhs.clone()];
5437 scalar_fn(func, args)
5438 };
5439
5440 let expr = match (token.id(), lhs_is_histogram, rhs_is_histogram) {
5441 (token::T_ADD, true, true) => scalar_fn(
5442 NativeHistogramAdd::scalar_udf_with_collector(promql_annotations.clone()),
5443 vec![lhs, rhs],
5444 ),
5445 (token::T_SUB, true, true) => scalar_fn(
5446 NativeHistogramSub::scalar_udf_with_collector(promql_annotations.clone()),
5447 vec![lhs, rhs],
5448 ),
5449 (token::T_MUL, true, false) => {
5450 scalar_fn(NativeHistogramMulScalar::scalar_udf(), vec![lhs, rhs])
5451 }
5452 (token::T_MUL, false, true) => {
5453 scalar_fn(NativeHistogramScalarMul::scalar_udf(), vec![lhs, rhs])
5454 }
5455 (token::T_DIV, true, false) => {
5456 scalar_fn(NativeHistogramDivScalar::scalar_udf(), vec![lhs, rhs])
5457 }
5458 (token::T_EQLC, true, true) => {
5459 scalar_fn(NativeHistogramEq::scalar_udf(), vec![lhs, rhs])
5460 }
5461 (token::T_NEQ, true, true) => {
5462 scalar_fn(NativeHistogramNotEq::scalar_udf(), vec![lhs, rhs])
5463 }
5464 _ => invalid_expr(),
5465 };
5466
5467 Ok(Some(expr))
5468 }
5469
5470 #[allow(clippy::type_complexity)]
5473 fn prom_token_to_binary_expr_builder(
5474 token: TokenType,
5475 ) -> Result<Box<dyn Fn(DfExpr, DfExpr) -> Result<DfExpr>>> {
5476 let cast_float = |expr| {
5477 if matches!(
5478 &expr,
5479 DfExpr::Cast(Cast {
5480 data_type: ArrowDataType::Float64,
5481 ..
5482 })
5483 ) || matches!(&expr, DfExpr::Literal(ScalarValue::Float64(_), _))
5484 {
5485 expr
5486 } else {
5487 DfExpr::Cast(Cast {
5488 expr: Box::new(expr),
5489 data_type: ArrowDataType::Float64,
5490 })
5491 }
5492 };
5493 match token.id() {
5494 token::T_ADD => Ok(Box::new(move |lhs, rhs| {
5495 Ok(cast_float(lhs) + cast_float(rhs))
5496 })),
5497 token::T_SUB => Ok(Box::new(move |lhs, rhs| {
5498 Ok(cast_float(lhs) - cast_float(rhs))
5499 })),
5500 token::T_MUL => Ok(Box::new(move |lhs, rhs| {
5501 Ok(cast_float(lhs) * cast_float(rhs))
5502 })),
5503 token::T_DIV => Ok(Box::new(move |lhs, rhs| {
5504 Ok(cast_float(lhs) / cast_float(rhs))
5505 })),
5506 token::T_MOD => Ok(Box::new(move |lhs: DfExpr, rhs| {
5507 Ok(cast_float(lhs) % cast_float(rhs))
5508 })),
5509 token::T_EQLC => Ok(Box::new(|lhs, rhs| Ok(lhs.eq(rhs)))),
5510 token::T_NEQ => Ok(Box::new(|lhs, rhs| Ok(lhs.not_eq(rhs)))),
5511 token::T_GTR => Ok(Box::new(|lhs, rhs| Ok(lhs.gt(rhs)))),
5512 token::T_LSS => Ok(Box::new(|lhs, rhs| Ok(lhs.lt(rhs)))),
5513 token::T_GTE => Ok(Box::new(|lhs, rhs| Ok(lhs.gt_eq(rhs)))),
5514 token::T_LTE => Ok(Box::new(|lhs, rhs| Ok(lhs.lt_eq(rhs)))),
5515 token::T_POW => Ok(Box::new(move |lhs, rhs| {
5516 Ok(DfExpr::ScalarFunction(ScalarFunction {
5517 func: datafusion_functions::math::power(),
5518 args: vec![cast_float(lhs), cast_float(rhs)],
5519 }))
5520 })),
5521 token::T_ATAN2 => Ok(Box::new(move |lhs, rhs| {
5522 Ok(DfExpr::ScalarFunction(ScalarFunction {
5523 func: datafusion_functions::math::atan2(),
5524 args: vec![cast_float(lhs), cast_float(rhs)],
5525 }))
5526 })),
5527 _ => UnexpectedTokenSnafu { token }.fail(),
5528 }
5529 }
5530
5531 fn is_token_a_comparison_op(token: TokenType) -> bool {
5533 matches!(
5534 token.id(),
5535 token::T_EQLC
5536 | token::T_NEQ
5537 | token::T_GTR
5538 | token::T_LSS
5539 | token::T_GTE
5540 | token::T_LTE
5541 )
5542 }
5543
5544 fn is_token_a_set_op(token: TokenType) -> bool {
5546 matches!(
5547 token.id(),
5548 token::T_LAND | token::T_LOR | token::T_LUNLESS )
5552 }
5553
5554 fn align_binary_field_columns<'a>(
5555 left_schema: &DFSchemaRef,
5556 right_schema: &DFSchemaRef,
5557 left_field_columns: &'a [String],
5558 right_field_columns: &'a [String],
5559 op: TokenType,
5560 left_is_scalar: bool,
5561 right_is_scalar: bool,
5562 ) -> (
5563 Vec<(String, Vec<BinaryFieldPair<'a>>)>,
5564 Vec<BinaryFieldPair<'a>>,
5565 ) {
5566 let left_alternative = Self::alternative_sample_columns(left_schema, left_field_columns);
5569 let right_alternative = Self::alternative_sample_columns(right_schema, right_field_columns);
5570 let alternative_alignment = match (left_alternative, right_alternative) {
5571 (Some(output_names), Some(_)) => Some((
5572 output_names,
5573 left_field_columns
5574 .iter()
5575 .flat_map(|left| right_field_columns.iter().map(move |right| (left, right)))
5576 .collect::<Vec<_>>(),
5577 )),
5578 (Some(output_names), None) if right_field_columns.len() == 1 => Some((
5579 output_names,
5580 left_field_columns
5581 .iter()
5582 .map(|left| (left, &right_field_columns[0]))
5583 .collect::<Vec<_>>(),
5584 )),
5585 (None, Some(output_names)) if left_field_columns.len() == 1 => Some((
5586 output_names,
5587 right_field_columns
5588 .iter()
5589 .map(|right| (&left_field_columns[0], right))
5590 .collect::<Vec<_>>(),
5591 )),
5592 _ => None,
5593 };
5594 let mut invalid_pairs = Vec::new();
5595 if let Some(((float_output, histogram_output), field_pairs)) = alternative_alignment {
5596 let mut float_pairs = Vec::new();
5597 let mut histogram_pairs = Vec::new();
5598 for (left, right) in field_pairs {
5599 let left_is_histogram = Self::field_column_is_native_histogram(left_schema, left);
5600 let right_is_histogram =
5601 Self::field_column_is_native_histogram(right_schema, right);
5602 match Self::binary_result_is_histogram(op, left_is_histogram, right_is_histogram) {
5603 Some(false) => float_pairs.push((left, right)),
5604 Some(true) => histogram_pairs.push((left, right)),
5605 None => invalid_pairs.push((left, right)),
5606 }
5607 }
5608 if !float_pairs.is_empty() || !histogram_pairs.is_empty() {
5609 return (
5610 [
5611 (!float_pairs.is_empty()).then(|| (float_output.to_string(), float_pairs)),
5612 (!histogram_pairs.is_empty())
5613 .then(|| (histogram_output.to_string(), histogram_pairs)),
5614 ]
5615 .into_iter()
5616 .flatten()
5617 .collect(),
5618 invalid_pairs,
5619 );
5620 }
5621 }
5622
5623 if left_is_scalar && !right_is_scalar && left_field_columns.len() == 1 {
5624 return (
5625 right_field_columns
5626 .iter()
5627 .map(|right| (right.clone(), vec![(&left_field_columns[0], right)]))
5628 .collect(),
5629 invalid_pairs,
5630 );
5631 }
5632 if right_is_scalar && !left_is_scalar && right_field_columns.len() == 1 {
5633 return (
5634 left_field_columns
5635 .iter()
5636 .map(|left| (left.clone(), vec![(left, &right_field_columns[0])]))
5637 .collect(),
5638 invalid_pairs,
5639 );
5640 }
5641
5642 (
5643 left_field_columns
5644 .iter()
5645 .zip(right_field_columns.iter())
5646 .map(|(left, right)| (left.clone(), vec![(left, right)]))
5647 .collect(),
5648 invalid_pairs,
5649 )
5650 }
5651
5652 fn binary_result_is_histogram(
5653 token: TokenType,
5654 lhs_is_histogram: bool,
5655 rhs_is_histogram: bool,
5656 ) -> Option<bool> {
5657 match (token.id(), lhs_is_histogram, rhs_is_histogram) {
5658 (_, false, false) => Some(false),
5659 (token::T_ADD | token::T_SUB, true, true)
5660 | (token::T_MUL, true, false)
5661 | (token::T_MUL, false, true)
5662 | (token::T_DIV, true, false) => Some(true),
5663 (token::T_EQLC | token::T_NEQ, true, true) => Some(false),
5664 _ => None,
5665 }
5666 }
5667
5668 fn plan_has_tsid_column(plan: &LogicalPlan) -> bool {
5669 plan.schema()
5670 .fields()
5671 .iter()
5672 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
5673 }
5674
5675 fn is_empty_metric(plan: &LogicalPlan) -> bool {
5676 matches!(plan, LogicalPlan::Extension(Extension { node }) if node.as_any().is::<EmptyMetric>())
5677 }
5678
5679 fn native_histogram_arrow_type() -> ArrowDataType {
5680 native_histogram_value_type().as_arrow_type()
5681 }
5682
5683 fn field_column_type<'a>(
5684 schema: &'a DFSchemaRef,
5685 field_column: &str,
5686 ) -> Option<&'a ArrowDataType> {
5687 schema
5688 .index_of_column_by_name(None, field_column)
5689 .map(|idx| schema.field(idx).data_type())
5690 }
5691
5692 fn field_column_is_native_histogram(schema: &DFSchemaRef, field_column: &str) -> bool {
5693 Self::field_column_type(schema, field_column)
5694 .is_some_and(|data_type| data_type == &Self::native_histogram_arrow_type())
5695 }
5696
5697 fn field_columns_contain_native_histogram(
5698 schema: &DFSchemaRef,
5699 field_columns: &[String],
5700 ) -> bool {
5701 field_columns
5702 .iter()
5703 .any(|field| Self::field_column_is_native_histogram(schema, field))
5704 }
5705
5706 fn field_column_is_float_range(schema: &DFSchemaRef, field_column: &str) -> bool {
5707 Self::field_column_type(schema, field_column).is_some_and(|data_type| {
5708 matches!(
5709 data_type,
5710 ArrowDataType::Dictionary(key_type, value_type)
5711 if key_type.as_ref() == &ArrowDataType::Int64
5712 && value_type.as_ref() == &ArrowDataType::Float64
5713 )
5714 })
5715 }
5716
5717 fn field_columns_are_alternative_samples(
5718 schema: &DFSchemaRef,
5719 field_columns: &[String],
5720 ) -> bool {
5721 Self::alternative_sample_columns(schema, field_columns).is_some()
5722 }
5723
5724 fn alternative_sample_columns<'a>(
5725 schema: &DFSchemaRef,
5726 field_columns: &'a [String],
5727 ) -> Option<(&'a str, &'a str)> {
5728 if field_columns.len() != 2 {
5729 return None;
5730 }
5731
5732 let canonical_float = field_columns.iter().find(|field| {
5733 field.as_str() == greptime_value()
5734 && (Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64)
5735 || Self::field_column_is_float_range(schema, field))
5736 });
5737 let canonical_histogram = field_columns.iter().find(|field| {
5738 field.as_str() == greptime_native_histogram()
5739 && (Self::field_column_is_native_histogram(schema, field)
5740 || Self::field_column_is_native_histogram_range(schema, field))
5741 });
5742 if let (Some(float), Some(histogram)) = (canonical_float, canonical_histogram) {
5743 return Some((float, histogram));
5744 }
5745
5746 let float = field_columns.iter().find(|field| {
5747 field.starts_with(OR_FLOAT_FIELD_PREFIX)
5748 && (Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64)
5749 || Self::field_column_is_float_range(schema, field))
5750 })?;
5751 let histogram = field_columns.iter().find(|field| {
5752 field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)
5753 && (Self::field_column_is_native_histogram(schema, field)
5754 || Self::field_column_is_native_histogram_range(schema, field))
5755 })?;
5756 Some((float, histogram))
5757 }
5758
5759 fn alternative_sample_range_columns<'a>(
5760 schema: &DFSchemaRef,
5761 field_columns: &'a [String],
5762 ) -> Option<(&'a str, &'a str)> {
5763 Self::alternative_sample_columns(schema, field_columns).filter(|(float, histogram)| {
5764 Self::field_column_is_float_range(schema, float)
5765 && Self::field_column_is_native_histogram_range(schema, histogram)
5766 })
5767 }
5768
5769 fn field_column_is_native_histogram_range(schema: &DFSchemaRef, field_column: &str) -> bool {
5770 Self::field_column_type(schema, field_column).is_some_and(|data_type| {
5771 matches!(
5772 data_type,
5773 ArrowDataType::Dictionary(key_type, value_type)
5774 if key_type.as_ref() == &ArrowDataType::Int64
5775 && value_type.as_ref() == &Self::native_histogram_arrow_type()
5776 )
5777 })
5778 }
5779
5780 fn all_field_columns_are_native_histograms(&self, schema: &DFSchemaRef) -> bool {
5781 !self.ctx.field_columns.is_empty()
5782 && self
5783 .ctx
5784 .field_columns
5785 .iter()
5786 .all(|field| Self::field_column_is_native_histogram(schema, field))
5787 }
5788
5789 fn all_field_columns_are_native_histogram_ranges(&self, schema: &DFSchemaRef) -> bool {
5790 !self.ctx.field_columns.is_empty()
5791 && self
5792 .ctx
5793 .field_columns
5794 .iter()
5795 .all(|field| Self::field_column_is_native_histogram_range(schema, field))
5796 }
5797
5798 fn optional_tsid_projection(
5799 schema: &DFSchemaRef,
5800 table_ref: Option<&TableReference>,
5801 keep_tsid: bool,
5802 ) -> Option<DfExpr> {
5803 keep_tsid.then_some(()).and_then(|_| {
5804 schema
5805 .qualified_field_with_name(table_ref, DATA_SCHEMA_TSID_COLUMN_NAME)
5806 .ok()
5807 .map(|field| DfExpr::Column(field.into()))
5808 })
5809 }
5810
5811 fn binary_join_key_columns(
5812 &self,
5813 left_schema: &DFSchemaRef,
5814 right_schema: &DFSchemaRef,
5815 left_context: &PromPlannerContext,
5816 right_context: &PromPlannerContext,
5817 only_join_time_index: bool,
5818 modifier: &Option<BinModifier>,
5819 ) -> Result<(BTreeSet<String>, BTreeSet<String>, bool)> {
5820 let has_tsid = |schema: &DFSchemaRef| {
5821 schema
5822 .fields()
5823 .iter()
5824 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
5825 };
5826 let use_tsid_join = !only_join_time_index
5827 && self.binary_modifier_preserves_tsid_join_key(left_context, right_context, modifier)
5828 && left_context.use_tsid
5829 && right_context.use_tsid
5830 && has_tsid(left_schema)
5831 && has_tsid(right_schema);
5832
5833 let (mut left_tag_columns, mut right_tag_columns) = if use_tsid_join {
5834 (
5835 BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]),
5836 BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]),
5837 )
5838 } else {
5839 if only_join_time_index {
5840 (BTreeSet::new(), BTreeSet::new())
5841 } else {
5842 (
5843 left_context
5844 .tag_columns
5845 .iter()
5846 .cloned()
5847 .collect::<BTreeSet<_>>(),
5848 right_context
5849 .tag_columns
5850 .iter()
5851 .cloned()
5852 .collect::<BTreeSet<_>>(),
5853 )
5854 }
5855 };
5856
5857 if !use_tsid_join
5858 && let Some(modifier) = modifier
5859 && let Some(matching) = &modifier.matching
5860 {
5861 match matching {
5862 LabelModifier::Include(on) => {
5863 let mask = on.labels.iter().cloned().collect::<BTreeSet<_>>();
5864 left_tag_columns = left_tag_columns.intersection(&mask).cloned().collect();
5865 right_tag_columns = right_tag_columns.intersection(&mask).cloned().collect();
5866 }
5867 LabelModifier::Exclude(ignoring) => {
5868 for label in &ignoring.labels {
5869 let _ = left_tag_columns.remove(label);
5870 let _ = right_tag_columns.remove(label);
5871 }
5872 }
5873 }
5874 }
5875
5876 let force_empty_join =
5877 !use_tsid_join && !only_join_time_index && left_tag_columns != right_tag_columns;
5878 if force_empty_join {
5879 let common_tag_columns = left_tag_columns
5880 .intersection(&right_tag_columns)
5881 .cloned()
5882 .collect::<BTreeSet<_>>();
5883 left_tag_columns = common_tag_columns.clone();
5884 right_tag_columns = common_tag_columns;
5885 }
5886
5887 Ok((left_tag_columns, right_tag_columns, force_empty_join))
5888 }
5889
5890 fn binary_modifier_preserves_tsid_join_key(
5891 &self,
5892 left_context: &PromPlannerContext,
5893 right_context: &PromPlannerContext,
5894 modifier: &Option<BinModifier>,
5895 ) -> bool {
5896 let Some(modifier) = modifier else {
5897 return true;
5898 };
5899
5900 if !matches!(modifier.card, VectorMatchCardinality::OneToOne) {
5901 return false;
5902 }
5903
5904 match &modifier.matching {
5905 None => true,
5906 Some(LabelModifier::Exclude(ignoring)) => ignoring.labels.iter().all(|label| {
5907 !left_context.tag_columns.contains(label)
5908 && !right_context.tag_columns.contains(label)
5909 }),
5910 Some(LabelModifier::Include(on)) => {
5911 let on_labels = on.labels.iter().cloned().collect::<BTreeSet<_>>();
5912 let left_labels = left_context
5913 .tag_columns
5914 .iter()
5915 .cloned()
5916 .collect::<BTreeSet<_>>();
5917 let right_labels = right_context
5918 .tag_columns
5919 .iter()
5920 .cloned()
5921 .collect::<BTreeSet<_>>();
5922
5923 on_labels == left_labels && on_labels == right_labels
5924 }
5925 }
5926 }
5927
5928 #[allow(clippy::too_many_arguments)]
5931 fn join_on_non_field_columns(
5932 &self,
5933 left: LogicalPlan,
5934 right: LogicalPlan,
5935 left_table_ref: TableReference,
5936 right_table_ref: TableReference,
5937 left_time_index_column: Option<String>,
5938 right_time_index_column: Option<String>,
5939 only_join_time_index: bool,
5940 modifier: &Option<BinModifier>,
5941 left_context: &PromPlannerContext,
5942 right_context: &PromPlannerContext,
5943 ) -> Result<LogicalPlan> {
5944 let (mut left_tag_columns, mut right_tag_columns, mut force_empty_join) = self
5945 .binary_join_key_columns(
5946 left.schema(),
5947 right.schema(),
5948 left_context,
5949 right_context,
5950 only_join_time_index,
5951 modifier,
5952 )?;
5953 let use_tsid_join = !only_join_time_index
5954 && !force_empty_join
5955 && left_tag_columns == BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()])
5956 && right_tag_columns == BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]);
5957 let (left, right) = if !only_join_time_index
5958 && !use_tsid_join
5959 && Self::only_temporality_match_label_mismatches(left_context, right_context, modifier)
5960 {
5961 let mut aligned_left_context = left_context.clone();
5962 let mut aligned_right_context = right_context.clone();
5963 let (left, right, _) = Self::align_temporality_match_column(
5964 left,
5965 right,
5966 &mut aligned_left_context,
5967 &mut aligned_right_context,
5968 )?;
5969 (left_tag_columns, right_tag_columns, force_empty_join) = self
5970 .binary_join_key_columns(
5971 left.schema(),
5972 right.schema(),
5973 &aligned_left_context,
5974 &aligned_right_context,
5975 false,
5976 modifier,
5977 )?;
5978 (left, right)
5979 } else {
5980 (left, right)
5981 };
5982
5983 if let (Some(left_time_index_column), Some(right_time_index_column)) =
5985 (left_time_index_column, right_time_index_column)
5986 {
5987 left_tag_columns.insert(left_time_index_column);
5988 right_tag_columns.insert(right_time_index_column);
5989 }
5990
5991 let right = LogicalPlanBuilder::from(right)
5992 .alias(right_table_ref)
5993 .context(DataFusionPlanningSnafu)?
5994 .build()
5995 .context(DataFusionPlanningSnafu)?;
5996
5997 LogicalPlanBuilder::from(left)
5999 .alias(left_table_ref)
6000 .context(DataFusionPlanningSnafu)?
6001 .join_detailed(
6002 right,
6003 JoinType::Inner,
6004 (
6005 left_tag_columns
6006 .into_iter()
6007 .map(Column::from_name)
6008 .collect::<Vec<_>>(),
6009 right_tag_columns
6010 .into_iter()
6011 .map(Column::from_name)
6012 .collect::<Vec<_>>(),
6013 ),
6014 force_empty_join.then_some(lit(false)),
6015 NullEquality::NullEqualsNull,
6016 )
6017 .context(DataFusionPlanningSnafu)?
6018 .build()
6019 .context(DataFusionPlanningSnafu)
6020 }
6021
6022 fn selected_binary_match_labels(
6023 left_context: &PromPlannerContext,
6024 right_context: &PromPlannerContext,
6025 modifier: &Option<BinModifier>,
6026 ) -> BTreeSet<String> {
6027 let mut labels = left_context
6028 .tag_columns
6029 .iter()
6030 .chain(&right_context.tag_columns)
6031 .cloned()
6032 .collect::<BTreeSet<_>>();
6033 if let Some(matching) = modifier
6034 .as_ref()
6035 .and_then(|modifier| modifier.matching.as_ref())
6036 {
6037 match matching {
6038 LabelModifier::Include(on) => {
6039 labels = on
6040 .labels
6041 .iter()
6042 .filter(|label| {
6043 left_context.tag_columns.contains(label)
6044 || right_context.tag_columns.contains(label)
6045 })
6046 .cloned()
6047 .collect();
6048 }
6049 LabelModifier::Exclude(ignoring) => {
6050 for label in &ignoring.labels {
6051 labels.remove(label);
6052 }
6053 }
6054 }
6055 }
6056 labels
6057 }
6058
6059 fn only_temporality_match_label_mismatches(
6060 left_context: &PromPlannerContext,
6061 right_context: &PromPlannerContext,
6062 modifier: &Option<BinModifier>,
6063 ) -> bool {
6064 let mut mismatches =
6065 Self::selected_binary_match_labels(left_context, right_context, modifier)
6066 .into_iter()
6067 .filter(|label| {
6068 left_context.tag_columns.contains(label)
6069 != right_context.tag_columns.contains(label)
6070 });
6071 matches!(
6072 (mismatches.next(), mismatches.next()),
6073 (Some(label), None) if label == OTLP_AGGREGATION_TEMPORALITY_LABEL
6074 )
6075 }
6076
6077 fn align_temporality_match_column(
6078 mut left: LogicalPlan,
6079 mut right: LogicalPlan,
6080 left_context: &mut PromPlannerContext,
6081 right_context: &mut PromPlannerContext,
6082 ) -> Result<(LogicalPlan, LogicalPlan, bool)> {
6083 let marker = OTLP_AGGREGATION_TEMPORALITY_LABEL;
6084 let left_has_marker = left_context.tag_columns.iter().any(|tag| tag == marker);
6085 let (present, add_to_left) = if left_has_marker {
6086 (&left, false)
6087 } else {
6088 (&right, true)
6089 };
6090 let data_type = present
6091 .schema()
6092 .fields()
6093 .iter()
6094 .find(|field| field.name() == marker)
6095 .map(|field| field.data_type().clone())
6096 .with_context(|| ColumnNotFoundSnafu {
6097 col: marker.to_string(),
6098 })?;
6099 let null = Self::string_scalar_value(&data_type, None).with_context(|| {
6100 UnexpectedPlanExprSnafu {
6101 desc: format!("temporality match label {marker} must be a string"),
6102 }
6103 })?;
6104 let add_marker = |plan: LogicalPlan| {
6105 let visible = plan
6106 .schema()
6107 .iter()
6108 .map(|(qualifier, field)| {
6109 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6110 })
6111 .collect::<Vec<_>>();
6112 LogicalPlanBuilder::from(plan)
6113 .project(
6114 visible
6115 .into_iter()
6116 .chain([DfExpr::Literal(null, None).alias(marker)]),
6117 )
6118 .context(DataFusionPlanningSnafu)?
6119 .build()
6120 .context(DataFusionPlanningSnafu)
6121 };
6122
6123 if add_to_left {
6124 left = add_marker(left)?;
6125 left_context.tag_columns.push(marker.to_string());
6126 } else {
6127 right = add_marker(right)?;
6128 right_context.tag_columns.push(marker.to_string());
6129 }
6130 Ok((left, right, add_to_left))
6131 }
6132
6133 fn normalized_match_key_expr(
6134 label: &str,
6135 field: Option<(Option<TableReference>, ArrowDataType)>,
6136 value_type: &ArrowDataType,
6137 internal_name: &str,
6138 ) -> DfExpr {
6139 let empty = Self::string_scalar_value(value_type, Some(String::new()))
6140 .expect("match label value type is a string");
6141 let expr = if let Some((qualifier, data_type)) = field {
6142 let column = DfExpr::Column(Column::new(qualifier, label));
6143 let column = if &data_type == value_type {
6144 column
6145 } else {
6146 DfExpr::Cast(Cast {
6147 expr: Box::new(column),
6148 data_type: value_type.clone(),
6149 })
6150 };
6151 DfExpr::ScalarFunction(ScalarFunction {
6152 func: coalesce(),
6153 args: vec![column, DfExpr::Literal(empty, None)],
6154 })
6155 } else {
6156 DfExpr::Literal(empty, None)
6157 };
6158 expr.alias(internal_name)
6159 }
6160
6161 fn is_zero_row_empty_relation(plan: &LogicalPlan) -> bool {
6162 matches!(plan, LogicalPlan::EmptyRelation(relation) if !relation.produce_one_row)
6165 }
6166
6167 fn set_op_on_non_field_columns(
6169 &mut self,
6170 mut left: LogicalPlan,
6171 mut right: LogicalPlan,
6172 left_context: PromPlannerContext,
6173 right_context: PromPlannerContext,
6174 op: TokenType,
6175 modifier: &Option<BinModifier>,
6176 ) -> Result<LogicalPlan> {
6177 let left_tag_col_set = left_context
6178 .tag_columns
6179 .iter()
6180 .cloned()
6181 .collect::<HashSet<_>>();
6182 let right_tag_col_set = right_context
6183 .tag_columns
6184 .iter()
6185 .cloned()
6186 .collect::<HashSet<_>>();
6187
6188 if matches!(op.id(), token::T_LOR) {
6189 return self.or_operator(
6190 left,
6191 right,
6192 left_tag_col_set,
6193 right_tag_col_set,
6194 left_context,
6195 right_context,
6196 modifier,
6197 );
6198 }
6199
6200 if let Some(modifier) = modifier {
6201 ensure!(
6202 matches!(
6203 modifier.card,
6204 VectorMatchCardinality::OneToOne | VectorMatchCardinality::ManyToMany
6205 ),
6206 UnsupportedVectorMatchSnafu {
6207 name: modifier.card.clone(),
6208 },
6209 );
6210 }
6211
6212 let output_context = left_context.clone();
6213 let visible_left_schema = left.schema().clone();
6214 let mut left_context = left_context;
6215 let mut right_context = right_context;
6216 let added_marker_to_left = if Self::only_temporality_match_label_mismatches(
6217 &left_context,
6218 &right_context,
6219 modifier,
6220 ) {
6221 let aligned = Self::align_temporality_match_column(
6222 left,
6223 right,
6224 &mut left_context,
6225 &mut right_context,
6226 )?;
6227 left = aligned.0;
6228 right = aligned.1;
6229 aligned.2
6230 } else {
6231 false
6232 };
6233
6234 let mut left_tag_col_set = left_context
6235 .tag_columns
6236 .iter()
6237 .cloned()
6238 .collect::<BTreeSet<_>>();
6239 let mut right_tag_col_set = right_context
6240 .tag_columns
6241 .iter()
6242 .cloned()
6243 .collect::<BTreeSet<_>>();
6244 if let Some(matching) = modifier
6245 .as_ref()
6246 .and_then(|modifier| modifier.matching.as_ref())
6247 {
6248 match matching {
6249 LabelModifier::Include(on) => {
6250 let mask = on.labels.iter().cloned().collect::<BTreeSet<_>>();
6251 left_tag_col_set = left_tag_col_set.intersection(&mask).cloned().collect();
6252 right_tag_col_set = right_tag_col_set.intersection(&mask).cloned().collect();
6253 }
6254 LabelModifier::Exclude(ignoring) => {
6255 for label in &ignoring.labels {
6256 let _ = left_tag_col_set.remove(label);
6257 let _ = right_tag_col_set.remove(label);
6258 }
6259 }
6260 }
6261 }
6262 ensure!(
6263 left_tag_col_set == right_tag_col_set,
6264 CombineTableColumnMismatchSnafu {
6265 left: left_tag_col_set.iter().cloned().collect::<Vec<_>>(),
6266 right: right_tag_col_set.iter().cloned().collect::<Vec<_>>(),
6267 }
6268 );
6269
6270 let left_time_index = left_context.time_index_column.clone().unwrap();
6271 let right_time_index = right_context.time_index_column.clone().unwrap();
6272
6273 if left_context.time_index_column != right_context.time_index_column {
6275 let right_project_exprs = right
6276 .schema()
6277 .fields()
6278 .iter()
6279 .map(|field| {
6280 if field.name() == &right_time_index {
6281 DfExpr::Column(Column::from_name(&right_time_index)).alias(&left_time_index)
6282 } else {
6283 DfExpr::Column(Column::from_name(field.name()))
6284 }
6285 })
6286 .collect::<Vec<_>>();
6287
6288 right = LogicalPlanBuilder::from(right)
6289 .project(right_project_exprs)
6290 .context(DataFusionPlanningSnafu)?
6291 .build()
6292 .context(DataFusionPlanningSnafu)?;
6293 }
6294
6295 let join_keys = left_tag_col_set
6296 .into_iter()
6297 .chain([left_time_index])
6298 .collect::<Vec<_>>();
6299
6300 ensure!(
6301 left_context.field_columns.len() == 1
6302 || Self::field_columns_are_alternative_samples(
6303 left.schema(),
6304 &left_context.field_columns,
6305 ),
6306 MultiFieldsNotSupportedSnafu {
6307 operator: "AND/UNLESS operator"
6308 }
6309 );
6310 let result = match op.id() {
6313 token::T_LAND => LogicalPlanBuilder::from(left)
6314 .distinct()
6315 .context(DataFusionPlanningSnafu)?
6316 .join_detailed(
6317 right,
6318 JoinType::LeftSemi,
6319 (join_keys.clone(), join_keys),
6320 None,
6321 NullEquality::NullEqualsNull,
6322 )
6323 .context(DataFusionPlanningSnafu)?
6324 .build()
6325 .context(DataFusionPlanningSnafu),
6326 token::T_LUNLESS => LogicalPlanBuilder::from(left)
6327 .distinct()
6328 .context(DataFusionPlanningSnafu)?
6329 .join_detailed(
6330 right,
6331 JoinType::LeftAnti,
6332 (join_keys.clone(), join_keys),
6333 None,
6334 NullEquality::NullEqualsNull,
6335 )
6336 .context(DataFusionPlanningSnafu)?
6337 .build()
6338 .context(DataFusionPlanningSnafu),
6339 token::T_LOR => {
6340 unreachable!()
6343 }
6344 _ => UnexpectedTokenSnafu { token: op }.fail(),
6345 }?;
6346 let result = if added_marker_to_left {
6347 LogicalPlanBuilder::from(result)
6348 .project(visible_left_schema.iter().map(|(qualifier, field)| {
6349 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6350 }))
6351 .context(DataFusionPlanningSnafu)?
6352 .build()
6353 .context(DataFusionPlanningSnafu)?
6354 } else {
6355 result
6356 };
6357
6358 self.ctx = output_context;
6360 Ok(result)
6361 }
6362
6363 fn string_value_data_type(data_type: &ArrowDataType) -> Option<&ArrowDataType> {
6364 match data_type {
6365 data_type if data_type.is_string() => Some(data_type),
6366 ArrowDataType::Dictionary(_, value_type) if value_type.is_string() => Some(value_type),
6367 _ => None,
6368 }
6369 }
6370
6371 fn string_scalar_value(
6372 data_type: &ArrowDataType,
6373 value: Option<String>,
6374 ) -> Option<ScalarValue> {
6375 match data_type {
6376 ArrowDataType::Utf8 => Some(ScalarValue::Utf8(value)),
6377 ArrowDataType::LargeUtf8 => Some(ScalarValue::LargeUtf8(value)),
6378 ArrowDataType::Utf8View => Some(ScalarValue::Utf8View(value)),
6379 ArrowDataType::Dictionary(key_type, value_type) => Some(ScalarValue::Dictionary(
6380 key_type.clone(),
6381 Box::new(Self::string_scalar_value(value_type, value)?),
6382 )),
6383 _ => None,
6384 }
6385 }
6386
6387 fn common_label_data_type(
6388 left: Option<&ArrowDataType>,
6389 right: Option<&ArrowDataType>,
6390 ) -> Option<ArrowDataType> {
6391 match (left, right) {
6392 (Some(left), Some(right)) if left == right => {
6393 Self::string_value_data_type(left).map(|_| left.clone())
6394 }
6395 (Some(left), Some(right)) => {
6396 let left_value_type = Self::string_value_data_type(left)?;
6397 let right_value_type = Self::string_value_data_type(right)?;
6398 match (left_value_type, right_value_type) {
6401 (left, right) if left == right => Some(left.clone()),
6402 (ArrowDataType::LargeUtf8, _) | (_, ArrowDataType::LargeUtf8) => {
6403 Some(ArrowDataType::LargeUtf8)
6404 }
6405 (ArrowDataType::Utf8View, ArrowDataType::Utf8View) => {
6406 Some(ArrowDataType::Utf8View)
6407 }
6408 _ => Some(ArrowDataType::Utf8),
6409 }
6410 }
6411 (Some(data_type), None) | (None, Some(data_type)) => {
6412 Self::string_value_data_type(data_type).cloned()
6413 }
6414 (None, None) => Some(ArrowDataType::Utf8),
6415 }
6416 }
6417
6418 #[allow(clippy::too_many_arguments)]
6420 fn or_operator(
6421 &mut self,
6422 left: LogicalPlan,
6423 right: LogicalPlan,
6424 left_tag_cols_set: HashSet<String>,
6425 right_tag_cols_set: HashSet<String>,
6426 left_context: PromPlannerContext,
6427 right_context: PromPlannerContext,
6428 modifier: &Option<BinModifier>,
6429 ) -> Result<LogicalPlan> {
6430 let left_is_empty = Self::is_zero_row_empty_relation(&left);
6431 let right_is_empty = Self::is_zero_row_empty_relation(&right);
6432 match (left_is_empty, right_is_empty) {
6433 (true, false) => {
6434 self.ctx = right_context;
6435 return Ok(right);
6436 }
6437 (false, true) => {
6438 self.ctx = left_context;
6439 return Ok(left);
6440 }
6441 (true, true) => {
6442 self.ctx = left_context;
6443 return Ok(left);
6444 }
6445 (false, false) => {}
6446 }
6447
6448 ensure!(
6449 !left.schema().fields().is_empty() && !right.schema().fields().is_empty(),
6450 UnexpectedPlanExprSnafu {
6451 desc: "OR operator input has zero columns",
6452 }
6453 );
6454 let left_has_alternative_samples =
6455 Self::field_columns_are_alternative_samples(left.schema(), &left_context.field_columns);
6456 let right_has_alternative_samples = Self::field_columns_are_alternative_samples(
6457 right.schema(),
6458 &right_context.field_columns,
6459 );
6460 ensure!(
6461 left_context.field_columns.len() == 1 || left_has_alternative_samples,
6462 MultiFieldsNotSupportedSnafu {
6463 operator: "OR operator"
6464 }
6465 );
6466 ensure!(
6467 right_context.field_columns.len() == 1 || right_has_alternative_samples,
6468 MultiFieldsNotSupportedSnafu {
6469 operator: "OR operator"
6470 }
6471 );
6472
6473 let all_tags = left_tag_cols_set
6475 .union(&right_tag_cols_set)
6476 .cloned()
6477 .collect::<HashSet<_>>();
6478 let left_qualifier = left.schema().qualified_field(0).0.cloned();
6479 let right_qualifier = right.schema().qualified_field(0).0.cloned();
6480 let left_qualifier_string = left_qualifier
6481 .as_ref()
6482 .map(|l| l.to_string())
6483 .unwrap_or_default();
6484 let right_qualifier_string = right_qualifier
6485 .as_ref()
6486 .map(|r| r.to_string())
6487 .unwrap_or_default();
6488 let left_time_index_column =
6489 left_context
6490 .time_index_column
6491 .clone()
6492 .with_context(|| TimeIndexNotFoundSnafu {
6493 table: left_qualifier_string.clone(),
6494 })?;
6495 let right_time_index_column =
6496 right_context
6497 .time_index_column
6498 .clone()
6499 .with_context(|| TimeIndexNotFoundSnafu {
6500 table: right_qualifier_string.clone(),
6501 })?;
6502 let native_histogram_type = Self::native_histogram_arrow_type();
6503 let is_numeric = |data_type: &ArrowDataType| {
6504 matches!(
6505 data_type,
6506 ArrowDataType::Int8
6507 | ArrowDataType::Int16
6508 | ArrowDataType::Int32
6509 | ArrowDataType::Int64
6510 | ArrowDataType::UInt8
6511 | ArrowDataType::UInt16
6512 | ArrowDataType::UInt32
6513 | ArrowDataType::UInt64
6514 | ArrowDataType::Float32
6515 | ArrowDataType::Float64
6516 )
6517 };
6518 let left_fields = left_context
6519 .field_columns
6520 .iter()
6521 .map(|name| {
6522 left.schema()
6523 .iter()
6524 .find(|(_, field)| field.name() == name)
6525 .map(|(qualifier, field)| {
6526 (name.clone(), qualifier.cloned(), field.data_type().clone())
6527 })
6528 .with_context(|| ColumnNotFoundSnafu { col: name.clone() })
6529 })
6530 .collect::<Result<Vec<_>>>()?;
6531 let right_fields = right_context
6532 .field_columns
6533 .iter()
6534 .map(|name| {
6535 right
6536 .schema()
6537 .iter()
6538 .find(|(_, field)| field.name() == name)
6539 .map(|(qualifier, field)| {
6540 (name.clone(), qualifier.cloned(), field.data_type().clone())
6541 })
6542 .with_context(|| ColumnNotFoundSnafu { col: name.clone() })
6543 })
6544 .collect::<Result<Vec<_>>>()?;
6545 let left_field = &left_fields[0];
6546 let right_field = &right_fields[0];
6547 let left_field_col = &left_field.0;
6548 let right_field_col = &right_field.0;
6549 let fields_are_samples = |fields: &[(String, Option<TableReference>, ArrowDataType)]| {
6550 fields.iter().all(|(_, _, data_type)| {
6551 is_numeric(data_type) || data_type == &native_histogram_type
6552 })
6553 };
6554 let mixed_sample_types = if left_has_alternative_samples || right_has_alternative_samples {
6555 if !fields_are_samples(&left_fields) || !fields_are_samples(&right_fields) {
6556 return UnexpectedPlanExprSnafu {
6557 desc: format!(
6558 "OR value fields have incompatible types: {:?} and {:?}",
6559 left_fields
6560 .iter()
6561 .map(|(_, _, data_type)| data_type)
6562 .collect::<Vec<_>>(),
6563 right_fields
6564 .iter()
6565 .map(|(_, _, data_type)| data_type)
6566 .collect::<Vec<_>>()
6567 ),
6568 }
6569 .fail();
6570 }
6571 true
6572 } else {
6573 (left_field.2 == native_histogram_type && is_numeric(&right_field.2))
6574 || (right_field.2 == native_histogram_type && is_numeric(&left_field.2))
6575 };
6576 let target_field_type = if mixed_sample_types {
6577 ArrowDataType::Float64
6580 } else if left_field.2 == right_field.2 {
6581 left_field.2.clone()
6582 } else if is_numeric(&left_field.2) && is_numeric(&right_field.2) {
6583 ArrowDataType::Float64
6584 } else {
6585 return UnexpectedPlanExprSnafu {
6586 desc: format!(
6587 "OR value fields have incompatible types: {:?} and {:?}",
6588 left_field.2, right_field.2
6589 ),
6590 }
6591 .fail();
6592 };
6593 let (mixed_float_field_col, mixed_histogram_field_col) = if mixed_sample_types {
6594 let mut reserved_names = left
6595 .schema()
6596 .fields()
6597 .iter()
6598 .chain(right.schema().fields().iter())
6599 .map(|field| field.name().clone())
6600 .collect::<HashSet<_>>();
6601 for (name, _, _) in left_fields.iter().chain(&right_fields) {
6602 reserved_names.remove(name);
6603 }
6604 reserved_names.extend(all_tags.iter().cloned());
6605 let unique_name = |prefix: &str, reserved_names: &mut HashSet<String>| {
6606 let mut index = 0;
6607 loop {
6608 let name = format!("{prefix}{index}");
6609 index += 1;
6610 if reserved_names.insert(name.clone()) {
6611 break name;
6612 }
6613 }
6614 };
6615 let float_field = unique_name(OR_FLOAT_FIELD_PREFIX, &mut reserved_names);
6616 let histogram_field = unique_name(OR_HISTOGRAM_FIELD_PREFIX, &mut reserved_names);
6617 (float_field, histogram_field)
6618 } else {
6619 (left_field_col.clone(), String::new())
6620 };
6621 let left_tag_types = left_tag_cols_set
6622 .iter()
6623 .map(|label| {
6624 left.schema()
6625 .fields()
6626 .iter()
6627 .find(|field| field.name() == label)
6628 .map(|field| (label.clone(), field.data_type().clone()))
6629 .with_context(|| ColumnNotFoundSnafu { col: label.clone() })
6630 })
6631 .collect::<Result<HashMap<_, _>>>()?;
6632 let right_tag_types = right_tag_cols_set
6633 .iter()
6634 .map(|label| {
6635 right
6636 .schema()
6637 .fields()
6638 .iter()
6639 .find(|field| field.name() == label)
6640 .map(|field| (label.clone(), field.data_type().clone()))
6641 .with_context(|| ColumnNotFoundSnafu { col: label.clone() })
6642 })
6643 .collect::<Result<HashMap<_, _>>>()?;
6644 let mut target_tag_types = HashMap::with_capacity(all_tags.len());
6645 for label in &all_tags {
6646 let Some(data_type) =
6647 Self::common_label_data_type(left_tag_types.get(label), right_tag_types.get(label))
6648 else {
6649 return UnexpectedPlanExprSnafu {
6650 desc: format!(
6651 "OR label {label} has incompatible types: {:?} and {:?}",
6652 left_tag_types.get(label),
6653 right_tag_types.get(label)
6654 ),
6655 }
6656 .fail();
6657 };
6658 target_tag_types.insert(label.clone(), data_type);
6659 }
6660 let left_has_tsid = left
6661 .schema()
6662 .fields()
6663 .iter()
6664 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME);
6665 let right_has_tsid = right
6666 .schema()
6667 .fields()
6668 .iter()
6669 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME);
6670
6671 let mut all_columns_set = left
6673 .schema()
6674 .fields()
6675 .iter()
6676 .chain(right.schema().fields().iter())
6677 .map(|field| field.name().clone())
6678 .collect::<HashSet<_>>();
6679 if !(left_has_tsid && right_has_tsid) {
6682 all_columns_set.remove(DATA_SCHEMA_TSID_COLUMN_NAME);
6683 }
6684 all_columns_set.remove(&left_time_index_column);
6686 all_columns_set.remove(&right_time_index_column);
6687 if mixed_sample_types {
6688 for (name, _, _) in left_fields.iter().chain(&right_fields) {
6689 all_columns_set.remove(name);
6690 }
6691 all_columns_set.extend(all_tags.iter().cloned());
6692 all_columns_set.insert(mixed_float_field_col.clone());
6693 all_columns_set.insert(mixed_histogram_field_col.clone());
6694 } else if left_field_col != right_field_col {
6695 all_columns_set.remove(right_field_col);
6697 }
6698 let mut all_columns = all_columns_set.into_iter().collect::<Vec<_>>();
6699 all_columns.sort_unstable();
6701 all_columns.insert(0, left_time_index_column.clone());
6703 let mut occupied_column_names = left
6704 .schema()
6705 .fields()
6706 .iter()
6707 .chain(right.schema().fields().iter())
6708 .map(|field| field.name().clone())
6709 .collect::<HashSet<_>>();
6710
6711 let aligned_label_expr = |col: &String, source_types: &HashMap<String, ArrowDataType>| {
6713 let target_type = &target_tag_types[col];
6714 if let Some(source_type) = source_types.get(col) {
6715 let expr = DfExpr::Column(Column::new(None::<String>, col));
6716 if source_type == target_type {
6717 expr
6718 } else {
6719 DfExpr::Cast(Cast {
6720 expr: Box::new(expr),
6721 data_type: target_type.clone(),
6722 })
6723 .alias(col.clone())
6724 }
6725 } else {
6726 DfExpr::Literal(
6727 Self::string_scalar_value(target_type, None)
6728 .expect("target label type is a string"),
6729 None,
6730 )
6731 .alias(col.clone())
6732 }
6733 };
6734 let null_histogram =
6735 ScalarValue::try_new_null(&native_histogram_type).context(DataFusionPlanningSnafu)?;
6736 let mixed_value_expr = |fields: &[(String, Option<TableReference>, ArrowDataType)],
6737 output_col: &String| {
6738 if output_col == &mixed_float_field_col {
6739 if let Some((name, qualifier, data_type)) = fields
6740 .iter()
6741 .find(|(_, _, data_type)| is_numeric(data_type))
6742 {
6743 let expr = DfExpr::Column(Column::new(qualifier.clone(), name));
6744 if data_type == &ArrowDataType::Float64 {
6745 expr.alias(output_col)
6746 } else {
6747 DfExpr::Cast(Cast {
6748 expr: Box::new(expr),
6749 data_type: ArrowDataType::Float64,
6750 })
6751 .alias(output_col)
6752 }
6753 } else {
6754 DfExpr::Literal(ScalarValue::Float64(None), None).alias(output_col)
6755 }
6756 } else {
6757 fields
6758 .iter()
6759 .find(|(_, _, data_type)| data_type == &native_histogram_type)
6760 .map(|(name, qualifier, _)| {
6761 DfExpr::Column(Column::new(qualifier.clone(), name)).alias(output_col)
6762 })
6763 .unwrap_or_else(|| {
6764 DfExpr::Literal(null_histogram.clone(), None).alias(output_col)
6765 })
6766 }
6767 };
6768 let left_proj_exprs = all_columns.iter().map(|col| {
6769 if mixed_sample_types
6770 && (col == &mixed_float_field_col || col == &mixed_histogram_field_col)
6771 {
6772 mixed_value_expr(&left_fields, col)
6773 } else if !mixed_sample_types
6774 && col == left_field_col
6775 && left_field.2 != target_field_type
6776 {
6777 DfExpr::Cast(Cast {
6778 expr: Box::new(DfExpr::Column(Column::new(
6779 left_field.1.clone(),
6780 left_field_col,
6781 ))),
6782 data_type: target_field_type.clone(),
6783 })
6784 .alias(left_field_col.clone())
6785 } else if target_tag_types.contains_key(col) {
6786 aligned_label_expr(col, &left_tag_types)
6787 } else {
6788 DfExpr::Column(Column::new(None::<String>, col))
6789 }
6790 });
6791 let right_time_index_expr = DfExpr::Column(Column::new(
6792 right_qualifier.clone(),
6793 right_time_index_column,
6794 ))
6795 .alias(left_time_index_column.clone());
6796 let right_proj_exprs_without_time_index = all_columns.iter().skip(1).map(|col| {
6800 if mixed_sample_types
6802 && (col == &mixed_float_field_col || col == &mixed_histogram_field_col)
6803 {
6804 mixed_value_expr(&right_fields, col)
6805 } else if !mixed_sample_types && col == left_field_col {
6806 let expr = DfExpr::Column(Column::new(right_field.1.clone(), right_field_col));
6807 if right_field.2 != target_field_type {
6808 DfExpr::Cast(Cast {
6809 expr: Box::new(expr),
6810 data_type: target_field_type.clone(),
6811 })
6812 .alias(left_field_col.clone())
6813 } else if left_field_col != right_field_col {
6814 expr.alias(left_field_col.clone())
6815 } else {
6816 expr
6817 }
6818 } else if target_tag_types.contains_key(col) {
6819 aligned_label_expr(col, &right_tag_types)
6820 } else {
6821 DfExpr::Column(Column::new(None::<String>, col))
6822 }
6823 });
6824 let right_proj_exprs = [right_time_index_expr]
6825 .into_iter()
6826 .chain(right_proj_exprs_without_time_index);
6827
6828 let left_projected = LogicalPlanBuilder::from(left)
6829 .project(left_proj_exprs)
6830 .context(DataFusionPlanningSnafu)?
6831 .alias(left_qualifier_string.clone())
6832 .context(DataFusionPlanningSnafu)?
6833 .build()
6834 .context(DataFusionPlanningSnafu)?;
6835 let right_projected = LogicalPlanBuilder::from(right)
6836 .project(right_proj_exprs)
6837 .context(DataFusionPlanningSnafu)?
6838 .alias(right_qualifier_string.clone())
6839 .context(DataFusionPlanningSnafu)?
6840 .build()
6841 .context(DataFusionPlanningSnafu)?;
6842
6843 let mut match_columns = if let Some(modifier) = modifier
6845 && let Some(matching) = &modifier.matching
6846 {
6847 match matching {
6848 LabelModifier::Include(on) => on.labels.clone(),
6850 LabelModifier::Exclude(ignoring) => {
6852 let ignoring = ignoring.labels.iter().cloned().collect::<HashSet<_>>();
6853 all_tags.difference(&ignoring).cloned().collect()
6854 }
6855 }
6856 } else {
6857 all_tags.iter().cloned().collect()
6858 };
6859 match_columns.sort_unstable();
6861 match_columns.dedup();
6862 occupied_column_names.extend(
6863 left_projected
6864 .schema()
6865 .fields()
6866 .iter()
6867 .chain(right_projected.schema().fields().iter())
6868 .map(|field| field.name().clone()),
6869 );
6870
6871 let visible_schema = left_projected.schema().clone();
6872 let visible_left_exprs = left_projected
6873 .schema()
6874 .iter()
6875 .map(|(qualifier, field)| {
6876 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6877 })
6878 .collect::<Vec<_>>();
6879 let visible_right_exprs = right_projected
6880 .schema()
6881 .iter()
6882 .map(|(qualifier, field)| {
6883 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6884 })
6885 .collect::<Vec<_>>();
6886 let mut left_match_exprs = Vec::with_capacity(match_columns.len());
6887 let mut right_match_exprs = Vec::with_capacity(match_columns.len());
6888 let mut next_internal_column = 0;
6889
6890 for label in &match_columns {
6891 let left_field = if left_tag_cols_set.contains(label) {
6892 Some(
6893 left_projected
6894 .schema()
6895 .iter()
6896 .find(|(_, field)| field.name() == label)
6897 .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone()))
6898 .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?,
6899 )
6900 } else {
6901 None
6902 };
6903 let right_field = if right_tag_cols_set.contains(label) {
6904 Some(
6905 right_projected
6906 .schema()
6907 .iter()
6908 .find(|(_, field)| field.name() == label)
6909 .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone()))
6910 .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?,
6911 )
6912 } else {
6913 None
6914 };
6915 let data_type = match (left_field.as_ref(), right_field.as_ref()) {
6916 (Some((_, left_type)), Some((_, right_type))) if left_type == right_type => {
6917 left_type.clone()
6918 }
6919 (Some((_, left_type)), Some((_, right_type))) => {
6920 return UnexpectedPlanExprSnafu {
6921 desc: format!(
6922 "OR match label {label} has incompatible types: {left_type:?} and {right_type:?}"
6923 ),
6924 }
6925 .fail();
6926 }
6927 (Some((_, data_type)), None) | (None, Some((_, data_type))) => data_type.clone(),
6928 (None, None) => ArrowDataType::Utf8,
6929 };
6930 let Some(value_type) = Self::string_value_data_type(&data_type).cloned() else {
6931 return UnexpectedPlanExprSnafu {
6932 desc: format!("OR match label {label} must be a string"),
6933 }
6934 .fail();
6935 };
6936 let internal_name = loop {
6937 let name = format!("__promql_or_match_{next_internal_column}");
6938 next_internal_column += 1;
6939 if occupied_column_names.insert(name.clone()) {
6940 break name;
6941 }
6942 };
6943 left_match_exprs.push(Self::normalized_match_key_expr(
6944 label,
6945 left_field,
6946 &value_type,
6947 &internal_name,
6948 ));
6949 right_match_exprs.push(Self::normalized_match_key_expr(
6950 label,
6951 right_field,
6952 &value_type,
6953 &internal_name,
6954 ));
6955 }
6956
6957 let left_augmented = LogicalPlanBuilder::from(left_projected)
6958 .project(visible_left_exprs.into_iter().chain(left_match_exprs))
6959 .context(DataFusionPlanningSnafu)?
6960 .build()
6961 .context(DataFusionPlanningSnafu)?;
6962 let right_augmented = LogicalPlanBuilder::from(right_projected)
6963 .project(visible_right_exprs.into_iter().chain(right_match_exprs))
6964 .context(DataFusionPlanningSnafu)?
6965 .build()
6966 .context(DataFusionPlanningSnafu)?;
6967
6968 let visible_field_count = visible_schema.fields().len();
6970 let compare_key_indices =
6971 (visible_field_count..visible_field_count + match_columns.len()).collect::<Vec<_>>();
6972 let (time_qualifier, _) = visible_schema
6973 .iter()
6974 .find(|(_, field)| field.name() == &left_time_index_column)
6975 .with_context(|| TimeIndexNotFoundSnafu {
6976 table: left_qualifier_string.clone(),
6977 })?;
6978 let ts_col_idx = left_augmented
6979 .schema()
6980 .iter()
6981 .position(|(qualifier, field)| {
6982 qualifier == time_qualifier && field.name() == &left_time_index_column
6983 })
6984 .with_context(|| TimeIndexNotFoundSnafu {
6985 table: left_qualifier_string.clone(),
6986 })?;
6987 let union_distinct_on = UnionDistinctOn::try_new(
6988 left_augmented,
6989 right_augmented,
6990 compare_key_indices,
6991 ts_col_idx,
6992 )
6993 .context(DataFusionPlanningSnafu)?;
6994 let augmented_result = LogicalPlan::Extension(Extension {
6995 node: Arc::new(union_distinct_on),
6996 });
6997 let result = LogicalPlanBuilder::from(augmented_result)
6998 .project(visible_schema.iter().map(|(qualifier, field)| {
6999 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
7000 }))
7001 .context(DataFusionPlanningSnafu)?
7002 .build()
7003 .context(DataFusionPlanningSnafu)?;
7004
7005 let output_field_col = left_field_col.clone();
7007 let mut output_context = left_context;
7008 let mut visible_tags = all_tags.into_iter().collect::<Vec<_>>();
7009 visible_tags.sort_unstable();
7010 output_context.time_index_column = Some(left_time_index_column);
7011 output_context.tag_columns = visible_tags;
7012 output_context.field_columns = if mixed_sample_types {
7013 vec![mixed_float_field_col, mixed_histogram_field_col]
7014 } else {
7015 vec![output_field_col]
7016 };
7017 output_context.use_tsid = left_has_tsid && right_has_tsid;
7018 self.ctx = output_context;
7019
7020 Ok(result)
7021 }
7022
7023 fn projection_for_each_field_column<F>(
7031 &mut self,
7032 input: LogicalPlan,
7033 name_to_expr: F,
7034 ) -> Result<LogicalPlan>
7035 where
7036 F: FnMut(&String) -> Result<DfExpr>,
7037 {
7038 let preserve_field_names =
7041 Self::field_columns_are_alternative_samples(input.schema(), &self.ctx.field_columns);
7042 let table_ref = self.ctx.table_name.clone().map(TableReference::bare);
7043 let input_schema = input.schema().clone();
7045 let non_field_columns_iter = self
7046 .ctx
7047 .tag_columns
7048 .iter()
7049 .chain(self.ctx.time_index_column.iter())
7050 .map(|col| {
7051 input_schema
7052 .qualified_field_with_name(table_ref.as_ref(), col)
7053 .or_else(|_| input_schema.qualified_field_with_unqualified_name(col))
7054 .map(|field| DfExpr::Column(field.into()))
7055 .context(DataFusionPlanningSnafu)
7056 });
7057 let tsid_iter =
7058 Self::optional_tsid_projection(input.schema(), table_ref.as_ref(), self.ctx.use_tsid)
7059 .into_iter()
7060 .map(Ok);
7061
7062 let result_field_columns = self
7064 .ctx
7065 .field_columns
7066 .iter()
7067 .map(name_to_expr)
7068 .collect::<Result<Vec<_>>>()?;
7069
7070 if !preserve_field_names {
7072 self.ctx.field_columns = result_field_columns
7073 .iter()
7074 .map(|expr| expr.schema_name().to_string())
7075 .collect();
7076 }
7077 let field_columns_iter = result_field_columns
7078 .into_iter()
7079 .zip(self.ctx.field_columns.iter())
7080 .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::<String>, name))));
7081
7082 let project_fields = non_field_columns_iter
7084 .chain(tsid_iter)
7085 .chain(field_columns_iter)
7086 .collect::<Result<Vec<_>>>()?;
7087
7088 LogicalPlanBuilder::from(input)
7089 .project(project_fields)
7090 .context(DataFusionPlanningSnafu)?
7091 .build()
7092 .context(DataFusionPlanningSnafu)
7093 }
7094
7095 fn filter_on_field_column<F>(&self, input: LogicalPlan, name_to_expr: F) -> Result<LogicalPlan>
7097 where
7098 F: FnMut(&String) -> Result<DfExpr>,
7099 {
7100 ensure!(
7101 self.ctx.field_columns.len() == 1
7102 || Self::field_columns_are_alternative_samples(
7103 input.schema(),
7104 &self.ctx.field_columns,
7105 ),
7106 UnsupportedExprSnafu {
7107 name: "filter on multi-value input"
7108 }
7109 );
7110
7111 let field_column_filters = self
7112 .ctx
7113 .field_columns
7114 .iter()
7115 .map(name_to_expr)
7116 .collect::<Result<Vec<_>>>()?;
7117 let field_column_filter =
7118 disjunction(field_column_filters).context(UnsupportedExprSnafu {
7119 name: "filter on empty input",
7120 })?;
7121
7122 LogicalPlanBuilder::from(input)
7123 .filter(field_column_filter)
7124 .context(DataFusionPlanningSnafu)?
7125 .build()
7126 .context(DataFusionPlanningSnafu)
7127 }
7128
7129 fn date_part_on_time_index(&self, date_part: &str) -> Result<DfExpr> {
7132 let input_expr = datafusion::logical_expr::col(
7133 self.ctx
7134 .time_index_column
7135 .as_ref()
7136 .with_context(|| TimeIndexNotFoundSnafu {
7138 table: "<doesn't matter>",
7139 })?,
7140 );
7141 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
7142 func: datafusion_functions::datetime::date_part(),
7143 args: vec![date_part.lit(), input_expr],
7144 });
7145 Ok(fn_expr)
7146 }
7147
7148 fn strip_tsid_column(&self, plan: LogicalPlan) -> Result<LogicalPlan> {
7149 let schema = plan.schema();
7150 if !schema
7151 .fields()
7152 .iter()
7153 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
7154 {
7155 return Ok(plan);
7156 }
7157
7158 let project_exprs = schema
7161 .iter()
7162 .filter(|(_, field)| field.name() != DATA_SCHEMA_TSID_COLUMN_NAME)
7163 .map(|(qualifier, field)| {
7164 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
7165 })
7166 .collect::<Vec<_>>();
7167
7168 LogicalPlanBuilder::from(plan)
7169 .project(project_exprs)
7170 .context(DataFusionPlanningSnafu)?
7171 .build()
7172 .context(DataFusionPlanningSnafu)
7173 }
7174
7175 fn apply_alias(&mut self, plan: LogicalPlan, alias_name: String) -> Result<LogicalPlan> {
7177 let fields_expr = self.create_field_column_exprs()?;
7178
7179 ensure!(
7181 fields_expr.len() == 1,
7182 UnsupportedExprSnafu {
7183 name: "alias on multi-value result"
7184 }
7185 );
7186
7187 let project_fields = fields_expr
7188 .into_iter()
7189 .map(|expr| expr.alias(&alias_name))
7190 .chain(self.create_tag_column_exprs()?)
7191 .chain(Some(self.create_time_index_column_expr()?));
7192
7193 LogicalPlanBuilder::from(plan)
7194 .project(project_fields)
7195 .context(DataFusionPlanningSnafu)?
7196 .build()
7197 .context(DataFusionPlanningSnafu)
7198 }
7199}
7200
7201#[derive(Default, Debug)]
7202struct FunctionArgs {
7203 input: Option<PromExpr>,
7204 literals: Vec<DfExpr>,
7205}
7206
7207#[derive(Debug, Clone)]
7210enum ScalarFunc {
7211 DataFusionBuiltin(Arc<ScalarUdfDef>),
7215 DataFusionUdf(Arc<ScalarUdfDef>),
7219 NativeHistogramUdf(Arc<ScalarUdfDef>),
7222 Udf(Arc<ScalarUdfDef>),
7227 ExtrapolateUdf(Arc<ScalarUdfDef>, i64),
7234 GeneratedExpr,
7238}
7239
7240#[cfg(test)]
7241mod test {
7242 use std::time::{Duration, UNIX_EPOCH};
7243
7244 use catalog::RegisterTableRequest;
7245 use catalog::memory::{MemoryCatalogManager, new_memory_catalog_manager};
7246 use common_base::Plugins;
7247 use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
7248 use common_query::native_histogram::{
7249 CUSTOM_BUCKETS_SCHEMA, CounterResetHint, NativeHistogram, build_histogram_array,
7250 };
7251 use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value};
7252 use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS;
7253 use common_query::test_util::DummyDecoder;
7254 use common_recordbatch::RecordBatch as GreptimeRecordBatch;
7255 use datafusion::arrow::array::{
7256 Array, Float64Array, Int64Array, StringArray, TimestampMillisecondArray,
7257 };
7258 use datafusion::arrow::datatypes::{Field, Schema as ArrowSchema};
7259 use datafusion::arrow::record_batch::RecordBatch;
7260 use datafusion::catalog::{CatalogProvider, MemoryCatalogProvider, MemorySchemaProvider};
7261 use datafusion::datasource::memory::MemorySourceConfig;
7262 use datafusion::datasource::source::DataSourceExec;
7263 use datafusion::datasource::{MemTable, provider_as_source};
7264 use datafusion::execution::context::SessionContext;
7265 use datafusion::logical_expr::Extension;
7266 use datatypes::prelude::ConcreteDataType;
7267 use datatypes::schema::{ColumnSchema, Schema};
7268 use promql_parser::label::Labels;
7269 use promql_parser::parser;
7270 use session::context::QueryContext;
7271 use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan};
7272 use table::Table;
7273 use table::metadata::{FilterPushDownType, TableInfoBuilder, TableMetaBuilder};
7274 use table::test_util::{EmptyTable, MemTable as GreptimeMemTable};
7275
7276 use super::*;
7277 use crate::QueryEngineContext;
7278 use crate::options::QueryOptions;
7279 use crate::parser::QueryLanguageParser;
7280 use crate::query_engine::DefaultSerializer;
7281
7282 mod delta;
7283
7284 fn find_instant_manipulate(plan: &LogicalPlan) -> Option<&InstantManipulate> {
7285 if let LogicalPlan::Extension(Extension { node }) = plan
7286 && let Some(instant_manipulate) = node.as_any().downcast_ref::<InstantManipulate>()
7287 {
7288 return Some(instant_manipulate);
7289 }
7290
7291 plan.inputs().into_iter().find_map(find_instant_manipulate)
7292 }
7293
7294 fn build_query_engine_state() -> QueryEngineState {
7295 QueryEngineState::new(
7296 new_memory_catalog_manager().unwrap(),
7297 None,
7298 None,
7299 None,
7300 None,
7301 None,
7302 false,
7303 Plugins::default(),
7304 QueryOptions::default(),
7305 )
7306 }
7307
7308 #[test]
7309 fn common_label_type_preserves_only_shared_dictionary_encoding() {
7310 let dictionary = ArrowDataType::Dictionary(
7311 Box::new(ArrowDataType::UInt32),
7312 Box::new(ArrowDataType::Utf8),
7313 );
7314 let other_dictionary = ArrowDataType::Dictionary(
7315 Box::new(ArrowDataType::Int32),
7316 Box::new(ArrowDataType::Utf8),
7317 );
7318
7319 assert_eq!(
7320 Some(dictionary.clone()),
7321 PromPlanner::common_label_data_type(Some(&dictionary), Some(&dictionary))
7322 );
7323 assert_eq!(
7324 Some(ArrowDataType::Utf8),
7325 PromPlanner::common_label_data_type(Some(&dictionary), Some(&ArrowDataType::Utf8))
7326 );
7327 assert_eq!(
7328 Some(ArrowDataType::Utf8),
7329 PromPlanner::common_label_data_type(Some(&dictionary), Some(&other_dictionary))
7330 );
7331 assert_eq!(
7332 Some(ArrowDataType::Utf8),
7333 PromPlanner::common_label_data_type(Some(&dictionary), None)
7334 );
7335 }
7336
7337 async fn build_optimized_promql_plan(
7338 table_provider: DfTableSourceProvider,
7339 eval_stmt: &EvalStmt,
7340 ) -> LogicalPlan {
7341 let state = build_query_engine_state();
7342 let raw_plan = PromPlanner::stmt_to_plan(table_provider, eval_stmt, &state)
7343 .await
7344 .unwrap();
7345 let context = QueryEngineContext::new(state.session_state(), QueryContext::arc());
7346 state
7347 .optimize_by_extension_rules(raw_plan, &context)
7348 .unwrap()
7349 }
7350
7351 async fn build_optimized_tsid_plan(
7352 query: &str,
7353 num_tag: usize,
7354 num_field: usize,
7355 end_secs: u64,
7356 lookback_secs: u64,
7357 ) -> String {
7358 let eval_stmt = EvalStmt {
7359 expr: parser::parse(query).unwrap(),
7360 start: UNIX_EPOCH,
7361 end: UNIX_EPOCH
7362 .checked_add(Duration::from_secs(end_secs))
7363 .unwrap(),
7364 interval: Duration::from_secs(5),
7365 lookback_delta: Duration::from_secs(lookback_secs),
7366 };
7367 let table_provider = build_test_table_provider_with_tsid(
7368 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
7369 num_tag,
7370 num_field,
7371 )
7372 .await;
7373
7374 build_optimized_promql_plan(table_provider, &eval_stmt)
7375 .await
7376 .display_indent_schema()
7377 .to_string()
7378 }
7379
7380 async fn assert_nested_count_rewrite_applies(query: &str, expected_outer_agg: &str) {
7381 let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await;
7382
7383 assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
7384 assert!(plan_str.contains("Projection: some_metric.timestamp, some_metric.tag_0"));
7385 assert!(plan_str.contains("Distinct:"));
7386 assert!(plan_str.contains(expected_outer_agg), "{plan_str}");
7387 assert!(!plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]"));
7388 }
7389
7390 async fn assert_nested_count_rewrite_missing(query: &str, num_tag: usize, lookback_secs: u64) {
7391 let plan_str = build_optimized_tsid_plan(query, num_tag, 1, 100_000, lookback_secs).await;
7392 assert!(!plan_str.contains("Distinct:"), "{plan_str}");
7393 }
7394
7395 fn build_eval_stmt(expr: &str) -> EvalStmt {
7396 EvalStmt {
7397 expr: parser::parse(expr).unwrap(),
7398 start: UNIX_EPOCH,
7399 end: UNIX_EPOCH
7400 .checked_add(Duration::from_secs(100_000))
7401 .unwrap(),
7402 interval: Duration::from_secs(5),
7403 lookback_delta: Duration::from_secs(1),
7404 }
7405 }
7406
7407 enum DirectOrValue {
7408 Float64(f64),
7409 Int64(i64),
7410 NativeHistogram(NativeHistogram),
7411 Utf8(&'static str),
7412 }
7413
7414 impl DirectOrValue {
7415 fn data_type(&self) -> ArrowDataType {
7416 match self {
7417 Self::Float64(_) => ArrowDataType::Float64,
7418 Self::Int64(_) => ArrowDataType::Int64,
7419 Self::NativeHistogram(_) => native_histogram_value_type().as_arrow_type(),
7420 Self::Utf8(_) => ArrowDataType::Utf8,
7421 }
7422 }
7423 fn array(&self) -> Arc<dyn Array> {
7424 match self {
7425 Self::Float64(v) => Arc::new(Float64Array::from(vec![*v])),
7426 Self::Int64(v) => Arc::new(Int64Array::from(vec![*v])),
7427 Self::NativeHistogram(v) => build_histogram_array(&[Some(v.clone())]),
7428 Self::Utf8(v) => Arc::new(StringArray::from(vec![*v])),
7429 }
7430 }
7431 }
7432
7433 fn direct_or_histogram() -> NativeHistogram {
7434 NativeHistogram {
7435 schema: 0,
7436 zero_threshold: 0.0,
7437 sum: 1.0,
7438 reset_hint: CounterResetHint::Unknown,
7439 start_timestamp: None,
7440 custom_values: vec![],
7441 positive_spans: vec![],
7442 negative_spans: vec![],
7443 count: 1.0,
7444 zero_count: 1.0,
7445 positive_buckets: vec![],
7446 negative_buckets: vec![],
7447 }
7448 }
7449
7450 fn operator_metric_table(
7451 name: &str,
7452 table_id: u32,
7453 tag: &str,
7454 le: Option<&str>,
7455 value: DirectOrValue,
7456 ) -> table::TableRef {
7457 let value_type = match &value {
7458 DirectOrValue::Float64(_) => ConcreteDataType::float64_datatype(),
7459 DirectOrValue::Int64(_) => ConcreteDataType::int64_datatype(),
7460 DirectOrValue::NativeHistogram(_) => native_histogram_value_type().clone(),
7461 DirectOrValue::Utf8(_) => ConcreteDataType::string_datatype(),
7462 };
7463 let tag_count = 1 + usize::from(le.is_some());
7464 let mut columns = vec![ColumnSchema::new(
7465 "tag".to_string(),
7466 ConcreteDataType::string_datatype(),
7467 false,
7468 )];
7469 if le.is_some() {
7470 columns.push(ColumnSchema::new(
7471 LE_COLUMN_NAME.to_string(),
7472 ConcreteDataType::string_datatype(),
7473 false,
7474 ));
7475 }
7476 columns.extend([
7477 ColumnSchema::new(
7478 "ts".to_string(),
7479 ConcreteDataType::timestamp_millisecond_datatype(),
7480 false,
7481 )
7482 .with_time_index(true),
7483 ColumnSchema::new("v".to_string(), value_type, true),
7484 ]);
7485 let schema = Arc::new(Schema::new(columns));
7486 let mut arrays = vec![Arc::new(StringArray::from(vec![tag])) as Arc<dyn Array>];
7487 if let Some(le) = le {
7488 arrays.push(Arc::new(StringArray::from(vec![le])));
7489 }
7490 arrays.extend([
7491 Arc::new(TimestampMillisecondArray::from(vec![1_000])) as Arc<dyn Array>,
7492 value.array(),
7493 ]);
7494 let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap();
7495 let backing = GreptimeMemTable::new_with_catalog(
7496 name,
7497 GreptimeRecordBatch::from_df_record_batch(schema.clone(), batch),
7498 table_id,
7499 DEFAULT_CATALOG_NAME.to_string(),
7500 DEFAULT_SCHEMA_NAME.to_string(),
7501 );
7502 let value_index = tag_count + 1;
7503 let meta = TableMetaBuilder::empty()
7504 .schema(schema)
7505 .primary_key_indices((0..tag_count).collect())
7506 .value_indices(vec![value_index])
7507 .next_column_id((value_index + 1) as u32)
7508 .build()
7509 .unwrap();
7510 let info = Arc::new(
7511 TableInfoBuilder::default()
7512 .table_id(table_id)
7513 .name(name)
7514 .meta(meta)
7515 .build()
7516 .unwrap(),
7517 );
7518 Arc::new(Table::new(
7519 info,
7520 FilterPushDownType::Unsupported,
7521 backing.data_source(),
7522 ))
7523 }
7524
7525 fn operator_table_provider() -> DfTableSourceProvider {
7526 let catalog = MemoryCatalogManager::with_default_setup();
7527 let tables = [
7528 operator_metric_table("lf", 2_001, "a", None, DirectOrValue::Float64(2.0)),
7529 operator_metric_table(
7530 "lh",
7531 2_002,
7532 "b",
7533 None,
7534 DirectOrValue::NativeHistogram(direct_or_histogram()),
7535 ),
7536 operator_metric_table("rf", 2_003, "b", None, DirectOrValue::Float64(3.0)),
7537 operator_metric_table(
7538 "rh",
7539 2_004,
7540 "a",
7541 None,
7542 DirectOrValue::NativeHistogram(direct_or_histogram()),
7543 ),
7544 operator_metric_table("fallback", 2_005, "c", None, DirectOrValue::Float64(7.0)),
7545 operator_metric_table(
7546 "bad_classic",
7547 2_006,
7548 "d",
7549 Some("broken"),
7550 DirectOrValue::Float64(1.0),
7551 ),
7552 operator_metric_table(
7553 "bad_native",
7554 2_007,
7555 "d",
7556 None,
7557 DirectOrValue::NativeHistogram(direct_or_histogram()),
7558 ),
7559 ];
7560 for table in tables {
7561 let info = table.table_info();
7562 catalog
7563 .register_table_sync(RegisterTableRequest {
7564 catalog: DEFAULT_CATALOG_NAME.to_string(),
7565 schema: DEFAULT_SCHEMA_NAME.to_string(),
7566 table_name: info.name.clone(),
7567 table_id: info.ident.table_id,
7568 table,
7569 })
7570 .unwrap();
7571 }
7572 DfTableSourceProvider::new(
7573 catalog,
7574 false,
7575 QueryContext::arc(),
7576 DummyDecoder::arc(),
7577 false,
7578 )
7579 }
7580
7581 fn operator_eval_stmt(expr: &str) -> EvalStmt {
7582 let time = UNIX_EPOCH.checked_add(Duration::from_secs(1)).unwrap();
7583 EvalStmt {
7584 expr: parser::parse(expr).unwrap(),
7585 start: time,
7586 end: time,
7587 interval: Duration::from_secs(1),
7588 lookback_delta: Duration::from_secs(5),
7589 }
7590 }
7591
7592 struct DirectOrSource {
7593 name: &'static str,
7594 empty: bool,
7595 timestamp: i64,
7596 tags: Vec<(&'static str, Option<&'static str>)>,
7597 value: DirectOrValue,
7598 }
7599
7600 fn source(
7601 name: &'static str,
7602 empty: bool,
7603 timestamp: i64,
7604 tags: Vec<(&'static str, Option<&'static str>)>,
7605 value: DirectOrValue,
7606 ) -> DirectOrSource {
7607 DirectOrSource {
7608 name,
7609 empty,
7610 timestamp,
7611 tags,
7612 value,
7613 }
7614 }
7615
7616 fn tagged_source(
7617 name: &'static str,
7618 empty: bool,
7619 tag: (&'static str, Option<&'static str>),
7620 value: DirectOrValue,
7621 ) -> DirectOrSource {
7622 source(name, empty, 1, vec![("job", Some("job")), tag], value)
7623 }
7624
7625 fn job_source(name: &'static str, value: DirectOrValue) -> DirectOrSource {
7626 source(name, true, 1, vec![("job", Some("job"))], value)
7627 }
7628
7629 fn table(source: &DirectOrSource) -> Arc<MemTable> {
7630 let mut fields = vec![Field::new(
7631 "ts",
7632 ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
7633 false,
7634 )];
7635 fields.extend(
7636 source
7637 .tags
7638 .iter()
7639 .map(|(name, _)| Field::new(*name, ArrowDataType::Utf8, true)),
7640 );
7641 fields.push(Field::new("v", source.value.data_type(), true));
7642 let schema = Arc::new(ArrowSchema::new(fields));
7643 let partitions = if source.empty {
7644 vec![vec![]]
7645 } else {
7646 let mut columns: Vec<Arc<dyn Array>> =
7647 vec![Arc::new(TimestampMillisecondArray::from(vec![
7648 source.timestamp,
7649 ]))];
7650 columns.extend(
7651 source
7652 .tags
7653 .iter()
7654 .map(|(_, value)| Arc::new(StringArray::from(vec![*value])) as Arc<dyn Array>),
7655 );
7656 columns.push(source.value.array());
7657 vec![vec![RecordBatch::try_new(schema.clone(), columns).unwrap()]]
7658 };
7659 Arc::new(MemTable::try_new(schema, partitions).unwrap())
7660 }
7661
7662 fn scan(source: &DirectOrSource) -> LogicalPlan {
7663 LogicalPlanBuilder::scan(source.name, provider_as_source(table(source)), None)
7664 .unwrap()
7665 .build()
7666 .unwrap()
7667 }
7668
7669 fn direct_or_context(qualifier: &str, tags: &[&str], field: &str) -> PromPlannerContext {
7670 PromPlannerContext {
7671 table_name: Some(qualifier.to_string()),
7672 time_index_column: Some("ts".to_string()),
7673 field_columns: vec![field.to_string()],
7674 tag_columns: tags.iter().map(|tag| (*tag).to_string()).collect(),
7675 ..Default::default()
7676 }
7677 }
7678
7679 fn or_modifier(expr: &str) -> Option<BinModifier> {
7680 let PromExpr::Binary(expr) = parser::parse(expr).unwrap() else {
7681 unreachable!()
7682 };
7683 expr.modifier
7684 }
7685
7686 async fn plan_direct_or(
7687 left: LogicalPlan,
7688 right: LogicalPlan,
7689 left_context: PromPlannerContext,
7690 right_context: PromPlannerContext,
7691 modifier: &Option<BinModifier>,
7692 ) -> LogicalPlan {
7693 let table_provider = build_test_table_provider_with_fields(
7694 &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
7695 &[],
7696 )
7697 .await;
7698 let mut planner = PromPlanner {
7699 table_provider,
7700 ctx: PromPlannerContext::default(),
7701 promql_annotations: None,
7702 };
7703 planner
7704 .or_operator(
7705 left,
7706 right,
7707 left_context.tag_columns.iter().cloned().collect(),
7708 right_context.tag_columns.iter().cloned().collect(),
7709 left_context,
7710 right_context,
7711 modifier,
7712 )
7713 .unwrap()
7714 }
7715
7716 async fn execute(
7717 plan: LogicalPlan,
7718 state: &QueryEngineState,
7719 ) -> (LogicalPlan, Vec<RecordBatch>) {
7720 let context = QueryEngineContext::new(state.session_state(), QueryContext::arc());
7721 let optimized = state.optimize_by_extension_rules(plan, &context).unwrap();
7722 let physical = state
7723 .session_state()
7724 .create_physical_plan(&optimized)
7725 .await
7726 .unwrap();
7727 let batches =
7728 datafusion::physical_plan::collect(physical, state.session_state().task_ctx())
7729 .await
7730 .unwrap();
7731 (optimized, batches)
7732 }
7733
7734 async fn run(
7735 left: &DirectOrSource,
7736 right: &DirectOrSource,
7737 left_context: PromPlannerContext,
7738 right_context: PromPlannerContext,
7739 modifier: &Option<BinModifier>,
7740 ) -> (LogicalPlan, Vec<RecordBatch>) {
7741 let plan = plan_direct_or(
7742 scan(left),
7743 scan(right),
7744 left_context,
7745 right_context,
7746 modifier,
7747 )
7748 .await;
7749 execute(plan, &build_query_engine_state()).await
7750 }
7751
7752 async fn mixed_direct_or(histogram_on_left: bool) -> (PromPlanner, LogicalPlan) {
7753 let sample = |histogram: bool| {
7754 if histogram {
7755 DirectOrValue::NativeHistogram(direct_or_histogram())
7756 } else {
7757 DirectOrValue::Float64(1.25)
7758 }
7759 };
7760 let left = tagged_source(
7761 "lhs",
7762 false,
7763 (
7764 "k",
7765 Some(if histogram_on_left {
7766 "histogram"
7767 } else {
7768 "float"
7769 }),
7770 ),
7771 sample(histogram_on_left),
7772 );
7773 let right = tagged_source(
7774 "rhs",
7775 false,
7776 (
7777 "k",
7778 Some(if histogram_on_left {
7779 "float"
7780 } else {
7781 "histogram"
7782 }),
7783 ),
7784 sample(!histogram_on_left),
7785 );
7786 let table_provider = build_test_table_provider_with_fields(
7787 &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
7788 &[],
7789 )
7790 .await;
7791 let mut planner = PromPlanner {
7792 table_provider,
7793 ctx: PromPlannerContext::default(),
7794 promql_annotations: None,
7795 };
7796 let left_context = direct_or_context("lhs", &["job", "k"], "v");
7797 let right_context = direct_or_context("rhs", &["job", "k"], "v");
7798 let plan = planner
7799 .or_operator(
7800 scan(&left),
7801 scan(&right),
7802 left_context.tag_columns.iter().cloned().collect(),
7803 right_context.tag_columns.iter().cloned().collect(),
7804 left_context,
7805 right_context,
7806 &or_modifier("lhs or on(k) rhs"),
7807 )
7808 .unwrap();
7809 (planner, plan)
7810 }
7811
7812 async fn mixed_aggregate_input(histograms: Vec<NativeHistogram>) -> (PromPlanner, LogicalPlan) {
7813 let float_field = format!("{OR_FLOAT_FIELD_PREFIX}0");
7814 let histogram_field = format!("{OR_HISTOGRAM_FIELD_PREFIX}0");
7815 let row_count = histograms.len() + 1;
7816 let schema = Arc::new(ArrowSchema::new(vec![
7817 Field::new(
7818 "ts",
7819 ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
7820 false,
7821 ),
7822 Field::new("k", ArrowDataType::Utf8, false),
7823 Field::new(&float_field, ArrowDataType::Float64, true),
7824 Field::new(
7825 &histogram_field,
7826 native_histogram_value_type().as_arrow_type(),
7827 true,
7828 ),
7829 ]));
7830 let mut histogram_values = Vec::with_capacity(row_count);
7831 histogram_values.push(None);
7832 histogram_values.extend(histograms.into_iter().map(Some));
7833 let batch = RecordBatch::try_new(
7834 schema.clone(),
7835 vec![
7836 Arc::new(TimestampMillisecondArray::from(vec![1; row_count])),
7837 Arc::new(StringArray::from_iter_values(
7838 (0..row_count).map(|row| format!("kind_{row}")),
7839 )),
7840 Arc::new(Float64Array::from_iter(
7841 (0..row_count).map(|row| (row == 0).then_some(1.25)),
7842 )),
7843 build_histogram_array(&histogram_values),
7844 ],
7845 )
7846 .unwrap();
7847 let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap());
7848 let plan = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None)
7849 .unwrap()
7850 .build()
7851 .unwrap();
7852 let table_provider = build_test_table_provider_with_fields(
7853 &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
7854 &[],
7855 )
7856 .await;
7857 let planner = PromPlanner {
7858 table_provider,
7859 ctx: PromPlannerContext {
7860 table_name: Some("mixed".to_string()),
7861 time_index_column: Some("ts".to_string()),
7862 field_columns: vec![float_field, histogram_field],
7863 tag_columns: vec!["k".to_string()],
7864 ..Default::default()
7865 },
7866 promql_annotations: None,
7867 };
7868 (planner, plan)
7869 }
7870
7871 fn assert_no_internal_or_keys(schema: &DFSchema) {
7872 assert!(
7873 schema
7874 .fields()
7875 .iter()
7876 .all(|field| !field.name().starts_with("__promql_or_match_")),
7877 "{schema:?}"
7878 );
7879 }
7880
7881 fn values(batches: &[RecordBatch], column: &str) -> Vec<f64> {
7882 batches
7883 .iter()
7884 .flat_map(|batch| {
7885 batch
7886 .column_by_name(column)
7887 .unwrap()
7888 .as_any()
7889 .downcast_ref::<Float64Array>()
7890 .unwrap()
7891 .iter()
7892 .flatten()
7893 })
7894 .collect()
7895 }
7896
7897 fn numeric_values(batches: &[RecordBatch], column: &str) -> Vec<f64> {
7898 batches
7899 .iter()
7900 .flat_map(|batch| {
7901 let values = datafusion::arrow::compute::cast(
7902 batch.column_by_name(column).unwrap(),
7903 &ArrowDataType::Float64,
7904 )
7905 .unwrap();
7906 values
7907 .as_any()
7908 .downcast_ref::<Float64Array>()
7909 .unwrap()
7910 .iter()
7911 .flatten()
7912 .collect::<Vec<_>>()
7913 })
7914 .collect()
7915 }
7916
7917 fn histograms(batches: &[RecordBatch], column: &str) -> Vec<NativeHistogram> {
7918 batches
7919 .iter()
7920 .flat_map(|batch| {
7921 let values = batch
7922 .column_by_name(column)
7923 .unwrap()
7924 .as_any()
7925 .downcast_ref::<datafusion::arrow::array::StructArray>()
7926 .unwrap();
7927 (0..values.len()).filter_map(|row| {
7928 common_query::native_histogram::read_histogram(values, row).unwrap()
7929 })
7930 })
7931 .collect()
7932 }
7933
7934 fn rows(batches: &[RecordBatch]) -> Vec<(f64, Option<String>)> {
7935 let mut rows = batches
7936 .iter()
7937 .flat_map(|batch| {
7938 let values = batch
7939 .column_by_name("v")
7940 .unwrap()
7941 .as_any()
7942 .downcast_ref::<Float64Array>()
7943 .unwrap();
7944 let labels = batch
7945 .column_by_name("k")
7946 .map(|column| column.as_any().downcast_ref::<StringArray>().unwrap());
7947 (0..batch.num_rows()).map(move |i| {
7948 (
7949 values.value(i),
7950 labels.and_then(|labels| {
7951 (!labels.is_null(i)).then(|| labels.value(i).to_string())
7952 }),
7953 )
7954 })
7955 })
7956 .collect::<Vec<_>>();
7957 rows.sort_by(|left, right| left.0.total_cmp(&right.0));
7958 rows
7959 }
7960
7961 fn matrix_source(
7962 name: &'static str,
7963 k: Option<Option<&'static str>>,
7964 timestamp: i64,
7965 value: f64,
7966 ) -> DirectOrSource {
7967 let mut tags = vec![("job", Some("job"))];
7968 if let Some(k) = k {
7969 tags.push(("k", k));
7970 }
7971 source(name, false, timestamp, tags, DirectOrValue::Float64(value))
7972 }
7973
7974 fn matrix_context(name: &str, k: Option<Option<&str>>) -> PromPlannerContext {
7975 direct_or_context(
7976 name,
7977 if k.is_some() { &["job", "k"] } else { &["job"] },
7978 "v",
7979 )
7980 }
7981
7982 async fn build_missing_le_or_normal_metric_table_provider() -> DfTableSourceProvider {
7983 build_test_table_provider_with_fields(
7984 &[
7985 (
7986 DEFAULT_SCHEMA_NAME.to_string(),
7987 "non_existent_histogram_bucket".to_string(),
7988 ),
7989 (DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string()),
7990 ],
7991 &["pod", "instance"],
7992 )
7993 .await
7994 }
7995
7996 fn assert_normal_metric_schema(plan: &LogicalPlan) {
7997 let fields = plan.schema().fields();
7998 assert_eq!(fields.len(), 4, "{fields:?}");
7999 assert!(
8000 fields.iter().any(|field| field.name() == "pod"),
8001 "{fields:?}"
8002 );
8003 assert!(
8004 fields.iter().any(|field| field.name() == "instance"),
8005 "{fields:?}"
8006 );
8007 assert!(
8008 fields
8009 .iter()
8010 .any(|field| field.name() == greptime_timestamp()),
8011 "{fields:?}"
8012 );
8013 assert!(
8014 fields.iter().any(|field| {
8015 field.name() == greptime_value() && field.data_type() == &ArrowDataType::Float64
8016 }),
8017 "{fields:?}"
8018 );
8019 }
8020
8021 async fn build_test_table_provider_with_distinct_tags(
8022 table_tags: &[(&str, &[&str])],
8023 ) -> DfTableSourceProvider {
8024 let catalog_list = MemoryCatalogManager::with_default_setup();
8025 for (table_name, tags) in table_tags {
8026 let mut columns = tags
8027 .iter()
8028 .map(|tag| {
8029 ColumnSchema::new(
8030 (*tag).to_string(),
8031 ConcreteDataType::string_datatype(),
8032 false,
8033 )
8034 })
8035 .collect::<Vec<_>>();
8036 columns.push(
8037 ColumnSchema::new(
8038 greptime_timestamp().to_string(),
8039 ConcreteDataType::timestamp_millisecond_datatype(),
8040 false,
8041 )
8042 .with_time_index(true),
8043 );
8044 columns.push(ColumnSchema::new(
8045 greptime_value().to_string(),
8046 ConcreteDataType::float64_datatype(),
8047 true,
8048 ));
8049 let table_meta = TableMetaBuilder::empty()
8050 .schema(Arc::new(Schema::new(columns)))
8051 .primary_key_indices((0..tags.len()).collect())
8052 .next_column_id(1024)
8053 .build()
8054 .unwrap();
8055 let table_info = TableInfoBuilder::default()
8056 .name((*table_name).to_string())
8057 .meta(table_meta)
8058 .build()
8059 .unwrap();
8060
8061 assert!(
8062 catalog_list
8063 .register_table_sync(RegisterTableRequest {
8064 catalog: DEFAULT_CATALOG_NAME.to_string(),
8065 schema: DEFAULT_SCHEMA_NAME.to_string(),
8066 table_name: (*table_name).to_string(),
8067 table_id: 1024,
8068 table: EmptyTable::from_table_info(&table_info),
8069 })
8070 .is_ok()
8071 );
8072 }
8073
8074 DfTableSourceProvider::new(
8075 catalog_list,
8076 false,
8077 QueryContext::arc(),
8078 DummyDecoder::arc(),
8079 false,
8080 )
8081 }
8082
8083 fn contains_histogram_fold(plan: &LogicalPlan) -> bool {
8084 matches!(plan, LogicalPlan::Extension(Extension { node }) if node.as_any().is::<HistogramFold>())
8085 || plan.inputs().into_iter().any(contains_histogram_fold)
8086 }
8087
8088 async fn build_set_op_context_table_provider() -> DfTableSourceProvider {
8089 build_test_table_provider_with_distinct_tags(&[
8090 ("bucket_metric", &["job", "le"]),
8091 ("normal_metric", &["job"]),
8092 ("fallback_metric", &["instance"]),
8093 ])
8094 .await
8095 }
8096
8097 async fn build_or_context_table_provider() -> DfTableSourceProvider {
8098 build_test_table_provider_with_distinct_tags(&[
8099 ("normal_metric", &["job"]),
8100 ("other_metric", &["instance"]),
8101 ("non_hist_metric", &["instance"]),
8102 ])
8103 .await
8104 }
8105
8106 async fn optimize_and_create_physical_plan(
8107 state: &QueryEngineState,
8108 plan: LogicalPlan,
8109 ) -> (
8110 LogicalPlan,
8111 Arc<dyn datafusion::physical_plan::ExecutionPlan>,
8112 ) {
8113 let context = QueryEngineContext::new(state.session_state(), QueryContext::arc());
8114 let optimized = state.optimize_by_extension_rules(plan, &context).unwrap();
8115 let physical = state
8116 .session_state()
8117 .create_physical_plan(&optimized)
8118 .await
8119 .unwrap();
8120 (optimized, physical)
8121 }
8122
8123 async fn build_test_table_provider(
8124 table_name_tuples: &[(String, String)],
8125 num_tag: usize,
8126 num_field: usize,
8127 ) -> DfTableSourceProvider {
8128 let catalog_list = MemoryCatalogManager::with_default_setup();
8129 for (schema_name, table_name) in table_name_tuples {
8130 let mut columns = vec![];
8131 for i in 0..num_tag {
8132 columns.push(ColumnSchema::new(
8133 format!("tag_{i}"),
8134 ConcreteDataType::string_datatype(),
8135 false,
8136 ));
8137 }
8138 columns.push(
8139 ColumnSchema::new(
8140 "timestamp".to_string(),
8141 ConcreteDataType::timestamp_millisecond_datatype(),
8142 false,
8143 )
8144 .with_time_index(true),
8145 );
8146 for i in 0..num_field {
8147 columns.push(ColumnSchema::new(
8148 format!("field_{i}"),
8149 ConcreteDataType::float64_datatype(),
8150 true,
8151 ));
8152 }
8153 let schema = Arc::new(Schema::new(columns));
8154 let table_meta = TableMetaBuilder::empty()
8155 .schema(schema)
8156 .primary_key_indices((0..num_tag).collect())
8157 .value_indices((num_tag + 1..num_tag + 1 + num_field).collect())
8158 .next_column_id(1024)
8159 .build()
8160 .unwrap();
8161 let table_info = TableInfoBuilder::default()
8162 .name(table_name.clone())
8163 .meta(table_meta)
8164 .build()
8165 .unwrap();
8166 let table = EmptyTable::from_table_info(&table_info);
8167
8168 assert!(
8169 catalog_list
8170 .register_table_sync(RegisterTableRequest {
8171 catalog: DEFAULT_CATALOG_NAME.to_string(),
8172 schema: schema_name.clone(),
8173 table_name: table_name.clone(),
8174 table_id: 1024,
8175 table,
8176 })
8177 .is_ok()
8178 );
8179 }
8180
8181 DfTableSourceProvider::new(
8182 catalog_list,
8183 false,
8184 QueryContext::arc(),
8185 DummyDecoder::arc(),
8186 false,
8187 )
8188 }
8189
8190 async fn build_test_native_histogram_table_provider(table_name: &str) -> DfTableSourceProvider {
8191 build_test_native_histogram_table_provider_with_marker(table_name, false).await
8192 }
8193
8194 async fn build_test_native_histogram_table_provider_with_marker(
8195 table_name: &str,
8196 temporality_marker: bool,
8197 ) -> DfTableSourceProvider {
8198 let catalog_list = MemoryCatalogManager::with_default_setup();
8199 let mut columns = vec![
8200 ColumnSchema::new(
8201 "tag_0".to_string(),
8202 ConcreteDataType::string_datatype(),
8203 false,
8204 ),
8205 ColumnSchema::new(
8206 LE_COLUMN_NAME.to_string(),
8207 ConcreteDataType::string_datatype(),
8208 true,
8209 ),
8210 ];
8211 if temporality_marker {
8212 columns.push(ColumnSchema::new(
8213 OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(),
8214 ConcreteDataType::string_datatype(),
8215 true,
8216 ));
8217 }
8218 let tag_count = columns.len();
8219 columns.extend([
8220 ColumnSchema::new(
8221 "timestamp".to_string(),
8222 ConcreteDataType::timestamp_millisecond_datatype(),
8223 false,
8224 )
8225 .with_time_index(true),
8226 ColumnSchema::new(
8227 greptime_native_histogram().to_string(),
8228 native_histogram_value_type().clone(),
8229 true,
8230 ),
8231 ]);
8232 let schema = Arc::new(Schema::new(columns));
8233 let table_meta = TableMetaBuilder::empty()
8234 .schema(schema)
8235 .primary_key_indices((0..tag_count).collect())
8236 .value_indices(vec![tag_count + 1])
8237 .next_column_id(1024)
8238 .build()
8239 .unwrap();
8240 let table_info = TableInfoBuilder::default()
8241 .name(table_name)
8242 .meta(table_meta)
8243 .build()
8244 .unwrap();
8245 let table = EmptyTable::from_table_info(&table_info);
8246
8247 assert!(
8248 catalog_list
8249 .register_table_sync(RegisterTableRequest {
8250 catalog: DEFAULT_CATALOG_NAME.to_string(),
8251 schema: DEFAULT_SCHEMA_NAME.to_string(),
8252 table_name: table_name.to_string(),
8253 table_id: 1024,
8254 table,
8255 })
8256 .is_ok()
8257 );
8258
8259 DfTableSourceProvider::new(
8260 catalog_list,
8261 false,
8262 QueryContext::arc(),
8263 DummyDecoder::arc(),
8264 false,
8265 )
8266 }
8267
8268 async fn build_test_multi_histogram_table_provider(table_name: &str) -> DfTableSourceProvider {
8269 let catalog_list = MemoryCatalogManager::with_default_setup();
8270 let columns = vec![
8271 ColumnSchema::new(
8272 "tag_0".to_string(),
8273 ConcreteDataType::string_datatype(),
8274 false,
8275 ),
8276 ColumnSchema::new(
8277 "timestamp".to_string(),
8278 ConcreteDataType::timestamp_millisecond_datatype(),
8279 false,
8280 )
8281 .with_time_index(true),
8282 ColumnSchema::new(
8283 greptime_native_histogram().to_string(),
8284 native_histogram_value_type().clone(),
8285 true,
8286 ),
8287 ColumnSchema::new(
8288 "native_histogram_2".to_string(),
8289 native_histogram_value_type().clone(),
8290 true,
8291 ),
8292 ];
8293 let schema = Arc::new(Schema::new(columns));
8294 let table_meta = TableMetaBuilder::empty()
8295 .schema(schema)
8296 .primary_key_indices(vec![0])
8297 .value_indices(vec![2, 3])
8298 .next_column_id(1024)
8299 .build()
8300 .unwrap();
8301 let table_info = TableInfoBuilder::default()
8302 .name(table_name)
8303 .meta(table_meta)
8304 .build()
8305 .unwrap();
8306 let table = EmptyTable::from_table_info(&table_info);
8307
8308 assert!(
8309 catalog_list
8310 .register_table_sync(RegisterTableRequest {
8311 catalog: DEFAULT_CATALOG_NAME.to_string(),
8312 schema: DEFAULT_SCHEMA_NAME.to_string(),
8313 table_name: table_name.to_string(),
8314 table_id: 1024,
8315 table,
8316 })
8317 .is_ok()
8318 );
8319
8320 DfTableSourceProvider::new(
8321 catalog_list,
8322 false,
8323 QueryContext::arc(),
8324 DummyDecoder::arc(),
8325 false,
8326 )
8327 }
8328
8329 async fn build_test_mixed_native_histogram_table_provider(
8330 table_name: &str,
8331 ) -> DfTableSourceProvider {
8332 build_test_mixed_native_histogram_table_provider_with_marker(table_name, false).await
8333 }
8334
8335 async fn build_test_mixed_native_histogram_table_provider_with_marker(
8336 table_name: &str,
8337 temporality_marker: bool,
8338 ) -> DfTableSourceProvider {
8339 let catalog_list = MemoryCatalogManager::with_default_setup();
8340 let mut columns = vec![ColumnSchema::new(
8341 "tag_0".to_string(),
8342 ConcreteDataType::string_datatype(),
8343 false,
8344 )];
8345 if temporality_marker {
8346 columns.push(ColumnSchema::new(
8347 OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(),
8348 ConcreteDataType::string_datatype(),
8349 true,
8350 ));
8351 }
8352 let tag_count = columns.len();
8353 columns.extend([
8354 ColumnSchema::new(
8355 "timestamp".to_string(),
8356 ConcreteDataType::timestamp_millisecond_datatype(),
8357 false,
8358 )
8359 .with_time_index(true),
8360 ColumnSchema::new(
8361 greptime_native_histogram().to_string(),
8362 native_histogram_value_type().clone(),
8363 true,
8364 ),
8365 ColumnSchema::new(
8366 greptime_value().to_string(),
8367 ConcreteDataType::float64_datatype(),
8368 true,
8369 ),
8370 ]);
8371 let schema = Arc::new(Schema::new(columns));
8372 let table_meta = TableMetaBuilder::empty()
8373 .schema(schema.clone())
8374 .primary_key_indices((0..tag_count).collect())
8375 .value_indices(vec![tag_count + 1, tag_count + 2])
8376 .next_column_id(1024)
8377 .build()
8378 .unwrap();
8379 let table_info = Arc::new(
8380 TableInfoBuilder::default()
8381 .name(table_name)
8382 .meta(table_meta)
8383 .build()
8384 .unwrap(),
8385 );
8386 let mut arrays: Vec<Arc<dyn Array>> =
8387 vec![Arc::new(StringArray::from(vec!["float", "histogram"]))];
8388 if temporality_marker {
8389 arrays.push(Arc::new(StringArray::from(vec![
8390 Some(GREPTIME_TEMPORALITY_DELTA),
8391 Some(GREPTIME_TEMPORALITY_DELTA),
8392 ])));
8393 }
8394 arrays.extend([
8395 Arc::new(TimestampMillisecondArray::from(vec![1_000, 1_000])) as Arc<dyn Array>,
8396 build_histogram_array(&[None, Some(direct_or_histogram())]),
8397 Arc::new(Float64Array::from(vec![Some(2.0), None])),
8398 ]);
8399 let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap();
8400 let backing = GreptimeMemTable::new_with_catalog(
8401 table_name,
8402 GreptimeRecordBatch::from_df_record_batch(schema, batch),
8403 1024,
8404 DEFAULT_CATALOG_NAME.to_string(),
8405 DEFAULT_SCHEMA_NAME.to_string(),
8406 );
8407 let table = Arc::new(Table::new(
8408 table_info,
8409 FilterPushDownType::Unsupported,
8410 backing.data_source(),
8411 ));
8412
8413 assert!(
8414 catalog_list
8415 .register_table_sync(RegisterTableRequest {
8416 catalog: DEFAULT_CATALOG_NAME.to_string(),
8417 schema: DEFAULT_SCHEMA_NAME.to_string(),
8418 table_name: table_name.to_string(),
8419 table_id: 1024,
8420 table,
8421 })
8422 .is_ok()
8423 );
8424
8425 DfTableSourceProvider::new(
8426 catalog_list,
8427 false,
8428 QueryContext::arc(),
8429 DummyDecoder::arc(),
8430 false,
8431 )
8432 }
8433
8434 fn classic_and_native_histogram_table_provider(
8435 native_tag: &str,
8436 native_le: Option<&str>,
8437 native_histogram: NativeHistogram,
8438 ) -> DfTableSourceProvider {
8439 let table_name = "mixed_histogram";
8440 let catalog = MemoryCatalogManager::with_default_setup();
8441 let schema = Arc::new(Schema::new(vec![
8442 ColumnSchema::new(
8443 "tag".to_string(),
8444 ConcreteDataType::string_datatype(),
8445 false,
8446 ),
8447 ColumnSchema::new(
8448 LE_COLUMN_NAME.to_string(),
8449 ConcreteDataType::string_datatype(),
8450 true,
8451 ),
8452 ColumnSchema::new(
8453 "timestamp".to_string(),
8454 ConcreteDataType::timestamp_millisecond_datatype(),
8455 false,
8456 )
8457 .with_time_index(true),
8458 ColumnSchema::new(
8459 greptime_native_histogram().to_string(),
8460 native_histogram_value_type().clone(),
8461 true,
8462 ),
8463 ColumnSchema::new(
8464 greptime_value().to_string(),
8465 ConcreteDataType::float64_datatype(),
8466 true,
8467 ),
8468 ]));
8469 let table_meta = TableMetaBuilder::empty()
8470 .schema(schema.clone())
8471 .primary_key_indices(vec![0, 1])
8472 .value_indices(vec![3, 4])
8473 .next_column_id(5)
8474 .build()
8475 .unwrap();
8476 let table_info = Arc::new(
8477 TableInfoBuilder::default()
8478 .name(table_name)
8479 .meta(table_meta)
8480 .build()
8481 .unwrap(),
8482 );
8483 let batch = RecordBatch::try_new(
8484 schema.arrow_schema().clone(),
8485 vec![
8486 Arc::new(StringArray::from(vec![
8487 "classic", "classic", native_tag, "classic", "classic", native_tag,
8488 ])),
8489 Arc::new(StringArray::from(vec![
8490 Some("1"),
8491 Some("+Inf"),
8492 native_le,
8493 Some("1"),
8494 Some("+Inf"),
8495 native_le,
8496 ])),
8497 Arc::new(TimestampMillisecondArray::from(vec![
8498 1_000, 1_000, 1_000, 2_000, 2_000, 2_000,
8499 ])),
8500 build_histogram_array(&[
8501 None,
8502 None,
8503 Some(native_histogram.clone()),
8504 None,
8505 None,
8506 Some(native_histogram),
8507 ]),
8508 Arc::new(Float64Array::from(vec![
8509 Some(2.0),
8510 Some(4.0),
8511 None,
8512 Some(2.0),
8513 Some(4.0),
8514 None,
8515 ])),
8516 ],
8517 )
8518 .unwrap();
8519 let backing = GreptimeMemTable::new_with_catalog(
8520 table_name,
8521 GreptimeRecordBatch::from_df_record_batch(schema, batch),
8522 2_200,
8523 DEFAULT_CATALOG_NAME.to_string(),
8524 DEFAULT_SCHEMA_NAME.to_string(),
8525 );
8526 let table = Arc::new(Table::new(
8527 table_info,
8528 FilterPushDownType::Unsupported,
8529 backing.data_source(),
8530 ));
8531 catalog
8532 .register_table_sync(RegisterTableRequest {
8533 catalog: DEFAULT_CATALOG_NAME.to_string(),
8534 schema: DEFAULT_SCHEMA_NAME.to_string(),
8535 table_name: table_name.to_string(),
8536 table_id: 2_200,
8537 table,
8538 })
8539 .unwrap();
8540
8541 DfTableSourceProvider::new(
8542 catalog,
8543 false,
8544 QueryContext::arc(),
8545 DummyDecoder::arc(),
8546 false,
8547 )
8548 }
8549
8550 async fn build_test_table_provider_with_tsid(
8551 table_name_tuples: &[(String, String)],
8552 num_tag: usize,
8553 num_field: usize,
8554 ) -> DfTableSourceProvider {
8555 let table_specs = table_name_tuples
8556 .iter()
8557 .map(|(schema_name, table_name)| ((schema_name.clone(), table_name.clone()), num_field))
8558 .collect::<Vec<_>>();
8559 build_test_table_provider_with_tsid_fields(&table_specs, num_tag).await
8560 }
8561
8562 async fn build_test_table_provider_with_tsid_fields(
8563 table_specs: &[((String, String), usize)],
8564 num_tag: usize,
8565 ) -> DfTableSourceProvider {
8566 let table_specs = table_specs
8567 .iter()
8568 .map(|(table_name_tuple, num_field)| (table_name_tuple.clone(), num_tag, *num_field))
8569 .collect::<Vec<_>>();
8570 build_test_table_provider_with_tsid_tag_fields(&table_specs).await
8571 }
8572
8573 async fn build_test_table_provider_with_tsid_tag_fields(
8574 table_specs: &[((String, String), usize, usize)],
8575 ) -> DfTableSourceProvider {
8576 let catalog_list = MemoryCatalogManager::with_default_setup();
8577
8578 let physical_table_name = "phy";
8579 let physical_table_id = 999u32;
8580 let physical_num_tag = table_specs
8581 .iter()
8582 .map(|(_, num_tag, _)| *num_tag)
8583 .max()
8584 .unwrap_or(0);
8585 let physical_num_field = table_specs
8586 .iter()
8587 .map(|(_, _, num_field)| *num_field)
8588 .max()
8589 .unwrap_or(0);
8590
8591 {
8593 let mut columns = vec![
8594 ColumnSchema::new(
8595 DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(),
8596 ConcreteDataType::uint32_datatype(),
8597 false,
8598 ),
8599 ColumnSchema::new(
8600 DATA_SCHEMA_TSID_COLUMN_NAME.to_string(),
8601 ConcreteDataType::uint64_datatype(),
8602 false,
8603 ),
8604 ];
8605 for i in 0..physical_num_tag {
8606 columns.push(ColumnSchema::new(
8607 format!("tag_{i}"),
8608 ConcreteDataType::string_datatype(),
8609 false,
8610 ));
8611 }
8612 columns.push(
8613 ColumnSchema::new(
8614 "timestamp".to_string(),
8615 ConcreteDataType::timestamp_millisecond_datatype(),
8616 false,
8617 )
8618 .with_time_index(true),
8619 );
8620 for i in 0..physical_num_field {
8621 columns.push(ColumnSchema::new(
8622 format!("field_{i}"),
8623 ConcreteDataType::float64_datatype(),
8624 true,
8625 ));
8626 }
8627
8628 let schema = Arc::new(Schema::new(columns));
8629 let primary_key_indices = (0..(2 + physical_num_tag)).collect::<Vec<_>>();
8630 let table_meta = TableMetaBuilder::empty()
8631 .schema(schema)
8632 .primary_key_indices(primary_key_indices)
8633 .value_indices(
8634 (2 + physical_num_tag..2 + physical_num_tag + 1 + physical_num_field).collect(),
8635 )
8636 .engine(METRIC_ENGINE_NAME.to_string())
8637 .next_column_id(1024)
8638 .build()
8639 .unwrap();
8640 let table_info = TableInfoBuilder::default()
8641 .table_id(physical_table_id)
8642 .name(physical_table_name)
8643 .meta(table_meta)
8644 .build()
8645 .unwrap();
8646 let table = EmptyTable::from_table_info(&table_info);
8647
8648 assert!(
8649 catalog_list
8650 .register_table_sync(RegisterTableRequest {
8651 catalog: DEFAULT_CATALOG_NAME.to_string(),
8652 schema: DEFAULT_SCHEMA_NAME.to_string(),
8653 table_name: physical_table_name.to_string(),
8654 table_id: physical_table_id,
8655 table,
8656 })
8657 .is_ok()
8658 );
8659 }
8660
8661 for (idx, ((schema_name, table_name), num_tag, num_field)) in table_specs.iter().enumerate()
8663 {
8664 let mut columns = vec![];
8665 for i in 0..*num_tag {
8666 columns.push(ColumnSchema::new(
8667 format!("tag_{i}"),
8668 ConcreteDataType::string_datatype(),
8669 false,
8670 ));
8671 }
8672 columns.push(
8673 ColumnSchema::new(
8674 "timestamp".to_string(),
8675 ConcreteDataType::timestamp_millisecond_datatype(),
8676 false,
8677 )
8678 .with_time_index(true),
8679 );
8680 for i in 0..*num_field {
8681 columns.push(ColumnSchema::new(
8682 format!("field_{i}"),
8683 ConcreteDataType::float64_datatype(),
8684 true,
8685 ));
8686 }
8687
8688 let schema = Arc::new(Schema::new(columns));
8689 let mut options = table::requests::TableOptions::default();
8690 options.extra_options.insert(
8691 LOGICAL_TABLE_METADATA_KEY.to_string(),
8692 physical_table_name.to_string(),
8693 );
8694 let table_id = 1024u32 + idx as u32;
8695 let table_meta = TableMetaBuilder::empty()
8696 .schema(schema)
8697 .primary_key_indices((0..*num_tag).collect())
8698 .value_indices((*num_tag + 1..*num_tag + 1 + *num_field).collect())
8699 .engine(METRIC_ENGINE_NAME.to_string())
8700 .options(options)
8701 .next_column_id(1024)
8702 .build()
8703 .unwrap();
8704 let table_info = TableInfoBuilder::default()
8705 .table_id(table_id)
8706 .name(table_name.clone())
8707 .meta(table_meta)
8708 .build()
8709 .unwrap();
8710 let table = EmptyTable::from_table_info(&table_info);
8711
8712 assert!(
8713 catalog_list
8714 .register_table_sync(RegisterTableRequest {
8715 catalog: DEFAULT_CATALOG_NAME.to_string(),
8716 schema: schema_name.clone(),
8717 table_name: table_name.clone(),
8718 table_id,
8719 table,
8720 })
8721 .is_ok()
8722 );
8723 }
8724
8725 DfTableSourceProvider::new(
8726 catalog_list,
8727 false,
8728 QueryContext::arc(),
8729 DummyDecoder::arc(),
8730 false,
8731 )
8732 }
8733
8734 async fn build_test_table_provider_with_fields(
8735 table_name_tuples: &[(String, String)],
8736 tags: &[&str],
8737 ) -> DfTableSourceProvider {
8738 let catalog_list = MemoryCatalogManager::with_default_setup();
8739 for (schema_name, table_name) in table_name_tuples {
8740 let mut columns = vec![];
8741 let num_tag = tags.len();
8742 for tag in tags {
8743 columns.push(ColumnSchema::new(
8744 tag.to_string(),
8745 ConcreteDataType::string_datatype(),
8746 false,
8747 ));
8748 }
8749 columns.push(
8750 ColumnSchema::new(
8751 greptime_timestamp().to_string(),
8752 ConcreteDataType::timestamp_millisecond_datatype(),
8753 false,
8754 )
8755 .with_time_index(true),
8756 );
8757 columns.push(ColumnSchema::new(
8758 greptime_value().to_string(),
8759 ConcreteDataType::float64_datatype(),
8760 true,
8761 ));
8762 let schema = Arc::new(Schema::new(columns));
8763 let table_meta = TableMetaBuilder::empty()
8764 .schema(schema)
8765 .primary_key_indices((0..num_tag).collect())
8766 .next_column_id(1024)
8767 .build()
8768 .unwrap();
8769 let table_info = TableInfoBuilder::default()
8770 .name(table_name.clone())
8771 .meta(table_meta)
8772 .build()
8773 .unwrap();
8774 let table = EmptyTable::from_table_info(&table_info);
8775
8776 assert!(
8777 catalog_list
8778 .register_table_sync(RegisterTableRequest {
8779 catalog: DEFAULT_CATALOG_NAME.to_string(),
8780 schema: schema_name.clone(),
8781 table_name: table_name.clone(),
8782 table_id: 1024,
8783 table,
8784 })
8785 .is_ok()
8786 );
8787 }
8788
8789 DfTableSourceProvider::new(
8790 catalog_list,
8791 false,
8792 QueryContext::arc(),
8793 DummyDecoder::arc(),
8794 false,
8795 )
8796 }
8797
8798 async fn do_single_instant_function_call(fn_name: &'static str, plan_name: &str) {
8814 let prom_expr =
8815 parser::parse(&format!("{fn_name}(some_metric{{tag_0!=\"bar\"}})")).unwrap();
8816 let eval_stmt = EvalStmt {
8817 expr: prom_expr,
8818 start: UNIX_EPOCH,
8819 end: UNIX_EPOCH
8820 .checked_add(Duration::from_secs(100_000))
8821 .unwrap(),
8822 interval: Duration::from_secs(5),
8823 lookback_delta: Duration::from_secs(1),
8824 };
8825
8826 let table_provider = build_test_table_provider(
8827 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
8828 1,
8829 1,
8830 )
8831 .await;
8832 let plan =
8833 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
8834 .await
8835 .unwrap();
8836
8837 let expected = String::from(
8838 "Filter: TEMPLATE(field_0) IS NOT NULL [timestamp:Timestamp(ms), TEMPLATE(field_0):Float64;N, tag_0:Utf8]\
8839 \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]\
8840 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
8841 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
8842 \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]\
8843 \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]\
8844 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"
8845 ).replace("TEMPLATE", plan_name);
8846
8847 assert_eq!(plan.display_indent_schema().to_string(), expected);
8848 }
8849
8850 #[tokio::test]
8851 async fn single_abs() {
8852 do_single_instant_function_call("abs", "abs").await;
8853 }
8854
8855 #[tokio::test]
8856 #[should_panic]
8857 async fn single_absent() {
8858 do_single_instant_function_call("absent", "").await;
8859 }
8860
8861 #[tokio::test]
8862 async fn single_ceil() {
8863 do_single_instant_function_call("ceil", "ceil").await;
8864 }
8865
8866 #[tokio::test]
8867 async fn single_exp() {
8868 do_single_instant_function_call("exp", "exp").await;
8869 }
8870
8871 #[tokio::test]
8872 async fn single_ln() {
8873 do_single_instant_function_call("ln", "ln").await;
8874 }
8875
8876 #[tokio::test]
8877 async fn single_log2() {
8878 do_single_instant_function_call("log2", "log2").await;
8879 }
8880
8881 #[tokio::test]
8882 async fn single_log10() {
8883 do_single_instant_function_call("log10", "log10").await;
8884 }
8885
8886 #[tokio::test]
8887 #[should_panic]
8888 async fn single_scalar() {
8889 do_single_instant_function_call("scalar", "").await;
8890 }
8891
8892 #[tokio::test]
8893 #[should_panic]
8894 async fn single_sgn() {
8895 do_single_instant_function_call("sgn", "").await;
8896 }
8897
8898 #[tokio::test]
8899 #[should_panic]
8900 async fn single_sort() {
8901 do_single_instant_function_call("sort", "").await;
8902 }
8903
8904 #[tokio::test]
8905 #[should_panic]
8906 async fn single_sort_desc() {
8907 do_single_instant_function_call("sort_desc", "").await;
8908 }
8909
8910 #[tokio::test]
8911 async fn single_sqrt() {
8912 do_single_instant_function_call("sqrt", "sqrt").await;
8913 }
8914
8915 #[tokio::test]
8916 async fn single_timestamp_plan_preserves_source_value() {
8917 let eval_stmt = build_eval_stmt(r#"timestamp(some_metric{tag_0!="bar"})"#);
8918 let table_provider = build_test_table_provider(
8919 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
8920 1,
8921 1,
8922 )
8923 .await;
8924 let plan =
8925 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
8926 .await
8927 .unwrap();
8928
8929 let expected = String::from(
8930 "Filter: value IS NOT NULL [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\
8931 \n Projection: some_metric.timestamp, value AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\
8932 \n Projection: some_metric.timestamp, __promql_timestamp_value_ AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\
8933 \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]\
8934 \n Projection: some_metric.tag_0, some_metric.timestamp, some_metric.field_0, CAST(CAST(CAST(CAST(some_metric.timestamp AS Int64) AS Decimal128(19, 0)) * Decimal128(Some(1),1,0) + Decimal128(Some(0),19,0) AS Int64) AS Float64) / Float64(1000) AS __promql_timestamp_value_ [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\
8935 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
8936 \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]\
8937 \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]\
8938 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
8939 );
8940
8941 assert_eq!(plan.display_indent_schema().to_string(), expected);
8942 }
8943
8944 #[tokio::test]
8945 async fn single_acos() {
8946 do_single_instant_function_call("acos", "acos").await;
8947 }
8948
8949 #[tokio::test]
8950 #[should_panic]
8951 async fn single_acosh() {
8952 do_single_instant_function_call("acosh", "").await;
8953 }
8954
8955 #[tokio::test]
8956 async fn single_asin() {
8957 do_single_instant_function_call("asin", "asin").await;
8958 }
8959
8960 #[tokio::test]
8961 #[should_panic]
8962 async fn single_asinh() {
8963 do_single_instant_function_call("asinh", "").await;
8964 }
8965
8966 #[tokio::test]
8967 async fn single_atan() {
8968 do_single_instant_function_call("atan", "atan").await;
8969 }
8970
8971 #[tokio::test]
8972 #[should_panic]
8973 async fn single_atanh() {
8974 do_single_instant_function_call("atanh", "").await;
8975 }
8976
8977 #[tokio::test]
8978 async fn single_cos() {
8979 do_single_instant_function_call("cos", "cos").await;
8980 }
8981
8982 #[tokio::test]
8983 #[should_panic]
8984 async fn single_cosh() {
8985 do_single_instant_function_call("cosh", "").await;
8986 }
8987
8988 #[tokio::test]
8989 async fn single_sin() {
8990 do_single_instant_function_call("sin", "sin").await;
8991 }
8992
8993 #[tokio::test]
8994 #[should_panic]
8995 async fn single_sinh() {
8996 do_single_instant_function_call("sinh", "").await;
8997 }
8998
8999 #[tokio::test]
9000 async fn single_tan() {
9001 do_single_instant_function_call("tan", "tan").await;
9002 }
9003
9004 #[tokio::test]
9005 #[should_panic]
9006 async fn single_tanh() {
9007 do_single_instant_function_call("tanh", "").await;
9008 }
9009
9010 #[tokio::test]
9011 #[should_panic]
9012 async fn single_deg() {
9013 do_single_instant_function_call("deg", "").await;
9014 }
9015
9016 #[tokio::test]
9017 #[should_panic]
9018 async fn single_rad() {
9019 do_single_instant_function_call("rad", "").await;
9020 }
9021
9022 async fn do_aggregate_expr_plan(fn_name: &str, plan_name: &str) {
9044 let prom_expr = parser::parse(&format!(
9045 "{fn_name} by (tag_1)(some_metric{{tag_0!=\"bar\"}})",
9046 ))
9047 .unwrap();
9048 let mut eval_stmt = EvalStmt {
9049 expr: prom_expr,
9050 start: UNIX_EPOCH,
9051 end: UNIX_EPOCH
9052 .checked_add(Duration::from_secs(100_000))
9053 .unwrap(),
9054 interval: Duration::from_secs(5),
9055 lookback_delta: Duration::from_secs(1),
9056 };
9057
9058 let table_provider = build_test_table_provider(
9060 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9061 2,
9062 2,
9063 )
9064 .await;
9065 let plan =
9066 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9067 .await
9068 .unwrap();
9069 let expected_no_without = String::from(
9070 "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]\
9071 \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]\
9072 \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]\
9073 \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
9074 \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]\
9075 \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]\
9076 \n TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]"
9077 ).replace("TEMPLATE", plan_name);
9078 assert_eq!(
9079 plan.display_indent_schema().to_string(),
9080 expected_no_without
9081 );
9082
9083 if let PromExpr::Aggregate(AggregateExpr { modifier, .. }) = &mut eval_stmt.expr {
9085 *modifier = Some(LabelModifier::Exclude(Labels {
9086 labels: vec![String::from("tag_1")].into_iter().collect(),
9087 }));
9088 }
9089 let table_provider = build_test_table_provider(
9090 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9091 2,
9092 2,
9093 )
9094 .await;
9095 let plan =
9096 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9097 .await
9098 .unwrap();
9099 let expected_without = String::from(
9100 "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]\
9101 \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]\
9102 \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]\
9103 \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
9104 \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]\
9105 \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]\
9106 \n TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]"
9107 ).replace("TEMPLATE", plan_name);
9108 assert_eq!(plan.display_indent_schema().to_string(), expected_without);
9109 }
9110
9111 #[tokio::test]
9112 async fn aggregate_sum() {
9113 do_aggregate_expr_plan("sum", "sum").await;
9114 }
9115
9116 #[tokio::test]
9117 async fn tsid_is_used_for_series_divide_when_available() {
9118 let prom_expr = parser::parse("some_metric").unwrap();
9119 let eval_stmt = EvalStmt {
9120 expr: prom_expr,
9121 start: UNIX_EPOCH,
9122 end: UNIX_EPOCH
9123 .checked_add(Duration::from_secs(100_000))
9124 .unwrap(),
9125 interval: Duration::from_secs(5),
9126 lookback_delta: Duration::from_secs(1),
9127 };
9128
9129 let table_provider = build_test_table_provider_with_tsid(
9130 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9131 1,
9132 1,
9133 )
9134 .await;
9135 let plan =
9136 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9137 .await
9138 .unwrap();
9139
9140 let plan_str = plan.display_indent_schema().to_string();
9141 assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
9142 assert!(plan_str.contains("__tsid ASC NULLS FIRST"));
9143 assert!(
9144 !plan
9145 .schema()
9146 .fields()
9147 .iter()
9148 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
9149 );
9150
9151 let manipulate = find_instant_manipulate(&plan).unwrap();
9152 let exec = manipulate.to_execution_plan(Arc::new(DataSourceExec::new(Arc::new(
9153 MemorySourceConfig::try_new(
9154 &[],
9155 Arc::new(
9156 datafusion_expr::UserDefinedLogicalNodeCore::inputs(manipulate)[0]
9157 .schema()
9158 .as_arrow()
9159 .clone(),
9160 ),
9161 None,
9162 )
9163 .unwrap(),
9164 ))));
9165 assert!(format!("{exec:?}").contains("reuse_tsid_column: true"));
9166 }
9167
9168 #[tokio::test]
9169 async fn default_binary_join_uses_tsid_when_available() {
9170 let eval_stmt = build_eval_stmt("some_metric / some_alt_metric");
9171
9172 let table_provider = build_test_table_provider_with_tsid(
9173 &[
9174 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9175 (
9176 DEFAULT_SCHEMA_NAME.to_string(),
9177 "some_alt_metric".to_string(),
9178 ),
9179 ],
9180 1,
9181 1,
9182 )
9183 .await;
9184 let plan =
9185 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9186 .await
9187 .unwrap();
9188
9189 let plan_str = plan.display_indent_schema().to_string();
9190 assert!(
9191 plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
9192 "{plan_str}"
9193 );
9194 assert!(
9195 !plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
9196 "{plan_str}"
9197 );
9198 }
9199
9200 #[tokio::test]
9201 async fn reject_binary_fill_modifiers() {
9202 let state = build_query_engine_state();
9203
9204 for query in [
9205 "some_metric + fill(0) some_alt_metric",
9206 "some_metric + fill_left(0) some_alt_metric",
9207 "some_metric + fill_right(0) some_alt_metric",
9208 "(some_metric + fill(0) some_alt_metric) + some_metric",
9209 ] {
9210 let eval_stmt = build_eval_stmt(query);
9211 let table_provider = build_test_table_provider(&[], 0, 0).await;
9212 let err = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &state)
9213 .await
9214 .unwrap_err();
9215
9216 assert!(
9217 matches!(
9218 &err,
9219 crate::promql::error::Error::UnsupportedExpr { name, .. }
9220 if name == "PromQL fill modifiers"
9221 ),
9222 "{err}"
9223 );
9224 }
9225 }
9226
9227 #[tokio::test]
9228 async fn timestamp_binary_join_falls_back_when_tsid_is_projected_out() {
9229 for query in [
9230 "timestamp(some_metric) / some_metric",
9231 "some_metric / timestamp(some_metric)",
9232 ] {
9233 let eval_stmt = build_eval_stmt(query);
9234
9235 let table_provider = build_test_table_provider_with_tsid(
9236 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9237 1,
9238 1,
9239 )
9240 .await;
9241 let plan =
9242 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9243 .await
9244 .unwrap();
9245
9246 let plan_str = plan.display_indent_schema().to_string();
9247 assert!(!plan_str.contains("__tsid ="), "{query}: {plan_str}");
9248 assert!(
9249 plan_str.contains("lhs.tag_0 = rhs.tag_0"),
9250 "{query}: {plan_str}"
9251 );
9252 assert!(
9253 !plan
9254 .schema()
9255 .fields()
9256 .iter()
9257 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME),
9258 "{query}: {plan_str}"
9259 );
9260 }
9261 }
9262
9263 #[tokio::test]
9264 async fn timestamp_binary_join_rejects_default_matching_on_mismatched_labels() {
9265 let eval_stmt = build_eval_stmt("timestamp(left_host_job) / right_by_job");
9266
9267 let table_provider = build_test_table_provider_with_tsid_tag_fields(&[
9268 (
9269 (DEFAULT_SCHEMA_NAME.to_string(), "left_host_job".to_string()),
9270 2,
9271 1,
9272 ),
9273 (
9274 (DEFAULT_SCHEMA_NAME.to_string(), "right_by_job".to_string()),
9275 1,
9276 1,
9277 ),
9278 ])
9279 .await;
9280 let plan =
9281 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9282 .await
9283 .unwrap();
9284 let plan_str = plan.display_indent_schema().to_string();
9285
9286 assert!(
9287 plan_str.contains("Boolean(false)") || plan_str.contains("false"),
9288 "{plan_str}"
9289 );
9290 }
9291
9292 #[tokio::test]
9293 async fn tsid_is_preserved_for_nested_default_binary_joins() {
9294 let eval_stmt = build_eval_stmt("(some_metric - some_alt_metric) / some_third_metric");
9295
9296 let table_provider = build_test_table_provider_with_tsid(
9297 &[
9298 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9299 (
9300 DEFAULT_SCHEMA_NAME.to_string(),
9301 "some_alt_metric".to_string(),
9302 ),
9303 (
9304 DEFAULT_SCHEMA_NAME.to_string(),
9305 "some_third_metric".to_string(),
9306 ),
9307 ],
9308 1,
9309 1,
9310 )
9311 .await;
9312 let plan =
9313 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9314 .await
9315 .unwrap();
9316
9317 let plan_str = plan.display_indent_schema().to_string();
9318 assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
9319 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9320 }
9321
9322 #[tokio::test]
9323 async fn repeated_tsid_binary_operand_reuses_leaf_plan() {
9324 let eval_stmt = build_eval_stmt("((some_metric - some_alt_metric) / some_metric) * 100");
9325
9326 let table_provider = build_test_table_provider_with_tsid(
9327 &[
9328 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9329 (
9330 DEFAULT_SCHEMA_NAME.to_string(),
9331 "some_alt_metric".to_string(),
9332 ),
9333 ],
9334 1,
9335 1,
9336 )
9337 .await;
9338 let plan =
9339 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9340 .await
9341 .unwrap();
9342
9343 let plan_str = plan.display_indent_schema().to_string();
9344 assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}");
9345 assert_eq!(
9346 plan_str
9347 .matches("Filter: phy.__table_id = UInt32(1024)")
9348 .count(),
9349 1,
9350 "{plan_str}"
9351 );
9352 assert_eq!(
9353 plan_str.matches("PromInstantManipulate").count(),
9354 2,
9355 "{plan_str}"
9356 );
9357 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9358 }
9359
9360 #[tokio::test]
9361 async fn repeated_tsid_binary_operand_reuses_shorter_field_side() {
9362 let eval_stmt =
9363 build_eval_stmt("((two_field_metric - one_field_metric) / one_field_metric) * 100");
9364
9365 let table_provider = build_test_table_provider_with_tsid_fields(
9366 &[
9367 (
9368 (
9369 DEFAULT_SCHEMA_NAME.to_string(),
9370 "two_field_metric".to_string(),
9371 ),
9372 2,
9373 ),
9374 (
9375 (
9376 DEFAULT_SCHEMA_NAME.to_string(),
9377 "one_field_metric".to_string(),
9378 ),
9379 1,
9380 ),
9381 ],
9382 1,
9383 )
9384 .await;
9385 let plan =
9386 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9387 .await
9388 .unwrap();
9389
9390 let field_names = plan
9391 .schema()
9392 .fields()
9393 .iter()
9394 .map(|field| field.name().clone())
9395 .collect::<Vec<_>>();
9396 let value_columns = field_names
9397 .iter()
9398 .filter(|name| {
9399 *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME
9400 })
9401 .count();
9402 assert_eq!(value_columns, 1, "{field_names:?}");
9403 let plan_str = plan.display_indent_schema().to_string();
9404 assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}");
9405 assert_eq!(
9406 plan_str
9407 .matches("Filter: phy.__table_id = UInt32(1025)")
9408 .count(),
9409 1,
9410 "{plan_str}"
9411 );
9412 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9413 }
9414
9415 #[tokio::test]
9416 async fn binary_island_reuses_self_operand_without_join() {
9417 let eval_stmt = build_eval_stmt("some_metric / some_metric");
9418
9419 let table_provider = build_test_table_provider_with_tsid(
9420 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
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(), 0, "{plan_str}");
9432 assert_eq!(
9433 plan_str
9434 .matches("Filter: phy.__table_id = UInt32(1024)")
9435 .count(),
9436 1,
9437 "{plan_str}"
9438 );
9439 assert_eq!(
9440 plan_str.matches("PromInstantManipulate").count(),
9441 1,
9442 "{plan_str}"
9443 );
9444 }
9445
9446 #[tokio::test]
9447 async fn binary_island_reuses_leaf_across_two_branches() {
9448 let eval_stmt =
9449 build_eval_stmt("(some_metric + some_alt_metric) / (some_metric + third_metric)");
9450
9451 let table_provider = build_test_table_provider_with_tsid(
9452 &[
9453 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9454 (
9455 DEFAULT_SCHEMA_NAME.to_string(),
9456 "some_alt_metric".to_string(),
9457 ),
9458 (DEFAULT_SCHEMA_NAME.to_string(), "third_metric".to_string()),
9459 ],
9460 1,
9461 1,
9462 )
9463 .await;
9464 let plan =
9465 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9466 .await
9467 .unwrap();
9468
9469 let plan_str = plan.display_indent_schema().to_string();
9470 assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
9471 assert_eq!(
9472 plan_str
9473 .matches("Filter: phy.__table_id = UInt32(1024)")
9474 .count(),
9475 1,
9476 "{plan_str}"
9477 );
9478 assert_eq!(
9479 plan_str.matches("PromInstantManipulate").count(),
9480 3,
9481 "{plan_str}"
9482 );
9483 }
9484
9485 #[tokio::test]
9486 async fn binary_island_generated_alias_avoids_user_column_names() {
9487 let eval_stmt = build_eval_stmt("(some_metric + some_alt_metric) / some_metric");
9488
9489 let table_provider = build_test_table_provider_with_fields(
9490 &[
9491 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9492 (
9493 DEFAULT_SCHEMA_NAME.to_string(),
9494 "some_alt_metric".to_string(),
9495 ),
9496 ],
9497 &["prom_v0", "__prom_v0"],
9498 )
9499 .await;
9500 let plan =
9501 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9502 .await
9503 .unwrap();
9504
9505 let field_names = plan.schema().field_names();
9506 assert!(field_names.iter().any(|name| name.ends_with(".prom_v0")));
9507 assert!(field_names.iter().any(|name| name.ends_with(".__prom_v0")));
9508
9509 let plan_str = plan.display_indent_schema().to_string();
9510 assert!(plan_str.contains("SubqueryAlias: __prom_v0"), "{plan_str}");
9511 assert_eq!(
9512 plan_str.matches("PromInstantManipulate").count(),
9513 2,
9514 "{plan_str}"
9515 );
9516 }
9517
9518 #[tokio::test]
9519 async fn binary_island_clears_qualifier_for_nested_unary_projection() {
9520 let eval_stmt = build_eval_stmt("-((some_metric + some_alt_metric) / some_metric)");
9521
9522 let table_provider = build_test_table_provider_with_tsid(
9523 &[
9524 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9525 (
9526 DEFAULT_SCHEMA_NAME.to_string(),
9527 "some_alt_metric".to_string(),
9528 ),
9529 ],
9530 1,
9531 1,
9532 )
9533 .await;
9534 let plan =
9535 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9536 .await
9537 .unwrap();
9538
9539 let plan_str = plan.display_indent_schema().to_string();
9540 assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}");
9541 assert_eq!(
9542 plan_str.matches("PromInstantManipulate").count(),
9543 2,
9544 "{plan_str}"
9545 );
9546 }
9547
9548 #[tokio::test]
9549 async fn binary_island_keeps_distinct_matcher_leaves() {
9550 let eval_stmt = build_eval_stmt(
9551 "(some_metric{tag_0=\"foo\"} + some_alt_metric) / some_metric{tag_0=\"bar\"}",
9552 );
9553
9554 let table_provider = build_test_table_provider_with_tsid(
9555 &[
9556 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9557 (
9558 DEFAULT_SCHEMA_NAME.to_string(),
9559 "some_alt_metric".to_string(),
9560 ),
9561 ],
9562 1,
9563 1,
9564 )
9565 .await;
9566 let plan =
9567 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9568 .await
9569 .unwrap();
9570
9571 let plan_str = plan.display_indent_schema().to_string();
9572 assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
9573 assert_eq!(
9574 plan_str.matches("PromInstantManipulate").count(),
9575 3,
9576 "{plan_str}"
9577 );
9578 }
9579
9580 #[tokio::test]
9581 async fn binary_island_keeps_offset_leaves_distinct() {
9582 let eval_stmt = build_eval_stmt("(some_metric offset 5m + some_alt_metric) / some_metric");
9583
9584 let table_provider = build_test_table_provider_with_tsid(
9585 &[
9586 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9587 (
9588 DEFAULT_SCHEMA_NAME.to_string(),
9589 "some_alt_metric".to_string(),
9590 ),
9591 ],
9592 1,
9593 1,
9594 )
9595 .await;
9596 let plan =
9597 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9598 .await
9599 .unwrap();
9600
9601 let plan_str = plan.display_indent_schema().to_string();
9602 assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
9603 assert_eq!(
9604 plan_str.matches("PromInstantManipulate").count(),
9605 3,
9606 "{plan_str}"
9607 );
9608 }
9609
9610 #[tokio::test]
9611 async fn binary_island_falls_back_for_group_modifier() {
9612 let eval_stmt = build_eval_stmt(
9613 "(some_metric + ignoring(tag_0) group_left some_alt_metric) / some_metric",
9614 );
9615
9616 let table_provider = build_test_table_provider_with_tsid(
9617 &[
9618 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9619 (
9620 DEFAULT_SCHEMA_NAME.to_string(),
9621 "some_alt_metric".to_string(),
9622 ),
9623 ],
9624 1,
9625 1,
9626 )
9627 .await;
9628 let plan =
9629 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9630 .await
9631 .unwrap();
9632
9633 let plan_str = plan.display_indent_schema().to_string();
9634 assert_eq!(
9635 plan_str.matches("PromInstantManipulate").count(),
9636 3,
9637 "{plan_str}"
9638 );
9639 }
9640
9641 #[tokio::test]
9642 async fn binary_island_falls_back_for_comparison_filter() {
9643 let eval_stmt = build_eval_stmt("(some_metric > some_alt_metric) / some_metric");
9644
9645 let table_provider = build_test_table_provider_with_tsid(
9646 &[
9647 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9648 (
9649 DEFAULT_SCHEMA_NAME.to_string(),
9650 "some_alt_metric".to_string(),
9651 ),
9652 ],
9653 1,
9654 1,
9655 )
9656 .await;
9657 let plan =
9658 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9659 .await
9660 .unwrap();
9661
9662 let plan_str = plan.display_indent_schema().to_string();
9663 assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
9664 assert_eq!(
9665 plan_str.matches("PromInstantManipulate").count(),
9666 3,
9667 "{plan_str}"
9668 );
9669 }
9670
9671 #[tokio::test]
9672 async fn tsid_binary_join_uses_shorter_field_side() {
9673 let eval_stmt = build_eval_stmt("one_field_metric / two_field_metric");
9674
9675 let table_provider = build_test_table_provider_with_tsid_fields(
9676 &[
9677 (
9678 (
9679 DEFAULT_SCHEMA_NAME.to_string(),
9680 "one_field_metric".to_string(),
9681 ),
9682 1,
9683 ),
9684 (
9685 (
9686 DEFAULT_SCHEMA_NAME.to_string(),
9687 "two_field_metric".to_string(),
9688 ),
9689 2,
9690 ),
9691 ],
9692 1,
9693 )
9694 .await;
9695 let plan =
9696 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9697 .await
9698 .unwrap();
9699
9700 let field_names = plan
9701 .schema()
9702 .fields()
9703 .iter()
9704 .map(|field| field.name().clone())
9705 .collect::<Vec<_>>();
9706 let value_columns = field_names
9707 .iter()
9708 .filter(|name| {
9709 *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME
9710 })
9711 .count();
9712 assert_eq!(value_columns, 1, "{field_names:?}");
9713 }
9714
9715 #[tokio::test]
9716 async fn comparison_binary_join_uses_shorter_field_side() {
9717 let eval_stmt = build_eval_stmt("two_field_metric > one_field_metric");
9718
9719 let table_provider = build_test_table_provider_with_tsid_fields(
9720 &[
9721 (
9722 (
9723 DEFAULT_SCHEMA_NAME.to_string(),
9724 "two_field_metric".to_string(),
9725 ),
9726 2,
9727 ),
9728 (
9729 (
9730 DEFAULT_SCHEMA_NAME.to_string(),
9731 "one_field_metric".to_string(),
9732 ),
9733 1,
9734 ),
9735 ],
9736 1,
9737 )
9738 .await;
9739 let plan =
9740 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9741 .await
9742 .unwrap();
9743
9744 let field_names = plan
9745 .schema()
9746 .fields()
9747 .iter()
9748 .map(|field| field.name().clone())
9749 .collect::<Vec<_>>();
9750 assert!(
9751 field_names.iter().any(|name| name == "field_0"),
9752 "{field_names:?}"
9753 );
9754 assert!(
9755 !field_names.iter().any(|name| name == "field_1"),
9756 "{field_names:?}"
9757 );
9758 }
9759
9760 #[tokio::test]
9761 async fn label_matching_modifier_disables_tsid_binary_join() {
9762 let eval_stmt = build_eval_stmt("some_metric / ignoring(tag_0) some_alt_metric");
9763
9764 let table_provider = build_test_table_provider_with_tsid(
9765 &[
9766 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9767 (
9768 DEFAULT_SCHEMA_NAME.to_string(),
9769 "some_alt_metric".to_string(),
9770 ),
9771 ],
9772 2,
9773 1,
9774 )
9775 .await;
9776 let plan =
9777 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9778 .await
9779 .unwrap();
9780
9781 let plan_str = plan.display_indent_schema().to_string();
9782 assert!(!plan_str.contains("__tsid ="), "{plan_str}");
9783 assert!(
9784 plan_str.contains("some_metric.tag_1 = some_alt_metric.tag_1"),
9785 "{plan_str}"
9786 );
9787 }
9788
9789 #[tokio::test]
9790 async fn ignoring_absent_label_keeps_tsid_binary_join() {
9791 let eval_stmt = build_eval_stmt("some_metric / ignoring(missing) some_alt_metric");
9792
9793 let table_provider = build_test_table_provider_with_tsid(
9794 &[
9795 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9796 (
9797 DEFAULT_SCHEMA_NAME.to_string(),
9798 "some_alt_metric".to_string(),
9799 ),
9800 ],
9801 2,
9802 1,
9803 )
9804 .await;
9805 let plan =
9806 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9807 .await
9808 .unwrap();
9809
9810 let plan_str = plan.display_indent_schema().to_string();
9811 assert!(
9812 plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
9813 "{plan_str}"
9814 );
9815 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9816 assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
9817 }
9818
9819 #[tokio::test]
9820 async fn range_function_keeps_tsid_for_absent_ignoring_binary_join() {
9821 let eval_stmt =
9822 build_eval_stmt("rate(some_metric[5m]) / ignoring(missing) some_alt_metric");
9823
9824 let table_provider = build_test_table_provider_with_tsid(
9825 &[
9826 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9827 (
9828 DEFAULT_SCHEMA_NAME.to_string(),
9829 "some_alt_metric".to_string(),
9830 ),
9831 ],
9832 2,
9833 1,
9834 )
9835 .await;
9836 let plan =
9837 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9838 .await
9839 .unwrap();
9840
9841 let plan_str = plan.display_indent_schema().to_string();
9842 assert!(
9843 plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
9844 "{plan_str}"
9845 );
9846 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9847 assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
9848 }
9849
9850 #[tokio::test]
9851 async fn on_full_label_set_keeps_tsid_binary_join() {
9852 let eval_stmt = build_eval_stmt("some_metric / on(tag_0, tag_1) some_alt_metric");
9853
9854 let table_provider = build_test_table_provider_with_tsid(
9855 &[
9856 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9857 (
9858 DEFAULT_SCHEMA_NAME.to_string(),
9859 "some_alt_metric".to_string(),
9860 ),
9861 ],
9862 2,
9863 1,
9864 )
9865 .await;
9866 let plan =
9867 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9868 .await
9869 .unwrap();
9870
9871 let plan_str = plan.display_indent_schema().to_string();
9872 assert!(
9873 plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
9874 "{plan_str}"
9875 );
9876 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9877 assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
9878 }
9879
9880 #[tokio::test]
9881 async fn on_partial_label_set_disables_tsid_binary_join() {
9882 let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric");
9883
9884 let table_provider = build_test_table_provider_with_tsid(
9885 &[
9886 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9887 (
9888 DEFAULT_SCHEMA_NAME.to_string(),
9889 "some_alt_metric".to_string(),
9890 ),
9891 ],
9892 2,
9893 1,
9894 )
9895 .await;
9896 let plan =
9897 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9898 .await
9899 .unwrap();
9900
9901 let plan_str = plan.display_indent_schema().to_string();
9902 assert!(!plan_str.contains("__tsid ="), "{plan_str}");
9903 assert!(
9904 plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
9905 "{plan_str}"
9906 );
9907 assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
9908 }
9909
9910 #[tokio::test]
9911 async fn on_label_set_must_cover_both_sides_to_use_tsid_binary_join() {
9912 let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric");
9913
9914 let table_provider = build_test_table_provider_with_tsid_tag_fields(&[
9915 (
9916 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9917 2,
9918 1,
9919 ),
9920 (
9921 (
9922 DEFAULT_SCHEMA_NAME.to_string(),
9923 "some_alt_metric".to_string(),
9924 ),
9925 1,
9926 1,
9927 ),
9928 ])
9929 .await;
9930 let plan =
9931 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9932 .await
9933 .unwrap();
9934
9935 let plan_str = plan.display_indent_schema().to_string();
9936 assert!(!plan_str.contains("__tsid ="), "{plan_str}");
9937 assert!(
9938 plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
9939 "{plan_str}"
9940 );
9941 assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
9942 }
9943
9944 #[tokio::test]
9945 async fn comparison_binary_join_uses_tsid_and_keeps_it_in_filtered_result() {
9946 let eval_stmt = build_eval_stmt("some_metric > some_alt_metric");
9947
9948 let table_provider = build_test_table_provider_with_tsid(
9949 &[
9950 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9951 (
9952 DEFAULT_SCHEMA_NAME.to_string(),
9953 "some_alt_metric".to_string(),
9954 ),
9955 ],
9956 2,
9957 1,
9958 )
9959 .await;
9960 let mut planner = PromPlanner {
9961 table_provider,
9962 ctx: PromPlannerContext::from_eval_stmt(&eval_stmt),
9963 promql_annotations: None,
9964 };
9965 let plan = planner
9966 .prom_expr_to_plan(&eval_stmt.expr, &build_query_engine_state())
9967 .await
9968 .unwrap();
9969
9970 let plan_str = plan.display_indent_schema().to_string();
9971 assert!(
9972 plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
9973 "{plan_str}"
9974 );
9975 assert!(
9976 plan.schema()
9977 .fields()
9978 .iter()
9979 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME),
9980 "{plan_str}"
9981 );
9982 assert!(planner.ctx.use_tsid, "{plan_str}");
9983 }
9984
9985 #[tokio::test]
9986 async fn comparison_bool_binary_join_uses_tsid_when_available() {
9987 let eval_stmt = build_eval_stmt("some_metric > bool some_alt_metric");
9988
9989 let table_provider = build_test_table_provider_with_tsid(
9990 &[
9991 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9992 (
9993 DEFAULT_SCHEMA_NAME.to_string(),
9994 "some_alt_metric".to_string(),
9995 ),
9996 ],
9997 2,
9998 1,
9999 )
10000 .await;
10001 let plan =
10002 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10003 .await
10004 .unwrap();
10005
10006 let plan_str = plan.display_indent_schema().to_string();
10007 assert!(
10008 plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
10009 "{plan_str}"
10010 );
10011 assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
10012 assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
10013 }
10014
10015 #[tokio::test]
10016 async fn scalar_count_count_range_keeps_full_window() {
10017 let plan_str = build_optimized_tsid_plan(
10018 "scalar(count(count(some_metric) by (tag_0)))",
10019 1,
10020 1,
10021 100_000,
10022 1,
10023 )
10024 .await;
10025 assert!(plan_str.contains("ScalarCalculate: tags=[]"));
10026 assert!(plan_str.contains("PromInstantManipulate: range=[0..100000000]"));
10027 assert!(!plan_str.contains("PromInstantManipulate: range=[99999000..99999000]"));
10028 }
10029
10030 #[tokio::test]
10031 async fn scalar_count_count_rewrite_applies_inside_binary_expr_for_tsid_input() {
10032 let plan_str = build_optimized_tsid_plan(
10033 "sum(irate(some_metric[1h])) / scalar(count(count(some_metric) by (tag_0)))",
10034 2,
10035 1,
10036 10,
10037 300,
10038 )
10039 .await;
10040 assert!(plan_str.contains("Distinct:"), "{plan_str}");
10041 }
10042
10043 #[tokio::test]
10044 async fn nested_count_rewrite_keeps_full_series_key_with_tsid_input() {
10045 assert_nested_count_rewrite_applies(
10046 "count(count(some_metric) by (tag_0))",
10047 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(count(some_metric.field_0))]]"
10048 )
10049 .await;
10050 }
10051
10052 #[tokio::test]
10053 async fn nested_sum_count_rewrite_keeps_full_series_key_with_tsid_input() {
10054 assert_nested_count_rewrite_applies(
10055 "count(sum(some_metric) by (tag_0))",
10056 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(sum(some_metric.field_0))]]"
10057 )
10058 .await;
10059 }
10060
10061 #[tokio::test]
10062 async fn nested_supported_inner_aggs_rewrite_apply_for_tsid_input() {
10063 for (query, expected_outer_agg) in [
10064 (
10065 "count(avg(some_metric) by (tag_0))",
10066 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(avg(some_metric.field_0))]]",
10067 ),
10068 (
10069 "count(min(some_metric) by (tag_0))",
10070 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(min(some_metric.field_0))]]",
10071 ),
10072 (
10073 "count(max(some_metric) by (tag_0))",
10074 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(max(some_metric.field_0))]]",
10075 ),
10076 (
10077 "count(stddev(some_metric) by (tag_0))",
10078 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(stddev_pop(some_metric.field_0))]]",
10079 ),
10080 (
10081 "count(stdvar(some_metric) by (tag_0))",
10082 "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(var_pop(some_metric.field_0))]]",
10083 ),
10084 ] {
10085 assert_nested_count_rewrite_applies(query, expected_outer_agg).await;
10086 }
10087 }
10088
10089 #[tokio::test]
10090 async fn nested_non_count_inner_aggs_rewrite_filter_null_values_for_tsid_input() {
10091 let count_plan =
10092 build_optimized_tsid_plan("count(count(some_metric) by (tag_0))", 2, 1, 100_000, 1)
10093 .await;
10094 assert!(
10095 !count_plan.contains("some_metric.field_0 IS NOT NULL"),
10096 "{count_plan}"
10097 );
10098
10099 for query in [
10100 "count(sum(some_metric) by (tag_0))",
10101 "count(avg(some_metric) by (tag_0))",
10102 "count(min(some_metric) by (tag_0))",
10103 "count(max(some_metric) by (tag_0))",
10104 "count(stddev(some_metric) by (tag_0))",
10105 "count(stdvar(some_metric) by (tag_0))",
10106 ] {
10107 let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await;
10108 assert!(
10109 plan_str.contains("Filter: some_metric.field_0 IS NOT NULL"),
10110 "{query}: {plan_str}"
10111 );
10112 }
10113 }
10114
10115 #[tokio::test]
10116 async fn nested_unsupported_or_non_direct_inner_aggs_do_not_rewrite() {
10117 assert_nested_count_rewrite_missing("count(group(some_metric) by (tag_0))", 2, 1).await;
10118 assert_nested_count_rewrite_missing(
10119 "count(sum(irate(some_metric[1h])) by (tag_0))",
10120 2,
10121 300,
10122 )
10123 .await;
10124 }
10125
10126 #[tokio::test]
10127 async fn physical_table_name_is_not_leaked_in_plan() {
10128 let prom_expr = parser::parse("some_metric").unwrap();
10129 let eval_stmt = EvalStmt {
10130 expr: prom_expr,
10131 start: UNIX_EPOCH,
10132 end: UNIX_EPOCH
10133 .checked_add(Duration::from_secs(100_000))
10134 .unwrap(),
10135 interval: Duration::from_secs(5),
10136 lookback_delta: Duration::from_secs(1),
10137 };
10138
10139 let table_provider = build_test_table_provider_with_tsid(
10140 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10141 1,
10142 1,
10143 )
10144 .await;
10145 let plan =
10146 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10147 .await
10148 .unwrap();
10149
10150 let plan_str = plan.display_indent_schema().to_string();
10151 assert!(plan_str.contains("TableScan: phy"), "{plan}");
10152 assert!(plan_str.contains("SubqueryAlias: some_metric"));
10153 assert!(plan_str.contains("Filter: phy.__table_id = UInt32(1024)"));
10154 assert!(!plan_str.contains("TableScan: some_metric"));
10155 }
10156
10157 #[tokio::test]
10158 async fn sum_without_does_not_group_by_tsid() {
10159 let prom_expr = parser::parse("sum without (tag_0) (some_metric)").unwrap();
10160 let eval_stmt = EvalStmt {
10161 expr: prom_expr,
10162 start: UNIX_EPOCH,
10163 end: UNIX_EPOCH
10164 .checked_add(Duration::from_secs(100_000))
10165 .unwrap(),
10166 interval: Duration::from_secs(5),
10167 lookback_delta: Duration::from_secs(1),
10168 };
10169
10170 let table_provider = build_test_table_provider_with_tsid(
10171 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10172 1,
10173 1,
10174 )
10175 .await;
10176 let plan =
10177 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10178 .await
10179 .unwrap();
10180
10181 let plan_str = plan.display_indent_schema().to_string();
10182 assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10183
10184 let aggr_line = plan_str
10185 .lines()
10186 .find(|line| line.contains("Aggregate: groupBy="))
10187 .unwrap();
10188 assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
10189 }
10190
10191 #[tokio::test]
10192 async fn topk_without_does_not_partition_by_tsid() {
10193 let prom_expr = parser::parse("topk without (tag_0) (1, some_metric)").unwrap();
10194 let eval_stmt = EvalStmt {
10195 expr: prom_expr,
10196 start: UNIX_EPOCH,
10197 end: UNIX_EPOCH
10198 .checked_add(Duration::from_secs(100_000))
10199 .unwrap(),
10200 interval: Duration::from_secs(5),
10201 lookback_delta: Duration::from_secs(1),
10202 };
10203
10204 let table_provider = build_test_table_provider_with_tsid(
10205 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10206 1,
10207 1,
10208 )
10209 .await;
10210 let plan =
10211 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10212 .await
10213 .unwrap();
10214
10215 let plan_str = plan.display_indent_schema().to_string();
10216 assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10217
10218 let window_line = plan_str
10219 .lines()
10220 .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()"))
10221 .unwrap();
10222 let partition_by = window_line
10223 .split("PARTITION BY [")
10224 .nth(1)
10225 .and_then(|s| s.split("] ORDER BY").next())
10226 .unwrap();
10227 assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
10228 }
10229
10230 #[tokio::test]
10231 async fn sum_by_does_not_group_by_tsid() {
10232 let prom_expr = parser::parse("sum by (__tsid) (some_metric)").unwrap();
10233 let eval_stmt = EvalStmt {
10234 expr: prom_expr,
10235 start: UNIX_EPOCH,
10236 end: UNIX_EPOCH
10237 .checked_add(Duration::from_secs(100_000))
10238 .unwrap(),
10239 interval: Duration::from_secs(5),
10240 lookback_delta: Duration::from_secs(1),
10241 };
10242
10243 let table_provider = build_test_table_provider_with_tsid(
10244 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10245 1,
10246 1,
10247 )
10248 .await;
10249 let plan =
10250 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10251 .await
10252 .unwrap();
10253
10254 let plan_str = plan.display_indent_schema().to_string();
10255 assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10256
10257 let aggr_line = plan_str
10258 .lines()
10259 .find(|line| line.contains("Aggregate: groupBy="))
10260 .unwrap();
10261 assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
10262 }
10263
10264 #[tokio::test]
10265 async fn aggregate_over_binary_time_function_expr() {
10266 for op in ["sum", "min", "max", "avg"] {
10267 let prom_expr = parser::parse(&format!(
10268 "{op} by (tag_0, tag_1, tag_2) (time() - some_metric)"
10269 ))
10270 .unwrap();
10271 let eval_stmt = EvalStmt {
10272 expr: prom_expr,
10273 start: UNIX_EPOCH,
10274 end: UNIX_EPOCH
10275 .checked_add(Duration::from_secs(100_000))
10276 .unwrap(),
10277 interval: Duration::from_secs(5),
10278 lookback_delta: Duration::from_secs(1),
10279 };
10280
10281 let table_provider = build_test_table_provider_with_tsid(
10282 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10283 3,
10284 1,
10285 )
10286 .await;
10287 let plan =
10288 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10289 .await
10290 .unwrap();
10291
10292 let plan_str = plan.display_indent_schema().to_string();
10293 let aggr_line = plan_str
10294 .lines()
10295 .find(|line| line.contains("Aggregate: groupBy="))
10296 .unwrap();
10297 assert!(aggr_line.contains(op), "{plan_str}");
10298 assert!(aggr_line.contains("first_value"), "{plan_str}");
10299 assert!(
10300 !plan
10301 .schema()
10302 .fields()
10303 .iter()
10304 .any(|field| { field.name() == DATA_SCHEMA_TSID_COLUMN_NAME })
10305 );
10306 }
10307 }
10308
10309 #[tokio::test]
10310 async fn topk_by_does_not_partition_by_tsid() {
10311 let prom_expr = parser::parse("topk by (__tsid) (1, some_metric)").unwrap();
10312 let eval_stmt = EvalStmt {
10313 expr: prom_expr,
10314 start: UNIX_EPOCH,
10315 end: UNIX_EPOCH
10316 .checked_add(Duration::from_secs(100_000))
10317 .unwrap(),
10318 interval: Duration::from_secs(5),
10319 lookback_delta: Duration::from_secs(1),
10320 };
10321
10322 let table_provider = build_test_table_provider_with_tsid(
10323 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10324 1,
10325 1,
10326 )
10327 .await;
10328 let plan =
10329 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10330 .await
10331 .unwrap();
10332
10333 let plan_str = plan.display_indent_schema().to_string();
10334 assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10335
10336 let window_line = plan_str
10337 .lines()
10338 .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()"))
10339 .unwrap();
10340 let partition_by = window_line
10341 .split("PARTITION BY [")
10342 .nth(1)
10343 .and_then(|s| s.split("] ORDER BY").next())
10344 .unwrap();
10345 assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
10346 }
10347
10348 #[tokio::test]
10349 async fn selector_matcher_on_tsid_does_not_use_internal_column() {
10350 let prom_expr = parser::parse(r#"some_metric{__tsid="123"}"#).unwrap();
10351 let eval_stmt = EvalStmt {
10352 expr: prom_expr,
10353 start: UNIX_EPOCH,
10354 end: UNIX_EPOCH
10355 .checked_add(Duration::from_secs(100_000))
10356 .unwrap(),
10357 interval: Duration::from_secs(5),
10358 lookback_delta: Duration::from_secs(1),
10359 };
10360
10361 let table_provider = build_test_table_provider_with_tsid(
10362 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10363 1,
10364 1,
10365 )
10366 .await;
10367 let plan =
10368 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10369 .await
10370 .unwrap();
10371
10372 fn collect_filter_cols(plan: &LogicalPlan, out: &mut HashSet<Column>) {
10373 if let LogicalPlan::Filter(filter) = plan {
10374 datafusion_expr::utils::expr_to_columns(&filter.predicate, out).unwrap();
10375 }
10376 for input in plan.inputs() {
10377 collect_filter_cols(input, out);
10378 }
10379 }
10380
10381 let mut filter_cols = HashSet::new();
10382 collect_filter_cols(&plan, &mut filter_cols);
10383 assert!(
10384 !filter_cols
10385 .iter()
10386 .any(|c| c.name == DATA_SCHEMA_TSID_COLUMN_NAME)
10387 );
10388 }
10389
10390 #[tokio::test]
10391 async fn tsid_is_not_used_when_physical_table_is_missing() {
10392 let prom_expr = parser::parse("some_metric").unwrap();
10393 let eval_stmt = EvalStmt {
10394 expr: prom_expr,
10395 start: UNIX_EPOCH,
10396 end: UNIX_EPOCH
10397 .checked_add(Duration::from_secs(100_000))
10398 .unwrap(),
10399 interval: Duration::from_secs(5),
10400 lookback_delta: Duration::from_secs(1),
10401 };
10402
10403 let catalog_list = MemoryCatalogManager::with_default_setup();
10404
10405 let mut columns = vec![ColumnSchema::new(
10407 "tag_0".to_string(),
10408 ConcreteDataType::string_datatype(),
10409 false,
10410 )];
10411 columns.push(
10412 ColumnSchema::new(
10413 "timestamp".to_string(),
10414 ConcreteDataType::timestamp_millisecond_datatype(),
10415 false,
10416 )
10417 .with_time_index(true),
10418 );
10419 columns.push(ColumnSchema::new(
10420 "field_0".to_string(),
10421 ConcreteDataType::float64_datatype(),
10422 true,
10423 ));
10424 let schema = Arc::new(Schema::new(columns));
10425 let mut options = table::requests::TableOptions::default();
10426 options
10427 .extra_options
10428 .insert(LOGICAL_TABLE_METADATA_KEY.to_string(), "phy".to_string());
10429 let table_meta = TableMetaBuilder::empty()
10430 .schema(schema)
10431 .primary_key_indices(vec![0])
10432 .value_indices(vec![2])
10433 .engine(METRIC_ENGINE_NAME.to_string())
10434 .options(options)
10435 .next_column_id(1024)
10436 .build()
10437 .unwrap();
10438 let table_info = TableInfoBuilder::default()
10439 .table_id(1024)
10440 .name("some_metric")
10441 .meta(table_meta)
10442 .build()
10443 .unwrap();
10444 let table = EmptyTable::from_table_info(&table_info);
10445 catalog_list
10446 .register_table_sync(RegisterTableRequest {
10447 catalog: DEFAULT_CATALOG_NAME.to_string(),
10448 schema: DEFAULT_SCHEMA_NAME.to_string(),
10449 table_name: "some_metric".to_string(),
10450 table_id: 1024,
10451 table,
10452 })
10453 .unwrap();
10454
10455 let table_provider = DfTableSourceProvider::new(
10456 catalog_list,
10457 false,
10458 QueryContext::arc(),
10459 DummyDecoder::arc(),
10460 false,
10461 );
10462
10463 let plan =
10464 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10465 .await
10466 .unwrap();
10467
10468 let plan_str = plan.display_indent_schema().to_string();
10469 assert!(plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]"));
10470 assert!(!plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10471 }
10472
10473 #[tokio::test]
10474 async fn tsid_is_carried_only_when_aggregate_preserves_label_set() {
10475 let prom_expr = parser::parse("sum by (tag_0) (some_metric)").unwrap();
10476 let eval_stmt = EvalStmt {
10477 expr: prom_expr,
10478 start: UNIX_EPOCH,
10479 end: UNIX_EPOCH
10480 .checked_add(Duration::from_secs(100_000))
10481 .unwrap(),
10482 interval: Duration::from_secs(5),
10483 lookback_delta: Duration::from_secs(1),
10484 };
10485
10486 let table_provider = build_test_table_provider_with_tsid(
10487 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10488 1,
10489 1,
10490 )
10491 .await;
10492 let plan =
10493 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10494 .await
10495 .unwrap();
10496
10497 let plan_str = plan.display_indent_schema().to_string();
10498 assert!(plan_str.contains("first_value") && plan_str.contains("__tsid"));
10499 assert!(
10500 !plan
10501 .schema()
10502 .fields()
10503 .iter()
10504 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
10505 );
10506
10507 let prom_expr = parser::parse("sum(some_metric)").unwrap();
10509 let eval_stmt = EvalStmt {
10510 expr: prom_expr,
10511 start: UNIX_EPOCH,
10512 end: UNIX_EPOCH
10513 .checked_add(Duration::from_secs(100_000))
10514 .unwrap(),
10515 interval: Duration::from_secs(5),
10516 lookback_delta: Duration::from_secs(1),
10517 };
10518 let table_provider = build_test_table_provider_with_tsid(
10519 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10520 1,
10521 1,
10522 )
10523 .await;
10524 let plan =
10525 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10526 .await
10527 .unwrap();
10528 let plan_str = plan.display_indent_schema().to_string();
10529 assert!(!plan_str.contains("first_value"));
10530 }
10531
10532 #[tokio::test]
10533 async fn or_operator_with_unknown_metric_does_not_require_tsid() {
10534 let prom_expr = parser::parse("unknown_metric or some_metric").unwrap();
10535 let eval_stmt = EvalStmt {
10536 expr: prom_expr,
10537 start: UNIX_EPOCH,
10538 end: UNIX_EPOCH
10539 .checked_add(Duration::from_secs(100_000))
10540 .unwrap(),
10541 interval: Duration::from_secs(5),
10542 lookback_delta: Duration::from_secs(1),
10543 };
10544
10545 let table_provider = build_test_table_provider_with_tsid(
10546 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10547 1,
10548 1,
10549 )
10550 .await;
10551
10552 let plan =
10553 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10554 .await
10555 .unwrap();
10556
10557 assert!(
10558 !plan
10559 .schema()
10560 .fields()
10561 .iter()
10562 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
10563 );
10564 }
10565
10566 #[tokio::test]
10567 async fn aggregate_avg() {
10568 do_aggregate_expr_plan("avg", "avg").await;
10569 }
10570
10571 #[tokio::test]
10572 #[should_panic] async fn aggregate_count() {
10574 do_aggregate_expr_plan("count", "count").await;
10575 }
10576
10577 #[tokio::test]
10578 async fn aggregate_min() {
10579 do_aggregate_expr_plan("min", "min").await;
10580 }
10581
10582 #[tokio::test]
10583 async fn aggregate_max() {
10584 do_aggregate_expr_plan("max", "max").await;
10585 }
10586
10587 #[tokio::test]
10588 async fn aggregate_group() {
10589 let prom_expr = parser::parse(
10593 "sum(group by (cluster)(kubernetes_build_info{service=\"kubernetes\",job=\"apiserver\"}))",
10594 )
10595 .unwrap();
10596 let eval_stmt = EvalStmt {
10597 expr: prom_expr,
10598 start: UNIX_EPOCH,
10599 end: UNIX_EPOCH
10600 .checked_add(Duration::from_secs(100_000))
10601 .unwrap(),
10602 interval: Duration::from_secs(5),
10603 lookback_delta: Duration::from_secs(1),
10604 };
10605
10606 let table_provider = build_test_table_provider_with_fields(
10607 &[(
10608 DEFAULT_SCHEMA_NAME.to_string(),
10609 "kubernetes_build_info".to_string(),
10610 )],
10611 &["cluster", "service", "job"],
10612 )
10613 .await;
10614 let plan =
10615 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10616 .await
10617 .unwrap();
10618
10619 let plan_str = plan.display_indent_schema().to_string();
10620 assert!(plan_str.contains("max(Float64(1"));
10621 }
10622
10623 #[tokio::test]
10624 async fn aggregate_stddev() {
10625 do_aggregate_expr_plan("stddev", "stddev_pop").await;
10626 }
10627
10628 #[tokio::test]
10629 async fn aggregate_stdvar() {
10630 do_aggregate_expr_plan("stdvar", "var_pop").await;
10631 }
10632
10633 #[tokio::test]
10657 async fn binary_op_column_column() {
10658 let prom_expr =
10659 parser::parse(r#"some_metric{tag_0="foo"} + some_metric{tag_0="bar"}"#).unwrap();
10660 let eval_stmt = EvalStmt {
10661 expr: prom_expr,
10662 start: UNIX_EPOCH,
10663 end: UNIX_EPOCH
10664 .checked_add(Duration::from_secs(100_000))
10665 .unwrap(),
10666 interval: Duration::from_secs(5),
10667 lookback_delta: Duration::from_secs(1),
10668 };
10669
10670 let table_provider = build_test_table_provider(
10671 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10672 1,
10673 1,
10674 )
10675 .await;
10676 let plan =
10677 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10678 .await
10679 .unwrap();
10680
10681 let expected = String::from(
10682 "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]\
10683 \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]\
10684 \n SubqueryAlias: lhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10685 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10686 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10687 \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]\
10688 \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]\
10689 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10690 \n SubqueryAlias: rhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10691 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10692 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10693 \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]\
10694 \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]\
10695 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
10696 );
10697
10698 assert_eq!(plan.display_indent_schema().to_string(), expected);
10699 }
10700
10701 async fn indie_query_plan_compare<T: AsRef<str>>(query: &str, expected: T) {
10702 let prom_expr = parser::parse(query).unwrap();
10703 let eval_stmt = EvalStmt {
10704 expr: prom_expr,
10705 start: UNIX_EPOCH,
10706 end: UNIX_EPOCH
10707 .checked_add(Duration::from_secs(100_000))
10708 .unwrap(),
10709 interval: Duration::from_secs(5),
10710 lookback_delta: Duration::from_secs(1),
10711 };
10712
10713 let table_provider = build_test_table_provider(
10714 &[
10715 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10716 (
10717 "greptime_private".to_string(),
10718 "some_alt_metric".to_string(),
10719 ),
10720 ],
10721 1,
10722 1,
10723 )
10724 .await;
10725 let plan =
10726 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10727 .await
10728 .unwrap();
10729
10730 assert_eq!(plan.display_indent_schema().to_string(), expected.as_ref());
10731 }
10732
10733 #[tokio::test]
10734 async fn binary_op_literal_column() {
10735 let query = r#"1 + some_metric{tag_0="bar"}"#;
10736 let expected = String::from(
10737 "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]\
10738 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10739 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10740 \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]\
10741 \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]\
10742 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
10743 );
10744
10745 indie_query_plan_compare(query, expected).await;
10746 }
10747
10748 #[tokio::test]
10749 async fn binary_op_literal_literal() {
10750 let query = r#"1 + 1"#;
10751 let expected = r#"EmptyMetric: range=[0..100000000], interval=[5000] [time:Timestamp(ms), value:Float64;N]
10752 TableScan: dummy [time:Timestamp(ms), value:Float64;N]"#;
10753 indie_query_plan_compare(query, expected).await;
10754 }
10755
10756 #[tokio::test]
10757 async fn simple_bool_grammar() {
10758 let query = "some_metric != bool 1.2345";
10759 let expected = String::from(
10760 "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]\
10761 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10762 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10763 \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]\
10764 \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]\
10765 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
10766 );
10767
10768 indie_query_plan_compare(query, expected).await;
10769 }
10770
10771 #[tokio::test]
10772 async fn bool_with_additional_arithmetic() {
10773 let query = "some_metric + (1 == bool 2)";
10774 let expected = String::from(
10775 "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]\
10776 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10777 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10778 \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]\
10779 \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]\
10780 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
10781 );
10782
10783 indie_query_plan_compare(query, expected).await;
10784 }
10785
10786 #[tokio::test]
10787 async fn simple_unary() {
10788 let query = "-some_metric";
10789 let expected = String::from(
10790 "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]\
10791 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10792 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10793 \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]\
10794 \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]\
10795 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
10796 );
10797
10798 indie_query_plan_compare(query, expected).await;
10799 }
10800
10801 #[tokio::test]
10802 async fn increase_aggr() {
10803 let query = "increase(some_metric[5m])";
10804 let expected = String::from(
10805 "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]\
10806 \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]\
10807 \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))]\
10808 \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10809 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10810 \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]\
10811 \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]\
10812 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
10813 );
10814
10815 indie_query_plan_compare(query, expected).await;
10816 }
10817
10818 async fn native_histogram_plan(query: &str) -> String {
10819 let table_provider = build_test_native_histogram_table_provider("some_metric").await;
10820 let plan = PromPlanner::stmt_to_plan(
10821 table_provider,
10822 &build_eval_stmt(query),
10823 &build_query_engine_state(),
10824 )
10825 .await
10826 .unwrap();
10827 plan.display_indent_schema().to_string()
10828 }
10829
10830 #[tokio::test]
10831 async fn native_histogram_count_uses_native_udf() {
10832 let plan = native_histogram_plan("histogram_count(some_metric)").await;
10833
10834 assert!(plan.contains("prom_native_histogram_count"), "{plan}");
10835 assert!(!plan.contains("HistogramFold:"), "{plan}");
10836 }
10837
10838 #[tokio::test]
10839 async fn timestamp_filters_native_histogram_stale_marker_before_projection() {
10840 let mut stale = direct_or_histogram();
10841 stale.sum = f64::from_bits(PROMETHEUS_STALE_NAN_BITS);
10842 let table = operator_metric_table(
10843 "stale_histogram",
10844 2_100,
10845 "a",
10846 None,
10847 DirectOrValue::NativeHistogram(stale),
10848 );
10849 let catalog = MemoryCatalogManager::with_default_setup();
10850 catalog
10851 .register_table_sync(RegisterTableRequest {
10852 catalog: DEFAULT_CATALOG_NAME.to_string(),
10853 schema: DEFAULT_SCHEMA_NAME.to_string(),
10854 table_name: "stale_histogram".to_string(),
10855 table_id: 2_100,
10856 table,
10857 })
10858 .unwrap();
10859 let provider = DfTableSourceProvider::new(
10860 catalog,
10861 false,
10862 QueryContext::arc(),
10863 DummyDecoder::arc(),
10864 false,
10865 );
10866 let state = build_query_engine_state();
10867 let plan = PromPlanner::stmt_to_plan(
10868 provider,
10869 &operator_eval_stmt("timestamp(stale_histogram)"),
10870 &state,
10871 )
10872 .await
10873 .unwrap();
10874 let plan_text = plan.display_indent_schema().to_string();
10875 assert!(plan_text.contains(TIMESTAMP_VALUE_PREFIX), "{plan_text}");
10876
10877 let (_, batches) = execute(plan, &state).await;
10878 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
10879 }
10880
10881 #[tokio::test]
10882 async fn timestamp_filters_stale_marker_from_mixed_sample_companion() {
10883 let histograms = build_histogram_array(&[None]);
10884 let schema = Arc::new(ArrowSchema::new(vec![
10885 Field::new(
10886 "timestamp",
10887 ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
10888 false,
10889 ),
10890 Field::new(
10891 greptime_native_histogram(),
10892 histograms.data_type().clone(),
10893 true,
10894 ),
10895 Field::new(greptime_value(), ArrowDataType::Float64, true),
10896 ]));
10897 let batch = RecordBatch::try_new(
10898 schema.clone(),
10899 vec![
10900 Arc::new(TimestampMillisecondArray::from(vec![1_000])),
10901 histograms,
10902 Arc::new(Float64Array::from(vec![f64::from_bits(
10903 PROMETHEUS_STALE_NAN_BITS,
10904 )])),
10905 ],
10906 )
10907 .unwrap();
10908 let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap());
10909 let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None)
10910 .unwrap()
10911 .build()
10912 .unwrap();
10913 let input = LogicalPlan::Extension(Extension {
10914 node: Arc::new(SeriesDivide::new(
10915 Vec::new(),
10916 "timestamp".to_string(),
10917 input,
10918 )),
10919 });
10920 let input = LogicalPlan::Extension(Extension {
10921 node: Arc::new(InstantManipulate::new(
10922 1_000,
10923 1_000,
10924 5_000,
10925 1_000,
10926 0,
10927 "timestamp".to_string(),
10928 Vec::new(),
10929 Some(greptime_native_histogram().to_string()),
10930 input,
10931 )),
10932 });
10933 let plan = LogicalPlanBuilder::from(input)
10935 .project([col("timestamp")])
10936 .unwrap()
10937 .build()
10938 .unwrap();
10939
10940 let (_, batches) = execute(plan, &build_query_engine_state()).await;
10941 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
10942 }
10943
10944 #[tokio::test]
10945 async fn native_histogram_rate_can_feed_count() {
10946 let plan = native_histogram_plan("histogram_count(rate(some_metric[5m]))").await;
10947
10948 assert!(plan.contains("prom_native_histogram_rate"), "{plan}");
10949 assert!(plan.contains("prom_native_histogram_count"), "{plan}");
10950 }
10951
10952 #[tokio::test]
10953 async fn native_histogram_quantile_skips_classic_fold() {
10954 let plan = native_histogram_plan("histogram_quantile(0.9, some_metric)").await;
10955
10956 assert!(plan.contains("prom_native_histogram_quantile"), "{plan}");
10957 assert!(!plan.contains("HistogramFold:"), "{plan}");
10958 assert!(plan.contains("some_metric.le"), "{plan}");
10959 assert!(plan.contains("Float64(0.9)"), "{plan}");
10961 assert!(plan.contains("IS NOT NULL"), "{plan}");
10964 }
10965
10966 #[tokio::test]
10967 async fn mixed_native_histogram_quantile_uses_histogram_field() {
10968 let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
10969 let plan = PromPlanner::stmt_to_plan(
10970 table_provider,
10971 &build_eval_stmt("histogram_quantile(0.9, some_metric)"),
10972 &build_query_engine_state(),
10973 )
10974 .await
10975 .unwrap()
10976 .display_indent_schema()
10977 .to_string();
10978
10979 assert!(
10980 plan.contains("prom_native_histogram_quantile(greptime_native_histogram"),
10981 "{plan}"
10982 );
10983 assert!(!plan.contains("EmptyRelation"), "{plan}");
10984 }
10985
10986 #[tokio::test]
10987 async fn mixed_histogram_helpers_execute_classic_and_native_samples() {
10988 let state = build_query_engine_state();
10989 for (query, expected) in [
10990 (
10991 "histogram_quantile(0.5, mixed_histogram)",
10992 vec![("classic", 1.0), ("native", 0.0)],
10993 ),
10994 (
10995 "histogram_fraction(-Inf, +Inf, mixed_histogram)",
10996 vec![("classic", 1.0), ("native", 1.0)],
10997 ),
10998 ] {
10999 let plan = PromPlanner::stmt_to_plan(
11000 classic_and_native_histogram_table_provider("native", None, direct_or_histogram()),
11001 &operator_eval_stmt(query),
11002 &state,
11003 )
11004 .await
11005 .unwrap();
11006 let plan_text = plan.display_indent_schema().to_string();
11007 assert!(plan_text.contains("HistogramFold:"), "{plan_text}");
11008 assert!(plan_text.contains("prom_native_histogram_"), "{plan_text}");
11009 let value_field = plan
11010 .schema()
11011 .fields()
11012 .iter()
11013 .find(|field| field.data_type() == &ArrowDataType::Float64)
11014 .unwrap()
11015 .name()
11016 .clone();
11017
11018 let (_, batches) = execute(plan, &state).await;
11019 let mut actual = batches
11020 .iter()
11021 .flat_map(|batch| {
11022 let tags = batch
11023 .column_by_name("tag")
11024 .unwrap()
11025 .as_any()
11026 .downcast_ref::<StringArray>()
11027 .unwrap();
11028 let values = batch
11029 .column_by_name(&value_field)
11030 .unwrap()
11031 .as_any()
11032 .downcast_ref::<Float64Array>()
11033 .unwrap();
11034 (0..batch.num_rows()).map(|row| (tags.value(row), values.value(row)))
11035 })
11036 .collect::<Vec<_>>();
11037 actual.sort_by_key(|(tag, _)| *tag);
11038 assert_eq!(actual, expected, "{query}");
11039 }
11040 }
11041
11042 #[tokio::test]
11043 async fn mixed_histogram_helpers_report_annotations() {
11044 let state = build_query_engine_state();
11045 let mut native_histogram = direct_or_histogram();
11046 native_histogram.count = 2.0;
11047 native_histogram.sum = f64::NAN;
11048 for (native_tag, expected_rows, expected_warnings, expected_infos) in [
11049 (
11050 "classic",
11051 0,
11052 vec!["vector contains a mix of classic and native histograms"],
11053 vec![],
11054 ),
11055 (
11056 "native",
11057 2,
11058 vec![],
11059 vec!["input to histogram_quantile has NaN observations, result is skewed higher"],
11060 ),
11061 ] {
11062 let collector = PromqlAnnotationCollector::default();
11063 let plan = PromPlanner::stmt_to_plan_with_annotations(
11064 classic_and_native_histogram_table_provider(
11065 native_tag,
11066 None,
11067 native_histogram.clone(),
11068 ),
11069 &operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"),
11070 &state,
11071 Some(collector.clone()),
11072 )
11073 .await
11074 .unwrap();
11075
11076 let (_, batches) = execute(plan, &state).await;
11077 assert_eq!(
11078 batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
11079 expected_rows
11080 );
11081 let mut warnings = vec![];
11082 let mut infos = vec![];
11083 collector.append_to(&mut warnings, &mut infos);
11084 assert_eq!(warnings, expected_warnings);
11085 assert_eq!(infos, expected_infos);
11086 }
11087 }
11088
11089 #[tokio::test]
11090 async fn mixed_histogram_helper_preserves_native_le_and_scans_once() {
11091 let state = build_query_engine_state();
11092 let mut stmt = operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)");
11093 stmt.end = UNIX_EPOCH.checked_add(Duration::from_secs(2)).unwrap();
11094 let plan = PromPlanner::stmt_to_plan(
11095 classic_and_native_histogram_table_provider(
11096 "classic",
11097 Some("native"),
11098 direct_or_histogram(),
11099 ),
11100 &stmt,
11101 &state,
11102 )
11103 .await
11104 .unwrap();
11105 let plan_text = plan.display_indent_schema().to_string();
11106 assert_eq!(
11107 plan_text.matches("TableScan: mixed_histogram").count(),
11108 1,
11109 "{plan_text}"
11110 );
11111
11112 let value_field = plan
11113 .schema()
11114 .fields()
11115 .iter()
11116 .find(|field| field.data_type() == &ArrowDataType::Float64)
11117 .unwrap()
11118 .name()
11119 .clone();
11120 let (_, batches) = execute(plan, &state).await;
11121 let mut actual = batches
11122 .iter()
11123 .flat_map(|batch| {
11124 let le = batch
11125 .column_by_name(LE_COLUMN_NAME)
11126 .unwrap()
11127 .as_any()
11128 .downcast_ref::<StringArray>()
11129 .unwrap();
11130 let timestamps = batch
11131 .column_by_name("timestamp")
11132 .unwrap()
11133 .as_any()
11134 .downcast_ref::<TimestampMillisecondArray>()
11135 .unwrap();
11136 let values = batch
11137 .column_by_name(&value_field)
11138 .unwrap()
11139 .as_any()
11140 .downcast_ref::<Float64Array>()
11141 .unwrap();
11142 (0..batch.num_rows()).map(|row| {
11143 (
11144 timestamps.value(row),
11145 (!le.is_null(row)).then(|| le.value(row).to_string()),
11146 values.value(row),
11147 )
11148 })
11149 })
11150 .collect::<Vec<_>>();
11151 actual.sort_by(|lhs, rhs| (lhs.0, &lhs.1).cmp(&(rhs.0, &rhs.1)));
11152 assert_eq!(
11153 actual,
11154 vec![
11155 (1_000, None, 1.0),
11156 (1_000, Some("native".to_string()), 0.0),
11157 (2_000, None, 1.0),
11158 (2_000, Some("native".to_string()), 0.0),
11159 ]
11160 );
11161 }
11162
11163 #[tokio::test]
11164 async fn nested_histogram_helpers_ignore_unparsable_bucket_labels() {
11165 let state = build_query_engine_state();
11166 for native_le in [None, Some("native")] {
11167 for query in [
11168 "histogram_quantile(0.5, histogram_quantile(0.5, mixed_histogram))",
11169 "histogram_fraction(-Inf, +Inf, histogram_fraction(-Inf, +Inf, mixed_histogram))",
11170 ] {
11171 let plan = PromPlanner::stmt_to_plan(
11172 classic_and_native_histogram_table_provider(
11173 "native",
11174 native_le,
11175 direct_or_histogram(),
11176 ),
11177 &operator_eval_stmt(query),
11178 &state,
11179 )
11180 .await
11181 .unwrap();
11182
11183 let (_, batches) = execute(plan, &state).await;
11184 assert_eq!(
11185 batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
11186 0,
11187 "native_le={native_le:?}, query={query}"
11188 );
11189 }
11190 }
11191 }
11192
11193 #[tokio::test]
11194 async fn native_histogram_quantile_rejects_multi_field_input() {
11195 let table_provider = build_test_multi_histogram_table_provider("some_metric").await;
11196 let result = PromPlanner::stmt_to_plan(
11197 table_provider,
11198 &build_eval_stmt("histogram_quantile(0.9, some_metric)"),
11199 &build_query_engine_state(),
11200 )
11201 .await;
11202
11203 let err = result.expect_err("histogram_quantile on two native histogram fields must fail");
11204 assert!(
11205 err.to_string()
11206 .contains("Multi fields calculation is not supported in histogram_quantile"),
11207 "{err}"
11208 );
11209 }
11210
11211 #[tokio::test]
11212 async fn native_histogram_topk_uses_drop_udf() {
11213 let plan = native_histogram_plan("topk(1, some_metric)").await;
11214
11215 assert!(plan.contains("prom_native_histogram_drop_float"), "{plan}");
11216 assert!(
11217 plan.contains("Filter: prom_native_histogram_drop_float")
11218 && plan.contains("IS NOT NULL"),
11219 "{plan}"
11220 );
11221 }
11222
11223 #[tokio::test]
11224 async fn mixed_or_topk_bottomk_ignore_native_histograms() {
11225 for op in ["topk", "bottomk"] {
11226 let collector = PromqlAnnotationCollector::default();
11227 let state = build_query_engine_state();
11228 let plan = PromPlanner::stmt_to_plan_with_annotations(
11229 operator_table_provider(),
11230 &operator_eval_stmt(&format!("{op}(1, lf or on(tag) lh)")),
11231 &state,
11232 Some(collector.clone()),
11233 )
11234 .await
11235 .unwrap();
11236 let float_field = plan
11237 .schema()
11238 .fields()
11239 .iter()
11240 .find(|field| field.data_type() == &ArrowDataType::Float64)
11241 .unwrap()
11242 .name()
11243 .clone();
11244 assert!(
11245 plan.schema()
11246 .fields()
11247 .iter()
11248 .all(|field| field.data_type() != &PromPlanner::native_histogram_arrow_type()),
11249 "{plan:?}"
11250 );
11251
11252 let (_, batches) = execute(plan, &state).await;
11253 assert_eq!(values(&batches, &float_field), vec![2.0], "{op}");
11254 let mut warnings = vec![];
11255 let mut infos = vec![];
11256 collector.append_to(&mut warnings, &mut infos);
11257 assert!(warnings.is_empty());
11258 assert_eq!(
11259 infos,
11260 vec![format!(
11261 "{op}: dropped native histogram samples because this aggregation is not supported for native histograms"
11262 )]
11263 );
11264 }
11265 }
11266
11267 #[tokio::test]
11268 async fn native_histogram_scalar_is_ignored_before_scalar_calculate() {
11269 let plan = native_histogram_plan("scalar(some_metric)").await;
11270
11271 assert!(plan.contains("ScalarCalculate"), "{plan}");
11272 assert!(plan.contains("Filter: Boolean(false)"), "{plan}");
11273 assert!(!plan.contains("prom_native_histogram_drop"), "{plan}");
11274 }
11275
11276 #[tokio::test]
11277 async fn native_histogram_value_sort_is_empty_but_label_sort_preserves_samples() {
11278 for function in ["sort", "sort_desc"] {
11279 let plan = native_histogram_plan(&format!("{function}(some_metric)")).await;
11280
11281 assert!(plan.contains("Float64(NULL) IS NOT NULL"), "{plan}");
11282 assert!(
11283 !plan.contains(&format!("Sort: {}", greptime_native_histogram())),
11284 "{plan}"
11285 );
11286 assert!(!plan.contains("prom_native_histogram_drop"), "{plan}");
11287 }
11288
11289 for (function, direction) in [("sort_by_label", "ASC"), ("sort_by_label_desc", "DESC")] {
11290 let plan = native_histogram_plan(&format!("{function}(some_metric, \"tag_0\")")).await;
11291
11292 assert!(plan.contains(&format!("tag_0 {direction}")), "{plan}");
11293 assert!(plan.contains(greptime_native_histogram()), "{plan}");
11294 assert!(!plan.contains("Float64(NULL) IS NOT NULL"), "{plan}");
11295 }
11296 }
11297
11298 #[tokio::test]
11299 async fn unsupported_native_histogram_functions_use_drop_udf() {
11300 for query in [
11301 "deriv(some_metric[5m])",
11302 "min_over_time(some_metric[5m])",
11303 "quantile_over_time(0.9, some_metric[5m])",
11304 "predict_linear(some_metric[5m], 60)",
11305 "round(some_metric)",
11306 "abs(some_metric)",
11307 ] {
11308 let plan = native_histogram_plan(query).await;
11309
11310 assert!(
11311 plan.contains("prom_native_histogram_drop_float"),
11312 "{query}\n{plan}"
11313 );
11314 }
11315 }
11316
11317 #[tokio::test]
11318 async fn native_histogram_absent_over_time_uses_native_udf() {
11319 let plan = native_histogram_plan("absent_over_time(some_metric[5m])").await;
11320
11321 assert!(
11322 plan.contains("prom_native_histogram_absent_over_time"),
11323 "{plan}"
11324 );
11325 }
11326
11327 #[tokio::test]
11328 async fn native_histogram_all_function_arms_route_correctly() {
11329 let cases = [
11334 (
11336 "increase(some_metric[5m])",
11337 "prom_native_histogram_increase",
11338 ),
11339 ("rate(some_metric[5m])", "prom_native_histogram_rate"),
11340 ("delta(some_metric[5m])", "prom_native_histogram_delta"),
11341 ("idelta(some_metric[5m])", "prom_native_histogram_idelta"),
11342 ("irate(some_metric[5m])", "prom_native_histogram_irate"),
11343 ("resets(some_metric[5m])", "prom_native_histogram_resets"),
11344 ("changes(some_metric[5m])", "prom_native_histogram_changes"),
11345 (
11346 "avg_over_time(some_metric[5m])",
11347 "prom_native_histogram_avg_over_time",
11348 ),
11349 (
11350 "sum_over_time(some_metric[5m])",
11351 "prom_native_histogram_sum_over_time",
11352 ),
11353 (
11354 "count_over_time(some_metric[5m])",
11355 "prom_native_histogram_count_over_time",
11356 ),
11357 (
11358 "last_over_time(some_metric[5m])",
11359 "prom_native_histogram_last_over_time",
11360 ),
11361 (
11362 "present_over_time(some_metric[5m])",
11363 "prom_native_histogram_present_over_time",
11364 ),
11365 ("deriv(some_metric[5m])", "prom_native_histogram_drop_float"),
11367 (
11368 "min_over_time(some_metric[5m])",
11369 "prom_native_histogram_drop_float",
11370 ),
11371 (
11372 "max_over_time(some_metric[5m])",
11373 "prom_native_histogram_drop_float",
11374 ),
11375 (
11376 "stddev_over_time(some_metric[5m])",
11377 "prom_native_histogram_drop_float",
11378 ),
11379 (
11380 "stdvar_over_time(some_metric[5m])",
11381 "prom_native_histogram_drop_float",
11382 ),
11383 (
11384 "quantile_over_time(0.9, some_metric[5m])",
11385 "prom_native_histogram_drop_float",
11386 ),
11387 (
11388 "predict_linear(some_metric[5m], 60)",
11389 "prom_native_histogram_drop_float",
11390 ),
11391 (
11392 "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)",
11393 "prom_native_histogram_drop_float",
11394 ),
11395 ("round(some_metric)", "prom_native_histogram_drop_float"),
11396 ("rad(some_metric)", "prom_native_histogram_drop_float"),
11397 ("deg(some_metric)", "prom_native_histogram_drop_float"),
11398 ("sgn(some_metric)", "prom_native_histogram_drop_float"),
11399 (
11401 "histogram_count(some_metric)",
11402 "prom_native_histogram_count",
11403 ),
11404 ("histogram_sum(some_metric)", "prom_native_histogram_sum"),
11405 ("histogram_avg(some_metric)", "prom_native_histogram_avg"),
11406 (
11407 "histogram_stddev(some_metric)",
11408 "prom_native_histogram_stddev",
11409 ),
11410 (
11411 "histogram_stdvar(some_metric)",
11412 "prom_native_histogram_stdvar",
11413 ),
11414 (
11415 "histogram_fraction(-2 + 1, 2 / 2, some_metric)",
11416 "prom_native_histogram_fraction",
11417 ),
11418 ];
11419
11420 for (query, expected_udf) in cases {
11421 let plan = native_histogram_plan(query).await;
11422 assert!(plan.contains(expected_udf), "{query}\n{plan}");
11423 if query.starts_with("histogram_fraction") {
11424 assert!(plan.contains("Float64(-1)"), "{query}\n{plan}");
11425 }
11426 }
11427 }
11428
11429 #[tokio::test]
11430 async fn mixed_native_histogram_ranges_use_coordinated_udfs() {
11431 let dual_output = [
11432 "increase(some_metric[5m])",
11433 "rate(some_metric[5m])",
11434 "delta(some_metric[5m])",
11435 "idelta(some_metric[5m])",
11436 "irate(some_metric[5m])",
11437 "avg_over_time(some_metric[5m])",
11438 "sum_over_time(some_metric[5m])",
11439 "last_over_time(some_metric[5m])",
11440 ];
11441 let float_output = [
11442 "resets(some_metric[5m])",
11443 "changes(some_metric[5m])",
11444 "deriv(some_metric[5m])",
11445 "min_over_time(some_metric[5m])",
11446 "max_over_time(some_metric[5m])",
11447 "count_over_time(some_metric[5m])",
11448 "absent_over_time(some_metric[5m])",
11449 "present_over_time(some_metric[5m])",
11450 "stddev_over_time(some_metric[5m])",
11451 "stdvar_over_time(some_metric[5m])",
11452 "quantile_over_time(0.9, some_metric[5m])",
11453 "predict_linear(some_metric[5m], 60)",
11454 "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)",
11455 ];
11456
11457 for query in dual_output.iter().chain(float_output.iter()) {
11458 let plan = PromPlanner::stmt_to_plan(
11459 build_test_mixed_native_histogram_table_provider("some_metric").await,
11460 &build_eval_stmt(query),
11461 &build_query_engine_state(),
11462 )
11463 .await
11464 .unwrap()
11465 .display_indent_schema()
11466 .to_string();
11467 assert!(plan.contains("prom_mixed_range_float"), "{query}\n{plan}");
11468 assert_eq!(
11469 plan.contains("prom_mixed_range_histogram"),
11470 dual_output.contains(query),
11471 "{query}\n{plan}"
11472 );
11473 }
11474
11475 let plan = PromPlanner::stmt_to_plan(
11476 build_test_mixed_native_histogram_table_provider("some_metric").await,
11477 &build_eval_stmt("sum_over_time(rate(some_metric[5m])[10m:1m])"),
11478 &build_query_engine_state(),
11479 )
11480 .await
11481 .unwrap()
11482 .display_indent_schema()
11483 .to_string();
11484 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]
11485 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]
11486 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))]
11487 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]
11488 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]
11489 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]
11490 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]
11491 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))]
11492 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]
11493 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]
11494 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]
11495 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]
11496 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]"#;
11497 assert_eq!(plan, expected);
11498 }
11499
11500 #[tokio::test]
11501 async fn mixed_native_histogram_rate_executes_real_ranges() {
11502 let schema = Arc::new(ArrowSchema::new(vec![
11503 Field::new(
11504 "timestamp",
11505 ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
11506 false,
11507 ),
11508 Field::new(greptime_value(), ArrowDataType::Float64, true),
11509 Field::new(
11510 greptime_native_histogram(),
11511 native_histogram_value_type().as_arrow_type(),
11512 true,
11513 ),
11514 ]));
11515 let batch = RecordBatch::try_new(
11516 schema.clone(),
11517 vec![
11518 Arc::new(TimestampMillisecondArray::from(vec![1000, 2000, 3000])),
11519 Arc::new(Float64Array::from(vec![Some(1.0), None, Some(3.0)])),
11520 build_histogram_array(&[None, Some(direct_or_histogram()), None]),
11521 ],
11522 )
11523 .unwrap();
11524 let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap());
11525 let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None)
11526 .unwrap()
11527 .build()
11528 .unwrap();
11529 let collector = PromqlAnnotationCollector::default();
11530 let mut planner = PromPlanner {
11531 table_provider: build_test_table_provider_with_fields(
11532 &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
11533 &[],
11534 )
11535 .await,
11536 ctx: PromPlannerContext {
11537 start: 3000,
11538 end: 3000,
11539 interval: 1000,
11540 range: Some(3000),
11541 time_index_column: Some("timestamp".to_string()),
11542 field_columns: vec![
11543 greptime_native_histogram().to_string(),
11544 greptime_value().to_string(),
11545 ],
11546 ..Default::default()
11547 },
11548 promql_annotations: Some(collector.clone()),
11549 };
11550 let input = LogicalPlan::Extension(Extension {
11551 node: Arc::new(
11552 RangeManipulate::new(
11553 3000,
11554 3000,
11555 1000,
11556 0,
11557 3000,
11558 "timestamp".to_string(),
11559 planner.ctx.field_columns.clone(),
11560 input,
11561 )
11562 .unwrap(),
11563 ),
11564 });
11565 let PromExpr::Call(call) = parser::parse("rate(mixed[3s])").unwrap() else {
11566 unreachable!()
11567 };
11568 let preserve_any_value = PromPlanner::field_columns_are_alternative_samples(
11569 input.schema(),
11570 &planner.ctx.field_columns,
11571 );
11572 let state = build_query_engine_state();
11573 let (mut exprs, _) = planner
11574 .create_function_expr(&call.func, vec![], input.schema(), &state)
11575 .unwrap();
11576 exprs.insert(0, planner.create_time_index_column_expr().unwrap());
11577 let plan = LogicalPlanBuilder::from(input)
11578 .project(exprs)
11579 .unwrap()
11580 .filter(
11581 planner
11582 .create_empty_values_filter_expr(preserve_any_value)
11583 .unwrap(),
11584 )
11585 .unwrap()
11586 .build()
11587 .unwrap();
11588 let (_, batches) = execute(plan, &state).await;
11589 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
11590 let mut warnings = Vec::new();
11591 collector.append_to(&mut warnings, &mut Vec::new());
11592 assert!(
11593 warnings
11594 .iter()
11595 .any(|warning| warning.contains("mix of float and native histogram"))
11596 );
11597 }
11598
11599 #[tokio::test]
11600 async fn native_histogram_mixed_field_table_behaves() {
11601 let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
11605 let plan = PromPlanner::stmt_to_plan(
11606 table_provider,
11607 &build_eval_stmt("histogram_count(some_metric)"),
11608 &build_query_engine_state(),
11609 )
11610 .await
11611 .unwrap();
11612 let plan_str = plan.display_indent_schema().to_string();
11613 assert!(
11614 plan_str.contains("prom_native_histogram_count"),
11615 "{plan_str}"
11616 );
11617 assert!(!plan_str.contains("Float64(NULL)"), "{plan_str}");
11618 assert!(
11619 plan_str.contains("prom_native_histogram_count(greptime_native_histogram) IS NOT NULL"),
11620 "{plan_str}"
11621 );
11622
11623 let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
11625 let plan = PromPlanner::stmt_to_plan(
11626 table_provider,
11627 &build_eval_stmt("sort(some_metric)"),
11628 &build_query_engine_state(),
11629 )
11630 .await
11631 .unwrap();
11632 let plan_str = plan.display_indent_schema().to_string();
11633 assert!(
11634 plan_str.contains("greptime_value ASC NULLS FIRST"),
11635 "{plan_str}"
11636 );
11637 assert!(
11638 !plan_str.contains("greptime_native_histogram ASC"),
11639 "{plan_str}"
11640 );
11641
11642 let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
11644 let plan = PromPlanner::stmt_to_plan(
11645 table_provider,
11646 &build_eval_stmt("scalar(some_metric)"),
11647 &build_query_engine_state(),
11648 )
11649 .await
11650 .unwrap();
11651 let plan_str = plan.display_indent_schema().to_string();
11652 assert!(plan_str.contains("ScalarCalculate"), "{plan_str}");
11653 assert!(
11654 plan_str.contains("greptime_value IS NOT NULL"),
11655 "{plan_str}"
11656 );
11657
11658 let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
11660 let plan = PromPlanner::stmt_to_plan(
11661 table_provider,
11662 &build_eval_stmt(r#"label_replace(some_metric, "copied", "$1", "tag_0", "(.*)")"#),
11663 &build_query_engine_state(),
11664 )
11665 .await
11666 .unwrap();
11667 let plan_str = plan.display_indent_schema().to_string();
11668 let filter = plan_str.lines().next().unwrap();
11669 assert!(
11670 filter.starts_with("Filter: ")
11671 && filter.contains("greptime_native_histogram IS NOT NULL")
11672 && filter.contains(" OR ")
11673 && filter.contains("greptime_value IS NOT NULL"),
11674 "{plan_str}"
11675 );
11676 }
11677
11678 #[tokio::test]
11679 async fn less_filter_on_value() {
11680 let query = "some_metric < 1.2345";
11681 let expected = String::from(
11682 "Filter: some_metric.field_0 < Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11683 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11684 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11685 \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]\
11686 \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]\
11687 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
11688 );
11689
11690 indie_query_plan_compare(query, expected).await;
11691 }
11692
11693 #[tokio::test]
11694 async fn count_over_time() {
11695 let query = "count_over_time(some_metric[5m])";
11696 let expected = String::from(
11697 "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]\
11698 \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]\
11699 \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))]\
11700 \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11701 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11702 \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]\
11703 \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]\
11704 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
11705 );
11706
11707 indie_query_plan_compare(query, expected).await;
11708 }
11709
11710 #[tokio::test]
11713 async fn count_over_time_subquery() {
11714 let query = "count_over_time(some_metric[10m:1m])";
11715 let expected = String::from(
11716 "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]\
11717 \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]\
11718 \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))]\
11719 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11720 \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]\
11721 \n PromInstantManipulate: range=[-540000..100000000], lookback=[1000], interval=[60000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11722 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11723 \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]\
11724 \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]\
11725 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
11726 );
11727 indie_query_plan_compare(query, expected).await;
11728 }
11729
11730 #[tokio::test]
11731 async fn test_hash_join() {
11732 let mut eval_stmt = EvalStmt {
11733 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11734 start: UNIX_EPOCH,
11735 end: UNIX_EPOCH
11736 .checked_add(Duration::from_secs(100_000))
11737 .unwrap(),
11738 interval: Duration::from_secs(5),
11739 lookback_delta: Duration::from_secs(1),
11740 };
11741
11742 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"}"#;
11743
11744 let prom_expr = parser::parse(case).unwrap();
11745 eval_stmt.expr = prom_expr;
11746 let table_provider = build_test_table_provider_with_fields(
11747 &[
11748 (
11749 DEFAULT_SCHEMA_NAME.to_string(),
11750 "http_server_requests_seconds_sum".to_string(),
11751 ),
11752 (
11753 DEFAULT_SCHEMA_NAME.to_string(),
11754 "http_server_requests_seconds_count".to_string(),
11755 ),
11756 ],
11757 &["uri", "kubernetes_namespace", "kubernetes_pod_name"],
11758 )
11759 .await;
11760 let plan =
11762 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11763 .await
11764 .unwrap();
11765 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\
11766 \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\
11767 \n SubqueryAlias: http_server_requests_seconds_sum\
11768 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\
11769 \n PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\
11770 \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\
11771 \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)\
11772 \n TableScan: http_server_requests_seconds_sum\
11773 \n SubqueryAlias: http_server_requests_seconds_count\
11774 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\
11775 \n PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\
11776 \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\
11777 \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)\
11778 \n TableScan: http_server_requests_seconds_count";
11779 assert_eq!(plan.to_string(), expected);
11780 }
11781
11782 #[tokio::test]
11783 async fn test_nested_histogram_quantile() {
11784 let mut eval_stmt = EvalStmt {
11785 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11786 start: UNIX_EPOCH,
11787 end: UNIX_EPOCH
11788 .checked_add(Duration::from_secs(100_000))
11789 .unwrap(),
11790 interval: Duration::from_secs(5),
11791 lookback_delta: Duration::from_secs(1),
11792 };
11793
11794 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]*-(.*)")"#;
11795
11796 let prom_expr = parser::parse(case).unwrap();
11797 eval_stmt.expr = prom_expr;
11798 let table_provider = build_test_table_provider_with_fields(
11799 &[(
11800 DEFAULT_SCHEMA_NAME.to_string(),
11801 "greptime_servers_grpc_requests_elapsed_bucket".to_string(),
11802 )],
11803 &["pod", "le", "path", "code", "container"],
11804 )
11805 .await;
11806 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11808 .await
11809 .unwrap();
11810 }
11811
11812 #[tokio::test]
11813 async fn test_histogram_quantile_binary_op() {
11814 let mut eval_stmt = EvalStmt {
11815 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11816 start: UNIX_EPOCH,
11817 end: UNIX_EPOCH
11818 .checked_add(Duration::from_secs(100_000))
11819 .unwrap(),
11820 interval: Duration::from_secs(5),
11821 lookback_delta: Duration::from_secs(1),
11822 };
11823
11824 let case = r#"histogram_quantile(0.5, sum by (le, pod) (rate(http_request_duration_seconds_bucket[5m]))) + 0"#;
11828
11829 let prom_expr = parser::parse(case).unwrap();
11830 eval_stmt.expr = prom_expr;
11831 let table_provider = build_test_table_provider_with_fields(
11832 &[(
11833 DEFAULT_SCHEMA_NAME.to_string(),
11834 "http_request_duration_seconds_bucket".to_string(),
11835 )],
11836 &["pod", "le"],
11837 )
11838 .await;
11839 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11841 .await
11842 .unwrap();
11843 }
11844
11845 #[tokio::test]
11846 async fn test_parse_and_operator() {
11847 let mut eval_stmt = EvalStmt {
11848 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11849 start: UNIX_EPOCH,
11850 end: UNIX_EPOCH
11851 .checked_add(Duration::from_secs(100_000))
11852 .unwrap(),
11853 interval: Duration::from_secs(5),
11854 lookback_delta: Duration::from_secs(1),
11855 };
11856
11857 let cases = [
11858 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)"#,
11859 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)"#,
11860 ];
11861
11862 for case in cases {
11863 let prom_expr = parser::parse(case).unwrap();
11864 eval_stmt.expr = prom_expr;
11865 let table_provider = build_test_table_provider_with_fields(
11866 &[
11867 (
11868 DEFAULT_SCHEMA_NAME.to_string(),
11869 "kubelet_volume_stats_used_bytes".to_string(),
11870 ),
11871 (
11872 DEFAULT_SCHEMA_NAME.to_string(),
11873 "kubelet_volume_stats_capacity_bytes".to_string(),
11874 ),
11875 ],
11876 &["namespace", "persistentvolumeclaim"],
11877 )
11878 .await;
11879 let _ =
11881 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11882 .await
11883 .unwrap();
11884 }
11885 }
11886
11887 #[tokio::test]
11888 async fn test_nested_binary_op() {
11889 let mut eval_stmt = EvalStmt {
11890 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11891 start: UNIX_EPOCH,
11892 end: UNIX_EPOCH
11893 .checked_add(Duration::from_secs(100_000))
11894 .unwrap(),
11895 interval: Duration::from_secs(5),
11896 lookback_delta: Duration::from_secs(1),
11897 };
11898
11899 let case = r#"sum(rate(nginx_ingress_controller_requests{job=~".*"}[2m])) -
11900 (
11901 sum(rate(nginx_ingress_controller_requests{namespace=~".*"}[2m]))
11902 or
11903 vector(0)
11904 )"#;
11905
11906 let prom_expr = parser::parse(case).unwrap();
11907 eval_stmt.expr = prom_expr;
11908 let table_provider = build_test_table_provider_with_fields(
11909 &[(
11910 DEFAULT_SCHEMA_NAME.to_string(),
11911 "nginx_ingress_controller_requests".to_string(),
11912 )],
11913 &["namespace", "job"],
11914 )
11915 .await;
11916 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11918 .await
11919 .unwrap();
11920 }
11921
11922 #[tokio::test]
11923 async fn test_parse_or_operator() {
11924 let mut eval_stmt = EvalStmt {
11925 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11926 start: UNIX_EPOCH,
11927 end: UNIX_EPOCH
11928 .checked_add(Duration::from_secs(100_000))
11929 .unwrap(),
11930 interval: Duration::from_secs(5),
11931 lookback_delta: Duration::from_secs(1),
11932 };
11933
11934 let case = r#"
11935 sum(rate(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}[120s])) by (cluster_name,tenant_name) /
11936 (sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) * 100)
11937 or
11938 200 * sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) /
11939 sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)"#;
11940
11941 let table_provider = build_test_table_provider_with_fields(
11942 &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())],
11943 &["tenant_name", "cluster_name"],
11944 )
11945 .await;
11946 eval_stmt.expr = parser::parse(case).unwrap();
11947 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11948 .await
11949 .unwrap();
11950
11951 let case = r#"sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
11952 (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) +
11953 sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
11954 (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0
11955 or
11956 sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
11957 (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0
11958 or
11959 sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
11960 (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0"#;
11961 let table_provider = build_test_table_provider_with_fields(
11962 &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())],
11963 &["tenant_name", "cluster_name"],
11964 )
11965 .await;
11966 eval_stmt.expr = parser::parse(case).unwrap();
11967 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11968 .await
11969 .unwrap();
11970
11971 let case = r#"(sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) +
11972 sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or
11973 (sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or
11974 (sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name))"#;
11975 let table_provider = build_test_table_provider_with_fields(
11976 &[
11977 (
11978 DEFAULT_SCHEMA_NAME.to_string(),
11979 "background_waitevent_cnt".to_string(),
11980 ),
11981 (
11982 DEFAULT_SCHEMA_NAME.to_string(),
11983 "foreground_waitevent_cnt".to_string(),
11984 ),
11985 ],
11986 &["tenant_name", "cluster_name"],
11987 )
11988 .await;
11989 eval_stmt.expr = parser::parse(case).unwrap();
11990 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11991 .await
11992 .unwrap();
11993
11994 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)"#;
11995 let table_provider = build_test_table_provider_with_fields(
11996 &[
11997 (DEFAULT_SCHEMA_NAME.to_string(), "node_load1".to_string()),
11998 (
11999 DEFAULT_SCHEMA_NAME.to_string(),
12000 "container_cpu_load_average_10s".to_string(),
12001 ),
12002 (
12003 DEFAULT_SCHEMA_NAME.to_string(),
12004 "container_spec_cpu_quota".to_string(),
12005 ),
12006 ],
12007 &["cluster_name", "host_name"],
12008 )
12009 .await;
12010 eval_stmt.expr = parser::parse(case).unwrap();
12011 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12012 .await
12013 .unwrap();
12014 }
12015
12016 #[tokio::test]
12017 async fn value_matcher() {
12018 let mut eval_stmt = EvalStmt {
12020 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12021 start: UNIX_EPOCH,
12022 end: UNIX_EPOCH
12023 .checked_add(Duration::from_secs(100_000))
12024 .unwrap(),
12025 interval: Duration::from_secs(5),
12026 lookback_delta: Duration::from_secs(1),
12027 };
12028
12029 let cases = [
12030 (
12032 r#"some_metric{__field__="field_1"}"#,
12033 vec![
12034 "some_metric.field_1",
12035 "some_metric.tag_0",
12036 "some_metric.tag_1",
12037 "some_metric.tag_2",
12038 "some_metric.timestamp",
12039 ],
12040 ),
12041 (
12043 r#"some_metric{__field__="field_1", __field__="field_0"}"#,
12044 vec![
12045 "some_metric.field_0",
12046 "some_metric.field_1",
12047 "some_metric.tag_0",
12048 "some_metric.tag_1",
12049 "some_metric.tag_2",
12050 "some_metric.timestamp",
12051 ],
12052 ),
12053 (
12055 r#"some_metric{__field__!="field_1"}"#,
12056 vec![
12057 "some_metric.field_0",
12058 "some_metric.field_2",
12059 "some_metric.tag_0",
12060 "some_metric.tag_1",
12061 "some_metric.tag_2",
12062 "some_metric.timestamp",
12063 ],
12064 ),
12065 (
12067 r#"some_metric{__field__!="field_1", __field__!="field_2"}"#,
12068 vec![
12069 "some_metric.field_0",
12070 "some_metric.tag_0",
12071 "some_metric.tag_1",
12072 "some_metric.tag_2",
12073 "some_metric.timestamp",
12074 ],
12075 ),
12076 (
12078 r#"some_metric{__field__="field_1", __field__!="field_0"}"#,
12079 vec![
12080 "some_metric.field_1",
12081 "some_metric.tag_0",
12082 "some_metric.tag_1",
12083 "some_metric.tag_2",
12084 "some_metric.timestamp",
12085 ],
12086 ),
12087 (
12089 r#"some_metric{__field__="field_2", __field__!="field_2"}"#,
12090 vec![
12091 "some_metric.tag_0",
12092 "some_metric.tag_1",
12093 "some_metric.tag_2",
12094 "some_metric.timestamp",
12095 ],
12096 ),
12097 (
12099 r#"some_metric{__field__=~"field_1|field_2"}"#,
12100 vec![
12101 "some_metric.field_1",
12102 "some_metric.field_2",
12103 "some_metric.tag_0",
12104 "some_metric.tag_1",
12105 "some_metric.tag_2",
12106 "some_metric.timestamp",
12107 ],
12108 ),
12109 (
12111 r#"some_metric{__field__!~"field_1|field_2"}"#,
12112 vec![
12113 "some_metric.field_0",
12114 "some_metric.tag_0",
12115 "some_metric.tag_1",
12116 "some_metric.tag_2",
12117 "some_metric.timestamp",
12118 ],
12119 ),
12120 ];
12121
12122 for case in cases {
12123 let prom_expr = parser::parse(case.0).unwrap();
12124 eval_stmt.expr = prom_expr;
12125 let table_provider = build_test_table_provider(
12126 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12127 3,
12128 3,
12129 )
12130 .await;
12131 let plan =
12132 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12133 .await
12134 .unwrap();
12135 let mut fields = plan.schema().field_names();
12136 let mut expected = case.1.into_iter().map(String::from).collect::<Vec<_>>();
12137 fields.sort();
12138 expected.sort();
12139 assert_eq!(fields, expected, "case: {:?}", case.0);
12140 }
12141
12142 let bad_cases = [
12143 r#"some_metric{__field__="nonexistent"}"#,
12144 r#"some_metric{__field__!="nonexistent"}"#,
12145 ];
12146
12147 for case in bad_cases {
12148 let prom_expr = parser::parse(case).unwrap();
12149 eval_stmt.expr = prom_expr;
12150 let table_provider = build_test_table_provider(
12151 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12152 3,
12153 3,
12154 )
12155 .await;
12156 let plan =
12157 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12158 .await;
12159 assert!(plan.is_err(), "case: {:?}", case);
12160 }
12161 }
12162
12163 #[tokio::test]
12164 async fn custom_schema() {
12165 let query = "some_alt_metric{__schema__=\"greptime_private\"}";
12166 let expected = String::from(
12167 "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12168 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12169 \n Sort: 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]\
12170 \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]\
12171 \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
12172 );
12173
12174 indie_query_plan_compare(query, expected).await;
12175
12176 let query = "some_alt_metric{__database__=\"greptime_private\"}";
12177 let expected = String::from(
12178 "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12179 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12180 \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]\
12181 \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]\
12182 \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
12183 );
12184
12185 indie_query_plan_compare(query, expected).await;
12186
12187 let query = "some_alt_metric{__schema__=\"greptime_private\"} / some_metric";
12188 let expected = String::from(
12189 "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]\
12190 \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]\
12191 \n SubqueryAlias: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12192 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12193 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12194 \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]\
12195 \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]\
12196 \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12197 \n SubqueryAlias: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12198 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12199 \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12200 \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]\
12201 \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]\
12202 \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
12203 );
12204
12205 indie_query_plan_compare(query, expected).await;
12206 }
12207
12208 #[tokio::test]
12209 async fn only_equals_is_supported_for_special_matcher() {
12210 let queries = &[
12211 "some_alt_metric{__schema__!=\"greptime_private\"}",
12212 "some_alt_metric{__schema__=~\"lalala\"}",
12213 "some_alt_metric{__database__!=\"greptime_private\"}",
12214 "some_alt_metric{__database__=~\"lalala\"}",
12215 ];
12216
12217 for query in queries {
12218 let prom_expr = parser::parse(query).unwrap();
12219 let eval_stmt = EvalStmt {
12220 expr: prom_expr,
12221 start: UNIX_EPOCH,
12222 end: UNIX_EPOCH
12223 .checked_add(Duration::from_secs(100_000))
12224 .unwrap(),
12225 interval: Duration::from_secs(5),
12226 lookback_delta: Duration::from_secs(1),
12227 };
12228
12229 let table_provider = build_test_table_provider(
12230 &[
12231 (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
12232 (
12233 "greptime_private".to_string(),
12234 "some_alt_metric".to_string(),
12235 ),
12236 ],
12237 1,
12238 1,
12239 )
12240 .await;
12241
12242 let plan =
12243 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12244 .await;
12245 assert!(plan.is_err(), "query: {:?}", query);
12246 }
12247 }
12248
12249 #[tokio::test]
12250 async fn native_scan_bounds_preserve_zero_lookback_and_overflow() {
12251 let table_provider = build_test_table_provider(
12252 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12253 1,
12254 1,
12255 )
12256 .await;
12257 let mut planner = PromPlanner {
12258 table_provider,
12259 ctx: PromPlannerContext::from_eval_stmt(&build_eval_stmt("some_metric")),
12260 promql_annotations: None,
12261 };
12262 planner.ctx.time_index_column = Some("timestamp".to_string());
12263 planner.ctx.start = 1_000;
12264 planner.ctx.lookback_delta = 0;
12265 let schema = Arc::new(
12266 DFSchema::try_from(ArrowSchema::new(vec![Field::new(
12267 "timestamp",
12268 ArrowDataType::Timestamp(ArrowTimeUnit::Nanosecond, None),
12269 false,
12270 )]))
12271 .unwrap(),
12272 );
12273 for (end, interval, windows) in [
12274 (1_000, 1_000, 1),
12275 (2_000, 1_000, 1),
12276 (7_201_000, 7_200_000, 2),
12277 ] {
12278 planner.ctx.end = end;
12279 planner.ctx.interval = interval;
12280 let filter = planner
12281 .build_time_index_filter(0, &schema)
12282 .unwrap()
12283 .unwrap()
12284 .to_string();
12285 assert_eq!(filter.matches(">=").count(), windows, "{filter}");
12286 assert!(
12287 filter.contains("TimestampNanosecond(1000000000, None)"),
12288 "{filter}"
12289 );
12290 }
12291 planner.ctx.end = i64::MAX;
12292 let filter = planner
12293 .build_time_index_filter(0, &schema)
12294 .unwrap()
12295 .unwrap()
12296 .to_string();
12297 assert!(
12298 filter.contains("timestamp >= TimestampNanosecond(1000000000, None)"),
12299 "{filter}"
12300 );
12301
12302 let ms_schema = Arc::new(
12305 DFSchema::try_from(ArrowSchema::new(vec![Field::new(
12306 "timestamp",
12307 ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
12308 false,
12309 )]))
12310 .unwrap(),
12311 );
12312 planner.ctx.start = i64::MIN + 100;
12313 planner.ctx.end = planner.ctx.start;
12314 planner.ctx.lookback_delta = 200;
12315 let filter = planner
12316 .build_time_index_filter(0, &ms_schema)
12317 .unwrap()
12318 .unwrap()
12319 .to_string();
12320 assert_eq!(
12321 filter,
12322 format!(
12323 "timestamp <= TimestampMillisecond({}, None)",
12324 i64::MIN + 100
12325 )
12326 );
12327
12328 planner.ctx.start = 0;
12331 planner.ctx.end = 0;
12332 planner.ctx.lookback_delta = 300_000;
12333 let filter = planner
12334 .build_time_index_filter(9_223_372_036_854, &schema)
12335 .unwrap()
12336 .unwrap()
12337 .to_string();
12338 assert_eq!(
12339 filter,
12340 "timestamp <= TimestampNanosecond(-9223372036854000000, None)"
12341 );
12342 }
12343
12344 #[tokio::test]
12345 async fn test_non_ms_precision() {
12346 let catalog_list = MemoryCatalogManager::with_default_setup();
12347 let columns = vec![
12348 ColumnSchema::new(
12349 "tag".to_string(),
12350 ConcreteDataType::string_datatype(),
12351 false,
12352 ),
12353 ColumnSchema::new(
12354 "timestamp".to_string(),
12355 ConcreteDataType::timestamp_nanosecond_datatype(),
12356 false,
12357 )
12358 .with_time_index(true),
12359 ColumnSchema::new(
12360 "field".to_string(),
12361 ConcreteDataType::float64_datatype(),
12362 true,
12363 ),
12364 ];
12365 let schema = Arc::new(Schema::new(columns));
12366 let table_meta = TableMetaBuilder::empty()
12367 .schema(schema)
12368 .primary_key_indices(vec![0])
12369 .value_indices(vec![2])
12370 .next_column_id(1024)
12371 .build()
12372 .unwrap();
12373 let table_info = TableInfoBuilder::default()
12374 .name("metrics".to_string())
12375 .meta(table_meta)
12376 .build()
12377 .unwrap();
12378 let table = EmptyTable::from_table_info(&table_info);
12379 assert!(
12380 catalog_list
12381 .register_table_sync(RegisterTableRequest {
12382 catalog: DEFAULT_CATALOG_NAME.to_string(),
12383 schema: DEFAULT_SCHEMA_NAME.to_string(),
12384 table_name: "metrics".to_string(),
12385 table_id: 1024,
12386 table,
12387 })
12388 .is_ok()
12389 );
12390
12391 let plan = PromPlanner::stmt_to_plan(
12392 DfTableSourceProvider::new(
12393 catalog_list.clone(),
12394 false,
12395 QueryContext::arc(),
12396 DummyDecoder::arc(),
12397 true,
12398 ),
12399 &EvalStmt {
12400 expr: parser::parse("metrics{tag = \"1\"}").unwrap(),
12401 start: UNIX_EPOCH,
12402 end: UNIX_EPOCH
12403 .checked_add(Duration::from_secs(100_000))
12404 .unwrap(),
12405 interval: Duration::from_secs(5),
12406 lookback_delta: Duration::from_secs(1),
12407 },
12408 &build_query_engine_state(),
12409 )
12410 .await
12411 .unwrap();
12412 assert_eq!(
12413 plan.display_indent_schema().to_string(),
12414 "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\n PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp > TimestampNanosecond(-1000000000, None) AND metrics.timestamp <= TimestampNanosecond(100000000000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Projection: metrics.field, metrics.tag, metrics.timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]"
12415 );
12416 let plan = PromPlanner::stmt_to_plan(
12417 DfTableSourceProvider::new(
12418 catalog_list.clone(),
12419 false,
12420 QueryContext::arc(),
12421 DummyDecoder::arc(),
12422 true,
12423 ),
12424 &EvalStmt {
12425 expr: parser::parse("avg_over_time(metrics{tag = \"1\"}[5s])").unwrap(),
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 &build_query_engine_state(),
12434 )
12435 .await
12436 .unwrap();
12437 assert_eq!(
12438 plan.display_indent_schema().to_string(),
12439 "Filter: prom_avg_over_time(timestamp_range,field) IS NOT NULL [timestamp:Timestamp(ms), prom_avg_over_time(timestamp_range,field):Float64;N, tag:Utf8]\n Projection: metrics.timestamp, prom_avg_over_time(timestamp_range, field) AS prom_avg_over_time(timestamp_range,field), metrics.tag [timestamp:Timestamp(ms), prom_avg_over_time(timestamp_range,field):Float64;N, tag:Utf8]\n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[5000], time index=[timestamp], values=[\"field\"] [field:Dictionary(Int64, Float64);N, tag:Utf8, timestamp:Timestamp(ms), timestamp_range:Dictionary(Int64, Timestamp(ms))]\n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp > TimestampNanosecond(-5000000000, None) AND metrics.timestamp <= TimestampNanosecond(100000000000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Projection: metrics.field, metrics.tag, metrics.timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]"
12440 );
12441 }
12442
12443 #[tokio::test]
12444 async fn test_nonexistent_label() {
12445 let mut eval_stmt = EvalStmt {
12447 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12448 start: UNIX_EPOCH,
12449 end: UNIX_EPOCH
12450 .checked_add(Duration::from_secs(100_000))
12451 .unwrap(),
12452 interval: Duration::from_secs(5),
12453 lookback_delta: Duration::from_secs(1),
12454 };
12455
12456 let case = r#"some_metric{nonexistent="hi"}"#;
12457 let prom_expr = parser::parse(case).unwrap();
12458 eval_stmt.expr = prom_expr;
12459 let table_provider = build_test_table_provider(
12460 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12461 3,
12462 3,
12463 )
12464 .await;
12465 let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12467 .await
12468 .unwrap();
12469 }
12470
12471 #[tokio::test]
12472 async fn test_label_join() {
12473 let prom_expr = parser::parse(
12474 "label_join(up{tag_0='api-server'}, 'foo', ',', 'tag_1', 'tag_2', 'tag_3')",
12475 )
12476 .unwrap();
12477 let eval_stmt = EvalStmt {
12478 expr: prom_expr,
12479 start: UNIX_EPOCH,
12480 end: UNIX_EPOCH
12481 .checked_add(Duration::from_secs(100_000))
12482 .unwrap(),
12483 interval: Duration::from_secs(5),
12484 lookback_delta: Duration::from_secs(1),
12485 };
12486
12487 let table_provider =
12488 build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 4, 1)
12489 .await;
12490 let plan =
12491 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12492 .await
12493 .unwrap();
12494
12495 let expected = r#"
12496Filter: 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]
12497 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]
12498 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]
12499 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]
12500 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]
12501 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]
12502 TableScan: up [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#;
12503
12504 let ret = plan.display_indent_schema().to_string();
12505 assert_eq!(format!("\n{ret}"), expected, "\n{}", ret);
12506 }
12507
12508 #[tokio::test]
12509 async fn test_label_replace() {
12510 let prom_expr = parser::parse(
12511 "label_replace(up{tag_0=\"a:c\"}, \"foo\", \"$1\", \"tag_0\", \"(.*):.*\")",
12512 )
12513 .unwrap();
12514 let eval_stmt = EvalStmt {
12515 expr: prom_expr,
12516 start: UNIX_EPOCH,
12517 end: UNIX_EPOCH
12518 .checked_add(Duration::from_secs(100_000))
12519 .unwrap(),
12520 interval: Duration::from_secs(5),
12521 lookback_delta: Duration::from_secs(1),
12522 };
12523
12524 let table_provider =
12525 build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 1, 1)
12526 .await;
12527 let plan =
12528 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12529 .await
12530 .unwrap();
12531
12532 let expected = r#"
12533Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8]
12534 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]
12535 PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
12536 PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
12537 Sort: up.tag_0 ASC NULLS FIRST, up.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
12538 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]
12539 TableScan: up [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#;
12540
12541 let ret = plan.display_indent_schema().to_string();
12542 assert_eq!(format!("\n{ret}"), expected, "\n{}", ret);
12543 }
12544
12545 #[tokio::test]
12546 async fn label_replace_aggregation_queries_plan_successfully() {
12547 let aggregate =
12548 r#"sum by (foo) (label_replace(some_metric, "foo", "$1", "tag_0", "(.*)"))"#;
12549 let queries = [
12550 aggregate.to_string(),
12551 format!("{aggregate} <= 10"),
12552 format!("{aggregate} * 0.8"),
12553 format!("0.8 * {aggregate}"),
12554 format!("{aggregate} <= {aggregate} * 0.8"),
12555 ];
12556 let state = build_query_engine_state();
12557 let mut failures = Vec::new();
12558
12559 for query in queries {
12560 let table_provider = build_test_table_provider(
12561 &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12562 1,
12563 1,
12564 )
12565 .await;
12566 if let Err(error) =
12567 PromPlanner::stmt_to_plan(table_provider, &build_eval_stmt(&query), &state).await
12568 {
12569 failures.push(format!("{query}: {error:?}"));
12570 }
12571 }
12572
12573 assert!(failures.is_empty(), "{}", failures.join("\n"));
12574 }
12575
12576 #[tokio::test]
12577 async fn test_matchers_to_expr() {
12578 let mut eval_stmt = EvalStmt {
12579 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12580 start: UNIX_EPOCH,
12581 end: UNIX_EPOCH
12582 .checked_add(Duration::from_secs(100_000))
12583 .unwrap(),
12584 interval: Duration::from_secs(5),
12585 lookback_delta: Duration::from_secs(1),
12586 };
12587 let case =
12588 r#"sum(prometheus_tsdb_head_series{tag_1=~"(10.0.160.237:8080|10.0.160.237:9090)"})"#;
12589
12590 let prom_expr = parser::parse(case).unwrap();
12591 eval_stmt.expr = prom_expr;
12592 let table_provider = build_test_table_provider(
12593 &[(
12594 DEFAULT_SCHEMA_NAME.to_string(),
12595 "prometheus_tsdb_head_series".to_string(),
12596 )],
12597 3,
12598 3,
12599 )
12600 .await;
12601 let plan =
12602 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12603 .await
12604 .unwrap();
12605 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]\
12606 \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]\
12607 \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]\
12608 \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]\
12609 \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]\
12610 \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]\
12611 \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]";
12612 assert_eq!(plan.display_indent_schema().to_string(), expected);
12613 }
12614
12615 #[tokio::test]
12616 async fn test_topk_expr() {
12617 let mut eval_stmt = EvalStmt {
12618 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12619 start: UNIX_EPOCH,
12620 end: UNIX_EPOCH
12621 .checked_add(Duration::from_secs(100_000))
12622 .unwrap(),
12623 interval: Duration::from_secs(5),
12624 lookback_delta: Duration::from_secs(1),
12625 };
12626 let case = r#"topk(10, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#;
12627
12628 let prom_expr = parser::parse(case).unwrap();
12629 eval_stmt.expr = prom_expr;
12630 let table_provider = build_test_table_provider_with_fields(
12631 &[
12632 (
12633 DEFAULT_SCHEMA_NAME.to_string(),
12634 "prometheus_tsdb_head_series".to_string(),
12635 ),
12636 (
12637 DEFAULT_SCHEMA_NAME.to_string(),
12638 "http_server_requests_seconds_count".to_string(),
12639 ),
12640 ],
12641 &["ip"],
12642 )
12643 .await;
12644
12645 let plan =
12646 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12647 .await
12648 .unwrap();
12649 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)]\
12650 \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]\
12651 \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]\
12652 \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]\
12653 \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]\
12654 \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]\
12655 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
12656 \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
12657 \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]\
12658 \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]\
12659 \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]";
12660
12661 assert_eq!(plan.display_indent_schema().to_string(), expected);
12662 }
12663
12664 #[tokio::test]
12665 async fn test_count_values_expr() {
12666 let mut eval_stmt = EvalStmt {
12667 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12668 start: UNIX_EPOCH,
12669 end: UNIX_EPOCH
12670 .checked_add(Duration::from_secs(100_000))
12671 .unwrap(),
12672 interval: Duration::from_secs(5),
12673 lookback_delta: Duration::from_secs(1),
12674 };
12675 let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#;
12676
12677 let prom_expr = parser::parse(case).unwrap();
12678 eval_stmt.expr = prom_expr;
12679 let table_provider = build_test_table_provider_with_fields(
12680 &[
12681 (
12682 DEFAULT_SCHEMA_NAME.to_string(),
12683 "prometheus_tsdb_head_series".to_string(),
12684 ),
12685 (
12686 DEFAULT_SCHEMA_NAME.to_string(),
12687 "http_server_requests_seconds_count".to_string(),
12688 ),
12689 ],
12690 &["ip"],
12691 )
12692 .await;
12693
12694 let plan =
12695 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12696 .await
12697 .unwrap();
12698 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]\
12699 \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]\
12700 \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]\
12701 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
12702 \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
12703 \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]\
12704 \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]\
12705 \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]";
12706
12707 assert_eq!(plan.display_indent_schema().to_string(), expected);
12708 }
12709
12710 #[tokio::test]
12711 async fn test_value_alias() {
12712 let mut eval_stmt = EvalStmt {
12713 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12714 start: UNIX_EPOCH,
12715 end: UNIX_EPOCH
12716 .checked_add(Duration::from_secs(100_000))
12717 .unwrap(),
12718 interval: Duration::from_secs(5),
12719 lookback_delta: Duration::from_secs(1),
12720 };
12721 let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#;
12722
12723 let prom_expr = parser::parse(case).unwrap();
12724 eval_stmt.expr = prom_expr;
12725 eval_stmt = QueryLanguageParser::apply_alias_extension(eval_stmt, "my_series");
12726 let table_provider = build_test_table_provider_with_fields(
12727 &[
12728 (
12729 DEFAULT_SCHEMA_NAME.to_string(),
12730 "prometheus_tsdb_head_series".to_string(),
12731 ),
12732 (
12733 DEFAULT_SCHEMA_NAME.to_string(),
12734 "http_server_requests_seconds_count".to_string(),
12735 ),
12736 ],
12737 &["ip"],
12738 )
12739 .await;
12740
12741 let plan =
12742 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12743 .await
12744 .unwrap();
12745 let expected = r#"
12746Projection: 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)]
12747 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]
12748 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]
12749 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]
12750 PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
12751 PromSeriesDivide: tags=["ip"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
12752 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]
12753 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]
12754 TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"#;
12755 assert_eq!(format!("\n{}", plan.display_indent_schema()), expected);
12756 }
12757
12758 #[tokio::test]
12759 async fn test_quantile_expr() {
12760 let mut eval_stmt = EvalStmt {
12761 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12762 start: UNIX_EPOCH,
12763 end: UNIX_EPOCH
12764 .checked_add(Duration::from_secs(100_000))
12765 .unwrap(),
12766 interval: Duration::from_secs(5),
12767 lookback_delta: Duration::from_secs(1),
12768 };
12769 let case = r#"quantile(0.3, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#;
12770
12771 let prom_expr = parser::parse(case).unwrap();
12772 eval_stmt.expr = prom_expr;
12773 let table_provider = build_test_table_provider_with_fields(
12774 &[
12775 (
12776 DEFAULT_SCHEMA_NAME.to_string(),
12777 "prometheus_tsdb_head_series".to_string(),
12778 ),
12779 (
12780 DEFAULT_SCHEMA_NAME.to_string(),
12781 "http_server_requests_seconds_count".to_string(),
12782 ),
12783 ],
12784 &["ip"],
12785 )
12786 .await;
12787
12788 let plan =
12789 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12790 .await
12791 .unwrap();
12792 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]\
12793 \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]\
12794 \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]\
12795 \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]\
12796 \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
12797 \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
12798 \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]\
12799 \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]\
12800 \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]";
12801
12802 assert_eq!(plan.display_indent_schema().to_string(), expected);
12803 }
12804
12805 #[tokio::test]
12806 async fn test_or_not_exists_table_label() {
12807 let state = build_query_engine_state();
12808 let provider = build_test_table_provider_with_fields(
12809 &[(DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string())],
12810 &["job"],
12811 )
12812 .await;
12813 let raw = PromPlanner::stmt_to_plan(
12814 provider,
12815 &build_eval_stmt(r#"missing_metric or on(absent_label) normal_metric"#),
12816 &state,
12817 )
12818 .await
12819 .unwrap();
12820 assert!(
12821 raw.display_indent_schema()
12822 .to_string()
12823 .contains("__promql_or_match_0@")
12824 );
12825 let (optimized, batches) = execute(raw, &state).await;
12826 assert_no_internal_or_keys(optimized.schema());
12827 assert!(batches.iter().all(|batch| {
12828 batch
12829 .schema()
12830 .fields()
12831 .iter()
12832 .all(|field| !field.name().starts_with("__promql_or_match_"))
12833 }));
12834 }
12835
12836 #[tokio::test]
12837 async fn test_histogram_quantile_missing_le_column() {
12838 let mut eval_stmt = EvalStmt {
12839 expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12840 start: UNIX_EPOCH,
12841 end: UNIX_EPOCH
12842 .checked_add(Duration::from_secs(100_000))
12843 .unwrap(),
12844 interval: Duration::from_secs(5),
12845 lookback_delta: Duration::from_secs(1),
12846 };
12847
12848 let case = r#"histogram_quantile(0.99, sum by(pod,instance,le) (rate(non_existent_histogram_bucket{instance=~"xxx"}[1m])))"#;
12850
12851 let prom_expr = parser::parse(case).unwrap();
12852 eval_stmt.expr = prom_expr;
12853
12854 let table_provider = build_test_table_provider_with_fields(
12856 &[(
12857 DEFAULT_SCHEMA_NAME.to_string(),
12858 "non_existent_histogram_bucket".to_string(),
12859 )],
12860 &["pod", "instance"], )
12862 .await;
12863
12864 let result =
12866 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12867 .await;
12868
12869 assert!(
12871 result.is_ok(),
12872 "Expected successful plan creation with empty result, but got error: {:?}",
12873 result.err()
12874 );
12875
12876 let plan = result.unwrap();
12878 match plan {
12879 LogicalPlan::EmptyRelation(_) => {
12880 }
12882 _ => panic!("Expected EmptyRelation, but got: {:?}", plan),
12883 }
12884 }
12885
12886 #[tokio::test]
12887 async fn test_direct_or_normalizes_missing_match_labels() {
12888 type Case<'a> = (
12889 Option<Option<&'a str>>,
12890 Option<Option<&'a str>>,
12891 i64,
12892 i64,
12893 &'a [(f64, Option<&'a str>)],
12894 );
12895
12896 let modifier = or_modifier("lhs or on(k) rhs");
12897 #[rustfmt::skip]
12898 let cases: &[Case<'_>] = &[
12899 (None, None, 1, 1, &[(1.0, None)]),
12900 (None, Some(Some("")), 1, 1, &[(1.0, None)]),
12901 (Some(Some("")), None, 1, 1, &[(1.0, Some(""))]),
12902 (None, Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]),
12903 (Some(Some("l")), None, 1, 1, &[(1.0, Some("l")), (2.0, None)]),
12904 (Some(None), Some(Some("")), 1, 1, &[(1.0, None)]),
12905 (Some(None), Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]),
12906 (Some(Some("same")), Some(Some("same")), 1, 2, &[(1.0, Some("same")), (2.0, Some("same"))]),
12907 ];
12908 for &(left, right, left_ts, right_ts, expected) in cases {
12909 let (optimized, batches) = run(
12910 &matrix_source("lhs", left, left_ts, 1.0),
12911 &matrix_source("rhs", right, right_ts, 2.0),
12912 matrix_context("lhs", left),
12913 matrix_context("rhs", right),
12914 &modifier,
12915 )
12916 .await;
12917 assert_no_internal_or_keys(optimized.schema());
12918 assert_eq!(
12919 rows(&batches),
12920 expected
12921 .iter()
12922 .map(|(value, label)| (*value, label.map(str::to_string)))
12923 .collect::<Vec<_>>()
12924 );
12925 }
12926 }
12927
12928 #[tokio::test]
12929 async fn test_direct_or_match_modifiers() {
12930 for (modifier, left, right, expected) in [
12931 (None, "left", "right", 2),
12932 (or_modifier("lhs or on(k) rhs"), "same", "same", 1),
12933 (or_modifier("lhs or on() rhs"), "left", "right", 1),
12934 (or_modifier("lhs or ignoring(k) rhs"), "left", "right", 1),
12935 ] {
12936 let (_, batches) = run(
12937 &matrix_source("lhs", Some(Some(left)), 1, 1.0),
12938 &matrix_source("rhs", Some(Some(right)), 1, 2.0),
12939 direct_or_context("lhs", &["job", "k"], "v"),
12940 direct_or_context("rhs", &["job", "k"], "v"),
12941 &modifier,
12942 )
12943 .await;
12944 assert_eq!(
12945 batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
12946 expected
12947 );
12948 }
12949 }
12950
12951 #[tokio::test]
12952 async fn test_direct_or_nested_projection_uses_left_context() {
12953 let left = matrix_source("lhs", Some(Some("k")), 1, 1.0);
12954 let right = matrix_source("rhs", Some(Some("k")), 1, 2.0);
12955 let raw = plan_direct_or(
12956 scan(&left),
12957 scan(&right),
12958 direct_or_context("lhs", &["job", "k"], "v"),
12959 direct_or_context("rhs", &["job", "k"], "v"),
12960 &or_modifier("lhs or on(k) rhs"),
12961 )
12962 .await;
12963 assert!(raw.schema().iter().any(|(qualifier, field)| {
12964 qualifier.as_ref().is_some_and(|q| q.to_string() == "lhs") && field.name() == "v"
12965 }));
12966 let nested = LogicalPlanBuilder::from(raw)
12967 .project(vec![
12968 DfExpr::BinaryExpr(BinaryExpr {
12969 left: Box::new(DfExpr::Column(Column::new(
12970 Some(TableReference::bare("lhs")),
12971 "v",
12972 ))),
12973 op: Operator::Plus,
12974 right: Box::new(lit(1.0)),
12975 })
12976 .alias("v_plus"),
12977 ])
12978 .unwrap()
12979 .build()
12980 .unwrap();
12981 let (_, batches) = execute(nested, &build_query_engine_state()).await;
12982 assert_eq!(values(&batches, "v_plus"), vec![2.0]);
12983 }
12984
12985 #[tokio::test]
12986 async fn test_direct_or_skips_user_internal_key_name() {
12987 const USER_TAG: &str = "__promql_or_match_0";
12988 let left = tagged_source(
12989 "lhs",
12990 false,
12991 (USER_TAG, Some("left")),
12992 DirectOrValue::Float64(1.0),
12993 );
12994 let right = tagged_source(
12995 "rhs",
12996 false,
12997 (USER_TAG, Some("right")),
12998 DirectOrValue::Float64(2.0),
12999 );
13000 let raw = plan_direct_or(
13001 scan(&left),
13002 scan(&right),
13003 direct_or_context("lhs", &["job", USER_TAG], "v"),
13004 direct_or_context("rhs", &["job", USER_TAG], "v"),
13005 &or_modifier("lhs or on(missing_label) rhs"),
13006 )
13007 .await;
13008 assert!(
13009 raw.display_indent_schema()
13010 .to_string()
13011 .contains("__promql_or_match_1@")
13012 );
13013 let (_, batches) = execute(raw, &build_query_engine_state()).await;
13014 assert!(
13015 batches
13016 .iter()
13017 .all(|batch| batch.column_by_name(USER_TAG).is_some())
13018 );
13019 }
13020
13021 #[tokio::test]
13022 async fn test_direct_or_substrait_round_trip_with_normalized_key() {
13023 let state = build_query_engine_state();
13024 let ctx = SessionContext::new_with_state(state.session_state());
13025 let catalog = Arc::new(MemoryCatalogProvider::new());
13026 catalog
13027 .register_schema("public", Arc::new(MemorySchemaProvider::new()))
13028 .unwrap();
13029 ctx.register_catalog("datafusion", catalog);
13030 let left = matrix_source("lhs", Some(Some("")), 1, 1.0);
13031 let right = matrix_source("rhs", None, 1, 2.0);
13032 ctx.register_table(
13033 TableReference::full("datafusion", "public", "lhs"),
13034 table(&left),
13035 )
13036 .unwrap();
13037 ctx.register_table(
13038 TableReference::full("datafusion", "public", "rhs"),
13039 table(&right),
13040 )
13041 .unwrap();
13042 let raw = plan_direct_or(
13043 ctx.table("datafusion.public.lhs")
13044 .await
13045 .unwrap()
13046 .into_unoptimized_plan(),
13047 ctx.table("datafusion.public.rhs")
13048 .await
13049 .unwrap()
13050 .into_unoptimized_plan(),
13051 direct_or_context("lhs", &["job", "k"], "v"),
13052 direct_or_context("rhs", &["job"], "v"),
13053 &or_modifier("lhs or on(k) rhs"),
13054 )
13055 .await;
13056 let decoded = DFLogicalSubstraitConvertor
13057 .decode(
13058 DFLogicalSubstraitConvertor
13059 .encode(&raw, DefaultSerializer)
13060 .unwrap(),
13061 ctx.state(),
13062 )
13063 .await
13064 .unwrap();
13065 let (optimized, batches) = execute(decoded, &state).await;
13066 assert_no_internal_or_keys(optimized.schema());
13067 assert!(batches.iter().all(|batch| {
13068 batch
13069 .schema()
13070 .fields()
13071 .iter()
13072 .all(|field| !field.name().starts_with("__promql_or_match_"))
13073 }));
13074 assert_eq!(values(&batches, "v"), vec![1.0]);
13075 }
13076
13077 #[tokio::test]
13078 async fn test_direct_or_numeric_value_types() {
13079 let left = tagged_source("lhs", true, ("k", Some("lhs")), DirectOrValue::Int64(0));
13080 let right = tagged_source(
13081 "rhs",
13082 false,
13083 ("k", Some("rhs")),
13084 DirectOrValue::Float64(0.5),
13085 );
13086 let (optimized, batches) = run(
13087 &left,
13088 &right,
13089 direct_or_context("lhs", &["job", "k"], "v"),
13090 direct_or_context("rhs", &["job", "k"], "v"),
13091 &or_modifier("lhs or on(k) rhs"),
13092 )
13093 .await;
13094 assert_eq!(
13095 optimized
13096 .schema()
13097 .field_with_name(None, "v")
13098 .unwrap()
13099 .data_type(),
13100 &ArrowDataType::Float64
13101 );
13102 assert_eq!(values(&batches, "v"), vec![0.5]);
13103 let provider = build_test_table_provider_with_fields(
13104 &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
13105 &[],
13106 )
13107 .await;
13108 let mut planner = PromPlanner {
13109 table_provider: provider,
13110 ctx: PromPlannerContext::default(),
13111 promql_annotations: None,
13112 };
13113 let left_context = direct_or_context("lhs", &["job"], "v");
13114 let right_context = direct_or_context("rhs", &["job"], "v");
13115 let error = planner
13116 .or_operator(
13117 scan(&job_source("lhs", DirectOrValue::Utf8("x"))),
13118 scan(&job_source("rhs", DirectOrValue::Float64(1.0))),
13119 left_context.tag_columns.iter().cloned().collect(),
13120 right_context.tag_columns.iter().cloned().collect(),
13121 left_context,
13122 right_context,
13123 &or_modifier("lhs or on() rhs"),
13124 )
13125 .unwrap_err();
13126 assert!(
13127 error
13128 .to_string()
13129 .contains("OR value fields have incompatible types")
13130 );
13131 }
13132
13133 #[tokio::test]
13134 async fn test_or_with_histogram_quantile_missing_le_column() {
13135 let case = r#"histogram_quantile(0.99, non_existent_histogram_bucket) or normal_metric"#;
13136 let eval_stmt = build_eval_stmt(case);
13137 let table_provider = build_missing_le_or_normal_metric_table_provider().await;
13138
13139 let plan =
13140 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13141 .await
13142 .unwrap();
13143 assert_normal_metric_schema(&plan);
13144 }
13145
13146 #[tokio::test]
13147 async fn test_or_with_right_empty_histogram_restores_left_context() {
13148 let eval_stmt = build_eval_stmt(
13149 r#"abs(sum by(instance) (normal_metric) or histogram_quantile(0.99, sum by(pod) (non_existent_histogram_bucket)))"#,
13150 );
13151 let table_provider = build_missing_le_or_normal_metric_table_provider().await;
13152
13153 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13154 .await
13155 .unwrap();
13156 }
13157
13158 #[tokio::test]
13159 async fn test_or_with_both_empty_histograms() {
13160 let eval_stmt = build_eval_stmt(
13161 r#"histogram_quantile(0.99, sum by(pod) (left_histogram_bucket)) or histogram_quantile(0.99, sum by(instance) (right_histogram_bucket))"#,
13162 );
13163 let table_provider = build_test_table_provider_with_fields(
13164 &[
13165 (
13166 DEFAULT_SCHEMA_NAME.to_string(),
13167 "left_histogram_bucket".to_string(),
13168 ),
13169 (
13170 DEFAULT_SCHEMA_NAME.to_string(),
13171 "right_histogram_bucket".to_string(),
13172 ),
13173 ],
13174 &["pod", "instance"],
13175 )
13176 .await;
13177
13178 let plan =
13179 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13180 .await
13181 .unwrap();
13182 match plan {
13183 LogicalPlan::EmptyRelation(relation) => {
13184 assert!(!relation.produce_one_row);
13185 assert!(!relation.schema.fields().is_empty());
13186 assert!(
13187 relation
13188 .schema
13189 .fields()
13190 .iter()
13191 .any(|field| field.data_type() == &ArrowDataType::Float64)
13192 );
13193 assert!(
13194 relation
13195 .schema
13196 .fields()
13197 .iter()
13198 .any(|field| field.name() == "pod")
13199 );
13200 assert!(
13201 !relation
13202 .schema
13203 .fields()
13204 .iter()
13205 .any(|field| field.name() == "instance")
13206 );
13207 }
13208 _ => panic!("Expected EmptyRelation, but got: {plan:?}"),
13209 }
13210 }
13211
13212 #[tokio::test]
13213 async fn test_nested_or_with_both_empty_histograms() {
13214 for case in [
13215 r#"abs(histogram_quantile(0.99, left_histogram_bucket) or histogram_quantile(0.99, right_histogram_bucket))"#,
13216 r#"(histogram_quantile(0.99, left_histogram_bucket) or histogram_quantile(0.99, right_histogram_bucket)) + 1"#,
13217 ] {
13218 let eval_stmt = build_eval_stmt(case);
13219 let table_provider = build_test_table_provider_with_fields(
13220 &[
13221 (
13222 DEFAULT_SCHEMA_NAME.to_string(),
13223 "left_histogram_bucket".to_string(),
13224 ),
13225 (
13226 DEFAULT_SCHEMA_NAME.to_string(),
13227 "right_histogram_bucket".to_string(),
13228 ),
13229 ],
13230 &["pod", "instance"],
13231 )
13232 .await;
13233
13234 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13235 .await
13236 .unwrap();
13237 }
13238 }
13239
13240 #[tokio::test]
13241 async fn test_or_with_empty_histogram_modifiers() {
13242 for case in [
13243 r#"histogram_quantile(0.99, non_existent_histogram_bucket) or on(pod) normal_metric"#,
13244 r#"normal_metric or ignoring(instance) histogram_quantile(0.99, non_existent_histogram_bucket)"#,
13245 ] {
13246 let eval_stmt = build_eval_stmt(case);
13247 let table_provider = build_missing_le_or_normal_metric_table_provider().await;
13248
13249 let plan =
13250 PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13251 .await
13252 .unwrap();
13253 assert_normal_metric_schema(&plan);
13254 }
13255 }
13256
13257 #[tokio::test]
13258 async fn test_unless_preserves_left_context_for_histogram() {
13259 let eval_stmt = build_eval_stmt(
13260 r#"histogram_quantile(0.99, bucket_metric unless on(job) normal_metric) or fallback_metric"#,
13261 );
13262 let state = build_query_engine_state();
13263 let plan = PromPlanner::stmt_to_plan(
13264 build_set_op_context_table_provider().await,
13265 &eval_stmt,
13266 &state,
13267 )
13268 .await
13269 .unwrap();
13270 assert!(contains_histogram_fold(&plan), "{plan:?}");
13271 let (optimized, physical) = optimize_and_create_physical_plan(&state, plan).await;
13272 assert!(contains_histogram_fold(&optimized), "{optimized:?}");
13273 let batches =
13274 datafusion::physical_plan::collect(physical, state.session_state().task_ctx())
13275 .await
13276 .unwrap();
13277 assert!(batches.iter().all(|batch| batch.num_rows() == 0));
13278 }
13279
13280 #[tokio::test]
13281 async fn test_and_preserves_left_context_for_histogram() {
13282 let eval_stmt = build_eval_stmt(
13283 r#"histogram_quantile(0.99, bucket_metric and on(job) normal_metric) or fallback_metric"#,
13284 );
13285 let plan = PromPlanner::stmt_to_plan(
13286 build_set_op_context_table_provider().await,
13287 &eval_stmt,
13288 &build_query_engine_state(),
13289 )
13290 .await
13291 .unwrap();
13292 assert!(contains_histogram_fold(&plan), "{plan:?}");
13293 }
13294
13295 #[tokio::test]
13296 async fn test_and_preserves_left_context_when_le_is_missing() {
13297 let eval_stmt =
13298 build_eval_stmt(r#"histogram_quantile(0.99, normal_metric and on(job) bucket_metric)"#);
13299 let plan = PromPlanner::stmt_to_plan(
13300 build_set_op_context_table_provider().await,
13301 &eval_stmt,
13302 &build_query_engine_state(),
13303 )
13304 .await
13305 .unwrap();
13306 assert!(matches!(&plan, LogicalPlan::EmptyRelation(_)), "{plan:?}");
13307 assert!(!plan.schema().fields().is_empty());
13308 assert!(!contains_histogram_fold(&plan), "{plan:?}");
13309 }
13310
13311 #[tokio::test]
13312 async fn test_or_context_uses_left_qualified_output() {
13313 let case = r#"(normal_metric or other_metric) + 1"#;
13314 let eval_stmt = build_eval_stmt(case);
13315 let state = build_query_engine_state();
13316 let plan =
13317 PromPlanner::stmt_to_plan(build_or_context_table_provider().await, &eval_stmt, &state)
13318 .await
13319 .unwrap();
13320 assert!(
13321 plan.schema()
13322 .fields()
13323 .iter()
13324 .any(|field| field.data_type() == &ArrowDataType::Float64),
13325 "{plan:?}"
13326 );
13327 let (_optimized, _physical) = optimize_and_create_physical_plan(&state, plan).await;
13328 }
13329
13330 #[tokio::test]
13331 async fn test_or_context_uses_left_qualified_empty_histogram_output() {
13332 let case = r#"(abs(histogram_quantile(0.99, non_hist_metric)) or normal_metric) + 1"#;
13333 let eval_stmt = build_eval_stmt(case);
13334 let plan = PromPlanner::stmt_to_plan(
13335 build_or_context_table_provider().await,
13336 &eval_stmt,
13337 &build_query_engine_state(),
13338 )
13339 .await
13340 .unwrap();
13341 assert!(
13342 plan.schema()
13343 .fields()
13344 .iter()
13345 .any(|field| field.data_type() == &ArrowDataType::Float64),
13346 "{plan:?}"
13347 );
13348 }
13349
13350 #[tokio::test]
13351 async fn test_direct_or_preserves_float_and_native_histogram_samples() {
13352 for histogram_on_left in [false, true] {
13353 let (planner, plan) = mixed_direct_or(histogram_on_left).await;
13354
13355 let float_field = &planner.ctx.field_columns[0];
13356 let histogram_field = &planner.ctx.field_columns[1];
13357 assert!(float_field.starts_with(OR_FLOAT_FIELD_PREFIX));
13358 assert!(histogram_field.starts_with(OR_HISTOGRAM_FIELD_PREFIX));
13359 assert_eq!(
13360 plan.schema()
13361 .field_with_name(None, float_field)
13362 .unwrap()
13363 .data_type(),
13364 &ArrowDataType::Float64
13365 );
13366 assert_eq!(
13367 plan.schema()
13368 .field_with_name(None, histogram_field)
13369 .unwrap()
13370 .data_type(),
13371 &native_histogram_value_type().as_arrow_type()
13372 );
13373
13374 let (optimized, batches) = execute(plan, &build_query_engine_state()).await;
13375 assert_no_internal_or_keys(optimized.schema());
13376 let mut sample_kinds = batches
13377 .iter()
13378 .flat_map(|batch| {
13379 let values = batch.column_by_name(float_field).unwrap();
13380 let histograms = batch.column_by_name(histogram_field).unwrap();
13381 (0..batch.num_rows())
13382 .map(|row| (values.is_valid(row), histograms.is_valid(row)))
13383 })
13384 .collect::<Vec<_>>();
13385 sample_kinds.sort_unstable();
13386 assert_eq!(sample_kinds, vec![(false, true), (true, false)]);
13387 }
13388 }
13389
13390 #[tokio::test]
13391 async fn malformed_classic_bucket_does_not_drop_native_histogram() {
13392 let state = build_query_engine_state();
13393 let collector = PromqlAnnotationCollector::default();
13394 let plan = PromPlanner::stmt_to_plan_with_annotations(
13395 operator_table_provider(),
13396 &operator_eval_stmt("histogram_quantile(0.5, bad_classic or bad_native)"),
13397 &state,
13398 Some(collector.clone()),
13399 )
13400 .await
13401 .unwrap();
13402 let value_field = plan
13403 .schema()
13404 .fields()
13405 .iter()
13406 .find(|field| field.data_type() == &ArrowDataType::Float64)
13407 .unwrap()
13408 .name()
13409 .clone();
13410
13411 let (_, batches) = execute(plan, &state).await;
13412 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
13413 assert_eq!(values(&batches, &value_field), vec![0.0]);
13414 let mut warnings = vec![];
13415 let mut infos = vec![];
13416 collector.append_to(&mut warnings, &mut infos);
13417 assert!(warnings.is_empty());
13418 assert!(infos.is_empty());
13419 }
13420
13421 #[tokio::test]
13422 async fn test_mixed_binary_operator_aligns_both_alternative_inputs() {
13423 let state = build_query_engine_state();
13424 let plan = PromPlanner::stmt_to_plan(
13425 operator_table_provider(),
13426 &operator_eval_stmt("(lf or on(tag) lh) * on(tag) (rf or on(tag) rh)"),
13427 &state,
13428 )
13429 .await
13430 .unwrap();
13431 let plan_text = plan.display_indent_schema().to_string();
13432 assert!(
13433 plan_text.contains("prom_native_histogram_mul_scalar"),
13434 "{plan_text}"
13435 );
13436 assert!(
13437 plan_text.contains("prom_native_histogram_scalar_mul"),
13438 "{plan_text}"
13439 );
13440 let float_field = plan
13441 .schema()
13442 .fields()
13443 .iter()
13444 .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX))
13445 .unwrap()
13446 .name()
13447 .clone();
13448 let histogram_field = plan
13449 .schema()
13450 .fields()
13451 .iter()
13452 .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX))
13453 .unwrap()
13454 .name()
13455 .clone();
13456
13457 let (_, batches) = execute(plan, &state).await;
13458 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
13459 assert!(values(&batches, &float_field).is_empty());
13460 let mut sums = histograms(&batches, &histogram_field)
13461 .into_iter()
13462 .map(|histogram| histogram.sum)
13463 .collect::<Vec<_>>();
13464 sums.sort_by(f64::total_cmp);
13465 assert_eq!(sums, vec![2.0, 3.0]);
13466 }
13467
13468 #[tokio::test]
13469 async fn test_mixed_binary_operator_reports_only_dropped_samples() {
13470 for (query, expected_rows, expected_infos) in [
13471 ("(lf or on(tag) lh) + on(tag) (rf or on(tag) rh)", 0, 1),
13472 ("(lf or on(tag) lh) + on(tag) (lf or on(tag) lh)", 2, 0),
13473 ("(lf or on(tag) lh) % on(tag) lh", 0, 1),
13474 ] {
13475 let state = build_query_engine_state();
13476 let annotations = PromqlAnnotationCollector::default();
13477 let plan = PromPlanner::stmt_to_plan_with_annotations(
13478 operator_table_provider(),
13479 &operator_eval_stmt(query),
13480 &state,
13481 Some(annotations.clone()),
13482 )
13483 .await
13484 .unwrap();
13485
13486 let (_, batches) = execute(plan, &state).await;
13487 assert_eq!(
13488 batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
13489 expected_rows,
13490 "{query}"
13491 );
13492 let mut warnings = vec![];
13493 let mut infos = vec![];
13494 annotations.append_to(&mut warnings, &mut infos);
13495 assert!(warnings.is_empty(), "{query}: {warnings:?}");
13496 assert_eq!(infos.len(), expected_infos, "{query}: {infos:?}");
13497 }
13498 }
13499
13500 #[tokio::test]
13501 async fn test_histogram_only_min_drops_empty_aggregate_group() {
13502 let state = build_query_engine_state();
13506 for query in ["min(lh)", "group(min(lh))"] {
13507 let plan = PromPlanner::stmt_to_plan(
13508 operator_table_provider(),
13509 &operator_eval_stmt(query),
13510 &state,
13511 )
13512 .await
13513 .unwrap();
13514 let (_, batches) = execute(plan, &state).await;
13515 assert_eq!(
13516 batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
13517 0,
13518 "{query}"
13519 );
13520 }
13521 }
13522
13523 #[tokio::test]
13524 async fn test_mixed_min_drops_histogram_only_group() {
13525 let state = build_query_engine_state();
13529 let plan = PromPlanner::stmt_to_plan(
13530 operator_table_provider(),
13531 &operator_eval_stmt("min by (tag) (lf or on(tag) lh)"),
13532 &state,
13533 )
13534 .await
13535 .unwrap();
13536 let float_field = plan
13537 .schema()
13538 .fields()
13539 .iter()
13540 .find(|field| field.data_type() == &ArrowDataType::Float64)
13541 .unwrap()
13542 .name()
13543 .clone();
13544 let (_, batches) = execute(plan, &state).await;
13545 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
13546 assert_eq!(values(&batches, &float_field), vec![2.0]);
13547 }
13548
13549 #[tokio::test]
13550 async fn test_mixed_or_can_feed_another_or() {
13551 let state = build_query_engine_state();
13552 let plan = PromPlanner::stmt_to_plan(
13553 operator_table_provider(),
13554 &operator_eval_stmt("lf or on(tag) lh or on(tag) fallback"),
13555 &state,
13556 )
13557 .await
13558 .unwrap();
13559 let float_field = plan
13560 .schema()
13561 .fields()
13562 .iter()
13563 .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX))
13564 .unwrap()
13565 .name()
13566 .clone();
13567 let histogram_field = plan
13568 .schema()
13569 .fields()
13570 .iter()
13571 .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX))
13572 .unwrap()
13573 .name()
13574 .clone();
13575
13576 let (_, batches) = execute(plan, &state).await;
13577 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 3);
13578 let mut float_values = values(&batches, &float_field);
13579 float_values.sort_by(f64::total_cmp);
13580 assert_eq!(float_values, vec![2.0, 7.0]);
13581 assert_eq!(histograms(&batches, &histogram_field).len(), 1);
13582 }
13583
13584 #[tokio::test]
13585 async fn test_mixed_fields_align_with_single_float_vector() {
13586 let (planner, mixed) = mixed_direct_or(false).await;
13587 let scale = tagged_source(
13588 "scale",
13589 false,
13590 ("k", Some("float")),
13591 DirectOrValue::Float64(2.0),
13592 );
13593 let scale = scan(&scale);
13594 let scale_fields = vec!["v".to_string()];
13595 let PromExpr::Binary(binary) = parser::parse("lhs * rhs").unwrap() else {
13596 unreachable!()
13597 };
13598
13599 let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns(
13600 mixed.schema(),
13601 scale.schema(),
13602 &planner.ctx.field_columns,
13603 &scale_fields,
13604 binary.op,
13605 false,
13606 false,
13607 );
13608 assert!(invalid_pairs.is_empty());
13609 assert_eq!(
13610 groups
13611 .iter()
13612 .map(|(output, _)| output.clone())
13613 .collect::<Vec<_>>(),
13614 planner.ctx.field_columns
13615 );
13616 assert_eq!(groups.len(), 2);
13617 assert!(
13618 groups
13619 .iter()
13620 .flat_map(|(_, pairs)| pairs)
13621 .all(|(_, right)| *right == &scale_fields[0])
13622 );
13623
13624 let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns(
13625 scale.schema(),
13626 mixed.schema(),
13627 &scale_fields,
13628 &planner.ctx.field_columns,
13629 binary.op,
13630 false,
13631 false,
13632 );
13633 assert!(invalid_pairs.is_empty());
13634 assert_eq!(
13635 groups
13636 .iter()
13637 .map(|(output, _)| output.clone())
13638 .collect::<Vec<_>>(),
13639 planner.ctx.field_columns
13640 );
13641 assert_eq!(groups.len(), 2);
13642 assert!(
13643 groups
13644 .iter()
13645 .flat_map(|(_, pairs)| pairs)
13646 .all(|(left, _)| *left == &scale_fields[0])
13647 );
13648 }
13649
13650 #[tokio::test]
13651 async fn test_non_bool_comparison_filters_mixed_sample_lanes() {
13652 let (planner, input) = mixed_direct_or(false).await;
13653 let input_schema = input.schema().clone();
13654 let plan = planner
13655 .filter_on_field_column(input, |field| {
13656 if PromPlanner::field_column_is_native_histogram(&input_schema, field) {
13657 Ok(lit(false))
13658 } else {
13659 Ok(col(field).gt(lit(0.0)))
13660 }
13661 })
13662 .unwrap();
13663 let float_field = planner.ctx.field_columns[0].clone();
13664
13665 let (_, batches) = execute(plan, &build_query_engine_state()).await;
13666 assert_eq!(values(&batches, &float_field), vec![1.25]);
13667 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
13668 }
13669
13670 #[tokio::test]
13671 async fn test_mixed_left_and_unless_preserve_sample_lanes() {
13672 for (expression, expected_sample_kind) in [
13673 ("lhs and on(k) mask", (false, true)),
13674 ("lhs unless on(k) mask", (true, false)),
13675 ] {
13676 let (mut planner, left) = mixed_direct_or(false).await;
13677 let left_context = planner.ctx.clone();
13678 let float_field = left_context.field_columns[0].clone();
13679 let histogram_field = left_context.field_columns[1].clone();
13680 let mask = tagged_source(
13681 "mask",
13682 false,
13683 ("k", Some("histogram")),
13684 DirectOrValue::Float64(1.0),
13685 );
13686 let PromExpr::Binary(binary) = parser::parse(expression).unwrap() else {
13687 unreachable!()
13688 };
13689 let plan = planner
13690 .set_op_on_non_field_columns(
13691 left,
13692 scan(&mask),
13693 left_context,
13694 direct_or_context("mask", &["job", "k"], "v"),
13695 binary.op,
13696 &binary.modifier,
13697 )
13698 .unwrap();
13699
13700 let (_, batches) = execute(plan, &build_query_engine_state()).await;
13701 let sample_kinds = batches
13702 .iter()
13703 .flat_map(|batch| {
13704 let floats = batch.column_by_name(&float_field).unwrap();
13705 let histograms = batch.column_by_name(&histogram_field).unwrap();
13706 (0..batch.num_rows())
13707 .map(|row| (floats.is_valid(row), histograms.is_valid(row)))
13708 })
13709 .collect::<Vec<_>>();
13710 assert_eq!(sample_kinds, vec![expected_sample_kind], "{expression}");
13711 }
13712 }
13713
13714 #[tokio::test]
13715 async fn test_mixed_fields_arithmetic_broadcasts_computed_scalar() {
13716 let plan = PromPlanner::stmt_to_plan(
13717 build_test_mixed_native_histogram_table_provider("some_metric").await,
13718 &build_eval_stmt("some_metric * scalar(vector(2))"),
13719 &build_query_engine_state(),
13720 )
13721 .await
13722 .unwrap();
13723 let schema = plan.schema();
13724 assert_eq!(
13725 schema
13726 .field_with_unqualified_name(greptime_value())
13727 .unwrap()
13728 .data_type(),
13729 &ArrowDataType::Float64
13730 );
13731 assert_eq!(
13732 schema
13733 .field_with_unqualified_name(greptime_native_histogram())
13734 .unwrap()
13735 .data_type(),
13736 &native_histogram_value_type().as_arrow_type()
13737 );
13738 assert!(
13739 plan.display_indent_schema()
13740 .to_string()
13741 .contains("prom_native_histogram_mul_scalar"),
13742 "{plan:?}"
13743 );
13744 }
13745
13746 #[tokio::test]
13747 async fn test_unsupported_histogram_binary_does_not_block_or_fallback() {
13748 let state = build_query_engine_state();
13749 let plan = PromPlanner::stmt_to_plan(
13750 operator_table_provider(),
13751 &operator_eval_stmt("((lf or on(tag) lh) % 2) or on(tag) lh"),
13752 &state,
13753 )
13754 .await
13755 .unwrap();
13756 let float_field = plan
13757 .schema()
13758 .fields()
13759 .iter()
13760 .find(|field| field.data_type() == &ArrowDataType::Float64)
13761 .unwrap()
13762 .name()
13763 .clone();
13764 let histogram_field = plan
13765 .schema()
13766 .fields()
13767 .iter()
13768 .find(|field| field.data_type() == &native_histogram_value_type().as_arrow_type())
13769 .unwrap()
13770 .name()
13771 .clone();
13772
13773 let (_, batches) = execute(plan, &state).await;
13774 assert_eq!(values(&batches, &float_field), vec![0.0]);
13775 assert_eq!(histograms(&batches, &histogram_field).len(), 1);
13776 }
13777
13778 #[tokio::test]
13779 async fn test_unary_negates_mixed_float_and_native_histogram_samples() {
13780 for histogram_on_left in [false, true] {
13781 let (mut planner, input) = mixed_direct_or(histogram_on_left).await;
13782 let plan = planner.negate_field_columns(input).unwrap();
13783 assert!(PromPlanner::field_columns_are_alternative_samples(
13784 plan.schema(),
13785 &planner.ctx.field_columns
13786 ));
13787 let float_field = planner
13788 .ctx
13789 .field_columns
13790 .iter()
13791 .find(|field| field.starts_with(OR_FLOAT_FIELD_PREFIX))
13792 .unwrap();
13793 let histogram_field = planner
13794 .ctx
13795 .field_columns
13796 .iter()
13797 .find(|field| field.starts_with(OR_HISTOGRAM_FIELD_PREFIX))
13798 .unwrap();
13799
13800 let (_, batches) = execute(plan, &build_query_engine_state()).await;
13801 assert_eq!(values(&batches, float_field), vec![-1.25]);
13802 let histogram = batches
13803 .iter()
13804 .find_map(|batch| {
13805 let values = batch
13806 .column_by_name(histogram_field)
13807 .unwrap()
13808 .as_any()
13809 .downcast_ref::<datafusion::arrow::array::StructArray>()
13810 .unwrap();
13811 (0..values.len()).find_map(|row| {
13812 common_query::native_histogram::read_histogram(values, row).unwrap()
13813 })
13814 })
13815 .unwrap();
13816 assert_eq!(histogram.count, -1.0);
13817 assert_eq!(histogram.sum, -1.0);
13818 assert_eq!(histogram.reset_hint, CounterResetHint::Gauge);
13819 }
13820 }
13821
13822 #[tokio::test]
13823 async fn test_native_histogram_sum_and_avg_execute_real_batches() {
13824 for op_name in ["sum", "avg"] {
13825 for incompatible in [false, true] {
13826 let mut second = direct_or_histogram();
13827 if incompatible {
13828 second.schema = CUSTOM_BUCKETS_SCHEMA;
13829 second.custom_values = vec![1.0];
13830 }
13831 let collector = PromqlAnnotationCollector::default();
13832 let (mut planner, input) =
13833 mixed_aggregate_input(vec![direct_or_histogram(), second]).await;
13834 planner.promql_annotations = Some(collector.clone());
13835 let histogram_column = planner.ctx.field_columns[1].clone();
13836 planner.ctx.field_columns = vec![histogram_column.clone()];
13837 let input = LogicalPlanBuilder::from(input)
13838 .project([col("ts"), col(&histogram_column)])
13839 .unwrap()
13840 .build()
13841 .unwrap();
13842 let PromExpr::Aggregate(AggregateExpr { op, param, .. }) =
13843 parser::parse(&format!("{op_name}(mixed)")).unwrap()
13844 else {
13845 unreachable!()
13846 };
13847 let (aggregate_exprs, _) =
13848 planner.create_aggregate_exprs(op, ¶m, &input).unwrap();
13849 let plan = LogicalPlanBuilder::from(input)
13850 .aggregate(vec![col("ts")], aggregate_exprs)
13851 .unwrap()
13852 .filter(planner.create_empty_values_filter_expr(false).unwrap())
13853 .unwrap()
13854 .build()
13855 .unwrap();
13856
13857 let (_, batches) = execute(plan, &build_query_engine_state()).await;
13858 let mut warnings = vec![];
13859 let mut infos = vec![];
13860 collector.append_to(&mut warnings, &mut infos);
13861 assert!(infos.is_empty());
13862 if incompatible {
13863 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
13864 assert!(warnings.iter().any(|warning| {
13865 warning
13866 == &format!(
13867 "prom_native_histogram_agg_{op_name}: dropped native histogram aggregate with incompatible schemas"
13868 )
13869 }));
13870 } else {
13871 let histograms = histograms(&batches, &histogram_column);
13872 assert_eq!(histograms.len(), 1);
13873 let expected = if op_name == "sum" { 2.0 } else { 1.0 };
13874 assert_eq!(histograms[0].count, expected);
13875 assert_eq!(histograms[0].sum, expected);
13876 assert!(warnings.is_empty());
13877 }
13878 }
13879 }
13880 }
13881
13882 #[tokio::test]
13883 async fn test_canonical_mixed_count_group_and_count_values_execute() {
13884 let state = build_query_engine_state();
13885 for (query, expected) in [
13886 ("count(some_metric)", vec![2.0]),
13887 ("group(some_metric)", vec![1.0]),
13888 (r#"count_values("sample", some_metric)"#, vec![1.0, 1.0]),
13889 ] {
13890 let plan = PromPlanner::stmt_to_plan(
13891 build_test_mixed_native_histogram_table_provider("some_metric").await,
13892 &operator_eval_stmt(query),
13893 &state,
13894 )
13895 .await
13896 .unwrap();
13897 assert!(
13898 plan.schema()
13899 .fields()
13900 .iter()
13901 .all(|field| !field.name().starts_with("__promql_sample_count")),
13902 "{query}: {plan:?}"
13903 );
13904 let value_fields = plan
13905 .schema()
13906 .fields()
13907 .iter()
13908 .filter(|field| {
13909 matches!(
13910 field.data_type(),
13911 ArrowDataType::Float64 | ArrowDataType::Int64 | ArrowDataType::UInt64
13912 ) || field.data_type() == &native_histogram_value_type().as_arrow_type()
13913 })
13914 .collect::<Vec<_>>();
13915 assert_eq!(value_fields.len(), 1, "{query}: {plan:?}");
13916 assert_ne!(
13917 value_fields[0].data_type(),
13918 &native_histogram_value_type().as_arrow_type(),
13919 "{query}: {plan:?}"
13920 );
13921 let value_column = value_fields[0].name().clone();
13922
13923 let (_, batches) = execute(plan, &state).await;
13924 let mut actual = numeric_values(&batches, &value_column);
13925 actual.sort_by(f64::total_cmp);
13926 assert_eq!(actual, expected, "{query}");
13927
13928 if query.starts_with("count_values") {
13929 let mut sample_labels = batches
13930 .iter()
13931 .flat_map(|batch| {
13932 batch
13933 .column_by_name("sample")
13934 .unwrap()
13935 .as_any()
13936 .downcast_ref::<StringArray>()
13937 .unwrap()
13938 .iter()
13939 .flatten()
13940 .map(str::to_string)
13941 })
13942 .collect::<Vec<_>>();
13943 sample_labels.sort();
13944 let mut expected_labels =
13945 vec!["2".to_string(), direct_or_histogram().promql_string()];
13946 expected_labels.sort();
13947 assert_eq!(sample_labels, expected_labels);
13948 }
13949 }
13950 }
13951
13952 #[tokio::test]
13953 async fn test_mixed_or_sum_aggregates_each_sample_type() {
13954 let PromExpr::Aggregate(AggregateExpr { op, param, .. }) =
13955 parser::parse("sum(lhs)").unwrap()
13956 else {
13957 unreachable!()
13958 };
13959
13960 let collector = PromqlAnnotationCollector::default();
13961 let (mut planner, input) = mixed_direct_or(false).await;
13962 planner.promql_annotations = Some(collector.clone());
13963 let float_column = planner.ctx.field_columns[0].clone();
13964 let histogram_column = planner.ctx.field_columns[1].clone();
13965 let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap();
13966 let plan = LogicalPlanBuilder::from(input)
13967 .aggregate(vec![col("ts"), col("k")], aggregate_exprs)
13968 .unwrap()
13969 .filter(
13970 planner
13971 .mixed_aggregate_filter_expr(op, &float_column, &histogram_column)
13972 .unwrap(),
13973 )
13974 .unwrap()
13975 .project([
13976 col(&float_column),
13977 col(&histogram_column),
13978 col("ts"),
13979 col("k"),
13980 ])
13981 .unwrap()
13982 .build()
13983 .unwrap();
13984
13985 let (_, batches) = execute(plan, &build_query_engine_state()).await;
13986 assert_eq!(values(&batches, &float_column), vec![1.25]);
13987 let histogram = batches
13988 .iter()
13989 .find_map(|batch| {
13990 let values = batch
13991 .column_by_name(&histogram_column)?
13992 .as_any()
13993 .downcast_ref::<datafusion::arrow::array::StructArray>()?;
13994 (0..values.len()).find_map(|row| {
13995 common_query::native_histogram::read_histogram(values, row).unwrap()
13996 })
13997 })
13998 .unwrap();
13999 assert_eq!(histogram.count, 1.0);
14000 let mut warnings = vec![];
14001 let mut infos = vec![];
14002 collector.append_to(&mut warnings, &mut infos);
14003 assert!(warnings.is_empty());
14004
14005 let collector = PromqlAnnotationCollector::default();
14006 let (mut planner, input) = mixed_direct_or(false).await;
14007 planner.promql_annotations = Some(collector.clone());
14008 let float_column = planner.ctx.field_columns[0].clone();
14009 let histogram_column = planner.ctx.field_columns[1].clone();
14010 let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap();
14011 let plan = LogicalPlanBuilder::from(input)
14012 .aggregate(vec![col("ts")], aggregate_exprs)
14013 .unwrap()
14014 .filter(
14015 planner
14016 .mixed_aggregate_filter_expr(op, &float_column, &histogram_column)
14017 .unwrap(),
14018 )
14019 .unwrap()
14020 .project([col(&float_column), col(&histogram_column), col("ts")])
14021 .unwrap()
14022 .build()
14023 .unwrap();
14024
14025 let (_, batches) = execute(plan, &build_query_engine_state()).await;
14026 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
14027 let mut warnings = vec![];
14028 let mut infos = vec![];
14029 collector.append_to(&mut warnings, &mut infos);
14030 assert_eq!(
14031 warnings,
14032 vec![
14033 "sum: dropped aggregation result containing both float and native histogram samples"
14034 ]
14035 );
14036 }
14037
14038 #[tokio::test]
14039 async fn test_mixed_or_sum_drops_incompatible_mixed_group() {
14040 let PromExpr::Aggregate(AggregateExpr { op, param, .. }) =
14041 parser::parse("sum(lhs)").unwrap()
14042 else {
14043 unreachable!()
14044 };
14045 let mut custom = direct_or_histogram();
14046 custom.schema = CUSTOM_BUCKETS_SCHEMA;
14047 custom.custom_values = vec![1.0];
14048 let collector = PromqlAnnotationCollector::default();
14049 let (mut planner, input) = mixed_aggregate_input(vec![direct_or_histogram(), custom]).await;
14050 planner.promql_annotations = Some(collector.clone());
14051 let float_column = planner.ctx.field_columns[0].clone();
14052 let histogram_column = planner.ctx.field_columns[1].clone();
14053 let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap();
14054 let plan = LogicalPlanBuilder::from(input)
14055 .aggregate(vec![col("ts")], aggregate_exprs)
14056 .unwrap()
14057 .filter(
14058 planner
14059 .mixed_aggregate_filter_expr(op, &float_column, &histogram_column)
14060 .unwrap(),
14061 )
14062 .unwrap()
14063 .project([col(&float_column), col(&histogram_column), col("ts")])
14064 .unwrap()
14065 .build()
14066 .unwrap();
14067
14068 let (_, batches) = execute(plan, &build_query_engine_state()).await;
14069 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
14070 let mut warnings = vec![];
14071 let mut infos = vec![];
14072 collector.append_to(&mut warnings, &mut infos);
14073 assert!(warnings.iter().any(|warning| {
14074 warning
14075 == "sum: dropped aggregation result containing both float and native histogram samples"
14076 }));
14077 }
14078
14079 #[tokio::test]
14080 async fn test_mixed_or_min_records_only_present_histograms() {
14081 let PromExpr::Aggregate(AggregateExpr { op, param, .. }) =
14082 parser::parse("min(lhs)").unwrap()
14083 else {
14084 unreachable!()
14085 };
14086 let expected_info = "min: dropped native histogram samples because this aggregation is not supported for native histograms";
14087
14088 for (histograms, expected_infos) in [
14089 (vec![], vec![]),
14090 (vec![direct_or_histogram()], vec![expected_info]),
14091 ] {
14092 let collector = PromqlAnnotationCollector::default();
14093 let (mut planner, input) = mixed_aggregate_input(histograms).await;
14094 planner.promql_annotations = Some(collector.clone());
14095 let float_column = planner.ctx.field_columns[0].clone();
14096 let histogram_column = planner.ctx.field_columns[1].clone();
14097 let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap();
14098 let plan = LogicalPlanBuilder::from(input)
14099 .aggregate(vec![col("ts")], aggregate_exprs)
14100 .unwrap()
14101 .filter(
14102 planner
14103 .mixed_ignored_histogram_filter_expr(op, &histogram_column)
14104 .unwrap(),
14105 )
14106 .unwrap()
14107 .project([col(&float_column), col("ts")])
14108 .unwrap()
14109 .build()
14110 .unwrap();
14111
14112 let (_, batches) = execute(plan, &build_query_engine_state()).await;
14113 assert_eq!(values(&batches, &float_column), vec![1.25]);
14114 let mut warnings = vec![];
14115 let mut infos = vec![];
14116 collector.append_to(&mut warnings, &mut infos);
14117 assert!(warnings.is_empty());
14118 assert_eq!(infos, expected_infos);
14119 }
14120 }
14121
14122 #[tokio::test]
14123 async fn test_mixed_or_value_aliases_do_not_replace_labels() {
14124 let left = source(
14125 "lhs",
14126 false,
14127 1,
14128 vec![("job", Some("job")), ("k", Some("float"))],
14129 DirectOrValue::Float64(1.0),
14130 );
14131 let right = source(
14132 "rhs",
14133 false,
14134 1,
14135 vec![
14136 ("job", Some("job")),
14137 ("k", Some("histogram")),
14138 (greptime_value(), Some("value-label")),
14139 ],
14140 DirectOrValue::NativeHistogram(direct_or_histogram()),
14141 );
14142 let table_provider = build_test_table_provider_with_fields(
14143 &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
14144 &[],
14145 )
14146 .await;
14147 let mut planner = PromPlanner {
14148 table_provider,
14149 ctx: PromPlannerContext::default(),
14150 promql_annotations: None,
14151 };
14152 let left = LogicalPlanBuilder::from(scan(&left))
14153 .project(vec![
14154 col("ts"),
14155 col("job"),
14156 col("k"),
14157 col("v").alias(greptime_value()),
14158 ])
14159 .unwrap()
14160 .build()
14161 .unwrap();
14162 let left_context = direct_or_context("lhs", &["job", "k"], greptime_value());
14163 let right_context = direct_or_context("rhs", &["job", "k", greptime_value()], "v");
14164 let plan = planner
14165 .or_operator(
14166 left,
14167 scan(&right),
14168 left_context.tag_columns.iter().cloned().collect(),
14169 right_context.tag_columns.iter().cloned().collect(),
14170 left_context,
14171 right_context,
14172 &or_modifier("lhs or on(k) rhs"),
14173 )
14174 .unwrap();
14175
14176 assert_eq!(
14177 plan.schema()
14178 .field_with_name(None, greptime_value())
14179 .unwrap()
14180 .data_type(),
14181 &ArrowDataType::Utf8
14182 );
14183 assert!(
14184 planner
14185 .ctx
14186 .field_columns
14187 .iter()
14188 .all(|field| { field != greptime_value() && field != greptime_native_histogram() })
14189 );
14190 assert!(PromPlanner::field_columns_are_alternative_samples(
14191 plan.schema(),
14192 &planner.ctx.field_columns
14193 ));
14194 let (_, batches) = execute(plan, &build_query_engine_state()).await;
14195 assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
14196 let labels = batches
14197 .iter()
14198 .flat_map(|batch| {
14199 batch
14200 .column_by_name(greptime_value())
14201 .unwrap()
14202 .as_any()
14203 .downcast_ref::<StringArray>()
14204 .unwrap()
14205 .iter()
14206 .flatten()
14207 })
14208 .collect::<Vec<_>>();
14209 assert_eq!(labels, vec!["value-label"]);
14210 }
14211
14212 #[tokio::test]
14213 async fn test_mixed_or_routes_float_histogram_and_label_functions() {
14214 for (function, expected) in [("abs", 1.25), ("round", 1.0), ("histogram_count", 1.0)] {
14215 let (mut planner, input) = mixed_direct_or(false).await;
14216 let preserve_any_value = PromPlanner::field_columns_are_alternative_samples(
14217 input.schema(),
14218 &planner.ctx.field_columns,
14219 );
14220 let PromExpr::Call(call) = parser::parse(&format!("{function}(lhs)")).unwrap() else {
14221 unreachable!()
14222 };
14223 let state = build_query_engine_state();
14224 let (mut exprs, _) = planner
14225 .create_function_expr(&call.func, vec![], input.schema(), &state)
14226 .unwrap();
14227 exprs.insert(0, planner.create_time_index_column_expr().unwrap());
14228 exprs.extend(planner.create_tag_column_exprs().unwrap());
14229 let plan = LogicalPlanBuilder::from(input)
14230 .project(exprs)
14231 .unwrap()
14232 .filter(
14233 planner
14234 .create_empty_values_filter_expr(preserve_any_value)
14235 .unwrap(),
14236 )
14237 .unwrap()
14238 .build()
14239 .unwrap();
14240 let (_, batches) = execute(plan, &state).await;
14241 let values = batches
14242 .iter()
14243 .flat_map(|batch| {
14244 batch
14245 .schema()
14246 .fields()
14247 .iter()
14248 .position(|field| field.data_type() == &ArrowDataType::Float64)
14249 .map(|index| {
14250 batch
14251 .column(index)
14252 .as_any()
14253 .downcast_ref::<Float64Array>()
14254 .unwrap()
14255 .iter()
14256 .flatten()
14257 })
14258 .into_iter()
14259 .flatten()
14260 })
14261 .collect::<Vec<_>>();
14262 assert_eq!(values, vec![expected], "{function}");
14263 }
14264
14265 let (mut planner, input) = mixed_direct_or(false).await;
14266 let preserve_any_value = PromPlanner::field_columns_are_alternative_samples(
14267 input.schema(),
14268 &planner.ctx.field_columns,
14269 );
14270 let PromExpr::Call(call) =
14271 parser::parse(r#"label_replace(lhs, "copy", "$1", "k", "(.*)")"#).unwrap()
14272 else {
14273 unreachable!()
14274 };
14275 let args = planner.create_function_args(&call.args.args).unwrap();
14276 let state = build_query_engine_state();
14277 let (mut exprs, _) = planner
14278 .create_function_expr(&call.func, args.literals, input.schema(), &state)
14279 .unwrap();
14280 exprs.insert(0, planner.create_time_index_column_expr().unwrap());
14281 exprs.extend(planner.create_tag_column_exprs().unwrap());
14282 let plan = LogicalPlanBuilder::from(input)
14283 .project(exprs)
14284 .unwrap()
14285 .filter(
14286 planner
14287 .create_empty_values_filter_expr(preserve_any_value)
14288 .unwrap(),
14289 )
14290 .unwrap()
14291 .build()
14292 .unwrap();
14293 let (_, batches) = execute(plan, &state).await;
14294 let sample_count = batches.iter().map(RecordBatch::num_rows).sum::<usize>();
14295 assert_eq!(sample_count, 2);
14296 }
14297}