Skip to main content

query/range_select/
plan.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::cmp::Ordering;
16use std::collections::btree_map::Entry;
17use std::collections::{BTreeMap, HashMap};
18use std::fmt::Display;
19use std::pin::Pin;
20use std::sync::Arc;
21use std::task::{Context, Poll};
22use std::time::Duration;
23
24use ahash::RandomState;
25use arrow::compute::{self, CastOptions, cast_with_options, take_arrays};
26use arrow_schema::{DataType, Field, Schema, SchemaRef, SortOptions, TimeUnit};
27use common_function::aggrs::aggr_wrapper::get_aggr_func;
28use common_recordbatch::DfSendableRecordBatchStream;
29use datafusion::common::Result as DataFusionResult;
30use datafusion::error::Result as DfResult;
31use datafusion::execution::TaskContext;
32use datafusion::execution::context::SessionState;
33use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
34use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
35use datafusion::physical_plan::{
36    DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, RecordBatchStream,
37    SendableRecordBatchStream,
38};
39use datafusion_common::hash_utils::create_hashes;
40use datafusion_common::{DFSchema, DFSchemaRef, DataFusionError, ScalarValue};
41use datafusion_expr::utils::{COUNT_STAR_EXPANSION, exprlist_to_fields};
42use datafusion_expr::{
43    Accumulator, Expr, ExprSchemable, LogicalPlan, UserDefinedLogicalNodeCore, lit,
44};
45use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr};
46use datafusion_physical_expr::{
47    Distribution, EquivalenceProperties, Partitioning, PhysicalExpr, PhysicalSortExpr,
48    create_physical_expr, create_physical_sort_expr,
49};
50use datatypes::arrow::array::{
51    Array, ArrayRef, TimestampMillisecondArray, TimestampMillisecondBuilder, UInt32Builder,
52};
53use datatypes::arrow::datatypes::{ArrowPrimitiveType, TimestampMillisecondType};
54use datatypes::arrow::record_batch::RecordBatch;
55use datatypes::arrow::row::{OwnedRow, RowConverter, SortField};
56use futures::{Stream, ready};
57use futures_util::StreamExt;
58use snafu::ensure;
59
60use crate::error::{RangeQuerySnafu, Result};
61
62type Millisecond = <TimestampMillisecondType as ArrowPrimitiveType>::Native;
63
64#[derive(PartialEq, Eq, Debug, Hash, Clone)]
65pub enum Fill {
66    Null,
67    Prev,
68    Linear,
69    Const(ScalarValue),
70}
71
72impl Display for Fill {
73    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74        match self {
75            Fill::Null => write!(f, "NULL"),
76            Fill::Prev => write!(f, "PREV"),
77            Fill::Linear => write!(f, "LINEAR"),
78            Fill::Const(x) => write!(f, "{}", x),
79        }
80    }
81}
82
83impl Fill {
84    pub fn try_from_str(value: &str, datatype: &DataType) -> DfResult<Option<Self>> {
85        let s = value.to_uppercase();
86        match s.as_str() {
87            "" => Ok(None),
88            "NULL" => Ok(Some(Self::Null)),
89            "PREV" => Ok(Some(Self::Prev)),
90            "LINEAR" => {
91                if datatype.is_numeric() {
92                    Ok(Some(Self::Linear))
93                } else {
94                    Err(DataFusionError::Plan(format!(
95                        "Use FILL LINEAR on Non-numeric DataType {}",
96                        datatype
97                    )))
98                }
99            }
100            _ => ScalarValue::try_from_string(s.clone(), datatype)
101                .map_err(|err| {
102                    DataFusionError::Plan(format!(
103                        "{} is not a valid fill option, fail to convert to a const value. {{ {} }}",
104                        s, err
105                    ))
106                })
107                .map(|x| Some(Fill::Const(x))),
108        }
109    }
110
111    /// The input `data` contains data on a complete time series.
112    /// If the filling strategy is `PREV` or `LINEAR`, caller must be ensured that the incoming `ts`&`data` is ascending time order.
113    pub fn apply_fill_strategy(&self, ts: &[i64], data: &mut [ScalarValue]) -> DfResult<()> {
114        // No calculation need in `Fill::Null`
115        if matches!(self, Fill::Null) {
116            return Ok(());
117        }
118        let len = data.len();
119        if *self == Fill::Linear {
120            return Self::fill_linear(ts, data);
121        }
122        for i in 0..len {
123            if data[i].is_null() {
124                match self {
125                    Fill::Prev => {
126                        if i != 0 {
127                            data[i] = data[i - 1].clone()
128                        }
129                    }
130                    // The calculation of linear interpolation is relatively complicated.
131                    // `Self::fill_linear` is used to dispose `Fill::Linear`.
132                    // No calculation need in `Fill::Null`
133                    Fill::Linear | Fill::Null => unreachable!(),
134                    Fill::Const(v) => data[i] = v.clone(),
135                }
136            }
137        }
138        Ok(())
139    }
140
141    fn fill_linear(ts: &[i64], data: &mut [ScalarValue]) -> DfResult<()> {
142        let not_null_num = data
143            .iter()
144            .fold(0, |acc, x| if x.is_null() { acc } else { acc + 1 });
145        // We need at least two non-empty data points to perform linear interpolation
146        if not_null_num < 2 {
147            return Ok(());
148        }
149        let mut index = 0;
150        let mut head: Option<usize> = None;
151        let mut tail: Option<usize> = None;
152        while index < data.len() {
153            // find null interval [start, end)
154            // start is null, end is not-null
155            let start = data[index..]
156                .iter()
157                .position(ScalarValue::is_null)
158                .unwrap_or(data.len() - index)
159                + index;
160            if start == data.len() {
161                break;
162            }
163            let end = data[start..]
164                .iter()
165                .position(|r| !r.is_null())
166                .unwrap_or(data.len() - start)
167                + start;
168            index = end + 1;
169            // head or tail null dispose later, record start/end first
170            if start == 0 {
171                head = Some(end);
172            } else if end == data.len() {
173                tail = Some(start);
174            } else {
175                linear_interpolation(ts, data, start - 1, end, start, end)?;
176            }
177        }
178        // dispose head null interval
179        if let Some(end) = head {
180            linear_interpolation(ts, data, end, end + 1, 0, end)?;
181        }
182        // dispose tail null interval
183        if let Some(start) = tail {
184            linear_interpolation(ts, data, start - 2, start - 1, start, data.len())?;
185        }
186        Ok(())
187    }
188}
189
190/// use `(ts[i1], data[i1])`, `(ts[i2], data[i2])` as endpoint, linearly interpolates element over the interval `[start, end)`
191fn linear_interpolation(
192    ts: &[i64],
193    data: &mut [ScalarValue],
194    i1: usize,
195    i2: usize,
196    start: usize,
197    end: usize,
198) -> DfResult<()> {
199    let (x0, x1) = (ts[i1] as f64, ts[i2] as f64);
200    let (y0, y1, is_float32) = match (&data[i1], &data[i2]) {
201        (ScalarValue::Float64(Some(y0)), ScalarValue::Float64(Some(y1))) => (*y0, *y1, false),
202        (ScalarValue::Float32(Some(y0)), ScalarValue::Float32(Some(y1))) => {
203            (*y0 as f64, *y1 as f64, true)
204        }
205        _ => {
206            return Err(DataFusionError::Execution(
207                "RangePlan: Apply Fill LINEAR strategy on Non-floating type".to_string(),
208            ));
209        }
210    };
211    // To avoid divide zero error, kind of defensive programming
212    if x1 == x0 {
213        return Err(DataFusionError::Execution(
214            "RangePlan: Linear interpolation using the same coordinate points".to_string(),
215        ));
216    }
217    for i in start..end {
218        let val = y0 + (y1 - y0) / (x1 - x0) * (ts[i] as f64 - x0);
219        data[i] = if is_float32 {
220            ScalarValue::Float32(Some(val as f32))
221        } else {
222            ScalarValue::Float64(Some(val))
223        }
224    }
225    Ok(())
226}
227
228#[derive(Eq, Clone, Debug)]
229pub struct RangeFn {
230    /// with format like `max(a) RANGE 300s [FILL NULL]`
231    pub name: String,
232    pub data_type: DataType,
233    pub expr: Expr,
234    pub range: Duration,
235    pub fill: Option<Fill>,
236    /// If the `FIll` strategy is `Linear` and the output is an integer,
237    /// it is possible to calculate a floating point number.
238    /// So for `FILL==LINEAR`, the entire data will be implicitly converted to Float type
239    /// If `need_cast==true`, `data_type` may not consist with type `expr` generated.
240    pub need_cast: bool,
241}
242
243impl PartialEq for RangeFn {
244    fn eq(&self, other: &Self) -> bool {
245        self.name == other.name
246    }
247}
248
249impl PartialOrd for RangeFn {
250    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
251        Some(self.cmp(other))
252    }
253}
254
255impl Ord for RangeFn {
256    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
257        self.name.cmp(&other.name)
258    }
259}
260
261impl std::hash::Hash for RangeFn {
262    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
263        self.name.hash(state);
264    }
265}
266
267impl Display for RangeFn {
268    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
269        write!(f, "{}", self.name)
270    }
271}
272
273#[derive(Debug, PartialEq, Eq, Hash)]
274pub struct RangeSelect {
275    /// The incoming logical plan
276    pub input: Arc<LogicalPlan>,
277    /// all range expressions
278    pub range_expr: Vec<RangeFn>,
279    pub align: Duration,
280    pub align_to: i64,
281    pub time_index: String,
282    pub time_expr: Expr,
283    pub by: Vec<Expr>,
284    pub schema: DFSchemaRef,
285    pub by_schema: DFSchemaRef,
286    /// If the `schema` of the `RangeSelect` happens to be the same as the content of the upper-level Projection Plan,
287    /// the final output needs to be `project` through `schema_project`,
288    /// so that we can omit the upper-level Projection Plan.
289    pub schema_project: Option<Vec<usize>>,
290    /// The schema before run projection, follow the order of `range expr | time index | by columns`
291    /// `schema_before_project  ----  schema_project ----> schema`
292    /// if `schema_project==None` then `schema_before_project==schema`
293    pub schema_before_project: DFSchemaRef,
294}
295
296impl PartialOrd for RangeSelect {
297    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
298        // Compare fields in order excluding `schema`, `by_schema`, `schema_before_project`.
299        match self.input.partial_cmp(&other.input) {
300            Some(Ordering::Equal) => {}
301            ord => return ord,
302        }
303        match self.range_expr.partial_cmp(&other.range_expr) {
304            Some(Ordering::Equal) => {}
305            ord => return ord,
306        }
307        match self.align.partial_cmp(&other.align) {
308            Some(Ordering::Equal) => {}
309            ord => return ord,
310        }
311        match self.align_to.partial_cmp(&other.align_to) {
312            Some(Ordering::Equal) => {}
313            ord => return ord,
314        }
315        match self.time_index.partial_cmp(&other.time_index) {
316            Some(Ordering::Equal) => {}
317            ord => return ord,
318        }
319        match self.time_expr.partial_cmp(&other.time_expr) {
320            Some(Ordering::Equal) => {}
321            ord => return ord,
322        }
323        match self.by.partial_cmp(&other.by) {
324            Some(Ordering::Equal) => {}
325            ord => return ord,
326        }
327        self.schema_project.partial_cmp(&other.schema_project)
328    }
329}
330
331impl RangeSelect {
332    pub fn try_new(
333        input: Arc<LogicalPlan>,
334        range_expr: Vec<RangeFn>,
335        align: Duration,
336        align_to: i64,
337        time_index: Expr,
338        by: Vec<Expr>,
339        projection_expr: &[Expr],
340    ) -> Result<Self> {
341        ensure!(
342            align.as_millis() != 0,
343            RangeQuerySnafu {
344                msg: "Can't use 0 as align in Range Query"
345            }
346        );
347        for expr in &range_expr {
348            ensure!(
349                expr.range.as_millis() != 0,
350                RangeQuerySnafu {
351                    msg: format!(
352                        "Invalid Range expr `{}`, Can't use 0 as range in Range Query",
353                        expr.name
354                    )
355                }
356            );
357        }
358        let mut fields = range_expr
359            .iter()
360            .map(
361                |RangeFn {
362                     name,
363                     data_type,
364                     fill,
365                     ..
366                 }| {
367                    let field = Field::new(
368                        name,
369                        data_type.clone(),
370                        // Only when data fill with Const option, the data can't be null
371                        !matches!(fill, Some(Fill::Const(..))),
372                    );
373                    Ok((None, Arc::new(field)))
374                },
375            )
376            .collect::<DfResult<Vec<_>>>()?;
377        // add align_ts
378        let ts_field = time_index.to_field(input.schema().as_ref())?;
379        let time_index_name = ts_field.1.name().clone();
380        fields.push(ts_field);
381        // add by
382        let by_fields = exprlist_to_fields(&by, &input)?;
383        fields.extend(by_fields.clone());
384        let schema_before_project = Arc::new(DFSchema::new_with_metadata(
385            fields,
386            input.schema().metadata().clone(),
387        )?);
388        let by_schema = Arc::new(DFSchema::new_with_metadata(
389            by_fields,
390            input.schema().metadata().clone(),
391        )?);
392        // If the results of project plan can be obtained directly from range plan without any additional
393        // calculations, no project plan is required. We can simply project the final output of the range
394        // plan to produce the final result.
395        let schema_project = projection_expr
396            .iter()
397            .map(|project_expr| {
398                if let Expr::Column(column) = project_expr {
399                    schema_before_project
400                        .index_of_column_by_name(column.relation.as_ref(), &column.name)
401                        .ok_or(())
402                } else {
403                    let (qualifier, field) = project_expr
404                        .to_field(input.schema().as_ref())
405                        .map_err(|_| ())?;
406                    schema_before_project
407                        .index_of_column_by_name(qualifier.as_ref(), field.name())
408                        .ok_or(())
409                }
410            })
411            .collect::<std::result::Result<Vec<usize>, ()>>()
412            .ok();
413        let schema = if let Some(project) = &schema_project {
414            let project_field = project
415                .iter()
416                .map(|i| {
417                    let f = schema_before_project.qualified_field(*i);
418                    (f.0.cloned(), f.1.clone())
419                })
420                .collect();
421            Arc::new(DFSchema::new_with_metadata(
422                project_field,
423                input.schema().metadata().clone(),
424            )?)
425        } else {
426            schema_before_project.clone()
427        };
428        Ok(Self {
429            input,
430            range_expr,
431            align,
432            align_to,
433            time_index: time_index_name,
434            time_expr: time_index,
435            schema,
436            by_schema,
437            by,
438            schema_project,
439            schema_before_project,
440        })
441    }
442}
443
444impl UserDefinedLogicalNodeCore for RangeSelect {
445    fn name(&self) -> &str {
446        "RangeSelect"
447    }
448
449    fn inputs(&self) -> Vec<&LogicalPlan> {
450        vec![&self.input]
451    }
452
453    fn schema(&self) -> &DFSchemaRef {
454        &self.schema
455    }
456
457    fn expressions(&self) -> Vec<Expr> {
458        self.range_expr
459            .iter()
460            .map(|expr| expr.expr.clone())
461            .chain([self.time_expr.clone()])
462            .chain(self.by.clone())
463            .collect()
464    }
465
466    fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
467        write!(
468            f,
469            "RangeSelect: range_exprs=[{}], align={}ms, align_to={}ms, align_by=[{}], time_index={}",
470            self.range_expr
471                .iter()
472                .map(ToString::to_string)
473                .collect::<Vec<_>>()
474                .join(", "),
475            self.align.as_millis(),
476            self.align_to,
477            self.by
478                .iter()
479                .map(ToString::to_string)
480                .collect::<Vec<_>>()
481                .join(", "),
482            self.time_index
483        )
484    }
485
486    fn with_exprs_and_inputs(
487        &self,
488        exprs: Vec<Expr>,
489        inputs: Vec<LogicalPlan>,
490    ) -> DataFusionResult<Self> {
491        if inputs.is_empty() {
492            return Err(DataFusionError::Plan(
493                "RangeSelect: inputs is empty".to_string(),
494            ));
495        }
496        if exprs.len() != self.range_expr.len() + self.by.len() + 1 {
497            return Err(DataFusionError::Plan(
498                "RangeSelect: exprs length not match".to_string(),
499            ));
500        }
501
502        let range_expr = exprs
503            .iter()
504            .zip(self.range_expr.iter())
505            .map(|(e, range)| RangeFn {
506                name: range.name.clone(),
507                data_type: range.data_type.clone(),
508                expr: e.clone(),
509                range: range.range,
510                fill: range.fill.clone(),
511                need_cast: range.need_cast,
512            })
513            .collect();
514        let time_expr = exprs[self.range_expr.len()].clone();
515        let by = exprs[self.range_expr.len() + 1..].to_vec();
516        Ok(Self {
517            align: self.align,
518            align_to: self.align_to,
519            range_expr,
520            input: Arc::new(inputs[0].clone()),
521            time_index: self.time_index.clone(),
522            time_expr,
523            schema: self.schema.clone(),
524            by,
525            by_schema: self.by_schema.clone(),
526            schema_project: self.schema_project.clone(),
527            schema_before_project: self.schema_before_project.clone(),
528        })
529    }
530}
531
532impl RangeSelect {
533    fn create_physical_expr_list(
534        &self,
535        is_count_aggr: bool,
536        exprs: &[Expr],
537        df_schema: &Arc<DFSchema>,
538        session_state: &SessionState,
539    ) -> DfResult<Vec<Arc<dyn PhysicalExpr>>> {
540        exprs
541            .iter()
542            .map(|e| match e {
543                // `count(*)` will be rewritten by `CountWildcardRule` into `count(1)` when optimizing logical plan.
544                // The modification occurs after range plan rewrite.
545                // At this time, aggregate plan has been replaced by a custom range plan,
546                // so `CountWildcardRule` has not been applied.
547                // We manually modify it when creating the physical plan.
548                #[expect(deprecated)]
549                Expr::Wildcard { .. } if is_count_aggr => create_physical_expr(
550                    &lit(COUNT_STAR_EXPANSION),
551                    df_schema.as_ref(),
552                    session_state.execution_props(),
553                ),
554                _ => create_physical_expr(e, df_schema.as_ref(), session_state.execution_props()),
555            })
556            .collect::<DfResult<Vec<_>>>()
557    }
558
559    pub fn to_execution_plan(
560        &self,
561        logical_input: &LogicalPlan,
562        exec_input: Arc<dyn ExecutionPlan>,
563        session_state: &SessionState,
564    ) -> DfResult<Arc<dyn ExecutionPlan>> {
565        let fields: Vec<_> = self
566            .schema_before_project
567            .fields()
568            .iter()
569            .map(|field| Field::new(field.name(), field.data_type().clone(), field.is_nullable()))
570            .collect();
571        let by_fields: Vec<_> = self
572            .by_schema
573            .fields()
574            .iter()
575            .map(|field| Field::new(field.name(), field.data_type().clone(), field.is_nullable()))
576            .collect();
577        let input_dfschema = logical_input.schema();
578        let input_schema = exec_input.schema();
579        let range_exec: Vec<RangeFnExec> = self
580            .range_expr
581            .iter()
582            .map(|range_fn| {
583                let name = range_fn.expr.schema_name().to_string();
584                let range_expr = match &range_fn.expr {
585                    Expr::Alias(expr) => expr.expr.as_ref(),
586                    others => others,
587                };
588
589                let expr = match get_aggr_func(range_expr) {
590                    Some(aggr)
591                        if (aggr.func.name() == "last_value"
592                            || aggr.func.name() == "first_value") =>
593                    {
594                        let order_by = if !aggr.params.order_by.is_empty() {
595                            aggr.params
596                                .order_by
597                                .iter()
598                                .map(|x| {
599                                    create_physical_sort_expr(
600                                        x,
601                                        input_dfschema.as_ref(),
602                                        session_state.execution_props(),
603                                    )
604                                })
605                                .collect::<DfResult<Vec<_>>>()?
606                        } else {
607                            // if user not assign order by, time index is needed as default ordering
608                            let time_index = create_physical_expr(
609                                &self.time_expr,
610                                input_dfschema.as_ref(),
611                                session_state.execution_props(),
612                            )?;
613                            vec![PhysicalSortExpr {
614                                expr: time_index,
615                                options: SortOptions {
616                                    descending: false,
617                                    nulls_first: false,
618                                },
619                            }]
620                        };
621                        let arg = self.create_physical_expr_list(
622                            false,
623                            &aggr.params.args,
624                            input_dfschema,
625                            session_state,
626                        )?;
627                        // first_value/last_value has only one param.
628                        // The param have been checked by datafusion in logical plan stage.
629                        // We can safely assume that there is only one element here.
630                        AggregateExprBuilder::new(aggr.func.clone(), arg)
631                            .schema(input_schema.clone())
632                            .order_by(order_by)
633                            .alias(name)
634                            .build()
635                    }
636                    Some(aggr) => {
637                        let order_by = if !aggr.params.order_by.is_empty() {
638                            aggr.params
639                                .order_by
640                                .iter()
641                                .map(|x| {
642                                    create_physical_sort_expr(
643                                        x,
644                                        input_dfschema.as_ref(),
645                                        session_state.execution_props(),
646                                    )
647                                })
648                                .collect::<DfResult<Vec<_>>>()?
649                        } else {
650                            vec![]
651                        };
652                        let distinct = aggr.params.distinct;
653                        // TODO(discord9): add default null treatment?
654
655                        let input_phy_exprs = self.create_physical_expr_list(
656                            aggr.func.name() == "count",
657                            &aggr.params.args,
658                            input_dfschema,
659                            session_state,
660                        )?;
661                        AggregateExprBuilder::new(aggr.func.clone(), input_phy_exprs)
662                            .schema(input_schema.clone())
663                            .order_by(order_by)
664                            .with_distinct(distinct)
665                            .alias(name)
666                            .build()
667                    }
668                    None => Err(DataFusionError::Plan(format!(
669                        "Unexpected Expr: {} in RangeSelect",
670                        range_fn.expr
671                    ))),
672                }?;
673                Ok(RangeFnExec {
674                    expr: Arc::new(expr),
675                    range: range_fn.range.as_millis() as Millisecond,
676                    fill: range_fn.fill.clone(),
677                    need_cast: if range_fn.need_cast {
678                        Some(range_fn.data_type.clone())
679                    } else {
680                        None
681                    },
682                })
683            })
684            .collect::<DfResult<Vec<_>>>()?;
685        let schema_before_project = Arc::new(Schema::new(fields));
686        let schema = if let Some(project) = &self.schema_project {
687            Arc::new(schema_before_project.project(project)?)
688        } else {
689            schema_before_project.clone()
690        };
691        let by = self.create_physical_expr_list(false, &self.by, input_dfschema, session_state)?;
692        let cache = Arc::new(PlanProperties::new(
693            EquivalenceProperties::new(schema.clone()),
694            Partitioning::UnknownPartitioning(1),
695            EmissionType::Incremental,
696            Boundedness::Bounded,
697        ));
698        Ok(Arc::new(RangeSelectExec {
699            input: exec_input,
700            range_exec,
701            align: self.align.as_millis() as Millisecond,
702            align_to: self.align_to,
703            by,
704            time_index: self.time_index.clone(),
705            schema,
706            by_schema: Arc::new(Schema::new(by_fields)),
707            metric: ExecutionPlanMetricsSet::new(),
708            schema_before_project,
709            schema_project: self.schema_project.clone(),
710            cache,
711        }))
712    }
713}
714
715/// Range function expression.
716#[derive(Debug, Clone)]
717struct RangeFnExec {
718    expr: Arc<AggregateFunctionExpr>,
719    range: Millisecond,
720    fill: Option<Fill>,
721    need_cast: Option<DataType>,
722}
723
724impl RangeFnExec {
725    /// Returns the expressions to pass to the aggregator.
726    /// It also adds the order by expressions to the list of expressions.
727    /// Order-sensitive aggregators, such as `FIRST_VALUE(x ORDER BY y)` requires this.
728    fn expressions(&self) -> Vec<Arc<dyn PhysicalExpr>> {
729        let mut exprs = self.expr.expressions();
730        exprs.extend(self.expr.order_bys().iter().map(|sort| sort.expr.clone()));
731        exprs
732    }
733}
734
735impl Display for RangeFnExec {
736    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
737        if let Some(fill) = &self.fill {
738            write!(
739                f,
740                "{} RANGE {}s FILL {}",
741                self.expr.name(),
742                self.range / 1000,
743                fill
744            )
745        } else {
746            write!(f, "{} RANGE {}s", self.expr.name(), self.range / 1000)
747        }
748    }
749}
750
751#[derive(Debug)]
752pub struct RangeSelectExec {
753    input: Arc<dyn ExecutionPlan>,
754    range_exec: Vec<RangeFnExec>,
755    align: Millisecond,
756    align_to: i64,
757    time_index: String,
758    by: Vec<Arc<dyn PhysicalExpr>>,
759    schema: SchemaRef,
760    by_schema: SchemaRef,
761    metric: ExecutionPlanMetricsSet,
762    schema_project: Option<Vec<usize>>,
763    schema_before_project: SchemaRef,
764    cache: Arc<PlanProperties>,
765}
766
767impl DisplayAs for RangeSelectExec {
768    fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
769        match t {
770            DisplayFormatType::Default
771            | DisplayFormatType::Verbose
772            | DisplayFormatType::TreeRender => {
773                write!(f, "RangeSelectExec: ")?;
774                let range_expr_strs: Vec<String> =
775                    self.range_exec.iter().map(RangeFnExec::to_string).collect();
776                let by: Vec<String> = self.by.iter().map(|e| e.to_string()).collect();
777                write!(
778                    f,
779                    "range_expr=[{}], align={}ms, align_to={}ms, align_by=[{}], time_index={}",
780                    range_expr_strs.join(", "),
781                    self.align,
782                    self.align_to,
783                    by.join(", "),
784                    self.time_index,
785                )?;
786            }
787        }
788        Ok(())
789    }
790}
791
792impl ExecutionPlan for RangeSelectExec {
793    fn as_any(&self) -> &dyn std::any::Any {
794        self
795    }
796
797    fn schema(&self) -> SchemaRef {
798        self.schema.clone()
799    }
800
801    fn required_input_distribution(&self) -> Vec<Distribution> {
802        vec![Distribution::SinglePartition]
803    }
804
805    fn properties(&self) -> &Arc<PlanProperties> {
806        &self.cache
807    }
808
809    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
810        vec![&self.input]
811    }
812
813    fn with_new_children(
814        self: Arc<Self>,
815        children: Vec<Arc<dyn ExecutionPlan>>,
816    ) -> datafusion_common::Result<Arc<dyn ExecutionPlan>> {
817        assert!(!children.is_empty());
818        Ok(Arc::new(Self {
819            input: children[0].clone(),
820            range_exec: self.range_exec.clone(),
821            time_index: self.time_index.clone(),
822            by: self.by.clone(),
823            align: self.align,
824            align_to: self.align_to,
825            schema: self.schema.clone(),
826            by_schema: self.by_schema.clone(),
827            metric: self.metric.clone(),
828            schema_before_project: self.schema_before_project.clone(),
829            schema_project: self.schema_project.clone(),
830            cache: self.cache.clone(),
831        }))
832    }
833
834    fn execute(
835        &self,
836        partition: usize,
837        context: Arc<TaskContext>,
838    ) -> DfResult<DfSendableRecordBatchStream> {
839        let baseline_metric = BaselineMetrics::new(&self.metric, partition);
840        let batch_size = context.session_config().batch_size();
841        let input = self.input.execute(partition, context)?;
842        let schema = input.schema();
843        let time_index = schema
844            .column_with_name(&self.time_index)
845            .ok_or(DataFusionError::Execution(
846                "time index column not found".into(),
847            ))?
848            .0;
849        let row_converter = RowConverter::new(
850            self.by_schema
851                .fields()
852                .iter()
853                .map(|f| SortField::new(f.data_type().clone()))
854                .collect(),
855        )?;
856        Ok(Box::pin(RangeSelectStream {
857            batch_size,
858            schema: self.schema.clone(),
859            range_exec: self.range_exec.clone(),
860            input,
861            random_state: RandomState::new(),
862            time_index,
863            align: self.align,
864            align_to: self.align_to,
865            by: self.by.clone(),
866            series_map: HashMap::new(),
867            exec_state: ExecutionState::ReadingInput,
868            num_not_null_rows: 0,
869            row_converter,
870            modify_map: HashMap::new(),
871            metric: baseline_metric,
872            schema_project: self.schema_project.clone(),
873            schema_before_project: self.schema_before_project.clone(),
874            output_batch: None,
875            output_batch_offset: 0,
876        }))
877    }
878
879    fn metrics(&self) -> Option<MetricsSet> {
880        Some(self.metric.clone_inner())
881    }
882
883    fn name(&self) -> &str {
884        "RanegSelectExec"
885    }
886}
887
888struct RangeSelectStream {
889    batch_size: usize,
890    /// the schema of output column
891    schema: SchemaRef,
892    range_exec: Vec<RangeFnExec>,
893    input: SendableRecordBatchStream,
894    /// Column index of TIME INDEX column's position in the input schema
895    time_index: usize,
896    /// the unit of `align` is millisecond
897    align: Millisecond,
898    align_to: i64,
899    by: Vec<Arc<dyn PhysicalExpr>>,
900    exec_state: ExecutionState,
901    /// Converter for the by values
902    row_converter: RowConverter,
903    random_state: RandomState,
904    /// key: time series's hash value
905    /// value: time series's state on different align_ts
906    series_map: HashMap<u64, SeriesState>,
907    /// key: `(hash of by rows, align_ts)`
908    /// value: `[row_ids]`
909    /// It is used to record the data that needs to be aggregated in each time slot during the data update process
910    modify_map: HashMap<(u64, Millisecond), Vec<u32>>,
911    /// The number of rows of not null rows in the final output
912    num_not_null_rows: usize,
913    metric: BaselineMetrics,
914    schema_project: Option<Vec<usize>>,
915    schema_before_project: SchemaRef,
916    output_batch: Option<RecordBatch>,
917    output_batch_offset: usize,
918}
919
920#[derive(Debug)]
921struct SeriesState {
922    /// by values written by `RowWriter`
923    row: OwnedRow,
924    /// key: align_ts
925    /// value: a vector, each element is a range_fn follow the order of `range_exec`
926    align_ts_accumulator: BTreeMap<Millisecond, Vec<Box<dyn Accumulator>>>,
927}
928
929/// Use `align_to` as time origin.
930/// According to `align` as time interval, produces aligned time.
931/// Combining the parameters related to the range query,
932/// determine for each `Accumulator` `(hash, align_ts)` define,
933/// which rows of data will be applied to it.
934fn produce_align_time(
935    align_to: i64,
936    range: Millisecond,
937    align: Millisecond,
938    ts_column: &TimestampMillisecondArray,
939    by_columns_hash: &[u64],
940    modify_map: &mut HashMap<(u64, Millisecond), Vec<u32>>,
941) {
942    modify_map.clear();
943    // make modify_map for range_fn[i]
944    for (row, hash) in by_columns_hash.iter().enumerate() {
945        let ts = ts_column.value(row);
946        let ith_slot = (ts - align_to).div_floor(align);
947        let mut align_ts = ith_slot * align + align_to;
948        while align_ts <= ts && ts < align_ts + range {
949            modify_map
950                .entry((*hash, align_ts))
951                .or_default()
952                .push(row as u32);
953            align_ts -= align;
954        }
955    }
956}
957
958fn cast_scalar_values(values: &mut [ScalarValue], data_type: &DataType) -> DfResult<()> {
959    let array = ScalarValue::iter_to_array(values.to_vec())?;
960    let cast_array = cast_with_options(&array, data_type, &CastOptions::default())?;
961    for (i, value) in values.iter_mut().enumerate() {
962        *value = ScalarValue::try_from_array(&cast_array, i)?;
963    }
964    Ok(())
965}
966
967impl RangeSelectStream {
968    fn evaluate_many(
969        &self,
970        batch: &RecordBatch,
971        exprs: &[Arc<dyn PhysicalExpr>],
972    ) -> DfResult<Vec<ArrayRef>> {
973        exprs
974            .iter()
975            .map(|expr| {
976                let value = expr.evaluate(batch)?;
977                value.into_array(batch.num_rows())
978            })
979            .collect::<DfResult<Vec<_>>>()
980    }
981
982    fn update_range_context(&mut self, batch: RecordBatch) -> DfResult<()> {
983        let _timer = self.metric.elapsed_compute().timer();
984        let num_rows = batch.num_rows();
985        let by_arrays = self.evaluate_many(&batch, &self.by)?;
986        let mut hashes = vec![0; num_rows];
987        create_hashes(&by_arrays, &self.random_state, &mut hashes)?;
988        let by_rows = self.row_converter.convert_columns(&by_arrays)?;
989        let mut ts_column = batch.column(self.time_index).clone();
990        if !matches!(
991            ts_column.data_type(),
992            DataType::Timestamp(TimeUnit::Millisecond, _)
993        ) {
994            ts_column = compute::cast(
995                ts_column.as_ref(),
996                &DataType::Timestamp(TimeUnit::Millisecond, None),
997            )?;
998        }
999        let ts_column_ref = ts_column
1000            .as_any()
1001            .downcast_ref::<TimestampMillisecondArray>()
1002            .ok_or_else(|| {
1003                DataFusionError::Execution(
1004                    "Time index Column downcast to TimestampMillisecondArray failed".into(),
1005                )
1006            })?;
1007        for i in 0..self.range_exec.len() {
1008            let args = self.evaluate_many(&batch, &self.range_exec[i].expressions())?;
1009            // use self.modify_map record (hash, align_ts) => [row_nums]
1010            produce_align_time(
1011                self.align_to,
1012                self.range_exec[i].range,
1013                self.align,
1014                ts_column_ref,
1015                &hashes,
1016                &mut self.modify_map,
1017            );
1018            // build modify_rows/modify_index/offsets for batch update
1019            let mut modify_rows = UInt32Builder::with_capacity(0);
1020            // (hash, align_ts, row_num)
1021            // row_num use to find a by value
1022            // So we just need to record the row_num of a modify row randomly, because they all have the same by value
1023            let mut modify_index = Vec::with_capacity(self.modify_map.len());
1024            let mut offsets = vec![0];
1025            let mut offset_so_far = 0;
1026            for ((hash, ts), modify) in &self.modify_map {
1027                modify_rows.append_slice(modify);
1028                offset_so_far += modify.len();
1029                offsets.push(offset_so_far);
1030                modify_index.push((*hash, *ts, modify[0]));
1031            }
1032            let modify_rows = modify_rows.finish();
1033            let args = take_arrays(&args, &modify_rows, None)?;
1034            modify_index.iter().zip(offsets.windows(2)).try_for_each(
1035                |((hash, ts, row), offset)| {
1036                    let (offset, length) = (offset[0], offset[1] - offset[0]);
1037                    let sliced_arrays: Vec<ArrayRef> = args
1038                        .iter()
1039                        .map(|array| array.slice(offset, length))
1040                        .collect();
1041                    let accumulators_map =
1042                        self.series_map.entry(*hash).or_insert_with(|| SeriesState {
1043                            row: by_rows.row(*row as usize).owned(),
1044                            align_ts_accumulator: BTreeMap::new(),
1045                        });
1046                    match accumulators_map.align_ts_accumulator.entry(*ts) {
1047                        Entry::Occupied(mut e) => {
1048                            let accumulators = e.get_mut();
1049                            accumulators[i].update_batch(&sliced_arrays)
1050                        }
1051                        Entry::Vacant(e) => {
1052                            self.num_not_null_rows += 1;
1053                            let mut accumulators = self
1054                                .range_exec
1055                                .iter()
1056                                .map(|range| range.expr.create_accumulator())
1057                                .collect::<DfResult<Vec<_>>>()?;
1058                            let result = accumulators[i].update_batch(&sliced_arrays);
1059                            e.insert(accumulators);
1060                            result
1061                        }
1062                    }
1063                },
1064            )?;
1065        }
1066        Ok(())
1067    }
1068
1069    fn generate_output(&mut self) -> DfResult<RecordBatch> {
1070        let _timer = self.metric.elapsed_compute().timer();
1071        if self.series_map.is_empty() {
1072            return Ok(RecordBatch::new_empty(self.schema.clone()));
1073        }
1074        // 1 for time index column
1075        let mut columns: Vec<Arc<dyn Array>> =
1076            Vec::with_capacity(1 + self.range_exec.len() + self.by.len());
1077        let mut ts_builder = TimestampMillisecondBuilder::with_capacity(self.num_not_null_rows);
1078        let mut all_scalar =
1079            vec![Vec::with_capacity(self.num_not_null_rows); self.range_exec.len()];
1080        let mut by_rows = Vec::with_capacity(self.num_not_null_rows);
1081        let mut start_index = 0;
1082        // If any range expr need fill, we need fill both the missing align_ts and null value.
1083        let need_fill_output = self.range_exec.iter().any(|range| range.fill.is_some());
1084        // The padding value for each accumulator
1085        let padding_values = self
1086            .range_exec
1087            .iter()
1088            .map(|e| e.expr.create_accumulator()?.evaluate())
1089            .collect::<DfResult<Vec<_>>>()?;
1090        for SeriesState {
1091            row,
1092            align_ts_accumulator,
1093        } in self.series_map.values_mut()
1094        {
1095            // skip empty time series
1096            if align_ts_accumulator.is_empty() {
1097                continue;
1098            }
1099            // find the first and last align_ts
1100            let begin_ts = *align_ts_accumulator.first_key_value().unwrap().0;
1101            let end_ts = *align_ts_accumulator.last_key_value().unwrap().0;
1102            let align_ts = if need_fill_output {
1103                // we need to fill empty align_ts which not data in that solt
1104                (begin_ts..=end_ts).step_by(self.align as usize).collect()
1105            } else {
1106                align_ts_accumulator.keys().copied().collect::<Vec<_>>()
1107            };
1108            for ts in &align_ts {
1109                if let Some(slot) = align_ts_accumulator.get_mut(ts) {
1110                    for (column, acc) in all_scalar.iter_mut().zip(slot.iter_mut()) {
1111                        column.push(acc.evaluate()?);
1112                    }
1113                } else {
1114                    // fill null in empty time solt
1115                    for (column, padding) in all_scalar.iter_mut().zip(padding_values.iter()) {
1116                        column.push(padding.clone())
1117                    }
1118                }
1119            }
1120            ts_builder.append_slice(&align_ts);
1121            // apply fill strategy on time series
1122            for (
1123                i,
1124                RangeFnExec {
1125                    fill, need_cast, ..
1126                },
1127            ) in self.range_exec.iter().enumerate()
1128            {
1129                let time_series_data =
1130                    &mut all_scalar[i][start_index..start_index + align_ts.len()];
1131                if let Some(data_type) = need_cast {
1132                    cast_scalar_values(time_series_data, data_type)?;
1133                }
1134                if let Some(fill) = fill {
1135                    fill.apply_fill_strategy(&align_ts, time_series_data)?;
1136                }
1137            }
1138            by_rows.resize(by_rows.len() + align_ts.len(), row.row());
1139            start_index += align_ts.len();
1140        }
1141        for column_scalar in all_scalar {
1142            columns.push(ScalarValue::iter_to_array(column_scalar)?);
1143        }
1144        let ts_column = ts_builder.finish();
1145        // output schema before project follow the order of range expr | time index | by columns
1146        let ts_column = compute::cast(
1147            &ts_column,
1148            self.schema_before_project.field(columns.len()).data_type(),
1149        )?;
1150        columns.push(ts_column);
1151        // RowConverter decodes dictionary sort fields to their value arrays. Re-encode them so
1152        // the physical batch continues to match the logical output schema.
1153        for by_column in self.row_converter.convert_rows(by_rows)? {
1154            let output_type = self.schema_before_project.field(columns.len()).data_type();
1155            if by_column.data_type() == output_type {
1156                columns.push(by_column);
1157            } else {
1158                columns.push(compute::cast(by_column.as_ref(), output_type)?);
1159            }
1160        }
1161        let output = RecordBatch::try_new(self.schema_before_project.clone(), columns)?;
1162        let project_output = if let Some(project) = &self.schema_project {
1163            output.project(project)?
1164        } else {
1165            output
1166        };
1167        Ok(project_output)
1168    }
1169
1170    fn next_output_batch(&mut self) -> DfResult<Option<RecordBatch>> {
1171        if self.output_batch.is_none() {
1172            self.output_batch = Some(self.generate_output()?);
1173            self.output_batch_offset = 0;
1174        }
1175
1176        let num_rows = self.output_batch.as_ref().unwrap().num_rows();
1177        if num_rows == 0 {
1178            self.output_batch = None;
1179            self.output_batch_offset = 0;
1180            return Ok(None);
1181        }
1182
1183        if self.output_batch_offset == 0 && num_rows <= self.batch_size {
1184            return Ok(self.output_batch.take());
1185        }
1186
1187        let offset = self.output_batch_offset;
1188        let len = (num_rows - offset).min(self.batch_size);
1189        let batch = self.output_batch.as_ref().unwrap().slice(offset, len);
1190        self.output_batch_offset += len;
1191
1192        if self.output_batch_offset >= num_rows {
1193            self.output_batch = None;
1194            self.output_batch_offset = 0;
1195        }
1196
1197        Ok(Some(batch))
1198    }
1199}
1200
1201enum ExecutionState {
1202    ReadingInput,
1203    ProducingOutput,
1204    Done,
1205}
1206
1207impl RecordBatchStream for RangeSelectStream {
1208    fn schema(&self) -> SchemaRef {
1209        self.schema.clone()
1210    }
1211}
1212
1213impl Stream for RangeSelectStream {
1214    type Item = DataFusionResult<RecordBatch>;
1215
1216    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
1217        loop {
1218            match self.exec_state {
1219                ExecutionState::ReadingInput => {
1220                    match ready!(self.input.poll_next_unpin(cx)) {
1221                        // new batch to aggregate
1222                        Some(Ok(batch)) => {
1223                            if let Err(e) = self.update_range_context(batch) {
1224                                common_telemetry::debug!(
1225                                    "RangeSelectStream cannot update range context, schema: {:?}, err: {:?}",
1226                                    self.schema,
1227                                    e
1228                                );
1229                                return Poll::Ready(Some(Err(e)));
1230                            }
1231                        }
1232                        // inner had error, return to caller
1233                        Some(Err(e)) => return Poll::Ready(Some(Err(e))),
1234                        // inner is done, producing output
1235                        None => {
1236                            self.exec_state = ExecutionState::ProducingOutput;
1237                        }
1238                    }
1239                }
1240                ExecutionState::ProducingOutput => {
1241                    let result = self.next_output_batch();
1242                    return match result {
1243                        // made output
1244                        Ok(Some(batch)) => {
1245                            if self.output_batch.is_none() {
1246                                self.exec_state = ExecutionState::Done;
1247                            }
1248                            Poll::Ready(Some(Ok(batch)))
1249                        }
1250                        Ok(None) => {
1251                            self.exec_state = ExecutionState::Done;
1252                            Poll::Ready(None)
1253                        }
1254                        // error making output
1255                        Err(error) => Poll::Ready(Some(Err(error))),
1256                    };
1257                }
1258                ExecutionState::Done => return Poll::Ready(None),
1259            }
1260        }
1261    }
1262}
1263
1264#[cfg(test)]
1265mod test {
1266    macro_rules! nullable_array {
1267        ($builder:ident,) => {
1268        };
1269        ($array_type:ident ; $($tail:tt)*) => {
1270            paste::item! {
1271                {
1272                    let mut builder = arrow::array::[<$array_type Builder>]::new();
1273                    nullable_array!(builder, $($tail)*);
1274                    builder.finish()
1275                }
1276            }
1277        };
1278        ($builder:ident, null) => {
1279            $builder.append_null();
1280        };
1281        ($builder:ident, null, $($tail:tt)*) => {
1282            $builder.append_null();
1283            nullable_array!($builder, $($tail)*);
1284        };
1285        ($builder:ident, $value:literal) => {
1286            $builder.append_value($value);
1287        };
1288        ($builder:ident, $value:literal, $($tail:tt)*) => {
1289            $builder.append_value($value);
1290            nullable_array!($builder, $($tail)*);
1291        };
1292    }
1293
1294    use std::sync::Arc;
1295
1296    use arrow_schema::SortOptions;
1297    use datafusion::arrow::datatypes::{
1298        ArrowPrimitiveType, DataType, Field, Schema, TimestampMillisecondType,
1299    };
1300    use datafusion::datasource::memory::MemorySourceConfig;
1301    use datafusion::datasource::source::DataSourceExec;
1302    use datafusion::functions_aggregate::min_max;
1303    use datafusion::physical_plan::sorts::sort::SortExec;
1304    use datafusion::prelude::SessionContext;
1305    use datafusion_physical_expr::PhysicalSortExpr;
1306    use datafusion_physical_expr::expressions::Column;
1307    use datatypes::arrow::array::{Float64Array, Int64Array, TimestampMillisecondArray};
1308    use datatypes::arrow_array::StringArray;
1309
1310    use super::*;
1311
1312    const TIME_INDEX_COLUMN: &str = "timestamp";
1313
1314    fn prepare_test_data(is_float: bool, is_gap: bool) -> DataSourceExec {
1315        let schema = Arc::new(Schema::new(vec![
1316            Field::new(TIME_INDEX_COLUMN, TimestampMillisecondType::DATA_TYPE, true),
1317            Field::new(
1318                "value",
1319                if is_float {
1320                    DataType::Float64
1321                } else {
1322                    DataType::Int64
1323                },
1324                true,
1325            ),
1326            Field::new("host", DataType::Utf8, true),
1327        ]));
1328        let timestamp_column: Arc<dyn Array> = if !is_gap {
1329            Arc::new(TimestampMillisecondArray::from(vec![
1330                0, 5_000, 10_000, 15_000, 20_000, // host 1 every 5s
1331                0, 5_000, 10_000, 15_000, 20_000, // host 2 every 5s
1332            ])) as _
1333        } else {
1334            Arc::new(TimestampMillisecondArray::from(vec![
1335                0, 15_000, // host 1 every 5s, missing data on 5_000, 10_000
1336                0, 15_000, // host 2 every 5s, missing data on 5_000, 10_000
1337            ])) as _
1338        };
1339        let mut host = vec!["host1"; timestamp_column.len() / 2];
1340        host.extend(vec!["host2"; timestamp_column.len() / 2]);
1341        let mut value_column: Arc<dyn Array> = if is_gap {
1342            Arc::new(nullable_array!(Int64;
1343                0, 6, // data for host 1
1344                6, 12 // data for host 2
1345            )) as _
1346        } else {
1347            Arc::new(nullable_array!(Int64;
1348                0, null, 1, null, 2, // data for host 1
1349                3, null, 4, null, 5 // data for host 2
1350            )) as _
1351        };
1352        if is_float {
1353            value_column =
1354                cast_with_options(&value_column, &DataType::Float64, &CastOptions::default())
1355                    .unwrap();
1356        }
1357        let host_column: Arc<dyn Array> = Arc::new(StringArray::from(host)) as _;
1358        let data = RecordBatch::try_new(
1359            schema.clone(),
1360            vec![timestamp_column, value_column, host_column],
1361        )
1362        .unwrap();
1363
1364        DataSourceExec::new(Arc::new(
1365            MemorySourceConfig::try_new(&[vec![data]], schema, None).unwrap(),
1366        ))
1367    }
1368
1369    fn prepare_empty_test_data(is_float: bool) -> DataSourceExec {
1370        let schema = Arc::new(Schema::new(vec![
1371            Field::new(TIME_INDEX_COLUMN, TimestampMillisecondType::DATA_TYPE, true),
1372            Field::new(
1373                "value",
1374                if is_float {
1375                    DataType::Float64
1376                } else {
1377                    DataType::Int64
1378                },
1379                true,
1380            ),
1381            Field::new("host", DataType::Utf8, true),
1382        ]));
1383        let timestamp_column: Arc<dyn Array> =
1384            Arc::new(TimestampMillisecondArray::from(Vec::<i64>::new())) as _;
1385        let value_column: Arc<dyn Array> = if is_float {
1386            Arc::new(Float64Array::from(Vec::<Option<f64>>::new())) as _
1387        } else {
1388            Arc::new(Int64Array::from(Vec::<Option<i64>>::new())) as _
1389        };
1390        let host_column: Arc<dyn Array> =
1391            Arc::new(StringArray::from(Vec::<Option<&str>>::new())) as _;
1392        let data = RecordBatch::try_new(
1393            schema.clone(),
1394            vec![timestamp_column, value_column, host_column],
1395        )
1396        .unwrap();
1397
1398        DataSourceExec::new(Arc::new(
1399            MemorySourceConfig::try_new(&[vec![data]], schema, None).unwrap(),
1400        ))
1401    }
1402
1403    async fn collect_range_select_test(
1404        range1: Millisecond,
1405        range2: Millisecond,
1406        align: Millisecond,
1407        fill: Option<Fill>,
1408        is_float: bool,
1409        is_gap: bool,
1410        batch_size: usize,
1411    ) -> Vec<RecordBatch> {
1412        let data_type = if is_float {
1413            DataType::Float64
1414        } else {
1415            DataType::Int64
1416        };
1417        let (need_cast, schema_data_type) = if !is_float && matches!(fill, Some(Fill::Linear)) {
1418            // data_type = DataType::Float64;
1419            (Some(DataType::Float64), DataType::Float64)
1420        } else {
1421            (None, data_type.clone())
1422        };
1423        let memory_exec = Arc::new(prepare_test_data(is_float, is_gap));
1424        let schema = Arc::new(Schema::new(vec![
1425            Field::new("MIN(value)", schema_data_type.clone(), true),
1426            Field::new("MAX(value)", schema_data_type, true),
1427            Field::new(TIME_INDEX_COLUMN, TimestampMillisecondType::DATA_TYPE, true),
1428            Field::new("host", DataType::Utf8, true),
1429        ]));
1430        let cache = Arc::new(PlanProperties::new(
1431            EquivalenceProperties::new(schema.clone()),
1432            Partitioning::UnknownPartitioning(1),
1433            EmissionType::Incremental,
1434            Boundedness::Bounded,
1435        ));
1436        let input_schema = memory_exec.schema().clone();
1437        let range_select_exec = Arc::new(RangeSelectExec {
1438            input: memory_exec,
1439            range_exec: vec![
1440                RangeFnExec {
1441                    expr: Arc::new(
1442                        AggregateExprBuilder::new(
1443                            min_max::min_udaf(),
1444                            vec![Arc::new(Column::new("value", 1))],
1445                        )
1446                        .schema(input_schema.clone())
1447                        .alias("MIN(value)")
1448                        .build()
1449                        .unwrap(),
1450                    ),
1451                    range: range1,
1452                    fill: fill.clone(),
1453                    need_cast: need_cast.clone(),
1454                },
1455                RangeFnExec {
1456                    expr: Arc::new(
1457                        AggregateExprBuilder::new(
1458                            min_max::max_udaf(),
1459                            vec![Arc::new(Column::new("value", 1))],
1460                        )
1461                        .schema(input_schema.clone())
1462                        .alias("MAX(value)")
1463                        .build()
1464                        .unwrap(),
1465                    ),
1466                    range: range2,
1467                    fill,
1468                    need_cast,
1469                },
1470            ],
1471            align,
1472            align_to: 0,
1473            by: vec![Arc::new(Column::new("host", 2))],
1474            time_index: TIME_INDEX_COLUMN.to_string(),
1475            schema: schema.clone(),
1476            schema_before_project: schema.clone(),
1477            schema_project: None,
1478            by_schema: Arc::new(Schema::new(vec![Field::new("host", DataType::Utf8, true)])),
1479            metric: ExecutionPlanMetricsSet::new(),
1480            cache,
1481        });
1482        let sort_exec = SortExec::new(
1483            [
1484                PhysicalSortExpr {
1485                    expr: Arc::new(Column::new("host", 3)),
1486                    options: SortOptions {
1487                        descending: false,
1488                        nulls_first: true,
1489                    },
1490                },
1491                PhysicalSortExpr {
1492                    expr: Arc::new(Column::new(TIME_INDEX_COLUMN, 2)),
1493                    options: SortOptions {
1494                        descending: false,
1495                        nulls_first: true,
1496                    },
1497                },
1498            ]
1499            .into(),
1500            range_select_exec,
1501        );
1502        let session_context = SessionContext::new_with_config(
1503            datafusion::execution::config::SessionConfig::new().with_batch_size(batch_size),
1504        );
1505        datafusion::physical_plan::collect(Arc::new(sort_exec), session_context.task_ctx())
1506            .await
1507            .unwrap()
1508    }
1509
1510    async fn do_range_select_test(
1511        range1: Millisecond,
1512        range2: Millisecond,
1513        align: Millisecond,
1514        fill: Option<Fill>,
1515        is_float: bool,
1516        is_gap: bool,
1517        expected: String,
1518    ) {
1519        let result =
1520            collect_range_select_test(range1, range2, align, fill, is_float, is_gap, 8192).await;
1521
1522        let result_literal = arrow::util::pretty::pretty_format_batches(&result)
1523            .unwrap()
1524            .to_string();
1525
1526        assert_eq!(result_literal, expected);
1527    }
1528
1529    #[tokio::test]
1530    async fn range_10s_align_1000s() {
1531        let expected = String::from(
1532            "+------------+------------+---------------------+-------+\
1533            \n| MIN(value) | MAX(value) | timestamp           | host  |\
1534            \n+------------+------------+---------------------+-------+\
1535            \n| 0.0        | 0.0        | 1970-01-01T00:00:00 | host1 |\
1536            \n| 3.0        | 3.0        | 1970-01-01T00:00:00 | host2 |\
1537            \n+------------+------------+---------------------+-------+",
1538        );
1539        do_range_select_test(
1540            10_000,
1541            10_000,
1542            1_000_000,
1543            Some(Fill::Null),
1544            true,
1545            false,
1546            expected,
1547        )
1548        .await;
1549    }
1550
1551    #[tokio::test]
1552    async fn range_fill_null() {
1553        let expected = String::from(
1554            "+------------+------------+---------------------+-------+\
1555            \n| MIN(value) | MAX(value) | timestamp           | host  |\
1556            \n+------------+------------+---------------------+-------+\
1557            \n| 0.0        |            | 1969-12-31T23:59:55 | host1 |\
1558            \n| 0.0        | 0.0        | 1970-01-01T00:00:00 | host1 |\
1559            \n| 1.0        |            | 1970-01-01T00:00:05 | host1 |\
1560            \n| 1.0        | 1.0        | 1970-01-01T00:00:10 | host1 |\
1561            \n| 2.0        |            | 1970-01-01T00:00:15 | host1 |\
1562            \n| 2.0        | 2.0        | 1970-01-01T00:00:20 | host1 |\
1563            \n| 3.0        |            | 1969-12-31T23:59:55 | host2 |\
1564            \n| 3.0        | 3.0        | 1970-01-01T00:00:00 | host2 |\
1565            \n| 4.0        |            | 1970-01-01T00:00:05 | host2 |\
1566            \n| 4.0        | 4.0        | 1970-01-01T00:00:10 | host2 |\
1567            \n| 5.0        |            | 1970-01-01T00:00:15 | host2 |\
1568            \n| 5.0        | 5.0        | 1970-01-01T00:00:20 | host2 |\
1569            \n+------------+------------+---------------------+-------+",
1570        );
1571        do_range_select_test(
1572            10_000,
1573            5_000,
1574            5_000,
1575            Some(Fill::Null),
1576            true,
1577            false,
1578            expected,
1579        )
1580        .await;
1581    }
1582
1583    #[tokio::test]
1584    async fn range_fill_prev() {
1585        let expected = String::from(
1586            "+------------+------------+---------------------+-------+\
1587            \n| MIN(value) | MAX(value) | timestamp           | host  |\
1588            \n+------------+------------+---------------------+-------+\
1589            \n| 0.0        |            | 1969-12-31T23:59:55 | host1 |\
1590            \n| 0.0        | 0.0        | 1970-01-01T00:00:00 | host1 |\
1591            \n| 1.0        | 0.0        | 1970-01-01T00:00:05 | host1 |\
1592            \n| 1.0        | 1.0        | 1970-01-01T00:00:10 | host1 |\
1593            \n| 2.0        | 1.0        | 1970-01-01T00:00:15 | host1 |\
1594            \n| 2.0        | 2.0        | 1970-01-01T00:00:20 | host1 |\
1595            \n| 3.0        |            | 1969-12-31T23:59:55 | host2 |\
1596            \n| 3.0        | 3.0        | 1970-01-01T00:00:00 | host2 |\
1597            \n| 4.0        | 3.0        | 1970-01-01T00:00:05 | host2 |\
1598            \n| 4.0        | 4.0        | 1970-01-01T00:00:10 | host2 |\
1599            \n| 5.0        | 4.0        | 1970-01-01T00:00:15 | host2 |\
1600            \n| 5.0        | 5.0        | 1970-01-01T00:00:20 | host2 |\
1601            \n+------------+------------+---------------------+-------+",
1602        );
1603        do_range_select_test(
1604            10_000,
1605            5_000,
1606            5_000,
1607            Some(Fill::Prev),
1608            true,
1609            false,
1610            expected,
1611        )
1612        .await;
1613    }
1614
1615    #[tokio::test]
1616    async fn range_fill_linear() {
1617        let expected = String::from(
1618            "+------------+------------+---------------------+-------+\
1619            \n| MIN(value) | MAX(value) | timestamp           | host  |\
1620            \n+------------+------------+---------------------+-------+\
1621            \n| 0.0        | -0.5       | 1969-12-31T23:59:55 | host1 |\
1622            \n| 0.0        | 0.0        | 1970-01-01T00:00:00 | host1 |\
1623            \n| 1.0        | 0.5        | 1970-01-01T00:00:05 | host1 |\
1624            \n| 1.0        | 1.0        | 1970-01-01T00:00:10 | host1 |\
1625            \n| 2.0        | 1.5        | 1970-01-01T00:00:15 | host1 |\
1626            \n| 2.0        | 2.0        | 1970-01-01T00:00:20 | host1 |\
1627            \n| 3.0        | 2.5        | 1969-12-31T23:59:55 | host2 |\
1628            \n| 3.0        | 3.0        | 1970-01-01T00:00:00 | host2 |\
1629            \n| 4.0        | 3.5        | 1970-01-01T00:00:05 | host2 |\
1630            \n| 4.0        | 4.0        | 1970-01-01T00:00:10 | host2 |\
1631            \n| 5.0        | 4.5        | 1970-01-01T00:00:15 | host2 |\
1632            \n| 5.0        | 5.0        | 1970-01-01T00:00:20 | host2 |\
1633            \n+------------+------------+---------------------+-------+",
1634        );
1635        do_range_select_test(
1636            10_000,
1637            5_000,
1638            5_000,
1639            Some(Fill::Linear),
1640            true,
1641            false,
1642            expected,
1643        )
1644        .await;
1645    }
1646
1647    #[tokio::test]
1648    async fn range_fill_integer_linear() {
1649        let expected = String::from(
1650            "+------------+------------+---------------------+-------+\
1651            \n| MIN(value) | MAX(value) | timestamp           | host  |\
1652            \n+------------+------------+---------------------+-------+\
1653            \n| 0.0        | -0.5       | 1969-12-31T23:59:55 | host1 |\
1654            \n| 0.0        | 0.0        | 1970-01-01T00:00:00 | host1 |\
1655            \n| 1.0        | 0.5        | 1970-01-01T00:00:05 | host1 |\
1656            \n| 1.0        | 1.0        | 1970-01-01T00:00:10 | host1 |\
1657            \n| 2.0        | 1.5        | 1970-01-01T00:00:15 | host1 |\
1658            \n| 2.0        | 2.0        | 1970-01-01T00:00:20 | host1 |\
1659            \n| 3.0        | 2.5        | 1969-12-31T23:59:55 | host2 |\
1660            \n| 3.0        | 3.0        | 1970-01-01T00:00:00 | host2 |\
1661            \n| 4.0        | 3.5        | 1970-01-01T00:00:05 | host2 |\
1662            \n| 4.0        | 4.0        | 1970-01-01T00:00:10 | host2 |\
1663            \n| 5.0        | 4.5        | 1970-01-01T00:00:15 | host2 |\
1664            \n| 5.0        | 5.0        | 1970-01-01T00:00:20 | host2 |\
1665            \n+------------+------------+---------------------+-------+",
1666        );
1667        do_range_select_test(
1668            10_000,
1669            5_000,
1670            5_000,
1671            Some(Fill::Linear),
1672            false,
1673            false,
1674            expected,
1675        )
1676        .await;
1677    }
1678
1679    #[tokio::test]
1680    async fn range_fill_const() {
1681        let expected = String::from(
1682            "+------------+------------+---------------------+-------+\
1683            \n| MIN(value) | MAX(value) | timestamp           | host  |\
1684            \n+------------+------------+---------------------+-------+\
1685            \n| 0.0        | 6.6        | 1969-12-31T23:59:55 | host1 |\
1686            \n| 0.0        | 0.0        | 1970-01-01T00:00:00 | host1 |\
1687            \n| 1.0        | 6.6        | 1970-01-01T00:00:05 | host1 |\
1688            \n| 1.0        | 1.0        | 1970-01-01T00:00:10 | host1 |\
1689            \n| 2.0        | 6.6        | 1970-01-01T00:00:15 | host1 |\
1690            \n| 2.0        | 2.0        | 1970-01-01T00:00:20 | host1 |\
1691            \n| 3.0        | 6.6        | 1969-12-31T23:59:55 | host2 |\
1692            \n| 3.0        | 3.0        | 1970-01-01T00:00:00 | host2 |\
1693            \n| 4.0        | 6.6        | 1970-01-01T00:00:05 | host2 |\
1694            \n| 4.0        | 4.0        | 1970-01-01T00:00:10 | host2 |\
1695            \n| 5.0        | 6.6        | 1970-01-01T00:00:15 | host2 |\
1696            \n| 5.0        | 5.0        | 1970-01-01T00:00:20 | host2 |\
1697            \n+------------+------------+---------------------+-------+",
1698        );
1699        do_range_select_test(
1700            10_000,
1701            5_000,
1702            5_000,
1703            Some(Fill::Const(ScalarValue::Float64(Some(6.6)))),
1704            true,
1705            false,
1706            expected,
1707        )
1708        .await;
1709    }
1710
1711    #[tokio::test]
1712    async fn range_fill_gap() {
1713        let expected = String::from(
1714            "+------------+------------+---------------------+-------+\
1715            \n| MIN(value) | MAX(value) | timestamp           | host  |\
1716            \n+------------+------------+---------------------+-------+\
1717            \n| 0.0        | 0.0        | 1970-01-01T00:00:00 | host1 |\
1718            \n| 6.0        | 6.0        | 1970-01-01T00:00:15 | host1 |\
1719            \n| 6.0        | 6.0        | 1970-01-01T00:00:00 | host2 |\
1720            \n| 12.0       | 12.0       | 1970-01-01T00:00:15 | host2 |\
1721            \n+------------+------------+---------------------+-------+",
1722        );
1723        do_range_select_test(5_000, 5_000, 5_000, None, true, true, expected).await;
1724        let expected = String::from(
1725            "+------------+------------+---------------------+-------+\
1726            \n| MIN(value) | MAX(value) | timestamp           | host  |\
1727            \n+------------+------------+---------------------+-------+\
1728            \n| 0.0        | 0.0        | 1970-01-01T00:00:00 | host1 |\
1729            \n|            |            | 1970-01-01T00:00:05 | host1 |\
1730            \n|            |            | 1970-01-01T00:00:10 | host1 |\
1731            \n| 6.0        | 6.0        | 1970-01-01T00:00:15 | host1 |\
1732            \n| 6.0        | 6.0        | 1970-01-01T00:00:00 | host2 |\
1733            \n|            |            | 1970-01-01T00:00:05 | host2 |\
1734            \n|            |            | 1970-01-01T00:00:10 | host2 |\
1735            \n| 12.0       | 12.0       | 1970-01-01T00:00:15 | host2 |\
1736            \n+------------+------------+---------------------+-------+",
1737        );
1738        do_range_select_test(5_000, 5_000, 5_000, Some(Fill::Null), true, true, expected).await;
1739        let expected = String::from(
1740            "+------------+------------+---------------------+-------+\
1741            \n| MIN(value) | MAX(value) | timestamp           | host  |\
1742            \n+------------+------------+---------------------+-------+\
1743            \n| 0.0        | 0.0        | 1970-01-01T00:00:00 | host1 |\
1744            \n| 0.0        | 0.0        | 1970-01-01T00:00:05 | host1 |\
1745            \n| 0.0        | 0.0        | 1970-01-01T00:00:10 | host1 |\
1746            \n| 6.0        | 6.0        | 1970-01-01T00:00:15 | host1 |\
1747            \n| 6.0        | 6.0        | 1970-01-01T00:00:00 | host2 |\
1748            \n| 6.0        | 6.0        | 1970-01-01T00:00:05 | host2 |\
1749            \n| 6.0        | 6.0        | 1970-01-01T00:00:10 | host2 |\
1750            \n| 12.0       | 12.0       | 1970-01-01T00:00:15 | host2 |\
1751            \n+------------+------------+---------------------+-------+",
1752        );
1753        do_range_select_test(5_000, 5_000, 5_000, Some(Fill::Prev), true, true, expected).await;
1754        let expected = String::from(
1755            "+------------+------------+---------------------+-------+\
1756            \n| MIN(value) | MAX(value) | timestamp           | host  |\
1757            \n+------------+------------+---------------------+-------+\
1758            \n| 0.0        | 0.0        | 1970-01-01T00:00:00 | host1 |\
1759            \n| 2.0        | 2.0        | 1970-01-01T00:00:05 | host1 |\
1760            \n| 4.0        | 4.0        | 1970-01-01T00:00:10 | host1 |\
1761            \n| 6.0        | 6.0        | 1970-01-01T00:00:15 | host1 |\
1762            \n| 6.0        | 6.0        | 1970-01-01T00:00:00 | host2 |\
1763            \n| 8.0        | 8.0        | 1970-01-01T00:00:05 | host2 |\
1764            \n| 10.0       | 10.0       | 1970-01-01T00:00:10 | host2 |\
1765            \n| 12.0       | 12.0       | 1970-01-01T00:00:15 | host2 |\
1766            \n+------------+------------+---------------------+-------+",
1767        );
1768        do_range_select_test(
1769            5_000,
1770            5_000,
1771            5_000,
1772            Some(Fill::Linear),
1773            true,
1774            true,
1775            expected,
1776        )
1777        .await;
1778        let expected = String::from(
1779            "+------------+------------+---------------------+-------+\
1780            \n| MIN(value) | MAX(value) | timestamp           | host  |\
1781            \n+------------+------------+---------------------+-------+\
1782            \n| 0.0        | 0.0        | 1970-01-01T00:00:00 | host1 |\
1783            \n| 6.0        | 6.0        | 1970-01-01T00:00:05 | host1 |\
1784            \n| 6.0        | 6.0        | 1970-01-01T00:00:10 | host1 |\
1785            \n| 6.0        | 6.0        | 1970-01-01T00:00:15 | host1 |\
1786            \n| 6.0        | 6.0        | 1970-01-01T00:00:00 | host2 |\
1787            \n| 6.0        | 6.0        | 1970-01-01T00:00:05 | host2 |\
1788            \n| 6.0        | 6.0        | 1970-01-01T00:00:10 | host2 |\
1789            \n| 12.0       | 12.0       | 1970-01-01T00:00:15 | host2 |\
1790            \n+------------+------------+---------------------+-------+",
1791        );
1792        do_range_select_test(
1793            5_000,
1794            5_000,
1795            5_000,
1796            Some(Fill::Const(ScalarValue::Float64(Some(6.0)))),
1797            true,
1798            true,
1799            expected,
1800        )
1801        .await;
1802    }
1803
1804    #[tokio::test]
1805    async fn range_select_respects_session_batch_size() {
1806        let result =
1807            collect_range_select_test(10_000, 5_000, 5_000, Some(Fill::Null), true, false, 3).await;
1808
1809        let row_counts = result
1810            .iter()
1811            .map(|batch| batch.num_rows())
1812            .collect::<Vec<_>>();
1813        assert_eq!(vec![3, 3, 3, 3], row_counts);
1814    }
1815
1816    #[tokio::test]
1817    async fn range_select_skips_empty_output_batch() {
1818        let memory_exec = Arc::new(prepare_empty_test_data(true));
1819        let schema = Arc::new(Schema::new(vec![
1820            Field::new("MIN(value)", DataType::Float64, true),
1821            Field::new("MAX(value)", DataType::Float64, true),
1822            Field::new(TIME_INDEX_COLUMN, TimestampMillisecondType::DATA_TYPE, true),
1823            Field::new("host", DataType::Utf8, true),
1824        ]));
1825        let cache = Arc::new(PlanProperties::new(
1826            EquivalenceProperties::new(schema.clone()),
1827            Partitioning::UnknownPartitioning(1),
1828            EmissionType::Incremental,
1829            Boundedness::Bounded,
1830        ));
1831        let input_schema = memory_exec.schema().clone();
1832        let range_select_exec = Arc::new(RangeSelectExec {
1833            input: memory_exec,
1834            range_exec: vec![
1835                RangeFnExec {
1836                    expr: Arc::new(
1837                        AggregateExprBuilder::new(
1838                            min_max::min_udaf(),
1839                            vec![Arc::new(Column::new("value", 1))],
1840                        )
1841                        .schema(input_schema.clone())
1842                        .alias("MIN(value)")
1843                        .build()
1844                        .unwrap(),
1845                    ),
1846                    range: 10_000,
1847                    fill: Some(Fill::Null),
1848                    need_cast: None,
1849                },
1850                RangeFnExec {
1851                    expr: Arc::new(
1852                        AggregateExprBuilder::new(
1853                            min_max::max_udaf(),
1854                            vec![Arc::new(Column::new("value", 1))],
1855                        )
1856                        .schema(input_schema)
1857                        .alias("MAX(value)")
1858                        .build()
1859                        .unwrap(),
1860                    ),
1861                    range: 5_000,
1862                    fill: Some(Fill::Null),
1863                    need_cast: None,
1864                },
1865            ],
1866            align: 5_000,
1867            align_to: 0,
1868            by: vec![Arc::new(Column::new("host", 2))],
1869            time_index: TIME_INDEX_COLUMN.to_string(),
1870            schema: schema.clone(),
1871            schema_before_project: schema.clone(),
1872            schema_project: None,
1873            by_schema: Arc::new(Schema::new(vec![Field::new("host", DataType::Utf8, true)])),
1874            metric: ExecutionPlanMetricsSet::new(),
1875            cache,
1876        });
1877        let session_context = SessionContext::new();
1878        let result =
1879            datafusion::physical_plan::collect(range_select_exec, session_context.task_ctx())
1880                .await
1881                .unwrap();
1882
1883        assert!(result.is_empty());
1884    }
1885
1886    #[test]
1887    fn fill_test() {
1888        assert!(Fill::try_from_str("", &DataType::UInt8).unwrap().is_none());
1889        assert!(Fill::try_from_str("Linear", &DataType::UInt8).unwrap() == Some(Fill::Linear));
1890        assert_eq!(
1891            Fill::try_from_str("Linear", &DataType::Boolean)
1892                .unwrap_err()
1893                .to_string(),
1894            "Error during planning: Use FILL LINEAR on Non-numeric DataType Boolean"
1895        );
1896        assert_eq!(
1897            Fill::try_from_str("WHAT", &DataType::UInt8)
1898                .unwrap_err()
1899                .to_string(),
1900            "Error during planning: WHAT is not a valid fill option, fail to convert to a const value. { Arrow error: Cast error: Cannot cast string 'WHAT' to value of UInt8 type }"
1901        );
1902        assert_eq!(
1903            Fill::try_from_str("8.0", &DataType::UInt8)
1904                .unwrap_err()
1905                .to_string(),
1906            "Error during planning: 8.0 is not a valid fill option, fail to convert to a const value. { Arrow error: Cast error: Cannot cast string '8.0' to value of UInt8 type }"
1907        );
1908        assert!(
1909            Fill::try_from_str("8", &DataType::UInt8).unwrap()
1910                == Some(Fill::Const(ScalarValue::UInt8(Some(8))))
1911        );
1912        let mut test1 = vec![
1913            ScalarValue::UInt8(Some(8)),
1914            ScalarValue::UInt8(None),
1915            ScalarValue::UInt8(Some(9)),
1916        ];
1917        Fill::Null.apply_fill_strategy(&[], &mut test1).unwrap();
1918        assert_eq!(test1[1], ScalarValue::UInt8(None));
1919        Fill::Prev.apply_fill_strategy(&[], &mut test1).unwrap();
1920        assert_eq!(test1[1], ScalarValue::UInt8(Some(8)));
1921        test1[1] = ScalarValue::UInt8(None);
1922        Fill::Const(ScalarValue::UInt8(Some(10)))
1923            .apply_fill_strategy(&[], &mut test1)
1924            .unwrap();
1925        assert_eq!(test1[1], ScalarValue::UInt8(Some(10)));
1926    }
1927
1928    #[test]
1929    fn test_fill_linear() {
1930        let ts = vec![1, 2, 3, 4, 5];
1931        let mut test = vec![
1932            ScalarValue::Float32(Some(1.0)),
1933            ScalarValue::Float32(None),
1934            ScalarValue::Float32(Some(3.0)),
1935            ScalarValue::Float32(None),
1936            ScalarValue::Float32(Some(5.0)),
1937        ];
1938        Fill::Linear.apply_fill_strategy(&ts, &mut test).unwrap();
1939        let mut test1 = vec![
1940            ScalarValue::Float32(None),
1941            ScalarValue::Float32(Some(2.0)),
1942            ScalarValue::Float32(None),
1943            ScalarValue::Float32(Some(4.0)),
1944            ScalarValue::Float32(None),
1945        ];
1946        Fill::Linear.apply_fill_strategy(&ts, &mut test1).unwrap();
1947        assert_eq!(test, test1);
1948        // test linear interpolation on irregularly spaced ts/data
1949        let ts = vec![
1950            1,   // None
1951            3,   // 1.0
1952            8,   // 11.0
1953            30,  // None
1954            88,  // 10.0
1955            108, // 5.0
1956            128, // None
1957        ];
1958        let mut test = vec![
1959            ScalarValue::Float64(None),
1960            ScalarValue::Float64(Some(1.0)),
1961            ScalarValue::Float64(Some(11.0)),
1962            ScalarValue::Float64(None),
1963            ScalarValue::Float64(Some(10.0)),
1964            ScalarValue::Float64(Some(5.0)),
1965            ScalarValue::Float64(None),
1966        ];
1967        Fill::Linear.apply_fill_strategy(&ts, &mut test).unwrap();
1968        let data: Vec<_> = test
1969            .into_iter()
1970            .map(|x| {
1971                let ScalarValue::Float64(Some(f)) = x else {
1972                    unreachable!()
1973                };
1974                f
1975            })
1976            .collect();
1977        assert_eq!(data, vec![-3.0, 1.0, 11.0, 10.725, 10.0, 5.0, 0.0]);
1978        // test corner case
1979        let ts = vec![1];
1980        let test = vec![ScalarValue::Float32(None)];
1981        let mut test1 = test.clone();
1982        Fill::Linear.apply_fill_strategy(&ts, &mut test1).unwrap();
1983        assert_eq!(test, test1);
1984    }
1985}