promql/extension_plan/

range_manipulate.rs

// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashSet;
use std::ops::{Add, Mul, Sub};
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use common_telemetry::{debug, warn};
use datafusion::arrow::array::{Array, ArrayRef, Int64Array, TimestampMillisecondArray};
use datafusion::arrow::compute;
use datafusion::arrow::datatypes::{DataType, Field, Int64Type, SchemaRef, TimeUnit};
use datafusion::arrow::error::ArrowError;
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::common::stats::Precision;
use datafusion::common::tree_node::TreeNodeRecursion;
use datafusion::common::{DFSchema, DFSchemaRef, TableReference};
use datafusion::error::{DataFusionError, Result as DataFusionResult};
use datafusion::execution::context::TaskContext;
use datafusion::logical_expr::{EmptyRelation, Expr, LogicalPlan, UserDefinedLogicalNodeCore};
use datafusion::physical_expr::EquivalenceProperties;
use datafusion::physical_plan::metrics::{
    BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricValue, MetricsSet,
};
use datafusion::physical_plan::{
    ChildStats, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan,
    InputDistributionRequirements, PhysicalExpr, PlanProperties, RecordBatchStream,
    SendableRecordBatchStream, Statistics, StatisticsArgs,
};
use datafusion_expr::ident;
use datatypes::timestamp::timestamp_array_to_primitive;
use futures::{Stream, StreamExt, ready};
use greptime_proto::substrait_extension as pb;
use prost::Message;
use snafu::ResultExt;

use crate::error::{DeserializeSnafu, Result};
use crate::extension_plan::{
    METRIC_NUM_SERIES, Millisecond, local_offset, nanoseconds_per_native_tick, resolve_column_name,
    serialize_column_index, timestamp_unit,
};
use crate::metrics::PROMQL_SERIES_COUNT;
use crate::range_array::RangeArray;

/// Time series manipulator for range function.
///
/// This plan will "fold" time index and value columns into [RangeArray]s, and truncate
/// other columns to the same length with the "folded" [RangeArray] column.
///
/// To pass runtime information to the execution plan (or the range function), This plan
/// will add those extra columns:
/// - timestamp range with type [RangeArray], which is the folded timestamp column.
/// - end of current range with the same type as the timestamp column. (todo)
#[derive(Debug, PartialEq, Eq, Hash)]
pub struct RangeManipulate {
    start: Millisecond,
    end: Millisecond,
    interval: Millisecond,
    range: Millisecond,
    offset: Millisecond,
    time_index: String,
    field_columns: Vec<String>,
    input: LogicalPlan,
    output_schema: DFSchemaRef,
    unfix: Option<UnfixIndices>,
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct UnfixIndices {
    pub time_index_idx: u64,
    pub tag_column_indices: Vec<u64>,
}

impl RangeManipulate {
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        start: Millisecond,
        end: Millisecond,
        interval: Millisecond,
        offset: Millisecond,
        range: Millisecond,
        time_index: String,
        field_columns: Vec<String>,
        input: LogicalPlan,
    ) -> DataFusionResult<Self> {
        let output_schema =
            Self::calculate_output_schema(input.schema(), &time_index, &field_columns)?;
        Ok(Self {
            start,
            end,
            interval,
            range,
            offset,
            time_index,
            field_columns,
            input,
            output_schema,
            unfix: None,
        })
    }

    pub const fn name() -> &'static str {
        "RangeManipulate"
    }

    pub fn build_timestamp_range_name(time_index: &str) -> String {
        format!("{time_index}_range")
    }

    pub fn internal_range_end_col_name() -> String {
        "__internal_range_end".to_string()
    }

    fn range_timestamp_name(&self) -> String {
        Self::build_timestamp_range_name(&self.time_index)
    }

    fn calculate_output_schema(
        input_schema: &DFSchemaRef,
        time_index: &str,
        field_columns: &[String],
    ) -> DataFusionResult<DFSchemaRef> {
        let columns = input_schema.fields();
        let mut new_columns = Vec::with_capacity(columns.len() + 1);
        for i in 0..columns.len() {
            let x = input_schema.qualified_field(i);
            new_columns.push((x.0.cloned(), x.1.clone()));
        }

        // process time index column
        // the raw timestamp field is preserved. And a new timestamp_range field is appended to the last.
        let Some(ts_col_index) = input_schema.index_of_column_by_name(None, time_index) else {
            return Err(datafusion::common::field_not_found(
                None::<TableReference>,
                time_index,
                input_schema.as_ref(),
            ));
        };
        let ts_col_field = &columns[ts_col_index];
        let output_time_field = Arc::new(
            ts_col_field
                .as_ref()
                .clone()
                .with_data_type(DataType::Timestamp(TimeUnit::Millisecond, None)),
        );
        new_columns[ts_col_index] = (
            input_schema.qualified_field(ts_col_index).0.cloned(),
            output_time_field.clone(),
        );
        let timestamp_range_field = Field::new(
            Self::build_timestamp_range_name(time_index),
            RangeArray::convert_field(output_time_field.as_ref())
                .data_type()
                .clone(),
            ts_col_field.is_nullable(),
        );
        new_columns.push((None, Arc::new(timestamp_range_field)));

        // process value columns
        for name in field_columns {
            let Some(index) = input_schema.index_of_column_by_name(None, name) else {
                return Err(datafusion::common::field_not_found(
                    None::<TableReference>,
                    name,
                    input_schema.as_ref(),
                ));
            };
            new_columns[index] = (None, Arc::new(RangeArray::convert_field(&columns[index])));
        }

        Ok(Arc::new(DFSchema::new_with_metadata(
            new_columns,
            input_schema.metadata().clone(),
        )?))
    }

    pub fn to_execution_plan(&self, exec_input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
        let output_schema: SchemaRef = self.output_schema.inner().clone();
        let properties = exec_input.properties();
        let properties = Arc::new(PlanProperties::new(
            EquivalenceProperties::new(output_schema.clone()),
            properties.partitioning.clone(),
            properties.emission_type,
            properties.boundedness,
        ));
        Arc::new(RangeManipulateExec {
            offset: self.offset,
            start: self.start,
            end: self.end,
            interval: self.interval,
            range: self.range,
            time_index_column: self.time_index.clone(),
            time_range_column: self.range_timestamp_name(),
            field_columns: self.field_columns.clone(),
            input: exec_input,
            output_schema,
            metric: ExecutionPlanMetricsSet::new(),
            properties,
        })
    }

    pub fn serialize(&self) -> Vec<u8> {
        let time_index_idx = serialize_column_index(self.input.schema(), &self.time_index);

        let tag_column_indices = self
            .field_columns
            .iter()
            .map(|name| serialize_column_index(self.input.schema(), name))
            .collect::<Vec<u64>>();

        pb::RangeManipulate {
            start: self.start,
            end: self.end,
            interval: self.interval,
            range: self.range,
            time_index_idx,
            tag_column_indices,
            ..Default::default()
        }
        .encode_to_vec()
    }

    pub fn deserialize(bytes: &[u8]) -> Result<Self> {
        let pb_range_manipulate = pb::RangeManipulate::decode(bytes).context(DeserializeSnafu)?;
        let empty_schema = Arc::new(DFSchema::empty());
        let placeholder_plan = LogicalPlan::EmptyRelation(EmptyRelation {
            produce_one_row: false,
            schema: empty_schema.clone(),
        });

        let unfix = UnfixIndices {
            time_index_idx: pb_range_manipulate.time_index_idx,
            tag_column_indices: pb_range_manipulate.tag_column_indices.clone(),
        };
        debug!("RangeManipulate deserialize unfix: {:?}", unfix);

        // Unlike `Self::new()`, this method doesn't check the input schema as it will fail
        // because the input schema is empty.
        // But this is Ok since datafusion guarantees to call `with_exprs_and_inputs` for the
        // deserialized plan.
        Ok(Self {
            start: pb_range_manipulate.start,
            end: pb_range_manipulate.end,
            interval: pb_range_manipulate.interval,
            range: pb_range_manipulate.range,
            offset: 0,
            time_index: String::new(),
            field_columns: Vec::new(),
            input: placeholder_plan,
            output_schema: empty_schema,
            unfix: Some(unfix),
        })
    }
}

impl PartialOrd for RangeManipulate {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        // Compare fields in order excluding output_schema
        match self.start.partial_cmp(&other.start) {
            Some(core::cmp::Ordering::Equal) => {}
            ord => return ord,
        }
        match self.end.partial_cmp(&other.end) {
            Some(core::cmp::Ordering::Equal) => {}
            ord => return ord,
        }
        match self.interval.partial_cmp(&other.interval) {
            Some(core::cmp::Ordering::Equal) => {}
            ord => return ord,
        }
        match self.range.partial_cmp(&other.range) {
            Some(core::cmp::Ordering::Equal) => {}
            ord => return ord,
        }
        match self.offset.partial_cmp(&other.offset) {
            Some(core::cmp::Ordering::Equal) => {}
            ord => return ord,
        }
        match self.time_index.partial_cmp(&other.time_index) {
            Some(core::cmp::Ordering::Equal) => {}
            ord => return ord,
        }
        match self.field_columns.partial_cmp(&other.field_columns) {
            Some(core::cmp::Ordering::Equal) => {}
            ord => return ord,
        }
        self.input.partial_cmp(&other.input)
    }
}

impl UserDefinedLogicalNodeCore for RangeManipulate {
    fn name(&self) -> &str {
        Self::name()
    }

    fn inputs(&self) -> Vec<&LogicalPlan> {
        vec![&self.input]
    }

    fn schema(&self) -> &DFSchemaRef {
        &self.output_schema
    }

    fn expressions(&self) -> Vec<Expr> {
        if self.unfix.is_some() {
            return vec![];
        }

        let mut exprs = Vec::with_capacity(1 + self.field_columns.len());
        exprs.push(ident(&self.time_index));
        exprs.extend(self.field_columns.iter().map(ident));
        exprs
    }

    fn necessary_children_exprs(&self, output_columns: &[usize]) -> Option<Vec<Vec<usize>>> {
        if self.unfix.is_some() {
            return None;
        }

        let input_schema = self.input.schema();
        let input_len = input_schema.fields().len();
        let time_index_idx = input_schema.index_of_column_by_name(None, &self.time_index)?;

        if output_columns.is_empty() {
            let indices = (0..input_len).collect::<Vec<_>>();
            return Some(vec![indices]);
        }

        let mut required = Vec::with_capacity(output_columns.len() + 1 + self.field_columns.len());
        required.push(time_index_idx);
        for value_column in &self.field_columns {
            required.push(input_schema.index_of_column_by_name(None, value_column)?);
        }
        for &idx in output_columns {
            if idx < input_len {
                required.push(idx);
            } else if idx == input_len {
                // Derived timestamp range column.
                required.push(time_index_idx);
            } else {
                warn!(
                    "Output column index {} is out of bounds for input schema with length {}",
                    idx, input_len
                );
                return None;
            }
        }

        required.sort_unstable();
        required.dedup();
        Some(vec![required])
    }

    fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(
            f,
            "PromRangeManipulate: req range=[{}..{}], interval=[{}], eval range=[{}], time index=[{}], values={:?}",
            self.start, self.end, self.interval, self.range, self.time_index, self.field_columns
        )
    }

    fn with_exprs_and_inputs(
        &self,
        _exprs: Vec<Expr>,
        mut inputs: Vec<LogicalPlan>,
    ) -> DataFusionResult<Self> {
        if inputs.len() != 1 {
            return Err(DataFusionError::Internal(
                "RangeManipulate should have at exact one input".to_string(),
            ));
        }

        let input: LogicalPlan = inputs.pop().unwrap();
        let input_schema = input.schema();

        if let Some(unfix) = &self.unfix {
            // transform indices to names
            let time_index = resolve_column_name(
                unfix.time_index_idx,
                input_schema,
                "RangeManipulate",
                "time index",
            )?;

            let field_columns = unfix
                .tag_column_indices
                .iter()
                .map(|idx| resolve_column_name(*idx, input_schema, "RangeManipulate", "tag"))
                .collect::<DataFusionResult<Vec<String>>>()?;

            let output_schema =
                Self::calculate_output_schema(input_schema, &time_index, &field_columns)?;

            Ok(Self {
                start: self.start,
                end: self.end,
                interval: self.interval,
                range: self.range,
                offset: local_offset(&input, &time_index),
                time_index,
                field_columns,
                input,
                output_schema,
                unfix: None,
            })
        } else {
            let output_schema =
                Self::calculate_output_schema(input_schema, &self.time_index, &self.field_columns)?;

            Ok(Self {
                start: self.start,
                end: self.end,
                interval: self.interval,
                range: self.range,
                offset: self.offset,
                time_index: self.time_index.clone(),
                field_columns: self.field_columns.clone(),
                input,
                output_schema,
                unfix: None,
            })
        }
    }
}

#[derive(Debug)]
pub struct RangeManipulateExec {
    offset: Millisecond,
    start: Millisecond,
    end: Millisecond,
    interval: Millisecond,
    range: Millisecond,
    time_index_column: String,
    time_range_column: String,
    field_columns: Vec<String>,

    input: Arc<dyn ExecutionPlan>,
    output_schema: SchemaRef,
    metric: ExecutionPlanMetricsSet,
    properties: Arc<PlanProperties>,
}

impl ExecutionPlan for RangeManipulateExec {
    fn apply_expressions(
        &self,
        _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> datafusion_common::Result<TreeNodeRecursion>,
    ) -> DataFusionResult<TreeNodeRecursion> {
        Ok(TreeNodeRecursion::Continue)
    }

    fn schema(&self) -> SchemaRef {
        self.output_schema.clone()
    }

    fn properties(&self) -> &Arc<PlanProperties> {
        &self.properties
    }

    fn maintains_input_order(&self) -> Vec<bool> {
        vec![true; self.children().len()]
    }

    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
        vec![&self.input]
    }

    fn input_distribution_requirements(&self) -> InputDistributionRequirements {
        let input_requirement = self
            .input
            .input_distribution_requirements()
            .into_per_child();
        if input_requirement.is_empty() {
            // if the input is EmptyMetric, its required_input_distribution() is empty so we can't
            // use its input distribution.
            InputDistributionRequirements::new(vec![Distribution::UnspecifiedDistribution])
        } else {
            InputDistributionRequirements::new(input_requirement)
        }
    }

    fn with_new_children(
        self: Arc<Self>,
        children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
        assert!(!children.is_empty());
        let exec_input = children[0].clone();
        let properties = exec_input.properties();
        let properties = Arc::new(PlanProperties::new(
            EquivalenceProperties::new(self.output_schema.clone()),
            properties.partitioning.clone(),
            properties.emission_type,
            properties.boundedness,
        ));
        Ok(Arc::new(Self {
            offset: self.offset,
            start: self.start,
            end: self.end,
            interval: self.interval,
            range: self.range,
            time_index_column: self.time_index_column.clone(),
            time_range_column: self.time_range_column.clone(),
            field_columns: self.field_columns.clone(),
            output_schema: self.output_schema.clone(),
            input: children[0].clone(),
            metric: self.metric.clone(),
            properties,
        }))
    }

    fn execute(
        &self,
        partition: usize,
        context: Arc<TaskContext>,
    ) -> DataFusionResult<SendableRecordBatchStream> {
        let baseline_metric = BaselineMetrics::new(&self.metric, partition);
        let metrics_builder = MetricBuilder::new(&self.metric);
        let num_series = Count::new();
        metrics_builder
            .with_partition(partition)
            .build(MetricValue::Count {
                name: METRIC_NUM_SERIES.into(),
                count: num_series.clone(),
            });

        let input = self.input.execute(partition, context)?;
        let schema = input.schema();
        let time_index = schema
            .column_with_name(&self.time_index_column)
            .unwrap_or_else(|| panic!("time index column {} not found", self.time_index_column))
            .0;
        let field_columns = self
            .field_columns
            .iter()
            .map(|value_col| {
                schema
                    .column_with_name(value_col)
                    .unwrap_or_else(|| panic!("value column {value_col} not found",))
                    .0
            })
            .collect();
        let time_unit = timestamp_unit(schema.field(time_index).data_type())?;
        let aligned_ts_array =
            RangeManipulateStream::build_aligned_ts_array(self.start, self.end, self.interval);
        Ok(Box::pin(RangeManipulateStream {
            offset: self.offset,
            start: self.start,
            end: self.end,
            interval: self.interval,
            range: self.range,
            time_index,
            time_unit,
            field_columns,
            aligned_ts_array,
            output_schema: self.output_schema.clone(),
            input,
            metric: baseline_metric,
            num_series,
        }))
    }

    fn metrics(&self) -> Option<MetricsSet> {
        Some(self.metric.clone_inner())
    }

    fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {
        vec![ChildStats::At(partition)]
    }

    fn statistics_from_inputs(
        &self,
        input_stats: &[Arc<Statistics>],
        _args: &StatisticsArgs,
    ) -> DataFusionResult<Arc<Statistics>> {
        let input_stats = &input_stats[0];

        let estimated_row_num = (self.end - self.start) as f64 / self.interval as f64;
        let estimated_total_bytes = input_stats
            .total_byte_size
            .get_value()
            .zip(input_stats.num_rows.get_value())
            .map(|(size, rows)| {
                Precision::Inexact(((*size as f64 / *rows as f64) * estimated_row_num).floor() as _)
            })
            .unwrap_or_default();

        Ok(Arc::new(Statistics {
            num_rows: Precision::Inexact(estimated_row_num as _),
            total_byte_size: estimated_total_bytes,
            // TODO(ruihang): support this column statistics
            column_statistics: Statistics::unknown_column(&self.schema()),
        }))
    }

    fn name(&self) -> &str {
        "RangeManipulateExec"
    }
}

impl DisplayAs for RangeManipulateExec {
    fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match t {
            DisplayFormatType::Default
            | DisplayFormatType::Verbose
            | DisplayFormatType::TreeRender => {
                write!(
                    f,
                    "PromRangeManipulateExec: req range=[{}..{}], interval=[{}], eval range=[{}], time index=[{}]",
                    self.start, self.end, self.interval, self.range, self.time_index_column
                )
            }
        }
    }
}

pub struct RangeManipulateStream {
    offset: Millisecond,
    start: Millisecond,
    end: Millisecond,
    interval: Millisecond,
    range: Millisecond,
    time_index: usize,
    time_unit: TimeUnit,
    field_columns: Vec<usize>,
    aligned_ts_array: ArrayRef,

    output_schema: SchemaRef,
    input: SendableRecordBatchStream,
    metric: BaselineMetrics,
    /// Number of series processed.
    num_series: Count,
}

impl RecordBatchStream for RangeManipulateStream {
    fn schema(&self) -> SchemaRef {
        self.output_schema.clone()
    }
}

impl Stream for RangeManipulateStream {
    type Item = DataFusionResult<RecordBatch>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let poll = loop {
            match ready!(self.input.poll_next_unpin(cx)) {
                Some(Ok(batch)) => {
                    let timer = std::time::Instant::now();
                    let result = self.manipulate(batch);
                    if let Ok(None) = result {
                        self.metric.elapsed_compute().add_elapsed(timer);
                        continue;
                    } else {
                        self.num_series.add(1);
                        self.metric.elapsed_compute().add_elapsed(timer);
                        break Poll::Ready(result.transpose());
                    }
                }
                None => {
                    PROMQL_SERIES_COUNT.observe(self.num_series.value() as f64);
                    break Poll::Ready(None);
                }
                Some(Err(e)) => break Poll::Ready(Some(Err(e))),
            }
        };
        self.metric.record_poll(poll)
    }
}

impl RangeManipulateStream {
    // Prometheus: https://github.com/prometheus/prometheus/blob/e934d0f01158a1d55fa0ebb035346b195fcc1260/promql/engine.go#L1113-L1198
    // But they are not exactly the same, because we don't eager-evaluate on the data in this plan.
    // And the generated timestamp is not aligned to the step. It's expected to do later.
    pub fn manipulate(&self, input: RecordBatch) -> DataFusionResult<Option<RecordBatch>> {
        // Prometheus sample identity is one sample per series per timestamp, and this operator's
        // input is one series per batch (the planner inserts a `SeriesDivide` for that). A read
        // path that does not deduplicate can still deliver two rows for one timestamp, and every
        // range function reads its samples from here, so collapse them once instead of
        // special-casing each function. Comparing native ticks rather than the millisecond
        // payload keeps distinct sub-millisecond samples apart.
        //
        // The first row of a tied group wins. A deduplicating scan also keeps the first row
        // after ordering equal primary keys and timestamps by descending sequence, so the values
        // agree when the scan presents those rows newest-first. This operator cannot establish
        // that order itself: unordered scans and scans whose sequence values were overwritten
        // can select a different tied value. The collapse fixes sample multiplicity, not global
        // newest-write conflict resolution.
        let (native_timestamps, _) = timestamp_array_to_primitive(input.column(self.time_index))
            .ok_or_else(|| {
                DataFusionError::Execution("Time index column is not a timestamp".into())
            })?;
        let native_values = native_timestamps.values().as_ref();
        let first_duplicate = if native_timestamps.null_count() == 0 {
            native_values
                .windows(2)
                .position(|pair| pair[0] == pair[1])
                .map(|index| index + 1)
        } else {
            (1..input.num_rows()).find(|&i| {
                !native_timestamps.is_null(i - 1)
                    && !native_timestamps.is_null(i)
                    && native_values[i - 1] == native_values[i]
            })
        };
        let indices = first_duplicate.map(|first_duplicate| {
            let mut indices = Vec::with_capacity(input.num_rows());
            indices.extend(0..first_duplicate);
            indices.extend((first_duplicate + 1..input.num_rows()).filter(|&i| {
                native_timestamps.is_null(i)
                    || native_timestamps.is_null(i - 1)
                    || native_values[i] != native_values[i - 1]
            }));
            Int64Array::from_iter_values(indices.into_iter().map(|index| index as i64))
        });
        // Only sample columns need collapsing. Tags below still come from row zero of the
        // original batch, and the aligned time index is replaced below. In particular, do not
        // build a RecordBatch with collapsed sample columns and uncollapsed tags: Arrow requires
        // all columns to have the same length.
        let ts_column = match indices.as_ref() {
            Some(indices) => compute::take(input.column(self.time_index), indices, None)?,
            None => input.column(self.time_index).clone(),
        };
        let mut other_columns = (0..input.columns().len()).collect::<HashSet<_>>();
        other_columns.remove(&self.time_index);
        // calculate the range
        let (ranges, (start, end)) = self.calculate_range_on_timestamps(&ts_column)?;
        // ignore this if all ranges are empty
        if ranges.iter().all(|(_, len)| *len == 0) {
            return Ok(None);
        }

        // transform columns
        let mut new_columns = input.columns().to_vec();
        for index in self.field_columns.iter() {
            let _ = other_columns.remove(index);
            let column = match indices.as_ref() {
                Some(indices) => compute::take(input.column(*index), indices, None)?,
                None => input.column(*index).clone(),
            };
            let new_column = Arc::new(
                RangeArray::from_ranges(column, ranges.clone())
                    .map_err(|e| ArrowError::InvalidArgumentError(e.to_string()))?
                    .into_dict(),
            );
            new_columns[*index] = new_column;
        }

        // The timestamp range payload is always millisecond ABI. A column is reusable as-is
        // only when its own data type already is that ABI: the output schema fixes the
        // dictionary value type, so a millisecond column carrying a timezone must be
        // converted like any other native unit.
        let timestamp_values: ArrayRef = if self.offset == 0
            && matches!(
                ts_column.data_type(),
                DataType::Timestamp(TimeUnit::Millisecond, None)
            ) {
            ts_column.clone()
        } else {
            // Shift in wide native precision before truncating toward zero, preserving
            // null validity.
            let scale = nanoseconds_per_native_tick(self.time_unit);
            let (timestamps, _) = timestamp_array_to_primitive(&ts_column).ok_or_else(|| {
                DataFusionError::Execution("Time index column is not a timestamp".into())
            })?;
            // Single pass over the timestamps: an i128 intermediate holds any native
            // timestamp shift; only the final millisecond result must fit i64. Null-free
            // batches (the dominant PromQL case) take arrow's infallible `unary` kernel,
            // which LLVM can vectorize, and report overflow through a shared cell. But
            // `unary` visits NULL slots' underlying values, so a NULL slot holding an
            // arbitrary overflowing value would raise a false batch-wide error; nullable
            // batches therefore take `try_unary`, which applies the closure to valid rows
            // only.
            let offset_ns = (self.offset as i128) * 1_000_000;
            let timestamp_values = if timestamps.null_count() == 0 {
                let overflow = std::cell::Cell::new(false);
                let shifted = compute::unary::<Int64Type, _, Int64Type>(&timestamps, |v| {
                    let shifted_ns = (v as i128) * scale + offset_ns;
                    let shifted_ms = shifted_ns / 1_000_000;
                    match i64::try_from(shifted_ms) {
                        Ok(ms) => ms,
                        Err(_) => {
                            overflow.set(true);
                            0
                        }
                    }
                });
                if overflow.get() {
                    return Err(DataFusionError::Execution(
                        "RangeManipulate timestamp payload overflow".into(),
                    ));
                }
                shifted
            } else {
                compute::try_unary::<Int64Type, _, Int64Type>(&timestamps, |v| {
                    let shifted_ns = (v as i128) * scale + offset_ns;
                    i64::try_from(shifted_ns / 1_000_000).map_err(|_| {
                        ArrowError::ComputeError(
                            "RangeManipulate timestamp payload overflow".into(),
                        )
                    })
                })?
            };
            Arc::new(TimestampMillisecondArray::new(
                timestamp_values.values().clone(),
                timestamp_values.nulls().cloned(),
            ))
        };
        let ts_range_column = RangeArray::from_ranges(timestamp_values, ranges.clone())
            .map_err(|e| ArrowError::InvalidArgumentError(e.to_string()))?
            .into_dict();
        new_columns.push(Arc::new(ts_range_column));

        // truncate other columns
        let take_indices = Int64Array::from(vec![0; ranges.len()]);
        for index in other_columns.into_iter() {
            new_columns[index] = compute::take(&input.column(index), &take_indices, None)?;
        }
        // replace timestamp with the aligned one
        let new_time_index = if ranges.len() != self.aligned_ts_array.len() {
            Self::build_aligned_ts_array(start, end, self.interval)
        } else {
            self.aligned_ts_array.clone()
        };
        new_columns[self.time_index] = new_time_index;

        RecordBatch::try_new(self.output_schema.clone(), new_columns)
            .map(Some)
            .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))
    }

    fn build_aligned_ts_array(start: i64, end: i64, interval: i64) -> ArrayRef {
        Arc::new(TimestampMillisecondArray::from_iter_values(
            (start..=end).step_by(interval as _),
        ))
    }

    /// Return values:
    /// - A vector of tuples where each tuple contains the start index and length of the range.
    /// - A tuple of the actual start/end timestamp used to calculate the range.
    #[cfg(test)]
    #[allow(clippy::type_complexity)]
    fn calculate_range(
        &self,
        input: &RecordBatch,
    ) -> DataFusionResult<(Vec<(u32, u32)>, (i64, i64))> {
        self.calculate_range_on_timestamps(input.column(self.time_index))
    }

    #[allow(clippy::type_complexity)]
    fn calculate_range_on_timestamps(
        &self,
        ts_column: &ArrayRef,
    ) -> DataFusionResult<(Vec<(u32, u32)>, (i64, i64))> {
        let scale = nanoseconds_per_native_tick(self.time_unit);
        let (timestamps, _) = timestamp_array_to_primitive(ts_column).ok_or_else(|| {
            DataFusionError::Execution("Time index column is not a timestamp".into())
        })?;
        let timestamp_null_count = timestamps.null_count();
        let timestamps = timestamps.values();
        let timestamp =
            |index| (timestamps[index] as i128) * scale + (self.offset as i128) * 1_000_000;
        let len = timestamps.len();
        if len == 0 {
            return Ok((vec![], (self.start, self.end)));
        }

        // Shorten the range using wide arithmetic so timestamps near the native
        // type limits retain every query-aligned evaluation point.
        let query_start = self.start as i128;
        let query_end = self.end as i128;
        let interval = self.interval as i128;
        let first_ts = timestamp(0).div_euclid(1_000_000);
        // Preserve the query's alignment pattern when optimizing start time.
        let remainder = (first_ts - query_start).rem_euclid(interval);
        let first_ts_aligned = first_ts + (interval - remainder).rem_euclid(interval);
        let last_ts_with_range =
            (timestamp(len - 1) + (self.range as i128) * 1_000_000).div_euclid(1_000_000);
        let remainder = (last_ts_with_range - query_start).rem_euclid(interval);
        let last_ts_aligned = last_ts_with_range - remainder;
        let start = query_start.max(first_ts_aligned);
        let end = query_end.min(last_ts_aligned);
        if start > end {
            return Ok((vec![], (self.start, self.end)));
        }
        // The intersection is within the declared i64 query bounds.
        let start = start as i64;
        let end = end as i64;

        // Range membership is decided on shifted native ticks, before the
        // timestamp-range payload is converted to its millisecond ABI. This
        // keeps sub-millisecond samples distinct in a range; equal millisecond
        // payload values are not a reason to deduplicate input samples.
        //
        // Calculate for every aligned timestamp (`curr_ts`), assuming ordered
        // timestamps: the scan compares shifted samples in `i64` nanoseconds
        // whenever they fit the batch, and in exact `i128` nanoseconds otherwise.
        let narrow = (timestamp_null_count == 0)
            .then(|| {
                narrow_scan(
                    self.time_unit,
                    self.offset,
                    self.range,
                    timestamps,
                    start,
                    end,
                )
            })
            .flatten();
        let ranges = match narrow {
            Some(scan) => scan_ranges::<i64>(self.interval, timestamps, start, end, &scan),
            None => {
                let scan = ScanNanoseconds::<i128>::wide(self.time_unit, self.offset, self.range);
                scan_ranges::<i128>(self.interval, timestamps, start, end, &scan)
            }
        };

        Ok((ranges, (start, end)))
    }

    /// Cursor scan of `timestamps` through the narrow `i64` carrier, for benchmarks.
    ///
    /// `bench_range_fn` times the fast path directly, next to the exact `i128` fallback in
    /// [`RangeManipulateStream::bench_scan_wide`]; production selects the carrier through
    /// [`narrow_scan`] inside `calculate_range`. Returns `None` when that carrier is
    /// rejected, which is exactly the condition that makes the operator fall back to `i128`.
    #[doc(hidden)]
    pub fn bench_scan_narrow(
        time_unit: TimeUnit,
        offset: Millisecond,
        window: Millisecond,
        interval: Millisecond,
        timestamps: &[i64],
        start: Millisecond,
        end: Millisecond,
    ) -> Option<Vec<(u32, u32)>> {
        let scan = narrow_scan(time_unit, offset, window, timestamps, start, end)?;
        Some(scan_ranges::<i64>(interval, timestamps, start, end, &scan))
    }

    /// Cursor scan of `timestamps` through the exact `i128` carrier, for benchmarks.
    ///
    /// The carrier is exact for every input, so this times the fallback that
    /// [`narrow_scan`] selects whenever the narrow baselines would overflow.
    #[doc(hidden)]
    pub fn bench_scan_wide(
        time_unit: TimeUnit,
        offset: Millisecond,
        window: Millisecond,
        interval: Millisecond,
        timestamps: &[i64],
        start: Millisecond,
        end: Millisecond,
    ) -> Vec<(u32, u32)> {
        scan_ranges::<i128>(
            interval,
            timestamps,
            start,
            end,
            &ScanNanoseconds::<i128>::wide(time_unit, offset, window),
        )
    }
}

/// Cursor scan over the aligned evaluation points, in `Ns` nanoseconds.
fn scan_ranges<Ns: Nanoseconds>(
    interval: Millisecond,
    timestamps: &[i64],
    start: Millisecond,
    end: Millisecond,
    scan: &ScanNanoseconds<Ns>,
) -> Vec<(u32, u32)> {
    let len = timestamps.len();
    let mut ranges = Vec::new();
    let mut left = 0usize;
    let mut right = 0usize;
    for curr_ts in (start..=end).step_by(interval as _) {
        let curr_ts_ns = scan.instant(curr_ts);
        let start_ts = scan.window_start(curr_ts_ns);

        while left < len && scan.shift(timestamps[left]) <= start_ts {
            left += 1;
        }
        right = right.max(left);
        while right < len && scan.shift(timestamps[right]) <= curr_ts_ns {
            right += 1;
        }

        if left == right {
            ranges.push((0, 0));
        } else {
            ranges.push((left as _, (right - left) as _));
        }
    }
    ranges
}

/// Cursor scan baselines in nanoseconds, shared with `InstantManipulate`.
///
/// The narrow `i64` baselines come from [`narrow_scan`], the exact ones from
/// [`ScanNanoseconds::wide`].
#[derive(Debug, Clone, Copy)]
pub(super) struct ScanNanoseconds<Ns> {
    /// Nanoseconds per timestamp tick.
    scale: Ns,
    /// The offset shift in nanoseconds.
    offset_ns: Ns,
    /// The range or lookback window in nanoseconds.
    window_ns: Ns,
}

/// Nanoseconds carrier of the cursor scans: `i64` where [`narrow_scan`] proved that
/// the batch fits, `i128` for the exact fallback.
pub(super) trait Nanoseconds:
    Copy + PartialOrd + From<i64> + Add<Output = Self> + Sub<Output = Self> + Mul<Output = Self>
{
}

impl<T> Nanoseconds for T where
    T: Copy + PartialOrd + From<i64> + Add<Output = Self> + Sub<Output = Self> + Mul<Output = Self>
{
}

impl ScanNanoseconds<i128> {
    /// The `i128` baselines, exact for every input.
    pub(super) fn wide(time_unit: TimeUnit, offset: Millisecond, window: Millisecond) -> Self {
        Self {
            scale: nanoseconds_per_native_tick(time_unit),
            offset_ns: i128::from(offset) * 1_000_000,
            window_ns: i128::from(window) * 1_000_000,
        }
    }
}

impl<Ns: Nanoseconds> ScanNanoseconds<Ns> {
    /// Shifted nanoseconds of an input timestamp.
    pub(super) fn shift(&self, timestamp: i64) -> Ns {
        Ns::from(timestamp) * self.scale + self.offset_ns
    }

    /// Nanoseconds of an evaluation instant given in milliseconds.
    pub(super) fn instant(&self, millis: Millisecond) -> Ns {
        Ns::from(millis) * Ns::from(1_000_000)
    }

    /// The exclusive lower boundary of the window ending at `instant`.
    pub(super) fn window_start(&self, instant: Ns) -> Ns {
        instant - self.window_ns
    }
}

/// The narrow `i64` baselines of a sorted batch, or `None` when they would overflow.
pub(super) fn narrow_scan(
    time_unit: TimeUnit,
    offset: Millisecond,
    window: Millisecond,
    timestamps: &[i64],
    start: Millisecond,
    end: Millisecond,
) -> Option<ScanNanoseconds<i64>> {
    let scale = i64::try_from(nanoseconds_per_native_tick(time_unit)).ok()?;
    let offset_ns = offset.checked_mul(1_000_000)?;
    let window_ns = window.checked_mul(1_000_000)?;
    // Sorted samples, so the checked endpoints bound every interior shift and product.
    let shift = |timestamp: i64| timestamp.checked_mul(scale)?.checked_add(offset_ns);
    shift(*timestamps.first()?)?;
    shift(*timestamps.last()?)?;
    // The checked `[start, end]` bounds every instant and every `instant - window`.
    start.checked_mul(1_000_000)?.checked_sub(window_ns)?;
    end.checked_mul(1_000_000)?.checked_sub(window_ns)?;
    Some(ScanNanoseconds {
        scale,
        offset_ns,
        window_ns,
    })
}

#[cfg(test)]
mod test {
    use datafusion::arrow::array::{
        ArrayRef, DictionaryArray, Float64Array, StringArray, TimestampMicrosecondArray,
        TimestampNanosecondArray, TimestampSecondArray,
    };
    use datafusion::arrow::buffer::NullBuffer;
    use datafusion::arrow::datatypes::{
        ArrowPrimitiveType, DataType, Field, Int64Type, Schema, TimestampMillisecondType,
    };
    use datafusion::common::ToDFSchema;
    use datafusion::datasource::memory::MemorySourceConfig;
    use datafusion::datasource::source::DataSourceExec;
    use datafusion::logical_expr::{
        EmptyRelation, Extension, LogicalPlan, UserDefinedLogicalNodeCore,
    };
    use datafusion::physical_expr::Partitioning;
    use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
    use datafusion::physical_plan::memory::MemoryStream;
    use datafusion::physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions};
    use datafusion::prelude::SessionContext;
    use datatypes::arrow::array::TimestampMillisecondArray;
    use futures::FutureExt;

    use super::*;

    const TIME_INDEX_COLUMN: &str = "timestamp";

    /// The `(offset, length)` selections of one cursor scan.
    type RangeSelections = Vec<(u32, u32)>;

    fn project_batch(batch: &RecordBatch, indices: &[usize]) -> RecordBatch {
        let fields = indices
            .iter()
            .map(|&idx| batch.schema().field(idx).clone())
            .collect::<Vec<_>>();
        let columns = indices
            .iter()
            .map(|&idx| batch.column(idx).clone())
            .collect::<Vec<_>>();
        let schema = Arc::new(Schema::new(fields));
        RecordBatch::try_new(schema, columns).unwrap()
    }

    fn prepare_test_data() -> DataSourceExec {
        let schema = Arc::new(Schema::new(vec![
            Field::new(TIME_INDEX_COLUMN, TimestampMillisecondType::DATA_TYPE, true),
            Field::new("value_1", DataType::Float64, true),
            Field::new("value_2", DataType::Float64, true),
            Field::new("path", DataType::Utf8, true),
        ]));
        let timestamp_column = Arc::new(TimestampMillisecondArray::from(vec![
            0, 30_000, 60_000, 90_000, 120_000, // every 30s
            180_000, 240_000, // every 60s
            241_000, 271_000, 291_000, // others
        ])) as _;
        let field_column: ArrayRef = Arc::new(Float64Array::from(vec![1.0; 10])) as _;
        let path_column = Arc::new(StringArray::from(vec!["foo"; 10])) as _;
        let data = RecordBatch::try_new(
            schema.clone(),
            vec![
                timestamp_column,
                field_column.clone(),
                field_column,
                path_column,
            ],
        )
        .unwrap();

        DataSourceExec::new(Arc::new(
            MemorySourceConfig::try_new(&[vec![data]], schema, None).unwrap(),
        ))
    }

    async fn do_normalize_test(
        start: Millisecond,
        end: Millisecond,
        interval: Millisecond,
        range: Millisecond,
        expected: String,
    ) {
        let memory_exec = Arc::new(prepare_test_data());
        let time_index = TIME_INDEX_COLUMN.to_string();
        let field_columns = vec!["value_1".to_string(), "value_2".to_string()];
        let manipulate_output_schema = SchemaRef::new(
            RangeManipulate::calculate_output_schema(
                &memory_exec.schema().to_dfschema_ref().unwrap(),
                &time_index,
                &field_columns,
            )
            .unwrap()
            .as_arrow()
            .clone(),
        );
        let properties = Arc::new(PlanProperties::new(
            EquivalenceProperties::new(manipulate_output_schema.clone()),
            Partitioning::UnknownPartitioning(1),
            EmissionType::Incremental,
            Boundedness::Bounded,
        ));
        let normalize_exec = Arc::new(RangeManipulateExec {
            offset: 0,
            start,
            end,
            interval,
            range,
            field_columns,
            output_schema: manipulate_output_schema,
            time_range_column: RangeManipulate::build_timestamp_range_name(&time_index),
            time_index_column: time_index,
            input: memory_exec,
            metric: ExecutionPlanMetricsSet::new(),
            properties,
        });
        let session_context = SessionContext::default();
        let result = datafusion::physical_plan::collect(normalize_exec, session_context.task_ctx())
            .await
            .unwrap();
        // DirectoryArray from RangeArray cannot be print as normal arrays.
        let result_literal: String = result
            .into_iter()
            .filter_map(|batch| {
                batch
                    .columns()
                    .iter()
                    .map(|array| {
                        if matches!(array.data_type(), &DataType::Dictionary(..)) {
                            let dict_array = array
                                .as_any()
                                .downcast_ref::<DictionaryArray<Int64Type>>()
                                .unwrap()
                                .clone();
                            format!("{:?}", RangeArray::try_new(dict_array).unwrap())
                        } else {
                            format!("{array:?}")
                        }
                    })
                    .reduce(|lhs, rhs| lhs + "\n" + &rhs)
            })
            .reduce(|lhs, rhs| lhs + "\n\n" + &rhs)
            .unwrap();

        assert_eq!(result_literal, expected);
    }

    #[tokio::test]
    async fn duplicate_native_timestamps_collapse_in_real_plan() {
        use datafusion::arrow::array::Float64Array;

        let schema = Arc::new(Schema::new(vec![
            Field::new(
                TIME_INDEX_COLUMN,
                DataType::Timestamp(TimeUnit::Millisecond, None),
                false,
            ),
            Field::new("value", DataType::Float64, false),
        ]));
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(TimestampMillisecondArray::from(vec![
                    29_600, 29_600, 59_600, 59_600,
                ])),
                Arc::new(Float64Array::from(vec![100.0, 100.0, 102.0, 102.0])),
            ],
        )
        .unwrap();
        let plan = RangeManipulate::new(
            60_000,
            60_000,
            60_000,
            0,
            60_000,
            TIME_INDEX_COLUMN.to_string(),
            vec!["value".to_string()],
            LogicalPlan::EmptyRelation(EmptyRelation {
                produce_one_row: false,
                schema: schema.clone().to_dfschema_ref().unwrap(),
            }),
        )
        .unwrap();
        let input = Arc::new(DataSourceExec::new(Arc::new(
            MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(),
        )));
        let batches = datafusion::physical_plan::collect(
            plan.to_execution_plan(input),
            SessionContext::default().task_ctx(),
        )
        .await
        .unwrap();
        let out = &batches[0];
        let values = RangeArray::try_new(
            out.column(1)
                .as_any()
                .downcast_ref::<DictionaryArray<Int64Type>>()
                .unwrap()
                .clone(),
        )
        .unwrap();
        let timestamps = RangeArray::try_new(
            out.column(2)
                .as_any()
                .downcast_ref::<DictionaryArray<Int64Type>>()
                .unwrap()
                .clone(),
        )
        .unwrap();
        assert_eq!(
            values
                .get(0)
                .unwrap()
                .as_any()
                .downcast_ref::<Float64Array>()
                .unwrap()
                .values(),
            &[100.0, 102.0]
        );
        assert_eq!(
            timestamps
                .get(0)
                .unwrap()
                .as_any()
                .downcast_ref::<TimestampMillisecondArray>()
                .unwrap()
                .values(),
            &[29_600, 59_600]
        );
    }

    #[tokio::test]
    async fn duplicate_timestamp_keeps_first_row_and_aligned_tags() {
        use datafusion::arrow::array::{Float64Array, StringArray};

        // Rows sharing a timestamp but carrying different values and tags show which row the
        // collapse keeps, and that every column follows the same index instead of drifting apart.
        let schema = Arc::new(Schema::new(vec![
            Field::new(
                TIME_INDEX_COLUMN,
                DataType::Timestamp(TimeUnit::Millisecond, None),
                false,
            ),
            Field::new("value", DataType::Float64, false),
            Field::new("tag", DataType::Utf8, false),
        ]));
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(TimestampMillisecondArray::from(vec![
                    5_000, 10_000, 10_000, 15_000,
                ])),
                Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0])),
                Arc::new(StringArray::from(vec!["a", "b", "c", "d"])),
            ],
        )
        .unwrap();
        let plan = RangeManipulate::new(
            20_000,
            20_000,
            20_000,
            0,
            20_000,
            TIME_INDEX_COLUMN.to_string(),
            vec!["value".to_string()],
            LogicalPlan::EmptyRelation(EmptyRelation {
                produce_one_row: false,
                schema: schema.clone().to_dfschema_ref().unwrap(),
            }),
        )
        .unwrap();
        let input = Arc::new(DataSourceExec::new(Arc::new(
            MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(),
        )));
        let batches = datafusion::physical_plan::collect(
            plan.to_execution_plan(input),
            SessionContext::default().task_ctx(),
        )
        .await
        .unwrap();
        let out = &batches[0];
        let values = RangeArray::try_new(
            out.column(1)
                .as_any()
                .downcast_ref::<DictionaryArray<Int64Type>>()
                .unwrap()
                .clone(),
        )
        .unwrap();
        let timestamps = RangeArray::try_new(
            out.column(3)
                .as_any()
                .downcast_ref::<DictionaryArray<Int64Type>>()
                .unwrap()
                .clone(),
        )
        .unwrap();
        // The first row of the tied pair is kept, matching what a deduplicating read returns, so
        // 2.0 survives and 3.0 is dropped.
        assert_eq!(
            values
                .get(0)
                .unwrap()
                .as_any()
                .downcast_ref::<Float64Array>()
                .unwrap()
                .values(),
            &[1.0, 2.0, 4.0]
        );
        assert_eq!(
            timestamps
                .get(0)
                .unwrap()
                .as_any()
                .downcast_ref::<TimestampMillisecondArray>()
                .unwrap()
                .values(),
            &[5_000, 10_000, 15_000]
        );
        // Tag columns are not range payloads, so the operator repeats the batch's first row for
        // every output row: a batch holds one series, so every row's tag is that series' tag.
        assert_eq!(
            out.column(2)
                .as_any()
                .downcast_ref::<StringArray>()
                .unwrap()
                .iter()
                .collect::<Vec<_>>(),
            vec![Some("a")]
        );
    }

    #[test]
    fn duplicate_sample_range_udfs_report_deduplicated_values() {
        use crate::functions::{CountOverTime, Delta, IDelta, Increase, Rate};
        use datafusion::arrow::array::Float64Array;
        use datafusion::physical_plan::ColumnarValue;
        use datafusion_common::config::ConfigOptions;
        use datafusion_expr::ScalarFunctionArgs;
        let evaluate = |udf: datafusion::logical_expr::ScalarUDF, range_ms: Option<i64>| {
            let ts = RangeArray::from_ranges(
                Arc::new(TimestampMillisecondArray::from(vec![29_600, 59_600])),
                [(0, 2)],
            )
            .unwrap();
            let values =
                RangeArray::from_ranges(Arc::new(Float64Array::from(vec![100.0, 102.0])), [(0, 2)])
                    .unwrap();
            let mut args = vec![
                ColumnarValue::Array(Arc::new(ts.into_dict())),
                ColumnarValue::Array(Arc::new(values.into_dict())),
            ];
            let mut arg_fields = args
                .iter()
                .enumerate()
                .map(|(i, arg)| {
                    Arc::new(Field::new(
                        if i == 0 { "ts" } else { "value" },
                        arg.data_type(),
                        false,
                    ))
                })
                .collect::<Vec<_>>();
            if let Some(range_ms) = range_ms {
                let eval_ts: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![60_000]));
                args.push(ColumnarValue::Array(eval_ts));
                args.push(ColumnarValue::Scalar(
                    datafusion_common::ScalarValue::Int64(Some(range_ms)),
                ));
                arg_fields.push(Arc::new(Field::new("eval_ts", args[2].data_type(), false)));
                arg_fields.push(Arc::new(Field::new("range", DataType::Int64, false)));
            }
            let result = udf
                .invoke_with_args(ScalarFunctionArgs {
                    args,
                    arg_fields,
                    number_rows: 1,
                    return_field: Arc::new(Field::new("result", DataType::Float64, false)),
                    config_options: Arc::new(ConfigOptions::default()),
                })
                .unwrap();
            let array = match result {
                ColumnarValue::Array(a) => a,
                ColumnarValue::Scalar(s) => s.to_array_of_size(1).unwrap(),
            };
            array
                .as_any()
                .downcast_ref::<Float64Array>()
                .unwrap()
                .value(0)
        };
        assert_eq!(evaluate(CountOverTime::scalar_udf(), None), 2.0);
        assert_eq!(evaluate(Delta::scalar_udf(), Some(60_000)), 4.0);
        assert_eq!(evaluate(Increase::scalar_udf(), Some(60_000)), 4.0);
        assert!((evaluate(Rate::scalar_udf(), Some(60_000)) - 4.0 / 60.0).abs() < 0.0001);
        assert_eq!(evaluate(IDelta::<false>::scalar_udf(), None), 2.0);
        assert!((evaluate(IDelta::<true>::scalar_udf(), None) - 2.0 / 30.0).abs() < 0.0001);
    }

    #[tokio::test]
    async fn distinct_submillisecond_timestamps_are_not_collapsed() {
        let schema = Arc::new(Schema::new(vec![
            Field::new(
                TIME_INDEX_COLUMN,
                DataType::Timestamp(TimeUnit::Microsecond, None),
                false,
            ),
            Field::new("value", DataType::Float64, false),
        ]));
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(TimestampMicrosecondArray::from(vec![1_000_001, 1_000_002])),
                Arc::new(Float64Array::from(vec![10.0, 20.0])),
            ],
        )
        .unwrap();
        let plan = RangeManipulate::new(
            1_001,
            1_001,
            1,
            0,
            1,
            TIME_INDEX_COLUMN.to_string(),
            vec!["value".to_string()],
            LogicalPlan::EmptyRelation(EmptyRelation {
                produce_one_row: false,
                schema: schema.clone().to_dfschema_ref().unwrap(),
            }),
        )
        .unwrap();
        let input = Arc::new(DataSourceExec::new(Arc::new(
            MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(),
        )));
        let batches = datafusion::physical_plan::collect(
            plan.to_execution_plan(input),
            SessionContext::default().task_ctx(),
        )
        .await
        .unwrap();
        let output = &batches[0];
        let values = RangeArray::try_new(
            output
                .column(1)
                .as_any()
                .downcast_ref::<DictionaryArray<Int64Type>>()
                .unwrap()
                .clone(),
        )
        .unwrap();
        let timestamps = RangeArray::try_new(
            output
                .column(2)
                .as_any()
                .downcast_ref::<DictionaryArray<Int64Type>>()
                .unwrap()
                .clone(),
        )
        .unwrap();
        assert_eq!(
            values
                .get(0)
                .unwrap()
                .as_any()
                .downcast_ref::<Float64Array>()
                .unwrap()
                .values(),
            &[10.0, 20.0]
        );
        assert_eq!(
            timestamps
                .get(0)
                .unwrap()
                .as_any()
                .downcast_ref::<TimestampMillisecondArray>()
                .unwrap()
                .values(),
            &[1_000, 1_000]
        );
    }

    #[tokio::test]
    async fn native_timestamps_preserve_range_membership_and_ms_payload() {
        for (unit, ticks_per_ms) in [
            (TimeUnit::Microsecond, 1_000_i64),
            (TimeUnit::Nanosecond, 1_000_000_i64),
        ] {
            let lower = 1_000 * ticks_per_ms;
            let upper = 1_001 * ticks_per_ms;
            // Exclude the lower boundary and future sample; retain both native
            // samples in the same millisecond bucket and the exact upper sample.
            let timestamps = vec![lower, lower + 1, lower + 2, upper, upper + 1];
            let time: ArrayRef = match unit {
                TimeUnit::Microsecond => Arc::new(TimestampMicrosecondArray::from(timestamps)),
                TimeUnit::Nanosecond => Arc::new(TimestampNanosecondArray::from(timestamps)),
                _ => unreachable!(),
            };
            let schema = Arc::new(Schema::new(vec![
                Field::new(TIME_INDEX_COLUMN, DataType::Timestamp(unit, None), false),
                Field::new("value", DataType::Float64, true),
            ]));
            let batch = RecordBatch::try_new(
                schema.clone(),
                vec![
                    time,
                    Arc::new(Float64Array::from(vec![10.0, 20.0, 30.0, 40.0, 50.0])),
                ],
            )
            .unwrap();
            let logical_input = LogicalPlan::EmptyRelation(EmptyRelation {
                produce_one_row: false,
                schema: schema.clone().to_dfschema_ref().unwrap(),
            });
            let plan = RangeManipulate::new(
                1_001,
                1_001,
                1,
                0,
                1,
                TIME_INDEX_COLUMN.to_string(),
                vec!["value".to_string()],
                logical_input.clone(),
            )
            .unwrap();
            let output_time = Field::new(
                TIME_INDEX_COLUMN,
                DataType::Timestamp(TimeUnit::Millisecond, None),
                false,
            );
            let output_schema = Arc::new(Schema::new(vec![
                output_time.clone(),
                RangeArray::convert_field(&Field::new("value", DataType::Float64, true)),
                Field::new(
                    RangeManipulate::build_timestamp_range_name(TIME_INDEX_COLUMN),
                    RangeArray::convert_field(&output_time).data_type().clone(),
                    false,
                ),
            ]));
            assert_eq!(plan.schema().as_arrow(), output_schema.as_ref());

            let rebuilt = RangeManipulate::deserialize(&plan.serialize())
                .unwrap()
                .with_exprs_and_inputs(vec![], vec![logical_input])
                .unwrap();
            assert_eq!(rebuilt.schema(), plan.schema());
            assert_eq!(rebuilt.input.schema().as_arrow(), schema.as_ref());

            let input = Arc::new(DataSourceExec::new(Arc::new(
                MemorySourceConfig::try_new(&[vec![batch]], schema.clone(), None).unwrap(),
            )));
            let exec = rebuilt.to_execution_plan(input);
            assert_eq!(exec.schema(), output_schema);
            assert_eq!(exec.children()[0].schema(), schema);

            let batches =
                datafusion::physical_plan::collect(exec, SessionContext::default().task_ctx())
                    .await
                    .unwrap();
            assert_eq!(batches.len(), 1, "{unit:?}");
            let output = &batches[0];
            assert_eq!(output.schema(), output_schema);
            assert_eq!(output.num_rows(), 1);
            assert_eq!(
                output
                    .column(0)
                    .as_any()
                    .downcast_ref::<TimestampMillisecondArray>()
                    .unwrap()
                    .values()
                    .as_ref(),
                &[1_001]
            );

            // RangeArray packs offset/length into dictionary keys; Arrow dictionary
            // equality treats those packed keys as indices and cannot compare them.
            let values = RangeArray::try_new(
                output
                    .column(1)
                    .as_any()
                    .downcast_ref::<DictionaryArray<Int64Type>>()
                    .unwrap()
                    .clone(),
            )
            .unwrap();
            assert_eq!(values.get_offset_length(0), Some((1, 3)));
            assert_eq!(
                values.get(0).unwrap().to_data(),
                Float64Array::from(vec![20.0, 30.0, 40.0]).to_data()
            );
            let timestamps = RangeArray::try_new(
                output
                    .column(2)
                    .as_any()
                    .downcast_ref::<DictionaryArray<Int64Type>>()
                    .unwrap()
                    .clone(),
            )
            .unwrap();
            assert_eq!(timestamps.get_offset_length(0), Some((1, 3)));
            assert_eq!(
                timestamps.get(0).unwrap().to_data(),
                TimestampMillisecondArray::from(vec![1_000, 1_000, 1_001]).to_data()
            );
        }
    }

    #[test]
    fn logical_offset_participates_in_ordering() {
        let input = LogicalPlan::EmptyRelation(EmptyRelation {
            produce_one_row: false,
            schema: prepare_test_data().schema().to_dfschema_ref().unwrap(),
        });
        let first = RangeManipulate::new(
            0,
            0,
            0,
            0,
            0,
            TIME_INDEX_COLUMN.to_string(),
            vec!["value_1".to_string()],
            input.clone(),
        )
        .unwrap();
        let second = RangeManipulate::new(
            0,
            0,
            0,
            1,
            0,
            TIME_INDEX_COLUMN.to_string(),
            vec!["value_1".to_string()],
            input,
        )
        .unwrap();
        assert_ne!(first, second);
        assert_eq!(first.partial_cmp(&second), Some(std::cmp::Ordering::Less));
    }

    #[tokio::test]
    async fn logical_normalize_offset_survives_rebuild_and_executes() {
        for (name, time_unit, raw, offset, start, range, expected_payload) in [
            (
                "millisecond offset",
                TimeUnit::Millisecond,
                0,
                1_000,
                1_000,
                1_000,
                1_000,
            ),
            (
                "negative native lower limit with positive window",
                TimeUnit::Nanosecond,
                -9_223_112_837_000_000_000,
                -259_200_000,
                -9_223_372_037_000,
                300_000,
                -9_223_372_037_000,
            ),
            (
                "second timestamp with negative fractional offset",
                TimeUnit::Second,
                1,
                -500,
                1_000,
                1_000,
                500,
            ),
        ] {
            let schema = Arc::new(Schema::new(vec![
                Field::new(
                    TIME_INDEX_COLUMN,
                    DataType::Timestamp(time_unit, None),
                    false,
                ),
                Field::new("value", DataType::Float64, true),
            ]));
            let input = LogicalPlan::EmptyRelation(EmptyRelation {
                produce_one_row: false,
                schema: schema.clone().to_dfschema_ref().unwrap(),
            });
            let normalize = crate::extension_plan::SeriesNormalize::new(
                offset,
                TIME_INDEX_COLUMN,
                false,
                Vec::new(),
                input.clone(),
            );
            let normalize =
                crate::extension_plan::SeriesNormalize::deserialize(&normalize.serialize())
                    .unwrap()
                    .with_exprs_and_inputs(vec![], vec![input.clone()])
                    .unwrap();
            let normalized = LogicalPlan::Extension(Extension {
                node: Arc::new(normalize),
            });
            let fresh = RangeManipulate::new(
                start,
                start,
                1,
                offset,
                range,
                TIME_INDEX_COLUMN.to_string(),
                vec!["value".to_string()],
                input.clone(),
            )
            .unwrap()
            .with_exprs_and_inputs(vec![], vec![input.clone()])
            .unwrap();
            let serialized = RangeManipulate::new(
                start,
                start,
                1,
                offset,
                range,
                TIME_INDEX_COLUMN.to_string(),
                vec!["value".to_string()],
                normalized.clone(),
            )
            .unwrap();
            let decoded = RangeManipulate::deserialize(&serialized.serialize())
                .unwrap()
                .with_exprs_and_inputs(vec![], vec![normalized])
                .unwrap();
            let timestamp: ArrayRef = match time_unit {
                TimeUnit::Millisecond => Arc::new(TimestampMillisecondArray::from(vec![raw])),
                TimeUnit::Nanosecond => Arc::new(TimestampNanosecondArray::from(vec![raw])),
                TimeUnit::Second => Arc::new(TimestampSecondArray::from(vec![raw])),
                _ => unreachable!(),
            };
            let batch = RecordBatch::try_new(
                schema.clone(),
                vec![timestamp, Arc::new(Float64Array::from(vec![7.0]))],
            )
            .unwrap();
            for (mode, rebuilt) in [("fresh", fresh), ("decoded", decoded)] {
                let rebuilt = rebuilt
                    .with_exprs_and_inputs(vec![], vec![input.clone()])
                    .unwrap();
                assert_eq!(rebuilt.offset, offset, "{name}: {mode}");
                assert_eq!(rebuilt.input.schema(), input.schema(), "{name}: {mode}");

                let empty_exec_input = Arc::new(DataSourceExec::new(Arc::new(
                    MemorySourceConfig::try_new(&[vec![]], schema.clone(), None).unwrap(),
                )));
                let exec_input = Arc::new(DataSourceExec::new(Arc::new(
                    MemorySourceConfig::try_new(&[vec![batch.clone()]], schema.clone(), None)
                        .unwrap(),
                )));
                let exec = rebuilt
                    .to_execution_plan(empty_exec_input)
                    .replace_children(
                        vec![exec_input],
                        ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
                    )
                    .unwrap();
                let output =
                    datafusion::physical_plan::collect(exec, SessionContext::default().task_ctx())
                        .await
                        .unwrap();
                assert_eq!(output.len(), 1, "{name}: {mode}");
                let output = &output[0];
                assert_eq!(output.num_rows(), 1, "{name}: {mode}");
                assert_eq!(
                    output
                        .column(0)
                        .as_any()
                        .downcast_ref::<TimestampMillisecondArray>()
                        .unwrap()
                        .value(0),
                    start,
                    "{name}: {mode}"
                );
                let values = RangeArray::try_new(
                    output
                        .column(1)
                        .as_any()
                        .downcast_ref::<DictionaryArray<Int64Type>>()
                        .unwrap()
                        .clone(),
                )
                .unwrap();
                assert_eq!(values.get_offset_length(0), Some((0, 1)), "{name}: {mode}");
                assert_eq!(
                    values.get(0).unwrap().to_data(),
                    Float64Array::from(vec![7.0]).to_data(),
                    "{name}: {mode}"
                );
                let timestamps = RangeArray::try_new(
                    output
                        .column(2)
                        .as_any()
                        .downcast_ref::<DictionaryArray<Int64Type>>()
                        .unwrap()
                        .clone(),
                )
                .unwrap();
                assert_eq!(
                    timestamps.get_offset_length(0),
                    Some((0, 1)),
                    "{name}: {mode}"
                );
                assert_eq!(
                    timestamps.get(0).unwrap().to_data(),
                    TimestampMillisecondArray::from(vec![expected_payload]).to_data(),
                    "{name}: {mode}"
                );
            }
        }
    }

    #[tokio::test]
    async fn range_payload_preserves_null_timestamp_and_rejects_offset_overflow() {
        let schema = Arc::new(Schema::new(vec![
            Field::new(TIME_INDEX_COLUMN, TimestampMillisecondType::DATA_TYPE, true),
            Field::new("value", DataType::Float64, true),
        ]));
        let null_timestamp =
            TimestampMillisecondArray::new(vec![1_000].into(), Some(NullBuffer::from(vec![false])));
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(null_timestamp),
                Arc::new(Float64Array::from(vec![7.0])),
            ],
        )
        .unwrap();
        let input = Arc::new(DataSourceExec::new(Arc::new(
            MemorySourceConfig::try_new(&[vec![batch]], schema.clone(), None).unwrap(),
        )));
        let plan = RangeManipulate::new(
            1_000,
            1_000,
            1,
            0,
            1,
            TIME_INDEX_COLUMN.to_string(),
            vec!["value".to_string()],
            LogicalPlan::EmptyRelation(EmptyRelation {
                produce_one_row: false,
                schema: schema.clone().to_dfschema_ref().unwrap(),
            }),
        )
        .unwrap();
        let output = datafusion::physical_plan::collect(
            plan.to_execution_plan(input),
            SessionContext::default().task_ctx(),
        )
        .await
        .unwrap();
        let timestamps = RangeArray::try_new(
            output[0]
                .column(2)
                .as_any()
                .downcast_ref::<DictionaryArray<Int64Type>>()
                .unwrap()
                .clone(),
        )
        .unwrap();
        let payload = timestamps.get(0).unwrap();
        let payload = payload
            .as_any()
            .downcast_ref::<TimestampMillisecondArray>()
            .unwrap();
        assert_eq!(payload.len(), 1);
        assert!(!payload.is_valid(0));

        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(TimestampMillisecondArray::from(vec![0, i64::MAX])),
                Arc::new(Float64Array::from(vec![7.0, 8.0])),
            ],
        )
        .unwrap();
        let input = Arc::new(DataSourceExec::new(Arc::new(
            MemorySourceConfig::try_new(&[vec![batch]], schema.clone(), None).unwrap(),
        )));
        let normalized = crate::extension_plan::SeriesNormalize::new(
            1,
            TIME_INDEX_COLUMN,
            false,
            Vec::new(),
            LogicalPlan::EmptyRelation(EmptyRelation {
                produce_one_row: false,
                schema: schema.to_dfschema_ref().unwrap(),
            }),
        );
        let plan = RangeManipulate::new(
            1,
            1,
            1,
            1,
            1,
            TIME_INDEX_COLUMN.to_string(),
            vec!["value".to_string()],
            LogicalPlan::Extension(Extension {
                node: Arc::new(normalized),
            }),
        )
        .unwrap();
        let error = datafusion::physical_plan::collect(
            plan.to_execution_plan(input),
            SessionContext::default().task_ctx(),
        )
        .await
        .unwrap_err();
        assert!(error.to_string().contains("timestamp payload overflow"));
    }

    #[tokio::test]
    async fn pruning_should_keep_time_and_value_columns_for_exec() {
        let schema = Arc::new(Schema::new(vec![
            Field::new(TIME_INDEX_COLUMN, TimestampMillisecondType::DATA_TYPE, true),
            Field::new("value_1", DataType::Float64, true),
            Field::new("value_2", DataType::Float64, true),
            Field::new("path", DataType::Utf8, true),
        ]));
        let df_schema = schema.clone().to_dfschema_ref().unwrap();
        let input = LogicalPlan::EmptyRelation(EmptyRelation {
            produce_one_row: false,
            schema: df_schema,
        });
        let plan = RangeManipulate::new(
            0,
            310_000,
            30_000,
            0,
            90_000,
            TIME_INDEX_COLUMN.to_string(),
            vec!["value_1".to_string(), "value_2".to_string()],
            input,
        )
        .unwrap();

        // Simulate a parent projection requesting only the `path` column.
        let output_columns = [3usize];
        let required = plan.necessary_children_exprs(&output_columns).unwrap();
        let required = &required[0];
        assert_eq!(required.as_slice(), &[0, 1, 2, 3]);

        let timestamp_column = Arc::new(TimestampMillisecondArray::from(vec![
            0, 30_000, 60_000, 90_000, 120_000, // every 30s
            180_000, 240_000, // every 60s
            241_000, 271_000, 291_000, // others
        ])) as _;
        let field_column: ArrayRef = Arc::new(Float64Array::from(vec![1.0; 10])) as _;
        let path_column = Arc::new(StringArray::from(vec!["foo"; 10])) as _;
        let input_batch = RecordBatch::try_new(
            schema,
            vec![
                timestamp_column,
                field_column.clone(),
                field_column,
                path_column,
            ],
        )
        .unwrap();

        let projected = project_batch(&input_batch, required);
        let projected_schema = projected.schema();
        let memory_exec = Arc::new(DataSourceExec::new(Arc::new(
            MemorySourceConfig::try_new(&[vec![projected]], projected_schema, None).unwrap(),
        )));
        let range_exec = plan.to_execution_plan(memory_exec);
        let session_context = SessionContext::default();
        let output_batches =
            datafusion::physical_plan::collect(range_exec, session_context.task_ctx())
                .await
                .unwrap();
        assert_eq!(output_batches.len(), 1);

        let output_batch = &output_batches[0];
        let path = output_batch
            .column(3)
            .as_any()
            .downcast_ref::<StringArray>()
            .unwrap();
        assert!(path.iter().all(|v| v == Some("foo")));

        // Simulate the pre-fix pruning behavior: omit the timestamp/value columns from the child.
        let broken_required = [3usize];
        let broken = project_batch(&input_batch, &broken_required);
        let broken_schema = broken.schema();
        let broken_exec = Arc::new(DataSourceExec::new(Arc::new(
            MemorySourceConfig::try_new(&[vec![broken]], broken_schema, None).unwrap(),
        )));
        let broken_range_exec = plan.to_execution_plan(broken_exec);
        let session_context = SessionContext::default();
        let broken_result = std::panic::AssertUnwindSafe(async {
            datafusion::physical_plan::collect(broken_range_exec, session_context.task_ctx()).await
        })
        .catch_unwind()
        .await;
        assert!(broken_result.is_err());
    }

    #[tokio::test]
    async fn interval_30s_range_90s() {
        let expected = String::from(
            "PrimitiveArray<Timestamp(ms)>\n[\n  \
                1970-01-01T00:00:00,\n  \
                1970-01-01T00:00:30,\n  \
                1970-01-01T00:01:00,\n  \
                1970-01-01T00:01:30,\n  \
                1970-01-01T00:02:00,\n  \
                1970-01-01T00:02:30,\n  \
                1970-01-01T00:03:00,\n  \
                1970-01-01T00:03:30,\n  \
                1970-01-01T00:04:00,\n  \
                1970-01-01T00:04:30,\n  \
                1970-01-01T00:05:00,\n\
            ]\nRangeArray { \
                base array: PrimitiveArray<Float64>\n[\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n], \
                ranges: [Some(0..1), Some(0..2), Some(0..3), Some(1..4), Some(2..5), Some(3..5), Some(4..6), Some(5..6), Some(5..7), Some(6..8), Some(6..10)] \
            }\nRangeArray { \
                base array: PrimitiveArray<Float64>\n[\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n], \
                ranges: [Some(0..1), Some(0..2), Some(0..3), Some(1..4), Some(2..5), Some(3..5), Some(4..6), Some(5..6), Some(5..7), Some(6..8), Some(6..10)] \
            }\nStringArray\n[\n  \"foo\",\n  \"foo\",\n  \"foo\",\n  \"foo\",\n  \"foo\",\n  \"foo\",\n  \"foo\",\n  \"foo\",\n  \"foo\",\n  \"foo\",\n  \"foo\",\n]\n\
            RangeArray { \
                base array: PrimitiveArray<Timestamp(ms)>\n[\n  1970-01-01T00:00:00,\n  1970-01-01T00:00:30,\n  1970-01-01T00:01:00,\n  1970-01-01T00:01:30,\n  1970-01-01T00:02:00,\n  1970-01-01T00:03:00,\n  1970-01-01T00:04:00,\n  1970-01-01T00:04:01,\n  1970-01-01T00:04:31,\n  1970-01-01T00:04:51,\n], \
                ranges: [Some(0..1), Some(0..2), Some(0..3), Some(1..4), Some(2..5), Some(3..5), Some(4..6), Some(5..6), Some(5..7), Some(6..8), Some(6..10)] \
            }",
        );
        do_normalize_test(0, 310_000, 30_000, 90_000, expected.clone()).await;

        // dump large range
        do_normalize_test(-300000, 310_000, 30_000, 90_000, expected).await;
    }

    #[tokio::test]
    async fn small_empty_range() {
        let expected = String::from(
            "PrimitiveArray<Timestamp(ms)>\n[\n  \
            1970-01-01T00:00:00.001,\n  \
            1970-01-01T00:00:03.001,\n  \
            1970-01-01T00:00:06.001,\n  \
            1970-01-01T00:00:09.001,\n\
        ]\nRangeArray { \
            base array: PrimitiveArray<Float64>\n[\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n], \
            ranges: [Some(0..1), Some(0..0), Some(0..0), Some(0..0)] \
        }\nRangeArray { \
            base array: PrimitiveArray<Float64>\n[\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n  1.0,\n], \
            ranges: [Some(0..1), Some(0..0), Some(0..0), Some(0..0)] \
        }\nStringArray\n[\n  \"foo\",\n  \"foo\",\n  \"foo\",\n  \"foo\",\n]\n\
        RangeArray { \
            base array: PrimitiveArray<Timestamp(ms)>\n[\n  1970-01-01T00:00:00,\n  1970-01-01T00:00:30,\n  1970-01-01T00:01:00,\n  1970-01-01T00:01:30,\n  1970-01-01T00:02:00,\n  1970-01-01T00:03:00,\n  1970-01-01T00:04:00,\n  1970-01-01T00:04:01,\n  1970-01-01T00:04:31,\n  1970-01-01T00:04:51,\n], \
            ranges: [Some(0..1), Some(0..0), Some(0..0), Some(0..0)] \
        }",
        );
        do_normalize_test(1, 10_001, 3_000, 1_000, expected).await;
    }

    #[test]
    fn test_calculate_range_preserves_alignment() {
        // Test case: query starts at timestamp ending in 4000, step is 30s
        // Data starts at different alignment - should preserve query's 4000 pattern
        let schema = Arc::new(Schema::new(vec![Field::new(
            "timestamp",
            TimestampMillisecondType::DATA_TYPE,
            false,
        )]));
        let empty_stream = MemoryStream::try_new(vec![], schema.clone(), None).unwrap();

        let stream = RangeManipulateStream {
            offset: 0,
            start: 1758093274000, // ends in 4000
            end: 1758093334000,   // ends in 4000
            interval: 30000,      // 30s step
            range: 60000,         // 60s lookback
            time_index: 0,
            time_unit: TimeUnit::Millisecond,
            field_columns: vec![],
            aligned_ts_array: Arc::new(TimestampMillisecondArray::from(vec![0i64; 0])),
            output_schema: schema.clone(),
            input: Box::pin(empty_stream),
            metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
            num_series: Count::new(),
        };

        // Create test data with timestamps not aligned to query pattern
        let test_timestamps = vec![
            1758093260000, // ends in 0000 (different alignment)
            1758093290000, // ends in 0000
            1758093320000, // ends in 0000
        ];
        let ts_array = TimestampMillisecondArray::from(test_timestamps);
        let test_schema = Arc::new(Schema::new(vec![Field::new(
            "timestamp",
            TimestampMillisecondType::DATA_TYPE,
            false,
        )]));
        let batch = RecordBatch::try_new(test_schema, vec![Arc::new(ts_array)]).unwrap();

        let (ranges, (start, end)) = stream.calculate_range(&batch).unwrap();

        // Verify the optimized start preserves query alignment (should end in 4000)
        assert_eq!(
            start % 30000,
            1758093274000 % 30000,
            "Optimized start should preserve query alignment pattern"
        );

        // Verify we generate correct number of ranges for the alignment
        let expected_timestamps: Vec<i64> = (start..=end).step_by(30000).collect();
        assert_eq!(ranges.len(), expected_timestamps.len());

        // Verify all generated timestamps maintain the same alignment pattern
        for ts in expected_timestamps {
            assert_eq!(
                ts % 30000,
                1758093274000 % 30000,
                "All timestamps should maintain query alignment pattern"
            );
        }
    }

    #[tokio::test]
    async fn no_intersection_batch_is_skipped_and_stream_continues() {
        let schema = Arc::new(Schema::new(vec![
            Field::new(
                TIME_INDEX_COLUMN,
                TimestampMillisecondType::DATA_TYPE,
                false,
            ),
            Field::new("value", DataType::Float64, false),
        ]));
        let input = LogicalPlan::EmptyRelation(EmptyRelation {
            produce_one_row: false,
            schema: schema.clone().to_dfschema_ref().unwrap(),
        });
        let plan = RangeManipulate::new(
            0,
            50,
            10,
            0,
            1,
            TIME_INDEX_COLUMN.to_string(),
            vec!["value".to_string()],
            input,
        )
        .unwrap();
        let no_intersection = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(TimestampMillisecondArray::from(vec![100])),
                Arc::new(Float64Array::from(vec![1.0])),
            ],
        )
        .unwrap();
        let intersection = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(TimestampMillisecondArray::from(vec![20])),
                Arc::new(Float64Array::from(vec![2.0])),
            ],
        )
        .unwrap();
        let input = Arc::new(DataSourceExec::new(Arc::new(
            MemorySourceConfig::try_new(&[vec![no_intersection, intersection]], schema, None)
                .unwrap(),
        )));

        let batches = datafusion::physical_plan::collect(
            plan.to_execution_plan(input),
            SessionContext::default().task_ctx(),
        )
        .await
        .unwrap();

        assert_eq!(batches.len(), 1);
        assert_eq!(batches[0].num_rows(), 1);
        let timestamps = batches[0]
            .column(0)
            .as_any()
            .downcast_ref::<TimestampMillisecondArray>()
            .unwrap();
        assert_eq!(timestamps.value(0), 20);
    }

    fn calculate_range_for_test(
        query_start: i64,
        query_end: i64,
        interval: i64,
        range: i64,
        timestamps: &[i64],
    ) -> (Vec<(u32, u32)>, (i64, i64)) {
        let schema = Arc::new(Schema::new(vec![Field::new(
            TIME_INDEX_COLUMN,
            TimestampMillisecondType::DATA_TYPE,
            false,
        )]));
        let empty_stream = MemoryStream::try_new(vec![], schema.clone(), None).unwrap();
        let stream = RangeManipulateStream {
            offset: 0,
            start: query_start,
            end: query_end,
            interval,
            range,
            time_index: 0,
            time_unit: TimeUnit::Millisecond,
            field_columns: vec![],
            aligned_ts_array: Arc::new(TimestampMillisecondArray::from(vec![0i64; 0])),
            output_schema: schema.clone(),
            input: Box::pin(empty_stream),
            metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
            num_series: Count::new(),
        };
        let batch = RecordBatch::try_new(
            schema,
            vec![Arc::new(TimestampMillisecondArray::from(
                timestamps.to_vec(),
            ))],
        )
        .unwrap();

        stream.calculate_range(&batch).unwrap()
    }

    #[test]
    fn calculate_range_nullable_interior_timestamp_matches_wide_scan() {
        let schema = Arc::new(Schema::new(vec![Field::new(
            TIME_INDEX_COLUMN,
            DataType::Timestamp(TimeUnit::Second, None),
            true,
        )]));
        let stream = scan_stream_for_test(TimeUnit::Second, 0, 0, 1_000, 1_000, 1_000);
        let timestamps = [0, i64::MAX, 1];
        let array = TimestampSecondArray::new(
            timestamps.to_vec().into(),
            Some(NullBuffer::from(vec![false, false, true])),
        );
        let batch = RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap();

        let (ranges, bounds) = stream.calculate_range(&batch).unwrap();
        let wide = RangeManipulateStream::bench_scan_wide(
            TimeUnit::Second,
            0,
            1_000,
            1_000,
            &timestamps,
            bounds.0,
            bounds.1,
        );
        assert_eq!(bounds, (0, 1_000));
        assert_eq!(ranges, vec![(0, 1), (0, 0)]);
        assert_eq!(ranges, wide);
    }

    #[test]
    fn calculate_range_keeps_query_aligned_tail() {
        let (ranges, bounds) = calculate_range_for_test(4, 94, 30, 15, &[20, 50, 80]);

        assert_eq!(bounds, (34, 94));
        assert_eq!(ranges, vec![(0, 1), (1, 1), (2, 1)]);
    }

    /// Calculates exact offsets directly from the range predicate used by PromQL.
    ///
    /// Input timestamps are sorted and non-null. Interval is positive, range is
    /// nonnegative, and test values are chosen to avoid `i64` overflow.
    fn calculate_range_oracle(
        timestamps: &[i64],
        start: i64,
        end: i64,
        interval: i64,
        range: i64,
    ) -> Vec<(u32, u32)> {
        // Match `calculate_range`'s explicit empty-input early return.
        if timestamps.is_empty() || start > end {
            return vec![];
        }

        (start..=end)
            .step_by(interval as usize)
            .map(|curr| {
                let mut offset = None;
                let mut length = 0;
                for (index, &ts) in timestamps.iter().enumerate() {
                    if ts > curr - range && ts <= curr {
                        offset.get_or_insert(index);
                        length += 1;
                    }
                }
                (offset.unwrap_or(0) as u32, length)
            })
            .collect()
    }

    #[test]
    fn calculate_range_characterizes_returned_bounds() {
        let cases = [
            (
                "positive non-aligned last timestamp plus range",
                0,
                100,
                10,
                9,
                vec![13, 26],
                (20, 30),
            ),
            (
                "negative non-aligned last timestamp plus range",
                -50,
                50,
                10,
                9,
                vec![-37, -26],
                (-30, -20),
            ),
            (
                "query alignment not based on epoch",
                4,
                94,
                30,
                15,
                vec![20, 50, 80],
                (34, 94),
            ),
            ("leading data", 0, 100, 10, 10, vec![-10, 15], (0, 20)),
            (
                "trailing data past query end",
                0,
                100,
                10,
                10,
                vec![35, 45, 110],
                (40, 100),
            ),
            (
                "optimized start after query end",
                0,
                50,
                10,
                0,
                vec![100],
                (0, 50),
            ),
        ];

        for (name, query_start, query_end, interval, range, timestamps, expected_bounds) in cases {
            let (_, bounds) =
                calculate_range_for_test(query_start, query_end, interval, range, &timestamps);
            assert_eq!(bounds, expected_bounds, "{name}");
        }
    }

    #[test]
    fn calculate_range_keeps_extreme_range_tail() {
        let (ranges, bounds) =
            calculate_range_for_test(i64::MAX - 1, i64::MAX, 1, i64::MAX, &[i64::MAX]);

        assert_eq!(bounds, (i64::MAX, i64::MAX));
        assert_eq!(ranges, vec![(0, 1)]);
    }

    #[test]
    fn calculate_range_matches_bruteforce_oracle_for_deterministic_cases() {
        let cases = vec![
            (
                "duplicate lower and upper bounds",
                vec![0, 10, 10, 20, 20, 30],
                10,
                20,
                10,
                10,
            ),
            (
                "zero range excludes duplicates at current timestamp",
                vec![10, 10, 10],
                10,
                10,
                1,
                0,
            ),
            (
                "consecutive nonempty empty nonempty ranges",
                vec![10, 30],
                10,
                30,
                10,
                5,
            ),
            ("step smaller than range", vec![0, 4, 8, 12], 0, 12, 3, 5),
            ("step equal to range", vec![0, 5, 10, 15], 0, 15, 5, 5),
            ("step greater than range", vec![0, 7, 14, 21], 0, 21, 7, 3),
            (
                "negative sparse/tail timestamps",
                vec![-30, -20, -10, 0],
                -25,
                5,
                5,
                7,
            ),
            ("one sample", vec![42], 0, 100, 10, 15),
            ("empty input", vec![], -20, 20, 5, 10),
        ];

        for (name, timestamps, query_start, query_end, interval, range) in cases {
            let (actual, (start, end)) =
                calculate_range_for_test(query_start, query_end, interval, range, &timestamps);
            let expected = calculate_range_oracle(&timestamps, start, end, interval, range);
            assert_eq!(actual, expected, "{name}");
        }
    }

    #[test]
    fn calculate_range_positive_time_translated_regression() {
        let timestamps = [0, 10, 20, 30];
        let expected = vec![(0, 1), (1, 1), (1, 1), (2, 1), (2, 1), (3, 1), (3, 1)];
        let (actual, bounds) = calculate_range_for_test(5, 35, 5, 7, &timestamps);

        assert_eq!(bounds, (5, 35));
        assert_eq!(
            calculate_range_oracle(&timestamps, bounds.0, bounds.1, 5, 7),
            expected
        );
        assert_eq!(actual, expected);
    }

    #[test]
    fn calculate_range_matches_oracle_for_dense_positive_time_windows() {
        let timestamps = (0..=3_600).step_by(15).collect::<Vec<i64>>();

        for (name, range) in [
            ("one minute", 60),
            ("five minutes", 300),
            ("one hour", 3_600),
        ] {
            let (actual, (start, end)) = calculate_range_for_test(0, 3_600, 15, range, &timestamps);
            assert_eq!((start, end), (0, 3_600), "{name} bounds");
            assert_eq!(
                actual,
                calculate_range_oracle(&timestamps, start, end, 15, range),
                "{name} window"
            );
        }
    }

    #[derive(Clone, Copy)]
    struct TinyPrng(u64);

    impl TinyPrng {
        fn next_u64(&mut self) -> u64 {
            self.0 ^= self.0 << 13;
            self.0 ^= self.0 >> 7;
            self.0 ^= self.0 << 17;
            self.0
        }

        fn next_i64(&mut self, min: i64, max: i64) -> i64 {
            min + (self.next_u64() % (max - min + 1) as u64) as i64
        }
    }

    #[test]
    fn calculate_range_matches_bruteforce_oracle_for_seeded_matrix() {
        let mut prng = TinyPrng(0x5eed_cafe_f00d_baad);

        for case in 0..512 {
            let interval = prng.next_i64(1, 11);
            let range = prng.next_i64(0, 25);
            let query_start = prng.next_i64(-200, 200);
            let query_end = query_start + interval * prng.next_i64(0, 20);
            let mut timestamps = Vec::new();
            let mut timestamp = prng.next_i64(-250, 250);
            for _ in 0..prng.next_i64(0, 20) {
                timestamp += prng.next_i64(0, 7);
                timestamps.push(timestamp);
            }

            let (actual, (start, end)) =
                calculate_range_for_test(query_start, query_end, interval, range, &timestamps);
            let expected = calculate_range_oracle(&timestamps, start, end, interval, range);
            let expected = if actual.is_empty() && !expected.is_empty() {
                assert!(
                    expected.iter().all(|(_, len)| *len == 0),
                    "case={case}, timestamps={timestamps:?}, query=({query_start}, {query_end}), \
                     interval={interval}, range={range}, bounds=({start}, {end}): \
                     no-intersection output must have no selected samples"
                );
                vec![]
            } else {
                expected
            };
            assert_eq!(
                actual, expected,
                "case={case}, timestamps={timestamps:?}, query=({query_start}, {query_end}), \
                 interval={interval}, range={range}, bounds=({start}, {end})"
            );
        }
    }

    /// Nanoseconds per native tick, as kept by the narrow carrier.
    fn native_scale(time_unit: TimeUnit) -> i64 {
        i64::try_from(nanoseconds_per_native_tick(time_unit)).unwrap()
    }

    /// Shifted nanoseconds of a sample, computed in the exact `i128` carrier.
    fn shifted_nanoseconds(time_unit: TimeUnit, offset: Millisecond, timestamp: i64) -> i128 {
        (timestamp as i128) * nanoseconds_per_native_tick(time_unit) + (offset as i128) * 1_000_000
    }

    /// Shifted milliseconds of the instant grid of a sample, flooring like `calculate_range`.
    fn shifted_floor_millis(
        time_unit: TimeUnit,
        offset: Millisecond,
        timestamp: i64,
    ) -> Millisecond {
        Millisecond::try_from(
            shifted_nanoseconds(time_unit, offset, timestamp).div_euclid(1_000_000),
        )
        .unwrap()
    }

    /// Shifted milliseconds kept by the range payload of a sample.
    ///
    /// The payload conversion truncates toward zero, so samples below the epoch may keep
    /// the millisecond above their own native tick.
    fn shifted_millis(time_unit: TimeUnit, offset: Millisecond, timestamp: i64) -> Millisecond {
        Millisecond::try_from(shifted_nanoseconds(time_unit, offset, timestamp) / 1_000_000)
            .unwrap()
    }

    /// Stream whose cursor scan the differential tests below drive directly.
    fn scan_stream_for_test(
        time_unit: TimeUnit,
        offset: Millisecond,
        start: Millisecond,
        end: Millisecond,
        interval: Millisecond,
        range: Millisecond,
    ) -> RangeManipulateStream {
        let schema = Arc::new(Schema::new(vec![Field::new(
            TIME_INDEX_COLUMN,
            DataType::Timestamp(time_unit, None),
            false,
        )]));
        let empty_stream = MemoryStream::try_new(vec![], schema.clone(), None).unwrap();
        RangeManipulateStream {
            offset,
            start,
            end,
            interval,
            range,
            time_index: 0,
            time_unit,
            field_columns: vec![],
            aligned_ts_array: Arc::new(TimestampMillisecondArray::from(vec![0i64; 0])),
            output_schema: schema,
            input: Box::pin(empty_stream),
            metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
            num_series: Count::new(),
        }
    }

    /// Exact expectations of the cursor scan, formulated outside the production loops.
    ///
    /// Every evaluation instant keeps the samples of `(instant - range, instant]` in
    /// shifted nanoseconds: the lower boundary is exclusive, the upper one inclusive, and
    /// native ticks below a millisecond stay distinct.
    fn scan_oracle(
        time_unit: TimeUnit,
        offset: Millisecond,
        range: Millisecond,
        timestamps: &[i64],
        start: Millisecond,
        end: Millisecond,
        interval: Millisecond,
    ) -> Vec<(u32, u32)> {
        let window_ns = (range as i128) * 1_000_000;
        let shifted = timestamps
            .iter()
            .map(|timestamp| shifted_nanoseconds(time_unit, offset, *timestamp))
            .collect::<Vec<_>>();
        (start..=end)
            .step_by(interval as usize)
            .map(|instant| {
                let instant_ns = (instant as i128) * 1_000_000;
                let mut offset = None;
                let mut length = 0u32;
                for (index, timestamp) in shifted.iter().enumerate() {
                    if *timestamp > instant_ns - window_ns && *timestamp <= instant_ns {
                        offset.get_or_insert(index as u32);
                        length += 1;
                    }
                }
                (offset.unwrap_or(0), length)
            })
            .collect()
    }

    /// Runs the cursor scan of one input through both carriers.
    fn scan_both_carriers(
        stream: &RangeManipulateStream,
        timestamps: &[i64],
        start: Millisecond,
        end: Millisecond,
    ) -> (Option<RangeSelections>, RangeSelections) {
        let exact = scan_ranges::<i128>(
            stream.interval,
            timestamps,
            start,
            end,
            &ScanNanoseconds::<i128>::wide(stream.time_unit, stream.offset, stream.range),
        );
        let narrow = narrow_scan(
            stream.time_unit,
            stream.offset,
            stream.range,
            timestamps,
            start,
            end,
        )
        .map(|scan| scan_ranges::<i64>(stream.interval, timestamps, start, end, &scan));
        // The benchmark hooks must measure the very scans the operator runs.
        assert_eq!(
            RangeManipulateStream::bench_scan_wide(
                stream.time_unit,
                stream.offset,
                stream.range,
                stream.interval,
                timestamps,
                start,
                end,
            ),
            exact,
            "the i128 benchmark hook diverged from the cursor scan"
        );
        if let Some(narrow) = &narrow {
            assert_eq!(
                RangeManipulateStream::bench_scan_narrow(
                    stream.time_unit,
                    stream.offset,
                    stream.range,
                    stream.interval,
                    timestamps,
                    start,
                    end,
                ),
                Some(narrow.clone()),
                "the i64 benchmark hook diverged from the cursor scan"
            );
        } else {
            assert_eq!(
                RangeManipulateStream::bench_scan_narrow(
                    stream.time_unit,
                    stream.offset,
                    stream.range,
                    stream.interval,
                    timestamps,
                    start,
                    end,
                ),
                None,
                "the i64 benchmark hook must report the fallback"
            );
        }
        (narrow, exact)
    }

    /// Asserts the narrow carrier was taken and matches the exact scan evaluation by evaluation.
    fn assert_narrow_carrier_agrees(
        narrow: Option<Vec<(u32, u32)>>,
        exact: &[(u32, u32)],
        context: &str,
    ) {
        let Some(narrow) = narrow else {
            panic!("the i64 carrier was expected, but narrow_scan rejected the input: {context}");
        };
        assert_eq!(
            narrow.len(),
            exact.len(),
            "carriers disagree on the number of evaluations: {context}"
        );
        for (evaluation, (narrow_range, exact_range)) in narrow.iter().zip(exact).enumerate() {
            assert_eq!(
                narrow_range, exact_range,
                "carriers disagree on evaluation {evaluation}: {context}"
            );
        }
    }

    /// Replays the cursor scan in checked `i64` arithmetic.
    ///
    /// Reports whether the narrow carrier overflows on an operation the scan evaluates,
    /// which is what makes the `i128` fallback unavoidable for an input.
    fn narrow_carrier_overflows(
        time_unit: TimeUnit,
        offset: Millisecond,
        range: Millisecond,
        timestamps: &[i64],
        start: Millisecond,
        end: Millisecond,
        interval: Millisecond,
    ) -> bool {
        let Ok(scale) = i64::try_from(nanoseconds_per_native_tick(time_unit)) else {
            return true;
        };
        let Some(offset_ns) = offset.checked_mul(1_000_000) else {
            return true;
        };
        let Some(window_ns) = range.checked_mul(1_000_000) else {
            return true;
        };
        let shift = |timestamp: i64| timestamp.checked_mul(scale)?.checked_add(offset_ns);
        let mut left = 0usize;
        let mut right = 0usize;
        for instant in (start..=end).step_by(interval as usize) {
            let Some(instant_ns) = instant.checked_mul(1_000_000) else {
                return true;
            };
            let Some(window_start) = instant_ns.checked_sub(window_ns) else {
                return true;
            };
            while left < timestamps.len() {
                let Some(shifted) = shift(timestamps[left]) else {
                    return true;
                };
                if shifted <= window_start {
                    left += 1;
                } else {
                    break;
                }
            }
            right = right.max(left);
            while right < timestamps.len() {
                let Some(shifted) = shift(timestamps[right]) else {
                    return true;
                };
                if shifted <= instant_ns {
                    right += 1;
                } else {
                    break;
                }
            }
        }
        false
    }

    /// The checked `i64` operation that makes [`narrow_scan`] reject an input, if any.
    ///
    /// Mirrors the rejection criteria of [`narrow_scan`] so the deterministic cases below
    /// can pin which boundary drives the `i128` fallback.
    fn narrow_carrier_rejection(
        time_unit: TimeUnit,
        offset: Millisecond,
        range: Millisecond,
        timestamps: &[i64],
        start: Millisecond,
        end: Millisecond,
    ) -> Option<&'static str> {
        let Ok(scale) = i64::try_from(nanoseconds_per_native_tick(time_unit)) else {
            return Some("nanoseconds per native tick do not fit i64");
        };
        let Some(offset_ns) = offset.checked_mul(1_000_000) else {
            return Some("offset in nanoseconds does not fit i64");
        };
        let Some(window_ns) = range.checked_mul(1_000_000) else {
            return Some("range in nanoseconds does not fit i64");
        };
        let shift = |timestamp: i64| timestamp.checked_mul(scale)?.checked_add(offset_ns);
        let (Some(first), Some(last)) = (timestamps.first(), timestamps.last()) else {
            return Some("empty batch has no endpoints");
        };
        if shift(*first).is_none() || shift(*last).is_none() {
            return Some("sample nanoseconds do not fit i64");
        }
        if start
            .checked_mul(1_000_000)
            .and_then(|nanoseconds| nanoseconds.checked_sub(window_ns))
            .is_none()
        {
            return Some("start bound in nanoseconds does not fit i64");
        }
        if end
            .checked_mul(1_000_000)
            .and_then(|nanoseconds| nanoseconds.checked_sub(window_ns))
            .is_none()
        {
            return Some("end bound in nanoseconds do not fit i64");
        }
        None
    }

    #[test]
    fn scan_carriers_agree_on_range_boundaries_within_one_millisecond() {
        // Membership is `(lower, upper]` on native ticks: the sample exactly on the lower
        // boundary is dropped, the sample exactly on the upper boundary is kept, and the
        // three samples sharing one millisecond payload stay distinct and survive together.
        let cases = [
            (
                TimeUnit::Microsecond,
                vec![1_000_000, 1_000_001, 1_000_999, 1_001_000, 1_001_001],
            ),
            (
                TimeUnit::Nanosecond,
                vec![
                    1_000_000_000,
                    1_000_000_001,
                    1_000_000_999,
                    1_001_000_000,
                    1_001_000_001,
                ],
            ),
        ];

        for (time_unit, timestamps) in cases {
            let (start, end, interval, range) = (1_001, 1_001, 1, 1);
            let context = format!("unit={time_unit:?}, timestamps={timestamps:?}");
            let stream = scan_stream_for_test(time_unit, 0, start, end, interval, range);

            let oracle = scan_oracle(time_unit, 0, range, &timestamps, start, end, interval);
            assert_eq!(
                oracle,
                vec![(1, 3)],
                "{context}: (lower, upper] expectations"
            );

            let (narrow, exact) = scan_both_carriers(&stream, &timestamps, start, end);
            assert_eq!(exact, oracle, "{context}: i128 carrier must stay exact");
            assert_narrow_carrier_agrees(narrow, &exact, &context);

            // Two distinct native ticks below the upper boundary share the millisecond
            // payload `1_000`; they must not be deduplicated by going through that ABI.
            let payloads = timestamps[1..4]
                .iter()
                .map(|timestamp| shifted_millis(time_unit, 0, *timestamp))
                .collect::<Vec<_>>();
            assert_eq!(
                payloads,
                vec![1_000, 1_000, 1_001],
                "{context}: kept samples in one millisecond keep both native ticks"
            );
            assert_ne!(
                timestamps[1], timestamps[2],
                "{context}: the two ticks sharing a payload must stay distinct"
            );
        }
    }

    #[test]
    fn scan_falls_back_to_i128_when_only_the_sample_tail_overflows() {
        // `narrow_scan` inspects the sorted endpoints, so a batch whose prefix fits the
        // `i64` carrier but whose tail does not has to take the exact carrier. The tail
        // sample stays outside every window, so the selections keep the prefix ones.
        struct Case {
            time_unit: TimeUnit,
            narrow_safe: Vec<i64>,
            overflowing: i64,
            start: Millisecond,
            end: Millisecond,
            interval: Millisecond,
            range: Millisecond,
        }

        let cases = [
            Case {
                time_unit: TimeUnit::Second,
                narrow_safe: vec![9_223_372_034, 9_223_372_035, 9_223_372_036],
                overflowing: 9_223_372_037,
                start: 9_223_372_034_000,
                end: 9_223_372_036_000,
                interval: 1_000,
                range: 1_000,
            },
            Case {
                time_unit: TimeUnit::Millisecond,
                narrow_safe: vec![9_223_372_036_853, 9_223_372_036_854],
                overflowing: 9_223_372_036_855,
                start: 9_223_372_036_853,
                end: 9_223_372_036_854,
                interval: 1,
                range: 1,
            },
            Case {
                time_unit: TimeUnit::Microsecond,
                narrow_safe: vec![9_223_372_036_853_999, 9_223_372_036_854_000],
                overflowing: 9_223_372_036_854_776,
                start: 9_223_372_036_854,
                end: 9_223_372_036_854,
                interval: 1,
                range: 1,
            },
        ];

        for case in cases {
            let Case {
                time_unit,
                narrow_safe,
                overflowing,
                start,
                end,
                interval,
                range,
            } = case;
            let context = format!("unit={time_unit:?}, timestamps={narrow_safe:?} + {overflowing}");
            let scale = native_scale(time_unit);
            // The prefix fits the narrow carrier, the tail sample does not...
            assert!(
                narrow_safe
                    .iter()
                    .all(|timestamp| timestamp.checked_mul(scale).is_some()),
                "{context}: prefix samples must fit i64 nanoseconds"
            );
            assert!(
                overflowing.checked_mul(scale).is_none(),
                "{context}: the tail sample must overflow i64 nanoseconds"
            );
            // ...even though the millisecond payload of the tail stays legal.
            let payload = shifted_millis(time_unit, 0, overflowing);
            assert!(
                payload.abs() < i64::MAX / 2,
                "{context}: the shifted payload {payload} must stay legal"
            );

            let stream = scan_stream_for_test(time_unit, 0, start, end, interval, range);

            let (narrow, prefix_exact) = scan_both_carriers(&stream, &narrow_safe, start, end);
            let prefix_oracle =
                scan_oracle(time_unit, 0, range, &narrow_safe, start, end, interval);
            assert_eq!(
                prefix_exact, prefix_oracle,
                "{context}: prefix i128 carrier"
            );
            assert_narrow_carrier_agrees(narrow, &prefix_exact, &format!("{context} prefix"));

            let mut timestamps = narrow_safe.clone();
            timestamps.push(overflowing);
            assert_eq!(
                narrow_carrier_rejection(time_unit, 0, range, &timestamps, start, end),
                Some("sample nanoseconds do not fit i64"),
                "{context}: the overflowing endpoint must reject the narrow carrier"
            );
            assert!(
                narrow_carrier_overflows(time_unit, 0, range, &timestamps, start, end, interval),
                "{context}: the scan evaluates the overflowing tail"
            );

            let (narrow, exact) = scan_both_carriers(&stream, &timestamps, start, end);
            assert!(
                narrow.is_none(),
                "{context}: the i128 carrier must be taken"
            );
            let oracle = scan_oracle(time_unit, 0, range, &timestamps, start, end, interval);
            assert_eq!(exact, oracle, "{context}: exact carrier on the full batch");
            assert_eq!(
                exact, prefix_exact,
                "{context}: the out-of-window tail keeps the prefix selections"
            );
        }
    }

    #[test]
    fn scan_carriers_split_at_the_nanosecond_boundaries() {
        // Nanoseconds are not scaled, so the `i64` carrier is rejected by the query bounds:
        // below the millisecond of `i64::MIN` the start bound underflows, while a batch near
        // `i64::MAX` still fits. Both keep a legal millisecond payload.
        let low = i64::MIN;
        let low_floor = shifted_floor_millis(TimeUnit::Nanosecond, 0, low);
        assert_eq!(low_floor, -9_223_372_036_855);
        assert!(low_floor.checked_mul(1_000_000).is_none());
        assert_eq!(
            shifted_millis(TimeUnit::Nanosecond, 0, low),
            -9_223_372_036_854
        );

        let timestamps = vec![low, low + 1_000_000];
        let (start, end, interval, range) = (low_floor, low_floor + 1, 1, 1);
        let context = format!("nanosecond floor, timestamps={timestamps:?}");
        let stream = scan_stream_for_test(TimeUnit::Nanosecond, 0, start, end, interval, range);
        assert_eq!(
            narrow_carrier_rejection(TimeUnit::Nanosecond, 0, range, &timestamps, start, end),
            Some("start bound in nanoseconds does not fit i64"),
            "{context}"
        );
        let (narrow, exact) = scan_both_carriers(&stream, &timestamps, start, end);
        assert!(
            narrow.is_none(),
            "{context}: the i128 carrier must be taken"
        );
        let oracle = scan_oracle(
            TimeUnit::Nanosecond,
            0,
            range,
            &timestamps,
            start,
            end,
            interval,
        );
        assert_eq!(exact, oracle, "{context}: i128 carrier must stay exact");

        // The ceiling side still fits the narrow carrier, so both carriers have to agree.
        let high = i64::MAX - 2_000_000;
        let high_floor = shifted_floor_millis(TimeUnit::Nanosecond, 0, high);
        assert!(high_floor.checked_mul(1_000_000).is_some());
        let timestamps = vec![high, high + 1_000_000, high + 2_000_000];
        let (start, end, interval, range) = (high_floor, high_floor + 2, 1, 1);
        let context = format!("nanosecond ceiling, timestamps={timestamps:?}");
        let stream = scan_stream_for_test(TimeUnit::Nanosecond, 0, start, end, interval, range);
        assert_eq!(
            narrow_carrier_rejection(TimeUnit::Nanosecond, 0, range, &timestamps, start, end),
            None,
            "{context}: the ceiling batch still fits the narrow carrier"
        );
        let (narrow, exact) = scan_both_carriers(&stream, &timestamps, start, end);
        let oracle = scan_oracle(
            TimeUnit::Nanosecond,
            0,
            range,
            &timestamps,
            start,
            end,
            interval,
        );
        assert_eq!(exact, oracle, "{context}: i128 carrier must stay exact");
        assert_narrow_carrier_agrees(narrow, &exact, &context);
    }

    #[test]
    fn scan_carriers_agree_for_seeded_native_precision_matrix() {
        // Seeded matrix over all four native precisions: positive and negative offsets,
        // negative times, sub-millisecond samples sharing a millisecond payload, and native
        // extremes whose intermediate product overflows while the shifted millisecond stays
        // legal.
        let units = [
            TimeUnit::Second,
            TimeUnit::Millisecond,
            TimeUnit::Microsecond,
            TimeUnit::Nanosecond,
        ];
        // `narrow_scan` rejects the two large offsets: `offset * 1ms` overflows `i64`.
        let offsets = [0, 1, -1, 250, -250, 9_223_372_036_855, -9_223_372_036_855];
        let extremes = |time_unit: TimeUnit| match time_unit {
            TimeUnit::Second => [-9_223_372_037, 9_223_372_037],
            TimeUnit::Millisecond => [-9_223_372_036_855, 9_223_372_036_855],
            TimeUnit::Microsecond => [-9_223_372_036_854_776, 9_223_372_036_854_776],
            TimeUnit::Nanosecond => [i64::MIN, i64::MAX],
        };
        let mut prng = TinyPrng(0x5ca1_ab1e_d1ff_0001);
        let mut narrow_cases = 0usize;
        let mut wide_cases = 0usize;
        let mut skipped = 0usize;

        for case in 0..256usize {
            let time_unit = units[case % units.len()];
            let offset = offsets[(case / units.len()) % offsets.len()];
            let [low, high] = extremes(time_unit);
            let anchor = match prng.next_i64(0, 3) {
                0 => prng.next_i64(-2_000, 2_000),
                1 => -prng.next_i64(1, 2_000),
                2 => low,
                _ => high,
            };
            // Samples stay within a few milliseconds, with a native granularity small
            // enough that sub-millisecond samples share a millisecond payload.
            let step = match time_unit {
                TimeUnit::Second => 1,
                TimeUnit::Millisecond => 1,
                TimeUnit::Microsecond => 1 + prng.next_i64(0, 9),
                TimeUnit::Nanosecond => 1 + prng.next_i64(0, 999_999),
            };
            let count = 1 + prng.next_i64(0, 4);
            let ascending = anchor != high;
            let mut timestamps = (0..count)
                .map(|index| {
                    let delta = index * step;
                    if ascending {
                        anchor + delta
                    } else {
                        anchor - delta
                    }
                })
                .collect::<Vec<_>>();
            if !ascending {
                timestamps.reverse();
            }
            // The review shape keeps a legal millisecond payload; skip the extremes that the
            // large offsets push out of the millisecond space.
            let nanoseconds = timestamps
                .iter()
                .map(|timestamp| shifted_nanoseconds(time_unit, offset, *timestamp))
                .collect::<Vec<_>>();
            if nanoseconds
                .iter()
                .any(|value| i64::try_from(value / 1_000_000).is_err())
            {
                skipped += 1;
                continue;
            }

            let first_millis = Millisecond::try_from(nanoseconds[0].div_euclid(1_000_000)).unwrap();
            let interval = prng.next_i64(1, 4);
            let points = 1 + prng.next_i64(0, 4);
            let start = first_millis - interval;
            let end = start + interval * (points - 1);
            let range = prng.next_i64(0, 3) * interval;
            let context = format!(
                "case={case}, unit={time_unit:?}, offset={offset}, start={start}, end={end}, \
                 interval={interval}, range={range}, timestamps={timestamps:?}"
            );

            let stream = scan_stream_for_test(time_unit, offset, start, end, interval, range);
            let oracle = scan_oracle(time_unit, offset, range, &timestamps, start, end, interval);
            let (narrow, exact) = scan_both_carriers(&stream, &timestamps, start, end);
            assert_eq!(exact, oracle, "{context}: i128 carrier must stay exact");

            match (
                &narrow,
                narrow_carrier_rejection(time_unit, offset, range, &timestamps, start, end),
            ) {
                (Some(_), None) => {
                    narrow_cases += 1;
                    assert!(
                        !narrow_carrier_overflows(
                            time_unit,
                            offset,
                            range,
                            &timestamps,
                            start,
                            end,
                            interval
                        ),
                        "{context}: every narrow operation must fit i64"
                    );
                    assert_narrow_carrier_agrees(narrow, &exact, &context);
                }
                (None, Some(_reason)) => {
                    wide_cases += 1;
                }
                (narrow, rejection) => panic!(
                    "{context}: narrow_scan and its mirror disagree, narrow={}, rejection={rejection:?}",
                    narrow.is_some()
                ),
            }
        }

        assert!(
            narrow_cases > 0 && wide_cases > 0,
            "both carriers must be covered: {narrow_cases} narrow, {wide_cases} wide \
             ({skipped} skipped)"
        );
        assert!(
            skipped < 64,
            "the generator skipped too many cases: {skipped}"
        );
    }

    /// Builds the single-column timestamp batch the cursor scans consume.
    fn timestamp_batch(time_unit: TimeUnit, timestamps: &[i64]) -> RecordBatch {
        let schema = Arc::new(Schema::new(vec![Field::new(
            TIME_INDEX_COLUMN,
            DataType::Timestamp(time_unit, None),
            false,
        )]));
        let array: ArrayRef = match time_unit {
            TimeUnit::Second => Arc::new(TimestampSecondArray::from(timestamps.to_vec())),
            TimeUnit::Millisecond => Arc::new(TimestampMillisecondArray::from(timestamps.to_vec())),
            TimeUnit::Microsecond => Arc::new(TimestampMicrosecondArray::from(timestamps.to_vec())),
            TimeUnit::Nanosecond => Arc::new(TimestampNanosecondArray::from(timestamps.to_vec())),
        };
        RecordBatch::try_new(schema, vec![array]).unwrap()
    }

    #[test]
    fn calculate_range_dispatches_carriers_and_agrees_pointwise() {
        // Production dispatch, not the hooks: `calculate_range` must scan through the narrow
        // carrier while the batch fits and through the exact one once the shifted tail does
        // not, and both must select the same ranges for the bounds it derives itself.
        let narrow_timestamps = [0i64, 15_000, 30_000];
        let wide_timestamps = [9_223_372_036_853, 9_223_372_036_854, 9_223_372_036_855];
        let cases: [(&str, &[i64], Millisecond, Millisecond, bool); 2] = [
            ("narrow carrier", &narrow_timestamps, 0, 30_000, true),
            (
                "i128 fallback",
                &wide_timestamps,
                wide_timestamps[0],
                wide_timestamps[1],
                false,
            ),
        ];

        for (name, timestamps, query_start, query_end, expect_narrow) in cases {
            let (interval, range) = if expect_narrow {
                (15_000, 15_000)
            } else {
                (1, 1)
            };
            // `scan_stream_for_test` panics on an empty batch when `narrow_scan` is asked for
            // the endpoints, so the batch below is never empty.
            let stream = scan_stream_for_test(
                TimeUnit::Millisecond,
                0,
                query_start,
                query_end,
                interval,
                range,
            );
            let batch = timestamp_batch(TimeUnit::Millisecond, timestamps);
            let (ranges, (start, end)) = stream.calculate_range(&batch).unwrap();

            let oracle = scan_oracle(
                TimeUnit::Millisecond,
                0,
                range,
                timestamps,
                start,
                end,
                interval,
            );
            assert_eq!(ranges, oracle, "{name}: calculate_range pointwise");

            let rejection =
                narrow_carrier_rejection(TimeUnit::Millisecond, 0, range, timestamps, start, end);
            assert_eq!(
                rejection.is_none(),
                expect_narrow,
                "{name}: carrier choice by narrow_scan, rejection={rejection:?}"
            );
            let (narrow, exact) = scan_both_carriers(&stream, timestamps, start, end);
            assert_eq!(exact, ranges, "{name}: i128 carrier is the oracle");
            assert_eq!(narrow.is_some(), expect_narrow, "{name}: narrow carrier");
            assert_eq!(
                RangeManipulateStream::bench_scan_wide(
                    stream.time_unit,
                    stream.offset,
                    stream.range,
                    interval,
                    timestamps,
                    start,
                    end,
                ),
                ranges,
                "{name}: the benchmark hook matches calculate_range"
            );
        }
    }

    /// Stream whose single millisecond timestamp column is folded into the range payload.
    fn payload_stream(
        offset: Millisecond,
        start: Millisecond,
        end: Millisecond,
        interval: Millisecond,
        range: Millisecond,
    ) -> RangeManipulateStream {
        let payload_type = DataType::Timestamp(TimeUnit::Millisecond, None);
        let schema = Arc::new(Schema::new(vec![
            Field::new(TIME_INDEX_COLUMN, payload_type.clone(), false),
            Field::new(
                RangeManipulate::build_timestamp_range_name(TIME_INDEX_COLUMN),
                RangeArray::convert_data_type(payload_type),
                false,
            ),
        ]));
        let mut stream =
            scan_stream_for_test(TimeUnit::Millisecond, offset, start, end, interval, range);
        stream.output_schema = schema;
        stream
    }

    /// Builds the single-column millisecond timestamp batch the payload tests consume.
    fn timestamp_batch_with_timezone(timezone: Option<&str>, timestamps: &[i64]) -> RecordBatch {
        let schema = Arc::new(Schema::new(vec![Field::new(
            TIME_INDEX_COLUMN,
            DataType::Timestamp(TimeUnit::Millisecond, timezone.map(Arc::from)),
            false,
        )]));
        let array: ArrayRef = Arc::new(
            TimestampMillisecondArray::from(timestamps.to_vec()).with_timezone_opt(timezone),
        );
        RecordBatch::try_new(schema, vec![array]).unwrap()
    }

    /// Millisecond payload values of the range timestamp column of a manipulated batch.
    fn payload_millis(batch: &RecordBatch) -> Vec<i64> {
        let dict = batch
            .column(batch.num_columns() - 1)
            .as_any()
            .downcast_ref::<DictionaryArray<Int64Type>>()
            .expect("the timestamp range payload must be a folded dictionary")
            .clone();
        let payload = RangeArray::try_new(dict).unwrap();
        assert_eq!(
            payload.values().data_type(),
            &DataType::Timestamp(TimeUnit::Millisecond, None),
            "the payload value type must agree with the output schema"
        );
        payload
            .values()
            .as_any()
            .downcast_ref::<TimestampMillisecondArray>()
            .unwrap()
            .values()
            .to_vec()
    }

    #[test]
    fn manipulate_timezone_time_index_skips_zero_copy_payload() {
        // Oracle-verified regression: `Timestamp(Millisecond, Some(tz))` passes a
        // `time_unit == Millisecond` check but is not the payload ABI. Cloning such a column
        // would stamp a dictionary value type disagreeing with the output schema and
        // `RecordBatch::try_new` would fail, so the reuse must be guarded on the data type.
        let timestamps = [0i64, 1_000, 2_000, 3_000];
        let stream = payload_stream(0, 0, 3_000, 1_000, 1_000);

        // With offset 0 the data type is the only reason to convert, and the `try_unary`
        // path is the identity for a millisecond index.
        let zoned = stream
            .manipulate(timestamp_batch_with_timezone(Some("+08:00"), &timestamps))
            .expect("a timezone-carrying index must still build the output batch")
            .expect("the ranges are not all empty");
        assert_eq!(
            payload_millis(&zoned),
            timestamps,
            "the converted payload must keep the try_unary values"
        );

        // The same logical samples without a timezone do take the zero-copy path; both paths
        // must agree on the payload.
        let plain = stream
            .manipulate(timestamp_batch(TimeUnit::Millisecond, &timestamps))
            .unwrap()
            .unwrap();
        assert_eq!(
            payload_millis(&plain),
            payload_millis(&zoned),
            "the zero-copy and the converted path must agree"
        );
    }

    /// A NULL slot's underlying value must never trigger the payload overflow
    /// error: a nullable batch is converted with arrow's `try_unary`, which applies
    /// the closure to valid rows only. The same value as a *valid* sample must still
    /// error, through the null-free `unary` + overflow-cell path.
    #[test]
    fn manipulate_null_slot_underlying_value_does_not_overflow() {
        // Second-slot underlying value i64::MAX seconds overflows on conversion,
        // but the slot is NULL so the conversion must not see it. offset=1ms
        // keeps the batch off the zero-copy path. The query range must be bounded
        // and contain the first sample: with `end == 0` no range is selected and
        // the conversion (and so the regression) is skipped, while an unbounded
        // `end` makes the range scan visit `end - start` evaluation points.
        let timestamps = [0i64, i64::MAX];
        let validity = NullBuffer::from(vec![true, false]);
        let array: ArrayRef = Arc::new(TimestampSecondArray::new(
            timestamps.to_vec().into(),
            Some(validity),
        ));
        let schema = Arc::new(Schema::new(vec![Field::new(
            TIME_INDEX_COLUMN,
            DataType::Timestamp(TimeUnit::Second, None),
            true,
        )]));
        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();

        let mut stream = payload_stream(1, 0, 1_000, 1, 1_000);
        stream.time_unit = TimeUnit::Second;
        let output = stream
            .manipulate(batch)
            .expect("a NULL slot's underlying value must not overflow the payload")
            .expect("the first sample (0s + 1ms) is inside the query range");
        assert_eq!(
            payload_millis(&output)[0],
            1,
            "the valid sample must be converted to its shifted millisecond payload"
        );

        // Same underlying value, but now valid: the conversion must reject it.
        let array: ArrayRef = Arc::new(TimestampSecondArray::from(timestamps.to_vec()));
        let schema = Arc::new(Schema::new(vec![Field::new(
            TIME_INDEX_COLUMN,
            DataType::Timestamp(TimeUnit::Second, None),
            false,
        )]));
        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
        let mut stream = payload_stream(1, 0, 1_000, 1, 1_000);
        stream.time_unit = TimeUnit::Second;
        let error = stream
            .manipulate(batch)
            .expect_err("a valid overflowing timestamp must still error");
        assert!(
            error.to_string().contains("timestamp payload overflow"),
            "unexpected error: {error}"
        );
    }
}