promql/extension_plan/

normalize.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::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use common_query::native_histogram::{START_TIMESTAMP_FIELD, native_histogram_arrow_type};
use datafusion::arrow::array::{Array, BooleanArray, StructArray};
use datafusion::arrow::compute;
use datafusion::common::tree_node::TreeNodeRecursion;
use datafusion::common::{Column, DFSchema, DFSchemaRef, Result as DataFusionResult, Statistics};
use datafusion::error::DataFusionError;
use datafusion::execution::context::TaskContext;
use datafusion::logical_expr::{EmptyRelation, Expr, LogicalPlan, UserDefinedLogicalNodeCore};
use datafusion::physical_plan::expressions::Column as ColumnExpr;
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, StatisticsArgs,
};
use datafusion_expr::ident;
use datatypes::arrow::array::TimestampMillisecondArray;
use datatypes::arrow::datatypes::{SchemaRef, TimestampMillisecondType};
use datatypes::arrow::record_batch::RecordBatch;
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, is_prometheus_stale_sample, prometheus_stale_sample_column,
    resolve_column_name, serialize_column_index,
};
use crate::metrics::PROMQL_SERIES_COUNT;

/// Normalizes a single-series input batch and optionally removes Prometheus stale markers.
///
/// This node remains the serialized carrier of the selector offset. Native sample timestamps
/// stay raw: applying an offset can overflow native `i64` ticks even when evaluation is valid.
/// Manipulators instead apply it in `i128` during selection and produce millisecond outputs.
/// Histogram start timestamps are already millisecond payloads used for rate/reset, so their
/// offsets are applied here, preserving unknown zero values and nulls.
#[derive(Debug, PartialEq, Eq, Hash, PartialOrd)]
pub struct SeriesNormalize {
    offset: Millisecond,
    time_index_column_name: String,
    filter_stale_markers: bool,
    tag_columns: Vec<String>,

    input: LogicalPlan,
    unfix: Option<UnfixIndices>,
}

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

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

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

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

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

        self.tag_columns
            .iter()
            .map(ident)
            .chain(std::iter::once(ident(&self.time_index_column_name)))
            .collect()
    }

    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();
        if output_columns.is_empty() {
            let indices = (0..input_schema.fields().len()).collect::<Vec<_>>();
            return Some(vec![indices]);
        }

        let mut required = Vec::with_capacity(output_columns.len() + 1 + self.tag_columns.len());
        required.extend_from_slice(output_columns);
        required.push(input_schema.index_of_column_by_name(None, &self.time_index_column_name)?);
        for tag in &self.tag_columns {
            required.push(input_schema.index_of_column_by_name(None, tag)?);
        }

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

    fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(
            f,
            "PromSeriesNormalize: offset=[{}], time index=[{}], filter NaN: [{}]",
            self.offset, self.time_index_column_name, self.filter_stale_markers
        )
    }

    fn with_exprs_and_inputs(
        &self,
        _exprs: Vec<Expr>,
        inputs: Vec<LogicalPlan>,
    ) -> DataFusionResult<Self> {
        if inputs.is_empty() {
            return Err(DataFusionError::Internal(
                "SeriesNormalize should have at least one input".to_string(),
            ));
        }

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

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

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

            Ok(Self {
                offset: self.offset,
                time_index_column_name,
                filter_stale_markers: self.filter_stale_markers,
                tag_columns,
                input,
                unfix: None,
            })
        } else {
            Ok(Self {
                offset: self.offset,
                time_index_column_name: self.time_index_column_name.clone(),
                filter_stale_markers: self.filter_stale_markers,
                tag_columns: self.tag_columns.clone(),
                input,
                unfix: None,
            })
        }
    }
}

impl SeriesNormalize {
    pub(crate) fn offset_for_time_index(&self, time_index: &Column) -> Option<Millisecond> {
        let index = self.input.schema().maybe_index_of_column(time_index)?;
        let (qualifier, field) = self.input.schema().qualified_field(index);
        (field.name() == &self.time_index_column_name
            && time_index == &Column::new(qualifier.cloned(), field.name().clone()))
            .then_some(self.offset)
    }

    pub fn new<N: AsRef<str>>(
        offset: Millisecond,
        time_index_column_name: N,
        filter_stale_markers: bool,
        tag_columns: Vec<String>,
        input: LogicalPlan,
    ) -> Self {
        Self {
            offset,
            time_index_column_name: time_index_column_name.as_ref().to_string(),
            filter_stale_markers,
            tag_columns,
            input,
            unfix: None,
        }
    }

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

    /// Returns whether this plan removes Prometheus stale markers.
    pub const fn filter_stale_markers(&self) -> bool {
        self.filter_stale_markers
    }

    pub fn to_execution_plan(&self, exec_input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
        Arc::new(SeriesNormalizeExec {
            offset: self.offset,
            time_index_column_name: self.time_index_column_name.clone(),
            filter_stale_markers: self.filter_stale_markers,
            input: exec_input,
            tag_columns: self.tag_columns.clone(),
            metric: ExecutionPlanMetricsSet::new(),
        })
    }

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

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

        pb::SeriesNormalize {
            offset: self.offset,
            time_index_idx,
            filter_nan: self.filter_stale_markers,
            tag_column_indices,
            ..Default::default()
        }
        .encode_to_vec()
    }

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

        let unfix = UnfixIndices {
            time_index_idx: pb_normalize.time_index_idx,
            tag_column_indices: pb_normalize.tag_column_indices.clone(),
        };

        Ok(Self {
            offset: pb_normalize.offset,
            time_index_column_name: String::new(),
            filter_stale_markers: pb_normalize.filter_nan,
            tag_columns: Vec::new(),
            input: placeholder_plan,
            unfix: Some(unfix),
        })
    }
}

#[derive(Debug)]
pub struct SeriesNormalizeExec {
    offset: Millisecond,
    time_index_column_name: String,
    filter_stale_markers: bool,
    tag_columns: Vec<String>,

    input: Arc<dyn ExecutionPlan>,
    metric: ExecutionPlanMetricsSet,
}

impl ExecutionPlan for SeriesNormalizeExec {
    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.input.schema()
    }

    fn input_distribution_requirements(&self) -> InputDistributionRequirements {
        if self.tag_columns.is_empty() {
            return InputDistributionRequirements::new(vec![Distribution::SinglePartition]);
        }

        let schema = self.input.schema();
        InputDistributionRequirements::new(vec![Distribution::KeyPartitioned(
            self.tag_columns
                .iter()
                // Safety: the tag column names is verified in the planning phase
                .map(|tag| Arc::new(ColumnExpr::new_with_schema(tag, &schema).unwrap()) as _)
                .collect(),
        )])
    }

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

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

    fn with_new_children(
        self: Arc<Self>,
        children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
        assert!(!children.is_empty());
        Ok(Arc::new(Self {
            offset: self.offset,
            time_index_column_name: self.time_index_column_name.clone(),
            filter_stale_markers: self.filter_stale_markers,
            input: children[0].clone(),
            tag_columns: self.tag_columns.clone(),
            metric: self.metric.clone(),
        }))
    }

    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();
        Ok(Box::pin(SeriesNormalizeStream {
            offset: self.offset,
            filter_stale_markers: self.filter_stale_markers,
            schema,
            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>> {
        Ok(Arc::clone(&input_stats[0]))
    }

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

impl DisplayAs for SeriesNormalizeExec {
    fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match t {
            DisplayFormatType::Default
            | DisplayFormatType::Verbose
            | DisplayFormatType::TreeRender => {
                write!(
                    f,
                    "PromSeriesNormalizeExec: offset=[{}], time index=[{}], filter NaN: [{}]",
                    self.offset, self.time_index_column_name, self.filter_stale_markers
                )
            }
        }
    }
}

pub struct SeriesNormalizeStream {
    offset: Millisecond,
    filter_stale_markers: bool,

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

impl SeriesNormalizeStream {
    pub fn normalize(&self, input: RecordBatch) -> DataFusionResult<RecordBatch> {
        // Native sample timestamps remain raw. Manipulators apply the selector offset
        // in wide nanosecond arithmetic, avoiding overflow in native Arrow storage.
        let mut columns = input.columns().to_vec();

        // Offset selectors move native histogram start timestamps onto the evaluation
        // timeline for rate and reset calculations. These payloads are milliseconds.
        if self.offset != 0 {
            let native_histogram_type = native_histogram_arrow_type();
            for column in &mut columns {
                let Some((histograms, start_timestamp_index, start_timestamps)) = column
                    .as_any()
                    .downcast_ref::<StructArray>()
                    .filter(|histograms| histograms.data_type() == &native_histogram_type)
                    .and_then(|histograms| {
                        let (index, _) = histograms.fields().find(START_TIMESTAMP_FIELD)?;
                        let timestamps = histograms
                            .column(index)
                            .as_any()
                            .downcast_ref::<TimestampMillisecondArray>()?;
                        Some((histograms, index, timestamps))
                    })
                else {
                    continue;
                };
                let start_timestamps = start_timestamps
                    .try_unary::<_, TimestampMillisecondType, _>(|timestamp| {
                        // Prometheus uses zero to mean that the start timestamp is unknown.
                        if timestamp == 0 {
                            Ok(0)
                        } else {
                            timestamp.checked_add(self.offset).ok_or_else(|| {
                                DataFusionError::Execution(
                                    "SeriesNormalize: histogram timestamp offset overflow".into(),
                                )
                            })
                        }
                    })?;
                // Struct arrays are immutable, so rebuild the physical histogram payload with
                // only its start-timestamp child replaced. Its logical schema stays unchanged:
                // the selector offset affects histogram reset/rate metadata, not the native
                // sample timestamp column consumed by later manipulators.
                let mut children = histograms.columns().to_vec();
                children[start_timestamp_index] = Arc::new(start_timestamps);
                *column = Arc::new(StructArray::new(
                    histograms.fields().clone(),
                    children,
                    histograms.nulls().cloned(),
                ));
            }
        }

        let result_batch = RecordBatch::try_new(input.schema(), columns)?;
        if !self.filter_stale_markers {
            return Ok(result_batch);
        }

        // Filter out Prometheus stale markers.
        let mut stale_marker_filter = vec![true; input.num_rows()];
        for column in result_batch.columns() {
            let Some(stale_sample_column) = prometheus_stale_sample_column(column.as_ref()) else {
                continue;
            };
            for (i, flag) in stale_marker_filter.iter_mut().enumerate() {
                if is_prometheus_stale_sample(stale_sample_column, i) {
                    *flag = false;
                }
            }
        }

        let result =
            compute::filter_record_batch(&result_batch, &BooleanArray::from(stale_marker_filter))
                .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?;
        Ok(result)
    }
}

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

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

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

#[cfg(test)]
mod test {
    use common_query::native_histogram::{build_histogram_array, read_histogram};
    use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS;
    use datafusion::arrow::array::{
        DictionaryArray, Float64Array, TimestampMicrosecondArray, TimestampNanosecondArray,
    };
    use datafusion::arrow::buffer::NullBuffer;
    use datafusion::arrow::datatypes::{
        ArrowPrimitiveType, DataType, Field, Int64Type, Schema, TimeUnit, TimestampMillisecondType,
    };
    use datafusion::common::ToDFSchema;
    use datafusion::datasource::memory::MemorySourceConfig;
    use datafusion::datasource::source::DataSourceExec;
    use datafusion::logical_expr::{EmptyRelation, LogicalPlan};
    use datafusion::prelude::SessionContext;
    use datatypes::arrow::array::TimestampMillisecondArray;
    use datatypes::arrow_array::StringArray;

    use super::*;
    use crate::extension_plan::RangeManipulate;
    use crate::extension_plan::test_util::native_histogram;
    use crate::range_array::RangeArray;

    const TIME_INDEX_COLUMN: &str = "timestamp";

    fn prepare_test_data() -> DataSourceExec {
        let schema = Arc::new(Schema::new(vec![
            Field::new(TIME_INDEX_COLUMN, TimestampMillisecondType::DATA_TYPE, true),
            Field::new("value", DataType::Float64, true),
            Field::new("path", DataType::Utf8, true),
        ]));
        let timestamp_column = Arc::new(TimestampMillisecondArray::from(vec![
            60_000, 120_000, 0, 30_000, 90_000,
        ])) as _;
        let field_column = Arc::new(Float64Array::from(vec![0.0, 1.0, 10.0, 100.0, 1000.0])) as _;
        let path_column = Arc::new(StringArray::from(vec!["foo", "foo", "foo", "foo", "foo"])) as _;
        let data = RecordBatch::try_new(
            schema.clone(),
            vec![timestamp_column, field_column, path_column],
        )
        .unwrap();

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

    #[test]
    fn pruning_should_keep_time_and_tag_columns_for_exec() {
        let df_schema = prepare_test_data().schema().to_dfschema_ref().unwrap();
        let input = LogicalPlan::EmptyRelation(EmptyRelation {
            produce_one_row: false,
            schema: df_schema,
        });
        let plan =
            SeriesNormalize::new(0, TIME_INDEX_COLUMN, true, vec!["path".to_string()], input);

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

    #[tokio::test]
    async fn test_sort_record_batch() {
        let memory_exec = Arc::new(prepare_test_data());
        let normalize_exec = Arc::new(SeriesNormalizeExec {
            offset: 0,
            time_index_column_name: TIME_INDEX_COLUMN.to_string(),
            filter_stale_markers: true,
            input: memory_exec,
            tag_columns: vec!["path".to_string()],
            metric: ExecutionPlanMetricsSet::new(),
        });
        let session_context = SessionContext::default();
        let result = datafusion::physical_plan::collect(normalize_exec, session_context.task_ctx())
            .await
            .unwrap();
        let result_literal = datatypes::arrow::util::pretty::pretty_format_batches(&result)
            .unwrap()
            .to_string();

        let expected = String::from(
            "+---------------------+--------+------+\
            \n| timestamp           | value  | path |\
            \n+---------------------+--------+------+\
            \n| 1970-01-01T00:01:00 | 0.0    | foo  |\
            \n| 1970-01-01T00:02:00 | 1.0    | foo  |\
            \n| 1970-01-01T00:00:00 | 10.0   | foo  |\
            \n| 1970-01-01T00:00:30 | 100.0  | foo  |\
            \n| 1970-01-01T00:01:30 | 1000.0 | foo  |\
            \n+---------------------+--------+------+",
        );

        assert_eq!(result_literal, expected);
    }

    #[tokio::test]
    async fn test_offset_record_batch() {
        let memory_exec = Arc::new(prepare_test_data());
        let normalize_exec = Arc::new(SeriesNormalizeExec {
            offset: 1_000,
            time_index_column_name: TIME_INDEX_COLUMN.to_string(),
            filter_stale_markers: true,
            input: memory_exec,
            metric: ExecutionPlanMetricsSet::new(),
            tag_columns: vec!["path".to_string()],
        });
        let session_context = SessionContext::default();
        let result = datafusion::physical_plan::collect(normalize_exec, session_context.task_ctx())
            .await
            .unwrap();
        let result_literal = datatypes::arrow::util::pretty::pretty_format_batches(&result)
            .unwrap()
            .to_string();

        let expected = String::from(
            "+---------------------+--------+------+\
            \n| timestamp           | value  | path |\
            \n+---------------------+--------+------+\
            \n| 1970-01-01T00:01:00 | 0.0    | foo  |\
            \n| 1970-01-01T00:02:00 | 1.0    | foo  |\
            \n| 1970-01-01T00:00:00 | 10.0   | foo  |\
            \n| 1970-01-01T00:00:30 | 100.0  | foo  |\
            \n| 1970-01-01T00:01:30 | 1000.0 | foo  |\
            \n+---------------------+--------+------+",
        );

        assert_eq!(result_literal, expected);
    }

    #[tokio::test]
    async fn filters_stale_markers_and_preserves_ordinary_nan() {
        let schema = Arc::new(Schema::new(vec![
            Field::new(
                TIME_INDEX_COLUMN,
                TimestampMillisecondType::DATA_TYPE,
                false,
            ),
            Field::new("value", DataType::Float64, true),
            Field::new("auxiliary", DataType::Float64, true),
        ]));
        let value_column = Float64Array::new(
            vec![
                42.0,
                f64::from_bits(0x7ff0_0000_0000_0002),
                24.0,
                f64::from_bits(0x7ff0_0000_0000_0002),
            ]
            .into(),
            Some(NullBuffer::from(vec![true, true, true, false])),
        );
        assert!(!value_column.is_valid(3));
        assert_eq!(value_column.value(3).to_bits(), 0x7ff0_0000_0000_0002);
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(TimestampMillisecondArray::from(vec![
                    1_000, 2_000, 3_000, 4_000,
                ])),
                Arc::new(value_column),
                Arc::new(Float64Array::from(vec![
                    f64::from_bits(0x7ff8_0000_0000_0000),
                    1.0,
                    f64::from_bits(0x7ff0_0000_0000_0002),
                    2.0,
                ])),
            ],
        )
        .unwrap();
        let input = Arc::new(DataSourceExec::new(Arc::new(
            MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(),
        )));
        let exec = Arc::new(SeriesNormalizeExec {
            offset: 0,
            time_index_column_name: TIME_INDEX_COLUMN.to_string(),
            filter_stale_markers: true,
            tag_columns: Vec::new(),
            input,
            metric: ExecutionPlanMetricsSet::new(),
        });

        let context = SessionContext::default();
        let batches = datafusion::physical_plan::collect(exec, context.task_ctx())
            .await
            .unwrap();
        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
        let batch = batches.iter().find(|batch| batch.num_rows() == 2).unwrap();
        let value = batch
            .column(1)
            .as_any()
            .downcast_ref::<Float64Array>()
            .unwrap();
        let auxiliary = batch
            .column(2)
            .as_any()
            .downcast_ref::<Float64Array>()
            .unwrap();

        assert_eq!(value.value(0), 42.0);
        assert_eq!(auxiliary.value(0).to_bits(), 0x7ff8_0000_0000_0000);
        assert!(!value.is_valid(1));
    }

    #[tokio::test]
    async fn offsets_native_histogram_timestamps_and_filters_stale_markers() {
        let mut regular = native_histogram(42.0);
        regular.start_timestamp = Some(500);
        let mut ordinary_nan = native_histogram(f64::NAN);
        ordinary_nan.start_timestamp = Some(0);
        let mut unknown_start = native_histogram(7.0);
        unknown_start.start_timestamp = None;
        let histograms = build_histogram_array(&[
            Some(regular),
            Some(native_histogram(f64::from_bits(PROMETHEUS_STALE_NAN_BITS))),
            Some(ordinary_nan),
            Some(unknown_start),
            None,
        ]);
        for (unit, ticks_per_ms) in [
            (TimeUnit::Millisecond, 1_i64),
            (TimeUnit::Microsecond, 1_000),
            (TimeUnit::Nanosecond, 1_000_000),
        ] {
            let timestamp_array = |values: Vec<i64>| -> Arc<dyn Array> {
                match unit {
                    TimeUnit::Millisecond => Arc::new(TimestampMillisecondArray::from(values)),
                    TimeUnit::Microsecond => Arc::new(TimestampMicrosecondArray::from(values)),
                    TimeUnit::Nanosecond => Arc::new(TimestampNanosecondArray::from(values)),
                    TimeUnit::Second => unreachable!(),
                }
            };
            for offset in [-1_i64, 1] {
                let timestamps = timestamp_array(
                    [1_000, 2_000, 3_000, 4_000, 5_000]
                        .into_iter()
                        .map(|timestamp| timestamp * ticks_per_ms)
                        .collect(),
                );
                let schema = Arc::new(Schema::new(vec![
                    Field::new(TIME_INDEX_COLUMN, timestamps.data_type().clone(), false),
                    Field::new("value", histograms.data_type().clone(), true),
                ]));
                let batch =
                    RecordBatch::try_new(schema.clone(), vec![timestamps, histograms.clone()])
                        .unwrap();
                let input = Arc::new(DataSourceExec::new(Arc::new(
                    MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(),
                )));
                let exec = Arc::new(SeriesNormalizeExec {
                    offset,
                    time_index_column_name: TIME_INDEX_COLUMN.to_string(),
                    filter_stale_markers: true,
                    tag_columns: Vec::new(),
                    input,
                    metric: ExecutionPlanMetricsSet::new(),
                });
                let context = SessionContext::default();
                let batches = datafusion::physical_plan::collect(exec, context.task_ctx())
                    .await
                    .unwrap();
                assert_eq!(
                    batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
                    4,
                    "unit={unit:?}, offset={offset}"
                );
                let batch = batches.iter().find(|batch| batch.num_rows() == 4).unwrap();
                let expected_timestamps = timestamp_array(
                    [1_000, 3_000, 4_000, 5_000]
                        .into_iter()
                        .map(|timestamp| timestamp * ticks_per_ms)
                        .collect(),
                );
                assert_eq!(
                    batch.column(0).to_data(),
                    expected_timestamps.to_data(),
                    "unit={unit:?}, offset={offset}"
                );
                let values = batch
                    .column(1)
                    .as_any()
                    .downcast_ref::<datafusion::arrow::array::StructArray>()
                    .unwrap();
                let regular = read_histogram(values, 0).unwrap().unwrap();
                assert_eq!(
                    (regular.sum, regular.start_timestamp),
                    (42.0, Some(500 + offset)),
                    "unit={unit:?}, offset={offset}"
                );
                let ordinary_nan = read_histogram(values, 1).unwrap().unwrap();
                assert!(ordinary_nan.sum.is_nan());
                assert_eq!(ordinary_nan.start_timestamp, Some(0));
                let unknown_start = read_histogram(values, 2).unwrap().unwrap();
                assert_eq!(
                    (unknown_start.sum, unknown_start.start_timestamp),
                    (7.0, None)
                );
                assert!(read_histogram(values, 3).unwrap().is_none());
            }
        }

        let mut known_start = native_histogram(42.0);
        known_start.start_timestamp = Some(500);
        let mut sentinel_start = native_histogram(8.0);
        sentinel_start.start_timestamp = Some(0);
        let unknown_start = native_histogram(7.0);
        let histograms =
            build_histogram_array(&[Some(known_start), Some(sentinel_start), Some(unknown_start)]);
        let schema = Arc::new(Schema::new(vec![
            Field::new(
                TIME_INDEX_COLUMN,
                TimestampMillisecondType::DATA_TYPE,
                false,
            ),
            Field::new("value", histograms.data_type().clone(), true),
        ]));
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(TimestampMillisecondArray::from(vec![1_000; 3])),
                histograms,
            ],
        )
        .unwrap();
        let logical_input = LogicalPlan::EmptyRelation(EmptyRelation {
            produce_one_row: false,
            schema: schema.clone().to_dfschema_ref().unwrap(),
        });
        let normalized =
            SeriesNormalize::new(1_000, TIME_INDEX_COLUMN, false, Vec::new(), logical_input);
        let range = RangeManipulate::new(
            2_000,
            2_000,
            1,
            1_000,
            1,
            TIME_INDEX_COLUMN.to_string(),
            vec!["value".to_string()],
            LogicalPlan::Extension(datafusion::logical_expr::Extension {
                node: Arc::new(normalized),
            }),
        )
        .unwrap();
        let input = Arc::new(DataSourceExec::new(Arc::new(
            MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(),
        )));
        let normalized_input = Arc::new(SeriesNormalizeExec {
            offset: 1_000,
            time_index_column_name: TIME_INDEX_COLUMN.to_string(),
            filter_stale_markers: false,
            tag_columns: Vec::new(),
            input,
            metric: ExecutionPlanMetricsSet::new(),
        });
        let output = datafusion::physical_plan::collect(
            range.to_execution_plan(normalized_input),
            SessionContext::default().task_ctx(),
        )
        .await
        .unwrap();
        let values = RangeArray::try_new(
            output[0]
                .column(1)
                .as_any()
                .downcast_ref::<DictionaryArray<Int64Type>>()
                .unwrap()
                .clone(),
        )
        .unwrap();
        let values = values.get(0).unwrap();
        let values = values
            .as_any()
            .downcast_ref::<datafusion::arrow::array::StructArray>()
            .unwrap();
        assert_eq!(
            read_histogram(values, 0).unwrap().unwrap().start_timestamp,
            Some(1_500)
        );
        assert_eq!(
            read_histogram(values, 1).unwrap().unwrap().start_timestamp,
            Some(0)
        );
        assert_eq!(
            read_histogram(values, 2).unwrap().unwrap().start_timestamp,
            None
        );
        let timestamps = RangeArray::try_new(
            output[0]
                .column(2)
                .as_any()
                .downcast_ref::<DictionaryArray<Int64Type>>()
                .unwrap()
                .clone(),
        )
        .unwrap();
        assert_eq!(
            timestamps.get(0).unwrap().to_data(),
            TimestampMillisecondArray::from(vec![2_000; 3]).to_data()
        );
    }
}