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