query/range_select/
plan_rewrite.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
15use std::collections::BTreeSet;
16use std::sync::Arc;
17use std::time::Duration;
18
19use arrow_schema::DataType;
20use async_recursion::async_recursion;
21use catalog::table_source::DfTableSourceProvider;
22use chrono::{DateTime, Utc};
23use common_time::interval::{MS_PER_DAY, NANOS_PER_MILLI};
24use common_time::timestamp::TimeUnit;
25use common_time::{IntervalDayTime, IntervalMonthDayNano, IntervalYearMonth, Timestamp, Timezone};
26use datafusion::datasource::DefaultTableSource;
27use datafusion::prelude::Column;
28use datafusion::scalar::ScalarValue;
29use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter};
30use datafusion_common::{DFSchema, DataFusionError, Result as DFResult};
31use datafusion_expr::expr::WildcardOptions;
32use datafusion_expr::simplify::SimplifyContext;
33use datafusion_expr::{
34    Aggregate, Analyze, Cast, Distinct, DistinctOn, Explain, Expr, ExprSchemable, Extension,
35    Literal, LogicalPlan, LogicalPlanBuilder, Projection,
36};
37use datafusion_optimizer::simplify_expressions::ExprSimplifier;
38use datatypes::prelude::ConcreteDataType;
39use datatypes::schema::TIME_INDEX_KEY;
40use promql_parser::util::parse_duration;
41use session::context::QueryContextRef;
42use snafu::{OptionExt, ResultExt, ensure};
43use table::table::adapter::DfTableProviderAdapter;
44
45use crate::error::{
46    CatalogSnafu, RangeQuerySnafu, Result, TimeIndexNotFoundSnafu, UnknownTableSnafu,
47};
48use crate::plan::ExtractExpr;
49use crate::range_select::plan::{Fill, RangeFn, RangeSelect};
50
51/// `RangeExprRewriter` will recursively search certain `Expr`, find all `range_fn` scalar udf contained in `Expr`,
52/// and collect the information required by the RangeSelect query,
53/// and finally modify the `range_fn` scalar udf to an ordinary column field.
54pub struct RangeExprRewriter<'a> {
55    input_plan: &'a Arc<LogicalPlan>,
56    align: Duration,
57    align_to: i64,
58    by: Vec<Expr>,
59    /// Use `BTreeSet` to avoid in case like `avg(a) RANGE '5m' + avg(a) RANGE '5m'`, duplicate range expr `avg(a) RANGE '5m'` be calculate twice
60    range_fn: BTreeSet<RangeFn>,
61    sub_aggr: &'a Aggregate,
62    query_ctx: &'a QueryContextRef,
63}
64
65impl RangeExprRewriter<'_> {
66    pub fn get_range_expr(&self, args: &[Expr], i: usize) -> DFResult<Expr> {
67        match args.get(i) {
68            Some(Expr::Column(column)) => {
69                let index = self.sub_aggr.schema.index_of_column(column)?;
70                let len = self.sub_aggr.group_expr.len();
71                self.sub_aggr
72                    .aggr_expr
73                    .get(index - len)
74                    .cloned()
75                    .ok_or(DataFusionError::Plan(
76                        "Range expr not found in underlying Aggregate Plan".into(),
77                    ))
78            }
79            Some(Expr::Alias(alias)) => {
80                self.get_range_expr(std::slice::from_ref(alias.expr.as_ref()), 0)
81            }
82            other => Err(dispose_parse_error(other)),
83        }
84    }
85}
86
87#[inline]
88fn dispose_parse_error(expr: Option<&Expr>) -> DataFusionError {
89    DataFusionError::Plan(
90        expr.map(|x| {
91            format!(
92                "Illegal argument `{}` in range select query",
93                x.schema_name()
94            )
95        })
96        .unwrap_or("Missing argument in range select query".into()),
97    )
98}
99
100fn parse_str_expr(args: &[Expr], i: usize) -> DFResult<&str> {
101    match args.get(i) {
102        Some(Expr::Literal(ScalarValue::Utf8(Some(str)), _)) => Ok(str.as_str()),
103        other => Err(dispose_parse_error(other)),
104    }
105}
106
107fn parse_expr_to_string(args: &[Expr], i: usize) -> DFResult<String> {
108    match args.get(i) {
109        Some(Expr::Literal(ScalarValue::Utf8(Some(str)), _)) => Ok(str.clone()),
110        Some(expr) => Ok(expr.schema_name().to_string()),
111        None => Err(dispose_parse_error(None)),
112    }
113}
114
115/// Parse a duraion expr:
116/// 1. duration string (e.g. `'1h'`)
117/// 2. Interval expr (e.g. `INTERVAL '1 year 3 hours 20 minutes'`)
118/// 3. An interval expr can be evaluated at the logical plan stage (e.g. `INTERVAL '2' day - INTERVAL '1' day`)
119fn parse_duration_expr(args: &[Expr], i: usize) -> DFResult<Duration> {
120    match args.get(i) {
121        Some(Expr::Literal(ScalarValue::Utf8(Some(str)), _)) => {
122            parse_duration(str).map_err(DataFusionError::Plan)
123        }
124        Some(expr) => {
125            let ms = evaluate_expr_to_millisecond(args, i, true, None)?;
126            if ms <= 0 {
127                return Err(dispose_parse_error(Some(expr)));
128            }
129            Ok(Duration::from_millis(ms as u64))
130        }
131        None => Err(dispose_parse_error(None)),
132    }
133}
134
135/// Evaluate a time calculation expr, case like:
136/// 1. `INTERVAL '1' day + INTERVAL '1 year 2 hours 3 minutes'`
137/// 2. `now() - INTERVAL '1' day` (when `interval_only==false`)
138///
139/// Output a millisecond timestamp
140///
141/// if `interval_only==true`, only accept expr with all interval type (case 2 will return a error)
142fn evaluate_expr_to_millisecond(
143    args: &[Expr],
144    i: usize,
145    interval_only: bool,
146    scheduled_time: Option<DateTime<Utc>>,
147) -> DFResult<i64> {
148    let Some(expr) = args.get(i) else {
149        return Err(dispose_parse_error(None));
150    };
151    if interval_only && !interval_only_in_expr(expr) {
152        return Err(dispose_parse_error(Some(expr)));
153    }
154    let info = match scheduled_time {
155        Some(dt) => SimplifyContext::default().with_query_execution_start_time(Some(dt)),
156        None => SimplifyContext::default().with_current_time(),
157    };
158    let simplify_expr = ExprSimplifier::new(info).simplify(expr.clone())?;
159    match simplify_expr {
160        Expr::Literal(ScalarValue::TimestampNanosecond(ts_nanos, _), _)
161        | Expr::Literal(ScalarValue::DurationNanosecond(ts_nanos), _) => {
162            ts_nanos.map(|v| v / 1_000_000)
163        }
164        Expr::Literal(ScalarValue::TimestampMicrosecond(ts_micros, _), _)
165        | Expr::Literal(ScalarValue::DurationMicrosecond(ts_micros), _) => {
166            ts_micros.map(|v| v / 1_000)
167        }
168        Expr::Literal(ScalarValue::TimestampMillisecond(ts_millis, _), _)
169        | Expr::Literal(ScalarValue::DurationMillisecond(ts_millis), _) => ts_millis,
170        Expr::Literal(ScalarValue::TimestampSecond(ts_secs, _), _)
171        | Expr::Literal(ScalarValue::DurationSecond(ts_secs), _) => ts_secs.map(|v| v * 1_000),
172        // We don't support interval with months as days in a month is unclear.
173        Expr::Literal(ScalarValue::IntervalYearMonth(interval), _) => interval
174            .map(|v| {
175                let interval = IntervalYearMonth::from_i32(v);
176                if interval.months != 0 {
177                    return Err(DataFusionError::Plan(format!(
178                        "Year or month interval is not allowed in range query: {}",
179                        expr.schema_name()
180                    )));
181                }
182
183                Ok(0)
184            })
185            .transpose()?,
186        Expr::Literal(ScalarValue::IntervalDayTime(interval), _) => interval.map(|v| {
187            let interval = IntervalDayTime::from(v);
188            interval.as_millis()
189        }),
190        Expr::Literal(ScalarValue::IntervalMonthDayNano(interval), _) => interval
191            .map(|v| {
192                let interval = IntervalMonthDayNano::from(v);
193                if interval.months != 0 {
194                    return Err(DataFusionError::Plan(format!(
195                        "Year or month interval is not allowed in range query: {}",
196                        expr.schema_name()
197                    )));
198                }
199
200                Ok(interval.days as i64 * MS_PER_DAY + interval.nanoseconds / NANOS_PER_MILLI)
201            })
202            .transpose()?,
203        _ => None,
204    }
205    .ok_or_else(|| {
206        DataFusionError::Plan(format!(
207            "{} is not a expr can be evaluate and use in range query",
208            expr.schema_name()
209        ))
210    })
211}
212
213/// Parse the `align to` clause and return a UTC timestamp with unit of millisecond,
214/// which is used as the basis for dividing time slot during the align operation.
215/// 1. NOW: align to current execute time
216/// 2. Timestamp string: align to specific timestamp
217/// 3. An expr can be evaluated at the logical plan stage (e.g. `now() - INTERVAL '1' day`)
218/// 4. leave empty (as Default Option): align to unix epoch 0 (timezone aware)
219fn parse_align_to(
220    args: &[Expr],
221    i: usize,
222    timezone: Option<&Timezone>,
223    scheduled_time: Option<DateTime<Utc>>,
224) -> DFResult<i64> {
225    let Ok(s) = parse_str_expr(args, i) else {
226        return evaluate_expr_to_millisecond(args, i, false, scheduled_time);
227    };
228    let upper = s.to_uppercase();
229    match upper.as_str() {
230        "NOW" => {
231            return Ok(scheduled_time
232                .map(|dt| dt.timestamp_millis())
233                .unwrap_or_else(|| Timestamp::current_millis().value()));
234        }
235        // default align to unix epoch 0 (timezone aware)
236        "" => return Ok(timezone.map(|tz| tz.local_minus_utc() * 1000).unwrap_or(0)),
237        _ => (),
238    }
239
240    Timestamp::from_str(s, timezone)
241        .map_err(|e| {
242            DataFusionError::Plan(format!(
243                "Illegal `align to` argument `{}` in range select query, can't be parse as NOW/CALENDAR/Timestamp, error: {}",
244                s, e
245            ))
246        })?.convert_to(TimeUnit::Millisecond).map(|x|x.value()).ok_or(DataFusionError::Plan(format!(
247            "Illegal `align to` argument `{}` in range select query, can't be convert to a valid Timestamp",
248            s
249        ))
250        )
251}
252
253fn parse_expr_list(args: &[Expr], start: usize, len: usize) -> DFResult<Vec<Expr>> {
254    let mut outs = Vec::with_capacity(len);
255    for i in start..start + len {
256        outs.push(match &args.get(i) {
257            Some(
258                Expr::Column(_)
259                | Expr::Literal(_, _)
260                | Expr::BinaryExpr(_)
261                | Expr::ScalarFunction(_),
262            ) => args[i].clone(),
263            Some(Expr::Alias(alias)) if matches!(*alias.expr, Expr::ScalarFunction(_)) => {
264                args[i].clone()
265            }
266            other => {
267                return Err(dispose_parse_error(*other));
268            }
269        });
270    }
271    Ok(outs)
272}
273
274macro_rules! inconsistent_check {
275    ($self: ident.$name: ident, $cond: expr) => {
276        if $cond && $self.$name != $name {
277            return Err(DataFusionError::Plan(
278                concat!(
279                    "Inconsistent ",
280                    stringify!($name),
281                    " given in Range Function Rewrite"
282                )
283                .into(),
284            ));
285        } else {
286            $self.$name = $name;
287        }
288    };
289}
290
291impl TreeNodeRewriter for RangeExprRewriter<'_> {
292    type Node = Expr;
293
294    fn f_down(&mut self, node: Expr) -> DFResult<Transformed<Expr>> {
295        if let Expr::ScalarFunction(func) = &node
296            && func.name() == "range_fn"
297        {
298            // `range_fn(func, range, fill, byc, [byv], align, to)`
299            // `[byv]` are variadic arguments, byc indicate the length of arguments
300            let range_expr = self.get_range_expr(&func.args, 0)?;
301            let range = parse_duration_expr(&func.args, 1)?;
302            let byc = str::parse::<usize>(parse_str_expr(&func.args, 3)?)
303                .map_err(|e| DataFusionError::Plan(e.to_string()))?;
304            let by = parse_expr_list(&func.args, 4, byc)?;
305            let align = parse_duration_expr(&func.args, byc + 4)?;
306            let scheduled_time =
307                crate::options::parse_scheduled_time_datetime(&self.query_ctx.extensions())
308                    .map_err(|err| DataFusionError::Plan(err.to_string()))?;
309            let align_to = parse_align_to(
310                &func.args,
311                byc + 5,
312                Some(&self.query_ctx.timezone()),
313                scheduled_time,
314            )?;
315            let mut data_type = range_expr.get_type(self.input_plan.schema())?;
316            let mut need_cast = false;
317            let fill = Fill::try_from_str(parse_str_expr(&func.args, 2)?, &data_type)?;
318            if matches!(fill, Some(Fill::Linear)) && data_type.is_integer() {
319                data_type = DataType::Float64;
320                need_cast = true;
321            }
322            inconsistent_check!(self.by, !self.by.is_empty());
323            inconsistent_check!(self.align, self.align != Duration::default());
324            inconsistent_check!(self.align_to, self.align_to != 0);
325            let range_fn = RangeFn {
326                name: if let Some(fill) = &fill {
327                    format!(
328                        "{} RANGE {} FILL {}",
329                        range_expr.schema_name(),
330                        parse_expr_to_string(&func.args, 1)?,
331                        fill
332                    )
333                } else {
334                    format!(
335                        "{} RANGE {}",
336                        range_expr.schema_name(),
337                        parse_expr_to_string(&func.args, 1)?,
338                    )
339                },
340                data_type,
341                expr: range_expr,
342                range,
343                fill,
344                need_cast,
345            };
346            let alias = Expr::Column(Column::from_name(range_fn.name.clone()));
347            self.range_fn.insert(range_fn);
348            return Ok(Transformed::yes(alias));
349        }
350        Ok(Transformed::no(node))
351    }
352}
353
354/// In order to implement RangeSelect query like `avg(field_0) RANGE '5m' FILL NULL`,
355/// All RangeSelect query items are converted into udf scalar function in sql parse stage, with format like `range_fn(avg(field_0), .....)`.
356/// `range_fn` contains all the parameters we need to execute RangeSelect.
357/// In order to correctly execute the query process of range select, we need to modify the query plan generated by datafusion.
358/// We need to recursively find the entire LogicalPlan, and find all `range_fn` scalar udf contained in the project plan,
359/// collecting info we need to generate RangeSelect Query LogicalPlan and rewrite th original LogicalPlan.
360pub struct RangePlanRewriter {
361    table_provider: DfTableSourceProvider,
362    query_ctx: QueryContextRef,
363}
364
365impl RangePlanRewriter {
366    pub fn new(table_provider: DfTableSourceProvider, query_ctx: QueryContextRef) -> Self {
367        Self {
368            table_provider,
369            query_ctx,
370        }
371    }
372
373    pub async fn rewrite(&mut self, plan: LogicalPlan) -> Result<LogicalPlan> {
374        match self.rewrite_logical_plan(&plan).await? {
375            Some(new_plan) => Ok(new_plan),
376            None => Ok(plan),
377        }
378    }
379
380    #[async_recursion]
381    async fn rewrite_logical_plan(&mut self, plan: &LogicalPlan) -> Result<Option<LogicalPlan>> {
382        let inputs = plan.inputs();
383        let mut new_inputs = Vec::with_capacity(inputs.len());
384        for input in &inputs {
385            new_inputs.push(self.rewrite_logical_plan(input).await?)
386        }
387        match plan {
388            LogicalPlan::Projection(Projection { expr, input, .. })
389                if have_range_in_exprs(expr) =>
390            {
391                let (aggr_plan, input) = if let LogicalPlan::Aggregate(aggr) = input.as_ref() {
392                    // Expr like `rate(max(a) RANGE '6m') RANGE '6m'` have legal syntax but illegal semantic.
393                    if have_range_in_exprs(&aggr.aggr_expr) {
394                        return RangeQuerySnafu {
395                            msg: "Nest Range Query is not allowed",
396                        }
397                        .fail();
398                    }
399                    (aggr, aggr.input.clone())
400                } else {
401                    return RangeQuerySnafu {
402                        msg: "Window functions is not allowed in Range Query",
403                    }
404                    .fail();
405                };
406                let query_ctx = self.query_ctx.clone();
407                let mut range_rewriter = RangeExprRewriter {
408                    input_plan: &input,
409                    align: Duration::default(),
410                    align_to: 0,
411                    by: vec![],
412                    range_fn: BTreeSet::new(),
413                    sub_aggr: aggr_plan,
414                    query_ctx: &query_ctx,
415                };
416                let new_expr = expr
417                    .iter()
418                    .map(|expr| expr.clone().rewrite(&mut range_rewriter).map(|x| x.data))
419                    .collect::<DFResult<Vec<_>>>()?;
420                let need_default_by = range_rewriter.by.is_empty();
421                let (time_index, default_by) =
422                    self.get_index_by(input.schema(), need_default_by).await?;
423                if need_default_by {
424                    range_rewriter.by = default_by;
425                }
426                let range_select = RangeSelect::try_new(
427                    input.clone(),
428                    range_rewriter.range_fn.into_iter().collect(),
429                    range_rewriter.align,
430                    range_rewriter.align_to,
431                    time_index,
432                    range_rewriter.by,
433                    &new_expr,
434                )?;
435                let no_additional_project = range_select.schema_project.is_some();
436                let range_plan = LogicalPlan::Extension(Extension {
437                    node: Arc::new(range_select),
438                });
439                if no_additional_project {
440                    Ok(Some(range_plan))
441                } else {
442                    let project_plan = LogicalPlanBuilder::from(range_plan)
443                        .project(new_expr)
444                        .and_then(|x| x.build())?;
445                    Ok(Some(project_plan))
446                }
447            }
448            _ => {
449                if new_inputs.iter().any(|x| x.is_some()) {
450                    let inputs: Vec<LogicalPlan> = new_inputs
451                        .into_iter()
452                        .zip(inputs)
453                        .map(|(x, y)| match x {
454                            Some(plan) => plan,
455                            None => y.clone(),
456                        })
457                        .collect();
458                    // Due to the limitations of Datafusion, for `LogicalPlan::Analyze` and `LogicalPlan::Explain`,
459                    // directly using the method `with_new_inputs` to rebuild a new `LogicalPlan` will cause an error,
460                    // so here we directly use the `LogicalPlanBuilder` to build a new plan.
461                    let plan = match plan {
462                        LogicalPlan::Analyze(Analyze { verbose, .. }) => {
463                            ensure!(
464                                inputs.len() == 1,
465                                RangeQuerySnafu {
466                                    msg: "Illegal subplan nums when rewrite Analyze logical plan",
467                                }
468                            );
469                            LogicalPlanBuilder::from(inputs[0].clone())
470                                .explain(*verbose, true)?
471                                .build()
472                        }
473                        LogicalPlan::Explain(Explain { verbose, .. }) => {
474                            ensure!(
475                                inputs.len() == 1,
476                                RangeQuerySnafu {
477                                    msg: "Illegal subplan nums when rewrite Explain logical plan",
478                                }
479                            );
480                            LogicalPlanBuilder::from(inputs[0].clone())
481                                .explain(*verbose, false)?
482                                .build()
483                        }
484                        LogicalPlan::Distinct(Distinct::On(DistinctOn {
485                            on_expr,
486                            select_expr,
487                            sort_expr,
488                            ..
489                        })) => {
490                            ensure!(
491                                inputs.len() == 1,
492                                RangeQuerySnafu {
493                                    msg: "Illegal subplan nums when rewrite DistinctOn logical plan",
494                                }
495                            );
496                            LogicalPlanBuilder::from(inputs[0].clone())
497                                .distinct_on(
498                                    on_expr.clone(),
499                                    select_expr.clone(),
500                                    sort_expr.clone(),
501                                )?
502                                .build()
503                        }
504                        _ => plan.with_new_exprs(plan.expressions_consider_join(), inputs),
505                    }?;
506                    Ok(Some(plan))
507                } else {
508                    Ok(None)
509                }
510            }
511        }
512    }
513
514    /// Finds the time index column and default row-key grouping from the input schema.
515    ///
516    /// Returns `(time_index, [row_columns])` to the rewriter. If the user omits `BY`,
517    /// `[row_columns]` is used as the default time-series grouping.
518    ///
519    /// For derived inputs such as subqueries, joins, or set operations, the source table
520    /// qualifier may no longer resolve back to a table provider. In that case we can still
521    /// recover the time index from column metadata, but we cannot safely reconstruct the
522    /// original row-key columns, so omitted `BY` must be rejected by the caller.
523    async fn get_index_by(
524        &mut self,
525        schema: &Arc<DFSchema>,
526        need_default_by: bool,
527    ) -> Result<(Expr, Vec<Expr>)> {
528        #[allow(deprecated)]
529        let mut time_index_expr = Expr::Wildcard {
530            qualifier: None,
531            options: Box::new(WildcardOptions::default()),
532        };
533        let mut default_by = vec![];
534        let metadata_time_index_expr = (0..schema.fields().len()).find_map(|i| {
535            let (qualifier, field) = schema.qualified_field(i);
536            if field.metadata().contains_key(TIME_INDEX_KEY)
537                && matches!(field.data_type(), DataType::Timestamp(_, _))
538            {
539                Some(Expr::Column(Column::new(
540                    qualifier.cloned(),
541                    field.name().clone(),
542                )))
543            } else {
544                None
545            }
546        });
547        for i in 0..schema.fields().len() {
548            let (qualifier, _) = schema.qualified_field(i);
549            if let Some(table_ref) = qualifier {
550                let table_source = match self.table_provider.resolve_table(table_ref.clone()).await
551                {
552                    Ok(table_source) => table_source,
553                    Err(error) => {
554                        // `TableNotExist` here usually means the qualifier now refers to a derived
555                        // input instead of a base table. We can still salvage the time index from
556                        // field metadata, but only when such metadata is present.
557                        if matches!(&error, catalog::error::Error::TableNotExist { .. })
558                            && metadata_time_index_expr.is_some()
559                        {
560                            continue;
561                        }
562                        return Err(error).context(CatalogSnafu);
563                    }
564                };
565                let table = table_source
566                    .as_any()
567                    .downcast_ref::<DefaultTableSource>()
568                    .context(UnknownTableSnafu)?
569                    .table_provider
570                    .as_any()
571                    .downcast_ref::<DfTableProviderAdapter>()
572                    .context(UnknownTableSnafu)?
573                    .table();
574                let schema = table.schema();
575                let time_index_column =
576                    schema
577                        .timestamp_column()
578                        .with_context(|| TimeIndexNotFoundSnafu {
579                            table: table_ref.to_string(),
580                        })?;
581                // assert time_index's datatype is timestamp
582                if let ConcreteDataType::Timestamp(_) = time_index_column.data_type {
583                    default_by = table
584                        .table_info()
585                        .meta
586                        .row_key_column_names()
587                        .map(|key| Expr::Column(Column::new(Some(table_ref.clone()), key)))
588                        .collect();
589                    // If the user does not specify a primary key when creating a table,
590                    // then by default all data will be aggregated into one time series,
591                    // which is equivalent to using `by(1)` in SQL
592                    if default_by.is_empty() {
593                        default_by = vec![1.lit()];
594                    }
595                    time_index_expr = Expr::Column(Column::new(
596                        Some(table_ref.clone()),
597                        time_index_column.name.clone(),
598                    ));
599                }
600            }
601        }
602        #[allow(deprecated)]
603        if matches!(time_index_expr, Expr::Wildcard { .. })
604            && let Some(expr) = metadata_time_index_expr
605        {
606            common_telemetry::debug!(
607                "Range query falling back to time-index metadata for derived input schema: {}",
608                schema
609            );
610            ensure!(
611                !need_default_by,
612                RangeQuerySnafu {
613                    msg: "Cannot infer default BY columns from derived range query input"
614                }
615            );
616            time_index_expr = expr;
617        }
618        #[allow(deprecated)]
619        if matches!(time_index_expr, Expr::Wildcard { .. }) {
620            TimeIndexNotFoundSnafu {
621                table: schema.to_string(),
622            }
623            .fail()
624        } else {
625            Ok((time_index_expr, default_by))
626        }
627    }
628}
629
630fn have_range_in_exprs(exprs: &[Expr]) -> bool {
631    exprs.iter().any(|expr| {
632        let mut find_range = false;
633        let _ = expr.apply(|expr| {
634            Ok(match expr {
635                Expr::ScalarFunction(func) if func.name() == "range_fn" => {
636                    find_range = true;
637                    TreeNodeRecursion::Stop
638                }
639                _ => TreeNodeRecursion::Continue,
640            })
641        });
642        find_range
643    })
644}
645
646fn interval_only_in_expr(expr: &Expr) -> bool {
647    let mut all_interval = true;
648    let _ = expr.apply(|expr| {
649        // A cast expression for an interval.
650        if matches!(
651            expr,
652            Expr::Cast(Cast{
653                expr,
654                data_type: DataType::Interval(_)
655            }) if matches!(&**expr, Expr::Literal(ScalarValue::Utf8(_), _))
656        ) {
657            // Stop checking the sub `expr`,
658            // which is a `Utf8` type and has already been tested above.
659            return Ok(TreeNodeRecursion::Stop);
660        }
661
662        if !matches!(
663            expr,
664            Expr::Literal(ScalarValue::IntervalDayTime(_), _)
665                | Expr::Literal(ScalarValue::IntervalMonthDayNano(_), _)
666                | Expr::Literal(ScalarValue::IntervalYearMonth(_), _)
667                | Expr::BinaryExpr(_)
668                | Expr::Cast(Cast {
669                    data_type: DataType::Interval(_),
670                    ..
671                })
672        ) {
673            all_interval = false;
674            Ok(TreeNodeRecursion::Stop)
675        } else {
676            Ok(TreeNodeRecursion::Continue)
677        }
678    });
679
680    all_interval
681}
682
683#[cfg(test)]
684mod test {
685
686    use arrow::datatypes::IntervalUnit;
687    use catalog::RegisterTableRequest;
688    use catalog::memory::MemoryCatalogManager;
689    use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
690    use common_time::IntervalYearMonth;
691    use datafusion_expr::{BinaryExpr, Literal, Operator};
692    use datatypes::prelude::ConcreteDataType;
693    use datatypes::schema::{ColumnSchema, Schema};
694    use session::context::{QueryContext, QueryContextBuilder};
695    use table::metadata::{TableInfoBuilder, TableMetaBuilder};
696    use table::table::TableRef;
697    use table::test_util::EmptyTable;
698
699    use super::*;
700    use crate::options::QueryOptions;
701    use crate::parser::QueryLanguageParser;
702    use crate::{QueryEngineFactory, QueryEngineRef};
703
704    async fn create_test_engine() -> QueryEngineRef {
705        create_test_engine_with_tables(&["test"], false).await
706    }
707
708    async fn create_union_test_engine() -> QueryEngineRef {
709        create_test_engine_with_tables(&["test_0", "test_1"], true).await
710    }
711
712    async fn create_test_engine_with_tables(
713        table_names: &[&str],
714        with_extra_timestamp: bool,
715    ) -> QueryEngineRef {
716        let catalog_list = MemoryCatalogManager::with_default_setup();
717        for (i, table_name) in table_names.iter().enumerate() {
718            let table = create_test_table(table_name, with_extra_timestamp);
719            assert!(
720                catalog_list
721                    .register_table_sync(RegisterTableRequest {
722                        catalog: DEFAULT_CATALOG_NAME.to_string(),
723                        schema: DEFAULT_SCHEMA_NAME.to_string(),
724                        table_name: (*table_name).to_string(),
725                        table_id: 1024 + i as u32,
726                        table,
727                    })
728                    .is_ok()
729            );
730        }
731        QueryEngineFactory::new(
732            catalog_list,
733            None,
734            None,
735            None,
736            None,
737            false,
738            QueryOptions::default(),
739        )
740        .query_engine()
741    }
742
743    fn create_test_table(table_name: &str, with_extra_timestamp: bool) -> TableRef {
744        let mut columns = vec![];
745        for i in 0..5 {
746            columns.push(ColumnSchema::new(
747                format!("tag_{i}"),
748                ConcreteDataType::string_datatype(),
749                false,
750            ));
751        }
752        columns.push(
753            ColumnSchema::new(
754                "timestamp".to_string(),
755                ConcreteDataType::timestamp_millisecond_datatype(),
756                false,
757            )
758            .with_time_index(true),
759        );
760        if with_extra_timestamp {
761            columns.push(ColumnSchema::new(
762                "timestamp_2".to_string(),
763                ConcreteDataType::timestamp_millisecond_datatype(),
764                true,
765            ));
766        }
767        for i in 0..5 {
768            columns.push(ColumnSchema::new(
769                format!("field_{i}"),
770                ConcreteDataType::float64_datatype(),
771                true,
772            ));
773        }
774        let schema = Arc::new(Schema::new(columns));
775        let table_meta = TableMetaBuilder::empty()
776            .schema(schema)
777            .primary_key_indices((0..5).collect())
778            .value_indices(if with_extra_timestamp {
779                (6..12).collect()
780            } else {
781                (6..11).collect()
782            })
783            .next_column_id(1024)
784            .build()
785            .unwrap();
786        let table_info = TableInfoBuilder::default()
787            .name(table_name)
788            .meta(table_meta)
789            .build()
790            .unwrap();
791        EmptyTable::from_table_info(&table_info)
792    }
793
794    async fn do_query(sql: &str) -> Result<LogicalPlan> {
795        let stmt = QueryLanguageParser::parse_sql(sql, &QueryContext::arc()).unwrap();
796        let engine = create_test_engine().await;
797        engine.planner().plan(&stmt, QueryContext::arc()).await
798    }
799
800    async fn do_query_with_ctx(sql: &str, query_ctx: QueryContextRef) -> Result<LogicalPlan> {
801        let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap();
802        let engine = create_test_engine().await;
803        engine.planner().plan(&stmt, query_ctx).await
804    }
805
806    async fn do_union_query(sql: &str) -> Result<LogicalPlan> {
807        let stmt = QueryLanguageParser::parse_sql(sql, &QueryContext::arc()).unwrap();
808        let engine = create_union_test_engine().await;
809        engine.planner().plan(&stmt, QueryContext::arc()).await
810    }
811
812    async fn query_plan_compare(sql: &str, expected: String) {
813        let plan = do_query(sql).await.unwrap();
814        assert_eq!(plan.display_indent_schema().to_string(), expected);
815    }
816
817    #[tokio::test]
818    async fn range_align_to_now_uses_scheduled_time_extension() {
819        let query_ctx = Arc::new(
820            QueryContextBuilder::default()
821                .set_extension(
822                    crate::options::FLOW_SCHEDULED_TIME_MILLIS.to_string(),
823                    "1700000000123".to_string(),
824                )
825                .build(),
826        );
827        let query = r#"SELECT timestamp, tag_0, tag_1, avg(field_0) RANGE '5m' FROM test ALIGN '1h' TO NOW by (tag_0,tag_1);"#;
828        let plan = do_query_with_ctx(query, query_ctx).await.unwrap();
829
830        assert!(
831            plan.display_indent_schema()
832                .to_string()
833                .contains("align_to=1700000000123ms")
834        );
835    }
836
837    #[tokio::test]
838    async fn range_no_project() {
839        let query = r#"SELECT timestamp, tag_0, tag_1, avg(field_0 + field_1) RANGE '5m' FROM test ALIGN '1h' by (tag_0,tag_1);"#;
840        let expected = String::from(
841            "RangeSelect: range_exprs=[avg(test.field_0 + test.field_1) RANGE 5m], align=3600000ms, align_to=0ms, align_by=[test.tag_0, test.tag_1], time_index=timestamp [timestamp:Timestamp(ms), tag_0:Utf8, tag_1:Utf8, avg(test.field_0 + test.field_1) RANGE 5m:Float64;N]\
842            \n  TableScan: test [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, tag_4:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N, field_3:Float64;N, field_4:Float64;N]",
843        );
844        query_plan_compare(query, expected).await;
845    }
846
847    #[tokio::test]
848    async fn range_expr_calculation() {
849        let query = r#"SELECT (avg(field_0 + field_1)/4) RANGE '5m' FROM test ALIGN '1h' by (tag_0,tag_1);"#;
850        let expected = String::from(
851            "Projection: avg(test.field_0 + test.field_1) RANGE 5m / Int64(4) [avg(test.field_0 + test.field_1) RANGE 5m / Int64(4):Float64;N]\
852            \n  RangeSelect: range_exprs=[avg(test.field_0 + test.field_1) RANGE 5m], align=3600000ms, align_to=0ms, align_by=[test.tag_0, test.tag_1], time_index=timestamp [avg(test.field_0 + test.field_1) RANGE 5m:Float64;N, timestamp:Timestamp(ms), tag_0:Utf8, tag_1:Utf8]\
853            \n    TableScan: test [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, tag_4:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N, field_3:Float64;N, field_4:Float64;N]",
854        );
855        query_plan_compare(query, expected).await;
856    }
857
858    #[tokio::test]
859    async fn range_multi_args() {
860        let query =
861            r#"SELECT (covar(field_0 + field_1, field_1)/4) RANGE '5m' FROM test ALIGN '1h';"#;
862        let expected = String::from(
863            "Projection: covar_samp(test.field_0 + test.field_1,test.field_1) RANGE 5m / Int64(4) [covar_samp(test.field_0 + test.field_1,test.field_1) RANGE 5m / Int64(4):Float64;N]\
864            \n  RangeSelect: range_exprs=[covar_samp(test.field_0 + test.field_1,test.field_1) RANGE 5m], align=3600000ms, align_to=0ms, align_by=[test.tag_0, test.tag_1, test.tag_2, test.tag_3, test.tag_4], time_index=timestamp [covar_samp(test.field_0 + test.field_1,test.field_1) RANGE 5m:Float64;N, timestamp:Timestamp(ms), tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, tag_4:Utf8]\
865            \n    TableScan: test [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, tag_4:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N, field_3:Float64;N, field_4:Float64;N]",
866        );
867        query_plan_compare(query, expected).await;
868    }
869
870    #[tokio::test]
871    async fn range_calculation() {
872        let query = r#"SELECT ((avg(field_0)+sum(field_1))/4) RANGE '5m' FROM test ALIGN '1h' by (tag_0,tag_1) FILL NULL;"#;
873        let expected = String::from(
874            "Projection: (avg(test.field_0) RANGE 5m FILL NULL + sum(test.field_1) RANGE 5m FILL NULL) / Int64(4) [avg(test.field_0) RANGE 5m FILL NULL + sum(test.field_1) RANGE 5m FILL NULL / Int64(4):Float64;N]\
875            \n  RangeSelect: range_exprs=[avg(test.field_0) RANGE 5m FILL NULL, sum(test.field_1) RANGE 5m FILL NULL], align=3600000ms, align_to=0ms, align_by=[test.tag_0, test.tag_1], time_index=timestamp [avg(test.field_0) RANGE 5m FILL NULL:Float64;N, sum(test.field_1) RANGE 5m FILL NULL:Float64;N, timestamp:Timestamp(ms), tag_0:Utf8, tag_1:Utf8]\
876            \n    TableScan: test [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, tag_4:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N, field_3:Float64;N, field_4:Float64;N]",
877        );
878        query_plan_compare(query, expected).await;
879    }
880
881    #[tokio::test]
882    async fn range_as_sub_query() {
883        let query = r#"SELECT foo + 1 from (SELECT ((avg(field_0)+sum(field_1))/4) RANGE '5m' as foo FROM test ALIGN '1h' by (tag_0,tag_1) FILL NULL) where foo > 1;"#;
884        let expected = String::from(
885            "Projection: foo + Int64(1) [foo + Int64(1):Float64;N]\
886            \n  Filter: foo > Int64(1) [foo:Float64;N]\
887            \n    Projection: (avg(test.field_0) RANGE 5m FILL NULL + sum(test.field_1) RANGE 5m FILL NULL) / Int64(4) AS foo [foo:Float64;N]\
888            \n      RangeSelect: range_exprs=[avg(test.field_0) RANGE 5m FILL NULL, sum(test.field_1) RANGE 5m FILL NULL], align=3600000ms, align_to=0ms, align_by=[test.tag_0, test.tag_1], time_index=timestamp [avg(test.field_0) RANGE 5m FILL NULL:Float64;N, sum(test.field_1) RANGE 5m FILL NULL:Float64;N, timestamp:Timestamp(ms), tag_0:Utf8, tag_1:Utf8]\
889            \n        TableScan: test [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, tag_4:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N, field_3:Float64;N, field_4:Float64;N]",
890        );
891        query_plan_compare(query, expected).await;
892    }
893
894    #[tokio::test]
895    async fn range_from_nest_query() {
896        let query = r#"SELECT ((avg(a)+sum(b))/4) RANGE '5m' FROM (SELECT field_0 as a, field_1 as b, tag_0 as c, tag_1 as d, timestamp from test where field_0 > 1.0) ALIGN '1h' by (c, d) FILL NULL;"#;
897        let expected = String::from(
898            "Projection: (avg(a) RANGE 5m FILL NULL + sum(b) RANGE 5m FILL NULL) / Int64(4) [avg(a) RANGE 5m FILL NULL + sum(b) RANGE 5m FILL NULL / Int64(4):Float64;N]\
899            \n  RangeSelect: range_exprs=[avg(a) RANGE 5m FILL NULL, sum(b) RANGE 5m FILL NULL], align=3600000ms, align_to=0ms, align_by=[c, d], time_index=timestamp [avg(a) RANGE 5m FILL NULL:Float64;N, sum(b) RANGE 5m FILL NULL:Float64;N, timestamp:Timestamp(ms), c:Utf8, d:Utf8]\
900            \n    Projection: test.field_0 AS a, test.field_1 AS b, test.tag_0 AS c, test.tag_1 AS d, test.timestamp [a:Float64;N, b:Float64;N, c:Utf8, d:Utf8, timestamp:Timestamp(ms)]\
901            \n      Filter: test.field_0 > Float64(1) [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, tag_4:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N, field_3:Float64;N, field_4:Float64;N]\
902            \n        TableScan: test [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, tag_4:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N, field_3:Float64;N, field_4:Float64;N]",
903        );
904        query_plan_compare(query, expected).await;
905    }
906
907    #[tokio::test]
908    async fn range_from_union_query() {
909        let queries = [
910            r#"SELECT timestamp, tag_0, avg(field_0) RANGE '5m'
911            FROM (
912                SELECT timestamp, tag_0, field_0, timestamp_2 FROM test_0
913                UNION ALL
914                SELECT timestamp, tag_0, field_0, timestamp_2 FROM test_1
915            )
916            WHERE timestamp >= '1970-01-01 00:00:00'
917            ALIGN '1h' by (tag_0)"#,
918            r#"SELECT tmp.timestamp, tmp.tag_0, avg(tmp.field_0) RANGE '5m'
919            FROM (
920                SELECT timestamp, tag_0, field_0, timestamp_2 FROM test_0
921                UNION ALL
922                SELECT timestamp, tag_0, field_0, timestamp_2 FROM test_1
923            ) AS tmp
924            WHERE tmp.timestamp >= '1970-01-01 00:00:00'
925            ALIGN '1h' by (tmp.tag_0)"#,
926        ];
927
928        for query in queries {
929            let plan = do_union_query(query)
930                .await
931                .unwrap()
932                .display_indent_schema()
933                .to_string();
934
935            assert!(plan.contains("RangeSelect"));
936            assert!(plan.contains("Union"));
937            assert!(plan.contains("time_index=timestamp"));
938        }
939    }
940
941    #[tokio::test]
942    async fn range_from_derived_query_without_by_err() {
943        let queries = [
944            r#"SELECT timestamp, tag_0, avg(field_0) RANGE '5m'
945            FROM (
946                SELECT timestamp, tag_0, field_0, timestamp_2 FROM test_0
947                UNION ALL
948                SELECT timestamp, tag_0, field_0, timestamp_2 FROM test_1
949            )
950            WHERE timestamp >= '1970-01-01 00:00:00'
951            ALIGN '1h'"#,
952            r#"SELECT tmp.timestamp, tmp.tag_0, avg(tmp.field_0) RANGE '5m'
953            FROM (
954                SELECT timestamp, tag_0, field_0, timestamp_2 FROM test_0
955                UNION ALL
956                SELECT timestamp, tag_0, field_0, timestamp_2 FROM test_1
957            ) AS tmp
958            WHERE tmp.timestamp >= '1970-01-01 00:00:00'
959            ALIGN '1h'"#,
960        ];
961
962        for query in queries {
963            assert_eq!(
964                do_union_query(query).await.unwrap_err().to_string(),
965                "Range Query: Cannot infer default BY columns from derived range query input"
966            );
967        }
968    }
969
970    #[tokio::test]
971    async fn range_in_expr() {
972        let query = r#"SELECT sin(avg(field_0 + field_1) RANGE '5m' + 1) FROM test ALIGN '1h' by (tag_0,tag_1);"#;
973        let expected = String::from(
974            "Projection: sin(avg(test.field_0 + test.field_1) RANGE 5m + Int64(1)) [sin(avg(test.field_0 + test.field_1) RANGE 5m + Int64(1)):Float64;N]\
975            \n  RangeSelect: range_exprs=[avg(test.field_0 + test.field_1) RANGE 5m], align=3600000ms, align_to=0ms, align_by=[test.tag_0, test.tag_1], time_index=timestamp [avg(test.field_0 + test.field_1) RANGE 5m:Float64;N, timestamp:Timestamp(ms), tag_0:Utf8, tag_1:Utf8]\
976            \n    TableScan: test [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, tag_4:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N, field_3:Float64;N, field_4:Float64;N]",
977        );
978        query_plan_compare(query, expected).await;
979    }
980
981    #[tokio::test]
982    async fn duplicate_range_expr() {
983        let query = r#"SELECT avg(field_0) RANGE '5m' FILL 6.0 + avg(field_0) RANGE '5m' FILL 6.0 FROM test ALIGN '1h' by (tag_0,tag_1);"#;
984        let expected = String::from(
985            "Projection: avg(test.field_0) RANGE 5m FILL 6 + avg(test.field_0) RANGE 5m FILL 6 [avg(test.field_0) RANGE 5m FILL 6 + avg(test.field_0) RANGE 5m FILL 6:Float64]\
986            \n  RangeSelect: range_exprs=[avg(test.field_0) RANGE 5m FILL 6], align=3600000ms, align_to=0ms, align_by=[test.tag_0, test.tag_1], time_index=timestamp [avg(test.field_0) RANGE 5m FILL 6:Float64, timestamp:Timestamp(ms), tag_0:Utf8, tag_1:Utf8]\
987            \n    TableScan: test [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, tag_4:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N, field_3:Float64;N, field_4:Float64;N]",
988        );
989        query_plan_compare(query, expected).await;
990    }
991
992    #[tokio::test]
993    async fn deep_nest_range_expr() {
994        let query = r#"SELECT round(sin(avg(field_0 + field_1) RANGE '5m' + 1)) FROM test ALIGN '1h' by (tag_0,tag_1);"#;
995        let expected = String::from(
996            "Projection: round(sin(avg(test.field_0 + test.field_1) RANGE 5m + Int64(1))) [round(sin(avg(test.field_0 + test.field_1) RANGE 5m + Int64(1))):Float64;N]\
997            \n  RangeSelect: range_exprs=[avg(test.field_0 + test.field_1) RANGE 5m], align=3600000ms, align_to=0ms, align_by=[test.tag_0, test.tag_1], time_index=timestamp [avg(test.field_0 + test.field_1) RANGE 5m:Float64;N, timestamp:Timestamp(ms), tag_0:Utf8, tag_1:Utf8]\
998            \n    TableScan: test [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, tag_4:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N, field_3:Float64;N, field_4:Float64;N]",
999        );
1000        query_plan_compare(query, expected).await;
1001    }
1002
1003    #[tokio::test]
1004    async fn complex_range_expr() {
1005        let query = r#"SELECT gcd(CAST(max(field_0 + 1) Range '5m' FILL NULL AS Int64), CAST(tag_0 AS Int64)) + round(max(field_2+1) Range '6m' FILL NULL + 1) + max(field_2+3) Range '10m' FILL NULL * CAST(tag_1 AS Float64) + 1 FROM test ALIGN '1h' by (tag_0, tag_1);"#;
1006        let expected = String::from(
1007            "Projection: gcd(arrow_cast(max(test.field_0 + Int64(1)) RANGE 5m FILL NULL, Utf8(\"Int64\")), arrow_cast(test.tag_0, Utf8(\"Int64\"))) + round(max(test.field_2 + Int64(1)) RANGE 6m FILL NULL + Int64(1)) + max(test.field_2 + Int64(3)) RANGE 10m FILL NULL * arrow_cast(test.tag_1, Utf8(\"Float64\")) + Int64(1) [gcd(arrow_cast(max(test.field_0 + Int64(1)) RANGE 5m FILL NULL,Utf8(\"Int64\")),arrow_cast(test.tag_0,Utf8(\"Int64\"))) + round(max(test.field_2 + Int64(1)) RANGE 6m FILL NULL + Int64(1)) + max(test.field_2 + Int64(3)) RANGE 10m FILL NULL * arrow_cast(test.tag_1,Utf8(\"Float64\")) + Int64(1):Float64;N]\
1008            \n  RangeSelect: range_exprs=[max(test.field_0 + Int64(1)) RANGE 5m FILL NULL, max(test.field_2 + Int64(1)) RANGE 6m FILL NULL, max(test.field_2 + Int64(3)) RANGE 10m FILL NULL], align=3600000ms, align_to=0ms, align_by=[test.tag_0, test.tag_1], time_index=timestamp [max(test.field_0 + Int64(1)) RANGE 5m FILL NULL:Float64;N, max(test.field_2 + Int64(1)) RANGE 6m FILL NULL:Float64;N, max(test.field_2 + Int64(3)) RANGE 10m FILL NULL:Float64;N, timestamp:Timestamp(ms), tag_0:Utf8, tag_1:Utf8]\
1009            \n    TableScan: test [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, tag_4:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N, field_3:Float64;N, field_4:Float64;N]",
1010        );
1011        query_plan_compare(query, expected).await;
1012    }
1013
1014    #[tokio::test]
1015    async fn range_linear_on_integer() {
1016        let query = r#"SELECT min(CAST(field_0 AS Int64) + CAST(field_1 AS Int64)) RANGE '5m' FILL LINEAR FROM test ALIGN '1h' by (tag_0,tag_1);"#;
1017        let expected = String::from(
1018            "RangeSelect: range_exprs=[min(arrow_cast(test.field_0,Utf8(\"Int64\")) + arrow_cast(test.field_1,Utf8(\"Int64\"))) RANGE 5m FILL LINEAR], align=3600000ms, align_to=0ms, align_by=[test.tag_0, test.tag_1], time_index=timestamp [min(arrow_cast(test.field_0,Utf8(\"Int64\")) + arrow_cast(test.field_1,Utf8(\"Int64\"))) RANGE 5m FILL LINEAR:Float64;N]\
1019            \n  TableScan: test [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, tag_4:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N, field_3:Float64;N, field_4:Float64;N]",
1020        );
1021        query_plan_compare(query, expected).await;
1022    }
1023
1024    #[tokio::test]
1025    async fn range_nest_range_err() {
1026        let query = r#"SELECT sum(avg(field_0 + field_1) RANGE '5m' + 1) RANGE '5m' + 1 FROM test ALIGN '1h' by (tag_0,tag_1);"#;
1027        assert_eq!(
1028            do_query(query).await.unwrap_err().to_string(),
1029            "Range Query: Nest Range Query is not allowed"
1030        )
1031    }
1032
1033    #[tokio::test]
1034    /// Start directly from the rewritten SQL and check whether the error reported by the range expression rewriting is as expected.
1035    /// the right argument is `range_fn(avg(field_0), '5m', 'NULL', '0', '1h')`
1036    async fn range_argument_err_1() {
1037        let query = r#"SELECT range_fn('5m', avg(field_0), 'NULL', '1', tag_0, '1h') FROM test group by tag_0;"#;
1038        let error = do_query(query).await.unwrap_err().to_string();
1039        assert_eq!(
1040            error,
1041            "Error during planning: Illegal argument `Utf8(\"5m\")` in range select query"
1042        )
1043    }
1044
1045    #[tokio::test]
1046    async fn range_argument_err_2() {
1047        let query = r#"SELECT range_fn(avg(field_0), 5, 'NULL', '1', tag_0, '1h') FROM test group by tag_0;"#;
1048        let error = do_query(query).await.unwrap_err().to_string();
1049        assert_eq!(
1050            error,
1051            "Error during planning: Illegal argument `Int64(5)` in range select query"
1052        )
1053    }
1054
1055    #[test]
1056    fn test_parse_duration_expr() {
1057        // test IntervalYearMonth
1058        let interval = IntervalYearMonth::new(10);
1059        let args = vec![ScalarValue::IntervalYearMonth(Some(interval.to_i32())).lit()];
1060        assert!(parse_duration_expr(&args, 0).is_err(),);
1061        // test IntervalDayTime
1062        let interval = IntervalDayTime::new(10, 10);
1063        let args = vec![ScalarValue::IntervalDayTime(Some(interval.into())).lit()];
1064        assert_eq!(
1065            parse_duration_expr(&args, 0).unwrap().as_millis() as i64,
1066            interval.as_millis()
1067        );
1068        // test IntervalMonthDayNano
1069        let interval = IntervalMonthDayNano::new(0, 10, 10);
1070        let args = vec![ScalarValue::IntervalMonthDayNano(Some(interval.into())).lit()];
1071        assert_eq!(
1072            parse_duration_expr(&args, 0).unwrap().as_millis() as i64,
1073            interval.days as i64 * MS_PER_DAY + interval.nanoseconds / NANOS_PER_MILLI,
1074        );
1075        // test Duration
1076        let args = vec!["1y4w".lit()];
1077        assert_eq!(
1078            parse_duration_expr(&args, 0).unwrap(),
1079            parse_duration("1y4w").unwrap()
1080        );
1081        // test cast expression
1082        let args = vec![Expr::Cast(Cast {
1083            expr: Box::new("15 minutes".lit()),
1084            data_type: DataType::Interval(IntervalUnit::MonthDayNano),
1085        })];
1086        assert_eq!(
1087            parse_duration_expr(&args, 0).unwrap(),
1088            parse_duration("15m").unwrap()
1089        );
1090        // test index err
1091        assert!(parse_duration_expr(&args, 10).is_err());
1092        // test evaluate expr
1093        let args = vec![Expr::BinaryExpr(BinaryExpr {
1094            left: Box::new(
1095                ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(0, 10).into())).lit(),
1096            ),
1097            op: Operator::Plus,
1098            right: Box::new(
1099                ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(0, 10).into())).lit(),
1100            ),
1101        })];
1102        assert_eq!(
1103            parse_duration_expr(&args, 0).unwrap(),
1104            Duration::from_millis(20)
1105        );
1106        let args = vec![Expr::BinaryExpr(BinaryExpr {
1107            left: Box::new(
1108                ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(0, 10).into())).lit(),
1109            ),
1110            op: Operator::Minus,
1111            right: Box::new(
1112                ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(0, 10).into())).lit(),
1113            ),
1114        })];
1115        // test zero interval error
1116        assert!(parse_duration_expr(&args, 0).is_err());
1117        // test must all be interval
1118        let args = vec![Expr::BinaryExpr(BinaryExpr {
1119            left: Box::new(
1120                ScalarValue::IntervalYearMonth(Some(IntervalYearMonth::new(10).to_i32())).lit(),
1121            ),
1122            op: Operator::Minus,
1123            right: Box::new(ScalarValue::Time64Microsecond(Some(0)).lit()),
1124        })];
1125        assert!(parse_duration_expr(&args, 0).is_err());
1126    }
1127
1128    #[test]
1129    fn test_parse_align_to() {
1130        // test NOW
1131        let args = vec!["NOW".lit()];
1132        let epsinon =
1133            parse_align_to(&args, 0, None, None).unwrap() - Timestamp::current_millis().value();
1134        assert!(epsinon.abs() < 100);
1135        let scheduled_time = DateTime::from_timestamp(1_700_000_000, 123_000_000).unwrap();
1136        assert_eq!(
1137            scheduled_time.timestamp_millis(),
1138            parse_align_to(&args, 0, None, Some(scheduled_time)).unwrap()
1139        );
1140        // test default
1141        let args = vec!["".lit()];
1142        assert_eq!(0, parse_align_to(&args, 0, None, None).unwrap());
1143        // test default with timezone
1144        let args = vec!["".lit()];
1145        assert_eq!(
1146            -36000 * 1000,
1147            parse_align_to(
1148                &args,
1149                0,
1150                Some(&Timezone::from_tz_string("HST").unwrap()),
1151                None
1152            )
1153            .unwrap()
1154        );
1155        assert_eq!(
1156            28800 * 1000,
1157            parse_align_to(
1158                &args,
1159                0,
1160                Some(&Timezone::from_tz_string("Asia/Shanghai").unwrap()),
1161                None
1162            )
1163            .unwrap()
1164        );
1165
1166        // test Timestamp
1167        let args = vec!["1970-01-01T00:00:00+08:00".lit()];
1168        assert_eq!(
1169            parse_align_to(&args, 0, None, None).unwrap(),
1170            -8 * 60 * 60 * 1000
1171        );
1172        // timezone
1173        let args = vec!["1970-01-01T00:00:00".lit()];
1174        assert_eq!(
1175            parse_align_to(
1176                &args,
1177                0,
1178                Some(&Timezone::from_tz_string("Asia/Shanghai").unwrap()),
1179                None
1180            )
1181            .unwrap(),
1182            -8 * 60 * 60 * 1000
1183        );
1184        // test evaluate expr
1185        let args = vec![Expr::BinaryExpr(BinaryExpr {
1186            left: Box::new(
1187                ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(0, 10).into())).lit(),
1188            ),
1189            op: Operator::Plus,
1190            right: Box::new(
1191                ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(0, 10).into())).lit(),
1192            ),
1193        })];
1194        assert_eq!(parse_align_to(&args, 0, None, None).unwrap(), 20);
1195    }
1196
1197    #[test]
1198    fn test_interval_only() {
1199        let expr = Expr::BinaryExpr(BinaryExpr {
1200            left: Box::new(ScalarValue::DurationMillisecond(Some(20)).lit()),
1201            op: Operator::Minus,
1202            right: Box::new(
1203                ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(10, 0).into())).lit(),
1204            ),
1205        });
1206        assert!(!interval_only_in_expr(&expr));
1207        let expr = Expr::BinaryExpr(BinaryExpr {
1208            left: Box::new(
1209                ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(10, 0).into())).lit(),
1210            ),
1211            op: Operator::Minus,
1212            right: Box::new(
1213                ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(10, 0).into())).lit(),
1214            ),
1215        });
1216        assert!(interval_only_in_expr(&expr));
1217
1218        let expr = Expr::BinaryExpr(BinaryExpr {
1219            left: Box::new(Expr::Cast(Cast {
1220                expr: Box::new("15 minute".lit()),
1221                data_type: DataType::Interval(IntervalUnit::MonthDayNano),
1222            })),
1223            op: Operator::Minus,
1224            right: Box::new(
1225                ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(10, 0).into())).lit(),
1226            ),
1227        });
1228        assert!(interval_only_in_expr(&expr));
1229
1230        let expr = Expr::Cast(Cast {
1231            expr: Box::new(Expr::BinaryExpr(BinaryExpr {
1232                left: Box::new(Expr::Cast(Cast {
1233                    expr: Box::new("15 minute".lit()),
1234                    data_type: DataType::Interval(IntervalUnit::MonthDayNano),
1235                })),
1236                op: Operator::Minus,
1237                right: Box::new(
1238                    ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(10, 0).into())).lit(),
1239                ),
1240            })),
1241            data_type: DataType::Interval(IntervalUnit::MonthDayNano),
1242        });
1243
1244        assert!(interval_only_in_expr(&expr));
1245    }
1246}
query/range_select/plan_rewrite.rs

query/range_select/
plan_rewrite.rs
