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 if window >= self.ctx.interval {
3027 let (lower, _) = bounds(start);
3028 let last = i128::from(start) + num_points * i128::from(self.ctx.interval);
3029 let upper = last - i128::from(offset_duration);
3030 return Ok(filter(lower, upper));
3031 }
3032 let mut filters = Vec::new();
3033 for timestamp in (start..=end).step_by(self.ctx.interval as usize) {
3034 let (lower, upper) = bounds(timestamp);
3035 let Some(filter) = filter(lower, upper) else {
3036 return Ok(None);
3039 };
3040 filters.push(filter);
3041 }
3042 while filters.len() > 1 {
3045 let mut next = Vec::with_capacity(filters.len().div_ceil(2));
3046 let mut iter = filters.into_iter();
3047 while let Some(left) = iter.next() {
3048 next.push(match iter.next() {
3049 Some(right) => left.or(right),
3050 None => left,
3051 });
3052 }
3053 filters = next;
3054 }
3055 Ok(filters.pop())
3056 }
3057
3058 async fn create_table_scan_plan(&mut self, table_ref: TableReference) -> Result<LogicalPlan> {
3063 let provider = self
3064 .table_provider
3065 .resolve_table(table_ref.clone())
3066 .await
3067 .context(CatalogSnafu)?;
3068
3069 let logical_table = self.table_from_source(&provider)?;
3070
3071 let mut maybe_phy_table_ref = table_ref.clone();
3073 let mut scan_provider = provider;
3074 let mut table_id_filter: Option<u32> = None;
3075
3076 if logical_table.table_info().meta.engine == METRIC_ENGINE_NAME
3079 && let Some(physical_table_name) = logical_table
3080 .table_info()
3081 .meta
3082 .options
3083 .extra_options
3084 .get(LOGICAL_TABLE_METADATA_KEY)
3085 {
3086 let physical_table_ref = if let Some(schema_name) = &self.ctx.schema_name {
3087 TableReference::partial(schema_name.as_str(), physical_table_name.as_str())
3088 } else {
3089 TableReference::bare(physical_table_name.as_str())
3090 };
3091
3092 let physical_provider = match self
3093 .table_provider
3094 .resolve_table(physical_table_ref.clone())
3095 .await
3096 {
3097 Ok(provider) => provider,
3098 Err(e) if e.status_code() == StatusCode::TableNotFound => {
3099 scan_provider.clone()
3102 }
3103 Err(e) => return Err(e).context(CatalogSnafu),
3104 };
3105
3106 if !Arc::ptr_eq(&physical_provider, &scan_provider) {
3107 let physical_table = self.table_from_source(&physical_provider)?;
3109
3110 let has_table_id = physical_table
3111 .schema()
3112 .column_schema_by_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME)
3113 .is_some();
3114 let has_tsid = physical_table
3115 .schema()
3116 .column_schema_by_name(DATA_SCHEMA_TSID_COLUMN_NAME)
3117 .is_some_and(|col| matches!(col.data_type, ConcreteDataType::UInt64(_)));
3118
3119 if has_table_id && has_tsid {
3120 scan_provider = physical_provider;
3121 maybe_phy_table_ref = physical_table_ref;
3122 table_id_filter = Some(logical_table.table_info().ident.table_id);
3123 }
3124 }
3125 }
3126
3127 let scan_table = self.table_from_source(&scan_provider)?;
3128
3129 let use_tsid = table_id_filter.is_some()
3130 && scan_table
3131 .schema()
3132 .column_schema_by_name(DATA_SCHEMA_TSID_COLUMN_NAME)
3133 .is_some_and(|col| matches!(col.data_type, ConcreteDataType::UInt64(_)));
3134 self.ctx.use_tsid = use_tsid;
3135
3136 let all_table_tags = self.ctx.tag_columns.clone();
3137
3138 let scan_tag_columns = if use_tsid {
3139 let mut scan_tags = self.ctx.tag_columns.clone();
3140 for matcher in &self.ctx.selector_matcher {
3141 if is_metric_engine_internal_column(&matcher.name) {
3142 continue;
3143 }
3144 if all_table_tags.iter().any(|tag| tag == &matcher.name) {
3145 scan_tags.push(matcher.name.clone());
3146 }
3147 }
3148 scan_tags.sort_unstable();
3149 scan_tags.dedup();
3150 scan_tags
3151 } else {
3152 self.ctx.tag_columns.clone()
3153 };
3154
3155 let time_index_data_type = scan_table
3156 .schema()
3157 .timestamp_column()
3158 .with_context(|| TimeIndexNotFoundSnafu {
3159 table: maybe_phy_table_ref.to_quoted_string(),
3160 })?
3161 .data_type
3162 .clone();
3163 let is_time_index_second =
3164 time_index_data_type == ConcreteDataType::timestamp_second_datatype();
3165
3166 let scan_projection = if table_id_filter.is_some() {
3167 let mut required_columns = HashSet::new();
3168 required_columns.insert(DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string());
3169 required_columns.insert(self.ctx.time_index_column.clone().with_context(|| {
3170 TimeIndexNotFoundSnafu {
3171 table: maybe_phy_table_ref.to_quoted_string(),
3172 }
3173 })?);
3174 for col in &scan_tag_columns {
3175 required_columns.insert(col.clone());
3176 }
3177 for col in &self.ctx.field_columns {
3178 required_columns.insert(col.clone());
3179 }
3180 if use_tsid {
3181 required_columns.insert(DATA_SCHEMA_TSID_COLUMN_NAME.to_string());
3182 }
3183
3184 let arrow_schema = scan_table.schema().arrow_schema().clone();
3185 Some(
3186 arrow_schema
3187 .fields()
3188 .iter()
3189 .enumerate()
3190 .filter(|(_, field)| required_columns.contains(field.name().as_str()))
3191 .map(|(idx, _)| idx)
3192 .collect::<Vec<_>>(),
3193 )
3194 } else {
3195 None
3196 };
3197
3198 let mut scan_plan =
3199 LogicalPlanBuilder::scan(maybe_phy_table_ref.clone(), scan_provider, scan_projection)
3200 .context(DataFusionPlanningSnafu)?
3201 .build()
3202 .context(DataFusionPlanningSnafu)?;
3203
3204 if let Some(table_id) = table_id_filter {
3205 scan_plan = LogicalPlanBuilder::from(scan_plan)
3206 .filter(
3207 DfExpr::Column(Column::from_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME))
3208 .eq(lit(table_id)),
3209 )
3210 .context(DataFusionPlanningSnafu)?
3211 .alias(table_ref.clone()) .context(DataFusionPlanningSnafu)?
3213 .build()
3214 .context(DataFusionPlanningSnafu)?;
3215 }
3216
3217 if is_time_index_second {
3218 let expr: Vec<_> = self
3223 .create_field_column_exprs()?
3224 .into_iter()
3225 .chain(
3226 scan_tag_columns
3227 .iter()
3228 .map(|tag| DfExpr::Column(Column::from_name(tag))),
3229 )
3230 .chain(self.ctx.use_tsid.then_some(DfExpr::Column(Column::new(
3231 Some(table_ref.clone()),
3232 DATA_SCHEMA_TSID_COLUMN_NAME.to_string(),
3233 ))))
3234 .chain(Some(DfExpr::Alias(Alias {
3235 expr: Box::new(DfExpr::Cast(Cast::new(
3236 Box::new(self.create_time_index_column_expr()?),
3237 ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
3238 ))),
3239 relation: Some(table_ref.clone()),
3240 name: self
3241 .ctx
3242 .time_index_column
3243 .as_ref()
3244 .with_context(|| TimeIndexNotFoundSnafu {
3245 table: table_ref.to_quoted_string(),
3246 })?
3247 .clone(),
3248 metadata: None,
3249 })))
3250 .collect::<Vec<_>>();
3251 scan_plan = LogicalPlanBuilder::from(scan_plan)
3252 .project(expr)
3253 .context(DataFusionPlanningSnafu)?
3254 .build()
3255 .context(DataFusionPlanningSnafu)?;
3256 } else if table_id_filter.is_some()
3257 || time_index_data_type == ConcreteDataType::timestamp_microsecond_datatype()
3258 || time_index_data_type == ConcreteDataType::timestamp_nanosecond_datatype()
3259 {
3260 let project_exprs = self
3266 .create_field_column_exprs()?
3267 .into_iter()
3268 .chain(
3269 scan_tag_columns
3270 .iter()
3271 .map(|tag| DfExpr::Column(Column::from_name(tag))),
3272 )
3273 .chain(
3274 self.ctx
3275 .use_tsid
3276 .then_some(DfExpr::Column(Column::from_name(
3277 DATA_SCHEMA_TSID_COLUMN_NAME,
3278 ))),
3279 )
3280 .chain(Some(self.create_time_index_column_expr()?))
3281 .collect::<Vec<_>>();
3282
3283 scan_plan = LogicalPlanBuilder::from(scan_plan)
3284 .project(project_exprs)
3285 .context(DataFusionPlanningSnafu)?
3286 .build()
3287 .context(DataFusionPlanningSnafu)?;
3288 }
3289
3290 let result = LogicalPlanBuilder::from(scan_plan)
3291 .build()
3292 .context(DataFusionPlanningSnafu)?;
3293 Ok(result)
3294 }
3295
3296 fn collect_row_key_tag_columns_from_plan(
3297 &self,
3298 plan: &LogicalPlan,
3299 ) -> Result<BTreeSet<String>> {
3300 fn walk(
3301 planner: &PromPlanner,
3302 plan: &LogicalPlan,
3303 out: &mut BTreeSet<String>,
3304 ) -> Result<()> {
3305 if let LogicalPlan::TableScan(scan) = plan
3307 && let Ok(table) = planner.table_from_source(&scan.source)
3308 {
3309 for col in table.table_info().meta.row_key_column_names() {
3310 if col != DATA_SCHEMA_TABLE_ID_COLUMN_NAME
3311 && col != DATA_SCHEMA_TSID_COLUMN_NAME
3312 && !is_metric_engine_internal_column(col)
3313 {
3314 out.insert(col.clone());
3315 }
3316 }
3317 }
3318
3319 for input in plan.inputs() {
3320 walk(planner, input, out)?;
3321 }
3322 Ok(())
3323 }
3324
3325 let mut out = BTreeSet::new();
3326 walk(self, plan, &mut out)?;
3327 Ok(out)
3328 }
3329
3330 async fn setup_context(&mut self) -> Result<Option<LogicalPlan>> {
3334 let table_ref = self.table_ref()?;
3335 let source = match self.table_provider.resolve_table(table_ref.clone()).await {
3336 Err(e) if e.status_code() == StatusCode::TableNotFound => {
3337 let plan = self.setup_context_for_empty_metric()?;
3338 return Ok(Some(plan));
3339 }
3340 res => res.context(CatalogSnafu)?,
3341 };
3342 let table = self.table_from_source(&source)?;
3343
3344 let time_index = table
3346 .schema()
3347 .timestamp_column()
3348 .with_context(|| TimeIndexNotFoundSnafu {
3349 table: table_ref.to_quoted_string(),
3350 })?
3351 .name
3352 .clone();
3353 self.ctx.time_index_column = Some(time_index);
3354
3355 let values = table
3357 .table_info()
3358 .meta
3359 .field_column_names()
3360 .cloned()
3361 .collect();
3362 self.ctx.field_columns = values;
3363
3364 let tags = table
3366 .table_info()
3367 .meta
3368 .row_key_column_names()
3369 .filter(|col| {
3370 col != &DATA_SCHEMA_TABLE_ID_COLUMN_NAME && col != &DATA_SCHEMA_TSID_COLUMN_NAME
3372 })
3373 .cloned()
3374 .collect();
3375 self.ctx.tag_columns = tags;
3376 self.ctx.aggregation_field_labels.clear();
3378
3379 self.ctx.use_tsid = false;
3380
3381 Ok(None)
3382 }
3383
3384 fn setup_context_for_empty_metric(&mut self) -> Result<LogicalPlan> {
3387 self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
3388 self.ctx.reset_table_name_and_schema();
3389 self.ctx.tag_columns = vec![];
3390 self.ctx.aggregation_field_labels.clear();
3391 self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
3392 self.ctx.use_tsid = false;
3393
3394 let plan = LogicalPlan::Extension(Extension {
3396 node: Arc::new(
3397 EmptyMetric::new(
3398 0,
3399 -1,
3400 self.ctx.interval,
3401 SPECIAL_TIME_FUNCTION.to_string(),
3402 DEFAULT_FIELD_COLUMN.to_string(),
3403 Some(lit(0.0f64)),
3404 )
3405 .context(DataFusionPlanningSnafu)?,
3406 ),
3407 });
3408 Ok(plan)
3409 }
3410
3411 fn create_function_args(&self, args: &[Box<PromExpr>]) -> Result<FunctionArgs> {
3413 let mut result = FunctionArgs::default();
3414
3415 for arg in args {
3416 if let Some(expr) = Self::try_build_literal_expr(arg) {
3418 result.literals.push(expr);
3419 } else {
3420 match arg.as_ref() {
3422 PromExpr::Subquery(_)
3423 | PromExpr::VectorSelector(_)
3424 | PromExpr::MatrixSelector(_)
3425 | PromExpr::Extension(_)
3426 | PromExpr::Aggregate(_)
3427 | PromExpr::Paren(_)
3428 | PromExpr::Call(_)
3429 | PromExpr::Binary(_)
3430 | PromExpr::Unary(_) => {
3431 if result.input.replace(*arg.clone()).is_some() {
3432 MultipleVectorSnafu { expr: *arg.clone() }.fail()?;
3433 }
3434 }
3435
3436 _ => {
3437 let expr = Self::get_param_as_literal_expr(Some(arg.as_ref()), None, None)?;
3438 result.literals.push(expr);
3439 }
3440 }
3441 }
3442 }
3443
3444 Ok(result)
3445 }
3446
3447 fn create_mixed_range_function_exprs(
3448 &mut self,
3449 func: &Function,
3450 mut other_input_exprs: VecDeque<DfExpr>,
3451 float_field: &str,
3452 histogram_field: &str,
3453 input_schema: &DFSchemaRef,
3454 range_fold_offset: Option<Millisecond>,
3455 ) -> Result<Option<Vec<DfExpr>>> {
3456 let returns_histogram = matches!(
3457 func.name,
3458 "rate"
3459 | "increase"
3460 | "delta"
3461 | "idelta"
3462 | "irate"
3463 | "avg_over_time"
3464 | "sum_over_time"
3465 | "last_over_time"
3466 );
3467 if !returns_histogram
3468 && !matches!(
3469 func.name,
3470 "changes"
3471 | "resets"
3472 | "deriv"
3473 | "min_over_time"
3474 | "max_over_time"
3475 | "count_over_time"
3476 | "absent_over_time"
3477 | "present_over_time"
3478 | "stddev_over_time"
3479 | "stdvar_over_time"
3480 | "quantile_over_time"
3481 | "predict_linear"
3482 | "double_exponential_smoothing"
3483 | "holt_winters"
3484 )
3485 {
3486 return Ok(None);
3487 }
3488
3489 if func.name == "predict_linear" {
3490 other_input_exprs[0] = DfExpr::Cast(Cast::new(
3491 Box::new(other_input_exprs[0].clone()),
3492 ArrowDataType::Int64,
3493 ));
3494 other_input_exprs.push_back(self.create_range_eval_ts_expr(
3499 range_fold_offset.context(ExpectRangeSelectorSnafu)?,
3500 input_schema,
3501 )?);
3502 }
3503
3504 let timestamp_range = DfExpr::Column(Column::from_name(
3505 RangeManipulate::build_timestamp_range_name(
3506 self.ctx.time_index_column.as_ref().unwrap(),
3507 ),
3508 ));
3509 let float_range = DfExpr::Column(Column::from_name(float_field));
3510 let histogram_range = DfExpr::Column(Column::from_name(histogram_field));
3511 let mut args = Vec::with_capacity(other_input_exprs.len() + 6);
3512 args.push(lit(func.name));
3513 args.push(timestamp_range.clone());
3514 args.push(float_range.clone());
3515 args.push(histogram_range.clone());
3516 args.extend(other_input_exprs);
3517 if matches!(func.name, "rate" | "increase" | "delta") {
3518 args.push(self.create_time_index_column_expr()?);
3519 args.push(lit(self.ctx.range.context(ExpectRangeSelectorSnafu)?));
3520 }
3521
3522 let mut float_expr = DfExpr::ScalarFunction(ScalarFunction {
3523 func: Arc::new(MixedRange::float_udf(self.promql_annotations.clone())),
3524 args: args.clone(),
3525 });
3526 if matches!(func.name, "rate" | "increase") {
3527 let raw_delta_function = if func.name == "rate" {
3528 "raw_delta_rate"
3529 } else {
3530 "raw_delta_increase"
3531 };
3532 let delta_sum = DfExpr::ScalarFunction(ScalarFunction {
3533 func: Arc::new(MixedRange::float_udf(self.promql_annotations.clone())),
3534 args: vec![
3535 lit(raw_delta_function),
3536 timestamp_range,
3537 float_range,
3538 histogram_range,
3539 ],
3540 });
3541 float_expr = self.select_delta_range_math(
3542 func.name,
3543 input_schema,
3544 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3545 delta_sum,
3546 float_expr,
3547 )?;
3548 }
3549 let exprs = if returns_histogram {
3550 self.ctx.field_columns = vec![float_field.to_string(), histogram_field.to_string()];
3551 vec![
3552 float_expr.alias(float_field),
3553 DfExpr::ScalarFunction(ScalarFunction {
3554 func: Arc::new(MixedRange::histogram_udf(self.promql_annotations.clone())),
3555 args,
3556 })
3557 .alias(histogram_field),
3558 ]
3559 } else {
3560 let display_name = if func.name == "predict_linear" {
3561 Self::name_without_last_arg(&float_expr)
3564 } else {
3565 float_expr.schema_name().to_string()
3566 };
3567 self.ctx.field_columns = vec![display_name.clone()];
3568 vec![float_expr.alias(display_name)]
3569 };
3570 Ok(Some(exprs))
3571 }
3572
3573 fn create_function_expr(
3579 &mut self,
3580 func: &Function,
3581 other_input_exprs: Vec<DfExpr>,
3582 input_schema: &DFSchemaRef,
3583 query_engine_state: &QueryEngineState,
3584 range_fold_offset: Option<Millisecond>,
3585 ) -> Result<(Vec<DfExpr>, Vec<String>)> {
3586 let mut other_input_exprs: VecDeque<DfExpr> = other_input_exprs.into();
3588 if let Some((float_field, histogram_field)) =
3589 Self::alternative_sample_range_columns(input_schema, &self.ctx.field_columns)
3590 .map(|(float, histogram)| (float.to_string(), histogram.to_string()))
3591 && let Some(exprs) = self.create_mixed_range_function_exprs(
3592 func,
3593 other_input_exprs.clone(),
3594 &float_field,
3595 &histogram_field,
3596 input_schema,
3597 range_fold_offset,
3598 )?
3599 {
3600 return Ok((exprs, vec![]));
3601 }
3602 let alternative_samples =
3603 Self::field_columns_are_alternative_samples(input_schema, &self.ctx.field_columns);
3604 let all_field_columns_are_native_histogram_ranges =
3605 self.all_field_columns_are_native_histogram_ranges(input_schema);
3606
3607 let field_column_pos = 0;
3609 let mut exprs = Vec::with_capacity(self.ctx.field_columns.len());
3610 let mut new_tags = vec![];
3612 let promql_annotations = self.promql_annotations.clone();
3613 let native_histogram_drop_udf = |name: &str| {
3614 Arc::new(NativeHistogramDrop::float_null_udf(
3615 format!(
3616 "{name}: dropped native histogram samples because this function is not supported for native histograms"
3617 ),
3618 promql_annotations.clone(),
3619 ))
3620 };
3621 let scalar_func = match func.name {
3622 "increase" => {
3623 if all_field_columns_are_native_histogram_ranges {
3624 ScalarFunc::ExtrapolateUdf(
3625 Arc::new(NativeHistogramIncrease::scalar_udf_with_collector(
3626 self.promql_annotations.clone(),
3627 )),
3628 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3629 )
3630 } else {
3631 ScalarFunc::ExtrapolateUdf(
3632 Arc::new(Increase::scalar_udf()),
3633 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3634 )
3635 }
3636 }
3637 "rate" => {
3638 if all_field_columns_are_native_histogram_ranges {
3639 ScalarFunc::ExtrapolateUdf(
3640 Arc::new(NativeHistogramRate::scalar_udf_with_collector(
3641 self.promql_annotations.clone(),
3642 )),
3643 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3644 )
3645 } else {
3646 ScalarFunc::ExtrapolateUdf(
3647 Arc::new(Rate::scalar_udf()),
3648 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3649 )
3650 }
3651 }
3652 "delta" => {
3653 if all_field_columns_are_native_histogram_ranges {
3654 ScalarFunc::ExtrapolateUdf(
3655 Arc::new(NativeHistogramDelta::scalar_udf_with_collector(
3656 self.promql_annotations.clone(),
3657 )),
3658 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3659 )
3660 } else {
3661 ScalarFunc::ExtrapolateUdf(
3662 Arc::new(Delta::scalar_udf()),
3663 self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3664 )
3665 }
3666 }
3667 "idelta" => {
3668 if all_field_columns_are_native_histogram_ranges {
3669 ScalarFunc::Udf(Arc::new(NativeHistogramIDelta::scalar_udf_with_collector(
3670 self.promql_annotations.clone(),
3671 )))
3672 } else {
3673 ScalarFunc::Udf(Arc::new(IDelta::<false>::scalar_udf()))
3674 }
3675 }
3676 "irate" => {
3677 if all_field_columns_are_native_histogram_ranges {
3678 ScalarFunc::Udf(Arc::new(NativeHistogramIRate::scalar_udf_with_collector(
3679 self.promql_annotations.clone(),
3680 )))
3681 } else {
3682 ScalarFunc::Udf(Arc::new(IDelta::<true>::scalar_udf()))
3683 }
3684 }
3685 "resets" => {
3686 if all_field_columns_are_native_histogram_ranges {
3687 ScalarFunc::Udf(Arc::new(NativeHistogramResets::scalar_udf()))
3688 } else {
3689 ScalarFunc::Udf(Arc::new(Resets::scalar_udf()))
3690 }
3691 }
3692 "changes" => {
3693 if all_field_columns_are_native_histogram_ranges {
3694 ScalarFunc::Udf(Arc::new(NativeHistogramChanges::scalar_udf()))
3695 } else {
3696 ScalarFunc::Udf(Arc::new(Changes::scalar_udf()))
3697 }
3698 }
3699 "deriv" => {
3700 if all_field_columns_are_native_histogram_ranges {
3701 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3702 } else {
3703 ScalarFunc::Udf(Arc::new(Deriv::scalar_udf()))
3704 }
3705 }
3706 "avg_over_time" => {
3707 if all_field_columns_are_native_histogram_ranges {
3708 ScalarFunc::Udf(Arc::new(
3709 NativeHistogramAvgOverTime::scalar_udf_with_collector(
3710 self.promql_annotations.clone(),
3711 ),
3712 ))
3713 } else {
3714 ScalarFunc::Udf(Arc::new(AvgOverTime::scalar_udf()))
3715 }
3716 }
3717 "min_over_time" => {
3718 if all_field_columns_are_native_histogram_ranges {
3719 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3720 } else {
3721 ScalarFunc::Udf(Arc::new(MinOverTime::scalar_udf()))
3722 }
3723 }
3724 "max_over_time" => {
3725 if all_field_columns_are_native_histogram_ranges {
3726 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3727 } else {
3728 ScalarFunc::Udf(Arc::new(MaxOverTime::scalar_udf()))
3729 }
3730 }
3731 "sum_over_time" => {
3732 if all_field_columns_are_native_histogram_ranges {
3733 ScalarFunc::Udf(Arc::new(
3734 NativeHistogramSumOverTime::scalar_udf_with_collector(
3735 self.promql_annotations.clone(),
3736 ),
3737 ))
3738 } else {
3739 ScalarFunc::Udf(Arc::new(SumOverTime::scalar_udf()))
3740 }
3741 }
3742 "count_over_time" => {
3743 if all_field_columns_are_native_histogram_ranges {
3744 ScalarFunc::Udf(Arc::new(NativeHistogramCountOverTime::scalar_udf()))
3745 } else {
3746 ScalarFunc::Udf(Arc::new(CountOverTime::scalar_udf()))
3747 }
3748 }
3749 "last_over_time" => {
3750 if all_field_columns_are_native_histogram_ranges {
3751 ScalarFunc::Udf(Arc::new(NativeHistogramLastOverTime::scalar_udf()))
3752 } else {
3753 ScalarFunc::Udf(Arc::new(LastOverTime::scalar_udf()))
3754 }
3755 }
3756 "absent_over_time" => {
3757 if all_field_columns_are_native_histogram_ranges {
3758 ScalarFunc::Udf(Arc::new(NativeHistogramAbsentOverTime::scalar_udf()))
3759 } else {
3760 ScalarFunc::Udf(Arc::new(AbsentOverTime::scalar_udf()))
3761 }
3762 }
3763 "present_over_time" => {
3764 if all_field_columns_are_native_histogram_ranges {
3765 ScalarFunc::Udf(Arc::new(NativeHistogramPresentOverTime::scalar_udf()))
3766 } else {
3767 ScalarFunc::Udf(Arc::new(PresentOverTime::scalar_udf()))
3768 }
3769 }
3770 "stddev_over_time" => {
3771 if all_field_columns_are_native_histogram_ranges {
3772 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3773 } else {
3774 ScalarFunc::Udf(Arc::new(StddevOverTime::scalar_udf()))
3775 }
3776 }
3777 "stdvar_over_time" => {
3778 if all_field_columns_are_native_histogram_ranges {
3779 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3780 } else {
3781 ScalarFunc::Udf(Arc::new(StdvarOverTime::scalar_udf()))
3782 }
3783 }
3784 "quantile_over_time" => {
3785 if all_field_columns_are_native_histogram_ranges {
3786 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3787 } else {
3788 ScalarFunc::Udf(Arc::new(QuantileOverTime::scalar_udf()))
3789 }
3790 }
3791 "predict_linear" => {
3792 if all_field_columns_are_native_histogram_ranges {
3793 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3794 } else {
3795 other_input_exprs[0] = DfExpr::Cast(Cast::new(
3796 Box::new(other_input_exprs[0].clone()),
3797 ArrowDataType::Int64,
3798 ));
3799 other_input_exprs.push_back(self.create_range_eval_ts_expr(
3806 range_fold_offset.context(ExpectRangeSelectorSnafu)?,
3807 input_schema,
3808 )?);
3809 ScalarFunc::Udf(Arc::new(PredictLinear::scalar_udf()))
3810 }
3811 }
3812 "double_exponential_smoothing" | "holt_winters" => {
3813 if all_field_columns_are_native_histogram_ranges {
3814 ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3815 } else {
3816 ScalarFunc::Udf(Arc::new(DoubleExponentialSmoothing::scalar_udf()))
3817 }
3818 }
3819 "histogram_count" => {
3820 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramCount::scalar_udf()))
3821 }
3822 "histogram_sum" => {
3823 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramSum::scalar_udf()))
3824 }
3825 "histogram_avg" => {
3826 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramAvg::scalar_udf()))
3827 }
3828 "histogram_stddev" => {
3829 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramStddev::scalar_udf()))
3830 }
3831 "histogram_stdvar" => {
3832 ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramStdvar::scalar_udf()))
3833 }
3834 "time" => {
3835 exprs.push(build_special_time_expr(
3836 self.ctx.time_index_column.as_ref().unwrap(),
3837 ));
3838 ScalarFunc::GeneratedExpr
3839 }
3840 "minute" => {
3841 let expr = self.date_part_on_time_index("minute")?;
3843 exprs.push(expr);
3844 ScalarFunc::GeneratedExpr
3845 }
3846 "hour" => {
3847 let expr = self.date_part_on_time_index("hour")?;
3849 exprs.push(expr);
3850 ScalarFunc::GeneratedExpr
3851 }
3852 "month" => {
3853 let expr = self.date_part_on_time_index("month")?;
3855 exprs.push(expr);
3856 ScalarFunc::GeneratedExpr
3857 }
3858 "year" => {
3859 let expr = self.date_part_on_time_index("year")?;
3861 exprs.push(expr);
3862 ScalarFunc::GeneratedExpr
3863 }
3864 "day_of_month" => {
3865 let expr = self.date_part_on_time_index("day")?;
3867 exprs.push(expr);
3868 ScalarFunc::GeneratedExpr
3869 }
3870 "day_of_week" => {
3871 let expr = self.date_part_on_time_index("dow")?;
3873 exprs.push(expr);
3874 ScalarFunc::GeneratedExpr
3875 }
3876 "day_of_year" => {
3877 let expr = self.date_part_on_time_index("doy")?;
3879 exprs.push(expr);
3880 ScalarFunc::GeneratedExpr
3881 }
3882 "days_in_month" => {
3883 let day_lit_expr = "day".lit();
3888 let month_lit_expr = "month".lit();
3889 let interval_1month_lit_expr =
3890 DfExpr::Literal(ScalarValue::IntervalYearMonth(Some(1)), None);
3891 let interval_1day_lit_expr = DfExpr::Literal(
3892 ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(1, 0))),
3893 None,
3894 );
3895 let the_1month_minus_1day_expr = DfExpr::BinaryExpr(BinaryExpr {
3896 left: Box::new(interval_1month_lit_expr),
3897 op: Operator::Minus,
3898 right: Box::new(interval_1day_lit_expr),
3899 });
3900 let date_trunc_expr = DfExpr::ScalarFunction(ScalarFunction {
3901 func: datafusion_functions::datetime::date_trunc(),
3902 args: vec![month_lit_expr, self.create_time_index_column_expr()?],
3903 });
3904 let date_trunc_plus_interval_expr = DfExpr::BinaryExpr(BinaryExpr {
3905 left: Box::new(date_trunc_expr),
3906 op: Operator::Plus,
3907 right: Box::new(the_1month_minus_1day_expr),
3908 });
3909 let date_part_expr = DfExpr::ScalarFunction(ScalarFunction {
3910 func: datafusion_functions::datetime::date_part(),
3911 args: vec![day_lit_expr, date_trunc_plus_interval_expr],
3912 });
3913
3914 exprs.push(date_part_expr);
3915 ScalarFunc::GeneratedExpr
3916 }
3917
3918 "label_join" => {
3919 self.ctx.use_tsid = false;
3920 let (concat_expr, dst_label) = Self::build_concat_labels_expr(
3921 &mut other_input_exprs,
3922 &self.ctx,
3923 input_schema,
3924 query_engine_state,
3925 )?;
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(concat_expr);
3940
3941 ScalarFunc::GeneratedExpr
3942 }
3943 "label_replace" => {
3944 self.ctx.use_tsid = false;
3945 if let Some((replace_expr, dst_label)) =
3946 self.build_regexp_replace_label_expr(&mut other_input_exprs, input_schema)?
3947 {
3948 for value in &self.ctx.field_columns {
3950 if *value != dst_label {
3951 let expr = DfExpr::Column(Column::from_name(value));
3952 exprs.push(expr);
3953 }
3954 }
3955
3956 self.ctx.tag_columns.retain(|tag| *tag != dst_label);
3958 new_tags.push(dst_label);
3959 exprs.push(replace_expr);
3961 } else {
3962 for value in &self.ctx.field_columns {
3964 let expr = DfExpr::Column(Column::from_name(value));
3965 exprs.push(expr);
3966 }
3967 }
3968
3969 ScalarFunc::GeneratedExpr
3970 }
3971 "sort" | "sort_desc" => {
3972 for value in &self.ctx.field_columns {
3974 if !Self::field_column_is_native_histogram(input_schema, value) {
3975 exprs.push(DfExpr::Column(Column::from_name(value)));
3976 }
3977 }
3978 if exprs.is_empty() {
3981 exprs.push(DfExpr::Literal(ScalarValue::Float64(None), None));
3982 }
3983
3984 ScalarFunc::GeneratedExpr
3985 }
3986 "sort_by_label" | "sort_by_label_desc" | "timestamp" => {
3987 for value in &self.ctx.field_columns {
3990 let expr = DfExpr::Column(Column::from_name(value));
3991 exprs.push(expr);
3992 }
3993
3994 ScalarFunc::GeneratedExpr
3995 }
3996 "round" if self.all_field_columns_are_native_histograms(input_schema) => {
3997 if other_input_exprs.is_empty() {
3998 other_input_exprs.push_front(0.0f64.lit());
3999 }
4000 ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
4001 }
4002 "round" => {
4003 if other_input_exprs.is_empty() {
4004 other_input_exprs.push_front(1.0f64.lit());
4005 }
4006 ScalarFunc::DataFusionUdf(Arc::new(Round::scalar_udf()))
4007 }
4008 "sqrt" if !self.all_field_columns_are_native_histograms(input_schema) => {
4010 ScalarFunc::DataFusionUdf(Arc::new(IeeeSqrt::scalar_udf()))
4011 }
4012 "rad" | "deg" | "sgn" if self.all_field_columns_are_native_histograms(input_schema) => {
4013 ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
4014 }
4015 "rad" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::radians()),
4016 "deg" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::degrees()),
4017 "sgn" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::signum()),
4018 "pi" => {
4019 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
4021 func: datafusion::functions::math::pi(),
4022 args: vec![],
4023 });
4024 exprs.push(fn_expr);
4025
4026 ScalarFunc::GeneratedExpr
4027 }
4028 _ => {
4029 if let Some(f) = query_engine_state
4030 .session_state()
4031 .scalar_functions()
4032 .get(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::DataFusionBuiltin(f.clone())
4038 }
4039 } else if let Some(factory) = query_engine_state.scalar_function(func.name) {
4040 if self.all_field_columns_are_native_histograms(input_schema) {
4041 ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
4042 } else {
4043 let func_state = query_engine_state.function_state();
4044 let query_ctx = self.table_provider.query_ctx();
4045
4046 ScalarFunc::DataFusionUdf(Arc::new(factory.provide(FunctionContext {
4047 state: func_state,
4048 query_ctx: query_ctx.clone(),
4049 })))
4050 }
4051 } else if let Some(f) = datafusion_functions::math::functions()
4052 .iter()
4053 .find(|f| f.name() == func.name)
4054 {
4055 if self.all_field_columns_are_native_histograms(input_schema) {
4056 ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
4057 } else {
4058 ScalarFunc::DataFusionUdf(f.clone())
4059 }
4060 } else {
4061 return UnsupportedExprSnafu {
4062 name: func.name.to_string(),
4063 }
4064 .fail();
4065 }
4066 }
4067 };
4068
4069 for value in &self.ctx.field_columns {
4070 let col_expr = DfExpr::Column(Column::from_name(value));
4071 let value_is_histogram = Self::field_column_is_native_histogram(input_schema, value);
4072
4073 match scalar_func.clone() {
4074 ScalarFunc::DataFusionBuiltin(func) => {
4075 if alternative_samples && value_is_histogram {
4076 continue;
4077 }
4078 other_input_exprs.insert(field_column_pos, col_expr);
4079 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
4080 func,
4081 args: other_input_exprs.clone().into(),
4082 });
4083 exprs.push(fn_expr);
4084 let _ = other_input_exprs.remove(field_column_pos);
4085 }
4086 ScalarFunc::DataFusionUdf(func) => {
4087 if alternative_samples && value_is_histogram {
4088 continue;
4089 }
4090 let args = itertools::chain!(
4091 other_input_exprs.iter().take(field_column_pos).cloned(),
4092 std::iter::once(col_expr),
4093 other_input_exprs.iter().skip(field_column_pos).cloned()
4094 )
4095 .collect_vec();
4096 exprs.push(DfExpr::ScalarFunction(ScalarFunction { func, args }))
4097 }
4098 ScalarFunc::NativeHistogramUdf(func) => {
4099 if value_is_histogram {
4100 let args = itertools::chain!(
4101 other_input_exprs.iter().take(field_column_pos).cloned(),
4102 std::iter::once(col_expr),
4103 other_input_exprs.iter().skip(field_column_pos).cloned()
4104 )
4105 .collect_vec();
4106 exprs.push(DfExpr::ScalarFunction(ScalarFunction { func, args }));
4107 } else if !alternative_samples {
4108 exprs.push(
4109 DfExpr::Literal(ScalarValue::Float64(None), None).alias(format!(
4110 "{}_{}",
4111 func.name(),
4112 value
4113 )),
4114 );
4115 }
4116 }
4117 ScalarFunc::Udf(func) => {
4118 let ts_range_expr = DfExpr::Column(Column::from_name(
4119 RangeManipulate::build_timestamp_range_name(
4120 self.ctx.time_index_column.as_ref().unwrap(),
4121 ),
4122 ));
4123 other_input_exprs.insert(field_column_pos, ts_range_expr);
4124 other_input_exprs.insert(field_column_pos + 1, col_expr);
4125 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
4126 func,
4127 args: other_input_exprs.clone().into(),
4128 });
4129 exprs.push(fn_expr);
4130 let _ = other_input_exprs.remove(field_column_pos + 1);
4131 let _ = other_input_exprs.remove(field_column_pos);
4132 }
4133 ScalarFunc::ExtrapolateUdf(udf, range_length) => {
4134 let ts_range_expr = DfExpr::Column(Column::from_name(
4135 RangeManipulate::build_timestamp_range_name(
4136 self.ctx.time_index_column.as_ref().unwrap(),
4137 ),
4138 ));
4139 other_input_exprs.insert(field_column_pos, ts_range_expr.clone());
4140 other_input_exprs.insert(field_column_pos + 1, col_expr.clone());
4141 other_input_exprs
4142 .insert(field_column_pos + 2, self.create_time_index_column_expr()?);
4143 other_input_exprs.push_back(lit(range_length));
4144 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
4145 func: udf,
4146 args: other_input_exprs.clone().into(),
4147 });
4148 let fn_expr = if matches!(func.name, "rate" | "increase")
4149 && !all_field_columns_are_native_histogram_ranges
4150 {
4151 let delta_sum = DfExpr::ScalarFunction(ScalarFunction {
4152 func: Arc::new(SumOverTime::scalar_udf()),
4153 args: vec![ts_range_expr, col_expr],
4154 });
4155 self.select_delta_range_math(
4156 func.name,
4157 input_schema,
4158 range_length,
4159 delta_sum,
4160 fn_expr,
4161 )?
4162 } else {
4163 fn_expr
4164 };
4165 exprs.push(fn_expr);
4166 let _ = other_input_exprs.pop_back();
4167 let _ = other_input_exprs.remove(field_column_pos + 2);
4168 let _ = other_input_exprs.remove(field_column_pos + 1);
4169 let _ = other_input_exprs.remove(field_column_pos);
4170 }
4171 ScalarFunc::GeneratedExpr => {}
4172 }
4173 }
4174
4175 if !matches!(func.name, "label_join" | "label_replace") {
4179 let mut new_field_columns = Vec::with_capacity(exprs.len());
4180
4181 exprs = exprs
4182 .into_iter()
4183 .map(|expr| {
4184 let display_name = if func.name == "predict_linear"
4189 && !all_field_columns_are_native_histogram_ranges
4190 {
4191 Self::name_without_last_arg(&expr)
4192 } else {
4193 expr.schema_name().to_string()
4194 };
4195 new_field_columns.push(display_name.clone());
4196 Ok(expr.alias(display_name))
4197 })
4198 .collect::<std::result::Result<Vec<_>, _>>()
4199 .context(DataFusionPlanningSnafu)?;
4200
4201 self.ctx.field_columns = new_field_columns;
4202 }
4203
4204 Ok((exprs, new_tags))
4205 }
4206
4207 fn select_delta_range_math(
4208 &self,
4209 function: &str,
4210 input_schema: &DFSchemaRef,
4211 range_length: Millisecond,
4212 delta_sum: DfExpr,
4213 cumulative: DfExpr,
4214 ) -> Result<DfExpr> {
4215 let marker_is_delta = if self
4216 .ctx
4217 .tag_columns
4218 .iter()
4219 .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL)
4220 {
4221 Self::field_column_type(input_schema, OTLP_AGGREGATION_TEMPORALITY_LABEL)
4222 .filter(|data_type| Self::string_value_data_type(data_type).is_some())
4223 .map(|_| {
4224 DfExpr::Column(Column::from_name(OTLP_AGGREGATION_TEMPORALITY_LABEL))
4225 .eq(lit(GREPTIME_TEMPORALITY_DELTA))
4226 })
4227 } else {
4228 None
4229 };
4230 let Some(marker_is_delta) = marker_is_delta else {
4231 return Ok(cumulative);
4232 };
4233
4234 let delta = if function == "rate" {
4235 DfExpr::BinaryExpr(BinaryExpr {
4236 left: Box::new(delta_sum),
4237 op: Operator::Divide,
4238 right: Box::new(lit(range_length as f64 / 1000.0)),
4239 })
4240 } else {
4241 delta_sum
4242 };
4243 let display_name = cumulative.schema_name().to_string();
4244 when(marker_is_delta, delta)
4245 .otherwise(cumulative)
4246 .context(DataFusionPlanningSnafu)
4247 .map(|expr| expr.alias(display_name))
4248 }
4249
4250 fn validate_label_name(label_name: &str) -> Result<()> {
4254 if label_name.starts_with("__") {
4256 return InvalidDestinationLabelNameSnafu { label_name }.fail();
4257 }
4258 if !LABEL_NAME_REGEX.is_match(label_name) {
4260 return InvalidDestinationLabelNameSnafu { label_name }.fail();
4261 }
4262
4263 Ok(())
4264 }
4265
4266 fn build_regexp_replace_label_expr(
4268 &self,
4269 other_input_exprs: &mut VecDeque<DfExpr>,
4270 input_schema: &DFSchemaRef,
4271 ) -> Result<Option<(DfExpr, String)>> {
4272 let dst_label = match other_input_exprs.pop_front() {
4274 Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4275 other => UnexpectedPlanExprSnafu {
4276 desc: format!("expected dst_label string literal, but found {:?}", other),
4277 }
4278 .fail()?,
4279 };
4280
4281 Self::validate_label_name(&dst_label)?;
4283 let replacement = match other_input_exprs.pop_front() {
4284 Some(DfExpr::Literal(ScalarValue::Utf8(Some(r)), _)) => r,
4285 other => UnexpectedPlanExprSnafu {
4286 desc: format!("expected replacement string literal, but found {:?}", other),
4287 }
4288 .fail()?,
4289 };
4290 let src_label = match other_input_exprs.pop_front() {
4291 Some(DfExpr::Literal(ScalarValue::Utf8(Some(s)), None)) => s,
4292 other => UnexpectedPlanExprSnafu {
4293 desc: format!("expected src_label string literal, but found {:?}", other),
4294 }
4295 .fail()?,
4296 };
4297
4298 let regex = match other_input_exprs.pop_front() {
4299 Some(DfExpr::Literal(ScalarValue::Utf8(Some(r)), None)) => r,
4300 other => UnexpectedPlanExprSnafu {
4301 desc: format!("expected regex string literal, but found {:?}", other),
4302 }
4303 .fail()?,
4304 };
4305
4306 let anchored = format!("^(?s:{regex})$");
4310 let compiled = regex::Regex::new(&anchored)
4311 .map_err(|_| InvalidRegularExpressionSnafu { regex }.build())?;
4312 let dst_exists = self.ctx.tag_columns.contains(&dst_label);
4313
4314 if !self.ctx.tag_columns.contains(&src_label) {
4317 let Some(captures) = compiled.captures("") else {
4318 return Ok(None);
4319 };
4320 let mut value = String::new();
4321 captures.expand(&replacement, &mut value);
4322 if value.is_empty() {
4323 if !dst_exists {
4326 return Ok(None);
4327 }
4328 return Ok(Some((
4329 lit(ScalarValue::Utf8(None)).alias(&dst_label),
4330 dst_label,
4331 )));
4332 }
4333 return Ok(Some((lit(value).alias(&dst_label), dst_label)));
4334 }
4335
4336 let src = Self::label_value_expr(&src_label, input_schema)?;
4337 let matched = DfExpr::ScalarFunction(ScalarFunction {
4338 func: datafusion_functions::regex::regexp_like(),
4339 args: vec![src.clone(), lit(anchored.clone())],
4340 });
4341 let replaced = Self::empty_label_to_null(DfExpr::ScalarFunction(ScalarFunction {
4342 func: datafusion_functions::regex::regexp_replace(),
4343 args: vec![src, lit(anchored), lit(replacement)],
4344 }));
4345 let unchanged = if dst_exists {
4346 DfExpr::Column(Column::from_name(&dst_label))
4347 .cast_to(&ArrowDataType::Utf8, input_schema)
4348 .context(DataFusionPlanningSnafu)?
4349 } else {
4350 lit(ScalarValue::Utf8(None))
4351 };
4352 let replace_expr = when(matched, replaced)
4353 .otherwise(unchanged)
4354 .context(DataFusionPlanningSnafu)?;
4355
4356 Ok(Some((replace_expr.alias(&dst_label), dst_label)))
4357 }
4358
4359 fn build_concat_labels_expr(
4361 other_input_exprs: &mut VecDeque<DfExpr>,
4362 ctx: &PromPlannerContext,
4363 input_schema: &DFSchemaRef,
4364 query_engine_state: &QueryEngineState,
4365 ) -> Result<(DfExpr, String)> {
4366 let dst_label = match other_input_exprs.pop_front() {
4369 Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4370 other => UnexpectedPlanExprSnafu {
4371 desc: format!("expected dst_label string literal, but found {:?}", other),
4372 }
4373 .fail()?,
4374 };
4375 let separator = match other_input_exprs.pop_front() {
4376 Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4377 other => UnexpectedPlanExprSnafu {
4378 desc: format!("expected separator string literal, but found {:?}", other),
4379 }
4380 .fail()?,
4381 };
4382
4383 let available_columns: HashSet<&str> = ctx
4385 .tag_columns
4386 .iter()
4387 .chain(ctx.field_columns.iter())
4388 .chain(ctx.time_index_column.as_ref())
4389 .map(|s| s.as_str())
4390 .collect();
4391
4392 let src_labels = other_input_exprs
4393 .iter()
4394 .map(|expr| {
4395 match expr {
4398 DfExpr::Literal(ScalarValue::Utf8(Some(label)), None) => {
4399 if label.is_empty() {
4400 FunctionInvalidArgumentSnafu {
4401 fn_name: "label_join",
4402 }
4403 .fail()
4404 } else if available_columns.contains(label.as_str()) {
4405 Self::label_value_expr(label, input_schema)
4406 } else {
4407 Ok(lit(""))
4408 }
4409 }
4410 other => UnexpectedPlanExprSnafu {
4411 desc: format!(
4412 "expected source label string literal, but found {:?}",
4413 other
4414 ),
4415 }
4416 .fail(),
4417 }
4418 })
4419 .collect::<Result<Vec<_>>>()?;
4420 if src_labels.is_empty() {
4422 return Ok((lit(ScalarValue::Utf8(None)).alias(&dst_label), dst_label));
4423 }
4424
4425 let session_state = query_engine_state.session_state();
4426 let func = session_state
4427 .scalar_functions()
4428 .get("concat_ws")
4429 .context(UnsupportedExprSnafu { name: "concat_ws" })?;
4430
4431 let mut args = Vec::with_capacity(1 + src_labels.len());
4433 args.push(DfExpr::Literal(ScalarValue::Utf8(Some(separator)), None));
4434 args.extend(src_labels);
4435
4436 Ok((
4437 Self::empty_label_to_null(DfExpr::ScalarFunction(ScalarFunction {
4438 func: func.clone(),
4439 args,
4440 }))
4441 .alias(&dst_label),
4442 dst_label,
4443 ))
4444 }
4445
4446 fn label_value_expr(label: &str, input_schema: &DFSchemaRef) -> Result<DfExpr> {
4449 let value = DfExpr::Column(Column::from_name(label))
4450 .cast_to(&ArrowDataType::Utf8, input_schema)
4451 .context(DataFusionPlanningSnafu)?;
4452 Ok(DfExpr::ScalarFunction(ScalarFunction {
4453 func: coalesce(),
4454 args: vec![value, lit("")],
4455 }))
4456 }
4457
4458 fn empty_label_to_null(value: DfExpr) -> DfExpr {
4462 DfExpr::ScalarFunction(ScalarFunction {
4463 func: datafusion_functions::core::nullif(),
4464 args: vec![value, lit("")],
4465 })
4466 }
4467
4468 fn create_time_index_column_expr(&self) -> Result<DfExpr> {
4469 Ok(DfExpr::Column(Column::from_name(
4470 self.ctx
4471 .time_index_column
4472 .clone()
4473 .with_context(|| TimeIndexNotFoundSnafu { table: "unknown" })?,
4474 )))
4475 }
4476
4477 fn create_range_eval_ts_expr(
4492 &self,
4493 fold_offset: Millisecond,
4494 input_schema: &DFSchemaRef,
4495 ) -> Result<DfExpr> {
4496 let eval_ts = self
4497 .create_time_index_column_expr()?
4498 .cast_to(
4499 &ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
4500 input_schema,
4501 )
4502 .context(DataFusionPlanningSnafu)?
4503 .cast_to(&ArrowDataType::Int64, input_schema)
4504 .context(DataFusionPlanningSnafu)?;
4505 DfExpr::BinaryExpr(BinaryExpr {
4506 left: Box::new(eval_ts),
4507 op: Operator::Plus,
4508 right: Box::new(lit(fold_offset)),
4509 })
4510 .cast_to(
4511 &ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
4512 input_schema,
4513 )
4514 .context(DataFusionPlanningSnafu)
4515 }
4516
4517 fn name_without_last_arg(expr: &DfExpr) -> String {
4524 if let DfExpr::ScalarFunction(ScalarFunction { func, args }) = expr
4525 && let Some((_, visible_args)) = args.split_last()
4526 {
4527 let visible = ScalarFunction {
4528 func: func.clone(),
4529 args: visible_args.to_vec(),
4530 };
4531 return DfExpr::ScalarFunction(visible).schema_name().to_string();
4532 }
4533
4534 expr.schema_name().to_string()
4535 }
4536
4537 fn create_tag_column_exprs(&self) -> Result<Vec<DfExpr>> {
4538 let mut result = Vec::with_capacity(self.ctx.tag_columns.len());
4539 for tag in &self.ctx.tag_columns {
4540 let expr = DfExpr::Column(Column::from_name(tag));
4541 result.push(expr);
4542 }
4543 Ok(result)
4544 }
4545
4546 fn create_field_column_exprs(&self) -> Result<Vec<DfExpr>> {
4547 let mut result = Vec::with_capacity(self.ctx.field_columns.len());
4548 for field in &self.ctx.field_columns {
4549 let expr = DfExpr::Column(Column::from_name(field));
4550 result.push(expr);
4551 }
4552 Ok(result)
4553 }
4554
4555 fn create_tag_and_time_index_column_sort_exprs(&self) -> Result<Vec<SortExpr>> {
4556 let mut result = self
4557 .ctx
4558 .tag_columns
4559 .iter()
4560 .map(|col| DfExpr::Column(Column::from_name(col)).sort(true, true))
4561 .collect::<Vec<_>>();
4562 result.push(self.create_time_index_column_expr()?.sort(true, true));
4563 Ok(result)
4564 }
4565
4566 fn create_field_columns_sort_exprs(&self, asc: bool) -> Vec<SortExpr> {
4567 self.ctx
4568 .field_columns
4569 .iter()
4570 .map(|col| DfExpr::Column(Column::from_name(col)).sort(asc, true))
4571 .collect::<Vec<_>>()
4572 }
4573
4574 fn create_sort_exprs_by_tags(
4575 func: &str,
4576 tags: Vec<DfExpr>,
4577 asc: bool,
4578 ) -> Result<Vec<SortExpr>> {
4579 ensure!(
4580 !tags.is_empty(),
4581 FunctionInvalidArgumentSnafu { fn_name: func }
4582 );
4583
4584 tags.iter()
4585 .map(|col| match col {
4586 DfExpr::Literal(ScalarValue::Utf8(Some(label)), _) => {
4587 Ok(DfExpr::Column(Column::from_name(label)).sort(asc, false))
4588 }
4589 other => UnexpectedPlanExprSnafu {
4590 desc: format!("expected label string literal, but found {:?}", other),
4591 }
4592 .fail(),
4593 })
4594 .collect::<Result<Vec<_>>>()
4595 }
4596
4597 fn create_empty_values_filter_expr(&self, preserve_any_value: bool) -> Result<DfExpr> {
4598 let mut exprs = Vec::with_capacity(self.ctx.field_columns.len());
4599 for value in &self.ctx.field_columns {
4600 let expr = DfExpr::Column(Column::from_name(value)).is_not_null();
4601 exprs.push(expr);
4602 }
4603
4604 let predicate = if preserve_any_value {
4609 disjunction(exprs)
4610 } else {
4611 conjunction(exprs)
4612 };
4613 predicate.with_context(|| ValueNotFoundSnafu {
4614 table: self
4615 .table_ref()
4616 .map(|t| t.to_quoted_string())
4617 .unwrap_or_else(|_| "unknown".to_string()),
4618 })
4619 }
4620
4621 fn create_aggregate_exprs(
4638 &mut self,
4639 op: TokenType,
4640 param: &Option<Box<PromExpr>>,
4641 input_plan: &LogicalPlan,
4642 ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4643 let mixed_sample_columns =
4644 Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns)
4645 .map(|(float, histogram)| (float.to_string(), histogram.to_string()));
4646 let is_group_agg = op.id() == token::T_GROUP;
4647 if is_group_agg && mixed_sample_columns.is_none() {
4648 ensure!(
4649 self.ctx.field_columns.len() == 1,
4650 MultiFieldsNotSupportedSnafu {
4651 operator: "group()"
4652 }
4653 );
4654 }
4655
4656 if let Some((float, histogram)) = mixed_sample_columns {
4657 return self.create_mixed_aggregate_exprs(op, param, &float, &histogram);
4658 }
4659
4660 if self.all_field_columns_are_native_histograms(input_plan.schema()) {
4661 return self.create_native_histogram_aggregate_exprs(op, input_plan);
4662 }
4663
4664 let exprs = self
4666 .ctx
4667 .field_columns
4668 .iter()
4669 .map(|col| {
4670 Self::create_numeric_aggregate_expr(
4671 op,
4672 param,
4673 DfExpr::Column(Column::from_name(col)),
4674 )
4675 })
4676 .collect::<Result<Vec<_>>>()?;
4677
4678 let prev_field_exprs = if op.id() == token::T_COUNT_VALUES {
4686 ensure!(
4687 self.ctx.field_columns.len() == 1,
4688 UnsupportedExprSnafu {
4689 name: "count_values on multi-value input"
4690 }
4691 );
4692
4693 self.ctx
4694 .field_columns
4695 .iter()
4696 .map(|col| {
4697 let value = DfExpr::Column(Column::from_name(col));
4698 let value = if Self::field_column_type(input_plan.schema(), col)
4703 == Some(&ArrowDataType::Float64)
4704 {
4705 value
4706 } else {
4707 DfExpr::Cast(Cast::new(Box::new(value), ArrowDataType::Float64))
4708 };
4709 DfExpr::ScalarFunction(ScalarFunction {
4710 func: Arc::new(PromqlFloatToString::scalar_udf()),
4711 args: vec![value],
4712 })
4713 })
4714 .collect()
4715 } else {
4716 vec![]
4717 };
4718
4719 let exprs = normalize_cols(exprs, input_plan)
4722 .context(DataFusionPlanningSnafu)?
4723 .into_iter()
4724 .map(Self::alias_extremum_as_builtin)
4725 .collect::<Vec<_>>();
4726
4727 self.ctx.field_columns = exprs
4729 .iter()
4730 .map(|expr| expr.schema_name().to_string())
4731 .collect();
4732
4733 Ok((exprs, prev_field_exprs))
4734 }
4735
4736 fn alias_extremum_as_builtin(expr: DfExpr) -> DfExpr {
4739 let DfExpr::AggregateFunction(aggr) = &expr else {
4740 return expr;
4741 };
4742 let builtin = match aggr.func.name() {
4743 Extremum::MIN_NAME => min_udaf(),
4744 Extremum::MAX_NAME => max_udaf(),
4745 _ => return expr,
4746 };
4747 let name = builtin
4748 .call(aggr.params.args.clone())
4749 .schema_name()
4750 .to_string();
4751 expr.alias(name)
4752 }
4753
4754 fn create_numeric_aggregate_expr(
4755 op: TokenType,
4756 param: &Option<Box<PromExpr>>,
4757 input: DfExpr,
4758 ) -> Result<DfExpr> {
4759 let expr = match op.id() {
4760 token::T_SUM => sum_udaf().call(vec![input]),
4761 token::T_QUANTILE => {
4762 let q = Self::get_param_as_literal_expr(
4763 param.as_deref(),
4764 Some(op),
4765 Some(ArrowDataType::Float64),
4766 )?;
4767 quantile_udaf().call(vec![q, input])
4768 }
4769 token::T_AVG => avg_udaf().call(vec![input]),
4770 token::T_COUNT_VALUES | token::T_COUNT => count_udaf().call(vec![input]),
4771 token::T_MIN => Extremum::min_udaf().call(vec![input]),
4772 token::T_MAX => Extremum::max_udaf().call(vec![input]),
4773 token::T_GROUP => max_udaf().call(vec![lit(1_f64)]),
4776 token::T_STDDEV => stddev_pop_udaf().call(vec![input]),
4777 token::T_STDVAR => var_pop_udaf().call(vec![input]),
4778 token::T_TOPK | token::T_BOTTOMK => {
4779 return UnsupportedExprSnafu {
4780 name: format!("{op:?}"),
4781 }
4782 .fail();
4783 }
4784 _ => return UnexpectedTokenSnafu { token: op }.fail(),
4785 };
4786 Ok(expr)
4787 }
4788
4789 fn create_mixed_aggregate_exprs(
4790 &mut self,
4791 op: TokenType,
4792 param: &Option<Box<PromExpr>>,
4793 float_column: &str,
4794 histogram_column: &str,
4795 ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4796 let float_input = DfExpr::Column(Column::from_name(float_column));
4797 let histogram_input = DfExpr::Column(Column::from_name(histogram_column));
4798 let float_count = count_udaf().call(vec![float_input.clone()]);
4799 let histogram_count = count_udaf().call(vec![histogram_input.clone()]);
4800 let mixed_sample_value = || {
4801 DfExpr::ScalarFunction(ScalarFunction {
4802 func: coalesce(),
4803 args: vec![
4804 DfExpr::ScalarFunction(ScalarFunction {
4805 func: Arc::new(PromqlFloatToString::scalar_udf()),
4806 args: vec![float_input.clone()],
4807 }),
4808 DfExpr::ScalarFunction(ScalarFunction {
4809 func: Arc::new(NativeHistogramToString::scalar_udf()),
4810 args: vec![histogram_input.clone()],
4811 }),
4812 ],
4813 })
4814 };
4815
4816 let (exprs, prev_field_exprs, field_columns) = match op.id() {
4817 token::T_SUM | token::T_AVG => (
4818 vec![
4819 Self::create_numeric_aggregate_expr(op, param, float_input)?
4820 .alias(float_column),
4821 self.create_native_histogram_aggregate_expr(op, histogram_column)?,
4822 float_count.alias(Self::mixed_sample_count_name(float_column)),
4823 histogram_count.alias(Self::mixed_sample_count_name(histogram_column)),
4824 ],
4825 vec![],
4826 vec![float_column.to_string(), histogram_column.to_string()],
4827 ),
4828 token::T_COUNT => {
4829 let present = when(
4830 float_input
4831 .clone()
4832 .is_not_null()
4833 .or(histogram_input.clone().is_not_null()),
4834 lit(1_i64),
4835 )
4836 .otherwise(lit(ScalarValue::Int64(None)))
4837 .context(DataFusionPlanningSnafu)?;
4838 (
4839 vec![count_udaf().call(vec![present]).alias(float_column)],
4840 vec![],
4841 vec![float_column.to_string()],
4842 )
4843 }
4844 token::T_GROUP => (
4845 vec![max_udaf().call(vec![lit(1_f64)]).alias(float_column)],
4846 vec![],
4847 vec![float_column.to_string()],
4848 ),
4849 token::T_COUNT_VALUES => {
4850 let value = mixed_sample_value();
4851 (
4852 vec![count_udaf().call(vec![value.clone()]).alias(float_column)],
4853 vec![value],
4854 vec![float_column.to_string()],
4855 )
4856 }
4857 token::T_MIN | token::T_MAX | token::T_STDDEV | token::T_STDVAR | token::T_QUANTILE => {
4858 (
4859 vec![
4860 Self::create_numeric_aggregate_expr(op, param, float_input)?
4861 .alias(float_column),
4862 histogram_count.alias(Self::mixed_sample_count_name(histogram_column)),
4863 ],
4864 vec![],
4865 vec![float_column.to_string()],
4866 )
4867 }
4868 token::T_TOPK | token::T_BOTTOMK => {
4869 return UnsupportedExprSnafu {
4870 name: format!("{op:?}"),
4871 }
4872 .fail();
4873 }
4874 _ => return UnexpectedTokenSnafu { token: op }.fail(),
4875 };
4876
4877 self.ctx.field_columns = field_columns;
4878 Ok((exprs, prev_field_exprs))
4879 }
4880
4881 fn mixed_sample_count_column(column: &str) -> DfExpr {
4882 DfExpr::Column(Column::from_name(Self::mixed_sample_count_name(column)))
4883 }
4884
4885 fn mixed_sample_count_name(column: &str) -> String {
4886 format!("__promql_sample_count({column})")
4887 }
4888
4889 fn mixed_aggregate_filter_expr(
4890 &self,
4891 op: TokenType,
4892 float_column: &str,
4893 histogram_column: &str,
4894 ) -> Result<DfExpr> {
4895 let float_count = Self::mixed_sample_count_column(float_column);
4896 let histogram_count = Self::mixed_sample_count_column(histogram_column);
4897 let mixed = float_count
4898 .clone()
4899 .gt(lit(0_i64))
4900 .and(histogram_count.clone().gt(lit(0_i64)));
4901 let drop_mixed = DfExpr::ScalarFunction(ScalarFunction {
4902 func: Arc::new(NativeHistogramDrop::warning_bool_false_udf(
4903 format!(
4904 "{op}: dropped aggregation result containing both float and native histogram samples"
4905 ),
4906 self.promql_annotations.clone(),
4907 )),
4908 args: vec![float_count, histogram_count],
4909 });
4910
4911 when(mixed, drop_mixed)
4912 .otherwise(lit(true))
4913 .context(DataFusionPlanningSnafu)
4914 }
4915
4916 fn mixed_ignored_histogram_filter_expr(
4917 &self,
4918 op: TokenType,
4919 histogram_column: &str,
4920 ) -> Result<DfExpr> {
4921 let histogram_count = Self::mixed_sample_count_column(histogram_column);
4922 let has_histograms = histogram_count.clone().gt(lit(0_i64));
4923 let record_info = DfExpr::ScalarFunction(ScalarFunction {
4924 func: Arc::new(NativeHistogramDrop::bool_true_udf(
4925 format!(
4926 "{op}: dropped native histogram samples because this aggregation is not supported for native histograms"
4927 ),
4928 self.promql_annotations.clone(),
4929 )),
4930 args: vec![histogram_count],
4931 });
4932
4933 when(has_histograms, record_info)
4934 .otherwise(lit(true))
4935 .context(DataFusionPlanningSnafu)
4936 }
4937
4938 fn create_native_histogram_aggregate_expr(
4939 &self,
4940 op: TokenType,
4941 column: &str,
4942 ) -> Result<DfExpr> {
4943 let input = DfExpr::Column(Column::from_name(column));
4944 let expr = match op.id() {
4945 token::T_SUM => Arc::new(NativeHistogramAggSum::aggregate_udf_with_collector(
4946 self.promql_annotations.clone(),
4947 ))
4948 .call(vec![input])
4949 .alias(column),
4950 token::T_AVG => Arc::new(NativeHistogramAggAvg::aggregate_udf_with_collector(
4951 self.promql_annotations.clone(),
4952 ))
4953 .call(vec![input])
4954 .alias(column),
4955 token::T_COUNT_VALUES | token::T_COUNT => {
4956 count_udaf().call(vec![input]).alias(column)
4957 }
4958 token::T_GROUP => max_udaf().call(vec![lit(1_f64)]).alias(column),
4959 token::T_MIN
4960 | token::T_MAX
4961 | token::T_STDDEV
4962 | token::T_STDVAR
4963 | token::T_QUANTILE
4964 | token::T_TOPK
4965 | token::T_BOTTOMK => sum_udaf()
4966 .call(vec![DfExpr::ScalarFunction(ScalarFunction {
4967 func: Arc::new(NativeHistogramDrop::float_null_udf(
4968 format!(
4969 "{op}: dropped native histogram samples because this aggregation is not supported for native histograms"
4970 ),
4971 self.promql_annotations.clone(),
4972 )),
4973 args: vec![input],
4974 })])
4975 .alias(column),
4976 _ => return UnexpectedTokenSnafu { token: op }.fail(),
4977 };
4978 Ok(expr)
4979 }
4980
4981 fn create_native_histogram_aggregate_exprs(
4982 &mut self,
4983 op: TokenType,
4984 input_plan: &LogicalPlan,
4985 ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4986 let prev_field_exprs = if op.id() == token::T_COUNT_VALUES {
4987 ensure!(
4988 self.ctx.field_columns.len() == 1,
4989 UnsupportedExprSnafu {
4990 name: "count_values on multi-value input"
4991 }
4992 );
4993 self.ctx
4994 .field_columns
4995 .iter()
4996 .map(|col| {
4997 DfExpr::ScalarFunction(ScalarFunction {
4998 func: Arc::new(NativeHistogramToString::scalar_udf()),
4999 args: vec![DfExpr::Column(Column::from_name(col))],
5000 })
5001 })
5002 .collect::<Vec<_>>()
5003 } else {
5004 vec![]
5005 };
5006
5007 let exprs = self
5008 .ctx
5009 .field_columns
5010 .iter()
5011 .map(|col| self.create_native_histogram_aggregate_expr(op, col))
5012 .collect::<Result<Vec<_>>>()?;
5013
5014 let normalized_exprs =
5015 normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
5016 self.ctx.field_columns = normalized_exprs
5017 .into_iter()
5018 .map(|expr| expr.schema_name().to_string())
5019 .collect();
5020
5021 Ok((exprs, prev_field_exprs))
5022 }
5023
5024 fn get_param_value_as_str(op: TokenType, param: &Option<Box<PromExpr>>) -> Result<&str> {
5025 let param = param
5026 .as_deref()
5027 .with_context(|| FunctionInvalidArgumentSnafu {
5028 fn_name: op.to_string(),
5029 })?;
5030 let PromExpr::StringLiteral(StringLiteral { val }) = param else {
5031 return FunctionInvalidArgumentSnafu {
5032 fn_name: op.to_string(),
5033 }
5034 .fail();
5035 };
5036
5037 Ok(val)
5038 }
5039
5040 fn get_param_as_literal_expr(
5041 param: Option<&PromExpr>,
5042 op: Option<TokenType>,
5043 expected_type: Option<ArrowDataType>,
5044 ) -> Result<DfExpr> {
5045 let prom_param = param.with_context(|| {
5046 if let Some(op) = op {
5047 FunctionInvalidArgumentSnafu {
5048 fn_name: op.to_string(),
5049 }
5050 } else {
5051 FunctionInvalidArgumentSnafu {
5052 fn_name: "unknown".to_string(),
5053 }
5054 }
5055 })?;
5056
5057 let expr = Self::try_build_literal_expr(prom_param).with_context(|| {
5058 if let Some(op) = op {
5059 FunctionInvalidArgumentSnafu {
5060 fn_name: op.to_string(),
5061 }
5062 } else {
5063 FunctionInvalidArgumentSnafu {
5064 fn_name: "unknown".to_string(),
5065 }
5066 }
5067 })?;
5068
5069 if let Some(expected_type) = expected_type {
5071 let expr_type = expr
5073 .get_type(&DFSchema::empty())
5074 .context(DataFusionPlanningSnafu)?;
5075 if expected_type != expr_type {
5076 return FunctionInvalidArgumentSnafu {
5077 fn_name: format!("expected {expected_type:?}, but found {expr_type:?}"),
5078 }
5079 .fail();
5080 }
5081 }
5082
5083 Ok(expr)
5084 }
5085
5086 fn create_window_exprs(
5089 &mut self,
5090 op: TokenType,
5091 group_exprs: Vec<DfExpr>,
5092 input_plan: &LogicalPlan,
5093 ) -> Result<Vec<DfExpr>> {
5094 ensure!(
5095 self.ctx.field_columns.len() == 1,
5096 UnsupportedExprSnafu {
5097 name: "topk or bottomk on multi-value input"
5098 }
5099 );
5100
5101 assert!(matches!(op.id(), token::T_TOPK | token::T_BOTTOMK));
5102
5103 let asc = matches!(op.id(), token::T_BOTTOMK);
5104
5105 let tag_sort_exprs = self
5106 .create_tag_column_exprs()?
5107 .into_iter()
5108 .map(|expr| expr.sort(asc, true));
5109
5110 let exprs: Vec<DfExpr> = self
5112 .ctx
5113 .field_columns
5114 .iter()
5115 .map(|col| {
5116 let mut sort_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 2);
5117 sort_exprs.push(
5120 datafusion_functions::math::expr_fn::isnan(DfExpr::Column(Column::from_name(
5121 col,
5122 )))
5123 .sort(true, true),
5124 );
5125 sort_exprs.push(DfExpr::Column(Column::from_name(col)).sort(asc, true));
5127 sort_exprs.extend(tag_sort_exprs.clone());
5130
5131 DfExpr::WindowFunction(Box::new(WindowFunction {
5132 fun: WindowFunctionDefinition::WindowUDF(Arc::new(RowNumber::new().into())),
5133 params: WindowFunctionParams {
5134 args: vec![],
5135 partition_by: group_exprs.clone(),
5136 order_by: sort_exprs,
5137 window_frame: WindowFrame::new(Some(true)),
5138 null_treatment: None,
5139 distinct: false,
5140 filter: None,
5141 },
5142 }))
5143 })
5144 .collect();
5145
5146 let normalized_exprs =
5147 normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
5148 Ok(normalized_exprs)
5149 }
5150
5151 fn try_build_literal_expr(expr: &PromExpr) -> Option<DfExpr> {
5154 match expr {
5155 PromExpr::NumberLiteral(NumberLiteral { val }) => Some(val.lit()),
5156 PromExpr::StringLiteral(StringLiteral { val }) => Some(val.lit()),
5157 PromExpr::VectorSelector(_)
5158 | PromExpr::MatrixSelector(_)
5159 | PromExpr::Extension(_)
5160 | PromExpr::Aggregate(_)
5161 | PromExpr::Subquery(_) => None,
5162 PromExpr::Call(Call { func, .. }) => {
5163 if func.name == SPECIAL_TIME_FUNCTION {
5164 None
5167 } else {
5168 None
5169 }
5170 }
5171 PromExpr::Paren(ParenExpr { expr }) => Self::try_build_literal_expr(expr),
5172 PromExpr::Unary(UnaryExpr { expr, .. }) => Some(DfExpr::Negative(Box::new(
5173 Self::try_build_literal_expr(expr)?,
5174 ))),
5175 PromExpr::Binary(PromBinaryExpr {
5176 lhs,
5177 rhs,
5178 op,
5179 modifier,
5180 }) => {
5181 let lhs = Self::try_build_literal_expr(lhs)?;
5182 let rhs = Self::try_build_literal_expr(rhs)?;
5183 let is_comparison_op = Self::is_token_a_comparison_op(*op);
5184 let expr_builder = Self::prom_token_to_binary_expr_builder(*op).ok()?;
5185 let expr = expr_builder(lhs, rhs).ok()?;
5186
5187 let should_return_bool = if let Some(m) = modifier {
5188 m.return_bool
5189 } else {
5190 false
5191 };
5192 if is_comparison_op && should_return_bool {
5193 Some(DfExpr::Cast(Cast::new(
5194 Box::new(expr),
5195 ArrowDataType::Float64,
5196 )))
5197 } else {
5198 Some(expr)
5199 }
5200 }
5201 }
5202 }
5203
5204 fn try_build_special_time_expr_with_context(&self, expr: &PromExpr) -> Option<DfExpr> {
5205 match expr {
5206 PromExpr::Call(Call { func, .. }) => {
5207 if func.name == SPECIAL_TIME_FUNCTION
5208 && let Some(time_index_col) = self.ctx.time_index_column.as_ref()
5209 {
5210 Some(build_special_time_expr(time_index_col))
5211 } else {
5212 None
5213 }
5214 }
5215 _ => None,
5216 }
5217 }
5218
5219 fn native_histogram_binary_expr(
5220 token: TokenType,
5221 lhs: DfExpr,
5222 lhs_is_histogram: bool,
5223 rhs: DfExpr,
5224 rhs_is_histogram: bool,
5225 filter_context: bool,
5226 promql_annotations: Option<PromqlAnnotationCollector>,
5227 ) -> Result<Option<DfExpr>> {
5228 if !lhs_is_histogram && !rhs_is_histogram {
5229 return Ok(None);
5230 }
5231
5232 let scalar_fn = |func: ScalarUdfDef, args| {
5233 DfExpr::ScalarFunction(ScalarFunction {
5234 func: Arc::new(func),
5235 args,
5236 })
5237 };
5238 let invalid_expr = || {
5239 let message = format!(
5240 "{}: dropped native histogram samples because this binary operation is not supported for native histograms",
5241 token
5242 );
5243 let func = if filter_context {
5244 NativeHistogramDrop::bool_false_udf(message, promql_annotations.clone())
5245 } else {
5246 NativeHistogramDrop::float_null_udf(message, promql_annotations.clone())
5247 };
5248 let args = vec![lhs.clone(), rhs.clone()];
5249 scalar_fn(func, args)
5250 };
5251
5252 let expr = match (token.id(), lhs_is_histogram, rhs_is_histogram) {
5253 (token::T_ADD, true, true) => scalar_fn(
5254 NativeHistogramAdd::scalar_udf_with_collector(promql_annotations.clone()),
5255 vec![lhs, rhs],
5256 ),
5257 (token::T_SUB, true, true) => scalar_fn(
5258 NativeHistogramSub::scalar_udf_with_collector(promql_annotations.clone()),
5259 vec![lhs, rhs],
5260 ),
5261 (token::T_MUL, true, false) => {
5262 scalar_fn(NativeHistogramMulScalar::scalar_udf(), vec![lhs, rhs])
5263 }
5264 (token::T_MUL, false, true) => {
5265 scalar_fn(NativeHistogramScalarMul::scalar_udf(), vec![lhs, rhs])
5266 }
5267 (token::T_DIV, true, false) => {
5268 scalar_fn(NativeHistogramDivScalar::scalar_udf(), vec![lhs, rhs])
5269 }
5270 (token::T_EQLC, true, true) => {
5271 scalar_fn(NativeHistogramEq::scalar_udf(), vec![lhs, rhs])
5272 }
5273 (token::T_NEQ, true, true) => {
5274 scalar_fn(NativeHistogramNotEq::scalar_udf(), vec![lhs, rhs])
5275 }
5276 _ => invalid_expr(),
5277 };
5278
5279 Ok(Some(expr))
5280 }
5281
5282 #[allow(clippy::type_complexity)]
5285 fn prom_token_to_binary_expr_builder(
5286 token: TokenType,
5287 ) -> Result<Box<dyn Fn(DfExpr, DfExpr) -> Result<DfExpr>>> {
5288 let cast_float = |expr| {
5289 if matches!(
5290 &expr,
5291 DfExpr::Cast(Cast { field, .. }) if field.data_type() == &ArrowDataType::Float64
5292 ) || matches!(&expr, DfExpr::Literal(ScalarValue::Float64(_), _))
5293 {
5294 expr
5295 } else {
5296 DfExpr::Cast(Cast::new(Box::new(expr), ArrowDataType::Float64))
5297 }
5298 };
5299 match token.id() {
5300 token::T_ADD => Ok(Box::new(move |lhs, rhs| {
5301 Ok(cast_float(lhs) + cast_float(rhs))
5302 })),
5303 token::T_SUB => Ok(Box::new(move |lhs, rhs| {
5304 Ok(cast_float(lhs) - cast_float(rhs))
5305 })),
5306 token::T_MUL => Ok(Box::new(move |lhs, rhs| {
5307 Ok(cast_float(lhs) * cast_float(rhs))
5308 })),
5309 token::T_DIV => Ok(Box::new(move |lhs, rhs| {
5310 Ok(cast_float(lhs) / cast_float(rhs))
5311 })),
5312 token::T_MOD => Ok(Box::new(move |lhs: DfExpr, rhs| {
5313 Ok(cast_float(lhs) % cast_float(rhs))
5314 })),
5315 token::T_EQLC => Ok(Box::new(|lhs, rhs| {
5316 Ok(Self::ieee_comparison(lhs, Operator::Eq, rhs))
5317 })),
5318 token::T_NEQ => Ok(Box::new(|lhs, rhs| {
5319 Ok(Self::ieee_comparison(lhs, Operator::NotEq, rhs))
5320 })),
5321 token::T_GTR => Ok(Box::new(|lhs, rhs| {
5322 Ok(Self::ieee_comparison(lhs, Operator::Gt, rhs))
5323 })),
5324 token::T_LSS => Ok(Box::new(|lhs, rhs| {
5325 Ok(Self::ieee_comparison(lhs, Operator::Lt, rhs))
5326 })),
5327 token::T_GTE => Ok(Box::new(|lhs, rhs| {
5328 Ok(Self::ieee_comparison(lhs, Operator::GtEq, rhs))
5329 })),
5330 token::T_LTE => Ok(Box::new(|lhs, rhs| {
5331 Ok(Self::ieee_comparison(lhs, Operator::LtEq, rhs))
5332 })),
5333 token::T_POW => Ok(Box::new(move |lhs, rhs| {
5334 Ok(DfExpr::ScalarFunction(ScalarFunction {
5335 func: datafusion_functions::math::power(),
5336 args: vec![cast_float(lhs), cast_float(rhs)],
5337 }))
5338 })),
5339 token::T_ATAN2 => Ok(Box::new(move |lhs, rhs| {
5340 Ok(DfExpr::ScalarFunction(ScalarFunction {
5341 func: datafusion_functions::math::atan2(),
5342 args: vec![cast_float(lhs), cast_float(rhs)],
5343 }))
5344 })),
5345 _ => UnexpectedTokenSnafu { token }.fail(),
5346 }
5347 }
5348
5349 fn ieee_comparison(lhs: DfExpr, op: Operator, rhs: DfExpr) -> DfExpr {
5351 DfExpr::ScalarFunction(ScalarFunction {
5352 func: Arc::new(IeeeComparison::scalar_udf(op)),
5353 args: vec![lhs, rhs],
5354 })
5355 }
5356
5357 fn is_token_a_comparison_op(token: TokenType) -> bool {
5359 matches!(
5360 token.id(),
5361 token::T_EQLC
5362 | token::T_NEQ
5363 | token::T_GTR
5364 | token::T_LSS
5365 | token::T_GTE
5366 | token::T_LTE
5367 )
5368 }
5369
5370 fn is_token_a_set_op(token: TokenType) -> bool {
5372 matches!(
5373 token.id(),
5374 token::T_LAND | token::T_LOR | token::T_LUNLESS )
5378 }
5379
5380 fn align_binary_field_columns<'a>(
5381 left_schema: &DFSchemaRef,
5382 right_schema: &DFSchemaRef,
5383 left_field_columns: &'a [String],
5384 right_field_columns: &'a [String],
5385 op: TokenType,
5386 left_is_scalar: bool,
5387 right_is_scalar: bool,
5388 ) -> (
5389 Vec<(String, Vec<BinaryFieldPair<'a>>)>,
5390 Vec<BinaryFieldPair<'a>>,
5391 ) {
5392 let left_alternative = Self::alternative_sample_columns(left_schema, left_field_columns);
5395 let right_alternative = Self::alternative_sample_columns(right_schema, right_field_columns);
5396 let alternative_alignment = match (left_alternative, right_alternative) {
5397 (Some(output_names), Some(_)) => Some((
5398 output_names,
5399 left_field_columns
5400 .iter()
5401 .flat_map(|left| right_field_columns.iter().map(move |right| (left, right)))
5402 .collect::<Vec<_>>(),
5403 )),
5404 (Some(output_names), None) if right_field_columns.len() == 1 => Some((
5405 output_names,
5406 left_field_columns
5407 .iter()
5408 .map(|left| (left, &right_field_columns[0]))
5409 .collect::<Vec<_>>(),
5410 )),
5411 (None, Some(output_names)) if left_field_columns.len() == 1 => Some((
5412 output_names,
5413 right_field_columns
5414 .iter()
5415 .map(|right| (&left_field_columns[0], right))
5416 .collect::<Vec<_>>(),
5417 )),
5418 _ => None,
5419 };
5420 let mut invalid_pairs = Vec::new();
5421 if let Some(((float_output, histogram_output), field_pairs)) = alternative_alignment {
5422 let mut float_pairs = Vec::new();
5423 let mut histogram_pairs = Vec::new();
5424 for (left, right) in field_pairs {
5425 let left_is_histogram = Self::field_column_is_native_histogram(left_schema, left);
5426 let right_is_histogram =
5427 Self::field_column_is_native_histogram(right_schema, right);
5428 match Self::binary_result_is_histogram(op, left_is_histogram, right_is_histogram) {
5429 Some(false) => float_pairs.push((left, right)),
5430 Some(true) => histogram_pairs.push((left, right)),
5431 None => invalid_pairs.push((left, right)),
5432 }
5433 }
5434 if !float_pairs.is_empty() || !histogram_pairs.is_empty() {
5435 return (
5436 [
5437 (!float_pairs.is_empty()).then(|| (float_output.to_string(), float_pairs)),
5438 (!histogram_pairs.is_empty())
5439 .then(|| (histogram_output.to_string(), histogram_pairs)),
5440 ]
5441 .into_iter()
5442 .flatten()
5443 .collect(),
5444 invalid_pairs,
5445 );
5446 }
5447 }
5448
5449 if left_is_scalar && !right_is_scalar && left_field_columns.len() == 1 {
5450 return (
5451 right_field_columns
5452 .iter()
5453 .map(|right| (right.clone(), vec![(&left_field_columns[0], right)]))
5454 .collect(),
5455 invalid_pairs,
5456 );
5457 }
5458 if right_is_scalar && !left_is_scalar && right_field_columns.len() == 1 {
5459 return (
5460 left_field_columns
5461 .iter()
5462 .map(|left| (left.clone(), vec![(left, &right_field_columns[0])]))
5463 .collect(),
5464 invalid_pairs,
5465 );
5466 }
5467
5468 (
5469 left_field_columns
5470 .iter()
5471 .zip(right_field_columns.iter())
5472 .map(|(left, right)| (left.clone(), vec![(left, right)]))
5473 .collect(),
5474 invalid_pairs,
5475 )
5476 }
5477
5478 fn binary_result_is_histogram(
5479 token: TokenType,
5480 lhs_is_histogram: bool,
5481 rhs_is_histogram: bool,
5482 ) -> Option<bool> {
5483 match (token.id(), lhs_is_histogram, rhs_is_histogram) {
5484 (_, false, false) => Some(false),
5485 (token::T_ADD | token::T_SUB, true, true)
5486 | (token::T_MUL, true, false)
5487 | (token::T_MUL, false, true)
5488 | (token::T_DIV, true, false) => Some(true),
5489 (token::T_EQLC | token::T_NEQ, true, true) => Some(false),
5490 _ => None,
5491 }
5492 }
5493
5494 fn plan_has_tsid_column(plan: &LogicalPlan) -> bool {
5495 plan.schema()
5496 .fields()
5497 .iter()
5498 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
5499 }
5500
5501 fn is_empty_metric(plan: &LogicalPlan) -> bool {
5502 matches!(plan, LogicalPlan::Extension(Extension { node }) if node.as_any().is::<EmptyMetric>())
5503 }
5504
5505 fn native_histogram_arrow_type() -> ArrowDataType {
5506 native_histogram_value_type().as_arrow_type()
5507 }
5508
5509 fn field_column_type<'a>(
5510 schema: &'a DFSchemaRef,
5511 field_column: &str,
5512 ) -> Option<&'a ArrowDataType> {
5513 schema
5514 .index_of_column_by_name(None, field_column)
5515 .map(|idx| schema.field(idx).data_type())
5516 }
5517
5518 fn field_column_is_native_histogram(schema: &DFSchemaRef, field_column: &str) -> bool {
5519 Self::field_column_type(schema, field_column)
5520 .is_some_and(|data_type| data_type == &Self::native_histogram_arrow_type())
5521 }
5522
5523 fn field_columns_contain_native_histogram(
5524 schema: &DFSchemaRef,
5525 field_columns: &[String],
5526 ) -> bool {
5527 field_columns
5528 .iter()
5529 .any(|field| Self::field_column_is_native_histogram(schema, field))
5530 }
5531
5532 fn field_column_is_float_range(schema: &DFSchemaRef, field_column: &str) -> bool {
5533 Self::field_column_type(schema, field_column).is_some_and(|data_type| {
5534 matches!(
5535 data_type,
5536 ArrowDataType::Dictionary(key_type, value_type)
5537 if key_type.as_ref() == &ArrowDataType::Int64
5538 && value_type.as_ref() == &ArrowDataType::Float64
5539 )
5540 })
5541 }
5542
5543 fn field_columns_are_alternative_samples(
5544 schema: &DFSchemaRef,
5545 field_columns: &[String],
5546 ) -> bool {
5547 Self::alternative_sample_columns(schema, field_columns).is_some()
5548 }
5549
5550 fn alternative_sample_columns<'a>(
5551 schema: &DFSchemaRef,
5552 field_columns: &'a [String],
5553 ) -> Option<(&'a str, &'a str)> {
5554 if field_columns.len() != 2 {
5555 return None;
5556 }
5557
5558 let canonical_float = field_columns.iter().find(|field| {
5559 field.as_str() == greptime_value()
5560 && (Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64)
5561 || Self::field_column_is_float_range(schema, field))
5562 });
5563 let canonical_histogram = field_columns.iter().find(|field| {
5564 field.as_str() == greptime_native_histogram()
5565 && (Self::field_column_is_native_histogram(schema, field)
5566 || Self::field_column_is_native_histogram_range(schema, field))
5567 });
5568 if let (Some(float), Some(histogram)) = (canonical_float, canonical_histogram) {
5569 return Some((float, histogram));
5570 }
5571
5572 let float = field_columns.iter().find(|field| {
5573 field.starts_with(OR_FLOAT_FIELD_PREFIX)
5574 && (Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64)
5575 || Self::field_column_is_float_range(schema, field))
5576 })?;
5577 let histogram = field_columns.iter().find(|field| {
5578 field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)
5579 && (Self::field_column_is_native_histogram(schema, field)
5580 || Self::field_column_is_native_histogram_range(schema, field))
5581 })?;
5582 Some((float, histogram))
5583 }
5584
5585 fn alternative_sample_range_columns<'a>(
5586 schema: &DFSchemaRef,
5587 field_columns: &'a [String],
5588 ) -> Option<(&'a str, &'a str)> {
5589 Self::alternative_sample_columns(schema, field_columns).filter(|(float, histogram)| {
5590 Self::field_column_is_float_range(schema, float)
5591 && Self::field_column_is_native_histogram_range(schema, histogram)
5592 })
5593 }
5594
5595 fn field_column_is_native_histogram_range(schema: &DFSchemaRef, field_column: &str) -> bool {
5596 Self::field_column_type(schema, field_column).is_some_and(|data_type| {
5597 matches!(
5598 data_type,
5599 ArrowDataType::Dictionary(key_type, value_type)
5600 if key_type.as_ref() == &ArrowDataType::Int64
5601 && value_type.as_ref() == &Self::native_histogram_arrow_type()
5602 )
5603 })
5604 }
5605
5606 fn all_field_columns_are_native_histograms(&self, schema: &DFSchemaRef) -> bool {
5607 !self.ctx.field_columns.is_empty()
5608 && self
5609 .ctx
5610 .field_columns
5611 .iter()
5612 .all(|field| Self::field_column_is_native_histogram(schema, field))
5613 }
5614
5615 fn all_field_columns_are_native_histogram_ranges(&self, schema: &DFSchemaRef) -> bool {
5616 !self.ctx.field_columns.is_empty()
5617 && self
5618 .ctx
5619 .field_columns
5620 .iter()
5621 .all(|field| Self::field_column_is_native_histogram_range(schema, field))
5622 }
5623
5624 fn optional_tsid_projection(
5625 schema: &DFSchemaRef,
5626 table_ref: Option<&TableReference>,
5627 keep_tsid: bool,
5628 ) -> Option<DfExpr> {
5629 keep_tsid.then_some(()).and_then(|_| {
5630 schema
5631 .qualified_field_with_name(table_ref, DATA_SCHEMA_TSID_COLUMN_NAME)
5632 .ok()
5633 .map(|field| DfExpr::Column(field.into()))
5634 })
5635 }
5636
5637 fn binary_join_key_columns(
5638 &self,
5639 left_schema: &DFSchemaRef,
5640 right_schema: &DFSchemaRef,
5641 left_context: &PromPlannerContext,
5642 right_context: &PromPlannerContext,
5643 only_join_time_index: bool,
5644 modifier: &Option<BinModifier>,
5645 ) -> Result<(BTreeSet<String>, BTreeSet<String>, bool)> {
5646 let has_tsid = |schema: &DFSchemaRef| {
5647 schema
5648 .fields()
5649 .iter()
5650 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
5651 };
5652 let use_tsid_join = !only_join_time_index
5653 && self.binary_modifier_preserves_tsid_join_key(left_context, right_context, modifier)
5654 && left_context.use_tsid
5655 && right_context.use_tsid
5656 && has_tsid(left_schema)
5657 && has_tsid(right_schema);
5658
5659 let (mut left_tag_columns, mut right_tag_columns) = if use_tsid_join {
5660 (
5661 BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]),
5662 BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]),
5663 )
5664 } else {
5665 if only_join_time_index {
5666 (BTreeSet::new(), BTreeSet::new())
5667 } else {
5668 (
5669 left_context
5670 .tag_columns
5671 .iter()
5672 .cloned()
5673 .collect::<BTreeSet<_>>(),
5674 right_context
5675 .tag_columns
5676 .iter()
5677 .cloned()
5678 .collect::<BTreeSet<_>>(),
5679 )
5680 }
5681 };
5682
5683 if !use_tsid_join
5684 && let Some(modifier) = modifier
5685 && let Some(matching) = &modifier.matching
5686 {
5687 match matching {
5688 LabelModifier::Include(on) => {
5689 let mask = on.labels.iter().cloned().collect::<BTreeSet<_>>();
5690 left_tag_columns = left_tag_columns.intersection(&mask).cloned().collect();
5691 right_tag_columns = right_tag_columns.intersection(&mask).cloned().collect();
5692 }
5693 LabelModifier::Exclude(ignoring) => {
5694 for label in &ignoring.labels {
5695 let _ = left_tag_columns.remove(label);
5696 let _ = right_tag_columns.remove(label);
5697 }
5698 }
5699 }
5700 }
5701
5702 let force_empty_join =
5703 !use_tsid_join && !only_join_time_index && left_tag_columns != right_tag_columns;
5704 if force_empty_join {
5705 let common_tag_columns = left_tag_columns
5706 .intersection(&right_tag_columns)
5707 .cloned()
5708 .collect::<BTreeSet<_>>();
5709 left_tag_columns = common_tag_columns.clone();
5710 right_tag_columns = common_tag_columns;
5711 }
5712
5713 Ok((left_tag_columns, right_tag_columns, force_empty_join))
5714 }
5715
5716 fn binary_result_labels(
5723 left_context: &PromPlannerContext,
5724 right_context: &PromPlannerContext,
5725 modifier: &Option<BinModifier>,
5726 ) -> Option<Vec<(bool, String)>> {
5727 let modifier = modifier.as_ref()?;
5728 let (many_is_left, include) = match &modifier.card {
5729 VectorMatchCardinality::OneToOne => (true, None),
5730 VectorMatchCardinality::ManyToOne(labels) => (true, Some(labels)),
5731 VectorMatchCardinality::OneToMany(labels) => (false, Some(labels)),
5732 VectorMatchCardinality::ManyToMany => return None,
5734 };
5735
5736 let Some(include) = include else {
5737 let matching = modifier.matching.as_ref()?;
5738 let reduced = |keep: bool, labels: &BTreeSet<&String>| {
5739 left_context
5740 .tag_columns
5741 .iter()
5742 .filter(|tag| labels.contains(tag) == keep)
5743 .map(|tag| (true, tag.clone()))
5744 .collect()
5745 };
5746 return Some(match matching {
5747 LabelModifier::Include(on) => reduced(true, &on.labels.iter().collect()),
5748 LabelModifier::Exclude(ignoring) => {
5749 reduced(false, &ignoring.labels.iter().collect())
5750 }
5751 });
5752 };
5753
5754 let (many_context, one_context) = if many_is_left {
5755 (left_context, right_context)
5756 } else {
5757 (right_context, left_context)
5758 };
5759 let include = include.labels.iter().collect::<BTreeSet<_>>();
5760 let mut labels = many_context
5763 .tag_columns
5764 .iter()
5765 .filter(|tag| !include.contains(tag))
5766 .map(|tag| (many_is_left, tag.clone()))
5767 .collect::<Vec<_>>();
5768 labels.extend(
5769 include
5770 .into_iter()
5771 .filter(|label| one_context.tag_columns.contains(label))
5772 .map(|label| (!many_is_left, label.clone())),
5773 );
5774 Some(labels)
5775 }
5776
5777 fn binary_result_label_projection(
5779 schema: &DFSchemaRef,
5780 left_table_ref: &TableReference,
5781 right_table_ref: &TableReference,
5782 left_context: &PromPlannerContext,
5783 right_context: &PromPlannerContext,
5784 labels: Vec<(bool, String)>,
5785 ) -> Result<BinaryResultLabels> {
5786 let mut exprs = Vec::with_capacity(labels.len());
5787 let mut names = Vec::with_capacity(labels.len());
5788 let mut aggregation_field_labels = Vec::new();
5789 let mut sources = HashSet::new();
5790 for (from_left, label) in labels {
5791 let (table_ref, context) = if from_left {
5792 (left_table_ref, left_context)
5793 } else {
5794 (right_table_ref, right_context)
5795 };
5796 let field = schema
5797 .qualified_field_with_name(Some(table_ref), &label)
5798 .context(DataFusionPlanningSnafu)?;
5799 exprs.push(DfExpr::Column(field.into()));
5800 if context.aggregation_field_labels.contains(&label) {
5801 aggregation_field_labels.push(label.clone());
5802 }
5803 sources.insert(from_left);
5804 names.push(label);
5805 }
5806
5807 let source_context = match sources.into_iter().collect::<Vec<_>>().as_slice() {
5811 [true] => Some((left_table_ref, left_context)),
5812 [false] => Some((right_table_ref, right_context)),
5813 _ => None,
5814 };
5815 let tsid = source_context
5816 .filter(|(_, context)| {
5817 context.use_tsid
5818 && context.tag_columns.len() == names.len()
5819 && context.tag_columns.iter().all(|tag| names.contains(tag))
5820 })
5821 .and_then(|(table_ref, _)| {
5822 Self::optional_tsid_projection(schema, Some(table_ref), true)
5823 });
5824
5825 Ok(BinaryResultLabels {
5826 exprs,
5827 names,
5828 aggregation_field_labels,
5829 tsid,
5830 })
5831 }
5832
5833 fn binary_result_labels_may_repeat(
5837 left_context: &PromPlannerContext,
5838 right_context: &PromPlannerContext,
5839 modifier: &Option<BinModifier>,
5840 ) -> bool {
5841 let Some(modifier) = modifier else {
5842 return false;
5843 };
5844 match &modifier.card {
5845 VectorMatchCardinality::OneToOne => match &modifier.matching {
5846 None => false,
5847 Some(LabelModifier::Include(on)) => {
5848 let on = on.labels.iter().collect::<BTreeSet<_>>();
5849 !left_context.tag_columns.iter().all(|tag| on.contains(tag))
5850 }
5851 Some(LabelModifier::Exclude(ignoring)) => ignoring
5852 .labels
5853 .iter()
5854 .any(|label| left_context.tag_columns.contains(label)),
5855 },
5856 VectorMatchCardinality::ManyToOne(include) => include
5857 .labels
5858 .iter()
5859 .any(|label| left_context.tag_columns.contains(label)),
5860 VectorMatchCardinality::OneToMany(include) => include
5861 .labels
5862 .iter()
5863 .any(|label| right_context.tag_columns.contains(label)),
5864 VectorMatchCardinality::ManyToMany => false,
5865 }
5866 }
5867
5868 fn assert_unique_match_group(
5871 plan: LogicalPlan,
5872 group_exprs: Vec<DfExpr>,
5873 group_labels: Vec<String>,
5874 time_index_expr: DfExpr,
5875 violation: MatchGroupViolation,
5876 ) -> Result<LogicalPlan> {
5877 let mut partition_by = group_exprs.clone();
5878 partition_by.push(time_index_expr);
5879 let occupied_column_names = plan
5882 .schema()
5883 .fields()
5884 .iter()
5885 .map(|field| field.name().as_str())
5886 .collect::<HashSet<_>>();
5887 let mut next_suffix = 0;
5888 let count_column = loop {
5889 let name = match next_suffix {
5890 0 => MATCH_GROUP_COUNT_COLUMN.to_string(),
5891 suffix => format!("{MATCH_GROUP_COUNT_COLUMN}_{suffix}"),
5892 };
5893 next_suffix += 1;
5894 if !occupied_column_names.contains(name.as_str()) {
5895 break name;
5896 }
5897 };
5898 let count = DfExpr::WindowFunction(Box::new(WindowFunction {
5899 fun: WindowFunctionDefinition::AggregateUDF(count_udaf()),
5900 params: WindowFunctionParams {
5901 args: vec![lit(1i64)],
5902 partition_by,
5903 order_by: vec![],
5904 window_frame: WindowFrame::new(None),
5905 null_treatment: None,
5906 distinct: false,
5907 filter: None,
5908 },
5909 }))
5910 .alias(count_column.as_str());
5911
5912 let output_exprs = plan
5913 .schema()
5914 .iter()
5915 .map(|(qualifier, field)| DfExpr::Column(Column::new(qualifier.cloned(), field.name())))
5916 .collect::<Vec<_>>();
5917 let assert_expr = DfExpr::ScalarFunction(ScalarFunction {
5918 func: Arc::new(UniqueMatchGroup::scalar_udf(group_labels, violation)),
5919 args: std::iter::once(ident(count_column.as_str()))
5920 .chain(group_exprs)
5921 .collect(),
5922 });
5923
5924 LogicalPlanBuilder::from(plan)
5925 .window(vec![count])
5926 .context(DataFusionPlanningSnafu)?
5927 .filter(assert_expr)
5928 .context(DataFusionPlanningSnafu)?
5929 .project(output_exprs)
5930 .context(DataFusionPlanningSnafu)?
5931 .build()
5932 .context(DataFusionPlanningSnafu)
5933 }
5934
5935 fn binary_modifier_preserves_tsid_join_key(
5936 &self,
5937 left_context: &PromPlannerContext,
5938 right_context: &PromPlannerContext,
5939 modifier: &Option<BinModifier>,
5940 ) -> bool {
5941 let Some(modifier) = modifier else {
5942 return true;
5943 };
5944
5945 if !matches!(modifier.card, VectorMatchCardinality::OneToOne) {
5946 return false;
5947 }
5948
5949 match &modifier.matching {
5950 None => true,
5951 Some(LabelModifier::Exclude(ignoring)) => ignoring.labels.iter().all(|label| {
5952 !left_context.tag_columns.contains(label)
5953 && !right_context.tag_columns.contains(label)
5954 }),
5955 Some(LabelModifier::Include(on)) => {
5956 let on_labels = on.labels.iter().cloned().collect::<BTreeSet<_>>();
5957 let left_labels = left_context
5958 .tag_columns
5959 .iter()
5960 .cloned()
5961 .collect::<BTreeSet<_>>();
5962 let right_labels = right_context
5963 .tag_columns
5964 .iter()
5965 .cloned()
5966 .collect::<BTreeSet<_>>();
5967
5968 on_labels == left_labels && on_labels == right_labels
5969 }
5970 }
5971 }
5972
5973 #[allow(clippy::too_many_arguments)]
5976 fn join_on_non_field_columns(
5977 &self,
5978 left: LogicalPlan,
5979 right: LogicalPlan,
5980 left_table_ref: TableReference,
5981 right_table_ref: TableReference,
5982 left_time_index_column: Option<String>,
5983 right_time_index_column: Option<String>,
5984 only_join_time_index: bool,
5985 modifier: &Option<BinModifier>,
5986 left_context: &PromPlannerContext,
5987 right_context: &PromPlannerContext,
5988 ) -> Result<LogicalPlan> {
5989 let (mut left_tag_columns, mut right_tag_columns, mut force_empty_join) = self
5990 .binary_join_key_columns(
5991 left.schema(),
5992 right.schema(),
5993 left_context,
5994 right_context,
5995 only_join_time_index,
5996 modifier,
5997 )?;
5998 let use_tsid_join = !only_join_time_index
5999 && !force_empty_join
6000 && left_tag_columns == BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()])
6001 && right_tag_columns == BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]);
6002 let (left, right, left_matched_tags, right_matched_tags) = if !only_join_time_index
6003 && !use_tsid_join
6004 && Self::only_temporality_match_label_mismatches(left_context, right_context, modifier)
6005 {
6006 let mut aligned_left_context = left_context.clone();
6007 let mut aligned_right_context = right_context.clone();
6008 let (left, right, _) = Self::align_temporality_match_column(
6009 left,
6010 right,
6011 &mut aligned_left_context,
6012 &mut aligned_right_context,
6013 )?;
6014 (left_tag_columns, right_tag_columns, force_empty_join) = self
6015 .binary_join_key_columns(
6016 left.schema(),
6017 right.schema(),
6018 &aligned_left_context,
6019 &aligned_right_context,
6020 false,
6021 modifier,
6022 )?;
6023 (
6024 left,
6025 right,
6026 aligned_left_context.tag_columns,
6027 aligned_right_context.tag_columns,
6028 )
6029 } else {
6030 (
6031 left,
6032 right,
6033 left_context.tag_columns.clone(),
6034 right_context.tag_columns.clone(),
6035 )
6036 };
6037
6038 let checked_matching = !only_join_time_index
6041 && !force_empty_join
6042 && !use_tsid_join
6043 && !matches!(
6044 modifier.as_ref().map(|modifier| &modifier.card),
6045 Some(VectorMatchCardinality::ManyToMany)
6046 );
6047 let one_side_is_left = matches!(
6049 modifier.as_ref().map(|modifier| &modifier.card),
6050 Some(VectorMatchCardinality::OneToMany(_))
6051 );
6052 let (left, right) = if !checked_matching {
6053 (left, right)
6054 } else if one_side_is_left {
6055 (
6056 Self::assert_unique_one_side(
6057 left,
6058 &left_tag_columns,
6059 &left_matched_tags,
6060 left_time_index_column.as_deref(),
6061 true,
6062 )?,
6063 right,
6064 )
6065 } else {
6066 (
6067 left,
6068 Self::assert_unique_one_side(
6069 right,
6070 &right_tag_columns,
6071 &right_matched_tags,
6072 right_time_index_column.as_deref(),
6073 false,
6074 )?,
6075 )
6076 };
6077
6078 if let (Some(left_time_index_column), Some(right_time_index_column)) = (
6080 left_time_index_column.clone(),
6081 right_time_index_column.clone(),
6082 ) {
6083 left_tag_columns.insert(left_time_index_column);
6084 right_tag_columns.insert(right_time_index_column);
6085 }
6086
6087 let right = LogicalPlanBuilder::from(right)
6088 .alias(right_table_ref.clone())
6089 .context(DataFusionPlanningSnafu)?
6090 .build()
6091 .context(DataFusionPlanningSnafu)?;
6092
6093 let join_plan = LogicalPlanBuilder::from(left)
6095 .alias(left_table_ref.clone())
6096 .context(DataFusionPlanningSnafu)?
6097 .join_detailed(
6098 right,
6099 JoinType::Inner,
6100 (
6101 left_tag_columns
6102 .into_iter()
6103 .map(Column::from_name)
6104 .collect::<Vec<_>>(),
6105 right_tag_columns
6106 .into_iter()
6107 .map(Column::from_name)
6108 .collect::<Vec<_>>(),
6109 ),
6110 force_empty_join.then_some(lit(false)),
6111 NullEquality::NullEqualsNull,
6112 )
6113 .context(DataFusionPlanningSnafu)?
6114 .build()
6115 .context(DataFusionPlanningSnafu)?;
6116
6117 let labels = checked_matching
6120 .then(|| Self::binary_result_labels(left_context, right_context, modifier))
6121 .flatten()
6122 .filter(|_| {
6123 Self::binary_result_labels_may_repeat(left_context, right_context, modifier)
6124 });
6125 let Some(labels) = labels else {
6126 return Ok(join_plan);
6127 };
6128 let schema = join_plan.schema().clone();
6129 let group_exprs = labels
6130 .iter()
6131 .map(|(from_left, label)| {
6132 let table_ref = if *from_left {
6133 &left_table_ref
6134 } else {
6135 &right_table_ref
6136 };
6137 schema
6138 .qualified_field_with_name(Some(table_ref), label)
6139 .map(|field| DfExpr::Column(field.into()))
6140 .context(DataFusionPlanningSnafu)
6141 })
6142 .collect::<Result<Vec<_>>>()?;
6143 let time_index_expr = left_time_index_column
6144 .as_deref()
6145 .map(|column| {
6146 schema
6147 .qualified_field_with_name(Some(&left_table_ref), column)
6148 .map(|field| DfExpr::Column(field.into()))
6149 .context(DataFusionPlanningSnafu)
6150 })
6151 .transpose()?
6152 .with_context(|| UnexpectedPlanExprSnafu {
6153 desc: "vector matching on a plan without a time index",
6154 })?;
6155 let violation = if matches!(
6156 modifier.as_ref().map(|modifier| &modifier.card),
6157 Some(VectorMatchCardinality::OneToOne)
6158 ) {
6159 MatchGroupViolation::ImplicitManyToOne
6160 } else {
6161 MatchGroupViolation::AmbiguousGroupLabels
6162 };
6163 Self::assert_unique_match_group(
6164 join_plan,
6165 group_exprs,
6166 labels.into_iter().map(|(_, label)| label).collect(),
6167 time_index_expr,
6168 violation,
6169 )
6170 }
6171
6172 fn assert_unique_one_side(
6174 plan: LogicalPlan,
6175 join_keys: &BTreeSet<String>,
6176 tag_columns: &[String],
6177 time_index_column: Option<&str>,
6178 one_side_is_left: bool,
6179 ) -> Result<LogicalPlan> {
6180 let Some(time_index_column) = time_index_column else {
6181 return Ok(plan);
6182 };
6183 if tag_columns.iter().all(|tag| join_keys.contains(tag)) {
6184 return Ok(plan);
6185 }
6186
6187 let group_labels = join_keys.iter().cloned().collect::<Vec<_>>();
6188 let group_exprs = group_labels
6189 .iter()
6190 .map(|label| DfExpr::Column(Column::from_name(label)))
6191 .collect();
6192 Self::assert_unique_match_group(
6193 plan,
6194 group_exprs,
6195 group_labels,
6196 DfExpr::Column(Column::from_name(time_index_column)),
6197 MatchGroupViolation::DuplicateOnOneSide { one_side_is_left },
6198 )
6199 }
6200
6201 fn selected_binary_match_labels(
6202 left_context: &PromPlannerContext,
6203 right_context: &PromPlannerContext,
6204 modifier: &Option<BinModifier>,
6205 ) -> BTreeSet<String> {
6206 let mut labels = left_context
6207 .tag_columns
6208 .iter()
6209 .chain(&right_context.tag_columns)
6210 .cloned()
6211 .collect::<BTreeSet<_>>();
6212 if let Some(matching) = modifier
6213 .as_ref()
6214 .and_then(|modifier| modifier.matching.as_ref())
6215 {
6216 match matching {
6217 LabelModifier::Include(on) => {
6218 labels = on
6219 .labels
6220 .iter()
6221 .filter(|label| {
6222 left_context.tag_columns.contains(label)
6223 || right_context.tag_columns.contains(label)
6224 })
6225 .cloned()
6226 .collect();
6227 }
6228 LabelModifier::Exclude(ignoring) => {
6229 for label in &ignoring.labels {
6230 labels.remove(label);
6231 }
6232 }
6233 }
6234 }
6235 labels
6236 }
6237
6238 fn only_temporality_match_label_mismatches(
6239 left_context: &PromPlannerContext,
6240 right_context: &PromPlannerContext,
6241 modifier: &Option<BinModifier>,
6242 ) -> bool {
6243 let mut mismatches =
6244 Self::selected_binary_match_labels(left_context, right_context, modifier)
6245 .into_iter()
6246 .filter(|label| {
6247 left_context.tag_columns.contains(label)
6248 != right_context.tag_columns.contains(label)
6249 });
6250 matches!(
6251 (mismatches.next(), mismatches.next()),
6252 (Some(label), None) if label == OTLP_AGGREGATION_TEMPORALITY_LABEL
6253 )
6254 }
6255
6256 fn align_temporality_match_column(
6257 mut left: LogicalPlan,
6258 mut right: LogicalPlan,
6259 left_context: &mut PromPlannerContext,
6260 right_context: &mut PromPlannerContext,
6261 ) -> Result<(LogicalPlan, LogicalPlan, bool)> {
6262 let marker = OTLP_AGGREGATION_TEMPORALITY_LABEL;
6263 let left_has_marker = left_context.tag_columns.iter().any(|tag| tag == marker);
6264 let (data_type, value_type, add_to_left) = {
6265 let (present, add_to_left) = if left_has_marker {
6266 (&left, false)
6267 } else {
6268 (&right, true)
6269 };
6270 let data_type = present
6271 .schema()
6272 .fields()
6273 .iter()
6274 .find(|field| field.name() == marker)
6275 .map(|field| field.data_type().clone())
6276 .with_context(|| ColumnNotFoundSnafu {
6277 col: marker.to_string(),
6278 })?;
6279 let value_type = Self::string_value_data_type(&data_type)
6280 .cloned()
6281 .with_context(|| UnexpectedPlanExprSnafu {
6282 desc: format!("temporality match label {marker} must be a string"),
6283 })?;
6284 (data_type, value_type, add_to_left)
6285 };
6286 let null = Self::string_scalar_value(&value_type, None).with_context(|| {
6287 UnexpectedPlanExprSnafu {
6288 desc: format!("temporality match label {marker} must be a string"),
6289 }
6290 })?;
6291 let add_marker = |plan: LogicalPlan| {
6292 let visible = plan
6293 .schema()
6294 .iter()
6295 .map(|(qualifier, field)| {
6296 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6297 })
6298 .collect::<Vec<_>>();
6299 LogicalPlanBuilder::from(plan)
6300 .project(
6301 visible
6302 .into_iter()
6303 .chain([DfExpr::Literal(null, None).alias(marker)]),
6304 )
6305 .context(DataFusionPlanningSnafu)?
6306 .build()
6307 .context(DataFusionPlanningSnafu)
6308 };
6309 if data_type != value_type {
6310 let present = if add_to_left { &mut right } else { &mut left };
6311 let visible = present
6312 .schema()
6313 .iter()
6314 .map(|(qualifier, field)| {
6315 let column =
6316 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()));
6317 if field.name() == marker {
6318 DfExpr::Cast(Cast::new(Box::new(column), value_type.clone()))
6319 .alias_qualified(qualifier.cloned(), field.name().clone())
6320 } else {
6321 column
6322 }
6323 })
6324 .collect::<Vec<_>>();
6325 *present = LogicalPlanBuilder::from(present.clone())
6326 .project(visible)
6327 .context(DataFusionPlanningSnafu)?
6328 .build()
6329 .context(DataFusionPlanningSnafu)?;
6330 }
6331 if add_to_left {
6332 left = add_marker(left)?;
6333 left_context.tag_columns.push(marker.to_string());
6334 } else {
6335 right = add_marker(right)?;
6336 right_context.tag_columns.push(marker.to_string());
6337 }
6338 Ok((left, right, add_to_left))
6339 }
6340
6341 fn normalized_match_key_expr(
6342 label: &str,
6343 field: Option<(Option<TableReference>, ArrowDataType)>,
6344 value_type: &ArrowDataType,
6345 internal_name: &str,
6346 ) -> DfExpr {
6347 let empty = Self::string_scalar_value(value_type, Some(String::new()))
6348 .expect("match label value type is a string");
6349 let expr = if let Some((qualifier, data_type)) = field {
6350 let column = DfExpr::Column(Column::new(qualifier, label));
6351 let column = if &data_type == value_type {
6352 column
6353 } else {
6354 DfExpr::Cast(Cast::new(Box::new(column), value_type.clone()))
6355 };
6356 DfExpr::ScalarFunction(ScalarFunction {
6357 func: coalesce(),
6358 args: vec![column, DfExpr::Literal(empty, None)],
6359 })
6360 } else {
6361 DfExpr::Literal(empty, None)
6362 };
6363 expr.alias(internal_name)
6364 }
6365
6366 fn is_zero_row_empty_relation(plan: &LogicalPlan) -> bool {
6367 matches!(plan, LogicalPlan::EmptyRelation(relation) if !relation.produce_one_row)
6370 }
6371
6372 fn string_value_data_type(data_type: &ArrowDataType) -> Option<&ArrowDataType> {
6373 match data_type {
6374 data_type if data_type.is_string() => Some(data_type),
6375 ArrowDataType::Dictionary(_, value_type) if value_type.is_string() => Some(value_type),
6376 _ => None,
6377 }
6378 }
6379
6380 fn string_scalar_value(
6381 data_type: &ArrowDataType,
6382 value: Option<String>,
6383 ) -> Option<ScalarValue> {
6384 match data_type {
6385 ArrowDataType::Utf8 => Some(ScalarValue::Utf8(value)),
6386 ArrowDataType::LargeUtf8 => Some(ScalarValue::LargeUtf8(value)),
6387 ArrowDataType::Utf8View => Some(ScalarValue::Utf8View(value)),
6388 ArrowDataType::Dictionary(key_type, value_type) => Some(ScalarValue::Dictionary(
6389 key_type.clone(),
6390 Box::new(Self::string_scalar_value(value_type, value)?),
6391 )),
6392 _ => None,
6393 }
6394 }
6395
6396 fn common_label_data_type(
6397 left: Option<&ArrowDataType>,
6398 right: Option<&ArrowDataType>,
6399 ) -> Option<ArrowDataType> {
6400 match (left, right) {
6401 (Some(left), Some(right)) if left == right => {
6402 Self::string_value_data_type(left).map(|_| left.clone())
6403 }
6404 (Some(left), Some(right)) => {
6405 let left_value_type = Self::string_value_data_type(left)?;
6406 let right_value_type = Self::string_value_data_type(right)?;
6407 match (left_value_type, right_value_type) {
6410 (left, right) if left == right => Some(left.clone()),
6411 (ArrowDataType::LargeUtf8, _) | (_, ArrowDataType::LargeUtf8) => {
6412 Some(ArrowDataType::LargeUtf8)
6413 }
6414 (ArrowDataType::Utf8View, ArrowDataType::Utf8View) => {
6415 Some(ArrowDataType::Utf8View)
6416 }
6417 _ => Some(ArrowDataType::Utf8),
6418 }
6419 }
6420 (Some(data_type), None) | (None, Some(data_type)) => {
6421 Self::string_value_data_type(data_type).cloned()
6422 }
6423 (None, None) => Some(ArrowDataType::Utf8),
6424 }
6425 }
6426
6427 fn projection_for_each_field_column<F>(
6435 &mut self,
6436 input: LogicalPlan,
6437 name_to_expr: F,
6438 ) -> Result<LogicalPlan>
6439 where
6440 F: FnMut(&String) -> Result<DfExpr>,
6441 {
6442 self.projection_for_each_field_column_with_labels(input, None, name_to_expr)
6443 }
6444
6445 fn projection_for_each_field_column_with_labels<F>(
6448 &mut self,
6449 input: LogicalPlan,
6450 result_labels: Option<&BinaryResultLabels>,
6451 name_to_expr: F,
6452 ) -> Result<LogicalPlan>
6453 where
6454 F: FnMut(&String) -> Result<DfExpr>,
6455 {
6456 let preserve_field_names =
6459 Self::field_columns_are_alternative_samples(input.schema(), &self.ctx.field_columns);
6460 let table_ref = self.ctx.table_name.clone().map(TableReference::bare);
6461 let input_schema = input.schema().clone();
6463 let lookup = |col: &String| {
6464 input_schema
6465 .qualified_field_with_name(table_ref.as_ref(), col)
6466 .or_else(|_| input_schema.qualified_field_with_unqualified_name(col))
6467 .map(|field| DfExpr::Column(field.into()))
6468 .context(DataFusionPlanningSnafu)
6469 };
6470 let tag_columns_iter = match result_labels {
6471 Some(labels) => labels.exprs.iter().cloned().map(Ok).collect::<Vec<_>>(),
6472 None => self.ctx.tag_columns.iter().map(lookup).collect::<Vec<_>>(),
6473 };
6474 let non_field_columns_iter = tag_columns_iter
6475 .into_iter()
6476 .chain(self.ctx.time_index_column.iter().map(lookup));
6477 let tsid_iter = match result_labels {
6478 Some(labels) => labels.tsid_projection(table_ref.clone()),
6479 None => Self::optional_tsid_projection(
6480 input.schema(),
6481 table_ref.as_ref(),
6482 self.ctx.use_tsid,
6483 ),
6484 }
6485 .into_iter()
6486 .map(Ok);
6487
6488 let result_field_columns = self
6490 .ctx
6491 .field_columns
6492 .iter()
6493 .map(name_to_expr)
6494 .collect::<Result<Vec<_>>>()?;
6495
6496 if !preserve_field_names {
6498 self.ctx.field_columns = result_field_columns
6499 .iter()
6500 .map(|expr| expr.schema_name().to_string())
6501 .collect();
6502 }
6503 let field_columns_iter = result_field_columns
6504 .into_iter()
6505 .zip(self.ctx.field_columns.iter())
6506 .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::<String>, name))));
6507
6508 let project_fields = non_field_columns_iter
6510 .chain(tsid_iter)
6511 .chain(field_columns_iter)
6512 .collect::<Result<Vec<_>>>()?;
6513
6514 LogicalPlanBuilder::from(input)
6515 .project(project_fields)
6516 .context(DataFusionPlanningSnafu)?
6517 .build()
6518 .context(DataFusionPlanningSnafu)
6519 }
6520
6521 fn filter_on_field_column<F>(&self, input: LogicalPlan, name_to_expr: F) -> Result<LogicalPlan>
6523 where
6524 F: FnMut(&String) -> Result<DfExpr>,
6525 {
6526 ensure!(
6527 self.ctx.field_columns.len() == 1
6528 || Self::field_columns_are_alternative_samples(
6529 input.schema(),
6530 &self.ctx.field_columns,
6531 ),
6532 UnsupportedExprSnafu {
6533 name: "filter on multi-value input"
6534 }
6535 );
6536
6537 let field_column_filters = self
6538 .ctx
6539 .field_columns
6540 .iter()
6541 .map(name_to_expr)
6542 .collect::<Result<Vec<_>>>()?;
6543 let field_column_filter =
6544 disjunction(field_column_filters).context(UnsupportedExprSnafu {
6545 name: "filter on empty input",
6546 })?;
6547
6548 LogicalPlanBuilder::from(input)
6549 .filter(field_column_filter)
6550 .context(DataFusionPlanningSnafu)?
6551 .build()
6552 .context(DataFusionPlanningSnafu)
6553 }
6554
6555 fn date_part_on_time_index(&self, date_part: &str) -> Result<DfExpr> {
6558 let input_expr = datafusion::logical_expr::ident(
6559 self.ctx
6560 .time_index_column
6561 .as_ref()
6562 .with_context(|| TimeIndexNotFoundSnafu {
6564 table: "<doesn't matter>",
6565 })?,
6566 );
6567 let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
6568 func: datafusion_functions::datetime::date_part(),
6569 args: vec![date_part.lit(), input_expr],
6570 });
6571 Ok(fn_expr)
6572 }
6573
6574 fn strip_tsid_column(&self, plan: LogicalPlan) -> Result<LogicalPlan> {
6575 let schema = plan.schema();
6576 if !schema
6577 .fields()
6578 .iter()
6579 .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
6580 {
6581 return Ok(plan);
6582 }
6583
6584 let project_exprs = schema
6587 .iter()
6588 .filter(|(_, field)| field.name() != DATA_SCHEMA_TSID_COLUMN_NAME)
6589 .map(|(qualifier, field)| {
6590 DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6591 })
6592 .collect::<Vec<_>>();
6593
6594 LogicalPlanBuilder::from(plan)
6595 .project(project_exprs)
6596 .context(DataFusionPlanningSnafu)?
6597 .build()
6598 .context(DataFusionPlanningSnafu)
6599 }
6600
6601 fn apply_alias(&mut self, plan: LogicalPlan, alias_name: String) -> Result<LogicalPlan> {
6603 let fields_expr = self.create_field_column_exprs()?;
6604
6605 ensure!(
6607 fields_expr.len() == 1,
6608 UnsupportedExprSnafu {
6609 name: "alias on multi-value result"
6610 }
6611 );
6612
6613 let project_fields = fields_expr
6614 .into_iter()
6615 .map(|expr| expr.alias(&alias_name))
6616 .chain(self.create_tag_column_exprs()?)
6617 .chain(Some(self.create_time_index_column_expr()?));
6618
6619 LogicalPlanBuilder::from(plan)
6620 .project(project_fields)
6621 .context(DataFusionPlanningSnafu)?
6622 .build()
6623 .context(DataFusionPlanningSnafu)
6624 }
6625}
6626
6627#[derive(Default, Debug)]
6628struct FunctionArgs {
6629 input: Option<PromExpr>,
6630 literals: Vec<DfExpr>,
6631}
6632
6633#[derive(Debug, Clone)]
6636enum ScalarFunc {
6637 DataFusionBuiltin(Arc<ScalarUdfDef>),
6641 DataFusionUdf(Arc<ScalarUdfDef>),
6645 NativeHistogramUdf(Arc<ScalarUdfDef>),
6648 Udf(Arc<ScalarUdfDef>),
6653 ExtrapolateUdf(Arc<ScalarUdfDef>, i64),
6660 GeneratedExpr,
6664}
6665
6666#[cfg(test)]
6667mod test;