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