promql/extension_plan/
empty_metric.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashMap;
16use std::ops::Div;
17use std::pin::Pin;
18use std::sync::Arc;
19use std::task::{Context, Poll};
20
21use datafusion::arrow::array::ArrayRef;
22use datafusion::arrow::datatypes::{DataType, TimeUnit};
23use datafusion::catalog::Session;
24use datafusion::common::arrow::datatypes::Field;
25use datafusion::common::stats::Precision;
26use datafusion::common::tree_node::TreeNodeRecursion;
27use datafusion::common::{
28    DFSchema, DFSchemaRef, Result as DataFusionResult, Statistics, TableReference,
29};
30use datafusion::datasource::{MemTable, provider_as_source};
31use datafusion::error::DataFusionError;
32use datafusion::execution::context::TaskContext;
33use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext;
34use datafusion::logical_expr::{ExprSchemable, LogicalPlan, UserDefinedLogicalNodeCore};
35use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr, PhysicalExprRef};
36use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
37use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
38use datafusion::physical_plan::{
39    DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, RecordBatchStream,
40    SendableRecordBatchStream, StatisticsArgs,
41};
42use datafusion::physical_planner::PhysicalPlanner;
43use datafusion::prelude::{Expr, ident, lit};
44use datafusion_expr::LogicalPlanBuilder;
45use datatypes::arrow::array::TimestampMillisecondArray;
46use datatypes::arrow::datatypes::SchemaRef;
47use datatypes::arrow::record_batch::RecordBatch;
48use futures::Stream;
49
50use crate::extension_plan::Millisecond;
51
52/// Empty source plan that generate record batch with two columns:
53/// - time index column, computed from start, end and interval
54/// - value column, generated by the input expr. The expr should not
55///   reference any column except the time index column.
56#[derive(Debug, Clone, PartialEq, Eq, Hash)]
57pub struct EmptyMetric {
58    start: Millisecond,
59    end: Millisecond,
60    interval: Millisecond,
61    expr: Option<Expr>,
62    /// Schema that only contains the time index column.
63    /// This is for intermediate result only.
64    time_index_schema: DFSchemaRef,
65    /// Schema of the output record batch
66    result_schema: DFSchemaRef,
67    // This dummy input's sole purpose is to provide a schema for use in DataFusion's
68    // `SimplifyExpressions`. Otherwise it may report a "no field name ..." error.
69    // The error is caused by an optimization that tries to rewrite "A = A", during
70    // which will find the field in plan's schema. However, the schema is empty if the
71    // plan does not have an input.
72    dummy_input: LogicalPlan,
73}
74
75impl EmptyMetric {
76    pub fn new(
77        start: Millisecond,
78        end: Millisecond,
79        interval: Millisecond,
80        time_index_column_name: String,
81        field_column_name: String,
82        field_expr: Option<Expr>,
83    ) -> DataFusionResult<Self> {
84        let qualifier = Some(TableReference::bare(""));
85        let ts_only_schema = build_ts_only_schema(&time_index_column_name);
86        let mut fields = vec![(qualifier.clone(), ts_only_schema.field(0).clone())];
87        if let Some(field_expr) = &field_expr {
88            let field_data_type = field_expr.get_type(&ts_only_schema)?;
89            fields.push((
90                qualifier.clone(),
91                Arc::new(Field::new(field_column_name, field_data_type, true)),
92            ));
93        }
94        let schema = Arc::new(DFSchema::new_with_metadata(fields, HashMap::new())?);
95
96        let table = MemTable::try_new(Arc::new(schema.as_arrow().clone()), vec![vec![]])?;
97        let source = provider_as_source(Arc::new(table));
98        let dummy_input =
99            LogicalPlanBuilder::scan("dummy", source, None).and_then(|x| x.build())?;
100
101        Ok(Self {
102            start,
103            end,
104            interval,
105            time_index_schema: Arc::new(ts_only_schema),
106            result_schema: schema,
107            expr: field_expr,
108            dummy_input,
109        })
110    }
111
112    /// Returns the expression used to generate metric values.
113    pub fn expr(&self) -> Option<&Expr> {
114        self.expr.as_ref()
115    }
116
117    pub const fn name() -> &'static str {
118        "EmptyMetric"
119    }
120
121    pub fn to_execution_plan(
122        &self,
123        session: &dyn Session,
124        physical_planner: &dyn PhysicalPlanner,
125        planning_ctx: &PhysicalPlanningContext,
126    ) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
127        let physical_expr = self
128            .expr
129            .as_ref()
130            .map(|expr| {
131                physical_planner.create_physical_expr(
132                    expr,
133                    &self.time_index_schema,
134                    session,
135                    planning_ctx,
136                )
137            })
138            .transpose()?;
139        let result_schema: SchemaRef = self.result_schema.inner().clone();
140        let properties = Arc::new(PlanProperties::new(
141            EquivalenceProperties::new(result_schema.clone()),
142            Partitioning::UnknownPartitioning(1),
143            EmissionType::Incremental,
144            Boundedness::Bounded,
145        ));
146        Ok(Arc::new(EmptyMetricExec {
147            start: self.start,
148            end: self.end,
149            interval: self.interval,
150            time_index_schema: self.time_index_schema.inner().clone(),
151            result_schema,
152            expr: physical_expr,
153            properties,
154            metric: ExecutionPlanMetricsSet::new(),
155        }))
156    }
157}
158
159impl UserDefinedLogicalNodeCore for EmptyMetric {
160    fn name(&self) -> &str {
161        Self::name()
162    }
163
164    fn inputs(&self) -> Vec<&LogicalPlan> {
165        vec![&self.dummy_input]
166    }
167
168    fn schema(&self) -> &DFSchemaRef {
169        &self.result_schema
170    }
171
172    fn expressions(&self) -> Vec<Expr> {
173        if let Some(expr) = &self.expr {
174            vec![expr.clone()]
175        } else {
176            vec![]
177        }
178    }
179
180    fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
181        write!(
182            f,
183            "EmptyMetric: range=[{}..{}], interval=[{}]",
184            self.start, self.end, self.interval,
185        )
186    }
187
188    fn with_exprs_and_inputs(
189        &self,
190        exprs: Vec<Expr>,
191        _inputs: Vec<LogicalPlan>,
192    ) -> DataFusionResult<Self> {
193        Ok(Self {
194            start: self.start,
195            end: self.end,
196            interval: self.interval,
197            expr: exprs.into_iter().next(),
198            time_index_schema: self.time_index_schema.clone(),
199            result_schema: self.result_schema.clone(),
200            dummy_input: self.dummy_input.clone(),
201        })
202    }
203}
204
205impl PartialOrd for EmptyMetric {
206    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
207        // Compare fields in order excluding schema fields
208        match self.start.partial_cmp(&other.start) {
209            Some(core::cmp::Ordering::Equal) => {}
210            ord => return ord,
211        }
212        match self.end.partial_cmp(&other.end) {
213            Some(core::cmp::Ordering::Equal) => {}
214            ord => return ord,
215        }
216        match self.interval.partial_cmp(&other.interval) {
217            Some(core::cmp::Ordering::Equal) => {}
218            ord => return ord,
219        }
220        self.expr.partial_cmp(&other.expr)
221    }
222}
223
224#[derive(Debug, Clone)]
225pub struct EmptyMetricExec {
226    start: Millisecond,
227    end: Millisecond,
228    interval: Millisecond,
229    /// Schema that only contains the time index column.
230    /// This is for intermediate result only.
231    time_index_schema: SchemaRef,
232    /// Schema of the output record batch
233    result_schema: SchemaRef,
234    expr: Option<PhysicalExprRef>,
235    properties: Arc<PlanProperties>,
236    metric: ExecutionPlanMetricsSet,
237}
238
239impl ExecutionPlan for EmptyMetricExec {
240    fn apply_expressions(
241        &self,
242        f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> datafusion_common::Result<TreeNodeRecursion>,
243    ) -> DataFusionResult<TreeNodeRecursion> {
244        datafusion::physical_plan::apply_expression_roots(self.expr.iter(), f)
245    }
246
247    fn schema(&self) -> SchemaRef {
248        self.result_schema.clone()
249    }
250
251    fn properties(&self) -> &Arc<PlanProperties> {
252        &self.properties
253    }
254
255    fn maintains_input_order(&self) -> Vec<bool> {
256        vec![]
257    }
258
259    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
260        vec![]
261    }
262
263    fn with_new_children(
264        self: Arc<Self>,
265        _children: Vec<Arc<dyn ExecutionPlan>>,
266    ) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
267        Ok(Arc::new(self.as_ref().clone()))
268    }
269
270    fn execute(
271        &self,
272        partition: usize,
273        _context: Arc<TaskContext>,
274    ) -> DataFusionResult<SendableRecordBatchStream> {
275        let baseline_metric = BaselineMetrics::new(&self.metric, partition);
276        Ok(Box::pin(EmptyMetricStream {
277            start: self.start,
278            end: self.end,
279            interval: self.interval,
280            expr: self.expr.clone(),
281            is_first_poll: true,
282            time_index_schema: self.time_index_schema.clone(),
283            result_schema: self.result_schema.clone(),
284            metric: baseline_metric,
285        }))
286    }
287
288    fn metrics(&self) -> Option<MetricsSet> {
289        Some(self.metric.clone_inner())
290    }
291
292    fn statistics_from_inputs(
293        &self,
294        _input_stats: &[Arc<Statistics>],
295        args: &StatisticsArgs,
296    ) -> DataFusionResult<Arc<Statistics>> {
297        let partition = args.partition();
298        if partition.is_some() {
299            return Ok(Arc::new(Statistics::new_unknown(self.schema().as_ref())));
300        }
301
302        let estimated_row_num = if self.end > self.start {
303            (self.end - self.start) as f64 / self.interval as f64
304        } else {
305            0.0
306        };
307        let total_byte_size = estimated_row_num * std::mem::size_of::<Millisecond>() as f64;
308
309        Ok(Arc::new(Statistics {
310            num_rows: Precision::Inexact(estimated_row_num.floor() as _),
311            total_byte_size: Precision::Inexact(total_byte_size.floor() as _),
312            column_statistics: Statistics::unknown_column(&self.schema()),
313        }))
314    }
315
316    fn name(&self) -> &str {
317        "EmptyMetricExec"
318    }
319}
320
321impl DisplayAs for EmptyMetricExec {
322    fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
323        match t {
324            DisplayFormatType::Default
325            | DisplayFormatType::Verbose
326            | DisplayFormatType::TreeRender => write!(
327                f,
328                "EmptyMetric: range=[{}..{}], interval=[{}]",
329                self.start, self.end, self.interval,
330            ),
331        }
332    }
333}
334
335pub struct EmptyMetricStream {
336    start: Millisecond,
337    end: Millisecond,
338    interval: Millisecond,
339    expr: Option<PhysicalExprRef>,
340    /// This stream only generate one record batch at the first poll
341    is_first_poll: bool,
342    /// Schema that only contains the time index column.
343    /// This is for intermediate result only.
344    time_index_schema: SchemaRef,
345    /// Schema of the output record batch
346    result_schema: SchemaRef,
347    metric: BaselineMetrics,
348}
349
350impl RecordBatchStream for EmptyMetricStream {
351    fn schema(&self) -> SchemaRef {
352        self.result_schema.clone()
353    }
354}
355
356impl Stream for EmptyMetricStream {
357    type Item = DataFusionResult<RecordBatch>;
358
359    fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
360        let result = if self.is_first_poll {
361            self.is_first_poll = false;
362            let _timer = self.metric.elapsed_compute().timer();
363
364            // build the time index array, and a record batch that
365            // only contains that array as the input of field expr
366            let time_array = (self.start..=self.end)
367                .step_by(self.interval as _)
368                .collect::<Vec<_>>();
369            let time_array = Arc::new(TimestampMillisecondArray::from(time_array));
370            let num_rows = time_array.len();
371            let input_record_batch =
372                RecordBatch::try_new(self.time_index_schema.clone(), vec![time_array.clone()])
373                    .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?;
374            let mut result_arrays: Vec<ArrayRef> = vec![time_array];
375
376            // evaluate the field expr and get the result
377            if let Some(field_expr) = &self.expr {
378                result_arrays.push(
379                    field_expr
380                        .evaluate(&input_record_batch)
381                        .and_then(|x| x.into_array(num_rows))?,
382                );
383            }
384
385            // assemble the output record batch
386            let batch = RecordBatch::try_new(self.result_schema.clone(), result_arrays)
387                .map_err(|e| DataFusionError::ArrowError(Box::new(e), None));
388
389            Poll::Ready(Some(batch))
390        } else {
391            Poll::Ready(None)
392        };
393        self.metric.record_poll(result)
394    }
395}
396
397/// Build a schema that only contains **millisecond** timestamp column
398fn build_ts_only_schema(column_name: &str) -> DFSchema {
399    let ts_field = Field::new(
400        column_name,
401        DataType::Timestamp(TimeUnit::Millisecond, None),
402        false,
403    );
404    // safety: should not fail (UT covers this)
405    DFSchema::new_with_metadata(
406        vec![(Some(TableReference::bare("")), Arc::new(ts_field))],
407        HashMap::new(),
408    )
409    .unwrap()
410}
411
412// Convert timestamp column to UNIX epoch second:
413// https://prometheus.io/docs/prometheus/latest/querying/functions/#time
414pub fn build_special_time_expr(time_index_column_name: &str) -> Expr {
415    let input_schema = build_ts_only_schema(time_index_column_name);
416    // safety: should not failed (UT covers this)
417    ident(time_index_column_name)
418        .cast_to(&DataType::Int64, &input_schema)
419        .unwrap()
420        .cast_to(&DataType::Float64, &input_schema)
421        .unwrap()
422        .div(lit(1000.0)) // cast to second will lost precision, so we cast to float64 first and manually divide by 1000
423}
424
425#[cfg(test)]
426mod test {
427    use datafusion::physical_planner::DefaultPhysicalPlanner;
428    use datafusion::prelude::SessionContext;
429
430    use super::*;
431
432    async fn do_empty_metric_test(
433        start: Millisecond,
434        end: Millisecond,
435        interval: Millisecond,
436        time_column_name: String,
437        field_column_name: String,
438        expected: String,
439    ) {
440        let session_context = SessionContext::default();
441        let df_default_physical_planner = DefaultPhysicalPlanner::default();
442        let time_expr = build_special_time_expr(&time_column_name);
443        let empty_metric = EmptyMetric::new(
444            start,
445            end,
446            interval,
447            time_column_name,
448            field_column_name,
449            Some(time_expr),
450        )
451        .unwrap();
452        let empty_metric_exec = empty_metric
453            .to_execution_plan(
454                &session_context.state(),
455                &df_default_physical_planner,
456                &PhysicalPlanningContext::default(),
457            )
458            .unwrap();
459
460        let result =
461            datafusion::physical_plan::collect(empty_metric_exec, session_context.task_ctx())
462                .await
463                .unwrap();
464        let result_literal = datatypes::arrow::util::pretty::pretty_format_batches(&result)
465            .unwrap()
466            .to_string();
467
468        assert_eq!(result_literal, expected);
469    }
470
471    #[tokio::test]
472    async fn normal_empty_metric_test() {
473        do_empty_metric_test(
474            0,
475            100,
476            10,
477            "time".to_string(),
478            "value".to_string(),
479            String::from(
480                "+-------------------------+-------+\
481                \n| time                    | value |\
482                \n+-------------------------+-------+\
483                \n| 1970-01-01T00:00:00     | 0.0   |\
484                \n| 1970-01-01T00:00:00.010 | 0.01  |\
485                \n| 1970-01-01T00:00:00.020 | 0.02  |\
486                \n| 1970-01-01T00:00:00.030 | 0.03  |\
487                \n| 1970-01-01T00:00:00.040 | 0.04  |\
488                \n| 1970-01-01T00:00:00.050 | 0.05  |\
489                \n| 1970-01-01T00:00:00.060 | 0.06  |\
490                \n| 1970-01-01T00:00:00.070 | 0.07  |\
491                \n| 1970-01-01T00:00:00.080 | 0.08  |\
492                \n| 1970-01-01T00:00:00.090 | 0.09  |\
493                \n| 1970-01-01T00:00:00.100 | 0.1   |\
494                \n+-------------------------+-------+",
495            ),
496        )
497        .await
498    }
499
500    #[tokio::test]
501    async fn unaligned_empty_metric_test() {
502        do_empty_metric_test(
503            0,
504            100,
505            11,
506            "time".to_string(),
507            "value".to_string(),
508            String::from(
509                "+-------------------------+-------+\
510                \n| time                    | value |\
511                \n+-------------------------+-------+\
512                \n| 1970-01-01T00:00:00     | 0.0   |\
513                \n| 1970-01-01T00:00:00.011 | 0.011 |\
514                \n| 1970-01-01T00:00:00.022 | 0.022 |\
515                \n| 1970-01-01T00:00:00.033 | 0.033 |\
516                \n| 1970-01-01T00:00:00.044 | 0.044 |\
517                \n| 1970-01-01T00:00:00.055 | 0.055 |\
518                \n| 1970-01-01T00:00:00.066 | 0.066 |\
519                \n| 1970-01-01T00:00:00.077 | 0.077 |\
520                \n| 1970-01-01T00:00:00.088 | 0.088 |\
521                \n| 1970-01-01T00:00:00.099 | 0.099 |\
522                \n+-------------------------+-------+",
523            ),
524        )
525        .await
526    }
527
528    #[tokio::test]
529    async fn one_row_empty_metric_test() {
530        do_empty_metric_test(
531            0,
532            100,
533            1000,
534            "time".to_string(),
535            "value".to_string(),
536            String::from(
537                "+---------------------+-------+\
538                \n| time                | value |\
539                \n+---------------------+-------+\
540                \n| 1970-01-01T00:00:00 | 0.0   |\
541                \n+---------------------+-------+",
542            ),
543        )
544        .await
545    }
546
547    #[tokio::test]
548    async fn negative_range_empty_metric_test() {
549        do_empty_metric_test(
550            1000,
551            -1000,
552            10,
553            "time".to_string(),
554            "value".to_string(),
555            String::from(
556                "+------+-------+\
557                \n| time | value |\
558                \n+------+-------+\
559                \n+------+-------+",
560            ),
561        )
562        .await
563    }
564
565    #[tokio::test]
566    async fn no_field_expr() {
567        let session_context = SessionContext::default();
568        let df_default_physical_planner = DefaultPhysicalPlanner::default();
569        let empty_metric =
570            EmptyMetric::new(0, 200, 1000, "time".to_string(), "value".to_string(), None).unwrap();
571        let empty_metric_exec = empty_metric
572            .to_execution_plan(
573                &session_context.state(),
574                &df_default_physical_planner,
575                &PhysicalPlanningContext::default(),
576            )
577            .unwrap();
578
579        let result =
580            datafusion::physical_plan::collect(empty_metric_exec, session_context.task_ctx())
581                .await
582                .unwrap();
583        let result_literal = datatypes::arrow::util::pretty::pretty_format_batches(&result)
584            .unwrap()
585            .to_string();
586
587        let expected = String::from(
588            "+---------------------+\
589            \n| time                |\
590            \n+---------------------+\
591            \n| 1970-01-01T00:00:00 |\
592            \n+---------------------+",
593        );
594        assert_eq!(result_literal, expected);
595    }
596}
promql/extension_plan/empty_metric.rs

promql/extension_plan/
empty_metric.rs
