1mod at_modifier;
16mod function_plans;
17mod island;
18mod matching_filters;
19
20mod set_operator;
21
22use std::collections::{BTreeSet, HashSet, VecDeque};
23use std::sync::Arc;
24
25use arrow::datatypes::IntervalDayTime;
26use async_recursion::async_recursion;
27use catalog::table_source::DfTableSourceProvider;
28use common_error::ext::ErrorExt;
29use common_error::status_code::StatusCode;
30use common_function::aggrs::extremum::Extremum;
31use common_function::function::FunctionContext;
32use common_query::native_histogram::native_histogram_value_type;
33use common_query::prelude::{
34 GREPTIME_TEMPORALITY_DELTA, OTLP_AGGREGATION_TEMPORALITY_LABEL, greptime_native_histogram,
35 greptime_value,
36};
37use common_query::promql_annotations::PromqlAnnotationCollector;
38use datafusion::common::DFSchemaRef;
39use datafusion::datasource::DefaultTableSource;
40use datafusion::functions_aggregate::average::avg_udaf;
41use datafusion::functions_aggregate::count::count_udaf;
42use datafusion::functions_aggregate::expr_fn::first_value;
43use datafusion::functions_aggregate::min_max::{max_udaf, min_udaf};
44use datafusion::functions_aggregate::stddev::stddev_pop_udaf;
45use datafusion::functions_aggregate::sum::sum_udaf;
46use datafusion::functions_aggregate::variance::var_pop_udaf;
47use datafusion::functions_window::row_number::RowNumber;
48use datafusion::logical_expr::expr::{Alias, ScalarFunction, WindowFunction};
49use datafusion::logical_expr::expr_rewriter::normalize_cols;
50use datafusion::logical_expr::{
51 BinaryExpr, Cast, EmptyRelation, Extension, LogicalPlan, LogicalPlanBuilder, Operator,
52 ScalarUDF as ScalarUdfDef, WindowFrame, WindowFunctionDefinition,
53};
54use datafusion::prelude as df_prelude;
55use datafusion::prelude::{Column, Expr as DfExpr, JoinType};
56use datafusion::scalar::ScalarValue;
57use datafusion_common::{DFSchema, NullEquality, TableReference};
58use datafusion_expr::expr::WindowFunctionParams;
59use datafusion_expr::expr_fn::when;
60use datafusion_expr::utils::{conjunction, disjunction};
61use datafusion_expr::{ExprSchemable, Literal, SortExpr, TableSource, ident, lit};
62use datafusion_functions::core::coalesce;
63use datatypes::arrow::datatypes::{DataType as ArrowDataType, TimeUnit as ArrowTimeUnit};
64use datatypes::data_type::{ConcreteDataType, DataType as GreptimeDataType};
65use itertools::Itertools;
66use once_cell::sync::Lazy;
67use promql::extension_plan::{
68 EmptyMetric, InstantManipulate, Millisecond, RangeManipulate, SeriesDivide, SeriesNormalize,
69 build_special_time_expr,
70};
71use promql::functions::{
72 AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing,
73 IDelta, IeeeComparison, IeeeSqrt, Increase, LastOverTime, MatchGroupViolation, MaxOverTime,
74 MinOverTime, MixedRange, NativeHistogramAbsentOverTime, NativeHistogramAdd,
75 NativeHistogramAggAvg, NativeHistogramAggSum, NativeHistogramAvg, NativeHistogramAvgOverTime,
76 NativeHistogramChanges, NativeHistogramCount, NativeHistogramCountOverTime,
77 NativeHistogramDelta, NativeHistogramDivScalar, NativeHistogramDrop, NativeHistogramEq,
78 NativeHistogramIDelta, NativeHistogramIRate, NativeHistogramIncrease,
79 NativeHistogramLastOverTime, NativeHistogramMulScalar, NativeHistogramNeg,
80 NativeHistogramNotEq, NativeHistogramPresentOverTime, NativeHistogramRate,
81 NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev, NativeHistogramStdvar,
82 NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime, NativeHistogramToString,
83 PredictLinear, PresentOverTime, PromqlFloatToString, QuantileOverTime, Rate, Resets, Round,
84 StddevOverTime, StdvarOverTime, SumOverTime, UniqueMatchGroup, 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, LabelModifier, MatrixSelector, NumberLiteral, Offset, ParenExpr, StringLiteral,
92 SubqueryExpr, UnaryExpr, VectorMatchCardinality, VectorSelector, token,
93};
94use regex::{self, Regex};
95use snafu::{OptionExt, ResultExt, ensure};
96use store_api::metric_engine_consts::{
97 DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, LOGICAL_TABLE_METADATA_KEY,
98 METRIC_ENGINE_NAME, is_metric_engine_internal_column,
99};
100use table::table::adapter::DfTableProviderAdapter;
101
102use crate::parser::{
103 ALIAS_NODE_NAME, ANALYZE_NODE_NAME, ANALYZE_VERBOSE_NODE_NAME, AliasExpr, EXPLAIN_NODE_NAME,
104 EXPLAIN_VERBOSE_NODE_NAME,
105};
106use crate::promql::error::{
107 CatalogSnafu, ColumnNotFoundSnafu, DataFusionPlanningSnafu, ExpectRangeSelectorSnafu,
108 FunctionInvalidArgumentSnafu, InvalidDestinationLabelNameSnafu, InvalidRegularExpressionSnafu,
109 InvalidTimeRangeSnafu, MultiFieldsNotSupportedSnafu, MultipleMetricMatchersSnafu,
110 MultipleVectorSnafu, NoMetricMatcherSnafu, RangeVectorInRangeQuerySnafu, Result,
111 SubqueryTimestampOutOfRangeSnafu, TableNameNotFoundSnafu, TimeIndexNotFoundSnafu,
112 UnexpectedPlanExprSnafu, UnexpectedTokenSnafu, UnknownTableSnafu, UnsupportedExprSnafu,
113 UnsupportedMatcherOpSnafu, ValueNotFoundSnafu, ZeroRangeSelectorSnafu, ZeroSubqueryStepSnafu,
114};
115use crate::promql::label_values::signed_millis_since_epoch;
116use crate::query_engine::QueryEngineState;
117
118const DEFAULT_SUBQUERY_STEP_MS: Millisecond = 60_000;
121const SPECIAL_TIME_FUNCTION: &str = "time";
123const SCALAR_FUNCTION: &str = "scalar";
125const SPECIAL_ABSENT_FUNCTION: &str = "absent";
127const SPECIAL_HISTOGRAM_QUANTILE: &str = "histogram_quantile";
129const SPECIAL_HISTOGRAM_FRACTION: &str = "histogram_fraction";
131const SPECIAL_VECTOR_FUNCTION: &str = "vector";
133const LE_COLUMN_NAME: &str = "le";
135
136static LABEL_NAME_REGEX: Lazy<Regex> =
139 Lazy::new(|| Regex::new(r"^[a-zA-Z_][a-zA-Z0-9_]*$").unwrap());
140
141const DEFAULT_TIME_INDEX_COLUMN: &str = "time";
142
143const DEFAULT_FIELD_COLUMN: &str = "value";
145
146const FIELD_COLUMN_MATCHER: &str = "__field__";
148
149const SCHEMA_COLUMN_MATCHER: &str = "__schema__";
151const DB_COLUMN_MATCHER: &str = "__database__";
152
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_";
156const MATCH_GROUP_COUNT_COLUMN: &str = "__promql_match_group_count";
158
159const MAX_SCATTER_POINTS: i64 = 400;
161
162const INTERVAL_1H: i64 = 60 * 60 * 1000;
164
165#[derive(Default, Debug, Clone)]
166struct PromPlannerContext {
167 start: Millisecond,
169 end: Millisecond,
170 interval: Millisecond,
171 lookback_delta: Millisecond,
172 stmt_start: Millisecond,
177 stmt_end: Millisecond,
178
179 table_name: Option<String>,
181 time_index_column: Option<String>,
182 field_columns: Vec<String>,
183 tag_columns: Vec<String>,
184 aggregation_field_labels: Vec<String>,
193 use_tsid: bool,
199 field_column_matcher: Option<Vec<Matcher>>,
201 selector_matcher: Vec<Matcher>,
203 schema_name: Option<String>,
204 range: Option<Millisecond>,
206 range_fold_offset: Option<Millisecond>,
216}
217
218#[derive(Debug)]
221struct BinaryResultLabels {
222 exprs: Vec<DfExpr>,
223 names: Vec<String>,
224 aggregation_field_labels: Vec<String>,
225 tsid: Option<DfExpr>,
228}
229
230impl BinaryResultLabels {
231 fn apply(&self, ctx: &mut PromPlannerContext) {
232 ctx.tag_columns = self.names.clone();
233 ctx.aggregation_field_labels = self.aggregation_field_labels.clone();
234 ctx.use_tsid = self.tsid.is_some();
235 }
236
237 fn tsid_projection(&self, table_ref: Option<TableReference>) -> Option<DfExpr> {
241 self.tsid
242 .clone()
243 .map(|tsid| tsid.alias_qualified(table_ref, DATA_SCHEMA_TSID_COLUMN_NAME))
244 }
245}
246
247impl PromPlannerContext {
248 fn from_eval_stmt(stmt: &EvalStmt) -> Result<Self> {
249 let start = signed_millis_since_epoch(stmt.start)?;
251 let end = signed_millis_since_epoch(stmt.end)?;
252 Ok(Self {
253 start,
254 end,
255 interval: stmt.interval.as_millis() as _,
256 lookback_delta: stmt.lookback_delta.as_millis() as _,
257 stmt_start: start,
258 stmt_end: end,
259 ..Default::default()
260 })
261 }
262
263 fn reset(&mut self) {
265 self.table_name = None;
266 self.time_index_column = None;
267 self.field_columns = vec![];
268 self.tag_columns = vec![];
269 self.use_tsid = false;
270 self.field_column_matcher = None;
271 self.selector_matcher.clear();
272 self.schema_name = None;
273 self.range = None;
274 self.range_fold_offset = None;
275 }
276
277 fn reset_table_name_and_schema(&mut self) {
279 self.table_name = Some(String::new());
280 self.schema_name = None;
281 self.use_tsid = false;
282 }
283
284 fn has_le_tag(&self) -> bool {
286 self.tag_columns.iter().any(|c| c.eq(&LE_COLUMN_NAME))
287 }
288}
289
290pub struct PromPlanner {
291 table_provider: DfTableSourceProvider,
292 ctx: PromPlannerContext,
293 promql_annotations: Option<PromqlAnnotationCollector>,
295}
296
297type BinaryFieldPair<'a> = (&'a String, &'a String);
298
299impl PromPlanner {
300 pub async fn stmt_to_plan(
301 table_provider: DfTableSourceProvider,
302 stmt: &EvalStmt,
303 query_engine_state: &QueryEngineState,
304 ) -> Result<LogicalPlan> {
305 Self::stmt_to_plan_with_annotations(table_provider, stmt, query_engine_state, None).await
306 }
307
308 pub async fn stmt_to_plan_with_annotations(
310 table_provider: DfTableSourceProvider,
311 stmt: &EvalStmt,
312 query_engine_state: &QueryEngineState,
313 promql_annotations: Option<PromqlAnnotationCollector>,
314 ) -> Result<LogicalPlan> {
315 let mut planner = Self {
316 table_provider,
317 ctx: PromPlannerContext::from_eval_stmt(stmt)?,
318 promql_annotations,
319 };
320
321 let plan = planner
322 .prom_root_expr_to_plan(&stmt.expr, query_engine_state)
323 .await?;
324
325 planner.strip_tsid_column(plan)
327 }
328
329 pub async fn prom_expr_to_plan(
330 &mut self,
331 prom_expr: &PromExpr,
332 query_engine_state: &QueryEngineState,
333 ) -> Result<LogicalPlan> {
334 self.prom_expr_to_plan_inner(prom_expr, false, query_engine_state)
335 .await
336 }
337
338 #[async_recursion]
348 async fn prom_expr_to_plan_inner(
349 &mut self,
350 prom_expr: &PromExpr,
351 timestamp_fn: bool,
352 query_engine_state: &QueryEngineState,
353 ) -> Result<LogicalPlan> {
354 if let Some(plan) = self
358 .promote_anchored_range_call(prom_expr, timestamp_fn, query_engine_state)
359 .await?
360 {
361 return Ok(plan);
362 }
363
364 let res = match prom_expr {
365 PromExpr::Aggregate(expr) => {
366 self.prom_aggr_expr_to_plan(query_engine_state, expr)
367 .await?
368 }
369 PromExpr::Unary(expr) => {
370 self.prom_unary_expr_to_plan(query_engine_state, expr)
371 .await?
372 }
373 PromExpr::Binary(expr) => {
374 self.prom_binary_expr_to_plan(query_engine_state, expr)
375 .await?
376 }
377 PromExpr::Paren(ParenExpr { expr }) => {
378 self.prom_expr_to_plan_inner(expr, timestamp_fn, query_engine_state)
379 .await?
380 }
381 PromExpr::Subquery(expr) => {
382 self.prom_subquery_expr_to_plan(query_engine_state, expr)
383 .await?
384 }
385 PromExpr::NumberLiteral(lit) => self.prom_number_lit_to_plan(lit)?,
386 PromExpr::StringLiteral(lit) => self.prom_string_lit_to_plan(lit)?,
387 PromExpr::VectorSelector(selector) => {
388 self.prom_vector_selector_to_plan(selector, timestamp_fn)
389 .await?
390 }
391 PromExpr::MatrixSelector(selector) => {
392 self.prom_matrix_selector_to_plan(selector).await?
393 }
394 PromExpr::Call(expr) => {
395 self.prom_call_expr_to_plan(query_engine_state, expr)
396 .await?
397 }
398 PromExpr::Extension(expr) => {
399 self.prom_ext_expr_to_plan(query_engine_state, expr).await?
400 }
401 };
402
403 Ok(res)
404 }
405
406 async fn prom_root_expr_to_plan(
412 &mut self,
413 prom_expr: &PromExpr,
414 query_engine_state: &QueryEngineState,
415 ) -> Result<LogicalPlan> {
416 if matches!(prom_expr, PromExpr::Extension(_))
417 || prom_expr.value_type() != ValueType::Matrix
418 {
419 return self.prom_expr_to_plan(prom_expr, query_engine_state).await;
420 }
421 ensure!(self.ctx.start == self.ctx.end, RangeVectorInRangeQuerySnafu);
422 self.prom_range_vector_result_to_plan(prom_expr, query_engine_state)
423 .await
424 }
425
426 async fn prom_range_vector_result_to_plan(
432 &mut self,
433 prom_expr: &PromExpr,
434 query_engine_state: &QueryEngineState,
435 ) -> Result<LogicalPlan> {
436 let mut expr = prom_expr;
437 while let PromExpr::Paren(ParenExpr { expr: inner }) = expr {
438 expr = inner;
439 }
440 match expr {
441 PromExpr::MatrixSelector(selector) => {
442 self.prom_range_selector_samples_to_plan(selector).await
443 }
444 PromExpr::Subquery(subquery) => {
445 self.prom_subquery_samples_to_plan(subquery, query_engine_state)
446 .await
447 }
448 _ => UnsupportedExprSnafu {
449 name: format!("range vector result of {expr:?}"),
450 }
451 .fail(),
452 }
453 }
454
455 async fn prom_range_selector_samples_to_plan(
458 &mut self,
459 selector: &MatrixSelector,
460 ) -> Result<LogicalPlan> {
461 let MatrixSelector { vs, range } = selector;
462 let VectorSelector {
463 name,
464 offset,
465 matchers,
466 at,
467 } = vs;
468 let matchers = self.preprocess_label_matchers(matchers, name)?;
469 ensure!(!range.is_zero(), ZeroRangeSelectorSnafu);
470 self.ctx.range = Some(range.as_millis() as _);
471 let offset_ms = match self.at_modifier_offset(at, offset)? {
472 Some(at_offset) => at_offset,
473 None => Self::offset_millis(offset),
474 };
475 if let Some(empty_plan) = self.setup_context().await? {
476 return Ok(empty_plan);
477 }
478 let plan = self
479 .selector_to_series_normalize_plan(offset_ms, matchers, true)
480 .await?;
481 self.millisecond_time_index_plan(plan)
482 }
483
484 fn millisecond_time_index_plan(&self, plan: LogicalPlan) -> Result<LogicalPlan> {
487 let time_index =
488 self.ctx
489 .time_index_column
490 .as_deref()
491 .with_context(|| TimeIndexNotFoundSnafu {
492 table: self.ctx.table_name.clone().unwrap_or_default(),
493 })?;
494 let (_, field) = plan
495 .schema()
496 .qualified_field_with_unqualified_name(time_index)
497 .context(DataFusionPlanningSnafu)?;
498 let millis = match field.data_type() {
499 ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, _) => return Ok(plan),
500 ArrowDataType::Timestamp(_, timezone) => {
501 ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, timezone.clone())
502 }
503 other => unreachable!("time index is a timestamp, got {other}"),
504 };
505 let exprs = plan
506 .schema()
507 .iter()
508 .map(|(qualifier, field)| {
509 let column = DfExpr::Column(Column::new(qualifier.cloned(), field.name()));
510 if field.name() == time_index {
511 column
512 .cast_to(&millis, plan.schema())
513 .map(|cast| cast.alias_qualified(qualifier.cloned(), field.name().clone()))
514 } else {
515 Ok(column)
516 }
517 })
518 .collect::<datafusion::error::Result<Vec<_>>>()
519 .context(DataFusionPlanningSnafu)?;
520 LogicalPlanBuilder::from(plan)
521 .project(exprs)
522 .context(DataFusionPlanningSnafu)?
523 .build()
524 .context(DataFusionPlanningSnafu)
525 }
526
527 async fn prom_subquery_samples_to_plan(
530 &mut self,
531 subquery: &SubqueryExpr,
532 query_engine_state: &QueryEngineState,
533 ) -> Result<LogicalPlan> {
534 let SubqueryExpr {
535 expr,
536 range,
537 step,
538 offset,
539 at,
540 ..
541 } = subquery;
542 ensure!(!range.is_zero(), ZeroRangeSelectorSnafu);
543 let end = match self.at_ref_time(at, offset)? {
544 Some(anchor) => anchor,
545 None => Self::anchor_sub(self.ctx.end, Self::offset_millis(offset))?,
546 };
547 self.ctx.interval = match step {
548 Some(step) => step.as_millis() as _,
549 None => DEFAULT_SUBQUERY_STEP_MS,
550 };
551 let interval = self.ctx.interval;
552 ensure!(interval > 0, ZeroSubqueryStepSnafu);
553 let window = Self::subquery_child_window(end, end, 0, 0, range.as_millis() as _, interval)?;
554 let Some((start, end)) = window else {
555 self.ctx.start = end;
558 self.ctx.end = end;
559 let plan = self.prom_expr_to_plan(expr, query_engine_state).await?;
560 return Ok(LogicalPlan::EmptyRelation(EmptyRelation {
561 produce_one_row: false,
562 schema: plan.schema().clone(),
563 }));
564 };
565 self.ctx.start = start;
566 self.ctx.end = end;
567 self.prom_expr_to_plan(expr, query_engine_state).await
568 }
569
570 async fn prom_subquery_expr_to_plan(
571 &mut self,
572 query_engine_state: &QueryEngineState,
573 subquery_expr: &SubqueryExpr,
574 ) -> Result<LogicalPlan> {
575 let SubqueryExpr {
576 expr,
577 range,
578 step,
579 offset,
580 ..
581 } = subquery_expr;
582
583 let offset_ms = match offset {
585 Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond,
586 Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond),
587 None => 0,
588 };
589
590 let current_interval = self.ctx.interval;
591 self.ctx.interval = match step {
592 Some(step) => step.as_millis() as _,
593 None => DEFAULT_SUBQUERY_STEP_MS,
594 };
595 ensure!(self.ctx.interval > 0, ZeroRangeSelectorSnafu);
596 let current_start = self.ctx.start;
597 let current_end = self.ctx.end;
598 ensure!(
600 current_start <= current_end,
601 InvalidTimeRangeSnafu {
602 start: current_start,
603 end: current_end,
604 }
605 );
606 let child_window = Self::subquery_child_window(
608 current_start,
609 current_end,
610 current_interval,
611 offset_ms,
612 range.as_millis() as Millisecond,
613 self.ctx.interval,
614 )?;
615 if let Some((child_start, child_end)) = child_window {
616 self.ctx.start = child_start;
617 self.ctx.end = child_end;
618 }
619 let input = self.prom_expr_to_plan(expr, query_engine_state).await?;
620 self.ctx.interval = current_interval;
621 self.ctx.start = current_start;
622 self.ctx.end = current_end;
623 let input = match child_window {
625 Some(_) => input,
626 None => LogicalPlan::EmptyRelation(EmptyRelation {
627 produce_one_row: false,
628 schema: input.schema().clone(),
629 }),
630 };
631
632 ensure!(!range.is_zero(), ZeroRangeSelectorSnafu);
633 let range_ms = range.as_millis() as _;
634 self.ctx.range = Some(range_ms);
635
636 let time_index_column =
637 self.ctx
638 .time_index_column
639 .clone()
640 .with_context(|| TimeIndexNotFoundSnafu {
641 table: self.ctx.table_name.clone().unwrap_or_default(),
642 })?;
643
644 let input_schema = input.schema();
650 let input_has_tsid = input_schema.fields().iter().any(|field| {
651 field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
652 && field.data_type() == &ArrowDataType::UInt64
653 });
654 let (series_key_columns, mut sort_exprs) = if input_has_tsid {
655 (
656 vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()],
657 vec![
658 DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME))
659 .sort(true, true),
660 ],
661 )
662 } else {
663 let key_columns: Vec<String> = self
666 .ctx
667 .tag_columns
668 .iter()
669 .filter(|name| input_schema.has_column_with_unqualified_name(name))
670 .cloned()
671 .collect();
672 let sort = key_columns
673 .iter()
674 .map(|name| DfExpr::Column(Column::from_name(name)).sort(true, true))
675 .collect::<Vec<_>>();
676 (key_columns, sort)
677 };
678 sort_exprs.push(DfExpr::Column(Column::from_name(&time_index_column)).sort(true, true));
679
680 let sort_plan = LogicalPlanBuilder::from(input)
681 .sort(sort_exprs)
682 .context(DataFusionPlanningSnafu)?
683 .build()
684 .context(DataFusionPlanningSnafu)?;
685 let divide_plan = LogicalPlan::Extension(Extension {
686 node: Arc::new(SeriesDivide::new(
687 series_key_columns.clone(),
688 time_index_column.clone(),
689 sort_plan,
690 )),
691 });
692
693 let divide_plan = if offset_ms != 0 {
696 LogicalPlan::Extension(Extension {
697 node: Arc::new(SeriesNormalize::new(
698 offset_ms,
699 time_index_column.clone(),
700 false,
701 series_key_columns,
702 divide_plan,
703 )),
704 })
705 } else {
706 divide_plan
707 };
708
709 let manipulate = RangeManipulate::new(
710 self.ctx.start,
711 self.ctx.end,
712 self.ctx.interval,
713 offset_ms,
714 range_ms,
715 time_index_column,
716 self.ctx.field_columns.clone(),
717 divide_plan,
718 )
719 .context(DataFusionPlanningSnafu)?;
720 self.ctx.range_fold_offset = Some(offset_ms);
723
724 Ok(LogicalPlan::Extension(Extension {
725 node: Arc::new(manipulate),
726 }))
727 }
728
729 async fn prom_aggr_expr_to_plan(
730 &mut self,
731 query_engine_state: &QueryEngineState,
732 aggr_expr: &AggregateExpr,
733 ) -> Result<LogicalPlan> {
734 let AggregateExpr {
735 op,
736 expr,
737 modifier,
738 param,
739 } = aggr_expr;
740
741 let input = self.prom_expr_to_plan(expr, query_engine_state).await?;
742 let input_has_tsid = input.schema().fields().iter().any(|field| {
743 field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
744 && field.data_type() == &ArrowDataType::UInt64
745 });
746
747 match (*op).id() {
748 token::T_TOPK | token::T_BOTTOMK => {
749 self.prom_topk_bottomk_to_plan(aggr_expr, input).await
750 }
751 _ => {
752 let input_tag_columns = if input_has_tsid {
756 self.collect_row_key_tag_columns_from_plan(&input)?
757 .into_iter()
758 .collect::<Vec<_>>()
759 } else {
760 self.ctx.tag_columns.clone()
761 };
762 let mut group_exprs = self.agg_modifier_to_col(input.schema(), modifier, true)?;
765 let mixed_sample_columns =
766 Self::alternative_sample_columns(input.schema(), &self.ctx.field_columns)
767 .map(|(float, histogram)| (float.to_string(), histogram.to_string()));
768 let preserve_any_value = mixed_sample_columns.is_some();
773 let has_native_histogram = preserve_any_value
774 || self.all_field_columns_are_native_histograms(input.schema());
775 let (mut aggr_exprs, prev_field_exprs) =
777 self.create_aggregate_exprs(*op, param, &input)?;
778 let prev_field_exprs =
779 normalize_cols(prev_field_exprs, &input).context(DataFusionPlanningSnafu)?;
780
781 let keep_tsid = op.id() != token::T_COUNT_VALUES
782 && input_has_tsid
783 && input_tag_columns.iter().collect::<HashSet<_>>()
784 == self.ctx.tag_columns.iter().collect::<HashSet<_>>();
785
786 if keep_tsid {
787 aggr_exprs.push(
788 first_value(
789 DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME)),
790 vec![],
791 )
792 .alias(DATA_SCHEMA_TSID_COLUMN_NAME),
793 );
794 }
795 self.ctx.use_tsid = keep_tsid;
796
797 let builder = LogicalPlanBuilder::from(input);
799 let builder = if op.id() == token::T_COUNT_VALUES {
800 let label = Self::get_param_value_as_str(*op, param)?;
801 self.ctx.tag_columns.retain(|tag| tag != label);
815 group_exprs.retain(
816 |expr| !matches!(expr, DfExpr::Column(column) if column.name == label),
817 );
818 let label_qualifier = self
823 .ctx
824 .time_index_column
825 .as_deref()
826 .and_then(|time_index| {
827 builder
828 .schema()
829 .qualified_field_with_unqualified_name(time_index)
830 .ok()
831 })
832 .and_then(|(qualifier, _)| qualifier.cloned());
833 let count_value_exprs = prev_field_exprs
840 .iter()
841 .map(|expr| {
842 let value = match expr {
843 DfExpr::Column(column) => {
844 let value = DfExpr::Column(column.clone());
845 let value = if matches!(
848 builder.schema().field_with_unqualified_name(&column.name),
849 Ok(field) if field.data_type() == &ArrowDataType::Float64
850 ) {
851 value
852 } else {
853 DfExpr::Cast(Cast::new(
854 Box::new(value),
855 ArrowDataType::Float64,
856 ))
857 };
858 DfExpr::ScalarFunction(ScalarFunction {
859 func: Arc::new(PromqlFloatToString::scalar_udf()),
860 args: vec![value],
861 })
862 }
863 _ => DfExpr::Column(Column::from_name(
867 expr.schema_name().to_string(),
868 )),
869 };
870 DfExpr::Alias(Alias::new(value, label_qualifier.clone(), label))
871 })
872 .collect::<Vec<_>>();
873 let aggregate_group_exprs = group_exprs
874 .iter()
875 .cloned()
876 .chain(prev_field_exprs.clone())
877 .collect::<Vec<_>>();
878 group_exprs.push(ident(label));
879 let project_fields = self
880 .create_field_column_exprs()?
881 .into_iter()
882 .chain(self.create_tag_column_exprs()?)
883 .chain(Some(self.create_time_index_column_expr()?))
884 .chain(count_value_exprs);
885
886 let builder = builder
887 .aggregate(aggregate_group_exprs, aggr_exprs)
888 .context(DataFusionPlanningSnafu)?
889 .project(project_fields)
890 .context(DataFusionPlanningSnafu)?;
891 self.ctx.tag_columns.push(label.to_string());
893 builder
894 } else {
895 builder
896 .aggregate(group_exprs.clone(), aggr_exprs)
897 .context(DataFusionPlanningSnafu)?
898 };
899
900 let builder = if let Some((float, histogram)) = mixed_sample_columns {
901 let builder = match op.id() {
902 token::T_SUM | token::T_AVG => builder
903 .filter(self.mixed_aggregate_filter_expr(*op, &float, &histogram)?)
904 .context(DataFusionPlanningSnafu)?,
905 token::T_MIN
906 | token::T_MAX
907 | token::T_STDDEV
908 | token::T_STDVAR
909 | token::T_QUANTILE => builder
910 .filter(self.mixed_ignored_histogram_filter_expr(*op, &histogram)?)
911 .context(DataFusionPlanningSnafu)?,
912 _ => builder,
913 };
914
915 match op.id() {
916 token::T_SUM
917 | token::T_AVG
918 | token::T_MIN
919 | token::T_MAX
920 | token::T_STDDEV
921 | token::T_STDVAR
922 | token::T_QUANTILE => {
923 let project_fields = self
924 .create_field_column_exprs()?
925 .into_iter()
926 .chain(self.create_tag_column_exprs()?)
927 .chain(self.ctx.use_tsid.then_some(DfExpr::Column(
928 Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME),
929 )))
930 .chain(Some(self.create_time_index_column_expr()?));
931 builder
932 .project(project_fields)
933 .context(DataFusionPlanningSnafu)?
934 }
935 _ => builder,
936 }
937 } else {
938 builder
939 };
940
941 let builder = if has_native_histogram {
946 builder
947 .filter(self.create_empty_values_filter_expr(preserve_any_value)?)
948 .context(DataFusionPlanningSnafu)?
949 } else {
950 builder
951 };
952
953 let sort_expr = group_exprs.into_iter().map(|expr| expr.sort(true, false));
954
955 builder
956 .sort(sort_expr)
957 .context(DataFusionPlanningSnafu)?
958 .build()
959 .context(DataFusionPlanningSnafu)
960 }
961 }
962 }
963
964 async fn prom_topk_bottomk_to_plan(
966 &mut self,
967 aggr_expr: &AggregateExpr,
968 input: LogicalPlan,
969 ) -> Result<LogicalPlan> {
970 let AggregateExpr {
971 op,
972 param,
973 modifier,
974 ..
975 } = aggr_expr;
976
977 let input_has_tsid = input.schema().fields().iter().any(|field| {
978 field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
979 && field.data_type() == &ArrowDataType::UInt64
980 });
981 self.ctx.use_tsid = input_has_tsid;
982
983 let group_exprs = self.agg_modifier_to_col(input.schema(), modifier, false)?;
984
985 let mut input = input;
986 if let Some((float_column, histogram_column)) =
987 Self::alternative_sample_columns(input.schema(), &self.ctx.field_columns)
988 .map(|(float, histogram)| (float.to_string(), histogram.to_string()))
989 {
990 let drop_histogram = DfExpr::ScalarFunction(ScalarFunction {
991 func: Arc::new(NativeHistogramDrop::bool_false_udf(
992 format!(
993 "{}: dropped native histogram samples because this aggregation is not supported for native histograms",
994 op
995 ),
996 self.promql_annotations.clone(),
997 )),
998 args: vec![ident(&histogram_column)],
999 });
1000 let keep_float = when(ident(&histogram_column).is_not_null(), drop_histogram)
1001 .otherwise(ident(&float_column).is_not_null())
1002 .context(DataFusionPlanningSnafu)?;
1003 input = LogicalPlanBuilder::from(input)
1004 .filter(keep_float)
1005 .context(DataFusionPlanningSnafu)?
1006 .build()
1007 .context(DataFusionPlanningSnafu)?;
1008 self.ctx.field_columns = vec![float_column];
1009 }
1010
1011 if self.all_field_columns_are_native_histograms(input.schema()) {
1012 let promql_annotations = self.promql_annotations.clone();
1013 let input = self.projection_for_each_field_column(input, |col| {
1014 Ok(DfExpr::ScalarFunction(ScalarFunction {
1015 func: Arc::new(NativeHistogramDrop::float_null_udf(
1016 format!(
1017 "{}: dropped native histogram samples because this aggregation is not supported for native histograms",
1018 op
1019 ),
1020 promql_annotations.clone(),
1021 )),
1022 args: vec![DfExpr::Column(Column::from_name(col))],
1023 }))
1024 })?;
1025 return LogicalPlanBuilder::from(input)
1026 .filter(self.create_empty_values_filter_expr(false)?)
1027 .context(DataFusionPlanningSnafu)?
1028 .build()
1029 .context(DataFusionPlanningSnafu);
1030 }
1031
1032 let val = Self::get_param_as_literal_expr(
1033 param.as_deref(),
1034 Some(*op),
1035 Some(ArrowDataType::Float64),
1036 )?;
1037
1038 let window_exprs = self.create_window_exprs(*op, group_exprs.clone(), &input)?;
1040
1041 let rank_columns: Vec<_> = window_exprs
1042 .iter()
1043 .map(|expr| expr.schema_name().to_string())
1044 .collect();
1045
1046 let filter: DfExpr = rank_columns
1049 .iter()
1050 .fold(None, |expr, rank| {
1051 let predicate = DfExpr::BinaryExpr(BinaryExpr {
1052 left: Box::new(ident(rank)),
1053 op: Operator::LtEq,
1054 right: Box::new(val.clone()),
1055 });
1056
1057 match expr {
1058 None => Some(predicate),
1059 Some(expr) => Some(DfExpr::BinaryExpr(BinaryExpr {
1060 left: Box::new(expr),
1061 op: Operator::Or,
1062 right: Box::new(predicate),
1063 })),
1064 }
1065 })
1066 .unwrap();
1067
1068 let rank_columns: Vec<_> = rank_columns.into_iter().map(ident).collect();
1069
1070 let mut new_group_exprs = group_exprs.clone();
1071 new_group_exprs.extend(rank_columns);
1073
1074 let group_sort_expr = new_group_exprs
1075 .into_iter()
1076 .map(|expr| expr.sort(true, false));
1077
1078 let project_fields = self
1079 .create_field_column_exprs()?
1080 .into_iter()
1081 .chain(self.create_tag_column_exprs()?)
1082 .chain(
1083 self.ctx
1084 .use_tsid
1085 .then_some(DfExpr::Column(Column::from_name(
1086 DATA_SCHEMA_TSID_COLUMN_NAME,
1087 ))),
1088 )
1089 .chain(Some(self.create_time_index_column_expr()?));
1090
1091 LogicalPlanBuilder::from(input)
1092 .window(window_exprs)
1093 .context(DataFusionPlanningSnafu)?
1094 .filter(filter)
1095 .context(DataFusionPlanningSnafu)?
1096 .sort(group_sort_expr)
1097 .context(DataFusionPlanningSnafu)?
1098 .project(project_fields)
1099 .context(DataFusionPlanningSnafu)?
1100 .build()
1101 .context(DataFusionPlanningSnafu)
1102 }
1103
1104 async fn prom_unary_expr_to_plan(
1105 &mut self,
1106 query_engine_state: &QueryEngineState,
1107 unary_expr: &UnaryExpr,
1108 ) -> Result<LogicalPlan> {
1109 let UnaryExpr { expr } = unary_expr;
1110 let input = self.prom_expr_to_plan(expr, query_engine_state).await?;
1112 self.negate_field_columns(input)
1113 }
1114
1115 fn negate_field_columns(&mut self, input: LogicalPlan) -> Result<LogicalPlan> {
1116 let input_schema = input.schema().clone();
1117 self.projection_for_each_field_column(input, |col| {
1118 if Self::field_column_is_native_histogram(&input_schema, col) {
1119 Ok(DfExpr::ScalarFunction(ScalarFunction {
1120 func: Arc::new(NativeHistogramNeg::scalar_udf()),
1121 args: vec![DfExpr::Column(Column::from_name(col))],
1122 }))
1123 } else {
1124 Ok(DfExpr::Negative(Box::new(DfExpr::Column(
1125 Column::from_name(col),
1126 ))))
1127 }
1128 })
1129 }
1130
1131 async fn prom_binary_expr_to_plan(
1132 &mut self,
1133 query_engine_state: &QueryEngineState,
1134 binary_expr: &PromBinaryExpr,
1135 ) -> Result<LogicalPlan> {
1136 if let Some(modifier) = &binary_expr.modifier {
1140 ensure!(
1141 modifier.fill_values.lhs.is_none() && modifier.fill_values.rhs.is_none(),
1142 UnsupportedExprSnafu {
1143 name: "PromQL fill modifiers"
1144 }
1145 );
1146 }
1147
1148 if let Some(plan) = self.try_plan_binary_island(binary_expr).await? {
1149 return Ok(plan);
1150 }
1151
1152 let PromBinaryExpr {
1153 lhs,
1154 rhs,
1155 op,
1156 modifier,
1157 } = binary_expr;
1158
1159 let should_return_bool = if let Some(m) = modifier {
1162 m.return_bool
1163 } else {
1164 false
1165 };
1166 let is_comparison_op = Self::is_token_a_comparison_op(*op);
1167
1168 match (
1171 Self::try_build_literal_expr(lhs),
1172 Self::try_build_literal_expr(rhs),
1173 ) {
1174 (Some(lhs), Some(rhs)) => {
1175 self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
1176 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1177 self.ctx.reset_table_name_and_schema();
1178 let field_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1179 let mut field_expr = field_expr_builder(lhs, rhs)?;
1180
1181 if is_comparison_op && should_return_bool {
1182 field_expr =
1183 DfExpr::Cast(Cast::new(Box::new(field_expr), ArrowDataType::Float64));
1184 }
1185
1186 Ok(LogicalPlan::Extension(Extension {
1187 node: Arc::new(
1188 EmptyMetric::new(
1189 self.ctx.start,
1190 self.ctx.end,
1191 self.ctx.interval,
1192 SPECIAL_TIME_FUNCTION.to_string(),
1193 DEFAULT_FIELD_COLUMN.to_string(),
1194 Some(field_expr),
1195 )
1196 .context(DataFusionPlanningSnafu)?,
1197 ),
1198 }))
1199 }
1200 (Some(mut expr), None) => {
1202 let input = self.prom_expr_to_plan(rhs, query_engine_state).await?;
1203 if let Some(time_expr) = self.try_build_special_time_expr_with_context(lhs) {
1207 expr = time_expr
1208 }
1209 let input_schema = input.schema().clone();
1210 let preserve_any_value = Self::field_columns_are_alternative_samples(
1211 &input_schema,
1212 &self.ctx.field_columns,
1213 );
1214 let has_native_histogram = Self::field_columns_contain_native_histogram(
1215 &input_schema,
1216 &self.ctx.field_columns,
1217 );
1218 let retain_field_columns = self
1219 .ctx
1220 .field_columns
1221 .iter()
1222 .map(|col| {
1223 Self::binary_result_is_histogram(
1224 *op,
1225 false,
1226 Self::field_column_is_native_histogram(&input_schema, col),
1227 )
1228 .is_some()
1229 })
1230 .collect();
1231 let promql_annotations = self.promql_annotations.clone();
1232 let bin_expr_builder = |col: &String| {
1233 let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1234 let rhs_is_histogram =
1235 Self::field_column_is_native_histogram(&input_schema, col);
1236 let rhs = DfExpr::Column(Column::from_name(col));
1237 let mut binary_expr = match Self::native_histogram_binary_expr(
1238 *op,
1239 expr.clone(),
1240 false,
1241 rhs.clone(),
1242 rhs_is_histogram,
1243 is_comparison_op && !should_return_bool,
1244 promql_annotations.clone(),
1245 )? {
1246 Some(expr) => expr,
1247 None => binary_expr_builder(expr.clone(), rhs)?,
1248 };
1249
1250 if is_comparison_op && should_return_bool {
1251 binary_expr =
1252 DfExpr::Cast(Cast::new(Box::new(binary_expr), ArrowDataType::Float64));
1253 }
1254 Ok(binary_expr)
1255 };
1256 if is_comparison_op && !should_return_bool {
1257 self.filter_on_field_column(input, bin_expr_builder)
1258 } else {
1259 let projected =
1260 self.projection_for_each_field_column(input, bin_expr_builder)?;
1261 self.filter_binary_projection(
1262 projected,
1263 has_native_histogram,
1264 preserve_any_value,
1265 retain_field_columns,
1266 )
1267 }
1268 }
1269 (None, Some(mut expr)) => {
1271 let input = self.prom_expr_to_plan(lhs, query_engine_state).await?;
1272 if let Some(time_expr) = self.try_build_special_time_expr_with_context(rhs) {
1274 expr = time_expr
1275 }
1276 let input_schema = input.schema().clone();
1277 let preserve_any_value = Self::field_columns_are_alternative_samples(
1278 &input_schema,
1279 &self.ctx.field_columns,
1280 );
1281 let has_native_histogram = Self::field_columns_contain_native_histogram(
1282 &input_schema,
1283 &self.ctx.field_columns,
1284 );
1285 let retain_field_columns = self
1286 .ctx
1287 .field_columns
1288 .iter()
1289 .map(|col| {
1290 Self::binary_result_is_histogram(
1291 *op,
1292 Self::field_column_is_native_histogram(&input_schema, col),
1293 false,
1294 )
1295 .is_some()
1296 })
1297 .collect();
1298 let promql_annotations = self.promql_annotations.clone();
1299 let bin_expr_builder = |col: &String| {
1300 let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1301 let lhs_is_histogram =
1302 Self::field_column_is_native_histogram(&input_schema, col);
1303 let lhs = DfExpr::Column(Column::from_name(col));
1304 let mut binary_expr = match Self::native_histogram_binary_expr(
1305 *op,
1306 lhs.clone(),
1307 lhs_is_histogram,
1308 expr.clone(),
1309 false,
1310 is_comparison_op && !should_return_bool,
1311 promql_annotations.clone(),
1312 )? {
1313 Some(expr) => expr,
1314 None => binary_expr_builder(lhs, expr.clone())?,
1315 };
1316
1317 if is_comparison_op && should_return_bool {
1318 binary_expr =
1319 DfExpr::Cast(Cast::new(Box::new(binary_expr), ArrowDataType::Float64));
1320 }
1321 Ok(binary_expr)
1322 };
1323 if is_comparison_op && !should_return_bool {
1324 self.filter_on_field_column(input, bin_expr_builder)
1325 } else {
1326 let projected =
1327 self.projection_for_each_field_column(input, bin_expr_builder)?;
1328 self.filter_binary_projection(
1329 projected,
1330 has_native_histogram,
1331 preserve_any_value,
1332 retain_field_columns,
1333 )
1334 }
1335 }
1336 (None, None) => {
1338 let mut left_input = self.prom_expr_to_plan(lhs, query_engine_state).await?;
1339 let left_field_columns = self.ctx.field_columns.clone();
1340 let left_time_index_column = self.ctx.time_index_column.clone();
1341 let mut left_table_ref = self
1342 .table_ref()
1343 .unwrap_or_else(|_| TableReference::bare(""));
1344 let mut left_context = self.ctx.clone();
1345
1346 let mut right_input = self.prom_expr_to_plan(rhs, query_engine_state).await?;
1347 let right_field_columns = self.ctx.field_columns.clone();
1348 let right_time_index_column = self.ctx.time_index_column.clone();
1349 let mut right_table_ref = self
1350 .table_ref()
1351 .unwrap_or_else(|_| TableReference::bare(""));
1352 let mut right_context = self.ctx.clone();
1353
1354 if let Some(rewritten) = matching_filters::propagate(
1358 binary_expr,
1359 &left_context.tag_columns,
1360 &right_context.tag_columns,
1361 &left_context.aggregation_field_labels,
1362 &right_context.aggregation_field_labels,
1363 ) {
1364 if rewritten.lhs.as_ref() != lhs.as_ref() {
1368 let selectors = std::mem::take(&mut left_context.selector_matcher);
1369 left_input = self
1370 .prom_expr_to_plan(&rewritten.lhs, query_engine_state)
1371 .await?;
1372 left_context = self.ctx.clone();
1373 left_context.selector_matcher = selectors;
1374 }
1375 if rewritten.rhs.as_ref() != rhs.as_ref() {
1376 let selectors = std::mem::take(&mut right_context.selector_matcher);
1377 right_input = self
1378 .prom_expr_to_plan(&rewritten.rhs, query_engine_state)
1379 .await?;
1380 right_context = self.ctx.clone();
1381 right_context.selector_matcher = selectors;
1382 }
1383 self.ctx = right_context.clone();
1385 }
1386
1387 let left_is_empty_metric = Self::is_empty_metric(&left_input);
1388 let right_is_empty_metric = Self::is_empty_metric(&right_input);
1389
1390 if Self::is_token_a_set_op(*op) {
1394 return self.set_op_on_non_field_columns(
1395 left_input,
1396 right_input,
1397 left_context,
1398 right_context,
1399 *op,
1400 modifier,
1401 );
1402 }
1403
1404 let has_native_histogram = Self::field_columns_contain_native_histogram(
1405 left_input.schema(),
1406 &left_field_columns,
1407 ) || Self::field_columns_contain_native_histogram(
1408 right_input.schema(),
1409 &right_field_columns,
1410 );
1411
1412 if left_table_ref == right_table_ref {
1414 left_table_ref = TableReference::bare("lhs");
1416 right_table_ref = TableReference::bare("rhs");
1417 if self.ctx.tag_columns.is_empty() {
1423 self.ctx = left_context.clone();
1424 self.ctx.table_name = Some("lhs".to_string());
1425 } else {
1426 self.ctx.table_name = Some("rhs".to_string());
1427 }
1428 } else if right_is_empty_metric && !left_is_empty_metric {
1429 self.ctx = left_context.clone();
1430 } else if rhs.value_type() == ValueType::Scalar
1431 && lhs.value_type() == ValueType::Vector
1432 {
1433 self.ctx = left_context.clone();
1436 }
1437 if self.ctx.table_name.is_none() {
1441 self.ctx.table_name = Some(String::new());
1442 }
1443 let broadcast_scalar = !is_comparison_op;
1446 let (field_groups, invalid_field_pairs) = Self::align_binary_field_columns(
1447 left_input.schema(),
1448 right_input.schema(),
1449 &left_field_columns,
1450 &right_field_columns,
1451 *op,
1452 broadcast_scalar && lhs.value_type() == ValueType::Scalar,
1453 broadcast_scalar && rhs.value_type() == ValueType::Scalar,
1454 );
1455 let left_aligned_field_columns = field_groups
1456 .iter()
1457 .flat_map(|(_, pairs)| {
1458 pairs
1459 .iter()
1460 .map(|(left_col_name, _)| (*left_col_name).clone())
1461 })
1462 .collect::<Vec<_>>();
1463 let right_aligned_field_columns = field_groups
1464 .iter()
1465 .flat_map(|(_, pairs)| {
1466 pairs
1467 .iter()
1468 .map(|(_, right_col_name)| (*right_col_name).clone())
1469 })
1470 .collect::<Vec<_>>();
1471 self.ctx.field_columns = field_groups
1474 .iter()
1475 .map(|(output, _)| output.clone())
1476 .collect();
1477 let mut field_groups = field_groups.into_iter();
1478 let has_empty_metric_operand = left_is_empty_metric || right_is_empty_metric;
1480
1481 let only_join_time_index = lhs.value_type() == ValueType::Scalar
1482 || rhs.value_type() == ValueType::Scalar
1483 || has_empty_metric_operand
1484 || ((left_context.tag_columns.is_empty()
1485 || right_context.tag_columns.is_empty())
1486 && !left_context
1487 .tag_columns
1488 .iter()
1489 .chain(&right_context.tag_columns)
1490 .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL));
1491 let join_plan = self.join_on_non_field_columns(
1492 left_input,
1493 right_input,
1494 left_table_ref.clone(),
1495 right_table_ref.clone(),
1496 left_time_index_column,
1497 right_time_index_column,
1498 only_join_time_index,
1499 modifier,
1500 &left_context,
1501 &right_context,
1502 )?;
1503 let join_plan_schema = join_plan.schema().clone();
1504 let result_labels = if only_join_time_index {
1508 None
1509 } else {
1510 Self::binary_result_labels(&left_context, &right_context, modifier)
1511 .map(|labels| {
1512 Self::binary_result_label_projection(
1513 &join_plan_schema,
1514 &left_table_ref,
1515 &right_table_ref,
1516 &left_context,
1517 &right_context,
1518 labels,
1519 )
1520 })
1521 .transpose()?
1522 };
1523 if let Some(labels) = &result_labels {
1524 labels.apply(&mut self.ctx);
1525 }
1526 let promql_annotations = self.promql_annotations.clone();
1527 let invalid_pair_predicates = invalid_field_pairs
1530 .into_iter()
1531 .filter(|_| promql_annotations.is_some())
1532 .map(|(left_col_name, right_col_name)| {
1533 let left_field = join_plan_schema
1534 .qualified_field_with_name(Some(&left_table_ref), left_col_name)
1535 .context(DataFusionPlanningSnafu)?;
1536 let right_field = join_plan_schema
1537 .qualified_field_with_name(Some(&right_table_ref), right_col_name)
1538 .context(DataFusionPlanningSnafu)?;
1539 let left_is_histogram =
1540 left_field.1.data_type() == &Self::native_histogram_arrow_type();
1541 let right_is_histogram =
1542 right_field.1.data_type() == &Self::native_histogram_arrow_type();
1543 let drop_expr = Self::native_histogram_binary_expr(
1544 *op,
1545 DfExpr::Column(left_field.into()),
1546 left_is_histogram,
1547 DfExpr::Column(right_field.into()),
1548 right_is_histogram,
1549 true,
1550 promql_annotations.clone(),
1551 )?
1552 .with_context(|| UnexpectedPlanExprSnafu {
1553 desc: "invalid native histogram pair produced no drop expression",
1554 })?;
1555 Ok(DfExpr::Not(Box::new(drop_expr)))
1556 })
1557 .collect::<Result<Vec<_>>>()?;
1558 let join_plan = if let Some(predicate) = conjunction(invalid_pair_predicates) {
1559 LogicalPlanBuilder::from(join_plan)
1560 .filter(predicate)
1561 .context(DataFusionPlanningSnafu)?
1562 .build()
1563 .context(DataFusionPlanningSnafu)?
1564 } else {
1565 join_plan
1566 };
1567
1568 let bin_expr_builder = |_: &String| {
1569 let (_, field_pairs) =
1570 field_groups
1571 .next()
1572 .with_context(|| UnexpectedPlanExprSnafu {
1573 desc: "missing binary field group",
1574 })?;
1575 let binary_exprs = field_pairs
1576 .into_iter()
1577 .map(|(left_col_name, right_col_name)| {
1578 let left_field = join_plan_schema
1579 .qualified_field_with_name(Some(&left_table_ref), left_col_name)
1580 .context(DataFusionPlanningSnafu)?;
1581 let right_field = join_plan_schema
1582 .qualified_field_with_name(Some(&right_table_ref), right_col_name)
1583 .context(DataFusionPlanningSnafu)?;
1584 let left_is_histogram =
1585 left_field.1.data_type() == &Self::native_histogram_arrow_type();
1586 let right_is_histogram =
1587 right_field.1.data_type() == &Self::native_histogram_arrow_type();
1588 let left_col = left_field.into();
1589 let right_col = right_field.into();
1590
1591 let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1592 let lhs = DfExpr::Column(left_col);
1593 let rhs = DfExpr::Column(right_col);
1594 let mut binary_expr = match Self::native_histogram_binary_expr(
1595 *op,
1596 lhs.clone(),
1597 left_is_histogram,
1598 rhs.clone(),
1599 right_is_histogram,
1600 is_comparison_op && !should_return_bool,
1601 promql_annotations.clone(),
1602 )? {
1603 Some(expr) => expr,
1604 None => binary_expr_builder(lhs, rhs)?,
1605 };
1606 if is_comparison_op && should_return_bool {
1607 binary_expr = DfExpr::Cast(Cast::new(
1608 Box::new(binary_expr),
1609 ArrowDataType::Float64,
1610 ));
1611 }
1612 Ok(binary_expr)
1613 })
1614 .collect::<Result<Vec<_>>>()?;
1615 if let [binary_expr] = binary_exprs.as_slice() {
1616 Ok(binary_expr.clone())
1617 } else {
1618 Ok(DfExpr::ScalarFunction(ScalarFunction {
1619 func: coalesce(),
1620 args: binary_exprs,
1621 }))
1622 }
1623 };
1624 if is_comparison_op && !should_return_bool {
1625 let filtered = self.filter_on_field_column(join_plan, bin_expr_builder)?;
1632 let (project_table_ref, mut project_context, project_field_columns) =
1633 match (lhs.value_type(), rhs.value_type()) {
1634 (ValueType::Scalar, ValueType::Vector) => (
1635 &right_table_ref,
1636 right_context.clone(),
1637 right_aligned_field_columns,
1638 ),
1639 _ => (
1640 &left_table_ref,
1641 left_context.clone(),
1642 left_aligned_field_columns,
1643 ),
1644 };
1645 project_context.field_columns = project_field_columns;
1646 self.project_binary_join_side(
1647 filtered,
1648 project_table_ref,
1649 &project_context,
1650 result_labels.as_ref(),
1651 )
1652 } else {
1653 let projected = self.projection_for_each_field_column_with_labels(
1654 join_plan,
1655 result_labels.as_ref(),
1656 bin_expr_builder,
1657 )?;
1658 let preserve_any_value = Self::field_columns_are_alternative_samples(
1659 projected.schema(),
1660 &self.ctx.field_columns,
1661 );
1662 let retain_field_columns = vec![true; self.ctx.field_columns.len()];
1663 self.filter_binary_projection(
1664 projected,
1665 has_native_histogram,
1666 preserve_any_value,
1667 retain_field_columns,
1668 )
1669 }
1670 }
1671 }
1672 }
1673
1674 fn filter_binary_projection(
1675 &mut self,
1676 input: LogicalPlan,
1677 has_native_histogram: bool,
1678 preserve_any_value: bool,
1679 retain_field_columns: Vec<bool>,
1680 ) -> Result<LogicalPlan> {
1681 if !has_native_histogram {
1682 return Ok(input);
1683 }
1684
1685 ensure!(
1686 retain_field_columns.len() == self.ctx.field_columns.len(),
1687 UnexpectedPlanExprSnafu {
1688 desc: "binary output field count changed unexpectedly",
1689 }
1690 );
1691
1692 let filtered = LogicalPlanBuilder::from(input)
1693 .filter(self.create_empty_values_filter_expr(preserve_any_value)?)
1694 .context(DataFusionPlanningSnafu)?
1695 .build()
1696 .context(DataFusionPlanningSnafu)?;
1697 if retain_field_columns.iter().all(|retain| *retain) {
1698 return Ok(filtered);
1699 }
1700
1701 let retained = self
1702 .ctx
1703 .field_columns
1704 .iter()
1705 .zip(retain_field_columns)
1706 .filter(|(_, retain)| *retain)
1707 .map(|(field, _)| field.clone())
1708 .collect::<Vec<_>>();
1709 if retained.is_empty() {
1710 return Ok(filtered);
1711 }
1712 self.ctx.field_columns = retained;
1713
1714 let mut output_columns = self
1715 .ctx
1716 .field_columns
1717 .iter()
1718 .chain(&self.ctx.tag_columns)
1719 .cloned()
1720 .collect::<HashSet<_>>();
1721 output_columns.extend(self.ctx.time_index_column.iter().cloned());
1722 if self.ctx.use_tsid {
1723 output_columns.insert(DATA_SCHEMA_TSID_COLUMN_NAME.to_string());
1724 }
1725 let project_exprs = filtered
1726 .schema()
1727 .iter()
1728 .filter(|(_, field)| output_columns.contains(field.name()))
1729 .map(|(qualifier, field)| {
1730 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
1731 })
1732 .collect::<Vec<_>>();
1733 LogicalPlanBuilder::from(filtered)
1734 .project(project_exprs)
1735 .context(DataFusionPlanningSnafu)?
1736 .build()
1737 .context(DataFusionPlanningSnafu)
1738 }
1739
1740 fn project_binary_join_side(
1741 &mut self,
1742 input: LogicalPlan,
1743 table_ref: &TableReference,
1744 context: &PromPlannerContext,
1745 result_labels: Option<&BinaryResultLabels>,
1746 ) -> Result<LogicalPlan> {
1747 let schema = input.schema();
1748
1749 let mut project_exprs =
1750 Vec::with_capacity(context.tag_columns.len() + context.field_columns.len() + 2);
1751
1752 if let Some(time_index_column) = &context.time_index_column {
1754 let time_index_col = schema
1755 .qualified_field_with_name(Some(table_ref), time_index_column)
1756 .context(DataFusionPlanningSnafu)?
1757 .into();
1758 project_exprs.push(DfExpr::Column(time_index_col));
1759 }
1760
1761 for field_column in &context.field_columns {
1763 let field_col = schema
1764 .qualified_field_with_name(Some(table_ref), field_column)
1765 .context(DataFusionPlanningSnafu)?
1766 .into();
1767 project_exprs.push(DfExpr::Column(field_col));
1768 }
1769
1770 match result_labels {
1772 Some(labels) => project_exprs.extend(labels.exprs.iter().cloned()),
1773 None => {
1774 for tag_column in &context.tag_columns {
1775 let tag_col = schema
1776 .qualified_field_with_name(Some(table_ref), tag_column)
1777 .context(DataFusionPlanningSnafu)?
1778 .into();
1779 project_exprs.push(DfExpr::Column(tag_col));
1780 }
1781 }
1782 }
1783
1784 let tsid_col = match result_labels {
1787 Some(labels) => labels.tsid_projection(Some(table_ref.clone())),
1788 None => Self::optional_tsid_projection(schema, Some(table_ref), context.use_tsid),
1789 };
1790 if let Some(tsid_col) = tsid_col {
1791 project_exprs.push(tsid_col);
1792 }
1793
1794 let plan = LogicalPlanBuilder::from(input)
1795 .project(project_exprs)
1796 .context(DataFusionPlanningSnafu)?
1797 .build()
1798 .context(DataFusionPlanningSnafu)?;
1799
1800 self.ctx = context.clone();
1803 self.ctx.table_name = None;
1804 self.ctx.schema_name = None;
1805 if let Some(labels) = result_labels {
1806 labels.apply(&mut self.ctx);
1807 }
1808
1809 Ok(plan)
1810 }
1811
1812 fn prom_number_lit_to_plan(&mut self, number_literal: &NumberLiteral) -> Result<LogicalPlan> {
1813 let NumberLiteral { val } = number_literal;
1814 self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
1815 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1816 self.ctx.reset_table_name_and_schema();
1817 let literal_expr = df_prelude::lit(*val);
1818
1819 let plan = LogicalPlan::Extension(Extension {
1820 node: Arc::new(
1821 EmptyMetric::new(
1822 self.ctx.start,
1823 self.ctx.end,
1824 self.ctx.interval,
1825 SPECIAL_TIME_FUNCTION.to_string(),
1826 DEFAULT_FIELD_COLUMN.to_string(),
1827 Some(literal_expr),
1828 )
1829 .context(DataFusionPlanningSnafu)?,
1830 ),
1831 });
1832 Ok(plan)
1833 }
1834
1835 fn prom_string_lit_to_plan(&mut self, string_literal: &StringLiteral) -> Result<LogicalPlan> {
1836 let StringLiteral { val } = string_literal;
1837 self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
1838 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1839 self.ctx.reset_table_name_and_schema();
1840 let literal_expr = df_prelude::lit(val.clone());
1841
1842 let plan = LogicalPlan::Extension(Extension {
1843 node: Arc::new(
1844 EmptyMetric::new(
1845 self.ctx.start,
1846 self.ctx.end,
1847 self.ctx.interval,
1848 SPECIAL_TIME_FUNCTION.to_string(),
1849 DEFAULT_FIELD_COLUMN.to_string(),
1850 Some(literal_expr),
1851 )
1852 .context(DataFusionPlanningSnafu)?,
1853 ),
1854 });
1855 Ok(plan)
1856 }
1857
1858 fn offset_millis(offset: &Option<Offset>) -> Millisecond {
1861 match offset {
1862 Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond,
1863 Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond),
1864 None => 0,
1865 }
1866 }
1867
1868 fn subquery_child_window(
1872 start: Millisecond,
1873 end: Millisecond,
1874 interval: Millisecond,
1875 offset: Millisecond,
1876 range: Millisecond,
1877 step: Millisecond,
1878 ) -> Result<Option<(Millisecond, Millisecond)>> {
1879 let window_start = i128::from(start) - i128::from(offset) - i128::from(range);
1880 let grid_start = (window_start.div_euclid(i128::from(step)) + 1) * i128::from(step);
1881 let grid_end = if interval > 0 {
1882 let steps = (i128::from(end) - i128::from(start)) / i128::from(interval);
1883 i128::from(start) + steps * i128::from(interval)
1884 } else {
1885 i128::from(end)
1886 } - i128::from(offset);
1887 if grid_start > grid_end {
1889 return Ok(None);
1890 }
1891 let bound = |name: &str, value: i128| {
1892 i64::try_from(value).map_err(|_| {
1893 SubqueryTimestampOutOfRangeSnafu {
1894 timestamp: format!(
1895 "subquery child window {name} = {value} ms (start = {start} ms, \
1896 end = {end} ms, interval = {interval} ms, offset = {offset} ms, \
1897 range = {range} ms, step = {step} ms)"
1898 ),
1899 }
1900 .build()
1901 })
1902 };
1903 Ok(Some((
1904 bound("grid start", grid_start)?,
1905 bound("grid end", grid_end)?,
1906 )))
1907 }
1908
1909 fn series_key_columns(&self) -> Vec<String> {
1912 if self.ctx.use_tsid {
1913 vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]
1914 } else {
1915 self.ctx.tag_columns.clone()
1916 }
1917 }
1918
1919 fn series_key_columns_for_schema(&self, schema: &DFSchemaRef) -> Vec<String> {
1921 let has_tsid = schema.fields().iter().any(|field| {
1922 field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
1923 && field.data_type() == &ArrowDataType::UInt64
1924 });
1925 if has_tsid {
1926 vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]
1927 } else {
1928 self.ctx
1929 .tag_columns
1930 .iter()
1931 .filter(|name| schema.has_column_with_unqualified_name(name))
1932 .cloned()
1933 .collect()
1934 }
1935 }
1936
1937 async fn prom_vector_selector_to_plan(
1938 &mut self,
1939 vector_selector: &VectorSelector,
1940 timestamp_fn: bool,
1941 ) -> Result<LogicalPlan> {
1942 let VectorSelector {
1943 name,
1944 offset,
1945 matchers,
1946 at,
1947 } = vector_selector;
1948 let matchers = self.preprocess_label_matchers(matchers, name)?;
1949 if let Some(empty_plan) = self.setup_context().await? {
1950 return Ok(empty_plan);
1951 }
1952 let offset_ms = Self::offset_millis(offset);
1953 let at_offset = self.at_modifier_offset(at, offset)?;
1957 let grid_start = self.ctx.start;
1958 let grid_end = self.ctx.end;
1959 let normalize = match at_offset {
1960 Some(at_offset) => {
1961 self.ctx.end = grid_start;
1966 let plan = self
1967 .selector_to_series_normalize_plan(at_offset, matchers, false)
1968 .await?;
1969 self.ctx.end = grid_end;
1970 plan
1971 }
1972 None => {
1973 self.selector_to_series_normalize_plan(offset_ms, matchers, false)
1974 .await?
1975 }
1976 };
1977 let time_index_column =
1978 self.ctx
1979 .time_index_column
1980 .clone()
1981 .with_context(|| TimeIndexNotFoundSnafu {
1982 table: self.ctx.table_name.clone().unwrap_or_default(),
1983 })?;
1984
1985 let (normalize, timestamp_value_column) = if timestamp_fn {
1986 let occupied = normalize
1989 .schema()
1990 .fields()
1991 .iter()
1992 .map(|field| field.name().as_str())
1993 .collect::<HashSet<_>>();
1994 let mut timestamp_value_column = TIMESTAMP_VALUE_PREFIX.to_string();
1995 while occupied.contains(timestamp_value_column.as_str()) {
1996 timestamp_value_column.push('_');
1997 }
1998 let mut project_exprs = normalize
1999 .schema()
2000 .iter()
2001 .map(|(qualifier, field)| {
2002 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
2003 })
2004 .collect::<Vec<_>>();
2005 let sample_time = Self::timestamp_seconds_expr(&time_index_column, normalize.schema())?;
2008 project_exprs.push(sample_time.alias(×tamp_value_column));
2009 let normalize = LogicalPlanBuilder::from(normalize)
2010 .project(project_exprs)
2011 .context(DataFusionPlanningSnafu)?
2012 .build()
2013 .context(DataFusionPlanningSnafu)?;
2014 (normalize, Some(timestamp_value_column))
2015 } else {
2016 (normalize, None)
2017 };
2018
2019 let field_column = self.ctx.field_columns.first().cloned();
2020 let series_key_columns = self.series_key_columns();
2021 let manipulate = match at_offset {
2022 Some(at_offset) => {
2023 let anchored = InstantManipulate::new(
2027 grid_start,
2028 grid_start,
2029 self.ctx.lookback_delta,
2030 self.ctx.interval,
2031 at_offset,
2032 time_index_column.clone(),
2033 series_key_columns,
2034 field_column,
2035 normalize,
2036 );
2037 self.replay_over_grid(
2038 LogicalPlan::Extension(Extension {
2039 node: Arc::new(anchored),
2040 }),
2041 grid_start,
2042 grid_end,
2043 time_index_column,
2044 )
2045 }
2046 None => LogicalPlan::Extension(Extension {
2047 node: Arc::new(InstantManipulate::new(
2048 grid_start,
2049 grid_end,
2050 self.ctx.lookback_delta,
2051 self.ctx.interval,
2052 offset_ms,
2053 time_index_column,
2054 series_key_columns,
2055 field_column,
2056 normalize,
2057 )),
2058 }),
2059 };
2060 if let Some(timestamp_value_column) = timestamp_value_column {
2061 self.create_timestamp_func_plan(manipulate, ident(timestamp_value_column))
2062 } else {
2063 Ok(manipulate)
2064 }
2065 }
2066
2067 fn timestamp_seconds_expr(column: &str, schema: &DFSchema) -> Result<DfExpr> {
2069 let column = DfExpr::Column(Column::from_name(column));
2070 let ArrowDataType::Timestamp(_, timezone) =
2071 column.get_type(schema).context(DataFusionPlanningSnafu)?
2072 else {
2073 unreachable!("time index is a timestamp")
2074 };
2075 let millis = column
2076 .cast_to(
2077 &ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, timezone),
2078 schema,
2079 )
2080 .context(DataFusionPlanningSnafu)?
2081 .cast_to(&ArrowDataType::Int64, schema)
2082 .context(DataFusionPlanningSnafu)?
2083 .cast_to(&ArrowDataType::Float64, schema)
2084 .context(DataFusionPlanningSnafu)?;
2085 Ok(DfExpr::BinaryExpr(BinaryExpr {
2086 left: Box::new(millis),
2087 op: Operator::Divide,
2088 right: Box::new(lit(1000.0)),
2089 }))
2090 }
2091
2092 fn create_timestamp_func_plan(
2098 &mut self,
2099 input: LogicalPlan,
2100 timestamp_value: DfExpr,
2101 ) -> Result<LogicalPlan> {
2102 let has_sample = self.create_empty_values_filter_expr(true)?;
2105 let time_expr = timestamp_value.alias(DEFAULT_FIELD_COLUMN);
2106 self.ctx.field_columns = vec![time_expr.schema_name().to_string()];
2107 let mut project_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 2);
2108 project_exprs.push(self.create_time_index_column_expr()?);
2109 project_exprs.push(time_expr);
2110 project_exprs.extend(self.create_tag_column_exprs()?);
2111
2112 LogicalPlanBuilder::from(input)
2113 .filter(has_sample)
2114 .context(DataFusionPlanningSnafu)?
2115 .project(project_exprs)
2116 .context(DataFusionPlanningSnafu)?
2117 .build()
2118 .context(DataFusionPlanningSnafu)
2119 }
2120
2121 async fn prom_matrix_selector_to_plan(
2122 &mut self,
2123 matrix_selector: &MatrixSelector,
2124 ) -> Result<LogicalPlan> {
2125 let MatrixSelector { vs, range } = matrix_selector;
2126 let VectorSelector {
2127 name,
2128 offset,
2129 matchers,
2130 at,
2131 } = vs;
2132 let matchers = self.preprocess_label_matchers(matchers, name)?;
2133 ensure!(!range.is_zero(), ZeroRangeSelectorSnafu);
2134 let range_ms = range.as_millis() as _;
2135 self.ctx.range = Some(range_ms);
2136 let offset_ms = Self::offset_millis(offset);
2137
2138 let at_offset = self.at_modifier_offset(at, offset)?;
2141 let grid_start = self.ctx.start;
2142 let grid_end = self.ctx.end;
2143
2144 let (normalize, at_offset) = match self.setup_context().await? {
2147 Some(empty_plan) => {
2151 self.ctx.range_fold_offset = Some(offset_ms);
2152 (empty_plan, None)
2153 }
2154 None => {
2155 let normalize = match at_offset {
2156 Some(at_offset) => {
2157 self.ctx.end = grid_start;
2161 let plan = self
2162 .selector_to_series_normalize_plan(at_offset, matchers, true)
2163 .await?;
2164 self.ctx.end = grid_end;
2165 plan
2166 }
2167 None => {
2168 self.selector_to_series_normalize_plan(offset_ms, matchers, true)
2169 .await?
2170 }
2171 };
2172 self.ctx.range_fold_offset = Some(at_offset.unwrap_or(offset_ms));
2177 (normalize, at_offset)
2178 }
2179 };
2180 let time_index_column = self
2181 .ctx
2182 .time_index_column
2183 .clone()
2184 .expect("time index should be set in `setup_context`");
2185 let manipulate = match at_offset {
2186 Some(at_offset) => {
2187 let anchored = RangeManipulate::new(
2191 grid_start,
2192 grid_start,
2193 self.ctx.interval,
2194 at_offset,
2195 range_ms,
2197 time_index_column.clone(),
2198 self.ctx.field_columns.clone(),
2199 normalize,
2200 )
2201 .context(DataFusionPlanningSnafu)?;
2202 self.replay_over_grid(
2203 LogicalPlan::Extension(Extension {
2204 node: Arc::new(anchored),
2205 }),
2206 grid_start,
2207 grid_end,
2208 time_index_column,
2209 )
2210 }
2211 None => {
2212 let manipulate = RangeManipulate::new(
2213 grid_start,
2214 grid_end,
2215 self.ctx.interval,
2216 offset_ms,
2217 range_ms,
2219 time_index_column,
2220 self.ctx.field_columns.clone(),
2221 normalize,
2222 )
2223 .context(DataFusionPlanningSnafu)?;
2224
2225 LogicalPlan::Extension(Extension {
2226 node: Arc::new(manipulate),
2227 })
2228 }
2229 };
2230
2231 Ok(manipulate)
2232 }
2233
2234 async fn prom_call_expr_to_plan(
2235 &mut self,
2236 query_engine_state: &QueryEngineState,
2237 call_expr: &Call,
2238 ) -> Result<LogicalPlan> {
2239 let Call { func, args } = call_expr;
2240 match func.name {
2242 SPECIAL_HISTOGRAM_QUANTILE | SPECIAL_HISTOGRAM_FRACTION => {
2243 return self
2244 .create_histogram_plan(func.name, args, query_engine_state)
2245 .await;
2246 }
2247 SPECIAL_VECTOR_FUNCTION => return self.create_vector_plan(args).await,
2248 SCALAR_FUNCTION => return self.create_scalar_plan(args, query_engine_state).await,
2249 SPECIAL_ABSENT_FUNCTION => {
2250 return self.create_absent_plan(args, query_engine_state).await;
2251 }
2252 _ => {}
2253 }
2254
2255 let args = self.create_function_args(&args.args)?;
2257 let mut timestamp_arg = args.input.as_ref();
2260 while let Some(PromExpr::Paren(ParenExpr { expr })) = timestamp_arg {
2261 timestamp_arg = Some(expr);
2262 }
2263 let timestamp_of_selector =
2264 func.name == "timestamp" && matches!(timestamp_arg, Some(PromExpr::VectorSelector(_)));
2265 let input = if let Some(prom_expr) = &args.input {
2266 let input = self
2267 .prom_expr_to_plan_inner(prom_expr, timestamp_of_selector, query_engine_state)
2268 .await?;
2269 if func.name == "timestamp" && !timestamp_of_selector {
2270 let time_index_column =
2271 self.ctx
2272 .time_index_column
2273 .clone()
2274 .with_context(|| TimeIndexNotFoundSnafu {
2275 table: self.ctx.table_name.clone().unwrap_or_default(),
2276 })?;
2277 let eval_time = Self::timestamp_seconds_expr(&time_index_column, input.schema())?;
2278 self.create_timestamp_func_plan(input, eval_time)?
2279 } else {
2280 input
2281 }
2282 } else {
2283 self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
2284 self.ctx.reset_table_name_and_schema();
2285 self.ctx.tag_columns = vec![];
2286 self.ctx.aggregation_field_labels.clear();
2287 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
2288 LogicalPlan::Extension(Extension {
2289 node: Arc::new(
2290 EmptyMetric::new(
2291 self.ctx.start,
2292 self.ctx.end,
2293 self.ctx.interval,
2294 SPECIAL_TIME_FUNCTION.to_string(),
2295 DEFAULT_FIELD_COLUMN.to_string(),
2296 None,
2297 )
2298 .context(DataFusionPlanningSnafu)?,
2299 ),
2300 })
2301 };
2302 let range_fold_offset = self.ctx.range_fold_offset.take();
2307 let input_tag_columns = self.ctx.tag_columns.clone();
2308 let (mut func_exprs, new_tags) = self.create_function_expr(
2309 func,
2310 args.literals.clone(),
2311 input.schema(),
2312 query_engine_state,
2313 range_fold_offset,
2314 )?;
2315 func_exprs.insert(0, self.create_time_index_column_expr()?);
2316 func_exprs.extend_from_slice(&self.create_tag_column_exprs()?);
2317 if let Some(tsid_col) =
2318 Self::optional_tsid_projection(input.schema(), None, self.ctx.use_tsid)
2319 {
2320 func_exprs.push(tsid_col);
2321 }
2322
2323 let builder = LogicalPlanBuilder::from(input)
2329 .project(func_exprs)
2330 .context(DataFusionPlanningSnafu)?
2331 .filter(self.create_empty_values_filter_expr(true)?)
2332 .context(DataFusionPlanningSnafu)?;
2333
2334 let builder = match func.name {
2335 "sort" => builder
2336 .sort(self.create_field_columns_sort_exprs(true))
2337 .context(DataFusionPlanningSnafu)?,
2338 "sort_desc" => builder
2339 .sort(self.create_field_columns_sort_exprs(false))
2340 .context(DataFusionPlanningSnafu)?,
2341 "sort_by_label" => builder
2342 .sort(Self::create_sort_exprs_by_tags(
2343 func.name,
2344 args.literals,
2345 true,
2346 )?)
2347 .context(DataFusionPlanningSnafu)?,
2348 "sort_by_label_desc" => builder
2349 .sort(Self::create_sort_exprs_by_tags(
2350 func.name,
2351 args.literals,
2352 false,
2353 )?)
2354 .context(DataFusionPlanningSnafu)?,
2355
2356 _ => builder,
2357 };
2358
2359 let may_duplicate_label_sets = matches!(func.name, "label_join" | "label_replace")
2362 && new_tags.iter().any(|tag| input_tag_columns.contains(tag));
2363
2364 for tag in new_tags {
2367 self.ctx.tag_columns.push(tag);
2368 }
2369
2370 let mut plan = builder.build().context(DataFusionPlanningSnafu)?;
2371 if may_duplicate_label_sets {
2372 let labels = self.ctx.tag_columns.clone();
2373 plan = Self::assert_unique_match_group(
2374 plan,
2375 labels
2376 .iter()
2377 .map(|label| DfExpr::Column(Column::from_name(label)))
2378 .collect(),
2379 labels,
2380 self.create_time_index_column_expr()?,
2381 MatchGroupViolation::DuplicateLabelSet,
2382 )?;
2383 }
2384 common_telemetry::debug!("Created PromQL function plan: {plan:?} for {call_expr:?}");
2385
2386 Ok(plan)
2387 }
2388
2389 async fn prom_ext_expr_to_plan(
2390 &mut self,
2391 query_engine_state: &QueryEngineState,
2392 ext_expr: &promql_parser::parser::ast::Extension,
2393 ) -> Result<LogicalPlan> {
2394 let expr = &ext_expr.expr;
2396 let children = expr.children();
2397 let plan = self
2398 .prom_root_expr_to_plan(&children[0], query_engine_state)
2399 .await?;
2400 match expr.name() {
2406 ANALYZE_NODE_NAME => LogicalPlanBuilder::from(plan)
2407 .explain(false, true)
2408 .unwrap()
2409 .build()
2410 .context(DataFusionPlanningSnafu),
2411 ANALYZE_VERBOSE_NODE_NAME => LogicalPlanBuilder::from(plan)
2412 .explain(true, true)
2413 .unwrap()
2414 .build()
2415 .context(DataFusionPlanningSnafu),
2416 EXPLAIN_NODE_NAME => LogicalPlanBuilder::from(plan)
2417 .explain(false, false)
2418 .unwrap()
2419 .build()
2420 .context(DataFusionPlanningSnafu),
2421 EXPLAIN_VERBOSE_NODE_NAME => LogicalPlanBuilder::from(plan)
2422 .explain(true, false)
2423 .unwrap()
2424 .build()
2425 .context(DataFusionPlanningSnafu),
2426 ALIAS_NODE_NAME => {
2427 let alias = expr
2428 .as_any()
2429 .downcast_ref::<AliasExpr>()
2430 .context(UnexpectedPlanExprSnafu {
2431 desc: "Expected AliasExpr",
2432 })?
2433 .alias
2434 .clone();
2435 self.apply_alias(plan, alias)
2436 }
2437 _ => LogicalPlanBuilder::empty(true)
2438 .build()
2439 .context(DataFusionPlanningSnafu),
2440 }
2441 }
2442
2443 #[allow(clippy::mutable_key_type)]
2453 fn preprocess_label_matchers(
2454 &mut self,
2455 label_matchers: &Matchers,
2456 name: &Option<String>,
2457 ) -> Result<Matchers> {
2458 self.ctx.reset();
2459
2460 let metric_name;
2461 if let Some(name) = name.clone() {
2462 metric_name = Some(name);
2463 ensure!(
2464 label_matchers.find_matchers(METRIC_NAME).is_empty(),
2465 MultipleMetricMatchersSnafu
2466 );
2467 } else {
2468 let mut matches = label_matchers.find_matchers(METRIC_NAME);
2469 ensure!(!matches.is_empty(), NoMetricMatcherSnafu);
2470 ensure!(matches.len() == 1, MultipleMetricMatchersSnafu);
2471 ensure!(
2472 matches[0].op == MatchOp::Equal,
2473 UnsupportedMatcherOpSnafu {
2474 matcher_op: matches[0].op.to_string(),
2475 matcher: METRIC_NAME
2476 }
2477 );
2478 metric_name = matches.pop().map(|m| m.value);
2479 }
2480
2481 self.ctx.table_name = metric_name;
2482
2483 let mut matchers: Vec<Matcher> = Vec::with_capacity(label_matchers.matchers.len());
2486 for matcher in &label_matchers.matchers {
2487 if matcher.name == FIELD_COLUMN_MATCHER {
2489 self.ctx
2490 .field_column_matcher
2491 .get_or_insert_default()
2492 .push(matcher.clone());
2493 } else if matcher.name == SCHEMA_COLUMN_MATCHER || matcher.name == DB_COLUMN_MATCHER {
2494 ensure!(
2495 matcher.op == MatchOp::Equal,
2496 UnsupportedMatcherOpSnafu {
2497 matcher: matcher.name.clone(),
2498 matcher_op: matcher.op.to_string(),
2499 }
2500 );
2501 self.ctx.schema_name = Some(matcher.value.clone());
2502 } else if matcher.name != METRIC_NAME {
2503 self.ctx.selector_matcher.push(matcher.clone());
2504 if !matchers.contains(matcher) {
2505 matchers.push(matcher.clone());
2506 }
2507 }
2508 }
2509
2510 Ok(Matchers::new(matchers))
2511 }
2512
2513 async fn selector_to_series_normalize_plan(
2514 &mut self,
2515 offset_duration: Millisecond,
2516 label_matchers: Matchers,
2517 is_range_selector: bool,
2518 ) -> Result<LogicalPlan> {
2519 let table_ref = self.table_ref()?;
2521 let mut table_scan = self.create_table_scan_plan(table_ref.clone()).await?;
2522 let table_schema = table_scan.schema();
2523
2524 let mut scan_filters = Self::matchers_to_expr(label_matchers.clone(), table_schema)?;
2526 if let Some(time_index_filter) =
2527 self.build_time_index_filter(offset_duration, table_schema)?
2528 {
2529 scan_filters.push(time_index_filter);
2530 }
2531 if let Some(filter) = conjunction(scan_filters) {
2532 table_scan = LogicalPlanBuilder::from(table_scan)
2533 .filter(filter)
2534 .context(DataFusionPlanningSnafu)?
2535 .build()
2536 .context(DataFusionPlanningSnafu)?;
2537 }
2538
2539 if let Some(field_matchers) = &self.ctx.field_column_matcher {
2541 let col_set = self.ctx.field_columns.iter().collect::<HashSet<_>>();
2542 let mut result_set = HashSet::new();
2544 let mut reverse_set = HashSet::new();
2546 for matcher in field_matchers {
2547 match &matcher.op {
2548 MatchOp::Equal => {
2549 if col_set.contains(&matcher.value) {
2550 let _ = result_set.insert(matcher.value.clone());
2551 } else {
2552 return Err(ColumnNotFoundSnafu {
2553 col: matcher.value.clone(),
2554 }
2555 .build());
2556 }
2557 }
2558 MatchOp::NotEqual => {
2559 if col_set.contains(&matcher.value) {
2560 let _ = reverse_set.insert(matcher.value.clone());
2561 } else {
2562 return Err(ColumnNotFoundSnafu {
2563 col: matcher.value.clone(),
2564 }
2565 .build());
2566 }
2567 }
2568 MatchOp::Re(regex) => {
2569 for col in &self.ctx.field_columns {
2570 if regex.is_match(col) {
2571 let _ = result_set.insert(col.clone());
2572 }
2573 }
2574 }
2575 MatchOp::NotRe(regex) => {
2576 for col in &self.ctx.field_columns {
2577 if regex.is_match(col) {
2578 let _ = reverse_set.insert(col.clone());
2579 }
2580 }
2581 }
2582 }
2583 }
2584 if result_set.is_empty() {
2586 result_set = col_set.into_iter().cloned().collect();
2587 }
2588 for col in reverse_set {
2589 let _ = result_set.remove(&col);
2590 }
2591
2592 self.ctx.field_columns = self
2594 .ctx
2595 .field_columns
2596 .drain(..)
2597 .filter(|col| result_set.contains(col))
2598 .collect();
2599
2600 let exprs = result_set
2601 .into_iter()
2602 .map(|col| DfExpr::Column(Column::new_unqualified(col)))
2603 .chain(self.create_tag_column_exprs()?)
2604 .chain(
2605 self.ctx
2606 .use_tsid
2607 .then_some(DfExpr::Column(Column::new_unqualified(
2608 DATA_SCHEMA_TSID_COLUMN_NAME,
2609 ))),
2610 )
2611 .chain(Some(self.create_time_index_column_expr()?))
2612 .collect::<Vec<_>>();
2613
2614 table_scan = LogicalPlanBuilder::from(table_scan)
2616 .project(exprs)
2617 .context(DataFusionPlanningSnafu)?
2618 .build()
2619 .context(DataFusionPlanningSnafu)?;
2620 }
2621
2622 let series_key_columns = self.series_key_columns();
2624
2625 let sort_exprs = if self.ctx.use_tsid {
2626 vec![
2627 DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME)).sort(true, true),
2628 self.create_time_index_column_expr()?.sort(true, true),
2629 ]
2630 } else {
2631 self.create_tag_and_time_index_column_sort_exprs()?
2632 };
2633
2634 let sort_plan = LogicalPlanBuilder::from(table_scan)
2635 .sort(sort_exprs)
2636 .context(DataFusionPlanningSnafu)?
2637 .build()
2638 .context(DataFusionPlanningSnafu)?;
2639
2640 let time_index_column =
2642 self.ctx
2643 .time_index_column
2644 .clone()
2645 .with_context(|| TimeIndexNotFoundSnafu {
2646 table: table_ref.to_string(),
2647 })?;
2648 let divide_plan = LogicalPlan::Extension(Extension {
2649 node: Arc::new(SeriesDivide::new(
2650 series_key_columns.clone(),
2651 time_index_column,
2652 sort_plan,
2653 )),
2654 });
2655
2656 if !is_range_selector && offset_duration == 0 {
2658 return Ok(divide_plan);
2659 }
2660 let series_normalize = SeriesNormalize::new(
2661 offset_duration,
2662 self.ctx
2663 .time_index_column
2664 .clone()
2665 .with_context(|| TimeIndexNotFoundSnafu {
2666 table: table_ref.to_quoted_string(),
2667 })?,
2668 is_range_selector,
2669 series_key_columns,
2670 divide_plan,
2671 );
2672 let logical_plan = LogicalPlan::Extension(Extension {
2673 node: Arc::new(series_normalize),
2674 });
2675
2676 Ok(logical_plan)
2677 }
2678
2679 fn agg_modifier_to_col(
2686 &mut self,
2687 input_schema: &DFSchemaRef,
2688 modifier: &Option<LabelModifier>,
2689 update_ctx: bool,
2690 ) -> Result<Vec<DfExpr>> {
2691 match modifier {
2692 None => {
2693 if update_ctx {
2694 self.ctx.tag_columns.clear();
2695 self.ctx.aggregation_field_labels.clear();
2696 }
2697 Ok(vec![self.create_time_index_column_expr()?])
2698 }
2699 Some(LabelModifier::Include(labels)) => {
2700 if update_ctx {
2701 self.ctx.aggregation_field_labels = labels
2707 .labels
2708 .iter()
2709 .filter(|label| {
2710 self.ctx.field_columns.contains(label)
2711 || !self.ctx.tag_columns.contains(label)
2712 || self.ctx.aggregation_field_labels.contains(label)
2713 })
2714 .cloned()
2715 .collect();
2716 self.ctx.tag_columns.clear();
2717 }
2718 let mut exprs = Vec::with_capacity(labels.labels.len());
2719 for label in &labels.labels {
2720 if is_metric_engine_internal_column(label) {
2721 continue;
2722 }
2723 if let Some(column_name) = Self::find_case_sensitive_column(input_schema, label)
2725 {
2726 exprs.push(DfExpr::Column(Column::from_name(column_name.clone())));
2727
2728 if update_ctx {
2729 self.ctx.tag_columns.push(column_name);
2731 }
2732 }
2733 }
2734 exprs.push(self.create_time_index_column_expr()?);
2736
2737 Ok(exprs)
2738 }
2739 Some(LabelModifier::Exclude(labels)) => {
2740 let mut all_fields = input_schema
2741 .fields()
2742 .iter()
2743 .map(|f| f.name())
2744 .collect::<BTreeSet<_>>();
2745
2746 all_fields.retain(|col| !is_metric_engine_internal_column(col.as_str()));
2749
2750 for label in &labels.labels {
2753 let _ = all_fields.remove(label);
2754 }
2755
2756 if let Some(time_index) = &self.ctx.time_index_column {
2758 let _ = all_fields.remove(time_index);
2759 }
2760 for value in &self.ctx.field_columns {
2761 let _ = all_fields.remove(value);
2762 }
2763
2764 if update_ctx {
2765 self.ctx.tag_columns = all_fields.iter().map(|col| (*col).clone()).collect();
2767 self.ctx
2771 .aggregation_field_labels
2772 .retain(|label| all_fields.iter().any(|col| *col == label));
2773 }
2774
2775 let mut exprs = all_fields
2777 .into_iter()
2778 .map(|c| DfExpr::Column(Column::from_name(c)))
2779 .collect::<Vec<_>>();
2780
2781 exprs.push(self.create_time_index_column_expr()?);
2783
2784 Ok(exprs)
2785 }
2786 }
2787 }
2788
2789 pub fn matchers_to_expr(
2791 label_matchers: Matchers,
2792 table_schema: &DFSchemaRef,
2793 ) -> Result<Vec<DfExpr>> {
2794 let mut exprs = Vec::with_capacity(label_matchers.matchers.len());
2795 for matcher in label_matchers.matchers {
2796 if matcher.name == SCHEMA_COLUMN_MATCHER
2797 || matcher.name == DB_COLUMN_MATCHER
2798 || matcher.name == FIELD_COLUMN_MATCHER
2799 {
2800 continue;
2801 }
2802
2803 let accepts_empty = matcher.is_match("");
2804 let column_name = Self::find_case_sensitive_column(table_schema, matcher.name.as_str());
2805 let mut null_label = None;
2811 let col = if let Some(column_name) = column_name {
2812 let column = DfExpr::Column(Column::from_name(&column_name));
2813 let field = table_schema
2814 .index_of_column_by_name(None, &column_name)
2815 .map(|index| table_schema.field(index));
2816 if accepts_empty
2817 && field.is_some_and(|field| {
2818 field.is_nullable()
2819 && Self::string_value_data_type(field.data_type()).is_some()
2820 })
2821 {
2822 null_label = Some(column.clone().is_null());
2823 }
2824 column
2825 } else {
2826 DfExpr::Literal(ScalarValue::Utf8(Some(String::new())), None)
2827 };
2828 let lit = DfExpr::Literal(ScalarValue::Utf8(Some(matcher.value)), None);
2829 let expr = match matcher.op {
2830 MatchOp::Equal => col.eq(lit),
2831 MatchOp::NotEqual => col.not_eq(lit),
2832 MatchOp::Re(re) => {
2833 if re.as_str() == "^(?:.*)$" {
2839 continue;
2840 }
2841 if re.as_str() == "^(?:.+)$" {
2842 col.not_eq(DfExpr::Literal(
2843 ScalarValue::Utf8(Some(String::new())),
2844 None,
2845 ))
2846 } else {
2847 DfExpr::BinaryExpr(BinaryExpr {
2848 left: Box::new(col),
2849 op: Operator::RegexMatch,
2850 right: Box::new(DfExpr::Literal(
2851 ScalarValue::Utf8(Some(re.as_str().to_string())),
2852 None,
2853 )),
2854 })
2855 }
2856 }
2857 MatchOp::NotRe(re) => {
2858 if re.as_str() == "^(?:.*)$" {
2859 DfExpr::Literal(ScalarValue::Boolean(Some(false)), None)
2860 } else if re.as_str() == "^(?:.+)$" {
2861 col.eq(DfExpr::Literal(
2862 ScalarValue::Utf8(Some(String::new())),
2863 None,
2864 ))
2865 } else {
2866 DfExpr::BinaryExpr(BinaryExpr {
2867 left: Box::new(col),
2868 op: Operator::RegexNotMatch,
2869 right: Box::new(DfExpr::Literal(
2870 ScalarValue::Utf8(Some(re.as_str().to_string())),
2871 None,
2872 )),
2873 })
2874 }
2875 }
2876 };
2877 let expr = match null_label {
2878 Some(null_label) => null_label.or(expr),
2879 None => expr,
2880 };
2881 exprs.push(expr);
2882 }
2883
2884 Ok(exprs)
2885 }
2886
2887 fn find_case_sensitive_column(schema: &DFSchemaRef, column: &str) -> Option<String> {
2888 if is_metric_engine_internal_column(column) {
2889 return None;
2890 }
2891 schema
2892 .fields()
2893 .iter()
2894 .find(|field| field.name() == column)
2895 .map(|field| field.name().clone())
2896 }
2897
2898 fn table_from_source(&self, source: &Arc<dyn TableSource>) -> Result<table::TableRef> {
2899 Ok(source
2900 .downcast_ref::<DefaultTableSource>()
2901 .context(UnknownTableSnafu)?
2902 .table_provider
2903 .downcast_ref::<DfTableProviderAdapter>()
2904 .context(UnknownTableSnafu)?
2905 .table())
2906 }
2907
2908 fn table_ref(&self) -> Result<TableReference> {
2909 let table_name = self
2910 .ctx
2911 .table_name
2912 .clone()
2913 .context(TableNameNotFoundSnafu)?;
2914
2915 let table_ref = if let Some(schema_name) = &self.ctx.schema_name {
2917 TableReference::partial(schema_name.as_str(), table_name.as_str())
2918 } else {
2919 TableReference::bare(table_name.as_str())
2920 };
2921
2922 Ok(table_ref)
2923 }
2924
2925 fn build_time_index_filter(
2926 &self,
2927 offset_duration: i64,
2928 schema: &DFSchemaRef,
2929 ) -> Result<Option<DfExpr>> {
2930 let start = self.ctx.start;
2931 let end = self.ctx.end;
2932 if end < start {
2933 return InvalidTimeRangeSnafu { start, end }.fail();
2934 }
2935 let time_index_expr = self.create_time_index_column_expr()?;
2936 let time_index_name = self.ctx.time_index_column.as_ref().unwrap();
2937 let unit = schema
2938 .index_of_column_by_name(None, time_index_name)
2939 .and_then(|index| match schema.field(index).data_type() {
2940 ArrowDataType::Timestamp(unit, _) => Some(*unit),
2941 _ => None,
2942 })
2943 .unwrap_or(ArrowTimeUnit::Millisecond);
2944 let native_value = |milliseconds: i128| match unit {
2945 ArrowTimeUnit::Second => milliseconds.div_euclid(1_000),
2946 ArrowTimeUnit::Millisecond => milliseconds,
2947 ArrowTimeUnit::Microsecond => milliseconds * 1_000,
2948 ArrowTimeUnit::Nanosecond => milliseconds * 1_000_000,
2949 };
2950 let scalar = |milliseconds: i128| -> Option<ScalarValue> {
2951 let value = i64::try_from(native_value(milliseconds)).ok()?;
2952 Some(match unit {
2953 ArrowTimeUnit::Second => ScalarValue::TimestampSecond(Some(value), None),
2954 ArrowTimeUnit::Millisecond => ScalarValue::TimestampMillisecond(Some(value), None),
2955 ArrowTimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(Some(value), None),
2956 ArrowTimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(Some(value), None),
2957 })
2958 };
2959 let window = self.ctx.range.unwrap_or(self.ctx.lookback_delta);
2960 let filter = |lower_ms: i128, upper_ms: i128| {
2961 let lower_value = native_value(lower_ms);
2962 let upper_value = native_value(upper_ms);
2963 if lower_value > i128::from(i64::MAX) || upper_value < i128::from(i64::MIN) {
2964 return Some(lit(false));
2965 }
2966 let lower_filter = (lower_value >= i128::from(i64::MIN)).then(|| {
2967 let lower = DfExpr::Literal(scalar(lower_ms).unwrap(), None);
2968 if window == 0 {
2969 time_index_expr.clone().gt_eq(lower)
2970 } else if unit == ArrowTimeUnit::Millisecond
2971 && let Some(inclusive_lower) = lower_ms
2972 .checked_add(1)
2973 .and_then(|lower| i64::try_from(lower).ok())
2974 .and_then(|lower| scalar(i128::from(lower)))
2975 {
2976 time_index_expr
2977 .clone()
2978 .gt_eq(DfExpr::Literal(inclusive_lower, None))
2979 } else {
2980 time_index_expr.clone().gt(lower)
2981 }
2982 });
2983 let upper_filter = (upper_value <= i128::from(i64::MAX)).then(|| {
2984 time_index_expr
2985 .clone()
2986 .lt_eq(DfExpr::Literal(scalar(upper_ms).unwrap(), None))
2987 });
2988
2989 match (lower_filter, upper_filter) {
2993 (Some(lower), Some(upper)) => Some(lower.and(upper)),
2994 (Some(filter), None) | (None, Some(filter)) => Some(filter),
2995 (None, None) => None,
2996 }
2997 };
2998 let bounds = |timestamp: i64| {
2999 let upper = i128::from(timestamp) - i128::from(offset_duration);
3000 (upper - i128::from(window), upper)
3001 };
3002 let num_points = (end as i128 - start as i128) / self.ctx.interval as i128;
3003 if num_points > MAX_SCATTER_POINTS as i128 || self.ctx.interval <= INTERVAL_1H {
3004 let (lower, _) = bounds(start);
3005 let (_, upper) = bounds(end);
3006 return Ok(filter(lower, upper));
3007 }
3008 let mut filters = Vec::new();
3009 for timestamp in (start..=end).step_by(self.ctx.interval as usize) {
3010 let (lower, upper) = bounds(timestamp);
3011 let Some(filter) = filter(lower, upper) else {
3012 return Ok(None);
3015 };
3016 filters.push(filter);
3017 }
3018 Ok(filters.into_iter().reduce(DfExpr::or))
3019 }
3020
3021 async fn create_table_scan_plan(&mut self, table_ref: TableReference) -> Result<LogicalPlan> {
3026 let provider = self
3027 .table_provider
3028 .resolve_table(table_ref.clone())
3029 .await
3030 .context(CatalogSnafu)?;
3031
3032 let logical_table = self.table_from_source(&provider)?;
3033
3034 let mut maybe_phy_table_ref = table_ref.clone();
3036 let mut scan_provider = provider;
3037 let mut table_id_filter: Option<u32> = None;
3038
3039 if logical_table.table_info().meta.engine == METRIC_ENGINE_NAME
3042 && let Some(physical_table_name) = logical_table
3043 .table_info()
3044 .meta
3045 .options
3046 .extra_options
3047 .get(LOGICAL_TABLE_METADATA_KEY)
3048 {
3049 let physical_table_ref = if let Some(schema_name) = &self.ctx.schema_name {
3050 TableReference::partial(schema_name.as_str(), physical_table_name.as_str())
3051 } else {
3052 TableReference::bare(physical_table_name.as_str())
3053 };
3054
3055 let physical_provider = match self
3056 .table_provider
3057 .resolve_table(physical_table_ref.clone())
3058 .await
3059 {
3060 Ok(provider) => provider,
3061 Err(e) if e.status_code() == StatusCode::TableNotFound => {
3062 scan_provider.clone()
3065 }
3066 Err(e) => return Err(e).context(CatalogSnafu),
3067 };
3068
3069 if !Arc::ptr_eq(&physical_provider, &scan_provider) {
3070 let physical_table = self.table_from_source(&physical_provider)?;
3072
3073 let has_table_id = physical_table
3074 .schema()
3075 .column_schema_by_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME)
3076 .is_some();
3077 let has_tsid = physical_table
3078 .schema()
3079 .column_schema_by_name(DATA_SCHEMA_TSID_COLUMN_NAME)
3080 .is_some_and(|col| matches!(col.data_type, ConcreteDataType::UInt64(_)));
3081
3082 if has_table_id && has_tsid {
3083 scan_provider = physical_provider;
3084 maybe_phy_table_ref = physical_table_ref;
3085 table_id_filter = Some(logical_table.table_info().ident.table_id);
3086 }
3087 }
3088 }
3089
3090 let scan_table = self.table_from_source(&scan_provider)?;
3091
3092 let use_tsid = table_id_filter.is_some()
3093 && scan_table
3094 .schema()
3095 .column_schema_by_name(DATA_SCHEMA_TSID_COLUMN_NAME)
3096 .is_some_and(|col| matches!(col.data_type, ConcreteDataType::UInt64(_)));
3097 self.ctx.use_tsid = use_tsid;
3098
3099 let all_table_tags = self.ctx.tag_columns.clone();
3100
3101 let scan_tag_columns = if use_tsid {
3102 let mut scan_tags = self.ctx.tag_columns.clone();
3103 for matcher in &self.ctx.selector_matcher {
3104 if is_metric_engine_internal_column(&matcher.name) {
3105 continue;
3106 }
3107 if all_table_tags.iter().any(|tag| tag == &matcher.name) {
3108 scan_tags.push(matcher.name.clone());
3109 }
3110 }
3111 scan_tags.sort_unstable();
3112 scan_tags.dedup();
3113 scan_tags
3114 } else {
3115 self.ctx.tag_columns.clone()
3116 };
3117
3118 let time_index_data_type = scan_table
3119 .schema()
3120 .timestamp_column()
3121 .with_context(|| TimeIndexNotFoundSnafu {
3122 table: maybe_phy_table_ref.to_quoted_string(),
3123 })?
3124 .data_type
3125 .clone();
3126 let is_time_index_second =
3127 time_index_data_type == ConcreteDataType::timestamp_second_datatype();
3128
3129 let scan_projection = if table_id_filter.is_some() {
3130 let mut required_columns = HashSet::new();
3131 required_columns.insert(DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string());
3132 required_columns.insert(self.ctx.time_index_column.clone().with_context(|| {
3133 TimeIndexNotFoundSnafu {
3134 table: maybe_phy_table_ref.to_quoted_string(),
3135 }
3136 })?);
3137 for col in &scan_tag_columns {
3138 required_columns.insert(col.clone());
3139 }
3140 for col in &self.ctx.field_columns {
3141 required_columns.insert(col.clone());
3142 }
3143 if use_tsid {
3144 required_columns.insert(DATA_SCHEMA_TSID_COLUMN_NAME.to_string());
3145 }
3146
3147 let arrow_schema = scan_table.schema().arrow_schema().clone();
3148 Some(
3149 arrow_schema
3150 .fields()
3151 .iter()
3152 .enumerate()
3153 .filter(|(_, field)| required_columns.contains(field.name().as_str()))
3154 .map(|(idx, _)| idx)
3155 .collect::<Vec<_>>(),
3156 )
3157 } else {
3158 None
3159 };
3160
3161 let mut scan_plan =
3162 LogicalPlanBuilder::scan(maybe_phy_table_ref.clone(), scan_provider, scan_projection)
3163 .context(DataFusionPlanningSnafu)?
3164 .build()
3165 .context(DataFusionPlanningSnafu)?;
3166
3167 if let Some(table_id) = table_id_filter {
3168 scan_plan = LogicalPlanBuilder::from(scan_plan)
3169 .filter(
3170 DfExpr::Column(Column::from_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME))
3171 .eq(lit(table_id)),
3172 )
3173 .context(DataFusionPlanningSnafu)?
3174 .alias(table_ref.clone()) .context(DataFusionPlanningSnafu)?
3176 .build()
3177 .context(DataFusionPlanningSnafu)?;
3178 }
3179
3180 if is_time_index_second {
3181 let expr: Vec<_> = self
3186 .create_field_column_exprs()?
3187 .into_iter()
3188 .chain(
3189 scan_tag_columns
3190 .iter()
3191 .map(|tag| DfExpr::Column(Column::from_name(tag))),
3192 )
3193 .chain(self.ctx.use_tsid.then_some(DfExpr::Column(Column::new(
3194 Some(table_ref.clone()),
3195 DATA_SCHEMA_TSID_COLUMN_NAME.to_string(),
3196 ))))
3197 .chain(Some(DfExpr::Alias(Alias {
3198 expr: Box::new(DfExpr::Cast(Cast::new(
3199 Box::new(self.create_time_index_column_expr()?),
3200 ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
3201 ))),
3202 relation: Some(table_ref.clone()),
3203 name: self
3204 .ctx
3205 .time_index_column
3206 .as_ref()
3207 .with_context(|| TimeIndexNotFoundSnafu {
3208 table: table_ref.to_quoted_string(),
3209 })?
3210 .clone(),
3211 metadata: None,
3212 })))
3213 .collect::<Vec<_>>();
3214 scan_plan = LogicalPlanBuilder::from(scan_plan)
3215 .project(expr)
3216 .context(DataFusionPlanningSnafu)?
3217 .build()
3218 .context(DataFusionPlanningSnafu)?;
3219 } else if table_id_filter.is_some()
3220 || time_index_data_type == ConcreteDataType::timestamp_microsecond_datatype()
3221 || time_index_data_type == ConcreteDataType::timestamp_nanosecond_datatype()
3222 {
3223 let project_exprs = self
3229 .create_field_column_exprs()?
3230 .into_iter()
3231 .chain(
3232 scan_tag_columns
3233 .iter()
3234 .map(|tag| DfExpr::Column(Column::from_name(tag))),
3235 )
3236 .chain(
3237 self.ctx
3238 .use_tsid
3239 .then_some(DfExpr::Column(Column::from_name(
3240 DATA_SCHEMA_TSID_COLUMN_NAME,
3241 ))),
3242 )
3243 .chain(Some(self.create_time_index_column_expr()?))
3244 .collect::<Vec<_>>();
3245
3246 scan_plan = LogicalPlanBuilder::from(scan_plan)
3247 .project(project_exprs)
3248 .context(DataFusionPlanningSnafu)?
3249 .build()
3250 .context(DataFusionPlanningSnafu)?;
3251 }
3252
3253 let result = LogicalPlanBuilder::from(scan_plan)
3254 .build()
3255 .context(DataFusionPlanningSnafu)?;
3256 Ok(result)
3257 }
3258
3259 fn collect_row_key_tag_columns_from_plan(
3260 &self,
3261 plan: &LogicalPlan,
3262 ) -> Result<BTreeSet<String>> {
3263 fn walk(
3264 planner: &PromPlanner,
3265 plan: &LogicalPlan,
3266 out: &mut BTreeSet<String>,
3267 ) -> Result<()> {
3268 if let LogicalPlan::TableScan(scan) = plan
3270 && let Ok(table) = planner.table_from_source(&scan.source)
3271 {
3272 for col in table.table_info().meta.row_key_column_names() {
3273 if col != DATA_SCHEMA_TABLE_ID_COLUMN_NAME
3274 && col != DATA_SCHEMA_TSID_COLUMN_NAME
3275 && !is_metric_engine_internal_column(col)
3276 {
3277 out.insert(col.clone());
3278 }
3279 }
3280 }
3281
3282 for input in plan.inputs() {
3283 walk(planner, input, out)?;
3284 }
3285 Ok(())
3286 }
3287
3288 let mut out = BTreeSet::new();
3289 walk(self, plan, &mut out)?;
3290 Ok(out)
3291 }
3292
3293 async fn setup_context(&mut self) -> Result<Option<LogicalPlan>> {
3297 let table_ref = self.table_ref()?;
3298 let source = match self.table_provider.resolve_table(table_ref.clone()).await {
3299 Err(e) if e.status_code() == StatusCode::TableNotFound => {
3300 let plan = self.setup_context_for_empty_metric()?;
3301 return Ok(Some(plan));
3302 }
3303 res => res.context(CatalogSnafu)?,
3304 };
3305 let table = self.table_from_source(&source)?;
3306
3307 let time_index = table
3309 .schema()
3310 .timestamp_column()
3311 .with_context(|| TimeIndexNotFoundSnafu {
3312 table: table_ref.to_quoted_string(),
3313 })?
3314 .name
3315 .clone();
3316 self.ctx.time_index_column = Some(time_index);
3317
3318 let values = table
3320 .table_info()
3321 .meta
3322 .field_column_names()
3323 .cloned()
3324 .collect();
3325 self.ctx.field_columns = values;
3326
3327 let tags = table
3329 .table_info()
3330 .meta
3331 .row_key_column_names()
3332 .filter(|col| {
3333 col != &DATA_SCHEMA_TABLE_ID_COLUMN_NAME && col != &DATA_SCHEMA_TSID_COLUMN_NAME
3335 })
3336 .cloned()
3337 .collect();
3338 self.ctx.tag_columns = tags;
3339 self.ctx.aggregation_field_labels.clear();
3341
3342 self.ctx.use_tsid = false;
3343
3344 Ok(None)
3345 }
3346
3347 fn setup_context_for_empty_metric(&mut self) -> Result<LogicalPlan> {
3350 self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
3351 self.ctx.reset_table_name_and_schema();
3352 self.ctx.tag_columns = vec![];
3353 self.ctx.aggregation_field_labels.clear();
3354 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
3355 self.ctx.use_tsid = false;
3356
3357 let plan = LogicalPlan::Extension(Extension {
3359 node: Arc::new(
3360 EmptyMetric::new(
3361 0,
3362 -1,
3363 self.ctx.interval,
3364 SPECIAL_TIME_FUNCTION.to_string(),
3365 DEFAULT_FIELD_COLUMN.to_string(),
3366 Some(lit(0.0f64)),
3367 )
3368 .context(DataFusionPlanningSnafu)?,
3369 ),
3370 });
3371 Ok(plan)
3372 }
3373
3374 fn create_function_args(&self, args: &[Box<PromExpr>]) -> Result<FunctionArgs> {
3376 let mut result = FunctionArgs::default();
3377
3378 for arg in args {
3379 if let Some(expr) = Self::try_build_literal_expr(arg) {
3381 result.literals.push(expr);
3382 } else {
3383 match arg.as_ref() {
3385 PromExpr::Subquery(_)
3386 | PromExpr::VectorSelector(_)
3387 | PromExpr::MatrixSelector(_)
3388 | PromExpr::Extension(_)
3389 | PromExpr::Aggregate(_)
3390 | PromExpr::Paren(_)
3391 | PromExpr::Call(_)
3392 | PromExpr::Binary(_)
3393 | PromExpr::Unary(_) => {
3394 if result.input.replace(*arg.clone()).is_some() {
3395 MultipleVectorSnafu { expr: *arg.clone() }.fail()?;
3396 }
3397 }
3398
3399 _ => {
3400 let expr = Self::get_param_as_literal_expr(Some(arg.as_ref()), None, None)?;
3401 result.literals.push(expr);
3402 }
3403 }
3404 }
3405 }
3406
3407 Ok(result)
3408 }
3409
3410 fn create_mixed_range_function_exprs(
3411 &mut self,
3412 func: &Function,
3413 mut other_input_exprs: VecDeque<DfExpr>,
3414 float_field: &str,
3415 histogram_field: &str,
3416 input_schema: &DFSchemaRef,
3417 range_fold_offset: Option<Millisecond>,
3418 ) -> Result<Option<Vec<DfExpr>>> {
3419 let returns_histogram = matches!(
3420 func.name,
3421 "rate"
3422 | "increase"
3423 | "delta"
3424 | "idelta"
3425 | "irate"
3426 | "avg_over_time"
3427 | "sum_over_time"
3428 | "last_over_time"
3429 );
3430 if !returns_histogram
3431 && !matches!(
3432 func.name,
3433 "changes"
3434 | "resets"
3435 | "deriv"
3436 | "min_over_time"
3437 | "max_over_time"
3438 | "count_over_time"
3439 | "absent_over_time"
3440 | "present_over_time"
3441 | "stddev_over_time"
3442 | "stdvar_over_time"
3443 | "quantile_over_time"
3444 | "predict_linear"
3445 | "double_exponential_smoothing"
3446 | "holt_winters"
3447 )
3448 {
3449 return Ok(None);
3450 }
3451
3452 if func.name == "predict_linear" {
3453 other_input_exprs[0] = DfExpr::Cast(Cast::new(
3454 Box::new(other_input_exprs[0].clone()),
3455 ArrowDataType::Int64,
3456 ));
3457 other_input_exprs.push_back(self.create_range_eval_ts_expr(
3462 range_fold_offset.context(ExpectRangeSelectorSnafu)?,
3463 input_schema,
3464 )?);
3465 }
3466
3467 let timestamp_range = DfExpr::Column(Column::from_name(
3468 RangeManipulate::build_timestamp_range_name(
3469 self.ctx.time_index_column.as_ref().unwrap(),
3470 ),
3471 ));
3472 let float_range = DfExpr::Column(Column::from_name(float_field));
3473 let histogram_range = DfExpr::Column(Column::from_name(histogram_field));
3474 let mut args = Vec::with_capacity(other_input_exprs.len() + 6);
3475 args.push(lit(func.name));
3476 args.push(timestamp_range.clone());
3477 args.push(float_range.clone());
3478 args.push(histogram_range.clone());
3479 args.extend(other_input_exprs);
3480 if matches!(func.name, "rate" | "increase" | "delta") {
3481 args.push(self.create_time_index_column_expr()?);
3482 args.push(lit(self.ctx.range.context(ExpectRangeSelectorSnafu)?));
3483 }
3484
3485 let mut float_expr = DfExpr::ScalarFunction(ScalarFunction {
3486 func: Arc::new(MixedRange::float_udf(self.promql_annotations.clone())),
3487 args: args.clone(),
3488 });
3489 if matches!(func.name, "rate" | "increase") {
3490 let raw_delta_function = if func.name == "rate" {
3491 "raw_delta_rate"
3492 } else {
3493 "raw_delta_increase"
3494 };
3495 let delta_sum = DfExpr::ScalarFunction(ScalarFunction {
3496 func: Arc::new(MixedRange::float_udf(self.promql_annotations.clone())),
3497 args: vec![
3498 lit(raw_delta_function),
3499 timestamp_range,
3500 float_range,
3501 histogram_range,
3502 ],
3503 });
3504 float_expr = self.select_delta_range_math(
3505 func.name,
3506 input_schema,
3507 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3508 delta_sum,
3509 float_expr,
3510 )?;
3511 }
3512 let exprs = if returns_histogram {
3513 self.ctx.field_columns = vec![float_field.to_string(), histogram_field.to_string()];
3514 vec![
3515 float_expr.alias(float_field),
3516 DfExpr::ScalarFunction(ScalarFunction {
3517 func: Arc::new(MixedRange::histogram_udf(self.promql_annotations.clone())),
3518 args,
3519 })
3520 .alias(histogram_field),
3521 ]
3522 } else {
3523 let display_name = if func.name == "predict_linear" {
3524 Self::name_without_last_arg(&float_expr)
3527 } else {
3528 float_expr.schema_name().to_string()
3529 };
3530 self.ctx.field_columns = vec![display_name.clone()];
3531 vec![float_expr.alias(display_name)]
3532 };
3533 Ok(Some(exprs))
3534 }
3535
3536 fn create_function_expr(
3542 &mut self,
3543 func: &Function,
3544 other_input_exprs: Vec<DfExpr>,
3545 input_schema: &DFSchemaRef,
3546 query_engine_state: &QueryEngineState,
3547 range_fold_offset: Option<Millisecond>,
3548 ) -> Result<(Vec<DfExpr>, Vec<String>)> {
3549 let mut other_input_exprs: VecDeque<DfExpr> = other_input_exprs.into();
3551 if let Some((float_field, histogram_field)) =
3552 Self::alternative_sample_range_columns(input_schema, &self.ctx.field_columns)
3553 .map(|(float, histogram)| (float.to_string(), histogram.to_string()))
3554 && let Some(exprs) = self.create_mixed_range_function_exprs(
3555 func,
3556 other_input_exprs.clone(),
3557 &float_field,
3558 &histogram_field,
3559 input_schema,
3560 range_fold_offset,
3561 )?
3562 {
3563 return Ok((exprs, vec![]));
3564 }
3565 let alternative_samples =
3566 Self::field_columns_are_alternative_samples(input_schema, &self.ctx.field_columns);
3567 let all_field_columns_are_native_histogram_ranges =
3568 self.all_field_columns_are_native_histogram_ranges(input_schema);
3569
3570 let field_column_pos = 0;
3572 let mut exprs = Vec::with_capacity(self.ctx.field_columns.len());
3573 let mut new_tags = vec![];
3575 let promql_annotations = self.promql_annotations.clone();
3576 let native_histogram_drop_udf = |name: &str| {
3577 Arc::new(NativeHistogramDrop::float_null_udf(
3578 format!(
3579 "{name}: dropped native histogram samples because this function is not supported for native histograms"
3580 ),
3581 promql_annotations.clone(),
3582 ))
3583 };
3584 let scalar_func = match func.name {
3585 "increase" => {
3586 if all_field_columns_are_native_histogram_ranges {
3587 ScalarFunc::ExtrapolateUdf(
3588 Arc::new(NativeHistogramIncrease::scalar_udf_with_collector(
3589 self.promql_annotations.clone(),
3590 )),
3591 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3592 )
3593 } else {
3594 ScalarFunc::ExtrapolateUdf(
3595 Arc::new(Increase::scalar_udf()),
3596 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3597 )
3598 }
3599 }
3600 "rate" => {
3601 if all_field_columns_are_native_histogram_ranges {
3602 ScalarFunc::ExtrapolateUdf(
3603 Arc::new(NativeHistogramRate::scalar_udf_with_collector(
3604 self.promql_annotations.clone(),
3605 )),
3606 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3607 )
3608 } else {
3609 ScalarFunc::ExtrapolateUdf(
3610 Arc::new(Rate::scalar_udf()),
3611 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3612 )
3613 }
3614 }
3615 "delta" => {
3616 if all_field_columns_are_native_histogram_ranges {
3617 ScalarFunc::ExtrapolateUdf(
3618 Arc::new(NativeHistogramDelta::scalar_udf_with_collector(
3619 self.promql_annotations.clone(),
3620 )),
3621 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3622 )
3623 } else {
3624 ScalarFunc::ExtrapolateUdf(
3625 Arc::new(Delta::scalar_udf()),
3626 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3627 )
3628 }
3629 }
3630 "idelta" => {
3631 if all_field_columns_are_native_histogram_ranges {
3632 ScalarFunc::Udf(Arc::new(NativeHistogramIDelta::scalar_udf_with_collector(
3633 self.promql_annotations.clone(),
3634 )))
3635 } else {
3636 ScalarFunc::Udf(Arc::new(IDelta::<false>::scalar_udf()))
3637 }
3638 }
3639 "irate" => {
3640 if all_field_columns_are_native_histogram_ranges {
3641 ScalarFunc::Udf(Arc::new(NativeHistogramIRate::scalar_udf_with_collector(
3642 self.promql_annotations.clone(),
3643 )))
3644 } else {
3645 ScalarFunc::Udf(Arc::new(IDelta::<true>::scalar_udf()))
3646 }
3647 }
3648 "resets" => {
3649 if all_field_columns_are_native_histogram_ranges {
3650 ScalarFunc::Udf(Arc::new(NativeHistogramResets::scalar_udf()))
3651 } else {
3652 ScalarFunc::Udf(Arc::new(Resets::scalar_udf()))
3653 }
3654 }
3655 "changes" => {
3656 if all_field_columns_are_native_histogram_ranges {
3657 ScalarFunc::Udf(Arc::new(NativeHistogramChanges::scalar_udf()))
3658 } else {
3659 ScalarFunc::Udf(Arc::new(Changes::scalar_udf()))
3660 }
3661 }
3662 "deriv" => {
3663 if all_field_columns_are_native_histogram_ranges {
3664 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3665 } else {
3666 ScalarFunc::Udf(Arc::new(Deriv::scalar_udf()))
3667 }
3668 }
3669 "avg_over_time" => {
3670 if all_field_columns_are_native_histogram_ranges {
3671 ScalarFunc::Udf(Arc::new(
3672 NativeHistogramAvgOverTime::scalar_udf_with_collector(
3673 self.promql_annotations.clone(),
3674 ),
3675 ))
3676 } else {
3677 ScalarFunc::Udf(Arc::new(AvgOverTime::scalar_udf()))
3678 }
3679 }
3680 "min_over_time" => {
3681 if all_field_columns_are_native_histogram_ranges {
3682 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3683 } else {
3684 ScalarFunc::Udf(Arc::new(MinOverTime::scalar_udf()))
3685 }
3686 }
3687 "max_over_time" => {
3688 if all_field_columns_are_native_histogram_ranges {
3689 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3690 } else {
3691 ScalarFunc::Udf(Arc::new(MaxOverTime::scalar_udf()))
3692 }
3693 }
3694 "sum_over_time" => {
3695 if all_field_columns_are_native_histogram_ranges {
3696 ScalarFunc::Udf(Arc::new(
3697 NativeHistogramSumOverTime::scalar_udf_with_collector(
3698 self.promql_annotations.clone(),
3699 ),
3700 ))
3701 } else {
3702 ScalarFunc::Udf(Arc::new(SumOverTime::scalar_udf()))
3703 }
3704 }
3705 "count_over_time" => {
3706 if all_field_columns_are_native_histogram_ranges {
3707 ScalarFunc::Udf(Arc::new(NativeHistogramCountOverTime::scalar_udf()))
3708 } else {
3709 ScalarFunc::Udf(Arc::new(CountOverTime::scalar_udf()))
3710 }
3711 }
3712 "last_over_time" => {
3713 if all_field_columns_are_native_histogram_ranges {
3714 ScalarFunc::Udf(Arc::new(NativeHistogramLastOverTime::scalar_udf()))
3715 } else {
3716 ScalarFunc::Udf(Arc::new(LastOverTime::scalar_udf()))
3717 }
3718 }
3719 "absent_over_time" => {
3720 if all_field_columns_are_native_histogram_ranges {
3721 ScalarFunc::Udf(Arc::new(NativeHistogramAbsentOverTime::scalar_udf()))
3722 } else {
3723 ScalarFunc::Udf(Arc::new(AbsentOverTime::scalar_udf()))
3724 }
3725 }
3726 "present_over_time" => {
3727 if all_field_columns_are_native_histogram_ranges {
3728 ScalarFunc::Udf(Arc::new(NativeHistogramPresentOverTime::scalar_udf()))
3729 } else {
3730 ScalarFunc::Udf(Arc::new(PresentOverTime::scalar_udf()))
3731 }
3732 }
3733 "stddev_over_time" => {
3734 if all_field_columns_are_native_histogram_ranges {
3735 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3736 } else {
3737 ScalarFunc::Udf(Arc::new(StddevOverTime::scalar_udf()))
3738 }
3739 }
3740 "stdvar_over_time" => {
3741 if all_field_columns_are_native_histogram_ranges {
3742 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3743 } else {
3744 ScalarFunc::Udf(Arc::new(StdvarOverTime::scalar_udf()))
3745 }
3746 }
3747 "quantile_over_time" => {
3748 if all_field_columns_are_native_histogram_ranges {
3749 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3750 } else {
3751 ScalarFunc::Udf(Arc::new(QuantileOverTime::scalar_udf()))
3752 }
3753 }
3754 "predict_linear" => {
3755 if all_field_columns_are_native_histogram_ranges {
3756 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3757 } else {
3758 other_input_exprs[0] = DfExpr::Cast(Cast::new(
3759 Box::new(other_input_exprs[0].clone()),
3760 ArrowDataType::Int64,
3761 ));
3762 other_input_exprs.push_back(self.create_range_eval_ts_expr(
3769 range_fold_offset.context(ExpectRangeSelectorSnafu)?,
3770 input_schema,
3771 )?);
3772 ScalarFunc::Udf(Arc::new(PredictLinear::scalar_udf()))
3773 }
3774 }
3775 "double_exponential_smoothing" | "holt_winters" => {
3776 if all_field_columns_are_native_histogram_ranges {
3777 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3778 } else {
3779 ScalarFunc::Udf(Arc::new(DoubleExponentialSmoothing::scalar_udf()))
3780 }
3781 }
3782 "histogram_count" => {
3783 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramCount::scalar_udf()))
3784 }
3785 "histogram_sum" => {
3786 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramSum::scalar_udf()))
3787 }
3788 "histogram_avg" => {
3789 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramAvg::scalar_udf()))
3790 }
3791 "histogram_stddev" => {
3792 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramStddev::scalar_udf()))
3793 }
3794 "histogram_stdvar" => {
3795 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramStdvar::scalar_udf()))
3796 }
3797 "time" => {
3798 exprs.push(build_special_time_expr(
3799 self.ctx.time_index_column.as_ref().unwrap(),
3800 ));
3801 ScalarFunc::GeneratedExpr
3802 }
3803 "minute" => {
3804 let expr = self.date_part_on_time_index("minute")?;
3806 exprs.push(expr);
3807 ScalarFunc::GeneratedExpr
3808 }
3809 "hour" => {
3810 let expr = self.date_part_on_time_index("hour")?;
3812 exprs.push(expr);
3813 ScalarFunc::GeneratedExpr
3814 }
3815 "month" => {
3816 let expr = self.date_part_on_time_index("month")?;
3818 exprs.push(expr);
3819 ScalarFunc::GeneratedExpr
3820 }
3821 "year" => {
3822 let expr = self.date_part_on_time_index("year")?;
3824 exprs.push(expr);
3825 ScalarFunc::GeneratedExpr
3826 }
3827 "day_of_month" => {
3828 let expr = self.date_part_on_time_index("day")?;
3830 exprs.push(expr);
3831 ScalarFunc::GeneratedExpr
3832 }
3833 "day_of_week" => {
3834 let expr = self.date_part_on_time_index("dow")?;
3836 exprs.push(expr);
3837 ScalarFunc::GeneratedExpr
3838 }
3839 "day_of_year" => {
3840 let expr = self.date_part_on_time_index("doy")?;
3842 exprs.push(expr);
3843 ScalarFunc::GeneratedExpr
3844 }
3845 "days_in_month" => {
3846 let day_lit_expr = "day".lit();
3851 let month_lit_expr = "month".lit();
3852 let interval_1month_lit_expr =
3853 DfExpr::Literal(ScalarValue::IntervalYearMonth(Some(1)), None);
3854 let interval_1day_lit_expr = DfExpr::Literal(
3855 ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(1, 0))),
3856 None,
3857 );
3858 let the_1month_minus_1day_expr = DfExpr::BinaryExpr(BinaryExpr {
3859 left: Box::new(interval_1month_lit_expr),
3860 op: Operator::Minus,
3861 right: Box::new(interval_1day_lit_expr),
3862 });
3863 let date_trunc_expr = DfExpr::ScalarFunction(ScalarFunction {
3864 func: datafusion_functions::datetime::date_trunc(),
3865 args: vec![month_lit_expr, self.create_time_index_column_expr()?],
3866 });
3867 let date_trunc_plus_interval_expr = DfExpr::BinaryExpr(BinaryExpr {
3868 left: Box::new(date_trunc_expr),
3869 op: Operator::Plus,
3870 right: Box::new(the_1month_minus_1day_expr),
3871 });
3872 let date_part_expr = DfExpr::ScalarFunction(ScalarFunction {
3873 func: datafusion_functions::datetime::date_part(),
3874 args: vec![day_lit_expr, date_trunc_plus_interval_expr],
3875 });
3876
3877 exprs.push(date_part_expr);
3878 ScalarFunc::GeneratedExpr
3879 }
3880
3881 "label_join" => {
3882 self.ctx.use_tsid = false;
3883 let (concat_expr, dst_label) = Self::build_concat_labels_expr(
3884 &mut other_input_exprs,
3885 &self.ctx,
3886 input_schema,
3887 query_engine_state,
3888 )?;
3889
3890 for value in &self.ctx.field_columns {
3892 if *value != dst_label {
3893 let expr = DfExpr::Column(Column::from_name(value));
3894 exprs.push(expr);
3895 }
3896 }
3897
3898 self.ctx.tag_columns.retain(|tag| *tag != dst_label);
3900 new_tags.push(dst_label);
3901 exprs.push(concat_expr);
3903
3904 ScalarFunc::GeneratedExpr
3905 }
3906 "label_replace" => {
3907 self.ctx.use_tsid = false;
3908 if let Some((replace_expr, dst_label)) =
3909 self.build_regexp_replace_label_expr(&mut other_input_exprs, input_schema)?
3910 {
3911 for value in &self.ctx.field_columns {
3913 if *value != dst_label {
3914 let expr = DfExpr::Column(Column::from_name(value));
3915 exprs.push(expr);
3916 }
3917 }
3918
3919 self.ctx.tag_columns.retain(|tag| *tag != dst_label);
3921 new_tags.push(dst_label);
3922 exprs.push(replace_expr);
3924 } else {
3925 for value in &self.ctx.field_columns {
3927 let expr = DfExpr::Column(Column::from_name(value));
3928 exprs.push(expr);
3929 }
3930 }
3931
3932 ScalarFunc::GeneratedExpr
3933 }
3934 "sort" | "sort_desc" => {
3935 for value in &self.ctx.field_columns {
3937 if !Self::field_column_is_native_histogram(input_schema, value) {
3938 exprs.push(DfExpr::Column(Column::from_name(value)));
3939 }
3940 }
3941 if exprs.is_empty() {
3944 exprs.push(DfExpr::Literal(ScalarValue::Float64(None), None));
3945 }
3946
3947 ScalarFunc::GeneratedExpr
3948 }
3949 "sort_by_label" | "sort_by_label_desc" | "timestamp" => {
3950 for value in &self.ctx.field_columns {
3953 let expr = DfExpr::Column(Column::from_name(value));
3954 exprs.push(expr);
3955 }
3956
3957 ScalarFunc::GeneratedExpr
3958 }
3959 "round" if self.all_field_columns_are_native_histograms(input_schema) => {
3960 if other_input_exprs.is_empty() {
3961 other_input_exprs.push_front(0.0f64.lit());
3962 }
3963 ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3964 }
3965 "round" => {
3966 if other_input_exprs.is_empty() {
3967 other_input_exprs.push_front(1.0f64.lit());
3968 }
3969 ScalarFunc::DataFusionUdf(Arc::new(Round::scalar_udf()))
3970 }
3971 "sqrt" if !self.all_field_columns_are_native_histograms(input_schema) => {
3973 ScalarFunc::DataFusionUdf(Arc::new(IeeeSqrt::scalar_udf()))
3974 }
3975 "rad" | "deg" | "sgn" if self.all_field_columns_are_native_histograms(input_schema) => {
3976 ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3977 }
3978 "rad" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::radians()),
3979 "deg" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::degrees()),
3980 "sgn" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::signum()),
3981 "pi" => {
3982 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
3984 func: datafusion::functions::math::pi(),
3985 args: vec![],
3986 });
3987 exprs.push(fn_expr);
3988
3989 ScalarFunc::GeneratedExpr
3990 }
3991 _ => {
3992 if let Some(f) = query_engine_state
3993 .session_state()
3994 .scalar_functions()
3995 .get(func.name)
3996 {
3997 if self.all_field_columns_are_native_histograms(input_schema) {
3998 ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3999 } else {
4000 ScalarFunc::DataFusionBuiltin(f.clone())
4001 }
4002 } else if let Some(factory) = query_engine_state.scalar_function(func.name) {
4003 if self.all_field_columns_are_native_histograms(input_schema) {
4004 ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
4005 } else {
4006 let func_state = query_engine_state.function_state();
4007 let query_ctx = self.table_provider.query_ctx();
4008
4009 ScalarFunc::DataFusionUdf(Arc::new(factory.provide(FunctionContext {
4010 state: func_state,
4011 query_ctx: query_ctx.clone(),
4012 })))
4013 }
4014 } else if let Some(f) = datafusion_functions::math::functions()
4015 .iter()
4016 .find(|f| f.name() == func.name)
4017 {
4018 if self.all_field_columns_are_native_histograms(input_schema) {
4019 ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
4020 } else {
4021 ScalarFunc::DataFusionUdf(f.clone())
4022 }
4023 } else {
4024 return UnsupportedExprSnafu {
4025 name: func.name.to_string(),
4026 }
4027 .fail();
4028 }
4029 }
4030 };
4031
4032 for value in &self.ctx.field_columns {
4033 let col_expr = DfExpr::Column(Column::from_name(value));
4034 let value_is_histogram = Self::field_column_is_native_histogram(input_schema, value);
4035
4036 match scalar_func.clone() {
4037 ScalarFunc::DataFusionBuiltin(func) => {
4038 if alternative_samples && value_is_histogram {
4039 continue;
4040 }
4041 other_input_exprs.insert(field_column_pos, col_expr);
4042 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
4043 func,
4044 args: other_input_exprs.clone().into(),
4045 });
4046 exprs.push(fn_expr);
4047 let _ = other_input_exprs.remove(field_column_pos);
4048 }
4049 ScalarFunc::DataFusionUdf(func) => {
4050 if alternative_samples && value_is_histogram {
4051 continue;
4052 }
4053 let args = itertools::chain!(
4054 other_input_exprs.iter().take(field_column_pos).cloned(),
4055 std::iter::once(col_expr),
4056 other_input_exprs.iter().skip(field_column_pos).cloned()
4057 )
4058 .collect_vec();
4059 exprs.push(DfExpr::ScalarFunction(ScalarFunction { func, args }))
4060 }
4061 ScalarFunc::NativeHistogramUdf(func) => {
4062 if value_is_histogram {
4063 let args = itertools::chain!(
4064 other_input_exprs.iter().take(field_column_pos).cloned(),
4065 std::iter::once(col_expr),
4066 other_input_exprs.iter().skip(field_column_pos).cloned()
4067 )
4068 .collect_vec();
4069 exprs.push(DfExpr::ScalarFunction(ScalarFunction { func, args }));
4070 } else if !alternative_samples {
4071 exprs.push(
4072 DfExpr::Literal(ScalarValue::Float64(None), None).alias(format!(
4073 "{}_{}",
4074 func.name(),
4075 value
4076 )),
4077 );
4078 }
4079 }
4080 ScalarFunc::Udf(func) => {
4081 let ts_range_expr = DfExpr::Column(Column::from_name(
4082 RangeManipulate::build_timestamp_range_name(
4083 self.ctx.time_index_column.as_ref().unwrap(),
4084 ),
4085 ));
4086 other_input_exprs.insert(field_column_pos, ts_range_expr);
4087 other_input_exprs.insert(field_column_pos + 1, col_expr);
4088 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
4089 func,
4090 args: other_input_exprs.clone().into(),
4091 });
4092 exprs.push(fn_expr);
4093 let _ = other_input_exprs.remove(field_column_pos + 1);
4094 let _ = other_input_exprs.remove(field_column_pos);
4095 }
4096 ScalarFunc::ExtrapolateUdf(udf, range_length) => {
4097 let ts_range_expr = DfExpr::Column(Column::from_name(
4098 RangeManipulate::build_timestamp_range_name(
4099 self.ctx.time_index_column.as_ref().unwrap(),
4100 ),
4101 ));
4102 other_input_exprs.insert(field_column_pos, ts_range_expr.clone());
4103 other_input_exprs.insert(field_column_pos + 1, col_expr.clone());
4104 other_input_exprs
4105 .insert(field_column_pos + 2, self.create_time_index_column_expr()?);
4106 other_input_exprs.push_back(lit(range_length));
4107 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
4108 func: udf,
4109 args: other_input_exprs.clone().into(),
4110 });
4111 let fn_expr = if matches!(func.name, "rate" | "increase")
4112 && !all_field_columns_are_native_histogram_ranges
4113 {
4114 let delta_sum = DfExpr::ScalarFunction(ScalarFunction {
4115 func: Arc::new(SumOverTime::scalar_udf()),
4116 args: vec![ts_range_expr, col_expr],
4117 });
4118 self.select_delta_range_math(
4119 func.name,
4120 input_schema,
4121 range_length,
4122 delta_sum,
4123 fn_expr,
4124 )?
4125 } else {
4126 fn_expr
4127 };
4128 exprs.push(fn_expr);
4129 let _ = other_input_exprs.pop_back();
4130 let _ = other_input_exprs.remove(field_column_pos + 2);
4131 let _ = other_input_exprs.remove(field_column_pos + 1);
4132 let _ = other_input_exprs.remove(field_column_pos);
4133 }
4134 ScalarFunc::GeneratedExpr => {}
4135 }
4136 }
4137
4138 if !matches!(func.name, "label_join" | "label_replace") {
4142 let mut new_field_columns = Vec::with_capacity(exprs.len());
4143
4144 exprs = exprs
4145 .into_iter()
4146 .map(|expr| {
4147 let display_name = if func.name == "predict_linear"
4152 && !all_field_columns_are_native_histogram_ranges
4153 {
4154 Self::name_without_last_arg(&expr)
4155 } else {
4156 expr.schema_name().to_string()
4157 };
4158 new_field_columns.push(display_name.clone());
4159 Ok(expr.alias(display_name))
4160 })
4161 .collect::<std::result::Result<Vec<_>, _>>()
4162 .context(DataFusionPlanningSnafu)?;
4163
4164 self.ctx.field_columns = new_field_columns;
4165 }
4166
4167 Ok((exprs, new_tags))
4168 }
4169
4170 fn select_delta_range_math(
4171 &self,
4172 function: &str,
4173 input_schema: &DFSchemaRef,
4174 range_length: Millisecond,
4175 delta_sum: DfExpr,
4176 cumulative: DfExpr,
4177 ) -> Result<DfExpr> {
4178 let marker_is_delta = if self
4179 .ctx
4180 .tag_columns
4181 .iter()
4182 .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL)
4183 {
4184 Self::field_column_type(input_schema, OTLP_AGGREGATION_TEMPORALITY_LABEL)
4185 .filter(|data_type| Self::string_value_data_type(data_type).is_some())
4186 .map(|_| {
4187 DfExpr::Column(Column::from_name(OTLP_AGGREGATION_TEMPORALITY_LABEL))
4188 .eq(lit(GREPTIME_TEMPORALITY_DELTA))
4189 })
4190 } else {
4191 None
4192 };
4193 let Some(marker_is_delta) = marker_is_delta else {
4194 return Ok(cumulative);
4195 };
4196
4197 let delta = if function == "rate" {
4198 DfExpr::BinaryExpr(BinaryExpr {
4199 left: Box::new(delta_sum),
4200 op: Operator::Divide,
4201 right: Box::new(lit(range_length as f64 / 1000.0)),
4202 })
4203 } else {
4204 delta_sum
4205 };
4206 let display_name = cumulative.schema_name().to_string();
4207 when(marker_is_delta, delta)
4208 .otherwise(cumulative)
4209 .context(DataFusionPlanningSnafu)
4210 .map(|expr| expr.alias(display_name))
4211 }
4212
4213 fn validate_label_name(label_name: &str) -> Result<()> {
4217 if label_name.starts_with("__") {
4219 return InvalidDestinationLabelNameSnafu { label_name }.fail();
4220 }
4221 if !LABEL_NAME_REGEX.is_match(label_name) {
4223 return InvalidDestinationLabelNameSnafu { label_name }.fail();
4224 }
4225
4226 Ok(())
4227 }
4228
4229 fn build_regexp_replace_label_expr(
4231 &self,
4232 other_input_exprs: &mut VecDeque<DfExpr>,
4233 input_schema: &DFSchemaRef,
4234 ) -> Result<Option<(DfExpr, String)>> {
4235 let dst_label = match other_input_exprs.pop_front() {
4237 Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4238 other => UnexpectedPlanExprSnafu {
4239 desc: format!("expected dst_label string literal, but found {:?}", other),
4240 }
4241 .fail()?,
4242 };
4243
4244 Self::validate_label_name(&dst_label)?;
4246 let replacement = match other_input_exprs.pop_front() {
4247 Some(DfExpr::Literal(ScalarValue::Utf8(Some(r)), _)) => r,
4248 other => UnexpectedPlanExprSnafu {
4249 desc: format!("expected replacement string literal, but found {:?}", other),
4250 }
4251 .fail()?,
4252 };
4253 let src_label = match other_input_exprs.pop_front() {
4254 Some(DfExpr::Literal(ScalarValue::Utf8(Some(s)), None)) => s,
4255 other => UnexpectedPlanExprSnafu {
4256 desc: format!("expected src_label string literal, but found {:?}", other),
4257 }
4258 .fail()?,
4259 };
4260
4261 let regex = match other_input_exprs.pop_front() {
4262 Some(DfExpr::Literal(ScalarValue::Utf8(Some(r)), None)) => r,
4263 other => UnexpectedPlanExprSnafu {
4264 desc: format!("expected regex string literal, but found {:?}", other),
4265 }
4266 .fail()?,
4267 };
4268
4269 let anchored = format!("^(?s:{regex})$");
4273 let compiled = regex::Regex::new(&anchored)
4274 .map_err(|_| InvalidRegularExpressionSnafu { regex }.build())?;
4275 let dst_exists = self.ctx.tag_columns.contains(&dst_label);
4276
4277 if !self.ctx.tag_columns.contains(&src_label) {
4280 let Some(captures) = compiled.captures("") else {
4281 return Ok(None);
4282 };
4283 let mut value = String::new();
4284 captures.expand(&replacement, &mut value);
4285 if value.is_empty() {
4286 if !dst_exists {
4289 return Ok(None);
4290 }
4291 return Ok(Some((
4292 lit(ScalarValue::Utf8(None)).alias(&dst_label),
4293 dst_label,
4294 )));
4295 }
4296 return Ok(Some((lit(value).alias(&dst_label), dst_label)));
4297 }
4298
4299 let src = Self::label_value_expr(&src_label, input_schema)?;
4300 let matched = DfExpr::ScalarFunction(ScalarFunction {
4301 func: datafusion_functions::regex::regexp_like(),
4302 args: vec![src.clone(), lit(anchored.clone())],
4303 });
4304 let replaced = Self::empty_label_to_null(DfExpr::ScalarFunction(ScalarFunction {
4305 func: datafusion_functions::regex::regexp_replace(),
4306 args: vec![src, lit(anchored), lit(replacement)],
4307 }));
4308 let unchanged = if dst_exists {
4309 DfExpr::Column(Column::from_name(&dst_label))
4310 .cast_to(&ArrowDataType::Utf8, input_schema)
4311 .context(DataFusionPlanningSnafu)?
4312 } else {
4313 lit(ScalarValue::Utf8(None))
4314 };
4315 let replace_expr = when(matched, replaced)
4316 .otherwise(unchanged)
4317 .context(DataFusionPlanningSnafu)?;
4318
4319 Ok(Some((replace_expr.alias(&dst_label), dst_label)))
4320 }
4321
4322 fn build_concat_labels_expr(
4324 other_input_exprs: &mut VecDeque<DfExpr>,
4325 ctx: &PromPlannerContext,
4326 input_schema: &DFSchemaRef,
4327 query_engine_state: &QueryEngineState,
4328 ) -> Result<(DfExpr, String)> {
4329 let dst_label = match other_input_exprs.pop_front() {
4332 Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4333 other => UnexpectedPlanExprSnafu {
4334 desc: format!("expected dst_label string literal, but found {:?}", other),
4335 }
4336 .fail()?,
4337 };
4338 let separator = match other_input_exprs.pop_front() {
4339 Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4340 other => UnexpectedPlanExprSnafu {
4341 desc: format!("expected separator string literal, but found {:?}", other),
4342 }
4343 .fail()?,
4344 };
4345
4346 let available_columns: HashSet<&str> = ctx
4348 .tag_columns
4349 .iter()
4350 .chain(ctx.field_columns.iter())
4351 .chain(ctx.time_index_column.as_ref())
4352 .map(|s| s.as_str())
4353 .collect();
4354
4355 let src_labels = other_input_exprs
4356 .iter()
4357 .map(|expr| {
4358 match expr {
4361 DfExpr::Literal(ScalarValue::Utf8(Some(label)), None) => {
4362 if label.is_empty() {
4363 FunctionInvalidArgumentSnafu {
4364 fn_name: "label_join",
4365 }
4366 .fail()
4367 } else if available_columns.contains(label.as_str()) {
4368 Self::label_value_expr(label, input_schema)
4369 } else {
4370 Ok(lit(""))
4371 }
4372 }
4373 other => UnexpectedPlanExprSnafu {
4374 desc: format!(
4375 "expected source label string literal, but found {:?}",
4376 other
4377 ),
4378 }
4379 .fail(),
4380 }
4381 })
4382 .collect::<Result<Vec<_>>>()?;
4383 if src_labels.is_empty() {
4385 return Ok((lit(ScalarValue::Utf8(None)).alias(&dst_label), dst_label));
4386 }
4387
4388 let session_state = query_engine_state.session_state();
4389 let func = session_state
4390 .scalar_functions()
4391 .get("concat_ws")
4392 .context(UnsupportedExprSnafu { name: "concat_ws" })?;
4393
4394 let mut args = Vec::with_capacity(1 + src_labels.len());
4396 args.push(DfExpr::Literal(ScalarValue::Utf8(Some(separator)), None));
4397 args.extend(src_labels);
4398
4399 Ok((
4400 Self::empty_label_to_null(DfExpr::ScalarFunction(ScalarFunction {
4401 func: func.clone(),
4402 args,
4403 }))
4404 .alias(&dst_label),
4405 dst_label,
4406 ))
4407 }
4408
4409 fn label_value_expr(label: &str, input_schema: &DFSchemaRef) -> Result<DfExpr> {
4412 let value = DfExpr::Column(Column::from_name(label))
4413 .cast_to(&ArrowDataType::Utf8, input_schema)
4414 .context(DataFusionPlanningSnafu)?;
4415 Ok(DfExpr::ScalarFunction(ScalarFunction {
4416 func: coalesce(),
4417 args: vec![value, lit("")],
4418 }))
4419 }
4420
4421 fn empty_label_to_null(value: DfExpr) -> DfExpr {
4425 DfExpr::ScalarFunction(ScalarFunction {
4426 func: datafusion_functions::core::nullif(),
4427 args: vec![value, lit("")],
4428 })
4429 }
4430
4431 fn create_time_index_column_expr(&self) -> Result<DfExpr> {
4432 Ok(DfExpr::Column(Column::from_name(
4433 self.ctx
4434 .time_index_column
4435 .clone()
4436 .with_context(|| TimeIndexNotFoundSnafu { table: "unknown" })?,
4437 )))
4438 }
4439
4440 fn create_range_eval_ts_expr(
4455 &self,
4456 fold_offset: Millisecond,
4457 input_schema: &DFSchemaRef,
4458 ) -> Result<DfExpr> {
4459 let eval_ts = self
4460 .create_time_index_column_expr()?
4461 .cast_to(
4462 &ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
4463 input_schema,
4464 )
4465 .context(DataFusionPlanningSnafu)?
4466 .cast_to(&ArrowDataType::Int64, input_schema)
4467 .context(DataFusionPlanningSnafu)?;
4468 DfExpr::BinaryExpr(BinaryExpr {
4469 left: Box::new(eval_ts),
4470 op: Operator::Plus,
4471 right: Box::new(lit(fold_offset)),
4472 })
4473 .cast_to(
4474 &ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
4475 input_schema,
4476 )
4477 .context(DataFusionPlanningSnafu)
4478 }
4479
4480 fn name_without_last_arg(expr: &DfExpr) -> String {
4487 if let DfExpr::ScalarFunction(ScalarFunction { func, args }) = expr
4488 && let Some((_, visible_args)) = args.split_last()
4489 {
4490 let visible = ScalarFunction {
4491 func: func.clone(),
4492 args: visible_args.to_vec(),
4493 };
4494 return DfExpr::ScalarFunction(visible).schema_name().to_string();
4495 }
4496
4497 expr.schema_name().to_string()
4498 }
4499
4500 fn create_tag_column_exprs(&self) -> Result<Vec<DfExpr>> {
4501 let mut result = Vec::with_capacity(self.ctx.tag_columns.len());
4502 for tag in &self.ctx.tag_columns {
4503 let expr = DfExpr::Column(Column::from_name(tag));
4504 result.push(expr);
4505 }
4506 Ok(result)
4507 }
4508
4509 fn create_field_column_exprs(&self) -> Result<Vec<DfExpr>> {
4510 let mut result = Vec::with_capacity(self.ctx.field_columns.len());
4511 for field in &self.ctx.field_columns {
4512 let expr = DfExpr::Column(Column::from_name(field));
4513 result.push(expr);
4514 }
4515 Ok(result)
4516 }
4517
4518 fn create_tag_and_time_index_column_sort_exprs(&self) -> Result<Vec<SortExpr>> {
4519 let mut result = self
4520 .ctx
4521 .tag_columns
4522 .iter()
4523 .map(|col| DfExpr::Column(Column::from_name(col)).sort(true, true))
4524 .collect::<Vec<_>>();
4525 result.push(self.create_time_index_column_expr()?.sort(true, true));
4526 Ok(result)
4527 }
4528
4529 fn create_field_columns_sort_exprs(&self, asc: bool) -> Vec<SortExpr> {
4530 self.ctx
4531 .field_columns
4532 .iter()
4533 .map(|col| DfExpr::Column(Column::from_name(col)).sort(asc, true))
4534 .collect::<Vec<_>>()
4535 }
4536
4537 fn create_sort_exprs_by_tags(
4538 func: &str,
4539 tags: Vec<DfExpr>,
4540 asc: bool,
4541 ) -> Result<Vec<SortExpr>> {
4542 ensure!(
4543 !tags.is_empty(),
4544 FunctionInvalidArgumentSnafu { fn_name: func }
4545 );
4546
4547 tags.iter()
4548 .map(|col| match col {
4549 DfExpr::Literal(ScalarValue::Utf8(Some(label)), _) => {
4550 Ok(DfExpr::Column(Column::from_name(label)).sort(asc, false))
4551 }
4552 other => UnexpectedPlanExprSnafu {
4553 desc: format!("expected label string literal, but found {:?}", other),
4554 }
4555 .fail(),
4556 })
4557 .collect::<Result<Vec<_>>>()
4558 }
4559
4560 fn create_empty_values_filter_expr(&self, preserve_any_value: bool) -> Result<DfExpr> {
4561 let mut exprs = Vec::with_capacity(self.ctx.field_columns.len());
4562 for value in &self.ctx.field_columns {
4563 let expr = DfExpr::Column(Column::from_name(value)).is_not_null();
4564 exprs.push(expr);
4565 }
4566
4567 let predicate = if preserve_any_value {
4572 disjunction(exprs)
4573 } else {
4574 conjunction(exprs)
4575 };
4576 predicate.with_context(|| ValueNotFoundSnafu {
4577 table: self
4578 .table_ref()
4579 .map(|t| t.to_quoted_string())
4580 .unwrap_or_else(|_| "unknown".to_string()),
4581 })
4582 }
4583
4584 fn create_aggregate_exprs(
4601 &mut self,
4602 op: TokenType,
4603 param: &Option<Box<PromExpr>>,
4604 input_plan: &LogicalPlan,
4605 ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4606 let mixed_sample_columns =
4607 Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns)
4608 .map(|(float, histogram)| (float.to_string(), histogram.to_string()));
4609 let is_group_agg = op.id() == token::T_GROUP;
4610 if is_group_agg && mixed_sample_columns.is_none() {
4611 ensure!(
4612 self.ctx.field_columns.len() == 1,
4613 MultiFieldsNotSupportedSnafu {
4614 operator: "group()"
4615 }
4616 );
4617 }
4618
4619 if let Some((float, histogram)) = mixed_sample_columns {
4620 return self.create_mixed_aggregate_exprs(op, param, &float, &histogram);
4621 }
4622
4623 if self.all_field_columns_are_native_histograms(input_plan.schema()) {
4624 return self.create_native_histogram_aggregate_exprs(op, input_plan);
4625 }
4626
4627 let exprs = self
4629 .ctx
4630 .field_columns
4631 .iter()
4632 .map(|col| {
4633 Self::create_numeric_aggregate_expr(
4634 op,
4635 param,
4636 DfExpr::Column(Column::from_name(col)),
4637 )
4638 })
4639 .collect::<Result<Vec<_>>>()?;
4640
4641 let prev_field_exprs = if op.id() == token::T_COUNT_VALUES {
4649 ensure!(
4650 self.ctx.field_columns.len() == 1,
4651 UnsupportedExprSnafu {
4652 name: "count_values on multi-value input"
4653 }
4654 );
4655
4656 self.ctx
4657 .field_columns
4658 .iter()
4659 .map(|col| {
4660 let value = DfExpr::Column(Column::from_name(col));
4661 let value = if Self::field_column_type(input_plan.schema(), col)
4666 == Some(&ArrowDataType::Float64)
4667 {
4668 value
4669 } else {
4670 DfExpr::Cast(Cast::new(Box::new(value), ArrowDataType::Float64))
4671 };
4672 DfExpr::ScalarFunction(ScalarFunction {
4673 func: Arc::new(PromqlFloatToString::scalar_udf()),
4674 args: vec![value],
4675 })
4676 })
4677 .collect()
4678 } else {
4679 vec![]
4680 };
4681
4682 let exprs = normalize_cols(exprs, input_plan)
4685 .context(DataFusionPlanningSnafu)?
4686 .into_iter()
4687 .map(Self::alias_extremum_as_builtin)
4688 .collect::<Vec<_>>();
4689
4690 self.ctx.field_columns = exprs
4692 .iter()
4693 .map(|expr| expr.schema_name().to_string())
4694 .collect();
4695
4696 Ok((exprs, prev_field_exprs))
4697 }
4698
4699 fn alias_extremum_as_builtin(expr: DfExpr) -> DfExpr {
4702 let DfExpr::AggregateFunction(aggr) = &expr else {
4703 return expr;
4704 };
4705 let builtin = match aggr.func.name() {
4706 Extremum::MIN_NAME => min_udaf(),
4707 Extremum::MAX_NAME => max_udaf(),
4708 _ => return expr,
4709 };
4710 let name = builtin
4711 .call(aggr.params.args.clone())
4712 .schema_name()
4713 .to_string();
4714 expr.alias(name)
4715 }
4716
4717 fn create_numeric_aggregate_expr(
4718 op: TokenType,
4719 param: &Option<Box<PromExpr>>,
4720 input: DfExpr,
4721 ) -> Result<DfExpr> {
4722 let expr = match op.id() {
4723 token::T_SUM => sum_udaf().call(vec![input]),
4724 token::T_QUANTILE => {
4725 let q = Self::get_param_as_literal_expr(
4726 param.as_deref(),
4727 Some(op),
4728 Some(ArrowDataType::Float64),
4729 )?;
4730 quantile_udaf().call(vec![q, input])
4731 }
4732 token::T_AVG => avg_udaf().call(vec![input]),
4733 token::T_COUNT_VALUES | token::T_COUNT => count_udaf().call(vec![input]),
4734 token::T_MIN => Extremum::min_udaf().call(vec![input]),
4735 token::T_MAX => Extremum::max_udaf().call(vec![input]),
4736 token::T_GROUP => max_udaf().call(vec![lit(1_f64)]),
4739 token::T_STDDEV => stddev_pop_udaf().call(vec![input]),
4740 token::T_STDVAR => var_pop_udaf().call(vec![input]),
4741 token::T_TOPK | token::T_BOTTOMK => {
4742 return UnsupportedExprSnafu {
4743 name: format!("{op:?}"),
4744 }
4745 .fail();
4746 }
4747 _ => return UnexpectedTokenSnafu { token: op }.fail(),
4748 };
4749 Ok(expr)
4750 }
4751
4752 fn create_mixed_aggregate_exprs(
4753 &mut self,
4754 op: TokenType,
4755 param: &Option<Box<PromExpr>>,
4756 float_column: &str,
4757 histogram_column: &str,
4758 ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4759 let float_input = DfExpr::Column(Column::from_name(float_column));
4760 let histogram_input = DfExpr::Column(Column::from_name(histogram_column));
4761 let float_count = count_udaf().call(vec![float_input.clone()]);
4762 let histogram_count = count_udaf().call(vec![histogram_input.clone()]);
4763 let mixed_sample_value = || {
4764 DfExpr::ScalarFunction(ScalarFunction {
4765 func: coalesce(),
4766 args: vec![
4767 DfExpr::ScalarFunction(ScalarFunction {
4768 func: Arc::new(PromqlFloatToString::scalar_udf()),
4769 args: vec![float_input.clone()],
4770 }),
4771 DfExpr::ScalarFunction(ScalarFunction {
4772 func: Arc::new(NativeHistogramToString::scalar_udf()),
4773 args: vec![histogram_input.clone()],
4774 }),
4775 ],
4776 })
4777 };
4778
4779 let (exprs, prev_field_exprs, field_columns) = match op.id() {
4780 token::T_SUM | token::T_AVG => (
4781 vec![
4782 Self::create_numeric_aggregate_expr(op, param, float_input)?
4783 .alias(float_column),
4784 self.create_native_histogram_aggregate_expr(op, histogram_column)?,
4785 float_count.alias(Self::mixed_sample_count_name(float_column)),
4786 histogram_count.alias(Self::mixed_sample_count_name(histogram_column)),
4787 ],
4788 vec![],
4789 vec![float_column.to_string(), histogram_column.to_string()],
4790 ),
4791 token::T_COUNT => {
4792 let present = when(
4793 float_input
4794 .clone()
4795 .is_not_null()
4796 .or(histogram_input.clone().is_not_null()),
4797 lit(1_i64),
4798 )
4799 .otherwise(lit(ScalarValue::Int64(None)))
4800 .context(DataFusionPlanningSnafu)?;
4801 (
4802 vec![count_udaf().call(vec![present]).alias(float_column)],
4803 vec![],
4804 vec![float_column.to_string()],
4805 )
4806 }
4807 token::T_GROUP => (
4808 vec![max_udaf().call(vec![lit(1_f64)]).alias(float_column)],
4809 vec![],
4810 vec![float_column.to_string()],
4811 ),
4812 token::T_COUNT_VALUES => {
4813 let value = mixed_sample_value();
4814 (
4815 vec![count_udaf().call(vec![value.clone()]).alias(float_column)],
4816 vec![value],
4817 vec![float_column.to_string()],
4818 )
4819 }
4820 token::T_MIN | token::T_MAX | token::T_STDDEV | token::T_STDVAR | token::T_QUANTILE => {
4821 (
4822 vec![
4823 Self::create_numeric_aggregate_expr(op, param, float_input)?
4824 .alias(float_column),
4825 histogram_count.alias(Self::mixed_sample_count_name(histogram_column)),
4826 ],
4827 vec![],
4828 vec![float_column.to_string()],
4829 )
4830 }
4831 token::T_TOPK | token::T_BOTTOMK => {
4832 return UnsupportedExprSnafu {
4833 name: format!("{op:?}"),
4834 }
4835 .fail();
4836 }
4837 _ => return UnexpectedTokenSnafu { token: op }.fail(),
4838 };
4839
4840 self.ctx.field_columns = field_columns;
4841 Ok((exprs, prev_field_exprs))
4842 }
4843
4844 fn mixed_sample_count_column(column: &str) -> DfExpr {
4845 DfExpr::Column(Column::from_name(Self::mixed_sample_count_name(column)))
4846 }
4847
4848 fn mixed_sample_count_name(column: &str) -> String {
4849 format!("__promql_sample_count({column})")
4850 }
4851
4852 fn mixed_aggregate_filter_expr(
4853 &self,
4854 op: TokenType,
4855 float_column: &str,
4856 histogram_column: &str,
4857 ) -> Result<DfExpr> {
4858 let float_count = Self::mixed_sample_count_column(float_column);
4859 let histogram_count = Self::mixed_sample_count_column(histogram_column);
4860 let mixed = float_count
4861 .clone()
4862 .gt(lit(0_i64))
4863 .and(histogram_count.clone().gt(lit(0_i64)));
4864 let drop_mixed = DfExpr::ScalarFunction(ScalarFunction {
4865 func: Arc::new(NativeHistogramDrop::warning_bool_false_udf(
4866 format!(
4867 "{op}: dropped aggregation result containing both float and native histogram samples"
4868 ),
4869 self.promql_annotations.clone(),
4870 )),
4871 args: vec![float_count, histogram_count],
4872 });
4873
4874 when(mixed, drop_mixed)
4875 .otherwise(lit(true))
4876 .context(DataFusionPlanningSnafu)
4877 }
4878
4879 fn mixed_ignored_histogram_filter_expr(
4880 &self,
4881 op: TokenType,
4882 histogram_column: &str,
4883 ) -> Result<DfExpr> {
4884 let histogram_count = Self::mixed_sample_count_column(histogram_column);
4885 let has_histograms = histogram_count.clone().gt(lit(0_i64));
4886 let record_info = DfExpr::ScalarFunction(ScalarFunction {
4887 func: Arc::new(NativeHistogramDrop::bool_true_udf(
4888 format!(
4889 "{op}: dropped native histogram samples because this aggregation is not supported for native histograms"
4890 ),
4891 self.promql_annotations.clone(),
4892 )),
4893 args: vec![histogram_count],
4894 });
4895
4896 when(has_histograms, record_info)
4897 .otherwise(lit(true))
4898 .context(DataFusionPlanningSnafu)
4899 }
4900
4901 fn create_native_histogram_aggregate_expr(
4902 &self,
4903 op: TokenType,
4904 column: &str,
4905 ) -> Result<DfExpr> {
4906 let input = DfExpr::Column(Column::from_name(column));
4907 let expr = match op.id() {
4908 token::T_SUM => Arc::new(NativeHistogramAggSum::aggregate_udf_with_collector(
4909 self.promql_annotations.clone(),
4910 ))
4911 .call(vec![input])
4912 .alias(column),
4913 token::T_AVG => Arc::new(NativeHistogramAggAvg::aggregate_udf_with_collector(
4914 self.promql_annotations.clone(),
4915 ))
4916 .call(vec![input])
4917 .alias(column),
4918 token::T_COUNT_VALUES | token::T_COUNT => {
4919 count_udaf().call(vec![input]).alias(column)
4920 }
4921 token::T_GROUP => max_udaf().call(vec![lit(1_f64)]).alias(column),
4922 token::T_MIN
4923 | token::T_MAX
4924 | token::T_STDDEV
4925 | token::T_STDVAR
4926 | token::T_QUANTILE
4927 | token::T_TOPK
4928 | token::T_BOTTOMK => sum_udaf()
4929 .call(vec![DfExpr::ScalarFunction(ScalarFunction {
4930 func: Arc::new(NativeHistogramDrop::float_null_udf(
4931 format!(
4932 "{op}: dropped native histogram samples because this aggregation is not supported for native histograms"
4933 ),
4934 self.promql_annotations.clone(),
4935 )),
4936 args: vec![input],
4937 })])
4938 .alias(column),
4939 _ => return UnexpectedTokenSnafu { token: op }.fail(),
4940 };
4941 Ok(expr)
4942 }
4943
4944 fn create_native_histogram_aggregate_exprs(
4945 &mut self,
4946 op: TokenType,
4947 input_plan: &LogicalPlan,
4948 ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4949 let prev_field_exprs = if op.id() == token::T_COUNT_VALUES {
4950 ensure!(
4951 self.ctx.field_columns.len() == 1,
4952 UnsupportedExprSnafu {
4953 name: "count_values on multi-value input"
4954 }
4955 );
4956 self.ctx
4957 .field_columns
4958 .iter()
4959 .map(|col| {
4960 DfExpr::ScalarFunction(ScalarFunction {
4961 func: Arc::new(NativeHistogramToString::scalar_udf()),
4962 args: vec![DfExpr::Column(Column::from_name(col))],
4963 })
4964 })
4965 .collect::<Vec<_>>()
4966 } else {
4967 vec![]
4968 };
4969
4970 let exprs = self
4971 .ctx
4972 .field_columns
4973 .iter()
4974 .map(|col| self.create_native_histogram_aggregate_expr(op, col))
4975 .collect::<Result<Vec<_>>>()?;
4976
4977 let normalized_exprs =
4978 normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
4979 self.ctx.field_columns = normalized_exprs
4980 .into_iter()
4981 .map(|expr| expr.schema_name().to_string())
4982 .collect();
4983
4984 Ok((exprs, prev_field_exprs))
4985 }
4986
4987 fn get_param_value_as_str(op: TokenType, param: &Option<Box<PromExpr>>) -> Result<&str> {
4988 let param = param
4989 .as_deref()
4990 .with_context(|| FunctionInvalidArgumentSnafu {
4991 fn_name: op.to_string(),
4992 })?;
4993 let PromExpr::StringLiteral(StringLiteral { val }) = param else {
4994 return FunctionInvalidArgumentSnafu {
4995 fn_name: op.to_string(),
4996 }
4997 .fail();
4998 };
4999
5000 Ok(val)
5001 }
5002
5003 fn get_param_as_literal_expr(
5004 param: Option<&PromExpr>,
5005 op: Option<TokenType>,
5006 expected_type: Option<ArrowDataType>,
5007 ) -> Result<DfExpr> {
5008 let prom_param = param.with_context(|| {
5009 if let Some(op) = op {
5010 FunctionInvalidArgumentSnafu {
5011 fn_name: op.to_string(),
5012 }
5013 } else {
5014 FunctionInvalidArgumentSnafu {
5015 fn_name: "unknown".to_string(),
5016 }
5017 }
5018 })?;
5019
5020 let expr = Self::try_build_literal_expr(prom_param).with_context(|| {
5021 if let Some(op) = op {
5022 FunctionInvalidArgumentSnafu {
5023 fn_name: op.to_string(),
5024 }
5025 } else {
5026 FunctionInvalidArgumentSnafu {
5027 fn_name: "unknown".to_string(),
5028 }
5029 }
5030 })?;
5031
5032 if let Some(expected_type) = expected_type {
5034 let expr_type = expr
5036 .get_type(&DFSchema::empty())
5037 .context(DataFusionPlanningSnafu)?;
5038 if expected_type != expr_type {
5039 return FunctionInvalidArgumentSnafu {
5040 fn_name: format!("expected {expected_type:?}, but found {expr_type:?}"),
5041 }
5042 .fail();
5043 }
5044 }
5045
5046 Ok(expr)
5047 }
5048
5049 fn create_window_exprs(
5052 &mut self,
5053 op: TokenType,
5054 group_exprs: Vec<DfExpr>,
5055 input_plan: &LogicalPlan,
5056 ) -> Result<Vec<DfExpr>> {
5057 ensure!(
5058 self.ctx.field_columns.len() == 1,
5059 UnsupportedExprSnafu {
5060 name: "topk or bottomk on multi-value input"
5061 }
5062 );
5063
5064 assert!(matches!(op.id(), token::T_TOPK | token::T_BOTTOMK));
5065
5066 let asc = matches!(op.id(), token::T_BOTTOMK);
5067
5068 let tag_sort_exprs = self
5069 .create_tag_column_exprs()?
5070 .into_iter()
5071 .map(|expr| expr.sort(asc, true));
5072
5073 let exprs: Vec<DfExpr> = self
5075 .ctx
5076 .field_columns
5077 .iter()
5078 .map(|col| {
5079 let mut sort_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 2);
5080 sort_exprs.push(
5083 datafusion_functions::math::expr_fn::isnan(DfExpr::Column(Column::from_name(
5084 col,
5085 )))
5086 .sort(true, true),
5087 );
5088 sort_exprs.push(DfExpr::Column(Column::from_name(col)).sort(asc, true));
5090 sort_exprs.extend(tag_sort_exprs.clone());
5093
5094 DfExpr::WindowFunction(Box::new(WindowFunction {
5095 fun: WindowFunctionDefinition::WindowUDF(Arc::new(RowNumber::new().into())),
5096 params: WindowFunctionParams {
5097 args: vec![],
5098 partition_by: group_exprs.clone(),
5099 order_by: sort_exprs,
5100 window_frame: WindowFrame::new(Some(true)),
5101 null_treatment: None,
5102 distinct: false,
5103 filter: None,
5104 },
5105 }))
5106 })
5107 .collect();
5108
5109 let normalized_exprs =
5110 normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
5111 Ok(normalized_exprs)
5112 }
5113
5114 fn try_build_literal_expr(expr: &PromExpr) -> Option<DfExpr> {
5117 match expr {
5118 PromExpr::NumberLiteral(NumberLiteral { val }) => Some(val.lit()),
5119 PromExpr::StringLiteral(StringLiteral { val }) => Some(val.lit()),
5120 PromExpr::VectorSelector(_)
5121 | PromExpr::MatrixSelector(_)
5122 | PromExpr::Extension(_)
5123 | PromExpr::Aggregate(_)
5124 | PromExpr::Subquery(_) => None,
5125 PromExpr::Call(Call { func, .. }) => {
5126 if func.name == SPECIAL_TIME_FUNCTION {
5127 None
5130 } else {
5131 None
5132 }
5133 }
5134 PromExpr::Paren(ParenExpr { expr }) => Self::try_build_literal_expr(expr),
5135 PromExpr::Unary(UnaryExpr { expr, .. }) => Some(DfExpr::Negative(Box::new(
5136 Self::try_build_literal_expr(expr)?,
5137 ))),
5138 PromExpr::Binary(PromBinaryExpr {
5139 lhs,
5140 rhs,
5141 op,
5142 modifier,
5143 }) => {
5144 let lhs = Self::try_build_literal_expr(lhs)?;
5145 let rhs = Self::try_build_literal_expr(rhs)?;
5146 let is_comparison_op = Self::is_token_a_comparison_op(*op);
5147 let expr_builder = Self::prom_token_to_binary_expr_builder(*op).ok()?;
5148 let expr = expr_builder(lhs, rhs).ok()?;
5149
5150 let should_return_bool = if let Some(m) = modifier {
5151 m.return_bool
5152 } else {
5153 false
5154 };
5155 if is_comparison_op && should_return_bool {
5156 Some(DfExpr::Cast(Cast::new(
5157 Box::new(expr),
5158 ArrowDataType::Float64,
5159 )))
5160 } else {
5161 Some(expr)
5162 }
5163 }
5164 }
5165 }
5166
5167 fn try_build_special_time_expr_with_context(&self, expr: &PromExpr) -> Option<DfExpr> {
5168 match expr {
5169 PromExpr::Call(Call { func, .. }) => {
5170 if func.name == SPECIAL_TIME_FUNCTION
5171 && let Some(time_index_col) = self.ctx.time_index_column.as_ref()
5172 {
5173 Some(build_special_time_expr(time_index_col))
5174 } else {
5175 None
5176 }
5177 }
5178 _ => None,
5179 }
5180 }
5181
5182 fn native_histogram_binary_expr(
5183 token: TokenType,
5184 lhs: DfExpr,
5185 lhs_is_histogram: bool,
5186 rhs: DfExpr,
5187 rhs_is_histogram: bool,
5188 filter_context: bool,
5189 promql_annotations: Option<PromqlAnnotationCollector>,
5190 ) -> Result<Option<DfExpr>> {
5191 if !lhs_is_histogram && !rhs_is_histogram {
5192 return Ok(None);
5193 }
5194
5195 let scalar_fn = |func: ScalarUdfDef, args| {
5196 DfExpr::ScalarFunction(ScalarFunction {
5197 func: Arc::new(func),
5198 args,
5199 })
5200 };
5201 let invalid_expr = || {
5202 let message = format!(
5203 "{}: dropped native histogram samples because this binary operation is not supported for native histograms",
5204 token
5205 );
5206 let func = if filter_context {
5207 NativeHistogramDrop::bool_false_udf(message, promql_annotations.clone())
5208 } else {
5209 NativeHistogramDrop::float_null_udf(message, promql_annotations.clone())
5210 };
5211 let args = vec![lhs.clone(), rhs.clone()];
5212 scalar_fn(func, args)
5213 };
5214
5215 let expr = match (token.id(), lhs_is_histogram, rhs_is_histogram) {
5216 (token::T_ADD, true, true) => scalar_fn(
5217 NativeHistogramAdd::scalar_udf_with_collector(promql_annotations.clone()),
5218 vec![lhs, rhs],
5219 ),
5220 (token::T_SUB, true, true) => scalar_fn(
5221 NativeHistogramSub::scalar_udf_with_collector(promql_annotations.clone()),
5222 vec![lhs, rhs],
5223 ),
5224 (token::T_MUL, true, false) => {
5225 scalar_fn(NativeHistogramMulScalar::scalar_udf(), vec![lhs, rhs])
5226 }
5227 (token::T_MUL, false, true) => {
5228 scalar_fn(NativeHistogramScalarMul::scalar_udf(), vec![lhs, rhs])
5229 }
5230 (token::T_DIV, true, false) => {
5231 scalar_fn(NativeHistogramDivScalar::scalar_udf(), vec![lhs, rhs])
5232 }
5233 (token::T_EQLC, true, true) => {
5234 scalar_fn(NativeHistogramEq::scalar_udf(), vec![lhs, rhs])
5235 }
5236 (token::T_NEQ, true, true) => {
5237 scalar_fn(NativeHistogramNotEq::scalar_udf(), vec![lhs, rhs])
5238 }
5239 _ => invalid_expr(),
5240 };
5241
5242 Ok(Some(expr))
5243 }
5244
5245 #[allow(clippy::type_complexity)]
5248 fn prom_token_to_binary_expr_builder(
5249 token: TokenType,
5250 ) -> Result<Box<dyn Fn(DfExpr, DfExpr) -> Result<DfExpr>>> {
5251 let cast_float = |expr| {
5252 if matches!(
5253 &expr,
5254 DfExpr::Cast(Cast { field, .. }) if field.data_type() == &ArrowDataType::Float64
5255 ) || matches!(&expr, DfExpr::Literal(ScalarValue::Float64(_), _))
5256 {
5257 expr
5258 } else {
5259 DfExpr::Cast(Cast::new(Box::new(expr), ArrowDataType::Float64))
5260 }
5261 };
5262 match token.id() {
5263 token::T_ADD => Ok(Box::new(move |lhs, rhs| {
5264 Ok(cast_float(lhs) + cast_float(rhs))
5265 })),
5266 token::T_SUB => Ok(Box::new(move |lhs, rhs| {
5267 Ok(cast_float(lhs) - cast_float(rhs))
5268 })),
5269 token::T_MUL => Ok(Box::new(move |lhs, rhs| {
5270 Ok(cast_float(lhs) * cast_float(rhs))
5271 })),
5272 token::T_DIV => Ok(Box::new(move |lhs, rhs| {
5273 Ok(cast_float(lhs) / cast_float(rhs))
5274 })),
5275 token::T_MOD => Ok(Box::new(move |lhs: DfExpr, rhs| {
5276 Ok(cast_float(lhs) % cast_float(rhs))
5277 })),
5278 token::T_EQLC => Ok(Box::new(|lhs, rhs| {
5279 Ok(Self::ieee_comparison(lhs, Operator::Eq, rhs))
5280 })),
5281 token::T_NEQ => Ok(Box::new(|lhs, rhs| {
5282 Ok(Self::ieee_comparison(lhs, Operator::NotEq, rhs))
5283 })),
5284 token::T_GTR => Ok(Box::new(|lhs, rhs| {
5285 Ok(Self::ieee_comparison(lhs, Operator::Gt, rhs))
5286 })),
5287 token::T_LSS => Ok(Box::new(|lhs, rhs| {
5288 Ok(Self::ieee_comparison(lhs, Operator::Lt, rhs))
5289 })),
5290 token::T_GTE => Ok(Box::new(|lhs, rhs| {
5291 Ok(Self::ieee_comparison(lhs, Operator::GtEq, rhs))
5292 })),
5293 token::T_LTE => Ok(Box::new(|lhs, rhs| {
5294 Ok(Self::ieee_comparison(lhs, Operator::LtEq, rhs))
5295 })),
5296 token::T_POW => Ok(Box::new(move |lhs, rhs| {
5297 Ok(DfExpr::ScalarFunction(ScalarFunction {
5298 func: datafusion_functions::math::power(),
5299 args: vec![cast_float(lhs), cast_float(rhs)],
5300 }))
5301 })),
5302 token::T_ATAN2 => Ok(Box::new(move |lhs, rhs| {
5303 Ok(DfExpr::ScalarFunction(ScalarFunction {
5304 func: datafusion_functions::math::atan2(),
5305 args: vec![cast_float(lhs), cast_float(rhs)],
5306 }))
5307 })),
5308 _ => UnexpectedTokenSnafu { token }.fail(),
5309 }
5310 }
5311
5312 fn ieee_comparison(lhs: DfExpr, op: Operator, rhs: DfExpr) -> DfExpr {
5314 DfExpr::ScalarFunction(ScalarFunction {
5315 func: Arc::new(IeeeComparison::scalar_udf(op)),
5316 args: vec![lhs, rhs],
5317 })
5318 }
5319
5320 fn is_token_a_comparison_op(token: TokenType) -> bool {
5322 matches!(
5323 token.id(),
5324 token::T_EQLC
5325 | token::T_NEQ
5326 | token::T_GTR
5327 | token::T_LSS
5328 | token::T_GTE
5329 | token::T_LTE
5330 )
5331 }
5332
5333 fn is_token_a_set_op(token: TokenType) -> bool {
5335 matches!(
5336 token.id(),
5337 token::T_LAND | token::T_LOR | token::T_LUNLESS )
5341 }
5342
5343 fn align_binary_field_columns<'a>(
5344 left_schema: &DFSchemaRef,
5345 right_schema: &DFSchemaRef,
5346 left_field_columns: &'a [String],
5347 right_field_columns: &'a [String],
5348 op: TokenType,
5349 left_is_scalar: bool,
5350 right_is_scalar: bool,
5351 ) -> (
5352 Vec<(String, Vec<BinaryFieldPair<'a>>)>,
5353 Vec<BinaryFieldPair<'a>>,
5354 ) {
5355 let left_alternative = Self::alternative_sample_columns(left_schema, left_field_columns);
5358 let right_alternative = Self::alternative_sample_columns(right_schema, right_field_columns);
5359 let alternative_alignment = match (left_alternative, right_alternative) {
5360 (Some(output_names), Some(_)) => Some((
5361 output_names,
5362 left_field_columns
5363 .iter()
5364 .flat_map(|left| right_field_columns.iter().map(move |right| (left, right)))
5365 .collect::<Vec<_>>(),
5366 )),
5367 (Some(output_names), None) if right_field_columns.len() == 1 => Some((
5368 output_names,
5369 left_field_columns
5370 .iter()
5371 .map(|left| (left, &right_field_columns[0]))
5372 .collect::<Vec<_>>(),
5373 )),
5374 (None, Some(output_names)) if left_field_columns.len() == 1 => Some((
5375 output_names,
5376 right_field_columns
5377 .iter()
5378 .map(|right| (&left_field_columns[0], right))
5379 .collect::<Vec<_>>(),
5380 )),
5381 _ => None,
5382 };
5383 let mut invalid_pairs = Vec::new();
5384 if let Some(((float_output, histogram_output), field_pairs)) = alternative_alignment {
5385 let mut float_pairs = Vec::new();
5386 let mut histogram_pairs = Vec::new();
5387 for (left, right) in field_pairs {
5388 let left_is_histogram = Self::field_column_is_native_histogram(left_schema, left);
5389 let right_is_histogram =
5390 Self::field_column_is_native_histogram(right_schema, right);
5391 match Self::binary_result_is_histogram(op, left_is_histogram, right_is_histogram) {
5392 Some(false) => float_pairs.push((left, right)),
5393 Some(true) => histogram_pairs.push((left, right)),
5394 None => invalid_pairs.push((left, right)),
5395 }
5396 }
5397 if !float_pairs.is_empty() || !histogram_pairs.is_empty() {
5398 return (
5399 [
5400 (!float_pairs.is_empty()).then(|| (float_output.to_string(), float_pairs)),
5401 (!histogram_pairs.is_empty())
5402 .then(|| (histogram_output.to_string(), histogram_pairs)),
5403 ]
5404 .into_iter()
5405 .flatten()
5406 .collect(),
5407 invalid_pairs,
5408 );
5409 }
5410 }
5411
5412 if left_is_scalar && !right_is_scalar && left_field_columns.len() == 1 {
5413 return (
5414 right_field_columns
5415 .iter()
5416 .map(|right| (right.clone(), vec![(&left_field_columns[0], right)]))
5417 .collect(),
5418 invalid_pairs,
5419 );
5420 }
5421 if right_is_scalar && !left_is_scalar && right_field_columns.len() == 1 {
5422 return (
5423 left_field_columns
5424 .iter()
5425 .map(|left| (left.clone(), vec![(left, &right_field_columns[0])]))
5426 .collect(),
5427 invalid_pairs,
5428 );
5429 }
5430
5431 (
5432 left_field_columns
5433 .iter()
5434 .zip(right_field_columns.iter())
5435 .map(|(left, right)| (left.clone(), vec![(left, right)]))
5436 .collect(),
5437 invalid_pairs,
5438 )
5439 }
5440
5441 fn binary_result_is_histogram(
5442 token: TokenType,
5443 lhs_is_histogram: bool,
5444 rhs_is_histogram: bool,
5445 ) -> Option<bool> {
5446 match (token.id(), lhs_is_histogram, rhs_is_histogram) {
5447 (_, false, false) => Some(false),
5448 (token::T_ADD | token::T_SUB, true, true)
5449 | (token::T_MUL, true, false)
5450 | (token::T_MUL, false, true)
5451 | (token::T_DIV, true, false) => Some(true),
5452 (token::T_EQLC | token::T_NEQ, true, true) => Some(false),
5453 _ => None,
5454 }
5455 }
5456
5457 fn plan_has_tsid_column(plan: &LogicalPlan) -> bool {
5458 plan.schema()
5459 .fields()
5460 .iter()
5461 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
5462 }
5463
5464 fn is_empty_metric(plan: &LogicalPlan) -> bool {
5465 matches!(plan, LogicalPlan::Extension(Extension { node }) if node.as_any().is::<EmptyMetric>())
5466 }
5467
5468 fn native_histogram_arrow_type() -> ArrowDataType {
5469 native_histogram_value_type().as_arrow_type()
5470 }
5471
5472 fn field_column_type<'a>(
5473 schema: &'a DFSchemaRef,
5474 field_column: &str,
5475 ) -> Option<&'a ArrowDataType> {
5476 schema
5477 .index_of_column_by_name(None, field_column)
5478 .map(|idx| schema.field(idx).data_type())
5479 }
5480
5481 fn field_column_is_native_histogram(schema: &DFSchemaRef, field_column: &str) -> bool {
5482 Self::field_column_type(schema, field_column)
5483 .is_some_and(|data_type| data_type == &Self::native_histogram_arrow_type())
5484 }
5485
5486 fn field_columns_contain_native_histogram(
5487 schema: &DFSchemaRef,
5488 field_columns: &[String],
5489 ) -> bool {
5490 field_columns
5491 .iter()
5492 .any(|field| Self::field_column_is_native_histogram(schema, field))
5493 }
5494
5495 fn field_column_is_float_range(schema: &DFSchemaRef, field_column: &str) -> bool {
5496 Self::field_column_type(schema, field_column).is_some_and(|data_type| {
5497 matches!(
5498 data_type,
5499 ArrowDataType::Dictionary(key_type, value_type)
5500 if key_type.as_ref() == &ArrowDataType::Int64
5501 && value_type.as_ref() == &ArrowDataType::Float64
5502 )
5503 })
5504 }
5505
5506 fn field_columns_are_alternative_samples(
5507 schema: &DFSchemaRef,
5508 field_columns: &[String],
5509 ) -> bool {
5510 Self::alternative_sample_columns(schema, field_columns).is_some()
5511 }
5512
5513 fn alternative_sample_columns<'a>(
5514 schema: &DFSchemaRef,
5515 field_columns: &'a [String],
5516 ) -> Option<(&'a str, &'a str)> {
5517 if field_columns.len() != 2 {
5518 return None;
5519 }
5520
5521 let canonical_float = field_columns.iter().find(|field| {
5522 field.as_str() == greptime_value()
5523 && (Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64)
5524 || Self::field_column_is_float_range(schema, field))
5525 });
5526 let canonical_histogram = field_columns.iter().find(|field| {
5527 field.as_str() == greptime_native_histogram()
5528 && (Self::field_column_is_native_histogram(schema, field)
5529 || Self::field_column_is_native_histogram_range(schema, field))
5530 });
5531 if let (Some(float), Some(histogram)) = (canonical_float, canonical_histogram) {
5532 return Some((float, histogram));
5533 }
5534
5535 let float = field_columns.iter().find(|field| {
5536 field.starts_with(OR_FLOAT_FIELD_PREFIX)
5537 && (Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64)
5538 || Self::field_column_is_float_range(schema, field))
5539 })?;
5540 let histogram = field_columns.iter().find(|field| {
5541 field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)
5542 && (Self::field_column_is_native_histogram(schema, field)
5543 || Self::field_column_is_native_histogram_range(schema, field))
5544 })?;
5545 Some((float, histogram))
5546 }
5547
5548 fn alternative_sample_range_columns<'a>(
5549 schema: &DFSchemaRef,
5550 field_columns: &'a [String],
5551 ) -> Option<(&'a str, &'a str)> {
5552 Self::alternative_sample_columns(schema, field_columns).filter(|(float, histogram)| {
5553 Self::field_column_is_float_range(schema, float)
5554 && Self::field_column_is_native_histogram_range(schema, histogram)
5555 })
5556 }
5557
5558 fn field_column_is_native_histogram_range(schema: &DFSchemaRef, field_column: &str) -> bool {
5559 Self::field_column_type(schema, field_column).is_some_and(|data_type| {
5560 matches!(
5561 data_type,
5562 ArrowDataType::Dictionary(key_type, value_type)
5563 if key_type.as_ref() == &ArrowDataType::Int64
5564 && value_type.as_ref() == &Self::native_histogram_arrow_type()
5565 )
5566 })
5567 }
5568
5569 fn all_field_columns_are_native_histograms(&self, schema: &DFSchemaRef) -> bool {
5570 !self.ctx.field_columns.is_empty()
5571 && self
5572 .ctx
5573 .field_columns
5574 .iter()
5575 .all(|field| Self::field_column_is_native_histogram(schema, field))
5576 }
5577
5578 fn all_field_columns_are_native_histogram_ranges(&self, schema: &DFSchemaRef) -> bool {
5579 !self.ctx.field_columns.is_empty()
5580 && self
5581 .ctx
5582 .field_columns
5583 .iter()
5584 .all(|field| Self::field_column_is_native_histogram_range(schema, field))
5585 }
5586
5587 fn optional_tsid_projection(
5588 schema: &DFSchemaRef,
5589 table_ref: Option<&TableReference>,
5590 keep_tsid: bool,
5591 ) -> Option<DfExpr> {
5592 keep_tsid.then_some(()).and_then(|_| {
5593 schema
5594 .qualified_field_with_name(table_ref, DATA_SCHEMA_TSID_COLUMN_NAME)
5595 .ok()
5596 .map(|field| DfExpr::Column(field.into()))
5597 })
5598 }
5599
5600 fn binary_join_key_columns(
5601 &self,
5602 left_schema: &DFSchemaRef,
5603 right_schema: &DFSchemaRef,
5604 left_context: &PromPlannerContext,
5605 right_context: &PromPlannerContext,
5606 only_join_time_index: bool,
5607 modifier: &Option<BinModifier>,
5608 ) -> Result<(BTreeSet<String>, BTreeSet<String>, bool)> {
5609 let has_tsid = |schema: &DFSchemaRef| {
5610 schema
5611 .fields()
5612 .iter()
5613 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
5614 };
5615 let use_tsid_join = !only_join_time_index
5616 && self.binary_modifier_preserves_tsid_join_key(left_context, right_context, modifier)
5617 && left_context.use_tsid
5618 && right_context.use_tsid
5619 && has_tsid(left_schema)
5620 && has_tsid(right_schema);
5621
5622 let (mut left_tag_columns, mut right_tag_columns) = if use_tsid_join {
5623 (
5624 BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]),
5625 BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]),
5626 )
5627 } else {
5628 if only_join_time_index {
5629 (BTreeSet::new(), BTreeSet::new())
5630 } else {
5631 (
5632 left_context
5633 .tag_columns
5634 .iter()
5635 .cloned()
5636 .collect::<BTreeSet<_>>(),
5637 right_context
5638 .tag_columns
5639 .iter()
5640 .cloned()
5641 .collect::<BTreeSet<_>>(),
5642 )
5643 }
5644 };
5645
5646 if !use_tsid_join
5647 && let Some(modifier) = modifier
5648 && let Some(matching) = &modifier.matching
5649 {
5650 match matching {
5651 LabelModifier::Include(on) => {
5652 let mask = on.labels.iter().cloned().collect::<BTreeSet<_>>();
5653 left_tag_columns = left_tag_columns.intersection(&mask).cloned().collect();
5654 right_tag_columns = right_tag_columns.intersection(&mask).cloned().collect();
5655 }
5656 LabelModifier::Exclude(ignoring) => {
5657 for label in &ignoring.labels {
5658 let _ = left_tag_columns.remove(label);
5659 let _ = right_tag_columns.remove(label);
5660 }
5661 }
5662 }
5663 }
5664
5665 let force_empty_join =
5666 !use_tsid_join && !only_join_time_index && left_tag_columns != right_tag_columns;
5667 if force_empty_join {
5668 let common_tag_columns = left_tag_columns
5669 .intersection(&right_tag_columns)
5670 .cloned()
5671 .collect::<BTreeSet<_>>();
5672 left_tag_columns = common_tag_columns.clone();
5673 right_tag_columns = common_tag_columns;
5674 }
5675
5676 Ok((left_tag_columns, right_tag_columns, force_empty_join))
5677 }
5678
5679 fn binary_result_labels(
5686 left_context: &PromPlannerContext,
5687 right_context: &PromPlannerContext,
5688 modifier: &Option<BinModifier>,
5689 ) -> Option<Vec<(bool, String)>> {
5690 let modifier = modifier.as_ref()?;
5691 let (many_is_left, include) = match &modifier.card {
5692 VectorMatchCardinality::OneToOne => (true, None),
5693 VectorMatchCardinality::ManyToOne(labels) => (true, Some(labels)),
5694 VectorMatchCardinality::OneToMany(labels) => (false, Some(labels)),
5695 VectorMatchCardinality::ManyToMany => return None,
5697 };
5698
5699 let Some(include) = include else {
5700 let matching = modifier.matching.as_ref()?;
5701 let reduced = |keep: bool, labels: &BTreeSet<&String>| {
5702 left_context
5703 .tag_columns
5704 .iter()
5705 .filter(|tag| labels.contains(tag) == keep)
5706 .map(|tag| (true, tag.clone()))
5707 .collect()
5708 };
5709 return Some(match matching {
5710 LabelModifier::Include(on) => reduced(true, &on.labels.iter().collect()),
5711 LabelModifier::Exclude(ignoring) => {
5712 reduced(false, &ignoring.labels.iter().collect())
5713 }
5714 });
5715 };
5716
5717 let (many_context, one_context) = if many_is_left {
5718 (left_context, right_context)
5719 } else {
5720 (right_context, left_context)
5721 };
5722 let include = include.labels.iter().collect::<BTreeSet<_>>();
5723 let mut labels = many_context
5726 .tag_columns
5727 .iter()
5728 .filter(|tag| !include.contains(tag))
5729 .map(|tag| (many_is_left, tag.clone()))
5730 .collect::<Vec<_>>();
5731 labels.extend(
5732 include
5733 .into_iter()
5734 .filter(|label| one_context.tag_columns.contains(label))
5735 .map(|label| (!many_is_left, label.clone())),
5736 );
5737 Some(labels)
5738 }
5739
5740 fn binary_result_label_projection(
5742 schema: &DFSchemaRef,
5743 left_table_ref: &TableReference,
5744 right_table_ref: &TableReference,
5745 left_context: &PromPlannerContext,
5746 right_context: &PromPlannerContext,
5747 labels: Vec<(bool, String)>,
5748 ) -> Result<BinaryResultLabels> {
5749 let mut exprs = Vec::with_capacity(labels.len());
5750 let mut names = Vec::with_capacity(labels.len());
5751 let mut aggregation_field_labels = Vec::new();
5752 let mut sources = HashSet::new();
5753 for (from_left, label) in labels {
5754 let (table_ref, context) = if from_left {
5755 (left_table_ref, left_context)
5756 } else {
5757 (right_table_ref, right_context)
5758 };
5759 let field = schema
5760 .qualified_field_with_name(Some(table_ref), &label)
5761 .context(DataFusionPlanningSnafu)?;
5762 exprs.push(DfExpr::Column(field.into()));
5763 if context.aggregation_field_labels.contains(&label) {
5764 aggregation_field_labels.push(label.clone());
5765 }
5766 sources.insert(from_left);
5767 names.push(label);
5768 }
5769
5770 let source_context = match sources.into_iter().collect::<Vec<_>>().as_slice() {
5774 [true] => Some((left_table_ref, left_context)),
5775 [false] => Some((right_table_ref, right_context)),
5776 _ => None,
5777 };
5778 let tsid = source_context
5779 .filter(|(_, context)| {
5780 context.use_tsid
5781 && context.tag_columns.len() == names.len()
5782 && context.tag_columns.iter().all(|tag| names.contains(tag))
5783 })
5784 .and_then(|(table_ref, _)| {
5785 Self::optional_tsid_projection(schema, Some(table_ref), true)
5786 });
5787
5788 Ok(BinaryResultLabels {
5789 exprs,
5790 names,
5791 aggregation_field_labels,
5792 tsid,
5793 })
5794 }
5795
5796 fn binary_result_labels_may_repeat(
5800 left_context: &PromPlannerContext,
5801 right_context: &PromPlannerContext,
5802 modifier: &Option<BinModifier>,
5803 ) -> bool {
5804 let Some(modifier) = modifier else {
5805 return false;
5806 };
5807 match &modifier.card {
5808 VectorMatchCardinality::OneToOne => match &modifier.matching {
5809 None => false,
5810 Some(LabelModifier::Include(on)) => {
5811 let on = on.labels.iter().collect::<BTreeSet<_>>();
5812 !left_context.tag_columns.iter().all(|tag| on.contains(tag))
5813 }
5814 Some(LabelModifier::Exclude(ignoring)) => ignoring
5815 .labels
5816 .iter()
5817 .any(|label| left_context.tag_columns.contains(label)),
5818 },
5819 VectorMatchCardinality::ManyToOne(include) => include
5820 .labels
5821 .iter()
5822 .any(|label| left_context.tag_columns.contains(label)),
5823 VectorMatchCardinality::OneToMany(include) => include
5824 .labels
5825 .iter()
5826 .any(|label| right_context.tag_columns.contains(label)),
5827 VectorMatchCardinality::ManyToMany => false,
5828 }
5829 }
5830
5831 fn assert_unique_match_group(
5834 plan: LogicalPlan,
5835 group_exprs: Vec<DfExpr>,
5836 group_labels: Vec<String>,
5837 time_index_expr: DfExpr,
5838 violation: MatchGroupViolation,
5839 ) -> Result<LogicalPlan> {
5840 let mut partition_by = group_exprs.clone();
5841 partition_by.push(time_index_expr);
5842 let occupied_column_names = plan
5845 .schema()
5846 .fields()
5847 .iter()
5848 .map(|field| field.name().as_str())
5849 .collect::<HashSet<_>>();
5850 let mut next_suffix = 0;
5851 let count_column = loop {
5852 let name = match next_suffix {
5853 0 => MATCH_GROUP_COUNT_COLUMN.to_string(),
5854 suffix => format!("{MATCH_GROUP_COUNT_COLUMN}_{suffix}"),
5855 };
5856 next_suffix += 1;
5857 if !occupied_column_names.contains(name.as_str()) {
5858 break name;
5859 }
5860 };
5861 let count = DfExpr::WindowFunction(Box::new(WindowFunction {
5862 fun: WindowFunctionDefinition::AggregateUDF(count_udaf()),
5863 params: WindowFunctionParams {
5864 args: vec![lit(1i64)],
5865 partition_by,
5866 order_by: vec![],
5867 window_frame: WindowFrame::new(None),
5868 null_treatment: None,
5869 distinct: false,
5870 filter: None,
5871 },
5872 }))
5873 .alias(count_column.as_str());
5874
5875 let output_exprs = plan
5876 .schema()
5877 .iter()
5878 .map(|(qualifier, field)| DfExpr::Column(Column::new(qualifier.cloned(), field.name())))
5879 .collect::<Vec<_>>();
5880 let assert_expr = DfExpr::ScalarFunction(ScalarFunction {
5881 func: Arc::new(UniqueMatchGroup::scalar_udf(group_labels, violation)),
5882 args: std::iter::once(ident(count_column.as_str()))
5883 .chain(group_exprs)
5884 .collect(),
5885 });
5886
5887 LogicalPlanBuilder::from(plan)
5888 .window(vec![count])
5889 .context(DataFusionPlanningSnafu)?
5890 .filter(assert_expr)
5891 .context(DataFusionPlanningSnafu)?
5892 .project(output_exprs)
5893 .context(DataFusionPlanningSnafu)?
5894 .build()
5895 .context(DataFusionPlanningSnafu)
5896 }
5897
5898 fn binary_modifier_preserves_tsid_join_key(
5899 &self,
5900 left_context: &PromPlannerContext,
5901 right_context: &PromPlannerContext,
5902 modifier: &Option<BinModifier>,
5903 ) -> bool {
5904 let Some(modifier) = modifier else {
5905 return true;
5906 };
5907
5908 if !matches!(modifier.card, VectorMatchCardinality::OneToOne) {
5909 return false;
5910 }
5911
5912 match &modifier.matching {
5913 None => true,
5914 Some(LabelModifier::Exclude(ignoring)) => ignoring.labels.iter().all(|label| {
5915 !left_context.tag_columns.contains(label)
5916 && !right_context.tag_columns.contains(label)
5917 }),
5918 Some(LabelModifier::Include(on)) => {
5919 let on_labels = on.labels.iter().cloned().collect::<BTreeSet<_>>();
5920 let left_labels = left_context
5921 .tag_columns
5922 .iter()
5923 .cloned()
5924 .collect::<BTreeSet<_>>();
5925 let right_labels = right_context
5926 .tag_columns
5927 .iter()
5928 .cloned()
5929 .collect::<BTreeSet<_>>();
5930
5931 on_labels == left_labels && on_labels == right_labels
5932 }
5933 }
5934 }
5935
5936 #[allow(clippy::too_many_arguments)]
5939 fn join_on_non_field_columns(
5940 &self,
5941 left: LogicalPlan,
5942 right: LogicalPlan,
5943 left_table_ref: TableReference,
5944 right_table_ref: TableReference,
5945 left_time_index_column: Option<String>,
5946 right_time_index_column: Option<String>,
5947 only_join_time_index: bool,
5948 modifier: &Option<BinModifier>,
5949 left_context: &PromPlannerContext,
5950 right_context: &PromPlannerContext,
5951 ) -> Result<LogicalPlan> {
5952 let (mut left_tag_columns, mut right_tag_columns, mut force_empty_join) = self
5953 .binary_join_key_columns(
5954 left.schema(),
5955 right.schema(),
5956 left_context,
5957 right_context,
5958 only_join_time_index,
5959 modifier,
5960 )?;
5961 let use_tsid_join = !only_join_time_index
5962 && !force_empty_join
5963 && left_tag_columns == BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()])
5964 && right_tag_columns == BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]);
5965 let (left, right, left_matched_tags, right_matched_tags) = if !only_join_time_index
5966 && !use_tsid_join
5967 && Self::only_temporality_match_label_mismatches(left_context, right_context, modifier)
5968 {
5969 let mut aligned_left_context = left_context.clone();
5970 let mut aligned_right_context = right_context.clone();
5971 let (left, right, _) = Self::align_temporality_match_column(
5972 left,
5973 right,
5974 &mut aligned_left_context,
5975 &mut aligned_right_context,
5976 )?;
5977 (left_tag_columns, right_tag_columns, force_empty_join) = self
5978 .binary_join_key_columns(
5979 left.schema(),
5980 right.schema(),
5981 &aligned_left_context,
5982 &aligned_right_context,
5983 false,
5984 modifier,
5985 )?;
5986 (
5987 left,
5988 right,
5989 aligned_left_context.tag_columns,
5990 aligned_right_context.tag_columns,
5991 )
5992 } else {
5993 (
5994 left,
5995 right,
5996 left_context.tag_columns.clone(),
5997 right_context.tag_columns.clone(),
5998 )
5999 };
6000
6001 let checked_matching = !only_join_time_index
6004 && !force_empty_join
6005 && !use_tsid_join
6006 && !matches!(
6007 modifier.as_ref().map(|modifier| &modifier.card),
6008 Some(VectorMatchCardinality::ManyToMany)
6009 );
6010 let one_side_is_left = matches!(
6012 modifier.as_ref().map(|modifier| &modifier.card),
6013 Some(VectorMatchCardinality::OneToMany(_))
6014 );
6015 let (left, right) = if !checked_matching {
6016 (left, right)
6017 } else if one_side_is_left {
6018 (
6019 Self::assert_unique_one_side(
6020 left,
6021 &left_tag_columns,
6022 &left_matched_tags,
6023 left_time_index_column.as_deref(),
6024 true,
6025 )?,
6026 right,
6027 )
6028 } else {
6029 (
6030 left,
6031 Self::assert_unique_one_side(
6032 right,
6033 &right_tag_columns,
6034 &right_matched_tags,
6035 right_time_index_column.as_deref(),
6036 false,
6037 )?,
6038 )
6039 };
6040
6041 if let (Some(left_time_index_column), Some(right_time_index_column)) = (
6043 left_time_index_column.clone(),
6044 right_time_index_column.clone(),
6045 ) {
6046 left_tag_columns.insert(left_time_index_column);
6047 right_tag_columns.insert(right_time_index_column);
6048 }
6049
6050 let right = LogicalPlanBuilder::from(right)
6051 .alias(right_table_ref.clone())
6052 .context(DataFusionPlanningSnafu)?
6053 .build()
6054 .context(DataFusionPlanningSnafu)?;
6055
6056 let join_plan = LogicalPlanBuilder::from(left)
6058 .alias(left_table_ref.clone())
6059 .context(DataFusionPlanningSnafu)?
6060 .join_detailed(
6061 right,
6062 JoinType::Inner,
6063 (
6064 left_tag_columns
6065 .into_iter()
6066 .map(Column::from_name)
6067 .collect::<Vec<_>>(),
6068 right_tag_columns
6069 .into_iter()
6070 .map(Column::from_name)
6071 .collect::<Vec<_>>(),
6072 ),
6073 force_empty_join.then_some(lit(false)),
6074 NullEquality::NullEqualsNull,
6075 )
6076 .context(DataFusionPlanningSnafu)?
6077 .build()
6078 .context(DataFusionPlanningSnafu)?;
6079
6080 let labels = checked_matching
6083 .then(|| Self::binary_result_labels(left_context, right_context, modifier))
6084 .flatten()
6085 .filter(|_| {
6086 Self::binary_result_labels_may_repeat(left_context, right_context, modifier)
6087 });
6088 let Some(labels) = labels else {
6089 return Ok(join_plan);
6090 };
6091 let schema = join_plan.schema().clone();
6092 let group_exprs = labels
6093 .iter()
6094 .map(|(from_left, label)| {
6095 let table_ref = if *from_left {
6096 &left_table_ref
6097 } else {
6098 &right_table_ref
6099 };
6100 schema
6101 .qualified_field_with_name(Some(table_ref), label)
6102 .map(|field| DfExpr::Column(field.into()))
6103 .context(DataFusionPlanningSnafu)
6104 })
6105 .collect::<Result<Vec<_>>>()?;
6106 let time_index_expr = left_time_index_column
6107 .as_deref()
6108 .map(|column| {
6109 schema
6110 .qualified_field_with_name(Some(&left_table_ref), column)
6111 .map(|field| DfExpr::Column(field.into()))
6112 .context(DataFusionPlanningSnafu)
6113 })
6114 .transpose()?
6115 .with_context(|| UnexpectedPlanExprSnafu {
6116 desc: "vector matching on a plan without a time index",
6117 })?;
6118 let violation = if matches!(
6119 modifier.as_ref().map(|modifier| &modifier.card),
6120 Some(VectorMatchCardinality::OneToOne)
6121 ) {
6122 MatchGroupViolation::ImplicitManyToOne
6123 } else {
6124 MatchGroupViolation::AmbiguousGroupLabels
6125 };
6126 Self::assert_unique_match_group(
6127 join_plan,
6128 group_exprs,
6129 labels.into_iter().map(|(_, label)| label).collect(),
6130 time_index_expr,
6131 violation,
6132 )
6133 }
6134
6135 fn assert_unique_one_side(
6137 plan: LogicalPlan,
6138 join_keys: &BTreeSet<String>,
6139 tag_columns: &[String],
6140 time_index_column: Option<&str>,
6141 one_side_is_left: bool,
6142 ) -> Result<LogicalPlan> {
6143 let Some(time_index_column) = time_index_column else {
6144 return Ok(plan);
6145 };
6146 if tag_columns.iter().all(|tag| join_keys.contains(tag)) {
6147 return Ok(plan);
6148 }
6149
6150 let group_labels = join_keys.iter().cloned().collect::<Vec<_>>();
6151 let group_exprs = group_labels
6152 .iter()
6153 .map(|label| DfExpr::Column(Column::from_name(label)))
6154 .collect();
6155 Self::assert_unique_match_group(
6156 plan,
6157 group_exprs,
6158 group_labels,
6159 DfExpr::Column(Column::from_name(time_index_column)),
6160 MatchGroupViolation::DuplicateOnOneSide { one_side_is_left },
6161 )
6162 }
6163
6164 fn selected_binary_match_labels(
6165 left_context: &PromPlannerContext,
6166 right_context: &PromPlannerContext,
6167 modifier: &Option<BinModifier>,
6168 ) -> BTreeSet<String> {
6169 let mut labels = left_context
6170 .tag_columns
6171 .iter()
6172 .chain(&right_context.tag_columns)
6173 .cloned()
6174 .collect::<BTreeSet<_>>();
6175 if let Some(matching) = modifier
6176 .as_ref()
6177 .and_then(|modifier| modifier.matching.as_ref())
6178 {
6179 match matching {
6180 LabelModifier::Include(on) => {
6181 labels = on
6182 .labels
6183 .iter()
6184 .filter(|label| {
6185 left_context.tag_columns.contains(label)
6186 || right_context.tag_columns.contains(label)
6187 })
6188 .cloned()
6189 .collect();
6190 }
6191 LabelModifier::Exclude(ignoring) => {
6192 for label in &ignoring.labels {
6193 labels.remove(label);
6194 }
6195 }
6196 }
6197 }
6198 labels
6199 }
6200
6201 fn only_temporality_match_label_mismatches(
6202 left_context: &PromPlannerContext,
6203 right_context: &PromPlannerContext,
6204 modifier: &Option<BinModifier>,
6205 ) -> bool {
6206 let mut mismatches =
6207 Self::selected_binary_match_labels(left_context, right_context, modifier)
6208 .into_iter()
6209 .filter(|label| {
6210 left_context.tag_columns.contains(label)
6211 != right_context.tag_columns.contains(label)
6212 });
6213 matches!(
6214 (mismatches.next(), mismatches.next()),
6215 (Some(label), None) if label == OTLP_AGGREGATION_TEMPORALITY_LABEL
6216 )
6217 }
6218
6219 fn align_temporality_match_column(
6220 mut left: LogicalPlan,
6221 mut right: LogicalPlan,
6222 left_context: &mut PromPlannerContext,
6223 right_context: &mut PromPlannerContext,
6224 ) -> Result<(LogicalPlan, LogicalPlan, bool)> {
6225 let marker = OTLP_AGGREGATION_TEMPORALITY_LABEL;
6226 let left_has_marker = left_context.tag_columns.iter().any(|tag| tag == marker);
6227 let (data_type, value_type, add_to_left) = {
6228 let (present, add_to_left) = if left_has_marker {
6229 (&left, false)
6230 } else {
6231 (&right, true)
6232 };
6233 let data_type = present
6234 .schema()
6235 .fields()
6236 .iter()
6237 .find(|field| field.name() == marker)
6238 .map(|field| field.data_type().clone())
6239 .with_context(|| ColumnNotFoundSnafu {
6240 col: marker.to_string(),
6241 })?;
6242 let value_type = Self::string_value_data_type(&data_type)
6243 .cloned()
6244 .with_context(|| UnexpectedPlanExprSnafu {
6245 desc: format!("temporality match label {marker} must be a string"),
6246 })?;
6247 (data_type, value_type, add_to_left)
6248 };
6249 let null = Self::string_scalar_value(&value_type, None).with_context(|| {
6250 UnexpectedPlanExprSnafu {
6251 desc: format!("temporality match label {marker} must be a string"),
6252 }
6253 })?;
6254 let add_marker = |plan: LogicalPlan| {
6255 let visible = plan
6256 .schema()
6257 .iter()
6258 .map(|(qualifier, field)| {
6259 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6260 })
6261 .collect::<Vec<_>>();
6262 LogicalPlanBuilder::from(plan)
6263 .project(
6264 visible
6265 .into_iter()
6266 .chain([DfExpr::Literal(null, None).alias(marker)]),
6267 )
6268 .context(DataFusionPlanningSnafu)?
6269 .build()
6270 .context(DataFusionPlanningSnafu)
6271 };
6272 if data_type != value_type {
6273 let present = if add_to_left { &mut right } else { &mut left };
6274 let visible = present
6275 .schema()
6276 .iter()
6277 .map(|(qualifier, field)| {
6278 let column =
6279 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()));
6280 if field.name() == marker {
6281 DfExpr::Cast(Cast::new(Box::new(column), value_type.clone()))
6282 .alias_qualified(qualifier.cloned(), field.name().clone())
6283 } else {
6284 column
6285 }
6286 })
6287 .collect::<Vec<_>>();
6288 *present = LogicalPlanBuilder::from(present.clone())
6289 .project(visible)
6290 .context(DataFusionPlanningSnafu)?
6291 .build()
6292 .context(DataFusionPlanningSnafu)?;
6293 }
6294 if add_to_left {
6295 left = add_marker(left)?;
6296 left_context.tag_columns.push(marker.to_string());
6297 } else {
6298 right = add_marker(right)?;
6299 right_context.tag_columns.push(marker.to_string());
6300 }
6301 Ok((left, right, add_to_left))
6302 }
6303
6304 fn normalized_match_key_expr(
6305 label: &str,
6306 field: Option<(Option<TableReference>, ArrowDataType)>,
6307 value_type: &ArrowDataType,
6308 internal_name: &str,
6309 ) -> DfExpr {
6310 let empty = Self::string_scalar_value(value_type, Some(String::new()))
6311 .expect("match label value type is a string");
6312 let expr = if let Some((qualifier, data_type)) = field {
6313 let column = DfExpr::Column(Column::new(qualifier, label));
6314 let column = if &data_type == value_type {
6315 column
6316 } else {
6317 DfExpr::Cast(Cast::new(Box::new(column), value_type.clone()))
6318 };
6319 DfExpr::ScalarFunction(ScalarFunction {
6320 func: coalesce(),
6321 args: vec![column, DfExpr::Literal(empty, None)],
6322 })
6323 } else {
6324 DfExpr::Literal(empty, None)
6325 };
6326 expr.alias(internal_name)
6327 }
6328
6329 fn is_zero_row_empty_relation(plan: &LogicalPlan) -> bool {
6330 matches!(plan, LogicalPlan::EmptyRelation(relation) if !relation.produce_one_row)
6333 }
6334
6335 fn string_value_data_type(data_type: &ArrowDataType) -> Option<&ArrowDataType> {
6336 match data_type {
6337 data_type if data_type.is_string() => Some(data_type),
6338 ArrowDataType::Dictionary(_, value_type) if value_type.is_string() => Some(value_type),
6339 _ => None,
6340 }
6341 }
6342
6343 fn string_scalar_value(
6344 data_type: &ArrowDataType,
6345 value: Option<String>,
6346 ) -> Option<ScalarValue> {
6347 match data_type {
6348 ArrowDataType::Utf8 => Some(ScalarValue::Utf8(value)),
6349 ArrowDataType::LargeUtf8 => Some(ScalarValue::LargeUtf8(value)),
6350 ArrowDataType::Utf8View => Some(ScalarValue::Utf8View(value)),
6351 ArrowDataType::Dictionary(key_type, value_type) => Some(ScalarValue::Dictionary(
6352 key_type.clone(),
6353 Box::new(Self::string_scalar_value(value_type, value)?),
6354 )),
6355 _ => None,
6356 }
6357 }
6358
6359 fn common_label_data_type(
6360 left: Option<&ArrowDataType>,
6361 right: Option<&ArrowDataType>,
6362 ) -> Option<ArrowDataType> {
6363 match (left, right) {
6364 (Some(left), Some(right)) if left == right => {
6365 Self::string_value_data_type(left).map(|_| left.clone())
6366 }
6367 (Some(left), Some(right)) => {
6368 let left_value_type = Self::string_value_data_type(left)?;
6369 let right_value_type = Self::string_value_data_type(right)?;
6370 match (left_value_type, right_value_type) {
6373 (left, right) if left == right => Some(left.clone()),
6374 (ArrowDataType::LargeUtf8, _) | (_, ArrowDataType::LargeUtf8) => {
6375 Some(ArrowDataType::LargeUtf8)
6376 }
6377 (ArrowDataType::Utf8View, ArrowDataType::Utf8View) => {
6378 Some(ArrowDataType::Utf8View)
6379 }
6380 _ => Some(ArrowDataType::Utf8),
6381 }
6382 }
6383 (Some(data_type), None) | (None, Some(data_type)) => {
6384 Self::string_value_data_type(data_type).cloned()
6385 }
6386 (None, None) => Some(ArrowDataType::Utf8),
6387 }
6388 }
6389
6390 fn projection_for_each_field_column<F>(
6398 &mut self,
6399 input: LogicalPlan,
6400 name_to_expr: F,
6401 ) -> Result<LogicalPlan>
6402 where
6403 F: FnMut(&String) -> Result<DfExpr>,
6404 {
6405 self.projection_for_each_field_column_with_labels(input, None, name_to_expr)
6406 }
6407
6408 fn projection_for_each_field_column_with_labels<F>(
6411 &mut self,
6412 input: LogicalPlan,
6413 result_labels: Option<&BinaryResultLabels>,
6414 name_to_expr: F,
6415 ) -> Result<LogicalPlan>
6416 where
6417 F: FnMut(&String) -> Result<DfExpr>,
6418 {
6419 let preserve_field_names =
6422 Self::field_columns_are_alternative_samples(input.schema(), &self.ctx.field_columns);
6423 let table_ref = self.ctx.table_name.clone().map(TableReference::bare);
6424 let input_schema = input.schema().clone();
6426 let lookup = |col: &String| {
6427 input_schema
6428 .qualified_field_with_name(table_ref.as_ref(), col)
6429 .or_else(|_| input_schema.qualified_field_with_unqualified_name(col))
6430 .map(|field| DfExpr::Column(field.into()))
6431 .context(DataFusionPlanningSnafu)
6432 };
6433 let tag_columns_iter = match result_labels {
6434 Some(labels) => labels.exprs.iter().cloned().map(Ok).collect::<Vec<_>>(),
6435 None => self.ctx.tag_columns.iter().map(lookup).collect::<Vec<_>>(),
6436 };
6437 let non_field_columns_iter = tag_columns_iter
6438 .into_iter()
6439 .chain(self.ctx.time_index_column.iter().map(lookup));
6440 let tsid_iter = match result_labels {
6441 Some(labels) => labels.tsid_projection(table_ref.clone()),
6442 None => Self::optional_tsid_projection(
6443 input.schema(),
6444 table_ref.as_ref(),
6445 self.ctx.use_tsid,
6446 ),
6447 }
6448 .into_iter()
6449 .map(Ok);
6450
6451 let result_field_columns = self
6453 .ctx
6454 .field_columns
6455 .iter()
6456 .map(name_to_expr)
6457 .collect::<Result<Vec<_>>>()?;
6458
6459 if !preserve_field_names {
6461 self.ctx.field_columns = result_field_columns
6462 .iter()
6463 .map(|expr| expr.schema_name().to_string())
6464 .collect();
6465 }
6466 let field_columns_iter = result_field_columns
6467 .into_iter()
6468 .zip(self.ctx.field_columns.iter())
6469 .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::<String>, name))));
6470
6471 let project_fields = non_field_columns_iter
6473 .chain(tsid_iter)
6474 .chain(field_columns_iter)
6475 .collect::<Result<Vec<_>>>()?;
6476
6477 LogicalPlanBuilder::from(input)
6478 .project(project_fields)
6479 .context(DataFusionPlanningSnafu)?
6480 .build()
6481 .context(DataFusionPlanningSnafu)
6482 }
6483
6484 fn filter_on_field_column<F>(&self, input: LogicalPlan, name_to_expr: F) -> Result<LogicalPlan>
6486 where
6487 F: FnMut(&String) -> Result<DfExpr>,
6488 {
6489 ensure!(
6490 self.ctx.field_columns.len() == 1
6491 || Self::field_columns_are_alternative_samples(
6492 input.schema(),
6493 &self.ctx.field_columns,
6494 ),
6495 UnsupportedExprSnafu {
6496 name: "filter on multi-value input"
6497 }
6498 );
6499
6500 let field_column_filters = self
6501 .ctx
6502 .field_columns
6503 .iter()
6504 .map(name_to_expr)
6505 .collect::<Result<Vec<_>>>()?;
6506 let field_column_filter =
6507 disjunction(field_column_filters).context(UnsupportedExprSnafu {
6508 name: "filter on empty input",
6509 })?;
6510
6511 LogicalPlanBuilder::from(input)
6512 .filter(field_column_filter)
6513 .context(DataFusionPlanningSnafu)?
6514 .build()
6515 .context(DataFusionPlanningSnafu)
6516 }
6517
6518 fn date_part_on_time_index(&self, date_part: &str) -> Result<DfExpr> {
6521 let input_expr = datafusion::logical_expr::ident(
6522 self.ctx
6523 .time_index_column
6524 .as_ref()
6525 .with_context(|| TimeIndexNotFoundSnafu {
6527 table: "<doesn't matter>",
6528 })?,
6529 );
6530 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
6531 func: datafusion_functions::datetime::date_part(),
6532 args: vec![date_part.lit(), input_expr],
6533 });
6534 Ok(fn_expr)
6535 }
6536
6537 fn strip_tsid_column(&self, plan: LogicalPlan) -> Result<LogicalPlan> {
6538 let schema = plan.schema();
6539 if !schema
6540 .fields()
6541 .iter()
6542 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
6543 {
6544 return Ok(plan);
6545 }
6546
6547 let project_exprs = schema
6550 .iter()
6551 .filter(|(_, field)| field.name() != DATA_SCHEMA_TSID_COLUMN_NAME)
6552 .map(|(qualifier, field)| {
6553 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6554 })
6555 .collect::<Vec<_>>();
6556
6557 LogicalPlanBuilder::from(plan)
6558 .project(project_exprs)
6559 .context(DataFusionPlanningSnafu)?
6560 .build()
6561 .context(DataFusionPlanningSnafu)
6562 }
6563
6564 fn apply_alias(&mut self, plan: LogicalPlan, alias_name: String) -> Result<LogicalPlan> {
6566 let fields_expr = self.create_field_column_exprs()?;
6567
6568 ensure!(
6570 fields_expr.len() == 1,
6571 UnsupportedExprSnafu {
6572 name: "alias on multi-value result"
6573 }
6574 );
6575
6576 let project_fields = fields_expr
6577 .into_iter()
6578 .map(|expr| expr.alias(&alias_name))
6579 .chain(self.create_tag_column_exprs()?)
6580 .chain(Some(self.create_time_index_column_expr()?));
6581
6582 LogicalPlanBuilder::from(plan)
6583 .project(project_fields)
6584 .context(DataFusionPlanningSnafu)?
6585 .build()
6586 .context(DataFusionPlanningSnafu)
6587 }
6588}
6589
6590#[derive(Default, Debug)]
6591struct FunctionArgs {
6592 input: Option<PromExpr>,
6593 literals: Vec<DfExpr>,
6594}
6595
6596#[derive(Debug, Clone)]
6599enum ScalarFunc {
6600 DataFusionBuiltin(Arc<ScalarUdfDef>),
6604 DataFusionUdf(Arc<ScalarUdfDef>),
6608 NativeHistogramUdf(Arc<ScalarUdfDef>),
6611 Udf(Arc<ScalarUdfDef>),
6616 ExtrapolateUdf(Arc<ScalarUdfDef>, i64),
6623 GeneratedExpr,
6627}
6628
6629#[cfg(test)]
6630mod test;