1use 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 pub fn apply_fill_strategy(&self, ts: &[i64], data: &mut [ScalarValue]) -> DfResult<()> {
114 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 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 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 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 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 if let Some(end) = head {
180 linear_interpolation(ts, data, end, end + 1, 0, end)?;
181 }
182 if let Some(start) = tail {
184 linear_interpolation(ts, data, start - 2, start - 1, start, data.len())?;
185 }
186 Ok(())
187 }
188}
189
190fn 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 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 pub name: String,
232 pub data_type: DataType,
233 pub expr: Expr,
234 pub range: Duration,
235 pub fill: Option<Fill>,
236 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 pub input: Arc<LogicalPlan>,
277 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 pub schema_project: Option<Vec<usize>>,
290 pub schema_before_project: DFSchemaRef,
294}
295
296impl PartialOrd for RangeSelect {
297 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
298 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 !matches!(fill, Some(Fill::Const(..))),
372 );
373 Ok((None, Arc::new(field)))
374 },
375 )
376 .collect::<DfResult<Vec<_>>>()?;
377 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 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 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 #[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 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 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 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#[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 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 schema: SchemaRef,
892 range_exec: Vec<RangeFnExec>,
893 input: SendableRecordBatchStream,
894 time_index: usize,
896 align: Millisecond,
898 align_to: i64,
899 by: Vec<Arc<dyn PhysicalExpr>>,
900 exec_state: ExecutionState,
901 row_converter: RowConverter,
903 random_state: RandomState,
904 series_map: HashMap<u64, SeriesState>,
907 modify_map: HashMap<(u64, Millisecond), Vec<u32>>,
911 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 row: OwnedRow,
924 align_ts_accumulator: BTreeMap<Millisecond, Vec<Box<dyn Accumulator>>>,
927}
928
929fn 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 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 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 let mut modify_rows = UInt32Builder::with_capacity(0);
1020 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 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 let need_fill_output = self.range_exec.iter().any(|range| range.fill.is_some());
1084 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 if align_ts_accumulator.is_empty() {
1097 continue;
1098 }
1099 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 (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 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 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 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 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 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 Some(Err(e)) => return Poll::Ready(Some(Err(e))),
1234 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 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 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, 0, 5_000, 10_000, 15_000, 20_000, ])) as _
1333 } else {
1334 Arc::new(TimestampMillisecondArray::from(vec![
1335 0, 15_000, 0, 15_000, ])) 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, 6, 12 )) as _
1346 } else {
1347 Arc::new(nullable_array!(Int64;
1348 0, null, 1, null, 2, 3, null, 4, null, 5 )) 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 (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 let ts = vec![
1950 1, 3, 8, 30, 88, 108, 128, ];
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 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}