Skip to main content

promql/extension_plan/
instant_manipulate.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::pin::Pin;
16use std::sync::Arc;
17use std::task::{Context, Poll};
18
19use common_query::prelude::{greptime_native_histogram, greptime_value};
20use datafusion::arrow::array::{Array, TimestampMillisecondArray, UInt64Array};
21use datafusion::arrow::datatypes::{DataType, SchemaRef};
22use datafusion::arrow::record_batch::RecordBatch;
23use datafusion::common::stats::Precision;
24use datafusion::common::tree_node::TreeNodeRecursion;
25use datafusion::common::{DFSchema, DFSchemaRef, ScalarValue};
26use datafusion::error::{DataFusionError, Result as DataFusionResult};
27use datafusion::execution::context::TaskContext;
28use datafusion::logical_expr::{
29    EmptyRelation, Expr, Extension, LogicalPlan, UserDefinedLogicalNodeCore,
30};
31use datafusion::physical_expr::EquivalenceProperties;
32use datafusion::physical_plan::metrics::{
33    BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricValue, MetricsSet,
34};
35use datafusion::physical_plan::{
36    ChildStats, DisplayAs, DisplayFormatType, ExecutionPlan, InputDistributionRequirements,
37    PhysicalExpr, PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics,
38    StatisticsArgs,
39};
40use datafusion_expr::ident;
41use datatypes::arrow::compute;
42use datatypes::timestamp::timestamp_array_to_primitive;
43use futures::{Stream, StreamExt, ready};
44use greptime_proto::substrait_extension as pb;
45use prost::Message;
46use snafu::ResultExt;
47
48use crate::error::{DeserializeSnafu, Result};
49use crate::extension_plan::range_manipulate::{Nanoseconds, ScanNanoseconds, narrow_scan};
50use crate::extension_plan::series_divide::SeriesDivide;
51use crate::extension_plan::{
52    METRIC_NUM_SERIES, Millisecond, is_prometheus_stale_sample, local_offset,
53    nanoseconds_per_native_tick, prometheus_stale_sample_column, resolve_column_name,
54    serialize_column_index, timestamp_unit,
55};
56use crate::metrics::PROMQL_SERIES_COUNT;
57
58const MAX_INSTANT_MANIPULATE_OUTPUT_POINTS: usize = 1_000_000;
59
60fn mixed_sample_fields(field: Option<&str>) -> [Option<&str>; 2] {
61    let companion = match field {
62        Some(field) if field == greptime_value() => Some(greptime_native_histogram()),
63        Some(field) if field == greptime_native_histogram() => Some(greptime_value()),
64        _ => None,
65    };
66    [field, companion]
67}
68
69/// Manipulate the input record batch to make it suitable for Instant Operator.
70///
71/// This plan will try to align the input time series, for every timestamp between
72/// `start` and `end` with step `interval`. Find in the `lookback` range if data
73/// is missing at the given timestamp.
74#[derive(Debug, PartialEq, Eq, Hash)]
75pub struct InstantManipulate {
76    start: Millisecond,
77    end: Millisecond,
78    lookback_delta: Millisecond,
79    interval: Millisecond,
80    offset: Millisecond,
81    time_index_column: String,
82    // Planner-provided tag-column hint for execution fast paths.
83    tag_columns: Vec<String>,
84    /// Primary sample column used to derive the columns checked for staleness.
85    field_column: Option<String>,
86    input: LogicalPlan,
87    output_schema: DFSchemaRef,
88    unfix: Option<UnfixIndices>,
89}
90
91impl PartialOrd for InstantManipulate {
92    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
93        (
94            self.start,
95            self.end,
96            self.lookback_delta,
97            self.interval,
98            self.offset,
99            &self.time_index_column,
100            &self.tag_columns,
101            &self.field_column,
102            &self.input,
103            &self.unfix,
104        )
105            .partial_cmp(&(
106                other.start,
107                other.end,
108                other.lookback_delta,
109                other.interval,
110                other.offset,
111                &other.time_index_column,
112                &other.tag_columns,
113                &other.field_column,
114                &other.input,
115                &other.unfix,
116            ))
117    }
118}
119
120#[derive(Debug, PartialEq, Eq, Hash, PartialOrd)]
121struct UnfixIndices {
122    pub time_index_idx: u64,
123    pub field_index_idx: u64,
124}
125
126impl UserDefinedLogicalNodeCore for InstantManipulate {
127    fn name(&self) -> &str {
128        Self::name()
129    }
130
131    fn inputs(&self) -> Vec<&LogicalPlan> {
132        vec![&self.input]
133    }
134
135    fn schema(&self) -> &DFSchemaRef {
136        &self.output_schema
137    }
138
139    fn expressions(&self) -> Vec<Expr> {
140        if self.unfix.is_some() {
141            return vec![];
142        }
143
144        let mut exprs = vec![ident(&self.time_index_column)];
145        exprs.extend(self.staleness_field_columns().map(ident));
146        exprs
147    }
148
149    fn necessary_children_exprs(&self, output_columns: &[usize]) -> Option<Vec<Vec<usize>>> {
150        if self.unfix.is_some() {
151            return None;
152        }
153
154        let input_schema = self.input.schema();
155        if output_columns.is_empty() {
156            let indices = (0..input_schema.fields().len()).collect::<Vec<_>>();
157            return Some(vec![indices]);
158        }
159
160        let mut required = output_columns.to_vec();
161        required.push(input_schema.index_of_column_by_name(None, &self.time_index_column)?);
162        for field in self.staleness_field_columns() {
163            required.push(input_schema.index_of_column_by_name(None, field)?);
164        }
165
166        required.sort_unstable();
167        required.dedup();
168        Some(vec![required])
169    }
170
171    fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
172        write!(
173            f,
174            "PromInstantManipulate: range=[{}..{}], lookback=[{}], interval=[{}], time index=[{}]",
175            self.start, self.end, self.lookback_delta, self.interval, self.time_index_column
176        )
177    }
178
179    fn with_exprs_and_inputs(
180        &self,
181        _exprs: Vec<Expr>,
182        inputs: Vec<LogicalPlan>,
183    ) -> DataFusionResult<Self> {
184        if inputs.len() != 1 {
185            return Err(DataFusionError::Internal(
186                "InstantManipulate should have exact one input".to_string(),
187            ));
188        }
189
190        let input: LogicalPlan = inputs.into_iter().next().unwrap();
191        let input_schema = input.schema();
192
193        if let Some(unfix) = &self.unfix {
194            // transform indices to names
195            let time_index_column = resolve_column_name(
196                unfix.time_index_idx,
197                input_schema,
198                "InstantManipulate",
199                "time index",
200            )?;
201
202            let field_column = if unfix.field_index_idx == u64::MAX {
203                None
204            } else {
205                Some(resolve_column_name(
206                    unfix.field_index_idx,
207                    input_schema,
208                    "InstantManipulate",
209                    "field",
210                )?)
211            };
212
213            Ok(Self {
214                start: self.start,
215                end: self.end,
216                lookback_delta: self.lookback_delta,
217                interval: self.interval,
218                offset: local_offset(&input, &time_index_column),
219                output_schema: Self::calculate_output_schema(&input, &time_index_column)?,
220                time_index_column,
221                tag_columns: Self::resolve_tag_columns(&input, &self.tag_columns),
222                field_column,
223                input,
224                unfix: None,
225            })
226        } else {
227            Ok(Self {
228                start: self.start,
229                end: self.end,
230                lookback_delta: self.lookback_delta,
231                interval: self.interval,
232                offset: self.offset,
233                time_index_column: self.time_index_column.clone(),
234                tag_columns: Self::resolve_tag_columns(&input, &self.tag_columns),
235                field_column: self.field_column.clone(),
236                output_schema: Self::calculate_output_schema(&input, &self.time_index_column)?,
237                input,
238                unfix: None,
239            })
240        }
241    }
242}
243
244impl InstantManipulate {
245    fn calculate_output_schema(
246        input: &LogicalPlan,
247        time_index_column: &str,
248    ) -> DataFusionResult<DFSchemaRef> {
249        let input_schema = input.schema();
250        let time_index = input_schema
251            .index_of_column_by_name(None, time_index_column)
252            .ok_or_else(|| {
253                DataFusionError::Internal(format!(
254                    "InstantManipulate time index {time_index_column} not found"
255                ))
256            })?;
257        let mut fields = (0..input_schema.fields().len())
258            .map(|index| {
259                let (qualifier, field) = input_schema.qualified_field(index);
260                (qualifier.cloned(), field.clone())
261            })
262            .collect::<Vec<_>>();
263        let (qualifier, field) = input_schema.qualified_field(time_index);
264        fields[time_index] = (
265            qualifier.cloned(),
266            Arc::new(field.as_ref().clone().with_data_type(DataType::Timestamp(
267                datafusion::arrow::datatypes::TimeUnit::Millisecond,
268                None,
269            ))),
270        );
271        Ok(Arc::new(DFSchema::new_with_metadata(
272            fields,
273            input_schema.metadata().clone(),
274        )?))
275    }
276
277    #[allow(clippy::too_many_arguments)]
278    pub fn new(
279        start: Millisecond,
280        end: Millisecond,
281        lookback_delta: Millisecond,
282        interval: Millisecond,
283        offset: Millisecond,
284        time_index_column: String,
285        tag_columns: Vec<String>,
286        field_column: Option<String>,
287        input: LogicalPlan,
288    ) -> Self {
289        Self {
290            start,
291            end,
292            lookback_delta,
293            interval,
294            offset,
295            output_schema: Self::calculate_output_schema(&input, &time_index_column)
296                .unwrap_or_else(|_| input.schema().clone()),
297            time_index_column,
298            tag_columns,
299            field_column,
300            input,
301            unfix: None,
302        }
303    }
304
305    pub const fn name() -> &'static str {
306        "InstantManipulate"
307    }
308
309    /// Returns whether this node evaluates a single timestamp rather than a range.
310    pub fn is_single_evaluation(&self) -> bool {
311        self.start == self.end
312    }
313
314    fn staleness_field_columns(&self) -> impl Iterator<Item = &str> {
315        let [field, companion] = mixed_sample_fields(self.field_column.as_deref());
316        [
317            field,
318            companion.filter(|companion| {
319                self.input
320                    .schema()
321                    .index_of_column_by_name(None, companion)
322                    .is_some()
323            }),
324        ]
325        .into_iter()
326        .flatten()
327    }
328
329    fn resolve_tag_columns(input: &LogicalPlan, tag_columns: &[String]) -> Vec<String> {
330        if !tag_columns.is_empty() {
331            return tag_columns.to_vec();
332        }
333
334        Self::find_series_divide_tags(input).unwrap_or_default()
335    }
336
337    fn find_series_divide_tags(plan: &LogicalPlan) -> Option<Vec<String>> {
338        if let LogicalPlan::Extension(Extension { node }) = plan
339            && let Some(series_divide) = node.as_any().downcast_ref::<SeriesDivide>()
340        {
341            return Some(series_divide.tags().to_vec());
342        }
343
344        plan.inputs()
345            .into_iter()
346            .find_map(Self::find_series_divide_tags)
347    }
348
349    pub fn to_execution_plan(&self, exec_input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
350        let reuse_tsid_column = matches!(self.tag_columns.as_slice(), [tag] if tag == "__tsid");
351
352        let mut fields = exec_input.schema().fields().to_vec();
353        let time_index = exec_input
354            .schema()
355            .index_of(&self.time_index_column)
356            .expect("time index column not found");
357        fields[time_index] = Arc::new(fields[time_index].as_ref().clone().with_data_type(
358            DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
359        ));
360        let output_schema = Arc::new(datafusion::arrow::datatypes::Schema::new_with_metadata(
361            fields,
362            exec_input.schema().metadata().clone(),
363        ));
364        let input_properties = exec_input.properties();
365        let properties = Arc::new(PlanProperties::new(
366            EquivalenceProperties::new(output_schema.clone()),
367            input_properties.partitioning.clone(),
368            input_properties.emission_type,
369            input_properties.boundedness,
370        ));
371        Arc::new(InstantManipulateExec {
372            offset: self.offset,
373            start: self.start,
374            end: self.end,
375            lookback_delta: self.lookback_delta,
376            interval: self.interval,
377            time_index_column: self.time_index_column.clone(),
378            field_column: self.field_column.clone(),
379            reuse_tsid_column,
380            input: exec_input,
381            output_schema,
382            properties,
383            metric: ExecutionPlanMetricsSet::new(),
384        })
385    }
386
387    pub fn serialize(&self) -> Vec<u8> {
388        let time_index_idx = serialize_column_index(self.input.schema(), &self.time_index_column);
389
390        let field_index_idx = self
391            .field_column
392            .as_ref()
393            .map(|name| serialize_column_index(self.input.schema(), name))
394            .unwrap_or(u64::MAX);
395
396        pb::InstantManipulate {
397            start: self.start,
398            end: self.end,
399            interval: self.interval,
400            lookback_delta: self.lookback_delta,
401            time_index_idx,
402            field_index_idx,
403            ..Default::default()
404        }
405        .encode_to_vec()
406    }
407
408    pub fn deserialize(bytes: &[u8]) -> Result<Self> {
409        let pb_instant_manipulate =
410            pb::InstantManipulate::decode(bytes).context(DeserializeSnafu)?;
411        let empty_schema = Arc::new(DFSchema::empty());
412        let placeholder_plan = LogicalPlan::EmptyRelation(EmptyRelation {
413            produce_one_row: false,
414            schema: empty_schema.clone(),
415        });
416
417        let unfix = UnfixIndices {
418            time_index_idx: pb_instant_manipulate.time_index_idx,
419            field_index_idx: pb_instant_manipulate.field_index_idx,
420        };
421
422        Ok(Self {
423            start: pb_instant_manipulate.start,
424            end: pb_instant_manipulate.end,
425            lookback_delta: pb_instant_manipulate.lookback_delta,
426            interval: pb_instant_manipulate.interval,
427            offset: 0,
428            time_index_column: String::new(),
429            tag_columns: Vec::new(),
430            field_column: None,
431            output_schema: empty_schema,
432            input: placeholder_plan,
433            unfix: Some(unfix),
434        })
435    }
436}
437
438#[derive(Debug)]
439pub struct InstantManipulateExec {
440    offset: Millisecond,
441    start: Millisecond,
442    end: Millisecond,
443    lookback_delta: Millisecond,
444    interval: Millisecond,
445    time_index_column: String,
446    field_column: Option<String>,
447    reuse_tsid_column: bool,
448
449    input: Arc<dyn ExecutionPlan>,
450    output_schema: SchemaRef,
451    properties: Arc<PlanProperties>,
452    metric: ExecutionPlanMetricsSet,
453}
454
455impl ExecutionPlan for InstantManipulateExec {
456    fn apply_expressions(
457        &self,
458        _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> datafusion_common::Result<TreeNodeRecursion>,
459    ) -> DataFusionResult<TreeNodeRecursion> {
460        Ok(TreeNodeRecursion::Continue)
461    }
462
463    fn schema(&self) -> SchemaRef {
464        self.output_schema.clone()
465    }
466
467    fn properties(&self) -> &Arc<PlanProperties> {
468        &self.properties
469    }
470
471    fn input_distribution_requirements(&self) -> InputDistributionRequirements {
472        self.input.input_distribution_requirements()
473    }
474
475    // Prevent reordering of input
476    fn maintains_input_order(&self) -> Vec<bool> {
477        vec![false; self.children().len()]
478    }
479
480    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
481        vec![&self.input]
482    }
483
484    fn with_new_children(
485        self: Arc<Self>,
486        children: Vec<Arc<dyn ExecutionPlan>>,
487    ) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
488        assert!(!children.is_empty());
489        let input = children[0].clone();
490        // The child may expose native timestamps, but our output schema uses ms.
491        // Retain its execution properties, not its schema equivalences or ordering.
492        let input_properties = input.properties();
493        let properties = Arc::new(PlanProperties::new(
494            EquivalenceProperties::new(self.output_schema.clone()),
495            input_properties.partitioning.clone(),
496            input_properties.emission_type,
497            input_properties.boundedness,
498        ));
499        Ok(Arc::new(Self {
500            offset: self.offset,
501            start: self.start,
502            end: self.end,
503            lookback_delta: self.lookback_delta,
504            interval: self.interval,
505            time_index_column: self.time_index_column.clone(),
506            field_column: self.field_column.clone(),
507            reuse_tsid_column: self.reuse_tsid_column,
508            input,
509            output_schema: self.output_schema.clone(),
510            properties,
511            metric: self.metric.clone(),
512        }))
513    }
514
515    fn execute(
516        &self,
517        partition: usize,
518        context: Arc<TaskContext>,
519    ) -> DataFusionResult<SendableRecordBatchStream> {
520        let baseline_metric = BaselineMetrics::new(&self.metric, partition);
521        let num_series = Count::new();
522        MetricBuilder::new(&self.metric)
523            .with_partition(partition)
524            .build(MetricValue::Count {
525                name: METRIC_NUM_SERIES.into(),
526                count: num_series.clone(),
527            });
528
529        let input = self.input.execute(partition, context)?;
530        let schema = input.schema();
531        let time_index = schema
532            .column_with_name(&self.time_index_column)
533            .expect("time index column not found")
534            .0;
535        let time_unit = timestamp_unit(schema.field(time_index).data_type())?;
536        let field_indices = mixed_sample_fields(self.field_column.as_deref()).map(|field| {
537            field.and_then(|field| schema.column_with_name(field).map(|(index, _)| index))
538        });
539        let tsid_index = schema
540            .column_with_name("__tsid")
541            .filter(|(_, field)| field.data_type() == &DataType::UInt64)
542            .map(|(index, _)| index);
543        Ok(Box::pin(InstantManipulateStream {
544            offset: self.offset,
545            start: self.start,
546            end: self.end,
547            lookback_delta: self.lookback_delta,
548            interval: self.interval,
549            time_index,
550            time_unit,
551            field_indices,
552            tsid_index,
553            reuse_tsid_column: self.reuse_tsid_column && tsid_index.is_some(),
554            schema: self.output_schema.clone(),
555            input,
556            metric: baseline_metric,
557            num_series,
558        }))
559    }
560
561    fn metrics(&self) -> Option<MetricsSet> {
562        Some(self.metric.clone_inner())
563    }
564
565    fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {
566        vec![ChildStats::At(partition)]
567    }
568
569    fn statistics_from_inputs(
570        &self,
571        input_stats: &[Arc<Statistics>],
572        _args: &StatisticsArgs,
573    ) -> DataFusionResult<Arc<Statistics>> {
574        let input_stats = &input_stats[0];
575
576        let estimated_row_num = (self.end - self.start) as f64 / self.interval as f64;
577        let estimated_total_bytes = input_stats
578            .total_byte_size
579            .get_value()
580            .zip(input_stats.num_rows.get_value())
581            .map(|(size, rows)| {
582                Precision::Inexact(((*size as f64 / *rows as f64) * estimated_row_num).floor() as _)
583            })
584            .unwrap_or(Precision::Absent);
585
586        Ok(Arc::new(Statistics {
587            num_rows: Precision::Inexact(estimated_row_num.floor() as _),
588            total_byte_size: estimated_total_bytes,
589            // TODO(ruihang): support this column statistics
590            column_statistics: Statistics::unknown_column(&self.schema()),
591        }))
592    }
593
594    fn name(&self) -> &str {
595        "InstantManipulateExec"
596    }
597}
598
599impl DisplayAs for InstantManipulateExec {
600    fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
601        match t {
602            DisplayFormatType::Default
603            | DisplayFormatType::Verbose
604            | DisplayFormatType::TreeRender => {
605                write!(
606                    f,
607                    "PromInstantManipulateExec: range=[{}..{}], lookback=[{}], interval=[{}], time index=[{}]",
608                    self.start,
609                    self.end,
610                    self.lookback_delta,
611                    self.interval,
612                    self.time_index_column
613                )
614            }
615        }
616    }
617}
618
619pub struct InstantManipulateStream {
620    offset: Millisecond,
621    start: Millisecond,
622    end: Millisecond,
623    lookback_delta: Millisecond,
624    interval: Millisecond,
625    // Column index of TIME INDEX column's position in schema
626    time_index: usize,
627    time_unit: datafusion::arrow::datatypes::TimeUnit,
628    field_indices: [Option<usize>; 2],
629    tsid_index: Option<usize>,
630    reuse_tsid_column: bool,
631
632    schema: SchemaRef,
633    input: SendableRecordBatchStream,
634    metric: BaselineMetrics,
635    /// Number of series processed.
636    num_series: Count,
637}
638
639impl RecordBatchStream for InstantManipulateStream {
640    fn schema(&self) -> SchemaRef {
641        self.schema.clone()
642    }
643}
644
645impl Stream for InstantManipulateStream {
646    type Item = DataFusionResult<RecordBatch>;
647
648    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
649        let poll = match ready!(self.input.poll_next_unpin(cx)) {
650            Some(Ok(batch)) => {
651                let timer = std::time::Instant::now();
652                if batch.num_rows() != 0 {
653                    self.num_series.add(1);
654                }
655                let result = Ok(batch).and_then(|batch| self.manipulate(batch));
656                self.metric.elapsed_compute().add_elapsed(timer);
657                Poll::Ready(Some(result))
658            }
659            None => {
660                PROMQL_SERIES_COUNT.observe(self.num_series.value() as f64);
661                Poll::Ready(None)
662            }
663            Some(Err(e)) => Poll::Ready(Some(Err(e))),
664        };
665        self.metric.record_poll(poll)
666    }
667}
668
669impl InstantManipulateStream {
670    /// Manipulates one complete series sorted by timestamp. The planner enforces
671    /// this input contract with a sort followed by [`SeriesDivide`].
672    ///
673    /// Prometheus `v3.9.1`'s `vectorSelectorSingle` uses a start-exclusive
674    /// lookback window `(eval_ts - lookback_delta, eval_ts]`; a sample at exactly
675    /// `eval_ts - lookback_delta` is too old.
676    pub fn manipulate(&self, input: RecordBatch) -> DataFusionResult<RecordBatch> {
677        let ts_column = input.column(self.time_index);
678        if ts_column.is_empty() {
679            // Returning Pending after consuming a ready batch without a wake can stall the stream.
680            return Ok(RecordBatch::new_empty(self.schema.clone()));
681        }
682        let scale = nanoseconds_per_native_tick(self.time_unit);
683        let stale_sample_columns = self.field_indices.map(|index| {
684            index.and_then(|index| prometheus_stale_sample_column(input.column(index).as_ref()))
685        });
686        let is_stale = |row| {
687            stale_sample_columns
688                .iter()
689                .flatten()
690                .any(|column| is_prometheus_stale_sample(*column, row))
691        };
692        let (timestamps, _) = timestamp_array_to_primitive(ts_column).ok_or_else(|| {
693            DataFusionError::Execution("Time index column is not a timestamp".into())
694        })?;
695        let timestamp_null_count = timestamps.null_count();
696        let timestamps = timestamps.values();
697        let len = timestamps.len();
698        // Shift the native-tick timeline in i128 before comparing samples. Doing
699        // this in the Arrow storage unit can overflow even when the shifted
700        // PromQL millisecond evaluation time is representable.
701        let to_nanoseconds =
702            |timestamp: i64| (timestamp as i128) * scale + (self.offset as i128) * 1_000_000;
703        let first_ns = to_nanoseconds(timestamps[0]);
704        let last_ns = to_nanoseconds(timestamps[len - 1]);
705        // An exact sample remains useful with zero lookback. Otherwise the lower
706        // boundary is exclusive, so subtract one nanosecond from its final window.
707        let last_useful = if self.lookback_delta == 0 {
708            last_ns
709        } else {
710            last_ns + (self.lookback_delta as i128) * 1_000_000 - 1
711        };
712        let first_ms = (first_ns + 999_999).div_euclid(1_000_000);
713        let last_ms = last_useful.div_euclid(1_000_000);
714        let query_start = self.start as i128;
715        let query_end = self.end as i128;
716        let interval = self.interval as i128;
717        let max_start = first_ms.max(query_start);
718        let min_end = last_ms.min(query_end);
719        let (aligned_start, aligned_end) = if max_start > min_end {
720            (1, 0)
721        } else {
722            (
723                query_start + (max_start - query_start) / interval * interval,
724                query_end - (query_end - min_end) / interval * interval,
725            )
726        };
727        let estimated_points = if aligned_end >= aligned_start {
728            (aligned_end - aligned_start) / interval + 1
729        } else {
730            0
731        };
732        if estimated_points > MAX_INSTANT_MANIPULATE_OUTPUT_POINTS as i128 {
733            return Err(DataFusionError::Execution(format!(
734                "InstantManipulate output points exceed limit: {estimated_points} > {MAX_INSTANT_MANIPULATE_OUTPUT_POINTS}"
735            )));
736        }
737        let estimated_points = estimated_points as usize;
738        let aligned_start = aligned_start as i64;
739        let aligned_end = aligned_end as i64;
740        // Narrow `i64` nanoseconds whenever the batch fits them, exact `i128` otherwise.
741        let narrow = (timestamp_null_count == 0)
742            .then(|| {
743                narrow_scan(
744                    self.time_unit,
745                    self.offset,
746                    self.lookback_delta,
747                    timestamps,
748                    aligned_start,
749                    aligned_end,
750                )
751            })
752            .flatten();
753        if let Some(scan) = &narrow {
754            debug_assert_eq!(i64::try_from(first_ns), Ok(scan.shift(timestamps[0])));
755            debug_assert_eq!(i64::try_from(last_ns), Ok(scan.shift(timestamps[len - 1])));
756        }
757        let (take_indices, aligned_ts) = match narrow {
758            Some(scan) => self.scan_instants::<i64>(
759                timestamps,
760                aligned_start,
761                aligned_end,
762                &scan,
763                estimated_points,
764                &is_stale,
765            ),
766            None => self.scan_instants::<i128>(
767                timestamps,
768                aligned_start,
769                aligned_end,
770                &ScanNanoseconds::<i128>::wide(self.time_unit, self.offset, self.lookback_delta),
771                estimated_points,
772                &is_stale,
773            ),
774        };
775        self.take_record_batch_optional(input, take_indices, aligned_ts)
776    }
777
778    /// Chooses the sample of every aligned evaluation instant, in `Ns` nanoseconds.
779    ///
780    /// Keeps the first row among exact timestamp ties and the latest preceding row
781    /// otherwise; zero lookback admits only exact samples. Staleness is tested after
782    /// choosing: a selected stale marker suppresses this evaluation rather than
783    /// falling back to an older finite value.
784    fn scan_instants<Ns: Nanoseconds>(
785        &self,
786        timestamps: &[i64],
787        start: Millisecond,
788        end: Millisecond,
789        scan: &ScanNanoseconds<Ns>,
790        capacity: usize,
791        is_stale: impl Fn(usize) -> bool,
792    ) -> (Vec<u64>, Vec<Millisecond>) {
793        let len = timestamps.len();
794        let mut cursor = 0usize;
795        let mut take_indices = Vec::with_capacity(capacity);
796        let mut aligned_ts = Vec::with_capacity(capacity);
797        for expected_ms in (start..=end).step_by(self.interval as usize) {
798            let expected = scan.instant(expected_ms);
799            let mut exact_candidate = None;
800            while cursor < len && scan.shift(timestamps[cursor]) <= expected {
801                if scan.shift(timestamps[cursor]) == expected && exact_candidate.is_none() {
802                    exact_candidate = Some(cursor);
803                }
804                cursor += 1;
805            }
806            let Some(candidate) = exact_candidate.or_else(|| cursor.checked_sub(1)) else {
807                continue;
808            };
809            let candidate_ts = scan.shift(timestamps[candidate]);
810            if (candidate_ts == expected || candidate_ts > scan.window_start(expected))
811                && candidate_ts <= expected
812                && !is_stale(candidate)
813            {
814                take_indices.push(candidate as u64);
815                aligned_ts.push(expected_ms);
816            }
817        }
818        (take_indices, aligned_ts)
819    }
820
821    /// Helper function to apply "take" on record batch.
822    fn take_record_batch_optional(
823        &self,
824        record_batch: RecordBatch,
825        take_indices: Vec<u64>,
826        aligned_ts: Vec<Millisecond>,
827    ) -> DataFusionResult<RecordBatch> {
828        assert_eq!(take_indices.len(), aligned_ts.len());
829
830        let output_len = aligned_ts.len();
831        let mut indices_array = None;
832        let mut arrays = Vec::with_capacity(record_batch.num_columns());
833        let aligned_ts = Arc::new(TimestampMillisecondArray::from(aligned_ts)) as Arc<dyn Array>;
834
835        for (index, array) in record_batch.columns().iter().enumerate() {
836            if index == self.time_index {
837                arrays.push(aligned_ts.clone());
838                continue;
839            }
840
841            if self.reuse_tsid_column && self.tsid_index == Some(index) {
842                arrays.push(reuse_constant_column(array, output_len)?);
843                continue;
844            }
845
846            let indices_array =
847                indices_array.get_or_insert_with(|| UInt64Array::from(take_indices.clone()));
848            arrays.push(compute::take(array, indices_array, None)?);
849        }
850
851        let result = RecordBatch::try_new(self.schema.clone(), arrays)
852            .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?;
853        Ok(result)
854    }
855}
856
857fn reuse_constant_column(array: &Arc<dyn Array>, len: usize) -> DataFusionResult<Arc<dyn Array>> {
858    if len <= array.len() {
859        return Ok(array.slice(0, len));
860    }
861
862    if array.is_empty() {
863        return Ok(array.slice(0, 0));
864    }
865
866    ScalarValue::try_from_array(array.as_ref(), 0)?.to_array_of_size(len)
867}
868
869#[cfg(test)]
870mod test {
871    use common_query::native_histogram::build_histogram_array;
872    use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS;
873    use datafusion::arrow::array::{
874        Float64Array, TimestampMicrosecondArray, TimestampNanosecondArray, TimestampSecondArray,
875    };
876    use datafusion::arrow::buffer::NullBuffer;
877    use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit};
878    use datafusion::common::ToDFSchema;
879    use datafusion::datasource::memory::MemorySourceConfig;
880    use datafusion::datasource::source::DataSourceExec;
881    use datafusion::logical_expr::{
882        EmptyRelation, Extension, LogicalPlan, Projection, UserDefinedLogicalNodeCore,
883    };
884    use datafusion::physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions};
885    use datafusion::prelude::SessionContext;
886    use datafusion_expr::col;
887
888    use super::*;
889    use crate::extension_plan::test_util::{
890        TIME_INDEX_COLUMN, native_histogram, prepare_test_data, prepare_test_data_with_stale_marker,
891    };
892
893    async fn do_normalize_test(
894        start: Millisecond,
895        end: Millisecond,
896        lookback_delta: Millisecond,
897        interval: Millisecond,
898        expected: String,
899        contains_stale_marker: bool,
900    ) {
901        let memory_exec = if contains_stale_marker {
902            Arc::new(prepare_test_data_with_stale_marker())
903        } else {
904            Arc::new(prepare_test_data())
905        };
906        let normalize_exec = Arc::new(InstantManipulateExec {
907            offset: 0,
908            start,
909            end,
910            lookback_delta,
911            interval,
912            time_index_column: TIME_INDEX_COLUMN.to_string(),
913            field_column: Some("value".to_string()),
914            reuse_tsid_column: false,
915            output_schema: memory_exec.schema(),
916            properties: memory_exec.properties().clone(),
917            input: memory_exec,
918            metric: ExecutionPlanMetricsSet::new(),
919        });
920        let session_context = SessionContext::default();
921        let result = datafusion::physical_plan::collect(normalize_exec, session_context.task_ctx())
922            .await
923            .unwrap();
924        let result_literal = datatypes::arrow::util::pretty::pretty_format_batches(&result)
925            .unwrap()
926            .to_string();
927
928        assert_eq!(result_literal, expected);
929    }
930
931    #[tokio::test]
932    async fn native_timestamps_select_exact_samples_and_keep_ms_output() {
933        for (unit, ticks_per_ms) in [
934            (TimeUnit::Microsecond, 1_000_i64),
935            (TimeUnit::Nanosecond, 1_000_000_i64),
936        ] {
937            let lower = 1_000 * ticks_per_ms;
938            let upper = 1_001 * ticks_per_ms;
939            let stale = f64::from_bits(PROMETHEUS_STALE_NAN_BITS);
940            for (name, timestamps, values, expected_timestamps, expected_values) in [
941                (
942                    "exact upper sample",
943                    vec![lower + 1, upper],
944                    vec![1.0, 2.0],
945                    vec![1_001],
946                    vec![2.0],
947                ),
948                (
949                    "exclusive lower boundary and future sample",
950                    vec![lower, upper + 1],
951                    vec![1.0, 2.0],
952                    vec![1_000],
953                    vec![1.0],
954                ),
955                (
956                    "one native tick above lower boundary",
957                    vec![lower + 1, upper + 1],
958                    vec![1.0, 2.0],
959                    vec![1_001],
960                    vec![1.0],
961                ),
962                (
963                    "future stale marker does not suppress",
964                    vec![lower + 1, upper + 1],
965                    vec![1.0, stale],
966                    vec![1_001],
967                    vec![1.0],
968                ),
969                (
970                    "latest in-window stale marker suppresses",
971                    vec![lower + 1, lower + 2, upper + 1],
972                    vec![1.0, stale, 3.0],
973                    vec![],
974                    vec![],
975                ),
976            ] {
977                let schema = Arc::new(Schema::new(vec![
978                    Field::new(TIME_INDEX_COLUMN, DataType::Timestamp(unit, None), false),
979                    Field::new("value", DataType::Float64, true),
980                ]));
981                let time: Arc<dyn Array> = match unit {
982                    TimeUnit::Microsecond => Arc::new(TimestampMicrosecondArray::from(timestamps)),
983                    TimeUnit::Nanosecond => Arc::new(TimestampNanosecondArray::from(timestamps)),
984                    _ => unreachable!(),
985                };
986                let batch = RecordBatch::try_new(
987                    schema.clone(),
988                    vec![time, Arc::new(Float64Array::from(values))],
989                )
990                .unwrap();
991                let logical_input = LogicalPlan::EmptyRelation(EmptyRelation {
992                    produce_one_row: false,
993                    schema: schema.clone().to_dfschema_ref().unwrap(),
994                });
995                let plan = InstantManipulate::new(
996                    1_000,
997                    1_001,
998                    1,
999                    1,
1000                    0,
1001                    TIME_INDEX_COLUMN.to_string(),
1002                    Vec::new(),
1003                    Some("value".to_string()),
1004                    logical_input.clone(),
1005                );
1006                let output_schema = Arc::new(Schema::new(vec![
1007                    Field::new(
1008                        TIME_INDEX_COLUMN,
1009                        DataType::Timestamp(TimeUnit::Millisecond, None),
1010                        false,
1011                    ),
1012                    Field::new("value", DataType::Float64, true),
1013                ]));
1014                assert_eq!(plan.schema().as_arrow(), output_schema.as_ref());
1015
1016                let rebuilt = InstantManipulate::deserialize(&plan.serialize())
1017                    .unwrap()
1018                    .with_exprs_and_inputs(vec![], vec![logical_input])
1019                    .unwrap();
1020                assert_eq!(rebuilt.schema(), plan.schema());
1021                assert_eq!(rebuilt.input.schema().as_arrow(), schema.as_ref());
1022
1023                let input = Arc::new(DataSourceExec::new(Arc::new(
1024                    MemorySourceConfig::try_new(&[vec![batch]], schema.clone(), None).unwrap(),
1025                )));
1026                let exec = rebuilt.to_execution_plan(input);
1027                assert_eq!(exec.schema(), output_schema);
1028                assert_eq!(exec.children()[0].schema(), schema);
1029
1030                let batches =
1031                    datafusion::physical_plan::collect(exec, SessionContext::default().task_ctx())
1032                        .await
1033                        .unwrap();
1034                assert_eq!(batches.len(), 1, "{unit:?}: {name}");
1035                let output = &batches[0];
1036                assert_eq!(output.schema(), output_schema);
1037                let timestamps = output
1038                    .column(0)
1039                    .as_any()
1040                    .downcast_ref::<TimestampMillisecondArray>()
1041                    .unwrap();
1042                let values = output
1043                    .column(1)
1044                    .as_any()
1045                    .downcast_ref::<Float64Array>()
1046                    .unwrap();
1047                assert_eq!(
1048                    timestamps.values().as_ref(),
1049                    expected_timestamps.as_slice(),
1050                    "{unit:?}: {name}"
1051                );
1052                assert_eq!(
1053                    values.values().as_ref(),
1054                    expected_values.as_slice(),
1055                    "{unit:?}: {name}"
1056                );
1057                assert_eq!(values.null_count(), 0, "{unit:?}: {name}");
1058            }
1059        }
1060    }
1061
1062    #[tokio::test]
1063    async fn logical_normalize_offset_survives_rebuild_and_executes() {
1064        for (name, time_unit, raw, offset, start, lookback_delta) in [
1065            (
1066                "millisecond offset",
1067                TimeUnit::Millisecond,
1068                0,
1069                1_000,
1070                1_000,
1071                0,
1072            ),
1073            (
1074                "second timestamp with negative fractional offset",
1075                TimeUnit::Second,
1076                1,
1077                -500,
1078                500,
1079                0,
1080            ),
1081        ] {
1082            let schema = Arc::new(Schema::new(vec![
1083                Field::new(
1084                    TIME_INDEX_COLUMN,
1085                    DataType::Timestamp(time_unit, None),
1086                    false,
1087                ),
1088                Field::new("value", DataType::Float64, true),
1089            ]));
1090            let input = LogicalPlan::EmptyRelation(EmptyRelation {
1091                produce_one_row: false,
1092                schema: schema.clone().to_dfschema_ref().unwrap(),
1093            });
1094            let normalize = crate::extension_plan::SeriesNormalize::new(
1095                offset,
1096                TIME_INDEX_COLUMN,
1097                false,
1098                Vec::new(),
1099                input.clone(),
1100            );
1101            let normalize =
1102                crate::extension_plan::SeriesNormalize::deserialize(&normalize.serialize())
1103                    .unwrap()
1104                    .with_exprs_and_inputs(vec![], vec![input.clone()])
1105                    .unwrap();
1106            let normalized = LogicalPlan::Projection(
1107                Projection::try_new(
1108                    vec![col(TIME_INDEX_COLUMN), col("value")],
1109                    Arc::new(LogicalPlan::Extension(Extension {
1110                        node: Arc::new(normalize),
1111                    })),
1112                )
1113                .unwrap(),
1114            );
1115            let fresh = InstantManipulate::new(
1116                start,
1117                start,
1118                lookback_delta,
1119                1,
1120                offset,
1121                TIME_INDEX_COLUMN.to_string(),
1122                Vec::new(),
1123                Some("value".to_string()),
1124                input.clone(),
1125            )
1126            .with_exprs_and_inputs(vec![], vec![input.clone()])
1127            .unwrap();
1128            let serialized = InstantManipulate::new(
1129                start,
1130                start,
1131                lookback_delta,
1132                1,
1133                offset,
1134                TIME_INDEX_COLUMN.to_string(),
1135                Vec::new(),
1136                Some("value".to_string()),
1137                normalized.clone(),
1138            );
1139            let decoded = InstantManipulate::deserialize(&serialized.serialize())
1140                .unwrap()
1141                .with_exprs_and_inputs(vec![], vec![normalized])
1142                .unwrap();
1143            let timestamp: Arc<dyn Array> = match time_unit {
1144                TimeUnit::Millisecond => Arc::new(TimestampMillisecondArray::from(vec![raw])),
1145                TimeUnit::Second => Arc::new(TimestampSecondArray::from(vec![raw])),
1146                _ => unreachable!(),
1147            };
1148            let batch = RecordBatch::try_new(
1149                schema.clone(),
1150                vec![timestamp, Arc::new(Float64Array::from(vec![7.0]))],
1151            )
1152            .unwrap();
1153            for (mode, rebuilt) in [("fresh", fresh), ("decoded", decoded)] {
1154                let rebuilt = rebuilt
1155                    .with_exprs_and_inputs(vec![], vec![input.clone()])
1156                    .unwrap();
1157                assert_eq!(rebuilt.offset, offset, "{name}: {mode}");
1158                assert_eq!(rebuilt.input.schema(), input.schema(), "{name}: {mode}");
1159
1160                let empty_exec_input = Arc::new(DataSourceExec::new(Arc::new(
1161                    MemorySourceConfig::try_new(&[vec![]], schema.clone(), None).unwrap(),
1162                )));
1163                let exec_input = Arc::new(DataSourceExec::new(Arc::new(
1164                    MemorySourceConfig::try_new(&[vec![batch.clone()]], schema.clone(), None)
1165                        .unwrap(),
1166                )));
1167                let exec = rebuilt
1168                    .to_execution_plan(empty_exec_input)
1169                    .replace_children(
1170                        vec![exec_input],
1171                        ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
1172                    )
1173                    .unwrap();
1174                let output =
1175                    datafusion::physical_plan::collect(exec, SessionContext::default().task_ctx())
1176                        .await
1177                        .unwrap();
1178                let output = &output[0];
1179                assert_eq!(output.num_rows(), 1, "{name}: {mode}");
1180                assert_eq!(
1181                    output
1182                        .column(0)
1183                        .as_any()
1184                        .downcast_ref::<TimestampMillisecondArray>()
1185                        .unwrap()
1186                        .value(0),
1187                    start,
1188                    "{name}: {mode}"
1189                );
1190                assert_eq!(
1191                    output
1192                        .column(1)
1193                        .as_any()
1194                        .downcast_ref::<Float64Array>()
1195                        .unwrap()
1196                        .value(0),
1197                    7.0,
1198                    "{name}: {mode}"
1199                );
1200            }
1201        }
1202    }
1203
1204    #[test]
1205    fn deserialized_ordering_preserves_column_indices() {
1206        let mut wire = pb::InstantManipulate::default();
1207        let first = InstantManipulate::deserialize(&wire.encode_to_vec()).unwrap();
1208        wire.time_index_idx = 1;
1209        let second = InstantManipulate::deserialize(&wire.encode_to_vec()).unwrap();
1210        assert_ne!(first, second);
1211        assert_eq!(first.partial_cmp(&second), Some(std::cmp::Ordering::Less));
1212        wire.field_index_idx = 2;
1213        let third = InstantManipulate::deserialize(&wire.encode_to_vec()).unwrap();
1214        assert_ne!(second, third);
1215        assert_eq!(second.partial_cmp(&third), Some(std::cmp::Ordering::Less));
1216
1217        let input = LogicalPlan::EmptyRelation(EmptyRelation {
1218            produce_one_row: false,
1219            schema: prepare_test_data().schema().to_dfschema_ref().unwrap(),
1220        });
1221        let first = InstantManipulate::new(
1222            0,
1223            0,
1224            0,
1225            0,
1226            0,
1227            TIME_INDEX_COLUMN.to_string(),
1228            Vec::new(),
1229            Some("value".to_string()),
1230            input.clone(),
1231        );
1232        let second = InstantManipulate::new(
1233            0,
1234            0,
1235            0,
1236            0,
1237            1,
1238            TIME_INDEX_COLUMN.to_string(),
1239            Vec::new(),
1240            Some("value".to_string()),
1241            input,
1242        );
1243        assert_ne!(first, second);
1244        assert_eq!(first.partial_cmp(&second), Some(std::cmp::Ordering::Less));
1245    }
1246
1247    #[test]
1248    fn pruning_should_keep_time_and_field_columns_for_exec() {
1249        let df_schema = prepare_test_data().schema().to_dfschema_ref().unwrap();
1250        let input = LogicalPlan::EmptyRelation(EmptyRelation {
1251            produce_one_row: false,
1252            schema: df_schema,
1253        });
1254        let plan = InstantManipulate::new(
1255            0,
1256            0,
1257            0,
1258            0,
1259            0,
1260            TIME_INDEX_COLUMN.to_string(),
1261            Vec::new(),
1262            Some("value".to_string()),
1263            input,
1264        );
1265
1266        // Simulate a parent projection requesting only the `path` column.
1267        let output_columns = [2usize];
1268        let required = plan.necessary_children_exprs(&output_columns).unwrap();
1269        let required = &required[0];
1270        assert_eq!(required.as_slice(), &[0, 1, 2]);
1271    }
1272
1273    #[test]
1274    fn rebuild_should_recover_tag_columns_from_series_divide_input() {
1275        let df_schema = prepare_test_data().schema().to_dfschema_ref().unwrap();
1276        let input = LogicalPlan::EmptyRelation(EmptyRelation {
1277            produce_one_row: false,
1278            schema: df_schema,
1279        });
1280        let series_divide = LogicalPlan::Extension(Extension {
1281            node: Arc::new(SeriesDivide::new(
1282                vec!["__tsid".to_string()],
1283                TIME_INDEX_COLUMN.to_string(),
1284                input,
1285            )),
1286        });
1287        let bytes = InstantManipulate::new(
1288            0,
1289            0,
1290            0,
1291            0,
1292            0,
1293            TIME_INDEX_COLUMN.to_string(),
1294            vec!["__tsid".to_string()],
1295            Some("value".to_string()),
1296            series_divide.clone(),
1297        )
1298        .serialize();
1299        let plan = InstantManipulate::deserialize(&bytes)
1300            .unwrap()
1301            .with_exprs_and_inputs(vec![], vec![series_divide])
1302            .unwrap();
1303
1304        assert_eq!(plan.tag_columns, vec!["__tsid".to_string()]);
1305    }
1306
1307    #[test]
1308    fn rebuild_should_recover_tag_columns_from_series_normalize_input() {
1309        let df_schema = prepare_test_data().schema().to_dfschema_ref().unwrap();
1310        let input = LogicalPlan::EmptyRelation(EmptyRelation {
1311            produce_one_row: false,
1312            schema: df_schema,
1313        });
1314        let series_divide = LogicalPlan::Extension(Extension {
1315            node: Arc::new(SeriesDivide::new(
1316                vec!["__tsid".to_string()],
1317                TIME_INDEX_COLUMN.to_string(),
1318                input,
1319            )),
1320        });
1321        let series_normalize = LogicalPlan::Extension(Extension {
1322            node: Arc::new(crate::extension_plan::SeriesNormalize::new(
1323                0,
1324                TIME_INDEX_COLUMN,
1325                false,
1326                vec!["__tsid".to_string()],
1327                series_divide,
1328            )),
1329        });
1330        let bytes = InstantManipulate::new(
1331            0,
1332            0,
1333            0,
1334            0,
1335            0,
1336            TIME_INDEX_COLUMN.to_string(),
1337            vec!["__tsid".to_string()],
1338            Some("value".to_string()),
1339            series_normalize.clone(),
1340        )
1341        .serialize();
1342        let plan = InstantManipulate::deserialize(&bytes)
1343            .unwrap()
1344            .with_exprs_and_inputs(vec![], vec![series_normalize])
1345            .unwrap();
1346
1347        assert_eq!(plan.tag_columns, vec!["__tsid".to_string()]);
1348    }
1349
1350    #[test]
1351    fn to_execution_plan_enables_tsid_fast_path() {
1352        let schema = Arc::new(Schema::new(vec![
1353            Field::new(
1354                TIME_INDEX_COLUMN,
1355                DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
1356                false,
1357            ),
1358            Field::new("value", DataType::Float64, true),
1359        ]));
1360        let exec_input: Arc<dyn ExecutionPlan> = Arc::new(DataSourceExec::new(Arc::new(
1361            MemorySourceConfig::try_new(&[], schema, None).unwrap(),
1362        )));
1363
1364        let exec = InstantManipulate::new(
1365            0,
1366            0,
1367            0,
1368            0,
1369            0,
1370            TIME_INDEX_COLUMN.to_string(),
1371            vec!["__tsid".to_string()],
1372            Some("value".to_string()),
1373            LogicalPlan::EmptyRelation(EmptyRelation {
1374                produce_one_row: false,
1375                schema: Arc::new(datafusion::common::DFSchema::empty()),
1376            }),
1377        )
1378        .to_execution_plan(exec_input);
1379
1380        assert!(format!("{exec:?}").contains("reuse_tsid_column: true"));
1381    }
1382
1383    #[tokio::test]
1384    async fn tsid_fast_path_reuses_tsid_column_when_output_grows() {
1385        let schema = Arc::new(Schema::new(vec![
1386            Field::new(
1387                TIME_INDEX_COLUMN,
1388                DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
1389                false,
1390            ),
1391            Field::new("value", DataType::Float64, true),
1392            Field::new("host", DataType::Utf8, true),
1393            Field::new("__tsid", DataType::UInt64, false),
1394        ]));
1395        let batch = RecordBatch::try_new(
1396            schema.clone(),
1397            vec![
1398                Arc::new(TimestampMillisecondArray::from(vec![0, 1_000])),
1399                Arc::new(Float64Array::from(vec![1.0, 2.0])),
1400                Arc::new(datafusion::arrow::array::StringArray::from(vec![
1401                    "foo", "foo",
1402                ])),
1403                Arc::new(UInt64Array::from(vec![42, 42])),
1404            ],
1405        )
1406        .unwrap();
1407        let input = Arc::new(DataSourceExec::new(Arc::new(
1408            MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(),
1409        )));
1410        let normalize_exec = Arc::new(InstantManipulateExec {
1411            offset: 0,
1412            start: 0,
1413            end: 1_500,
1414            lookback_delta: 1_000,
1415            interval: 500,
1416            time_index_column: TIME_INDEX_COLUMN.to_string(),
1417            field_column: Some("value".to_string()),
1418            reuse_tsid_column: true,
1419            output_schema: input.schema(),
1420            properties: input.properties().clone(),
1421            input,
1422            metric: ExecutionPlanMetricsSet::new(),
1423        });
1424        let session_context = SessionContext::default();
1425        let result = datafusion::physical_plan::collect(normalize_exec, session_context.task_ctx())
1426            .await
1427            .unwrap();
1428        let result_literal = datatypes::arrow::util::pretty::pretty_format_batches(&result)
1429            .unwrap()
1430            .to_string();
1431
1432        assert_eq!(
1433            result_literal,
1434            "+-------------------------+-------+------+--------+\
1435            \n| timestamp               | value | host | __tsid |\
1436            \n+-------------------------+-------+------+--------+\
1437            \n| 1970-01-01T00:00:00     | 1.0   | foo  | 42     |\
1438            \n| 1970-01-01T00:00:00.500 | 1.0   | foo  | 42     |\
1439            \n| 1970-01-01T00:00:01     | 2.0   | foo  | 42     |\
1440            \n| 1970-01-01T00:00:01.500 | 2.0   | foo  | 42     |\
1441            \n+-------------------------+-------+------+--------+"
1442        );
1443    }
1444
1445    #[tokio::test]
1446    async fn tsid_fast_path_still_takes_additional_field_columns() {
1447        let schema = Arc::new(Schema::new(vec![
1448            Field::new(
1449                TIME_INDEX_COLUMN,
1450                DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
1451                false,
1452            ),
1453            Field::new("value", DataType::Float64, true),
1454            Field::new("value_2", DataType::Float64, true),
1455            Field::new("host", DataType::Utf8, true),
1456            Field::new("__tsid", DataType::UInt64, false),
1457        ]));
1458        let batch = RecordBatch::try_new(
1459            schema.clone(),
1460            vec![
1461                Arc::new(TimestampMillisecondArray::from(vec![0, 1_000])),
1462                Arc::new(Float64Array::from(vec![1.0, 2.0])),
1463                Arc::new(Float64Array::from(vec![10.0, 20.0])),
1464                Arc::new(datafusion::arrow::array::StringArray::from(vec![
1465                    "foo", "foo",
1466                ])),
1467                Arc::new(UInt64Array::from(vec![42, 42])),
1468            ],
1469        )
1470        .unwrap();
1471        let input = Arc::new(DataSourceExec::new(Arc::new(
1472            MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(),
1473        )));
1474        let normalize_exec = Arc::new(InstantManipulateExec {
1475            offset: 0,
1476            start: 0,
1477            end: 1_500,
1478            lookback_delta: 1_000,
1479            interval: 500,
1480            time_index_column: TIME_INDEX_COLUMN.to_string(),
1481            field_column: Some("value".to_string()),
1482            reuse_tsid_column: true,
1483            output_schema: input.schema(),
1484            properties: input.properties().clone(),
1485            input,
1486            metric: ExecutionPlanMetricsSet::new(),
1487        });
1488        let session_context = SessionContext::default();
1489        let result = datafusion::physical_plan::collect(normalize_exec, session_context.task_ctx())
1490            .await
1491            .unwrap();
1492        let result_literal = datatypes::arrow::util::pretty::pretty_format_batches(&result)
1493            .unwrap()
1494            .to_string();
1495
1496        assert_eq!(
1497            result_literal,
1498            "+-------------------------+-------+---------+------+--------+\
1499            \n| timestamp               | value | value_2 | host | __tsid |\
1500            \n+-------------------------+-------+---------+------+--------+\
1501            \n| 1970-01-01T00:00:00     | 1.0   | 10.0    | foo  | 42     |\
1502            \n| 1970-01-01T00:00:00.500 | 1.0   | 10.0    | foo  | 42     |\
1503            \n| 1970-01-01T00:00:01     | 2.0   | 20.0    | foo  | 42     |\
1504            \n| 1970-01-01T00:00:01.500 | 2.0   | 20.0    | foo  | 42     |\
1505            \n+-------------------------+-------+---------+------+--------+"
1506        );
1507    }
1508
1509    #[tokio::test]
1510    async fn manipulate_should_reject_too_many_output_points() {
1511        let schema = Arc::new(Schema::new(vec![
1512            Field::new(
1513                TIME_INDEX_COLUMN,
1514                DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
1515                false,
1516            ),
1517            Field::new("value", DataType::Float64, true),
1518        ]));
1519        let batch = RecordBatch::try_new(
1520            schema.clone(),
1521            vec![
1522                Arc::new(TimestampMillisecondArray::from(vec![0])),
1523                Arc::new(Float64Array::from(vec![1.0])),
1524            ],
1525        )
1526        .unwrap();
1527        let input = Arc::new(DataSourceExec::new(Arc::new(
1528            MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(),
1529        )));
1530        let too_many_points = MAX_INSTANT_MANIPULATE_OUTPUT_POINTS as Millisecond + 1;
1531        let normalize_exec = Arc::new(InstantManipulateExec {
1532            offset: 0,
1533            start: 0,
1534            end: too_many_points,
1535            lookback_delta: too_many_points + 1,
1536            interval: 1,
1537            time_index_column: TIME_INDEX_COLUMN.to_string(),
1538            field_column: Some("value".to_string()),
1539            reuse_tsid_column: false,
1540            output_schema: input.schema(),
1541            properties: input.properties().clone(),
1542            input,
1543            metric: ExecutionPlanMetricsSet::new(),
1544        });
1545        let session_context = SessionContext::default();
1546        let err = datafusion::physical_plan::collect(normalize_exec, session_context.task_ctx())
1547            .await
1548            .unwrap_err();
1549
1550        assert!(
1551            err.to_string()
1552                .contains("InstantManipulate output points exceed limit")
1553        );
1554    }
1555
1556    #[tokio::test]
1557    async fn lookback_10s_interval_30s() {
1558        let expected = String::from(
1559            "+---------------------+-------+------+\
1560            \n| timestamp           | value | path |\
1561            \n+---------------------+-------+------+\
1562            \n| 1970-01-01T00:00:00 | 1.0   | foo  |\
1563            \n| 1970-01-01T00:00:30 | 1.0   | foo  |\
1564            \n| 1970-01-01T00:01:00 | 1.0   | foo  |\
1565            \n| 1970-01-01T00:01:30 | 1.0   | foo  |\
1566            \n| 1970-01-01T00:02:00 | 1.0   | foo  |\
1567            \n| 1970-01-01T00:03:00 | 1.0   | foo  |\
1568            \n| 1970-01-01T00:04:00 | 1.0   | foo  |\
1569            \n| 1970-01-01T00:05:00 | 1.0   | foo  |\
1570            \n+---------------------+-------+------+",
1571        );
1572        do_normalize_test(0, 310_000, 10_000, 30_000, expected, false).await;
1573    }
1574
1575    #[tokio::test]
1576    async fn lookback_10s_interval_10s() {
1577        let expected = String::from(
1578            "+---------------------+-------+------+\
1579            \n| timestamp           | value | path |\
1580            \n+---------------------+-------+------+\
1581            \n| 1970-01-01T00:00:00 | 1.0   | foo  |\
1582            \n| 1970-01-01T00:00:30 | 1.0   | foo  |\
1583            \n| 1970-01-01T00:01:00 | 1.0   | foo  |\
1584            \n| 1970-01-01T00:01:30 | 1.0   | foo  |\
1585            \n| 1970-01-01T00:02:00 | 1.0   | foo  |\
1586            \n| 1970-01-01T00:03:00 | 1.0   | foo  |\
1587            \n| 1970-01-01T00:04:00 | 1.0   | foo  |\
1588            \n| 1970-01-01T00:04:10 | 1.0   | foo  |\
1589            \n| 1970-01-01T00:04:40 | 1.0   | foo  |\
1590            \n| 1970-01-01T00:05:00 | 1.0   | foo  |\
1591            \n+---------------------+-------+------+",
1592        );
1593        do_normalize_test(0, 300_000, 10_000, 10_000, expected, false).await;
1594    }
1595
1596    #[tokio::test]
1597    async fn lookback_30s_interval_30s() {
1598        let expected = String::from(
1599            "+---------------------+-------+------+\
1600            \n| timestamp           | value | path |\
1601            \n+---------------------+-------+------+\
1602            \n| 1970-01-01T00:00:00 | 1.0   | foo  |\
1603            \n| 1970-01-01T00:00:30 | 1.0   | foo  |\
1604            \n| 1970-01-01T00:01:00 | 1.0   | foo  |\
1605            \n| 1970-01-01T00:01:30 | 1.0   | foo  |\
1606            \n| 1970-01-01T00:02:00 | 1.0   | foo  |\
1607            \n| 1970-01-01T00:03:00 | 1.0   | foo  |\
1608            \n| 1970-01-01T00:04:00 | 1.0   | foo  |\
1609            \n| 1970-01-01T00:04:30 | 1.0   | foo  |\
1610            \n| 1970-01-01T00:05:00 | 1.0   | foo  |\
1611            \n+---------------------+-------+------+",
1612        );
1613        do_normalize_test(0, 300_000, 30_000, 30_000, expected, false).await;
1614    }
1615
1616    #[tokio::test]
1617    async fn lookback_30s_interval_10s() {
1618        let expected = String::from(
1619            "+---------------------+-------+------+\
1620            \n| timestamp           | value | path |\
1621            \n+---------------------+-------+------+\
1622            \n| 1970-01-01T00:00:00 | 1.0   | foo  |\
1623            \n| 1970-01-01T00:00:10 | 1.0   | foo  |\
1624            \n| 1970-01-01T00:00:20 | 1.0   | foo  |\
1625            \n| 1970-01-01T00:00:30 | 1.0   | foo  |\
1626            \n| 1970-01-01T00:00:40 | 1.0   | foo  |\
1627            \n| 1970-01-01T00:00:50 | 1.0   | foo  |\
1628            \n| 1970-01-01T00:01:00 | 1.0   | foo  |\
1629            \n| 1970-01-01T00:01:10 | 1.0   | foo  |\
1630            \n| 1970-01-01T00:01:20 | 1.0   | foo  |\
1631            \n| 1970-01-01T00:01:30 | 1.0   | foo  |\
1632            \n| 1970-01-01T00:01:40 | 1.0   | foo  |\
1633            \n| 1970-01-01T00:01:50 | 1.0   | foo  |\
1634            \n| 1970-01-01T00:02:00 | 1.0   | foo  |\
1635            \n| 1970-01-01T00:02:10 | 1.0   | foo  |\
1636            \n| 1970-01-01T00:02:20 | 1.0   | foo  |\
1637            \n| 1970-01-01T00:03:00 | 1.0   | foo  |\
1638            \n| 1970-01-01T00:03:10 | 1.0   | foo  |\
1639            \n| 1970-01-01T00:03:20 | 1.0   | foo  |\
1640            \n| 1970-01-01T00:04:00 | 1.0   | foo  |\
1641            \n| 1970-01-01T00:04:10 | 1.0   | foo  |\
1642            \n| 1970-01-01T00:04:20 | 1.0   | foo  |\
1643            \n| 1970-01-01T00:04:30 | 1.0   | foo  |\
1644            \n| 1970-01-01T00:04:40 | 1.0   | foo  |\
1645            \n| 1970-01-01T00:04:50 | 1.0   | foo  |\
1646            \n| 1970-01-01T00:05:00 | 1.0   | foo  |\
1647            \n+---------------------+-------+------+",
1648        );
1649        do_normalize_test(0, 300_000, 30_000, 10_000, expected, false).await;
1650    }
1651
1652    #[tokio::test]
1653    async fn lookback_60s_interval_10s() {
1654        let expected = String::from(
1655            "+---------------------+-------+------+\
1656            \n| timestamp           | value | path |\
1657            \n+---------------------+-------+------+\
1658            \n| 1970-01-01T00:00:00 | 1.0   | foo  |\
1659            \n| 1970-01-01T00:00:10 | 1.0   | foo  |\
1660            \n| 1970-01-01T00:00:20 | 1.0   | foo  |\
1661            \n| 1970-01-01T00:00:30 | 1.0   | foo  |\
1662            \n| 1970-01-01T00:00:40 | 1.0   | foo  |\
1663            \n| 1970-01-01T00:00:50 | 1.0   | foo  |\
1664            \n| 1970-01-01T00:01:00 | 1.0   | foo  |\
1665            \n| 1970-01-01T00:01:10 | 1.0   | foo  |\
1666            \n| 1970-01-01T00:01:20 | 1.0   | foo  |\
1667            \n| 1970-01-01T00:01:30 | 1.0   | foo  |\
1668            \n| 1970-01-01T00:01:40 | 1.0   | foo  |\
1669            \n| 1970-01-01T00:01:50 | 1.0   | foo  |\
1670            \n| 1970-01-01T00:02:00 | 1.0   | foo  |\
1671            \n| 1970-01-01T00:02:10 | 1.0   | foo  |\
1672            \n| 1970-01-01T00:02:20 | 1.0   | foo  |\
1673            \n| 1970-01-01T00:02:30 | 1.0   | foo  |\
1674            \n| 1970-01-01T00:02:40 | 1.0   | foo  |\
1675            \n| 1970-01-01T00:02:50 | 1.0   | foo  |\
1676            \n| 1970-01-01T00:03:00 | 1.0   | foo  |\
1677            \n| 1970-01-01T00:03:10 | 1.0   | foo  |\
1678            \n| 1970-01-01T00:03:20 | 1.0   | foo  |\
1679            \n| 1970-01-01T00:03:30 | 1.0   | foo  |\
1680            \n| 1970-01-01T00:03:40 | 1.0   | foo  |\
1681            \n| 1970-01-01T00:03:50 | 1.0   | foo  |\
1682            \n| 1970-01-01T00:04:00 | 1.0   | foo  |\
1683            \n| 1970-01-01T00:04:10 | 1.0   | foo  |\
1684            \n| 1970-01-01T00:04:20 | 1.0   | foo  |\
1685            \n| 1970-01-01T00:04:30 | 1.0   | foo  |\
1686            \n| 1970-01-01T00:04:40 | 1.0   | foo  |\
1687            \n| 1970-01-01T00:04:50 | 1.0   | foo  |\
1688            \n| 1970-01-01T00:05:00 | 1.0   | foo  |\
1689            \n+---------------------+-------+------+",
1690        );
1691        do_normalize_test(0, 300_000, 60_000, 10_000, expected, false).await;
1692    }
1693
1694    #[tokio::test]
1695    async fn lookback_60s_interval_30s() {
1696        let expected = String::from(
1697            "+---------------------+-------+------+\
1698            \n| timestamp           | value | path |\
1699            \n+---------------------+-------+------+\
1700            \n| 1970-01-01T00:00:00 | 1.0   | foo  |\
1701            \n| 1970-01-01T00:00:30 | 1.0   | foo  |\
1702            \n| 1970-01-01T00:01:00 | 1.0   | foo  |\
1703            \n| 1970-01-01T00:01:30 | 1.0   | foo  |\
1704            \n| 1970-01-01T00:02:00 | 1.0   | foo  |\
1705            \n| 1970-01-01T00:02:30 | 1.0   | foo  |\
1706            \n| 1970-01-01T00:03:00 | 1.0   | foo  |\
1707            \n| 1970-01-01T00:03:30 | 1.0   | foo  |\
1708            \n| 1970-01-01T00:04:00 | 1.0   | foo  |\
1709            \n| 1970-01-01T00:04:30 | 1.0   | foo  |\
1710            \n| 1970-01-01T00:05:00 | 1.0   | foo  |\
1711            \n+---------------------+-------+------+",
1712        );
1713        do_normalize_test(0, 300_000, 60_000, 30_000, expected, false).await;
1714    }
1715
1716    #[tokio::test]
1717    async fn small_range_lookback_0s_interval_1s() {
1718        let expected = String::from(
1719            "+---------------------+-------+------+\
1720            \n| timestamp           | value | path |\
1721            \n+---------------------+-------+------+\
1722            \n| 1970-01-01T00:04:00 | 1.0   | foo  |\
1723            \n| 1970-01-01T00:04:01 | 1.0   | foo  |\
1724            \n+---------------------+-------+------+",
1725        );
1726        do_normalize_test(230_000, 245_000, 0, 1_000, expected, false).await;
1727    }
1728
1729    #[tokio::test]
1730    async fn small_range_lookback_10s_interval_10s() {
1731        let expected = String::from(
1732            "+---------------------+-------+------+\
1733            \n| timestamp           | value | path |\
1734            \n+---------------------+-------+------+\
1735            \n| 1970-01-01T00:00:00 | 1.0   | foo  |\
1736            \n| 1970-01-01T00:00:30 | 1.0   | foo  |\
1737            \n+---------------------+-------+------+",
1738        );
1739        do_normalize_test(0, 30_000, 10_000, 10_000, expected, false).await;
1740    }
1741
1742    #[tokio::test]
1743    async fn large_range_lookback_30s_interval_60s() {
1744        let expected = String::from(
1745            "+---------------------+-------+------+\
1746            \n| timestamp           | value | path |\
1747            \n+---------------------+-------+------+\
1748            \n| 1970-01-01T00:00:00 | 1.0   | foo  |\
1749            \n| 1970-01-01T00:01:00 | 1.0   | foo  |\
1750            \n| 1970-01-01T00:02:00 | 1.0   | foo  |\
1751            \n| 1970-01-01T00:03:00 | 1.0   | foo  |\
1752            \n| 1970-01-01T00:04:00 | 1.0   | foo  |\
1753            \n| 1970-01-01T00:05:00 | 1.0   | foo  |\
1754            \n+---------------------+-------+------+",
1755        );
1756        do_normalize_test(-900_000, 900_000, 30_000, 60_000, expected, false).await;
1757    }
1758
1759    #[tokio::test]
1760    async fn small_range_lookback_30s_interval_30s() {
1761        let expected = String::from(
1762            "+---------------------+-------+------+\
1763            \n| timestamp           | value | path |\
1764            \n+---------------------+-------+------+\
1765            \n| 1970-01-01T00:03:10 | 1.0   | foo  |\
1766            \n| 1970-01-01T00:03:20 | 1.0   | foo  |\
1767            \n| 1970-01-01T00:04:00 | 1.0   | foo  |\
1768            \n| 1970-01-01T00:04:10 | 1.0   | foo  |\
1769            \n| 1970-01-01T00:04:20 | 1.0   | foo  |\
1770            \n| 1970-01-01T00:04:30 | 1.0   | foo  |\
1771            \n| 1970-01-01T00:04:40 | 1.0   | foo  |\
1772            \n| 1970-01-01T00:04:50 | 1.0   | foo  |\
1773            \n| 1970-01-01T00:05:00 | 1.0   | foo  |\
1774            \n+---------------------+-------+------+",
1775        );
1776        do_normalize_test(190_000, 300_000, 30_000, 10_000, expected, false).await;
1777    }
1778
1779    #[tokio::test]
1780    async fn lookback_10s_interval_10s_with_stale_marker() {
1781        let expected = String::from(
1782            "+---------------------+-------+\
1783            \n| timestamp           | value |\
1784            \n+---------------------+-------+\
1785            \n| 1970-01-01T00:00:00 | 0.0   |\
1786            \n| 1970-01-01T00:01:00 | 6.0   |\
1787            \n| 1970-01-01T00:02:00 | 12.0  |\
1788            \n+---------------------+-------+",
1789        );
1790        do_normalize_test(0, 300_000, 10_000, 10_000, expected, true).await;
1791    }
1792
1793    #[tokio::test]
1794    async fn lookback_10s_interval_10s_with_stale_marker_unaligned() {
1795        let expected = String::from(
1796            "+-------------------------+-------+\
1797            \n| timestamp               | value |\
1798            \n+-------------------------+-------+\
1799            \n| 1970-01-01T00:00:00.001 | 0.0   |\
1800            \n| 1970-01-01T00:01:00.001 | 6.0   |\
1801            \n| 1970-01-01T00:02:00.001 | 12.0  |\
1802            \n+-------------------------+-------+",
1803        );
1804        do_normalize_test(1, 300_001, 10_000, 10_000, expected, true).await;
1805    }
1806
1807    #[tokio::test]
1808    async fn ultra_large_range() {
1809        let expected = String::from(
1810            "+-------------------------+-------+\
1811            \n| timestamp               | value |\
1812            \n+-------------------------+-------+\
1813            \n| 1970-01-01T00:00:00.001 | 0.0   |\
1814            \n| 1970-01-01T00:01:00.001 | 6.0   |\
1815            \n| 1970-01-01T00:02:00.001 | 12.0  |\
1816            \n+-------------------------+-------+",
1817        );
1818        do_normalize_test(
1819            -900_000_000_000_000 + 1,
1820            900_000_000_000_000,
1821            10_000,
1822            10_000,
1823            expected,
1824            true,
1825        )
1826        .await;
1827    }
1828
1829    #[test]
1830    fn nullable_interior_timestamp_matches_wide_scan() {
1831        let schema = Arc::new(Schema::new(vec![
1832            Field::new(
1833                TIME_INDEX_COLUMN,
1834                DataType::Timestamp(TimeUnit::Second, None),
1835                true,
1836            ),
1837            Field::new("value", DataType::Float64, false),
1838        ]));
1839        let timestamps = [0, i64::MAX, 1];
1840        let array = TimestampSecondArray::new(
1841            timestamps.to_vec().into(),
1842            Some(NullBuffer::from(vec![false, false, true])),
1843        );
1844        let input = RecordBatch::try_new(
1845            schema.clone(),
1846            vec![
1847                Arc::new(array),
1848                Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])),
1849            ],
1850        )
1851        .unwrap();
1852        let stream = InstantManipulateStream {
1853            offset: 0,
1854            start: 0,
1855            end: 1_000,
1856            lookback_delta: 1_000,
1857            interval: 1_000,
1858            time_index: 0,
1859            time_unit: TimeUnit::Second,
1860            field_indices: [Some(1), None],
1861            tsid_index: None,
1862            reuse_tsid_column: false,
1863            schema: Arc::new(Schema::new(vec![
1864                Field::new(
1865                    TIME_INDEX_COLUMN,
1866                    DataType::Timestamp(TimeUnit::Millisecond, None),
1867                    false,
1868                ),
1869                Field::new("value", DataType::Float64, false),
1870            ])),
1871            input: Box::pin(
1872                datafusion::physical_plan::memory::MemoryStream::try_new(vec![], schema, None)
1873                    .unwrap(),
1874            ),
1875            metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
1876            num_series: Count::new(),
1877        };
1878
1879        let output = stream.manipulate(input).unwrap();
1880        let output_timestamps = output
1881            .column(0)
1882            .as_any()
1883            .downcast_ref::<TimestampMillisecondArray>()
1884            .unwrap();
1885        let output_values = output
1886            .column(1)
1887            .as_any()
1888            .downcast_ref::<Float64Array>()
1889            .unwrap();
1890        let wide = stream.scan_instants::<i128>(
1891            &timestamps,
1892            0,
1893            1_000,
1894            &ScanNanoseconds::<i128>::wide(TimeUnit::Second, 0, 1_000),
1895            2,
1896            |_| false,
1897        );
1898        assert_eq!(output_timestamps.values(), &[0]);
1899        assert_eq!(output_values.values(), &[1.0]);
1900        assert_eq!(wide, (vec![0], vec![0]));
1901        assert_eq!(wide.1, output_timestamps.values().to_vec());
1902    }
1903
1904    #[test]
1905    fn exact_ties_select_first_and_lookback_uses_latest() {
1906        for (values, expected_timestamp) in [
1907            (vec![42.0, f64::from_bits(PROMETHEUS_STALE_NAN_BITS)], 1_000),
1908            (vec![f64::from_bits(PROMETHEUS_STALE_NAN_BITS), 42.0], 1_050),
1909        ] {
1910            let schema = Arc::new(Schema::new(vec![
1911                Field::new(
1912                    TIME_INDEX_COLUMN,
1913                    DataType::Timestamp(TimeUnit::Millisecond, None),
1914                    false,
1915                ),
1916                Field::new("value", DataType::Float64, true),
1917            ]));
1918            let input = RecordBatch::try_new(
1919                schema.clone(),
1920                vec![
1921                    Arc::new(TimestampMillisecondArray::from(vec![1_000, 1_000])),
1922                    Arc::new(Float64Array::from(values)),
1923                ],
1924            )
1925            .unwrap();
1926            let stream = InstantManipulateStream {
1927                offset: 0,
1928                start: 1_000,
1929                end: 1_050,
1930                lookback_delta: 100,
1931                interval: 50,
1932                time_index: 0,
1933                time_unit: TimeUnit::Millisecond,
1934                field_indices: [Some(1), None],
1935                tsid_index: None,
1936                reuse_tsid_column: false,
1937                schema: schema.clone(),
1938                input: Box::pin(
1939                    datafusion::physical_plan::memory::MemoryStream::try_new(vec![], schema, None)
1940                        .unwrap(),
1941                ),
1942                metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
1943                num_series: Count::new(),
1944            };
1945
1946            let output = stream.manipulate(input).unwrap();
1947            let timestamps = output
1948                .column(0)
1949                .as_any()
1950                .downcast_ref::<TimestampMillisecondArray>()
1951                .unwrap();
1952            let values = output
1953                .column(1)
1954                .as_any()
1955                .downcast_ref::<Float64Array>()
1956                .unwrap();
1957            assert_eq!(timestamps.values(), &[expected_timestamp]);
1958            assert_eq!(values.values(), &[42.0]);
1959        }
1960    }
1961
1962    #[test]
1963    fn empty_batches_preserve_stream_progression_and_output_schema() {
1964        use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
1965        use futures::stream;
1966        use futures::task::noop_waker_ref;
1967
1968        let input_schema = Arc::new(Schema::new(vec![Field::new(
1969            TIME_INDEX_COLUMN,
1970            DataType::Timestamp(TimeUnit::Second, None),
1971            false,
1972        )]));
1973        let output_schema = Arc::new(Schema::new(vec![Field::new(
1974            TIME_INDEX_COLUMN,
1975            DataType::Timestamp(TimeUnit::Millisecond, None),
1976            false,
1977        )]));
1978        let empty = RecordBatch::new_empty(input_schema.clone());
1979        let valid = RecordBatch::try_new(
1980            input_schema.clone(),
1981            vec![Arc::new(TimestampSecondArray::from(vec![1]))],
1982        )
1983        .unwrap();
1984        let input = RecordBatchStreamAdapter::new(
1985            input_schema,
1986            stream::iter(vec![
1987                Ok(empty.clone()),
1988                Ok(empty.clone()),
1989                Ok(valid),
1990                Ok(empty),
1991                Err(DataFusionError::Execution("injected input error".into())),
1992            ]),
1993        );
1994        let mut stream = InstantManipulateStream {
1995            offset: 0,
1996            start: 1_000,
1997            end: 1_000,
1998            lookback_delta: 0,
1999            interval: 1,
2000            time_index: 0,
2001            time_unit: TimeUnit::Second,
2002            field_indices: [None, None],
2003            tsid_index: None,
2004            reuse_tsid_column: false,
2005            schema: output_schema.clone(),
2006            input: Box::pin(input),
2007            metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
2008            num_series: Count::new(),
2009        };
2010        let waker = noop_waker_ref();
2011        let mut cx = Context::from_waker(waker);
2012
2013        for _ in 0..2 {
2014            let Poll::Ready(Some(Ok(batch))) = Pin::new(&mut stream).poll_next(&mut cx) else {
2015                panic!("empty batch must be returned immediately");
2016            };
2017            assert_eq!(batch.num_rows(), 0);
2018            assert_eq!(batch.schema(), output_schema);
2019        }
2020
2021        let Poll::Ready(Some(Ok(valid))) = Pin::new(&mut stream).poll_next(&mut cx) else {
2022            panic!("valid batch must follow empty batches");
2023        };
2024        assert_eq!(valid.schema(), output_schema);
2025        assert_eq!(
2026            valid.column(0).data_type(),
2027            &DataType::Timestamp(TimeUnit::Millisecond, None)
2028        );
2029        assert_eq!(
2030            valid
2031                .column(0)
2032                .as_any()
2033                .downcast_ref::<TimestampMillisecondArray>()
2034                .unwrap()
2035                .values(),
2036            &[1_000]
2037        );
2038
2039        let Poll::Ready(Some(Ok(empty))) = Pin::new(&mut stream).poll_next(&mut cx) else {
2040            panic!("empty batch after valid batch must be returned immediately");
2041        };
2042        assert_eq!(empty.num_rows(), 0);
2043        assert_eq!(empty.schema(), output_schema);
2044
2045        let Poll::Ready(Some(Err(error))) = Pin::new(&mut stream).poll_next(&mut cx) else {
2046            panic!("input error must propagate");
2047        };
2048        assert!(error.to_string().contains("injected input error"));
2049        assert!(matches!(
2050            Pin::new(&mut stream).poll_next(&mut cx),
2051            Poll::Ready(None)
2052        ));
2053        assert_eq!(stream.num_series.value(), 1);
2054    }
2055
2056    #[test]
2057    fn extreme_alignment_retains_exact_sample() {
2058        let schema = Arc::new(Schema::new(vec![
2059            Field::new(
2060                TIME_INDEX_COLUMN,
2061                DataType::Timestamp(TimeUnit::Millisecond, None),
2062                false,
2063            ),
2064            Field::new("value", DataType::Float64, true),
2065        ]));
2066        let input = RecordBatch::try_new(
2067            schema.clone(),
2068            vec![
2069                Arc::new(TimestampMillisecondArray::from(vec![i64::MAX])),
2070                Arc::new(Float64Array::from(vec![7.0])),
2071            ],
2072        )
2073        .unwrap();
2074        let stream = InstantManipulateStream {
2075            offset: 0,
2076            start: i64::MIN + 1,
2077            end: i64::MAX,
2078            lookback_delta: 0,
2079            interval: i64::MAX,
2080            time_index: 0,
2081            time_unit: TimeUnit::Millisecond,
2082            field_indices: [Some(1), None],
2083            tsid_index: None,
2084            reuse_tsid_column: false,
2085            schema: schema.clone(),
2086            input: Box::pin(
2087                datafusion::physical_plan::memory::MemoryStream::try_new(vec![], schema, None)
2088                    .unwrap(),
2089            ),
2090            metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
2091            num_series: Count::new(),
2092        };
2093
2094        let output = stream.manipulate(input).unwrap();
2095        let timestamps = output
2096            .column(0)
2097            .as_any()
2098            .downcast_ref::<TimestampMillisecondArray>()
2099            .unwrap();
2100        let values = output
2101            .column(1)
2102            .as_any()
2103            .downcast_ref::<Float64Array>()
2104            .unwrap();
2105        assert_eq!(timestamps.values(), &[i64::MAX]);
2106        assert_eq!(values.values(), &[7.0]);
2107    }
2108
2109    #[test]
2110    fn native_nanosecond_offset_uses_wide_shifted_timeline() {
2111        for (raw, offset, eval) in [
2112            (
2113                9_223_112_837_000_000_000_i64,
2114                259_200_000,
2115                9_223_372_037_000,
2116            ),
2117            (
2118                -9_223_112_837_000_000_000_i64,
2119                -259_200_000,
2120                -9_223_372_037_000,
2121            ),
2122        ] {
2123            let schema = Arc::new(Schema::new(vec![
2124                Field::new(
2125                    TIME_INDEX_COLUMN,
2126                    DataType::Timestamp(TimeUnit::Nanosecond, None),
2127                    false,
2128                ),
2129                Field::new("value", DataType::Float64, true),
2130            ]));
2131            let input = RecordBatch::try_new(
2132                schema.clone(),
2133                vec![
2134                    Arc::new(TimestampNanosecondArray::from(vec![raw])),
2135                    Arc::new(Float64Array::from(vec![7.0])),
2136                ],
2137            )
2138            .unwrap();
2139            let stream = InstantManipulateStream {
2140                offset,
2141                start: eval,
2142                end: eval,
2143                lookback_delta: 300_000,
2144                interval: 1,
2145                time_index: 0,
2146                time_unit: TimeUnit::Nanosecond,
2147                field_indices: [Some(1), None],
2148                tsid_index: None,
2149                reuse_tsid_column: false,
2150                schema: Arc::new(Schema::new(vec![
2151                    Field::new(
2152                        TIME_INDEX_COLUMN,
2153                        DataType::Timestamp(TimeUnit::Millisecond, None),
2154                        false,
2155                    ),
2156                    Field::new("value", DataType::Float64, true),
2157                ])),
2158                input: Box::pin(
2159                    datafusion::physical_plan::memory::MemoryStream::try_new(vec![], schema, None)
2160                        .unwrap(),
2161                ),
2162                metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
2163                num_series: Count::new(),
2164            };
2165            let output = stream.manipulate(input).unwrap();
2166            assert_eq!(output.num_rows(), 1);
2167            assert_eq!(
2168                output
2169                    .column(0)
2170                    .as_any()
2171                    .downcast_ref::<TimestampMillisecondArray>()
2172                    .unwrap()
2173                    .value(0),
2174                eval
2175            );
2176            assert_eq!(
2177                output
2178                    .column(1)
2179                    .as_any()
2180                    .downcast_ref::<Float64Array>()
2181                    .unwrap()
2182                    .value(0),
2183                7.0
2184            );
2185        }
2186    }
2187
2188    #[tokio::test]
2189    async fn ordinary_nan_is_selected_for_exact_and_lookback() {
2190        let schema = Arc::new(Schema::new(vec![
2191            Field::new(
2192                TIME_INDEX_COLUMN,
2193                DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
2194                false,
2195            ),
2196            Field::new("value", DataType::Float64, true),
2197        ]));
2198        let batch = RecordBatch::try_new(
2199            schema.clone(),
2200            vec![
2201                Arc::new(TimestampMillisecondArray::from(vec![1_000])),
2202                Arc::new(Float64Array::from(vec![f64::NAN])),
2203            ],
2204        )
2205        .unwrap();
2206        let input = Arc::new(DataSourceExec::new(Arc::new(
2207            MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(),
2208        )));
2209        let exec = Arc::new(InstantManipulateExec {
2210            offset: 0,
2211            start: 1_000,
2212            end: 1_500,
2213            lookback_delta: 1_000,
2214            interval: 500,
2215            time_index_column: TIME_INDEX_COLUMN.to_string(),
2216            field_column: Some("value".to_string()),
2217            reuse_tsid_column: false,
2218            output_schema: input.schema(),
2219            properties: input.properties().clone(),
2220            input,
2221            metric: ExecutionPlanMetricsSet::new(),
2222        });
2223
2224        let context = SessionContext::default();
2225        let batches = datafusion::physical_plan::collect(exec, context.task_ctx())
2226            .await
2227            .unwrap();
2228        let values = batches
2229            .iter()
2230            .flat_map(|batch| {
2231                batch
2232                    .column(1)
2233                    .as_any()
2234                    .downcast_ref::<Float64Array>()
2235                    .unwrap()
2236                    .values()
2237                    .iter()
2238                    .copied()
2239            })
2240            .collect::<Vec<_>>();
2241        let timestamps = batches
2242            .iter()
2243            .flat_map(|batch| {
2244                batch
2245                    .column(0)
2246                    .as_any()
2247                    .downcast_ref::<TimestampMillisecondArray>()
2248                    .unwrap()
2249                    .values()
2250                    .iter()
2251                    .copied()
2252            })
2253            .collect::<Vec<_>>();
2254
2255        assert_eq!(values.len(), 2);
2256        assert_eq!(timestamps, vec![1_000, 1_500]);
2257        assert!(values.iter().all(|value| value.is_nan()));
2258        assert_eq!(
2259            values
2260                .iter()
2261                .map(|value| value.to_bits())
2262                .collect::<Vec<_>>(),
2263            vec![f64::NAN.to_bits(); 2]
2264        );
2265    }
2266
2267    #[tokio::test]
2268    async fn prometheus_stale_nan_selects_before_and_suppresses_after_marker() {
2269        let schema = Arc::new(Schema::new(vec![
2270            Field::new(
2271                TIME_INDEX_COLUMN,
2272                DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
2273                false,
2274            ),
2275            Field::new("value", DataType::Float64, true),
2276        ]));
2277        let batch = RecordBatch::try_new(
2278            schema.clone(),
2279            vec![
2280                Arc::new(TimestampMillisecondArray::from(vec![500, 1_000])),
2281                Arc::new(Float64Array::from(vec![
2282                    42.0,
2283                    f64::from_bits(0x7ff0_0000_0000_0002),
2284                ])),
2285            ],
2286        )
2287        .unwrap();
2288        let input = Arc::new(DataSourceExec::new(Arc::new(
2289            MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(),
2290        )));
2291        let exec = Arc::new(InstantManipulateExec {
2292            offset: 0,
2293            start: 750,
2294            end: 1_500,
2295            lookback_delta: 1_001,
2296            interval: 250,
2297            time_index_column: TIME_INDEX_COLUMN.to_string(),
2298            field_column: Some("value".to_string()),
2299            reuse_tsid_column: false,
2300            output_schema: input.schema(),
2301            properties: input.properties().clone(),
2302            input,
2303            metric: ExecutionPlanMetricsSet::new(),
2304        });
2305
2306        let context = SessionContext::default();
2307        let batches = datafusion::physical_plan::collect(exec, context.task_ctx())
2308            .await
2309            .unwrap();
2310
2311        let row_count = batches.iter().map(RecordBatch::num_rows).sum::<usize>();
2312        let batch = batches.iter().find(|batch| batch.num_rows() > 0).unwrap();
2313        let timestamp = batch
2314            .column(0)
2315            .as_any()
2316            .downcast_ref::<TimestampMillisecondArray>()
2317            .unwrap()
2318            .value(0);
2319        let value = batch
2320            .column(1)
2321            .as_any()
2322            .downcast_ref::<Float64Array>()
2323            .unwrap()
2324            .value(0);
2325        assert_eq!(
2326            (row_count, timestamp, value),
2327            (1, 750, 42.0),
2328            "only the evaluation before the stale marker should select 42.0"
2329        );
2330    }
2331
2332    #[tokio::test]
2333    async fn native_histogram_stale_nan_suppresses_exact_and_lookback() {
2334        let histograms = build_histogram_array(&[
2335            Some(native_histogram(42.0)),
2336            Some(native_histogram(f64::from_bits(PROMETHEUS_STALE_NAN_BITS))),
2337        ]);
2338        let schema = Arc::new(Schema::new(vec![
2339            Field::new(
2340                TIME_INDEX_COLUMN,
2341                DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
2342                false,
2343            ),
2344            Field::new("value", histograms.data_type().clone(), true),
2345        ]));
2346        let batch = RecordBatch::try_new(
2347            schema.clone(),
2348            vec![
2349                Arc::new(TimestampMillisecondArray::from(vec![500, 1_000])),
2350                histograms,
2351            ],
2352        )
2353        .unwrap();
2354        let input = Arc::new(DataSourceExec::new(Arc::new(
2355            MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(),
2356        )));
2357        let exec = Arc::new(InstantManipulateExec {
2358            offset: 0,
2359            start: 1_000,
2360            end: 1_500,
2361            lookback_delta: 1_001,
2362            interval: 500,
2363            time_index_column: TIME_INDEX_COLUMN.to_string(),
2364            field_column: Some("value".to_string()),
2365            reuse_tsid_column: false,
2366            output_schema: input.schema(),
2367            properties: input.properties().clone(),
2368            input,
2369            metric: ExecutionPlanMetricsSet::new(),
2370        });
2371
2372        let context = SessionContext::default();
2373        let batches = datafusion::physical_plan::collect(exec, context.task_ctx())
2374            .await
2375            .unwrap();
2376
2377        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
2378    }
2379
2380    #[tokio::test]
2381    async fn null_value_backed_by_stale_bits_is_selected_for_exact_and_lookback() {
2382        let schema = Arc::new(Schema::new(vec![
2383            Field::new(
2384                TIME_INDEX_COLUMN,
2385                DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
2386                false,
2387            ),
2388            Field::new("value", DataType::Float64, true),
2389        ]));
2390        let field_column = Float64Array::new(
2391            vec![f64::from_bits(0x7ff0_0000_0000_0002)].into(),
2392            Some(NullBuffer::from(vec![false])),
2393        );
2394        assert!(!field_column.is_valid(0));
2395        assert_eq!(field_column.value(0).to_bits(), 0x7ff0_0000_0000_0002);
2396        let batch = RecordBatch::try_new(
2397            schema.clone(),
2398            vec![
2399                Arc::new(TimestampMillisecondArray::from(vec![1_000])),
2400                Arc::new(field_column),
2401            ],
2402        )
2403        .unwrap();
2404        let input = Arc::new(DataSourceExec::new(Arc::new(
2405            MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(),
2406        )));
2407        let exec = Arc::new(InstantManipulateExec {
2408            offset: 0,
2409            start: 1_000,
2410            end: 1_500,
2411            lookback_delta: 1_000,
2412            interval: 500,
2413            time_index_column: TIME_INDEX_COLUMN.to_string(),
2414            field_column: Some("value".to_string()),
2415            reuse_tsid_column: false,
2416            output_schema: input.schema(),
2417            properties: input.properties().clone(),
2418            input,
2419            metric: ExecutionPlanMetricsSet::new(),
2420        });
2421
2422        let context = SessionContext::default();
2423        let batches = datafusion::physical_plan::collect(exec, context.task_ctx())
2424            .await
2425            .unwrap();
2426        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
2427        let batch = batches.iter().find(|batch| batch.num_rows() == 2).unwrap();
2428        let timestamps = batch
2429            .column(0)
2430            .as_any()
2431            .downcast_ref::<TimestampMillisecondArray>()
2432            .unwrap();
2433        let values = batch
2434            .column(1)
2435            .as_any()
2436            .downcast_ref::<Float64Array>()
2437            .unwrap();
2438
2439        assert_eq!(timestamps.values(), &[1_000, 1_500]);
2440        assert!(!values.is_valid(0));
2441        assert!(!values.is_valid(1));
2442    }
2443}