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 matching_filters;
16
17use std::collections::{BTreeSet, HashMap, HashSet, VecDeque};
18use std::sync::Arc;
19use std::time::UNIX_EPOCH;
20
21use arrow::datatypes::IntervalDayTime;
22use async_recursion::async_recursion;
23use catalog::table_source::DfTableSourceProvider;
24use common_error::ext::ErrorExt;
25use common_error::status_code::StatusCode;
26use common_function::function::FunctionContext;
27use common_query::native_histogram::native_histogram_value_type;
28use common_query::prelude::{
29    GREPTIME_TEMPORALITY_DELTA, OTLP_AGGREGATION_TEMPORALITY_LABEL, greptime_native_histogram,
30    greptime_value,
31};
32use common_query::promql_annotations::PromqlAnnotationCollector;
33use datafusion::common::DFSchemaRef;
34use datafusion::datasource::DefaultTableSource;
35use datafusion::functions_aggregate::average::avg_udaf;
36use datafusion::functions_aggregate::count::count_udaf;
37use datafusion::functions_aggregate::expr_fn::first_value;
38use datafusion::functions_aggregate::min_max::{max_udaf, min_udaf};
39use datafusion::functions_aggregate::stddev::stddev_pop_udaf;
40use datafusion::functions_aggregate::sum::sum_udaf;
41use datafusion::functions_aggregate::variance::var_pop_udaf;
42use datafusion::functions_window::row_number::RowNumber;
43use datafusion::logical_expr::expr::{Alias, ScalarFunction, WindowFunction};
44use datafusion::logical_expr::expr_rewriter::normalize_cols;
45use datafusion::logical_expr::{
46    BinaryExpr, Cast, Extension, LogicalPlan, LogicalPlanBuilder, Operator,
47    ScalarUDF as ScalarUdfDef, WindowFrame, WindowFunctionDefinition,
48};
49use datafusion::optimizer::simplify_expressions::ExprSimplifier;
50use datafusion::prelude as df_prelude;
51use datafusion::prelude::{Column, Expr as DfExpr, JoinType};
52use datafusion::scalar::ScalarValue;
53use datafusion_common::{DFSchema, NullEquality, TableReference};
54use datafusion_expr::expr::WindowFunctionParams;
55use datafusion_expr::expr_fn::when;
56use datafusion_expr::simplify::SimplifyContext;
57use datafusion_expr::utils::{conjunction, disjunction};
58use datafusion_expr::{ExprSchemable, Literal, SortExpr, TableSource, col, lit};
59use datafusion_functions::core::coalesce;
60use datatypes::arrow::datatypes::{DataType as ArrowDataType, TimeUnit as ArrowTimeUnit};
61use datatypes::data_type::{ConcreteDataType, DataType as GreptimeDataType};
62use itertools::Itertools;
63use once_cell::sync::Lazy;
64use promql::extension_plan::{
65    Absent, EmptyMetric, HistogramFold, HistogramFoldOperation, InstantManipulate, Millisecond,
66    RangeManipulate, ScalarCalculate, SeriesDivide, SeriesNormalize, UnionDistinctOn,
67    build_special_time_expr,
68};
69use promql::functions::{
70    AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing,
71    IDelta, Increase, LastOverTime, MatchGroupViolation, MaxOverTime, MinOverTime, MixedRange,
72    NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAggAvg,
73    NativeHistogramAggSum, NativeHistogramAvg, NativeHistogramAvgOverTime, NativeHistogramChanges,
74    NativeHistogramCount, NativeHistogramCountOverTime, NativeHistogramDelta,
75    NativeHistogramDivScalar, NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction,
76    NativeHistogramIDelta, NativeHistogramIRate, NativeHistogramIncrease,
77    NativeHistogramLastOverTime, NativeHistogramMulScalar, NativeHistogramNeg,
78    NativeHistogramNotEq, NativeHistogramPresentOverTime, NativeHistogramQuantile,
79    NativeHistogramRate, NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev,
80    NativeHistogramStdvar, NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime,
81    NativeHistogramToString, PredictLinear, PresentOverTime, PromqlFloatToString, QuantileOverTime,
82    Rate, Resets, Round, StddevOverTime, StdvarOverTime, SumOverTime, UniqueMatchGroup,
83    quantile_udaf,
84};
85use promql_parser::label::{METRIC_NAME, MatchOp, Matcher, Matchers};
86use promql_parser::parser::token::TokenType;
87use promql_parser::parser::value::ValueType;
88use promql_parser::parser::{
89    AggregateExpr, BinModifier, BinaryExpr as PromBinaryExpr, Call, EvalStmt, Expr as PromExpr,
90    Function, FunctionArgs as PromFunctionArgs, LabelModifier, MatrixSelector, NumberLiteral,
91    Offset, ParenExpr, StringLiteral, SubqueryExpr, UnaryExpr, VectorMatchCardinality,
92    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, CombineTableColumnMismatchSnafu, DataFusionPlanningSnafu,
108    ExpectRangeSelectorSnafu, FunctionInvalidArgumentSnafu, InvalidDestinationLabelNameSnafu,
109    InvalidRegularExpressionSnafu, InvalidTimeRangeSnafu, MultiFieldsNotSupportedSnafu,
110    MultipleMetricMatchersSnafu, MultipleVectorSnafu, NoMetricMatcherSnafu, PromqlPlanNodeSnafu,
111    Result, SameLabelSetSnafu, TableNameNotFoundSnafu, TimeIndexNotFoundSnafu,
112    UnexpectedPlanExprSnafu, UnexpectedTokenSnafu, UnknownTableSnafu, UnsupportedExprSnafu,
113    UnsupportedMatcherOpSnafu, UnsupportedVectorMatchSnafu, ValueNotFoundSnafu,
114    ZeroRangeSelectorSnafu,
115};
116use crate::query_engine::QueryEngineState;
117
118/// `time()` function in PromQL.
119const SPECIAL_TIME_FUNCTION: &str = "time";
120/// `scalar()` function in PromQL.
121const SCALAR_FUNCTION: &str = "scalar";
122/// `absent()` function in PromQL
123const SPECIAL_ABSENT_FUNCTION: &str = "absent";
124/// `histogram_quantile` function in PromQL
125const SPECIAL_HISTOGRAM_QUANTILE: &str = "histogram_quantile";
126/// `histogram_fraction` function in PromQL
127const SPECIAL_HISTOGRAM_FRACTION: &str = "histogram_fraction";
128/// `vector` function in PromQL
129const SPECIAL_VECTOR_FUNCTION: &str = "vector";
130/// `le` column for conventional histogram.
131const LE_COLUMN_NAME: &str = "le";
132
133/// Static regex for validating label names according to Prometheus specification.
134/// Label names must match the regex: [a-zA-Z_][a-zA-Z0-9_]*
135static LABEL_NAME_REGEX: Lazy<Regex> =
136    Lazy::new(|| Regex::new(r"^[a-zA-Z_][a-zA-Z0-9_]*$").unwrap());
137
138const DEFAULT_TIME_INDEX_COLUMN: &str = "time";
139
140/// default value column name for empty metric
141const DEFAULT_FIELD_COLUMN: &str = "value";
142
143/// Special modifier to project field columns under multi-field mode
144const FIELD_COLUMN_MATCHER: &str = "__field__";
145
146/// Special modifier for cross schema query
147const SCHEMA_COLUMN_MATCHER: &str = "__schema__";
148const DB_COLUMN_MATCHER: &str = "__database__";
149
150/// Prefix for generated binary island leaf aliases.
151const BINARY_ISLAND_LEAF_ALIAS_PREFIX: &str = "__prom_v";
152const OR_FLOAT_FIELD_PREFIX: &str = "__promql_or_float_";
153const OR_HISTOGRAM_FIELD_PREFIX: &str = "__promql_or_histogram_";
154const TIMESTAMP_VALUE_PREFIX: &str = "__promql_timestamp_value_";
155/// Per-match-group row count the vector matching cardinality check is built on.
156const MATCH_GROUP_COUNT_COLUMN: &str = "__promql_match_group_count";
157
158/// Threshold for scatter scan mode
159const MAX_SCATTER_POINTS: i64 = 400;
160
161/// Interval 1 hour in millisecond
162const INTERVAL_1H: i64 = 60 * 60 * 1000;
163
164#[derive(Default, Debug, Clone)]
165struct PromPlannerContext {
166    // query parameters
167    start: Millisecond,
168    end: Millisecond,
169    interval: Millisecond,
170    lookback_delta: Millisecond,
171
172    // planner states
173    table_name: Option<String>,
174    time_index_column: Option<String>,
175    field_columns: Vec<String>,
176    tag_columns: Vec<String>,
177    /// `by(...)` labels of the aggregation that produced this operand that are not series tags of
178    /// its input, i.e. value fields (or a label an inner aggregation already reported as one).
179    ///
180    /// An aggregation reports every `by(...)` label it finds in its input schema among its tag
181    /// columns, and a value field named there is a group key of the aggregate rather than a
182    /// property of a series: its value varies between the samples of one series. Lowering a
183    /// matcher on it into the scan would change which samples are selected (#9242), so
184    /// [`matching_filters`] refuses to propagate such a matcher.
185    aggregation_field_labels: Vec<String>,
186    /// Use metric engine internal series identifier column (`__tsid`) as series key.
187    ///
188    /// This is enabled only when the underlying scan can provide `__tsid` (`UInt64`). The planner
189    /// uses it internally (e.g. as the series key for [`SeriesDivide`]) and strips it from the
190    /// final output.
191    use_tsid: bool,
192    /// The matcher for field columns `__field__`.
193    field_column_matcher: Option<Vec<Matcher>>,
194    /// The matcher for selectors (normal matchers).
195    selector_matcher: Vec<Matcher>,
196    schema_name: Option<String>,
197    /// The range in millisecond of range selector. None if there is no range selector.
198    range: Option<Millisecond>,
199}
200
201#[derive(Debug, Clone, PartialEq, Eq, Hash)]
202struct VectorLeafKey {
203    metric_name: String,
204    matchers: Vec<(String, String, String)>,
205    or_matchers: Vec<Vec<(String, String, String)>>,
206    offset_ms: i128,
207    at: String,
208}
209
210#[derive(Debug, Clone)]
211struct IslandLeaf {
212    selector: VectorSelector,
213    display_table: String,
214}
215
216#[derive(Debug, Clone)]
217enum IslandExpr {
218    VectorLeaf(usize),
219    Scalar(DfExpr),
220    Unary {
221        input: Box<IslandExpr>,
222    },
223    Binary {
224        op: TokenType,
225        lhs: Box<IslandExpr>,
226        rhs: Box<IslandExpr>,
227    },
228}
229
230impl IslandExpr {
231    fn try_new(expr: &PromExpr, env: &mut IslandCollectEnv) -> Option<Self> {
232        if let Some(expr) = PromPlanner::try_build_literal_expr(expr) {
233            return Some(Self::Scalar(expr));
234        }
235
236        match expr {
237            PromExpr::Paren(ParenExpr { expr }) => Self::try_new(expr, env),
238            PromExpr::VectorSelector(selector) => {
239                let leaf = env.intern_leaf(selector)?;
240                Some(Self::VectorLeaf(leaf))
241            }
242            PromExpr::Unary(UnaryExpr { expr }) => {
243                let input = Self::try_new(expr, env)?;
244                Some(Self::Unary {
245                    input: Box::new(input),
246                })
247            }
248            PromExpr::Binary(PromBinaryExpr {
249                lhs,
250                rhs,
251                op,
252                modifier,
253            }) if matches!(
254                op.id(),
255                token::T_ADD
256                    | token::T_SUB
257                    | token::T_MUL
258                    | token::T_DIV
259                    | token::T_MOD
260                    | token::T_POW
261                    | token::T_ATAN2
262            ) && modifier.as_ref().is_none_or(|modifier| {
263                !modifier.return_bool
264                    && modifier.matching.is_none()
265                    && matches!(modifier.card, VectorMatchCardinality::OneToOne)
266                    && modifier.fill_values.lhs.is_none()
267                    && modifier.fill_values.rhs.is_none()
268            }) =>
269            {
270                let lhs = Self::try_new(lhs, env)?;
271                let rhs = Self::try_new(rhs, env)?;
272                Some(Self::Binary {
273                    op: *op,
274                    lhs: Box::new(lhs),
275                    rhs: Box::new(rhs),
276                })
277            }
278            _ => None,
279        }
280    }
281}
282
283#[derive(Debug, Default)]
284struct IslandCollectEnv {
285    leaf_by_key: HashMap<VectorLeafKey, usize>,
286    leaves: Vec<IslandLeaf>,
287    vector_occurrences: usize,
288}
289
290#[derive(Debug)]
291struct PlannedIslandLeaf {
292    plan: LogicalPlan,
293    ctx: PromPlannerContext,
294    alias: TableReference,
295    display_table: String,
296}
297
298/// Result labels a vector-vector binary operation derives from its matching modifier, projected
299/// from the operand each label belongs to.
300#[derive(Debug)]
301struct BinaryResultLabels {
302    exprs: Vec<DfExpr>,
303    names: Vec<String>,
304    aggregation_field_labels: Vec<String>,
305    /// `__tsid` column of the operand the labels come from, when that operand contributes its
306    /// whole tag set and the column still identifies the result series.
307    tsid: Option<DfExpr>,
308}
309
310impl BinaryResultLabels {
311    fn apply(&self, ctx: &mut PromPlannerContext) {
312        ctx.tag_columns = self.names.clone();
313        ctx.aggregation_field_labels = self.aggregation_field_labels.clone();
314        ctx.use_tsid = self.tsid.is_some();
315    }
316
317    /// `__tsid` is projected from whichever operand the labels come from, so it carries that
318    /// operand's qualifier. Re-qualify it as the result's own, which is what the context names
319    /// and what the enclosing expression looks the column up by.
320    fn tsid_projection(&self, table_ref: Option<TableReference>) -> Option<DfExpr> {
321        self.tsid
322            .clone()
323            .map(|tsid| tsid.alias_qualified(table_ref, DATA_SCHEMA_TSID_COLUMN_NAME))
324    }
325}
326
327#[derive(Debug)]
328struct IslandFieldExprs {
329    exprs: Vec<DfExpr>,
330    names: Vec<String>,
331    scalar: bool,
332}
333
334impl VectorLeafKey {
335    fn from_selector(selector: &VectorSelector) -> Option<Self> {
336        let mut metric_name = selector.name.clone();
337        let mut matchers = Vec::with_capacity(selector.matchers.matchers.len());
338        let matcher_key = |matcher: &Matcher| {
339            (
340                matcher.name.clone(),
341                matcher.op.to_string(),
342                matcher.value.clone(),
343            )
344        };
345
346        for matcher in &selector.matchers.matchers {
347            if matcher.name == METRIC_NAME {
348                if matcher.op != MatchOp::Equal || metric_name.is_some() {
349                    return None;
350                }
351                metric_name = Some(matcher.value.clone());
352            } else {
353                matchers.push(matcher_key(matcher));
354            }
355        }
356        matchers.sort();
357
358        let mut or_matchers = selector
359            .matchers
360            .or_matchers
361            .iter()
362            .map(|group| {
363                let mut group = group.iter().map(matcher_key).collect::<Vec<_>>();
364                group.sort();
365                group
366            })
367            .collect::<Vec<_>>();
368        or_matchers.sort();
369
370        Some(Self {
371            metric_name: metric_name?,
372            matchers,
373            or_matchers,
374            offset_ms: match &selector.offset {
375                Some(Offset::Pos(duration)) => duration.as_millis() as i128,
376                Some(Offset::Neg(duration)) => -(duration.as_millis() as i128),
377                None => 0,
378            },
379            at: format!("{:?}", selector.at),
380        })
381    }
382}
383
384impl IslandCollectEnv {
385    fn intern_leaf(&mut self, selector: &VectorSelector) -> Option<usize> {
386        self.vector_occurrences += 1;
387        let key = VectorLeafKey::from_selector(selector)?;
388        if let Some(id) = self.leaf_by_key.get(&key) {
389            return Some(*id);
390        }
391
392        let id = self.leaves.len();
393        self.leaves.push(IslandLeaf {
394            selector: selector.clone(),
395            display_table: key.metric_name.clone(),
396        });
397        self.leaf_by_key.insert(key, id);
398        Some(id)
399    }
400}
401
402impl PromPlannerContext {
403    fn from_eval_stmt(stmt: &EvalStmt) -> Self {
404        Self {
405            start: stmt.start.duration_since(UNIX_EPOCH).unwrap().as_millis() as _,
406            end: stmt.end.duration_since(UNIX_EPOCH).unwrap().as_millis() as _,
407            interval: stmt.interval.as_millis() as _,
408            lookback_delta: stmt.lookback_delta.as_millis() as _,
409            ..Default::default()
410        }
411    }
412
413    /// Reset all planner states
414    fn reset(&mut self) {
415        self.table_name = None;
416        self.time_index_column = None;
417        self.field_columns = vec![];
418        self.tag_columns = vec![];
419        self.use_tsid = false;
420        self.field_column_matcher = None;
421        self.selector_matcher.clear();
422        self.schema_name = None;
423        self.range = None;
424    }
425
426    /// Reset table name and schema to empty
427    fn reset_table_name_and_schema(&mut self) {
428        self.table_name = Some(String::new());
429        self.schema_name = None;
430        self.use_tsid = false;
431    }
432
433    /// Check if `le` is present in tag columns
434    fn has_le_tag(&self) -> bool {
435        self.tag_columns.iter().any(|c| c.eq(&LE_COLUMN_NAME))
436    }
437}
438
439pub struct PromPlanner {
440    table_provider: DfTableSourceProvider,
441    ctx: PromPlannerContext,
442    /// Optional collector passed to native histogram UDFs.
443    promql_annotations: Option<PromqlAnnotationCollector>,
444}
445
446type BinaryFieldPair<'a> = (&'a String, &'a String);
447
448impl PromPlanner {
449    pub async fn stmt_to_plan(
450        table_provider: DfTableSourceProvider,
451        stmt: &EvalStmt,
452        query_engine_state: &QueryEngineState,
453    ) -> Result<LogicalPlan> {
454        Self::stmt_to_plan_with_annotations(table_provider, stmt, query_engine_state, None).await
455    }
456
457    /// Plans a PromQL statement and passes the optional collector to histogram UDFs.
458    pub async fn stmt_to_plan_with_annotations(
459        table_provider: DfTableSourceProvider,
460        stmt: &EvalStmt,
461        query_engine_state: &QueryEngineState,
462        promql_annotations: Option<PromqlAnnotationCollector>,
463    ) -> Result<LogicalPlan> {
464        let mut planner = Self {
465            table_provider,
466            ctx: PromPlannerContext::from_eval_stmt(stmt),
467            promql_annotations,
468        };
469
470        let plan = planner
471            .prom_expr_to_plan(&stmt.expr, query_engine_state)
472            .await?;
473
474        // Never leak internal series identifier to output.
475        planner.strip_tsid_column(plan)
476    }
477
478    pub async fn prom_expr_to_plan(
479        &mut self,
480        prom_expr: &PromExpr,
481        query_engine_state: &QueryEngineState,
482    ) -> Result<LogicalPlan> {
483        self.prom_expr_to_plan_inner(prom_expr, false, query_engine_state)
484            .await
485    }
486
487    /**
488    Converts a PromQL expression to a logical plan.
489
490    NOTE:
491        The `timestamp_fn` indicates whether the PromQL `timestamp()` function is being evaluated in the current context.
492        If `true`, the planner generates a logical plan that projects the timestamp (time index) column
493        as the value column for each input row, implementing the PromQL `timestamp()` function semantics.
494        If `false`, the planner generates the standard logical plan for the given PromQL expression.
495    */
496    #[async_recursion]
497    async fn prom_expr_to_plan_inner(
498        &mut self,
499        prom_expr: &PromExpr,
500        timestamp_fn: bool,
501        query_engine_state: &QueryEngineState,
502    ) -> Result<LogicalPlan> {
503        let res = match prom_expr {
504            PromExpr::Aggregate(expr) => {
505                self.prom_aggr_expr_to_plan(query_engine_state, expr)
506                    .await?
507            }
508            PromExpr::Unary(expr) => {
509                self.prom_unary_expr_to_plan(query_engine_state, expr)
510                    .await?
511            }
512            PromExpr::Binary(expr) => {
513                self.prom_binary_expr_to_plan(query_engine_state, expr)
514                    .await?
515            }
516            PromExpr::Paren(ParenExpr { expr }) => {
517                self.prom_expr_to_plan_inner(expr, timestamp_fn, query_engine_state)
518                    .await?
519            }
520            PromExpr::Subquery(expr) => {
521                self.prom_subquery_expr_to_plan(query_engine_state, expr)
522                    .await?
523            }
524            PromExpr::NumberLiteral(lit) => self.prom_number_lit_to_plan(lit)?,
525            PromExpr::StringLiteral(lit) => self.prom_string_lit_to_plan(lit)?,
526            PromExpr::VectorSelector(selector) => {
527                self.prom_vector_selector_to_plan(selector, timestamp_fn)
528                    .await?
529            }
530            PromExpr::MatrixSelector(selector) => {
531                self.prom_matrix_selector_to_plan(selector).await?
532            }
533            PromExpr::Call(expr) => {
534                self.prom_call_expr_to_plan(query_engine_state, expr)
535                    .await?
536            }
537            PromExpr::Extension(expr) => {
538                self.prom_ext_expr_to_plan(query_engine_state, expr).await?
539            }
540        };
541
542        Ok(res)
543    }
544
545    async fn prom_subquery_expr_to_plan(
546        &mut self,
547        query_engine_state: &QueryEngineState,
548        subquery_expr: &SubqueryExpr,
549    ) -> Result<LogicalPlan> {
550        let SubqueryExpr {
551            expr, range, step, ..
552        } = subquery_expr;
553
554        let current_interval = self.ctx.interval;
555        if let Some(step) = step {
556            self.ctx.interval = step.as_millis() as _;
557        }
558        let current_start = self.ctx.start;
559        self.ctx.start -= range.as_millis() as i64 - self.ctx.interval;
560        let input = self.prom_expr_to_plan(expr, query_engine_state).await?;
561        self.ctx.interval = current_interval;
562        self.ctx.start = current_start;
563
564        ensure!(!range.is_zero(), ZeroRangeSelectorSnafu);
565        let range_ms = range.as_millis() as _;
566        self.ctx.range = Some(range_ms);
567
568        let time_index_column =
569            self.ctx
570                .time_index_column
571                .clone()
572                .with_context(|| TimeIndexNotFoundSnafu {
573                    table: self.ctx.table_name.clone().unwrap_or_default(),
574                })?;
575
576        // `RangeManipulate` assumes each input batch holds exactly one series
577        // (it takes tag column values from row 0 and applies them to every
578        // output row). The inner expression may emit batches that mix series,
579        // so sort by series key + time index and split into per-series batches
580        // with a `SeriesDivide` first.
581        let input_schema = input.schema();
582        let input_has_tsid = input_schema.fields().iter().any(|field| {
583            field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
584                && field.data_type() == &ArrowDataType::UInt64
585        });
586        let (series_key_columns, mut sort_exprs) = if input_has_tsid {
587            (
588                vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()],
589                vec![
590                    DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME))
591                        .sort(true, true),
592                ],
593            )
594        } else {
595            // Only use tag columns that survive in the inner plan's schema —
596            // `ctx.tag_columns` can drift from the actual output.
597            let key_columns: Vec<String> = self
598                .ctx
599                .tag_columns
600                .iter()
601                .filter(|name| input_schema.has_column_with_unqualified_name(name))
602                .cloned()
603                .collect();
604            let sort = key_columns
605                .iter()
606                .map(|name| DfExpr::Column(Column::from_name(name)).sort(true, true))
607                .collect::<Vec<_>>();
608            (key_columns, sort)
609        };
610        sort_exprs.push(DfExpr::Column(Column::from_name(&time_index_column)).sort(true, true));
611
612        let sort_plan = LogicalPlanBuilder::from(input)
613            .sort(sort_exprs)
614            .context(DataFusionPlanningSnafu)?
615            .build()
616            .context(DataFusionPlanningSnafu)?;
617        let divide_plan = LogicalPlan::Extension(Extension {
618            node: Arc::new(SeriesDivide::new(
619                series_key_columns,
620                time_index_column.clone(),
621                sort_plan,
622            )),
623        });
624
625        let manipulate = RangeManipulate::new(
626            self.ctx.start,
627            self.ctx.end,
628            self.ctx.interval,
629            0,
630            range_ms,
631            time_index_column,
632            self.ctx.field_columns.clone(),
633            divide_plan,
634        )
635        .context(DataFusionPlanningSnafu)?;
636
637        Ok(LogicalPlan::Extension(Extension {
638            node: Arc::new(manipulate),
639        }))
640    }
641
642    async fn prom_aggr_expr_to_plan(
643        &mut self,
644        query_engine_state: &QueryEngineState,
645        aggr_expr: &AggregateExpr,
646    ) -> Result<LogicalPlan> {
647        let AggregateExpr {
648            op,
649            expr,
650            modifier,
651            param,
652        } = aggr_expr;
653
654        let input = self.prom_expr_to_plan(expr, query_engine_state).await?;
655        let input_has_tsid = input.schema().fields().iter().any(|field| {
656            field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
657                && field.data_type() == &ArrowDataType::UInt64
658        });
659
660        match (*op).id() {
661            token::T_TOPK | token::T_BOTTOMK => {
662                self.prom_topk_bottomk_to_plan(aggr_expr, input).await
663            }
664            _ => {
665                // When `__tsid` is available, tag columns may have been pruned from the input plan.
666                // For `keep_tsid` decision we should compare against the full row-key label set,
667                // otherwise we may incorrectly treat label-reducing aggregates as preserving labels.
668                let input_tag_columns = if input_has_tsid {
669                    self.collect_row_key_tag_columns_from_plan(&input)?
670                        .into_iter()
671                        .collect::<Vec<_>>()
672                } else {
673                    self.ctx.tag_columns.clone()
674                };
675                // calculate columns to group by
676                // Need to append time index column into group by columns
677                let mut group_exprs = self.agg_modifier_to_col(input.schema(), modifier, true)?;
678                let mixed_sample_columns =
679                    Self::alternative_sample_columns(input.schema(), &self.ctx.field_columns)
680                        .map(|(float, histogram)| (float.to_string(), histogram.to_string()));
681                // Aggregates over native histogram inputs may drop every sample in a group
682                // (e.g. `min` over histogram-only samples, or `sum` over histograms with
683                // incompatible schemas) and leave a NULL-valued group row behind. Compute this
684                // before `create_aggregate_exprs` mutates `ctx.field_columns`.
685                let preserve_any_value = mixed_sample_columns.is_some();
686                let has_native_histogram = preserve_any_value
687                    || self.all_field_columns_are_native_histograms(input.schema());
688                // convert op and value columns to aggregate exprs
689                let (mut aggr_exprs, prev_field_exprs) =
690                    self.create_aggregate_exprs(*op, param, &input)?;
691                let prev_field_exprs =
692                    normalize_cols(prev_field_exprs, &input).context(DataFusionPlanningSnafu)?;
693
694                let keep_tsid = op.id() != token::T_COUNT_VALUES
695                    && input_has_tsid
696                    && input_tag_columns.iter().collect::<HashSet<_>>()
697                        == self.ctx.tag_columns.iter().collect::<HashSet<_>>();
698
699                if keep_tsid {
700                    aggr_exprs.push(
701                        first_value(
702                            DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME)),
703                            vec![],
704                        )
705                        .alias(DATA_SCHEMA_TSID_COLUMN_NAME),
706                    );
707                }
708                self.ctx.use_tsid = keep_tsid;
709
710                // create plan
711                let builder = LogicalPlanBuilder::from(input);
712                let builder = if op.id() == token::T_COUNT_VALUES {
713                    let label = Self::get_param_value_as_str(*op, param)?;
714                    // `count_values` must be grouped by fields,
715                    // and project the fields to the new label.
716                    //
717                    // The generated label is a real label (column) of the output, so it must be
718                    // registered in `ctx.tag_columns` below. Otherwise enclosing expressions
719                    // rebuild their projections from `ctx.tag_columns` and silently drop it.
720                    //
721                    // PromQL sets the generated label *before* the grouping key is built, so an
722                    // input label with the same name is overwritten by the sample value and must
723                    // not remain a grouping key either: samples are grouped by the generated
724                    // label only. Dropping it from the projected tag columns is required as well:
725                    // projecting both would emit two columns with the same name (rejected as an
726                    // ambiguous reference).
727                    self.ctx.tag_columns.retain(|tag| tag != label);
728                    group_exprs.retain(
729                        |expr| !matches!(expr, DfExpr::Column(column) if column.name == label),
730                    );
731                    // The tag columns projected below are unqualified `Column` references, so they
732                    // end up qualified with whatever qualifier they carry in the input plan. Give
733                    // the generated label the same qualifier, otherwise qualified references to it
734                    // (e.g. from an enclosing binary expression or a vector join) fail to resolve.
735                    let label_qualifier = self
736                        .ctx
737                        .time_index_column
738                        .as_deref()
739                        .and_then(|time_index| {
740                            builder
741                                .schema()
742                                .qualified_field_with_unqualified_name(time_index)
743                                .ok()
744                        })
745                        .and_then(|(qualifier, _)| qualifier.cloned());
746                    // The generated label carries the *sample value*, so it must be materialized
747                    // as a string using Prometheus' format
748                    // (`strconv.FormatFloat(value, 'f', -1, 64)`: shortest decimal form without
749                    // an exponent, `1.0` becomes "1"). Labels are inferred from string columns
750                    // when a result is converted into the Prometheus HTTP API JSON format, so a
751                    // numeric label column would be mistaken for the sample value of the series.
752                    let count_value_exprs = prev_field_exprs
753                        .iter()
754                        .map(|expr| {
755                            let value = match expr {
756                                DfExpr::Column(column) => {
757                                    let value = DfExpr::Column(column.clone());
758                                    // `prom_float_to_string` formats exactly like Prometheus,
759                                    // while arrow's `Float64 -> Utf8` cast would render `1.0`.
760                                    let value = if matches!(
761                                        builder.schema().field_with_unqualified_name(&column.name),
762                                        Ok(field) if field.data_type() == &ArrowDataType::Float64
763                                    ) {
764                                        value
765                                    } else {
766                                        DfExpr::Cast(Cast::new(
767                                            Box::new(value),
768                                            ArrowDataType::Float64,
769                                        ))
770                                    };
771                                    DfExpr::ScalarFunction(ScalarFunction {
772                                        func: Arc::new(PromqlFloatToString::scalar_udf()),
773                                        args: vec![value],
774                                    })
775                                }
776                                // The value is already formatted (e.g. a native histogram is
777                                // converted to its string form by the aggregation), and the
778                                // aggregate output names it by its schema name.
779                                _ => DfExpr::Column(Column::from_name(
780                                    expr.schema_name().to_string(),
781                                )),
782                            };
783                            DfExpr::Alias(Alias::new(value, label_qualifier.clone(), label))
784                        })
785                        .collect::<Vec<_>>();
786                    let aggregate_group_exprs = group_exprs
787                        .iter()
788                        .cloned()
789                        .chain(prev_field_exprs.clone())
790                        .collect::<Vec<_>>();
791                    group_exprs.push(col(label));
792                    let project_fields = self
793                        .create_field_column_exprs()?
794                        .into_iter()
795                        .chain(self.create_tag_column_exprs()?)
796                        .chain(Some(self.create_time_index_column_expr()?))
797                        .chain(count_value_exprs);
798
799                    let builder = builder
800                        .aggregate(aggregate_group_exprs, aggr_exprs)
801                        .context(DataFusionPlanningSnafu)?
802                        .project(project_fields)
803                        .context(DataFusionPlanningSnafu)?;
804                    // The label only exists in the output schema from here on.
805                    self.ctx.tag_columns.push(label.to_string());
806                    builder
807                } else {
808                    builder
809                        .aggregate(group_exprs.clone(), aggr_exprs)
810                        .context(DataFusionPlanningSnafu)?
811                };
812
813                let builder = if let Some((float, histogram)) = mixed_sample_columns {
814                    let builder = match op.id() {
815                        token::T_SUM | token::T_AVG => builder
816                            .filter(self.mixed_aggregate_filter_expr(*op, &float, &histogram)?)
817                            .context(DataFusionPlanningSnafu)?,
818                        token::T_MIN
819                        | token::T_MAX
820                        | token::T_STDDEV
821                        | token::T_STDVAR
822                        | token::T_QUANTILE => builder
823                            .filter(self.mixed_ignored_histogram_filter_expr(*op, &histogram)?)
824                            .context(DataFusionPlanningSnafu)?,
825                        _ => builder,
826                    };
827
828                    match op.id() {
829                        token::T_SUM
830                        | token::T_AVG
831                        | token::T_MIN
832                        | token::T_MAX
833                        | token::T_STDDEV
834                        | token::T_STDVAR
835                        | token::T_QUANTILE => {
836                            let project_fields = self
837                                .create_field_column_exprs()?
838                                .into_iter()
839                                .chain(self.create_tag_column_exprs()?)
840                                .chain(self.ctx.use_tsid.then_some(DfExpr::Column(
841                                    Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME),
842                                )))
843                                .chain(Some(self.create_time_index_column_expr()?));
844                            builder
845                                .project(project_fields)
846                                .context(DataFusionPlanningSnafu)?
847                        }
848                        _ => builder,
849                    }
850                } else {
851                    builder
852                };
853
854                // Drop group rows whose every aggregated sample was discarded (NULL), so that
855                // e.g. `group(min(native_histogram))` doesn't resurrect groups Prometheus
856                // considers unseen. For alternative float/histogram fields keep the row if any
857                // field survived.
858                let builder = if has_native_histogram {
859                    builder
860                        .filter(self.create_empty_values_filter_expr(preserve_any_value)?)
861                        .context(DataFusionPlanningSnafu)?
862                } else {
863                    builder
864                };
865
866                let sort_expr = group_exprs.into_iter().map(|expr| expr.sort(true, false));
867
868                builder
869                    .sort(sort_expr)
870                    .context(DataFusionPlanningSnafu)?
871                    .build()
872                    .context(DataFusionPlanningSnafu)
873            }
874        }
875    }
876
877    /// Create logical plan for PromQL topk and bottomk expr.
878    async fn prom_topk_bottomk_to_plan(
879        &mut self,
880        aggr_expr: &AggregateExpr,
881        input: LogicalPlan,
882    ) -> Result<LogicalPlan> {
883        let AggregateExpr {
884            op,
885            param,
886            modifier,
887            ..
888        } = aggr_expr;
889
890        let input_has_tsid = input.schema().fields().iter().any(|field| {
891            field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
892                && field.data_type() == &ArrowDataType::UInt64
893        });
894        self.ctx.use_tsid = input_has_tsid;
895
896        let group_exprs = self.agg_modifier_to_col(input.schema(), modifier, false)?;
897
898        let mut input = input;
899        if let Some((float_column, histogram_column)) =
900            Self::alternative_sample_columns(input.schema(), &self.ctx.field_columns)
901                .map(|(float, histogram)| (float.to_string(), histogram.to_string()))
902        {
903            let drop_histogram = DfExpr::ScalarFunction(ScalarFunction {
904                func: Arc::new(NativeHistogramDrop::bool_false_udf(
905                    format!(
906                        "{}: dropped native histogram samples because this aggregation is not supported for native histograms",
907                        op
908                    ),
909                    self.promql_annotations.clone(),
910                )),
911                args: vec![col(&histogram_column)],
912            });
913            let keep_float = when(col(&histogram_column).is_not_null(), drop_histogram)
914                .otherwise(col(&float_column).is_not_null())
915                .context(DataFusionPlanningSnafu)?;
916            input = LogicalPlanBuilder::from(input)
917                .filter(keep_float)
918                .context(DataFusionPlanningSnafu)?
919                .build()
920                .context(DataFusionPlanningSnafu)?;
921            self.ctx.field_columns = vec![float_column];
922        }
923
924        if self.all_field_columns_are_native_histograms(input.schema()) {
925            let promql_annotations = self.promql_annotations.clone();
926            let input = self.projection_for_each_field_column(input, |col| {
927                Ok(DfExpr::ScalarFunction(ScalarFunction {
928                    func: Arc::new(NativeHistogramDrop::float_null_udf(
929                        format!(
930                            "{}: dropped native histogram samples because this aggregation is not supported for native histograms",
931                            op
932                        ),
933                        promql_annotations.clone(),
934                    )),
935                    args: vec![DfExpr::Column(Column::from_name(col))],
936                }))
937            })?;
938            return LogicalPlanBuilder::from(input)
939                .filter(self.create_empty_values_filter_expr(false)?)
940                .context(DataFusionPlanningSnafu)?
941                .build()
942                .context(DataFusionPlanningSnafu);
943        }
944
945        let val = Self::get_param_as_literal_expr(
946            param.as_deref(),
947            Some(*op),
948            Some(ArrowDataType::Float64),
949        )?;
950
951        // convert op and value columns to window exprs.
952        let window_exprs = self.create_window_exprs(*op, group_exprs.clone(), &input)?;
953
954        let rank_columns: Vec<_> = window_exprs
955            .iter()
956            .map(|expr| expr.schema_name().to_string())
957            .collect();
958
959        // Create ranks filter with `Operator::Or`.
960        // Safety: at least one rank column
961        let filter: DfExpr = rank_columns
962            .iter()
963            .fold(None, |expr, rank| {
964                let predicate = DfExpr::BinaryExpr(BinaryExpr {
965                    left: Box::new(col(rank)),
966                    op: Operator::LtEq,
967                    right: Box::new(val.clone()),
968                });
969
970                match expr {
971                    None => Some(predicate),
972                    Some(expr) => Some(DfExpr::BinaryExpr(BinaryExpr {
973                        left: Box::new(expr),
974                        op: Operator::Or,
975                        right: Box::new(predicate),
976                    })),
977                }
978            })
979            .unwrap();
980
981        let rank_columns: Vec<_> = rank_columns.into_iter().map(col).collect();
982
983        let mut new_group_exprs = group_exprs.clone();
984        // Order by ranks
985        new_group_exprs.extend(rank_columns);
986
987        let group_sort_expr = new_group_exprs
988            .into_iter()
989            .map(|expr| expr.sort(true, false));
990
991        let project_fields = self
992            .create_field_column_exprs()?
993            .into_iter()
994            .chain(self.create_tag_column_exprs()?)
995            .chain(
996                self.ctx
997                    .use_tsid
998                    .then_some(DfExpr::Column(Column::from_name(
999                        DATA_SCHEMA_TSID_COLUMN_NAME,
1000                    ))),
1001            )
1002            .chain(Some(self.create_time_index_column_expr()?));
1003
1004        LogicalPlanBuilder::from(input)
1005            .window(window_exprs)
1006            .context(DataFusionPlanningSnafu)?
1007            .filter(filter)
1008            .context(DataFusionPlanningSnafu)?
1009            .sort(group_sort_expr)
1010            .context(DataFusionPlanningSnafu)?
1011            .project(project_fields)
1012            .context(DataFusionPlanningSnafu)?
1013            .build()
1014            .context(DataFusionPlanningSnafu)
1015    }
1016
1017    async fn prom_unary_expr_to_plan(
1018        &mut self,
1019        query_engine_state: &QueryEngineState,
1020        unary_expr: &UnaryExpr,
1021    ) -> Result<LogicalPlan> {
1022        let UnaryExpr { expr } = unary_expr;
1023        // Unary Expr in PromQL implys the `-` operator
1024        let input = self.prom_expr_to_plan(expr, query_engine_state).await?;
1025        self.negate_field_columns(input)
1026    }
1027
1028    fn negate_field_columns(&mut self, input: LogicalPlan) -> Result<LogicalPlan> {
1029        let input_schema = input.schema().clone();
1030        self.projection_for_each_field_column(input, |col| {
1031            if Self::field_column_is_native_histogram(&input_schema, col) {
1032                Ok(DfExpr::ScalarFunction(ScalarFunction {
1033                    func: Arc::new(NativeHistogramNeg::scalar_udf()),
1034                    args: vec![DfExpr::Column(col.into())],
1035                }))
1036            } else {
1037                Ok(DfExpr::Negative(Box::new(DfExpr::Column(col.into()))))
1038            }
1039        })
1040    }
1041
1042    async fn try_plan_binary_island(
1043        &mut self,
1044        binary_expr: &PromBinaryExpr,
1045    ) -> Result<Option<LogicalPlan>> {
1046        let original_ctx = self.ctx.clone();
1047        let mut collect_env = IslandCollectEnv::default();
1048        let Some(island_expr) =
1049            IslandExpr::try_new(&PromExpr::Binary(binary_expr.clone()), &mut collect_env)
1050        else {
1051            return Ok(None);
1052        };
1053
1054        if collect_env.leaves.is_empty()
1055            || collect_env.vector_occurrences <= collect_env.leaves.len()
1056        {
1057            return Ok(None);
1058        }
1059
1060        let mut planned_leaves = Vec::with_capacity(collect_env.leaves.len());
1061        for (idx, leaf) in collect_env.leaves.iter().enumerate() {
1062            let plan = self
1063                .prom_vector_selector_to_plan(&leaf.selector, false)
1064                .await?;
1065            let ctx = self.ctx.clone();
1066            let alias = TableReference::bare(format!("{BINARY_ISLAND_LEAF_ALIAS_PREFIX}{idx}"));
1067            let plan = LogicalPlanBuilder::from(plan)
1068                .alias(alias.clone())
1069                .context(DataFusionPlanningSnafu)?
1070                .build()
1071                .context(DataFusionPlanningSnafu)?;
1072            planned_leaves.push(PlannedIslandLeaf {
1073                plan,
1074                ctx,
1075                alias,
1076                display_table: leaf.display_table.clone(),
1077            });
1078        }
1079
1080        if planned_leaves.iter().any(|leaf| {
1081            Self::field_columns_contain_native_histogram(
1082                leaf.plan.schema(),
1083                &leaf.ctx.field_columns,
1084            )
1085        }) {
1086            self.ctx = original_ctx;
1087            return Ok(None);
1088        }
1089
1090        if !Self::binary_island_join_contexts_supported(&planned_leaves) {
1091            self.ctx = original_ctx;
1092            return Ok(None);
1093        }
1094
1095        let mut input = planned_leaves[0].plan.clone();
1096        for right_idx in 1..planned_leaves.len() {
1097            input = self.join_binary_island_leaf(
1098                input,
1099                &planned_leaves[0],
1100                &planned_leaves[right_idx],
1101            )?;
1102        }
1103
1104        let field_exprs =
1105            Self::build_binary_island_field_exprs(&island_expr, &planned_leaves, input.schema())?;
1106        if field_exprs.scalar || field_exprs.exprs.is_empty() {
1107            self.ctx = original_ctx;
1108            return Ok(None);
1109        }
1110
1111        let plan = self.project_binary_island(
1112            input,
1113            &planned_leaves[0].alias,
1114            &planned_leaves[0].ctx,
1115            field_exprs,
1116        )?;
1117        Ok(Some(plan))
1118    }
1119
1120    fn binary_island_join_contexts_supported(leaves: &[PlannedIslandLeaf]) -> bool {
1121        if leaves
1122            .iter()
1123            .any(|leaf| leaf.ctx.time_index_column.is_none())
1124        {
1125            return false;
1126        }
1127
1128        if leaves.len() <= 1 {
1129            return true;
1130        }
1131
1132        let first_tags = leaves[0].ctx.tag_columns.iter().collect::<BTreeSet<_>>();
1133
1134        leaves.iter().skip(1).all(|leaf| {
1135            (Self::plan_has_tsid_column(&leaves[0].plan) && Self::plan_has_tsid_column(&leaf.plan))
1136                || leaf.ctx.tag_columns.iter().collect::<BTreeSet<_>>() == first_tags
1137        })
1138    }
1139
1140    fn join_binary_island_leaf(
1141        &self,
1142        left: LogicalPlan,
1143        first_leaf: &PlannedIslandLeaf,
1144        right_leaf: &PlannedIslandLeaf,
1145    ) -> Result<LogicalPlan> {
1146        let only_join_time_index = (first_leaf.ctx.tag_columns.is_empty()
1147            || right_leaf.ctx.tag_columns.is_empty())
1148            && !first_leaf
1149                .ctx
1150                .tag_columns
1151                .iter()
1152                .chain(&right_leaf.ctx.tag_columns)
1153                .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL);
1154        let (mut left_keys, mut right_keys, force_empty_join) = self.binary_join_key_columns(
1155            left.schema(),
1156            right_leaf.plan.schema(),
1157            &first_leaf.ctx,
1158            &right_leaf.ctx,
1159            only_join_time_index,
1160            &None,
1161        )?;
1162
1163        if let (Some(left_time_index_column), Some(right_time_index_column)) = (
1164            first_leaf.ctx.time_index_column.clone(),
1165            right_leaf.ctx.time_index_column.clone(),
1166        ) {
1167            left_keys.insert(left_time_index_column);
1168            right_keys.insert(right_time_index_column);
1169        }
1170
1171        LogicalPlanBuilder::from(left)
1172            .join_detailed(
1173                right_leaf.plan.clone(),
1174                JoinType::Inner,
1175                (
1176                    left_keys
1177                        .into_iter()
1178                        .map(|name| Column::new(Some(first_leaf.alias.clone()), name))
1179                        .collect::<Vec<_>>(),
1180                    right_keys
1181                        .into_iter()
1182                        .map(|name| Column::new(Some(right_leaf.alias.clone()), name))
1183                        .collect::<Vec<_>>(),
1184                ),
1185                force_empty_join.then_some(lit(false)),
1186                NullEquality::NullEqualsNull,
1187            )
1188            .context(DataFusionPlanningSnafu)?
1189            .build()
1190            .context(DataFusionPlanningSnafu)
1191    }
1192
1193    fn build_binary_island_field_exprs(
1194        expr: &IslandExpr,
1195        leaves: &[PlannedIslandLeaf],
1196        schema: &DFSchemaRef,
1197    ) -> Result<IslandFieldExprs> {
1198        match expr {
1199            IslandExpr::VectorLeaf(id) => {
1200                let leaf = &leaves[*id];
1201                let exprs = leaf
1202                    .ctx
1203                    .field_columns
1204                    .iter()
1205                    .map(|field| {
1206                        schema
1207                            .qualified_field_with_name(Some(&leaf.alias), field)
1208                            .context(DataFusionPlanningSnafu)
1209                            .map(|field| DfExpr::Column(field.into()))
1210                    })
1211                    .collect::<Result<Vec<_>>>()?;
1212                let names = leaf
1213                    .ctx
1214                    .field_columns
1215                    .iter()
1216                    .map(|field| format!("{}.{}", leaf.display_table, field))
1217                    .collect();
1218                Ok(IslandFieldExprs {
1219                    exprs,
1220                    names,
1221                    scalar: false,
1222                })
1223            }
1224            IslandExpr::Scalar(expr) => Ok(IslandFieldExprs {
1225                exprs: vec![expr.clone()],
1226                names: vec![expr.schema_name().to_string()],
1227                scalar: true,
1228            }),
1229            IslandExpr::Unary { input } => {
1230                let input = Self::build_binary_island_field_exprs(input, leaves, schema)?;
1231                let mut exprs = Vec::with_capacity(input.exprs.len());
1232                let mut names = Vec::with_capacity(input.names.len());
1233                for (expr, name) in input.exprs.into_iter().zip(input.names) {
1234                    exprs.push(DfExpr::Negative(Box::new(expr)));
1235                    names.push(format!("-{name}"));
1236                }
1237                Ok(IslandFieldExprs {
1238                    exprs,
1239                    names,
1240                    scalar: input.scalar,
1241                })
1242            }
1243            IslandExpr::Binary { op, lhs, rhs } => {
1244                let same_leaf = match (&**lhs, &**rhs) {
1245                    (IslandExpr::VectorLeaf(left), IslandExpr::VectorLeaf(right))
1246                        if left == right =>
1247                    {
1248                        Some(*left)
1249                    }
1250                    _ => None,
1251                };
1252                let lhs = Self::build_binary_island_field_exprs(lhs, leaves, schema)?;
1253                let rhs = Self::build_binary_island_field_exprs(rhs, leaves, schema)?;
1254                let expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1255                let scalar = lhs.scalar && rhs.scalar;
1256                let op = op.to_string();
1257
1258                let (exprs, names) = match (lhs.scalar, rhs.scalar) {
1259                    (true, true) => {
1260                        let expr = expr_builder(lhs.exprs[0].clone(), rhs.exprs[0].clone())?;
1261                        let name = format!("{} {op} {}", lhs.names[0], rhs.names[0]);
1262                        (vec![expr], vec![name])
1263                    }
1264                    (true, false) => {
1265                        let mut exprs = Vec::with_capacity(rhs.exprs.len());
1266                        let mut names = Vec::with_capacity(rhs.names.len());
1267                        for (rhs_expr, rhs_name) in rhs.exprs.into_iter().zip(rhs.names) {
1268                            exprs.push(expr_builder(lhs.exprs[0].clone(), rhs_expr)?);
1269                            names.push(format!("{} {op} {rhs_name}", lhs.names[0]));
1270                        }
1271                        (exprs, names)
1272                    }
1273                    (false, true) => {
1274                        let mut exprs = Vec::with_capacity(lhs.exprs.len());
1275                        let mut names = Vec::with_capacity(lhs.names.len());
1276                        for (lhs_expr, lhs_name) in lhs.exprs.into_iter().zip(lhs.names) {
1277                            exprs.push(expr_builder(lhs_expr, rhs.exprs[0].clone())?);
1278                            names.push(format!("{lhs_name} {op} {}", rhs.names[0]));
1279                        }
1280                        (exprs, names)
1281                    }
1282                    (false, false) => {
1283                        let mut exprs = Vec::new();
1284                        let mut names = Vec::new();
1285                        for (idx, ((lhs_expr, rhs_expr), (mut lhs_name, mut rhs_name))) in lhs
1286                            .exprs
1287                            .into_iter()
1288                            .zip(rhs.exprs)
1289                            .zip(lhs.names.into_iter().zip(rhs.names))
1290                            .enumerate()
1291                        {
1292                            if let Some(leaf) = same_leaf {
1293                                let field = leaves[leaf]
1294                                    .ctx
1295                                    .field_columns
1296                                    .get(idx)
1297                                    .cloned()
1298                                    .unwrap_or_else(|| lhs_name.clone());
1299                                lhs_name = format!("lhs.{field}");
1300                                rhs_name = format!("rhs.{field}");
1301                            }
1302                            exprs.push(expr_builder(lhs_expr, rhs_expr)?);
1303                            names.push(format!("{lhs_name} {op} {rhs_name}"));
1304                        }
1305                        (exprs, names)
1306                    }
1307                };
1308
1309                Ok(IslandFieldExprs {
1310                    exprs,
1311                    names,
1312                    scalar,
1313                })
1314            }
1315        }
1316    }
1317
1318    fn project_binary_island(
1319        &mut self,
1320        input: LogicalPlan,
1321        base_alias: &TableReference,
1322        base_ctx: &PromPlannerContext,
1323        field_exprs: IslandFieldExprs,
1324    ) -> Result<LogicalPlan> {
1325        self.ctx = base_ctx.clone();
1326
1327        let schema = input.schema();
1328        let non_field_exprs = base_ctx
1329            .tag_columns
1330            .iter()
1331            .chain(base_ctx.time_index_column.iter())
1332            .map(|column| {
1333                schema
1334                    .qualified_field_with_name(Some(base_alias), column)
1335                    .context(DataFusionPlanningSnafu)
1336                    .map(|field| DfExpr::Column(field.into()))
1337            });
1338        let tsid_expr = Self::optional_tsid_projection(schema, Some(base_alias), base_ctx.use_tsid)
1339            .into_iter()
1340            .map(Ok);
1341
1342        self.ctx.field_columns = field_exprs.names;
1343        let field_exprs = field_exprs
1344            .exprs
1345            .into_iter()
1346            .zip(self.ctx.field_columns.iter())
1347            .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::<String>, name))));
1348
1349        let project_exprs = non_field_exprs
1350            .chain(tsid_expr)
1351            .chain(field_exprs)
1352            .collect::<Result<Vec<_>>>()?;
1353
1354        let plan = LogicalPlanBuilder::from(input)
1355            .project(project_exprs)
1356            .context(DataFusionPlanningSnafu)?
1357            .build()
1358            .context(DataFusionPlanningSnafu)?;
1359
1360        self.ctx.table_name = None;
1361        self.ctx.schema_name = None;
1362
1363        Ok(plan)
1364    }
1365
1366    async fn prom_binary_expr_to_plan(
1367        &mut self,
1368        query_engine_state: &QueryEngineState,
1369        binary_expr: &PromBinaryExpr,
1370    ) -> Result<LogicalPlan> {
1371        // promql-parser accepts fill modifiers, but Greptime does not implement the
1372        // required outer joins and missing-value substitution. Reject them before the
1373        // binary-island fast path so they cannot silently behave like normal inner joins.
1374        if let Some(modifier) = &binary_expr.modifier {
1375            ensure!(
1376                modifier.fill_values.lhs.is_none() && modifier.fill_values.rhs.is_none(),
1377                UnsupportedExprSnafu {
1378                    name: "PromQL fill modifiers"
1379                }
1380            );
1381        }
1382
1383        if let Some(plan) = self.try_plan_binary_island(binary_expr).await? {
1384            return Ok(plan);
1385        }
1386
1387        let PromBinaryExpr {
1388            lhs,
1389            rhs,
1390            op,
1391            modifier,
1392        } = binary_expr;
1393
1394        // if set to true, comparison operator will return 0/1 (for true/false) instead of
1395        // filter on the result column
1396        let should_return_bool = if let Some(m) = modifier {
1397            m.return_bool
1398        } else {
1399            false
1400        };
1401        let is_comparison_op = Self::is_token_a_comparison_op(*op);
1402
1403        // we should build a filter plan here if the op is comparison op and need not
1404        // to return 0/1. Otherwise, we should build a projection plan
1405        match (
1406            Self::try_build_literal_expr(lhs),
1407            Self::try_build_literal_expr(rhs),
1408        ) {
1409            (Some(lhs), Some(rhs)) => {
1410                self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
1411                self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1412                self.ctx.reset_table_name_and_schema();
1413                let field_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1414                let mut field_expr = field_expr_builder(lhs, rhs)?;
1415
1416                if is_comparison_op && should_return_bool {
1417                    field_expr =
1418                        DfExpr::Cast(Cast::new(Box::new(field_expr), ArrowDataType::Float64));
1419                }
1420
1421                Ok(LogicalPlan::Extension(Extension {
1422                    node: Arc::new(
1423                        EmptyMetric::new(
1424                            self.ctx.start,
1425                            self.ctx.end,
1426                            self.ctx.interval,
1427                            SPECIAL_TIME_FUNCTION.to_string(),
1428                            DEFAULT_FIELD_COLUMN.to_string(),
1429                            Some(field_expr),
1430                        )
1431                        .context(DataFusionPlanningSnafu)?,
1432                    ),
1433                }))
1434            }
1435            // lhs is a literal, rhs is a column
1436            (Some(mut expr), None) => {
1437                let input = self.prom_expr_to_plan(rhs, query_engine_state).await?;
1438                // Arithmetic against a literal preserves the series labels, so
1439                // `aggregation_field_labels` passes through with `tag_columns` unchanged.
1440                // check if the literal is a special time expr
1441                if let Some(time_expr) = self.try_build_special_time_expr_with_context(lhs) {
1442                    expr = time_expr
1443                }
1444                let input_schema = input.schema().clone();
1445                let preserve_any_value = Self::field_columns_are_alternative_samples(
1446                    &input_schema,
1447                    &self.ctx.field_columns,
1448                );
1449                let has_native_histogram = Self::field_columns_contain_native_histogram(
1450                    &input_schema,
1451                    &self.ctx.field_columns,
1452                );
1453                let retain_field_columns = self
1454                    .ctx
1455                    .field_columns
1456                    .iter()
1457                    .map(|col| {
1458                        Self::binary_result_is_histogram(
1459                            *op,
1460                            false,
1461                            Self::field_column_is_native_histogram(&input_schema, col),
1462                        )
1463                        .is_some()
1464                    })
1465                    .collect();
1466                let promql_annotations = self.promql_annotations.clone();
1467                let bin_expr_builder = |col: &String| {
1468                    let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1469                    let rhs_is_histogram =
1470                        Self::field_column_is_native_histogram(&input_schema, col);
1471                    let rhs = DfExpr::Column(col.into());
1472                    let mut binary_expr = match Self::native_histogram_binary_expr(
1473                        *op,
1474                        expr.clone(),
1475                        false,
1476                        rhs.clone(),
1477                        rhs_is_histogram,
1478                        is_comparison_op && !should_return_bool,
1479                        promql_annotations.clone(),
1480                    )? {
1481                        Some(expr) => expr,
1482                        None => binary_expr_builder(expr.clone(), rhs)?,
1483                    };
1484
1485                    if is_comparison_op && should_return_bool {
1486                        binary_expr =
1487                            DfExpr::Cast(Cast::new(Box::new(binary_expr), ArrowDataType::Float64));
1488                    }
1489                    Ok(binary_expr)
1490                };
1491                if is_comparison_op && !should_return_bool {
1492                    self.filter_on_field_column(input, bin_expr_builder)
1493                } else {
1494                    let projected =
1495                        self.projection_for_each_field_column(input, bin_expr_builder)?;
1496                    self.filter_binary_projection(
1497                        projected,
1498                        has_native_histogram,
1499                        preserve_any_value,
1500                        retain_field_columns,
1501                    )
1502                }
1503            }
1504            // lhs is a column, rhs is a literal
1505            (None, Some(mut expr)) => {
1506                let input = self.prom_expr_to_plan(lhs, query_engine_state).await?;
1507                // check if the literal is a special time expr
1508                if let Some(time_expr) = self.try_build_special_time_expr_with_context(rhs) {
1509                    expr = time_expr
1510                }
1511                let input_schema = input.schema().clone();
1512                let preserve_any_value = Self::field_columns_are_alternative_samples(
1513                    &input_schema,
1514                    &self.ctx.field_columns,
1515                );
1516                let has_native_histogram = Self::field_columns_contain_native_histogram(
1517                    &input_schema,
1518                    &self.ctx.field_columns,
1519                );
1520                let retain_field_columns = self
1521                    .ctx
1522                    .field_columns
1523                    .iter()
1524                    .map(|col| {
1525                        Self::binary_result_is_histogram(
1526                            *op,
1527                            Self::field_column_is_native_histogram(&input_schema, col),
1528                            false,
1529                        )
1530                        .is_some()
1531                    })
1532                    .collect();
1533                let promql_annotations = self.promql_annotations.clone();
1534                let bin_expr_builder = |col: &String| {
1535                    let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1536                    let lhs_is_histogram =
1537                        Self::field_column_is_native_histogram(&input_schema, col);
1538                    let lhs = DfExpr::Column(col.into());
1539                    let mut binary_expr = match Self::native_histogram_binary_expr(
1540                        *op,
1541                        lhs.clone(),
1542                        lhs_is_histogram,
1543                        expr.clone(),
1544                        false,
1545                        is_comparison_op && !should_return_bool,
1546                        promql_annotations.clone(),
1547                    )? {
1548                        Some(expr) => expr,
1549                        None => binary_expr_builder(lhs, expr.clone())?,
1550                    };
1551
1552                    if is_comparison_op && should_return_bool {
1553                        binary_expr =
1554                            DfExpr::Cast(Cast::new(Box::new(binary_expr), ArrowDataType::Float64));
1555                    }
1556                    Ok(binary_expr)
1557                };
1558                if is_comparison_op && !should_return_bool {
1559                    self.filter_on_field_column(input, bin_expr_builder)
1560                } else {
1561                    let projected =
1562                        self.projection_for_each_field_column(input, bin_expr_builder)?;
1563                    self.filter_binary_projection(
1564                        projected,
1565                        has_native_histogram,
1566                        preserve_any_value,
1567                        retain_field_columns,
1568                    )
1569                }
1570            }
1571            // both are columns. join them on time index
1572            (None, None) => {
1573                let mut left_input = self.prom_expr_to_plan(lhs, query_engine_state).await?;
1574                let left_field_columns = self.ctx.field_columns.clone();
1575                let left_time_index_column = self.ctx.time_index_column.clone();
1576                let mut left_table_ref = self
1577                    .table_ref()
1578                    .unwrap_or_else(|_| TableReference::bare(""));
1579                let mut left_context = self.ctx.clone();
1580
1581                let mut right_input = self.prom_expr_to_plan(rhs, query_engine_state).await?;
1582                let right_field_columns = self.ctx.field_columns.clone();
1583                let right_time_index_column = self.ctx.time_index_column.clone();
1584                let mut right_table_ref = self
1585                    .table_ref()
1586                    .unwrap_or_else(|_| TableReference::bare(""));
1587                let mut right_context = self.ctx.clone();
1588
1589                // Both operands are planned first because only the planned contexts tell a tag
1590                // from a value field. The rewrite merely adds matchers to a selector, so the
1591                // table reference, time index and field columns captured above stay valid.
1592                if let Some(rewritten) = matching_filters::propagate(
1593                    binary_expr,
1594                    &left_context.tag_columns,
1595                    &right_context.tag_columns,
1596                    &left_context.aggregation_field_labels,
1597                    &right_context.aggregation_field_labels,
1598                ) {
1599                    // A copied matcher belongs to the scan, not to the operand's identity:
1600                    // `absent()` turns `selector_matcher` into the labels it reports, so the
1601                    // re-planned context keeps the matchers the operand was written with.
1602                    if rewritten.lhs.as_ref() != lhs.as_ref() {
1603                        let selectors = std::mem::take(&mut left_context.selector_matcher);
1604                        left_input = self
1605                            .prom_expr_to_plan(&rewritten.lhs, query_engine_state)
1606                            .await?;
1607                        left_context = self.ctx.clone();
1608                        left_context.selector_matcher = selectors;
1609                    }
1610                    if rewritten.rhs.as_ref() != rhs.as_ref() {
1611                        let selectors = std::mem::take(&mut right_context.selector_matcher);
1612                        right_input = self
1613                            .prom_expr_to_plan(&rewritten.rhs, query_engine_state)
1614                            .await?;
1615                        right_context = self.ctx.clone();
1616                        right_context.selector_matcher = selectors;
1617                    }
1618                    // The code below reads `self.ctx` as the right operand's context.
1619                    self.ctx = right_context.clone();
1620                }
1621
1622                let left_is_empty_metric = Self::is_empty_metric(&left_input);
1623                let right_is_empty_metric = Self::is_empty_metric(&right_input);
1624
1625                // TODO(ruihang): avoid join if left and right are the same table
1626
1627                // set op has "special" join semantics
1628                if Self::is_token_a_set_op(*op) {
1629                    return self.set_op_on_non_field_columns(
1630                        left_input,
1631                        right_input,
1632                        left_context,
1633                        right_context,
1634                        *op,
1635                        modifier,
1636                    );
1637                }
1638
1639                let has_native_histogram = Self::field_columns_contain_native_histogram(
1640                    left_input.schema(),
1641                    &left_field_columns,
1642                ) || Self::field_columns_contain_native_histogram(
1643                    right_input.schema(),
1644                    &right_field_columns,
1645                );
1646
1647                // normal join
1648                if left_table_ref == right_table_ref {
1649                    // rename table references to avoid ambiguity
1650                    left_table_ref = TableReference::bare("lhs");
1651                    right_table_ref = TableReference::bare("rhs");
1652                    // `self.ctx` have ctx in right plan, if right plan have no tag,
1653                    // we use left plan ctx as the ctx for subsequent calculations,
1654                    // to avoid case like `host + scalar(...)`
1655                    // we need preserve tag column on `host` table in subsequent projection,
1656                    // which only show in left plan ctx.
1657                    if self.ctx.tag_columns.is_empty() {
1658                        self.ctx = left_context.clone();
1659                        self.ctx.table_name = Some("lhs".to_string());
1660                    } else {
1661                        self.ctx.table_name = Some("rhs".to_string());
1662                    }
1663                } else if right_is_empty_metric && !left_is_empty_metric {
1664                    self.ctx = left_context.clone();
1665                }
1666                // Computed scalars reach this join path instead of the literal projection paths.
1667                // Broadcast them for arithmetic in the same way as literal scalars.
1668                let broadcast_scalar = !is_comparison_op;
1669                let (field_groups, invalid_field_pairs) = Self::align_binary_field_columns(
1670                    left_input.schema(),
1671                    right_input.schema(),
1672                    &left_field_columns,
1673                    &right_field_columns,
1674                    *op,
1675                    broadcast_scalar && lhs.value_type() == ValueType::Scalar,
1676                    broadcast_scalar && rhs.value_type() == ValueType::Scalar,
1677                );
1678                let left_aligned_field_columns = field_groups
1679                    .iter()
1680                    .flat_map(|(_, pairs)| {
1681                        pairs
1682                            .iter()
1683                            .map(|(left_col_name, _)| (*left_col_name).clone())
1684                    })
1685                    .collect::<Vec<_>>();
1686                let right_aligned_field_columns = field_groups
1687                    .iter()
1688                    .flat_map(|(_, pairs)| {
1689                        pairs
1690                            .iter()
1691                            .map(|(_, right_col_name)| (*right_col_name).clone())
1692                    })
1693                    .collect::<Vec<_>>();
1694                // Regular multi-field vectors combine their shared prefix. Alternative
1695                // float/histogram lanes instead align by valid PromQL sample combinations.
1696                self.ctx.field_columns = field_groups
1697                    .iter()
1698                    .map(|(output, _)| output.clone())
1699                    .collect();
1700                let mut field_groups = field_groups.into_iter();
1701                // `vector()` uses EmptyMetric and keeps GreptimeDB's timestamp broadcast.
1702                let has_empty_metric_operand = left_is_empty_metric || right_is_empty_metric;
1703
1704                let only_join_time_index = lhs.value_type() == ValueType::Scalar
1705                    || rhs.value_type() == ValueType::Scalar
1706                    || has_empty_metric_operand
1707                    || ((left_context.tag_columns.is_empty()
1708                        || right_context.tag_columns.is_empty())
1709                        && !left_context
1710                            .tag_columns
1711                            .iter()
1712                            .chain(&right_context.tag_columns)
1713                            .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL));
1714                let join_plan = self.join_on_non_field_columns(
1715                    left_input,
1716                    right_input,
1717                    left_table_ref.clone(),
1718                    right_table_ref.clone(),
1719                    left_time_index_column,
1720                    right_time_index_column,
1721                    only_join_time_index,
1722                    modifier,
1723                    &left_context,
1724                    &right_context,
1725                )?;
1726                let join_plan_schema = join_plan.schema().clone();
1727                // The matching modifier derives the result labels instead of the operands keeping
1728                // their own tag set. An operand that is broadcast rather than matched (a scalar,
1729                // or a vector without tags) keeps the other side's labels as before.
1730                let result_labels = if only_join_time_index {
1731                    None
1732                } else {
1733                    Self::binary_result_labels(&left_context, &right_context, modifier)
1734                        .map(|labels| {
1735                            Self::binary_result_label_projection(
1736                                &join_plan_schema,
1737                                &left_table_ref,
1738                                &right_table_ref,
1739                                &left_context,
1740                                &right_context,
1741                                labels,
1742                            )
1743                        })
1744                        .transpose()?
1745                };
1746                if let Some(labels) = &result_labels {
1747                    labels.apply(&mut self.ctx);
1748                }
1749                let promql_annotations = self.promql_annotations.clone();
1750                // These predicates always pass; they only evaluate otherwise-discarded pairs
1751                // while collecting annotations.
1752                let invalid_pair_predicates = invalid_field_pairs
1753                    .into_iter()
1754                    .filter(|_| promql_annotations.is_some())
1755                    .map(|(left_col_name, right_col_name)| {
1756                        let left_field = join_plan_schema
1757                            .qualified_field_with_name(Some(&left_table_ref), left_col_name)
1758                            .context(DataFusionPlanningSnafu)?;
1759                        let right_field = join_plan_schema
1760                            .qualified_field_with_name(Some(&right_table_ref), right_col_name)
1761                            .context(DataFusionPlanningSnafu)?;
1762                        let left_is_histogram =
1763                            left_field.1.data_type() == &Self::native_histogram_arrow_type();
1764                        let right_is_histogram =
1765                            right_field.1.data_type() == &Self::native_histogram_arrow_type();
1766                        let drop_expr = Self::native_histogram_binary_expr(
1767                            *op,
1768                            DfExpr::Column(left_field.into()),
1769                            left_is_histogram,
1770                            DfExpr::Column(right_field.into()),
1771                            right_is_histogram,
1772                            true,
1773                            promql_annotations.clone(),
1774                        )?
1775                        .with_context(|| UnexpectedPlanExprSnafu {
1776                            desc: "invalid native histogram pair produced no drop expression",
1777                        })?;
1778                        Ok(DfExpr::Not(Box::new(drop_expr)))
1779                    })
1780                    .collect::<Result<Vec<_>>>()?;
1781                let join_plan = if let Some(predicate) = conjunction(invalid_pair_predicates) {
1782                    LogicalPlanBuilder::from(join_plan)
1783                        .filter(predicate)
1784                        .context(DataFusionPlanningSnafu)?
1785                        .build()
1786                        .context(DataFusionPlanningSnafu)?
1787                } else {
1788                    join_plan
1789                };
1790
1791                let bin_expr_builder = |_: &String| {
1792                    let (_, field_pairs) =
1793                        field_groups
1794                            .next()
1795                            .with_context(|| UnexpectedPlanExprSnafu {
1796                                desc: "missing binary field group",
1797                            })?;
1798                    let binary_exprs = field_pairs
1799                        .into_iter()
1800                        .map(|(left_col_name, right_col_name)| {
1801                            let left_field = join_plan_schema
1802                                .qualified_field_with_name(Some(&left_table_ref), left_col_name)
1803                                .context(DataFusionPlanningSnafu)?;
1804                            let right_field = join_plan_schema
1805                                .qualified_field_with_name(Some(&right_table_ref), right_col_name)
1806                                .context(DataFusionPlanningSnafu)?;
1807                            let left_is_histogram =
1808                                left_field.1.data_type() == &Self::native_histogram_arrow_type();
1809                            let right_is_histogram =
1810                                right_field.1.data_type() == &Self::native_histogram_arrow_type();
1811                            let left_col = left_field.into();
1812                            let right_col = right_field.into();
1813
1814                            let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1815                            let lhs = DfExpr::Column(left_col);
1816                            let rhs = DfExpr::Column(right_col);
1817                            let mut binary_expr = match Self::native_histogram_binary_expr(
1818                                *op,
1819                                lhs.clone(),
1820                                left_is_histogram,
1821                                rhs.clone(),
1822                                right_is_histogram,
1823                                is_comparison_op && !should_return_bool,
1824                                promql_annotations.clone(),
1825                            )? {
1826                                Some(expr) => expr,
1827                                None => binary_expr_builder(lhs, rhs)?,
1828                            };
1829                            if is_comparison_op && should_return_bool {
1830                                binary_expr = DfExpr::Cast(Cast::new(
1831                                    Box::new(binary_expr),
1832                                    ArrowDataType::Float64,
1833                                ));
1834                            }
1835                            Ok(binary_expr)
1836                        })
1837                        .collect::<Result<Vec<_>>>()?;
1838                    if let [binary_expr] = binary_exprs.as_slice() {
1839                        Ok(binary_expr.clone())
1840                    } else {
1841                        Ok(DfExpr::ScalarFunction(ScalarFunction {
1842                            func: coalesce(),
1843                            args: binary_exprs,
1844                        }))
1845                    }
1846                };
1847                if is_comparison_op && !should_return_bool {
1848                    // PromQL comparison operators without `bool` are filters:
1849                    //   - keep the instant-vector side sample values
1850                    //   - drop samples where the comparison is false
1851                    //
1852                    // So we filter on the join result and then project only the side that should
1853                    // be preserved according to PromQL semantics.
1854                    let filtered = self.filter_on_field_column(join_plan, bin_expr_builder)?;
1855                    let (project_table_ref, mut project_context, project_field_columns) =
1856                        match (lhs.value_type(), rhs.value_type()) {
1857                            (ValueType::Scalar, ValueType::Vector) => (
1858                                &right_table_ref,
1859                                right_context.clone(),
1860                                right_aligned_field_columns,
1861                            ),
1862                            _ => (
1863                                &left_table_ref,
1864                                left_context.clone(),
1865                                left_aligned_field_columns,
1866                            ),
1867                        };
1868                    project_context.field_columns = project_field_columns;
1869                    self.project_binary_join_side(
1870                        filtered,
1871                        project_table_ref,
1872                        &project_context,
1873                        result_labels.as_ref(),
1874                    )
1875                } else {
1876                    let projected = self.projection_for_each_field_column_with_labels(
1877                        join_plan,
1878                        result_labels.as_ref(),
1879                        bin_expr_builder,
1880                    )?;
1881                    let preserve_any_value = Self::field_columns_are_alternative_samples(
1882                        projected.schema(),
1883                        &self.ctx.field_columns,
1884                    );
1885                    let retain_field_columns = vec![true; self.ctx.field_columns.len()];
1886                    self.filter_binary_projection(
1887                        projected,
1888                        has_native_histogram,
1889                        preserve_any_value,
1890                        retain_field_columns,
1891                    )
1892                }
1893            }
1894        }
1895    }
1896
1897    fn filter_binary_projection(
1898        &mut self,
1899        input: LogicalPlan,
1900        has_native_histogram: bool,
1901        preserve_any_value: bool,
1902        retain_field_columns: Vec<bool>,
1903    ) -> Result<LogicalPlan> {
1904        if !has_native_histogram {
1905            return Ok(input);
1906        }
1907
1908        ensure!(
1909            retain_field_columns.len() == self.ctx.field_columns.len(),
1910            UnexpectedPlanExprSnafu {
1911                desc: "binary output field count changed unexpectedly",
1912            }
1913        );
1914
1915        let filtered = LogicalPlanBuilder::from(input)
1916            .filter(self.create_empty_values_filter_expr(preserve_any_value)?)
1917            .context(DataFusionPlanningSnafu)?
1918            .build()
1919            .context(DataFusionPlanningSnafu)?;
1920        if retain_field_columns.iter().all(|retain| *retain) {
1921            return Ok(filtered);
1922        }
1923
1924        let retained = self
1925            .ctx
1926            .field_columns
1927            .iter()
1928            .zip(retain_field_columns)
1929            .filter(|(_, retain)| *retain)
1930            .map(|(field, _)| field.clone())
1931            .collect::<Vec<_>>();
1932        if retained.is_empty() {
1933            return Ok(filtered);
1934        }
1935        self.ctx.field_columns = retained;
1936
1937        let mut output_columns = self
1938            .ctx
1939            .field_columns
1940            .iter()
1941            .chain(&self.ctx.tag_columns)
1942            .cloned()
1943            .collect::<HashSet<_>>();
1944        output_columns.extend(self.ctx.time_index_column.iter().cloned());
1945        if self.ctx.use_tsid {
1946            output_columns.insert(DATA_SCHEMA_TSID_COLUMN_NAME.to_string());
1947        }
1948        let project_exprs = filtered
1949            .schema()
1950            .iter()
1951            .filter(|(_, field)| output_columns.contains(field.name()))
1952            .map(|(qualifier, field)| {
1953                DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
1954            })
1955            .collect::<Vec<_>>();
1956        LogicalPlanBuilder::from(filtered)
1957            .project(project_exprs)
1958            .context(DataFusionPlanningSnafu)?
1959            .build()
1960            .context(DataFusionPlanningSnafu)
1961    }
1962
1963    fn project_binary_join_side(
1964        &mut self,
1965        input: LogicalPlan,
1966        table_ref: &TableReference,
1967        context: &PromPlannerContext,
1968        result_labels: Option<&BinaryResultLabels>,
1969    ) -> Result<LogicalPlan> {
1970        let schema = input.schema();
1971
1972        let mut project_exprs =
1973            Vec::with_capacity(context.tag_columns.len() + context.field_columns.len() + 2);
1974
1975        // Project time index from the chosen side.
1976        if let Some(time_index_column) = &context.time_index_column {
1977            let time_index_col = schema
1978                .qualified_field_with_name(Some(table_ref), time_index_column)
1979                .context(DataFusionPlanningSnafu)?
1980                .into();
1981            project_exprs.push(DfExpr::Column(time_index_col));
1982        }
1983
1984        // Project field columns from the chosen side.
1985        for field_column in &context.field_columns {
1986            let field_col = schema
1987                .qualified_field_with_name(Some(table_ref), field_column)
1988                .context(DataFusionPlanningSnafu)?
1989                .into();
1990            project_exprs.push(DfExpr::Column(field_col));
1991        }
1992
1993        // Project tag columns: the labels the matching modifier derived, or the chosen side's.
1994        match result_labels {
1995            Some(labels) => project_exprs.extend(labels.exprs.iter().cloned()),
1996            None => {
1997                for tag_column in &context.tag_columns {
1998                    let tag_col = schema
1999                        .qualified_field_with_name(Some(table_ref), tag_column)
2000                        .context(DataFusionPlanningSnafu)?
2001                        .into();
2002                    project_exprs.push(DfExpr::Column(tag_col));
2003                }
2004            }
2005        }
2006
2007        // Preserve `__tsid` if present, so it can still be used internally downstream. It's
2008        // stripped from the final output anyway.
2009        let tsid_col = match result_labels {
2010            Some(labels) => labels.tsid_projection(Some(table_ref.clone())),
2011            None => Self::optional_tsid_projection(schema, Some(table_ref), context.use_tsid),
2012        };
2013        if let Some(tsid_col) = tsid_col {
2014            project_exprs.push(tsid_col);
2015        }
2016
2017        let plan = LogicalPlanBuilder::from(input)
2018            .project(project_exprs)
2019            .context(DataFusionPlanningSnafu)?
2020            .build()
2021            .context(DataFusionPlanningSnafu)?;
2022
2023        // Update context to reflect the projected schema. Don't keep a table qualifier since
2024        // the result is a derived expression.
2025        self.ctx = context.clone();
2026        self.ctx.table_name = None;
2027        self.ctx.schema_name = None;
2028        if let Some(labels) = result_labels {
2029            labels.apply(&mut self.ctx);
2030        }
2031
2032        Ok(plan)
2033    }
2034
2035    fn prom_number_lit_to_plan(&mut self, number_literal: &NumberLiteral) -> Result<LogicalPlan> {
2036        let NumberLiteral { val } = number_literal;
2037        self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
2038        self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
2039        self.ctx.reset_table_name_and_schema();
2040        let literal_expr = df_prelude::lit(*val);
2041
2042        let plan = LogicalPlan::Extension(Extension {
2043            node: Arc::new(
2044                EmptyMetric::new(
2045                    self.ctx.start,
2046                    self.ctx.end,
2047                    self.ctx.interval,
2048                    SPECIAL_TIME_FUNCTION.to_string(),
2049                    DEFAULT_FIELD_COLUMN.to_string(),
2050                    Some(literal_expr),
2051                )
2052                .context(DataFusionPlanningSnafu)?,
2053            ),
2054        });
2055        Ok(plan)
2056    }
2057
2058    fn prom_string_lit_to_plan(&mut self, string_literal: &StringLiteral) -> Result<LogicalPlan> {
2059        let StringLiteral { val } = string_literal;
2060        self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
2061        self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
2062        self.ctx.reset_table_name_and_schema();
2063        let literal_expr = df_prelude::lit(val.clone());
2064
2065        let plan = LogicalPlan::Extension(Extension {
2066            node: Arc::new(
2067                EmptyMetric::new(
2068                    self.ctx.start,
2069                    self.ctx.end,
2070                    self.ctx.interval,
2071                    SPECIAL_TIME_FUNCTION.to_string(),
2072                    DEFAULT_FIELD_COLUMN.to_string(),
2073                    Some(literal_expr),
2074                )
2075                .context(DataFusionPlanningSnafu)?,
2076            ),
2077        });
2078        Ok(plan)
2079    }
2080
2081    async fn prom_vector_selector_to_plan(
2082        &mut self,
2083        vector_selector: &VectorSelector,
2084        timestamp_fn: bool,
2085    ) -> Result<LogicalPlan> {
2086        let VectorSelector {
2087            name,
2088            offset,
2089            matchers,
2090            at: _,
2091        } = vector_selector;
2092        let matchers = self.preprocess_label_matchers(matchers, name)?;
2093        if let Some(empty_plan) = self.setup_context().await? {
2094            return Ok(empty_plan);
2095        }
2096        let offset_ms = match offset {
2097            Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond,
2098            Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond),
2099            None => 0,
2100        };
2101        let normalize = self
2102            .selector_to_series_normalize_plan(offset, matchers, false)
2103            .await?;
2104        let time_index_column =
2105            self.ctx
2106                .time_index_column
2107                .clone()
2108                .with_context(|| TimeIndexNotFoundSnafu {
2109                    table: self.ctx.table_name.clone().unwrap_or_default(),
2110                })?;
2111
2112        let (normalize, timestamp_value_column) = if timestamp_fn {
2113            // Keep the original sample for stale-marker detection while carrying
2114            // its timestamp through InstantManipulate in a private value column.
2115            let occupied = normalize
2116                .schema()
2117                .fields()
2118                .iter()
2119                .map(|field| field.name().as_str())
2120                .collect::<HashSet<_>>();
2121            let mut timestamp_value_column = TIMESTAMP_VALUE_PREFIX.to_string();
2122            while occupied.contains(timestamp_value_column.as_str()) {
2123                timestamp_value_column.push('_');
2124            }
2125            let mut project_exprs = normalize
2126                .schema()
2127                .iter()
2128                .map(|(qualifier, field)| {
2129                    DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
2130                })
2131                .collect::<Vec<_>>();
2132            // `timestamp()` preserves the shifted selector timeline even though
2133            // SeriesNormalize now retains raw native timestamp storage. Decimal
2134            // arithmetic shifts before truncating to milliseconds.
2135            let unit_factor = match col(&time_index_column)
2136                .get_type(normalize.schema())
2137                .context(DataFusionPlanningSnafu)?
2138            {
2139                ArrowDataType::Timestamp(ArrowTimeUnit::Second, _) => (1_000_i128, 4, 0),
2140                ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, _) => (1, 1, 0),
2141                ArrowDataType::Timestamp(ArrowTimeUnit::Microsecond, _) => (1, 4, 3),
2142                ArrowDataType::Timestamp(ArrowTimeUnit::Nanosecond, _) => (1, 7, 6),
2143                _ => unreachable!("time index is a timestamp"),
2144            };
2145            let sample_time = col(&time_index_column)
2146                .cast_to(&ArrowDataType::Int64, normalize.schema())
2147                .context(DataFusionPlanningSnafu)?
2148                .cast_to(&ArrowDataType::Decimal128(19, 0), normalize.schema())
2149                .context(DataFusionPlanningSnafu)?;
2150            let sample_time = DfExpr::BinaryExpr(BinaryExpr {
2151                left: Box::new(sample_time),
2152                op: Operator::Multiply,
2153                right: Box::new(lit(ScalarValue::Decimal128(
2154                    Some(unit_factor.0),
2155                    unit_factor.1,
2156                    unit_factor.2,
2157                ))),
2158            });
2159            let sample_time = DfExpr::BinaryExpr(BinaryExpr {
2160                left: Box::new(sample_time),
2161                op: Operator::Plus,
2162                right: Box::new(lit(ScalarValue::Decimal128(Some(offset_ms as i128), 19, 0))),
2163            })
2164            .cast_to(&ArrowDataType::Int64, normalize.schema())
2165            .context(DataFusionPlanningSnafu)?
2166            .cast_to(&ArrowDataType::Float64, normalize.schema())
2167            .context(DataFusionPlanningSnafu)?;
2168            let sample_time = DfExpr::BinaryExpr(BinaryExpr {
2169                left: Box::new(sample_time),
2170                op: Operator::Divide,
2171                right: Box::new(lit(1000.0)),
2172            });
2173            project_exprs.push(sample_time.alias(&timestamp_value_column));
2174            let normalize = LogicalPlanBuilder::from(normalize)
2175                .project(project_exprs)
2176                .context(DataFusionPlanningSnafu)?
2177                .build()
2178                .context(DataFusionPlanningSnafu)?;
2179            (normalize, Some(timestamp_value_column))
2180        } else {
2181            (normalize, None)
2182        };
2183
2184        let field_column = self.ctx.field_columns.first().cloned();
2185        let manipulate = InstantManipulate::new(
2186            self.ctx.start,
2187            self.ctx.end,
2188            self.ctx.lookback_delta,
2189            self.ctx.interval,
2190            offset_ms,
2191            time_index_column,
2192            if self.ctx.use_tsid {
2193                vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]
2194            } else {
2195                self.ctx.tag_columns.clone()
2196            },
2197            field_column,
2198            normalize,
2199        );
2200        let manipulate = LogicalPlan::Extension(Extension {
2201            node: Arc::new(manipulate),
2202        });
2203        if let Some(timestamp_value_column) = timestamp_value_column {
2204            self.create_timestamp_func_plan(manipulate, &timestamp_value_column)
2205        } else {
2206            Ok(manipulate)
2207        }
2208    }
2209
2210    /// Builds a projection plan for the PromQL `timestamp()` function.
2211    /// Projects the time index column as the value column for each row.
2212    ///
2213    /// # Arguments
2214    /// * `input` - Input [`LogicalPlan`] after instant-vector selection.
2215    /// * `timestamp_value_column` - Private column containing each selected sample's timestamp.
2216    ///
2217    /// # Returns
2218    /// Returns a [`Result<LogicalPlan>`] where the resulting logical plan projects the timestamp
2219    /// column as the value column, along with the original tag and time index columns.
2220    ///
2221    /// # Timestamp vs. Time Function
2222    ///
2223    /// - **Timestamp Function (`timestamp()`)**: In PromQL, the `timestamp()` function returns the
2224    ///   timestamp (time index) of each sample as the value column.
2225    ///
2226    /// - **Time Function (`time()`)**: The `time()` function returns the evaluation time of the query
2227    ///   as a scalar value.
2228    ///
2229    /// # Side Effects
2230    /// Updates the planner context's field columns to the timestamp column name.
2231    ///
2232    fn create_timestamp_func_plan(
2233        &mut self,
2234        input: LogicalPlan,
2235        timestamp_value_column: &str,
2236    ) -> Result<LogicalPlan> {
2237        let time_expr = col(timestamp_value_column).alias(DEFAULT_FIELD_COLUMN);
2238        self.ctx.field_columns = vec![time_expr.schema_name().to_string()];
2239        let mut project_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 2);
2240        project_exprs.push(self.create_time_index_column_expr()?);
2241        project_exprs.push(time_expr);
2242        project_exprs.extend(self.create_tag_column_exprs()?);
2243
2244        LogicalPlanBuilder::from(input)
2245            .project(project_exprs)
2246            .context(DataFusionPlanningSnafu)?
2247            .build()
2248            .context(DataFusionPlanningSnafu)
2249    }
2250
2251    async fn prom_matrix_selector_to_plan(
2252        &mut self,
2253        matrix_selector: &MatrixSelector,
2254    ) -> Result<LogicalPlan> {
2255        let MatrixSelector { vs, range } = matrix_selector;
2256        let VectorSelector {
2257            name,
2258            offset,
2259            matchers,
2260            ..
2261        } = vs;
2262        let matchers = self.preprocess_label_matchers(matchers, name)?;
2263        ensure!(!range.is_zero(), ZeroRangeSelectorSnafu);
2264        let range_ms = range.as_millis() as _;
2265        self.ctx.range = Some(range_ms);
2266        let offset_ms = match offset {
2267            Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond,
2268            Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond),
2269            None => 0,
2270        };
2271
2272        // Some functions like rate may require special fields in the RangeManipulate plan
2273        // so we can't skip RangeManipulate.
2274        let normalize = match self.setup_context().await? {
2275            Some(empty_plan) => empty_plan,
2276            None => {
2277                self.selector_to_series_normalize_plan(offset, matchers, true)
2278                    .await?
2279            }
2280        };
2281        let manipulate = RangeManipulate::new(
2282            self.ctx.start,
2283            self.ctx.end,
2284            self.ctx.interval,
2285            offset_ms,
2286            // TODO(ruihang): convert via Timestamp datatypes to support different time units
2287            range_ms,
2288            self.ctx
2289                .time_index_column
2290                .clone()
2291                .expect("time index should be set in `setup_context`"),
2292            self.ctx.field_columns.clone(),
2293            normalize,
2294        )
2295        .context(DataFusionPlanningSnafu)?;
2296
2297        Ok(LogicalPlan::Extension(Extension {
2298            node: Arc::new(manipulate),
2299        }))
2300    }
2301
2302    async fn prom_call_expr_to_plan(
2303        &mut self,
2304        query_engine_state: &QueryEngineState,
2305        call_expr: &Call,
2306    ) -> Result<LogicalPlan> {
2307        let Call { func, args } = call_expr;
2308        // some special functions that are not expression but a plan
2309        match func.name {
2310            SPECIAL_HISTOGRAM_QUANTILE | SPECIAL_HISTOGRAM_FRACTION => {
2311                return self
2312                    .create_histogram_plan(func.name, args, query_engine_state)
2313                    .await;
2314            }
2315            SPECIAL_VECTOR_FUNCTION => return self.create_vector_plan(args).await,
2316            SCALAR_FUNCTION => return self.create_scalar_plan(args, query_engine_state).await,
2317            SPECIAL_ABSENT_FUNCTION => {
2318                return self.create_absent_plan(args, query_engine_state).await;
2319            }
2320            _ => {}
2321        }
2322
2323        // transform function arguments
2324        let args = self.create_function_args(&args.args)?;
2325        let input = if let Some(prom_expr) = &args.input {
2326            self.prom_expr_to_plan_inner(prom_expr, func.name == "timestamp", query_engine_state)
2327                .await?
2328        } else {
2329            self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
2330            self.ctx.reset_table_name_and_schema();
2331            self.ctx.tag_columns = vec![];
2332            self.ctx.aggregation_field_labels.clear();
2333            self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
2334            LogicalPlan::Extension(Extension {
2335                node: Arc::new(
2336                    EmptyMetric::new(
2337                        self.ctx.start,
2338                        self.ctx.end,
2339                        self.ctx.interval,
2340                        SPECIAL_TIME_FUNCTION.to_string(),
2341                        DEFAULT_FIELD_COLUMN.to_string(),
2342                        None,
2343                    )
2344                    .context(DataFusionPlanningSnafu)?,
2345                ),
2346            })
2347        };
2348        let (mut func_exprs, new_tags) = self.create_function_expr(
2349            func,
2350            args.literals.clone(),
2351            input.schema(),
2352            query_engine_state,
2353        )?;
2354        func_exprs.insert(0, self.create_time_index_column_expr()?);
2355        func_exprs.extend_from_slice(&self.create_tag_column_exprs()?);
2356        if let Some(tsid_col) =
2357            Self::optional_tsid_projection(input.schema(), None, self.ctx.use_tsid)
2358        {
2359            func_exprs.push(tsid_col);
2360        }
2361
2362        // A row survives as long as one field column produced a sample, and the fields without
2363        // one stay NULL, which is the shape a selector already emits. Requiring every field to
2364        // be non-NULL would drop one field's samples because another field has none in the same
2365        // window — the reason alternative float/histogram columns already needed this form. A
2366        // single field column reduces to the same predicate either way.
2367        let builder = LogicalPlanBuilder::from(input)
2368            .project(func_exprs)
2369            .context(DataFusionPlanningSnafu)?
2370            .filter(self.create_empty_values_filter_expr(true)?)
2371            .context(DataFusionPlanningSnafu)?;
2372
2373        let builder = match func.name {
2374            "sort" => builder
2375                .sort(self.create_field_columns_sort_exprs(true))
2376                .context(DataFusionPlanningSnafu)?,
2377            "sort_desc" => builder
2378                .sort(self.create_field_columns_sort_exprs(false))
2379                .context(DataFusionPlanningSnafu)?,
2380            "sort_by_label" => builder
2381                .sort(Self::create_sort_exprs_by_tags(
2382                    func.name,
2383                    args.literals,
2384                    true,
2385                )?)
2386                .context(DataFusionPlanningSnafu)?,
2387            "sort_by_label_desc" => builder
2388                .sort(Self::create_sort_exprs_by_tags(
2389                    func.name,
2390                    args.literals,
2391                    false,
2392                )?)
2393                .context(DataFusionPlanningSnafu)?,
2394
2395            _ => builder,
2396        };
2397
2398        // Update context tags after building plan
2399        // We can't push them before planning, because they won't exist until projection.
2400        for tag in new_tags {
2401            self.ctx.tag_columns.push(tag);
2402        }
2403
2404        let plan = builder.build().context(DataFusionPlanningSnafu)?;
2405        common_telemetry::debug!("Created PromQL function plan: {plan:?} for {call_expr:?}");
2406
2407        Ok(plan)
2408    }
2409
2410    async fn prom_ext_expr_to_plan(
2411        &mut self,
2412        query_engine_state: &QueryEngineState,
2413        ext_expr: &promql_parser::parser::ast::Extension,
2414    ) -> Result<LogicalPlan> {
2415        // let promql_parser::parser::ast::Extension { expr } = ext_expr;
2416        let expr = &ext_expr.expr;
2417        let children = expr.children();
2418        let plan = self
2419            .prom_expr_to_plan(&children[0], query_engine_state)
2420            .await?;
2421        // Wrapper for the explanation/analyze of the existing plan
2422        // https://docs.rs/datafusion-expr/latest/datafusion_expr/logical_plan/builder/struct.LogicalPlanBuilder.html#method.explain
2423        // if `analyze` is true, runs the actual plan and produces
2424        // information about metrics during run.
2425        // if `verbose` is true, prints out additional details when VERBOSE keyword is specified
2426        match expr.name() {
2427            ANALYZE_NODE_NAME => LogicalPlanBuilder::from(plan)
2428                .explain(false, true)
2429                .unwrap()
2430                .build()
2431                .context(DataFusionPlanningSnafu),
2432            ANALYZE_VERBOSE_NODE_NAME => LogicalPlanBuilder::from(plan)
2433                .explain(true, true)
2434                .unwrap()
2435                .build()
2436                .context(DataFusionPlanningSnafu),
2437            EXPLAIN_NODE_NAME => LogicalPlanBuilder::from(plan)
2438                .explain(false, false)
2439                .unwrap()
2440                .build()
2441                .context(DataFusionPlanningSnafu),
2442            EXPLAIN_VERBOSE_NODE_NAME => LogicalPlanBuilder::from(plan)
2443                .explain(true, false)
2444                .unwrap()
2445                .build()
2446                .context(DataFusionPlanningSnafu),
2447            ALIAS_NODE_NAME => {
2448                let alias = expr
2449                    .as_any()
2450                    .downcast_ref::<AliasExpr>()
2451                    .context(UnexpectedPlanExprSnafu {
2452                        desc: "Expected AliasExpr",
2453                    })?
2454                    .alias
2455                    .clone();
2456                self.apply_alias(plan, alias)
2457            }
2458            _ => LogicalPlanBuilder::empty(true)
2459                .build()
2460                .context(DataFusionPlanningSnafu),
2461        }
2462    }
2463
2464    /// Extract metric name from `__name__` matcher and set it into [PromPlannerContext].
2465    /// Returns a new [Matchers] that doesn't contain metric name matcher.
2466    ///
2467    /// Each call to this function means new selector is started. Thus, the context will be reset
2468    /// at first.
2469    ///
2470    /// Name rule:
2471    /// - if `name` is some, then the matchers MUST NOT contain `__name__` matcher.
2472    /// - if `name` is none, then the matchers MAY contain NONE OR MULTIPLE `__name__` matchers.
2473    #[allow(clippy::mutable_key_type)]
2474    fn preprocess_label_matchers(
2475        &mut self,
2476        label_matchers: &Matchers,
2477        name: &Option<String>,
2478    ) -> Result<Matchers> {
2479        self.ctx.reset();
2480
2481        let metric_name;
2482        if let Some(name) = name.clone() {
2483            metric_name = Some(name);
2484            ensure!(
2485                label_matchers.find_matchers(METRIC_NAME).is_empty(),
2486                MultipleMetricMatchersSnafu
2487            );
2488        } else {
2489            let mut matches = label_matchers.find_matchers(METRIC_NAME);
2490            ensure!(!matches.is_empty(), NoMetricMatcherSnafu);
2491            ensure!(matches.len() == 1, MultipleMetricMatchersSnafu);
2492            ensure!(
2493                matches[0].op == MatchOp::Equal,
2494                UnsupportedMatcherOpSnafu {
2495                    matcher_op: matches[0].op.to_string(),
2496                    matcher: METRIC_NAME
2497                }
2498            );
2499            metric_name = matches.pop().map(|m| m.value);
2500        }
2501
2502        self.ctx.table_name = metric_name;
2503
2504        // Deduplicate in place instead of through a `HashSet`: the scan filter is built in
2505        // this order, and a hashed order makes the plan of one query vary between runs.
2506        let mut matchers: Vec<Matcher> = Vec::with_capacity(label_matchers.matchers.len());
2507        for matcher in &label_matchers.matchers {
2508            // TODO(ruihang): support other metric match ops
2509            if matcher.name == FIELD_COLUMN_MATCHER {
2510                self.ctx
2511                    .field_column_matcher
2512                    .get_or_insert_default()
2513                    .push(matcher.clone());
2514            } else if matcher.name == SCHEMA_COLUMN_MATCHER || matcher.name == DB_COLUMN_MATCHER {
2515                ensure!(
2516                    matcher.op == MatchOp::Equal,
2517                    UnsupportedMatcherOpSnafu {
2518                        matcher: matcher.name.clone(),
2519                        matcher_op: matcher.op.to_string(),
2520                    }
2521                );
2522                self.ctx.schema_name = Some(matcher.value.clone());
2523            } else if matcher.name != METRIC_NAME {
2524                self.ctx.selector_matcher.push(matcher.clone());
2525                if !matchers.contains(matcher) {
2526                    matchers.push(matcher.clone());
2527                }
2528            }
2529        }
2530
2531        Ok(Matchers::new(matchers))
2532    }
2533
2534    async fn selector_to_series_normalize_plan(
2535        &mut self,
2536        offset: &Option<Offset>,
2537        label_matchers: Matchers,
2538        is_range_selector: bool,
2539    ) -> Result<LogicalPlan> {
2540        // make table scan plan
2541        let table_ref = self.table_ref()?;
2542        let mut table_scan = self.create_table_scan_plan(table_ref.clone()).await?;
2543        let table_schema = table_scan.schema();
2544
2545        // make filter exprs
2546        let offset_duration = match offset {
2547            Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond,
2548            Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond),
2549            None => 0,
2550        };
2551        let mut scan_filters = Self::matchers_to_expr(label_matchers.clone(), table_schema)?;
2552        if let Some(time_index_filter) =
2553            self.build_time_index_filter(offset_duration, table_schema)?
2554        {
2555            scan_filters.push(time_index_filter);
2556        }
2557        if let Some(filter) = conjunction(scan_filters) {
2558            table_scan = LogicalPlanBuilder::from(table_scan)
2559                .filter(filter)
2560                .context(DataFusionPlanningSnafu)?
2561                .build()
2562                .context(DataFusionPlanningSnafu)?;
2563        }
2564
2565        // make a projection plan if there is any `__field__` matcher
2566        if let Some(field_matchers) = &self.ctx.field_column_matcher {
2567            let col_set = self.ctx.field_columns.iter().collect::<HashSet<_>>();
2568            // opt-in set
2569            let mut result_set = HashSet::new();
2570            // opt-out set
2571            let mut reverse_set = HashSet::new();
2572            for matcher in field_matchers {
2573                match &matcher.op {
2574                    MatchOp::Equal => {
2575                        if col_set.contains(&matcher.value) {
2576                            let _ = result_set.insert(matcher.value.clone());
2577                        } else {
2578                            return Err(ColumnNotFoundSnafu {
2579                                col: matcher.value.clone(),
2580                            }
2581                            .build());
2582                        }
2583                    }
2584                    MatchOp::NotEqual => {
2585                        if col_set.contains(&matcher.value) {
2586                            let _ = reverse_set.insert(matcher.value.clone());
2587                        } else {
2588                            return Err(ColumnNotFoundSnafu {
2589                                col: matcher.value.clone(),
2590                            }
2591                            .build());
2592                        }
2593                    }
2594                    MatchOp::Re(regex) => {
2595                        for col in &self.ctx.field_columns {
2596                            if regex.is_match(col) {
2597                                let _ = result_set.insert(col.clone());
2598                            }
2599                        }
2600                    }
2601                    MatchOp::NotRe(regex) => {
2602                        for col in &self.ctx.field_columns {
2603                            if regex.is_match(col) {
2604                                let _ = reverse_set.insert(col.clone());
2605                            }
2606                        }
2607                    }
2608                }
2609            }
2610            // merge two set
2611            if result_set.is_empty() {
2612                result_set = col_set.into_iter().cloned().collect();
2613            }
2614            for col in reverse_set {
2615                let _ = result_set.remove(&col);
2616            }
2617
2618            // mask the field columns in context using computed result set
2619            self.ctx.field_columns = self
2620                .ctx
2621                .field_columns
2622                .drain(..)
2623                .filter(|col| result_set.contains(col))
2624                .collect();
2625
2626            let exprs = result_set
2627                .into_iter()
2628                .map(|col| DfExpr::Column(Column::new_unqualified(col)))
2629                .chain(self.create_tag_column_exprs()?)
2630                .chain(
2631                    self.ctx
2632                        .use_tsid
2633                        .then_some(DfExpr::Column(Column::new_unqualified(
2634                            DATA_SCHEMA_TSID_COLUMN_NAME,
2635                        ))),
2636                )
2637                .chain(Some(self.create_time_index_column_expr()?))
2638                .collect::<Vec<_>>();
2639
2640            // reuse this variable for simplicity
2641            table_scan = LogicalPlanBuilder::from(table_scan)
2642                .project(exprs)
2643                .context(DataFusionPlanningSnafu)?
2644                .build()
2645                .context(DataFusionPlanningSnafu)?;
2646        }
2647
2648        // make sort plan
2649        let series_key_columns = if self.ctx.use_tsid {
2650            vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]
2651        } else {
2652            self.ctx.tag_columns.clone()
2653        };
2654
2655        let sort_exprs = if self.ctx.use_tsid {
2656            vec![
2657                DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME)).sort(true, true),
2658                self.create_time_index_column_expr()?.sort(true, true),
2659            ]
2660        } else {
2661            self.create_tag_and_time_index_column_sort_exprs()?
2662        };
2663
2664        let sort_plan = LogicalPlanBuilder::from(table_scan)
2665            .sort(sort_exprs)
2666            .context(DataFusionPlanningSnafu)?
2667            .build()
2668            .context(DataFusionPlanningSnafu)?;
2669
2670        // make divide plan
2671        let time_index_column =
2672            self.ctx
2673                .time_index_column
2674                .clone()
2675                .with_context(|| TimeIndexNotFoundSnafu {
2676                    table: table_ref.to_string(),
2677                })?;
2678        let divide_plan = LogicalPlan::Extension(Extension {
2679            node: Arc::new(SeriesDivide::new(
2680                series_key_columns.clone(),
2681                time_index_column,
2682                sort_plan,
2683            )),
2684        });
2685
2686        // make series_normalize plan
2687        if !is_range_selector && offset_duration == 0 {
2688            return Ok(divide_plan);
2689        }
2690        let series_normalize = SeriesNormalize::new(
2691            offset_duration,
2692            self.ctx
2693                .time_index_column
2694                .clone()
2695                .with_context(|| TimeIndexNotFoundSnafu {
2696                    table: table_ref.to_quoted_string(),
2697                })?,
2698            is_range_selector,
2699            series_key_columns,
2700            divide_plan,
2701        );
2702        let logical_plan = LogicalPlan::Extension(Extension {
2703            node: Arc::new(series_normalize),
2704        });
2705
2706        Ok(logical_plan)
2707    }
2708
2709    /// Convert [LabelModifier] to [Column] exprs for aggregation.
2710    /// Timestamp column and tag columns will be included.
2711    ///
2712    /// # Side effect
2713    ///
2714    /// This method will also change the tag columns in ctx if `update_ctx` is true.
2715    fn agg_modifier_to_col(
2716        &mut self,
2717        input_schema: &DFSchemaRef,
2718        modifier: &Option<LabelModifier>,
2719        update_ctx: bool,
2720    ) -> Result<Vec<DfExpr>> {
2721        match modifier {
2722            None => {
2723                if update_ctx {
2724                    self.ctx.tag_columns.clear();
2725                    self.ctx.aggregation_field_labels.clear();
2726                }
2727                Ok(vec![self.create_time_index_column_expr()?])
2728            }
2729            Some(LabelModifier::Include(labels)) => {
2730                if update_ctx {
2731                    // A `by(...)` label can name a value field of the input instead of a tag. The
2732                    // aggregate still reports it among its tag columns below, but unlike a tag it
2733                    // is a group key rather than a property of a series: its value varies between
2734                    // the samples of one series, so a matcher on it must stay above sample
2735                    // selection (#9242). Record it, before the tag columns are overwritten.
2736                    self.ctx.aggregation_field_labels = labels
2737                        .labels
2738                        .iter()
2739                        .filter(|label| {
2740                            self.ctx.field_columns.contains(label)
2741                                || !self.ctx.tag_columns.contains(label)
2742                                || self.ctx.aggregation_field_labels.contains(label)
2743                        })
2744                        .cloned()
2745                        .collect();
2746                    self.ctx.tag_columns.clear();
2747                }
2748                let mut exprs = Vec::with_capacity(labels.labels.len());
2749                for label in &labels.labels {
2750                    if is_metric_engine_internal_column(label) {
2751                        continue;
2752                    }
2753                    // nonexistence label will be ignored
2754                    if let Some(column_name) = Self::find_case_sensitive_column(input_schema, label)
2755                    {
2756                        exprs.push(DfExpr::Column(Column::from_name(column_name.clone())));
2757
2758                        if update_ctx {
2759                            // update the tag columns in context
2760                            self.ctx.tag_columns.push(column_name);
2761                        }
2762                    }
2763                }
2764                // add timestamp column
2765                exprs.push(self.create_time_index_column_expr()?);
2766
2767                Ok(exprs)
2768            }
2769            Some(LabelModifier::Exclude(labels)) => {
2770                let mut all_fields = input_schema
2771                    .fields()
2772                    .iter()
2773                    .map(|f| f.name())
2774                    .collect::<BTreeSet<_>>();
2775
2776                // Exclude metric engine internal columns (not PromQL labels) from the implicit
2777                // "without" label set.
2778                all_fields.retain(|col| !is_metric_engine_internal_column(col.as_str()));
2779
2780                // remove "without"-ed fields
2781                // nonexistence label will be ignored
2782                for label in &labels.labels {
2783                    let _ = all_fields.remove(label);
2784                }
2785
2786                // remove time index and value fields
2787                if let Some(time_index) = &self.ctx.time_index_column {
2788                    let _ = all_fields.remove(time_index);
2789                }
2790                for value in &self.ctx.field_columns {
2791                    let _ = all_fields.remove(value);
2792                }
2793
2794                if update_ctx {
2795                    // change the tag columns in context
2796                    self.ctx.tag_columns = all_fields.iter().map(|col| (*col).clone()).collect();
2797                    // `without(...)` drops the value fields of the input from its grouping labels,
2798                    // so a label stays a grouping label in name only if it survives in the input
2799                    // schema (e.g. an inner aggregation that grouped by it).
2800                    self.ctx
2801                        .aggregation_field_labels
2802                        .retain(|label| all_fields.iter().any(|col| *col == label));
2803                }
2804
2805                // collect remaining fields and convert to col expr
2806                let mut exprs = all_fields
2807                    .into_iter()
2808                    .map(|c| DfExpr::Column(Column::from(c)))
2809                    .collect::<Vec<_>>();
2810
2811                // add timestamp column
2812                exprs.push(self.create_time_index_column_expr()?);
2813
2814                Ok(exprs)
2815            }
2816        }
2817    }
2818
2819    // TODO(ruihang): ignore `MetricNameLabel` (`__name__`) matcher
2820    pub fn matchers_to_expr(
2821        label_matchers: Matchers,
2822        table_schema: &DFSchemaRef,
2823    ) -> Result<Vec<DfExpr>> {
2824        let mut exprs = Vec::with_capacity(label_matchers.matchers.len());
2825        for matcher in label_matchers.matchers {
2826            if matcher.name == SCHEMA_COLUMN_MATCHER
2827                || matcher.name == DB_COLUMN_MATCHER
2828                || matcher.name == FIELD_COLUMN_MATCHER
2829            {
2830                continue;
2831            }
2832
2833            let accepts_empty = matcher.is_match("");
2834            let column_name = Self::find_case_sensitive_column(table_schema, matcher.name.as_str());
2835            // Prometheus reads a label a series does not carry as the empty
2836            // string. A row can miss a label two ways: the table has no column
2837            // for it, or the column exists but is NULL on that row — the latter
2838            // is the norm for logical metrics sharing a physical table, which
2839            // holds the union of their label columns.
2840            let col = if let Some(column_name) = column_name {
2841                let column = DfExpr::Column(Column::from_name(&column_name));
2842                let field = table_schema
2843                    .index_of_column_by_name(None, &column_name)
2844                    .map(|index| table_schema.field(index));
2845                if accepts_empty
2846                    && let Some(data_type) = field
2847                        .filter(|field| {
2848                            field.is_nullable()
2849                                && Self::string_value_data_type(field.data_type()).is_some()
2850                        })
2851                        .map(|field| field.data_type())
2852                {
2853                    let empty = Self::string_scalar_value(data_type, Some(String::new()))
2854                        .expect("nullable label has a string type");
2855                    DfExpr::ScalarFunction(ScalarFunction {
2856                        func: coalesce(),
2857                        args: vec![column, DfExpr::Literal(empty, None)],
2858                    })
2859                } else {
2860                    column
2861                }
2862            } else {
2863                DfExpr::Literal(ScalarValue::Utf8(Some(String::new())), None)
2864            };
2865            let lit = DfExpr::Literal(ScalarValue::Utf8(Some(matcher.value)), None);
2866            let expr = match matcher.op {
2867                MatchOp::Equal => col.eq(lit),
2868                MatchOp::NotEqual => col.not_eq(lit),
2869                MatchOp::Re(re) => {
2870                    // TODO(ruihang): a more programmatic way to handle this in datafusion
2871
2872                    // This is a hack to handle `.+` and `.*`, and is not strictly correct
2873                    // `.` doesn't match newline (`\n`). Given this is in PromQL context,
2874                    // most of the time it's fine.
2875                    if re.as_str() == "^(?:.*)$" {
2876                        continue;
2877                    }
2878                    if re.as_str() == "^(?:.+)$" {
2879                        col.not_eq(DfExpr::Literal(
2880                            ScalarValue::Utf8(Some(String::new())),
2881                            None,
2882                        ))
2883                    } else {
2884                        DfExpr::BinaryExpr(BinaryExpr {
2885                            left: Box::new(col),
2886                            op: Operator::RegexMatch,
2887                            right: Box::new(DfExpr::Literal(
2888                                ScalarValue::Utf8(Some(re.as_str().to_string())),
2889                                None,
2890                            )),
2891                        })
2892                    }
2893                }
2894                MatchOp::NotRe(re) => {
2895                    if re.as_str() == "^(?:.*)$" {
2896                        DfExpr::Literal(ScalarValue::Boolean(Some(false)), None)
2897                    } else if re.as_str() == "^(?:.+)$" {
2898                        col.eq(DfExpr::Literal(
2899                            ScalarValue::Utf8(Some(String::new())),
2900                            None,
2901                        ))
2902                    } else {
2903                        DfExpr::BinaryExpr(BinaryExpr {
2904                            left: Box::new(col),
2905                            op: Operator::RegexNotMatch,
2906                            right: Box::new(DfExpr::Literal(
2907                                ScalarValue::Utf8(Some(re.as_str().to_string())),
2908                                None,
2909                            )),
2910                        })
2911                    }
2912                }
2913            };
2914            exprs.push(expr);
2915        }
2916
2917        Ok(exprs)
2918    }
2919
2920    fn find_case_sensitive_column(schema: &DFSchemaRef, column: &str) -> Option<String> {
2921        if is_metric_engine_internal_column(column) {
2922            return None;
2923        }
2924        schema
2925            .fields()
2926            .iter()
2927            .find(|field| field.name() == column)
2928            .map(|field| field.name().clone())
2929    }
2930
2931    fn table_from_source(&self, source: &Arc<dyn TableSource>) -> Result<table::TableRef> {
2932        Ok(source
2933            .downcast_ref::<DefaultTableSource>()
2934            .context(UnknownTableSnafu)?
2935            .table_provider
2936            .downcast_ref::<DfTableProviderAdapter>()
2937            .context(UnknownTableSnafu)?
2938            .table())
2939    }
2940
2941    fn table_ref(&self) -> Result<TableReference> {
2942        let table_name = self
2943            .ctx
2944            .table_name
2945            .clone()
2946            .context(TableNameNotFoundSnafu)?;
2947
2948        // set schema name if `__schema__` is given
2949        let table_ref = if let Some(schema_name) = &self.ctx.schema_name {
2950            TableReference::partial(schema_name.as_str(), table_name.as_str())
2951        } else {
2952            TableReference::bare(table_name.as_str())
2953        };
2954
2955        Ok(table_ref)
2956    }
2957
2958    fn build_time_index_filter(
2959        &self,
2960        offset_duration: i64,
2961        schema: &DFSchemaRef,
2962    ) -> Result<Option<DfExpr>> {
2963        let start = self.ctx.start;
2964        let end = self.ctx.end;
2965        if end < start {
2966            return InvalidTimeRangeSnafu { start, end }.fail();
2967        }
2968        let time_index_expr = self.create_time_index_column_expr()?;
2969        let time_index_name = self.ctx.time_index_column.as_ref().unwrap();
2970        let unit = schema
2971            .index_of_column_by_name(None, time_index_name)
2972            .and_then(|index| match schema.field(index).data_type() {
2973                ArrowDataType::Timestamp(unit, _) => Some(*unit),
2974                _ => None,
2975            })
2976            .unwrap_or(ArrowTimeUnit::Millisecond);
2977        let native_value = |milliseconds: i128| match unit {
2978            ArrowTimeUnit::Second => milliseconds.div_euclid(1_000),
2979            ArrowTimeUnit::Millisecond => milliseconds,
2980            ArrowTimeUnit::Microsecond => milliseconds * 1_000,
2981            ArrowTimeUnit::Nanosecond => milliseconds * 1_000_000,
2982        };
2983        let scalar = |milliseconds: i128| -> Option<ScalarValue> {
2984            let value = i64::try_from(native_value(milliseconds)).ok()?;
2985            Some(match unit {
2986                ArrowTimeUnit::Second => ScalarValue::TimestampSecond(Some(value), None),
2987                ArrowTimeUnit::Millisecond => ScalarValue::TimestampMillisecond(Some(value), None),
2988                ArrowTimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(Some(value), None),
2989                ArrowTimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(Some(value), None),
2990            })
2991        };
2992        let window = self.ctx.range.unwrap_or(self.ctx.lookback_delta);
2993        let filter = |lower_ms: i128, upper_ms: i128| {
2994            let lower_value = native_value(lower_ms);
2995            let upper_value = native_value(upper_ms);
2996            if lower_value > i128::from(i64::MAX) || upper_value < i128::from(i64::MIN) {
2997                return Some(lit(false));
2998            }
2999            let lower_filter = (lower_value >= i128::from(i64::MIN)).then(|| {
3000                let lower = DfExpr::Literal(scalar(lower_ms).unwrap(), None);
3001                if window == 0 {
3002                    time_index_expr.clone().gt_eq(lower)
3003                } else if unit == ArrowTimeUnit::Millisecond
3004                    && let Some(inclusive_lower) = lower_ms
3005                        .checked_add(1)
3006                        .and_then(|lower| i64::try_from(lower).ok())
3007                        .and_then(|lower| scalar(i128::from(lower)))
3008                {
3009                    time_index_expr
3010                        .clone()
3011                        .gt_eq(DfExpr::Literal(inclusive_lower, None))
3012                } else {
3013                    time_index_expr.clone().gt(lower)
3014                }
3015            });
3016            let upper_filter = (upper_value <= i128::from(i64::MAX)).then(|| {
3017                time_index_expr
3018                    .clone()
3019                    .lt_eq(DfExpr::Literal(scalar(upper_ms).unwrap(), None))
3020            });
3021
3022            // An underflowing lower bound must not discard a representable upper
3023            // bound: without it, LastRow could retain a future row and discard the
3024            // older eligible sample before the manipulator can check its time.
3025            match (lower_filter, upper_filter) {
3026                (Some(lower), Some(upper)) => Some(lower.and(upper)),
3027                (Some(filter), None) | (None, Some(filter)) => Some(filter),
3028                (None, None) => None,
3029            }
3030        };
3031        let bounds = |timestamp: i64| {
3032            let upper = i128::from(timestamp) - i128::from(offset_duration);
3033            (upper - i128::from(window), upper)
3034        };
3035        let num_points = (end as i128 - start as i128) / self.ctx.interval as i128;
3036        if num_points > MAX_SCATTER_POINTS as i128 || self.ctx.interval <= INTERVAL_1H {
3037            let (lower, _) = bounds(start);
3038            let (_, upper) = bounds(end);
3039            return Ok(filter(lower, upper));
3040        }
3041        let mut filters = Vec::new();
3042        for timestamp in (start..=end).step_by(self.ctx.interval as usize) {
3043            let (lower, upper) = bounds(timestamp);
3044            let Some(filter) = filter(lower, upper) else {
3045                // A point whose native bounds cannot be represented may cover the whole native
3046                // time domain, so its disjunct cannot be omitted.
3047                return Ok(None);
3048            };
3049            filters.push(filter);
3050        }
3051        Ok(filters.into_iter().reduce(DfExpr::or))
3052    }
3053
3054    /// Create a table scan plan and a filter plan with given filter.
3055    ///
3056    /// # Panic
3057    /// If the filter is empty
3058    async fn create_table_scan_plan(&mut self, table_ref: TableReference) -> Result<LogicalPlan> {
3059        let provider = self
3060            .table_provider
3061            .resolve_table(table_ref.clone())
3062            .await
3063            .context(CatalogSnafu)?;
3064
3065        let logical_table = self.table_from_source(&provider)?;
3066
3067        // Try to rewrite the table scan to physical table scan if possible.
3068        let mut maybe_phy_table_ref = table_ref.clone();
3069        let mut scan_provider = provider;
3070        let mut table_id_filter: Option<u32> = None;
3071
3072        // If it's a metric engine logical table, scan its physical table directly and filter by
3073        // `__table_id = logical_table_id` to get access to internal columns like `__tsid`.
3074        if logical_table.table_info().meta.engine == METRIC_ENGINE_NAME
3075            && let Some(physical_table_name) = logical_table
3076                .table_info()
3077                .meta
3078                .options
3079                .extra_options
3080                .get(LOGICAL_TABLE_METADATA_KEY)
3081        {
3082            let physical_table_ref = if let Some(schema_name) = &self.ctx.schema_name {
3083                TableReference::partial(schema_name.as_str(), physical_table_name.as_str())
3084            } else {
3085                TableReference::bare(physical_table_name.as_str())
3086            };
3087
3088            let physical_provider = match self
3089                .table_provider
3090                .resolve_table(physical_table_ref.clone())
3091                .await
3092            {
3093                Ok(provider) => provider,
3094                Err(e) if e.status_code() == StatusCode::TableNotFound => {
3095                    // Fall back to scanning the logical table. It still works, but without
3096                    // `__tsid` optimization.
3097                    scan_provider.clone()
3098                }
3099                Err(e) => return Err(e).context(CatalogSnafu),
3100            };
3101
3102            if !Arc::ptr_eq(&physical_provider, &scan_provider) {
3103                // Only rewrite when internal columns exist in physical schema.
3104                let physical_table = self.table_from_source(&physical_provider)?;
3105
3106                let has_table_id = physical_table
3107                    .schema()
3108                    .column_schema_by_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME)
3109                    .is_some();
3110                let has_tsid = physical_table
3111                    .schema()
3112                    .column_schema_by_name(DATA_SCHEMA_TSID_COLUMN_NAME)
3113                    .is_some_and(|col| matches!(col.data_type, ConcreteDataType::UInt64(_)));
3114
3115                if has_table_id && has_tsid {
3116                    scan_provider = physical_provider;
3117                    maybe_phy_table_ref = physical_table_ref;
3118                    table_id_filter = Some(logical_table.table_info().ident.table_id);
3119                }
3120            }
3121        }
3122
3123        let scan_table = self.table_from_source(&scan_provider)?;
3124
3125        let use_tsid = table_id_filter.is_some()
3126            && scan_table
3127                .schema()
3128                .column_schema_by_name(DATA_SCHEMA_TSID_COLUMN_NAME)
3129                .is_some_and(|col| matches!(col.data_type, ConcreteDataType::UInt64(_)));
3130        self.ctx.use_tsid = use_tsid;
3131
3132        let all_table_tags = self.ctx.tag_columns.clone();
3133
3134        let scan_tag_columns = if use_tsid {
3135            let mut scan_tags = self.ctx.tag_columns.clone();
3136            for matcher in &self.ctx.selector_matcher {
3137                if is_metric_engine_internal_column(&matcher.name) {
3138                    continue;
3139                }
3140                if all_table_tags.iter().any(|tag| tag == &matcher.name) {
3141                    scan_tags.push(matcher.name.clone());
3142                }
3143            }
3144            scan_tags.sort_unstable();
3145            scan_tags.dedup();
3146            scan_tags
3147        } else {
3148            self.ctx.tag_columns.clone()
3149        };
3150
3151        let time_index_data_type = scan_table
3152            .schema()
3153            .timestamp_column()
3154            .with_context(|| TimeIndexNotFoundSnafu {
3155                table: maybe_phy_table_ref.to_quoted_string(),
3156            })?
3157            .data_type
3158            .clone();
3159        let is_time_index_second =
3160            time_index_data_type == ConcreteDataType::timestamp_second_datatype();
3161
3162        let scan_projection = if table_id_filter.is_some() {
3163            let mut required_columns = HashSet::new();
3164            required_columns.insert(DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string());
3165            required_columns.insert(self.ctx.time_index_column.clone().with_context(|| {
3166                TimeIndexNotFoundSnafu {
3167                    table: maybe_phy_table_ref.to_quoted_string(),
3168                }
3169            })?);
3170            for col in &scan_tag_columns {
3171                required_columns.insert(col.clone());
3172            }
3173            for col in &self.ctx.field_columns {
3174                required_columns.insert(col.clone());
3175            }
3176            if use_tsid {
3177                required_columns.insert(DATA_SCHEMA_TSID_COLUMN_NAME.to_string());
3178            }
3179
3180            let arrow_schema = scan_table.schema().arrow_schema().clone();
3181            Some(
3182                arrow_schema
3183                    .fields()
3184                    .iter()
3185                    .enumerate()
3186                    .filter(|(_, field)| required_columns.contains(field.name().as_str()))
3187                    .map(|(idx, _)| idx)
3188                    .collect::<Vec<_>>(),
3189            )
3190        } else {
3191            None
3192        };
3193
3194        let mut scan_plan =
3195            LogicalPlanBuilder::scan(maybe_phy_table_ref.clone(), scan_provider, scan_projection)
3196                .context(DataFusionPlanningSnafu)?
3197                .build()
3198                .context(DataFusionPlanningSnafu)?;
3199
3200        if let Some(table_id) = table_id_filter {
3201            scan_plan = LogicalPlanBuilder::from(scan_plan)
3202                .filter(
3203                    DfExpr::Column(Column::from_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME))
3204                        .eq(lit(table_id)),
3205                )
3206                .context(DataFusionPlanningSnafu)?
3207                .alias(table_ref.clone()) // rename the relation back to logical table's name after filtering
3208                .context(DataFusionPlanningSnafu)?
3209                .build()
3210                .context(DataFusionPlanningSnafu)?;
3211        }
3212
3213        if is_time_index_second {
3214            // Promote seconds so millisecond offsets remain exact; retain finer precision.
3215            // Later manipulators compare native sample ticks, while PromQL evaluation and
3216            // emitted timestamps remain millisecond-based, so this projection must not
3217            // silently truncate a finer-grained time index.
3218            let expr: Vec<_> = self
3219                .create_field_column_exprs()?
3220                .into_iter()
3221                .chain(
3222                    scan_tag_columns
3223                        .iter()
3224                        .map(|tag| DfExpr::Column(Column::from_name(tag))),
3225                )
3226                .chain(self.ctx.use_tsid.then_some(DfExpr::Column(Column::new(
3227                    Some(table_ref.clone()),
3228                    DATA_SCHEMA_TSID_COLUMN_NAME.to_string(),
3229                ))))
3230                .chain(Some(DfExpr::Alias(Alias {
3231                    expr: Box::new(DfExpr::Cast(Cast::new(
3232                        Box::new(self.create_time_index_column_expr()?),
3233                        ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
3234                    ))),
3235                    relation: Some(table_ref.clone()),
3236                    name: self
3237                        .ctx
3238                        .time_index_column
3239                        .as_ref()
3240                        .with_context(|| TimeIndexNotFoundSnafu {
3241                            table: table_ref.to_quoted_string(),
3242                        })?
3243                        .clone(),
3244                    metadata: None,
3245                })))
3246                .collect::<Vec<_>>();
3247            scan_plan = LogicalPlanBuilder::from(scan_plan)
3248                .project(expr)
3249                .context(DataFusionPlanningSnafu)?
3250                .build()
3251                .context(DataFusionPlanningSnafu)?;
3252        } else if table_id_filter.is_some()
3253            || time_index_data_type == ConcreteDataType::timestamp_microsecond_datatype()
3254            || time_index_data_type == ConcreteDataType::timestamp_nanosecond_datatype()
3255        {
3256            // Drop the internal `__table_id` column after filtering and preserve PromQL's
3257            // field/tag/timestamp column order for native microsecond/nanosecond timestamps.
3258            // Keeping the original time column also lets the existing ordering hints
3259            // use PerSeries scans without a cast, repartition, and sort. This benefits
3260            // multi-evaluation selectors too; only a single evaluation can use LastRow.
3261            let project_exprs = self
3262                .create_field_column_exprs()?
3263                .into_iter()
3264                .chain(
3265                    scan_tag_columns
3266                        .iter()
3267                        .map(|tag| DfExpr::Column(Column::from_name(tag))),
3268                )
3269                .chain(
3270                    self.ctx
3271                        .use_tsid
3272                        .then_some(DfExpr::Column(Column::from_name(
3273                            DATA_SCHEMA_TSID_COLUMN_NAME,
3274                        ))),
3275                )
3276                .chain(Some(self.create_time_index_column_expr()?))
3277                .collect::<Vec<_>>();
3278
3279            scan_plan = LogicalPlanBuilder::from(scan_plan)
3280                .project(project_exprs)
3281                .context(DataFusionPlanningSnafu)?
3282                .build()
3283                .context(DataFusionPlanningSnafu)?;
3284        }
3285
3286        let result = LogicalPlanBuilder::from(scan_plan)
3287            .build()
3288            .context(DataFusionPlanningSnafu)?;
3289        Ok(result)
3290    }
3291
3292    fn collect_row_key_tag_columns_from_plan(
3293        &self,
3294        plan: &LogicalPlan,
3295    ) -> Result<BTreeSet<String>> {
3296        fn walk(
3297            planner: &PromPlanner,
3298            plan: &LogicalPlan,
3299            out: &mut BTreeSet<String>,
3300        ) -> Result<()> {
3301            // Derived PromQL plans may contain non-Greptime scans without row-key metadata.
3302            if let LogicalPlan::TableScan(scan) = plan
3303                && let Ok(table) = planner.table_from_source(&scan.source)
3304            {
3305                for col in table.table_info().meta.row_key_column_names() {
3306                    if col != DATA_SCHEMA_TABLE_ID_COLUMN_NAME
3307                        && col != DATA_SCHEMA_TSID_COLUMN_NAME
3308                        && !is_metric_engine_internal_column(col)
3309                    {
3310                        out.insert(col.clone());
3311                    }
3312                }
3313            }
3314
3315            for input in plan.inputs() {
3316                walk(planner, input, out)?;
3317            }
3318            Ok(())
3319        }
3320
3321        let mut out = BTreeSet::new();
3322        walk(self, plan, &mut out)?;
3323        Ok(out)
3324    }
3325
3326    /// Setup [PromPlannerContext]'s state fields.
3327    ///
3328    /// Returns a logical plan for an empty metric.
3329    async fn setup_context(&mut self) -> Result<Option<LogicalPlan>> {
3330        let table_ref = self.table_ref()?;
3331        let source = match self.table_provider.resolve_table(table_ref.clone()).await {
3332            Err(e) if e.status_code() == StatusCode::TableNotFound => {
3333                let plan = self.setup_context_for_empty_metric()?;
3334                return Ok(Some(plan));
3335            }
3336            res => res.context(CatalogSnafu)?,
3337        };
3338        let table = self.table_from_source(&source)?;
3339
3340        // set time index column name
3341        let time_index = table
3342            .schema()
3343            .timestamp_column()
3344            .with_context(|| TimeIndexNotFoundSnafu {
3345                table: table_ref.to_quoted_string(),
3346            })?
3347            .name
3348            .clone();
3349        self.ctx.time_index_column = Some(time_index);
3350
3351        // set values columns
3352        let values = table
3353            .table_info()
3354            .meta
3355            .field_column_names()
3356            .cloned()
3357            .collect();
3358        self.ctx.field_columns = values;
3359
3360        // set primary key (tag) columns
3361        let tags = table
3362            .table_info()
3363            .meta
3364            .row_key_column_names()
3365            .filter(|col| {
3366                // remove metric engine's internal columns
3367                col != &DATA_SCHEMA_TABLE_ID_COLUMN_NAME && col != &DATA_SCHEMA_TSID_COLUMN_NAME
3368            })
3369            .cloned()
3370            .collect();
3371        self.ctx.tag_columns = tags;
3372        // The operand is a plain selector: its tag columns are the table's.
3373        self.ctx.aggregation_field_labels.clear();
3374
3375        self.ctx.use_tsid = false;
3376
3377        Ok(None)
3378    }
3379
3380    /// Setup [PromPlannerContext]'s state fields for a non existent table
3381    /// without any rows.
3382    fn setup_context_for_empty_metric(&mut self) -> Result<LogicalPlan> {
3383        self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
3384        self.ctx.reset_table_name_and_schema();
3385        self.ctx.tag_columns = vec![];
3386        self.ctx.aggregation_field_labels.clear();
3387        self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
3388        self.ctx.use_tsid = false;
3389
3390        // The table doesn't have any data, so we set start to 0 and end to -1.
3391        let plan = LogicalPlan::Extension(Extension {
3392            node: Arc::new(
3393                EmptyMetric::new(
3394                    0,
3395                    -1,
3396                    self.ctx.interval,
3397                    SPECIAL_TIME_FUNCTION.to_string(),
3398                    DEFAULT_FIELD_COLUMN.to_string(),
3399                    Some(lit(0.0f64)),
3400                )
3401                .context(DataFusionPlanningSnafu)?,
3402            ),
3403        });
3404        Ok(plan)
3405    }
3406
3407    // TODO(ruihang): insert column expr
3408    fn create_function_args(&self, args: &[Box<PromExpr>]) -> Result<FunctionArgs> {
3409        let mut result = FunctionArgs::default();
3410
3411        for arg in args {
3412            // First try to parse as literal expression (including binary expressions like 100.0 + 3.0)
3413            if let Some(expr) = Self::try_build_literal_expr(arg) {
3414                result.literals.push(expr);
3415            } else {
3416                // If not a literal, treat as vector input
3417                match arg.as_ref() {
3418                    PromExpr::Subquery(_)
3419                    | PromExpr::VectorSelector(_)
3420                    | PromExpr::MatrixSelector(_)
3421                    | PromExpr::Extension(_)
3422                    | PromExpr::Aggregate(_)
3423                    | PromExpr::Paren(_)
3424                    | PromExpr::Call(_)
3425                    | PromExpr::Binary(_)
3426                    | PromExpr::Unary(_) => {
3427                        if result.input.replace(*arg.clone()).is_some() {
3428                            MultipleVectorSnafu { expr: *arg.clone() }.fail()?;
3429                        }
3430                    }
3431
3432                    _ => {
3433                        let expr = Self::get_param_as_literal_expr(Some(arg.as_ref()), None, None)?;
3434                        result.literals.push(expr);
3435                    }
3436                }
3437            }
3438        }
3439
3440        Ok(result)
3441    }
3442
3443    fn create_mixed_range_function_exprs(
3444        &mut self,
3445        func: &Function,
3446        mut other_input_exprs: VecDeque<DfExpr>,
3447        float_field: &str,
3448        histogram_field: &str,
3449        input_schema: &DFSchemaRef,
3450    ) -> Result<Option<Vec<DfExpr>>> {
3451        let returns_histogram = matches!(
3452            func.name,
3453            "rate"
3454                | "increase"
3455                | "delta"
3456                | "idelta"
3457                | "irate"
3458                | "avg_over_time"
3459                | "sum_over_time"
3460                | "last_over_time"
3461        );
3462        if !returns_histogram
3463            && !matches!(
3464                func.name,
3465                "changes"
3466                    | "resets"
3467                    | "deriv"
3468                    | "min_over_time"
3469                    | "max_over_time"
3470                    | "count_over_time"
3471                    | "absent_over_time"
3472                    | "present_over_time"
3473                    | "stddev_over_time"
3474                    | "stdvar_over_time"
3475                    | "quantile_over_time"
3476                    | "predict_linear"
3477                    | "double_exponential_smoothing"
3478                    | "holt_winters"
3479            )
3480        {
3481            return Ok(None);
3482        }
3483
3484        if func.name == "predict_linear" {
3485            other_input_exprs[0] = DfExpr::Cast(Cast::new(
3486                Box::new(other_input_exprs[0].clone()),
3487                ArrowDataType::Int64,
3488            ));
3489        }
3490
3491        let timestamp_range = DfExpr::Column(Column::from_name(
3492            RangeManipulate::build_timestamp_range_name(
3493                self.ctx.time_index_column.as_ref().unwrap(),
3494            ),
3495        ));
3496        let float_range = DfExpr::Column(Column::from_name(float_field));
3497        let histogram_range = DfExpr::Column(Column::from_name(histogram_field));
3498        let mut args = Vec::with_capacity(other_input_exprs.len() + 6);
3499        args.push(lit(func.name));
3500        args.push(timestamp_range.clone());
3501        args.push(float_range.clone());
3502        args.push(histogram_range.clone());
3503        args.extend(other_input_exprs);
3504        if matches!(func.name, "rate" | "increase" | "delta") {
3505            args.push(self.create_time_index_column_expr()?);
3506            args.push(lit(self.ctx.range.context(ExpectRangeSelectorSnafu)?));
3507        }
3508
3509        let mut float_expr = DfExpr::ScalarFunction(ScalarFunction {
3510            func: Arc::new(MixedRange::float_udf(self.promql_annotations.clone())),
3511            args: args.clone(),
3512        });
3513        if matches!(func.name, "rate" | "increase") {
3514            let raw_delta_function = if func.name == "rate" {
3515                "raw_delta_rate"
3516            } else {
3517                "raw_delta_increase"
3518            };
3519            let delta_sum = DfExpr::ScalarFunction(ScalarFunction {
3520                func: Arc::new(MixedRange::float_udf(self.promql_annotations.clone())),
3521                args: vec![
3522                    lit(raw_delta_function),
3523                    timestamp_range,
3524                    float_range,
3525                    histogram_range,
3526                ],
3527            });
3528            float_expr = self.select_delta_range_math(
3529                func.name,
3530                input_schema,
3531                self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3532                delta_sum,
3533                float_expr,
3534            )?;
3535        }
3536        let exprs = if returns_histogram {
3537            self.ctx.field_columns = vec![float_field.to_string(), histogram_field.to_string()];
3538            vec![
3539                float_expr.alias(float_field),
3540                DfExpr::ScalarFunction(ScalarFunction {
3541                    func: Arc::new(MixedRange::histogram_udf(self.promql_annotations.clone())),
3542                    args,
3543                })
3544                .alias(histogram_field),
3545            ]
3546        } else {
3547            let display_name = float_expr.schema_name().to_string();
3548            self.ctx.field_columns = vec![display_name.clone()];
3549            vec![float_expr.alias(display_name)]
3550        };
3551        Ok(Some(exprs))
3552    }
3553
3554    /// Creates function expressions for projection and returns the expressions and new tags.
3555    ///
3556    /// # Side Effects
3557    ///
3558    /// This method will update [PromPlannerContext]'s fields and tags if needed.
3559    fn create_function_expr(
3560        &mut self,
3561        func: &Function,
3562        other_input_exprs: Vec<DfExpr>,
3563        input_schema: &DFSchemaRef,
3564        query_engine_state: &QueryEngineState,
3565    ) -> Result<(Vec<DfExpr>, Vec<String>)> {
3566        // TODO(ruihang): check function args list
3567        let mut other_input_exprs: VecDeque<DfExpr> = other_input_exprs.into();
3568        if let Some((float_field, histogram_field)) =
3569            Self::alternative_sample_range_columns(input_schema, &self.ctx.field_columns)
3570                .map(|(float, histogram)| (float.to_string(), histogram.to_string()))
3571            && let Some(exprs) = self.create_mixed_range_function_exprs(
3572                func,
3573                other_input_exprs.clone(),
3574                &float_field,
3575                &histogram_field,
3576                input_schema,
3577            )?
3578        {
3579            return Ok((exprs, vec![]));
3580        }
3581        let alternative_samples =
3582            Self::field_columns_are_alternative_samples(input_schema, &self.ctx.field_columns);
3583        let all_field_columns_are_native_histogram_ranges =
3584            self.all_field_columns_are_native_histogram_ranges(input_schema);
3585
3586        // TODO(ruihang): set this according to in-param list
3587        let field_column_pos = 0;
3588        let mut exprs = Vec::with_capacity(self.ctx.field_columns.len());
3589        // New labels after executing the function, e.g. `label_replace` etc.
3590        let mut new_tags = vec![];
3591        let promql_annotations = self.promql_annotations.clone();
3592        let native_histogram_drop_udf = |name: &str| {
3593            Arc::new(NativeHistogramDrop::float_null_udf(
3594                format!(
3595                    "{name}: dropped native histogram samples because this function is not supported for native histograms"
3596                ),
3597                promql_annotations.clone(),
3598            ))
3599        };
3600        let scalar_func = match func.name {
3601            "increase" => {
3602                if all_field_columns_are_native_histogram_ranges {
3603                    ScalarFunc::ExtrapolateUdf(
3604                        Arc::new(NativeHistogramIncrease::scalar_udf_with_collector(
3605                            self.promql_annotations.clone(),
3606                        )),
3607                        self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3608                    )
3609                } else {
3610                    ScalarFunc::ExtrapolateUdf(
3611                        Arc::new(Increase::scalar_udf()),
3612                        self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3613                    )
3614                }
3615            }
3616            "rate" => {
3617                if all_field_columns_are_native_histogram_ranges {
3618                    ScalarFunc::ExtrapolateUdf(
3619                        Arc::new(NativeHistogramRate::scalar_udf_with_collector(
3620                            self.promql_annotations.clone(),
3621                        )),
3622                        self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3623                    )
3624                } else {
3625                    ScalarFunc::ExtrapolateUdf(
3626                        Arc::new(Rate::scalar_udf()),
3627                        self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3628                    )
3629                }
3630            }
3631            "delta" => {
3632                if all_field_columns_are_native_histogram_ranges {
3633                    ScalarFunc::ExtrapolateUdf(
3634                        Arc::new(NativeHistogramDelta::scalar_udf_with_collector(
3635                            self.promql_annotations.clone(),
3636                        )),
3637                        self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3638                    )
3639                } else {
3640                    ScalarFunc::ExtrapolateUdf(
3641                        Arc::new(Delta::scalar_udf()),
3642                        self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3643                    )
3644                }
3645            }
3646            "idelta" => {
3647                if all_field_columns_are_native_histogram_ranges {
3648                    ScalarFunc::Udf(Arc::new(NativeHistogramIDelta::scalar_udf_with_collector(
3649                        self.promql_annotations.clone(),
3650                    )))
3651                } else {
3652                    ScalarFunc::Udf(Arc::new(IDelta::<false>::scalar_udf()))
3653                }
3654            }
3655            "irate" => {
3656                if all_field_columns_are_native_histogram_ranges {
3657                    ScalarFunc::Udf(Arc::new(NativeHistogramIRate::scalar_udf_with_collector(
3658                        self.promql_annotations.clone(),
3659                    )))
3660                } else {
3661                    ScalarFunc::Udf(Arc::new(IDelta::<true>::scalar_udf()))
3662                }
3663            }
3664            "resets" => {
3665                if all_field_columns_are_native_histogram_ranges {
3666                    ScalarFunc::Udf(Arc::new(NativeHistogramResets::scalar_udf()))
3667                } else {
3668                    ScalarFunc::Udf(Arc::new(Resets::scalar_udf()))
3669                }
3670            }
3671            "changes" => {
3672                if all_field_columns_are_native_histogram_ranges {
3673                    ScalarFunc::Udf(Arc::new(NativeHistogramChanges::scalar_udf()))
3674                } else {
3675                    ScalarFunc::Udf(Arc::new(Changes::scalar_udf()))
3676                }
3677            }
3678            "deriv" => {
3679                if all_field_columns_are_native_histogram_ranges {
3680                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3681                } else {
3682                    ScalarFunc::Udf(Arc::new(Deriv::scalar_udf()))
3683                }
3684            }
3685            "avg_over_time" => {
3686                if all_field_columns_are_native_histogram_ranges {
3687                    ScalarFunc::Udf(Arc::new(
3688                        NativeHistogramAvgOverTime::scalar_udf_with_collector(
3689                            self.promql_annotations.clone(),
3690                        ),
3691                    ))
3692                } else {
3693                    ScalarFunc::Udf(Arc::new(AvgOverTime::scalar_udf()))
3694                }
3695            }
3696            "min_over_time" => {
3697                if all_field_columns_are_native_histogram_ranges {
3698                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3699                } else {
3700                    ScalarFunc::Udf(Arc::new(MinOverTime::scalar_udf()))
3701                }
3702            }
3703            "max_over_time" => {
3704                if all_field_columns_are_native_histogram_ranges {
3705                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3706                } else {
3707                    ScalarFunc::Udf(Arc::new(MaxOverTime::scalar_udf()))
3708                }
3709            }
3710            "sum_over_time" => {
3711                if all_field_columns_are_native_histogram_ranges {
3712                    ScalarFunc::Udf(Arc::new(
3713                        NativeHistogramSumOverTime::scalar_udf_with_collector(
3714                            self.promql_annotations.clone(),
3715                        ),
3716                    ))
3717                } else {
3718                    ScalarFunc::Udf(Arc::new(SumOverTime::scalar_udf()))
3719                }
3720            }
3721            "count_over_time" => {
3722                if all_field_columns_are_native_histogram_ranges {
3723                    ScalarFunc::Udf(Arc::new(NativeHistogramCountOverTime::scalar_udf()))
3724                } else {
3725                    ScalarFunc::Udf(Arc::new(CountOverTime::scalar_udf()))
3726                }
3727            }
3728            "last_over_time" => {
3729                if all_field_columns_are_native_histogram_ranges {
3730                    ScalarFunc::Udf(Arc::new(NativeHistogramLastOverTime::scalar_udf()))
3731                } else {
3732                    ScalarFunc::Udf(Arc::new(LastOverTime::scalar_udf()))
3733                }
3734            }
3735            "absent_over_time" => {
3736                if all_field_columns_are_native_histogram_ranges {
3737                    ScalarFunc::Udf(Arc::new(NativeHistogramAbsentOverTime::scalar_udf()))
3738                } else {
3739                    ScalarFunc::Udf(Arc::new(AbsentOverTime::scalar_udf()))
3740                }
3741            }
3742            "present_over_time" => {
3743                if all_field_columns_are_native_histogram_ranges {
3744                    ScalarFunc::Udf(Arc::new(NativeHistogramPresentOverTime::scalar_udf()))
3745                } else {
3746                    ScalarFunc::Udf(Arc::new(PresentOverTime::scalar_udf()))
3747                }
3748            }
3749            "stddev_over_time" => {
3750                if all_field_columns_are_native_histogram_ranges {
3751                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3752                } else {
3753                    ScalarFunc::Udf(Arc::new(StddevOverTime::scalar_udf()))
3754                }
3755            }
3756            "stdvar_over_time" => {
3757                if all_field_columns_are_native_histogram_ranges {
3758                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3759                } else {
3760                    ScalarFunc::Udf(Arc::new(StdvarOverTime::scalar_udf()))
3761                }
3762            }
3763            "quantile_over_time" => {
3764                if all_field_columns_are_native_histogram_ranges {
3765                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3766                } else {
3767                    ScalarFunc::Udf(Arc::new(QuantileOverTime::scalar_udf()))
3768                }
3769            }
3770            "predict_linear" => {
3771                if all_field_columns_are_native_histogram_ranges {
3772                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3773                } else {
3774                    other_input_exprs[0] = DfExpr::Cast(Cast::new(
3775                        Box::new(other_input_exprs[0].clone()),
3776                        ArrowDataType::Int64,
3777                    ));
3778                    ScalarFunc::Udf(Arc::new(PredictLinear::scalar_udf()))
3779                }
3780            }
3781            "double_exponential_smoothing" | "holt_winters" => {
3782                if all_field_columns_are_native_histogram_ranges {
3783                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3784                } else {
3785                    ScalarFunc::Udf(Arc::new(DoubleExponentialSmoothing::scalar_udf()))
3786                }
3787            }
3788            "histogram_count" => {
3789                ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramCount::scalar_udf()))
3790            }
3791            "histogram_sum" => {
3792                ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramSum::scalar_udf()))
3793            }
3794            "histogram_avg" => {
3795                ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramAvg::scalar_udf()))
3796            }
3797            "histogram_stddev" => {
3798                ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramStddev::scalar_udf()))
3799            }
3800            "histogram_stdvar" => {
3801                ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramStdvar::scalar_udf()))
3802            }
3803            "time" => {
3804                exprs.push(build_special_time_expr(
3805                    self.ctx.time_index_column.as_ref().unwrap(),
3806                ));
3807                ScalarFunc::GeneratedExpr
3808            }
3809            "minute" => {
3810                // date_part('minute', time_index)
3811                let expr = self.date_part_on_time_index("minute")?;
3812                exprs.push(expr);
3813                ScalarFunc::GeneratedExpr
3814            }
3815            "hour" => {
3816                // date_part('hour', time_index)
3817                let expr = self.date_part_on_time_index("hour")?;
3818                exprs.push(expr);
3819                ScalarFunc::GeneratedExpr
3820            }
3821            "month" => {
3822                // date_part('month', time_index)
3823                let expr = self.date_part_on_time_index("month")?;
3824                exprs.push(expr);
3825                ScalarFunc::GeneratedExpr
3826            }
3827            "year" => {
3828                // date_part('year', time_index)
3829                let expr = self.date_part_on_time_index("year")?;
3830                exprs.push(expr);
3831                ScalarFunc::GeneratedExpr
3832            }
3833            "day_of_month" => {
3834                // date_part('day', time_index)
3835                let expr = self.date_part_on_time_index("day")?;
3836                exprs.push(expr);
3837                ScalarFunc::GeneratedExpr
3838            }
3839            "day_of_week" => {
3840                // date_part('dow', time_index)
3841                let expr = self.date_part_on_time_index("dow")?;
3842                exprs.push(expr);
3843                ScalarFunc::GeneratedExpr
3844            }
3845            "day_of_year" => {
3846                // date_part('doy', time_index)
3847                let expr = self.date_part_on_time_index("doy")?;
3848                exprs.push(expr);
3849                ScalarFunc::GeneratedExpr
3850            }
3851            "days_in_month" => {
3852                // date_part(
3853                //     'days',
3854                //     (date_trunc('month', <TIME INDEX>::date) + interval '1 month - 1 day')
3855                // );
3856                let day_lit_expr = "day".lit();
3857                let month_lit_expr = "month".lit();
3858                let interval_1month_lit_expr =
3859                    DfExpr::Literal(ScalarValue::IntervalYearMonth(Some(1)), None);
3860                let interval_1day_lit_expr = DfExpr::Literal(
3861                    ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(1, 0))),
3862                    None,
3863                );
3864                let the_1month_minus_1day_expr = DfExpr::BinaryExpr(BinaryExpr {
3865                    left: Box::new(interval_1month_lit_expr),
3866                    op: Operator::Minus,
3867                    right: Box::new(interval_1day_lit_expr),
3868                });
3869                let date_trunc_expr = DfExpr::ScalarFunction(ScalarFunction {
3870                    func: datafusion_functions::datetime::date_trunc(),
3871                    args: vec![month_lit_expr, self.create_time_index_column_expr()?],
3872                });
3873                let date_trunc_plus_interval_expr = DfExpr::BinaryExpr(BinaryExpr {
3874                    left: Box::new(date_trunc_expr),
3875                    op: Operator::Plus,
3876                    right: Box::new(the_1month_minus_1day_expr),
3877                });
3878                let date_part_expr = DfExpr::ScalarFunction(ScalarFunction {
3879                    func: datafusion_functions::datetime::date_part(),
3880                    args: vec![day_lit_expr, date_trunc_plus_interval_expr],
3881                });
3882
3883                exprs.push(date_part_expr);
3884                ScalarFunc::GeneratedExpr
3885            }
3886
3887            "label_join" => {
3888                self.ctx.use_tsid = false;
3889                let (concat_expr, dst_label) = Self::build_concat_labels_expr(
3890                    &mut other_input_exprs,
3891                    &self.ctx,
3892                    query_engine_state,
3893                )?;
3894
3895                // Reserve the current field columns except the `dst_label`.
3896                for value in &self.ctx.field_columns {
3897                    if *value != dst_label {
3898                        let expr = DfExpr::Column(Column::from_name(value));
3899                        exprs.push(expr);
3900                    }
3901                }
3902
3903                // Remove it from tag columns if exists to avoid duplicated column names
3904                self.ctx.tag_columns.retain(|tag| *tag != dst_label);
3905                new_tags.push(dst_label);
3906                // Add the new label expr to evaluate
3907                exprs.push(concat_expr);
3908
3909                ScalarFunc::GeneratedExpr
3910            }
3911            "label_replace" => {
3912                self.ctx.use_tsid = false;
3913                if let Some((replace_expr, dst_label)) = self
3914                    .build_regexp_replace_label_expr(&mut other_input_exprs, query_engine_state)?
3915                {
3916                    // Reserve the current field columns except the `dst_label`.
3917                    for value in &self.ctx.field_columns {
3918                        if *value != dst_label {
3919                            let expr = DfExpr::Column(Column::from_name(value));
3920                            exprs.push(expr);
3921                        }
3922                    }
3923
3924                    ensure!(
3925                        !self.ctx.tag_columns.contains(&dst_label),
3926                        SameLabelSetSnafu
3927                    );
3928                    new_tags.push(dst_label);
3929                    // Add the new label expr to evaluate
3930                    exprs.push(replace_expr);
3931                } else {
3932                    // Keep the current field columns
3933                    for value in &self.ctx.field_columns {
3934                        let expr = DfExpr::Column(Column::from_name(value));
3935                        exprs.push(expr);
3936                    }
3937                }
3938
3939                ScalarFunc::GeneratedExpr
3940            }
3941            "sort" | "sort_desc" => {
3942                // Value sorting silently ignores native histogram samples.
3943                for value in &self.ctx.field_columns {
3944                    if !Self::field_column_is_native_histogram(input_schema, value) {
3945                        exprs.push(DfExpr::Column(Column::from_name(value)));
3946                    }
3947                }
3948                // Keep a nullable float field so the normal empty-value filter produces an
3949                // empty vector when the input contains only histograms.
3950                if exprs.is_empty() {
3951                    exprs.push(DfExpr::Literal(ScalarValue::Float64(None), None));
3952                }
3953
3954                ScalarFunc::GeneratedExpr
3955            }
3956            "sort_by_label" | "sort_by_label_desc" | "timestamp" => {
3957                // These functions are not expression but a part of plan,
3958                // they are processed by `prom_call_expr_to_plan`.
3959                for value in &self.ctx.field_columns {
3960                    let expr = DfExpr::Column(Column::from_name(value));
3961                    exprs.push(expr);
3962                }
3963
3964                ScalarFunc::GeneratedExpr
3965            }
3966            "round" if self.all_field_columns_are_native_histograms(input_schema) => {
3967                if other_input_exprs.is_empty() {
3968                    other_input_exprs.push_front(0.0f64.lit());
3969                }
3970                ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3971            }
3972            "round" => {
3973                if other_input_exprs.is_empty() {
3974                    other_input_exprs.push_front(0.0f64.lit());
3975                }
3976                ScalarFunc::DataFusionUdf(Arc::new(Round::scalar_udf()))
3977            }
3978            "rad" | "deg" | "sgn" if self.all_field_columns_are_native_histograms(input_schema) => {
3979                ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3980            }
3981            "rad" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::radians()),
3982            "deg" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::degrees()),
3983            "sgn" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::signum()),
3984            "pi" => {
3985                // pi functions doesn't accepts any arguments, needs special processing
3986                let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
3987                    func: datafusion::functions::math::pi(),
3988                    args: vec![],
3989                });
3990                exprs.push(fn_expr);
3991
3992                ScalarFunc::GeneratedExpr
3993            }
3994            _ => {
3995                if let Some(f) = query_engine_state
3996                    .session_state()
3997                    .scalar_functions()
3998                    .get(func.name)
3999                {
4000                    if self.all_field_columns_are_native_histograms(input_schema) {
4001                        ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
4002                    } else {
4003                        ScalarFunc::DataFusionBuiltin(f.clone())
4004                    }
4005                } else if let Some(factory) = query_engine_state.scalar_function(func.name) {
4006                    if self.all_field_columns_are_native_histograms(input_schema) {
4007                        ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
4008                    } else {
4009                        let func_state = query_engine_state.function_state();
4010                        let query_ctx = self.table_provider.query_ctx();
4011
4012                        ScalarFunc::DataFusionUdf(Arc::new(factory.provide(FunctionContext {
4013                            state: func_state,
4014                            query_ctx: query_ctx.clone(),
4015                        })))
4016                    }
4017                } else if let Some(f) = datafusion_functions::math::functions()
4018                    .iter()
4019                    .find(|f| f.name() == func.name)
4020                {
4021                    if self.all_field_columns_are_native_histograms(input_schema) {
4022                        ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
4023                    } else {
4024                        ScalarFunc::DataFusionUdf(f.clone())
4025                    }
4026                } else {
4027                    return UnsupportedExprSnafu {
4028                        name: func.name.to_string(),
4029                    }
4030                    .fail();
4031                }
4032            }
4033        };
4034
4035        for value in &self.ctx.field_columns {
4036            let col_expr = DfExpr::Column(Column::from_name(value));
4037            let value_is_histogram = Self::field_column_is_native_histogram(input_schema, value);
4038
4039            match scalar_func.clone() {
4040                ScalarFunc::DataFusionBuiltin(func) => {
4041                    if alternative_samples && value_is_histogram {
4042                        continue;
4043                    }
4044                    other_input_exprs.insert(field_column_pos, col_expr);
4045                    let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
4046                        func,
4047                        args: other_input_exprs.clone().into(),
4048                    });
4049                    exprs.push(fn_expr);
4050                    let _ = other_input_exprs.remove(field_column_pos);
4051                }
4052                ScalarFunc::DataFusionUdf(func) => {
4053                    if alternative_samples && value_is_histogram {
4054                        continue;
4055                    }
4056                    let args = itertools::chain!(
4057                        other_input_exprs.iter().take(field_column_pos).cloned(),
4058                        std::iter::once(col_expr),
4059                        other_input_exprs.iter().skip(field_column_pos).cloned()
4060                    )
4061                    .collect_vec();
4062                    exprs.push(DfExpr::ScalarFunction(ScalarFunction { func, args }))
4063                }
4064                ScalarFunc::NativeHistogramUdf(func) => {
4065                    if value_is_histogram {
4066                        let args = itertools::chain!(
4067                            other_input_exprs.iter().take(field_column_pos).cloned(),
4068                            std::iter::once(col_expr),
4069                            other_input_exprs.iter().skip(field_column_pos).cloned()
4070                        )
4071                        .collect_vec();
4072                        exprs.push(DfExpr::ScalarFunction(ScalarFunction { func, args }));
4073                    } else if !alternative_samples {
4074                        exprs.push(
4075                            DfExpr::Literal(ScalarValue::Float64(None), None).alias(format!(
4076                                "{}_{}",
4077                                func.name(),
4078                                value
4079                            )),
4080                        );
4081                    }
4082                }
4083                ScalarFunc::Udf(func) => {
4084                    let ts_range_expr = DfExpr::Column(Column::from_name(
4085                        RangeManipulate::build_timestamp_range_name(
4086                            self.ctx.time_index_column.as_ref().unwrap(),
4087                        ),
4088                    ));
4089                    other_input_exprs.insert(field_column_pos, ts_range_expr);
4090                    other_input_exprs.insert(field_column_pos + 1, col_expr);
4091                    let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
4092                        func,
4093                        args: other_input_exprs.clone().into(),
4094                    });
4095                    exprs.push(fn_expr);
4096                    let _ = other_input_exprs.remove(field_column_pos + 1);
4097                    let _ = other_input_exprs.remove(field_column_pos);
4098                }
4099                ScalarFunc::ExtrapolateUdf(udf, range_length) => {
4100                    let ts_range_expr = DfExpr::Column(Column::from_name(
4101                        RangeManipulate::build_timestamp_range_name(
4102                            self.ctx.time_index_column.as_ref().unwrap(),
4103                        ),
4104                    ));
4105                    other_input_exprs.insert(field_column_pos, ts_range_expr.clone());
4106                    other_input_exprs.insert(field_column_pos + 1, col_expr.clone());
4107                    other_input_exprs
4108                        .insert(field_column_pos + 2, self.create_time_index_column_expr()?);
4109                    other_input_exprs.push_back(lit(range_length));
4110                    let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
4111                        func: udf,
4112                        args: other_input_exprs.clone().into(),
4113                    });
4114                    let fn_expr = if matches!(func.name, "rate" | "increase")
4115                        && !all_field_columns_are_native_histogram_ranges
4116                    {
4117                        let delta_sum = DfExpr::ScalarFunction(ScalarFunction {
4118                            func: Arc::new(SumOverTime::scalar_udf()),
4119                            args: vec![ts_range_expr, col_expr],
4120                        });
4121                        self.select_delta_range_math(
4122                            func.name,
4123                            input_schema,
4124                            range_length,
4125                            delta_sum,
4126                            fn_expr,
4127                        )?
4128                    } else {
4129                        fn_expr
4130                    };
4131                    exprs.push(fn_expr);
4132                    let _ = other_input_exprs.pop_back();
4133                    let _ = other_input_exprs.remove(field_column_pos + 2);
4134                    let _ = other_input_exprs.remove(field_column_pos + 1);
4135                    let _ = other_input_exprs.remove(field_column_pos);
4136                }
4137                ScalarFunc::GeneratedExpr => {}
4138            }
4139        }
4140
4141        // Update value columns' name, and alias them to remove qualifiers
4142        // For label functions such as `label_join`, `label_replace`, etc.,
4143        // we keep the fields unchanged.
4144        if !matches!(func.name, "label_join" | "label_replace") {
4145            let mut new_field_columns = Vec::with_capacity(exprs.len());
4146
4147            exprs = exprs
4148                .into_iter()
4149                .map(|expr| {
4150                    let display_name = expr.schema_name().to_string();
4151                    new_field_columns.push(display_name.clone());
4152                    Ok(expr.alias(display_name))
4153                })
4154                .collect::<std::result::Result<Vec<_>, _>>()
4155                .context(DataFusionPlanningSnafu)?;
4156
4157            self.ctx.field_columns = new_field_columns;
4158        }
4159
4160        Ok((exprs, new_tags))
4161    }
4162
4163    fn select_delta_range_math(
4164        &self,
4165        function: &str,
4166        input_schema: &DFSchemaRef,
4167        range_length: Millisecond,
4168        delta_sum: DfExpr,
4169        cumulative: DfExpr,
4170    ) -> Result<DfExpr> {
4171        let marker_is_delta = if self
4172            .ctx
4173            .tag_columns
4174            .iter()
4175            .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL)
4176        {
4177            Self::field_column_type(input_schema, OTLP_AGGREGATION_TEMPORALITY_LABEL)
4178                .filter(|data_type| Self::string_value_data_type(data_type).is_some())
4179                .map(|_| {
4180                    DfExpr::Column(Column::from_name(OTLP_AGGREGATION_TEMPORALITY_LABEL))
4181                        .eq(lit(GREPTIME_TEMPORALITY_DELTA))
4182                })
4183        } else {
4184            None
4185        };
4186        let Some(marker_is_delta) = marker_is_delta else {
4187            return Ok(cumulative);
4188        };
4189
4190        let delta = if function == "rate" {
4191            DfExpr::BinaryExpr(BinaryExpr {
4192                left: Box::new(delta_sum),
4193                op: Operator::Divide,
4194                right: Box::new(lit(range_length as f64 / 1000.0)),
4195            })
4196        } else {
4197            delta_sum
4198        };
4199        let display_name = cumulative.schema_name().to_string();
4200        when(marker_is_delta, delta)
4201            .otherwise(cumulative)
4202            .context(DataFusionPlanningSnafu)
4203            .map(|expr| expr.alias(display_name))
4204    }
4205
4206    /// Validate label name according to Prometheus specification.
4207    /// Label names must match the regex: [a-zA-Z_][a-zA-Z0-9_]*
4208    /// Additionally, label names starting with double underscores are reserved for internal use.
4209    fn validate_label_name(label_name: &str) -> Result<()> {
4210        // Check if label name starts with double underscores (reserved)
4211        if label_name.starts_with("__") {
4212            return InvalidDestinationLabelNameSnafu { label_name }.fail();
4213        }
4214        // Check if label name matches the required pattern
4215        if !LABEL_NAME_REGEX.is_match(label_name) {
4216            return InvalidDestinationLabelNameSnafu { label_name }.fail();
4217        }
4218
4219        Ok(())
4220    }
4221
4222    /// Build expr for `label_replace` function
4223    fn build_regexp_replace_label_expr(
4224        &self,
4225        other_input_exprs: &mut VecDeque<DfExpr>,
4226        query_engine_state: &QueryEngineState,
4227    ) -> Result<Option<(DfExpr, String)>> {
4228        // label_replace(vector, dst_label, replacement, src_label, regex)
4229        let dst_label = match other_input_exprs.pop_front() {
4230            Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4231            other => UnexpectedPlanExprSnafu {
4232                desc: format!("expected dst_label string literal, but found {:?}", other),
4233            }
4234            .fail()?,
4235        };
4236
4237        // Validate the destination label name
4238        Self::validate_label_name(&dst_label)?;
4239        let replacement = match other_input_exprs.pop_front() {
4240            Some(DfExpr::Literal(ScalarValue::Utf8(Some(r)), _)) => r,
4241            other => UnexpectedPlanExprSnafu {
4242                desc: format!("expected replacement string literal, but found {:?}", other),
4243            }
4244            .fail()?,
4245        };
4246        let src_label = match other_input_exprs.pop_front() {
4247            Some(DfExpr::Literal(ScalarValue::Utf8(Some(s)), None)) => s,
4248            other => UnexpectedPlanExprSnafu {
4249                desc: format!("expected src_label string literal, but found {:?}", other),
4250            }
4251            .fail()?,
4252        };
4253
4254        let regex = match other_input_exprs.pop_front() {
4255            Some(DfExpr::Literal(ScalarValue::Utf8(Some(r)), None)) => r,
4256            other => UnexpectedPlanExprSnafu {
4257                desc: format!("expected regex string literal, but found {:?}", other),
4258            }
4259            .fail()?,
4260        };
4261
4262        // Validate the regex before using it
4263        // doc: https://prometheus.io/docs/prometheus/latest/querying/functions/#label_replace
4264        regex::Regex::new(&regex).map_err(|_| {
4265            InvalidRegularExpressionSnafu {
4266                regex: regex.clone(),
4267            }
4268            .build()
4269        })?;
4270
4271        // If the src_label exists and regex is empty, keep everything unchanged.
4272        if self.ctx.tag_columns.contains(&src_label) && regex.is_empty() {
4273            return Ok(None);
4274        }
4275
4276        // If the src_label doesn't exists, and
4277        if !self.ctx.tag_columns.contains(&src_label) {
4278            if replacement.is_empty() {
4279                // the replacement is empty, keep everything unchanged.
4280                return Ok(None);
4281            } else {
4282                // the replacement is not empty, always adds dst_label with replacement value.
4283                return Ok(Some((
4284                    // alias literal `replacement` as dst_label
4285                    lit(replacement).alias(&dst_label),
4286                    dst_label,
4287                )));
4288            }
4289        }
4290
4291        // Preprocess the regex:
4292        // https://github.com/prometheus/prometheus/blob/d902abc50d6652ba8fe9a81ff8e5cce936114eba/promql/functions.go#L1575C32-L1575C37
4293        let regex = format!("^(?s:{regex})$");
4294
4295        let session_state = query_engine_state.session_state();
4296        let func = session_state
4297            .scalar_functions()
4298            .get("regexp_replace")
4299            .context(UnsupportedExprSnafu {
4300                name: "regexp_replace",
4301            })?;
4302
4303        // regexp_replace(src_label, regex, replacement)
4304        let args = vec![
4305            if src_label.is_empty() {
4306                DfExpr::Literal(ScalarValue::Utf8(Some(String::new())), None)
4307            } else {
4308                DfExpr::Column(Column::from_name(src_label))
4309            },
4310            DfExpr::Literal(ScalarValue::Utf8(Some(regex)), None),
4311            DfExpr::Literal(ScalarValue::Utf8(Some(replacement)), None),
4312        ];
4313
4314        Ok(Some((
4315            DfExpr::ScalarFunction(ScalarFunction {
4316                func: func.clone(),
4317                args,
4318            })
4319            .alias(&dst_label),
4320            dst_label,
4321        )))
4322    }
4323
4324    /// Build expr for `label_join` function
4325    fn build_concat_labels_expr(
4326        other_input_exprs: &mut VecDeque<DfExpr>,
4327        ctx: &PromPlannerContext,
4328        query_engine_state: &QueryEngineState,
4329    ) -> Result<(DfExpr, String)> {
4330        // label_join(vector, dst_label, separator, src_label_1, src_label_2, ...)
4331
4332        let dst_label = match other_input_exprs.pop_front() {
4333            Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4334            other => UnexpectedPlanExprSnafu {
4335                desc: format!("expected dst_label string literal, but found {:?}", other),
4336            }
4337            .fail()?,
4338        };
4339        let separator = match other_input_exprs.pop_front() {
4340            Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4341            other => UnexpectedPlanExprSnafu {
4342                desc: format!("expected separator string literal, but found {:?}", other),
4343            }
4344            .fail()?,
4345        };
4346
4347        // Create a set of available columns (tag columns + field columns + time index column)
4348        let available_columns: HashSet<&str> = ctx
4349            .tag_columns
4350            .iter()
4351            .chain(ctx.field_columns.iter())
4352            .chain(ctx.time_index_column.as_ref())
4353            .map(|s| s.as_str())
4354            .collect();
4355
4356        let src_labels = other_input_exprs
4357            .iter()
4358            .map(|expr| {
4359                // Cast source label into column or null literal
4360                match expr {
4361                    DfExpr::Literal(ScalarValue::Utf8(Some(label)), None) => {
4362                        if label.is_empty() {
4363                            Ok(DfExpr::Literal(ScalarValue::Null, None))
4364                        } else if available_columns.contains(label.as_str()) {
4365                            // Label exists in the table schema
4366                            Ok(DfExpr::Column(Column::from_name(label)))
4367                        } else {
4368                            // Label doesn't exist, treat as empty string (null)
4369                            Ok(DfExpr::Literal(ScalarValue::Null, None))
4370                        }
4371                    }
4372                    other => UnexpectedPlanExprSnafu {
4373                        desc: format!(
4374                            "expected source label string literal, but found {:?}",
4375                            other
4376                        ),
4377                    }
4378                    .fail(),
4379                }
4380            })
4381            .collect::<Result<Vec<_>>>()?;
4382        ensure!(
4383            !src_labels.is_empty(),
4384            FunctionInvalidArgumentSnafu {
4385                fn_name: "label_join"
4386            }
4387        );
4388
4389        let session_state = query_engine_state.session_state();
4390        let func = session_state
4391            .scalar_functions()
4392            .get("concat_ws")
4393            .context(UnsupportedExprSnafu { name: "concat_ws" })?;
4394
4395        // concat_ws(separator, src_label_1, src_label_2, ...) as dst_label
4396        let mut args = Vec::with_capacity(1 + src_labels.len());
4397        args.push(DfExpr::Literal(ScalarValue::Utf8(Some(separator)), None));
4398        args.extend(src_labels);
4399
4400        Ok((
4401            DfExpr::ScalarFunction(ScalarFunction {
4402                func: func.clone(),
4403                args,
4404            })
4405            .alias(&dst_label),
4406            dst_label,
4407        ))
4408    }
4409
4410    fn create_time_index_column_expr(&self) -> Result<DfExpr> {
4411        Ok(DfExpr::Column(Column::from_name(
4412            self.ctx
4413                .time_index_column
4414                .clone()
4415                .with_context(|| TimeIndexNotFoundSnafu { table: "unknown" })?,
4416        )))
4417    }
4418
4419    fn create_tag_column_exprs(&self) -> Result<Vec<DfExpr>> {
4420        let mut result = Vec::with_capacity(self.ctx.tag_columns.len());
4421        for tag in &self.ctx.tag_columns {
4422            let expr = DfExpr::Column(Column::from_name(tag));
4423            result.push(expr);
4424        }
4425        Ok(result)
4426    }
4427
4428    fn create_field_column_exprs(&self) -> Result<Vec<DfExpr>> {
4429        let mut result = Vec::with_capacity(self.ctx.field_columns.len());
4430        for field in &self.ctx.field_columns {
4431            let expr = DfExpr::Column(Column::from_name(field));
4432            result.push(expr);
4433        }
4434        Ok(result)
4435    }
4436
4437    fn create_tag_and_time_index_column_sort_exprs(&self) -> Result<Vec<SortExpr>> {
4438        let mut result = self
4439            .ctx
4440            .tag_columns
4441            .iter()
4442            .map(|col| DfExpr::Column(Column::from_name(col)).sort(true, true))
4443            .collect::<Vec<_>>();
4444        result.push(self.create_time_index_column_expr()?.sort(true, true));
4445        Ok(result)
4446    }
4447
4448    fn create_field_columns_sort_exprs(&self, asc: bool) -> Vec<SortExpr> {
4449        self.ctx
4450            .field_columns
4451            .iter()
4452            .map(|col| DfExpr::Column(Column::from_name(col)).sort(asc, true))
4453            .collect::<Vec<_>>()
4454    }
4455
4456    fn create_sort_exprs_by_tags(
4457        func: &str,
4458        tags: Vec<DfExpr>,
4459        asc: bool,
4460    ) -> Result<Vec<SortExpr>> {
4461        ensure!(
4462            !tags.is_empty(),
4463            FunctionInvalidArgumentSnafu { fn_name: func }
4464        );
4465
4466        tags.iter()
4467            .map(|col| match col {
4468                DfExpr::Literal(ScalarValue::Utf8(Some(label)), _) => {
4469                    Ok(DfExpr::Column(Column::from_name(label)).sort(asc, false))
4470                }
4471                other => UnexpectedPlanExprSnafu {
4472                    desc: format!("expected label string literal, but found {:?}", other),
4473                }
4474                .fail(),
4475            })
4476            .collect::<Result<Vec<_>>>()
4477    }
4478
4479    fn create_empty_values_filter_expr(&self, preserve_any_value: bool) -> Result<DfExpr> {
4480        let mut exprs = Vec::with_capacity(self.ctx.field_columns.len());
4481        for value in &self.ctx.field_columns {
4482            let expr = DfExpr::Column(Column::from_name(value)).is_not_null();
4483            exprs.push(expr);
4484        }
4485
4486        // This error context should be computed lazily: the planner may set `ctx.table_name` to
4487        // `None` for derived expressions (e.g. after projecting the LHS of a vector-vector
4488        // comparison filter). Eagerly calling `table_ref()?` here can turn a valid plan into
4489        // a `TableNameNotFound` error even when predicate construction succeeds.
4490        let predicate = if preserve_any_value {
4491            disjunction(exprs)
4492        } else {
4493            conjunction(exprs)
4494        };
4495        predicate.with_context(|| ValueNotFoundSnafu {
4496            table: self
4497                .table_ref()
4498                .map(|t| t.to_quoted_string())
4499                .unwrap_or_else(|_| "unknown".to_string()),
4500        })
4501    }
4502
4503    /// Creates a set of DataFusion `DfExpr::AggregateFunction` expressions for each value column using the specified aggregate function.
4504    ///
4505    /// # Side Effects
4506    ///
4507    /// This method modifies the value columns in the context by replacing them with the new columns
4508    /// created by the aggregate function application.
4509    ///
4510    /// # Returns
4511    ///
4512    /// Returns a tuple of `(aggregate_expressions, previous_field_expressions)` where:
4513    /// - `aggregate_expressions`: Expressions that apply the aggregate function to the original fields
4514    /// - `previous_field_expressions`: Field expressions naming the pre-aggregation values. This is
4515    ///   non-empty only when the operation is `count_values`, which groups by the sample value and
4516    ///   projects it as the generated label, so these expressions are passed through the same
4517    ///   formatting as that label (`prom_float_to_string`).
4518    ///
4519    fn create_aggregate_exprs(
4520        &mut self,
4521        op: TokenType,
4522        param: &Option<Box<PromExpr>>,
4523        input_plan: &LogicalPlan,
4524    ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4525        let mixed_sample_columns =
4526            Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns)
4527                .map(|(float, histogram)| (float.to_string(), histogram.to_string()));
4528        let is_group_agg = op.id() == token::T_GROUP;
4529        if is_group_agg && mixed_sample_columns.is_none() {
4530            ensure!(
4531                self.ctx.field_columns.len() == 1,
4532                MultiFieldsNotSupportedSnafu {
4533                    operator: "group()"
4534                }
4535            );
4536        }
4537
4538        if let Some((float, histogram)) = mixed_sample_columns {
4539            return self.create_mixed_aggregate_exprs(op, param, &float, &histogram);
4540        }
4541
4542        if self.all_field_columns_are_native_histograms(input_plan.schema()) {
4543            return self.create_native_histogram_aggregate_exprs(op, input_plan);
4544        }
4545
4546        // perform aggregate operation to each value column
4547        let exprs = self
4548            .ctx
4549            .field_columns
4550            .iter()
4551            .map(|col| {
4552                Self::create_numeric_aggregate_expr(
4553                    op,
4554                    param,
4555                    DfExpr::Column(Column::from_name(col)),
4556                )
4557            })
4558            .collect::<Result<Vec<_>>>()?;
4559
4560        // if the aggregator is `count_values`, it must be grouped by current fields.
4561        //
4562        // The grouping key is the *formatted* sample value, i.e. the same expression that
4563        // produces the generated label below: PromQL groups by the value, and the label is
4564        // that value in Prometheus' textual form (`strconv.FormatFloat(value, 'f', -1, 64)`),
4565        // so grouping by the raw value would split samples that render to one label into
4566        // several groups, each emitting the same label set for one timestamp.
4567        let prev_field_exprs = if op.id() == token::T_COUNT_VALUES {
4568            ensure!(
4569                self.ctx.field_columns.len() == 1,
4570                UnsupportedExprSnafu {
4571                    name: "count_values on multi-value input"
4572                }
4573            );
4574
4575            self.ctx
4576                .field_columns
4577                .iter()
4578                .map(|col| {
4579                    let value = DfExpr::Column(Column::from_name(col));
4580                    // Normalize non `Float64` inputs the same way the label projection does,
4581                    // so both sides agree on the formatted value: `prom_float_to_string`
4582                    // formats exactly like Prometheus, while arrow's `Float64 -> Utf8` cast
4583                    // would render `1.0`.
4584                    let value = if Self::field_column_type(input_plan.schema(), col)
4585                        == Some(&ArrowDataType::Float64)
4586                    {
4587                        value
4588                    } else {
4589                        DfExpr::Cast(Cast::new(Box::new(value), ArrowDataType::Float64))
4590                    };
4591                    DfExpr::ScalarFunction(ScalarFunction {
4592                        func: Arc::new(PromqlFloatToString::scalar_udf()),
4593                        args: vec![value],
4594                    })
4595                })
4596                .collect()
4597        } else {
4598            vec![]
4599        };
4600
4601        // update value column name according to the aggregators,
4602        let mut new_field_columns = Vec::with_capacity(self.ctx.field_columns.len());
4603
4604        let normalized_exprs =
4605            normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
4606        for expr in normalized_exprs {
4607            new_field_columns.push(expr.schema_name().to_string());
4608        }
4609        self.ctx.field_columns = new_field_columns;
4610
4611        Ok((exprs, prev_field_exprs))
4612    }
4613
4614    fn create_numeric_aggregate_expr(
4615        op: TokenType,
4616        param: &Option<Box<PromExpr>>,
4617        input: DfExpr,
4618    ) -> Result<DfExpr> {
4619        let expr = match op.id() {
4620            token::T_SUM => sum_udaf().call(vec![input]),
4621            token::T_QUANTILE => {
4622                let q = Self::get_param_as_literal_expr(
4623                    param.as_deref(),
4624                    Some(op),
4625                    Some(ArrowDataType::Float64),
4626                )?;
4627                quantile_udaf().call(vec![q, input])
4628            }
4629            token::T_AVG => avg_udaf().call(vec![input]),
4630            token::T_COUNT_VALUES | token::T_COUNT => count_udaf().call(vec![input]),
4631            token::T_MIN => min_udaf().call(vec![input]),
4632            token::T_MAX => max_udaf().call(vec![input]),
4633            // PromQL's `group()` aggregator produces 1 for each group.
4634            // Use `max(1.0)` (per-group) to match semantics and output type (Float64).
4635            token::T_GROUP => max_udaf().call(vec![lit(1_f64)]),
4636            token::T_STDDEV => stddev_pop_udaf().call(vec![input]),
4637            token::T_STDVAR => var_pop_udaf().call(vec![input]),
4638            token::T_TOPK | token::T_BOTTOMK => {
4639                return UnsupportedExprSnafu {
4640                    name: format!("{op:?}"),
4641                }
4642                .fail();
4643            }
4644            _ => return UnexpectedTokenSnafu { token: op }.fail(),
4645        };
4646        Ok(expr)
4647    }
4648
4649    fn create_mixed_aggregate_exprs(
4650        &mut self,
4651        op: TokenType,
4652        param: &Option<Box<PromExpr>>,
4653        float_column: &str,
4654        histogram_column: &str,
4655    ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4656        let float_input = DfExpr::Column(Column::from_name(float_column));
4657        let histogram_input = DfExpr::Column(Column::from_name(histogram_column));
4658        let float_count = count_udaf().call(vec![float_input.clone()]);
4659        let histogram_count = count_udaf().call(vec![histogram_input.clone()]);
4660        let mixed_sample_value = || {
4661            DfExpr::ScalarFunction(ScalarFunction {
4662                func: coalesce(),
4663                args: vec![
4664                    DfExpr::ScalarFunction(ScalarFunction {
4665                        func: Arc::new(PromqlFloatToString::scalar_udf()),
4666                        args: vec![float_input.clone()],
4667                    }),
4668                    DfExpr::ScalarFunction(ScalarFunction {
4669                        func: Arc::new(NativeHistogramToString::scalar_udf()),
4670                        args: vec![histogram_input.clone()],
4671                    }),
4672                ],
4673            })
4674        };
4675
4676        let (exprs, prev_field_exprs, field_columns) = match op.id() {
4677            token::T_SUM | token::T_AVG => (
4678                vec![
4679                    Self::create_numeric_aggregate_expr(op, param, float_input)?
4680                        .alias(float_column),
4681                    self.create_native_histogram_aggregate_expr(op, histogram_column)?,
4682                    float_count.alias(Self::mixed_sample_count_name(float_column)),
4683                    histogram_count.alias(Self::mixed_sample_count_name(histogram_column)),
4684                ],
4685                vec![],
4686                vec![float_column.to_string(), histogram_column.to_string()],
4687            ),
4688            token::T_COUNT => {
4689                let present = when(
4690                    float_input
4691                        .clone()
4692                        .is_not_null()
4693                        .or(histogram_input.clone().is_not_null()),
4694                    lit(1_i64),
4695                )
4696                .otherwise(lit(ScalarValue::Int64(None)))
4697                .context(DataFusionPlanningSnafu)?;
4698                (
4699                    vec![count_udaf().call(vec![present]).alias(float_column)],
4700                    vec![],
4701                    vec![float_column.to_string()],
4702                )
4703            }
4704            token::T_GROUP => (
4705                vec![max_udaf().call(vec![lit(1_f64)]).alias(float_column)],
4706                vec![],
4707                vec![float_column.to_string()],
4708            ),
4709            token::T_COUNT_VALUES => {
4710                let value = mixed_sample_value();
4711                (
4712                    vec![count_udaf().call(vec![value.clone()]).alias(float_column)],
4713                    vec![value],
4714                    vec![float_column.to_string()],
4715                )
4716            }
4717            token::T_MIN | token::T_MAX | token::T_STDDEV | token::T_STDVAR | token::T_QUANTILE => {
4718                (
4719                    vec![
4720                        Self::create_numeric_aggregate_expr(op, param, float_input)?
4721                            .alias(float_column),
4722                        histogram_count.alias(Self::mixed_sample_count_name(histogram_column)),
4723                    ],
4724                    vec![],
4725                    vec![float_column.to_string()],
4726                )
4727            }
4728            token::T_TOPK | token::T_BOTTOMK => {
4729                return UnsupportedExprSnafu {
4730                    name: format!("{op:?}"),
4731                }
4732                .fail();
4733            }
4734            _ => return UnexpectedTokenSnafu { token: op }.fail(),
4735        };
4736
4737        self.ctx.field_columns = field_columns;
4738        Ok((exprs, prev_field_exprs))
4739    }
4740
4741    fn mixed_sample_count_column(column: &str) -> DfExpr {
4742        DfExpr::Column(Column::from_name(Self::mixed_sample_count_name(column)))
4743    }
4744
4745    fn mixed_sample_count_name(column: &str) -> String {
4746        format!("__promql_sample_count({column})")
4747    }
4748
4749    fn mixed_aggregate_filter_expr(
4750        &self,
4751        op: TokenType,
4752        float_column: &str,
4753        histogram_column: &str,
4754    ) -> Result<DfExpr> {
4755        let float_count = Self::mixed_sample_count_column(float_column);
4756        let histogram_count = Self::mixed_sample_count_column(histogram_column);
4757        let mixed = float_count
4758            .clone()
4759            .gt(lit(0_i64))
4760            .and(histogram_count.clone().gt(lit(0_i64)));
4761        let drop_mixed = DfExpr::ScalarFunction(ScalarFunction {
4762            func: Arc::new(NativeHistogramDrop::warning_bool_false_udf(
4763                format!(
4764                    "{op}: dropped aggregation result containing both float and native histogram samples"
4765                ),
4766                self.promql_annotations.clone(),
4767            )),
4768            args: vec![float_count, histogram_count],
4769        });
4770
4771        when(mixed, drop_mixed)
4772            .otherwise(lit(true))
4773            .context(DataFusionPlanningSnafu)
4774    }
4775
4776    fn mixed_ignored_histogram_filter_expr(
4777        &self,
4778        op: TokenType,
4779        histogram_column: &str,
4780    ) -> Result<DfExpr> {
4781        let histogram_count = Self::mixed_sample_count_column(histogram_column);
4782        let has_histograms = histogram_count.clone().gt(lit(0_i64));
4783        let record_info = DfExpr::ScalarFunction(ScalarFunction {
4784            func: Arc::new(NativeHistogramDrop::bool_true_udf(
4785                format!(
4786                    "{op}: dropped native histogram samples because this aggregation is not supported for native histograms"
4787                ),
4788                self.promql_annotations.clone(),
4789            )),
4790            args: vec![histogram_count],
4791        });
4792
4793        when(has_histograms, record_info)
4794            .otherwise(lit(true))
4795            .context(DataFusionPlanningSnafu)
4796    }
4797
4798    fn create_native_histogram_aggregate_expr(
4799        &self,
4800        op: TokenType,
4801        column: &str,
4802    ) -> Result<DfExpr> {
4803        let input = DfExpr::Column(Column::from_name(column));
4804        let expr = match op.id() {
4805            token::T_SUM => Arc::new(NativeHistogramAggSum::aggregate_udf_with_collector(
4806                self.promql_annotations.clone(),
4807            ))
4808            .call(vec![input])
4809            .alias(column),
4810            token::T_AVG => Arc::new(NativeHistogramAggAvg::aggregate_udf_with_collector(
4811                self.promql_annotations.clone(),
4812            ))
4813            .call(vec![input])
4814            .alias(column),
4815            token::T_COUNT_VALUES | token::T_COUNT => {
4816                count_udaf().call(vec![input]).alias(column)
4817            }
4818            token::T_GROUP => max_udaf().call(vec![lit(1_f64)]).alias(column),
4819            token::T_MIN
4820            | token::T_MAX
4821            | token::T_STDDEV
4822            | token::T_STDVAR
4823            | token::T_QUANTILE
4824            | token::T_TOPK
4825            | token::T_BOTTOMK => sum_udaf()
4826                .call(vec![DfExpr::ScalarFunction(ScalarFunction {
4827                    func: Arc::new(NativeHistogramDrop::float_null_udf(
4828                        format!(
4829                            "{op}: dropped native histogram samples because this aggregation is not supported for native histograms"
4830                        ),
4831                        self.promql_annotations.clone(),
4832                    )),
4833                    args: vec![input],
4834                })])
4835                .alias(column),
4836            _ => return UnexpectedTokenSnafu { token: op }.fail(),
4837        };
4838        Ok(expr)
4839    }
4840
4841    fn create_native_histogram_aggregate_exprs(
4842        &mut self,
4843        op: TokenType,
4844        input_plan: &LogicalPlan,
4845    ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4846        let prev_field_exprs = if op.id() == token::T_COUNT_VALUES {
4847            ensure!(
4848                self.ctx.field_columns.len() == 1,
4849                UnsupportedExprSnafu {
4850                    name: "count_values on multi-value input"
4851                }
4852            );
4853            self.ctx
4854                .field_columns
4855                .iter()
4856                .map(|col| {
4857                    DfExpr::ScalarFunction(ScalarFunction {
4858                        func: Arc::new(NativeHistogramToString::scalar_udf()),
4859                        args: vec![DfExpr::Column(Column::from_name(col))],
4860                    })
4861                })
4862                .collect::<Vec<_>>()
4863        } else {
4864            vec![]
4865        };
4866
4867        let exprs = self
4868            .ctx
4869            .field_columns
4870            .iter()
4871            .map(|col| self.create_native_histogram_aggregate_expr(op, col))
4872            .collect::<Result<Vec<_>>>()?;
4873
4874        let normalized_exprs =
4875            normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
4876        self.ctx.field_columns = normalized_exprs
4877            .into_iter()
4878            .map(|expr| expr.schema_name().to_string())
4879            .collect();
4880
4881        Ok((exprs, prev_field_exprs))
4882    }
4883
4884    fn get_param_value_as_str(op: TokenType, param: &Option<Box<PromExpr>>) -> Result<&str> {
4885        let param = param
4886            .as_deref()
4887            .with_context(|| FunctionInvalidArgumentSnafu {
4888                fn_name: op.to_string(),
4889            })?;
4890        let PromExpr::StringLiteral(StringLiteral { val }) = param else {
4891            return FunctionInvalidArgumentSnafu {
4892                fn_name: op.to_string(),
4893            }
4894            .fail();
4895        };
4896
4897        Ok(val)
4898    }
4899
4900    fn get_param_as_literal_expr(
4901        param: Option<&PromExpr>,
4902        op: Option<TokenType>,
4903        expected_type: Option<ArrowDataType>,
4904    ) -> Result<DfExpr> {
4905        let prom_param = param.with_context(|| {
4906            if let Some(op) = op {
4907                FunctionInvalidArgumentSnafu {
4908                    fn_name: op.to_string(),
4909                }
4910            } else {
4911                FunctionInvalidArgumentSnafu {
4912                    fn_name: "unknown".to_string(),
4913                }
4914            }
4915        })?;
4916
4917        let expr = Self::try_build_literal_expr(prom_param).with_context(|| {
4918            if let Some(op) = op {
4919                FunctionInvalidArgumentSnafu {
4920                    fn_name: op.to_string(),
4921                }
4922            } else {
4923                FunctionInvalidArgumentSnafu {
4924                    fn_name: "unknown".to_string(),
4925                }
4926            }
4927        })?;
4928
4929        // check if the type is expected
4930        if let Some(expected_type) = expected_type {
4931            // literal should not have reference to column
4932            let expr_type = expr
4933                .get_type(&DFSchema::empty())
4934                .context(DataFusionPlanningSnafu)?;
4935            if expected_type != expr_type {
4936                return FunctionInvalidArgumentSnafu {
4937                    fn_name: format!("expected {expected_type:?}, but found {expr_type:?}"),
4938                }
4939                .fail();
4940            }
4941        }
4942
4943        Ok(expr)
4944    }
4945
4946    /// Create [DfExpr::WindowFunction] expr for each value column with given window function.
4947    ///
4948    fn create_window_exprs(
4949        &mut self,
4950        op: TokenType,
4951        group_exprs: Vec<DfExpr>,
4952        input_plan: &LogicalPlan,
4953    ) -> Result<Vec<DfExpr>> {
4954        ensure!(
4955            self.ctx.field_columns.len() == 1,
4956            UnsupportedExprSnafu {
4957                name: "topk or bottomk on multi-value input"
4958            }
4959        );
4960
4961        assert!(matches!(op.id(), token::T_TOPK | token::T_BOTTOMK));
4962
4963        let asc = matches!(op.id(), token::T_BOTTOMK);
4964
4965        let tag_sort_exprs = self
4966            .create_tag_column_exprs()?
4967            .into_iter()
4968            .map(|expr| expr.sort(asc, true));
4969
4970        // perform window operation to each value column
4971        let exprs: Vec<DfExpr> = self
4972            .ctx
4973            .field_columns
4974            .iter()
4975            .map(|col| {
4976                let mut sort_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 1);
4977                // Order by value in the specific order
4978                sort_exprs.push(DfExpr::Column(Column::from(col)).sort(asc, true));
4979                // Then tags if the values are equal,
4980                // Try to ensure the relative stability of the output results.
4981                sort_exprs.extend(tag_sort_exprs.clone());
4982
4983                DfExpr::WindowFunction(Box::new(WindowFunction {
4984                    fun: WindowFunctionDefinition::WindowUDF(Arc::new(RowNumber::new().into())),
4985                    params: WindowFunctionParams {
4986                        args: vec![],
4987                        partition_by: group_exprs.clone(),
4988                        order_by: sort_exprs,
4989                        window_frame: WindowFrame::new(Some(true)),
4990                        null_treatment: None,
4991                        distinct: false,
4992                        filter: None,
4993                    },
4994                }))
4995            })
4996            .collect();
4997
4998        let normalized_exprs =
4999            normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
5000        Ok(normalized_exprs)
5001    }
5002
5003    /// Create a classic, native, or mixed histogram helper plan.
5004    async fn create_histogram_plan(
5005        &mut self,
5006        function_name: &str,
5007        args: &PromFunctionArgs,
5008        query_engine_state: &QueryEngineState,
5009    ) -> Result<LogicalPlan> {
5010        let float_literal = |param: &PromExpr| -> Result<f64> {
5011            let value = (|| {
5012                let expr = Self::get_param_as_literal_expr(
5013                    Some(param),
5014                    None,
5015                    Some(ArrowDataType::Float64),
5016                )
5017                .ok()?;
5018                let simplifier = ExprSimplifier::new(SimplifyContext::default());
5019                let expr = simplifier.coerce(expr, &DFSchema::empty()).ok()?;
5020                let DfExpr::Literal(value, _) = simplifier.simplify(expr).ok()? else {
5021                    return None;
5022                };
5023                let ScalarValue::Float64(Some(value)) =
5024                    value.cast_to(&ArrowDataType::Float64).ok()?
5025                else {
5026                    return None;
5027                };
5028                Some(value)
5029            })()
5030            .with_context(|| FunctionInvalidArgumentSnafu {
5031                fn_name: function_name.to_string(),
5032            })?;
5033            Ok(value)
5034        };
5035        let (function, input) = match (function_name, args.args.as_slice()) {
5036            (SPECIAL_HISTOGRAM_QUANTILE, [quantile, input]) => (
5037                HistogramFoldOperation::Quantile(float_literal(quantile)?.into()),
5038                input.as_ref().clone(),
5039            ),
5040            (SPECIAL_HISTOGRAM_FRACTION, [lower, upper, input]) => (
5041                HistogramFoldOperation::Fraction {
5042                    lower: float_literal(lower)?.into(),
5043                    upper: float_literal(upper)?.into(),
5044                },
5045                input.as_ref().clone(),
5046            ),
5047            _ => {
5048                return FunctionInvalidArgumentSnafu {
5049                    fn_name: function_name.to_string(),
5050                }
5051                .fail();
5052            }
5053        };
5054
5055        let input_plan = self.prom_expr_to_plan(&input, query_engine_state).await?;
5056        // Histogram helpers fold buckets across `le`, so `__tsid` (which includes `le`) is not a
5057        // stable series identifier anymore. HistogramFold must not treat it as a label column.
5058        let input_plan = self.strip_tsid_column(input_plan)?;
5059        self.ctx.use_tsid = false;
5060
5061        if let Some((float_field, histogram_field)) =
5062            Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns)
5063                .map(|(float, histogram)| (float.to_string(), histogram.to_string()))
5064        {
5065            if self.ctx.has_le_tag() {
5066                return self.create_mixed_histogram_plan(
5067                    function,
5068                    input_plan,
5069                    float_field,
5070                    histogram_field,
5071                );
5072            }
5073            self.ctx.field_columns = vec![histogram_field];
5074        }
5075        if self.all_field_columns_are_native_histograms(input_plan.schema()) {
5076            return self.create_native_histogram_plan(function, input_plan);
5077        }
5078
5079        if !self.ctx.has_le_tag() {
5080            // Return empty result instead of error when 'le' column is not found
5081            // This handles the case when histogram metrics don't exist
5082            return Ok(LogicalPlan::EmptyRelation(
5083                datafusion::logical_expr::EmptyRelation {
5084                    produce_one_row: false,
5085                    schema: input_plan.schema().clone(),
5086                },
5087            ));
5088        }
5089        let time_index_column =
5090            self.ctx
5091                .time_index_column
5092                .clone()
5093                .with_context(|| TimeIndexNotFoundSnafu {
5094                    table: self.ctx.table_name.clone().unwrap_or_default(),
5095                })?;
5096        // FIXME(ruihang): support multi fields
5097        let field_column = self
5098            .ctx
5099            .field_columns
5100            .first()
5101            .with_context(|| FunctionInvalidArgumentSnafu {
5102                fn_name: function.function_name().to_string(),
5103            })?
5104            .clone();
5105        // remove le column from tag columns
5106        self.ctx.tag_columns.retain(|col| col != LE_COLUMN_NAME);
5107
5108        let fold = HistogramFold::new_with_operation(
5109            LE_COLUMN_NAME.to_string(),
5110            field_column,
5111            time_index_column,
5112            function,
5113            None,
5114            input_plan,
5115        )
5116        .context(DataFusionPlanningSnafu)?;
5117        Ok(LogicalPlan::Extension(Extension {
5118            node: Arc::new(fold),
5119        }))
5120    }
5121
5122    fn create_native_histogram_expr(
5123        &self,
5124        function: HistogramFoldOperation,
5125        field_column: &str,
5126    ) -> DfExpr {
5127        let field = DfExpr::Column(Column::from_name(field_column));
5128        let (func, args) = match function {
5129            HistogramFoldOperation::Quantile(quantile) => (
5130                Arc::new(NativeHistogramQuantile::scalar_udf_with_collector(
5131                    self.promql_annotations.clone(),
5132                )),
5133                vec![field, lit(f64::from(quantile))],
5134            ),
5135            HistogramFoldOperation::Fraction { lower, upper } => (
5136                Arc::new(NativeHistogramFraction::scalar_udf_with_collector(
5137                    self.promql_annotations.clone(),
5138                )),
5139                vec![field, lit(f64::from(lower)), lit(f64::from(upper))],
5140            ),
5141        };
5142        DfExpr::ScalarFunction(ScalarFunction { func, args })
5143    }
5144
5145    fn create_native_histogram_plan(
5146        &mut self,
5147        function: HistogramFoldOperation,
5148        input_plan: LogicalPlan,
5149    ) -> Result<LogicalPlan> {
5150        ensure!(
5151            self.ctx.field_columns.len() == 1,
5152            MultiFieldsNotSupportedSnafu {
5153                operator: function.function_name()
5154            },
5155        );
5156
5157        let field_column = self.ctx.field_columns[0].clone();
5158        let function_expr = self.create_native_histogram_expr(function, &field_column);
5159        let display_name = function_expr.schema_name().to_string();
5160        self.ctx.field_columns = vec![display_name.clone()];
5161
5162        let project_exprs = std::iter::once(self.create_time_index_column_expr()?)
5163            .chain(std::iter::once(function_expr.alias(display_name)))
5164            .chain(self.create_tag_column_exprs()?)
5165            .collect::<Vec<_>>();
5166
5167        LogicalPlanBuilder::from(input_plan)
5168            .project(project_exprs)
5169            .context(DataFusionPlanningSnafu)?
5170            .filter(self.create_empty_values_filter_expr(false)?)
5171            .context(DataFusionPlanningSnafu)?
5172            .build()
5173            .context(DataFusionPlanningSnafu)
5174    }
5175
5176    fn create_mixed_histogram_plan(
5177        &mut self,
5178        function: HistogramFoldOperation,
5179        input_plan: LogicalPlan,
5180        float_field: String,
5181        histogram_field: String,
5182    ) -> Result<LogicalPlan> {
5183        let time_index_column =
5184            self.ctx
5185                .time_index_column
5186                .clone()
5187                .with_context(|| TimeIndexNotFoundSnafu {
5188                    table: self.ctx.table_name.clone().unwrap_or_default(),
5189                })?;
5190        let tag_columns = self.ctx.tag_columns.clone();
5191        let folded = HistogramFold::new_with_operation(
5192            LE_COLUMN_NAME.to_string(),
5193            float_field.clone(),
5194            time_index_column.clone(),
5195            function,
5196            Some(histogram_field.clone()),
5197            input_plan,
5198        )
5199        .context(DataFusionPlanningSnafu)?;
5200        let record_collision = DfExpr::ScalarFunction(ScalarFunction {
5201            func: Arc::new(NativeHistogramDrop::warning_bool_false_udf(
5202                "vector contains a mix of classic and native histograms".to_string(),
5203                self.promql_annotations.clone(),
5204            )),
5205            args: vec![col(&float_field), col(&histogram_field)],
5206        });
5207        let keep = when(
5208            col(&float_field)
5209                .is_not_null()
5210                .and(col(&histogram_field).is_not_null()),
5211            record_collision,
5212        )
5213        .otherwise(lit(true))
5214        .context(DataFusionPlanningSnafu)?;
5215
5216        let native_expr = self.create_native_histogram_expr(function, &histogram_field);
5217        let output_field = native_expr.schema_name().to_string();
5218        let value = DfExpr::ScalarFunction(ScalarFunction {
5219            func: coalesce(),
5220            args: vec![col(&float_field), native_expr],
5221        });
5222        self.ctx.field_columns = vec![output_field.clone()];
5223        LogicalPlanBuilder::from(LogicalPlan::Extension(Extension {
5224            node: Arc::new(folded),
5225        }))
5226        .filter(keep)
5227        .context(DataFusionPlanningSnafu)?
5228        .project(
5229            std::iter::once(col(&time_index_column))
5230                .chain(std::iter::once(value.alias(output_field)))
5231                .chain(tag_columns.iter().map(col)),
5232        )
5233        .context(DataFusionPlanningSnafu)?
5234        .build()
5235        .context(DataFusionPlanningSnafu)
5236    }
5237
5238    /// Create a [SPECIAL_VECTOR_FUNCTION] plan
5239    async fn create_vector_plan(&mut self, args: &PromFunctionArgs) -> Result<LogicalPlan> {
5240        if args.args.len() != 1 {
5241            return FunctionInvalidArgumentSnafu {
5242                fn_name: SPECIAL_VECTOR_FUNCTION.to_string(),
5243            }
5244            .fail();
5245        }
5246        let lit = Self::get_param_as_literal_expr(Some(args.args[0].as_ref()), None, None)?;
5247
5248        // reuse `SPECIAL_TIME_FUNCTION` as name of time index column
5249        self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
5250        self.ctx.reset_table_name_and_schema();
5251        self.ctx.tag_columns = vec![];
5252        self.ctx.aggregation_field_labels.clear();
5253        self.ctx.field_columns = vec![greptime_value().to_string()];
5254        Ok(LogicalPlan::Extension(Extension {
5255            node: Arc::new(
5256                EmptyMetric::new(
5257                    self.ctx.start,
5258                    self.ctx.end,
5259                    self.ctx.interval,
5260                    SPECIAL_TIME_FUNCTION.to_string(),
5261                    greptime_value().to_string(),
5262                    Some(lit),
5263                )
5264                .context(DataFusionPlanningSnafu)?,
5265            ),
5266        }))
5267    }
5268
5269    /// Create a [SCALAR_FUNCTION] plan
5270    async fn create_scalar_plan(
5271        &mut self,
5272        args: &PromFunctionArgs,
5273        query_engine_state: &QueryEngineState,
5274    ) -> Result<LogicalPlan> {
5275        ensure!(
5276            args.len() == 1,
5277            FunctionInvalidArgumentSnafu {
5278                fn_name: SCALAR_FUNCTION
5279            }
5280        );
5281        let input = self
5282            .prom_expr_to_plan(&args.args[0], query_engine_state)
5283            .await?;
5284        let input_schema = input.schema().clone();
5285        let alternative_samples =
5286            Self::field_columns_are_alternative_samples(&input_schema, &self.ctx.field_columns);
5287        let histogram_fields = self
5288            .ctx
5289            .field_columns
5290            .iter()
5291            .filter(|field| Self::field_column_is_native_histogram(&input_schema, field))
5292            .count();
5293        ensure!(
5294            self.ctx.field_columns.len() == 1 || alternative_samples,
5295            MultiFieldsNotSupportedSnafu {
5296                operator: SCALAR_FUNCTION
5297            },
5298        );
5299        let scalar_field = self
5300            .ctx
5301            .field_columns
5302            .iter()
5303            .find(|field| !Self::field_column_is_native_histogram(&input_schema, field))
5304            .or_else(|| self.ctx.field_columns.first())
5305            .cloned()
5306            .with_context(|| FunctionInvalidArgumentSnafu {
5307                fn_name: SCALAR_FUNCTION,
5308            })?;
5309        let input = if histogram_fields == self.ctx.field_columns.len() {
5310            // scalar() ignores histogram samples. An empty input makes ScalarCalculate emit NaN
5311            // for every evaluation timestamp without attempting a Struct-to-Float64 cast.
5312            LogicalPlanBuilder::from(input)
5313                .filter(lit(false))
5314                .context(DataFusionPlanningSnafu)?
5315                .build()
5316                .context(DataFusionPlanningSnafu)?
5317        } else if histogram_fields > 0 {
5318            // A mixed vector contributes only its float samples to scalar().
5319            LogicalPlanBuilder::from(input)
5320                .filter(DfExpr::Column(Column::from_name(&scalar_field)).is_not_null())
5321                .context(DataFusionPlanningSnafu)?
5322                .build()
5323                .context(DataFusionPlanningSnafu)?
5324        } else {
5325            input
5326        };
5327        let scalar_plan = LogicalPlan::Extension(Extension {
5328            node: Arc::new(
5329                ScalarCalculate::new(
5330                    self.ctx.start,
5331                    self.ctx.end,
5332                    self.ctx.interval,
5333                    input,
5334                    self.ctx.time_index_column.as_ref().unwrap(),
5335                    &self.ctx.tag_columns,
5336                    &scalar_field,
5337                    self.ctx.table_name.as_deref(),
5338                )
5339                .context(PromqlPlanNodeSnafu)?,
5340            ),
5341        });
5342        // scalar plan have no tag columns
5343        self.ctx.tag_columns.clear();
5344        self.ctx.aggregation_field_labels.clear();
5345        self.ctx.field_columns.clear();
5346        self.ctx
5347            .field_columns
5348            .push(scalar_plan.schema().field(1).name().clone());
5349        Ok(scalar_plan)
5350    }
5351
5352    /// Create a [SPECIAL_ABSENT_FUNCTION] plan
5353    async fn create_absent_plan(
5354        &mut self,
5355        args: &PromFunctionArgs,
5356        query_engine_state: &QueryEngineState,
5357    ) -> Result<LogicalPlan> {
5358        if args.args.len() != 1 {
5359            return FunctionInvalidArgumentSnafu {
5360                fn_name: SPECIAL_ABSENT_FUNCTION.to_string(),
5361            }
5362            .fail();
5363        }
5364        let input = self
5365            .prom_expr_to_plan(&args.args[0], query_engine_state)
5366            .await?;
5367
5368        let time_index_expr = self.create_time_index_column_expr()?;
5369        let first_field_expr =
5370            self.create_field_column_exprs()?
5371                .pop()
5372                .with_context(|| ValueNotFoundSnafu {
5373                    table: self.ctx.table_name.clone().unwrap_or_default(),
5374                })?;
5375        let first_value_expr = first_value(first_field_expr, vec![]);
5376
5377        let ordered_aggregated_input = LogicalPlanBuilder::from(input)
5378            .aggregate(
5379                vec![time_index_expr.clone()],
5380                vec![first_value_expr.clone()],
5381            )
5382            .context(DataFusionPlanningSnafu)?
5383            .sort(vec![time_index_expr.sort(true, false)])
5384            .context(DataFusionPlanningSnafu)?
5385            .build()
5386            .context(DataFusionPlanningSnafu)?;
5387
5388        let fake_labels = self
5389            .ctx
5390            .selector_matcher
5391            .iter()
5392            .filter_map(|matcher| match matcher.op {
5393                MatchOp::Equal => Some((matcher.name.clone(), matcher.value.clone())),
5394                _ => None,
5395            })
5396            .collect::<Vec<_>>();
5397
5398        // Create the absent plan
5399        let absent_plan = LogicalPlan::Extension(Extension {
5400            node: Arc::new(
5401                Absent::try_new(
5402                    self.ctx.start,
5403                    self.ctx.end,
5404                    self.ctx.interval,
5405                    self.ctx.time_index_column.as_ref().unwrap().clone(),
5406                    self.ctx.field_columns[0].clone(),
5407                    fake_labels,
5408                    ordered_aggregated_input,
5409                )
5410                .context(DataFusionPlanningSnafu)?,
5411            ),
5412        });
5413
5414        // The absent series carries the equality matchers as labels, not the input's
5415        // tags or value fields, so the input's field grouping labels no longer apply.
5416        self.ctx.aggregation_field_labels.clear();
5417        Ok(absent_plan)
5418    }
5419
5420    /// Try to build a DataFusion Literal Expression from PromQL Expr, return
5421    /// `None` if the input is not a literal expression.
5422    fn try_build_literal_expr(expr: &PromExpr) -> Option<DfExpr> {
5423        match expr {
5424            PromExpr::NumberLiteral(NumberLiteral { val }) => Some(val.lit()),
5425            PromExpr::StringLiteral(StringLiteral { val }) => Some(val.lit()),
5426            PromExpr::VectorSelector(_)
5427            | PromExpr::MatrixSelector(_)
5428            | PromExpr::Extension(_)
5429            | PromExpr::Aggregate(_)
5430            | PromExpr::Subquery(_) => None,
5431            PromExpr::Call(Call { func, .. }) => {
5432                if func.name == SPECIAL_TIME_FUNCTION {
5433                    // For time() function, don't treat it as a literal
5434                    // Let it be handled as a regular function call
5435                    None
5436                } else {
5437                    None
5438                }
5439            }
5440            PromExpr::Paren(ParenExpr { expr }) => Self::try_build_literal_expr(expr),
5441            PromExpr::Unary(UnaryExpr { expr, .. }) => Some(DfExpr::Negative(Box::new(
5442                Self::try_build_literal_expr(expr)?,
5443            ))),
5444            PromExpr::Binary(PromBinaryExpr {
5445                lhs,
5446                rhs,
5447                op,
5448                modifier,
5449            }) => {
5450                let lhs = Self::try_build_literal_expr(lhs)?;
5451                let rhs = Self::try_build_literal_expr(rhs)?;
5452                let is_comparison_op = Self::is_token_a_comparison_op(*op);
5453                let expr_builder = Self::prom_token_to_binary_expr_builder(*op).ok()?;
5454                let expr = expr_builder(lhs, rhs).ok()?;
5455
5456                let should_return_bool = if let Some(m) = modifier {
5457                    m.return_bool
5458                } else {
5459                    false
5460                };
5461                if is_comparison_op && should_return_bool {
5462                    Some(DfExpr::Cast(Cast::new(
5463                        Box::new(expr),
5464                        ArrowDataType::Float64,
5465                    )))
5466                } else {
5467                    Some(expr)
5468                }
5469            }
5470        }
5471    }
5472
5473    fn try_build_special_time_expr_with_context(&self, expr: &PromExpr) -> Option<DfExpr> {
5474        match expr {
5475            PromExpr::Call(Call { func, .. }) => {
5476                if func.name == SPECIAL_TIME_FUNCTION
5477                    && let Some(time_index_col) = self.ctx.time_index_column.as_ref()
5478                {
5479                    Some(build_special_time_expr(time_index_col))
5480                } else {
5481                    None
5482                }
5483            }
5484            _ => None,
5485        }
5486    }
5487
5488    fn native_histogram_binary_expr(
5489        token: TokenType,
5490        lhs: DfExpr,
5491        lhs_is_histogram: bool,
5492        rhs: DfExpr,
5493        rhs_is_histogram: bool,
5494        filter_context: bool,
5495        promql_annotations: Option<PromqlAnnotationCollector>,
5496    ) -> Result<Option<DfExpr>> {
5497        if !lhs_is_histogram && !rhs_is_histogram {
5498            return Ok(None);
5499        }
5500
5501        let scalar_fn = |func: ScalarUdfDef, args| {
5502            DfExpr::ScalarFunction(ScalarFunction {
5503                func: Arc::new(func),
5504                args,
5505            })
5506        };
5507        let invalid_expr = || {
5508            let message = format!(
5509                "{}: dropped native histogram samples because this binary operation is not supported for native histograms",
5510                token
5511            );
5512            let func = if filter_context {
5513                NativeHistogramDrop::bool_false_udf(message, promql_annotations.clone())
5514            } else {
5515                NativeHistogramDrop::float_null_udf(message, promql_annotations.clone())
5516            };
5517            let args = vec![lhs.clone(), rhs.clone()];
5518            scalar_fn(func, args)
5519        };
5520
5521        let expr = match (token.id(), lhs_is_histogram, rhs_is_histogram) {
5522            (token::T_ADD, true, true) => scalar_fn(
5523                NativeHistogramAdd::scalar_udf_with_collector(promql_annotations.clone()),
5524                vec![lhs, rhs],
5525            ),
5526            (token::T_SUB, true, true) => scalar_fn(
5527                NativeHistogramSub::scalar_udf_with_collector(promql_annotations.clone()),
5528                vec![lhs, rhs],
5529            ),
5530            (token::T_MUL, true, false) => {
5531                scalar_fn(NativeHistogramMulScalar::scalar_udf(), vec![lhs, rhs])
5532            }
5533            (token::T_MUL, false, true) => {
5534                scalar_fn(NativeHistogramScalarMul::scalar_udf(), vec![lhs, rhs])
5535            }
5536            (token::T_DIV, true, false) => {
5537                scalar_fn(NativeHistogramDivScalar::scalar_udf(), vec![lhs, rhs])
5538            }
5539            (token::T_EQLC, true, true) => {
5540                scalar_fn(NativeHistogramEq::scalar_udf(), vec![lhs, rhs])
5541            }
5542            (token::T_NEQ, true, true) => {
5543                scalar_fn(NativeHistogramNotEq::scalar_udf(), vec![lhs, rhs])
5544            }
5545            _ => invalid_expr(),
5546        };
5547
5548        Ok(Some(expr))
5549    }
5550
5551    /// Return a lambda to build binary expression from token.
5552    /// Because some binary operator are function in DataFusion like `atan2` or `^`.
5553    #[allow(clippy::type_complexity)]
5554    fn prom_token_to_binary_expr_builder(
5555        token: TokenType,
5556    ) -> Result<Box<dyn Fn(DfExpr, DfExpr) -> Result<DfExpr>>> {
5557        let cast_float = |expr| {
5558            if matches!(
5559                &expr,
5560                DfExpr::Cast(Cast { field, .. }) if field.data_type() == &ArrowDataType::Float64
5561            ) || matches!(&expr, DfExpr::Literal(ScalarValue::Float64(_), _))
5562            {
5563                expr
5564            } else {
5565                DfExpr::Cast(Cast::new(Box::new(expr), ArrowDataType::Float64))
5566            }
5567        };
5568        match token.id() {
5569            token::T_ADD => Ok(Box::new(move |lhs, rhs| {
5570                Ok(cast_float(lhs) + cast_float(rhs))
5571            })),
5572            token::T_SUB => Ok(Box::new(move |lhs, rhs| {
5573                Ok(cast_float(lhs) - cast_float(rhs))
5574            })),
5575            token::T_MUL => Ok(Box::new(move |lhs, rhs| {
5576                Ok(cast_float(lhs) * cast_float(rhs))
5577            })),
5578            token::T_DIV => Ok(Box::new(move |lhs, rhs| {
5579                Ok(cast_float(lhs) / cast_float(rhs))
5580            })),
5581            token::T_MOD => Ok(Box::new(move |lhs: DfExpr, rhs| {
5582                Ok(cast_float(lhs) % cast_float(rhs))
5583            })),
5584            token::T_EQLC => Ok(Box::new(|lhs, rhs| Ok(lhs.eq(rhs)))),
5585            token::T_NEQ => Ok(Box::new(|lhs, rhs| Ok(lhs.not_eq(rhs)))),
5586            token::T_GTR => Ok(Box::new(|lhs, rhs| Ok(lhs.gt(rhs)))),
5587            token::T_LSS => Ok(Box::new(|lhs, rhs| Ok(lhs.lt(rhs)))),
5588            token::T_GTE => Ok(Box::new(|lhs, rhs| Ok(lhs.gt_eq(rhs)))),
5589            token::T_LTE => Ok(Box::new(|lhs, rhs| Ok(lhs.lt_eq(rhs)))),
5590            token::T_POW => Ok(Box::new(move |lhs, rhs| {
5591                Ok(DfExpr::ScalarFunction(ScalarFunction {
5592                    func: datafusion_functions::math::power(),
5593                    args: vec![cast_float(lhs), cast_float(rhs)],
5594                }))
5595            })),
5596            token::T_ATAN2 => Ok(Box::new(move |lhs, rhs| {
5597                Ok(DfExpr::ScalarFunction(ScalarFunction {
5598                    func: datafusion_functions::math::atan2(),
5599                    args: vec![cast_float(lhs), cast_float(rhs)],
5600                }))
5601            })),
5602            _ => UnexpectedTokenSnafu { token }.fail(),
5603        }
5604    }
5605
5606    /// Check if the given op is a [comparison operator](https://prometheus.io/docs/prometheus/latest/querying/operators/#comparison-binary-operators).
5607    fn is_token_a_comparison_op(token: TokenType) -> bool {
5608        matches!(
5609            token.id(),
5610            token::T_EQLC
5611                | token::T_NEQ
5612                | token::T_GTR
5613                | token::T_LSS
5614                | token::T_GTE
5615                | token::T_LTE
5616        )
5617    }
5618
5619    /// Check if the given op is a set operator (UNION, INTERSECT and EXCEPT in SQL).
5620    fn is_token_a_set_op(token: TokenType) -> bool {
5621        matches!(
5622            token.id(),
5623            token::T_LAND // INTERSECT
5624                | token::T_LOR // UNION
5625                | token::T_LUNLESS // EXCEPT
5626        )
5627    }
5628
5629    fn align_binary_field_columns<'a>(
5630        left_schema: &DFSchemaRef,
5631        right_schema: &DFSchemaRef,
5632        left_field_columns: &'a [String],
5633        right_field_columns: &'a [String],
5634        op: TokenType,
5635        left_is_scalar: bool,
5636        right_is_scalar: bool,
5637    ) -> (
5638        Vec<(String, Vec<BinaryFieldPair<'a>>)>,
5639        Vec<BinaryFieldPair<'a>>,
5640    ) {
5641        // Mixed vectors store mutually exclusive float and histogram samples in two columns.
5642        // Retain each valid sample combination and group expressions by their output lane.
5643        let left_alternative = Self::alternative_sample_columns(left_schema, left_field_columns);
5644        let right_alternative = Self::alternative_sample_columns(right_schema, right_field_columns);
5645        let alternative_alignment = match (left_alternative, right_alternative) {
5646            (Some(output_names), Some(_)) => Some((
5647                output_names,
5648                left_field_columns
5649                    .iter()
5650                    .flat_map(|left| right_field_columns.iter().map(move |right| (left, right)))
5651                    .collect::<Vec<_>>(),
5652            )),
5653            (Some(output_names), None) if right_field_columns.len() == 1 => Some((
5654                output_names,
5655                left_field_columns
5656                    .iter()
5657                    .map(|left| (left, &right_field_columns[0]))
5658                    .collect::<Vec<_>>(),
5659            )),
5660            (None, Some(output_names)) if left_field_columns.len() == 1 => Some((
5661                output_names,
5662                right_field_columns
5663                    .iter()
5664                    .map(|right| (&left_field_columns[0], right))
5665                    .collect::<Vec<_>>(),
5666            )),
5667            _ => None,
5668        };
5669        let mut invalid_pairs = Vec::new();
5670        if let Some(((float_output, histogram_output), field_pairs)) = alternative_alignment {
5671            let mut float_pairs = Vec::new();
5672            let mut histogram_pairs = Vec::new();
5673            for (left, right) in field_pairs {
5674                let left_is_histogram = Self::field_column_is_native_histogram(left_schema, left);
5675                let right_is_histogram =
5676                    Self::field_column_is_native_histogram(right_schema, right);
5677                match Self::binary_result_is_histogram(op, left_is_histogram, right_is_histogram) {
5678                    Some(false) => float_pairs.push((left, right)),
5679                    Some(true) => histogram_pairs.push((left, right)),
5680                    None => invalid_pairs.push((left, right)),
5681                }
5682            }
5683            if !float_pairs.is_empty() || !histogram_pairs.is_empty() {
5684                return (
5685                    [
5686                        (!float_pairs.is_empty()).then(|| (float_output.to_string(), float_pairs)),
5687                        (!histogram_pairs.is_empty())
5688                            .then(|| (histogram_output.to_string(), histogram_pairs)),
5689                    ]
5690                    .into_iter()
5691                    .flatten()
5692                    .collect(),
5693                    invalid_pairs,
5694                );
5695            }
5696        }
5697
5698        if left_is_scalar && !right_is_scalar && left_field_columns.len() == 1 {
5699            return (
5700                right_field_columns
5701                    .iter()
5702                    .map(|right| (right.clone(), vec![(&left_field_columns[0], right)]))
5703                    .collect(),
5704                invalid_pairs,
5705            );
5706        }
5707        if right_is_scalar && !left_is_scalar && right_field_columns.len() == 1 {
5708            return (
5709                left_field_columns
5710                    .iter()
5711                    .map(|left| (left.clone(), vec![(left, &right_field_columns[0])]))
5712                    .collect(),
5713                invalid_pairs,
5714            );
5715        }
5716
5717        (
5718            left_field_columns
5719                .iter()
5720                .zip(right_field_columns.iter())
5721                .map(|(left, right)| (left.clone(), vec![(left, right)]))
5722                .collect(),
5723            invalid_pairs,
5724        )
5725    }
5726
5727    fn binary_result_is_histogram(
5728        token: TokenType,
5729        lhs_is_histogram: bool,
5730        rhs_is_histogram: bool,
5731    ) -> Option<bool> {
5732        match (token.id(), lhs_is_histogram, rhs_is_histogram) {
5733            (_, false, false) => Some(false),
5734            (token::T_ADD | token::T_SUB, true, true)
5735            | (token::T_MUL, true, false)
5736            | (token::T_MUL, false, true)
5737            | (token::T_DIV, true, false) => Some(true),
5738            (token::T_EQLC | token::T_NEQ, true, true) => Some(false),
5739            _ => None,
5740        }
5741    }
5742
5743    fn plan_has_tsid_column(plan: &LogicalPlan) -> bool {
5744        plan.schema()
5745            .fields()
5746            .iter()
5747            .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
5748    }
5749
5750    fn is_empty_metric(plan: &LogicalPlan) -> bool {
5751        matches!(plan, LogicalPlan::Extension(Extension { node }) if node.as_any().is::<EmptyMetric>())
5752    }
5753
5754    fn native_histogram_arrow_type() -> ArrowDataType {
5755        native_histogram_value_type().as_arrow_type()
5756    }
5757
5758    fn field_column_type<'a>(
5759        schema: &'a DFSchemaRef,
5760        field_column: &str,
5761    ) -> Option<&'a ArrowDataType> {
5762        schema
5763            .index_of_column_by_name(None, field_column)
5764            .map(|idx| schema.field(idx).data_type())
5765    }
5766
5767    fn field_column_is_native_histogram(schema: &DFSchemaRef, field_column: &str) -> bool {
5768        Self::field_column_type(schema, field_column)
5769            .is_some_and(|data_type| data_type == &Self::native_histogram_arrow_type())
5770    }
5771
5772    fn field_columns_contain_native_histogram(
5773        schema: &DFSchemaRef,
5774        field_columns: &[String],
5775    ) -> bool {
5776        field_columns
5777            .iter()
5778            .any(|field| Self::field_column_is_native_histogram(schema, field))
5779    }
5780
5781    fn field_column_is_float_range(schema: &DFSchemaRef, field_column: &str) -> bool {
5782        Self::field_column_type(schema, field_column).is_some_and(|data_type| {
5783            matches!(
5784                data_type,
5785                ArrowDataType::Dictionary(key_type, value_type)
5786                    if key_type.as_ref() == &ArrowDataType::Int64
5787                        && value_type.as_ref() == &ArrowDataType::Float64
5788            )
5789        })
5790    }
5791
5792    fn field_columns_are_alternative_samples(
5793        schema: &DFSchemaRef,
5794        field_columns: &[String],
5795    ) -> bool {
5796        Self::alternative_sample_columns(schema, field_columns).is_some()
5797    }
5798
5799    fn alternative_sample_columns<'a>(
5800        schema: &DFSchemaRef,
5801        field_columns: &'a [String],
5802    ) -> Option<(&'a str, &'a str)> {
5803        if field_columns.len() != 2 {
5804            return None;
5805        }
5806
5807        let canonical_float = field_columns.iter().find(|field| {
5808            field.as_str() == greptime_value()
5809                && (Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64)
5810                    || Self::field_column_is_float_range(schema, field))
5811        });
5812        let canonical_histogram = field_columns.iter().find(|field| {
5813            field.as_str() == greptime_native_histogram()
5814                && (Self::field_column_is_native_histogram(schema, field)
5815                    || Self::field_column_is_native_histogram_range(schema, field))
5816        });
5817        if let (Some(float), Some(histogram)) = (canonical_float, canonical_histogram) {
5818            return Some((float, histogram));
5819        }
5820
5821        let float = field_columns.iter().find(|field| {
5822            field.starts_with(OR_FLOAT_FIELD_PREFIX)
5823                && (Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64)
5824                    || Self::field_column_is_float_range(schema, field))
5825        })?;
5826        let histogram = field_columns.iter().find(|field| {
5827            field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)
5828                && (Self::field_column_is_native_histogram(schema, field)
5829                    || Self::field_column_is_native_histogram_range(schema, field))
5830        })?;
5831        Some((float, histogram))
5832    }
5833
5834    fn alternative_sample_range_columns<'a>(
5835        schema: &DFSchemaRef,
5836        field_columns: &'a [String],
5837    ) -> Option<(&'a str, &'a str)> {
5838        Self::alternative_sample_columns(schema, field_columns).filter(|(float, histogram)| {
5839            Self::field_column_is_float_range(schema, float)
5840                && Self::field_column_is_native_histogram_range(schema, histogram)
5841        })
5842    }
5843
5844    fn field_column_is_native_histogram_range(schema: &DFSchemaRef, field_column: &str) -> bool {
5845        Self::field_column_type(schema, field_column).is_some_and(|data_type| {
5846            matches!(
5847                data_type,
5848                ArrowDataType::Dictionary(key_type, value_type)
5849                    if key_type.as_ref() == &ArrowDataType::Int64
5850                        && value_type.as_ref() == &Self::native_histogram_arrow_type()
5851            )
5852        })
5853    }
5854
5855    fn all_field_columns_are_native_histograms(&self, schema: &DFSchemaRef) -> bool {
5856        !self.ctx.field_columns.is_empty()
5857            && self
5858                .ctx
5859                .field_columns
5860                .iter()
5861                .all(|field| Self::field_column_is_native_histogram(schema, field))
5862    }
5863
5864    fn all_field_columns_are_native_histogram_ranges(&self, schema: &DFSchemaRef) -> bool {
5865        !self.ctx.field_columns.is_empty()
5866            && self
5867                .ctx
5868                .field_columns
5869                .iter()
5870                .all(|field| Self::field_column_is_native_histogram_range(schema, field))
5871    }
5872
5873    fn optional_tsid_projection(
5874        schema: &DFSchemaRef,
5875        table_ref: Option<&TableReference>,
5876        keep_tsid: bool,
5877    ) -> Option<DfExpr> {
5878        keep_tsid.then_some(()).and_then(|_| {
5879            schema
5880                .qualified_field_with_name(table_ref, DATA_SCHEMA_TSID_COLUMN_NAME)
5881                .ok()
5882                .map(|field| DfExpr::Column(field.into()))
5883        })
5884    }
5885
5886    fn binary_join_key_columns(
5887        &self,
5888        left_schema: &DFSchemaRef,
5889        right_schema: &DFSchemaRef,
5890        left_context: &PromPlannerContext,
5891        right_context: &PromPlannerContext,
5892        only_join_time_index: bool,
5893        modifier: &Option<BinModifier>,
5894    ) -> Result<(BTreeSet<String>, BTreeSet<String>, bool)> {
5895        let has_tsid = |schema: &DFSchemaRef| {
5896            schema
5897                .fields()
5898                .iter()
5899                .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
5900        };
5901        let use_tsid_join = !only_join_time_index
5902            && self.binary_modifier_preserves_tsid_join_key(left_context, right_context, modifier)
5903            && left_context.use_tsid
5904            && right_context.use_tsid
5905            && has_tsid(left_schema)
5906            && has_tsid(right_schema);
5907
5908        let (mut left_tag_columns, mut right_tag_columns) = if use_tsid_join {
5909            (
5910                BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]),
5911                BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]),
5912            )
5913        } else {
5914            if only_join_time_index {
5915                (BTreeSet::new(), BTreeSet::new())
5916            } else {
5917                (
5918                    left_context
5919                        .tag_columns
5920                        .iter()
5921                        .cloned()
5922                        .collect::<BTreeSet<_>>(),
5923                    right_context
5924                        .tag_columns
5925                        .iter()
5926                        .cloned()
5927                        .collect::<BTreeSet<_>>(),
5928                )
5929            }
5930        };
5931
5932        if !use_tsid_join
5933            && let Some(modifier) = modifier
5934            && let Some(matching) = &modifier.matching
5935        {
5936            match matching {
5937                LabelModifier::Include(on) => {
5938                    let mask = on.labels.iter().cloned().collect::<BTreeSet<_>>();
5939                    left_tag_columns = left_tag_columns.intersection(&mask).cloned().collect();
5940                    right_tag_columns = right_tag_columns.intersection(&mask).cloned().collect();
5941                }
5942                LabelModifier::Exclude(ignoring) => {
5943                    for label in &ignoring.labels {
5944                        let _ = left_tag_columns.remove(label);
5945                        let _ = right_tag_columns.remove(label);
5946                    }
5947                }
5948            }
5949        }
5950
5951        let force_empty_join =
5952            !use_tsid_join && !only_join_time_index && left_tag_columns != right_tag_columns;
5953        if force_empty_join {
5954            let common_tag_columns = left_tag_columns
5955                .intersection(&right_tag_columns)
5956                .cloned()
5957                .collect::<BTreeSet<_>>();
5958            left_tag_columns = common_tag_columns.clone();
5959            right_tag_columns = common_tag_columns;
5960        }
5961
5962        Ok((left_tag_columns, right_tag_columns, force_empty_join))
5963    }
5964
5965    /// Result labels of a vector-vector binary operation, following Prometheus `resultMetric`:
5966    /// `on(...)` keeps only the matching labels, `ignoring(...)` drops them, and a group modifier
5967    /// keeps the "many" side's labels plus the `group_x(...)` labels taken from the "one" side.
5968    ///
5969    /// The flag of each entry tells which operand the label is projected from. `None` means the
5970    /// operation keeps a whole operand tag set, which the default projection already does.
5971    fn binary_result_labels(
5972        left_context: &PromPlannerContext,
5973        right_context: &PromPlannerContext,
5974        modifier: &Option<BinModifier>,
5975    ) -> Option<Vec<(bool, String)>> {
5976        let modifier = modifier.as_ref()?;
5977        let (many_is_left, include) = match &modifier.card {
5978            VectorMatchCardinality::OneToOne => (true, None),
5979            VectorMatchCardinality::ManyToOne(labels) => (true, Some(labels)),
5980            VectorMatchCardinality::OneToMany(labels) => (false, Some(labels)),
5981            // Set operators keep their operands' labels and don't reach this path.
5982            VectorMatchCardinality::ManyToMany => return None,
5983        };
5984
5985        let Some(include) = include else {
5986            let matching = modifier.matching.as_ref()?;
5987            let reduced = |keep: bool, labels: &BTreeSet<&String>| {
5988                left_context
5989                    .tag_columns
5990                    .iter()
5991                    .filter(|tag| labels.contains(tag) == keep)
5992                    .map(|tag| (true, tag.clone()))
5993                    .collect()
5994            };
5995            return Some(match matching {
5996                LabelModifier::Include(on) => reduced(true, &on.labels.iter().collect()),
5997                LabelModifier::Exclude(ignoring) => {
5998                    reduced(false, &ignoring.labels.iter().collect())
5999                }
6000            });
6001        };
6002
6003        let (many_context, one_context) = if many_is_left {
6004            (left_context, right_context)
6005        } else {
6006            (right_context, left_context)
6007        };
6008        let include = include.labels.iter().collect::<BTreeSet<_>>();
6009        // An included label the "one" side doesn't carry is deleted from the result, so it is
6010        // dropped from the "many" side as well.
6011        let mut labels = many_context
6012            .tag_columns
6013            .iter()
6014            .filter(|tag| !include.contains(tag))
6015            .map(|tag| (many_is_left, tag.clone()))
6016            .collect::<Vec<_>>();
6017        labels.extend(
6018            include
6019                .into_iter()
6020                .filter(|label| one_context.tag_columns.contains(label))
6021                .map(|label| (!many_is_left, label.clone())),
6022        );
6023        Some(labels)
6024    }
6025
6026    /// Resolve [`Self::binary_result_labels`] against the join output.
6027    fn binary_result_label_projection(
6028        schema: &DFSchemaRef,
6029        left_table_ref: &TableReference,
6030        right_table_ref: &TableReference,
6031        left_context: &PromPlannerContext,
6032        right_context: &PromPlannerContext,
6033        labels: Vec<(bool, String)>,
6034    ) -> Result<BinaryResultLabels> {
6035        let mut exprs = Vec::with_capacity(labels.len());
6036        let mut names = Vec::with_capacity(labels.len());
6037        let mut aggregation_field_labels = Vec::new();
6038        let mut sources = HashSet::new();
6039        for (from_left, label) in labels {
6040            let (table_ref, context) = if from_left {
6041                (left_table_ref, left_context)
6042            } else {
6043                (right_table_ref, right_context)
6044            };
6045            let field = schema
6046                .qualified_field_with_name(Some(table_ref), &label)
6047                .context(DataFusionPlanningSnafu)?;
6048            exprs.push(DfExpr::Column(field.into()));
6049            if context.aggregation_field_labels.contains(&label) {
6050                aggregation_field_labels.push(label.clone());
6051            }
6052            sources.insert(from_left);
6053            names.push(label);
6054        }
6055
6056        // One operand contributing every one of its tags as the whole result label set keeps the
6057        // one-to-one correspondence between its `__tsid` and a result series: no other operand
6058        // value reaches the labels, and the matching gives each of its rows a single partner.
6059        let source_context = match sources.into_iter().collect::<Vec<_>>().as_slice() {
6060            [true] => Some((left_table_ref, left_context)),
6061            [false] => Some((right_table_ref, right_context)),
6062            _ => None,
6063        };
6064        let tsid = source_context
6065            .filter(|(_, context)| {
6066                context.use_tsid
6067                    && context.tag_columns.len() == names.len()
6068                    && context.tag_columns.iter().all(|tag| names.contains(tag))
6069            })
6070            .and_then(|(table_ref, _)| {
6071                Self::optional_tsid_projection(schema, Some(table_ref), true)
6072            });
6073
6074        Ok(BinaryResultLabels {
6075            exprs,
6076            names,
6077            aggregation_field_labels,
6078            tsid,
6079        })
6080    }
6081
6082    /// Whether the result of a binary operation can hold two series with the same labels, which
6083    /// only [`Self::binary_result_labels`] can introduce: one-to-one matching on a subset of the
6084    /// tags, or a group modifier that overwrites a label of the "many" side.
6085    fn binary_result_labels_may_repeat(
6086        left_context: &PromPlannerContext,
6087        right_context: &PromPlannerContext,
6088        modifier: &Option<BinModifier>,
6089    ) -> bool {
6090        let Some(modifier) = modifier else {
6091            return false;
6092        };
6093        match &modifier.card {
6094            VectorMatchCardinality::OneToOne => match &modifier.matching {
6095                None => false,
6096                Some(LabelModifier::Include(on)) => {
6097                    let on = on.labels.iter().collect::<BTreeSet<_>>();
6098                    !left_context.tag_columns.iter().all(|tag| on.contains(tag))
6099                }
6100                Some(LabelModifier::Exclude(ignoring)) => ignoring
6101                    .labels
6102                    .iter()
6103                    .any(|label| left_context.tag_columns.contains(label)),
6104            },
6105            VectorMatchCardinality::ManyToOne(include) => include
6106                .labels
6107                .iter()
6108                .any(|label| left_context.tag_columns.contains(label)),
6109            VectorMatchCardinality::OneToMany(include) => include
6110                .labels
6111                .iter()
6112                .any(|label| right_context.tag_columns.contains(label)),
6113            VectorMatchCardinality::ManyToMany => false,
6114        }
6115    }
6116
6117    /// Wrap `plan` in a check that fails the query when a match group holds more than one row at
6118    /// a timestamp. `group_exprs` are the label columns of the group, resolved against `plan`.
6119    fn assert_unique_match_group(
6120        plan: LogicalPlan,
6121        group_exprs: Vec<DfExpr>,
6122        group_labels: Vec<String>,
6123        time_index_expr: DfExpr,
6124        violation: MatchGroupViolation,
6125    ) -> Result<LogicalPlan> {
6126        let mut partition_by = group_exprs.clone();
6127        partition_by.push(time_index_expr);
6128        // A label may carry the generated name, and the count column has to stay unambiguous
6129        // against every column the operand already has.
6130        let occupied_column_names = plan
6131            .schema()
6132            .fields()
6133            .iter()
6134            .map(|field| field.name().as_str())
6135            .collect::<HashSet<_>>();
6136        let mut next_suffix = 0;
6137        let count_column = loop {
6138            let name = match next_suffix {
6139                0 => MATCH_GROUP_COUNT_COLUMN.to_string(),
6140                suffix => format!("{MATCH_GROUP_COUNT_COLUMN}_{suffix}"),
6141            };
6142            next_suffix += 1;
6143            if !occupied_column_names.contains(name.as_str()) {
6144                break name;
6145            }
6146        };
6147        let count = DfExpr::WindowFunction(Box::new(WindowFunction {
6148            fun: WindowFunctionDefinition::AggregateUDF(count_udaf()),
6149            params: WindowFunctionParams {
6150                args: vec![lit(1i64)],
6151                partition_by,
6152                order_by: vec![],
6153                window_frame: WindowFrame::new(None),
6154                null_treatment: None,
6155                distinct: false,
6156                filter: None,
6157            },
6158        }))
6159        .alias(count_column.as_str());
6160
6161        let output_exprs = plan
6162            .schema()
6163            .iter()
6164            .map(|(qualifier, field)| DfExpr::Column(Column::new(qualifier.cloned(), field.name())))
6165            .collect::<Vec<_>>();
6166        let assert_expr = DfExpr::ScalarFunction(ScalarFunction {
6167            func: Arc::new(UniqueMatchGroup::scalar_udf(group_labels, violation)),
6168            args: std::iter::once(col(count_column.as_str()))
6169                .chain(group_exprs)
6170                .collect(),
6171        });
6172
6173        LogicalPlanBuilder::from(plan)
6174            .window(vec![count])
6175            .context(DataFusionPlanningSnafu)?
6176            .filter(assert_expr)
6177            .context(DataFusionPlanningSnafu)?
6178            .project(output_exprs)
6179            .context(DataFusionPlanningSnafu)?
6180            .build()
6181            .context(DataFusionPlanningSnafu)
6182    }
6183
6184    fn binary_modifier_preserves_tsid_join_key(
6185        &self,
6186        left_context: &PromPlannerContext,
6187        right_context: &PromPlannerContext,
6188        modifier: &Option<BinModifier>,
6189    ) -> bool {
6190        let Some(modifier) = modifier else {
6191            return true;
6192        };
6193
6194        if !matches!(modifier.card, VectorMatchCardinality::OneToOne) {
6195            return false;
6196        }
6197
6198        match &modifier.matching {
6199            None => true,
6200            Some(LabelModifier::Exclude(ignoring)) => ignoring.labels.iter().all(|label| {
6201                !left_context.tag_columns.contains(label)
6202                    && !right_context.tag_columns.contains(label)
6203            }),
6204            Some(LabelModifier::Include(on)) => {
6205                let on_labels = on.labels.iter().cloned().collect::<BTreeSet<_>>();
6206                let left_labels = left_context
6207                    .tag_columns
6208                    .iter()
6209                    .cloned()
6210                    .collect::<BTreeSet<_>>();
6211                let right_labels = right_context
6212                    .tag_columns
6213                    .iter()
6214                    .cloned()
6215                    .collect::<BTreeSet<_>>();
6216
6217                on_labels == left_labels && on_labels == right_labels
6218            }
6219        }
6220    }
6221
6222    /// Build a inner join on time index column and tag columns to concat two logical plans.
6223    /// When `only_join_time_index == true` we only join on the time index, because these two plan may not have the same tag columns
6224    #[allow(clippy::too_many_arguments)]
6225    fn join_on_non_field_columns(
6226        &self,
6227        left: LogicalPlan,
6228        right: LogicalPlan,
6229        left_table_ref: TableReference,
6230        right_table_ref: TableReference,
6231        left_time_index_column: Option<String>,
6232        right_time_index_column: Option<String>,
6233        only_join_time_index: bool,
6234        modifier: &Option<BinModifier>,
6235        left_context: &PromPlannerContext,
6236        right_context: &PromPlannerContext,
6237    ) -> Result<LogicalPlan> {
6238        let (mut left_tag_columns, mut right_tag_columns, mut force_empty_join) = self
6239            .binary_join_key_columns(
6240                left.schema(),
6241                right.schema(),
6242                left_context,
6243                right_context,
6244                only_join_time_index,
6245                modifier,
6246            )?;
6247        let use_tsid_join = !only_join_time_index
6248            && !force_empty_join
6249            && left_tag_columns == BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()])
6250            && right_tag_columns == BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]);
6251        let (left, right, left_matched_tags, right_matched_tags) = if !only_join_time_index
6252            && !use_tsid_join
6253            && Self::only_temporality_match_label_mismatches(left_context, right_context, modifier)
6254        {
6255            let mut aligned_left_context = left_context.clone();
6256            let mut aligned_right_context = right_context.clone();
6257            let (left, right, _) = Self::align_temporality_match_column(
6258                left,
6259                right,
6260                &mut aligned_left_context,
6261                &mut aligned_right_context,
6262            )?;
6263            (left_tag_columns, right_tag_columns, force_empty_join) = self
6264                .binary_join_key_columns(
6265                    left.schema(),
6266                    right.schema(),
6267                    &aligned_left_context,
6268                    &aligned_right_context,
6269                    false,
6270                    modifier,
6271                )?;
6272            (
6273                left,
6274                right,
6275                aligned_left_context.tag_columns,
6276                aligned_right_context.tag_columns,
6277            )
6278        } else {
6279            (
6280                left,
6281                right,
6282                left_context.tag_columns.clone(),
6283                right_context.tag_columns.clone(),
6284            )
6285        };
6286
6287        // A join key that covers the whole tag set of a side already makes that side's match
6288        // groups unique, and so does a join on `__tsid`. Only the reduced keys need the check.
6289        let checked_matching = !only_join_time_index
6290            && !force_empty_join
6291            && !use_tsid_join
6292            && !matches!(
6293                modifier.as_ref().map(|modifier| &modifier.card),
6294                Some(VectorMatchCardinality::ManyToMany)
6295            );
6296        // The right operand is the "one" side of the matching, unless `group_right` swaps them.
6297        let one_side_is_left = matches!(
6298            modifier.as_ref().map(|modifier| &modifier.card),
6299            Some(VectorMatchCardinality::OneToMany(_))
6300        );
6301        let (left, right) = if !checked_matching {
6302            (left, right)
6303        } else if one_side_is_left {
6304            (
6305                Self::assert_unique_one_side(
6306                    left,
6307                    &left_tag_columns,
6308                    &left_matched_tags,
6309                    left_time_index_column.as_deref(),
6310                    true,
6311                )?,
6312                right,
6313            )
6314        } else {
6315            (
6316                left,
6317                Self::assert_unique_one_side(
6318                    right,
6319                    &right_tag_columns,
6320                    &right_matched_tags,
6321                    right_time_index_column.as_deref(),
6322                    false,
6323                )?,
6324            )
6325        };
6326
6327        // push time index column if it exists
6328        if let (Some(left_time_index_column), Some(right_time_index_column)) = (
6329            left_time_index_column.clone(),
6330            right_time_index_column.clone(),
6331        ) {
6332            left_tag_columns.insert(left_time_index_column);
6333            right_tag_columns.insert(right_time_index_column);
6334        }
6335
6336        let right = LogicalPlanBuilder::from(right)
6337            .alias(right_table_ref.clone())
6338            .context(DataFusionPlanningSnafu)?
6339            .build()
6340            .context(DataFusionPlanningSnafu)?;
6341
6342        // Inner Join on time index column to concat two operator
6343        let join_plan = LogicalPlanBuilder::from(left)
6344            .alias(left_table_ref.clone())
6345            .context(DataFusionPlanningSnafu)?
6346            .join_detailed(
6347                right,
6348                JoinType::Inner,
6349                (
6350                    left_tag_columns
6351                        .into_iter()
6352                        .map(Column::from_name)
6353                        .collect::<Vec<_>>(),
6354                    right_tag_columns
6355                        .into_iter()
6356                        .map(Column::from_name)
6357                        .collect::<Vec<_>>(),
6358                ),
6359                force_empty_join.then_some(lit(false)),
6360                NullEquality::NullEqualsNull,
6361            )
6362            .context(DataFusionPlanningSnafu)?
6363            .build()
6364            .context(DataFusionPlanningSnafu)?;
6365
6366        // The result labels no longer identify a series on their own: the "many" side can hold
6367        // several series per result label set, which PromQL cannot represent.
6368        let labels = checked_matching
6369            .then(|| Self::binary_result_labels(left_context, right_context, modifier))
6370            .flatten()
6371            .filter(|_| {
6372                Self::binary_result_labels_may_repeat(left_context, right_context, modifier)
6373            });
6374        let Some(labels) = labels else {
6375            return Ok(join_plan);
6376        };
6377        let schema = join_plan.schema().clone();
6378        let group_exprs = labels
6379            .iter()
6380            .map(|(from_left, label)| {
6381                let table_ref = if *from_left {
6382                    &left_table_ref
6383                } else {
6384                    &right_table_ref
6385                };
6386                schema
6387                    .qualified_field_with_name(Some(table_ref), label)
6388                    .map(|field| DfExpr::Column(field.into()))
6389                    .context(DataFusionPlanningSnafu)
6390            })
6391            .collect::<Result<Vec<_>>>()?;
6392        let time_index_expr = left_time_index_column
6393            .as_deref()
6394            .map(|column| {
6395                schema
6396                    .qualified_field_with_name(Some(&left_table_ref), column)
6397                    .map(|field| DfExpr::Column(field.into()))
6398                    .context(DataFusionPlanningSnafu)
6399            })
6400            .transpose()?
6401            .with_context(|| UnexpectedPlanExprSnafu {
6402                desc: "vector matching on a plan without a time index",
6403            })?;
6404        let violation = if matches!(
6405            modifier.as_ref().map(|modifier| &modifier.card),
6406            Some(VectorMatchCardinality::OneToOne)
6407        ) {
6408            MatchGroupViolation::ImplicitManyToOne
6409        } else {
6410            MatchGroupViolation::AmbiguousGroupLabels
6411        };
6412        Self::assert_unique_match_group(
6413            join_plan,
6414            group_exprs,
6415            labels.into_iter().map(|(_, label)| label).collect(),
6416            time_index_expr,
6417            violation,
6418        )
6419    }
6420
6421    /// Guard the side of a vector matching that must hold one series per match group.
6422    fn assert_unique_one_side(
6423        plan: LogicalPlan,
6424        join_keys: &BTreeSet<String>,
6425        tag_columns: &[String],
6426        time_index_column: Option<&str>,
6427        one_side_is_left: bool,
6428    ) -> Result<LogicalPlan> {
6429        let Some(time_index_column) = time_index_column else {
6430            return Ok(plan);
6431        };
6432        if tag_columns.iter().all(|tag| join_keys.contains(tag)) {
6433            return Ok(plan);
6434        }
6435
6436        let group_labels = join_keys.iter().cloned().collect::<Vec<_>>();
6437        let group_exprs = group_labels
6438            .iter()
6439            .map(|label| DfExpr::Column(Column::from_name(label)))
6440            .collect();
6441        Self::assert_unique_match_group(
6442            plan,
6443            group_exprs,
6444            group_labels,
6445            DfExpr::Column(Column::from_name(time_index_column)),
6446            MatchGroupViolation::DuplicateOnOneSide { one_side_is_left },
6447        )
6448    }
6449
6450    fn selected_binary_match_labels(
6451        left_context: &PromPlannerContext,
6452        right_context: &PromPlannerContext,
6453        modifier: &Option<BinModifier>,
6454    ) -> BTreeSet<String> {
6455        let mut labels = left_context
6456            .tag_columns
6457            .iter()
6458            .chain(&right_context.tag_columns)
6459            .cloned()
6460            .collect::<BTreeSet<_>>();
6461        if let Some(matching) = modifier
6462            .as_ref()
6463            .and_then(|modifier| modifier.matching.as_ref())
6464        {
6465            match matching {
6466                LabelModifier::Include(on) => {
6467                    labels = on
6468                        .labels
6469                        .iter()
6470                        .filter(|label| {
6471                            left_context.tag_columns.contains(label)
6472                                || right_context.tag_columns.contains(label)
6473                        })
6474                        .cloned()
6475                        .collect();
6476                }
6477                LabelModifier::Exclude(ignoring) => {
6478                    for label in &ignoring.labels {
6479                        labels.remove(label);
6480                    }
6481                }
6482            }
6483        }
6484        labels
6485    }
6486
6487    fn only_temporality_match_label_mismatches(
6488        left_context: &PromPlannerContext,
6489        right_context: &PromPlannerContext,
6490        modifier: &Option<BinModifier>,
6491    ) -> bool {
6492        let mut mismatches =
6493            Self::selected_binary_match_labels(left_context, right_context, modifier)
6494                .into_iter()
6495                .filter(|label| {
6496                    left_context.tag_columns.contains(label)
6497                        != right_context.tag_columns.contains(label)
6498                });
6499        matches!(
6500            (mismatches.next(), mismatches.next()),
6501            (Some(label), None) if label == OTLP_AGGREGATION_TEMPORALITY_LABEL
6502        )
6503    }
6504
6505    fn align_temporality_match_column(
6506        mut left: LogicalPlan,
6507        mut right: LogicalPlan,
6508        left_context: &mut PromPlannerContext,
6509        right_context: &mut PromPlannerContext,
6510    ) -> Result<(LogicalPlan, LogicalPlan, bool)> {
6511        let marker = OTLP_AGGREGATION_TEMPORALITY_LABEL;
6512        let left_has_marker = left_context.tag_columns.iter().any(|tag| tag == marker);
6513        let (data_type, value_type, add_to_left) = {
6514            let (present, add_to_left) = if left_has_marker {
6515                (&left, false)
6516            } else {
6517                (&right, true)
6518            };
6519            let data_type = present
6520                .schema()
6521                .fields()
6522                .iter()
6523                .find(|field| field.name() == marker)
6524                .map(|field| field.data_type().clone())
6525                .with_context(|| ColumnNotFoundSnafu {
6526                    col: marker.to_string(),
6527                })?;
6528            let value_type = Self::string_value_data_type(&data_type)
6529                .cloned()
6530                .with_context(|| UnexpectedPlanExprSnafu {
6531                    desc: format!("temporality match label {marker} must be a string"),
6532                })?;
6533            (data_type, value_type, add_to_left)
6534        };
6535        let null = Self::string_scalar_value(&value_type, None).with_context(|| {
6536            UnexpectedPlanExprSnafu {
6537                desc: format!("temporality match label {marker} must be a string"),
6538            }
6539        })?;
6540        let add_marker = |plan: LogicalPlan| {
6541            let visible = plan
6542                .schema()
6543                .iter()
6544                .map(|(qualifier, field)| {
6545                    DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6546                })
6547                .collect::<Vec<_>>();
6548            LogicalPlanBuilder::from(plan)
6549                .project(
6550                    visible
6551                        .into_iter()
6552                        .chain([DfExpr::Literal(null, None).alias(marker)]),
6553                )
6554                .context(DataFusionPlanningSnafu)?
6555                .build()
6556                .context(DataFusionPlanningSnafu)
6557        };
6558        if data_type != value_type {
6559            let present = if add_to_left { &mut right } else { &mut left };
6560            let visible = present
6561                .schema()
6562                .iter()
6563                .map(|(qualifier, field)| {
6564                    let column =
6565                        DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()));
6566                    if field.name() == marker {
6567                        DfExpr::Cast(Cast::new(Box::new(column), value_type.clone()))
6568                            .alias_qualified(qualifier.cloned(), field.name().clone())
6569                    } else {
6570                        column
6571                    }
6572                })
6573                .collect::<Vec<_>>();
6574            *present = LogicalPlanBuilder::from(present.clone())
6575                .project(visible)
6576                .context(DataFusionPlanningSnafu)?
6577                .build()
6578                .context(DataFusionPlanningSnafu)?;
6579        }
6580        if add_to_left {
6581            left = add_marker(left)?;
6582            left_context.tag_columns.push(marker.to_string());
6583        } else {
6584            right = add_marker(right)?;
6585            right_context.tag_columns.push(marker.to_string());
6586        }
6587        Ok((left, right, add_to_left))
6588    }
6589
6590    fn normalized_match_key_expr(
6591        label: &str,
6592        field: Option<(Option<TableReference>, ArrowDataType)>,
6593        value_type: &ArrowDataType,
6594        internal_name: &str,
6595    ) -> DfExpr {
6596        let empty = Self::string_scalar_value(value_type, Some(String::new()))
6597            .expect("match label value type is a string");
6598        let expr = if let Some((qualifier, data_type)) = field {
6599            let column = DfExpr::Column(Column::new(qualifier, label));
6600            let column = if &data_type == value_type {
6601                column
6602            } else {
6603                DfExpr::Cast(Cast::new(Box::new(column), value_type.clone()))
6604            };
6605            DfExpr::ScalarFunction(ScalarFunction {
6606                func: coalesce(),
6607                args: vec![column, DfExpr::Literal(empty, None)],
6608            })
6609        } else {
6610            DfExpr::Literal(empty, None)
6611        };
6612        expr.alias(internal_name)
6613    }
6614
6615    fn is_zero_row_empty_relation(plan: &LogicalPlan) -> bool {
6616        // `produce_one_row` is used for input-free plans that still emit one row;
6617        // only the false case is a statically proven empty vector.
6618        matches!(plan, LogicalPlan::EmptyRelation(relation) if !relation.produce_one_row)
6619    }
6620
6621    /// Build a set operator (AND/OR/UNLESS)
6622    fn set_op_on_non_field_columns(
6623        &mut self,
6624        mut left: LogicalPlan,
6625        mut right: LogicalPlan,
6626        left_context: PromPlannerContext,
6627        right_context: PromPlannerContext,
6628        op: TokenType,
6629        modifier: &Option<BinModifier>,
6630    ) -> Result<LogicalPlan> {
6631        let left_tag_col_set = left_context
6632            .tag_columns
6633            .iter()
6634            .cloned()
6635            .collect::<HashSet<_>>();
6636        let right_tag_col_set = right_context
6637            .tag_columns
6638            .iter()
6639            .cloned()
6640            .collect::<HashSet<_>>();
6641
6642        if matches!(op.id(), token::T_LOR) {
6643            return self.or_operator(
6644                left,
6645                right,
6646                left_tag_col_set,
6647                right_tag_col_set,
6648                left_context,
6649                right_context,
6650                modifier,
6651            );
6652        }
6653
6654        if let Some(modifier) = modifier {
6655            ensure!(
6656                matches!(
6657                    modifier.card,
6658                    VectorMatchCardinality::OneToOne | VectorMatchCardinality::ManyToMany
6659                ),
6660                UnsupportedVectorMatchSnafu {
6661                    name: modifier.card.clone(),
6662                },
6663            );
6664        }
6665
6666        let output_context = left_context.clone();
6667        let visible_left_schema = left.schema().clone();
6668        let mut left_context = left_context;
6669        let mut right_context = right_context;
6670        let added_marker_to_left = if Self::only_temporality_match_label_mismatches(
6671            &left_context,
6672            &right_context,
6673            modifier,
6674        ) {
6675            let aligned = Self::align_temporality_match_column(
6676                left,
6677                right,
6678                &mut left_context,
6679                &mut right_context,
6680            )?;
6681            left = aligned.0;
6682            right = aligned.1;
6683            aligned.2
6684        } else {
6685            false
6686        };
6687
6688        let mut left_tag_col_set = left_context
6689            .tag_columns
6690            .iter()
6691            .cloned()
6692            .collect::<BTreeSet<_>>();
6693        let mut right_tag_col_set = right_context
6694            .tag_columns
6695            .iter()
6696            .cloned()
6697            .collect::<BTreeSet<_>>();
6698        if let Some(matching) = modifier
6699            .as_ref()
6700            .and_then(|modifier| modifier.matching.as_ref())
6701        {
6702            match matching {
6703                LabelModifier::Include(on) => {
6704                    let mask = on.labels.iter().cloned().collect::<BTreeSet<_>>();
6705                    left_tag_col_set = left_tag_col_set.intersection(&mask).cloned().collect();
6706                    right_tag_col_set = right_tag_col_set.intersection(&mask).cloned().collect();
6707                }
6708                LabelModifier::Exclude(ignoring) => {
6709                    for label in &ignoring.labels {
6710                        let _ = left_tag_col_set.remove(label);
6711                        let _ = right_tag_col_set.remove(label);
6712                    }
6713                }
6714            }
6715        }
6716        ensure!(
6717            left_tag_col_set == right_tag_col_set,
6718            CombineTableColumnMismatchSnafu {
6719                left: left_tag_col_set.iter().cloned().collect::<Vec<_>>(),
6720                right: right_tag_col_set.iter().cloned().collect::<Vec<_>>(),
6721            }
6722        );
6723
6724        let left_time_index = left_context.time_index_column.clone().unwrap();
6725        let right_time_index = right_context.time_index_column.clone().unwrap();
6726
6727        // alias right time index column if necessary
6728        if left_context.time_index_column != right_context.time_index_column {
6729            let right_project_exprs = right
6730                .schema()
6731                .fields()
6732                .iter()
6733                .map(|field| {
6734                    if field.name() == &right_time_index {
6735                        DfExpr::Column(Column::from_name(&right_time_index)).alias(&left_time_index)
6736                    } else {
6737                        DfExpr::Column(Column::from_name(field.name()))
6738                    }
6739                })
6740                .collect::<Vec<_>>();
6741
6742            right = LogicalPlanBuilder::from(right)
6743                .project(right_project_exprs)
6744                .context(DataFusionPlanningSnafu)?
6745                .build()
6746                .context(DataFusionPlanningSnafu)?;
6747        }
6748
6749        let join_keys = left_tag_col_set
6750            .into_iter()
6751            .chain([left_time_index])
6752            .collect::<Vec<_>>();
6753
6754        ensure!(
6755            left_context.field_columns.len() == 1
6756                || Self::field_columns_are_alternative_samples(
6757                    left.schema(),
6758                    &left_context.field_columns,
6759                ),
6760            MultiFieldsNotSupportedSnafu {
6761                operator: "AND/UNLESS operator"
6762            }
6763        );
6764        // Generate join plan.
6765        // All set operations in PromQL are "distinct"
6766        let result = match op.id() {
6767            token::T_LAND => LogicalPlanBuilder::from(left)
6768                .distinct()
6769                .context(DataFusionPlanningSnafu)?
6770                .join_detailed(
6771                    right,
6772                    JoinType::LeftSemi,
6773                    (join_keys.clone(), join_keys),
6774                    None,
6775                    NullEquality::NullEqualsNull,
6776                )
6777                .context(DataFusionPlanningSnafu)?
6778                .build()
6779                .context(DataFusionPlanningSnafu),
6780            token::T_LUNLESS => LogicalPlanBuilder::from(left)
6781                .distinct()
6782                .context(DataFusionPlanningSnafu)?
6783                .join_detailed(
6784                    right,
6785                    JoinType::LeftAnti,
6786                    (join_keys.clone(), join_keys),
6787                    None,
6788                    NullEquality::NullEqualsNull,
6789                )
6790                .context(DataFusionPlanningSnafu)?
6791                .build()
6792                .context(DataFusionPlanningSnafu),
6793            token::T_LOR => {
6794                // OR is handled at the beginning of this function, as it cannot
6795                // be expressed using JOIN like AND and UNLESS.
6796                unreachable!()
6797            }
6798            _ => UnexpectedTokenSnafu { token: op }.fail(),
6799        }?;
6800        let result = if added_marker_to_left {
6801            LogicalPlanBuilder::from(result)
6802                .project(visible_left_schema.iter().map(|(qualifier, field)| {
6803                    DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6804                }))
6805                .context(DataFusionPlanningSnafu)?
6806                .build()
6807                .context(DataFusionPlanningSnafu)?
6808        } else {
6809            result
6810        };
6811
6812        // AND/UNLESS preserve the complete left operand's visible columns and values; encoded
6813        // markers are decoded.
6814        self.ctx = output_context;
6815        Ok(result)
6816    }
6817
6818    fn string_value_data_type(data_type: &ArrowDataType) -> Option<&ArrowDataType> {
6819        match data_type {
6820            data_type if data_type.is_string() => Some(data_type),
6821            ArrowDataType::Dictionary(_, value_type) if value_type.is_string() => Some(value_type),
6822            _ => None,
6823        }
6824    }
6825
6826    fn string_scalar_value(
6827        data_type: &ArrowDataType,
6828        value: Option<String>,
6829    ) -> Option<ScalarValue> {
6830        match data_type {
6831            ArrowDataType::Utf8 => Some(ScalarValue::Utf8(value)),
6832            ArrowDataType::LargeUtf8 => Some(ScalarValue::LargeUtf8(value)),
6833            ArrowDataType::Utf8View => Some(ScalarValue::Utf8View(value)),
6834            ArrowDataType::Dictionary(key_type, value_type) => Some(ScalarValue::Dictionary(
6835                key_type.clone(),
6836                Box::new(Self::string_scalar_value(value_type, value)?),
6837            )),
6838            _ => None,
6839        }
6840    }
6841
6842    fn common_label_data_type(
6843        left: Option<&ArrowDataType>,
6844        right: Option<&ArrowDataType>,
6845    ) -> Option<ArrowDataType> {
6846        match (left, right) {
6847            (Some(left), Some(right)) if left == right => {
6848                Self::string_value_data_type(left).map(|_| left.clone())
6849            }
6850            (Some(left), Some(right)) => {
6851                let left_value_type = Self::string_value_data_type(left)?;
6852                let right_value_type = Self::string_value_data_type(right)?;
6853                // DataFusion projections can decode dictionaries, but do not encode plain strings
6854                // as dictionaries. Preserve the encoding only when both inputs already share it.
6855                match (left_value_type, right_value_type) {
6856                    (left, right) if left == right => Some(left.clone()),
6857                    (ArrowDataType::LargeUtf8, _) | (_, ArrowDataType::LargeUtf8) => {
6858                        Some(ArrowDataType::LargeUtf8)
6859                    }
6860                    (ArrowDataType::Utf8View, ArrowDataType::Utf8View) => {
6861                        Some(ArrowDataType::Utf8View)
6862                    }
6863                    _ => Some(ArrowDataType::Utf8),
6864                }
6865            }
6866            (Some(data_type), None) | (None, Some(data_type)) => {
6867                Self::string_value_data_type(data_type).cloned()
6868            }
6869            (None, None) => Some(ArrowDataType::Utf8),
6870        }
6871    }
6872
6873    // TODO(ruihang): change function name
6874    #[allow(clippy::too_many_arguments)]
6875    fn or_operator(
6876        &mut self,
6877        left: LogicalPlan,
6878        right: LogicalPlan,
6879        left_tag_cols_set: HashSet<String>,
6880        right_tag_cols_set: HashSet<String>,
6881        left_context: PromPlannerContext,
6882        right_context: PromPlannerContext,
6883        modifier: &Option<BinModifier>,
6884    ) -> Result<LogicalPlan> {
6885        let left_is_empty = Self::is_zero_row_empty_relation(&left);
6886        let right_is_empty = Self::is_zero_row_empty_relation(&right);
6887        match (left_is_empty, right_is_empty) {
6888            (true, false) => {
6889                self.ctx = right_context;
6890                return Ok(right);
6891            }
6892            (false, true) => {
6893                self.ctx = left_context;
6894                return Ok(left);
6895            }
6896            (true, true) => {
6897                self.ctx = left_context;
6898                return Ok(left);
6899            }
6900            (false, false) => {}
6901        }
6902
6903        ensure!(
6904            !left.schema().fields().is_empty() && !right.schema().fields().is_empty(),
6905            UnexpectedPlanExprSnafu {
6906                desc: "OR operator input has zero columns",
6907            }
6908        );
6909        let left_has_alternative_samples =
6910            Self::field_columns_are_alternative_samples(left.schema(), &left_context.field_columns);
6911        let right_has_alternative_samples = Self::field_columns_are_alternative_samples(
6912            right.schema(),
6913            &right_context.field_columns,
6914        );
6915        ensure!(
6916            left_context.field_columns.len() == 1 || left_has_alternative_samples,
6917            MultiFieldsNotSupportedSnafu {
6918                operator: "OR operator"
6919            }
6920        );
6921        ensure!(
6922            right_context.field_columns.len() == 1 || right_has_alternative_samples,
6923            MultiFieldsNotSupportedSnafu {
6924                operator: "OR operator"
6925            }
6926        );
6927
6928        // prepare hash sets
6929        let all_tags = left_tag_cols_set
6930            .union(&right_tag_cols_set)
6931            .cloned()
6932            .collect::<HashSet<_>>();
6933        let left_qualifier = left.schema().qualified_field(0).0.cloned();
6934        let right_qualifier = right.schema().qualified_field(0).0.cloned();
6935        let left_qualifier_string = left_qualifier
6936            .as_ref()
6937            .map(|l| l.to_string())
6938            .unwrap_or_default();
6939        let right_qualifier_string = right_qualifier
6940            .as_ref()
6941            .map(|r| r.to_string())
6942            .unwrap_or_default();
6943        let left_time_index_column =
6944            left_context
6945                .time_index_column
6946                .clone()
6947                .with_context(|| TimeIndexNotFoundSnafu {
6948                    table: left_qualifier_string.clone(),
6949                })?;
6950        let right_time_index_column =
6951            right_context
6952                .time_index_column
6953                .clone()
6954                .with_context(|| TimeIndexNotFoundSnafu {
6955                    table: right_qualifier_string.clone(),
6956                })?;
6957        let native_histogram_type = Self::native_histogram_arrow_type();
6958        let is_numeric = |data_type: &ArrowDataType| {
6959            matches!(
6960                data_type,
6961                ArrowDataType::Int8
6962                    | ArrowDataType::Int16
6963                    | ArrowDataType::Int32
6964                    | ArrowDataType::Int64
6965                    | ArrowDataType::UInt8
6966                    | ArrowDataType::UInt16
6967                    | ArrowDataType::UInt32
6968                    | ArrowDataType::UInt64
6969                    | ArrowDataType::Float32
6970                    | ArrowDataType::Float64
6971            )
6972        };
6973        let left_fields = left_context
6974            .field_columns
6975            .iter()
6976            .map(|name| {
6977                left.schema()
6978                    .iter()
6979                    .find(|(_, field)| field.name() == name)
6980                    .map(|(qualifier, field)| {
6981                        (name.clone(), qualifier.cloned(), field.data_type().clone())
6982                    })
6983                    .with_context(|| ColumnNotFoundSnafu { col: name.clone() })
6984            })
6985            .collect::<Result<Vec<_>>>()?;
6986        let right_fields = right_context
6987            .field_columns
6988            .iter()
6989            .map(|name| {
6990                right
6991                    .schema()
6992                    .iter()
6993                    .find(|(_, field)| field.name() == name)
6994                    .map(|(qualifier, field)| {
6995                        (name.clone(), qualifier.cloned(), field.data_type().clone())
6996                    })
6997                    .with_context(|| ColumnNotFoundSnafu { col: name.clone() })
6998            })
6999            .collect::<Result<Vec<_>>>()?;
7000        let left_field = &left_fields[0];
7001        let right_field = &right_fields[0];
7002        let left_field_col = &left_field.0;
7003        let right_field_col = &right_field.0;
7004        let fields_are_samples = |fields: &[(String, Option<TableReference>, ArrowDataType)]| {
7005            fields.iter().all(|(_, _, data_type)| {
7006                is_numeric(data_type) || data_type == &native_histogram_type
7007            })
7008        };
7009        let mixed_sample_types = if left_has_alternative_samples || right_has_alternative_samples {
7010            if !fields_are_samples(&left_fields) || !fields_are_samples(&right_fields) {
7011                return UnexpectedPlanExprSnafu {
7012                    desc: format!(
7013                        "OR value fields have incompatible types: {:?} and {:?}",
7014                        left_fields
7015                            .iter()
7016                            .map(|(_, _, data_type)| data_type)
7017                            .collect::<Vec<_>>(),
7018                        right_fields
7019                            .iter()
7020                            .map(|(_, _, data_type)| data_type)
7021                            .collect::<Vec<_>>()
7022                    ),
7023                }
7024                .fail();
7025            }
7026            true
7027        } else {
7028            (left_field.2 == native_histogram_type && is_numeric(&right_field.2))
7029                || (right_field.2 == native_histogram_type && is_numeric(&left_field.2))
7030        };
7031        let target_field_type = if mixed_sample_types {
7032            // Mixed vectors use the existing response representation: one nullable float column
7033            // and one nullable native-histogram column.
7034            ArrowDataType::Float64
7035        } else if left_field.2 == right_field.2 {
7036            left_field.2.clone()
7037        } else if is_numeric(&left_field.2) && is_numeric(&right_field.2) {
7038            ArrowDataType::Float64
7039        } else {
7040            return UnexpectedPlanExprSnafu {
7041                desc: format!(
7042                    "OR value fields have incompatible types: {:?} and {:?}",
7043                    left_field.2, right_field.2
7044                ),
7045            }
7046            .fail();
7047        };
7048        let (mixed_float_field_col, mixed_histogram_field_col) = if mixed_sample_types {
7049            let mut reserved_names = left
7050                .schema()
7051                .fields()
7052                .iter()
7053                .chain(right.schema().fields().iter())
7054                .map(|field| field.name().clone())
7055                .collect::<HashSet<_>>();
7056            for (name, _, _) in left_fields.iter().chain(&right_fields) {
7057                reserved_names.remove(name);
7058            }
7059            reserved_names.extend(all_tags.iter().cloned());
7060            let unique_name = |prefix: &str, reserved_names: &mut HashSet<String>| {
7061                let mut index = 0;
7062                loop {
7063                    let name = format!("{prefix}{index}");
7064                    index += 1;
7065                    if reserved_names.insert(name.clone()) {
7066                        break name;
7067                    }
7068                }
7069            };
7070            let float_field = unique_name(OR_FLOAT_FIELD_PREFIX, &mut reserved_names);
7071            let histogram_field = unique_name(OR_HISTOGRAM_FIELD_PREFIX, &mut reserved_names);
7072            (float_field, histogram_field)
7073        } else {
7074            (left_field_col.clone(), String::new())
7075        };
7076        let left_tag_types = left_tag_cols_set
7077            .iter()
7078            .map(|label| {
7079                left.schema()
7080                    .fields()
7081                    .iter()
7082                    .find(|field| field.name() == label)
7083                    .map(|field| (label.clone(), field.data_type().clone()))
7084                    .with_context(|| ColumnNotFoundSnafu { col: label.clone() })
7085            })
7086            .collect::<Result<HashMap<_, _>>>()?;
7087        let right_tag_types = right_tag_cols_set
7088            .iter()
7089            .map(|label| {
7090                right
7091                    .schema()
7092                    .fields()
7093                    .iter()
7094                    .find(|field| field.name() == label)
7095                    .map(|field| (label.clone(), field.data_type().clone()))
7096                    .with_context(|| ColumnNotFoundSnafu { col: label.clone() })
7097            })
7098            .collect::<Result<HashMap<_, _>>>()?;
7099        let mut target_tag_types = HashMap::with_capacity(all_tags.len());
7100        for label in &all_tags {
7101            let Some(data_type) =
7102                Self::common_label_data_type(left_tag_types.get(label), right_tag_types.get(label))
7103            else {
7104                return UnexpectedPlanExprSnafu {
7105                    desc: format!(
7106                        "OR label {label} has incompatible types: {:?} and {:?}",
7107                        left_tag_types.get(label),
7108                        right_tag_types.get(label)
7109                    ),
7110                }
7111                .fail();
7112            };
7113            target_tag_types.insert(label.clone(), data_type);
7114        }
7115        let left_has_tsid = left
7116            .schema()
7117            .fields()
7118            .iter()
7119            .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME);
7120        let right_has_tsid = right
7121            .schema()
7122            .fields()
7123            .iter()
7124            .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME);
7125
7126        // step 0: fill all columns in output schema
7127        let mut all_columns_set = left
7128            .schema()
7129            .fields()
7130            .iter()
7131            .chain(right.schema().fields().iter())
7132            .map(|field| field.name().clone())
7133            .collect::<HashSet<_>>();
7134        // Keep `__tsid` only when both sides contain it, otherwise it may break schema alignment
7135        // (e.g. `unknown_metric or some_metric`).
7136        if !(left_has_tsid && right_has_tsid) {
7137            all_columns_set.remove(DATA_SCHEMA_TSID_COLUMN_NAME);
7138        }
7139        // remove time index column
7140        all_columns_set.remove(&left_time_index_column);
7141        all_columns_set.remove(&right_time_index_column);
7142        if mixed_sample_types {
7143            for (name, _, _) in left_fields.iter().chain(&right_fields) {
7144                all_columns_set.remove(name);
7145            }
7146            all_columns_set.extend(all_tags.iter().cloned());
7147            all_columns_set.insert(mixed_float_field_col.clone());
7148            all_columns_set.insert(mixed_histogram_field_col.clone());
7149        } else if left_field_col != right_field_col {
7150            // remove field column in the right
7151            all_columns_set.remove(right_field_col);
7152        }
7153        let mut all_columns = all_columns_set.into_iter().collect::<Vec<_>>();
7154        // sort to ensure the generated schema is not volatile
7155        all_columns.sort_unstable();
7156        // use left time index column name as the result time index column name
7157        all_columns.insert(0, left_time_index_column.clone());
7158        let mut occupied_column_names = left
7159            .schema()
7160            .fields()
7161            .iter()
7162            .chain(right.schema().fields().iter())
7163            .map(|field| field.name().clone())
7164            .collect::<HashSet<_>>();
7165
7166        // step 1: align schema using project, fill non-exist columns with null
7167        let aligned_label_expr = |col: &String, source_types: &HashMap<String, ArrowDataType>| {
7168            let target_type = &target_tag_types[col];
7169            if let Some(source_type) = source_types.get(col) {
7170                let expr = DfExpr::Column(Column::new(None::<String>, col));
7171                if source_type == target_type {
7172                    expr
7173                } else {
7174                    DfExpr::Cast(Cast::new(Box::new(expr), target_type.clone())).alias(col.clone())
7175                }
7176            } else {
7177                DfExpr::Literal(
7178                    Self::string_scalar_value(target_type, None)
7179                        .expect("target label type is a string"),
7180                    None,
7181                )
7182                .alias(col.clone())
7183            }
7184        };
7185        let null_histogram =
7186            ScalarValue::try_new_null(&native_histogram_type).context(DataFusionPlanningSnafu)?;
7187        let mixed_value_expr = |fields: &[(String, Option<TableReference>, ArrowDataType)],
7188                                output_col: &String| {
7189            if output_col == &mixed_float_field_col {
7190                if let Some((name, qualifier, data_type)) = fields
7191                    .iter()
7192                    .find(|(_, _, data_type)| is_numeric(data_type))
7193                {
7194                    let expr = DfExpr::Column(Column::new(qualifier.clone(), name));
7195                    if data_type == &ArrowDataType::Float64 {
7196                        expr.alias(output_col)
7197                    } else {
7198                        DfExpr::Cast(Cast::new(Box::new(expr), ArrowDataType::Float64))
7199                            .alias(output_col)
7200                    }
7201                } else {
7202                    DfExpr::Literal(ScalarValue::Float64(None), None).alias(output_col)
7203                }
7204            } else {
7205                fields
7206                    .iter()
7207                    .find(|(_, _, data_type)| data_type == &native_histogram_type)
7208                    .map(|(name, qualifier, _)| {
7209                        DfExpr::Column(Column::new(qualifier.clone(), name)).alias(output_col)
7210                    })
7211                    .unwrap_or_else(|| {
7212                        DfExpr::Literal(null_histogram.clone(), None).alias(output_col)
7213                    })
7214            }
7215        };
7216        let left_proj_exprs = all_columns.iter().map(|col| {
7217            if mixed_sample_types
7218                && (col == &mixed_float_field_col || col == &mixed_histogram_field_col)
7219            {
7220                mixed_value_expr(&left_fields, col)
7221            } else if !mixed_sample_types
7222                && col == left_field_col
7223                && left_field.2 != target_field_type
7224            {
7225                DfExpr::Cast(Cast::new(
7226                    Box::new(DfExpr::Column(Column::new(
7227                        left_field.1.clone(),
7228                        left_field_col,
7229                    ))),
7230                    target_field_type.clone(),
7231                ))
7232                .alias(left_field_col.clone())
7233            } else if target_tag_types.contains_key(col) {
7234                aligned_label_expr(col, &left_tag_types)
7235            } else {
7236                DfExpr::Column(Column::new(None::<String>, col))
7237            }
7238        });
7239        let right_time_index_expr = DfExpr::Column(Column::new(
7240            right_qualifier.clone(),
7241            right_time_index_column,
7242        ))
7243        .alias(left_time_index_column.clone());
7244        // The field column in right side may not have qualifier (it may be removed by join operation),
7245        // so we need to find it from the schema.
7246        // `skip(1)` to skip the time index column
7247        let right_proj_exprs_without_time_index = all_columns.iter().skip(1).map(|col| {
7248            // expr
7249            if mixed_sample_types
7250                && (col == &mixed_float_field_col || col == &mixed_histogram_field_col)
7251            {
7252                mixed_value_expr(&right_fields, col)
7253            } else if !mixed_sample_types && col == left_field_col {
7254                let expr = DfExpr::Column(Column::new(right_field.1.clone(), right_field_col));
7255                if right_field.2 != target_field_type {
7256                    DfExpr::Cast(Cast::new(Box::new(expr), target_field_type.clone()))
7257                        .alias(left_field_col.clone())
7258                } else if left_field_col != right_field_col {
7259                    expr.alias(left_field_col.clone())
7260                } else {
7261                    expr
7262                }
7263            } else if target_tag_types.contains_key(col) {
7264                aligned_label_expr(col, &right_tag_types)
7265            } else {
7266                DfExpr::Column(Column::new(None::<String>, col))
7267            }
7268        });
7269        let right_proj_exprs = [right_time_index_expr]
7270            .into_iter()
7271            .chain(right_proj_exprs_without_time_index);
7272
7273        let left_projected = LogicalPlanBuilder::from(left)
7274            .project(left_proj_exprs)
7275            .context(DataFusionPlanningSnafu)?
7276            .alias(left_qualifier_string.clone())
7277            .context(DataFusionPlanningSnafu)?
7278            .build()
7279            .context(DataFusionPlanningSnafu)?;
7280        let right_projected = LogicalPlanBuilder::from(right)
7281            .project(right_proj_exprs)
7282            .context(DataFusionPlanningSnafu)?
7283            .alias(right_qualifier_string.clone())
7284            .context(DataFusionPlanningSnafu)?
7285            .build()
7286            .context(DataFusionPlanningSnafu)?;
7287
7288        // step 2: compute match columns
7289        let mut match_columns = if let Some(modifier) = modifier
7290            && let Some(matching) = &modifier.matching
7291        {
7292            match matching {
7293                // keeps columns mentioned in `on`
7294                LabelModifier::Include(on) => on.labels.clone(),
7295                // removes columns memtioned in `ignoring`
7296                LabelModifier::Exclude(ignoring) => {
7297                    let ignoring = ignoring.labels.iter().cloned().collect::<HashSet<_>>();
7298                    all_tags.difference(&ignoring).cloned().collect()
7299                }
7300            }
7301        } else {
7302            all_tags.iter().cloned().collect()
7303        };
7304        // sort to ensure the generated plan is not volatile
7305        match_columns.sort_unstable();
7306        match_columns.dedup();
7307        occupied_column_names.extend(
7308            left_projected
7309                .schema()
7310                .fields()
7311                .iter()
7312                .chain(right_projected.schema().fields().iter())
7313                .map(|field| field.name().clone()),
7314        );
7315
7316        let visible_schema = left_projected.schema().clone();
7317        let visible_left_exprs = left_projected
7318            .schema()
7319            .iter()
7320            .map(|(qualifier, field)| {
7321                DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
7322            })
7323            .collect::<Vec<_>>();
7324        let visible_right_exprs = right_projected
7325            .schema()
7326            .iter()
7327            .map(|(qualifier, field)| {
7328                DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
7329            })
7330            .collect::<Vec<_>>();
7331        let mut left_match_exprs = Vec::with_capacity(match_columns.len());
7332        let mut right_match_exprs = Vec::with_capacity(match_columns.len());
7333        let mut next_internal_column = 0;
7334
7335        for label in &match_columns {
7336            let left_field = if left_tag_cols_set.contains(label) {
7337                Some(
7338                    left_projected
7339                        .schema()
7340                        .iter()
7341                        .find(|(_, field)| field.name() == label)
7342                        .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone()))
7343                        .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?,
7344                )
7345            } else {
7346                None
7347            };
7348            let right_field = if right_tag_cols_set.contains(label) {
7349                Some(
7350                    right_projected
7351                        .schema()
7352                        .iter()
7353                        .find(|(_, field)| field.name() == label)
7354                        .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone()))
7355                        .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?,
7356                )
7357            } else {
7358                None
7359            };
7360            let data_type = match (left_field.as_ref(), right_field.as_ref()) {
7361                (Some((_, left_type)), Some((_, right_type))) if left_type == right_type => {
7362                    left_type.clone()
7363                }
7364                (Some((_, left_type)), Some((_, right_type))) => {
7365                    return UnexpectedPlanExprSnafu {
7366                        desc: format!(
7367                            "OR match label {label} has incompatible types: {left_type:?} and {right_type:?}"
7368                        ),
7369                    }
7370                    .fail();
7371                }
7372                (Some((_, data_type)), None) | (None, Some((_, data_type))) => data_type.clone(),
7373                (None, None) => ArrowDataType::Utf8,
7374            };
7375            let Some(value_type) = Self::string_value_data_type(&data_type).cloned() else {
7376                return UnexpectedPlanExprSnafu {
7377                    desc: format!("OR match label {label} must be a string"),
7378                }
7379                .fail();
7380            };
7381            let internal_name = loop {
7382                let name = format!("__promql_or_match_{next_internal_column}");
7383                next_internal_column += 1;
7384                if occupied_column_names.insert(name.clone()) {
7385                    break name;
7386                }
7387            };
7388            left_match_exprs.push(Self::normalized_match_key_expr(
7389                label,
7390                left_field,
7391                &value_type,
7392                &internal_name,
7393            ));
7394            right_match_exprs.push(Self::normalized_match_key_expr(
7395                label,
7396                right_field,
7397                &value_type,
7398                &internal_name,
7399            ));
7400        }
7401
7402        let left_augmented = LogicalPlanBuilder::from(left_projected)
7403            .project(visible_left_exprs.into_iter().chain(left_match_exprs))
7404            .context(DataFusionPlanningSnafu)?
7405            .build()
7406            .context(DataFusionPlanningSnafu)?;
7407        let right_augmented = LogicalPlanBuilder::from(right_projected)
7408            .project(visible_right_exprs.into_iter().chain(right_match_exprs))
7409            .context(DataFusionPlanningSnafu)?
7410            .build()
7411            .context(DataFusionPlanningSnafu)?;
7412
7413        // step 3: build `UnionDistinctOn` with normalized internal match keys.
7414        let visible_field_count = visible_schema.fields().len();
7415        let compare_key_indices =
7416            (visible_field_count..visible_field_count + match_columns.len()).collect::<Vec<_>>();
7417        let (time_qualifier, _) = visible_schema
7418            .iter()
7419            .find(|(_, field)| field.name() == &left_time_index_column)
7420            .with_context(|| TimeIndexNotFoundSnafu {
7421                table: left_qualifier_string.clone(),
7422            })?;
7423        let ts_col_idx = left_augmented
7424            .schema()
7425            .iter()
7426            .position(|(qualifier, field)| {
7427                qualifier == time_qualifier && field.name() == &left_time_index_column
7428            })
7429            .with_context(|| TimeIndexNotFoundSnafu {
7430                table: left_qualifier_string.clone(),
7431            })?;
7432        let union_distinct_on = UnionDistinctOn::try_new(
7433            left_augmented,
7434            right_augmented,
7435            compare_key_indices,
7436            ts_col_idx,
7437        )
7438        .context(DataFusionPlanningSnafu)?;
7439        let augmented_result = LogicalPlan::Extension(Extension {
7440            node: Arc::new(union_distinct_on),
7441        });
7442        let result = LogicalPlanBuilder::from(augmented_result)
7443            .project(visible_schema.iter().map(|(qualifier, field)| {
7444                DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
7445            }))
7446            .context(DataFusionPlanningSnafu)?
7447            .build()
7448            .context(DataFusionPlanningSnafu)?;
7449
7450        // step 4: update context
7451        let output_field_col = left_field_col.clone();
7452        let mut output_context = left_context;
7453        let mut visible_tags = all_tags.into_iter().collect::<Vec<_>>();
7454        visible_tags.sort_unstable();
7455        output_context.time_index_column = Some(left_time_index_column);
7456        output_context.tag_columns = visible_tags;
7457        output_context.field_columns = if mixed_sample_types {
7458            vec![mixed_float_field_col, mixed_histogram_field_col]
7459        } else {
7460            vec![output_field_col]
7461        };
7462        output_context.use_tsid = left_has_tsid && right_has_tsid;
7463        self.ctx = output_context;
7464
7465        Ok(result)
7466    }
7467
7468    /// Build a projection that project and perform operation expr for every value columns.
7469    /// Non-value columns (tag and timestamp) will be preserved in the projection.
7470    ///
7471    /// # Side effect
7472    ///
7473    /// This function will update the value columns in the context. Those new column names
7474    /// don't contains qualifier.
7475    fn projection_for_each_field_column<F>(
7476        &mut self,
7477        input: LogicalPlan,
7478        name_to_expr: F,
7479    ) -> Result<LogicalPlan>
7480    where
7481        F: FnMut(&String) -> Result<DfExpr>,
7482    {
7483        self.projection_for_each_field_column_with_labels(input, None, name_to_expr)
7484    }
7485
7486    /// Like [`Self::projection_for_each_field_column`], but projects `result_labels` instead of
7487    /// the context tag columns when a binary operation derived its own result label set.
7488    fn projection_for_each_field_column_with_labels<F>(
7489        &mut self,
7490        input: LogicalPlan,
7491        result_labels: Option<&BinaryResultLabels>,
7492        name_to_expr: F,
7493    ) -> Result<LogicalPlan>
7494    where
7495        F: FnMut(&String) -> Result<DfExpr>,
7496    {
7497        // Keep the generated float/histogram lane names while an element-wise operation
7498        // preserves both sample types, so downstream operators still recognize the pair.
7499        let preserve_field_names =
7500            Self::field_columns_are_alternative_samples(input.schema(), &self.ctx.field_columns);
7501        let table_ref = self.ctx.table_name.clone().map(TableReference::bare);
7502        // Derived labels can be unqualified even when the context still names the source table.
7503        let input_schema = input.schema().clone();
7504        let lookup = |col: &String| {
7505            input_schema
7506                .qualified_field_with_name(table_ref.as_ref(), col)
7507                .or_else(|_| input_schema.qualified_field_with_unqualified_name(col))
7508                .map(|field| DfExpr::Column(field.into()))
7509                .context(DataFusionPlanningSnafu)
7510        };
7511        let tag_columns_iter = match result_labels {
7512            Some(labels) => labels.exprs.iter().cloned().map(Ok).collect::<Vec<_>>(),
7513            None => self.ctx.tag_columns.iter().map(lookup).collect::<Vec<_>>(),
7514        };
7515        let non_field_columns_iter = tag_columns_iter
7516            .into_iter()
7517            .chain(self.ctx.time_index_column.iter().map(lookup));
7518        let tsid_iter = match result_labels {
7519            Some(labels) => labels.tsid_projection(table_ref.clone()),
7520            None => Self::optional_tsid_projection(
7521                input.schema(),
7522                table_ref.as_ref(),
7523                self.ctx.use_tsid,
7524            ),
7525        }
7526        .into_iter()
7527        .map(Ok);
7528
7529        // build computation exprs
7530        let result_field_columns = self
7531            .ctx
7532            .field_columns
7533            .iter()
7534            .map(name_to_expr)
7535            .collect::<Result<Vec<_>>>()?;
7536
7537        // alias the computation exprs to remove qualifier
7538        if !preserve_field_names {
7539            self.ctx.field_columns = result_field_columns
7540                .iter()
7541                .map(|expr| expr.schema_name().to_string())
7542                .collect();
7543        }
7544        let field_columns_iter = result_field_columns
7545            .into_iter()
7546            .zip(self.ctx.field_columns.iter())
7547            .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::<String>, name))));
7548
7549        // chain non-field columns (unchanged) and field columns (applied computation then alias)
7550        let project_fields = non_field_columns_iter
7551            .chain(tsid_iter)
7552            .chain(field_columns_iter)
7553            .collect::<Result<Vec<_>>>()?;
7554
7555        LogicalPlanBuilder::from(input)
7556            .project(project_fields)
7557            .context(DataFusionPlanningSnafu)?
7558            .build()
7559            .context(DataFusionPlanningSnafu)
7560    }
7561
7562    /// Build a filter plan on one value column or a float/histogram alternative pair.
7563    fn filter_on_field_column<F>(&self, input: LogicalPlan, name_to_expr: F) -> Result<LogicalPlan>
7564    where
7565        F: FnMut(&String) -> Result<DfExpr>,
7566    {
7567        ensure!(
7568            self.ctx.field_columns.len() == 1
7569                || Self::field_columns_are_alternative_samples(
7570                    input.schema(),
7571                    &self.ctx.field_columns,
7572                ),
7573            UnsupportedExprSnafu {
7574                name: "filter on multi-value input"
7575            }
7576        );
7577
7578        let field_column_filters = self
7579            .ctx
7580            .field_columns
7581            .iter()
7582            .map(name_to_expr)
7583            .collect::<Result<Vec<_>>>()?;
7584        let field_column_filter =
7585            disjunction(field_column_filters).context(UnsupportedExprSnafu {
7586                name: "filter on empty input",
7587            })?;
7588
7589        LogicalPlanBuilder::from(input)
7590            .filter(field_column_filter)
7591            .context(DataFusionPlanningSnafu)?
7592            .build()
7593            .context(DataFusionPlanningSnafu)
7594    }
7595
7596    /// Generate an expr like `date_part("hour", <TIME_INDEX>)`. Caller should ensure the
7597    /// time index column in context is set
7598    fn date_part_on_time_index(&self, date_part: &str) -> Result<DfExpr> {
7599        let input_expr = datafusion::logical_expr::col(
7600            self.ctx
7601                .time_index_column
7602                .as_ref()
7603                // table name doesn't matters here
7604                .with_context(|| TimeIndexNotFoundSnafu {
7605                    table: "<doesn't matter>",
7606                })?,
7607        );
7608        let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
7609            func: datafusion_functions::datetime::date_part(),
7610            args: vec![date_part.lit(), input_expr],
7611        });
7612        Ok(fn_expr)
7613    }
7614
7615    fn strip_tsid_column(&self, plan: LogicalPlan) -> Result<LogicalPlan> {
7616        let schema = plan.schema();
7617        if !schema
7618            .fields()
7619            .iter()
7620            .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
7621        {
7622            return Ok(plan);
7623        }
7624
7625        // Preserve column qualifiers so downstream plan nodes can keep referencing
7626        // the columns by their original qualified names.
7627        let project_exprs = schema
7628            .iter()
7629            .filter(|(_, field)| field.name() != DATA_SCHEMA_TSID_COLUMN_NAME)
7630            .map(|(qualifier, field)| {
7631                DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
7632            })
7633            .collect::<Vec<_>>();
7634
7635        LogicalPlanBuilder::from(plan)
7636            .project(project_exprs)
7637            .context(DataFusionPlanningSnafu)?
7638            .build()
7639            .context(DataFusionPlanningSnafu)
7640    }
7641
7642    /// Apply an alias to the query result by adding a projection with the alias name
7643    fn apply_alias(&mut self, plan: LogicalPlan, alias_name: String) -> Result<LogicalPlan> {
7644        let fields_expr = self.create_field_column_exprs()?;
7645
7646        // TODO(dennis): how to support multi-value aliasing?
7647        ensure!(
7648            fields_expr.len() == 1,
7649            UnsupportedExprSnafu {
7650                name: "alias on multi-value result"
7651            }
7652        );
7653
7654        let project_fields = fields_expr
7655            .into_iter()
7656            .map(|expr| expr.alias(&alias_name))
7657            .chain(self.create_tag_column_exprs()?)
7658            .chain(Some(self.create_time_index_column_expr()?));
7659
7660        LogicalPlanBuilder::from(plan)
7661            .project(project_fields)
7662            .context(DataFusionPlanningSnafu)?
7663            .build()
7664            .context(DataFusionPlanningSnafu)
7665    }
7666}
7667
7668#[derive(Default, Debug)]
7669struct FunctionArgs {
7670    input: Option<PromExpr>,
7671    literals: Vec<DfExpr>,
7672}
7673
7674/// Represents different types of scalar functions supported in PromQL expressions.
7675/// Each variant defines how the function should be processed and what arguments it expects.
7676#[derive(Debug, Clone)]
7677enum ScalarFunc {
7678    /// DataFusion's registered(including built-in) scalar functions (e.g., abs, sqrt, round, clamp).
7679    /// These are passed through directly to DataFusion's execution engine.
7680    /// Processing: Simple argument insertion at the specified position.
7681    DataFusionBuiltin(Arc<ScalarUdfDef>),
7682    /// User-defined functions registered in DataFusion's function registry.
7683    /// Similar to DataFusionBuiltin but for custom functions not built into DataFusion.
7684    /// Processing: Direct pass-through with argument positioning.
7685    DataFusionUdf(Arc<ScalarUdfDef>),
7686    /// Native histogram helper UDFs. Non-histogram inputs are projected as NULL
7687    /// so the normal PromQL empty-value filter drops them.
7688    NativeHistogramUdf(Arc<ScalarUdfDef>),
7689    /// PromQL-specific functions that operate on time series data with temporal context.
7690    /// These functions require both timestamp ranges and values to perform calculations.
7691    /// Processing: Automatically injects timestamp_range and value columns as first arguments.
7692    /// Examples: idelta, irate, resets, changes, deriv, *_over_time function
7693    Udf(Arc<ScalarUdfDef>),
7694    /// PromQL functions requiring extrapolation calculations with explicit range information.
7695    /// These functions need to know the time range length to perform rate calculations.
7696    /// The second field contains the range length in milliseconds.
7697    /// Processing: Injects timestamp_range, value, time_index columns and appends range_length.
7698    /// Examples: increase, rate, delta
7699    // TODO(ruihang): maybe merge with Udf later
7700    ExtrapolateUdf(Arc<ScalarUdfDef>, i64),
7701    /// Functions that generate expressions directly without external UDF calls.
7702    /// The expression is constructed during function matching and requires no additional processing.
7703    /// Examples: time(), minute(), hour(), month(), year() and other date/time extractors
7704    GeneratedExpr,
7705}
7706
7707#[cfg(test)]
7708mod test {
7709    use std::time::{Duration, UNIX_EPOCH};
7710
7711    use catalog::RegisterTableRequest;
7712    use catalog::memory::{MemoryCatalogManager, new_memory_catalog_manager};
7713    use common_base::Plugins;
7714    use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
7715    use common_query::native_histogram::{
7716        CUSTOM_BUCKETS_SCHEMA, CounterResetHint, NativeHistogram, build_histogram_array,
7717    };
7718    use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value};
7719    use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS;
7720    use common_query::test_util::DummyDecoder;
7721    use common_recordbatch::RecordBatch as GreptimeRecordBatch;
7722    use datafusion::arrow::array::{
7723        Array, ArrayRef, Float64Array, Int64Array, StringArray, TimestampMillisecondArray,
7724    };
7725    use datafusion::arrow::datatypes::{Field, Schema as ArrowSchema};
7726    use datafusion::arrow::record_batch::RecordBatch;
7727    use datafusion::catalog::{CatalogProvider, MemoryCatalogProvider, MemorySchemaProvider};
7728    use datafusion::datasource::memory::MemorySourceConfig;
7729    use datafusion::datasource::source::DataSourceExec;
7730    use datafusion::datasource::{MemTable, provider_as_source};
7731    use datafusion::execution::context::SessionContext;
7732    use datafusion::logical_expr::Extension;
7733    use datatypes::prelude::ConcreteDataType;
7734    use datatypes::schema::{ColumnSchema, Schema};
7735    use promql_parser::label::Labels;
7736    use promql_parser::parser;
7737    use session::context::QueryContext;
7738    use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan};
7739    use table::Table;
7740    use table::metadata::{FilterPushDownType, TableInfoBuilder, TableMetaBuilder};
7741    use table::test_util::{EmptyTable, MemTable as GreptimeMemTable};
7742
7743    use super::*;
7744    use crate::QueryEngineContext;
7745    use crate::options::QueryOptions;
7746    use crate::parser::QueryLanguageParser;
7747    use crate::query_engine::DefaultSerializer;
7748
7749    mod delta;
7750
7751    fn find_instant_manipulate(plan: &LogicalPlan) -> Option<&InstantManipulate> {
7752        if let LogicalPlan::Extension(Extension { node }) = plan
7753            && let Some(instant_manipulate) = node.as_any().downcast_ref::<InstantManipulate>()
7754        {
7755            return Some(instant_manipulate);
7756        }
7757
7758        plan.inputs().into_iter().find_map(find_instant_manipulate)
7759    }
7760
7761    fn build_query_engine_state() -> QueryEngineState {
7762        QueryEngineState::new(
7763            new_memory_catalog_manager().unwrap(),
7764            None,
7765            None,
7766            None,
7767            None,
7768            None,
7769            false,
7770            Plugins::default(),
7771            QueryOptions::default(),
7772        )
7773    }
7774
7775    #[test]
7776    fn common_label_type_preserves_only_shared_dictionary_encoding() {
7777        let dictionary = ArrowDataType::Dictionary(
7778            Box::new(ArrowDataType::UInt32),
7779            Box::new(ArrowDataType::Utf8),
7780        );
7781        let other_dictionary = ArrowDataType::Dictionary(
7782            Box::new(ArrowDataType::Int32),
7783            Box::new(ArrowDataType::Utf8),
7784        );
7785
7786        assert_eq!(
7787            Some(dictionary.clone()),
7788            PromPlanner::common_label_data_type(Some(&dictionary), Some(&dictionary))
7789        );
7790        assert_eq!(
7791            Some(ArrowDataType::Utf8),
7792            PromPlanner::common_label_data_type(Some(&dictionary), Some(&ArrowDataType::Utf8))
7793        );
7794        assert_eq!(
7795            Some(ArrowDataType::Utf8),
7796            PromPlanner::common_label_data_type(Some(&dictionary), Some(&other_dictionary))
7797        );
7798        assert_eq!(
7799            Some(ArrowDataType::Utf8),
7800            PromPlanner::common_label_data_type(Some(&dictionary), None)
7801        );
7802    }
7803
7804    async fn build_optimized_promql_plan(
7805        table_provider: DfTableSourceProvider,
7806        eval_stmt: &EvalStmt,
7807    ) -> LogicalPlan {
7808        let state = build_query_engine_state();
7809        let raw_plan = PromPlanner::stmt_to_plan(table_provider, eval_stmt, &state)
7810            .await
7811            .unwrap();
7812        let context = QueryEngineContext::new(state.session_state(), QueryContext::arc());
7813        state
7814            .optimize_by_extension_rules(raw_plan, &context)
7815            .unwrap()
7816    }
7817
7818    async fn build_optimized_tsid_plan(
7819        query: &str,
7820        num_tag: usize,
7821        num_field: usize,
7822        end_secs: u64,
7823        lookback_secs: u64,
7824    ) -> String {
7825        let eval_stmt = EvalStmt {
7826            expr: parser::parse(query).unwrap(),
7827            start: UNIX_EPOCH,
7828            end: UNIX_EPOCH
7829                .checked_add(Duration::from_secs(end_secs))
7830                .unwrap(),
7831            interval: Duration::from_secs(5),
7832            lookback_delta: Duration::from_secs(lookback_secs),
7833        };
7834        let table_provider = build_test_table_provider_with_tsid(
7835            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
7836            num_tag,
7837            num_field,
7838        )
7839        .await;
7840
7841        build_optimized_promql_plan(table_provider, &eval_stmt)
7842            .await
7843            .display_indent_schema()
7844            .to_string()
7845    }
7846
7847    async fn assert_nested_count_rewrite_applies(query: &str, expected_outer_agg: &str) {
7848        let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await;
7849
7850        assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
7851        assert!(plan_str.contains("Projection: some_metric.timestamp, some_metric.tag_0"));
7852        assert!(plan_str.contains("Distinct:"));
7853        assert!(plan_str.contains(expected_outer_agg), "{plan_str}");
7854        assert!(!plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]"));
7855    }
7856
7857    async fn assert_nested_count_rewrite_missing(query: &str, num_tag: usize, lookback_secs: u64) {
7858        let plan_str = build_optimized_tsid_plan(query, num_tag, 1, 100_000, lookback_secs).await;
7859        assert!(!plan_str.contains("Distinct:"), "{plan_str}");
7860    }
7861
7862    fn build_eval_stmt(expr: &str) -> EvalStmt {
7863        EvalStmt {
7864            expr: parser::parse(expr).unwrap(),
7865            start: UNIX_EPOCH,
7866            end: UNIX_EPOCH
7867                .checked_add(Duration::from_secs(100_000))
7868                .unwrap(),
7869            interval: Duration::from_secs(5),
7870            lookback_delta: Duration::from_secs(1),
7871        }
7872    }
7873
7874    enum DirectOrValue {
7875        Float64(f64),
7876        Int64(i64),
7877        NativeHistogram(NativeHistogram),
7878        Utf8(&'static str),
7879    }
7880
7881    impl DirectOrValue {
7882        fn data_type(&self) -> ArrowDataType {
7883            match self {
7884                Self::Float64(_) => ArrowDataType::Float64,
7885                Self::Int64(_) => ArrowDataType::Int64,
7886                Self::NativeHistogram(_) => native_histogram_value_type().as_arrow_type(),
7887                Self::Utf8(_) => ArrowDataType::Utf8,
7888            }
7889        }
7890        fn array(&self) -> Arc<dyn Array> {
7891            match self {
7892                Self::Float64(v) => Arc::new(Float64Array::from(vec![*v])),
7893                Self::Int64(v) => Arc::new(Int64Array::from(vec![*v])),
7894                Self::NativeHistogram(v) => build_histogram_array(&[Some(v.clone())]),
7895                Self::Utf8(v) => Arc::new(StringArray::from(vec![*v])),
7896            }
7897        }
7898    }
7899
7900    fn direct_or_histogram() -> NativeHistogram {
7901        NativeHistogram {
7902            schema: 0,
7903            zero_threshold: 0.0,
7904            sum: 1.0,
7905            reset_hint: CounterResetHint::Unknown,
7906            start_timestamp: None,
7907            custom_values: vec![],
7908            positive_spans: vec![],
7909            negative_spans: vec![],
7910            count: 1.0,
7911            zero_count: 1.0,
7912            positive_buckets: vec![],
7913            negative_buckets: vec![],
7914        }
7915    }
7916
7917    fn operator_metric_table(
7918        name: &str,
7919        table_id: u32,
7920        tag: &str,
7921        le: Option<&str>,
7922        value: DirectOrValue,
7923    ) -> table::TableRef {
7924        let value_type = match &value {
7925            DirectOrValue::Float64(_) => ConcreteDataType::float64_datatype(),
7926            DirectOrValue::Int64(_) => ConcreteDataType::int64_datatype(),
7927            DirectOrValue::NativeHistogram(_) => native_histogram_value_type().clone(),
7928            DirectOrValue::Utf8(_) => ConcreteDataType::string_datatype(),
7929        };
7930        let tag_count = 1 + usize::from(le.is_some());
7931        let mut columns = vec![ColumnSchema::new(
7932            "tag".to_string(),
7933            ConcreteDataType::string_datatype(),
7934            false,
7935        )];
7936        if le.is_some() {
7937            columns.push(ColumnSchema::new(
7938                LE_COLUMN_NAME.to_string(),
7939                ConcreteDataType::string_datatype(),
7940                false,
7941            ));
7942        }
7943        columns.extend([
7944            ColumnSchema::new(
7945                "ts".to_string(),
7946                ConcreteDataType::timestamp_millisecond_datatype(),
7947                false,
7948            )
7949            .with_time_index(true),
7950            ColumnSchema::new("v".to_string(), value_type, true),
7951        ]);
7952        let schema = Arc::new(Schema::new(columns));
7953        let mut arrays = vec![Arc::new(StringArray::from(vec![tag])) as Arc<dyn Array>];
7954        if let Some(le) = le {
7955            arrays.push(Arc::new(StringArray::from(vec![le])));
7956        }
7957        arrays.extend([
7958            Arc::new(TimestampMillisecondArray::from(vec![1_000])) as Arc<dyn Array>,
7959            value.array(),
7960        ]);
7961        let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap();
7962        let backing = GreptimeMemTable::new_with_catalog(
7963            name,
7964            GreptimeRecordBatch::from_df_record_batch(schema.clone(), batch),
7965            table_id,
7966            DEFAULT_CATALOG_NAME.to_string(),
7967            DEFAULT_SCHEMA_NAME.to_string(),
7968        );
7969        let value_index = tag_count + 1;
7970        let meta = TableMetaBuilder::empty()
7971            .schema(schema)
7972            .primary_key_indices((0..tag_count).collect())
7973            .value_indices(vec![value_index])
7974            .next_column_id((value_index + 1) as u32)
7975            .build()
7976            .unwrap();
7977        let info = Arc::new(
7978            TableInfoBuilder::default()
7979                .table_id(table_id)
7980                .name(name)
7981                .meta(meta)
7982                .build()
7983                .unwrap(),
7984        );
7985        Arc::new(Table::new(
7986            info,
7987            FilterPushDownType::Unsupported,
7988            backing.data_source(),
7989        ))
7990    }
7991
7992    fn operator_table_provider() -> DfTableSourceProvider {
7993        let catalog = MemoryCatalogManager::with_default_setup();
7994        let tables = [
7995            operator_metric_table("lf", 2_001, "a", None, DirectOrValue::Float64(2.0)),
7996            operator_metric_table(
7997                "lh",
7998                2_002,
7999                "b",
8000                None,
8001                DirectOrValue::NativeHistogram(direct_or_histogram()),
8002            ),
8003            operator_metric_table("rf", 2_003, "b", None, DirectOrValue::Float64(3.0)),
8004            operator_metric_table(
8005                "rh",
8006                2_004,
8007                "a",
8008                None,
8009                DirectOrValue::NativeHistogram(direct_or_histogram()),
8010            ),
8011            operator_metric_table("fallback", 2_005, "c", None, DirectOrValue::Float64(7.0)),
8012            operator_metric_table(
8013                "bad_classic",
8014                2_006,
8015                "d",
8016                Some("broken"),
8017                DirectOrValue::Float64(1.0),
8018            ),
8019            operator_metric_table(
8020                "bad_native",
8021                2_007,
8022                "d",
8023                None,
8024                DirectOrValue::NativeHistogram(direct_or_histogram()),
8025            ),
8026        ];
8027        for table in tables {
8028            let info = table.table_info();
8029            catalog
8030                .register_table_sync(RegisterTableRequest {
8031                    catalog: DEFAULT_CATALOG_NAME.to_string(),
8032                    schema: DEFAULT_SCHEMA_NAME.to_string(),
8033                    table_name: info.name.clone(),
8034                    table_id: info.ident.table_id,
8035                    table,
8036                })
8037                .unwrap();
8038        }
8039        DfTableSourceProvider::new(
8040            catalog,
8041            false,
8042            QueryContext::arc(),
8043            DummyDecoder::arc(),
8044            false,
8045        )
8046    }
8047
8048    fn operator_eval_stmt(expr: &str) -> EvalStmt {
8049        let time = UNIX_EPOCH.checked_add(Duration::from_secs(1)).unwrap();
8050        EvalStmt {
8051            expr: parser::parse(expr).unwrap(),
8052            start: time,
8053            end: time,
8054            interval: Duration::from_secs(1),
8055            lookback_delta: Duration::from_secs(5),
8056        }
8057    }
8058
8059    struct DirectOrSource {
8060        name: &'static str,
8061        empty: bool,
8062        timestamp: i64,
8063        tags: Vec<(&'static str, Option<&'static str>)>,
8064        value: DirectOrValue,
8065    }
8066
8067    fn source(
8068        name: &'static str,
8069        empty: bool,
8070        timestamp: i64,
8071        tags: Vec<(&'static str, Option<&'static str>)>,
8072        value: DirectOrValue,
8073    ) -> DirectOrSource {
8074        DirectOrSource {
8075            name,
8076            empty,
8077            timestamp,
8078            tags,
8079            value,
8080        }
8081    }
8082
8083    fn tagged_source(
8084        name: &'static str,
8085        empty: bool,
8086        tag: (&'static str, Option<&'static str>),
8087        value: DirectOrValue,
8088    ) -> DirectOrSource {
8089        source(name, empty, 1, vec![("job", Some("job")), tag], value)
8090    }
8091
8092    fn job_source(name: &'static str, value: DirectOrValue) -> DirectOrSource {
8093        source(name, true, 1, vec![("job", Some("job"))], value)
8094    }
8095
8096    fn table(source: &DirectOrSource) -> Arc<MemTable> {
8097        let mut fields = vec![Field::new(
8098            "ts",
8099            ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
8100            false,
8101        )];
8102        fields.extend(
8103            source
8104                .tags
8105                .iter()
8106                .map(|(name, _)| Field::new(*name, ArrowDataType::Utf8, true)),
8107        );
8108        fields.push(Field::new("v", source.value.data_type(), true));
8109        let schema = Arc::new(ArrowSchema::new(fields));
8110        let partitions = if source.empty {
8111            vec![vec![]]
8112        } else {
8113            let mut columns: Vec<Arc<dyn Array>> =
8114                vec![Arc::new(TimestampMillisecondArray::from(vec![
8115                    source.timestamp,
8116                ]))];
8117            columns.extend(
8118                source
8119                    .tags
8120                    .iter()
8121                    .map(|(_, value)| Arc::new(StringArray::from(vec![*value])) as Arc<dyn Array>),
8122            );
8123            columns.push(source.value.array());
8124            vec![vec![RecordBatch::try_new(schema.clone(), columns).unwrap()]]
8125        };
8126        Arc::new(MemTable::try_new(schema, partitions).unwrap())
8127    }
8128
8129    fn scan(source: &DirectOrSource) -> LogicalPlan {
8130        LogicalPlanBuilder::scan(source.name, provider_as_source(table(source)), None)
8131            .unwrap()
8132            .build()
8133            .unwrap()
8134    }
8135
8136    fn direct_or_context(qualifier: &str, tags: &[&str], field: &str) -> PromPlannerContext {
8137        PromPlannerContext {
8138            table_name: Some(qualifier.to_string()),
8139            time_index_column: Some("ts".to_string()),
8140            field_columns: vec![field.to_string()],
8141            tag_columns: tags.iter().map(|tag| (*tag).to_string()).collect(),
8142            ..Default::default()
8143        }
8144    }
8145
8146    fn or_modifier(expr: &str) -> Option<BinModifier> {
8147        let PromExpr::Binary(expr) = parser::parse(expr).unwrap() else {
8148            unreachable!()
8149        };
8150        expr.modifier
8151    }
8152
8153    async fn plan_direct_or(
8154        left: LogicalPlan,
8155        right: LogicalPlan,
8156        left_context: PromPlannerContext,
8157        right_context: PromPlannerContext,
8158        modifier: &Option<BinModifier>,
8159    ) -> LogicalPlan {
8160        let table_provider = build_test_table_provider_with_fields(
8161            &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
8162            &[],
8163        )
8164        .await;
8165        let mut planner = PromPlanner {
8166            table_provider,
8167            ctx: PromPlannerContext::default(),
8168            promql_annotations: None,
8169        };
8170        planner
8171            .or_operator(
8172                left,
8173                right,
8174                left_context.tag_columns.iter().cloned().collect(),
8175                right_context.tag_columns.iter().cloned().collect(),
8176                left_context,
8177                right_context,
8178                modifier,
8179            )
8180            .unwrap()
8181    }
8182
8183    async fn execute(
8184        plan: LogicalPlan,
8185        state: &QueryEngineState,
8186    ) -> (LogicalPlan, Vec<RecordBatch>) {
8187        let context = QueryEngineContext::new(state.session_state(), QueryContext::arc());
8188        let optimized = state.optimize_by_extension_rules(plan, &context).unwrap();
8189        let physical = state
8190            .session_state()
8191            .create_physical_plan(&optimized)
8192            .await
8193            .unwrap();
8194        let batches =
8195            datafusion::physical_plan::collect(physical, state.session_state().task_ctx())
8196                .await
8197                .unwrap();
8198        (optimized, batches)
8199    }
8200
8201    async fn run(
8202        left: &DirectOrSource,
8203        right: &DirectOrSource,
8204        left_context: PromPlannerContext,
8205        right_context: PromPlannerContext,
8206        modifier: &Option<BinModifier>,
8207    ) -> (LogicalPlan, Vec<RecordBatch>) {
8208        let plan = plan_direct_or(
8209            scan(left),
8210            scan(right),
8211            left_context,
8212            right_context,
8213            modifier,
8214        )
8215        .await;
8216        execute(plan, &build_query_engine_state()).await
8217    }
8218
8219    async fn mixed_direct_or(histogram_on_left: bool) -> (PromPlanner, LogicalPlan) {
8220        let sample = |histogram: bool| {
8221            if histogram {
8222                DirectOrValue::NativeHistogram(direct_or_histogram())
8223            } else {
8224                DirectOrValue::Float64(1.25)
8225            }
8226        };
8227        let left = tagged_source(
8228            "lhs",
8229            false,
8230            (
8231                "k",
8232                Some(if histogram_on_left {
8233                    "histogram"
8234                } else {
8235                    "float"
8236                }),
8237            ),
8238            sample(histogram_on_left),
8239        );
8240        let right = tagged_source(
8241            "rhs",
8242            false,
8243            (
8244                "k",
8245                Some(if histogram_on_left {
8246                    "float"
8247                } else {
8248                    "histogram"
8249                }),
8250            ),
8251            sample(!histogram_on_left),
8252        );
8253        let table_provider = build_test_table_provider_with_fields(
8254            &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
8255            &[],
8256        )
8257        .await;
8258        let mut planner = PromPlanner {
8259            table_provider,
8260            ctx: PromPlannerContext::default(),
8261            promql_annotations: None,
8262        };
8263        let left_context = direct_or_context("lhs", &["job", "k"], "v");
8264        let right_context = direct_or_context("rhs", &["job", "k"], "v");
8265        let plan = planner
8266            .or_operator(
8267                scan(&left),
8268                scan(&right),
8269                left_context.tag_columns.iter().cloned().collect(),
8270                right_context.tag_columns.iter().cloned().collect(),
8271                left_context,
8272                right_context,
8273                &or_modifier("lhs or on(k) rhs"),
8274            )
8275            .unwrap();
8276        (planner, plan)
8277    }
8278
8279    async fn mixed_aggregate_input(histograms: Vec<NativeHistogram>) -> (PromPlanner, LogicalPlan) {
8280        let float_field = format!("{OR_FLOAT_FIELD_PREFIX}0");
8281        let histogram_field = format!("{OR_HISTOGRAM_FIELD_PREFIX}0");
8282        let row_count = histograms.len() + 1;
8283        let schema = Arc::new(ArrowSchema::new(vec![
8284            Field::new(
8285                "ts",
8286                ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
8287                false,
8288            ),
8289            Field::new("k", ArrowDataType::Utf8, false),
8290            Field::new(&float_field, ArrowDataType::Float64, true),
8291            Field::new(
8292                &histogram_field,
8293                native_histogram_value_type().as_arrow_type(),
8294                true,
8295            ),
8296        ]));
8297        let mut histogram_values = Vec::with_capacity(row_count);
8298        histogram_values.push(None);
8299        histogram_values.extend(histograms.into_iter().map(Some));
8300        let batch = RecordBatch::try_new(
8301            schema.clone(),
8302            vec![
8303                Arc::new(TimestampMillisecondArray::from(vec![1; row_count])),
8304                Arc::new(StringArray::from_iter_values(
8305                    (0..row_count).map(|row| format!("kind_{row}")),
8306                )),
8307                Arc::new(Float64Array::from_iter(
8308                    (0..row_count).map(|row| (row == 0).then_some(1.25)),
8309                )),
8310                build_histogram_array(&histogram_values),
8311            ],
8312        )
8313        .unwrap();
8314        let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap());
8315        let plan = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None)
8316            .unwrap()
8317            .build()
8318            .unwrap();
8319        let table_provider = build_test_table_provider_with_fields(
8320            &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
8321            &[],
8322        )
8323        .await;
8324        let planner = PromPlanner {
8325            table_provider,
8326            ctx: PromPlannerContext {
8327                table_name: Some("mixed".to_string()),
8328                time_index_column: Some("ts".to_string()),
8329                field_columns: vec![float_field, histogram_field],
8330                tag_columns: vec!["k".to_string()],
8331                ..Default::default()
8332            },
8333            promql_annotations: None,
8334        };
8335        (planner, plan)
8336    }
8337
8338    fn assert_no_internal_or_keys(schema: &DFSchema) {
8339        assert!(
8340            schema
8341                .fields()
8342                .iter()
8343                .all(|field| !field.name().starts_with("__promql_or_match_")),
8344            "{schema:?}"
8345        );
8346    }
8347
8348    fn values(batches: &[RecordBatch], column: &str) -> Vec<f64> {
8349        batches
8350            .iter()
8351            .flat_map(|batch| {
8352                batch
8353                    .column_by_name(column)
8354                    .unwrap()
8355                    .as_any()
8356                    .downcast_ref::<Float64Array>()
8357                    .unwrap()
8358                    .iter()
8359                    .flatten()
8360            })
8361            .collect()
8362    }
8363
8364    fn numeric_values(batches: &[RecordBatch], column: &str) -> Vec<f64> {
8365        batches
8366            .iter()
8367            .flat_map(|batch| {
8368                let values = datafusion::arrow::compute::cast(
8369                    batch.column_by_name(column).unwrap(),
8370                    &ArrowDataType::Float64,
8371                )
8372                .unwrap();
8373                values
8374                    .as_any()
8375                    .downcast_ref::<Float64Array>()
8376                    .unwrap()
8377                    .iter()
8378                    .flatten()
8379                    .collect::<Vec<_>>()
8380            })
8381            .collect()
8382    }
8383
8384    fn histograms(batches: &[RecordBatch], column: &str) -> Vec<NativeHistogram> {
8385        batches
8386            .iter()
8387            .flat_map(|batch| {
8388                let values = batch
8389                    .column_by_name(column)
8390                    .unwrap()
8391                    .as_any()
8392                    .downcast_ref::<datafusion::arrow::array::StructArray>()
8393                    .unwrap();
8394                (0..values.len()).filter_map(|row| {
8395                    common_query::native_histogram::read_histogram(values, row).unwrap()
8396                })
8397            })
8398            .collect()
8399    }
8400
8401    fn rows(batches: &[RecordBatch]) -> Vec<(f64, Option<String>)> {
8402        let mut rows = batches
8403            .iter()
8404            .flat_map(|batch| {
8405                let values = batch
8406                    .column_by_name("v")
8407                    .unwrap()
8408                    .as_any()
8409                    .downcast_ref::<Float64Array>()
8410                    .unwrap();
8411                let labels = batch
8412                    .column_by_name("k")
8413                    .map(|column| column.as_any().downcast_ref::<StringArray>().unwrap());
8414                (0..batch.num_rows()).map(move |i| {
8415                    (
8416                        values.value(i),
8417                        labels.and_then(|labels| {
8418                            (!labels.is_null(i)).then(|| labels.value(i).to_string())
8419                        }),
8420                    )
8421                })
8422            })
8423            .collect::<Vec<_>>();
8424        rows.sort_by(|left, right| left.0.total_cmp(&right.0));
8425        rows
8426    }
8427
8428    fn matrix_source(
8429        name: &'static str,
8430        k: Option<Option<&'static str>>,
8431        timestamp: i64,
8432        value: f64,
8433    ) -> DirectOrSource {
8434        let mut tags = vec![("job", Some("job"))];
8435        if let Some(k) = k {
8436            tags.push(("k", k));
8437        }
8438        source(name, false, timestamp, tags, DirectOrValue::Float64(value))
8439    }
8440
8441    fn matrix_context(name: &str, k: Option<Option<&str>>) -> PromPlannerContext {
8442        direct_or_context(
8443            name,
8444            if k.is_some() { &["job", "k"] } else { &["job"] },
8445            "v",
8446        )
8447    }
8448
8449    async fn build_missing_le_or_normal_metric_table_provider() -> DfTableSourceProvider {
8450        build_test_table_provider_with_fields(
8451            &[
8452                (
8453                    DEFAULT_SCHEMA_NAME.to_string(),
8454                    "non_existent_histogram_bucket".to_string(),
8455                ),
8456                (DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string()),
8457            ],
8458            &["pod", "instance"],
8459        )
8460        .await
8461    }
8462
8463    fn assert_normal_metric_schema(plan: &LogicalPlan) {
8464        let fields = plan.schema().fields();
8465        assert_eq!(fields.len(), 4, "{fields:?}");
8466        assert!(
8467            fields.iter().any(|field| field.name() == "pod"),
8468            "{fields:?}"
8469        );
8470        assert!(
8471            fields.iter().any(|field| field.name() == "instance"),
8472            "{fields:?}"
8473        );
8474        assert!(
8475            fields
8476                .iter()
8477                .any(|field| field.name() == greptime_timestamp()),
8478            "{fields:?}"
8479        );
8480        assert!(
8481            fields.iter().any(|field| {
8482                field.name() == greptime_value() && field.data_type() == &ArrowDataType::Float64
8483            }),
8484            "{fields:?}"
8485        );
8486    }
8487
8488    async fn build_test_table_provider_with_distinct_tags(
8489        table_tags: &[(&str, &[&str])],
8490    ) -> DfTableSourceProvider {
8491        let catalog_list = MemoryCatalogManager::with_default_setup();
8492        for (table_name, tags) in table_tags {
8493            let mut columns = tags
8494                .iter()
8495                .map(|tag| {
8496                    ColumnSchema::new(
8497                        (*tag).to_string(),
8498                        ConcreteDataType::string_datatype(),
8499                        false,
8500                    )
8501                })
8502                .collect::<Vec<_>>();
8503            columns.push(
8504                ColumnSchema::new(
8505                    greptime_timestamp().to_string(),
8506                    ConcreteDataType::timestamp_millisecond_datatype(),
8507                    false,
8508                )
8509                .with_time_index(true),
8510            );
8511            columns.push(ColumnSchema::new(
8512                greptime_value().to_string(),
8513                ConcreteDataType::float64_datatype(),
8514                true,
8515            ));
8516            let table_meta = TableMetaBuilder::empty()
8517                .schema(Arc::new(Schema::new(columns)))
8518                .primary_key_indices((0..tags.len()).collect())
8519                .next_column_id(1024)
8520                .build()
8521                .unwrap();
8522            let table_info = TableInfoBuilder::default()
8523                .name((*table_name).to_string())
8524                .meta(table_meta)
8525                .build()
8526                .unwrap();
8527
8528            assert!(
8529                catalog_list
8530                    .register_table_sync(RegisterTableRequest {
8531                        catalog: DEFAULT_CATALOG_NAME.to_string(),
8532                        schema: DEFAULT_SCHEMA_NAME.to_string(),
8533                        table_name: (*table_name).to_string(),
8534                        table_id: 1024,
8535                        table: EmptyTable::from_table_info(&table_info),
8536                    })
8537                    .is_ok()
8538            );
8539        }
8540
8541        DfTableSourceProvider::new(
8542            catalog_list,
8543            false,
8544            QueryContext::arc(),
8545            DummyDecoder::arc(),
8546            false,
8547        )
8548    }
8549
8550    fn contains_histogram_fold(plan: &LogicalPlan) -> bool {
8551        matches!(plan, LogicalPlan::Extension(Extension { node }) if node.as_any().is::<HistogramFold>())
8552            || plan.inputs().into_iter().any(contains_histogram_fold)
8553    }
8554
8555    async fn build_set_op_context_table_provider() -> DfTableSourceProvider {
8556        build_test_table_provider_with_distinct_tags(&[
8557            ("bucket_metric", &["job", "le"]),
8558            ("normal_metric", &["job"]),
8559            ("fallback_metric", &["instance"]),
8560        ])
8561        .await
8562    }
8563
8564    async fn build_or_context_table_provider() -> DfTableSourceProvider {
8565        build_test_table_provider_with_distinct_tags(&[
8566            ("normal_metric", &["job"]),
8567            ("other_metric", &["instance"]),
8568            ("non_hist_metric", &["instance"]),
8569        ])
8570        .await
8571    }
8572
8573    async fn optimize_and_create_physical_plan(
8574        state: &QueryEngineState,
8575        plan: LogicalPlan,
8576    ) -> (
8577        LogicalPlan,
8578        Arc<dyn datafusion::physical_plan::ExecutionPlan>,
8579    ) {
8580        let context = QueryEngineContext::new(state.session_state(), QueryContext::arc());
8581        let optimized = state.optimize_by_extension_rules(plan, &context).unwrap();
8582        let physical = state
8583            .session_state()
8584            .create_physical_plan(&optimized)
8585            .await
8586            .unwrap();
8587        (optimized, physical)
8588    }
8589
8590    async fn build_test_table_provider(
8591        table_name_tuples: &[(String, String)],
8592        num_tag: usize,
8593        num_field: usize,
8594    ) -> DfTableSourceProvider {
8595        let catalog_list = MemoryCatalogManager::with_default_setup();
8596        for (schema_name, table_name) in table_name_tuples {
8597            let mut columns = vec![];
8598            for i in 0..num_tag {
8599                columns.push(ColumnSchema::new(
8600                    format!("tag_{i}"),
8601                    ConcreteDataType::string_datatype(),
8602                    false,
8603                ));
8604            }
8605            columns.push(
8606                ColumnSchema::new(
8607                    "timestamp".to_string(),
8608                    ConcreteDataType::timestamp_millisecond_datatype(),
8609                    false,
8610                )
8611                .with_time_index(true),
8612            );
8613            for i in 0..num_field {
8614                columns.push(ColumnSchema::new(
8615                    format!("field_{i}"),
8616                    ConcreteDataType::float64_datatype(),
8617                    true,
8618                ));
8619            }
8620            let schema = Arc::new(Schema::new(columns));
8621            let table_meta = TableMetaBuilder::empty()
8622                .schema(schema)
8623                .primary_key_indices((0..num_tag).collect())
8624                .value_indices((num_tag + 1..num_tag + 1 + num_field).collect())
8625                .next_column_id(1024)
8626                .build()
8627                .unwrap();
8628            let table_info = TableInfoBuilder::default()
8629                .name(table_name.clone())
8630                .meta(table_meta)
8631                .build()
8632                .unwrap();
8633            let table = EmptyTable::from_table_info(&table_info);
8634
8635            assert!(
8636                catalog_list
8637                    .register_table_sync(RegisterTableRequest {
8638                        catalog: DEFAULT_CATALOG_NAME.to_string(),
8639                        schema: schema_name.clone(),
8640                        table_name: table_name.clone(),
8641                        table_id: 1024,
8642                        table,
8643                    })
8644                    .is_ok()
8645            );
8646        }
8647
8648        DfTableSourceProvider::new(
8649            catalog_list,
8650            false,
8651            QueryContext::arc(),
8652            DummyDecoder::arc(),
8653            false,
8654        )
8655    }
8656
8657    async fn build_test_native_histogram_table_provider(table_name: &str) -> DfTableSourceProvider {
8658        build_test_native_histogram_table_provider_with_marker(table_name, false).await
8659    }
8660
8661    async fn build_test_native_histogram_table_provider_with_marker(
8662        table_name: &str,
8663        temporality_marker: bool,
8664    ) -> DfTableSourceProvider {
8665        let catalog_list = MemoryCatalogManager::with_default_setup();
8666        let mut columns = vec![
8667            ColumnSchema::new(
8668                "tag_0".to_string(),
8669                ConcreteDataType::string_datatype(),
8670                false,
8671            ),
8672            ColumnSchema::new(
8673                LE_COLUMN_NAME.to_string(),
8674                ConcreteDataType::string_datatype(),
8675                true,
8676            ),
8677        ];
8678        if temporality_marker {
8679            columns.push(ColumnSchema::new(
8680                OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(),
8681                ConcreteDataType::string_datatype(),
8682                true,
8683            ));
8684        }
8685        let tag_count = columns.len();
8686        columns.extend([
8687            ColumnSchema::new(
8688                "timestamp".to_string(),
8689                ConcreteDataType::timestamp_millisecond_datatype(),
8690                false,
8691            )
8692            .with_time_index(true),
8693            ColumnSchema::new(
8694                greptime_native_histogram().to_string(),
8695                native_histogram_value_type().clone(),
8696                true,
8697            ),
8698        ]);
8699        let schema = Arc::new(Schema::new(columns));
8700        let table_meta = TableMetaBuilder::empty()
8701            .schema(schema)
8702            .primary_key_indices((0..tag_count).collect())
8703            .value_indices(vec![tag_count + 1])
8704            .next_column_id(1024)
8705            .build()
8706            .unwrap();
8707        let table_info = TableInfoBuilder::default()
8708            .name(table_name)
8709            .meta(table_meta)
8710            .build()
8711            .unwrap();
8712        let table = EmptyTable::from_table_info(&table_info);
8713
8714        assert!(
8715            catalog_list
8716                .register_table_sync(RegisterTableRequest {
8717                    catalog: DEFAULT_CATALOG_NAME.to_string(),
8718                    schema: DEFAULT_SCHEMA_NAME.to_string(),
8719                    table_name: table_name.to_string(),
8720                    table_id: 1024,
8721                    table,
8722                })
8723                .is_ok()
8724        );
8725
8726        DfTableSourceProvider::new(
8727            catalog_list,
8728            false,
8729            QueryContext::arc(),
8730            DummyDecoder::arc(),
8731            false,
8732        )
8733    }
8734
8735    async fn build_test_multi_histogram_table_provider(table_name: &str) -> DfTableSourceProvider {
8736        let catalog_list = MemoryCatalogManager::with_default_setup();
8737        let columns = vec![
8738            ColumnSchema::new(
8739                "tag_0".to_string(),
8740                ConcreteDataType::string_datatype(),
8741                false,
8742            ),
8743            ColumnSchema::new(
8744                "timestamp".to_string(),
8745                ConcreteDataType::timestamp_millisecond_datatype(),
8746                false,
8747            )
8748            .with_time_index(true),
8749            ColumnSchema::new(
8750                greptime_native_histogram().to_string(),
8751                native_histogram_value_type().clone(),
8752                true,
8753            ),
8754            ColumnSchema::new(
8755                "native_histogram_2".to_string(),
8756                native_histogram_value_type().clone(),
8757                true,
8758            ),
8759        ];
8760        let schema = Arc::new(Schema::new(columns));
8761        let table_meta = TableMetaBuilder::empty()
8762            .schema(schema)
8763            .primary_key_indices(vec![0])
8764            .value_indices(vec![2, 3])
8765            .next_column_id(1024)
8766            .build()
8767            .unwrap();
8768        let table_info = TableInfoBuilder::default()
8769            .name(table_name)
8770            .meta(table_meta)
8771            .build()
8772            .unwrap();
8773        let table = EmptyTable::from_table_info(&table_info);
8774
8775        assert!(
8776            catalog_list
8777                .register_table_sync(RegisterTableRequest {
8778                    catalog: DEFAULT_CATALOG_NAME.to_string(),
8779                    schema: DEFAULT_SCHEMA_NAME.to_string(),
8780                    table_name: table_name.to_string(),
8781                    table_id: 1024,
8782                    table,
8783                })
8784                .is_ok()
8785        );
8786
8787        DfTableSourceProvider::new(
8788            catalog_list,
8789            false,
8790            QueryContext::arc(),
8791            DummyDecoder::arc(),
8792            false,
8793        )
8794    }
8795
8796    async fn build_test_mixed_native_histogram_table_provider(
8797        table_name: &str,
8798    ) -> DfTableSourceProvider {
8799        build_test_mixed_native_histogram_table_provider_with_marker(table_name, false).await
8800    }
8801
8802    async fn build_test_mixed_native_histogram_table_provider_with_marker(
8803        table_name: &str,
8804        temporality_marker: bool,
8805    ) -> DfTableSourceProvider {
8806        let catalog_list = MemoryCatalogManager::with_default_setup();
8807        let mut columns = vec![ColumnSchema::new(
8808            "tag_0".to_string(),
8809            ConcreteDataType::string_datatype(),
8810            false,
8811        )];
8812        if temporality_marker {
8813            columns.push(ColumnSchema::new(
8814                OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(),
8815                ConcreteDataType::string_datatype(),
8816                true,
8817            ));
8818        }
8819        let tag_count = columns.len();
8820        columns.extend([
8821            ColumnSchema::new(
8822                "timestamp".to_string(),
8823                ConcreteDataType::timestamp_millisecond_datatype(),
8824                false,
8825            )
8826            .with_time_index(true),
8827            ColumnSchema::new(
8828                greptime_native_histogram().to_string(),
8829                native_histogram_value_type().clone(),
8830                true,
8831            ),
8832            ColumnSchema::new(
8833                greptime_value().to_string(),
8834                ConcreteDataType::float64_datatype(),
8835                true,
8836            ),
8837        ]);
8838        let schema = Arc::new(Schema::new(columns));
8839        let table_meta = TableMetaBuilder::empty()
8840            .schema(schema.clone())
8841            .primary_key_indices((0..tag_count).collect())
8842            .value_indices(vec![tag_count + 1, tag_count + 2])
8843            .next_column_id(1024)
8844            .build()
8845            .unwrap();
8846        let table_info = Arc::new(
8847            TableInfoBuilder::default()
8848                .name(table_name)
8849                .meta(table_meta)
8850                .build()
8851                .unwrap(),
8852        );
8853        let mut arrays: Vec<Arc<dyn Array>> =
8854            vec![Arc::new(StringArray::from(vec!["float", "histogram"]))];
8855        if temporality_marker {
8856            arrays.push(Arc::new(StringArray::from(vec![
8857                Some(GREPTIME_TEMPORALITY_DELTA),
8858                Some(GREPTIME_TEMPORALITY_DELTA),
8859            ])));
8860        }
8861        arrays.extend([
8862            Arc::new(TimestampMillisecondArray::from(vec![1_000, 1_000])) as Arc<dyn Array>,
8863            build_histogram_array(&[None, Some(direct_or_histogram())]),
8864            Arc::new(Float64Array::from(vec![Some(2.0), None])),
8865        ]);
8866        let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap();
8867        let backing = GreptimeMemTable::new_with_catalog(
8868            table_name,
8869            GreptimeRecordBatch::from_df_record_batch(schema, batch),
8870            1024,
8871            DEFAULT_CATALOG_NAME.to_string(),
8872            DEFAULT_SCHEMA_NAME.to_string(),
8873        );
8874        let table = Arc::new(Table::new(
8875            table_info,
8876            FilterPushDownType::Unsupported,
8877            backing.data_source(),
8878        ));
8879
8880        assert!(
8881            catalog_list
8882                .register_table_sync(RegisterTableRequest {
8883                    catalog: DEFAULT_CATALOG_NAME.to_string(),
8884                    schema: DEFAULT_SCHEMA_NAME.to_string(),
8885                    table_name: table_name.to_string(),
8886                    table_id: 1024,
8887                    table,
8888                })
8889                .is_ok()
8890        );
8891
8892        DfTableSourceProvider::new(
8893            catalog_list,
8894            false,
8895            QueryContext::arc(),
8896            DummyDecoder::arc(),
8897            false,
8898        )
8899    }
8900
8901    fn classic_and_native_histogram_table_provider(
8902        native_tag: &str,
8903        native_le: Option<&str>,
8904        native_histogram: NativeHistogram,
8905    ) -> DfTableSourceProvider {
8906        let table_name = "mixed_histogram";
8907        let catalog = MemoryCatalogManager::with_default_setup();
8908        let schema = Arc::new(Schema::new(vec![
8909            ColumnSchema::new(
8910                "tag".to_string(),
8911                ConcreteDataType::string_datatype(),
8912                false,
8913            ),
8914            ColumnSchema::new(
8915                LE_COLUMN_NAME.to_string(),
8916                ConcreteDataType::string_datatype(),
8917                true,
8918            ),
8919            ColumnSchema::new(
8920                "timestamp".to_string(),
8921                ConcreteDataType::timestamp_millisecond_datatype(),
8922                false,
8923            )
8924            .with_time_index(true),
8925            ColumnSchema::new(
8926                greptime_native_histogram().to_string(),
8927                native_histogram_value_type().clone(),
8928                true,
8929            ),
8930            ColumnSchema::new(
8931                greptime_value().to_string(),
8932                ConcreteDataType::float64_datatype(),
8933                true,
8934            ),
8935        ]));
8936        let table_meta = TableMetaBuilder::empty()
8937            .schema(schema.clone())
8938            .primary_key_indices(vec![0, 1])
8939            .value_indices(vec![3, 4])
8940            .next_column_id(5)
8941            .build()
8942            .unwrap();
8943        let table_info = Arc::new(
8944            TableInfoBuilder::default()
8945                .name(table_name)
8946                .meta(table_meta)
8947                .build()
8948                .unwrap(),
8949        );
8950        let batch = RecordBatch::try_new(
8951            schema.arrow_schema().clone(),
8952            vec![
8953                Arc::new(StringArray::from(vec![
8954                    "classic", "classic", native_tag, "classic", "classic", native_tag,
8955                ])),
8956                Arc::new(StringArray::from(vec![
8957                    Some("1"),
8958                    Some("+Inf"),
8959                    native_le,
8960                    Some("1"),
8961                    Some("+Inf"),
8962                    native_le,
8963                ])),
8964                Arc::new(TimestampMillisecondArray::from(vec![
8965                    1_000, 1_000, 1_000, 2_000, 2_000, 2_000,
8966                ])),
8967                build_histogram_array(&[
8968                    None,
8969                    None,
8970                    Some(native_histogram.clone()),
8971                    None,
8972                    None,
8973                    Some(native_histogram),
8974                ]),
8975                Arc::new(Float64Array::from(vec![
8976                    Some(2.0),
8977                    Some(4.0),
8978                    None,
8979                    Some(2.0),
8980                    Some(4.0),
8981                    None,
8982                ])),
8983            ],
8984        )
8985        .unwrap();
8986        let backing = GreptimeMemTable::new_with_catalog(
8987            table_name,
8988            GreptimeRecordBatch::from_df_record_batch(schema, batch),
8989            2_200,
8990            DEFAULT_CATALOG_NAME.to_string(),
8991            DEFAULT_SCHEMA_NAME.to_string(),
8992        );
8993        let table = Arc::new(Table::new(
8994            table_info,
8995            FilterPushDownType::Unsupported,
8996            backing.data_source(),
8997        ));
8998        catalog
8999            .register_table_sync(RegisterTableRequest {
9000                catalog: DEFAULT_CATALOG_NAME.to_string(),
9001                schema: DEFAULT_SCHEMA_NAME.to_string(),
9002                table_name: table_name.to_string(),
9003                table_id: 2_200,
9004                table,
9005            })
9006            .unwrap();
9007
9008        DfTableSourceProvider::new(
9009            catalog,
9010            false,
9011            QueryContext::arc(),
9012            DummyDecoder::arc(),
9013            false,
9014        )
9015    }
9016
9017    async fn build_test_table_provider_with_tsid(
9018        table_name_tuples: &[(String, String)],
9019        num_tag: usize,
9020        num_field: usize,
9021    ) -> DfTableSourceProvider {
9022        let table_specs = table_name_tuples
9023            .iter()
9024            .map(|(schema_name, table_name)| ((schema_name.clone(), table_name.clone()), num_field))
9025            .collect::<Vec<_>>();
9026        build_test_table_provider_with_tsid_fields(&table_specs, num_tag).await
9027    }
9028
9029    async fn build_test_table_provider_with_tsid_fields(
9030        table_specs: &[((String, String), usize)],
9031        num_tag: usize,
9032    ) -> DfTableSourceProvider {
9033        let table_specs = table_specs
9034            .iter()
9035            .map(|(table_name_tuple, num_field)| (table_name_tuple.clone(), num_tag, *num_field))
9036            .collect::<Vec<_>>();
9037        build_test_table_provider_with_tsid_tag_fields(&table_specs).await
9038    }
9039
9040    async fn build_test_table_provider_with_tsid_tag_fields(
9041        table_specs: &[((String, String), usize, usize)],
9042    ) -> DfTableSourceProvider {
9043        let catalog_list = MemoryCatalogManager::with_default_setup();
9044
9045        let physical_table_name = "phy";
9046        let physical_table_id = 999u32;
9047        let physical_num_tag = table_specs
9048            .iter()
9049            .map(|(_, num_tag, _)| *num_tag)
9050            .max()
9051            .unwrap_or(0);
9052        let physical_num_field = table_specs
9053            .iter()
9054            .map(|(_, _, num_field)| *num_field)
9055            .max()
9056            .unwrap_or(0);
9057
9058        // Register a metric engine physical table with internal columns.
9059        {
9060            let mut columns = vec![
9061                ColumnSchema::new(
9062                    DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(),
9063                    ConcreteDataType::uint32_datatype(),
9064                    false,
9065                ),
9066                ColumnSchema::new(
9067                    DATA_SCHEMA_TSID_COLUMN_NAME.to_string(),
9068                    ConcreteDataType::uint64_datatype(),
9069                    false,
9070                ),
9071            ];
9072            for i in 0..physical_num_tag {
9073                columns.push(ColumnSchema::new(
9074                    format!("tag_{i}"),
9075                    ConcreteDataType::string_datatype(),
9076                    false,
9077                ));
9078            }
9079            columns.push(
9080                ColumnSchema::new(
9081                    "timestamp".to_string(),
9082                    ConcreteDataType::timestamp_millisecond_datatype(),
9083                    false,
9084                )
9085                .with_time_index(true),
9086            );
9087            for i in 0..physical_num_field {
9088                columns.push(ColumnSchema::new(
9089                    format!("field_{i}"),
9090                    ConcreteDataType::float64_datatype(),
9091                    true,
9092                ));
9093            }
9094
9095            let schema = Arc::new(Schema::new(columns));
9096            let primary_key_indices = (0..(2 + physical_num_tag)).collect::<Vec<_>>();
9097            let table_meta = TableMetaBuilder::empty()
9098                .schema(schema)
9099                .primary_key_indices(primary_key_indices)
9100                .value_indices(
9101                    (2 + physical_num_tag..2 + physical_num_tag + 1 + physical_num_field).collect(),
9102                )
9103                .engine(METRIC_ENGINE_NAME.to_string())
9104                .next_column_id(1024)
9105                .build()
9106                .unwrap();
9107            let table_info = TableInfoBuilder::default()
9108                .table_id(physical_table_id)
9109                .name(physical_table_name)
9110                .meta(table_meta)
9111                .build()
9112                .unwrap();
9113            let table = EmptyTable::from_table_info(&table_info);
9114
9115            assert!(
9116                catalog_list
9117                    .register_table_sync(RegisterTableRequest {
9118                        catalog: DEFAULT_CATALOG_NAME.to_string(),
9119                        schema: DEFAULT_SCHEMA_NAME.to_string(),
9120                        table_name: physical_table_name.to_string(),
9121                        table_id: physical_table_id,
9122                        table,
9123                    })
9124                    .is_ok()
9125            );
9126        }
9127
9128        // Register metric engine logical tables without `__tsid`, referencing the physical table.
9129        for (idx, ((schema_name, table_name), num_tag, num_field)) in table_specs.iter().enumerate()
9130        {
9131            let mut columns = vec![];
9132            for i in 0..*num_tag {
9133                columns.push(ColumnSchema::new(
9134                    format!("tag_{i}"),
9135                    ConcreteDataType::string_datatype(),
9136                    false,
9137                ));
9138            }
9139            columns.push(
9140                ColumnSchema::new(
9141                    "timestamp".to_string(),
9142                    ConcreteDataType::timestamp_millisecond_datatype(),
9143                    false,
9144                )
9145                .with_time_index(true),
9146            );
9147            for i in 0..*num_field {
9148                columns.push(ColumnSchema::new(
9149                    format!("field_{i}"),
9150                    ConcreteDataType::float64_datatype(),
9151                    true,
9152                ));
9153            }
9154
9155            let schema = Arc::new(Schema::new(columns));
9156            let mut options = table::requests::TableOptions::default();
9157            options.extra_options.insert(
9158                LOGICAL_TABLE_METADATA_KEY.to_string(),
9159                physical_table_name.to_string(),
9160            );
9161            let table_id = 1024u32 + idx as u32;
9162            let table_meta = TableMetaBuilder::empty()
9163                .schema(schema)
9164                .primary_key_indices((0..*num_tag).collect())
9165                .value_indices((*num_tag + 1..*num_tag + 1 + *num_field).collect())
9166                .engine(METRIC_ENGINE_NAME.to_string())
9167                .options(options)
9168                .next_column_id(1024)
9169                .build()
9170                .unwrap();
9171            let table_info = TableInfoBuilder::default()
9172                .table_id(table_id)
9173                .name(table_name.clone())
9174                .meta(table_meta)
9175                .build()
9176                .unwrap();
9177            let table = EmptyTable::from_table_info(&table_info);
9178
9179            assert!(
9180                catalog_list
9181                    .register_table_sync(RegisterTableRequest {
9182                        catalog: DEFAULT_CATALOG_NAME.to_string(),
9183                        schema: schema_name.clone(),
9184                        table_name: table_name.clone(),
9185                        table_id,
9186                        table,
9187                    })
9188                    .is_ok()
9189            );
9190        }
9191
9192        DfTableSourceProvider::new(
9193            catalog_list,
9194            false,
9195            QueryContext::arc(),
9196            DummyDecoder::arc(),
9197            false,
9198        )
9199    }
9200
9201    async fn build_test_table_provider_with_fields(
9202        table_name_tuples: &[(String, String)],
9203        tags: &[&str],
9204    ) -> DfTableSourceProvider {
9205        let catalog_list = MemoryCatalogManager::with_default_setup();
9206        for (schema_name, table_name) in table_name_tuples {
9207            let mut columns = vec![];
9208            let num_tag = tags.len();
9209            for tag in tags {
9210                columns.push(ColumnSchema::new(
9211                    tag.to_string(),
9212                    ConcreteDataType::string_datatype(),
9213                    false,
9214                ));
9215            }
9216            columns.push(
9217                ColumnSchema::new(
9218                    greptime_timestamp().to_string(),
9219                    ConcreteDataType::timestamp_millisecond_datatype(),
9220                    false,
9221                )
9222                .with_time_index(true),
9223            );
9224            columns.push(ColumnSchema::new(
9225                greptime_value().to_string(),
9226                ConcreteDataType::float64_datatype(),
9227                true,
9228            ));
9229            let schema = Arc::new(Schema::new(columns));
9230            let table_meta = TableMetaBuilder::empty()
9231                .schema(schema)
9232                .primary_key_indices((0..num_tag).collect())
9233                .next_column_id(1024)
9234                .build()
9235                .unwrap();
9236            let table_info = TableInfoBuilder::default()
9237                .name(table_name.clone())
9238                .meta(table_meta)
9239                .build()
9240                .unwrap();
9241            let table = EmptyTable::from_table_info(&table_info);
9242
9243            assert!(
9244                catalog_list
9245                    .register_table_sync(RegisterTableRequest {
9246                        catalog: DEFAULT_CATALOG_NAME.to_string(),
9247                        schema: schema_name.clone(),
9248                        table_name: table_name.clone(),
9249                        table_id: 1024,
9250                        table,
9251                    })
9252                    .is_ok()
9253            );
9254        }
9255
9256        DfTableSourceProvider::new(
9257            catalog_list,
9258            false,
9259            QueryContext::arc(),
9260            DummyDecoder::arc(),
9261            false,
9262        )
9263    }
9264
9265    // {
9266    //     input: `abs(some_metric{foo!="bar"})`,
9267    //     expected: &Call{
9268    //         Func: MustGetFunction("abs"),
9269    //         Args: Expressions{
9270    //             &VectorSelector{
9271    //                 Name: "some_metric",
9272    //                 LabelMatchers: []*labels.Matcher{
9273    //                     MustLabelMatcher(labels.MatchNotEqual, "foo", "bar"),
9274    //                     MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"),
9275    //                 },
9276    //             },
9277    //         },
9278    //     },
9279    // },
9280    async fn do_single_instant_function_call(fn_name: &'static str, plan_name: &str) {
9281        let prom_expr =
9282            parser::parse(&format!("{fn_name}(some_metric{{tag_0!=\"bar\"}})")).unwrap();
9283        let eval_stmt = EvalStmt {
9284            expr: prom_expr,
9285            start: UNIX_EPOCH,
9286            end: UNIX_EPOCH
9287                .checked_add(Duration::from_secs(100_000))
9288                .unwrap(),
9289            interval: Duration::from_secs(5),
9290            lookback_delta: Duration::from_secs(1),
9291        };
9292
9293        let table_provider = build_test_table_provider(
9294            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9295            1,
9296            1,
9297        )
9298        .await;
9299        let plan =
9300            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9301                .await
9302                .unwrap();
9303
9304        let expected = String::from(
9305            "Filter: TEMPLATE(field_0) IS NOT NULL [timestamp:Timestamp(ms), TEMPLATE(field_0):Float64;N, tag_0:Utf8]\
9306            \n  Projection: some_metric.timestamp, TEMPLATE(some_metric.field_0) AS TEMPLATE(field_0), some_metric.tag_0 [timestamp:Timestamp(ms), TEMPLATE(field_0):Float64;N, tag_0:Utf8]\
9307            \n    PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
9308            \n      PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
9309            \n        Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
9310	            \n          Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
9311            \n            TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"
9312        ).replace("TEMPLATE", plan_name);
9313
9314        assert_eq!(plan.display_indent_schema().to_string(), expected);
9315    }
9316
9317    #[tokio::test]
9318    async fn single_abs() {
9319        do_single_instant_function_call("abs", "abs").await;
9320    }
9321
9322    #[tokio::test]
9323    #[should_panic]
9324    async fn single_absent() {
9325        do_single_instant_function_call("absent", "").await;
9326    }
9327
9328    #[tokio::test]
9329    async fn single_ceil() {
9330        do_single_instant_function_call("ceil", "ceil").await;
9331    }
9332
9333    #[tokio::test]
9334    async fn single_exp() {
9335        do_single_instant_function_call("exp", "exp").await;
9336    }
9337
9338    #[tokio::test]
9339    async fn single_ln() {
9340        do_single_instant_function_call("ln", "ln").await;
9341    }
9342
9343    #[tokio::test]
9344    async fn single_log2() {
9345        do_single_instant_function_call("log2", "log2").await;
9346    }
9347
9348    #[tokio::test]
9349    async fn single_log10() {
9350        do_single_instant_function_call("log10", "log10").await;
9351    }
9352
9353    #[tokio::test]
9354    #[should_panic]
9355    async fn single_scalar() {
9356        do_single_instant_function_call("scalar", "").await;
9357    }
9358
9359    #[tokio::test]
9360    #[should_panic]
9361    async fn single_sgn() {
9362        do_single_instant_function_call("sgn", "").await;
9363    }
9364
9365    #[tokio::test]
9366    #[should_panic]
9367    async fn single_sort() {
9368        do_single_instant_function_call("sort", "").await;
9369    }
9370
9371    #[tokio::test]
9372    #[should_panic]
9373    async fn single_sort_desc() {
9374        do_single_instant_function_call("sort_desc", "").await;
9375    }
9376
9377    #[tokio::test]
9378    async fn single_sqrt() {
9379        do_single_instant_function_call("sqrt", "sqrt").await;
9380    }
9381
9382    #[tokio::test]
9383    async fn single_timestamp_plan_preserves_source_value() {
9384        let eval_stmt = build_eval_stmt(r#"timestamp(some_metric{tag_0!="bar"})"#);
9385        let table_provider = build_test_table_provider(
9386            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9387            1,
9388            1,
9389        )
9390        .await;
9391        let plan =
9392            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9393                .await
9394                .unwrap();
9395
9396        let expected = String::from(
9397            "Filter: value IS NOT NULL [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\
9398            \n  Projection: some_metric.timestamp, value AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\
9399            \n    Projection: some_metric.timestamp, __promql_timestamp_value_ AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\
9400            \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\
9401            \n        Projection: some_metric.tag_0, some_metric.timestamp, some_metric.field_0, CAST(CAST(CAST(CAST(some_metric.timestamp AS Int64) AS Decimal128(19, 0)) * Decimal128(1,1,0) + Decimal128(0,19,0) AS Int64) AS Float64) / Float64(1000) AS __promql_timestamp_value_ [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\
9402            \n          PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
9403            \n            Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
9404            \n              Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
9405            \n                TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
9406        );
9407
9408        assert_eq!(plan.display_indent_schema().to_string(), expected);
9409    }
9410
9411    #[tokio::test]
9412    async fn single_acos() {
9413        do_single_instant_function_call("acos", "acos").await;
9414    }
9415
9416    #[tokio::test]
9417    #[should_panic]
9418    async fn single_acosh() {
9419        do_single_instant_function_call("acosh", "").await;
9420    }
9421
9422    #[tokio::test]
9423    async fn single_asin() {
9424        do_single_instant_function_call("asin", "asin").await;
9425    }
9426
9427    #[tokio::test]
9428    #[should_panic]
9429    async fn single_asinh() {
9430        do_single_instant_function_call("asinh", "").await;
9431    }
9432
9433    #[tokio::test]
9434    async fn single_atan() {
9435        do_single_instant_function_call("atan", "atan").await;
9436    }
9437
9438    #[tokio::test]
9439    #[should_panic]
9440    async fn single_atanh() {
9441        do_single_instant_function_call("atanh", "").await;
9442    }
9443
9444    #[tokio::test]
9445    async fn single_cos() {
9446        do_single_instant_function_call("cos", "cos").await;
9447    }
9448
9449    #[tokio::test]
9450    #[should_panic]
9451    async fn single_cosh() {
9452        do_single_instant_function_call("cosh", "").await;
9453    }
9454
9455    #[tokio::test]
9456    async fn single_sin() {
9457        do_single_instant_function_call("sin", "sin").await;
9458    }
9459
9460    #[tokio::test]
9461    #[should_panic]
9462    async fn single_sinh() {
9463        do_single_instant_function_call("sinh", "").await;
9464    }
9465
9466    #[tokio::test]
9467    async fn single_tan() {
9468        do_single_instant_function_call("tan", "tan").await;
9469    }
9470
9471    #[tokio::test]
9472    #[should_panic]
9473    async fn single_tanh() {
9474        do_single_instant_function_call("tanh", "").await;
9475    }
9476
9477    #[tokio::test]
9478    #[should_panic]
9479    async fn single_deg() {
9480        do_single_instant_function_call("deg", "").await;
9481    }
9482
9483    #[tokio::test]
9484    #[should_panic]
9485    async fn single_rad() {
9486        do_single_instant_function_call("rad", "").await;
9487    }
9488
9489    // {
9490    //     input: "avg by (foo)(some_metric)",
9491    //     expected: &AggregateExpr{
9492    //         Op: AVG,
9493    //         Expr: &VectorSelector{
9494    //             Name: "some_metric",
9495    //             LabelMatchers: []*labels.Matcher{
9496    //                 MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"),
9497    //             },
9498    //             PosRange: PositionRange{
9499    //                 Start: 13,
9500    //                 End:   24,
9501    //             },
9502    //         },
9503    //         Grouping: []string{"foo"},
9504    //         PosRange: PositionRange{
9505    //             Start: 0,
9506    //             End:   25,
9507    //         },
9508    //     },
9509    // },
9510    async fn do_aggregate_expr_plan(fn_name: &str, plan_name: &str) {
9511        let prom_expr = parser::parse(&format!(
9512            "{fn_name} by (tag_1)(some_metric{{tag_0!=\"bar\"}})",
9513        ))
9514        .unwrap();
9515        let mut eval_stmt = EvalStmt {
9516            expr: prom_expr,
9517            start: UNIX_EPOCH,
9518            end: UNIX_EPOCH
9519                .checked_add(Duration::from_secs(100_000))
9520                .unwrap(),
9521            interval: Duration::from_secs(5),
9522            lookback_delta: Duration::from_secs(1),
9523        };
9524
9525        // test group by
9526        let table_provider = build_test_table_provider(
9527            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9528            2,
9529            2,
9530        )
9531        .await;
9532        let plan =
9533            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9534                .await
9535                .unwrap();
9536        let expected_no_without = String::from(
9537            "Sort: some_metric.tag_1 ASC NULLS LAST, some_metric.timestamp ASC NULLS LAST [tag_1:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\
9538            \n  Aggregate: groupBy=[[some_metric.tag_1, some_metric.timestamp]], aggr=[[TEMPLATE(some_metric.field_0), TEMPLATE(some_metric.field_1)]] [tag_1:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\
9539            \n    PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
9540            \n      PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
9541            \n        Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.tag_1 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
9542            \n          Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
9543            \n            TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]"
9544        ).replace("TEMPLATE", plan_name);
9545        assert_eq!(
9546            plan.display_indent_schema().to_string(),
9547            expected_no_without
9548        );
9549
9550        // test group without
9551        if let PromExpr::Aggregate(AggregateExpr { modifier, .. }) = &mut eval_stmt.expr {
9552            *modifier = Some(LabelModifier::Exclude(Labels {
9553                labels: vec![String::from("tag_1")].into_iter().collect(),
9554            }));
9555        }
9556        let table_provider = build_test_table_provider(
9557            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9558            2,
9559            2,
9560        )
9561        .await;
9562        let plan =
9563            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9564                .await
9565                .unwrap();
9566        let expected_without = String::from(
9567            "Sort: some_metric.tag_0 ASC NULLS LAST, some_metric.timestamp ASC NULLS LAST [tag_0:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\
9568            \n  Aggregate: groupBy=[[some_metric.tag_0, some_metric.timestamp]], aggr=[[TEMPLATE(some_metric.field_0), TEMPLATE(some_metric.field_1)]] [tag_0:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\
9569            \n    PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
9570            \n      PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
9571            \n        Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.tag_1 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
9572            \n          Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
9573            \n            TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]"
9574        ).replace("TEMPLATE", plan_name);
9575        assert_eq!(plan.display_indent_schema().to_string(), expected_without);
9576    }
9577
9578    #[tokio::test]
9579    async fn aggregate_sum() {
9580        do_aggregate_expr_plan("sum", "sum").await;
9581    }
9582
9583    #[tokio::test]
9584    async fn tsid_is_used_for_series_divide_when_available() {
9585        let prom_expr = parser::parse("some_metric").unwrap();
9586        let eval_stmt = EvalStmt {
9587            expr: prom_expr,
9588            start: UNIX_EPOCH,
9589            end: UNIX_EPOCH
9590                .checked_add(Duration::from_secs(100_000))
9591                .unwrap(),
9592            interval: Duration::from_secs(5),
9593            lookback_delta: Duration::from_secs(1),
9594        };
9595
9596        let table_provider = build_test_table_provider_with_tsid(
9597            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9598            1,
9599            1,
9600        )
9601        .await;
9602        let plan =
9603            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9604                .await
9605                .unwrap();
9606
9607        let plan_str = plan.display_indent_schema().to_string();
9608        assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
9609        assert!(plan_str.contains("__tsid ASC NULLS FIRST"));
9610        assert!(
9611            !plan
9612                .schema()
9613                .fields()
9614                .iter()
9615                .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
9616        );
9617
9618        let manipulate = find_instant_manipulate(&plan).unwrap();
9619        let exec = manipulate.to_execution_plan(Arc::new(DataSourceExec::new(Arc::new(
9620            MemorySourceConfig::try_new(
9621                &[],
9622                Arc::new(
9623                    datafusion_expr::UserDefinedLogicalNodeCore::inputs(manipulate)[0]
9624                        .schema()
9625                        .as_arrow()
9626                        .clone(),
9627                ),
9628                None,
9629            )
9630            .unwrap(),
9631        ))));
9632        assert!(format!("{exec:?}").contains("reuse_tsid_column: true"));
9633    }
9634
9635    #[tokio::test]
9636    async fn default_binary_join_uses_tsid_when_available() {
9637        let eval_stmt = build_eval_stmt("some_metric / some_alt_metric");
9638
9639        let table_provider = build_test_table_provider_with_tsid(
9640            &[
9641                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9642                (
9643                    DEFAULT_SCHEMA_NAME.to_string(),
9644                    "some_alt_metric".to_string(),
9645                ),
9646            ],
9647            1,
9648            1,
9649        )
9650        .await;
9651        let plan =
9652            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9653                .await
9654                .unwrap();
9655
9656        let plan_str = plan.display_indent_schema().to_string();
9657        assert!(
9658            plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
9659            "{plan_str}"
9660        );
9661        assert!(
9662            !plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
9663            "{plan_str}"
9664        );
9665    }
9666
9667    #[tokio::test]
9668    async fn reject_binary_fill_modifiers() {
9669        let state = build_query_engine_state();
9670
9671        for query in [
9672            "some_metric + fill(0) some_alt_metric",
9673            "some_metric + fill_left(0) some_alt_metric",
9674            "some_metric + fill_right(0) some_alt_metric",
9675            "(some_metric + fill(0) some_alt_metric) + some_metric",
9676        ] {
9677            let eval_stmt = build_eval_stmt(query);
9678            let table_provider = build_test_table_provider(&[], 0, 0).await;
9679            let err = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &state)
9680                .await
9681                .unwrap_err();
9682
9683            assert!(
9684                matches!(
9685                    &err,
9686                    crate::promql::error::Error::UnsupportedExpr { name, .. }
9687                        if name == "PromQL fill modifiers"
9688                ),
9689                "{err}"
9690            );
9691        }
9692    }
9693
9694    #[tokio::test]
9695    async fn timestamp_binary_join_falls_back_when_tsid_is_projected_out() {
9696        for query in [
9697            "timestamp(some_metric) / some_metric",
9698            "some_metric / timestamp(some_metric)",
9699        ] {
9700            let eval_stmt = build_eval_stmt(query);
9701
9702            let table_provider = build_test_table_provider_with_tsid(
9703                &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9704                1,
9705                1,
9706            )
9707            .await;
9708            let plan =
9709                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9710                    .await
9711                    .unwrap();
9712
9713            let plan_str = plan.display_indent_schema().to_string();
9714            assert!(!plan_str.contains("__tsid ="), "{query}: {plan_str}");
9715            assert!(
9716                plan_str.contains("lhs.tag_0 = rhs.tag_0"),
9717                "{query}: {plan_str}"
9718            );
9719            assert!(
9720                !plan
9721                    .schema()
9722                    .fields()
9723                    .iter()
9724                    .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME),
9725                "{query}: {plan_str}"
9726            );
9727        }
9728    }
9729
9730    #[tokio::test]
9731    async fn timestamp_binary_join_rejects_default_matching_on_mismatched_labels() {
9732        let eval_stmt = build_eval_stmt("timestamp(left_host_job) / right_by_job");
9733
9734        let table_provider = build_test_table_provider_with_tsid_tag_fields(&[
9735            (
9736                (DEFAULT_SCHEMA_NAME.to_string(), "left_host_job".to_string()),
9737                2,
9738                1,
9739            ),
9740            (
9741                (DEFAULT_SCHEMA_NAME.to_string(), "right_by_job".to_string()),
9742                1,
9743                1,
9744            ),
9745        ])
9746        .await;
9747        let plan =
9748            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9749                .await
9750                .unwrap();
9751        let plan_str = plan.display_indent_schema().to_string();
9752
9753        assert!(
9754            plan_str.contains("Boolean(false)") || plan_str.contains("false"),
9755            "{plan_str}"
9756        );
9757    }
9758
9759    #[tokio::test]
9760    async fn tsid_is_preserved_for_nested_default_binary_joins() {
9761        let eval_stmt = build_eval_stmt("(some_metric - some_alt_metric) / some_third_metric");
9762
9763        let table_provider = build_test_table_provider_with_tsid(
9764            &[
9765                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9766                (
9767                    DEFAULT_SCHEMA_NAME.to_string(),
9768                    "some_alt_metric".to_string(),
9769                ),
9770                (
9771                    DEFAULT_SCHEMA_NAME.to_string(),
9772                    "some_third_metric".to_string(),
9773                ),
9774            ],
9775            1,
9776            1,
9777        )
9778        .await;
9779        let plan =
9780            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9781                .await
9782                .unwrap();
9783
9784        let plan_str = plan.display_indent_schema().to_string();
9785        assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
9786        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9787    }
9788
9789    #[tokio::test]
9790    async fn repeated_tsid_binary_operand_reuses_leaf_plan() {
9791        let eval_stmt = build_eval_stmt("((some_metric - some_alt_metric) / some_metric) * 100");
9792
9793        let table_provider = build_test_table_provider_with_tsid(
9794            &[
9795                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9796                (
9797                    DEFAULT_SCHEMA_NAME.to_string(),
9798                    "some_alt_metric".to_string(),
9799                ),
9800            ],
9801            1,
9802            1,
9803        )
9804        .await;
9805        let plan =
9806            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9807                .await
9808                .unwrap();
9809
9810        let plan_str = plan.display_indent_schema().to_string();
9811        assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}");
9812        assert_eq!(
9813            plan_str
9814                .matches("Filter: phy.__table_id = UInt32(1024)")
9815                .count(),
9816            1,
9817            "{plan_str}"
9818        );
9819        assert_eq!(
9820            plan_str.matches("PromInstantManipulate").count(),
9821            2,
9822            "{plan_str}"
9823        );
9824        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9825    }
9826
9827    #[tokio::test]
9828    async fn repeated_tsid_binary_operand_reuses_shorter_field_side() {
9829        let eval_stmt =
9830            build_eval_stmt("((two_field_metric - one_field_metric) / one_field_metric) * 100");
9831
9832        let table_provider = build_test_table_provider_with_tsid_fields(
9833            &[
9834                (
9835                    (
9836                        DEFAULT_SCHEMA_NAME.to_string(),
9837                        "two_field_metric".to_string(),
9838                    ),
9839                    2,
9840                ),
9841                (
9842                    (
9843                        DEFAULT_SCHEMA_NAME.to_string(),
9844                        "one_field_metric".to_string(),
9845                    ),
9846                    1,
9847                ),
9848            ],
9849            1,
9850        )
9851        .await;
9852        let plan =
9853            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9854                .await
9855                .unwrap();
9856
9857        let field_names = plan
9858            .schema()
9859            .fields()
9860            .iter()
9861            .map(|field| field.name().clone())
9862            .collect::<Vec<_>>();
9863        let value_columns = field_names
9864            .iter()
9865            .filter(|name| {
9866                *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME
9867            })
9868            .count();
9869        assert_eq!(value_columns, 1, "{field_names:?}");
9870        let plan_str = plan.display_indent_schema().to_string();
9871        assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}");
9872        assert_eq!(
9873            plan_str
9874                .matches("Filter: phy.__table_id = UInt32(1025)")
9875                .count(),
9876            1,
9877            "{plan_str}"
9878        );
9879        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9880    }
9881
9882    #[tokio::test]
9883    async fn binary_island_reuses_self_operand_without_join() {
9884        let eval_stmt = build_eval_stmt("some_metric / some_metric");
9885
9886        let table_provider = build_test_table_provider_with_tsid(
9887            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9888            1,
9889            1,
9890        )
9891        .await;
9892        let plan =
9893            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9894                .await
9895                .unwrap();
9896
9897        let plan_str = plan.display_indent_schema().to_string();
9898        assert_eq!(plan_str.matches("__tsid =").count(), 0, "{plan_str}");
9899        assert_eq!(
9900            plan_str
9901                .matches("Filter: phy.__table_id = UInt32(1024)")
9902                .count(),
9903            1,
9904            "{plan_str}"
9905        );
9906        assert_eq!(
9907            plan_str.matches("PromInstantManipulate").count(),
9908            1,
9909            "{plan_str}"
9910        );
9911    }
9912
9913    #[tokio::test]
9914    async fn binary_island_reuses_leaf_across_two_branches() {
9915        let eval_stmt =
9916            build_eval_stmt("(some_metric + some_alt_metric) / (some_metric + third_metric)");
9917
9918        let table_provider = build_test_table_provider_with_tsid(
9919            &[
9920                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9921                (
9922                    DEFAULT_SCHEMA_NAME.to_string(),
9923                    "some_alt_metric".to_string(),
9924                ),
9925                (DEFAULT_SCHEMA_NAME.to_string(), "third_metric".to_string()),
9926            ],
9927            1,
9928            1,
9929        )
9930        .await;
9931        let plan =
9932            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9933                .await
9934                .unwrap();
9935
9936        let plan_str = plan.display_indent_schema().to_string();
9937        assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
9938        assert_eq!(
9939            plan_str
9940                .matches("Filter: phy.__table_id = UInt32(1024)")
9941                .count(),
9942            1,
9943            "{plan_str}"
9944        );
9945        assert_eq!(
9946            plan_str.matches("PromInstantManipulate").count(),
9947            3,
9948            "{plan_str}"
9949        );
9950    }
9951
9952    #[tokio::test]
9953    async fn binary_island_generated_alias_avoids_user_column_names() {
9954        let eval_stmt = build_eval_stmt("(some_metric + some_alt_metric) / some_metric");
9955
9956        let table_provider = build_test_table_provider_with_fields(
9957            &[
9958                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9959                (
9960                    DEFAULT_SCHEMA_NAME.to_string(),
9961                    "some_alt_metric".to_string(),
9962                ),
9963            ],
9964            &["prom_v0", "__prom_v0"],
9965        )
9966        .await;
9967        let plan =
9968            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9969                .await
9970                .unwrap();
9971
9972        let field_names = plan.schema().field_names();
9973        assert!(field_names.iter().any(|name| name.ends_with(".prom_v0")));
9974        assert!(field_names.iter().any(|name| name.ends_with(".__prom_v0")));
9975
9976        let plan_str = plan.display_indent_schema().to_string();
9977        assert!(plan_str.contains("SubqueryAlias: __prom_v0"), "{plan_str}");
9978        assert_eq!(
9979            plan_str.matches("PromInstantManipulate").count(),
9980            2,
9981            "{plan_str}"
9982        );
9983    }
9984
9985    #[tokio::test]
9986    async fn binary_island_clears_qualifier_for_nested_unary_projection() {
9987        let eval_stmt = build_eval_stmt("-((some_metric + some_alt_metric) / some_metric)");
9988
9989        let table_provider = build_test_table_provider_with_tsid(
9990            &[
9991                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9992                (
9993                    DEFAULT_SCHEMA_NAME.to_string(),
9994                    "some_alt_metric".to_string(),
9995                ),
9996            ],
9997            1,
9998            1,
9999        )
10000        .await;
10001        let plan =
10002            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10003                .await
10004                .unwrap();
10005
10006        let plan_str = plan.display_indent_schema().to_string();
10007        assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}");
10008        assert_eq!(
10009            plan_str.matches("PromInstantManipulate").count(),
10010            2,
10011            "{plan_str}"
10012        );
10013    }
10014
10015    #[tokio::test]
10016    async fn binary_island_keeps_distinct_matcher_leaves() {
10017        let eval_stmt = build_eval_stmt(
10018            "(some_metric{tag_0=\"foo\"} + some_alt_metric) / some_metric{tag_0=\"bar\"}",
10019        );
10020
10021        let table_provider = build_test_table_provider_with_tsid(
10022            &[
10023                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10024                (
10025                    DEFAULT_SCHEMA_NAME.to_string(),
10026                    "some_alt_metric".to_string(),
10027                ),
10028            ],
10029            1,
10030            1,
10031        )
10032        .await;
10033        let plan =
10034            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10035                .await
10036                .unwrap();
10037
10038        let plan_str = plan.display_indent_schema().to_string();
10039        assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
10040        assert_eq!(
10041            plan_str.matches("PromInstantManipulate").count(),
10042            3,
10043            "{plan_str}"
10044        );
10045    }
10046
10047    #[tokio::test]
10048    async fn binary_island_keeps_offset_leaves_distinct() {
10049        let eval_stmt = build_eval_stmt("(some_metric offset 5m + some_alt_metric) / some_metric");
10050
10051        let table_provider = build_test_table_provider_with_tsid(
10052            &[
10053                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10054                (
10055                    DEFAULT_SCHEMA_NAME.to_string(),
10056                    "some_alt_metric".to_string(),
10057                ),
10058            ],
10059            1,
10060            1,
10061        )
10062        .await;
10063        let plan =
10064            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10065                .await
10066                .unwrap();
10067
10068        let plan_str = plan.display_indent_schema().to_string();
10069        assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
10070        assert_eq!(
10071            plan_str.matches("PromInstantManipulate").count(),
10072            3,
10073            "{plan_str}"
10074        );
10075    }
10076
10077    #[tokio::test]
10078    async fn binary_island_falls_back_for_group_modifier() {
10079        let eval_stmt = build_eval_stmt(
10080            "(some_metric + ignoring(tag_0) group_left some_alt_metric) / some_metric",
10081        );
10082
10083        let table_provider = build_test_table_provider_with_tsid(
10084            &[
10085                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10086                (
10087                    DEFAULT_SCHEMA_NAME.to_string(),
10088                    "some_alt_metric".to_string(),
10089                ),
10090            ],
10091            1,
10092            1,
10093        )
10094        .await;
10095        let plan =
10096            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10097                .await
10098                .unwrap();
10099
10100        let plan_str = plan.display_indent_schema().to_string();
10101        assert_eq!(
10102            plan_str.matches("PromInstantManipulate").count(),
10103            3,
10104            "{plan_str}"
10105        );
10106    }
10107
10108    #[tokio::test]
10109    async fn binary_island_falls_back_for_comparison_filter() {
10110        let eval_stmt = build_eval_stmt("(some_metric > some_alt_metric) / some_metric");
10111
10112        let table_provider = build_test_table_provider_with_tsid(
10113            &[
10114                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10115                (
10116                    DEFAULT_SCHEMA_NAME.to_string(),
10117                    "some_alt_metric".to_string(),
10118                ),
10119            ],
10120            1,
10121            1,
10122        )
10123        .await;
10124        let plan =
10125            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10126                .await
10127                .unwrap();
10128
10129        let plan_str = plan.display_indent_schema().to_string();
10130        assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
10131        assert_eq!(
10132            plan_str.matches("PromInstantManipulate").count(),
10133            3,
10134            "{plan_str}"
10135        );
10136    }
10137
10138    #[tokio::test]
10139    async fn tsid_binary_join_uses_shorter_field_side() {
10140        let eval_stmt = build_eval_stmt("one_field_metric / two_field_metric");
10141
10142        let table_provider = build_test_table_provider_with_tsid_fields(
10143            &[
10144                (
10145                    (
10146                        DEFAULT_SCHEMA_NAME.to_string(),
10147                        "one_field_metric".to_string(),
10148                    ),
10149                    1,
10150                ),
10151                (
10152                    (
10153                        DEFAULT_SCHEMA_NAME.to_string(),
10154                        "two_field_metric".to_string(),
10155                    ),
10156                    2,
10157                ),
10158            ],
10159            1,
10160        )
10161        .await;
10162        let plan =
10163            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10164                .await
10165                .unwrap();
10166
10167        let field_names = plan
10168            .schema()
10169            .fields()
10170            .iter()
10171            .map(|field| field.name().clone())
10172            .collect::<Vec<_>>();
10173        let value_columns = field_names
10174            .iter()
10175            .filter(|name| {
10176                *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME
10177            })
10178            .count();
10179        assert_eq!(value_columns, 1, "{field_names:?}");
10180    }
10181
10182    #[tokio::test]
10183    async fn comparison_binary_join_uses_shorter_field_side() {
10184        let eval_stmt = build_eval_stmt("two_field_metric > one_field_metric");
10185
10186        let table_provider = build_test_table_provider_with_tsid_fields(
10187            &[
10188                (
10189                    (
10190                        DEFAULT_SCHEMA_NAME.to_string(),
10191                        "two_field_metric".to_string(),
10192                    ),
10193                    2,
10194                ),
10195                (
10196                    (
10197                        DEFAULT_SCHEMA_NAME.to_string(),
10198                        "one_field_metric".to_string(),
10199                    ),
10200                    1,
10201                ),
10202            ],
10203            1,
10204        )
10205        .await;
10206        let plan =
10207            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10208                .await
10209                .unwrap();
10210
10211        let field_names = plan
10212            .schema()
10213            .fields()
10214            .iter()
10215            .map(|field| field.name().clone())
10216            .collect::<Vec<_>>();
10217        assert!(
10218            field_names.iter().any(|name| name == "field_0"),
10219            "{field_names:?}"
10220        );
10221        assert!(
10222            !field_names.iter().any(|name| name == "field_1"),
10223            "{field_names:?}"
10224        );
10225    }
10226
10227    #[tokio::test]
10228    async fn label_matching_modifier_disables_tsid_binary_join() {
10229        let eval_stmt = build_eval_stmt("some_metric / ignoring(tag_0) some_alt_metric");
10230
10231        let table_provider = build_test_table_provider_with_tsid(
10232            &[
10233                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10234                (
10235                    DEFAULT_SCHEMA_NAME.to_string(),
10236                    "some_alt_metric".to_string(),
10237                ),
10238            ],
10239            2,
10240            1,
10241        )
10242        .await;
10243        let plan =
10244            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10245                .await
10246                .unwrap();
10247
10248        let plan_str = plan.display_indent_schema().to_string();
10249        assert!(!plan_str.contains("__tsid ="), "{plan_str}");
10250        assert!(
10251            plan_str.contains("some_metric.tag_1 = some_alt_metric.tag_1"),
10252            "{plan_str}"
10253        );
10254    }
10255
10256    #[tokio::test]
10257    async fn ignoring_absent_label_keeps_tsid_binary_join() {
10258        let eval_stmt = build_eval_stmt("some_metric / ignoring(missing) some_alt_metric");
10259
10260        let table_provider = build_test_table_provider_with_tsid(
10261            &[
10262                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10263                (
10264                    DEFAULT_SCHEMA_NAME.to_string(),
10265                    "some_alt_metric".to_string(),
10266                ),
10267            ],
10268            2,
10269            1,
10270        )
10271        .await;
10272        let plan =
10273            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10274                .await
10275                .unwrap();
10276
10277        let plan_str = plan.display_indent_schema().to_string();
10278        assert!(
10279            plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
10280            "{plan_str}"
10281        );
10282        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
10283        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
10284    }
10285
10286    #[tokio::test]
10287    async fn range_function_keeps_tsid_for_absent_ignoring_binary_join() {
10288        let eval_stmt =
10289            build_eval_stmt("rate(some_metric[5m]) / ignoring(missing) some_alt_metric");
10290
10291        let table_provider = build_test_table_provider_with_tsid(
10292            &[
10293                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10294                (
10295                    DEFAULT_SCHEMA_NAME.to_string(),
10296                    "some_alt_metric".to_string(),
10297                ),
10298            ],
10299            2,
10300            1,
10301        )
10302        .await;
10303        let plan =
10304            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10305                .await
10306                .unwrap();
10307
10308        let plan_str = plan.display_indent_schema().to_string();
10309        assert!(
10310            plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
10311            "{plan_str}"
10312        );
10313        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
10314        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
10315    }
10316
10317    #[tokio::test]
10318    async fn on_full_label_set_keeps_tsid_binary_join() {
10319        let eval_stmt = build_eval_stmt("some_metric / on(tag_0, tag_1) some_alt_metric");
10320
10321        let table_provider = build_test_table_provider_with_tsid(
10322            &[
10323                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10324                (
10325                    DEFAULT_SCHEMA_NAME.to_string(),
10326                    "some_alt_metric".to_string(),
10327                ),
10328            ],
10329            2,
10330            1,
10331        )
10332        .await;
10333        let plan =
10334            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10335                .await
10336                .unwrap();
10337
10338        let plan_str = plan.display_indent_schema().to_string();
10339        assert!(
10340            plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
10341            "{plan_str}"
10342        );
10343        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
10344        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
10345    }
10346
10347    #[tokio::test]
10348    async fn on_partial_label_set_disables_tsid_binary_join() {
10349        let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric");
10350
10351        let table_provider = build_test_table_provider_with_tsid(
10352            &[
10353                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10354                (
10355                    DEFAULT_SCHEMA_NAME.to_string(),
10356                    "some_alt_metric".to_string(),
10357                ),
10358            ],
10359            2,
10360            1,
10361        )
10362        .await;
10363        let plan =
10364            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10365                .await
10366                .unwrap();
10367
10368        let plan_str = plan.display_indent_schema().to_string();
10369        assert!(!plan_str.contains("__tsid ="), "{plan_str}");
10370        assert!(
10371            plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
10372            "{plan_str}"
10373        );
10374        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
10375    }
10376
10377    #[tokio::test]
10378    async fn on_label_set_must_cover_both_sides_to_use_tsid_binary_join() {
10379        let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric");
10380
10381        let table_provider = build_test_table_provider_with_tsid_tag_fields(&[
10382            (
10383                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10384                2,
10385                1,
10386            ),
10387            (
10388                (
10389                    DEFAULT_SCHEMA_NAME.to_string(),
10390                    "some_alt_metric".to_string(),
10391                ),
10392                1,
10393                1,
10394            ),
10395        ])
10396        .await;
10397        let plan =
10398            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10399                .await
10400                .unwrap();
10401
10402        let plan_str = plan.display_indent_schema().to_string();
10403        assert!(!plan_str.contains("__tsid ="), "{plan_str}");
10404        assert!(
10405            plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
10406            "{plan_str}"
10407        );
10408        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
10409    }
10410
10411    #[tokio::test]
10412    async fn comparison_binary_join_uses_tsid_and_keeps_it_in_filtered_result() {
10413        let eval_stmt = build_eval_stmt("some_metric > some_alt_metric");
10414
10415        let table_provider = build_test_table_provider_with_tsid(
10416            &[
10417                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10418                (
10419                    DEFAULT_SCHEMA_NAME.to_string(),
10420                    "some_alt_metric".to_string(),
10421                ),
10422            ],
10423            2,
10424            1,
10425        )
10426        .await;
10427        let mut planner = PromPlanner {
10428            table_provider,
10429            ctx: PromPlannerContext::from_eval_stmt(&eval_stmt),
10430            promql_annotations: None,
10431        };
10432        let plan = planner
10433            .prom_expr_to_plan(&eval_stmt.expr, &build_query_engine_state())
10434            .await
10435            .unwrap();
10436
10437        let plan_str = plan.display_indent_schema().to_string();
10438        assert!(
10439            plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
10440            "{plan_str}"
10441        );
10442        assert!(
10443            plan.schema()
10444                .fields()
10445                .iter()
10446                .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME),
10447            "{plan_str}"
10448        );
10449        assert!(planner.ctx.use_tsid, "{plan_str}");
10450    }
10451
10452    #[tokio::test]
10453    async fn comparison_bool_binary_join_uses_tsid_when_available() {
10454        let eval_stmt = build_eval_stmt("some_metric > bool some_alt_metric");
10455
10456        let table_provider = build_test_table_provider_with_tsid(
10457            &[
10458                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10459                (
10460                    DEFAULT_SCHEMA_NAME.to_string(),
10461                    "some_alt_metric".to_string(),
10462                ),
10463            ],
10464            2,
10465            1,
10466        )
10467        .await;
10468        let plan =
10469            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10470                .await
10471                .unwrap();
10472
10473        let plan_str = plan.display_indent_schema().to_string();
10474        assert!(
10475            plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
10476            "{plan_str}"
10477        );
10478        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
10479        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
10480    }
10481
10482    #[tokio::test]
10483    async fn scalar_count_count_range_keeps_full_window() {
10484        let plan_str = build_optimized_tsid_plan(
10485            "scalar(count(count(some_metric) by (tag_0)))",
10486            1,
10487            1,
10488            100_000,
10489            1,
10490        )
10491        .await;
10492        assert!(plan_str.contains("ScalarCalculate: tags=[]"));
10493        assert!(plan_str.contains("PromInstantManipulate: range=[0..100000000]"));
10494        assert!(!plan_str.contains("PromInstantManipulate: range=[99999000..99999000]"));
10495    }
10496
10497    #[tokio::test]
10498    async fn scalar_count_count_rewrite_applies_inside_binary_expr_for_tsid_input() {
10499        let plan_str = build_optimized_tsid_plan(
10500            "sum(irate(some_metric[1h])) / scalar(count(count(some_metric) by (tag_0)))",
10501            2,
10502            1,
10503            10,
10504            300,
10505        )
10506        .await;
10507        assert!(plan_str.contains("Distinct:"), "{plan_str}");
10508    }
10509
10510    #[tokio::test]
10511    async fn nested_count_rewrite_keeps_full_series_key_with_tsid_input() {
10512        assert_nested_count_rewrite_applies(
10513            "count(count(some_metric) by (tag_0))",
10514            "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(count(some_metric.field_0))]]"
10515        )
10516        .await;
10517    }
10518
10519    #[tokio::test]
10520    async fn nested_sum_count_rewrite_keeps_full_series_key_with_tsid_input() {
10521        assert_nested_count_rewrite_applies(
10522            "count(sum(some_metric) by (tag_0))",
10523            "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(sum(some_metric.field_0))]]"
10524        )
10525        .await;
10526    }
10527
10528    #[tokio::test]
10529    async fn nested_supported_inner_aggs_rewrite_apply_for_tsid_input() {
10530        for (query, expected_outer_agg) in [
10531            (
10532                "count(avg(some_metric) by (tag_0))",
10533                "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(avg(some_metric.field_0))]]",
10534            ),
10535            (
10536                "count(min(some_metric) by (tag_0))",
10537                "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(min(some_metric.field_0))]]",
10538            ),
10539            (
10540                "count(max(some_metric) by (tag_0))",
10541                "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(max(some_metric.field_0))]]",
10542            ),
10543            (
10544                "count(stddev(some_metric) by (tag_0))",
10545                "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(stddev_pop(some_metric.field_0))]]",
10546            ),
10547            (
10548                "count(stdvar(some_metric) by (tag_0))",
10549                "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(var_pop(some_metric.field_0))]]",
10550            ),
10551        ] {
10552            assert_nested_count_rewrite_applies(query, expected_outer_agg).await;
10553        }
10554    }
10555
10556    #[tokio::test]
10557    async fn nested_non_count_inner_aggs_rewrite_filter_null_values_for_tsid_input() {
10558        let count_plan =
10559            build_optimized_tsid_plan("count(count(some_metric) by (tag_0))", 2, 1, 100_000, 1)
10560                .await;
10561        assert!(
10562            !count_plan.contains("some_metric.field_0 IS NOT NULL"),
10563            "{count_plan}"
10564        );
10565
10566        for query in [
10567            "count(sum(some_metric) by (tag_0))",
10568            "count(avg(some_metric) by (tag_0))",
10569            "count(min(some_metric) by (tag_0))",
10570            "count(max(some_metric) by (tag_0))",
10571            "count(stddev(some_metric) by (tag_0))",
10572            "count(stdvar(some_metric) by (tag_0))",
10573        ] {
10574            let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await;
10575            assert!(
10576                plan_str.contains("Filter: some_metric.field_0 IS NOT NULL"),
10577                "{query}: {plan_str}"
10578            );
10579        }
10580    }
10581
10582    #[tokio::test]
10583    async fn nested_unsupported_or_non_direct_inner_aggs_do_not_rewrite() {
10584        assert_nested_count_rewrite_missing("count(group(some_metric) by (tag_0))", 2, 1).await;
10585        assert_nested_count_rewrite_missing(
10586            "count(sum(irate(some_metric[1h])) by (tag_0))",
10587            2,
10588            300,
10589        )
10590        .await;
10591    }
10592
10593    #[tokio::test]
10594    async fn physical_table_name_is_not_leaked_in_plan() {
10595        let prom_expr = parser::parse("some_metric").unwrap();
10596        let eval_stmt = EvalStmt {
10597            expr: prom_expr,
10598            start: UNIX_EPOCH,
10599            end: UNIX_EPOCH
10600                .checked_add(Duration::from_secs(100_000))
10601                .unwrap(),
10602            interval: Duration::from_secs(5),
10603            lookback_delta: Duration::from_secs(1),
10604        };
10605
10606        let table_provider = build_test_table_provider_with_tsid(
10607            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10608            1,
10609            1,
10610        )
10611        .await;
10612        let plan =
10613            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10614                .await
10615                .unwrap();
10616
10617        let plan_str = plan.display_indent_schema().to_string();
10618        assert!(plan_str.contains("TableScan: phy"), "{plan}");
10619        assert!(plan_str.contains("SubqueryAlias: some_metric"));
10620        assert!(plan_str.contains("Filter: phy.__table_id = UInt32(1024)"));
10621        assert!(!plan_str.contains("TableScan: some_metric"));
10622    }
10623
10624    #[tokio::test]
10625    async fn sum_without_does_not_group_by_tsid() {
10626        let prom_expr = parser::parse("sum without (tag_0) (some_metric)").unwrap();
10627        let eval_stmt = EvalStmt {
10628            expr: prom_expr,
10629            start: UNIX_EPOCH,
10630            end: UNIX_EPOCH
10631                .checked_add(Duration::from_secs(100_000))
10632                .unwrap(),
10633            interval: Duration::from_secs(5),
10634            lookback_delta: Duration::from_secs(1),
10635        };
10636
10637        let table_provider = build_test_table_provider_with_tsid(
10638            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10639            1,
10640            1,
10641        )
10642        .await;
10643        let plan =
10644            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10645                .await
10646                .unwrap();
10647
10648        let plan_str = plan.display_indent_schema().to_string();
10649        assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10650
10651        let aggr_line = plan_str
10652            .lines()
10653            .find(|line| line.contains("Aggregate: groupBy="))
10654            .unwrap();
10655        assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
10656    }
10657
10658    #[tokio::test]
10659    async fn topk_without_does_not_partition_by_tsid() {
10660        let prom_expr = parser::parse("topk without (tag_0) (1, some_metric)").unwrap();
10661        let eval_stmt = EvalStmt {
10662            expr: prom_expr,
10663            start: UNIX_EPOCH,
10664            end: UNIX_EPOCH
10665                .checked_add(Duration::from_secs(100_000))
10666                .unwrap(),
10667            interval: Duration::from_secs(5),
10668            lookback_delta: Duration::from_secs(1),
10669        };
10670
10671        let table_provider = build_test_table_provider_with_tsid(
10672            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10673            1,
10674            1,
10675        )
10676        .await;
10677        let plan =
10678            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10679                .await
10680                .unwrap();
10681
10682        let plan_str = plan.display_indent_schema().to_string();
10683        assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10684
10685        let window_line = plan_str
10686            .lines()
10687            .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()"))
10688            .unwrap();
10689        let partition_by = window_line
10690            .split("PARTITION BY [")
10691            .nth(1)
10692            .and_then(|s| s.split("] ORDER BY").next())
10693            .unwrap();
10694        assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
10695    }
10696
10697    #[tokio::test]
10698    async fn sum_by_does_not_group_by_tsid() {
10699        let prom_expr = parser::parse("sum by (__tsid) (some_metric)").unwrap();
10700        let eval_stmt = EvalStmt {
10701            expr: prom_expr,
10702            start: UNIX_EPOCH,
10703            end: UNIX_EPOCH
10704                .checked_add(Duration::from_secs(100_000))
10705                .unwrap(),
10706            interval: Duration::from_secs(5),
10707            lookback_delta: Duration::from_secs(1),
10708        };
10709
10710        let table_provider = build_test_table_provider_with_tsid(
10711            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10712            1,
10713            1,
10714        )
10715        .await;
10716        let plan =
10717            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10718                .await
10719                .unwrap();
10720
10721        let plan_str = plan.display_indent_schema().to_string();
10722        assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10723
10724        let aggr_line = plan_str
10725            .lines()
10726            .find(|line| line.contains("Aggregate: groupBy="))
10727            .unwrap();
10728        assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
10729    }
10730
10731    #[tokio::test]
10732    async fn aggregate_over_binary_time_function_expr() {
10733        for op in ["sum", "min", "max", "avg"] {
10734            let prom_expr = parser::parse(&format!(
10735                "{op} by (tag_0, tag_1, tag_2) (time() - some_metric)"
10736            ))
10737            .unwrap();
10738            let eval_stmt = EvalStmt {
10739                expr: prom_expr,
10740                start: UNIX_EPOCH,
10741                end: UNIX_EPOCH
10742                    .checked_add(Duration::from_secs(100_000))
10743                    .unwrap(),
10744                interval: Duration::from_secs(5),
10745                lookback_delta: Duration::from_secs(1),
10746            };
10747
10748            let table_provider = build_test_table_provider_with_tsid(
10749                &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10750                3,
10751                1,
10752            )
10753            .await;
10754            let plan =
10755                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10756                    .await
10757                    .unwrap();
10758
10759            let plan_str = plan.display_indent_schema().to_string();
10760            let aggr_line = plan_str
10761                .lines()
10762                .find(|line| line.contains("Aggregate: groupBy="))
10763                .unwrap();
10764            assert!(aggr_line.contains(op), "{plan_str}");
10765            assert!(aggr_line.contains("first_value"), "{plan_str}");
10766            assert!(
10767                !plan
10768                    .schema()
10769                    .fields()
10770                    .iter()
10771                    .any(|field| { field.name() == DATA_SCHEMA_TSID_COLUMN_NAME })
10772            );
10773        }
10774    }
10775
10776    #[tokio::test]
10777    async fn topk_by_does_not_partition_by_tsid() {
10778        let prom_expr = parser::parse("topk by (__tsid) (1, some_metric)").unwrap();
10779        let eval_stmt = EvalStmt {
10780            expr: prom_expr,
10781            start: UNIX_EPOCH,
10782            end: UNIX_EPOCH
10783                .checked_add(Duration::from_secs(100_000))
10784                .unwrap(),
10785            interval: Duration::from_secs(5),
10786            lookback_delta: Duration::from_secs(1),
10787        };
10788
10789        let table_provider = build_test_table_provider_with_tsid(
10790            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10791            1,
10792            1,
10793        )
10794        .await;
10795        let plan =
10796            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10797                .await
10798                .unwrap();
10799
10800        let plan_str = plan.display_indent_schema().to_string();
10801        assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10802
10803        let window_line = plan_str
10804            .lines()
10805            .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()"))
10806            .unwrap();
10807        let partition_by = window_line
10808            .split("PARTITION BY [")
10809            .nth(1)
10810            .and_then(|s| s.split("] ORDER BY").next())
10811            .unwrap();
10812        assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
10813    }
10814
10815    #[tokio::test]
10816    async fn selector_matcher_on_tsid_does_not_use_internal_column() {
10817        let prom_expr = parser::parse(r#"some_metric{__tsid="123"}"#).unwrap();
10818        let eval_stmt = EvalStmt {
10819            expr: prom_expr,
10820            start: UNIX_EPOCH,
10821            end: UNIX_EPOCH
10822                .checked_add(Duration::from_secs(100_000))
10823                .unwrap(),
10824            interval: Duration::from_secs(5),
10825            lookback_delta: Duration::from_secs(1),
10826        };
10827
10828        let table_provider = build_test_table_provider_with_tsid(
10829            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10830            1,
10831            1,
10832        )
10833        .await;
10834        let plan =
10835            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10836                .await
10837                .unwrap();
10838
10839        fn collect_filter_cols(plan: &LogicalPlan, out: &mut HashSet<Column>) {
10840            if let LogicalPlan::Filter(filter) = plan {
10841                datafusion_expr::utils::expr_to_columns(&filter.predicate, out).unwrap();
10842            }
10843            for input in plan.inputs() {
10844                collect_filter_cols(input, out);
10845            }
10846        }
10847
10848        let mut filter_cols = HashSet::new();
10849        collect_filter_cols(&plan, &mut filter_cols);
10850        assert!(
10851            !filter_cols
10852                .iter()
10853                .any(|c| c.name == DATA_SCHEMA_TSID_COLUMN_NAME)
10854        );
10855    }
10856
10857    #[tokio::test]
10858    async fn tsid_is_not_used_when_physical_table_is_missing() {
10859        let prom_expr = parser::parse("some_metric").unwrap();
10860        let eval_stmt = EvalStmt {
10861            expr: prom_expr,
10862            start: UNIX_EPOCH,
10863            end: UNIX_EPOCH
10864                .checked_add(Duration::from_secs(100_000))
10865                .unwrap(),
10866            interval: Duration::from_secs(5),
10867            lookback_delta: Duration::from_secs(1),
10868        };
10869
10870        let catalog_list = MemoryCatalogManager::with_default_setup();
10871
10872        // Register a metric engine logical table referencing a missing physical table.
10873        let mut columns = vec![ColumnSchema::new(
10874            "tag_0".to_string(),
10875            ConcreteDataType::string_datatype(),
10876            false,
10877        )];
10878        columns.push(
10879            ColumnSchema::new(
10880                "timestamp".to_string(),
10881                ConcreteDataType::timestamp_millisecond_datatype(),
10882                false,
10883            )
10884            .with_time_index(true),
10885        );
10886        columns.push(ColumnSchema::new(
10887            "field_0".to_string(),
10888            ConcreteDataType::float64_datatype(),
10889            true,
10890        ));
10891        let schema = Arc::new(Schema::new(columns));
10892        let mut options = table::requests::TableOptions::default();
10893        options
10894            .extra_options
10895            .insert(LOGICAL_TABLE_METADATA_KEY.to_string(), "phy".to_string());
10896        let table_meta = TableMetaBuilder::empty()
10897            .schema(schema)
10898            .primary_key_indices(vec![0])
10899            .value_indices(vec![2])
10900            .engine(METRIC_ENGINE_NAME.to_string())
10901            .options(options)
10902            .next_column_id(1024)
10903            .build()
10904            .unwrap();
10905        let table_info = TableInfoBuilder::default()
10906            .table_id(1024)
10907            .name("some_metric")
10908            .meta(table_meta)
10909            .build()
10910            .unwrap();
10911        let table = EmptyTable::from_table_info(&table_info);
10912        catalog_list
10913            .register_table_sync(RegisterTableRequest {
10914                catalog: DEFAULT_CATALOG_NAME.to_string(),
10915                schema: DEFAULT_SCHEMA_NAME.to_string(),
10916                table_name: "some_metric".to_string(),
10917                table_id: 1024,
10918                table,
10919            })
10920            .unwrap();
10921
10922        let table_provider = DfTableSourceProvider::new(
10923            catalog_list,
10924            false,
10925            QueryContext::arc(),
10926            DummyDecoder::arc(),
10927            false,
10928        );
10929
10930        let plan =
10931            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10932                .await
10933                .unwrap();
10934
10935        let plan_str = plan.display_indent_schema().to_string();
10936        assert!(plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]"));
10937        assert!(!plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10938    }
10939
10940    #[tokio::test]
10941    async fn tsid_is_carried_only_when_aggregate_preserves_label_set() {
10942        let prom_expr = parser::parse("sum by (tag_0) (some_metric)").unwrap();
10943        let eval_stmt = EvalStmt {
10944            expr: prom_expr,
10945            start: UNIX_EPOCH,
10946            end: UNIX_EPOCH
10947                .checked_add(Duration::from_secs(100_000))
10948                .unwrap(),
10949            interval: Duration::from_secs(5),
10950            lookback_delta: Duration::from_secs(1),
10951        };
10952
10953        let table_provider = build_test_table_provider_with_tsid(
10954            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10955            1,
10956            1,
10957        )
10958        .await;
10959        let plan =
10960            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10961                .await
10962                .unwrap();
10963
10964        let plan_str = plan.display_indent_schema().to_string();
10965        assert!(plan_str.contains("first_value") && plan_str.contains("__tsid"));
10966        assert!(
10967            !plan
10968                .schema()
10969                .fields()
10970                .iter()
10971                .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
10972        );
10973
10974        // Merging aggregate: label set is reduced, tsid should not be carried.
10975        let prom_expr = parser::parse("sum(some_metric)").unwrap();
10976        let eval_stmt = EvalStmt {
10977            expr: prom_expr,
10978            start: UNIX_EPOCH,
10979            end: UNIX_EPOCH
10980                .checked_add(Duration::from_secs(100_000))
10981                .unwrap(),
10982            interval: Duration::from_secs(5),
10983            lookback_delta: Duration::from_secs(1),
10984        };
10985        let table_provider = build_test_table_provider_with_tsid(
10986            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10987            1,
10988            1,
10989        )
10990        .await;
10991        let plan =
10992            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10993                .await
10994                .unwrap();
10995        let plan_str = plan.display_indent_schema().to_string();
10996        assert!(!plan_str.contains("first_value"));
10997    }
10998
10999    #[tokio::test]
11000    async fn or_operator_with_unknown_metric_does_not_require_tsid() {
11001        let prom_expr = parser::parse("unknown_metric or some_metric").unwrap();
11002        let eval_stmt = EvalStmt {
11003            expr: prom_expr,
11004            start: UNIX_EPOCH,
11005            end: UNIX_EPOCH
11006                .checked_add(Duration::from_secs(100_000))
11007                .unwrap(),
11008            interval: Duration::from_secs(5),
11009            lookback_delta: Duration::from_secs(1),
11010        };
11011
11012        let table_provider = build_test_table_provider_with_tsid(
11013            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
11014            1,
11015            1,
11016        )
11017        .await;
11018
11019        let plan =
11020            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11021                .await
11022                .unwrap();
11023
11024        assert!(
11025            !plan
11026                .schema()
11027                .fields()
11028                .iter()
11029                .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
11030        );
11031    }
11032
11033    #[tokio::test]
11034    async fn aggregate_avg() {
11035        do_aggregate_expr_plan("avg", "avg").await;
11036    }
11037
11038    #[tokio::test]
11039    #[should_panic] // output type doesn't match
11040    async fn aggregate_count() {
11041        do_aggregate_expr_plan("count", "count").await;
11042    }
11043
11044    #[tokio::test]
11045    async fn aggregate_min() {
11046        do_aggregate_expr_plan("min", "min").await;
11047    }
11048
11049    #[tokio::test]
11050    async fn aggregate_max() {
11051        do_aggregate_expr_plan("max", "max").await;
11052    }
11053
11054    #[tokio::test]
11055    async fn aggregate_group() {
11056        // Regression test for `group()` aggregator.
11057        // PromQL: sum(group by (cluster)(kubernetes_build_info{service="kubernetes",job="apiserver"}))
11058        // should be plannable, and `group()` should produce constant 1 for each group.
11059        let prom_expr = parser::parse(
11060            "sum(group by (cluster)(kubernetes_build_info{service=\"kubernetes\",job=\"apiserver\"}))",
11061        )
11062        .unwrap();
11063        let eval_stmt = EvalStmt {
11064            expr: prom_expr,
11065            start: UNIX_EPOCH,
11066            end: UNIX_EPOCH
11067                .checked_add(Duration::from_secs(100_000))
11068                .unwrap(),
11069            interval: Duration::from_secs(5),
11070            lookback_delta: Duration::from_secs(1),
11071        };
11072
11073        let table_provider = build_test_table_provider_with_fields(
11074            &[(
11075                DEFAULT_SCHEMA_NAME.to_string(),
11076                "kubernetes_build_info".to_string(),
11077            )],
11078            &["cluster", "service", "job"],
11079        )
11080        .await;
11081        let plan =
11082            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11083                .await
11084                .unwrap();
11085
11086        let plan_str = plan.display_indent_schema().to_string();
11087        assert!(plan_str.contains("max(Float64(1"));
11088    }
11089
11090    #[tokio::test]
11091    async fn aggregate_stddev() {
11092        do_aggregate_expr_plan("stddev", "stddev_pop").await;
11093    }
11094
11095    #[tokio::test]
11096    async fn aggregate_stdvar() {
11097        do_aggregate_expr_plan("stdvar", "var_pop").await;
11098    }
11099
11100    // TODO(ruihang): add range fn tests once exprs are ready.
11101
11102    // {
11103    //     input: "some_metric{tag_0="foo"} + some_metric{tag_0="bar"}",
11104    //     expected: &BinaryExpr{
11105    //         Op: ADD,
11106    //         LHS: &VectorSelector{
11107    //             Name: "a",
11108    //             LabelMatchers: []*labels.Matcher{
11109    //                     MustLabelMatcher(labels.MatchEqual, "tag_0", "foo"),
11110    //                     MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"),
11111    //             },
11112    //         },
11113    //         RHS: &VectorSelector{
11114    //             Name: "sum",
11115    //             LabelMatchers: []*labels.Matcher{
11116    //                     MustLabelMatcher(labels.MatchxEqual, "tag_0", "bar"),
11117    //                     MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"),
11118    //             },
11119    //         },
11120    //         VectorMatching: &VectorMatching{},
11121    //     },
11122    // },
11123    #[tokio::test]
11124    async fn binary_op_column_column() {
11125        let prom_expr =
11126            parser::parse(r#"some_metric{tag_0="foo"} + some_metric{tag_0="bar"}"#).unwrap();
11127        let eval_stmt = EvalStmt {
11128            expr: prom_expr,
11129            start: UNIX_EPOCH,
11130            end: UNIX_EPOCH
11131                .checked_add(Duration::from_secs(100_000))
11132                .unwrap(),
11133            interval: Duration::from_secs(5),
11134            lookback_delta: Duration::from_secs(1),
11135        };
11136
11137        let table_provider = build_test_table_provider(
11138            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
11139            1,
11140            1,
11141        )
11142        .await;
11143        let plan =
11144            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11145                .await
11146                .unwrap();
11147
11148        let expected = String::from(
11149            "Projection: rhs.tag_0, rhs.timestamp, CAST(lhs.field_0 AS Float64) + CAST(rhs.field_0 AS Float64) AS lhs.field_0 + rhs.field_0 [tag_0:Utf8, timestamp:Timestamp(ms), lhs.field_0 + rhs.field_0:Float64;N]\
11150            \n  Inner Join: lhs.tag_0 = rhs.tag_0, lhs.timestamp = rhs.timestamp [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11151            \n    SubqueryAlias: lhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11152            \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11153            \n        PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11154            \n          Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11155            \n            Filter: some_metric.tag_0 = Utf8(\"foo\") AND some_metric.tag_0 = Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11156            \n              TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11157            \n    SubqueryAlias: rhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11158            \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11159            \n        PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11160            \n          Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11161            \n            Filter: some_metric.tag_0 = Utf8(\"bar\") AND some_metric.tag_0 = Utf8(\"foo\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11162            \n              TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
11163        );
11164
11165        assert_eq!(plan.display_indent_schema().to_string(), expected);
11166    }
11167
11168    async fn indie_query_plan_compare<T: AsRef<str>>(query: &str, expected: T) {
11169        let prom_expr = parser::parse(query).unwrap();
11170        let eval_stmt = EvalStmt {
11171            expr: prom_expr,
11172            start: UNIX_EPOCH,
11173            end: UNIX_EPOCH
11174                .checked_add(Duration::from_secs(100_000))
11175                .unwrap(),
11176            interval: Duration::from_secs(5),
11177            lookback_delta: Duration::from_secs(1),
11178        };
11179
11180        let table_provider = build_test_table_provider(
11181            &[
11182                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
11183                (
11184                    "greptime_private".to_string(),
11185                    "some_alt_metric".to_string(),
11186                ),
11187            ],
11188            1,
11189            1,
11190        )
11191        .await;
11192        let plan =
11193            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11194                .await
11195                .unwrap();
11196
11197        assert_eq!(plan.display_indent_schema().to_string(), expected.as_ref());
11198    }
11199
11200    #[tokio::test]
11201    async fn binary_op_literal_column() {
11202        let query = r#"1 + some_metric{tag_0="bar"}"#;
11203        let expected = String::from(
11204            "Projection: some_metric.tag_0, some_metric.timestamp, Float64(1) + CAST(some_metric.field_0 AS Float64) AS Float64(1) + field_0 [tag_0:Utf8, timestamp:Timestamp(ms), Float64(1) + field_0:Float64;N]\
11205            \n  PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11206            \n    PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11207            \n      Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11208            \n        Filter: some_metric.tag_0 = Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11209            \n          TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
11210        );
11211
11212        indie_query_plan_compare(query, expected).await;
11213    }
11214
11215    #[tokio::test]
11216    async fn binary_op_literal_literal() {
11217        let query = r#"1 + 1"#;
11218        let expected = r#"EmptyMetric: range=[0..100000000], interval=[5000] [time:Timestamp(ms), value:Float64;N]
11219  TableScan: dummy [time:Timestamp(ms), value:Float64;N]"#;
11220        indie_query_plan_compare(query, expected).await;
11221    }
11222
11223    #[tokio::test]
11224    async fn simple_bool_grammar() {
11225        let query = "some_metric != bool 1.2345";
11226        let expected = String::from(
11227            "Projection: some_metric.tag_0, some_metric.timestamp, CAST(some_metric.field_0 != Float64(1.2345) AS Float64) AS field_0 != Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0 != Float64(1.2345):Float64;N]\
11228            \n  PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11229            \n    PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11230            \n      Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11231            \n        Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11232            \n          TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
11233        );
11234
11235        indie_query_plan_compare(query, expected).await;
11236    }
11237
11238    #[tokio::test]
11239    async fn bool_with_additional_arithmetic() {
11240        let query = "some_metric + (1 == bool 2)";
11241        let expected = String::from(
11242            "Projection: some_metric.tag_0, some_metric.timestamp, CAST(some_metric.field_0 AS Float64) + CAST(Float64(1) = Float64(2) AS Float64) AS field_0 + Float64(1) = Float64(2) [tag_0:Utf8, timestamp:Timestamp(ms), field_0 + Float64(1) = Float64(2):Float64;N]\
11243            \n  PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11244            \n    PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11245            \n      Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11246            \n        Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11247            \n          TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
11248        );
11249
11250        indie_query_plan_compare(query, expected).await;
11251    }
11252
11253    #[tokio::test]
11254    async fn simple_unary() {
11255        let query = "-some_metric";
11256        let expected = String::from(
11257            "Projection: some_metric.tag_0, some_metric.timestamp, (- some_metric.field_0) AS (- field_0) [tag_0:Utf8, timestamp:Timestamp(ms), (- field_0):Float64;N]\
11258            \n  PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11259            \n    PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11260            \n      Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11261            \n        Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11262            \n          TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
11263        );
11264
11265        indie_query_plan_compare(query, expected).await;
11266    }
11267
11268    #[tokio::test]
11269    async fn increase_aggr() {
11270        let query = "increase(some_metric[5m])";
11271        let expected = String::from(
11272            "Filter: prom_increase(timestamp_range,field_0,timestamp,Int64(300000)) IS NOT NULL [timestamp:Timestamp(ms), prom_increase(timestamp_range,field_0,timestamp,Int64(300000)):Float64;N, tag_0:Utf8]\
11273            \n  Projection: some_metric.timestamp, prom_increase(timestamp_range, field_0, some_metric.timestamp, Int64(300000)) AS prom_increase(timestamp_range,field_0,timestamp,Int64(300000)), some_metric.tag_0 [timestamp:Timestamp(ms), prom_increase(timestamp_range,field_0,timestamp,Int64(300000)):Float64;N, tag_0:Utf8]\
11274            \n    PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[300000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\
11275            \n      PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11276            \n        PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11277            \n          Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11278            \n            Filter: some_metric.timestamp >= TimestampMillisecond(-299999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11279            \n              TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
11280        );
11281
11282        indie_query_plan_compare(query, expected).await;
11283    }
11284
11285    async fn native_histogram_plan(query: &str) -> String {
11286        let table_provider = build_test_native_histogram_table_provider("some_metric").await;
11287        let plan = PromPlanner::stmt_to_plan(
11288            table_provider,
11289            &build_eval_stmt(query),
11290            &build_query_engine_state(),
11291        )
11292        .await
11293        .unwrap();
11294        plan.display_indent_schema().to_string()
11295    }
11296
11297    #[tokio::test]
11298    async fn native_histogram_count_uses_native_udf() {
11299        let plan = native_histogram_plan("histogram_count(some_metric)").await;
11300
11301        assert!(plan.contains("prom_native_histogram_count"), "{plan}");
11302        assert!(!plan.contains("HistogramFold:"), "{plan}");
11303    }
11304
11305    #[tokio::test]
11306    async fn timestamp_filters_native_histogram_stale_marker_before_projection() {
11307        let mut stale = direct_or_histogram();
11308        stale.sum = f64::from_bits(PROMETHEUS_STALE_NAN_BITS);
11309        let table = operator_metric_table(
11310            "stale_histogram",
11311            2_100,
11312            "a",
11313            None,
11314            DirectOrValue::NativeHistogram(stale),
11315        );
11316        let catalog = MemoryCatalogManager::with_default_setup();
11317        catalog
11318            .register_table_sync(RegisterTableRequest {
11319                catalog: DEFAULT_CATALOG_NAME.to_string(),
11320                schema: DEFAULT_SCHEMA_NAME.to_string(),
11321                table_name: "stale_histogram".to_string(),
11322                table_id: 2_100,
11323                table,
11324            })
11325            .unwrap();
11326        let provider = DfTableSourceProvider::new(
11327            catalog,
11328            false,
11329            QueryContext::arc(),
11330            DummyDecoder::arc(),
11331            false,
11332        );
11333        let state = build_query_engine_state();
11334        let plan = PromPlanner::stmt_to_plan(
11335            provider,
11336            &operator_eval_stmt("timestamp(stale_histogram)"),
11337            &state,
11338        )
11339        .await
11340        .unwrap();
11341        let plan_text = plan.display_indent_schema().to_string();
11342        assert!(plan_text.contains(TIMESTAMP_VALUE_PREFIX), "{plan_text}");
11343
11344        let (_, batches) = execute(plan, &state).await;
11345        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
11346    }
11347
11348    #[tokio::test]
11349    async fn timestamp_filters_stale_marker_from_mixed_sample_companion() {
11350        let histograms = build_histogram_array(&[None]);
11351        let schema = Arc::new(ArrowSchema::new(vec![
11352            Field::new(
11353                "timestamp",
11354                ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
11355                false,
11356            ),
11357            Field::new(
11358                greptime_native_histogram(),
11359                histograms.data_type().clone(),
11360                true,
11361            ),
11362            Field::new(greptime_value(), ArrowDataType::Float64, true),
11363        ]));
11364        let batch = RecordBatch::try_new(
11365            schema.clone(),
11366            vec![
11367                Arc::new(TimestampMillisecondArray::from(vec![1_000])),
11368                histograms,
11369                Arc::new(Float64Array::from(vec![f64::from_bits(
11370                    PROMETHEUS_STALE_NAN_BITS,
11371                )])),
11372            ],
11373        )
11374        .unwrap();
11375        let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap());
11376        let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None)
11377            .unwrap()
11378            .build()
11379            .unwrap();
11380        let input = LogicalPlan::Extension(Extension {
11381            node: Arc::new(SeriesDivide::new(
11382                Vec::new(),
11383                "timestamp".to_string(),
11384                input,
11385            )),
11386        });
11387        let input = LogicalPlan::Extension(Extension {
11388            node: Arc::new(InstantManipulate::new(
11389                1_000,
11390                1_000,
11391                5_000,
11392                1_000,
11393                0,
11394                "timestamp".to_string(),
11395                Vec::new(),
11396                Some(greptime_native_histogram().to_string()),
11397                input,
11398            )),
11399        });
11400        // Match timestamp()'s parent projection, which otherwise prunes the companion lane.
11401        let plan = LogicalPlanBuilder::from(input)
11402            .project([col("timestamp")])
11403            .unwrap()
11404            .build()
11405            .unwrap();
11406
11407        let (_, batches) = execute(plan, &build_query_engine_state()).await;
11408        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
11409    }
11410
11411    #[tokio::test]
11412    async fn native_histogram_rate_can_feed_count() {
11413        let plan = native_histogram_plan("histogram_count(rate(some_metric[5m]))").await;
11414
11415        assert!(plan.contains("prom_native_histogram_rate"), "{plan}");
11416        assert!(plan.contains("prom_native_histogram_count"), "{plan}");
11417    }
11418
11419    #[tokio::test]
11420    async fn native_histogram_quantile_skips_classic_fold() {
11421        let plan = native_histogram_plan("histogram_quantile(0.9, some_metric)").await;
11422
11423        assert!(plan.contains("prom_native_histogram_quantile"), "{plan}");
11424        assert!(!plan.contains("HistogramFold:"), "{plan}");
11425        assert!(plan.contains("some_metric.le"), "{plan}");
11426        // The phi literal is threaded into the native quantile UDF as its second argument.
11427        assert!(plan.contains("Float64(0.9)"), "{plan}");
11428        // The empty-values filter drops NULL quantile results so the output is empty
11429        // when all native histogram samples are dropped.
11430        assert!(plan.contains("IS NOT NULL"), "{plan}");
11431    }
11432
11433    #[tokio::test]
11434    async fn mixed_native_histogram_quantile_uses_histogram_field() {
11435        let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
11436        let plan = PromPlanner::stmt_to_plan(
11437            table_provider,
11438            &build_eval_stmt("histogram_quantile(0.9, some_metric)"),
11439            &build_query_engine_state(),
11440        )
11441        .await
11442        .unwrap()
11443        .display_indent_schema()
11444        .to_string();
11445
11446        assert!(
11447            plan.contains("prom_native_histogram_quantile(greptime_native_histogram"),
11448            "{plan}"
11449        );
11450        assert!(!plan.contains("EmptyRelation"), "{plan}");
11451    }
11452
11453    #[tokio::test]
11454    async fn mixed_histogram_helpers_execute_classic_and_native_samples() {
11455        let state = build_query_engine_state();
11456        for (query, expected) in [
11457            (
11458                "histogram_quantile(0.5, mixed_histogram)",
11459                vec![("classic", 1.0), ("native", 0.0)],
11460            ),
11461            (
11462                "histogram_fraction(-Inf, +Inf, mixed_histogram)",
11463                vec![("classic", 1.0), ("native", 1.0)],
11464            ),
11465        ] {
11466            let plan = PromPlanner::stmt_to_plan(
11467                classic_and_native_histogram_table_provider("native", None, direct_or_histogram()),
11468                &operator_eval_stmt(query),
11469                &state,
11470            )
11471            .await
11472            .unwrap();
11473            let plan_text = plan.display_indent_schema().to_string();
11474            assert!(plan_text.contains("HistogramFold:"), "{plan_text}");
11475            assert!(plan_text.contains("prom_native_histogram_"), "{plan_text}");
11476            let value_field = plan
11477                .schema()
11478                .fields()
11479                .iter()
11480                .find(|field| field.data_type() == &ArrowDataType::Float64)
11481                .unwrap()
11482                .name()
11483                .clone();
11484
11485            let (_, batches) = execute(plan, &state).await;
11486            let mut actual = batches
11487                .iter()
11488                .flat_map(|batch| {
11489                    let tags = batch
11490                        .column_by_name("tag")
11491                        .unwrap()
11492                        .as_any()
11493                        .downcast_ref::<StringArray>()
11494                        .unwrap();
11495                    let values = batch
11496                        .column_by_name(&value_field)
11497                        .unwrap()
11498                        .as_any()
11499                        .downcast_ref::<Float64Array>()
11500                        .unwrap();
11501                    (0..batch.num_rows()).map(|row| (tags.value(row), values.value(row)))
11502                })
11503                .collect::<Vec<_>>();
11504            actual.sort_by_key(|(tag, _)| *tag);
11505            assert_eq!(actual, expected, "{query}");
11506        }
11507    }
11508
11509    #[tokio::test]
11510    async fn mixed_histogram_helpers_report_annotations() {
11511        let state = build_query_engine_state();
11512        let mut native_histogram = direct_or_histogram();
11513        native_histogram.count = 2.0;
11514        native_histogram.sum = f64::NAN;
11515        for (native_tag, expected_rows, expected_warnings, expected_infos) in [
11516            (
11517                "classic",
11518                0,
11519                vec!["vector contains a mix of classic and native histograms"],
11520                vec![],
11521            ),
11522            (
11523                "native",
11524                2,
11525                vec![],
11526                vec!["input to histogram_quantile has NaN observations, result is skewed higher"],
11527            ),
11528        ] {
11529            let collector = PromqlAnnotationCollector::default();
11530            let plan = PromPlanner::stmt_to_plan_with_annotations(
11531                classic_and_native_histogram_table_provider(
11532                    native_tag,
11533                    None,
11534                    native_histogram.clone(),
11535                ),
11536                &operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"),
11537                &state,
11538                Some(collector.clone()),
11539            )
11540            .await
11541            .unwrap();
11542
11543            let (_, batches) = execute(plan, &state).await;
11544            assert_eq!(
11545                batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
11546                expected_rows
11547            );
11548            let mut warnings = vec![];
11549            let mut infos = vec![];
11550            collector.append_to(&mut warnings, &mut infos);
11551            assert_eq!(warnings, expected_warnings);
11552            assert_eq!(infos, expected_infos);
11553        }
11554    }
11555
11556    #[tokio::test]
11557    async fn mixed_histogram_helper_preserves_native_le_and_scans_once() {
11558        let state = build_query_engine_state();
11559        let mut stmt = operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)");
11560        stmt.end = UNIX_EPOCH.checked_add(Duration::from_secs(2)).unwrap();
11561        let plan = PromPlanner::stmt_to_plan(
11562            classic_and_native_histogram_table_provider(
11563                "classic",
11564                Some("native"),
11565                direct_or_histogram(),
11566            ),
11567            &stmt,
11568            &state,
11569        )
11570        .await
11571        .unwrap();
11572        let plan_text = plan.display_indent_schema().to_string();
11573        assert_eq!(
11574            plan_text.matches("TableScan: mixed_histogram").count(),
11575            1,
11576            "{plan_text}"
11577        );
11578
11579        let value_field = plan
11580            .schema()
11581            .fields()
11582            .iter()
11583            .find(|field| field.data_type() == &ArrowDataType::Float64)
11584            .unwrap()
11585            .name()
11586            .clone();
11587        let (_, batches) = execute(plan, &state).await;
11588        let mut actual = batches
11589            .iter()
11590            .flat_map(|batch| {
11591                let le = batch
11592                    .column_by_name(LE_COLUMN_NAME)
11593                    .unwrap()
11594                    .as_any()
11595                    .downcast_ref::<StringArray>()
11596                    .unwrap();
11597                let timestamps = batch
11598                    .column_by_name("timestamp")
11599                    .unwrap()
11600                    .as_any()
11601                    .downcast_ref::<TimestampMillisecondArray>()
11602                    .unwrap();
11603                let values = batch
11604                    .column_by_name(&value_field)
11605                    .unwrap()
11606                    .as_any()
11607                    .downcast_ref::<Float64Array>()
11608                    .unwrap();
11609                (0..batch.num_rows()).map(|row| {
11610                    (
11611                        timestamps.value(row),
11612                        (!le.is_null(row)).then(|| le.value(row).to_string()),
11613                        values.value(row),
11614                    )
11615                })
11616            })
11617            .collect::<Vec<_>>();
11618        actual.sort_by(|lhs, rhs| (lhs.0, &lhs.1).cmp(&(rhs.0, &rhs.1)));
11619        assert_eq!(
11620            actual,
11621            vec![
11622                (1_000, None, 1.0),
11623                (1_000, Some("native".to_string()), 0.0),
11624                (2_000, None, 1.0),
11625                (2_000, Some("native".to_string()), 0.0),
11626            ]
11627        );
11628    }
11629
11630    #[tokio::test]
11631    async fn nested_histogram_helpers_ignore_unparsable_bucket_labels() {
11632        let state = build_query_engine_state();
11633        for native_le in [None, Some("native")] {
11634            for query in [
11635                "histogram_quantile(0.5, histogram_quantile(0.5, mixed_histogram))",
11636                "histogram_fraction(-Inf, +Inf, histogram_fraction(-Inf, +Inf, mixed_histogram))",
11637            ] {
11638                let plan = PromPlanner::stmt_to_plan(
11639                    classic_and_native_histogram_table_provider(
11640                        "native",
11641                        native_le,
11642                        direct_or_histogram(),
11643                    ),
11644                    &operator_eval_stmt(query),
11645                    &state,
11646                )
11647                .await
11648                .unwrap();
11649
11650                let (_, batches) = execute(plan, &state).await;
11651                assert_eq!(
11652                    batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
11653                    0,
11654                    "native_le={native_le:?}, query={query}"
11655                );
11656            }
11657        }
11658    }
11659
11660    #[tokio::test]
11661    async fn native_histogram_quantile_rejects_multi_field_input() {
11662        let table_provider = build_test_multi_histogram_table_provider("some_metric").await;
11663        let result = PromPlanner::stmt_to_plan(
11664            table_provider,
11665            &build_eval_stmt("histogram_quantile(0.9, some_metric)"),
11666            &build_query_engine_state(),
11667        )
11668        .await;
11669
11670        let err = result.expect_err("histogram_quantile on two native histogram fields must fail");
11671        assert!(
11672            err.to_string()
11673                .contains("Multi fields calculation is not supported in histogram_quantile"),
11674            "{err}"
11675        );
11676    }
11677
11678    #[tokio::test]
11679    async fn native_histogram_topk_uses_drop_udf() {
11680        let plan = native_histogram_plan("topk(1, some_metric)").await;
11681
11682        assert!(plan.contains("prom_native_histogram_drop_float"), "{plan}");
11683        assert!(
11684            plan.contains("Filter: prom_native_histogram_drop_float")
11685                && plan.contains("IS NOT NULL"),
11686            "{plan}"
11687        );
11688    }
11689
11690    #[tokio::test]
11691    async fn mixed_or_topk_bottomk_ignore_native_histograms() {
11692        for op in ["topk", "bottomk"] {
11693            let collector = PromqlAnnotationCollector::default();
11694            let state = build_query_engine_state();
11695            let plan = PromPlanner::stmt_to_plan_with_annotations(
11696                operator_table_provider(),
11697                &operator_eval_stmt(&format!("{op}(1, lf or on(tag) lh)")),
11698                &state,
11699                Some(collector.clone()),
11700            )
11701            .await
11702            .unwrap();
11703            let float_field = plan
11704                .schema()
11705                .fields()
11706                .iter()
11707                .find(|field| field.data_type() == &ArrowDataType::Float64)
11708                .unwrap()
11709                .name()
11710                .clone();
11711            assert!(
11712                plan.schema()
11713                    .fields()
11714                    .iter()
11715                    .all(|field| field.data_type() != &PromPlanner::native_histogram_arrow_type()),
11716                "{plan:?}"
11717            );
11718
11719            let (_, batches) = execute(plan, &state).await;
11720            assert_eq!(values(&batches, &float_field), vec![2.0], "{op}");
11721            let mut warnings = vec![];
11722            let mut infos = vec![];
11723            collector.append_to(&mut warnings, &mut infos);
11724            assert!(warnings.is_empty());
11725            assert_eq!(
11726                infos,
11727                vec![format!(
11728                    "{op}: dropped native histogram samples because this aggregation is not supported for native histograms"
11729                )]
11730            );
11731        }
11732    }
11733
11734    #[tokio::test]
11735    async fn native_histogram_scalar_is_ignored_before_scalar_calculate() {
11736        let plan = native_histogram_plan("scalar(some_metric)").await;
11737
11738        assert!(plan.contains("ScalarCalculate"), "{plan}");
11739        assert!(plan.contains("Filter: Boolean(false)"), "{plan}");
11740        assert!(!plan.contains("prom_native_histogram_drop"), "{plan}");
11741    }
11742
11743    #[tokio::test]
11744    async fn native_histogram_value_sort_is_empty_but_label_sort_preserves_samples() {
11745        for function in ["sort", "sort_desc"] {
11746            let plan = native_histogram_plan(&format!("{function}(some_metric)")).await;
11747
11748            assert!(plan.contains("Float64(NULL) IS NOT NULL"), "{plan}");
11749            assert!(
11750                !plan.contains(&format!("Sort: {}", greptime_native_histogram())),
11751                "{plan}"
11752            );
11753            assert!(!plan.contains("prom_native_histogram_drop"), "{plan}");
11754        }
11755
11756        for (function, direction) in [("sort_by_label", "ASC"), ("sort_by_label_desc", "DESC")] {
11757            let plan = native_histogram_plan(&format!("{function}(some_metric, \"tag_0\")")).await;
11758
11759            assert!(plan.contains(&format!("tag_0 {direction}")), "{plan}");
11760            assert!(plan.contains(greptime_native_histogram()), "{plan}");
11761            assert!(!plan.contains("Float64(NULL) IS NOT NULL"), "{plan}");
11762        }
11763    }
11764
11765    #[tokio::test]
11766    async fn unsupported_native_histogram_functions_use_drop_udf() {
11767        for query in [
11768            "deriv(some_metric[5m])",
11769            "min_over_time(some_metric[5m])",
11770            "quantile_over_time(0.9, some_metric[5m])",
11771            "predict_linear(some_metric[5m], 60)",
11772            "round(some_metric)",
11773            "abs(some_metric)",
11774        ] {
11775            let plan = native_histogram_plan(query).await;
11776
11777            assert!(
11778                plan.contains("prom_native_histogram_drop_float"),
11779                "{query}\n{plan}"
11780            );
11781        }
11782    }
11783
11784    #[tokio::test]
11785    async fn native_histogram_absent_over_time_uses_native_udf() {
11786        let plan = native_histogram_plan("absent_over_time(some_metric[5m])").await;
11787
11788        assert!(
11789            plan.contains("prom_native_histogram_absent_over_time"),
11790            "{plan}"
11791        );
11792    }
11793
11794    #[tokio::test]
11795    async fn native_histogram_all_function_arms_route_correctly() {
11796        // Every native-histogram match arm in `create_function_expr` must route to the
11797        // expected UDF when all field columns are native histograms. `holt_winters` shares
11798        // the `double_exponential_smoothing` arm but is not registered in the promql
11799        // parser (0.10), so it cannot be exercised through a query string.
11800        let cases = [
11801            // Range functions routed to native histogram UDFs.
11802            (
11803                "increase(some_metric[5m])",
11804                "prom_native_histogram_increase",
11805            ),
11806            ("rate(some_metric[5m])", "prom_native_histogram_rate"),
11807            ("delta(some_metric[5m])", "prom_native_histogram_delta"),
11808            ("idelta(some_metric[5m])", "prom_native_histogram_idelta"),
11809            ("irate(some_metric[5m])", "prom_native_histogram_irate"),
11810            ("resets(some_metric[5m])", "prom_native_histogram_resets"),
11811            ("changes(some_metric[5m])", "prom_native_histogram_changes"),
11812            (
11813                "avg_over_time(some_metric[5m])",
11814                "prom_native_histogram_avg_over_time",
11815            ),
11816            (
11817                "sum_over_time(some_metric[5m])",
11818                "prom_native_histogram_sum_over_time",
11819            ),
11820            (
11821                "count_over_time(some_metric[5m])",
11822                "prom_native_histogram_count_over_time",
11823            ),
11824            (
11825                "last_over_time(some_metric[5m])",
11826                "prom_native_histogram_last_over_time",
11827            ),
11828            (
11829                "present_over_time(some_metric[5m])",
11830                "prom_native_histogram_present_over_time",
11831            ),
11832            // Unsupported functions dropped with the float-null UDF.
11833            ("deriv(some_metric[5m])", "prom_native_histogram_drop_float"),
11834            (
11835                "min_over_time(some_metric[5m])",
11836                "prom_native_histogram_drop_float",
11837            ),
11838            (
11839                "max_over_time(some_metric[5m])",
11840                "prom_native_histogram_drop_float",
11841            ),
11842            (
11843                "stddev_over_time(some_metric[5m])",
11844                "prom_native_histogram_drop_float",
11845            ),
11846            (
11847                "stdvar_over_time(some_metric[5m])",
11848                "prom_native_histogram_drop_float",
11849            ),
11850            (
11851                "quantile_over_time(0.9, some_metric[5m])",
11852                "prom_native_histogram_drop_float",
11853            ),
11854            (
11855                "predict_linear(some_metric[5m], 60)",
11856                "prom_native_histogram_drop_float",
11857            ),
11858            (
11859                "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)",
11860                "prom_native_histogram_drop_float",
11861            ),
11862            ("round(some_metric)", "prom_native_histogram_drop_float"),
11863            ("rad(some_metric)", "prom_native_histogram_drop_float"),
11864            ("deg(some_metric)", "prom_native_histogram_drop_float"),
11865            ("sgn(some_metric)", "prom_native_histogram_drop_float"),
11866            // Instant helper functions routed to native histogram UDFs.
11867            (
11868                "histogram_count(some_metric)",
11869                "prom_native_histogram_count",
11870            ),
11871            ("histogram_sum(some_metric)", "prom_native_histogram_sum"),
11872            ("histogram_avg(some_metric)", "prom_native_histogram_avg"),
11873            (
11874                "histogram_stddev(some_metric)",
11875                "prom_native_histogram_stddev",
11876            ),
11877            (
11878                "histogram_stdvar(some_metric)",
11879                "prom_native_histogram_stdvar",
11880            ),
11881            (
11882                "histogram_fraction(-2 + 1, 2 / 2, some_metric)",
11883                "prom_native_histogram_fraction",
11884            ),
11885        ];
11886
11887        for (query, expected_udf) in cases {
11888            let plan = native_histogram_plan(query).await;
11889            assert!(plan.contains(expected_udf), "{query}\n{plan}");
11890            if query.starts_with("histogram_fraction") {
11891                assert!(plan.contains("Float64(-1)"), "{query}\n{plan}");
11892            }
11893        }
11894    }
11895
11896    #[tokio::test]
11897    async fn mixed_native_histogram_ranges_use_coordinated_udfs() {
11898        let dual_output = [
11899            "increase(some_metric[5m])",
11900            "rate(some_metric[5m])",
11901            "delta(some_metric[5m])",
11902            "idelta(some_metric[5m])",
11903            "irate(some_metric[5m])",
11904            "avg_over_time(some_metric[5m])",
11905            "sum_over_time(some_metric[5m])",
11906            "last_over_time(some_metric[5m])",
11907        ];
11908        let float_output = [
11909            "resets(some_metric[5m])",
11910            "changes(some_metric[5m])",
11911            "deriv(some_metric[5m])",
11912            "min_over_time(some_metric[5m])",
11913            "max_over_time(some_metric[5m])",
11914            "count_over_time(some_metric[5m])",
11915            "absent_over_time(some_metric[5m])",
11916            "present_over_time(some_metric[5m])",
11917            "stddev_over_time(some_metric[5m])",
11918            "stdvar_over_time(some_metric[5m])",
11919            "quantile_over_time(0.9, some_metric[5m])",
11920            "predict_linear(some_metric[5m], 60)",
11921            "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)",
11922        ];
11923
11924        for query in dual_output.iter().chain(float_output.iter()) {
11925            let plan = PromPlanner::stmt_to_plan(
11926                build_test_mixed_native_histogram_table_provider("some_metric").await,
11927                &build_eval_stmt(query),
11928                &build_query_engine_state(),
11929            )
11930            .await
11931            .unwrap()
11932            .display_indent_schema()
11933            .to_string();
11934            assert!(plan.contains("prom_mixed_range_float"), "{query}\n{plan}");
11935            assert_eq!(
11936                plan.contains("prom_mixed_range_histogram"),
11937                dual_output.contains(query),
11938                "{query}\n{plan}"
11939            );
11940        }
11941
11942        let plan = PromPlanner::stmt_to_plan(
11943            build_test_mixed_native_histogram_table_provider("some_metric").await,
11944            &build_eval_stmt("sum_over_time(rate(some_metric[5m])[10m:1m])"),
11945            &build_query_engine_state(),
11946        )
11947        .await
11948        .unwrap()
11949        .display_indent_schema()
11950        .to_string();
11951        let expected = r#"Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8]
11952  Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_value, prom_mixed_range_histogram(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8]
11953    PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=["greptime_value", "greptime_native_histogram"] [timestamp:Timestamp(ms), greptime_value:Dictionary(Int64, Float64);N, greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, tag_0:Utf8, timestamp_range:Dictionary(Int64, Timestamp(ms))]
11954      PromSeriesDivide: tags=["tag_0"] [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8]
11955        Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8]
11956          Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8]
11957            Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_value, prom_mixed_range_histogram(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8]
11958              PromRangeManipulate: req range=[-540000..100000000], interval=[60000], eval range=[300000], time index=[timestamp], values=["greptime_native_histogram", "greptime_value"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, greptime_value:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]
11959                PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]
11960                  PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]
11961                    Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]
11962                      Filter: some_metric.timestamp >= TimestampMillisecond(-839999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]
11963                        TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]"#;
11964        assert_eq!(plan, expected);
11965    }
11966
11967    #[tokio::test]
11968    async fn mixed_native_histogram_rate_executes_real_ranges() {
11969        let schema = Arc::new(ArrowSchema::new(vec![
11970            Field::new(
11971                "timestamp",
11972                ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
11973                false,
11974            ),
11975            Field::new(greptime_value(), ArrowDataType::Float64, true),
11976            Field::new(
11977                greptime_native_histogram(),
11978                native_histogram_value_type().as_arrow_type(),
11979                true,
11980            ),
11981        ]));
11982        let batch = RecordBatch::try_new(
11983            schema.clone(),
11984            vec![
11985                Arc::new(TimestampMillisecondArray::from(vec![1000, 2000, 3000])),
11986                Arc::new(Float64Array::from(vec![Some(1.0), None, Some(3.0)])),
11987                build_histogram_array(&[None, Some(direct_or_histogram()), None]),
11988            ],
11989        )
11990        .unwrap();
11991        let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap());
11992        let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None)
11993            .unwrap()
11994            .build()
11995            .unwrap();
11996        let collector = PromqlAnnotationCollector::default();
11997        let mut planner = PromPlanner {
11998            table_provider: build_test_table_provider_with_fields(
11999                &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
12000                &[],
12001            )
12002            .await,
12003            ctx: PromPlannerContext {
12004                start: 3000,
12005                end: 3000,
12006                interval: 1000,
12007                range: Some(3000),
12008                time_index_column: Some("timestamp".to_string()),
12009                field_columns: vec![
12010                    greptime_native_histogram().to_string(),
12011                    greptime_value().to_string(),
12012                ],
12013                ..Default::default()
12014            },
12015            promql_annotations: Some(collector.clone()),
12016        };
12017        let input = LogicalPlan::Extension(Extension {
12018            node: Arc::new(
12019                RangeManipulate::new(
12020                    3000,
12021                    3000,
12022                    1000,
12023                    0,
12024                    3000,
12025                    "timestamp".to_string(),
12026                    planner.ctx.field_columns.clone(),
12027                    input,
12028                )
12029                .unwrap(),
12030            ),
12031        });
12032        let PromExpr::Call(call) = parser::parse("rate(mixed[3s])").unwrap() else {
12033            unreachable!()
12034        };
12035        let preserve_any_value = PromPlanner::field_columns_are_alternative_samples(
12036            input.schema(),
12037            &planner.ctx.field_columns,
12038        );
12039        let state = build_query_engine_state();
12040        let (mut exprs, _) = planner
12041            .create_function_expr(&call.func, vec![], input.schema(), &state)
12042            .unwrap();
12043        exprs.insert(0, planner.create_time_index_column_expr().unwrap());
12044        let plan = LogicalPlanBuilder::from(input)
12045            .project(exprs)
12046            .unwrap()
12047            .filter(
12048                planner
12049                    .create_empty_values_filter_expr(preserve_any_value)
12050                    .unwrap(),
12051            )
12052            .unwrap()
12053            .build()
12054            .unwrap();
12055        let (_, batches) = execute(plan, &state).await;
12056        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
12057        let mut warnings = Vec::new();
12058        collector.append_to(&mut warnings, &mut Vec::new());
12059        assert!(
12060            warnings
12061                .iter()
12062                .any(|warning| warning.contains("mix of float and native histogram"))
12063        );
12064    }
12065
12066    #[tokio::test]
12067    async fn native_histogram_mixed_field_table_behaves() {
12068        // Exercise function planning after float and histogram samples have already been
12069        // represented as alternative nullable fields. Histogram functions must select the
12070        // histogram field without adding a NULL float field that would reject every row.
12071        let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
12072        let plan = PromPlanner::stmt_to_plan(
12073            table_provider,
12074            &build_eval_stmt("histogram_count(some_metric)"),
12075            &build_query_engine_state(),
12076        )
12077        .await
12078        .unwrap();
12079        let plan_str = plan.display_indent_schema().to_string();
12080        assert!(
12081            plan_str.contains("prom_native_histogram_count"),
12082            "{plan_str}"
12083        );
12084        assert!(!plan_str.contains("Float64(NULL)"), "{plan_str}");
12085        assert!(
12086            plan_str.contains("prom_native_histogram_count(greptime_native_histogram) IS NOT NULL"),
12087            "{plan_str}"
12088        );
12089
12090        // Value sorting keeps the float column and never sorts by the histogram column.
12091        let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
12092        let plan = PromPlanner::stmt_to_plan(
12093            table_provider,
12094            &build_eval_stmt("sort(some_metric)"),
12095            &build_query_engine_state(),
12096        )
12097        .await
12098        .unwrap();
12099        let plan_str = plan.display_indent_schema().to_string();
12100        assert!(
12101            plan_str.contains("greptime_value ASC NULLS FIRST"),
12102            "{plan_str}"
12103        );
12104        assert!(
12105            !plan_str.contains("greptime_native_histogram ASC"),
12106            "{plan_str}"
12107        );
12108
12109        // scalar() ignores histogram samples and evaluates only the float field.
12110        let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
12111        let plan = PromPlanner::stmt_to_plan(
12112            table_provider,
12113            &build_eval_stmt("scalar(some_metric)"),
12114            &build_query_engine_state(),
12115        )
12116        .await
12117        .unwrap();
12118        let plan_str = plan.display_indent_schema().to_string();
12119        assert!(plan_str.contains("ScalarCalculate"), "{plan_str}");
12120        assert!(
12121            plan_str.contains("greptime_value IS NOT NULL"),
12122            "{plan_str}"
12123        );
12124
12125        // Functions that preserve both alternative fields keep rows with either sample type.
12126        let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
12127        let plan = PromPlanner::stmt_to_plan(
12128            table_provider,
12129            &build_eval_stmt(r#"label_replace(some_metric, "copied", "$1", "tag_0", "(.*)")"#),
12130            &build_query_engine_state(),
12131        )
12132        .await
12133        .unwrap();
12134        let plan_str = plan.display_indent_schema().to_string();
12135        let filter = plan_str.lines().next().unwrap();
12136        assert!(
12137            filter.starts_with("Filter: ")
12138                && filter.contains("greptime_native_histogram IS NOT NULL")
12139                && filter.contains(" OR ")
12140                && filter.contains("greptime_value IS NOT NULL"),
12141            "{plan_str}"
12142        );
12143    }
12144
12145    #[tokio::test]
12146    async fn less_filter_on_value() {
12147        let query = "some_metric < 1.2345";
12148        let expected = String::from(
12149            "Filter: some_metric.field_0 < Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12150            \n  PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12151            \n    PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12152            \n      Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12153            \n        Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12154            \n          TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
12155        );
12156
12157        indie_query_plan_compare(query, expected).await;
12158    }
12159
12160    #[tokio::test]
12161    async fn count_over_time() {
12162        let query = "count_over_time(some_metric[5m])";
12163        let expected = String::from(
12164            "Filter: prom_count_over_time(timestamp_range,field_0) IS NOT NULL [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\
12165            \n  Projection: some_metric.timestamp, prom_count_over_time(timestamp_range, field_0) AS prom_count_over_time(timestamp_range,field_0), some_metric.tag_0 [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\
12166            \n    PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[300000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\
12167            \n      PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12168            \n        PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12169            \n          Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12170            \n            Filter: some_metric.timestamp >= TimestampMillisecond(-299999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12171            \n              TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
12172        );
12173
12174        indie_query_plan_compare(query, expected).await;
12175    }
12176
12177    /// The outer `PromRangeManipulate` from a subquery must be preceded by
12178    /// `Sort` + `PromSeriesDivide`.
12179    #[tokio::test]
12180    async fn count_over_time_subquery() {
12181        let query = "count_over_time(some_metric[10m:1m])";
12182        let expected = String::from(
12183            "Filter: prom_count_over_time(timestamp_range,field_0) IS NOT NULL [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\
12184            \n  Projection: some_metric.timestamp, prom_count_over_time(timestamp_range, field_0) AS prom_count_over_time(timestamp_range,field_0), some_metric.tag_0 [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\
12185            \n    PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\
12186            \n      PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12187            \n        Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12188            \n          PromInstantManipulate: range=[-540000..100000000], lookback=[1000], interval=[60000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12189            \n            PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12190            \n              Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12191            \n                Filter: some_metric.timestamp >= TimestampMillisecond(-540999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12192            \n                  TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
12193        );
12194        indie_query_plan_compare(query, expected).await;
12195    }
12196
12197    #[tokio::test]
12198    async fn test_hash_join() {
12199        let mut eval_stmt = EvalStmt {
12200            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12201            start: UNIX_EPOCH,
12202            end: UNIX_EPOCH
12203                .checked_add(Duration::from_secs(100_000))
12204                .unwrap(),
12205            interval: Duration::from_secs(5),
12206            lookback_delta: Duration::from_secs(1),
12207        };
12208
12209        let case = r#"http_server_requests_seconds_sum{uri="/accounts/login"} / ignoring(kubernetes_pod_name,kubernetes_namespace) http_server_requests_seconds_count{uri="/accounts/login"}"#;
12210
12211        let prom_expr = parser::parse(case).unwrap();
12212        eval_stmt.expr = prom_expr;
12213        let table_provider = build_test_table_provider_with_fields(
12214            &[
12215                (
12216                    DEFAULT_SCHEMA_NAME.to_string(),
12217                    "http_server_requests_seconds_sum".to_string(),
12218                ),
12219                (
12220                    DEFAULT_SCHEMA_NAME.to_string(),
12221                    "http_server_requests_seconds_count".to_string(),
12222                ),
12223            ],
12224            &["uri", "kubernetes_namespace", "kubernetes_pod_name"],
12225        )
12226        .await;
12227        // Should be ok
12228        let plan =
12229            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12230                .await
12231                .unwrap();
12232        let expected = "Projection: http_server_requests_seconds_sum.uri, http_server_requests_seconds_count.greptime_timestamp, CAST(http_server_requests_seconds_sum.greptime_value AS Float64) / CAST(http_server_requests_seconds_count.greptime_value AS Float64) AS http_server_requests_seconds_sum.greptime_value / http_server_requests_seconds_count.greptime_value\
12233            \n  Projection: http_server_requests_seconds_sum.uri, http_server_requests_seconds_sum.kubernetes_namespace, http_server_requests_seconds_sum.kubernetes_pod_name, http_server_requests_seconds_sum.greptime_timestamp, http_server_requests_seconds_sum.greptime_value, http_server_requests_seconds_count.uri, http_server_requests_seconds_count.kubernetes_namespace, http_server_requests_seconds_count.kubernetes_pod_name, http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_count.greptime_value\
12234            \n    Filter: prom_assert_unique_match_group(__promql_match_group_count, http_server_requests_seconds_sum.uri)\
12235            \n      WindowAggr: windowExpr=[[count(Int64(1)) PARTITION BY [http_server_requests_seconds_sum.uri, http_server_requests_seconds_sum.greptime_timestamp] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS __promql_match_group_count]]\
12236            \n        Inner Join: http_server_requests_seconds_sum.greptime_timestamp = http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_sum.uri = http_server_requests_seconds_count.uri\
12237            \n          SubqueryAlias: http_server_requests_seconds_sum\
12238            \n            PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\
12239            \n              PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\
12240            \n                Sort: http_server_requests_seconds_sum.uri ASC NULLS FIRST, http_server_requests_seconds_sum.kubernetes_namespace ASC NULLS FIRST, http_server_requests_seconds_sum.kubernetes_pod_name ASC NULLS FIRST, http_server_requests_seconds_sum.greptime_timestamp ASC NULLS FIRST\
12241            \n                  Filter: http_server_requests_seconds_sum.uri = Utf8(\"/accounts/login\") AND http_server_requests_seconds_sum.greptime_timestamp >= TimestampMillisecond(-999, None) AND http_server_requests_seconds_sum.greptime_timestamp <= TimestampMillisecond(100000000, None)\
12242            \n                    TableScan: http_server_requests_seconds_sum\
12243            \n          SubqueryAlias: http_server_requests_seconds_count\
12244            \n            Projection: http_server_requests_seconds_count.uri, http_server_requests_seconds_count.kubernetes_namespace, http_server_requests_seconds_count.kubernetes_pod_name, http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_count.greptime_value\
12245            \n              Filter: prom_assert_unique_match_group(__promql_match_group_count, http_server_requests_seconds_count.uri)\
12246            \n                WindowAggr: windowExpr=[[count(Int64(1)) PARTITION BY [http_server_requests_seconds_count.uri, http_server_requests_seconds_count.greptime_timestamp] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS __promql_match_group_count]]\
12247            \n                  PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\
12248            \n                    PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\
12249            \n                      Sort: http_server_requests_seconds_count.uri ASC NULLS FIRST, http_server_requests_seconds_count.kubernetes_namespace ASC NULLS FIRST, http_server_requests_seconds_count.kubernetes_pod_name ASC NULLS FIRST, http_server_requests_seconds_count.greptime_timestamp ASC NULLS FIRST\
12250            \n                        Filter: http_server_requests_seconds_count.uri = Utf8(\"/accounts/login\") AND http_server_requests_seconds_count.greptime_timestamp >= TimestampMillisecond(-999, None) AND http_server_requests_seconds_count.greptime_timestamp <= TimestampMillisecond(100000000, None)\
12251            \n                          TableScan: http_server_requests_seconds_count";
12252        assert_eq!(plan.to_string(), expected);
12253    }
12254
12255    #[tokio::test]
12256    async fn test_nested_histogram_quantile() {
12257        let mut eval_stmt = EvalStmt {
12258            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12259            start: UNIX_EPOCH,
12260            end: UNIX_EPOCH
12261                .checked_add(Duration::from_secs(100_000))
12262                .unwrap(),
12263            interval: Duration::from_secs(5),
12264            lookback_delta: Duration::from_secs(1),
12265        };
12266
12267        let case = r#"label_replace(histogram_quantile(0.99, sum by(pod, le, path, code) (rate(greptime_servers_grpc_requests_elapsed_bucket{container="frontend"}[1m0s]))), "pod_new", "$1", "pod", "greptimedb-frontend-[0-9a-z]*-(.*)")"#;
12268
12269        let prom_expr = parser::parse(case).unwrap();
12270        eval_stmt.expr = prom_expr;
12271        let table_provider = build_test_table_provider_with_fields(
12272            &[(
12273                DEFAULT_SCHEMA_NAME.to_string(),
12274                "greptime_servers_grpc_requests_elapsed_bucket".to_string(),
12275            )],
12276            &["pod", "le", "path", "code", "container"],
12277        )
12278        .await;
12279        // Should be ok
12280        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12281            .await
12282            .unwrap();
12283    }
12284
12285    #[tokio::test]
12286    async fn test_histogram_quantile_binary_op() {
12287        let mut eval_stmt = EvalStmt {
12288            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12289            start: UNIX_EPOCH,
12290            end: UNIX_EPOCH
12291                .checked_add(Duration::from_secs(100_000))
12292                .unwrap(),
12293            interval: Duration::from_secs(5),
12294            lookback_delta: Duration::from_secs(1),
12295        };
12296
12297        // Arithmetic applied to a histogram_quantile() result. Regression for #8144:
12298        // HistogramFold used to drop the input column qualifiers, so the binary-op
12299        // projection failed to resolve the qualified tag column.
12300        let case = r#"histogram_quantile(0.5, sum by (le, pod) (rate(http_request_duration_seconds_bucket[5m]))) + 0"#;
12301
12302        let prom_expr = parser::parse(case).unwrap();
12303        eval_stmt.expr = prom_expr;
12304        let table_provider = build_test_table_provider_with_fields(
12305            &[(
12306                DEFAULT_SCHEMA_NAME.to_string(),
12307                "http_request_duration_seconds_bucket".to_string(),
12308            )],
12309            &["pod", "le"],
12310        )
12311        .await;
12312        // Should plan without a "No field named ..." error.
12313        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12314            .await
12315            .unwrap();
12316    }
12317
12318    #[tokio::test]
12319    async fn test_parse_and_operator() {
12320        let mut eval_stmt = EvalStmt {
12321            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12322            start: UNIX_EPOCH,
12323            end: UNIX_EPOCH
12324                .checked_add(Duration::from_secs(100_000))
12325                .unwrap(),
12326            interval: Duration::from_secs(5),
12327            lookback_delta: Duration::from_secs(1),
12328        };
12329
12330        let cases = [
12331            r#"count (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} ) and (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} )) / (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_capacity_bytes{namespace=~".+"} )) >= (80 / 100)) or vector (0)"#,
12332            r#"count (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} ) unless (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} )) / (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_capacity_bytes{namespace=~".+"} )) >= (80 / 100)) or vector (0)"#,
12333        ];
12334
12335        for case in cases {
12336            let prom_expr = parser::parse(case).unwrap();
12337            eval_stmt.expr = prom_expr;
12338            let table_provider = build_test_table_provider_with_fields(
12339                &[
12340                    (
12341                        DEFAULT_SCHEMA_NAME.to_string(),
12342                        "kubelet_volume_stats_used_bytes".to_string(),
12343                    ),
12344                    (
12345                        DEFAULT_SCHEMA_NAME.to_string(),
12346                        "kubelet_volume_stats_capacity_bytes".to_string(),
12347                    ),
12348                ],
12349                &["namespace", "persistentvolumeclaim"],
12350            )
12351            .await;
12352            // Should be ok
12353            let _ =
12354                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12355                    .await
12356                    .unwrap();
12357        }
12358    }
12359
12360    #[tokio::test]
12361    async fn test_nested_binary_op() {
12362        let mut eval_stmt = EvalStmt {
12363            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12364            start: UNIX_EPOCH,
12365            end: UNIX_EPOCH
12366                .checked_add(Duration::from_secs(100_000))
12367                .unwrap(),
12368            interval: Duration::from_secs(5),
12369            lookback_delta: Duration::from_secs(1),
12370        };
12371
12372        let case = r#"sum(rate(nginx_ingress_controller_requests{job=~".*"}[2m])) -
12373        (
12374            sum(rate(nginx_ingress_controller_requests{namespace=~".*"}[2m]))
12375            or
12376            vector(0)
12377        )"#;
12378
12379        let prom_expr = parser::parse(case).unwrap();
12380        eval_stmt.expr = prom_expr;
12381        let table_provider = build_test_table_provider_with_fields(
12382            &[(
12383                DEFAULT_SCHEMA_NAME.to_string(),
12384                "nginx_ingress_controller_requests".to_string(),
12385            )],
12386            &["namespace", "job"],
12387        )
12388        .await;
12389        // Should be ok
12390        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12391            .await
12392            .unwrap();
12393    }
12394
12395    #[tokio::test]
12396    async fn test_parse_or_operator() {
12397        let mut eval_stmt = EvalStmt {
12398            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12399            start: UNIX_EPOCH,
12400            end: UNIX_EPOCH
12401                .checked_add(Duration::from_secs(100_000))
12402                .unwrap(),
12403            interval: Duration::from_secs(5),
12404            lookback_delta: Duration::from_secs(1),
12405        };
12406
12407        let case = r#"
12408        sum(rate(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}[120s])) by (cluster_name,tenant_name) /
12409        (sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) * 100)
12410            or
12411        200 * sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) /
12412        sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)"#;
12413
12414        let table_provider = build_test_table_provider_with_fields(
12415            &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())],
12416            &["tenant_name", "cluster_name"],
12417        )
12418        .await;
12419        eval_stmt.expr = parser::parse(case).unwrap();
12420        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12421            .await
12422            .unwrap();
12423
12424        let case = r#"sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
12425            (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) +
12426            sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
12427            (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0
12428            or
12429            sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
12430            (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0
12431            or
12432            sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
12433            (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0"#;
12434        let table_provider = build_test_table_provider_with_fields(
12435            &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())],
12436            &["tenant_name", "cluster_name"],
12437        )
12438        .await;
12439        eval_stmt.expr = parser::parse(case).unwrap();
12440        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12441            .await
12442            .unwrap();
12443
12444        let case = r#"(sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) +
12445            sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or
12446            (sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or
12447            (sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name))"#;
12448        let table_provider = build_test_table_provider_with_fields(
12449            &[
12450                (
12451                    DEFAULT_SCHEMA_NAME.to_string(),
12452                    "background_waitevent_cnt".to_string(),
12453                ),
12454                (
12455                    DEFAULT_SCHEMA_NAME.to_string(),
12456                    "foreground_waitevent_cnt".to_string(),
12457                ),
12458            ],
12459            &["tenant_name", "cluster_name"],
12460        )
12461        .await;
12462        eval_stmt.expr = parser::parse(case).unwrap();
12463        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12464            .await
12465            .unwrap();
12466
12467        let case = r#"avg(node_load1{cluster_name=~"cluster1"}) by (cluster_name,host_name) or max(container_cpu_load_average_10s{cluster_name=~"cluster1"}) by (cluster_name,host_name) * 100 / max(container_spec_cpu_quota{cluster_name=~"cluster1"}) by (cluster_name,host_name)"#;
12468        let table_provider = build_test_table_provider_with_fields(
12469            &[
12470                (DEFAULT_SCHEMA_NAME.to_string(), "node_load1".to_string()),
12471                (
12472                    DEFAULT_SCHEMA_NAME.to_string(),
12473                    "container_cpu_load_average_10s".to_string(),
12474                ),
12475                (
12476                    DEFAULT_SCHEMA_NAME.to_string(),
12477                    "container_spec_cpu_quota".to_string(),
12478                ),
12479            ],
12480            &["cluster_name", "host_name"],
12481        )
12482        .await;
12483        eval_stmt.expr = parser::parse(case).unwrap();
12484        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12485            .await
12486            .unwrap();
12487    }
12488
12489    #[tokio::test]
12490    async fn value_matcher() {
12491        // template
12492        let mut eval_stmt = EvalStmt {
12493            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12494            start: UNIX_EPOCH,
12495            end: UNIX_EPOCH
12496                .checked_add(Duration::from_secs(100_000))
12497                .unwrap(),
12498            interval: Duration::from_secs(5),
12499            lookback_delta: Duration::from_secs(1),
12500        };
12501
12502        let cases = [
12503            // single equal matcher
12504            (
12505                r#"some_metric{__field__="field_1"}"#,
12506                vec![
12507                    "some_metric.field_1",
12508                    "some_metric.tag_0",
12509                    "some_metric.tag_1",
12510                    "some_metric.tag_2",
12511                    "some_metric.timestamp",
12512                ],
12513            ),
12514            // two equal matchers
12515            (
12516                r#"some_metric{__field__="field_1", __field__="field_0"}"#,
12517                vec![
12518                    "some_metric.field_0",
12519                    "some_metric.field_1",
12520                    "some_metric.tag_0",
12521                    "some_metric.tag_1",
12522                    "some_metric.tag_2",
12523                    "some_metric.timestamp",
12524                ],
12525            ),
12526            // single not_eq matcher
12527            (
12528                r#"some_metric{__field__!="field_1"}"#,
12529                vec![
12530                    "some_metric.field_0",
12531                    "some_metric.field_2",
12532                    "some_metric.tag_0",
12533                    "some_metric.tag_1",
12534                    "some_metric.tag_2",
12535                    "some_metric.timestamp",
12536                ],
12537            ),
12538            // two not_eq matchers
12539            (
12540                r#"some_metric{__field__!="field_1", __field__!="field_2"}"#,
12541                vec![
12542                    "some_metric.field_0",
12543                    "some_metric.tag_0",
12544                    "some_metric.tag_1",
12545                    "some_metric.tag_2",
12546                    "some_metric.timestamp",
12547                ],
12548            ),
12549            // equal and not_eq matchers (no conflict)
12550            (
12551                r#"some_metric{__field__="field_1", __field__!="field_0"}"#,
12552                vec![
12553                    "some_metric.field_1",
12554                    "some_metric.tag_0",
12555                    "some_metric.tag_1",
12556                    "some_metric.tag_2",
12557                    "some_metric.timestamp",
12558                ],
12559            ),
12560            // equal and not_eq matchers (conflict)
12561            (
12562                r#"some_metric{__field__="field_2", __field__!="field_2"}"#,
12563                vec![
12564                    "some_metric.tag_0",
12565                    "some_metric.tag_1",
12566                    "some_metric.tag_2",
12567                    "some_metric.timestamp",
12568                ],
12569            ),
12570            // single regex eq matcher
12571            (
12572                r#"some_metric{__field__=~"field_1|field_2"}"#,
12573                vec![
12574                    "some_metric.field_1",
12575                    "some_metric.field_2",
12576                    "some_metric.tag_0",
12577                    "some_metric.tag_1",
12578                    "some_metric.tag_2",
12579                    "some_metric.timestamp",
12580                ],
12581            ),
12582            // single regex not_eq matcher
12583            (
12584                r#"some_metric{__field__!~"field_1|field_2"}"#,
12585                vec![
12586                    "some_metric.field_0",
12587                    "some_metric.tag_0",
12588                    "some_metric.tag_1",
12589                    "some_metric.tag_2",
12590                    "some_metric.timestamp",
12591                ],
12592            ),
12593        ];
12594
12595        for case in cases {
12596            let prom_expr = parser::parse(case.0).unwrap();
12597            eval_stmt.expr = prom_expr;
12598            let table_provider = build_test_table_provider(
12599                &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12600                3,
12601                3,
12602            )
12603            .await;
12604            let plan =
12605                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12606                    .await
12607                    .unwrap();
12608            let mut fields = plan.schema().field_names();
12609            let mut expected = case.1.into_iter().map(String::from).collect::<Vec<_>>();
12610            fields.sort();
12611            expected.sort();
12612            assert_eq!(fields, expected, "case: {:?}", case.0);
12613        }
12614
12615        let bad_cases = [
12616            r#"some_metric{__field__="nonexistent"}"#,
12617            r#"some_metric{__field__!="nonexistent"}"#,
12618        ];
12619
12620        for case in bad_cases {
12621            let prom_expr = parser::parse(case).unwrap();
12622            eval_stmt.expr = prom_expr;
12623            let table_provider = build_test_table_provider(
12624                &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12625                3,
12626                3,
12627            )
12628            .await;
12629            let plan =
12630                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12631                    .await;
12632            assert!(plan.is_err(), "case: {:?}", case);
12633        }
12634    }
12635
12636    #[tokio::test]
12637    async fn custom_schema() {
12638        let query = "some_alt_metric{__schema__=\"greptime_private\"}";
12639        let expected = String::from(
12640            "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12641            \n  PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12642            \n    Sort: greptime_private.some_alt_metric.tag_0 ASC NULLS FIRST, greptime_private.some_alt_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12643            \n      Filter: greptime_private.some_alt_metric.timestamp >= TimestampMillisecond(-999, None) AND greptime_private.some_alt_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12644            \n        TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
12645        );
12646
12647        indie_query_plan_compare(query, expected).await;
12648
12649        let query = "some_alt_metric{__database__=\"greptime_private\"}";
12650        let expected = String::from(
12651            "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12652            \n  PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12653            \n    Sort: greptime_private.some_alt_metric.tag_0 ASC NULLS FIRST, greptime_private.some_alt_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12654            \n      Filter: greptime_private.some_alt_metric.timestamp >= TimestampMillisecond(-999, None) AND greptime_private.some_alt_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12655            \n        TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
12656        );
12657
12658        indie_query_plan_compare(query, expected).await;
12659
12660        let query = "some_alt_metric{__schema__=\"greptime_private\"} / some_metric";
12661        let expected = String::from(
12662            "Projection: some_metric.tag_0, some_metric.timestamp, CAST(greptime_private.some_alt_metric.field_0 AS Float64) / CAST(some_metric.field_0 AS Float64) AS greptime_private.some_alt_metric.field_0 / some_metric.field_0 [tag_0:Utf8, timestamp:Timestamp(ms), greptime_private.some_alt_metric.field_0 / some_metric.field_0:Float64;N]\
12663            \n  Inner Join: greptime_private.some_alt_metric.tag_0 = some_metric.tag_0, greptime_private.some_alt_metric.timestamp = some_metric.timestamp [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12664            \n    SubqueryAlias: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12665            \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12666            \n        PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12667            \n          Sort: greptime_private.some_alt_metric.tag_0 ASC NULLS FIRST, greptime_private.some_alt_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12668            \n            Filter: greptime_private.some_alt_metric.timestamp >= TimestampMillisecond(-999, None) AND greptime_private.some_alt_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12669            \n              TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12670            \n    SubqueryAlias: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12671            \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12672            \n        PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12673            \n          Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12674            \n            Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12675            \n              TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
12676        );
12677
12678        indie_query_plan_compare(query, expected).await;
12679    }
12680
12681    #[tokio::test]
12682    async fn only_equals_is_supported_for_special_matcher() {
12683        let queries = &[
12684            "some_alt_metric{__schema__!=\"greptime_private\"}",
12685            "some_alt_metric{__schema__=~\"lalala\"}",
12686            "some_alt_metric{__database__!=\"greptime_private\"}",
12687            "some_alt_metric{__database__=~\"lalala\"}",
12688        ];
12689
12690        for query in queries {
12691            let prom_expr = parser::parse(query).unwrap();
12692            let eval_stmt = EvalStmt {
12693                expr: prom_expr,
12694                start: UNIX_EPOCH,
12695                end: UNIX_EPOCH
12696                    .checked_add(Duration::from_secs(100_000))
12697                    .unwrap(),
12698                interval: Duration::from_secs(5),
12699                lookback_delta: Duration::from_secs(1),
12700            };
12701
12702            let table_provider = build_test_table_provider(
12703                &[
12704                    (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
12705                    (
12706                        "greptime_private".to_string(),
12707                        "some_alt_metric".to_string(),
12708                    ),
12709                ],
12710                1,
12711                1,
12712            )
12713            .await;
12714
12715            let plan =
12716                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12717                    .await;
12718            assert!(plan.is_err(), "query: {:?}", query);
12719        }
12720    }
12721
12722    #[tokio::test]
12723    async fn native_scan_bounds_preserve_zero_lookback_and_overflow() {
12724        let table_provider = build_test_table_provider(
12725            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12726            1,
12727            1,
12728        )
12729        .await;
12730        let mut planner = PromPlanner {
12731            table_provider,
12732            ctx: PromPlannerContext::from_eval_stmt(&build_eval_stmt("some_metric")),
12733            promql_annotations: None,
12734        };
12735        planner.ctx.time_index_column = Some("timestamp".to_string());
12736        planner.ctx.start = 1_000;
12737        planner.ctx.lookback_delta = 0;
12738        let schema = Arc::new(
12739            DFSchema::try_from(ArrowSchema::new(vec![Field::new(
12740                "timestamp",
12741                ArrowDataType::Timestamp(ArrowTimeUnit::Nanosecond, None),
12742                false,
12743            )]))
12744            .unwrap(),
12745        );
12746        for (end, interval, windows) in [
12747            (1_000, 1_000, 1),
12748            (2_000, 1_000, 1),
12749            (7_201_000, 7_200_000, 2),
12750        ] {
12751            planner.ctx.end = end;
12752            planner.ctx.interval = interval;
12753            let filter = planner
12754                .build_time_index_filter(0, &schema)
12755                .unwrap()
12756                .unwrap()
12757                .to_string();
12758            assert_eq!(filter.matches(">=").count(), windows, "{filter}");
12759            assert!(
12760                filter.contains("TimestampNanosecond(1000000000, None)"),
12761                "{filter}"
12762            );
12763        }
12764        planner.ctx.end = i64::MAX;
12765        let filter = planner
12766            .build_time_index_filter(0, &schema)
12767            .unwrap()
12768            .unwrap()
12769            .to_string();
12770        assert!(
12771            filter.contains("timestamp >= TimestampNanosecond(1000000000, None)"),
12772            "{filter}"
12773        );
12774
12775        // A lookback subtraction can underflow milliseconds while the upper bound remains
12776        // representable. Keep that upper bound so LastRow cannot select a future sample.
12777        let ms_schema = Arc::new(
12778            DFSchema::try_from(ArrowSchema::new(vec![Field::new(
12779                "timestamp",
12780                ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
12781                false,
12782            )]))
12783            .unwrap(),
12784        );
12785        planner.ctx.start = i64::MIN + 100;
12786        planner.ctx.end = planner.ctx.start;
12787        planner.ctx.lookback_delta = 200;
12788        let filter = planner
12789            .build_time_index_filter(0, &ms_schema)
12790            .unwrap()
12791            .unwrap()
12792            .to_string();
12793        assert_eq!(
12794            filter,
12795            format!(
12796                "timestamp <= TimestampMillisecond({}, None)",
12797                i64::MIN + 100
12798            )
12799        );
12800
12801        // The lower bound can also overflow while converting milliseconds to native nanoseconds.
12802        // Its representable upper bound still has to reach the scan.
12803        planner.ctx.start = 0;
12804        planner.ctx.end = 0;
12805        planner.ctx.lookback_delta = 300_000;
12806        let filter = planner
12807            .build_time_index_filter(9_223_372_036_854, &schema)
12808            .unwrap()
12809            .unwrap()
12810            .to_string();
12811        assert_eq!(
12812            filter,
12813            "timestamp <= TimestampNanosecond(-9223372036854000000, None)"
12814        );
12815    }
12816
12817    #[tokio::test]
12818    async fn test_non_ms_precision() {
12819        let catalog_list = MemoryCatalogManager::with_default_setup();
12820        let columns = vec![
12821            ColumnSchema::new(
12822                "tag".to_string(),
12823                ConcreteDataType::string_datatype(),
12824                false,
12825            ),
12826            ColumnSchema::new(
12827                "timestamp".to_string(),
12828                ConcreteDataType::timestamp_nanosecond_datatype(),
12829                false,
12830            )
12831            .with_time_index(true),
12832            ColumnSchema::new(
12833                "field".to_string(),
12834                ConcreteDataType::float64_datatype(),
12835                true,
12836            ),
12837        ];
12838        let schema = Arc::new(Schema::new(columns));
12839        let table_meta = TableMetaBuilder::empty()
12840            .schema(schema)
12841            .primary_key_indices(vec![0])
12842            .value_indices(vec![2])
12843            .next_column_id(1024)
12844            .build()
12845            .unwrap();
12846        let table_info = TableInfoBuilder::default()
12847            .name("metrics".to_string())
12848            .meta(table_meta)
12849            .build()
12850            .unwrap();
12851        let table = EmptyTable::from_table_info(&table_info);
12852        assert!(
12853            catalog_list
12854                .register_table_sync(RegisterTableRequest {
12855                    catalog: DEFAULT_CATALOG_NAME.to_string(),
12856                    schema: DEFAULT_SCHEMA_NAME.to_string(),
12857                    table_name: "metrics".to_string(),
12858                    table_id: 1024,
12859                    table,
12860                })
12861                .is_ok()
12862        );
12863
12864        let plan = PromPlanner::stmt_to_plan(
12865            DfTableSourceProvider::new(
12866                catalog_list.clone(),
12867                false,
12868                QueryContext::arc(),
12869                DummyDecoder::arc(),
12870                true,
12871            ),
12872            &EvalStmt {
12873                expr: parser::parse("metrics{tag = \"1\"}").unwrap(),
12874                start: UNIX_EPOCH,
12875                end: UNIX_EPOCH
12876                    .checked_add(Duration::from_secs(100_000))
12877                    .unwrap(),
12878                interval: Duration::from_secs(5),
12879                lookback_delta: Duration::from_secs(1),
12880            },
12881            &build_query_engine_state(),
12882        )
12883        .await
12884        .unwrap();
12885        assert_eq!(
12886            plan.display_indent_schema().to_string(),
12887            "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\n  PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n    Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n      Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp > TimestampNanosecond(-1000000000, None) AND metrics.timestamp <= TimestampNanosecond(100000000000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n        Projection: metrics.field, metrics.tag, metrics.timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n          TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]"
12888        );
12889        let plan = PromPlanner::stmt_to_plan(
12890            DfTableSourceProvider::new(
12891                catalog_list.clone(),
12892                false,
12893                QueryContext::arc(),
12894                DummyDecoder::arc(),
12895                true,
12896            ),
12897            &EvalStmt {
12898                expr: parser::parse("avg_over_time(metrics{tag = \"1\"}[5s])").unwrap(),
12899                start: UNIX_EPOCH,
12900                end: UNIX_EPOCH
12901                    .checked_add(Duration::from_secs(100_000))
12902                    .unwrap(),
12903                interval: Duration::from_secs(5),
12904                lookback_delta: Duration::from_secs(1),
12905            },
12906            &build_query_engine_state(),
12907        )
12908        .await
12909        .unwrap();
12910        assert_eq!(
12911            plan.display_indent_schema().to_string(),
12912            "Filter: prom_avg_over_time(timestamp_range,field) IS NOT NULL [timestamp:Timestamp(ms), prom_avg_over_time(timestamp_range,field):Float64;N, tag:Utf8]\n  Projection: metrics.timestamp, prom_avg_over_time(timestamp_range, field) AS prom_avg_over_time(timestamp_range,field), metrics.tag [timestamp:Timestamp(ms), prom_avg_over_time(timestamp_range,field):Float64;N, tag:Utf8]\n    PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[5000], time index=[timestamp], values=[\"field\"] [field:Dictionary(Int64, Float64);N, tag:Utf8, timestamp:Timestamp(ms), timestamp_range:Dictionary(Int64, Timestamp(ms))]\n      PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n        PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n          Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n            Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp > TimestampNanosecond(-5000000000, None) AND metrics.timestamp <= TimestampNanosecond(100000000000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n              Projection: metrics.field, metrics.tag, metrics.timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n                TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]"
12913        );
12914    }
12915
12916    #[tokio::test]
12917    async fn test_nonexistent_label() {
12918        // template
12919        let mut eval_stmt = EvalStmt {
12920            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12921            start: UNIX_EPOCH,
12922            end: UNIX_EPOCH
12923                .checked_add(Duration::from_secs(100_000))
12924                .unwrap(),
12925            interval: Duration::from_secs(5),
12926            lookback_delta: Duration::from_secs(1),
12927        };
12928
12929        let case = r#"some_metric{nonexistent="hi"}"#;
12930        let prom_expr = parser::parse(case).unwrap();
12931        eval_stmt.expr = prom_expr;
12932        let table_provider = build_test_table_provider(
12933            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12934            3,
12935            3,
12936        )
12937        .await;
12938        // Should be ok
12939        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12940            .await
12941            .unwrap();
12942    }
12943
12944    #[tokio::test]
12945    async fn test_label_join() {
12946        let prom_expr = parser::parse(
12947            "label_join(up{tag_0='api-server'}, 'foo', ',', 'tag_1', 'tag_2', 'tag_3')",
12948        )
12949        .unwrap();
12950        let eval_stmt = EvalStmt {
12951            expr: prom_expr,
12952            start: UNIX_EPOCH,
12953            end: UNIX_EPOCH
12954                .checked_add(Duration::from_secs(100_000))
12955                .unwrap(),
12956            interval: Duration::from_secs(5),
12957            lookback_delta: Duration::from_secs(1),
12958        };
12959
12960        let table_provider =
12961            build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 4, 1)
12962                .await;
12963        let plan =
12964            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12965                .await
12966                .unwrap();
12967
12968        let expected = r#"
12969Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8]
12970  Projection: up.timestamp, up.field_0, concat_ws(Utf8(","), up.tag_1, up.tag_2, up.tag_3) AS foo, up.tag_0, up.tag_1, up.tag_2, up.tag_3 [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8]
12971    PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
12972      PromSeriesDivide: tags=["tag_0", "tag_1", "tag_2", "tag_3"] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
12973        Sort: up.tag_0 ASC NULLS FIRST, up.tag_1 ASC NULLS FIRST, up.tag_2 ASC NULLS FIRST, up.tag_3 ASC NULLS FIRST, up.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
12974          Filter: up.tag_0 = Utf8("api-server") AND up.timestamp >= TimestampMillisecond(-999, None) AND up.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
12975            TableScan: up [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#;
12976
12977        let ret = plan.display_indent_schema().to_string();
12978        assert_eq!(format!("\n{ret}"), expected, "\n{}", ret);
12979    }
12980
12981    #[tokio::test]
12982    async fn test_label_replace() {
12983        let prom_expr = parser::parse(
12984            "label_replace(up{tag_0=\"a:c\"}, \"foo\", \"$1\", \"tag_0\", \"(.*):.*\")",
12985        )
12986        .unwrap();
12987        let eval_stmt = EvalStmt {
12988            expr: prom_expr,
12989            start: UNIX_EPOCH,
12990            end: UNIX_EPOCH
12991                .checked_add(Duration::from_secs(100_000))
12992                .unwrap(),
12993            interval: Duration::from_secs(5),
12994            lookback_delta: Duration::from_secs(1),
12995        };
12996
12997        let table_provider =
12998            build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 1, 1)
12999                .await;
13000        let plan =
13001            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13002                .await
13003                .unwrap();
13004
13005        let expected = r#"
13006Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8]
13007  Projection: up.timestamp, up.field_0, regexp_replace(up.tag_0, Utf8("^(?s:(.*):.*)$"), Utf8("$1")) AS foo, up.tag_0 [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8]
13008    PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
13009      PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
13010        Sort: up.tag_0 ASC NULLS FIRST, up.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
13011          Filter: up.tag_0 = Utf8("a:c") AND up.timestamp >= TimestampMillisecond(-999, None) AND up.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
13012            TableScan: up [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#;
13013
13014        let ret = plan.display_indent_schema().to_string();
13015        assert_eq!(format!("\n{ret}"), expected, "\n{}", ret);
13016    }
13017
13018    #[tokio::test]
13019    async fn label_replace_aggregation_queries_plan_successfully() {
13020        let aggregate =
13021            r#"sum by (foo) (label_replace(some_metric, "foo", "$1", "tag_0", "(.*)"))"#;
13022        let queries = [
13023            aggregate.to_string(),
13024            format!("{aggregate} <= 10"),
13025            format!("{aggregate} * 0.8"),
13026            format!("0.8 * {aggregate}"),
13027            format!("{aggregate} <= {aggregate} * 0.8"),
13028        ];
13029        let state = build_query_engine_state();
13030        let mut failures = Vec::new();
13031
13032        for query in queries {
13033            let table_provider = build_test_table_provider(
13034                &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
13035                1,
13036                1,
13037            )
13038            .await;
13039            if let Err(error) =
13040                PromPlanner::stmt_to_plan(table_provider, &build_eval_stmt(&query), &state).await
13041            {
13042                failures.push(format!("{query}: {error:?}"));
13043            }
13044        }
13045
13046        assert!(failures.is_empty(), "{}", failures.join("\n"));
13047    }
13048
13049    #[tokio::test]
13050    async fn test_matchers_to_expr() {
13051        let mut eval_stmt = EvalStmt {
13052            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
13053            start: UNIX_EPOCH,
13054            end: UNIX_EPOCH
13055                .checked_add(Duration::from_secs(100_000))
13056                .unwrap(),
13057            interval: Duration::from_secs(5),
13058            lookback_delta: Duration::from_secs(1),
13059        };
13060        let case =
13061            r#"sum(prometheus_tsdb_head_series{tag_1=~"(10.0.160.237:8080|10.0.160.237:9090)"})"#;
13062
13063        let prom_expr = parser::parse(case).unwrap();
13064        eval_stmt.expr = prom_expr;
13065        let table_provider = build_test_table_provider(
13066            &[(
13067                DEFAULT_SCHEMA_NAME.to_string(),
13068                "prometheus_tsdb_head_series".to_string(),
13069            )],
13070            3,
13071            3,
13072        )
13073        .await;
13074        let plan =
13075            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13076                .await
13077                .unwrap();
13078        let expected = "Sort: prometheus_tsdb_head_series.timestamp ASC NULLS LAST [timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.field_0):Float64;N, sum(prometheus_tsdb_head_series.field_1):Float64;N, sum(prometheus_tsdb_head_series.field_2):Float64;N]\
13079        \n  Aggregate: groupBy=[[prometheus_tsdb_head_series.timestamp]], aggr=[[sum(prometheus_tsdb_head_series.field_0), sum(prometheus_tsdb_head_series.field_1), sum(prometheus_tsdb_head_series.field_2)]] [timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.field_0):Float64;N, sum(prometheus_tsdb_head_series.field_1):Float64;N, sum(prometheus_tsdb_head_series.field_2):Float64;N]\
13080        \n    PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\
13081        \n      PromSeriesDivide: tags=[\"tag_0\", \"tag_1\", \"tag_2\"] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\
13082        \n        Sort: prometheus_tsdb_head_series.tag_0 ASC NULLS FIRST, prometheus_tsdb_head_series.tag_1 ASC NULLS FIRST, prometheus_tsdb_head_series.tag_2 ASC NULLS FIRST, prometheus_tsdb_head_series.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\
13083        \n          Filter: prometheus_tsdb_head_series.tag_1 ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\
13084        \n            TableScan: prometheus_tsdb_head_series [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]";
13085        assert_eq!(plan.display_indent_schema().to_string(), expected);
13086    }
13087
13088    #[tokio::test]
13089    async fn test_topk_expr() {
13090        let mut eval_stmt = EvalStmt {
13091            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
13092            start: UNIX_EPOCH,
13093            end: UNIX_EPOCH
13094                .checked_add(Duration::from_secs(100_000))
13095                .unwrap(),
13096            interval: Duration::from_secs(5),
13097            lookback_delta: Duration::from_secs(1),
13098        };
13099        let case = r#"topk(10, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#;
13100
13101        let prom_expr = parser::parse(case).unwrap();
13102        eval_stmt.expr = prom_expr;
13103        let table_provider = build_test_table_provider_with_fields(
13104            &[
13105                (
13106                    DEFAULT_SCHEMA_NAME.to_string(),
13107                    "prometheus_tsdb_head_series".to_string(),
13108                ),
13109                (
13110                    DEFAULT_SCHEMA_NAME.to_string(),
13111                    "http_server_requests_seconds_count".to_string(),
13112                ),
13113            ],
13114            &["ip"],
13115        )
13116        .await;
13117
13118        let plan =
13119            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13120                .await
13121                .unwrap();
13122        let expected = "Projection: sum(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp [sum(prometheus_tsdb_head_series.greptime_value):Float64;N, ip:Utf8, greptime_timestamp:Timestamp(ms)]\
13123        \n  Sort: prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\
13124        \n    Filter: row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= Float64(10) [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\
13125        \n      WindowAggr: windowExpr=[[row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\
13126        \n        Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\
13127        \n          Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[sum(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\
13128        \n            PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
13129        \n              PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
13130        \n                Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
13131        \n                  Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
13132        \n                    TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]";
13133
13134        assert_eq!(plan.display_indent_schema().to_string(), expected);
13135    }
13136
13137    #[tokio::test]
13138    async fn test_count_values_expr() {
13139        let mut eval_stmt = EvalStmt {
13140            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
13141            start: UNIX_EPOCH,
13142            end: UNIX_EPOCH
13143                .checked_add(Duration::from_secs(100_000))
13144                .unwrap(),
13145            interval: Duration::from_secs(5),
13146            lookback_delta: Duration::from_secs(1),
13147        };
13148        let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#;
13149
13150        let prom_expr = parser::parse(case).unwrap();
13151        eval_stmt.expr = prom_expr;
13152        let table_provider = build_test_table_provider_with_fields(
13153            &[
13154                (
13155                    DEFAULT_SCHEMA_NAME.to_string(),
13156                    "prometheus_tsdb_head_series".to_string(),
13157                ),
13158                (
13159                    DEFAULT_SCHEMA_NAME.to_string(),
13160                    "http_server_requests_seconds_count".to_string(),
13161                ),
13162            ],
13163            &["ip"],
13164        )
13165        .await;
13166
13167        let plan =
13168            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13169                .await
13170                .unwrap();
13171        let expected = "Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N]\
13172        \n  Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value) AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N]\
13173        \n    Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value)]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), prom_float_to_string(prometheus_tsdb_head_series.greptime_value):Utf8;N, count(prometheus_tsdb_head_series.greptime_value):Int64]\
13174        \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
13175        \n        PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
13176        \n          Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
13177        \n            Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
13178        \n              TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]";
13179
13180        assert_eq!(plan.display_indent_schema().to_string(), expected);
13181    }
13182
13183    #[tokio::test]
13184    async fn test_value_alias() {
13185        let mut eval_stmt = EvalStmt {
13186            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
13187            start: UNIX_EPOCH,
13188            end: UNIX_EPOCH
13189                .checked_add(Duration::from_secs(100_000))
13190                .unwrap(),
13191            interval: Duration::from_secs(5),
13192            lookback_delta: Duration::from_secs(1),
13193        };
13194        let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#;
13195
13196        let prom_expr = parser::parse(case).unwrap();
13197        eval_stmt.expr = prom_expr;
13198        eval_stmt = QueryLanguageParser::apply_alias_extension(eval_stmt, "my_series");
13199        let table_provider = build_test_table_provider_with_fields(
13200            &[
13201                (
13202                    DEFAULT_SCHEMA_NAME.to_string(),
13203                    "prometheus_tsdb_head_series".to_string(),
13204                ),
13205                (
13206                    DEFAULT_SCHEMA_NAME.to_string(),
13207                    "http_server_requests_seconds_count".to_string(),
13208                ),
13209            ],
13210            &["ip"],
13211        )
13212        .await;
13213
13214        let plan =
13215            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13216                .await
13217                .unwrap();
13218        let expected = r#"
13219Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.series, prometheus_tsdb_head_series.greptime_timestamp [my_series:Int64, ip:Utf8, series:Utf8;N, greptime_timestamp:Timestamp(ms)]
13220  Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N]
13221    Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value) AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N]
13222      Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value)]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), prom_float_to_string(prometheus_tsdb_head_series.greptime_value):Utf8;N, count(prometheus_tsdb_head_series.greptime_value):Int64]
13223        PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
13224          PromSeriesDivide: tags=["ip"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
13225            Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
13226              Filter: prometheus_tsdb_head_series.ip ~ Utf8("^(?:(10.0.160.237:8080|10.0.160.237:9090))$") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
13227                TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"#;
13228        assert_eq!(format!("\n{}", plan.display_indent_schema()), expected);
13229    }
13230
13231    #[tokio::test]
13232    async fn test_quantile_expr() {
13233        let mut eval_stmt = EvalStmt {
13234            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
13235            start: UNIX_EPOCH,
13236            end: UNIX_EPOCH
13237                .checked_add(Duration::from_secs(100_000))
13238                .unwrap(),
13239            interval: Duration::from_secs(5),
13240            lookback_delta: Duration::from_secs(1),
13241        };
13242        let case = r#"quantile(0.3, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#;
13243
13244        let prom_expr = parser::parse(case).unwrap();
13245        eval_stmt.expr = prom_expr;
13246        let table_provider = build_test_table_provider_with_fields(
13247            &[
13248                (
13249                    DEFAULT_SCHEMA_NAME.to_string(),
13250                    "prometheus_tsdb_head_series".to_string(),
13251                ),
13252                (
13253                    DEFAULT_SCHEMA_NAME.to_string(),
13254                    "http_server_requests_seconds_count".to_string(),
13255                ),
13256            ],
13257            &["ip"],
13258        )
13259        .await;
13260
13261        let plan =
13262            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13263                .await
13264                .unwrap();
13265        let expected = "Sort: prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [greptime_timestamp:Timestamp(ms), quantile(Float64(0.3),sum(prometheus_tsdb_head_series.greptime_value)):Float64;N]\
13266        \n  Aggregate: groupBy=[[prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[quantile(Float64(0.3), sum(prometheus_tsdb_head_series.greptime_value))]] [greptime_timestamp:Timestamp(ms), quantile(Float64(0.3),sum(prometheus_tsdb_head_series.greptime_value)):Float64;N]\
13267        \n    Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\
13268        \n      Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[sum(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\
13269        \n        PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
13270        \n          PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
13271        \n            Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
13272        \n              Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
13273        \n                TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]";
13274
13275        assert_eq!(plan.display_indent_schema().to_string(), expected);
13276    }
13277
13278    #[tokio::test]
13279    async fn test_or_not_exists_table_label() {
13280        let state = build_query_engine_state();
13281        let provider = build_test_table_provider_with_fields(
13282            &[(DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string())],
13283            &["job"],
13284        )
13285        .await;
13286        let raw = PromPlanner::stmt_to_plan(
13287            provider,
13288            &build_eval_stmt(r#"missing_metric or on(absent_label) normal_metric"#),
13289            &state,
13290        )
13291        .await
13292        .unwrap();
13293        assert!(
13294            raw.display_indent_schema()
13295                .to_string()
13296                .contains("__promql_or_match_0@")
13297        );
13298        let (optimized, batches) = execute(raw, &state).await;
13299        assert_no_internal_or_keys(optimized.schema());
13300        assert!(batches.iter().all(|batch| {
13301            batch
13302                .schema()
13303                .fields()
13304                .iter()
13305                .all(|field| !field.name().starts_with("__promql_or_match_"))
13306        }));
13307    }
13308
13309    #[tokio::test]
13310    async fn test_histogram_quantile_missing_le_column() {
13311        let mut eval_stmt = EvalStmt {
13312            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
13313            start: UNIX_EPOCH,
13314            end: UNIX_EPOCH
13315                .checked_add(Duration::from_secs(100_000))
13316                .unwrap(),
13317            interval: Duration::from_secs(5),
13318            lookback_delta: Duration::from_secs(1),
13319        };
13320
13321        // Test case: histogram_quantile with a table that doesn't have 'le' column
13322        let case = r#"histogram_quantile(0.99, sum by(pod,instance,le) (rate(non_existent_histogram_bucket{instance=~"xxx"}[1m])))"#;
13323
13324        let prom_expr = parser::parse(case).unwrap();
13325        eval_stmt.expr = prom_expr;
13326
13327        // Create a table provider with a table that doesn't have 'le' column
13328        let table_provider = build_test_table_provider_with_fields(
13329            &[(
13330                DEFAULT_SCHEMA_NAME.to_string(),
13331                "non_existent_histogram_bucket".to_string(),
13332            )],
13333            &["pod", "instance"], // Note: no 'le' column
13334        )
13335        .await;
13336
13337        // Should return empty result instead of error
13338        let result =
13339            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13340                .await;
13341
13342        // This should succeed now (returning empty result) instead of failing with "Cannot find column le"
13343        assert!(
13344            result.is_ok(),
13345            "Expected successful plan creation with empty result, but got error: {:?}",
13346            result.err()
13347        );
13348
13349        // Verify that the result is an EmptyRelation
13350        let plan = result.unwrap();
13351        match plan {
13352            LogicalPlan::EmptyRelation(_) => {
13353                // This is what we expect
13354            }
13355            _ => panic!("Expected EmptyRelation, but got: {:?}", plan),
13356        }
13357    }
13358
13359    #[tokio::test]
13360    async fn test_direct_or_normalizes_missing_match_labels() {
13361        type Case<'a> = (
13362            Option<Option<&'a str>>,
13363            Option<Option<&'a str>>,
13364            i64,
13365            i64,
13366            &'a [(f64, Option<&'a str>)],
13367        );
13368
13369        let modifier = or_modifier("lhs or on(k) rhs");
13370        #[rustfmt::skip]
13371        let cases: &[Case<'_>] = &[
13372            (None, None, 1, 1, &[(1.0, None)]),
13373            (None, Some(Some("")), 1, 1, &[(1.0, None)]),
13374            (Some(Some("")), None, 1, 1, &[(1.0, Some(""))]),
13375            (None, Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]),
13376            (Some(Some("l")), None, 1, 1, &[(1.0, Some("l")), (2.0, None)]),
13377            (Some(None), Some(Some("")), 1, 1, &[(1.0, None)]),
13378            (Some(None), Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]),
13379            (Some(Some("same")), Some(Some("same")), 1, 2, &[(1.0, Some("same")), (2.0, Some("same"))]),
13380        ];
13381        for &(left, right, left_ts, right_ts, expected) in cases {
13382            let (optimized, batches) = run(
13383                &matrix_source("lhs", left, left_ts, 1.0),
13384                &matrix_source("rhs", right, right_ts, 2.0),
13385                matrix_context("lhs", left),
13386                matrix_context("rhs", right),
13387                &modifier,
13388            )
13389            .await;
13390            assert_no_internal_or_keys(optimized.schema());
13391            assert_eq!(
13392                rows(&batches),
13393                expected
13394                    .iter()
13395                    .map(|(value, label)| (*value, label.map(str::to_string)))
13396                    .collect::<Vec<_>>()
13397            );
13398        }
13399    }
13400
13401    #[tokio::test]
13402    async fn test_direct_or_match_modifiers() {
13403        for (modifier, left, right, expected) in [
13404            (None, "left", "right", 2),
13405            (or_modifier("lhs or on(k) rhs"), "same", "same", 1),
13406            (or_modifier("lhs or on() rhs"), "left", "right", 1),
13407            (or_modifier("lhs or ignoring(k) rhs"), "left", "right", 1),
13408        ] {
13409            let (_, batches) = run(
13410                &matrix_source("lhs", Some(Some(left)), 1, 1.0),
13411                &matrix_source("rhs", Some(Some(right)), 1, 2.0),
13412                direct_or_context("lhs", &["job", "k"], "v"),
13413                direct_or_context("rhs", &["job", "k"], "v"),
13414                &modifier,
13415            )
13416            .await;
13417            assert_eq!(
13418                batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
13419                expected
13420            );
13421        }
13422    }
13423
13424    #[tokio::test]
13425    async fn test_direct_or_nested_projection_uses_left_context() {
13426        let left = matrix_source("lhs", Some(Some("k")), 1, 1.0);
13427        let right = matrix_source("rhs", Some(Some("k")), 1, 2.0);
13428        let raw = plan_direct_or(
13429            scan(&left),
13430            scan(&right),
13431            direct_or_context("lhs", &["job", "k"], "v"),
13432            direct_or_context("rhs", &["job", "k"], "v"),
13433            &or_modifier("lhs or on(k) rhs"),
13434        )
13435        .await;
13436        assert!(raw.schema().iter().any(|(qualifier, field)| {
13437            qualifier.as_ref().is_some_and(|q| q.to_string() == "lhs") && field.name() == "v"
13438        }));
13439        let nested = LogicalPlanBuilder::from(raw)
13440            .project(vec![
13441                DfExpr::BinaryExpr(BinaryExpr {
13442                    left: Box::new(DfExpr::Column(Column::new(
13443                        Some(TableReference::bare("lhs")),
13444                        "v",
13445                    ))),
13446                    op: Operator::Plus,
13447                    right: Box::new(lit(1.0)),
13448                })
13449                .alias("v_plus"),
13450            ])
13451            .unwrap()
13452            .build()
13453            .unwrap();
13454        let (_, batches) = execute(nested, &build_query_engine_state()).await;
13455        assert_eq!(values(&batches, "v_plus"), vec![2.0]);
13456    }
13457
13458    #[tokio::test]
13459    async fn test_direct_or_skips_user_internal_key_name() {
13460        const USER_TAG: &str = "__promql_or_match_0";
13461        let left = tagged_source(
13462            "lhs",
13463            false,
13464            (USER_TAG, Some("left")),
13465            DirectOrValue::Float64(1.0),
13466        );
13467        let right = tagged_source(
13468            "rhs",
13469            false,
13470            (USER_TAG, Some("right")),
13471            DirectOrValue::Float64(2.0),
13472        );
13473        let raw = plan_direct_or(
13474            scan(&left),
13475            scan(&right),
13476            direct_or_context("lhs", &["job", USER_TAG], "v"),
13477            direct_or_context("rhs", &["job", USER_TAG], "v"),
13478            &or_modifier("lhs or on(missing_label) rhs"),
13479        )
13480        .await;
13481        assert!(
13482            raw.display_indent_schema()
13483                .to_string()
13484                .contains("__promql_or_match_1@")
13485        );
13486        let (_, batches) = execute(raw, &build_query_engine_state()).await;
13487        assert!(
13488            batches
13489                .iter()
13490                .all(|batch| batch.column_by_name(USER_TAG).is_some())
13491        );
13492    }
13493
13494    #[tokio::test]
13495    async fn test_direct_or_substrait_round_trip_with_normalized_key() {
13496        let state = build_query_engine_state();
13497        let ctx = SessionContext::new_with_state(state.session_state());
13498        let catalog = Arc::new(MemoryCatalogProvider::new());
13499        catalog
13500            .register_schema("public", Arc::new(MemorySchemaProvider::new()))
13501            .unwrap();
13502        ctx.register_catalog("datafusion", catalog);
13503        let left = matrix_source("lhs", Some(Some("")), 1, 1.0);
13504        let right = matrix_source("rhs", None, 1, 2.0);
13505        ctx.register_table(
13506            TableReference::full("datafusion", "public", "lhs"),
13507            table(&left),
13508        )
13509        .unwrap();
13510        ctx.register_table(
13511            TableReference::full("datafusion", "public", "rhs"),
13512            table(&right),
13513        )
13514        .unwrap();
13515        let raw = plan_direct_or(
13516            ctx.table("datafusion.public.lhs")
13517                .await
13518                .unwrap()
13519                .into_unoptimized_plan(),
13520            ctx.table("datafusion.public.rhs")
13521                .await
13522                .unwrap()
13523                .into_unoptimized_plan(),
13524            direct_or_context("lhs", &["job", "k"], "v"),
13525            direct_or_context("rhs", &["job"], "v"),
13526            &or_modifier("lhs or on(k) rhs"),
13527        )
13528        .await;
13529        let decoded = DFLogicalSubstraitConvertor
13530            .decode(
13531                DFLogicalSubstraitConvertor
13532                    .encode(&raw, DefaultSerializer)
13533                    .unwrap(),
13534                ctx.state(),
13535            )
13536            .await
13537            .unwrap();
13538        let (optimized, batches) = execute(decoded, &state).await;
13539        assert_no_internal_or_keys(optimized.schema());
13540        assert!(batches.iter().all(|batch| {
13541            batch
13542                .schema()
13543                .fields()
13544                .iter()
13545                .all(|field| !field.name().starts_with("__promql_or_match_"))
13546        }));
13547        assert_eq!(values(&batches, "v"), vec![1.0]);
13548    }
13549
13550    #[tokio::test]
13551    async fn test_direct_or_numeric_value_types() {
13552        let left = tagged_source("lhs", true, ("k", Some("lhs")), DirectOrValue::Int64(0));
13553        let right = tagged_source(
13554            "rhs",
13555            false,
13556            ("k", Some("rhs")),
13557            DirectOrValue::Float64(0.5),
13558        );
13559        let (optimized, batches) = run(
13560            &left,
13561            &right,
13562            direct_or_context("lhs", &["job", "k"], "v"),
13563            direct_or_context("rhs", &["job", "k"], "v"),
13564            &or_modifier("lhs or on(k) rhs"),
13565        )
13566        .await;
13567        assert_eq!(
13568            optimized
13569                .schema()
13570                .field_with_name(None, "v")
13571                .unwrap()
13572                .data_type(),
13573            &ArrowDataType::Float64
13574        );
13575        assert_eq!(values(&batches, "v"), vec![0.5]);
13576        let provider = build_test_table_provider_with_fields(
13577            &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
13578            &[],
13579        )
13580        .await;
13581        let mut planner = PromPlanner {
13582            table_provider: provider,
13583            ctx: PromPlannerContext::default(),
13584            promql_annotations: None,
13585        };
13586        let left_context = direct_or_context("lhs", &["job"], "v");
13587        let right_context = direct_or_context("rhs", &["job"], "v");
13588        let error = planner
13589            .or_operator(
13590                scan(&job_source("lhs", DirectOrValue::Utf8("x"))),
13591                scan(&job_source("rhs", DirectOrValue::Float64(1.0))),
13592                left_context.tag_columns.iter().cloned().collect(),
13593                right_context.tag_columns.iter().cloned().collect(),
13594                left_context,
13595                right_context,
13596                &or_modifier("lhs or on() rhs"),
13597            )
13598            .unwrap_err();
13599        assert!(
13600            error
13601                .to_string()
13602                .contains("OR value fields have incompatible types")
13603        );
13604    }
13605
13606    #[tokio::test]
13607    async fn test_or_with_histogram_quantile_missing_le_column() {
13608        let case = r#"histogram_quantile(0.99, non_existent_histogram_bucket) or normal_metric"#;
13609        let eval_stmt = build_eval_stmt(case);
13610        let table_provider = build_missing_le_or_normal_metric_table_provider().await;
13611
13612        let plan =
13613            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13614                .await
13615                .unwrap();
13616        assert_normal_metric_schema(&plan);
13617    }
13618
13619    #[tokio::test]
13620    async fn test_or_with_right_empty_histogram_restores_left_context() {
13621        let eval_stmt = build_eval_stmt(
13622            r#"abs(sum by(instance) (normal_metric) or histogram_quantile(0.99, sum by(pod) (non_existent_histogram_bucket)))"#,
13623        );
13624        let table_provider = build_missing_le_or_normal_metric_table_provider().await;
13625
13626        PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13627            .await
13628            .unwrap();
13629    }
13630
13631    #[tokio::test]
13632    async fn test_or_with_both_empty_histograms() {
13633        let eval_stmt = build_eval_stmt(
13634            r#"histogram_quantile(0.99, sum by(pod) (left_histogram_bucket)) or histogram_quantile(0.99, sum by(instance) (right_histogram_bucket))"#,
13635        );
13636        let table_provider = build_test_table_provider_with_fields(
13637            &[
13638                (
13639                    DEFAULT_SCHEMA_NAME.to_string(),
13640                    "left_histogram_bucket".to_string(),
13641                ),
13642                (
13643                    DEFAULT_SCHEMA_NAME.to_string(),
13644                    "right_histogram_bucket".to_string(),
13645                ),
13646            ],
13647            &["pod", "instance"],
13648        )
13649        .await;
13650
13651        let plan =
13652            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13653                .await
13654                .unwrap();
13655        match plan {
13656            LogicalPlan::EmptyRelation(relation) => {
13657                assert!(!relation.produce_one_row);
13658                assert!(!relation.schema.fields().is_empty());
13659                assert!(
13660                    relation
13661                        .schema
13662                        .fields()
13663                        .iter()
13664                        .any(|field| field.data_type() == &ArrowDataType::Float64)
13665                );
13666                assert!(
13667                    relation
13668                        .schema
13669                        .fields()
13670                        .iter()
13671                        .any(|field| field.name() == "pod")
13672                );
13673                assert!(
13674                    !relation
13675                        .schema
13676                        .fields()
13677                        .iter()
13678                        .any(|field| field.name() == "instance")
13679                );
13680            }
13681            _ => panic!("Expected EmptyRelation, but got: {plan:?}"),
13682        }
13683    }
13684
13685    #[tokio::test]
13686    async fn test_nested_or_with_both_empty_histograms() {
13687        for case in [
13688            r#"abs(histogram_quantile(0.99, left_histogram_bucket) or histogram_quantile(0.99, right_histogram_bucket))"#,
13689            r#"(histogram_quantile(0.99, left_histogram_bucket) or histogram_quantile(0.99, right_histogram_bucket)) + 1"#,
13690        ] {
13691            let eval_stmt = build_eval_stmt(case);
13692            let table_provider = build_test_table_provider_with_fields(
13693                &[
13694                    (
13695                        DEFAULT_SCHEMA_NAME.to_string(),
13696                        "left_histogram_bucket".to_string(),
13697                    ),
13698                    (
13699                        DEFAULT_SCHEMA_NAME.to_string(),
13700                        "right_histogram_bucket".to_string(),
13701                    ),
13702                ],
13703                &["pod", "instance"],
13704            )
13705            .await;
13706
13707            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13708                .await
13709                .unwrap();
13710        }
13711    }
13712
13713    #[tokio::test]
13714    async fn test_or_with_empty_histogram_modifiers() {
13715        for case in [
13716            r#"histogram_quantile(0.99, non_existent_histogram_bucket) or on(pod) normal_metric"#,
13717            r#"normal_metric or ignoring(instance) histogram_quantile(0.99, non_existent_histogram_bucket)"#,
13718        ] {
13719            let eval_stmt = build_eval_stmt(case);
13720            let table_provider = build_missing_le_or_normal_metric_table_provider().await;
13721
13722            let plan =
13723                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13724                    .await
13725                    .unwrap();
13726            assert_normal_metric_schema(&plan);
13727        }
13728    }
13729
13730    #[tokio::test]
13731    async fn test_unless_preserves_left_context_for_histogram() {
13732        let eval_stmt = build_eval_stmt(
13733            r#"histogram_quantile(0.99, bucket_metric unless on(job) normal_metric) or fallback_metric"#,
13734        );
13735        let state = build_query_engine_state();
13736        let plan = PromPlanner::stmt_to_plan(
13737            build_set_op_context_table_provider().await,
13738            &eval_stmt,
13739            &state,
13740        )
13741        .await
13742        .unwrap();
13743        assert!(contains_histogram_fold(&plan), "{plan:?}");
13744        let (optimized, physical) = optimize_and_create_physical_plan(&state, plan).await;
13745        assert!(contains_histogram_fold(&optimized), "{optimized:?}");
13746        let batches =
13747            datafusion::physical_plan::collect(physical, state.session_state().task_ctx())
13748                .await
13749                .unwrap();
13750        assert!(batches.iter().all(|batch| batch.num_rows() == 0));
13751    }
13752
13753    #[tokio::test]
13754    async fn test_and_preserves_left_context_for_histogram() {
13755        let eval_stmt = build_eval_stmt(
13756            r#"histogram_quantile(0.99, bucket_metric and on(job) normal_metric) or fallback_metric"#,
13757        );
13758        let plan = PromPlanner::stmt_to_plan(
13759            build_set_op_context_table_provider().await,
13760            &eval_stmt,
13761            &build_query_engine_state(),
13762        )
13763        .await
13764        .unwrap();
13765        assert!(contains_histogram_fold(&plan), "{plan:?}");
13766    }
13767
13768    #[tokio::test]
13769    async fn test_and_preserves_left_context_when_le_is_missing() {
13770        let eval_stmt =
13771            build_eval_stmt(r#"histogram_quantile(0.99, normal_metric and on(job) bucket_metric)"#);
13772        let plan = PromPlanner::stmt_to_plan(
13773            build_set_op_context_table_provider().await,
13774            &eval_stmt,
13775            &build_query_engine_state(),
13776        )
13777        .await
13778        .unwrap();
13779        assert!(matches!(&plan, LogicalPlan::EmptyRelation(_)), "{plan:?}");
13780        assert!(!plan.schema().fields().is_empty());
13781        assert!(!contains_histogram_fold(&plan), "{plan:?}");
13782    }
13783
13784    async fn build_matching_filter_plan(query: &str) -> String {
13785        let table_provider = build_test_table_provider_with_distinct_tags(&[
13786            ("metric_a", &["host", "device"]),
13787            ("metric_b", &["host", "device"]),
13788        ])
13789        .await;
13790        let plan = PromPlanner::stmt_to_plan(
13791            table_provider,
13792            &build_eval_stmt(query),
13793            &build_query_engine_state(),
13794        )
13795        .await
13796        .unwrap();
13797        plan.display_indent().to_string()
13798    }
13799
13800    /// [`build_test_table_provider_with_distinct_tags`] plus a `status` string column: a value
13801    /// field that is neither a primary key nor a value column of the metric, so
13802    /// `count by(status) (...)` still reports it among the aggregation's tag columns.
13803    async fn build_test_table_provider_with_string_field(
13804        table_tags: &[(&str, &[&str])],
13805    ) -> DfTableSourceProvider {
13806        let catalog_list = MemoryCatalogManager::with_default_setup();
13807        for (table_name, tags) in table_tags {
13808            let mut columns = tags
13809                .iter()
13810                .map(|tag| {
13811                    ColumnSchema::new(
13812                        (*tag).to_string(),
13813                        ConcreteDataType::string_datatype(),
13814                        false,
13815                    )
13816                })
13817                .collect::<Vec<_>>();
13818            columns.push(
13819                ColumnSchema::new(
13820                    greptime_timestamp().to_string(),
13821                    ConcreteDataType::timestamp_millisecond_datatype(),
13822                    false,
13823                )
13824                .with_time_index(true),
13825            );
13826            columns.push(ColumnSchema::new(
13827                greptime_value().to_string(),
13828                ConcreteDataType::float64_datatype(),
13829                true,
13830            ));
13831            columns.push(ColumnSchema::new(
13832                "status".to_string(),
13833                ConcreteDataType::string_datatype(),
13834                true,
13835            ));
13836            let table_meta = TableMetaBuilder::empty()
13837                .schema(Arc::new(Schema::new(columns)))
13838                .primary_key_indices((0..tags.len()).collect())
13839                .next_column_id(1024)
13840                .build()
13841                .unwrap();
13842            let table_info = TableInfoBuilder::default()
13843                .name((*table_name).to_string())
13844                .meta(table_meta)
13845                .build()
13846                .unwrap();
13847
13848            assert!(
13849                catalog_list
13850                    .register_table_sync(RegisterTableRequest {
13851                        catalog: DEFAULT_CATALOG_NAME.to_string(),
13852                        schema: DEFAULT_SCHEMA_NAME.to_string(),
13853                        table_name: (*table_name).to_string(),
13854                        table_id: 1024,
13855                        table: EmptyTable::from_table_info(&table_info),
13856                    })
13857                    .is_ok()
13858            );
13859        }
13860
13861        DfTableSourceProvider::new(
13862            catalog_list,
13863            false,
13864            QueryContext::arc(),
13865            DummyDecoder::arc(),
13866            false,
13867        )
13868    }
13869
13870    async fn build_matching_filter_plan_with_string_field(query: &str) -> String {
13871        let table_provider = build_test_table_provider_with_string_field(&[
13872            ("metric_a", &["host", "device"]),
13873            ("metric_b", &["host", "device"]),
13874        ])
13875        .await;
13876        let plan = PromPlanner::stmt_to_plan(
13877            table_provider,
13878            &build_eval_stmt(query),
13879            &build_query_engine_state(),
13880        )
13881        .await
13882        .unwrap();
13883        plan.display_indent().to_string()
13884    }
13885
13886    #[tokio::test]
13887    async fn binary_matching_label_filter_reaches_both_operands() {
13888        for query in [
13889            r#"metric_a / metric_b{host="foo"}"#,
13890            r#"metric_a / on(host, device) metric_b{host="foo"}"#,
13891            r#"count_over_time(metric_a[1m]) / on(host) count_over_time(metric_b{host="foo"}[1m])"#,
13892            r#"metric_a / ignoring(device) metric_b{host="foo"}"#,
13893            r#"sum by(host) (metric_a) / on(host) sum by(host) (metric_b{host="foo"})"#,
13894            r#"metric_a / on(host) avg without(device) (metric_b{host="foo"})"#,
13895        ] {
13896            let plan = build_matching_filter_plan(query).await;
13897            assert_eq!(
13898                plan.matches(r#"host = Utf8("foo")"#).count(),
13899                2,
13900                "{query}\n{plan}"
13901            );
13902        }
13903    }
13904
13905    #[tokio::test]
13906    async fn binary_matching_label_filter_reaches_scalar_ranking_and_grouped_operands() {
13907        for query in [
13908            r#"(8 * metric_a{host="foo"}) / on(host) metric_b"#,
13909            r#"topk(1, metric_a{host="foo"}) / on(host, device) metric_b"#,
13910            r#"(8 * metric_a{host="foo"}) / on(host) group_left topk by(host)(1, max by(host)(metric_b))"#,
13911        ] {
13912            let plan = build_matching_filter_plan(query).await;
13913            assert_eq!(
13914                plan.matches(r#"host = Utf8("foo")"#).count(),
13915                2,
13916                "{query}\n{plan}"
13917            );
13918        }
13919        // A global ranking one-side must see every host, so the matcher stays put.
13920        let query = r#"metric_a{host="foo"} / on(host) group_left topk(1, max by(host)(metric_b))"#;
13921        let plan = build_matching_filter_plan(query).await;
13922        assert_eq!(
13923            plan.matches(r#"host = Utf8("foo")"#).count(),
13924            1,
13925            "{query}\n{plan}"
13926        );
13927    }
13928
13929    #[tokio::test]
13930    async fn binary_matching_label_filter_skips_selecting_aggregations() {
13931        // `topk` ranks its input, so filtering before it changes the candidate set.
13932        let query = r#"topk(1, metric_a) / on(host, device) metric_b{host="foo"}"#;
13933        let plan = build_matching_filter_plan(query).await;
13934        assert_eq!(plan.matches("foo").count(), 1, "{query}\n{plan}");
13935    }
13936
13937    #[tokio::test]
13938    async fn binary_value_field_matcher_stays_on_its_own_operand() {
13939        let value = greptime_value();
13940        for query in [
13941            format!(r#"metric_a / metric_b{{{value}="2"}}"#),
13942            format!(r#"metric_a / on(host, device, {value}) metric_b{{{value}="2"}}"#),
13943        ] {
13944            let plan = build_matching_filter_plan(&query).await;
13945            assert_eq!(plan.matches(r#"Utf8("2")"#).count(), 1, "{query}\n{plan}");
13946        }
13947    }
13948
13949    #[tokio::test]
13950    async fn binary_value_field_matcher_is_not_copied_across_aggregations() {
13951        // `status` varies between the samples of one series, so filtering the other operand by it
13952        // would drop the newest sample before sample selection (#9242).
13953        for query in [
13954            r#"count by(status) (metric_a) / on(status) count by(status) (metric_b{status="ready"})"#,
13955            r#"(8 * count by(status)(metric_a{__field__="status"})) / on(status) topk by(status)(1, count by(status)(metric_b{__field__="status",status="ready"}))"#,
13956            r#"count by(status)(metric_a{__field__="status"}) / on(status) group_left topk by(status)(1, count by(status)(metric_b{__field__="status",status="ready"}))"#,
13957        ] {
13958            let plan = build_matching_filter_plan_with_string_field(query).await;
13959            assert_eq!(
13960                plan.matches(r#"Utf8("ready")"#).count(),
13961                1,
13962                "{query}\n{plan}"
13963            );
13964        }
13965    }
13966
13967    #[tokio::test]
13968    async fn binary_matching_label_filter_reaches_aggregations_grouping_by_tags() {
13969        // `count`, not `sum`: the string value field is not summable.
13970        let query = r#"count by(host) (metric_a) / on(host) count by(host) (metric_b{host="foo"})"#;
13971        let plan = build_matching_filter_plan_with_string_field(query).await;
13972        assert_eq!(plan.matches(r#"Utf8("foo")"#).count(), 2, "{query}\n{plan}");
13973    }
13974
13975    #[tokio::test]
13976    async fn binary_matching_label_filter_skips_unproven_expressions() {
13977        for query in [
13978            r#"metric_a > on(host, device) metric_b{host="foo"}"#,
13979            r#"metric_a / on(host) group_left metric_b{host="foo"}"#,
13980            r#"metric_a / on(host) label_replace(metric_b{host="foo"},"extra","e","host",".*")"#,
13981            r#"metric_a / on(device) metric_b{host="foo"}"#,
13982        ] {
13983            let plan = build_matching_filter_plan(query).await;
13984            assert_eq!(plan.matches("foo").count(), 1, "{query}\n{plan}");
13985        }
13986    }
13987
13988    #[tokio::test]
13989    async fn test_or_context_uses_left_qualified_output() {
13990        let case = r#"(normal_metric or other_metric) + 1"#;
13991        let eval_stmt = build_eval_stmt(case);
13992        let state = build_query_engine_state();
13993        let plan =
13994            PromPlanner::stmt_to_plan(build_or_context_table_provider().await, &eval_stmt, &state)
13995                .await
13996                .unwrap();
13997        assert!(
13998            plan.schema()
13999                .fields()
14000                .iter()
14001                .any(|field| field.data_type() == &ArrowDataType::Float64),
14002            "{plan:?}"
14003        );
14004        let (_optimized, _physical) = optimize_and_create_physical_plan(&state, plan).await;
14005    }
14006
14007    #[tokio::test]
14008    async fn test_or_context_uses_left_qualified_empty_histogram_output() {
14009        let case = r#"(abs(histogram_quantile(0.99, non_hist_metric)) or normal_metric) + 1"#;
14010        let eval_stmt = build_eval_stmt(case);
14011        let plan = PromPlanner::stmt_to_plan(
14012            build_or_context_table_provider().await,
14013            &eval_stmt,
14014            &build_query_engine_state(),
14015        )
14016        .await
14017        .unwrap();
14018        assert!(
14019            plan.schema()
14020                .fields()
14021                .iter()
14022                .any(|field| field.data_type() == &ArrowDataType::Float64),
14023            "{plan:?}"
14024        );
14025    }
14026
14027    #[tokio::test]
14028    async fn test_direct_or_preserves_float_and_native_histogram_samples() {
14029        for histogram_on_left in [false, true] {
14030            let (planner, plan) = mixed_direct_or(histogram_on_left).await;
14031
14032            let float_field = &planner.ctx.field_columns[0];
14033            let histogram_field = &planner.ctx.field_columns[1];
14034            assert!(float_field.starts_with(OR_FLOAT_FIELD_PREFIX));
14035            assert!(histogram_field.starts_with(OR_HISTOGRAM_FIELD_PREFIX));
14036            assert_eq!(
14037                plan.schema()
14038                    .field_with_name(None, float_field)
14039                    .unwrap()
14040                    .data_type(),
14041                &ArrowDataType::Float64
14042            );
14043            assert_eq!(
14044                plan.schema()
14045                    .field_with_name(None, histogram_field)
14046                    .unwrap()
14047                    .data_type(),
14048                &native_histogram_value_type().as_arrow_type()
14049            );
14050
14051            let (optimized, batches) = execute(plan, &build_query_engine_state()).await;
14052            assert_no_internal_or_keys(optimized.schema());
14053            let mut sample_kinds = batches
14054                .iter()
14055                .flat_map(|batch| {
14056                    let values = batch.column_by_name(float_field).unwrap();
14057                    let histograms = batch.column_by_name(histogram_field).unwrap();
14058                    (0..batch.num_rows())
14059                        .map(|row| (values.is_valid(row), histograms.is_valid(row)))
14060                })
14061                .collect::<Vec<_>>();
14062            sample_kinds.sort_unstable();
14063            assert_eq!(sample_kinds, vec![(false, true), (true, false)]);
14064        }
14065    }
14066
14067    #[tokio::test]
14068    async fn malformed_classic_bucket_does_not_drop_native_histogram() {
14069        let state = build_query_engine_state();
14070        let collector = PromqlAnnotationCollector::default();
14071        let plan = PromPlanner::stmt_to_plan_with_annotations(
14072            operator_table_provider(),
14073            &operator_eval_stmt("histogram_quantile(0.5, bad_classic or bad_native)"),
14074            &state,
14075            Some(collector.clone()),
14076        )
14077        .await
14078        .unwrap();
14079        let value_field = plan
14080            .schema()
14081            .fields()
14082            .iter()
14083            .find(|field| field.data_type() == &ArrowDataType::Float64)
14084            .unwrap()
14085            .name()
14086            .clone();
14087
14088        let (_, batches) = execute(plan, &state).await;
14089        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
14090        assert_eq!(values(&batches, &value_field), vec![0.0]);
14091        let mut warnings = vec![];
14092        let mut infos = vec![];
14093        collector.append_to(&mut warnings, &mut infos);
14094        assert!(warnings.is_empty());
14095        assert!(infos.is_empty());
14096    }
14097
14098    #[tokio::test]
14099    async fn test_mixed_binary_operator_aligns_both_alternative_inputs() {
14100        let state = build_query_engine_state();
14101        let plan = PromPlanner::stmt_to_plan(
14102            operator_table_provider(),
14103            &operator_eval_stmt("(lf or on(tag) lh) * on(tag) (rf or on(tag) rh)"),
14104            &state,
14105        )
14106        .await
14107        .unwrap();
14108        let plan_text = plan.display_indent_schema().to_string();
14109        assert!(
14110            plan_text.contains("prom_native_histogram_mul_scalar"),
14111            "{plan_text}"
14112        );
14113        assert!(
14114            plan_text.contains("prom_native_histogram_scalar_mul"),
14115            "{plan_text}"
14116        );
14117        let float_field = plan
14118            .schema()
14119            .fields()
14120            .iter()
14121            .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX))
14122            .unwrap()
14123            .name()
14124            .clone();
14125        let histogram_field = plan
14126            .schema()
14127            .fields()
14128            .iter()
14129            .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX))
14130            .unwrap()
14131            .name()
14132            .clone();
14133
14134        let (_, batches) = execute(plan, &state).await;
14135        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
14136        assert!(values(&batches, &float_field).is_empty());
14137        let mut sums = histograms(&batches, &histogram_field)
14138            .into_iter()
14139            .map(|histogram| histogram.sum)
14140            .collect::<Vec<_>>();
14141        sums.sort_by(f64::total_cmp);
14142        assert_eq!(sums, vec![2.0, 3.0]);
14143    }
14144
14145    #[tokio::test]
14146    async fn test_mixed_binary_operator_reports_only_dropped_samples() {
14147        for (query, expected_rows, expected_infos) in [
14148            ("(lf or on(tag) lh) + on(tag) (rf or on(tag) rh)", 0, 1),
14149            ("(lf or on(tag) lh) + on(tag) (lf or on(tag) lh)", 2, 0),
14150            ("(lf or on(tag) lh) % on(tag) lh", 0, 1),
14151        ] {
14152            let state = build_query_engine_state();
14153            let annotations = PromqlAnnotationCollector::default();
14154            let plan = PromPlanner::stmt_to_plan_with_annotations(
14155                operator_table_provider(),
14156                &operator_eval_stmt(query),
14157                &state,
14158                Some(annotations.clone()),
14159            )
14160            .await
14161            .unwrap();
14162
14163            let (_, batches) = execute(plan, &state).await;
14164            assert_eq!(
14165                batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
14166                expected_rows,
14167                "{query}"
14168            );
14169            let mut warnings = vec![];
14170            let mut infos = vec![];
14171            annotations.append_to(&mut warnings, &mut infos);
14172            assert!(warnings.is_empty(), "{query}: {warnings:?}");
14173            assert_eq!(infos.len(), expected_infos, "{query}: {infos:?}");
14174        }
14175    }
14176
14177    #[tokio::test]
14178    async fn test_histogram_only_min_drops_empty_aggregate_group() {
14179        // `min` over native-histogram-only input drops every sample in the group, so the
14180        // NULL-valued aggregate row must be filtered out. Otherwise an outer expression
14181        // like `group()` resurrects the group Prometheus considers unseen.
14182        let state = build_query_engine_state();
14183        for query in ["min(lh)", "group(min(lh))"] {
14184            let plan = PromPlanner::stmt_to_plan(
14185                operator_table_provider(),
14186                &operator_eval_stmt(query),
14187                &state,
14188            )
14189            .await
14190            .unwrap();
14191            let (_, batches) = execute(plan, &state).await;
14192            assert_eq!(
14193                batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
14194                0,
14195                "{query}"
14196            );
14197        }
14198    }
14199
14200    #[tokio::test]
14201    async fn test_mixed_min_drops_histogram_only_group() {
14202        // With alternative float/histogram fields, `min by (tag)` keeps float-only groups
14203        // (tag=a from `lf`) and drops histogram-only groups (tag=b from `lh`) instead of
14204        // emitting a NULL-valued row for them.
14205        let state = build_query_engine_state();
14206        let plan = PromPlanner::stmt_to_plan(
14207            operator_table_provider(),
14208            &operator_eval_stmt("min by (tag) (lf or on(tag) lh)"),
14209            &state,
14210        )
14211        .await
14212        .unwrap();
14213        let float_field = plan
14214            .schema()
14215            .fields()
14216            .iter()
14217            .find(|field| field.data_type() == &ArrowDataType::Float64)
14218            .unwrap()
14219            .name()
14220            .clone();
14221        let (_, batches) = execute(plan, &state).await;
14222        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
14223        assert_eq!(values(&batches, &float_field), vec![2.0]);
14224    }
14225
14226    #[tokio::test]
14227    async fn test_mixed_or_can_feed_another_or() {
14228        let state = build_query_engine_state();
14229        let plan = PromPlanner::stmt_to_plan(
14230            operator_table_provider(),
14231            &operator_eval_stmt("lf or on(tag) lh or on(tag) fallback"),
14232            &state,
14233        )
14234        .await
14235        .unwrap();
14236        let float_field = plan
14237            .schema()
14238            .fields()
14239            .iter()
14240            .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX))
14241            .unwrap()
14242            .name()
14243            .clone();
14244        let histogram_field = plan
14245            .schema()
14246            .fields()
14247            .iter()
14248            .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX))
14249            .unwrap()
14250            .name()
14251            .clone();
14252
14253        let (_, batches) = execute(plan, &state).await;
14254        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 3);
14255        let mut float_values = values(&batches, &float_field);
14256        float_values.sort_by(f64::total_cmp);
14257        assert_eq!(float_values, vec![2.0, 7.0]);
14258        assert_eq!(histograms(&batches, &histogram_field).len(), 1);
14259    }
14260
14261    #[tokio::test]
14262    async fn test_mixed_fields_align_with_single_float_vector() {
14263        let (planner, mixed) = mixed_direct_or(false).await;
14264        let scale = tagged_source(
14265            "scale",
14266            false,
14267            ("k", Some("float")),
14268            DirectOrValue::Float64(2.0),
14269        );
14270        let scale = scan(&scale);
14271        let scale_fields = vec!["v".to_string()];
14272        let PromExpr::Binary(binary) = parser::parse("lhs * rhs").unwrap() else {
14273            unreachable!()
14274        };
14275
14276        let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns(
14277            mixed.schema(),
14278            scale.schema(),
14279            &planner.ctx.field_columns,
14280            &scale_fields,
14281            binary.op,
14282            false,
14283            false,
14284        );
14285        assert!(invalid_pairs.is_empty());
14286        assert_eq!(
14287            groups
14288                .iter()
14289                .map(|(output, _)| output.clone())
14290                .collect::<Vec<_>>(),
14291            planner.ctx.field_columns
14292        );
14293        assert_eq!(groups.len(), 2);
14294        assert!(
14295            groups
14296                .iter()
14297                .flat_map(|(_, pairs)| pairs)
14298                .all(|(_, right)| *right == &scale_fields[0])
14299        );
14300
14301        let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns(
14302            scale.schema(),
14303            mixed.schema(),
14304            &scale_fields,
14305            &planner.ctx.field_columns,
14306            binary.op,
14307            false,
14308            false,
14309        );
14310        assert!(invalid_pairs.is_empty());
14311        assert_eq!(
14312            groups
14313                .iter()
14314                .map(|(output, _)| output.clone())
14315                .collect::<Vec<_>>(),
14316            planner.ctx.field_columns
14317        );
14318        assert_eq!(groups.len(), 2);
14319        assert!(
14320            groups
14321                .iter()
14322                .flat_map(|(_, pairs)| pairs)
14323                .all(|(left, _)| *left == &scale_fields[0])
14324        );
14325    }
14326
14327    #[tokio::test]
14328    async fn test_non_bool_comparison_filters_mixed_sample_lanes() {
14329        let (planner, input) = mixed_direct_or(false).await;
14330        let input_schema = input.schema().clone();
14331        let plan = planner
14332            .filter_on_field_column(input, |field| {
14333                if PromPlanner::field_column_is_native_histogram(&input_schema, field) {
14334                    Ok(lit(false))
14335                } else {
14336                    Ok(col(field).gt(lit(0.0)))
14337                }
14338            })
14339            .unwrap();
14340        let float_field = planner.ctx.field_columns[0].clone();
14341
14342        let (_, batches) = execute(plan, &build_query_engine_state()).await;
14343        assert_eq!(values(&batches, &float_field), vec![1.25]);
14344        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
14345    }
14346
14347    #[tokio::test]
14348    async fn test_mixed_left_and_unless_preserve_sample_lanes() {
14349        for (expression, expected_sample_kind) in [
14350            ("lhs and on(k) mask", (false, true)),
14351            ("lhs unless on(k) mask", (true, false)),
14352        ] {
14353            let (mut planner, left) = mixed_direct_or(false).await;
14354            let left_context = planner.ctx.clone();
14355            let float_field = left_context.field_columns[0].clone();
14356            let histogram_field = left_context.field_columns[1].clone();
14357            let mask = tagged_source(
14358                "mask",
14359                false,
14360                ("k", Some("histogram")),
14361                DirectOrValue::Float64(1.0),
14362            );
14363            let PromExpr::Binary(binary) = parser::parse(expression).unwrap() else {
14364                unreachable!()
14365            };
14366            let plan = planner
14367                .set_op_on_non_field_columns(
14368                    left,
14369                    scan(&mask),
14370                    left_context,
14371                    direct_or_context("mask", &["job", "k"], "v"),
14372                    binary.op,
14373                    &binary.modifier,
14374                )
14375                .unwrap();
14376
14377            let (_, batches) = execute(plan, &build_query_engine_state()).await;
14378            let sample_kinds = batches
14379                .iter()
14380                .flat_map(|batch| {
14381                    let floats = batch.column_by_name(&float_field).unwrap();
14382                    let histograms = batch.column_by_name(&histogram_field).unwrap();
14383                    (0..batch.num_rows())
14384                        .map(|row| (floats.is_valid(row), histograms.is_valid(row)))
14385                })
14386                .collect::<Vec<_>>();
14387            assert_eq!(sample_kinds, vec![expected_sample_kind], "{expression}");
14388        }
14389    }
14390
14391    #[tokio::test]
14392    async fn test_mixed_fields_arithmetic_broadcasts_computed_scalar() {
14393        let plan = PromPlanner::stmt_to_plan(
14394            build_test_mixed_native_histogram_table_provider("some_metric").await,
14395            &build_eval_stmt("some_metric * scalar(vector(2))"),
14396            &build_query_engine_state(),
14397        )
14398        .await
14399        .unwrap();
14400        let schema = plan.schema();
14401        assert_eq!(
14402            schema
14403                .field_with_unqualified_name(greptime_value())
14404                .unwrap()
14405                .data_type(),
14406            &ArrowDataType::Float64
14407        );
14408        assert_eq!(
14409            schema
14410                .field_with_unqualified_name(greptime_native_histogram())
14411                .unwrap()
14412                .data_type(),
14413            &native_histogram_value_type().as_arrow_type()
14414        );
14415        assert!(
14416            plan.display_indent_schema()
14417                .to_string()
14418                .contains("prom_native_histogram_mul_scalar"),
14419            "{plan:?}"
14420        );
14421    }
14422
14423    #[tokio::test]
14424    async fn test_unsupported_histogram_binary_does_not_block_or_fallback() {
14425        let state = build_query_engine_state();
14426        let plan = PromPlanner::stmt_to_plan(
14427            operator_table_provider(),
14428            &operator_eval_stmt("((lf or on(tag) lh) % 2) or on(tag) lh"),
14429            &state,
14430        )
14431        .await
14432        .unwrap();
14433        let float_field = plan
14434            .schema()
14435            .fields()
14436            .iter()
14437            .find(|field| field.data_type() == &ArrowDataType::Float64)
14438            .unwrap()
14439            .name()
14440            .clone();
14441        let histogram_field = plan
14442            .schema()
14443            .fields()
14444            .iter()
14445            .find(|field| field.data_type() == &native_histogram_value_type().as_arrow_type())
14446            .unwrap()
14447            .name()
14448            .clone();
14449
14450        let (_, batches) = execute(plan, &state).await;
14451        assert_eq!(values(&batches, &float_field), vec![0.0]);
14452        assert_eq!(histograms(&batches, &histogram_field).len(), 1);
14453    }
14454
14455    #[tokio::test]
14456    async fn test_unary_negates_mixed_float_and_native_histogram_samples() {
14457        for histogram_on_left in [false, true] {
14458            let (mut planner, input) = mixed_direct_or(histogram_on_left).await;
14459            let plan = planner.negate_field_columns(input).unwrap();
14460            assert!(PromPlanner::field_columns_are_alternative_samples(
14461                plan.schema(),
14462                &planner.ctx.field_columns
14463            ));
14464            let float_field = planner
14465                .ctx
14466                .field_columns
14467                .iter()
14468                .find(|field| field.starts_with(OR_FLOAT_FIELD_PREFIX))
14469                .unwrap();
14470            let histogram_field = planner
14471                .ctx
14472                .field_columns
14473                .iter()
14474                .find(|field| field.starts_with(OR_HISTOGRAM_FIELD_PREFIX))
14475                .unwrap();
14476
14477            let (_, batches) = execute(plan, &build_query_engine_state()).await;
14478            assert_eq!(values(&batches, float_field), vec![-1.25]);
14479            let histogram = batches
14480                .iter()
14481                .find_map(|batch| {
14482                    let values = batch
14483                        .column_by_name(histogram_field)
14484                        .unwrap()
14485                        .as_any()
14486                        .downcast_ref::<datafusion::arrow::array::StructArray>()
14487                        .unwrap();
14488                    (0..values.len()).find_map(|row| {
14489                        common_query::native_histogram::read_histogram(values, row).unwrap()
14490                    })
14491                })
14492                .unwrap();
14493            assert_eq!(histogram.count, -1.0);
14494            assert_eq!(histogram.sum, -1.0);
14495            assert_eq!(histogram.reset_hint, CounterResetHint::Gauge);
14496        }
14497    }
14498
14499    #[tokio::test]
14500    async fn test_native_histogram_sum_and_avg_execute_real_batches() {
14501        for op_name in ["sum", "avg"] {
14502            for incompatible in [false, true] {
14503                let mut second = direct_or_histogram();
14504                if incompatible {
14505                    second.schema = CUSTOM_BUCKETS_SCHEMA;
14506                    second.custom_values = vec![1.0];
14507                }
14508                let collector = PromqlAnnotationCollector::default();
14509                let (mut planner, input) =
14510                    mixed_aggregate_input(vec![direct_or_histogram(), second]).await;
14511                planner.promql_annotations = Some(collector.clone());
14512                let histogram_column = planner.ctx.field_columns[1].clone();
14513                planner.ctx.field_columns = vec![histogram_column.clone()];
14514                let input = LogicalPlanBuilder::from(input)
14515                    .project([col("ts"), col(&histogram_column)])
14516                    .unwrap()
14517                    .build()
14518                    .unwrap();
14519                let PromExpr::Aggregate(AggregateExpr { op, param, .. }) =
14520                    parser::parse(&format!("{op_name}(mixed)")).unwrap()
14521                else {
14522                    unreachable!()
14523                };
14524                let (aggregate_exprs, _) =
14525                    planner.create_aggregate_exprs(op, &param, &input).unwrap();
14526                let plan = LogicalPlanBuilder::from(input)
14527                    .aggregate(vec![col("ts")], aggregate_exprs)
14528                    .unwrap()
14529                    .filter(planner.create_empty_values_filter_expr(false).unwrap())
14530                    .unwrap()
14531                    .build()
14532                    .unwrap();
14533
14534                let (_, batches) = execute(plan, &build_query_engine_state()).await;
14535                let mut warnings = vec![];
14536                let mut infos = vec![];
14537                collector.append_to(&mut warnings, &mut infos);
14538                assert!(infos.is_empty());
14539                if incompatible {
14540                    assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
14541                    assert!(warnings.iter().any(|warning| {
14542                        warning
14543                            == &format!(
14544                                "prom_native_histogram_agg_{op_name}: dropped native histogram aggregate with incompatible schemas"
14545                            )
14546                    }));
14547                } else {
14548                    let histograms = histograms(&batches, &histogram_column);
14549                    assert_eq!(histograms.len(), 1);
14550                    let expected = if op_name == "sum" { 2.0 } else { 1.0 };
14551                    assert_eq!(histograms[0].count, expected);
14552                    assert_eq!(histograms[0].sum, expected);
14553                    assert!(warnings.is_empty());
14554                }
14555            }
14556        }
14557    }
14558
14559    #[tokio::test]
14560    async fn test_canonical_mixed_count_group_and_count_values_execute() {
14561        let state = build_query_engine_state();
14562        for (query, expected) in [
14563            ("count(some_metric)", vec![2.0]),
14564            ("group(some_metric)", vec![1.0]),
14565            (r#"count_values("sample", some_metric)"#, vec![1.0, 1.0]),
14566        ] {
14567            let plan = PromPlanner::stmt_to_plan(
14568                build_test_mixed_native_histogram_table_provider("some_metric").await,
14569                &operator_eval_stmt(query),
14570                &state,
14571            )
14572            .await
14573            .unwrap();
14574            assert!(
14575                plan.schema()
14576                    .fields()
14577                    .iter()
14578                    .all(|field| !field.name().starts_with("__promql_sample_count")),
14579                "{query}: {plan:?}"
14580            );
14581            let value_fields = plan
14582                .schema()
14583                .fields()
14584                .iter()
14585                .filter(|field| {
14586                    matches!(
14587                        field.data_type(),
14588                        ArrowDataType::Float64 | ArrowDataType::Int64 | ArrowDataType::UInt64
14589                    ) || field.data_type() == &native_histogram_value_type().as_arrow_type()
14590                })
14591                .collect::<Vec<_>>();
14592            assert_eq!(value_fields.len(), 1, "{query}: {plan:?}");
14593            assert_ne!(
14594                value_fields[0].data_type(),
14595                &native_histogram_value_type().as_arrow_type(),
14596                "{query}: {plan:?}"
14597            );
14598            let value_column = value_fields[0].name().clone();
14599
14600            let (_, batches) = execute(plan, &state).await;
14601            let mut actual = numeric_values(&batches, &value_column);
14602            actual.sort_by(f64::total_cmp);
14603            assert_eq!(actual, expected, "{query}");
14604
14605            if query.starts_with("count_values") {
14606                let mut sample_labels = batches
14607                    .iter()
14608                    .flat_map(|batch| {
14609                        batch
14610                            .column_by_name("sample")
14611                            .unwrap()
14612                            .as_any()
14613                            .downcast_ref::<StringArray>()
14614                            .unwrap()
14615                            .iter()
14616                            .flatten()
14617                            .map(str::to_string)
14618                    })
14619                    .collect::<Vec<_>>();
14620                sample_labels.sort();
14621                let mut expected_labels =
14622                    vec!["2".to_string(), direct_or_histogram().promql_string()];
14623                expected_labels.sort();
14624                assert_eq!(sample_labels, expected_labels);
14625            }
14626        }
14627    }
14628
14629    #[tokio::test]
14630    async fn test_mixed_or_sum_aggregates_each_sample_type() {
14631        let PromExpr::Aggregate(AggregateExpr { op, param, .. }) =
14632            parser::parse("sum(lhs)").unwrap()
14633        else {
14634            unreachable!()
14635        };
14636
14637        let collector = PromqlAnnotationCollector::default();
14638        let (mut planner, input) = mixed_direct_or(false).await;
14639        planner.promql_annotations = Some(collector.clone());
14640        let float_column = planner.ctx.field_columns[0].clone();
14641        let histogram_column = planner.ctx.field_columns[1].clone();
14642        let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, &param, &input).unwrap();
14643        let plan = LogicalPlanBuilder::from(input)
14644            .aggregate(vec![col("ts"), col("k")], aggregate_exprs)
14645            .unwrap()
14646            .filter(
14647                planner
14648                    .mixed_aggregate_filter_expr(op, &float_column, &histogram_column)
14649                    .unwrap(),
14650            )
14651            .unwrap()
14652            .project([
14653                col(&float_column),
14654                col(&histogram_column),
14655                col("ts"),
14656                col("k"),
14657            ])
14658            .unwrap()
14659            .build()
14660            .unwrap();
14661
14662        let (_, batches) = execute(plan, &build_query_engine_state()).await;
14663        assert_eq!(values(&batches, &float_column), vec![1.25]);
14664        let histogram = batches
14665            .iter()
14666            .find_map(|batch| {
14667                let values = batch
14668                    .column_by_name(&histogram_column)?
14669                    .as_any()
14670                    .downcast_ref::<datafusion::arrow::array::StructArray>()?;
14671                (0..values.len()).find_map(|row| {
14672                    common_query::native_histogram::read_histogram(values, row).unwrap()
14673                })
14674            })
14675            .unwrap();
14676        assert_eq!(histogram.count, 1.0);
14677        let mut warnings = vec![];
14678        let mut infos = vec![];
14679        collector.append_to(&mut warnings, &mut infos);
14680        assert!(warnings.is_empty());
14681
14682        let collector = PromqlAnnotationCollector::default();
14683        let (mut planner, input) = mixed_direct_or(false).await;
14684        planner.promql_annotations = Some(collector.clone());
14685        let float_column = planner.ctx.field_columns[0].clone();
14686        let histogram_column = planner.ctx.field_columns[1].clone();
14687        let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, &param, &input).unwrap();
14688        let plan = LogicalPlanBuilder::from(input)
14689            .aggregate(vec![col("ts")], aggregate_exprs)
14690            .unwrap()
14691            .filter(
14692                planner
14693                    .mixed_aggregate_filter_expr(op, &float_column, &histogram_column)
14694                    .unwrap(),
14695            )
14696            .unwrap()
14697            .project([col(&float_column), col(&histogram_column), col("ts")])
14698            .unwrap()
14699            .build()
14700            .unwrap();
14701
14702        let (_, batches) = execute(plan, &build_query_engine_state()).await;
14703        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
14704        let mut warnings = vec![];
14705        let mut infos = vec![];
14706        collector.append_to(&mut warnings, &mut infos);
14707        assert_eq!(
14708            warnings,
14709            vec![
14710                "sum: dropped aggregation result containing both float and native histogram samples"
14711            ]
14712        );
14713    }
14714
14715    #[tokio::test]
14716    async fn test_mixed_or_sum_drops_incompatible_mixed_group() {
14717        let PromExpr::Aggregate(AggregateExpr { op, param, .. }) =
14718            parser::parse("sum(lhs)").unwrap()
14719        else {
14720            unreachable!()
14721        };
14722        let mut custom = direct_or_histogram();
14723        custom.schema = CUSTOM_BUCKETS_SCHEMA;
14724        custom.custom_values = vec![1.0];
14725        let collector = PromqlAnnotationCollector::default();
14726        let (mut planner, input) = mixed_aggregate_input(vec![direct_or_histogram(), custom]).await;
14727        planner.promql_annotations = Some(collector.clone());
14728        let float_column = planner.ctx.field_columns[0].clone();
14729        let histogram_column = planner.ctx.field_columns[1].clone();
14730        let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, &param, &input).unwrap();
14731        let plan = LogicalPlanBuilder::from(input)
14732            .aggregate(vec![col("ts")], aggregate_exprs)
14733            .unwrap()
14734            .filter(
14735                planner
14736                    .mixed_aggregate_filter_expr(op, &float_column, &histogram_column)
14737                    .unwrap(),
14738            )
14739            .unwrap()
14740            .project([col(&float_column), col(&histogram_column), col("ts")])
14741            .unwrap()
14742            .build()
14743            .unwrap();
14744
14745        let (_, batches) = execute(plan, &build_query_engine_state()).await;
14746        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
14747        let mut warnings = vec![];
14748        let mut infos = vec![];
14749        collector.append_to(&mut warnings, &mut infos);
14750        assert!(warnings.iter().any(|warning| {
14751            warning
14752                == "sum: dropped aggregation result containing both float and native histogram samples"
14753        }));
14754    }
14755
14756    #[tokio::test]
14757    async fn test_mixed_or_min_records_only_present_histograms() {
14758        let PromExpr::Aggregate(AggregateExpr { op, param, .. }) =
14759            parser::parse("min(lhs)").unwrap()
14760        else {
14761            unreachable!()
14762        };
14763        let expected_info = "min: dropped native histogram samples because this aggregation is not supported for native histograms";
14764
14765        for (histograms, expected_infos) in [
14766            (vec![], vec![]),
14767            (vec![direct_or_histogram()], vec![expected_info]),
14768        ] {
14769            let collector = PromqlAnnotationCollector::default();
14770            let (mut planner, input) = mixed_aggregate_input(histograms).await;
14771            planner.promql_annotations = Some(collector.clone());
14772            let float_column = planner.ctx.field_columns[0].clone();
14773            let histogram_column = planner.ctx.field_columns[1].clone();
14774            let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, &param, &input).unwrap();
14775            let plan = LogicalPlanBuilder::from(input)
14776                .aggregate(vec![col("ts")], aggregate_exprs)
14777                .unwrap()
14778                .filter(
14779                    planner
14780                        .mixed_ignored_histogram_filter_expr(op, &histogram_column)
14781                        .unwrap(),
14782                )
14783                .unwrap()
14784                .project([col(&float_column), col("ts")])
14785                .unwrap()
14786                .build()
14787                .unwrap();
14788
14789            let (_, batches) = execute(plan, &build_query_engine_state()).await;
14790            assert_eq!(values(&batches, &float_column), vec![1.25]);
14791            let mut warnings = vec![];
14792            let mut infos = vec![];
14793            collector.append_to(&mut warnings, &mut infos);
14794            assert!(warnings.is_empty());
14795            assert_eq!(infos, expected_infos);
14796        }
14797    }
14798
14799    #[tokio::test]
14800    async fn test_mixed_or_value_aliases_do_not_replace_labels() {
14801        let left = source(
14802            "lhs",
14803            false,
14804            1,
14805            vec![("job", Some("job")), ("k", Some("float"))],
14806            DirectOrValue::Float64(1.0),
14807        );
14808        let right = source(
14809            "rhs",
14810            false,
14811            1,
14812            vec![
14813                ("job", Some("job")),
14814                ("k", Some("histogram")),
14815                (greptime_value(), Some("value-label")),
14816            ],
14817            DirectOrValue::NativeHistogram(direct_or_histogram()),
14818        );
14819        let table_provider = build_test_table_provider_with_fields(
14820            &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
14821            &[],
14822        )
14823        .await;
14824        let mut planner = PromPlanner {
14825            table_provider,
14826            ctx: PromPlannerContext::default(),
14827            promql_annotations: None,
14828        };
14829        let left = LogicalPlanBuilder::from(scan(&left))
14830            .project(vec![
14831                col("ts"),
14832                col("job"),
14833                col("k"),
14834                col("v").alias(greptime_value()),
14835            ])
14836            .unwrap()
14837            .build()
14838            .unwrap();
14839        let left_context = direct_or_context("lhs", &["job", "k"], greptime_value());
14840        let right_context = direct_or_context("rhs", &["job", "k", greptime_value()], "v");
14841        let plan = planner
14842            .or_operator(
14843                left,
14844                scan(&right),
14845                left_context.tag_columns.iter().cloned().collect(),
14846                right_context.tag_columns.iter().cloned().collect(),
14847                left_context,
14848                right_context,
14849                &or_modifier("lhs or on(k) rhs"),
14850            )
14851            .unwrap();
14852
14853        assert_eq!(
14854            plan.schema()
14855                .field_with_name(None, greptime_value())
14856                .unwrap()
14857                .data_type(),
14858            &ArrowDataType::Utf8
14859        );
14860        assert!(
14861            planner
14862                .ctx
14863                .field_columns
14864                .iter()
14865                .all(|field| { field != greptime_value() && field != greptime_native_histogram() })
14866        );
14867        assert!(PromPlanner::field_columns_are_alternative_samples(
14868            plan.schema(),
14869            &planner.ctx.field_columns
14870        ));
14871        let (_, batches) = execute(plan, &build_query_engine_state()).await;
14872        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
14873        let labels = batches
14874            .iter()
14875            .flat_map(|batch| {
14876                batch
14877                    .column_by_name(greptime_value())
14878                    .unwrap()
14879                    .as_any()
14880                    .downcast_ref::<StringArray>()
14881                    .unwrap()
14882                    .iter()
14883                    .flatten()
14884            })
14885            .collect::<Vec<_>>();
14886        assert_eq!(labels, vec!["value-label"]);
14887    }
14888
14889    #[tokio::test]
14890    async fn test_mixed_or_routes_float_histogram_and_label_functions() {
14891        for (function, expected) in [("abs", 1.25), ("round", 1.0), ("histogram_count", 1.0)] {
14892            let (mut planner, input) = mixed_direct_or(false).await;
14893            let preserve_any_value = PromPlanner::field_columns_are_alternative_samples(
14894                input.schema(),
14895                &planner.ctx.field_columns,
14896            );
14897            let PromExpr::Call(call) = parser::parse(&format!("{function}(lhs)")).unwrap() else {
14898                unreachable!()
14899            };
14900            let state = build_query_engine_state();
14901            let (mut exprs, _) = planner
14902                .create_function_expr(&call.func, vec![], input.schema(), &state)
14903                .unwrap();
14904            exprs.insert(0, planner.create_time_index_column_expr().unwrap());
14905            exprs.extend(planner.create_tag_column_exprs().unwrap());
14906            let plan = LogicalPlanBuilder::from(input)
14907                .project(exprs)
14908                .unwrap()
14909                .filter(
14910                    planner
14911                        .create_empty_values_filter_expr(preserve_any_value)
14912                        .unwrap(),
14913                )
14914                .unwrap()
14915                .build()
14916                .unwrap();
14917            let (_, batches) = execute(plan, &state).await;
14918            let values = batches
14919                .iter()
14920                .flat_map(|batch| {
14921                    batch
14922                        .schema()
14923                        .fields()
14924                        .iter()
14925                        .position(|field| field.data_type() == &ArrowDataType::Float64)
14926                        .map(|index| {
14927                            batch
14928                                .column(index)
14929                                .as_any()
14930                                .downcast_ref::<Float64Array>()
14931                                .unwrap()
14932                                .iter()
14933                                .flatten()
14934                        })
14935                        .into_iter()
14936                        .flatten()
14937                })
14938                .collect::<Vec<_>>();
14939            assert_eq!(values, vec![expected], "{function}");
14940        }
14941
14942        let (mut planner, input) = mixed_direct_or(false).await;
14943        let preserve_any_value = PromPlanner::field_columns_are_alternative_samples(
14944            input.schema(),
14945            &planner.ctx.field_columns,
14946        );
14947        let PromExpr::Call(call) =
14948            parser::parse(r#"label_replace(lhs, "copy", "$1", "k", "(.*)")"#).unwrap()
14949        else {
14950            unreachable!()
14951        };
14952        let args = planner.create_function_args(&call.args.args).unwrap();
14953        let state = build_query_engine_state();
14954        let (mut exprs, _) = planner
14955            .create_function_expr(&call.func, args.literals, input.schema(), &state)
14956            .unwrap();
14957        exprs.insert(0, planner.create_time_index_column_expr().unwrap());
14958        exprs.extend(planner.create_tag_column_exprs().unwrap());
14959        let plan = LogicalPlanBuilder::from(input)
14960            .project(exprs)
14961            .unwrap()
14962            .filter(
14963                planner
14964                    .create_empty_values_filter_expr(preserve_any_value)
14965                    .unwrap(),
14966            )
14967            .unwrap()
14968            .build()
14969            .unwrap();
14970        let (_, batches) = execute(plan, &state).await;
14971        let sample_count = batches.iter().map(RecordBatch::num_rows).sum::<usize>();
14972        assert_eq!(sample_count, 2);
14973    }
14974
14975    /// Table provider with a single metric `cv_metric` holding three series at `ts=1000`:
14976    /// `k="a"` and `k="c"` carry the value 1.0, `k="b"` carries 2.0.
14977    async fn build_count_values_table_provider() -> DfTableSourceProvider {
14978        build_count_values_table_provider_with_values(&[1.0, 2.0, 1.0]).await
14979    }
14980
14981    /// Table provider with a single metric `cv_metric` holding one series per given value at
14982    /// `ts=1000` (`k` names the series, the sample value is the given value).
14983    async fn build_count_values_table_provider_with_values(
14984        values: &[f64],
14985    ) -> DfTableSourceProvider {
14986        build_count_values_table_provider_with_value_array(Arc::new(Float64Array::from(
14987            values.to_vec(),
14988        )))
14989        .await
14990    }
14991
14992    /// Like [`build_count_values_table_provider_with_values`], but with a caller provided value
14993    /// column, so tests can cover value columns that are not `Float64` (e.g. `BIGINT`).
14994    async fn build_count_values_table_provider_with_value_array(
14995        values: ArrayRef,
14996    ) -> DfTableSourceProvider {
14997        let value_data_type = ConcreteDataType::from_arrow_type(values.data_type());
14998        let catalog_list = MemoryCatalogManager::with_default_setup();
14999        let columns = vec![
15000            ColumnSchema::new("k".to_string(), ConcreteDataType::string_datatype(), false),
15001            ColumnSchema::new(
15002                "timestamp".to_string(),
15003                ConcreteDataType::timestamp_millisecond_datatype(),
15004                false,
15005            )
15006            .with_time_index(true),
15007            ColumnSchema::new(greptime_value().to_string(), value_data_type, true),
15008        ];
15009        let schema = Arc::new(Schema::new(columns));
15010        let table_meta = TableMetaBuilder::empty()
15011            .schema(schema.clone())
15012            .primary_key_indices(vec![0])
15013            .value_indices(vec![2])
15014            .next_column_id(1024)
15015            .build()
15016            .unwrap();
15017        let table_info = Arc::new(
15018            TableInfoBuilder::default()
15019                .table_id(3_001)
15020                .name("cv_metric")
15021                .meta(table_meta)
15022                .build()
15023                .unwrap(),
15024        );
15025        let batch = RecordBatch::try_new(
15026            schema.arrow_schema().clone(),
15027            vec![
15028                Arc::new(StringArray::from(
15029                    (0..values.len())
15030                        .map(|index| format!("k{index}"))
15031                        .collect::<Vec<_>>(),
15032                )),
15033                Arc::new(TimestampMillisecondArray::from(vec![1_000; values.len()])),
15034                values.clone(),
15035            ],
15036        )
15037        .unwrap();
15038        let backing = GreptimeMemTable::new_with_catalog(
15039            "cv_metric",
15040            GreptimeRecordBatch::from_df_record_batch(schema, batch),
15041            3_001,
15042            DEFAULT_CATALOG_NAME.to_string(),
15043            DEFAULT_SCHEMA_NAME.to_string(),
15044        );
15045        let table = Arc::new(Table::new(
15046            table_info,
15047            FilterPushDownType::Unsupported,
15048            backing.data_source(),
15049        ));
15050
15051        assert!(
15052            catalog_list
15053                .register_table_sync(RegisterTableRequest {
15054                    catalog: DEFAULT_CATALOG_NAME.to_string(),
15055                    schema: DEFAULT_SCHEMA_NAME.to_string(),
15056                    table_name: "cv_metric".to_string(),
15057                    table_id: 3_001,
15058                    table,
15059                })
15060                .is_ok()
15061        );
15062
15063        DfTableSourceProvider::new(
15064            catalog_list,
15065            false,
15066            QueryContext::arc(),
15067            DummyDecoder::arc(),
15068            false,
15069        )
15070    }
15071
15072    /// Collects `(label, value)` pairs of a `count_values` result, where `label` is the
15073    /// PromQL label generated by `count_values` and `value` is the aggregated sample value.
15074    ///
15075    /// The generated label holds the original sample value in PromQL's textual form, so it
15076    /// is asserted as a string: comparing it as a number would not catch formatting bugs
15077    /// (`1.0` instead of `1`, scientific notation, ...).
15078    fn count_values_rows<'a>(batches: &'a [RecordBatch], label: &str) -> Vec<(&'a str, f64)> {
15079        let mut rows = batches
15080            .iter()
15081            .flat_map(|batch| {
15082                // The aggregated value is the only numeric column that is not the generated label.
15083                let value_index = batch
15084                    .schema()
15085                    .fields()
15086                    .iter()
15087                    .position(|field| {
15088                        field.name() != label
15089                            && matches!(
15090                                field.data_type(),
15091                                ArrowDataType::Float64
15092                                    | ArrowDataType::Int64
15093                                    | ArrowDataType::UInt64
15094                            )
15095                    })
15096                    .expect("no aggregated value column");
15097                let labels = batch
15098                    .column_by_name(label)
15099                    .expect("no generated label column")
15100                    .as_any()
15101                    .downcast_ref::<StringArray>()
15102                    .expect("the generated label must be a string column");
15103                let values = datafusion::arrow::compute::cast(
15104                    batch.column(value_index),
15105                    &ArrowDataType::Float64,
15106                )
15107                .unwrap();
15108                let values = values.as_any().downcast_ref::<Float64Array>().unwrap();
15109                labels
15110                    .iter()
15111                    .zip(values.iter())
15112                    .map(|(label, value)| (label.unwrap(), value.unwrap()))
15113                    .collect::<Vec<_>>()
15114            })
15115            .collect::<Vec<_>>();
15116        rows.sort_by(|left, right| left.0.cmp(right.0).then(left.1.total_cmp(&right.1)));
15117        rows
15118    }
15119
15120    /// Asserts that a `count_values` result holds one sample per label set and evaluation
15121    /// timestamp: Prometheus groups by the generated label, so a timestamp must never repeat
15122    /// a label set (that would mean the samples were still grouped by the overwritten input
15123    /// label).
15124    fn assert_unique_label_set_per_timestamp(batches: &[RecordBatch], label: &str) {
15125        let mut seen = HashMap::<i64, HashSet<String>>::new();
15126        for batch in batches {
15127            let timestamp_index = batch
15128                .schema()
15129                .fields()
15130                .iter()
15131                .position(|field| matches!(field.data_type(), ArrowDataType::Timestamp(..)))
15132                .expect("no timestamp column");
15133            let timestamps = batch
15134                .column(timestamp_index)
15135                .as_any()
15136                .downcast_ref::<TimestampMillisecondArray>()
15137                .expect("timestamp column is not a millisecond timestamp");
15138            let labels = batch
15139                .column_by_name(label)
15140                .expect("no generated label column")
15141                .as_any()
15142                .downcast_ref::<StringArray>()
15143                .expect("the generated label must be a string column");
15144            for (timestamp, label) in timestamps.iter().zip(labels.iter()) {
15145                let timestamp = timestamp.unwrap();
15146                let label = label.unwrap();
15147                assert!(
15148                    seen.entry(timestamp).or_default().insert(label.to_string()),
15149                    "duplicated label set `{label}` at timestamp {timestamp}"
15150                );
15151            }
15152        }
15153    }
15154
15155    #[tokio::test]
15156    async fn test_count_values_generated_label_survives_enclosing_expr() {
15157        // https://github.com/GreptimeTeam/greptimedb/issues/9181
15158        for (case, label) in [
15159            (r#"count_values("v", prometheus_tsdb_head_series)"#, "v"),
15160            (
15161                r#"abs(count_values("v", prometheus_tsdb_head_series))"#,
15162                "v",
15163            ),
15164            (
15165                r#"round(count_values("v", prometheus_tsdb_head_series))"#,
15166                "v",
15167            ),
15168            (r#"count_values("v", prometheus_tsdb_head_series) + 1"#, "v"),
15169            (
15170                r#"topk(1, count_values("v", prometheus_tsdb_head_series))"#,
15171                "v",
15172            ),
15173            (
15174                r#"sum by (v) (count_values("v", prometheus_tsdb_head_series))"#,
15175                "v",
15176            ),
15177            (
15178                r#"label_replace(count_values("v", prometheus_tsdb_head_series), "vcopy", "$1", "v", "(.*)")"#,
15179                "v",
15180            ),
15181            (
15182                r#"count_values("v", prometheus_tsdb_head_series) by (ip) + 1"#,
15183                "v",
15184            ),
15185            // The generated label overwrites an input label with the same name.
15186            (
15187                r#"count_values("ip", prometheus_tsdb_head_series) by (ip)"#,
15188                "ip",
15189            ),
15190            (
15191                r#"count_values("ip", prometheus_tsdb_head_series) by (ip) + 1"#,
15192                "ip",
15193            ),
15194        ] {
15195            let plan = PromPlanner::stmt_to_plan(
15196                build_test_table_provider_with_fields(
15197                    &[(
15198                        DEFAULT_SCHEMA_NAME.to_string(),
15199                        "prometheus_tsdb_head_series".to_string(),
15200                    )],
15201                    &["ip"],
15202                )
15203                .await,
15204                &build_eval_stmt(case),
15205                &build_query_engine_state(),
15206            )
15207            .await
15208            .unwrap();
15209
15210            let label_columns = plan
15211                .schema()
15212                .fields()
15213                .iter()
15214                .filter(|field| field.name() == label)
15215                .count();
15216            assert_eq!(
15217                label_columns,
15218                1,
15219                "{case}: the `{label}` label must survive: {}",
15220                plan.display_indent()
15221            );
15222        }
15223    }
15224
15225    #[tokio::test]
15226    async fn test_count_values_generated_label_in_enclosing_expr_execute() {
15227        // https://github.com/GreptimeTeam/greptimedb/issues/9181
15228        let state = build_query_engine_state();
15229        // (query, generated label, expected `(label value, aggregated value)` pairs)
15230        for (query, label, expected) in [
15231            (
15232                r#"count_values("v", cv_metric)"#,
15233                "v",
15234                vec![("1", 2.0), ("2", 1.0)],
15235            ),
15236            (
15237                r#"abs(count_values("v", cv_metric))"#,
15238                "v",
15239                vec![("1", 2.0), ("2", 1.0)],
15240            ),
15241            (
15242                r#"round(count_values("v", cv_metric))"#,
15243                "v",
15244                vec![("1", 2.0), ("2", 1.0)],
15245            ),
15246            (
15247                r#"count_values("v", cv_metric) + 1"#,
15248                "v",
15249                vec![("1", 3.0), ("2", 2.0)],
15250            ),
15251            (
15252                r#"sum by (v) (count_values("v", cv_metric))"#,
15253                "v",
15254                vec![("1", 2.0), ("2", 1.0)],
15255            ),
15256            (
15257                r#"topk(10, count_values("v", cv_metric))"#,
15258                "v",
15259                vec![("1", 2.0), ("2", 1.0)],
15260            ),
15261            // The generated label overwrites the input label with the same name, and the
15262            // samples are grouped by the generated label only: `{k="1"}` holds the two
15263            // samples of value `1.0` instead of one row per (overwritten label, value).
15264            (
15265                r#"count_values("k", cv_metric) by (k)"#,
15266                "k",
15267                vec![("1", 2.0), ("2", 1.0)],
15268            ),
15269        ] {
15270            let plan = PromPlanner::stmt_to_plan(
15271                build_count_values_table_provider().await,
15272                &operator_eval_stmt(query),
15273                &state,
15274            )
15275            .await
15276            .unwrap_or_else(|err| panic!("{query}: {err}"));
15277
15278            assert_eq!(
15279                plan.schema()
15280                    .fields()
15281                    .iter()
15282                    .filter(|field| field.name() == label)
15283                    .count(),
15284                1,
15285                "{query}: {}",
15286                plan.display_indent()
15287            );
15288
15289            let (_, batches) = execute(plan, &state).await;
15290            assert_eq!(count_values_rows(&batches, label), expected, "{query}");
15291            assert_unique_label_set_per_timestamp(&batches, label);
15292        }
15293    }
15294
15295    #[tokio::test]
15296    async fn test_count_values_label_is_prometheus_formatted_value() {
15297        // PromQL materializes the generated label with `strconv.FormatFloat(value, 'f', -1, 64)`:
15298        // the shortest decimal form of the sample value without an exponent. The label is a
15299        // label, so it must be a string column holding exactly that text: arrow's
15300        // `Float64 -> Utf8` cast would render `1`/`200`/`1e21` as `1.0`/`200.0`/`1e21`.
15301        let state = build_query_engine_state();
15302        // Samples are grouped by that formatted text, exactly like Prometheus groups by the
15303        // generated label: `-0.0` and `0.0` are two series (`-0` and `0`), while values that
15304        // round to the same text share one group. `0.0` formats as "0".
15305        let plan = PromPlanner::stmt_to_plan(
15306            build_count_values_table_provider_with_values(&[
15307                -0.0, 0.0, 1.0, 0.5, 200.0, 1e21, 1e-7, 2.5,
15308            ])
15309            .await,
15310            &operator_eval_stmt(r#"count_values("v", cv_metric)"#),
15311            &state,
15312        )
15313        .await
15314        .unwrap();
15315
15316        let (_, batches) = execute(plan, &state).await;
15317        assert_unique_label_set_per_timestamp(&batches, "v");
15318        let mut labels = count_values_rows(&batches, "v")
15319            .into_iter()
15320            .map(|(label, _)| label)
15321            .collect::<Vec<_>>();
15322        labels.sort();
15323        assert_eq!(
15324            labels,
15325            vec![
15326                "-0",
15327                "0",
15328                "0.0000001",
15329                "0.5",
15330                "1",
15331                "1000000000000000000000",
15332                "2.5",
15333                "200",
15334            ]
15335        );
15336    }
15337
15338    #[tokio::test]
15339    async fn test_count_values_groups_by_formatted_value_for_bigint_input() {
15340        // The grouping key of `count_values` is the formatted sample value, not the raw input
15341        // value. Two `BIGINT` values that differ below the `Float64` precision (`2^53` and
15342        // `2^53 + 1`) cast and format to the same label, so they must share one group and one
15343        // count, exactly like Prometheus, which groups by the generated label text.
15344        let state = build_query_engine_state();
15345        let plan = PromPlanner::stmt_to_plan(
15346            build_count_values_table_provider_with_value_array(Arc::new(Int64Array::from(vec![
15347                9_007_199_254_740_992_i64,
15348                9_007_199_254_740_993_i64,
15349            ])))
15350            .await,
15351            &operator_eval_stmt(r#"count_values("v", cv_metric)"#),
15352            &state,
15353        )
15354        .await
15355        .unwrap();
15356
15357        let (_, batches) = execute(plan, &state).await;
15358        // One timestamp must never carry the same label set twice.
15359        assert_unique_label_set_per_timestamp(&batches, "v");
15360        assert_eq!(
15361            count_values_rows(&batches, "v"),
15362            vec![("9007199254740992", 2.0)]
15363        );
15364    }
15365}