Skip to main content

query/promql/
planner.rs

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