Skip to main content

query/promql/
planner.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15mod 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
118/// Step of a subquery without an explicit one (`x[5m:]`). Prometheus uses the global
119/// evaluation interval, whose default is one minute, rather than the query step.
120const DEFAULT_SUBQUERY_STEP_MS: Millisecond = 60_000;
121/// `time()` function in PromQL.
122const SPECIAL_TIME_FUNCTION: &str = "time";
123/// `scalar()` function in PromQL.
124const SCALAR_FUNCTION: &str = "scalar";
125/// `absent()` function in PromQL
126const SPECIAL_ABSENT_FUNCTION: &str = "absent";
127/// `histogram_quantile` function in PromQL
128const SPECIAL_HISTOGRAM_QUANTILE: &str = "histogram_quantile";
129/// `histogram_fraction` function in PromQL
130const SPECIAL_HISTOGRAM_FRACTION: &str = "histogram_fraction";
131/// `vector` function in PromQL
132const SPECIAL_VECTOR_FUNCTION: &str = "vector";
133/// `le` column for conventional histogram.
134const LE_COLUMN_NAME: &str = "le";
135
136/// Static regex for validating label names according to Prometheus specification.
137/// Label names must match the regex: [a-zA-Z_][a-zA-Z0-9_]*
138static 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
143/// default value column name for empty metric
144const DEFAULT_FIELD_COLUMN: &str = "value";
145
146/// Special modifier to project field columns under multi-field mode
147const FIELD_COLUMN_MATCHER: &str = "__field__";
148
149/// Special modifier for cross schema query
150const 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_";
156/// Per-match-group row count the vector matching cardinality check is built on.
157const MATCH_GROUP_COUNT_COLUMN: &str = "__promql_match_group_count";
158
159/// Threshold for scatter scan mode
160const MAX_SCATTER_POINTS: i64 = 400;
161
162/// Interval 1 hour in millisecond
163const INTERVAL_1H: i64 = 60 * 60 * 1000;
164
165#[derive(Default, Debug, Clone)]
166struct PromPlannerContext {
167    // query parameters
168    start: Millisecond,
169    end: Millisecond,
170    interval: Millisecond,
171    lookback_delta: Millisecond,
172    /// Evaluation range of the whole statement, which `@ start()` and `@ end()` refer to.
173    ///
174    /// Unlike [`Self::start`] and [`Self::end`], these are never rewritten while planning, so a
175    /// selector inside a subquery still resolves `@ start()` / `@ end()` against the statement.
176    stmt_start: Millisecond,
177    stmt_end: Millisecond,
178
179    // planner states
180    table_name: Option<String>,
181    time_index_column: Option<String>,
182    field_columns: Vec<String>,
183    tag_columns: Vec<String>,
184    /// `by(...)` labels of the aggregation that produced this operand that are not series tags of
185    /// its input, i.e. value fields (or a label an inner aggregation already reported as one).
186    ///
187    /// An aggregation reports every `by(...)` label it finds in its input schema among its tag
188    /// columns, and a value field named there is a group key of the aggregate rather than a
189    /// property of a series: its value varies between the samples of one series. Lowering a
190    /// matcher on it into the scan would change which samples are selected (#9242), so
191    /// [`matching_filters`] refuses to propagate such a matcher.
192    aggregation_field_labels: Vec<String>,
193    /// Use metric engine internal series identifier column (`__tsid`) as series key.
194    ///
195    /// This is enabled only when the underlying scan can provide `__tsid` (`UInt64`). The planner
196    /// uses it internally (e.g. as the series key for [`SeriesDivide`]) and strips it from the
197    /// final output.
198    use_tsid: bool,
199    /// The matcher for field columns `__field__`.
200    field_column_matcher: Option<Vec<Matcher>>,
201    /// The matcher for selectors (normal matchers).
202    selector_matcher: Vec<Matcher>,
203    schema_name: Option<String>,
204    /// The range in millisecond of range selector. None if there is no range selector.
205    range: Option<Millisecond>,
206    /// The offset in milliseconds the window of the last planned range selector is folded with,
207    /// or `None` when no range selector has been planned since the last read.
208    ///
209    /// [`Self::start`] and the sample timestamps are compared on the shifted timeline: a range
210    /// payload carries `sample_timestamp + offset`, while the time index column of a folded row
211    /// stays the evaluation timestamp of its step. A function that reads it right after its input
212    /// plan is built, like `predict_linear`, consumes it instead of reading the state, so the
213    /// offset cannot leak from one input to another; see
214    /// [`PromPlanner::create_function_expr`].
215    range_fold_offset: Option<Millisecond>,
216}
217
218/// Result labels a vector-vector binary operation derives from its matching modifier, projected
219/// from the operand each label belongs to.
220#[derive(Debug)]
221struct BinaryResultLabels {
222    exprs: Vec<DfExpr>,
223    names: Vec<String>,
224    aggregation_field_labels: Vec<String>,
225    /// `__tsid` column of the operand the labels come from, when that operand contributes its
226    /// whole tag set and the column still identifies the result series.
227    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    /// `__tsid` is projected from whichever operand the labels come from, so it carries that
238    /// operand's qualifier. Re-qualify it as the result's own, which is what the context names
239    /// and what the enclosing expression looks the column up by.
240    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        // Evaluation times before the Unix epoch are valid in Prometheus.
250        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    /// Reset all planner states
264    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    /// Reset table name and schema to empty
278    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    /// Check if `le` is present in tag columns
285    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    /// Optional collector passed to native histogram UDFs.
294    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    /// Plans a PromQL statement and passes the optional collector to histogram UDFs.
309    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        // Never leak internal series identifier to output.
326        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    /**
339    Converts a PromQL expression to a logical plan.
340
341    NOTE:
342        The `timestamp_fn` indicates whether the PromQL `timestamp()` function is being evaluated in the current context.
343        If `true`, the planner generates a logical plan that projects the timestamp (time index) column
344        as the value column for each input row, implementing the PromQL `timestamp()` function semantics.
345        If `false`, the planner generates the standard logical plan for the given PromQL expression.
346    */
347    #[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        // An anchored range call is step-invariant: evaluate it once, at the start of the
355        // evaluation, and report its result at every step; see
356        // [`Self::promote_anchored_range_call`].
357        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    /// Plans the root expression of a statement, or the expression wrapped by a root extension
407    /// such as `EXPLAIN` or an alias.
408    ///
409    /// Prometheus only accepts a range-vector expression in an instant query, whose result is
410    /// the matrix of samples the expression selects.
411    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    /// Plans the result of a range-vector expression in an instant query: the raw samples of a
427    /// range selector, or the evaluation points of a subquery, with their own timestamps.
428    ///
429    /// Range functions consume the same expressions through [`RangeManipulate`], whose range
430    /// arrays are an internal representation that must not reach the query output.
431    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    /// Selects the samples of `selector` in `(t - offset - range, t - offset]`, where `t` is the
456    /// evaluation time or the `@` anchor. Stale markers are dropped and timestamps are kept.
457    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    /// Casts a microsecond or nanosecond time index to milliseconds, the precision of every
485    /// other PromQL result and the only one the Prometheus HTTP API encodes.
486    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    /// Evaluates the inner expression of `subquery` at the multiples of its step in
528    /// `(t - offset - range, t - offset]`, where `t` is the evaluation time or the `@` anchor.
529    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            // No step multiple falls in the window; plan the inner expression only for its
556            // schema.
557            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        // Shift the child window back; `RangeManipulate` restores the outer timeline.
584        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        // Reject an invalid parent window rather than treating it as empty.
599        ensure!(
600            current_start <= current_end,
601            InvalidTimeRangeSnafu {
602                start: current_start,
603                end: current_end,
604            }
605        );
606        // `@` makes the subquery step-invariant: like Prometheus (`setOffsetForAtModifier`), fold
607        // the anchor into the offset and evaluate once at `current_start`, then replay the result
608        // over the outer grid. See [`Self::at_modifier_offset`].
609        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        // An empty child needs only its schema; plan it over the valid caller window.
615        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        // Return no rows with the child schema so enclosing fallbacks still work.
632        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        // `RangeManipulate` assumes each input batch holds exactly one series
653        // (it takes tag column values from row 0 and applies them to every
654        // output row). The inner expression may emit batches that mix series,
655        // so sort by series key + time index and split into per-series batches
656        // with a `SeriesDivide` first.
657        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            // Only use tag columns that survive in the inner plan's schema —
672            // `ctx.tag_columns` can drift from the actual output.
673            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        // `RangeManipulate` has no offset in its protobuf message; decoding recovers it from the
702        // `SeriesNormalize` directly below. Stale markers are not filtered: the input is computed.
703        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        // The payload timestamps are shifted by the subquery offset (the rewritten one under `@`);
729        // see [`Self::create_range_eval_ts_expr`].
730        self.ctx.range_fold_offset = Some(offset_ms);
731        let manipulate = LogicalPlan::Extension(Extension {
732            node: Arc::new(manipulate),
733        });
734
735        // An anchored subquery was folded once, at `current_start`; report that window at every
736        // step of the outer grid, like an anchored range selector.
737        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                // When `__tsid` is available, tag columns may have been pruned from the input plan.
769                // For `keep_tsid` decision we should compare against the full row-key label set,
770                // otherwise we may incorrectly treat label-reducing aggregates as preserving labels.
771                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                // calculate columns to group by
779                // Need to append time index column into group by columns
780                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                // Aggregates over native histogram inputs may drop every sample in a group
785                // (e.g. `min` over histogram-only samples, or `sum` over histograms with
786                // incompatible schemas) and leave a NULL-valued group row behind. Compute this
787                // before `create_aggregate_exprs` mutates `ctx.field_columns`.
788                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                // convert op and value columns to aggregate exprs
792                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                // create plan
814                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                    // `count_values` must be grouped by fields,
818                    // and project the fields to the new label.
819                    //
820                    // The generated label is a real label (column) of the output, so it must be
821                    // registered in `ctx.tag_columns` below. Otherwise enclosing expressions
822                    // rebuild their projections from `ctx.tag_columns` and silently drop it.
823                    //
824                    // PromQL sets the generated label *before* the grouping key is built, so an
825                    // input label with the same name is overwritten by the sample value and must
826                    // not remain a grouping key either: samples are grouped by the generated
827                    // label only. Dropping it from the projected tag columns is required as well:
828                    // projecting both would emit two columns with the same name (rejected as an
829                    // ambiguous reference).
830                    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                    // The tag columns projected below are unqualified `Column` references, so they
835                    // end up qualified with whatever qualifier they carry in the input plan. Give
836                    // the generated label the same qualifier, otherwise qualified references to it
837                    // (e.g. from an enclosing binary expression or a vector join) fail to resolve.
838                    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                    // The generated label carries the *sample value*, so it must be materialized
850                    // as a string using Prometheus' format
851                    // (`strconv.FormatFloat(value, 'f', -1, 64)`: shortest decimal form without
852                    // an exponent, `1.0` becomes "1"). Labels are inferred from string columns
853                    // when a result is converted into the Prometheus HTTP API JSON format, so a
854                    // numeric label column would be mistaken for the sample value of the series.
855                    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                                    // `prom_float_to_string` formats exactly like Prometheus,
862                                    // while arrow's `Float64 -> Utf8` cast would render `1.0`.
863                                    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                                // The value is already formatted (e.g. a native histogram is
880                                // converted to its string form by the aggregation), and the
881                                // aggregate output names it by its schema name.
882                                _ => 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                    // The label only exists in the output schema from here on.
908                    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                // Drop group rows whose every aggregated sample was discarded (NULL), so that
958                // e.g. `group(min(native_histogram))` doesn't resurrect groups Prometheus
959                // considers unseen. For alternative float/histogram fields keep the row if any
960                // field survived.
961                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    /// Create logical plan for PromQL topk and bottomk expr.
981    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        // convert op and value columns to window exprs.
1055        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        // Create ranks filter with `Operator::Or`.
1063        // Safety: at least one rank column
1064        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        // Order by ranks
1088        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        // Unary Expr in PromQL implys the `-` operator
1127        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        // promql-parser accepts fill modifiers, but Greptime does not implement the
1153        // required outer joins and missing-value substitution. Reject them before the
1154        // binary-island fast path so they cannot silently behave like normal inner joins.
1155        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        // if set to true, comparison operator will return 0/1 (for true/false) instead of
1176        // filter on the result column
1177        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        // we should build a filter plan here if the op is comparison op and need not
1185        // to return 0/1. Otherwise, we should build a projection plan
1186        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            // lhs is a literal, rhs is a column
1217            (Some(mut expr), None) => {
1218                let input = self.prom_expr_to_plan(rhs, query_engine_state).await?;
1219                // Arithmetic against a literal preserves the series labels, so
1220                // `aggregation_field_labels` passes through with `tag_columns` unchanged.
1221                // check if the literal is a special time expr
1222                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            // lhs is a column, rhs is a literal
1286            (None, Some(mut expr)) => {
1287                let input = self.prom_expr_to_plan(lhs, query_engine_state).await?;
1288                // check if the literal is a special time expr
1289                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            // both are columns. join them on time index
1353            (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                // Both operands are planned first because only the planned contexts tell a tag
1371                // from a value field. The rewrite merely adds matchers to a selector, so the
1372                // table reference, time index and field columns captured above stay valid.
1373                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                    // A copied matcher belongs to the scan, not to the operand's identity:
1381                    // `absent()` turns `selector_matcher` into the labels it reports, so the
1382                    // re-planned context keeps the matchers the operand was written with.
1383                    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                    // The code below reads `self.ctx` as the right operand's context.
1400                    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                // TODO(ruihang): avoid join if left and right are the same table
1407
1408                // set op has "special" join semantics
1409                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                // normal join
1429                if left_table_ref == right_table_ref {
1430                    // rename table references to avoid ambiguity
1431                    left_table_ref = TableReference::bare("lhs");
1432                    right_table_ref = TableReference::bare("rhs");
1433                    // `self.ctx` have ctx in right plan, if right plan have no tag,
1434                    // we use left plan ctx as the ctx for subsequent calculations,
1435                    // to avoid case like `host + scalar(...)`
1436                    // we need preserve tag column on `host` table in subsequent projection,
1437                    // which only show in left plan ctx.
1438                    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                    // A computed scalar on the right (`time()`, `scalar(...)`) has no labels;
1450                    // the result keeps the labels of the vector on the left.
1451                    self.ctx = left_context.clone();
1452                }
1453                // An operand without a table name, such as a comparison filter's projection, is
1454                // joined under the `""` qualifier; resolve the projected columns against it.
1455                // Otherwise a time index present on both sides is ambiguous.
1456                if self.ctx.table_name.is_none() {
1457                    self.ctx.table_name = Some(String::new());
1458                }
1459                // Computed scalars reach this join path instead of the literal projection paths.
1460                // Broadcast them for arithmetic in the same way as literal scalars.
1461                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                // Regular multi-field vectors combine their shared prefix. Alternative
1488                // float/histogram lanes instead align by valid PromQL sample combinations.
1489                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                // `vector()` uses EmptyMetric and keeps GreptimeDB's timestamp broadcast.
1495                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                // The matching modifier derives the result labels instead of the operands keeping
1521                // their own tag set. An operand that is broadcast rather than matched (a scalar,
1522                // or a vector without tags) keeps the other side's labels as before.
1523                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                // These predicates always pass; they only evaluate otherwise-discarded pairs
1544                // while collecting annotations.
1545                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                    // PromQL comparison operators without `bool` are filters:
1642                    //   - keep the instant-vector side sample values
1643                    //   - drop samples where the comparison is false
1644                    //
1645                    // So we filter on the join result and then project only the side that should
1646                    // be preserved according to PromQL semantics.
1647                    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        // Project time index from the chosen side.
1769        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        // Project field columns from the chosen side.
1778        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        // Project tag columns: the labels the matching modifier derived, or the chosen side's.
1787        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        // Preserve `__tsid` if present, so it can still be used internally downstream. It's
1801        // stripped from the final output anyway.
1802        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        // Update context to reflect the projected schema. Don't keep a table qualifier since
1817        // the result is a derived expression.
1818        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    /// The offset of a selector in milliseconds. A positive offset selects samples from an earlier
1875    /// time and moves them forward into the evaluation timeline.
1876    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    /// Returns the first absolute step multiple strictly after `start - offset - range`
1885    /// through the last parent evaluation instant minus `offset`. Requires a positive step.
1886    /// Returns `None` for an empty window; an unrepresentable nonempty bound is an error.
1887    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        // Check emptiness before converting bounds to i64.
1904        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    /// The columns that identify one series, which is the series key expected by the PromQL plan
1926    /// nodes that hold exactly one series per input batch.
1927    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    /// Keep replay and series division on the same effective keys when a call rewrites labels.
1936    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        // `@` anchors the sample window at a fixed timestamp: the selector selects its samples
1970        // around the anchor once, instead of following the outer evaluation grid. See
1971        // [`Self::at_modifier_offset`].
1972        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                // Select the anchored samples at the start of the evaluation, with the offset that
1978                // re-anchors the selector.
1979                // The planner is single-use (one `EvalStmt` produces one plan), so an error
1980                // below aborts the whole planning and `ctx.end` needs no restore-on-error.
1981                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            // Keep the original sample for stale-marker detection while carrying
2003            // its timestamp through InstantManipulate in a private value column.
2004            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            // The time index still holds the raw sample timestamp here, which is what
2022            // `timestamp()` reports regardless of `offset` and `@`.
2023            let sample_time = Self::timestamp_seconds_expr(&time_index_column, normalize.schema())?;
2024            project_exprs.push(sample_time.alias(&timestamp_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                // Select the anchored sample once, then report it at every step of the outer
2040                // grid. The samples keep their native anchor-time timestamps, so the manipulate
2041                // must shift them onto the evaluation timeline with the rewritten offset.
2042                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    /// Converts the timestamp column `column` into PromQL seconds, truncated to milliseconds.
2084    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    /// Builds a projection plan for the PromQL `timestamp()` function, which reports
2109    /// `timestamp_value` as the value of each row, along with the original tag and time index
2110    /// columns.
2111    ///
2112    /// Updates the planner context's field columns to the timestamp column name.
2113    fn create_timestamp_func_plan(
2114        &mut self,
2115        input: LogicalPlan,
2116        timestamp_value: DfExpr,
2117    ) -> Result<LogicalPlan> {
2118        // A row whose fields are all NULL holds no sample, so it must not get a timestamp. The
2119        // check reads the input fields, which the projection below replaces.
2120        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        // `@` anchors the range selector's window at a fixed timestamp, so the same window is fed
2155        // to the enclosing function at every evaluation step. See [`Self::at_modifier_offset`].
2156        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        // Some functions like rate may require special fields in the RangeManipulate plan
2161        // so we can't skip RangeManipulate.
2162        let (normalize, at_offset) = match self.setup_context().await? {
2163            // An empty metric does not contain any sample, so anchoring cannot change the result.
2164            // The manipulate below folds the empty input with `offset_ms`, and the recorded fold
2165            // offset agrees with it instead of with the anchor the window cannot use.
2166            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                        // Fold the anchored window once, at the start of the evaluation.
2174                        // Single-use planner: an error below aborts planning, so `ctx.end`
2175                        // needs no restore-on-error.
2176                        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                // Samples are shifted onto the evaluation timeline with the very same offset while
2189                // the window is folded. Record it so that the function above the selector can
2190                // recover the evaluation instant of the folded window; see
2191                // [`Self::create_range_eval_ts_expr`].
2192                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                // Fold the anchored window once, then report it at every step of the outer
2204                // grid. The samples keep their native anchor-time timestamps, so the manipulate
2205                // must shift them onto the evaluation timeline with the rewritten offset.
2206                let anchored = RangeManipulate::new(
2207                    grid_start,
2208                    grid_start,
2209                    self.ctx.interval,
2210                    at_offset,
2211                    // TODO(ruihang): convert via Timestamp datatypes to support different time units
2212                    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                    // TODO(ruihang): convert via Timestamp datatypes to support different time units
2234                    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        // some special functions that are not expression but a plan
2257        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        // transform function arguments
2272        let args = self.create_function_args(&args.args)?;
2273        // Only a vector selector keeps the timestamps of its samples. Any other expression
2274        // produces samples at the evaluation time, which is what `timestamp()` reports for it.
2275        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        // The input plan records the fold offset of the range selector it is built from. Take it
2319        // here, so that the offset of one input cannot leak into another call, and pass it to
2320        // `create_function_expr`: the function that reads it (`predict_linear`) then depends on an
2321        // argument instead of on planner state written by the selector below it.
2322        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        // A row survives as long as one field column produced a sample, and the fields without
2340        // one stay NULL, which is the shape a selector already emits. Requiring every field to
2341        // be non-NULL would drop one field's samples because another field has none in the same
2342        // window — the reason alternative float/histogram columns already needed this form. A
2343        // single field column reduces to the same predicate either way.
2344        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        // Rewriting a label the input series already have can map several of them onto the same
2376        // label set, which PromQL rejects. A new label keeps the series distinct.
2377        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        // Update context tags after building plan
2381        // We can't push them before planning, because they won't exist until projection.
2382        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 promql_parser::parser::ast::Extension { expr } = ext_expr;
2411        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        // Wrapper for the explanation/analyze of the existing plan
2417        // https://docs.rs/datafusion-expr/latest/datafusion_expr/logical_plan/builder/struct.LogicalPlanBuilder.html#method.explain
2418        // if `analyze` is true, runs the actual plan and produces
2419        // information about metrics during run.
2420        // if `verbose` is true, prints out additional details when VERBOSE keyword is specified
2421        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    /// Extract metric name from `__name__` matcher and set it into [PromPlannerContext].
2460    /// Returns a new [Matchers] that doesn't contain metric name matcher.
2461    ///
2462    /// Each call to this function means new selector is started. Thus, the context will be reset
2463    /// at first.
2464    ///
2465    /// Name rule:
2466    /// - if `name` is some, then the matchers MUST NOT contain `__name__` matcher.
2467    /// - if `name` is none, then the matchers MAY contain NONE OR MULTIPLE `__name__` matchers.
2468    #[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        // Deduplicate in place instead of through a `HashSet`: the scan filter is built in
2500        // this order, and a hashed order makes the plan of one query vary between runs.
2501        let mut matchers: Vec<Matcher> = Vec::with_capacity(label_matchers.matchers.len());
2502        for matcher in &label_matchers.matchers {
2503            // TODO(ruihang): support other metric match ops
2504            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        // make table scan plan
2536        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        // make filter exprs
2541        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        // make a projection plan if there is any `__field__` matcher
2556        if let Some(field_matchers) = &self.ctx.field_column_matcher {
2557            let col_set = self.ctx.field_columns.iter().collect::<HashSet<_>>();
2558            // opt-in set
2559            let mut result_set = HashSet::new();
2560            // opt-out set
2561            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            // merge two set
2601            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            // mask the field columns in context using computed result set
2609            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            // reuse this variable for simplicity
2631            table_scan = LogicalPlanBuilder::from(table_scan)
2632                .project(exprs)
2633                .context(DataFusionPlanningSnafu)?
2634                .build()
2635                .context(DataFusionPlanningSnafu)?;
2636        }
2637
2638        // make sort plan
2639        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        // make divide plan
2657        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        // make series_normalize plan
2673        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    /// Convert [LabelModifier] to [Column] exprs for aggregation.
2696    /// Timestamp column and tag columns will be included.
2697    ///
2698    /// # Side effect
2699    ///
2700    /// This method will also change the tag columns in ctx if `update_ctx` is true.
2701    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                    // A `by(...)` label can name a value field of the input instead of a tag. The
2718                    // aggregate still reports it among its tag columns below, but unlike a tag it
2719                    // is a group key rather than a property of a series: its value varies between
2720                    // the samples of one series, so a matcher on it must stay above sample
2721                    // selection (#9242). Record it, before the tag columns are overwritten.
2722                    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                    // nonexistence label will be ignored
2740                    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                            // update the tag columns in context
2746                            self.ctx.tag_columns.push(column_name);
2747                        }
2748                    }
2749                }
2750                // add timestamp column
2751                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                // Exclude metric engine internal columns (not PromQL labels) from the implicit
2763                // "without" label set.
2764                all_fields.retain(|col| !is_metric_engine_internal_column(col.as_str()));
2765
2766                // remove "without"-ed fields
2767                // nonexistence label will be ignored
2768                for label in &labels.labels {
2769                    let _ = all_fields.remove(label);
2770                }
2771
2772                // remove time index and value fields
2773                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                    // change the tag columns in context
2782                    self.ctx.tag_columns = all_fields.iter().map(|col| (*col).clone()).collect();
2783                    // `without(...)` drops the value fields of the input from its grouping labels,
2784                    // so a label stays a grouping label in name only if it survives in the input
2785                    // schema (e.g. an inner aggregation that grouped by it).
2786                    self.ctx
2787                        .aggregation_field_labels
2788                        .retain(|label| all_fields.iter().any(|col| *col == label));
2789                }
2790
2791                // collect remaining fields and convert to col expr
2792                let mut exprs = all_fields
2793                    .into_iter()
2794                    .map(|c| DfExpr::Column(Column::from_name(c)))
2795                    .collect::<Vec<_>>();
2796
2797                // add timestamp column
2798                exprs.push(self.create_time_index_column_expr()?);
2799
2800                Ok(exprs)
2801            }
2802        }
2803    }
2804
2805    // TODO(ruihang): ignore `MetricNameLabel` (`__name__`) matcher
2806    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            // Prometheus reads a label a series does not carry as the empty
2822            // string. A row can miss a label two ways: the table has no column
2823            // for it, or the column exists but is NULL on that row — the latter
2824            // is the norm for logical metrics sharing a physical table, which
2825            // holds the union of their label columns.
2826            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                    // TODO(ruihang): a more programmatic way to handle this in datafusion
2850
2851                    // This is a hack to handle `.+` and `.*`, and is not strictly correct
2852                    // `.` doesn't match newline (`\n`). Given this is in PromQL context,
2853                    // most of the time it's fine.
2854                    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        // set schema name if `__schema__` is given
2932        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            // An underflowing lower bound must not discard a representable upper
3006            // bound: without it, LastRow could retain a future row and discard the
3007            // older eligible sample before the manipulator can check its time.
3008            match (lower_filter, upper_filter) {
3009                (Some(lower), Some(upper)) => Some(lower.and(upper)),
3010                (Some(filter), None) | (None, Some(filter)) => Some(filter),
3011                (None, None) => None,
3012            }
3013        };
3014        let bounds = |timestamp: i64| {
3015            let upper = i128::from(timestamp) - i128::from(offset_duration);
3016            (upper - i128::from(window), upper)
3017        };
3018        let num_points = (end as i128 - start as i128) / self.ctx.interval as i128;
3019        if num_points > MAX_SCATTER_POINTS as i128 || self.ctx.interval <= INTERVAL_1H {
3020            let (lower, _) = bounds(start);
3021            let (_, upper) = bounds(end);
3022            return Ok(filter(lower, upper));
3023        }
3024        let mut filters = Vec::new();
3025        for timestamp in (start..=end).step_by(self.ctx.interval as usize) {
3026            let (lower, upper) = bounds(timestamp);
3027            let Some(filter) = filter(lower, upper) else {
3028                // A point whose native bounds cannot be represented may cover the whole native
3029                // time domain, so its disjunct cannot be omitted.
3030                return Ok(None);
3031            };
3032            filters.push(filter);
3033        }
3034        Ok(filters.into_iter().reduce(DfExpr::or))
3035    }
3036
3037    /// Create a table scan plan and a filter plan with given filter.
3038    ///
3039    /// # Panic
3040    /// If the filter is empty
3041    async fn create_table_scan_plan(&mut self, table_ref: TableReference) -> Result<LogicalPlan> {
3042        let provider = self
3043            .table_provider
3044            .resolve_table(table_ref.clone())
3045            .await
3046            .context(CatalogSnafu)?;
3047
3048        let logical_table = self.table_from_source(&provider)?;
3049
3050        // Try to rewrite the table scan to physical table scan if possible.
3051        let mut maybe_phy_table_ref = table_ref.clone();
3052        let mut scan_provider = provider;
3053        let mut table_id_filter: Option<u32> = None;
3054
3055        // If it's a metric engine logical table, scan its physical table directly and filter by
3056        // `__table_id = logical_table_id` to get access to internal columns like `__tsid`.
3057        if logical_table.table_info().meta.engine == METRIC_ENGINE_NAME
3058            && let Some(physical_table_name) = logical_table
3059                .table_info()
3060                .meta
3061                .options
3062                .extra_options
3063                .get(LOGICAL_TABLE_METADATA_KEY)
3064        {
3065            let physical_table_ref = if let Some(schema_name) = &self.ctx.schema_name {
3066                TableReference::partial(schema_name.as_str(), physical_table_name.as_str())
3067            } else {
3068                TableReference::bare(physical_table_name.as_str())
3069            };
3070
3071            let physical_provider = match self
3072                .table_provider
3073                .resolve_table(physical_table_ref.clone())
3074                .await
3075            {
3076                Ok(provider) => provider,
3077                Err(e) if e.status_code() == StatusCode::TableNotFound => {
3078                    // Fall back to scanning the logical table. It still works, but without
3079                    // `__tsid` optimization.
3080                    scan_provider.clone()
3081                }
3082                Err(e) => return Err(e).context(CatalogSnafu),
3083            };
3084
3085            if !Arc::ptr_eq(&physical_provider, &scan_provider) {
3086                // Only rewrite when internal columns exist in physical schema.
3087                let physical_table = self.table_from_source(&physical_provider)?;
3088
3089                let has_table_id = physical_table
3090                    .schema()
3091                    .column_schema_by_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME)
3092                    .is_some();
3093                let has_tsid = physical_table
3094                    .schema()
3095                    .column_schema_by_name(DATA_SCHEMA_TSID_COLUMN_NAME)
3096                    .is_some_and(|col| matches!(col.data_type, ConcreteDataType::UInt64(_)));
3097
3098                if has_table_id && has_tsid {
3099                    scan_provider = physical_provider;
3100                    maybe_phy_table_ref = physical_table_ref;
3101                    table_id_filter = Some(logical_table.table_info().ident.table_id);
3102                }
3103            }
3104        }
3105
3106        let scan_table = self.table_from_source(&scan_provider)?;
3107
3108        let use_tsid = table_id_filter.is_some()
3109            && scan_table
3110                .schema()
3111                .column_schema_by_name(DATA_SCHEMA_TSID_COLUMN_NAME)
3112                .is_some_and(|col| matches!(col.data_type, ConcreteDataType::UInt64(_)));
3113        self.ctx.use_tsid = use_tsid;
3114
3115        let all_table_tags = self.ctx.tag_columns.clone();
3116
3117        let scan_tag_columns = if use_tsid {
3118            let mut scan_tags = self.ctx.tag_columns.clone();
3119            for matcher in &self.ctx.selector_matcher {
3120                if is_metric_engine_internal_column(&matcher.name) {
3121                    continue;
3122                }
3123                if all_table_tags.iter().any(|tag| tag == &matcher.name) {
3124                    scan_tags.push(matcher.name.clone());
3125                }
3126            }
3127            scan_tags.sort_unstable();
3128            scan_tags.dedup();
3129            scan_tags
3130        } else {
3131            self.ctx.tag_columns.clone()
3132        };
3133
3134        let time_index_data_type = scan_table
3135            .schema()
3136            .timestamp_column()
3137            .with_context(|| TimeIndexNotFoundSnafu {
3138                table: maybe_phy_table_ref.to_quoted_string(),
3139            })?
3140            .data_type
3141            .clone();
3142        let is_time_index_second =
3143            time_index_data_type == ConcreteDataType::timestamp_second_datatype();
3144
3145        let scan_projection = if table_id_filter.is_some() {
3146            let mut required_columns = HashSet::new();
3147            required_columns.insert(DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string());
3148            required_columns.insert(self.ctx.time_index_column.clone().with_context(|| {
3149                TimeIndexNotFoundSnafu {
3150                    table: maybe_phy_table_ref.to_quoted_string(),
3151                }
3152            })?);
3153            for col in &scan_tag_columns {
3154                required_columns.insert(col.clone());
3155            }
3156            for col in &self.ctx.field_columns {
3157                required_columns.insert(col.clone());
3158            }
3159            if use_tsid {
3160                required_columns.insert(DATA_SCHEMA_TSID_COLUMN_NAME.to_string());
3161            }
3162
3163            let arrow_schema = scan_table.schema().arrow_schema().clone();
3164            Some(
3165                arrow_schema
3166                    .fields()
3167                    .iter()
3168                    .enumerate()
3169                    .filter(|(_, field)| required_columns.contains(field.name().as_str()))
3170                    .map(|(idx, _)| idx)
3171                    .collect::<Vec<_>>(),
3172            )
3173        } else {
3174            None
3175        };
3176
3177        let mut scan_plan =
3178            LogicalPlanBuilder::scan(maybe_phy_table_ref.clone(), scan_provider, scan_projection)
3179                .context(DataFusionPlanningSnafu)?
3180                .build()
3181                .context(DataFusionPlanningSnafu)?;
3182
3183        if let Some(table_id) = table_id_filter {
3184            scan_plan = LogicalPlanBuilder::from(scan_plan)
3185                .filter(
3186                    DfExpr::Column(Column::from_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME))
3187                        .eq(lit(table_id)),
3188                )
3189                .context(DataFusionPlanningSnafu)?
3190                .alias(table_ref.clone()) // rename the relation back to logical table's name after filtering
3191                .context(DataFusionPlanningSnafu)?
3192                .build()
3193                .context(DataFusionPlanningSnafu)?;
3194        }
3195
3196        if is_time_index_second {
3197            // Promote seconds so millisecond offsets remain exact; retain finer precision.
3198            // Later manipulators compare native sample ticks, while PromQL evaluation and
3199            // emitted timestamps remain millisecond-based, so this projection must not
3200            // silently truncate a finer-grained time index.
3201            let expr: Vec<_> = self
3202                .create_field_column_exprs()?
3203                .into_iter()
3204                .chain(
3205                    scan_tag_columns
3206                        .iter()
3207                        .map(|tag| DfExpr::Column(Column::from_name(tag))),
3208                )
3209                .chain(self.ctx.use_tsid.then_some(DfExpr::Column(Column::new(
3210                    Some(table_ref.clone()),
3211                    DATA_SCHEMA_TSID_COLUMN_NAME.to_string(),
3212                ))))
3213                .chain(Some(DfExpr::Alias(Alias {
3214                    expr: Box::new(DfExpr::Cast(Cast::new(
3215                        Box::new(self.create_time_index_column_expr()?),
3216                        ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
3217                    ))),
3218                    relation: Some(table_ref.clone()),
3219                    name: self
3220                        .ctx
3221                        .time_index_column
3222                        .as_ref()
3223                        .with_context(|| TimeIndexNotFoundSnafu {
3224                            table: table_ref.to_quoted_string(),
3225                        })?
3226                        .clone(),
3227                    metadata: None,
3228                })))
3229                .collect::<Vec<_>>();
3230            scan_plan = LogicalPlanBuilder::from(scan_plan)
3231                .project(expr)
3232                .context(DataFusionPlanningSnafu)?
3233                .build()
3234                .context(DataFusionPlanningSnafu)?;
3235        } else if table_id_filter.is_some()
3236            || time_index_data_type == ConcreteDataType::timestamp_microsecond_datatype()
3237            || time_index_data_type == ConcreteDataType::timestamp_nanosecond_datatype()
3238        {
3239            // Drop the internal `__table_id` column after filtering and preserve PromQL's
3240            // field/tag/timestamp column order for native microsecond/nanosecond timestamps.
3241            // Keeping the original time column also lets the existing ordering hints
3242            // use PerSeries scans without a cast, repartition, and sort. This benefits
3243            // multi-evaluation selectors too; only a single evaluation can use LastRow.
3244            let project_exprs = self
3245                .create_field_column_exprs()?
3246                .into_iter()
3247                .chain(
3248                    scan_tag_columns
3249                        .iter()
3250                        .map(|tag| DfExpr::Column(Column::from_name(tag))),
3251                )
3252                .chain(
3253                    self.ctx
3254                        .use_tsid
3255                        .then_some(DfExpr::Column(Column::from_name(
3256                            DATA_SCHEMA_TSID_COLUMN_NAME,
3257                        ))),
3258                )
3259                .chain(Some(self.create_time_index_column_expr()?))
3260                .collect::<Vec<_>>();
3261
3262            scan_plan = LogicalPlanBuilder::from(scan_plan)
3263                .project(project_exprs)
3264                .context(DataFusionPlanningSnafu)?
3265                .build()
3266                .context(DataFusionPlanningSnafu)?;
3267        }
3268
3269        let result = LogicalPlanBuilder::from(scan_plan)
3270            .build()
3271            .context(DataFusionPlanningSnafu)?;
3272        Ok(result)
3273    }
3274
3275    fn collect_row_key_tag_columns_from_plan(
3276        &self,
3277        plan: &LogicalPlan,
3278    ) -> Result<BTreeSet<String>> {
3279        fn walk(
3280            planner: &PromPlanner,
3281            plan: &LogicalPlan,
3282            out: &mut BTreeSet<String>,
3283        ) -> Result<()> {
3284            // Derived PromQL plans may contain non-Greptime scans without row-key metadata.
3285            if let LogicalPlan::TableScan(scan) = plan
3286                && let Ok(table) = planner.table_from_source(&scan.source)
3287            {
3288                for col in table.table_info().meta.row_key_column_names() {
3289                    if col != DATA_SCHEMA_TABLE_ID_COLUMN_NAME
3290                        && col != DATA_SCHEMA_TSID_COLUMN_NAME
3291                        && !is_metric_engine_internal_column(col)
3292                    {
3293                        out.insert(col.clone());
3294                    }
3295                }
3296            }
3297
3298            for input in plan.inputs() {
3299                walk(planner, input, out)?;
3300            }
3301            Ok(())
3302        }
3303
3304        let mut out = BTreeSet::new();
3305        walk(self, plan, &mut out)?;
3306        Ok(out)
3307    }
3308
3309    /// Setup [PromPlannerContext]'s state fields.
3310    ///
3311    /// Returns a logical plan for an empty metric.
3312    async fn setup_context(&mut self) -> Result<Option<LogicalPlan>> {
3313        let table_ref = self.table_ref()?;
3314        let source = match self.table_provider.resolve_table(table_ref.clone()).await {
3315            Err(e) if e.status_code() == StatusCode::TableNotFound => {
3316                let plan = self.setup_context_for_empty_metric()?;
3317                return Ok(Some(plan));
3318            }
3319            res => res.context(CatalogSnafu)?,
3320        };
3321        let table = self.table_from_source(&source)?;
3322
3323        // set time index column name
3324        let time_index = table
3325            .schema()
3326            .timestamp_column()
3327            .with_context(|| TimeIndexNotFoundSnafu {
3328                table: table_ref.to_quoted_string(),
3329            })?
3330            .name
3331            .clone();
3332        self.ctx.time_index_column = Some(time_index);
3333
3334        // set values columns
3335        let values = table
3336            .table_info()
3337            .meta
3338            .field_column_names()
3339            .cloned()
3340            .collect();
3341        self.ctx.field_columns = values;
3342
3343        // set primary key (tag) columns
3344        let tags = table
3345            .table_info()
3346            .meta
3347            .row_key_column_names()
3348            .filter(|col| {
3349                // remove metric engine's internal columns
3350                col != &DATA_SCHEMA_TABLE_ID_COLUMN_NAME && col != &DATA_SCHEMA_TSID_COLUMN_NAME
3351            })
3352            .cloned()
3353            .collect();
3354        self.ctx.tag_columns = tags;
3355        // The operand is a plain selector: its tag columns are the table's.
3356        self.ctx.aggregation_field_labels.clear();
3357
3358        self.ctx.use_tsid = false;
3359
3360        Ok(None)
3361    }
3362
3363    /// Setup [PromPlannerContext]'s state fields for a non existent table
3364    /// without any rows.
3365    fn setup_context_for_empty_metric(&mut self) -> Result<LogicalPlan> {
3366        self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
3367        self.ctx.reset_table_name_and_schema();
3368        self.ctx.tag_columns = vec![];
3369        self.ctx.aggregation_field_labels.clear();
3370        self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
3371        self.ctx.use_tsid = false;
3372
3373        // The table doesn't have any data, so we set start to 0 and end to -1.
3374        let plan = LogicalPlan::Extension(Extension {
3375            node: Arc::new(
3376                EmptyMetric::new(
3377                    0,
3378                    -1,
3379                    self.ctx.interval,
3380                    SPECIAL_TIME_FUNCTION.to_string(),
3381                    DEFAULT_FIELD_COLUMN.to_string(),
3382                    Some(lit(0.0f64)),
3383                )
3384                .context(DataFusionPlanningSnafu)?,
3385            ),
3386        });
3387        Ok(plan)
3388    }
3389
3390    // TODO(ruihang): insert column expr
3391    fn create_function_args(&self, args: &[Box<PromExpr>]) -> Result<FunctionArgs> {
3392        let mut result = FunctionArgs::default();
3393
3394        for arg in args {
3395            // First try to parse as literal expression (including binary expressions like 100.0 + 3.0)
3396            if let Some(expr) = Self::try_build_literal_expr(arg) {
3397                result.literals.push(expr);
3398            } else {
3399                // If not a literal, treat as vector input
3400                match arg.as_ref() {
3401                    PromExpr::Subquery(_)
3402                    | PromExpr::VectorSelector(_)
3403                    | PromExpr::MatrixSelector(_)
3404                    | PromExpr::Extension(_)
3405                    | PromExpr::Aggregate(_)
3406                    | PromExpr::Paren(_)
3407                    | PromExpr::Call(_)
3408                    | PromExpr::Binary(_)
3409                    | PromExpr::Unary(_) => {
3410                        if result.input.replace(*arg.clone()).is_some() {
3411                            MultipleVectorSnafu { expr: *arg.clone() }.fail()?;
3412                        }
3413                    }
3414
3415                    _ => {
3416                        let expr = Self::get_param_as_literal_expr(Some(arg.as_ref()), None, None)?;
3417                        result.literals.push(expr);
3418                    }
3419                }
3420            }
3421        }
3422
3423        Ok(result)
3424    }
3425
3426    fn create_mixed_range_function_exprs(
3427        &mut self,
3428        func: &Function,
3429        mut other_input_exprs: VecDeque<DfExpr>,
3430        float_field: &str,
3431        histogram_field: &str,
3432        input_schema: &DFSchemaRef,
3433        range_fold_offset: Option<Millisecond>,
3434    ) -> Result<Option<Vec<DfExpr>>> {
3435        let returns_histogram = matches!(
3436            func.name,
3437            "rate"
3438                | "increase"
3439                | "delta"
3440                | "idelta"
3441                | "irate"
3442                | "avg_over_time"
3443                | "sum_over_time"
3444                | "last_over_time"
3445        );
3446        if !returns_histogram
3447            && !matches!(
3448                func.name,
3449                "changes"
3450                    | "resets"
3451                    | "deriv"
3452                    | "min_over_time"
3453                    | "max_over_time"
3454                    | "count_over_time"
3455                    | "absent_over_time"
3456                    | "present_over_time"
3457                    | "stddev_over_time"
3458                    | "stdvar_over_time"
3459                    | "quantile_over_time"
3460                    | "predict_linear"
3461                    | "double_exponential_smoothing"
3462                    | "holt_winters"
3463            )
3464        {
3465            return Ok(None);
3466        }
3467
3468        if func.name == "predict_linear" {
3469            other_input_exprs[0] = DfExpr::Cast(Cast::new(
3470                Box::new(other_input_exprs[0].clone()),
3471                ArrowDataType::Int64,
3472            ));
3473            // Same evaluation instant as in the non-mixed path; it follows the other inputs so the
3474            // float UDF below is called as `predict_linear(ts_range, value_range, t, eval_ts)`.
3475            // The regression is only defined for a folded window, so a missing fold offset is an
3476            // error instead of a default.
3477            other_input_exprs.push_back(self.create_range_eval_ts_expr(
3478                range_fold_offset.context(ExpectRangeSelectorSnafu)?,
3479                input_schema,
3480            )?);
3481        }
3482
3483        let timestamp_range = DfExpr::Column(Column::from_name(
3484            RangeManipulate::build_timestamp_range_name(
3485                self.ctx.time_index_column.as_ref().unwrap(),
3486            ),
3487        ));
3488        let float_range = DfExpr::Column(Column::from_name(float_field));
3489        let histogram_range = DfExpr::Column(Column::from_name(histogram_field));
3490        let mut args = Vec::with_capacity(other_input_exprs.len() + 6);
3491        args.push(lit(func.name));
3492        args.push(timestamp_range.clone());
3493        args.push(float_range.clone());
3494        args.push(histogram_range.clone());
3495        args.extend(other_input_exprs);
3496        if matches!(func.name, "rate" | "increase" | "delta") {
3497            args.push(self.create_time_index_column_expr()?);
3498            args.push(lit(self.ctx.range.context(ExpectRangeSelectorSnafu)?));
3499        }
3500
3501        let mut float_expr = DfExpr::ScalarFunction(ScalarFunction {
3502            func: Arc::new(MixedRange::float_udf(self.promql_annotations.clone())),
3503            args: args.clone(),
3504        });
3505        if matches!(func.name, "rate" | "increase") {
3506            let raw_delta_function = if func.name == "rate" {
3507                "raw_delta_rate"
3508            } else {
3509                "raw_delta_increase"
3510            };
3511            let delta_sum = DfExpr::ScalarFunction(ScalarFunction {
3512                func: Arc::new(MixedRange::float_udf(self.promql_annotations.clone())),
3513                args: vec![
3514                    lit(raw_delta_function),
3515                    timestamp_range,
3516                    float_range,
3517                    histogram_range,
3518                ],
3519            });
3520            float_expr = self.select_delta_range_math(
3521                func.name,
3522                input_schema,
3523                self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3524                delta_sum,
3525                float_expr,
3526            )?;
3527        }
3528        let exprs = if returns_histogram {
3529            self.ctx.field_columns = vec![float_field.to_string(), histogram_field.to_string()];
3530            vec![
3531                float_expr.alias(float_field),
3532                DfExpr::ScalarFunction(ScalarFunction {
3533                    func: Arc::new(MixedRange::histogram_udf(self.promql_annotations.clone())),
3534                    args,
3535                })
3536                .alias(histogram_field),
3537            ]
3538        } else {
3539            let display_name = if func.name == "predict_linear" {
3540                // The evaluation instant is the private last argument of the mixed float UDF
3541                // call; keep it out of the output column name like on the non-mixed path.
3542                Self::name_without_last_arg(&float_expr)
3543            } else {
3544                float_expr.schema_name().to_string()
3545            };
3546            self.ctx.field_columns = vec![display_name.clone()];
3547            vec![float_expr.alias(display_name)]
3548        };
3549        Ok(Some(exprs))
3550    }
3551
3552    /// Creates function expressions for projection and returns the expressions and new tags.
3553    ///
3554    /// # Side Effects
3555    ///
3556    /// This method will update [PromPlannerContext]'s fields and tags if needed.
3557    fn create_function_expr(
3558        &mut self,
3559        func: &Function,
3560        other_input_exprs: Vec<DfExpr>,
3561        input_schema: &DFSchemaRef,
3562        query_engine_state: &QueryEngineState,
3563        range_fold_offset: Option<Millisecond>,
3564    ) -> Result<(Vec<DfExpr>, Vec<String>)> {
3565        // TODO(ruihang): check function args list
3566        let mut other_input_exprs: VecDeque<DfExpr> = other_input_exprs.into();
3567        if let Some((float_field, histogram_field)) =
3568            Self::alternative_sample_range_columns(input_schema, &self.ctx.field_columns)
3569                .map(|(float, histogram)| (float.to_string(), histogram.to_string()))
3570            && let Some(exprs) = self.create_mixed_range_function_exprs(
3571                func,
3572                other_input_exprs.clone(),
3573                &float_field,
3574                &histogram_field,
3575                input_schema,
3576                range_fold_offset,
3577            )?
3578        {
3579            return Ok((exprs, vec![]));
3580        }
3581        let alternative_samples =
3582            Self::field_columns_are_alternative_samples(input_schema, &self.ctx.field_columns);
3583        let all_field_columns_are_native_histogram_ranges =
3584            self.all_field_columns_are_native_histogram_ranges(input_schema);
3585
3586        // TODO(ruihang): set this according to in-param list
3587        let field_column_pos = 0;
3588        let mut exprs = Vec::with_capacity(self.ctx.field_columns.len());
3589        // New labels after executing the function, e.g. `label_replace` etc.
3590        let mut new_tags = vec![];
3591        let promql_annotations = self.promql_annotations.clone();
3592        let native_histogram_drop_udf = |name: &str| {
3593            Arc::new(NativeHistogramDrop::float_null_udf(
3594                format!(
3595                    "{name}: dropped native histogram samples because this function is not supported for native histograms"
3596                ),
3597                promql_annotations.clone(),
3598            ))
3599        };
3600        let scalar_func = match func.name {
3601            "increase" => {
3602                if all_field_columns_are_native_histogram_ranges {
3603                    ScalarFunc::ExtrapolateUdf(
3604                        Arc::new(NativeHistogramIncrease::scalar_udf_with_collector(
3605                            self.promql_annotations.clone(),
3606                        )),
3607                        self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3608                    )
3609                } else {
3610                    ScalarFunc::ExtrapolateUdf(
3611                        Arc::new(Increase::scalar_udf()),
3612                        self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3613                    )
3614                }
3615            }
3616            "rate" => {
3617                if all_field_columns_are_native_histogram_ranges {
3618                    ScalarFunc::ExtrapolateUdf(
3619                        Arc::new(NativeHistogramRate::scalar_udf_with_collector(
3620                            self.promql_annotations.clone(),
3621                        )),
3622                        self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3623                    )
3624                } else {
3625                    ScalarFunc::ExtrapolateUdf(
3626                        Arc::new(Rate::scalar_udf()),
3627                        self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3628                    )
3629                }
3630            }
3631            "delta" => {
3632                if all_field_columns_are_native_histogram_ranges {
3633                    ScalarFunc::ExtrapolateUdf(
3634                        Arc::new(NativeHistogramDelta::scalar_udf_with_collector(
3635                            self.promql_annotations.clone(),
3636                        )),
3637                        self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3638                    )
3639                } else {
3640                    ScalarFunc::ExtrapolateUdf(
3641                        Arc::new(Delta::scalar_udf()),
3642                        self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3643                    )
3644                }
3645            }
3646            "idelta" => {
3647                if all_field_columns_are_native_histogram_ranges {
3648                    ScalarFunc::Udf(Arc::new(NativeHistogramIDelta::scalar_udf_with_collector(
3649                        self.promql_annotations.clone(),
3650                    )))
3651                } else {
3652                    ScalarFunc::Udf(Arc::new(IDelta::<false>::scalar_udf()))
3653                }
3654            }
3655            "irate" => {
3656                if all_field_columns_are_native_histogram_ranges {
3657                    ScalarFunc::Udf(Arc::new(NativeHistogramIRate::scalar_udf_with_collector(
3658                        self.promql_annotations.clone(),
3659                    )))
3660                } else {
3661                    ScalarFunc::Udf(Arc::new(IDelta::<true>::scalar_udf()))
3662                }
3663            }
3664            "resets" => {
3665                if all_field_columns_are_native_histogram_ranges {
3666                    ScalarFunc::Udf(Arc::new(NativeHistogramResets::scalar_udf()))
3667                } else {
3668                    ScalarFunc::Udf(Arc::new(Resets::scalar_udf()))
3669                }
3670            }
3671            "changes" => {
3672                if all_field_columns_are_native_histogram_ranges {
3673                    ScalarFunc::Udf(Arc::new(NativeHistogramChanges::scalar_udf()))
3674                } else {
3675                    ScalarFunc::Udf(Arc::new(Changes::scalar_udf()))
3676                }
3677            }
3678            "deriv" => {
3679                if all_field_columns_are_native_histogram_ranges {
3680                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3681                } else {
3682                    ScalarFunc::Udf(Arc::new(Deriv::scalar_udf()))
3683                }
3684            }
3685            "avg_over_time" => {
3686                if all_field_columns_are_native_histogram_ranges {
3687                    ScalarFunc::Udf(Arc::new(
3688                        NativeHistogramAvgOverTime::scalar_udf_with_collector(
3689                            self.promql_annotations.clone(),
3690                        ),
3691                    ))
3692                } else {
3693                    ScalarFunc::Udf(Arc::new(AvgOverTime::scalar_udf()))
3694                }
3695            }
3696            "min_over_time" => {
3697                if all_field_columns_are_native_histogram_ranges {
3698                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3699                } else {
3700                    ScalarFunc::Udf(Arc::new(MinOverTime::scalar_udf()))
3701                }
3702            }
3703            "max_over_time" => {
3704                if all_field_columns_are_native_histogram_ranges {
3705                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3706                } else {
3707                    ScalarFunc::Udf(Arc::new(MaxOverTime::scalar_udf()))
3708                }
3709            }
3710            "sum_over_time" => {
3711                if all_field_columns_are_native_histogram_ranges {
3712                    ScalarFunc::Udf(Arc::new(
3713                        NativeHistogramSumOverTime::scalar_udf_with_collector(
3714                            self.promql_annotations.clone(),
3715                        ),
3716                    ))
3717                } else {
3718                    ScalarFunc::Udf(Arc::new(SumOverTime::scalar_udf()))
3719                }
3720            }
3721            "count_over_time" => {
3722                if all_field_columns_are_native_histogram_ranges {
3723                    ScalarFunc::Udf(Arc::new(NativeHistogramCountOverTime::scalar_udf()))
3724                } else {
3725                    ScalarFunc::Udf(Arc::new(CountOverTime::scalar_udf()))
3726                }
3727            }
3728            "last_over_time" => {
3729                if all_field_columns_are_native_histogram_ranges {
3730                    ScalarFunc::Udf(Arc::new(NativeHistogramLastOverTime::scalar_udf()))
3731                } else {
3732                    ScalarFunc::Udf(Arc::new(LastOverTime::scalar_udf()))
3733                }
3734            }
3735            "absent_over_time" => {
3736                if all_field_columns_are_native_histogram_ranges {
3737                    ScalarFunc::Udf(Arc::new(NativeHistogramAbsentOverTime::scalar_udf()))
3738                } else {
3739                    ScalarFunc::Udf(Arc::new(AbsentOverTime::scalar_udf()))
3740                }
3741            }
3742            "present_over_time" => {
3743                if all_field_columns_are_native_histogram_ranges {
3744                    ScalarFunc::Udf(Arc::new(NativeHistogramPresentOverTime::scalar_udf()))
3745                } else {
3746                    ScalarFunc::Udf(Arc::new(PresentOverTime::scalar_udf()))
3747                }
3748            }
3749            "stddev_over_time" => {
3750                if all_field_columns_are_native_histogram_ranges {
3751                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3752                } else {
3753                    ScalarFunc::Udf(Arc::new(StddevOverTime::scalar_udf()))
3754                }
3755            }
3756            "stdvar_over_time" => {
3757                if all_field_columns_are_native_histogram_ranges {
3758                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3759                } else {
3760                    ScalarFunc::Udf(Arc::new(StdvarOverTime::scalar_udf()))
3761                }
3762            }
3763            "quantile_over_time" => {
3764                if all_field_columns_are_native_histogram_ranges {
3765                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3766                } else {
3767                    ScalarFunc::Udf(Arc::new(QuantileOverTime::scalar_udf()))
3768                }
3769            }
3770            "predict_linear" => {
3771                if all_field_columns_are_native_histogram_ranges {
3772                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3773                } else {
3774                    other_input_exprs[0] = DfExpr::Cast(Cast::new(
3775                        Box::new(other_input_exprs[0].clone()),
3776                        ArrowDataType::Int64,
3777                    ));
3778                    // The prediction starts at the evaluation instant of the step, which the
3779                    // window's fold offset recovers from the row's time index; see
3780                    // [`Self::create_range_eval_ts_expr`]. It is appended last, so the UDF is
3781                    // called as `predict_linear(ts_range, value_range, t, eval_ts)`. The
3782                    // regression is only defined for a folded window, so a missing fold offset
3783                    // is an error instead of a default.
3784                    other_input_exprs.push_back(self.create_range_eval_ts_expr(
3785                        range_fold_offset.context(ExpectRangeSelectorSnafu)?,
3786                        input_schema,
3787                    )?);
3788                    ScalarFunc::Udf(Arc::new(PredictLinear::scalar_udf()))
3789                }
3790            }
3791            "double_exponential_smoothing" | "holt_winters" => {
3792                if all_field_columns_are_native_histogram_ranges {
3793                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3794                } else {
3795                    ScalarFunc::Udf(Arc::new(DoubleExponentialSmoothing::scalar_udf()))
3796                }
3797            }
3798            "histogram_count" => {
3799                ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramCount::scalar_udf()))
3800            }
3801            "histogram_sum" => {
3802                ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramSum::scalar_udf()))
3803            }
3804            "histogram_avg" => {
3805                ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramAvg::scalar_udf()))
3806            }
3807            "histogram_stddev" => {
3808                ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramStddev::scalar_udf()))
3809            }
3810            "histogram_stdvar" => {
3811                ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramStdvar::scalar_udf()))
3812            }
3813            "time" => {
3814                exprs.push(build_special_time_expr(
3815                    self.ctx.time_index_column.as_ref().unwrap(),
3816                ));
3817                ScalarFunc::GeneratedExpr
3818            }
3819            "minute" => {
3820                // date_part('minute', time_index)
3821                let expr = self.date_part_on_time_index("minute")?;
3822                exprs.push(expr);
3823                ScalarFunc::GeneratedExpr
3824            }
3825            "hour" => {
3826                // date_part('hour', time_index)
3827                let expr = self.date_part_on_time_index("hour")?;
3828                exprs.push(expr);
3829                ScalarFunc::GeneratedExpr
3830            }
3831            "month" => {
3832                // date_part('month', time_index)
3833                let expr = self.date_part_on_time_index("month")?;
3834                exprs.push(expr);
3835                ScalarFunc::GeneratedExpr
3836            }
3837            "year" => {
3838                // date_part('year', time_index)
3839                let expr = self.date_part_on_time_index("year")?;
3840                exprs.push(expr);
3841                ScalarFunc::GeneratedExpr
3842            }
3843            "day_of_month" => {
3844                // date_part('day', time_index)
3845                let expr = self.date_part_on_time_index("day")?;
3846                exprs.push(expr);
3847                ScalarFunc::GeneratedExpr
3848            }
3849            "day_of_week" => {
3850                // date_part('dow', time_index)
3851                let expr = self.date_part_on_time_index("dow")?;
3852                exprs.push(expr);
3853                ScalarFunc::GeneratedExpr
3854            }
3855            "day_of_year" => {
3856                // date_part('doy', time_index)
3857                let expr = self.date_part_on_time_index("doy")?;
3858                exprs.push(expr);
3859                ScalarFunc::GeneratedExpr
3860            }
3861            "days_in_month" => {
3862                // date_part(
3863                //     'days',
3864                //     (date_trunc('month', <TIME INDEX>::date) + interval '1 month - 1 day')
3865                // );
3866                let day_lit_expr = "day".lit();
3867                let month_lit_expr = "month".lit();
3868                let interval_1month_lit_expr =
3869                    DfExpr::Literal(ScalarValue::IntervalYearMonth(Some(1)), None);
3870                let interval_1day_lit_expr = DfExpr::Literal(
3871                    ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(1, 0))),
3872                    None,
3873                );
3874                let the_1month_minus_1day_expr = DfExpr::BinaryExpr(BinaryExpr {
3875                    left: Box::new(interval_1month_lit_expr),
3876                    op: Operator::Minus,
3877                    right: Box::new(interval_1day_lit_expr),
3878                });
3879                let date_trunc_expr = DfExpr::ScalarFunction(ScalarFunction {
3880                    func: datafusion_functions::datetime::date_trunc(),
3881                    args: vec![month_lit_expr, self.create_time_index_column_expr()?],
3882                });
3883                let date_trunc_plus_interval_expr = DfExpr::BinaryExpr(BinaryExpr {
3884                    left: Box::new(date_trunc_expr),
3885                    op: Operator::Plus,
3886                    right: Box::new(the_1month_minus_1day_expr),
3887                });
3888                let date_part_expr = DfExpr::ScalarFunction(ScalarFunction {
3889                    func: datafusion_functions::datetime::date_part(),
3890                    args: vec![day_lit_expr, date_trunc_plus_interval_expr],
3891                });
3892
3893                exprs.push(date_part_expr);
3894                ScalarFunc::GeneratedExpr
3895            }
3896
3897            "label_join" => {
3898                self.ctx.use_tsid = false;
3899                let (concat_expr, dst_label) = Self::build_concat_labels_expr(
3900                    &mut other_input_exprs,
3901                    &self.ctx,
3902                    input_schema,
3903                    query_engine_state,
3904                )?;
3905
3906                // Reserve the current field columns except the `dst_label`.
3907                for value in &self.ctx.field_columns {
3908                    if *value != dst_label {
3909                        let expr = DfExpr::Column(Column::from_name(value));
3910                        exprs.push(expr);
3911                    }
3912                }
3913
3914                // Remove it from tag columns if exists to avoid duplicated column names
3915                self.ctx.tag_columns.retain(|tag| *tag != dst_label);
3916                new_tags.push(dst_label);
3917                // Add the new label expr to evaluate
3918                exprs.push(concat_expr);
3919
3920                ScalarFunc::GeneratedExpr
3921            }
3922            "label_replace" => {
3923                self.ctx.use_tsid = false;
3924                if let Some((replace_expr, dst_label)) =
3925                    self.build_regexp_replace_label_expr(&mut other_input_exprs, input_schema)?
3926                {
3927                    // Reserve the current field columns except the `dst_label`.
3928                    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                    // Remove it from tag columns if exists to avoid duplicated column names
3936                    self.ctx.tag_columns.retain(|tag| *tag != dst_label);
3937                    new_tags.push(dst_label);
3938                    // Add the new label expr to evaluate
3939                    exprs.push(replace_expr);
3940                } else {
3941                    // Keep the current field columns
3942                    for value in &self.ctx.field_columns {
3943                        let expr = DfExpr::Column(Column::from_name(value));
3944                        exprs.push(expr);
3945                    }
3946                }
3947
3948                ScalarFunc::GeneratedExpr
3949            }
3950            "sort" | "sort_desc" => {
3951                // Value sorting silently ignores native histogram samples.
3952                for value in &self.ctx.field_columns {
3953                    if !Self::field_column_is_native_histogram(input_schema, value) {
3954                        exprs.push(DfExpr::Column(Column::from_name(value)));
3955                    }
3956                }
3957                // Keep a nullable float field so the normal empty-value filter produces an
3958                // empty vector when the input contains only histograms.
3959                if exprs.is_empty() {
3960                    exprs.push(DfExpr::Literal(ScalarValue::Float64(None), None));
3961                }
3962
3963                ScalarFunc::GeneratedExpr
3964            }
3965            "sort_by_label" | "sort_by_label_desc" | "timestamp" => {
3966                // These functions are not expression but a part of plan,
3967                // they are processed by `prom_call_expr_to_plan`.
3968                for value in &self.ctx.field_columns {
3969                    let expr = DfExpr::Column(Column::from_name(value));
3970                    exprs.push(expr);
3971                }
3972
3973                ScalarFunc::GeneratedExpr
3974            }
3975            "round" if self.all_field_columns_are_native_histograms(input_schema) => {
3976                if other_input_exprs.is_empty() {
3977                    other_input_exprs.push_front(0.0f64.lit());
3978                }
3979                ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3980            }
3981            "round" => {
3982                if other_input_exprs.is_empty() {
3983                    other_input_exprs.push_front(1.0f64.lit());
3984                }
3985                ScalarFunc::DataFusionUdf(Arc::new(Round::scalar_udf()))
3986            }
3987            // DataFusion's `sqrt` rejects negative input; PromQL returns NaN for it.
3988            "sqrt" if !self.all_field_columns_are_native_histograms(input_schema) => {
3989                ScalarFunc::DataFusionUdf(Arc::new(IeeeSqrt::scalar_udf()))
3990            }
3991            "rad" | "deg" | "sgn" if self.all_field_columns_are_native_histograms(input_schema) => {
3992                ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3993            }
3994            "rad" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::radians()),
3995            "deg" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::degrees()),
3996            "sgn" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::signum()),
3997            "pi" => {
3998                // pi functions doesn't accepts any arguments, needs special processing
3999                let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
4000                    func: datafusion::functions::math::pi(),
4001                    args: vec![],
4002                });
4003                exprs.push(fn_expr);
4004
4005                ScalarFunc::GeneratedExpr
4006            }
4007            _ => {
4008                if let Some(f) = query_engine_state
4009                    .session_state()
4010                    .scalar_functions()
4011                    .get(func.name)
4012                {
4013                    if self.all_field_columns_are_native_histograms(input_schema) {
4014                        ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
4015                    } else {
4016                        ScalarFunc::DataFusionBuiltin(f.clone())
4017                    }
4018                } else if let Some(factory) = query_engine_state.scalar_function(func.name) {
4019                    if self.all_field_columns_are_native_histograms(input_schema) {
4020                        ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
4021                    } else {
4022                        let func_state = query_engine_state.function_state();
4023                        let query_ctx = self.table_provider.query_ctx();
4024
4025                        ScalarFunc::DataFusionUdf(Arc::new(factory.provide(FunctionContext {
4026                            state: func_state,
4027                            query_ctx: query_ctx.clone(),
4028                        })))
4029                    }
4030                } else if let Some(f) = datafusion_functions::math::functions()
4031                    .iter()
4032                    .find(|f| f.name() == func.name)
4033                {
4034                    if self.all_field_columns_are_native_histograms(input_schema) {
4035                        ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
4036                    } else {
4037                        ScalarFunc::DataFusionUdf(f.clone())
4038                    }
4039                } else {
4040                    return UnsupportedExprSnafu {
4041                        name: func.name.to_string(),
4042                    }
4043                    .fail();
4044                }
4045            }
4046        };
4047
4048        for value in &self.ctx.field_columns {
4049            let col_expr = DfExpr::Column(Column::from_name(value));
4050            let value_is_histogram = Self::field_column_is_native_histogram(input_schema, value);
4051
4052            match scalar_func.clone() {
4053                ScalarFunc::DataFusionBuiltin(func) => {
4054                    if alternative_samples && value_is_histogram {
4055                        continue;
4056                    }
4057                    other_input_exprs.insert(field_column_pos, col_expr);
4058                    let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
4059                        func,
4060                        args: other_input_exprs.clone().into(),
4061                    });
4062                    exprs.push(fn_expr);
4063                    let _ = other_input_exprs.remove(field_column_pos);
4064                }
4065                ScalarFunc::DataFusionUdf(func) => {
4066                    if alternative_samples && value_is_histogram {
4067                        continue;
4068                    }
4069                    let args = itertools::chain!(
4070                        other_input_exprs.iter().take(field_column_pos).cloned(),
4071                        std::iter::once(col_expr),
4072                        other_input_exprs.iter().skip(field_column_pos).cloned()
4073                    )
4074                    .collect_vec();
4075                    exprs.push(DfExpr::ScalarFunction(ScalarFunction { func, args }))
4076                }
4077                ScalarFunc::NativeHistogramUdf(func) => {
4078                    if value_is_histogram {
4079                        let args = itertools::chain!(
4080                            other_input_exprs.iter().take(field_column_pos).cloned(),
4081                            std::iter::once(col_expr),
4082                            other_input_exprs.iter().skip(field_column_pos).cloned()
4083                        )
4084                        .collect_vec();
4085                        exprs.push(DfExpr::ScalarFunction(ScalarFunction { func, args }));
4086                    } else if !alternative_samples {
4087                        exprs.push(
4088                            DfExpr::Literal(ScalarValue::Float64(None), None).alias(format!(
4089                                "{}_{}",
4090                                func.name(),
4091                                value
4092                            )),
4093                        );
4094                    }
4095                }
4096                ScalarFunc::Udf(func) => {
4097                    let ts_range_expr = DfExpr::Column(Column::from_name(
4098                        RangeManipulate::build_timestamp_range_name(
4099                            self.ctx.time_index_column.as_ref().unwrap(),
4100                        ),
4101                    ));
4102                    other_input_exprs.insert(field_column_pos, ts_range_expr);
4103                    other_input_exprs.insert(field_column_pos + 1, col_expr);
4104                    let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
4105                        func,
4106                        args: other_input_exprs.clone().into(),
4107                    });
4108                    exprs.push(fn_expr);
4109                    let _ = other_input_exprs.remove(field_column_pos + 1);
4110                    let _ = other_input_exprs.remove(field_column_pos);
4111                }
4112                ScalarFunc::ExtrapolateUdf(udf, range_length) => {
4113                    let ts_range_expr = DfExpr::Column(Column::from_name(
4114                        RangeManipulate::build_timestamp_range_name(
4115                            self.ctx.time_index_column.as_ref().unwrap(),
4116                        ),
4117                    ));
4118                    other_input_exprs.insert(field_column_pos, ts_range_expr.clone());
4119                    other_input_exprs.insert(field_column_pos + 1, col_expr.clone());
4120                    other_input_exprs
4121                        .insert(field_column_pos + 2, self.create_time_index_column_expr()?);
4122                    other_input_exprs.push_back(lit(range_length));
4123                    let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
4124                        func: udf,
4125                        args: other_input_exprs.clone().into(),
4126                    });
4127                    let fn_expr = if matches!(func.name, "rate" | "increase")
4128                        && !all_field_columns_are_native_histogram_ranges
4129                    {
4130                        let delta_sum = DfExpr::ScalarFunction(ScalarFunction {
4131                            func: Arc::new(SumOverTime::scalar_udf()),
4132                            args: vec![ts_range_expr, col_expr],
4133                        });
4134                        self.select_delta_range_math(
4135                            func.name,
4136                            input_schema,
4137                            range_length,
4138                            delta_sum,
4139                            fn_expr,
4140                        )?
4141                    } else {
4142                        fn_expr
4143                    };
4144                    exprs.push(fn_expr);
4145                    let _ = other_input_exprs.pop_back();
4146                    let _ = other_input_exprs.remove(field_column_pos + 2);
4147                    let _ = other_input_exprs.remove(field_column_pos + 1);
4148                    let _ = other_input_exprs.remove(field_column_pos);
4149                }
4150                ScalarFunc::GeneratedExpr => {}
4151            }
4152        }
4153
4154        // Update value columns' name, and alias them to remove qualifiers
4155        // For label functions such as `label_join`, `label_replace`, etc.,
4156        // we keep the fields unchanged.
4157        if !matches!(func.name, "label_join" | "label_replace") {
4158            let mut new_field_columns = Vec::with_capacity(exprs.len());
4159
4160            exprs = exprs
4161                .into_iter()
4162                .map(|expr| {
4163                    // `predict_linear` appends its private evaluation instant as the last argument
4164                    // of the UDF call; the output column is named after the call without it, so
4165                    // the injected expression stays out of the user-visible schema. The
4166                    // native-histogram drop UDF takes no such argument.
4167                    let display_name = if func.name == "predict_linear"
4168                        && !all_field_columns_are_native_histogram_ranges
4169                    {
4170                        Self::name_without_last_arg(&expr)
4171                    } else {
4172                        expr.schema_name().to_string()
4173                    };
4174                    new_field_columns.push(display_name.clone());
4175                    Ok(expr.alias(display_name))
4176                })
4177                .collect::<std::result::Result<Vec<_>, _>>()
4178                .context(DataFusionPlanningSnafu)?;
4179
4180            self.ctx.field_columns = new_field_columns;
4181        }
4182
4183        Ok((exprs, new_tags))
4184    }
4185
4186    fn select_delta_range_math(
4187        &self,
4188        function: &str,
4189        input_schema: &DFSchemaRef,
4190        range_length: Millisecond,
4191        delta_sum: DfExpr,
4192        cumulative: DfExpr,
4193    ) -> Result<DfExpr> {
4194        let marker_is_delta = if self
4195            .ctx
4196            .tag_columns
4197            .iter()
4198            .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL)
4199        {
4200            Self::field_column_type(input_schema, OTLP_AGGREGATION_TEMPORALITY_LABEL)
4201                .filter(|data_type| Self::string_value_data_type(data_type).is_some())
4202                .map(|_| {
4203                    DfExpr::Column(Column::from_name(OTLP_AGGREGATION_TEMPORALITY_LABEL))
4204                        .eq(lit(GREPTIME_TEMPORALITY_DELTA))
4205                })
4206        } else {
4207            None
4208        };
4209        let Some(marker_is_delta) = marker_is_delta else {
4210            return Ok(cumulative);
4211        };
4212
4213        let delta = if function == "rate" {
4214            DfExpr::BinaryExpr(BinaryExpr {
4215                left: Box::new(delta_sum),
4216                op: Operator::Divide,
4217                right: Box::new(lit(range_length as f64 / 1000.0)),
4218            })
4219        } else {
4220            delta_sum
4221        };
4222        let display_name = cumulative.schema_name().to_string();
4223        when(marker_is_delta, delta)
4224            .otherwise(cumulative)
4225            .context(DataFusionPlanningSnafu)
4226            .map(|expr| expr.alias(display_name))
4227    }
4228
4229    /// Validate label name according to Prometheus specification.
4230    /// Label names must match the regex: [a-zA-Z_][a-zA-Z0-9_]*
4231    /// Additionally, label names starting with double underscores are reserved for internal use.
4232    fn validate_label_name(label_name: &str) -> Result<()> {
4233        // Check if label name starts with double underscores (reserved)
4234        if label_name.starts_with("__") {
4235            return InvalidDestinationLabelNameSnafu { label_name }.fail();
4236        }
4237        // Check if label name matches the required pattern
4238        if !LABEL_NAME_REGEX.is_match(label_name) {
4239            return InvalidDestinationLabelNameSnafu { label_name }.fail();
4240        }
4241
4242        Ok(())
4243    }
4244
4245    /// Build expr for `label_replace` function
4246    fn build_regexp_replace_label_expr(
4247        &self,
4248        other_input_exprs: &mut VecDeque<DfExpr>,
4249        input_schema: &DFSchemaRef,
4250    ) -> Result<Option<(DfExpr, String)>> {
4251        // label_replace(vector, dst_label, replacement, src_label, regex)
4252        let dst_label = match other_input_exprs.pop_front() {
4253            Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4254            other => UnexpectedPlanExprSnafu {
4255                desc: format!("expected dst_label string literal, but found {:?}", other),
4256            }
4257            .fail()?,
4258        };
4259
4260        // Validate the destination label name
4261        Self::validate_label_name(&dst_label)?;
4262        let replacement = match other_input_exprs.pop_front() {
4263            Some(DfExpr::Literal(ScalarValue::Utf8(Some(r)), _)) => r,
4264            other => UnexpectedPlanExprSnafu {
4265                desc: format!("expected replacement string literal, but found {:?}", other),
4266            }
4267            .fail()?,
4268        };
4269        let src_label = match other_input_exprs.pop_front() {
4270            Some(DfExpr::Literal(ScalarValue::Utf8(Some(s)), None)) => s,
4271            other => UnexpectedPlanExprSnafu {
4272                desc: format!("expected src_label string literal, but found {:?}", other),
4273            }
4274            .fail()?,
4275        };
4276
4277        let regex = match other_input_exprs.pop_front() {
4278            Some(DfExpr::Literal(ScalarValue::Utf8(Some(r)), None)) => r,
4279            other => UnexpectedPlanExprSnafu {
4280                desc: format!("expected regex string literal, but found {:?}", other),
4281            }
4282            .fail()?,
4283        };
4284
4285        // Like Prometheus, match the whole source value. A series whose source value matches gets
4286        // `dst_label` set to the expanded replacement; any other series is left unchanged.
4287        // doc: https://prometheus.io/docs/prometheus/latest/querying/functions/#label_replace
4288        let anchored = format!("^(?s:{regex})$");
4289        let compiled = regex::Regex::new(&anchored)
4290            .map_err(|_| InvalidRegularExpressionSnafu { regex }.build())?;
4291        let dst_exists = self.ctx.tag_columns.contains(&dst_label);
4292
4293        // A missing source label reads as the empty string for every series, so the result is
4294        // the same for all of them and is decided here.
4295        if !self.ctx.tag_columns.contains(&src_label) {
4296            let Some(captures) = compiled.captures("") else {
4297                return Ok(None);
4298            };
4299            let mut value = String::new();
4300            captures.expand(&replacement, &mut value);
4301            if value.is_empty() {
4302                // Setting a label to the empty string removes it, which is a no-op for a new
4303                // label.
4304                if !dst_exists {
4305                    return Ok(None);
4306                }
4307                return Ok(Some((
4308                    lit(ScalarValue::Utf8(None)).alias(&dst_label),
4309                    dst_label,
4310                )));
4311            }
4312            return Ok(Some((lit(value).alias(&dst_label), dst_label)));
4313        }
4314
4315        let src = Self::label_value_expr(&src_label, input_schema)?;
4316        let matched = DfExpr::ScalarFunction(ScalarFunction {
4317            func: datafusion_functions::regex::regexp_like(),
4318            args: vec![src.clone(), lit(anchored.clone())],
4319        });
4320        let replaced = Self::empty_label_to_null(DfExpr::ScalarFunction(ScalarFunction {
4321            func: datafusion_functions::regex::regexp_replace(),
4322            args: vec![src, lit(anchored), lit(replacement)],
4323        }));
4324        let unchanged = if dst_exists {
4325            DfExpr::Column(Column::from_name(&dst_label))
4326                .cast_to(&ArrowDataType::Utf8, input_schema)
4327                .context(DataFusionPlanningSnafu)?
4328        } else {
4329            lit(ScalarValue::Utf8(None))
4330        };
4331        let replace_expr = when(matched, replaced)
4332            .otherwise(unchanged)
4333            .context(DataFusionPlanningSnafu)?;
4334
4335        Ok(Some((replace_expr.alias(&dst_label), dst_label)))
4336    }
4337
4338    /// Build expr for `label_join` function
4339    fn build_concat_labels_expr(
4340        other_input_exprs: &mut VecDeque<DfExpr>,
4341        ctx: &PromPlannerContext,
4342        input_schema: &DFSchemaRef,
4343        query_engine_state: &QueryEngineState,
4344    ) -> Result<(DfExpr, String)> {
4345        // label_join(vector, dst_label, separator, src_label_1, src_label_2, ...)
4346
4347        let dst_label = match other_input_exprs.pop_front() {
4348            Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4349            other => UnexpectedPlanExprSnafu {
4350                desc: format!("expected dst_label string literal, but found {:?}", other),
4351            }
4352            .fail()?,
4353        };
4354        let separator = match other_input_exprs.pop_front() {
4355            Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4356            other => UnexpectedPlanExprSnafu {
4357                desc: format!("expected separator string literal, but found {:?}", other),
4358            }
4359            .fail()?,
4360        };
4361
4362        // Create a set of available columns (tag columns + field columns + time index column)
4363        let available_columns: HashSet<&str> = ctx
4364            .tag_columns
4365            .iter()
4366            .chain(ctx.field_columns.iter())
4367            .chain(ctx.time_index_column.as_ref())
4368            .map(|s| s.as_str())
4369            .collect();
4370
4371        let src_labels = other_input_exprs
4372            .iter()
4373            .map(|expr| {
4374                // `concat_ws` skips NULL arguments together with their separator, while an
4375                // absent label joins as the empty string.
4376                match expr {
4377                    DfExpr::Literal(ScalarValue::Utf8(Some(label)), None) => {
4378                        if label.is_empty() {
4379                            FunctionInvalidArgumentSnafu {
4380                                fn_name: "label_join",
4381                            }
4382                            .fail()
4383                        } else if available_columns.contains(label.as_str()) {
4384                            Self::label_value_expr(label, input_schema)
4385                        } else {
4386                            Ok(lit(""))
4387                        }
4388                    }
4389                    other => UnexpectedPlanExprSnafu {
4390                        desc: format!(
4391                            "expected source label string literal, but found {:?}",
4392                            other
4393                        ),
4394                    }
4395                    .fail(),
4396                }
4397            })
4398            .collect::<Result<Vec<_>>>()?;
4399        // Joining no labels yields the empty string, i.e. removes `dst_label`.
4400        if src_labels.is_empty() {
4401            return Ok((lit(ScalarValue::Utf8(None)).alias(&dst_label), dst_label));
4402        }
4403
4404        let session_state = query_engine_state.session_state();
4405        let func = session_state
4406            .scalar_functions()
4407            .get("concat_ws")
4408            .context(UnsupportedExprSnafu { name: "concat_ws" })?;
4409
4410        // concat_ws(separator, src_label_1, src_label_2, ...) as dst_label
4411        let mut args = Vec::with_capacity(1 + src_labels.len());
4412        args.push(DfExpr::Literal(ScalarValue::Utf8(Some(separator)), None));
4413        args.extend(src_labels);
4414
4415        Ok((
4416            Self::empty_label_to_null(DfExpr::ScalarFunction(ScalarFunction {
4417                func: func.clone(),
4418                args,
4419            }))
4420            .alias(&dst_label),
4421            dst_label,
4422        ))
4423    }
4424
4425    /// The value of `label` as a string, where NULL (the series has no such label) reads as the
4426    /// empty string, as in PromQL.
4427    fn label_value_expr(label: &str, input_schema: &DFSchemaRef) -> Result<DfExpr> {
4428        let value = DfExpr::Column(Column::from_name(label))
4429            .cast_to(&ArrowDataType::Utf8, input_schema)
4430            .context(DataFusionPlanningSnafu)?;
4431        Ok(DfExpr::ScalarFunction(ScalarFunction {
4432            func: coalesce(),
4433            args: vec![value, lit("")],
4434        }))
4435    }
4436
4437    /// An empty label value means the label is absent in PromQL. Label functions represent it
4438    /// as NULL, like a series that never had the label, so that both compare equal when label
4439    /// sets are matched and neither is reported as a label.
4440    fn empty_label_to_null(value: DfExpr) -> DfExpr {
4441        DfExpr::ScalarFunction(ScalarFunction {
4442            func: datafusion_functions::core::nullif(),
4443            args: vec![value, lit("")],
4444        })
4445    }
4446
4447    fn create_time_index_column_expr(&self) -> Result<DfExpr> {
4448        Ok(DfExpr::Column(Column::from_name(
4449            self.ctx
4450                .time_index_column
4451                .clone()
4452                .with_context(|| TimeIndexNotFoundSnafu { table: "unknown" })?,
4453        )))
4454    }
4455
4456    /// Builds the evaluation instant the window of the last planned range selector is folded for,
4457    /// as a `Timestamp(Millisecond)` expression.
4458    ///
4459    /// The timestamp payload of a folded window is shifted onto the evaluation timeline by the
4460    /// offset the window is folded with (`fold_offset`), while the time index column of a folded
4461    /// row keeps the evaluation timestamp of its step. Adding the offset back yields the
4462    /// evaluation instant on the payload timeline, which is where the regression of
4463    /// `predict_linear` is centered: neither a plain window (which may end before the step, and
4464    /// is additionally shifted by `offset` on the payload timeline) nor an `@`-anchored one
4465    /// (whose end is the anchor, while the payload is shifted by `at_offset`) ends at the step it
4466    /// is evaluated at.
4467    ///
4468    /// The sum is computed on the millisecond representation and cast back, so that the result
4469    /// keeps the `Timestamp(Millisecond)` type the range functions declare for it.
4470    fn create_range_eval_ts_expr(
4471        &self,
4472        fold_offset: Millisecond,
4473        input_schema: &DFSchemaRef,
4474    ) -> Result<DfExpr> {
4475        let eval_ts = self
4476            .create_time_index_column_expr()?
4477            .cast_to(
4478                &ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
4479                input_schema,
4480            )
4481            .context(DataFusionPlanningSnafu)?
4482            .cast_to(&ArrowDataType::Int64, input_schema)
4483            .context(DataFusionPlanningSnafu)?;
4484        DfExpr::BinaryExpr(BinaryExpr {
4485            left: Box::new(eval_ts),
4486            op: Operator::Plus,
4487            right: Box::new(lit(fold_offset)),
4488        })
4489        .cast_to(
4490            &ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
4491            input_schema,
4492        )
4493        .context(DataFusionPlanningSnafu)
4494    }
4495
4496    /// The name of `expr` without its last argument.
4497    ///
4498    /// A `predict_linear` call appends the evaluation instant of its window as a private last
4499    /// argument ([`Self::create_range_eval_ts_expr`]). Naming the output column after the call
4500    /// the user wrote, without that argument, keeps the injected expression out of the
4501    /// user-visible schema; the expression itself keeps every argument it needs.
4502    fn name_without_last_arg(expr: &DfExpr) -> String {
4503        if let DfExpr::ScalarFunction(ScalarFunction { func, args }) = expr
4504            && let Some((_, visible_args)) = args.split_last()
4505        {
4506            let visible = ScalarFunction {
4507                func: func.clone(),
4508                args: visible_args.to_vec(),
4509            };
4510            return DfExpr::ScalarFunction(visible).schema_name().to_string();
4511        }
4512
4513        expr.schema_name().to_string()
4514    }
4515
4516    fn create_tag_column_exprs(&self) -> Result<Vec<DfExpr>> {
4517        let mut result = Vec::with_capacity(self.ctx.tag_columns.len());
4518        for tag in &self.ctx.tag_columns {
4519            let expr = DfExpr::Column(Column::from_name(tag));
4520            result.push(expr);
4521        }
4522        Ok(result)
4523    }
4524
4525    fn create_field_column_exprs(&self) -> Result<Vec<DfExpr>> {
4526        let mut result = Vec::with_capacity(self.ctx.field_columns.len());
4527        for field in &self.ctx.field_columns {
4528            let expr = DfExpr::Column(Column::from_name(field));
4529            result.push(expr);
4530        }
4531        Ok(result)
4532    }
4533
4534    fn create_tag_and_time_index_column_sort_exprs(&self) -> Result<Vec<SortExpr>> {
4535        let mut result = self
4536            .ctx
4537            .tag_columns
4538            .iter()
4539            .map(|col| DfExpr::Column(Column::from_name(col)).sort(true, true))
4540            .collect::<Vec<_>>();
4541        result.push(self.create_time_index_column_expr()?.sort(true, true));
4542        Ok(result)
4543    }
4544
4545    fn create_field_columns_sort_exprs(&self, asc: bool) -> Vec<SortExpr> {
4546        self.ctx
4547            .field_columns
4548            .iter()
4549            .map(|col| DfExpr::Column(Column::from_name(col)).sort(asc, true))
4550            .collect::<Vec<_>>()
4551    }
4552
4553    fn create_sort_exprs_by_tags(
4554        func: &str,
4555        tags: Vec<DfExpr>,
4556        asc: bool,
4557    ) -> Result<Vec<SortExpr>> {
4558        ensure!(
4559            !tags.is_empty(),
4560            FunctionInvalidArgumentSnafu { fn_name: func }
4561        );
4562
4563        tags.iter()
4564            .map(|col| match col {
4565                DfExpr::Literal(ScalarValue::Utf8(Some(label)), _) => {
4566                    Ok(DfExpr::Column(Column::from_name(label)).sort(asc, false))
4567                }
4568                other => UnexpectedPlanExprSnafu {
4569                    desc: format!("expected label string literal, but found {:?}", other),
4570                }
4571                .fail(),
4572            })
4573            .collect::<Result<Vec<_>>>()
4574    }
4575
4576    fn create_empty_values_filter_expr(&self, preserve_any_value: bool) -> Result<DfExpr> {
4577        let mut exprs = Vec::with_capacity(self.ctx.field_columns.len());
4578        for value in &self.ctx.field_columns {
4579            let expr = DfExpr::Column(Column::from_name(value)).is_not_null();
4580            exprs.push(expr);
4581        }
4582
4583        // This error context should be computed lazily: the planner may set `ctx.table_name` to
4584        // `None` for derived expressions (e.g. after projecting the LHS of a vector-vector
4585        // comparison filter). Eagerly calling `table_ref()?` here can turn a valid plan into
4586        // a `TableNameNotFound` error even when predicate construction succeeds.
4587        let predicate = if preserve_any_value {
4588            disjunction(exprs)
4589        } else {
4590            conjunction(exprs)
4591        };
4592        predicate.with_context(|| ValueNotFoundSnafu {
4593            table: self
4594                .table_ref()
4595                .map(|t| t.to_quoted_string())
4596                .unwrap_or_else(|_| "unknown".to_string()),
4597        })
4598    }
4599
4600    /// Creates a set of DataFusion `DfExpr::AggregateFunction` expressions for each value column using the specified aggregate function.
4601    ///
4602    /// # Side Effects
4603    ///
4604    /// This method modifies the value columns in the context by replacing them with the new columns
4605    /// created by the aggregate function application.
4606    ///
4607    /// # Returns
4608    ///
4609    /// Returns a tuple of `(aggregate_expressions, previous_field_expressions)` where:
4610    /// - `aggregate_expressions`: Expressions that apply the aggregate function to the original fields
4611    /// - `previous_field_expressions`: Field expressions naming the pre-aggregation values. This is
4612    ///   non-empty only when the operation is `count_values`, which groups by the sample value and
4613    ///   projects it as the generated label, so these expressions are passed through the same
4614    ///   formatting as that label (`prom_float_to_string`).
4615    ///
4616    fn create_aggregate_exprs(
4617        &mut self,
4618        op: TokenType,
4619        param: &Option<Box<PromExpr>>,
4620        input_plan: &LogicalPlan,
4621    ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4622        let mixed_sample_columns =
4623            Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns)
4624                .map(|(float, histogram)| (float.to_string(), histogram.to_string()));
4625        let is_group_agg = op.id() == token::T_GROUP;
4626        if is_group_agg && mixed_sample_columns.is_none() {
4627            ensure!(
4628                self.ctx.field_columns.len() == 1,
4629                MultiFieldsNotSupportedSnafu {
4630                    operator: "group()"
4631                }
4632            );
4633        }
4634
4635        if let Some((float, histogram)) = mixed_sample_columns {
4636            return self.create_mixed_aggregate_exprs(op, param, &float, &histogram);
4637        }
4638
4639        if self.all_field_columns_are_native_histograms(input_plan.schema()) {
4640            return self.create_native_histogram_aggregate_exprs(op, input_plan);
4641        }
4642
4643        // perform aggregate operation to each value column
4644        let exprs = self
4645            .ctx
4646            .field_columns
4647            .iter()
4648            .map(|col| {
4649                Self::create_numeric_aggregate_expr(
4650                    op,
4651                    param,
4652                    DfExpr::Column(Column::from_name(col)),
4653                )
4654            })
4655            .collect::<Result<Vec<_>>>()?;
4656
4657        // if the aggregator is `count_values`, it must be grouped by current fields.
4658        //
4659        // The grouping key is the *formatted* sample value, i.e. the same expression that
4660        // produces the generated label below: PromQL groups by the value, and the label is
4661        // that value in Prometheus' textual form (`strconv.FormatFloat(value, 'f', -1, 64)`),
4662        // so grouping by the raw value would split samples that render to one label into
4663        // several groups, each emitting the same label set for one timestamp.
4664        let prev_field_exprs = if op.id() == token::T_COUNT_VALUES {
4665            ensure!(
4666                self.ctx.field_columns.len() == 1,
4667                UnsupportedExprSnafu {
4668                    name: "count_values on multi-value input"
4669                }
4670            );
4671
4672            self.ctx
4673                .field_columns
4674                .iter()
4675                .map(|col| {
4676                    let value = DfExpr::Column(Column::from_name(col));
4677                    // Normalize non `Float64` inputs the same way the label projection does,
4678                    // so both sides agree on the formatted value: `prom_float_to_string`
4679                    // formats exactly like Prometheus, while arrow's `Float64 -> Utf8` cast
4680                    // would render `1.0`.
4681                    let value = if Self::field_column_type(input_plan.schema(), col)
4682                        == Some(&ArrowDataType::Float64)
4683                    {
4684                        value
4685                    } else {
4686                        DfExpr::Cast(Cast::new(Box::new(value), ArrowDataType::Float64))
4687                    };
4688                    DfExpr::ScalarFunction(ScalarFunction {
4689                        func: Arc::new(PromqlFloatToString::scalar_udf()),
4690                        args: vec![value],
4691                    })
4692                })
4693                .collect()
4694        } else {
4695            vec![]
4696        };
4697
4698        // The alias is built from normalized columns so it matches the builtin's
4699        // qualified name.
4700        let exprs = normalize_cols(exprs, input_plan)
4701            .context(DataFusionPlanningSnafu)?
4702            .into_iter()
4703            .map(Self::alias_extremum_as_builtin)
4704            .collect::<Vec<_>>();
4705
4706        // update value column name according to the aggregators,
4707        self.ctx.field_columns = exprs
4708            .iter()
4709            .map(|expr| expr.schema_name().to_string())
4710            .collect();
4711
4712        Ok((exprs, prev_field_exprs))
4713    }
4714
4715    /// `prom_min`/`prom_max` keep their own names so SQL can use them next to `min`/`max`
4716    /// in one aggregate. PromQL output columns keep the builtin names.
4717    fn alias_extremum_as_builtin(expr: DfExpr) -> DfExpr {
4718        let DfExpr::AggregateFunction(aggr) = &expr else {
4719            return expr;
4720        };
4721        let builtin = match aggr.func.name() {
4722            Extremum::MIN_NAME => min_udaf(),
4723            Extremum::MAX_NAME => max_udaf(),
4724            _ => return expr,
4725        };
4726        let name = builtin
4727            .call(aggr.params.args.clone())
4728            .schema_name()
4729            .to_string();
4730        expr.alias(name)
4731    }
4732
4733    fn create_numeric_aggregate_expr(
4734        op: TokenType,
4735        param: &Option<Box<PromExpr>>,
4736        input: DfExpr,
4737    ) -> Result<DfExpr> {
4738        let expr = match op.id() {
4739            token::T_SUM => sum_udaf().call(vec![input]),
4740            token::T_QUANTILE => {
4741                let q = Self::get_param_as_literal_expr(
4742                    param.as_deref(),
4743                    Some(op),
4744                    Some(ArrowDataType::Float64),
4745                )?;
4746                quantile_udaf().call(vec![q, input])
4747            }
4748            token::T_AVG => avg_udaf().call(vec![input]),
4749            token::T_COUNT_VALUES | token::T_COUNT => count_udaf().call(vec![input]),
4750            token::T_MIN => Extremum::min_udaf().call(vec![input]),
4751            token::T_MAX => Extremum::max_udaf().call(vec![input]),
4752            // PromQL's `group()` aggregator produces 1 for each group.
4753            // Use `max(1.0)` (per-group) to match semantics and output type (Float64).
4754            token::T_GROUP => max_udaf().call(vec![lit(1_f64)]),
4755            token::T_STDDEV => stddev_pop_udaf().call(vec![input]),
4756            token::T_STDVAR => var_pop_udaf().call(vec![input]),
4757            token::T_TOPK | token::T_BOTTOMK => {
4758                return UnsupportedExprSnafu {
4759                    name: format!("{op:?}"),
4760                }
4761                .fail();
4762            }
4763            _ => return UnexpectedTokenSnafu { token: op }.fail(),
4764        };
4765        Ok(expr)
4766    }
4767
4768    fn create_mixed_aggregate_exprs(
4769        &mut self,
4770        op: TokenType,
4771        param: &Option<Box<PromExpr>>,
4772        float_column: &str,
4773        histogram_column: &str,
4774    ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4775        let float_input = DfExpr::Column(Column::from_name(float_column));
4776        let histogram_input = DfExpr::Column(Column::from_name(histogram_column));
4777        let float_count = count_udaf().call(vec![float_input.clone()]);
4778        let histogram_count = count_udaf().call(vec![histogram_input.clone()]);
4779        let mixed_sample_value = || {
4780            DfExpr::ScalarFunction(ScalarFunction {
4781                func: coalesce(),
4782                args: vec![
4783                    DfExpr::ScalarFunction(ScalarFunction {
4784                        func: Arc::new(PromqlFloatToString::scalar_udf()),
4785                        args: vec![float_input.clone()],
4786                    }),
4787                    DfExpr::ScalarFunction(ScalarFunction {
4788                        func: Arc::new(NativeHistogramToString::scalar_udf()),
4789                        args: vec![histogram_input.clone()],
4790                    }),
4791                ],
4792            })
4793        };
4794
4795        let (exprs, prev_field_exprs, field_columns) = match op.id() {
4796            token::T_SUM | token::T_AVG => (
4797                vec![
4798                    Self::create_numeric_aggregate_expr(op, param, float_input)?
4799                        .alias(float_column),
4800                    self.create_native_histogram_aggregate_expr(op, histogram_column)?,
4801                    float_count.alias(Self::mixed_sample_count_name(float_column)),
4802                    histogram_count.alias(Self::mixed_sample_count_name(histogram_column)),
4803                ],
4804                vec![],
4805                vec![float_column.to_string(), histogram_column.to_string()],
4806            ),
4807            token::T_COUNT => {
4808                let present = when(
4809                    float_input
4810                        .clone()
4811                        .is_not_null()
4812                        .or(histogram_input.clone().is_not_null()),
4813                    lit(1_i64),
4814                )
4815                .otherwise(lit(ScalarValue::Int64(None)))
4816                .context(DataFusionPlanningSnafu)?;
4817                (
4818                    vec![count_udaf().call(vec![present]).alias(float_column)],
4819                    vec![],
4820                    vec![float_column.to_string()],
4821                )
4822            }
4823            token::T_GROUP => (
4824                vec![max_udaf().call(vec![lit(1_f64)]).alias(float_column)],
4825                vec![],
4826                vec![float_column.to_string()],
4827            ),
4828            token::T_COUNT_VALUES => {
4829                let value = mixed_sample_value();
4830                (
4831                    vec![count_udaf().call(vec![value.clone()]).alias(float_column)],
4832                    vec![value],
4833                    vec![float_column.to_string()],
4834                )
4835            }
4836            token::T_MIN | token::T_MAX | token::T_STDDEV | token::T_STDVAR | token::T_QUANTILE => {
4837                (
4838                    vec![
4839                        Self::create_numeric_aggregate_expr(op, param, float_input)?
4840                            .alias(float_column),
4841                        histogram_count.alias(Self::mixed_sample_count_name(histogram_column)),
4842                    ],
4843                    vec![],
4844                    vec![float_column.to_string()],
4845                )
4846            }
4847            token::T_TOPK | token::T_BOTTOMK => {
4848                return UnsupportedExprSnafu {
4849                    name: format!("{op:?}"),
4850                }
4851                .fail();
4852            }
4853            _ => return UnexpectedTokenSnafu { token: op }.fail(),
4854        };
4855
4856        self.ctx.field_columns = field_columns;
4857        Ok((exprs, prev_field_exprs))
4858    }
4859
4860    fn mixed_sample_count_column(column: &str) -> DfExpr {
4861        DfExpr::Column(Column::from_name(Self::mixed_sample_count_name(column)))
4862    }
4863
4864    fn mixed_sample_count_name(column: &str) -> String {
4865        format!("__promql_sample_count({column})")
4866    }
4867
4868    fn mixed_aggregate_filter_expr(
4869        &self,
4870        op: TokenType,
4871        float_column: &str,
4872        histogram_column: &str,
4873    ) -> Result<DfExpr> {
4874        let float_count = Self::mixed_sample_count_column(float_column);
4875        let histogram_count = Self::mixed_sample_count_column(histogram_column);
4876        let mixed = float_count
4877            .clone()
4878            .gt(lit(0_i64))
4879            .and(histogram_count.clone().gt(lit(0_i64)));
4880        let drop_mixed = DfExpr::ScalarFunction(ScalarFunction {
4881            func: Arc::new(NativeHistogramDrop::warning_bool_false_udf(
4882                format!(
4883                    "{op}: dropped aggregation result containing both float and native histogram samples"
4884                ),
4885                self.promql_annotations.clone(),
4886            )),
4887            args: vec![float_count, histogram_count],
4888        });
4889
4890        when(mixed, drop_mixed)
4891            .otherwise(lit(true))
4892            .context(DataFusionPlanningSnafu)
4893    }
4894
4895    fn mixed_ignored_histogram_filter_expr(
4896        &self,
4897        op: TokenType,
4898        histogram_column: &str,
4899    ) -> Result<DfExpr> {
4900        let histogram_count = Self::mixed_sample_count_column(histogram_column);
4901        let has_histograms = histogram_count.clone().gt(lit(0_i64));
4902        let record_info = DfExpr::ScalarFunction(ScalarFunction {
4903            func: Arc::new(NativeHistogramDrop::bool_true_udf(
4904                format!(
4905                    "{op}: dropped native histogram samples because this aggregation is not supported for native histograms"
4906                ),
4907                self.promql_annotations.clone(),
4908            )),
4909            args: vec![histogram_count],
4910        });
4911
4912        when(has_histograms, record_info)
4913            .otherwise(lit(true))
4914            .context(DataFusionPlanningSnafu)
4915    }
4916
4917    fn create_native_histogram_aggregate_expr(
4918        &self,
4919        op: TokenType,
4920        column: &str,
4921    ) -> Result<DfExpr> {
4922        let input = DfExpr::Column(Column::from_name(column));
4923        let expr = match op.id() {
4924            token::T_SUM => Arc::new(NativeHistogramAggSum::aggregate_udf_with_collector(
4925                self.promql_annotations.clone(),
4926            ))
4927            .call(vec![input])
4928            .alias(column),
4929            token::T_AVG => Arc::new(NativeHistogramAggAvg::aggregate_udf_with_collector(
4930                self.promql_annotations.clone(),
4931            ))
4932            .call(vec![input])
4933            .alias(column),
4934            token::T_COUNT_VALUES | token::T_COUNT => {
4935                count_udaf().call(vec![input]).alias(column)
4936            }
4937            token::T_GROUP => max_udaf().call(vec![lit(1_f64)]).alias(column),
4938            token::T_MIN
4939            | token::T_MAX
4940            | token::T_STDDEV
4941            | token::T_STDVAR
4942            | token::T_QUANTILE
4943            | token::T_TOPK
4944            | token::T_BOTTOMK => sum_udaf()
4945                .call(vec![DfExpr::ScalarFunction(ScalarFunction {
4946                    func: Arc::new(NativeHistogramDrop::float_null_udf(
4947                        format!(
4948                            "{op}: dropped native histogram samples because this aggregation is not supported for native histograms"
4949                        ),
4950                        self.promql_annotations.clone(),
4951                    )),
4952                    args: vec![input],
4953                })])
4954                .alias(column),
4955            _ => return UnexpectedTokenSnafu { token: op }.fail(),
4956        };
4957        Ok(expr)
4958    }
4959
4960    fn create_native_histogram_aggregate_exprs(
4961        &mut self,
4962        op: TokenType,
4963        input_plan: &LogicalPlan,
4964    ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4965        let prev_field_exprs = if op.id() == token::T_COUNT_VALUES {
4966            ensure!(
4967                self.ctx.field_columns.len() == 1,
4968                UnsupportedExprSnafu {
4969                    name: "count_values on multi-value input"
4970                }
4971            );
4972            self.ctx
4973                .field_columns
4974                .iter()
4975                .map(|col| {
4976                    DfExpr::ScalarFunction(ScalarFunction {
4977                        func: Arc::new(NativeHistogramToString::scalar_udf()),
4978                        args: vec![DfExpr::Column(Column::from_name(col))],
4979                    })
4980                })
4981                .collect::<Vec<_>>()
4982        } else {
4983            vec![]
4984        };
4985
4986        let exprs = self
4987            .ctx
4988            .field_columns
4989            .iter()
4990            .map(|col| self.create_native_histogram_aggregate_expr(op, col))
4991            .collect::<Result<Vec<_>>>()?;
4992
4993        let normalized_exprs =
4994            normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
4995        self.ctx.field_columns = normalized_exprs
4996            .into_iter()
4997            .map(|expr| expr.schema_name().to_string())
4998            .collect();
4999
5000        Ok((exprs, prev_field_exprs))
5001    }
5002
5003    fn get_param_value_as_str(op: TokenType, param: &Option<Box<PromExpr>>) -> Result<&str> {
5004        let param = param
5005            .as_deref()
5006            .with_context(|| FunctionInvalidArgumentSnafu {
5007                fn_name: op.to_string(),
5008            })?;
5009        let PromExpr::StringLiteral(StringLiteral { val }) = param else {
5010            return FunctionInvalidArgumentSnafu {
5011                fn_name: op.to_string(),
5012            }
5013            .fail();
5014        };
5015
5016        Ok(val)
5017    }
5018
5019    fn get_param_as_literal_expr(
5020        param: Option<&PromExpr>,
5021        op: Option<TokenType>,
5022        expected_type: Option<ArrowDataType>,
5023    ) -> Result<DfExpr> {
5024        let prom_param = param.with_context(|| {
5025            if let Some(op) = op {
5026                FunctionInvalidArgumentSnafu {
5027                    fn_name: op.to_string(),
5028                }
5029            } else {
5030                FunctionInvalidArgumentSnafu {
5031                    fn_name: "unknown".to_string(),
5032                }
5033            }
5034        })?;
5035
5036        let expr = Self::try_build_literal_expr(prom_param).with_context(|| {
5037            if let Some(op) = op {
5038                FunctionInvalidArgumentSnafu {
5039                    fn_name: op.to_string(),
5040                }
5041            } else {
5042                FunctionInvalidArgumentSnafu {
5043                    fn_name: "unknown".to_string(),
5044                }
5045            }
5046        })?;
5047
5048        // check if the type is expected
5049        if let Some(expected_type) = expected_type {
5050            // literal should not have reference to column
5051            let expr_type = expr
5052                .get_type(&DFSchema::empty())
5053                .context(DataFusionPlanningSnafu)?;
5054            if expected_type != expr_type {
5055                return FunctionInvalidArgumentSnafu {
5056                    fn_name: format!("expected {expected_type:?}, but found {expr_type:?}"),
5057                }
5058                .fail();
5059            }
5060        }
5061
5062        Ok(expr)
5063    }
5064
5065    /// Create [DfExpr::WindowFunction] expr for each value column with given window function.
5066    ///
5067    fn create_window_exprs(
5068        &mut self,
5069        op: TokenType,
5070        group_exprs: Vec<DfExpr>,
5071        input_plan: &LogicalPlan,
5072    ) -> Result<Vec<DfExpr>> {
5073        ensure!(
5074            self.ctx.field_columns.len() == 1,
5075            UnsupportedExprSnafu {
5076                name: "topk or bottomk on multi-value input"
5077            }
5078        );
5079
5080        assert!(matches!(op.id(), token::T_TOPK | token::T_BOTTOMK));
5081
5082        let asc = matches!(op.id(), token::T_BOTTOMK);
5083
5084        let tag_sort_exprs = self
5085            .create_tag_column_exprs()?
5086            .into_iter()
5087            .map(|expr| expr.sort(asc, true));
5088
5089        // perform window operation to each value column
5090        let exprs: Vec<DfExpr> = self
5091            .ctx
5092            .field_columns
5093            .iter()
5094            .map(|col| {
5095                let mut sort_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 2);
5096                // PromQL ranks NaN last for both `topk` and `bottomk`, while DataFusion
5097                // orders NaN above every number.
5098                sort_exprs.push(
5099                    datafusion_functions::math::expr_fn::isnan(DfExpr::Column(Column::from_name(
5100                        col,
5101                    )))
5102                    .sort(true, true),
5103                );
5104                // Order by value in the specific order
5105                sort_exprs.push(DfExpr::Column(Column::from_name(col)).sort(asc, true));
5106                // Then tags if the values are equal,
5107                // Try to ensure the relative stability of the output results.
5108                sort_exprs.extend(tag_sort_exprs.clone());
5109
5110                DfExpr::WindowFunction(Box::new(WindowFunction {
5111                    fun: WindowFunctionDefinition::WindowUDF(Arc::new(RowNumber::new().into())),
5112                    params: WindowFunctionParams {
5113                        args: vec![],
5114                        partition_by: group_exprs.clone(),
5115                        order_by: sort_exprs,
5116                        window_frame: WindowFrame::new(Some(true)),
5117                        null_treatment: None,
5118                        distinct: false,
5119                        filter: None,
5120                    },
5121                }))
5122            })
5123            .collect();
5124
5125        let normalized_exprs =
5126            normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
5127        Ok(normalized_exprs)
5128    }
5129
5130    /// Try to build a DataFusion Literal Expression from PromQL Expr, return
5131    /// `None` if the input is not a literal expression.
5132    fn try_build_literal_expr(expr: &PromExpr) -> Option<DfExpr> {
5133        match expr {
5134            PromExpr::NumberLiteral(NumberLiteral { val }) => Some(val.lit()),
5135            PromExpr::StringLiteral(StringLiteral { val }) => Some(val.lit()),
5136            PromExpr::VectorSelector(_)
5137            | PromExpr::MatrixSelector(_)
5138            | PromExpr::Extension(_)
5139            | PromExpr::Aggregate(_)
5140            | PromExpr::Subquery(_) => None,
5141            PromExpr::Call(Call { func, .. }) => {
5142                if func.name == SPECIAL_TIME_FUNCTION {
5143                    // For time() function, don't treat it as a literal
5144                    // Let it be handled as a regular function call
5145                    None
5146                } else {
5147                    None
5148                }
5149            }
5150            PromExpr::Paren(ParenExpr { expr }) => Self::try_build_literal_expr(expr),
5151            PromExpr::Unary(UnaryExpr { expr, .. }) => Some(DfExpr::Negative(Box::new(
5152                Self::try_build_literal_expr(expr)?,
5153            ))),
5154            PromExpr::Binary(PromBinaryExpr {
5155                lhs,
5156                rhs,
5157                op,
5158                modifier,
5159            }) => {
5160                let lhs = Self::try_build_literal_expr(lhs)?;
5161                let rhs = Self::try_build_literal_expr(rhs)?;
5162                let is_comparison_op = Self::is_token_a_comparison_op(*op);
5163                let expr_builder = Self::prom_token_to_binary_expr_builder(*op).ok()?;
5164                let expr = expr_builder(lhs, rhs).ok()?;
5165
5166                let should_return_bool = if let Some(m) = modifier {
5167                    m.return_bool
5168                } else {
5169                    false
5170                };
5171                if is_comparison_op && should_return_bool {
5172                    Some(DfExpr::Cast(Cast::new(
5173                        Box::new(expr),
5174                        ArrowDataType::Float64,
5175                    )))
5176                } else {
5177                    Some(expr)
5178                }
5179            }
5180        }
5181    }
5182
5183    fn try_build_special_time_expr_with_context(&self, expr: &PromExpr) -> Option<DfExpr> {
5184        match expr {
5185            PromExpr::Call(Call { func, .. }) => {
5186                if func.name == SPECIAL_TIME_FUNCTION
5187                    && let Some(time_index_col) = self.ctx.time_index_column.as_ref()
5188                {
5189                    Some(build_special_time_expr(time_index_col))
5190                } else {
5191                    None
5192                }
5193            }
5194            _ => None,
5195        }
5196    }
5197
5198    fn native_histogram_binary_expr(
5199        token: TokenType,
5200        lhs: DfExpr,
5201        lhs_is_histogram: bool,
5202        rhs: DfExpr,
5203        rhs_is_histogram: bool,
5204        filter_context: bool,
5205        promql_annotations: Option<PromqlAnnotationCollector>,
5206    ) -> Result<Option<DfExpr>> {
5207        if !lhs_is_histogram && !rhs_is_histogram {
5208            return Ok(None);
5209        }
5210
5211        let scalar_fn = |func: ScalarUdfDef, args| {
5212            DfExpr::ScalarFunction(ScalarFunction {
5213                func: Arc::new(func),
5214                args,
5215            })
5216        };
5217        let invalid_expr = || {
5218            let message = format!(
5219                "{}: dropped native histogram samples because this binary operation is not supported for native histograms",
5220                token
5221            );
5222            let func = if filter_context {
5223                NativeHistogramDrop::bool_false_udf(message, promql_annotations.clone())
5224            } else {
5225                NativeHistogramDrop::float_null_udf(message, promql_annotations.clone())
5226            };
5227            let args = vec![lhs.clone(), rhs.clone()];
5228            scalar_fn(func, args)
5229        };
5230
5231        let expr = match (token.id(), lhs_is_histogram, rhs_is_histogram) {
5232            (token::T_ADD, true, true) => scalar_fn(
5233                NativeHistogramAdd::scalar_udf_with_collector(promql_annotations.clone()),
5234                vec![lhs, rhs],
5235            ),
5236            (token::T_SUB, true, true) => scalar_fn(
5237                NativeHistogramSub::scalar_udf_with_collector(promql_annotations.clone()),
5238                vec![lhs, rhs],
5239            ),
5240            (token::T_MUL, true, false) => {
5241                scalar_fn(NativeHistogramMulScalar::scalar_udf(), vec![lhs, rhs])
5242            }
5243            (token::T_MUL, false, true) => {
5244                scalar_fn(NativeHistogramScalarMul::scalar_udf(), vec![lhs, rhs])
5245            }
5246            (token::T_DIV, true, false) => {
5247                scalar_fn(NativeHistogramDivScalar::scalar_udf(), vec![lhs, rhs])
5248            }
5249            (token::T_EQLC, true, true) => {
5250                scalar_fn(NativeHistogramEq::scalar_udf(), vec![lhs, rhs])
5251            }
5252            (token::T_NEQ, true, true) => {
5253                scalar_fn(NativeHistogramNotEq::scalar_udf(), vec![lhs, rhs])
5254            }
5255            _ => invalid_expr(),
5256        };
5257
5258        Ok(Some(expr))
5259    }
5260
5261    /// Return a lambda to build binary expression from token.
5262    /// Because some binary operator are function in DataFusion like `atan2` or `^`.
5263    #[allow(clippy::type_complexity)]
5264    fn prom_token_to_binary_expr_builder(
5265        token: TokenType,
5266    ) -> Result<Box<dyn Fn(DfExpr, DfExpr) -> Result<DfExpr>>> {
5267        let cast_float = |expr| {
5268            if matches!(
5269                &expr,
5270                DfExpr::Cast(Cast { field, .. }) if field.data_type() == &ArrowDataType::Float64
5271            ) || matches!(&expr, DfExpr::Literal(ScalarValue::Float64(_), _))
5272            {
5273                expr
5274            } else {
5275                DfExpr::Cast(Cast::new(Box::new(expr), ArrowDataType::Float64))
5276            }
5277        };
5278        match token.id() {
5279            token::T_ADD => Ok(Box::new(move |lhs, rhs| {
5280                Ok(cast_float(lhs) + cast_float(rhs))
5281            })),
5282            token::T_SUB => Ok(Box::new(move |lhs, rhs| {
5283                Ok(cast_float(lhs) - cast_float(rhs))
5284            })),
5285            token::T_MUL => Ok(Box::new(move |lhs, rhs| {
5286                Ok(cast_float(lhs) * cast_float(rhs))
5287            })),
5288            token::T_DIV => Ok(Box::new(move |lhs, rhs| {
5289                Ok(cast_float(lhs) / cast_float(rhs))
5290            })),
5291            token::T_MOD => Ok(Box::new(move |lhs: DfExpr, rhs| {
5292                Ok(cast_float(lhs) % cast_float(rhs))
5293            })),
5294            token::T_EQLC => Ok(Box::new(|lhs, rhs| {
5295                Ok(Self::ieee_comparison(lhs, Operator::Eq, rhs))
5296            })),
5297            token::T_NEQ => Ok(Box::new(|lhs, rhs| {
5298                Ok(Self::ieee_comparison(lhs, Operator::NotEq, rhs))
5299            })),
5300            token::T_GTR => Ok(Box::new(|lhs, rhs| {
5301                Ok(Self::ieee_comparison(lhs, Operator::Gt, rhs))
5302            })),
5303            token::T_LSS => Ok(Box::new(|lhs, rhs| {
5304                Ok(Self::ieee_comparison(lhs, Operator::Lt, rhs))
5305            })),
5306            token::T_GTE => Ok(Box::new(|lhs, rhs| {
5307                Ok(Self::ieee_comparison(lhs, Operator::GtEq, rhs))
5308            })),
5309            token::T_LTE => Ok(Box::new(|lhs, rhs| {
5310                Ok(Self::ieee_comparison(lhs, Operator::LtEq, rhs))
5311            })),
5312            token::T_POW => Ok(Box::new(move |lhs, rhs| {
5313                Ok(DfExpr::ScalarFunction(ScalarFunction {
5314                    func: datafusion_functions::math::power(),
5315                    args: vec![cast_float(lhs), cast_float(rhs)],
5316                }))
5317            })),
5318            token::T_ATAN2 => Ok(Box::new(move |lhs, rhs| {
5319                Ok(DfExpr::ScalarFunction(ScalarFunction {
5320                    func: datafusion_functions::math::atan2(),
5321                    args: vec![cast_float(lhs), cast_float(rhs)],
5322                }))
5323            })),
5324            _ => UnexpectedTokenSnafu { token }.fail(),
5325        }
5326    }
5327
5328    /// Compares two expressions with IEEE 754 semantics, as PromQL does; see [`IeeeComparison`].
5329    fn ieee_comparison(lhs: DfExpr, op: Operator, rhs: DfExpr) -> DfExpr {
5330        DfExpr::ScalarFunction(ScalarFunction {
5331            func: Arc::new(IeeeComparison::scalar_udf(op)),
5332            args: vec![lhs, rhs],
5333        })
5334    }
5335
5336    /// Check if the given op is a [comparison operator](https://prometheus.io/docs/prometheus/latest/querying/operators/#comparison-binary-operators).
5337    fn is_token_a_comparison_op(token: TokenType) -> bool {
5338        matches!(
5339            token.id(),
5340            token::T_EQLC
5341                | token::T_NEQ
5342                | token::T_GTR
5343                | token::T_LSS
5344                | token::T_GTE
5345                | token::T_LTE
5346        )
5347    }
5348
5349    /// Check if the given op is a set operator (UNION, INTERSECT and EXCEPT in SQL).
5350    fn is_token_a_set_op(token: TokenType) -> bool {
5351        matches!(
5352            token.id(),
5353            token::T_LAND // INTERSECT
5354                | token::T_LOR // UNION
5355                | token::T_LUNLESS // EXCEPT
5356        )
5357    }
5358
5359    fn align_binary_field_columns<'a>(
5360        left_schema: &DFSchemaRef,
5361        right_schema: &DFSchemaRef,
5362        left_field_columns: &'a [String],
5363        right_field_columns: &'a [String],
5364        op: TokenType,
5365        left_is_scalar: bool,
5366        right_is_scalar: bool,
5367    ) -> (
5368        Vec<(String, Vec<BinaryFieldPair<'a>>)>,
5369        Vec<BinaryFieldPair<'a>>,
5370    ) {
5371        // Mixed vectors store mutually exclusive float and histogram samples in two columns.
5372        // Retain each valid sample combination and group expressions by their output lane.
5373        let left_alternative = Self::alternative_sample_columns(left_schema, left_field_columns);
5374        let right_alternative = Self::alternative_sample_columns(right_schema, right_field_columns);
5375        let alternative_alignment = match (left_alternative, right_alternative) {
5376            (Some(output_names), Some(_)) => Some((
5377                output_names,
5378                left_field_columns
5379                    .iter()
5380                    .flat_map(|left| right_field_columns.iter().map(move |right| (left, right)))
5381                    .collect::<Vec<_>>(),
5382            )),
5383            (Some(output_names), None) if right_field_columns.len() == 1 => Some((
5384                output_names,
5385                left_field_columns
5386                    .iter()
5387                    .map(|left| (left, &right_field_columns[0]))
5388                    .collect::<Vec<_>>(),
5389            )),
5390            (None, Some(output_names)) if left_field_columns.len() == 1 => Some((
5391                output_names,
5392                right_field_columns
5393                    .iter()
5394                    .map(|right| (&left_field_columns[0], right))
5395                    .collect::<Vec<_>>(),
5396            )),
5397            _ => None,
5398        };
5399        let mut invalid_pairs = Vec::new();
5400        if let Some(((float_output, histogram_output), field_pairs)) = alternative_alignment {
5401            let mut float_pairs = Vec::new();
5402            let mut histogram_pairs = Vec::new();
5403            for (left, right) in field_pairs {
5404                let left_is_histogram = Self::field_column_is_native_histogram(left_schema, left);
5405                let right_is_histogram =
5406                    Self::field_column_is_native_histogram(right_schema, right);
5407                match Self::binary_result_is_histogram(op, left_is_histogram, right_is_histogram) {
5408                    Some(false) => float_pairs.push((left, right)),
5409                    Some(true) => histogram_pairs.push((left, right)),
5410                    None => invalid_pairs.push((left, right)),
5411                }
5412            }
5413            if !float_pairs.is_empty() || !histogram_pairs.is_empty() {
5414                return (
5415                    [
5416                        (!float_pairs.is_empty()).then(|| (float_output.to_string(), float_pairs)),
5417                        (!histogram_pairs.is_empty())
5418                            .then(|| (histogram_output.to_string(), histogram_pairs)),
5419                    ]
5420                    .into_iter()
5421                    .flatten()
5422                    .collect(),
5423                    invalid_pairs,
5424                );
5425            }
5426        }
5427
5428        if left_is_scalar && !right_is_scalar && left_field_columns.len() == 1 {
5429            return (
5430                right_field_columns
5431                    .iter()
5432                    .map(|right| (right.clone(), vec![(&left_field_columns[0], right)]))
5433                    .collect(),
5434                invalid_pairs,
5435            );
5436        }
5437        if right_is_scalar && !left_is_scalar && right_field_columns.len() == 1 {
5438            return (
5439                left_field_columns
5440                    .iter()
5441                    .map(|left| (left.clone(), vec![(left, &right_field_columns[0])]))
5442                    .collect(),
5443                invalid_pairs,
5444            );
5445        }
5446
5447        (
5448            left_field_columns
5449                .iter()
5450                .zip(right_field_columns.iter())
5451                .map(|(left, right)| (left.clone(), vec![(left, right)]))
5452                .collect(),
5453            invalid_pairs,
5454        )
5455    }
5456
5457    fn binary_result_is_histogram(
5458        token: TokenType,
5459        lhs_is_histogram: bool,
5460        rhs_is_histogram: bool,
5461    ) -> Option<bool> {
5462        match (token.id(), lhs_is_histogram, rhs_is_histogram) {
5463            (_, false, false) => Some(false),
5464            (token::T_ADD | token::T_SUB, true, true)
5465            | (token::T_MUL, true, false)
5466            | (token::T_MUL, false, true)
5467            | (token::T_DIV, true, false) => Some(true),
5468            (token::T_EQLC | token::T_NEQ, true, true) => Some(false),
5469            _ => None,
5470        }
5471    }
5472
5473    fn plan_has_tsid_column(plan: &LogicalPlan) -> bool {
5474        plan.schema()
5475            .fields()
5476            .iter()
5477            .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
5478    }
5479
5480    fn is_empty_metric(plan: &LogicalPlan) -> bool {
5481        matches!(plan, LogicalPlan::Extension(Extension { node }) if node.as_any().is::<EmptyMetric>())
5482    }
5483
5484    fn native_histogram_arrow_type() -> ArrowDataType {
5485        native_histogram_value_type().as_arrow_type()
5486    }
5487
5488    fn field_column_type<'a>(
5489        schema: &'a DFSchemaRef,
5490        field_column: &str,
5491    ) -> Option<&'a ArrowDataType> {
5492        schema
5493            .index_of_column_by_name(None, field_column)
5494            .map(|idx| schema.field(idx).data_type())
5495    }
5496
5497    fn field_column_is_native_histogram(schema: &DFSchemaRef, field_column: &str) -> bool {
5498        Self::field_column_type(schema, field_column)
5499            .is_some_and(|data_type| data_type == &Self::native_histogram_arrow_type())
5500    }
5501
5502    fn field_columns_contain_native_histogram(
5503        schema: &DFSchemaRef,
5504        field_columns: &[String],
5505    ) -> bool {
5506        field_columns
5507            .iter()
5508            .any(|field| Self::field_column_is_native_histogram(schema, field))
5509    }
5510
5511    fn field_column_is_float_range(schema: &DFSchemaRef, field_column: &str) -> bool {
5512        Self::field_column_type(schema, field_column).is_some_and(|data_type| {
5513            matches!(
5514                data_type,
5515                ArrowDataType::Dictionary(key_type, value_type)
5516                    if key_type.as_ref() == &ArrowDataType::Int64
5517                        && value_type.as_ref() == &ArrowDataType::Float64
5518            )
5519        })
5520    }
5521
5522    fn field_columns_are_alternative_samples(
5523        schema: &DFSchemaRef,
5524        field_columns: &[String],
5525    ) -> bool {
5526        Self::alternative_sample_columns(schema, field_columns).is_some()
5527    }
5528
5529    fn alternative_sample_columns<'a>(
5530        schema: &DFSchemaRef,
5531        field_columns: &'a [String],
5532    ) -> Option<(&'a str, &'a str)> {
5533        if field_columns.len() != 2 {
5534            return None;
5535        }
5536
5537        let canonical_float = field_columns.iter().find(|field| {
5538            field.as_str() == greptime_value()
5539                && (Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64)
5540                    || Self::field_column_is_float_range(schema, field))
5541        });
5542        let canonical_histogram = field_columns.iter().find(|field| {
5543            field.as_str() == greptime_native_histogram()
5544                && (Self::field_column_is_native_histogram(schema, field)
5545                    || Self::field_column_is_native_histogram_range(schema, field))
5546        });
5547        if let (Some(float), Some(histogram)) = (canonical_float, canonical_histogram) {
5548            return Some((float, histogram));
5549        }
5550
5551        let float = field_columns.iter().find(|field| {
5552            field.starts_with(OR_FLOAT_FIELD_PREFIX)
5553                && (Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64)
5554                    || Self::field_column_is_float_range(schema, field))
5555        })?;
5556        let histogram = field_columns.iter().find(|field| {
5557            field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)
5558                && (Self::field_column_is_native_histogram(schema, field)
5559                    || Self::field_column_is_native_histogram_range(schema, field))
5560        })?;
5561        Some((float, histogram))
5562    }
5563
5564    fn alternative_sample_range_columns<'a>(
5565        schema: &DFSchemaRef,
5566        field_columns: &'a [String],
5567    ) -> Option<(&'a str, &'a str)> {
5568        Self::alternative_sample_columns(schema, field_columns).filter(|(float, histogram)| {
5569            Self::field_column_is_float_range(schema, float)
5570                && Self::field_column_is_native_histogram_range(schema, histogram)
5571        })
5572    }
5573
5574    fn field_column_is_native_histogram_range(schema: &DFSchemaRef, field_column: &str) -> bool {
5575        Self::field_column_type(schema, field_column).is_some_and(|data_type| {
5576            matches!(
5577                data_type,
5578                ArrowDataType::Dictionary(key_type, value_type)
5579                    if key_type.as_ref() == &ArrowDataType::Int64
5580                        && value_type.as_ref() == &Self::native_histogram_arrow_type()
5581            )
5582        })
5583    }
5584
5585    fn all_field_columns_are_native_histograms(&self, schema: &DFSchemaRef) -> bool {
5586        !self.ctx.field_columns.is_empty()
5587            && self
5588                .ctx
5589                .field_columns
5590                .iter()
5591                .all(|field| Self::field_column_is_native_histogram(schema, field))
5592    }
5593
5594    fn all_field_columns_are_native_histogram_ranges(&self, schema: &DFSchemaRef) -> bool {
5595        !self.ctx.field_columns.is_empty()
5596            && self
5597                .ctx
5598                .field_columns
5599                .iter()
5600                .all(|field| Self::field_column_is_native_histogram_range(schema, field))
5601    }
5602
5603    fn optional_tsid_projection(
5604        schema: &DFSchemaRef,
5605        table_ref: Option<&TableReference>,
5606        keep_tsid: bool,
5607    ) -> Option<DfExpr> {
5608        keep_tsid.then_some(()).and_then(|_| {
5609            schema
5610                .qualified_field_with_name(table_ref, DATA_SCHEMA_TSID_COLUMN_NAME)
5611                .ok()
5612                .map(|field| DfExpr::Column(field.into()))
5613        })
5614    }
5615
5616    fn binary_join_key_columns(
5617        &self,
5618        left_schema: &DFSchemaRef,
5619        right_schema: &DFSchemaRef,
5620        left_context: &PromPlannerContext,
5621        right_context: &PromPlannerContext,
5622        only_join_time_index: bool,
5623        modifier: &Option<BinModifier>,
5624    ) -> Result<(BTreeSet<String>, BTreeSet<String>, bool)> {
5625        let has_tsid = |schema: &DFSchemaRef| {
5626            schema
5627                .fields()
5628                .iter()
5629                .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
5630        };
5631        let use_tsid_join = !only_join_time_index
5632            && self.binary_modifier_preserves_tsid_join_key(left_context, right_context, modifier)
5633            && left_context.use_tsid
5634            && right_context.use_tsid
5635            && has_tsid(left_schema)
5636            && has_tsid(right_schema);
5637
5638        let (mut left_tag_columns, mut right_tag_columns) = if use_tsid_join {
5639            (
5640                BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]),
5641                BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]),
5642            )
5643        } else {
5644            if only_join_time_index {
5645                (BTreeSet::new(), BTreeSet::new())
5646            } else {
5647                (
5648                    left_context
5649                        .tag_columns
5650                        .iter()
5651                        .cloned()
5652                        .collect::<BTreeSet<_>>(),
5653                    right_context
5654                        .tag_columns
5655                        .iter()
5656                        .cloned()
5657                        .collect::<BTreeSet<_>>(),
5658                )
5659            }
5660        };
5661
5662        if !use_tsid_join
5663            && let Some(modifier) = modifier
5664            && let Some(matching) = &modifier.matching
5665        {
5666            match matching {
5667                LabelModifier::Include(on) => {
5668                    let mask = on.labels.iter().cloned().collect::<BTreeSet<_>>();
5669                    left_tag_columns = left_tag_columns.intersection(&mask).cloned().collect();
5670                    right_tag_columns = right_tag_columns.intersection(&mask).cloned().collect();
5671                }
5672                LabelModifier::Exclude(ignoring) => {
5673                    for label in &ignoring.labels {
5674                        let _ = left_tag_columns.remove(label);
5675                        let _ = right_tag_columns.remove(label);
5676                    }
5677                }
5678            }
5679        }
5680
5681        let force_empty_join =
5682            !use_tsid_join && !only_join_time_index && left_tag_columns != right_tag_columns;
5683        if force_empty_join {
5684            let common_tag_columns = left_tag_columns
5685                .intersection(&right_tag_columns)
5686                .cloned()
5687                .collect::<BTreeSet<_>>();
5688            left_tag_columns = common_tag_columns.clone();
5689            right_tag_columns = common_tag_columns;
5690        }
5691
5692        Ok((left_tag_columns, right_tag_columns, force_empty_join))
5693    }
5694
5695    /// Result labels of a vector-vector binary operation, following Prometheus `resultMetric`:
5696    /// `on(...)` keeps only the matching labels, `ignoring(...)` drops them, and a group modifier
5697    /// keeps the "many" side's labels plus the `group_x(...)` labels taken from the "one" side.
5698    ///
5699    /// The flag of each entry tells which operand the label is projected from. `None` means the
5700    /// operation keeps a whole operand tag set, which the default projection already does.
5701    fn binary_result_labels(
5702        left_context: &PromPlannerContext,
5703        right_context: &PromPlannerContext,
5704        modifier: &Option<BinModifier>,
5705    ) -> Option<Vec<(bool, String)>> {
5706        let modifier = modifier.as_ref()?;
5707        let (many_is_left, include) = match &modifier.card {
5708            VectorMatchCardinality::OneToOne => (true, None),
5709            VectorMatchCardinality::ManyToOne(labels) => (true, Some(labels)),
5710            VectorMatchCardinality::OneToMany(labels) => (false, Some(labels)),
5711            // Set operators keep their operands' labels and don't reach this path.
5712            VectorMatchCardinality::ManyToMany => return None,
5713        };
5714
5715        let Some(include) = include else {
5716            let matching = modifier.matching.as_ref()?;
5717            let reduced = |keep: bool, labels: &BTreeSet<&String>| {
5718                left_context
5719                    .tag_columns
5720                    .iter()
5721                    .filter(|tag| labels.contains(tag) == keep)
5722                    .map(|tag| (true, tag.clone()))
5723                    .collect()
5724            };
5725            return Some(match matching {
5726                LabelModifier::Include(on) => reduced(true, &on.labels.iter().collect()),
5727                LabelModifier::Exclude(ignoring) => {
5728                    reduced(false, &ignoring.labels.iter().collect())
5729                }
5730            });
5731        };
5732
5733        let (many_context, one_context) = if many_is_left {
5734            (left_context, right_context)
5735        } else {
5736            (right_context, left_context)
5737        };
5738        let include = include.labels.iter().collect::<BTreeSet<_>>();
5739        // An included label the "one" side doesn't carry is deleted from the result, so it is
5740        // dropped from the "many" side as well.
5741        let mut labels = many_context
5742            .tag_columns
5743            .iter()
5744            .filter(|tag| !include.contains(tag))
5745            .map(|tag| (many_is_left, tag.clone()))
5746            .collect::<Vec<_>>();
5747        labels.extend(
5748            include
5749                .into_iter()
5750                .filter(|label| one_context.tag_columns.contains(label))
5751                .map(|label| (!many_is_left, label.clone())),
5752        );
5753        Some(labels)
5754    }
5755
5756    /// Resolve [`Self::binary_result_labels`] against the join output.
5757    fn binary_result_label_projection(
5758        schema: &DFSchemaRef,
5759        left_table_ref: &TableReference,
5760        right_table_ref: &TableReference,
5761        left_context: &PromPlannerContext,
5762        right_context: &PromPlannerContext,
5763        labels: Vec<(bool, String)>,
5764    ) -> Result<BinaryResultLabels> {
5765        let mut exprs = Vec::with_capacity(labels.len());
5766        let mut names = Vec::with_capacity(labels.len());
5767        let mut aggregation_field_labels = Vec::new();
5768        let mut sources = HashSet::new();
5769        for (from_left, label) in labels {
5770            let (table_ref, context) = if from_left {
5771                (left_table_ref, left_context)
5772            } else {
5773                (right_table_ref, right_context)
5774            };
5775            let field = schema
5776                .qualified_field_with_name(Some(table_ref), &label)
5777                .context(DataFusionPlanningSnafu)?;
5778            exprs.push(DfExpr::Column(field.into()));
5779            if context.aggregation_field_labels.contains(&label) {
5780                aggregation_field_labels.push(label.clone());
5781            }
5782            sources.insert(from_left);
5783            names.push(label);
5784        }
5785
5786        // One operand contributing every one of its tags as the whole result label set keeps the
5787        // one-to-one correspondence between its `__tsid` and a result series: no other operand
5788        // value reaches the labels, and the matching gives each of its rows a single partner.
5789        let source_context = match sources.into_iter().collect::<Vec<_>>().as_slice() {
5790            [true] => Some((left_table_ref, left_context)),
5791            [false] => Some((right_table_ref, right_context)),
5792            _ => None,
5793        };
5794        let tsid = source_context
5795            .filter(|(_, context)| {
5796                context.use_tsid
5797                    && context.tag_columns.len() == names.len()
5798                    && context.tag_columns.iter().all(|tag| names.contains(tag))
5799            })
5800            .and_then(|(table_ref, _)| {
5801                Self::optional_tsid_projection(schema, Some(table_ref), true)
5802            });
5803
5804        Ok(BinaryResultLabels {
5805            exprs,
5806            names,
5807            aggregation_field_labels,
5808            tsid,
5809        })
5810    }
5811
5812    /// Whether the result of a binary operation can hold two series with the same labels, which
5813    /// only [`Self::binary_result_labels`] can introduce: one-to-one matching on a subset of the
5814    /// tags, or a group modifier that overwrites a label of the "many" side.
5815    fn binary_result_labels_may_repeat(
5816        left_context: &PromPlannerContext,
5817        right_context: &PromPlannerContext,
5818        modifier: &Option<BinModifier>,
5819    ) -> bool {
5820        let Some(modifier) = modifier else {
5821            return false;
5822        };
5823        match &modifier.card {
5824            VectorMatchCardinality::OneToOne => match &modifier.matching {
5825                None => false,
5826                Some(LabelModifier::Include(on)) => {
5827                    let on = on.labels.iter().collect::<BTreeSet<_>>();
5828                    !left_context.tag_columns.iter().all(|tag| on.contains(tag))
5829                }
5830                Some(LabelModifier::Exclude(ignoring)) => ignoring
5831                    .labels
5832                    .iter()
5833                    .any(|label| left_context.tag_columns.contains(label)),
5834            },
5835            VectorMatchCardinality::ManyToOne(include) => include
5836                .labels
5837                .iter()
5838                .any(|label| left_context.tag_columns.contains(label)),
5839            VectorMatchCardinality::OneToMany(include) => include
5840                .labels
5841                .iter()
5842                .any(|label| right_context.tag_columns.contains(label)),
5843            VectorMatchCardinality::ManyToMany => false,
5844        }
5845    }
5846
5847    /// Wrap `plan` in a check that fails the query when a match group holds more than one row at
5848    /// a timestamp. `group_exprs` are the label columns of the group, resolved against `plan`.
5849    fn assert_unique_match_group(
5850        plan: LogicalPlan,
5851        group_exprs: Vec<DfExpr>,
5852        group_labels: Vec<String>,
5853        time_index_expr: DfExpr,
5854        violation: MatchGroupViolation,
5855    ) -> Result<LogicalPlan> {
5856        let mut partition_by = group_exprs.clone();
5857        partition_by.push(time_index_expr);
5858        // A label may carry the generated name, and the count column has to stay unambiguous
5859        // against every column the operand already has.
5860        let occupied_column_names = plan
5861            .schema()
5862            .fields()
5863            .iter()
5864            .map(|field| field.name().as_str())
5865            .collect::<HashSet<_>>();
5866        let mut next_suffix = 0;
5867        let count_column = loop {
5868            let name = match next_suffix {
5869                0 => MATCH_GROUP_COUNT_COLUMN.to_string(),
5870                suffix => format!("{MATCH_GROUP_COUNT_COLUMN}_{suffix}"),
5871            };
5872            next_suffix += 1;
5873            if !occupied_column_names.contains(name.as_str()) {
5874                break name;
5875            }
5876        };
5877        let count = DfExpr::WindowFunction(Box::new(WindowFunction {
5878            fun: WindowFunctionDefinition::AggregateUDF(count_udaf()),
5879            params: WindowFunctionParams {
5880                args: vec![lit(1i64)],
5881                partition_by,
5882                order_by: vec![],
5883                window_frame: WindowFrame::new(None),
5884                null_treatment: None,
5885                distinct: false,
5886                filter: None,
5887            },
5888        }))
5889        .alias(count_column.as_str());
5890
5891        let output_exprs = plan
5892            .schema()
5893            .iter()
5894            .map(|(qualifier, field)| DfExpr::Column(Column::new(qualifier.cloned(), field.name())))
5895            .collect::<Vec<_>>();
5896        let assert_expr = DfExpr::ScalarFunction(ScalarFunction {
5897            func: Arc::new(UniqueMatchGroup::scalar_udf(group_labels, violation)),
5898            args: std::iter::once(ident(count_column.as_str()))
5899                .chain(group_exprs)
5900                .collect(),
5901        });
5902
5903        LogicalPlanBuilder::from(plan)
5904            .window(vec![count])
5905            .context(DataFusionPlanningSnafu)?
5906            .filter(assert_expr)
5907            .context(DataFusionPlanningSnafu)?
5908            .project(output_exprs)
5909            .context(DataFusionPlanningSnafu)?
5910            .build()
5911            .context(DataFusionPlanningSnafu)
5912    }
5913
5914    fn binary_modifier_preserves_tsid_join_key(
5915        &self,
5916        left_context: &PromPlannerContext,
5917        right_context: &PromPlannerContext,
5918        modifier: &Option<BinModifier>,
5919    ) -> bool {
5920        let Some(modifier) = modifier else {
5921            return true;
5922        };
5923
5924        if !matches!(modifier.card, VectorMatchCardinality::OneToOne) {
5925            return false;
5926        }
5927
5928        match &modifier.matching {
5929            None => true,
5930            Some(LabelModifier::Exclude(ignoring)) => ignoring.labels.iter().all(|label| {
5931                !left_context.tag_columns.contains(label)
5932                    && !right_context.tag_columns.contains(label)
5933            }),
5934            Some(LabelModifier::Include(on)) => {
5935                let on_labels = on.labels.iter().cloned().collect::<BTreeSet<_>>();
5936                let left_labels = left_context
5937                    .tag_columns
5938                    .iter()
5939                    .cloned()
5940                    .collect::<BTreeSet<_>>();
5941                let right_labels = right_context
5942                    .tag_columns
5943                    .iter()
5944                    .cloned()
5945                    .collect::<BTreeSet<_>>();
5946
5947                on_labels == left_labels && on_labels == right_labels
5948            }
5949        }
5950    }
5951
5952    /// Build a inner join on time index column and tag columns to concat two logical plans.
5953    /// When `only_join_time_index == true` we only join on the time index, because these two plan may not have the same tag columns
5954    #[allow(clippy::too_many_arguments)]
5955    fn join_on_non_field_columns(
5956        &self,
5957        left: LogicalPlan,
5958        right: LogicalPlan,
5959        left_table_ref: TableReference,
5960        right_table_ref: TableReference,
5961        left_time_index_column: Option<String>,
5962        right_time_index_column: Option<String>,
5963        only_join_time_index: bool,
5964        modifier: &Option<BinModifier>,
5965        left_context: &PromPlannerContext,
5966        right_context: &PromPlannerContext,
5967    ) -> Result<LogicalPlan> {
5968        let (mut left_tag_columns, mut right_tag_columns, mut force_empty_join) = self
5969            .binary_join_key_columns(
5970                left.schema(),
5971                right.schema(),
5972                left_context,
5973                right_context,
5974                only_join_time_index,
5975                modifier,
5976            )?;
5977        let use_tsid_join = !only_join_time_index
5978            && !force_empty_join
5979            && left_tag_columns == BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()])
5980            && right_tag_columns == BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]);
5981        let (left, right, left_matched_tags, right_matched_tags) = if !only_join_time_index
5982            && !use_tsid_join
5983            && Self::only_temporality_match_label_mismatches(left_context, right_context, modifier)
5984        {
5985            let mut aligned_left_context = left_context.clone();
5986            let mut aligned_right_context = right_context.clone();
5987            let (left, right, _) = Self::align_temporality_match_column(
5988                left,
5989                right,
5990                &mut aligned_left_context,
5991                &mut aligned_right_context,
5992            )?;
5993            (left_tag_columns, right_tag_columns, force_empty_join) = self
5994                .binary_join_key_columns(
5995                    left.schema(),
5996                    right.schema(),
5997                    &aligned_left_context,
5998                    &aligned_right_context,
5999                    false,
6000                    modifier,
6001                )?;
6002            (
6003                left,
6004                right,
6005                aligned_left_context.tag_columns,
6006                aligned_right_context.tag_columns,
6007            )
6008        } else {
6009            (
6010                left,
6011                right,
6012                left_context.tag_columns.clone(),
6013                right_context.tag_columns.clone(),
6014            )
6015        };
6016
6017        // A join key that covers the whole tag set of a side already makes that side's match
6018        // groups unique, and so does a join on `__tsid`. Only the reduced keys need the check.
6019        let checked_matching = !only_join_time_index
6020            && !force_empty_join
6021            && !use_tsid_join
6022            && !matches!(
6023                modifier.as_ref().map(|modifier| &modifier.card),
6024                Some(VectorMatchCardinality::ManyToMany)
6025            );
6026        // The right operand is the "one" side of the matching, unless `group_right` swaps them.
6027        let one_side_is_left = matches!(
6028            modifier.as_ref().map(|modifier| &modifier.card),
6029            Some(VectorMatchCardinality::OneToMany(_))
6030        );
6031        let (left, right) = if !checked_matching {
6032            (left, right)
6033        } else if one_side_is_left {
6034            (
6035                Self::assert_unique_one_side(
6036                    left,
6037                    &left_tag_columns,
6038                    &left_matched_tags,
6039                    left_time_index_column.as_deref(),
6040                    true,
6041                )?,
6042                right,
6043            )
6044        } else {
6045            (
6046                left,
6047                Self::assert_unique_one_side(
6048                    right,
6049                    &right_tag_columns,
6050                    &right_matched_tags,
6051                    right_time_index_column.as_deref(),
6052                    false,
6053                )?,
6054            )
6055        };
6056
6057        // push time index column if it exists
6058        if let (Some(left_time_index_column), Some(right_time_index_column)) = (
6059            left_time_index_column.clone(),
6060            right_time_index_column.clone(),
6061        ) {
6062            left_tag_columns.insert(left_time_index_column);
6063            right_tag_columns.insert(right_time_index_column);
6064        }
6065
6066        let right = LogicalPlanBuilder::from(right)
6067            .alias(right_table_ref.clone())
6068            .context(DataFusionPlanningSnafu)?
6069            .build()
6070            .context(DataFusionPlanningSnafu)?;
6071
6072        // Inner Join on time index column to concat two operator
6073        let join_plan = LogicalPlanBuilder::from(left)
6074            .alias(left_table_ref.clone())
6075            .context(DataFusionPlanningSnafu)?
6076            .join_detailed(
6077                right,
6078                JoinType::Inner,
6079                (
6080                    left_tag_columns
6081                        .into_iter()
6082                        .map(Column::from_name)
6083                        .collect::<Vec<_>>(),
6084                    right_tag_columns
6085                        .into_iter()
6086                        .map(Column::from_name)
6087                        .collect::<Vec<_>>(),
6088                ),
6089                force_empty_join.then_some(lit(false)),
6090                NullEquality::NullEqualsNull,
6091            )
6092            .context(DataFusionPlanningSnafu)?
6093            .build()
6094            .context(DataFusionPlanningSnafu)?;
6095
6096        // The result labels no longer identify a series on their own: the "many" side can hold
6097        // several series per result label set, which PromQL cannot represent.
6098        let labels = checked_matching
6099            .then(|| Self::binary_result_labels(left_context, right_context, modifier))
6100            .flatten()
6101            .filter(|_| {
6102                Self::binary_result_labels_may_repeat(left_context, right_context, modifier)
6103            });
6104        let Some(labels) = labels else {
6105            return Ok(join_plan);
6106        };
6107        let schema = join_plan.schema().clone();
6108        let group_exprs = labels
6109            .iter()
6110            .map(|(from_left, label)| {
6111                let table_ref = if *from_left {
6112                    &left_table_ref
6113                } else {
6114                    &right_table_ref
6115                };
6116                schema
6117                    .qualified_field_with_name(Some(table_ref), label)
6118                    .map(|field| DfExpr::Column(field.into()))
6119                    .context(DataFusionPlanningSnafu)
6120            })
6121            .collect::<Result<Vec<_>>>()?;
6122        let time_index_expr = left_time_index_column
6123            .as_deref()
6124            .map(|column| {
6125                schema
6126                    .qualified_field_with_name(Some(&left_table_ref), column)
6127                    .map(|field| DfExpr::Column(field.into()))
6128                    .context(DataFusionPlanningSnafu)
6129            })
6130            .transpose()?
6131            .with_context(|| UnexpectedPlanExprSnafu {
6132                desc: "vector matching on a plan without a time index",
6133            })?;
6134        let violation = if matches!(
6135            modifier.as_ref().map(|modifier| &modifier.card),
6136            Some(VectorMatchCardinality::OneToOne)
6137        ) {
6138            MatchGroupViolation::ImplicitManyToOne
6139        } else {
6140            MatchGroupViolation::AmbiguousGroupLabels
6141        };
6142        Self::assert_unique_match_group(
6143            join_plan,
6144            group_exprs,
6145            labels.into_iter().map(|(_, label)| label).collect(),
6146            time_index_expr,
6147            violation,
6148        )
6149    }
6150
6151    /// Guard the side of a vector matching that must hold one series per match group.
6152    fn assert_unique_one_side(
6153        plan: LogicalPlan,
6154        join_keys: &BTreeSet<String>,
6155        tag_columns: &[String],
6156        time_index_column: Option<&str>,
6157        one_side_is_left: bool,
6158    ) -> Result<LogicalPlan> {
6159        let Some(time_index_column) = time_index_column else {
6160            return Ok(plan);
6161        };
6162        if tag_columns.iter().all(|tag| join_keys.contains(tag)) {
6163            return Ok(plan);
6164        }
6165
6166        let group_labels = join_keys.iter().cloned().collect::<Vec<_>>();
6167        let group_exprs = group_labels
6168            .iter()
6169            .map(|label| DfExpr::Column(Column::from_name(label)))
6170            .collect();
6171        Self::assert_unique_match_group(
6172            plan,
6173            group_exprs,
6174            group_labels,
6175            DfExpr::Column(Column::from_name(time_index_column)),
6176            MatchGroupViolation::DuplicateOnOneSide { one_side_is_left },
6177        )
6178    }
6179
6180    fn selected_binary_match_labels(
6181        left_context: &PromPlannerContext,
6182        right_context: &PromPlannerContext,
6183        modifier: &Option<BinModifier>,
6184    ) -> BTreeSet<String> {
6185        let mut labels = left_context
6186            .tag_columns
6187            .iter()
6188            .chain(&right_context.tag_columns)
6189            .cloned()
6190            .collect::<BTreeSet<_>>();
6191        if let Some(matching) = modifier
6192            .as_ref()
6193            .and_then(|modifier| modifier.matching.as_ref())
6194        {
6195            match matching {
6196                LabelModifier::Include(on) => {
6197                    labels = on
6198                        .labels
6199                        .iter()
6200                        .filter(|label| {
6201                            left_context.tag_columns.contains(label)
6202                                || right_context.tag_columns.contains(label)
6203                        })
6204                        .cloned()
6205                        .collect();
6206                }
6207                LabelModifier::Exclude(ignoring) => {
6208                    for label in &ignoring.labels {
6209                        labels.remove(label);
6210                    }
6211                }
6212            }
6213        }
6214        labels
6215    }
6216
6217    fn only_temporality_match_label_mismatches(
6218        left_context: &PromPlannerContext,
6219        right_context: &PromPlannerContext,
6220        modifier: &Option<BinModifier>,
6221    ) -> bool {
6222        let mut mismatches =
6223            Self::selected_binary_match_labels(left_context, right_context, modifier)
6224                .into_iter()
6225                .filter(|label| {
6226                    left_context.tag_columns.contains(label)
6227                        != right_context.tag_columns.contains(label)
6228                });
6229        matches!(
6230            (mismatches.next(), mismatches.next()),
6231            (Some(label), None) if label == OTLP_AGGREGATION_TEMPORALITY_LABEL
6232        )
6233    }
6234
6235    fn align_temporality_match_column(
6236        mut left: LogicalPlan,
6237        mut right: LogicalPlan,
6238        left_context: &mut PromPlannerContext,
6239        right_context: &mut PromPlannerContext,
6240    ) -> Result<(LogicalPlan, LogicalPlan, bool)> {
6241        let marker = OTLP_AGGREGATION_TEMPORALITY_LABEL;
6242        let left_has_marker = left_context.tag_columns.iter().any(|tag| tag == marker);
6243        let (data_type, value_type, add_to_left) = {
6244            let (present, add_to_left) = if left_has_marker {
6245                (&left, false)
6246            } else {
6247                (&right, true)
6248            };
6249            let data_type = present
6250                .schema()
6251                .fields()
6252                .iter()
6253                .find(|field| field.name() == marker)
6254                .map(|field| field.data_type().clone())
6255                .with_context(|| ColumnNotFoundSnafu {
6256                    col: marker.to_string(),
6257                })?;
6258            let value_type = Self::string_value_data_type(&data_type)
6259                .cloned()
6260                .with_context(|| UnexpectedPlanExprSnafu {
6261                    desc: format!("temporality match label {marker} must be a string"),
6262                })?;
6263            (data_type, value_type, add_to_left)
6264        };
6265        let null = Self::string_scalar_value(&value_type, None).with_context(|| {
6266            UnexpectedPlanExprSnafu {
6267                desc: format!("temporality match label {marker} must be a string"),
6268            }
6269        })?;
6270        let add_marker = |plan: LogicalPlan| {
6271            let visible = plan
6272                .schema()
6273                .iter()
6274                .map(|(qualifier, field)| {
6275                    DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6276                })
6277                .collect::<Vec<_>>();
6278            LogicalPlanBuilder::from(plan)
6279                .project(
6280                    visible
6281                        .into_iter()
6282                        .chain([DfExpr::Literal(null, None).alias(marker)]),
6283                )
6284                .context(DataFusionPlanningSnafu)?
6285                .build()
6286                .context(DataFusionPlanningSnafu)
6287        };
6288        if data_type != value_type {
6289            let present = if add_to_left { &mut right } else { &mut left };
6290            let visible = present
6291                .schema()
6292                .iter()
6293                .map(|(qualifier, field)| {
6294                    let column =
6295                        DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()));
6296                    if field.name() == marker {
6297                        DfExpr::Cast(Cast::new(Box::new(column), value_type.clone()))
6298                            .alias_qualified(qualifier.cloned(), field.name().clone())
6299                    } else {
6300                        column
6301                    }
6302                })
6303                .collect::<Vec<_>>();
6304            *present = LogicalPlanBuilder::from(present.clone())
6305                .project(visible)
6306                .context(DataFusionPlanningSnafu)?
6307                .build()
6308                .context(DataFusionPlanningSnafu)?;
6309        }
6310        if add_to_left {
6311            left = add_marker(left)?;
6312            left_context.tag_columns.push(marker.to_string());
6313        } else {
6314            right = add_marker(right)?;
6315            right_context.tag_columns.push(marker.to_string());
6316        }
6317        Ok((left, right, add_to_left))
6318    }
6319
6320    fn normalized_match_key_expr(
6321        label: &str,
6322        field: Option<(Option<TableReference>, ArrowDataType)>,
6323        value_type: &ArrowDataType,
6324        internal_name: &str,
6325    ) -> DfExpr {
6326        let empty = Self::string_scalar_value(value_type, Some(String::new()))
6327            .expect("match label value type is a string");
6328        let expr = if let Some((qualifier, data_type)) = field {
6329            let column = DfExpr::Column(Column::new(qualifier, label));
6330            let column = if &data_type == value_type {
6331                column
6332            } else {
6333                DfExpr::Cast(Cast::new(Box::new(column), value_type.clone()))
6334            };
6335            DfExpr::ScalarFunction(ScalarFunction {
6336                func: coalesce(),
6337                args: vec![column, DfExpr::Literal(empty, None)],
6338            })
6339        } else {
6340            DfExpr::Literal(empty, None)
6341        };
6342        expr.alias(internal_name)
6343    }
6344
6345    fn is_zero_row_empty_relation(plan: &LogicalPlan) -> bool {
6346        // `produce_one_row` is used for input-free plans that still emit one row;
6347        // only the false case is a statically proven empty vector.
6348        matches!(plan, LogicalPlan::EmptyRelation(relation) if !relation.produce_one_row)
6349    }
6350
6351    fn string_value_data_type(data_type: &ArrowDataType) -> Option<&ArrowDataType> {
6352        match data_type {
6353            data_type if data_type.is_string() => Some(data_type),
6354            ArrowDataType::Dictionary(_, value_type) if value_type.is_string() => Some(value_type),
6355            _ => None,
6356        }
6357    }
6358
6359    fn string_scalar_value(
6360        data_type: &ArrowDataType,
6361        value: Option<String>,
6362    ) -> Option<ScalarValue> {
6363        match data_type {
6364            ArrowDataType::Utf8 => Some(ScalarValue::Utf8(value)),
6365            ArrowDataType::LargeUtf8 => Some(ScalarValue::LargeUtf8(value)),
6366            ArrowDataType::Utf8View => Some(ScalarValue::Utf8View(value)),
6367            ArrowDataType::Dictionary(key_type, value_type) => Some(ScalarValue::Dictionary(
6368                key_type.clone(),
6369                Box::new(Self::string_scalar_value(value_type, value)?),
6370            )),
6371            _ => None,
6372        }
6373    }
6374
6375    fn common_label_data_type(
6376        left: Option<&ArrowDataType>,
6377        right: Option<&ArrowDataType>,
6378    ) -> Option<ArrowDataType> {
6379        match (left, right) {
6380            (Some(left), Some(right)) if left == right => {
6381                Self::string_value_data_type(left).map(|_| left.clone())
6382            }
6383            (Some(left), Some(right)) => {
6384                let left_value_type = Self::string_value_data_type(left)?;
6385                let right_value_type = Self::string_value_data_type(right)?;
6386                // DataFusion projections can decode dictionaries, but do not encode plain strings
6387                // as dictionaries. Preserve the encoding only when both inputs already share it.
6388                match (left_value_type, right_value_type) {
6389                    (left, right) if left == right => Some(left.clone()),
6390                    (ArrowDataType::LargeUtf8, _) | (_, ArrowDataType::LargeUtf8) => {
6391                        Some(ArrowDataType::LargeUtf8)
6392                    }
6393                    (ArrowDataType::Utf8View, ArrowDataType::Utf8View) => {
6394                        Some(ArrowDataType::Utf8View)
6395                    }
6396                    _ => Some(ArrowDataType::Utf8),
6397                }
6398            }
6399            (Some(data_type), None) | (None, Some(data_type)) => {
6400                Self::string_value_data_type(data_type).cloned()
6401            }
6402            (None, None) => Some(ArrowDataType::Utf8),
6403        }
6404    }
6405
6406    /// Build a projection that project and perform operation expr for every value columns.
6407    /// Non-value columns (tag and timestamp) will be preserved in the projection.
6408    ///
6409    /// # Side effect
6410    ///
6411    /// This function will update the value columns in the context. Those new column names
6412    /// don't contains qualifier.
6413    fn projection_for_each_field_column<F>(
6414        &mut self,
6415        input: LogicalPlan,
6416        name_to_expr: F,
6417    ) -> Result<LogicalPlan>
6418    where
6419        F: FnMut(&String) -> Result<DfExpr>,
6420    {
6421        self.projection_for_each_field_column_with_labels(input, None, name_to_expr)
6422    }
6423
6424    /// Like [`Self::projection_for_each_field_column`], but projects `result_labels` instead of
6425    /// the context tag columns when a binary operation derived its own result label set.
6426    fn projection_for_each_field_column_with_labels<F>(
6427        &mut self,
6428        input: LogicalPlan,
6429        result_labels: Option<&BinaryResultLabels>,
6430        name_to_expr: F,
6431    ) -> Result<LogicalPlan>
6432    where
6433        F: FnMut(&String) -> Result<DfExpr>,
6434    {
6435        // Keep the generated float/histogram lane names while an element-wise operation
6436        // preserves both sample types, so downstream operators still recognize the pair.
6437        let preserve_field_names =
6438            Self::field_columns_are_alternative_samples(input.schema(), &self.ctx.field_columns);
6439        let table_ref = self.ctx.table_name.clone().map(TableReference::bare);
6440        // Derived labels can be unqualified even when the context still names the source table.
6441        let input_schema = input.schema().clone();
6442        let lookup = |col: &String| {
6443            input_schema
6444                .qualified_field_with_name(table_ref.as_ref(), col)
6445                .or_else(|_| input_schema.qualified_field_with_unqualified_name(col))
6446                .map(|field| DfExpr::Column(field.into()))
6447                .context(DataFusionPlanningSnafu)
6448        };
6449        let tag_columns_iter = match result_labels {
6450            Some(labels) => labels.exprs.iter().cloned().map(Ok).collect::<Vec<_>>(),
6451            None => self.ctx.tag_columns.iter().map(lookup).collect::<Vec<_>>(),
6452        };
6453        let non_field_columns_iter = tag_columns_iter
6454            .into_iter()
6455            .chain(self.ctx.time_index_column.iter().map(lookup));
6456        let tsid_iter = match result_labels {
6457            Some(labels) => labels.tsid_projection(table_ref.clone()),
6458            None => Self::optional_tsid_projection(
6459                input.schema(),
6460                table_ref.as_ref(),
6461                self.ctx.use_tsid,
6462            ),
6463        }
6464        .into_iter()
6465        .map(Ok);
6466
6467        // build computation exprs
6468        let result_field_columns = self
6469            .ctx
6470            .field_columns
6471            .iter()
6472            .map(name_to_expr)
6473            .collect::<Result<Vec<_>>>()?;
6474
6475        // alias the computation exprs to remove qualifier
6476        if !preserve_field_names {
6477            self.ctx.field_columns = result_field_columns
6478                .iter()
6479                .map(|expr| expr.schema_name().to_string())
6480                .collect();
6481        }
6482        let field_columns_iter = result_field_columns
6483            .into_iter()
6484            .zip(self.ctx.field_columns.iter())
6485            .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::<String>, name))));
6486
6487        // chain non-field columns (unchanged) and field columns (applied computation then alias)
6488        let project_fields = non_field_columns_iter
6489            .chain(tsid_iter)
6490            .chain(field_columns_iter)
6491            .collect::<Result<Vec<_>>>()?;
6492
6493        LogicalPlanBuilder::from(input)
6494            .project(project_fields)
6495            .context(DataFusionPlanningSnafu)?
6496            .build()
6497            .context(DataFusionPlanningSnafu)
6498    }
6499
6500    /// Build a filter plan on one value column or a float/histogram alternative pair.
6501    fn filter_on_field_column<F>(&self, input: LogicalPlan, name_to_expr: F) -> Result<LogicalPlan>
6502    where
6503        F: FnMut(&String) -> Result<DfExpr>,
6504    {
6505        ensure!(
6506            self.ctx.field_columns.len() == 1
6507                || Self::field_columns_are_alternative_samples(
6508                    input.schema(),
6509                    &self.ctx.field_columns,
6510                ),
6511            UnsupportedExprSnafu {
6512                name: "filter on multi-value input"
6513            }
6514        );
6515
6516        let field_column_filters = self
6517            .ctx
6518            .field_columns
6519            .iter()
6520            .map(name_to_expr)
6521            .collect::<Result<Vec<_>>>()?;
6522        let field_column_filter =
6523            disjunction(field_column_filters).context(UnsupportedExprSnafu {
6524                name: "filter on empty input",
6525            })?;
6526
6527        LogicalPlanBuilder::from(input)
6528            .filter(field_column_filter)
6529            .context(DataFusionPlanningSnafu)?
6530            .build()
6531            .context(DataFusionPlanningSnafu)
6532    }
6533
6534    /// Generate an expr like `date_part("hour", <TIME_INDEX>)`. Caller should ensure the
6535    /// time index column in context is set
6536    fn date_part_on_time_index(&self, date_part: &str) -> Result<DfExpr> {
6537        let input_expr = datafusion::logical_expr::ident(
6538            self.ctx
6539                .time_index_column
6540                .as_ref()
6541                // table name doesn't matters here
6542                .with_context(|| TimeIndexNotFoundSnafu {
6543                    table: "<doesn't matter>",
6544                })?,
6545        );
6546        let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
6547            func: datafusion_functions::datetime::date_part(),
6548            args: vec![date_part.lit(), input_expr],
6549        });
6550        Ok(fn_expr)
6551    }
6552
6553    fn strip_tsid_column(&self, plan: LogicalPlan) -> Result<LogicalPlan> {
6554        let schema = plan.schema();
6555        if !schema
6556            .fields()
6557            .iter()
6558            .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
6559        {
6560            return Ok(plan);
6561        }
6562
6563        // Preserve column qualifiers so downstream plan nodes can keep referencing
6564        // the columns by their original qualified names.
6565        let project_exprs = schema
6566            .iter()
6567            .filter(|(_, field)| field.name() != DATA_SCHEMA_TSID_COLUMN_NAME)
6568            .map(|(qualifier, field)| {
6569                DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6570            })
6571            .collect::<Vec<_>>();
6572
6573        LogicalPlanBuilder::from(plan)
6574            .project(project_exprs)
6575            .context(DataFusionPlanningSnafu)?
6576            .build()
6577            .context(DataFusionPlanningSnafu)
6578    }
6579
6580    /// Apply an alias to the query result by adding a projection with the alias name
6581    fn apply_alias(&mut self, plan: LogicalPlan, alias_name: String) -> Result<LogicalPlan> {
6582        let fields_expr = self.create_field_column_exprs()?;
6583
6584        // TODO(dennis): how to support multi-value aliasing?
6585        ensure!(
6586            fields_expr.len() == 1,
6587            UnsupportedExprSnafu {
6588                name: "alias on multi-value result"
6589            }
6590        );
6591
6592        let project_fields = fields_expr
6593            .into_iter()
6594            .map(|expr| expr.alias(&alias_name))
6595            .chain(self.create_tag_column_exprs()?)
6596            .chain(Some(self.create_time_index_column_expr()?));
6597
6598        LogicalPlanBuilder::from(plan)
6599            .project(project_fields)
6600            .context(DataFusionPlanningSnafu)?
6601            .build()
6602            .context(DataFusionPlanningSnafu)
6603    }
6604}
6605
6606#[derive(Default, Debug)]
6607struct FunctionArgs {
6608    input: Option<PromExpr>,
6609    literals: Vec<DfExpr>,
6610}
6611
6612/// Represents different types of scalar functions supported in PromQL expressions.
6613/// Each variant defines how the function should be processed and what arguments it expects.
6614#[derive(Debug, Clone)]
6615enum ScalarFunc {
6616    /// DataFusion's registered(including built-in) scalar functions (e.g., abs, sqrt, round, clamp).
6617    /// These are passed through directly to DataFusion's execution engine.
6618    /// Processing: Simple argument insertion at the specified position.
6619    DataFusionBuiltin(Arc<ScalarUdfDef>),
6620    /// User-defined functions registered in DataFusion's function registry.
6621    /// Similar to DataFusionBuiltin but for custom functions not built into DataFusion.
6622    /// Processing: Direct pass-through with argument positioning.
6623    DataFusionUdf(Arc<ScalarUdfDef>),
6624    /// Native histogram helper UDFs. Non-histogram inputs are projected as NULL
6625    /// so the normal PromQL empty-value filter drops them.
6626    NativeHistogramUdf(Arc<ScalarUdfDef>),
6627    /// PromQL-specific functions that operate on time series data with temporal context.
6628    /// These functions require both timestamp ranges and values to perform calculations.
6629    /// Processing: Automatically injects timestamp_range and value columns as first arguments.
6630    /// Examples: idelta, irate, resets, changes, deriv, *_over_time function
6631    Udf(Arc<ScalarUdfDef>),
6632    /// PromQL functions requiring extrapolation calculations with explicit range information.
6633    /// These functions need to know the time range length to perform rate calculations.
6634    /// The second field contains the range length in milliseconds.
6635    /// Processing: Injects timestamp_range, value, time_index columns and appends range_length.
6636    /// Examples: increase, rate, delta
6637    // TODO(ruihang): maybe merge with Udf later
6638    ExtrapolateUdf(Arc<ScalarUdfDef>, i64),
6639    /// Functions that generate expressions directly without external UDF calls.
6640    /// The expression is constructed during function matching and requires no additional processing.
6641    /// Examples: time(), minute(), hour(), month(), year() and other date/time extractors
6642    GeneratedExpr,
6643}
6644
6645#[cfg(test)]
6646mod test;