feat(promql): add native instant manipulate execution over raw series (#9660)

I hereby agree to the terms of the [GreptimeDB
CLA](https://github.com/GreptimeTeam/.github/blob/main/CLA.md).

## Refer to a related PR or issue link (optional)

**Stacked on #9654** (source materializer). Also requires **#9649**
(instant eval) and **#9655** (frame ordering).

## What's changed and what's your intention?

Add native instant manipulate execution over raw series streams:

- `NativeInstantMapping`: preflighted mapping from typed raw series to
query output schema.
- `NativeInstantManipulateExec`: ExecutionPlan that consumes
`SendableRawSeriesStream`, materializes labels, selects instants, and
emits aligned RecordBatches.
- `InstantManipulate::try_to_native_execution_plan()`: validates the raw
frame schema and constructs the native plan.

**5 files, +1830/-0 lines.** No callers wired yet.

Note: diff currently includes cherry-picked commits from #9649 and
#9655. Once those merge into narthex and the base is updated, the diff
will shrink to only the 5 native exec files.

### Verification

- `cargo check --all-targets` ✅, `cargo nextest run -p promql` →
**357/357 passed** ✅, `cargo clippy -- -D warnings` ✅

## PR Checklist

- [x] I have written the necessary rustdoc comments.
- [x] I have added the necessary unit tests and integration tests.
- [ ] This PR requires documentation updates.
- [x] API changes are backward compatible.
- [x] Schema or data changes are backward compatible.
- [ ] This PR needs to be backported to release branches.

src/promql/src/extension_plan/instant_manipulate.rs

@@ -12,6 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+mod native;
+mod native_exec;
+
 use std::pin::Pin;
 use std::sync::Arc;
 use std::task::{Context, Poll};

src/promql/src/extension_plan/instant_manipulate/native.rs

@@ -0,0 +1,302 @@
+// 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::sync::Arc;
+
+use datafusion::arrow::array::{ArrayRef, Float64Array, StringArray, TimestampMillisecondArray};
+use datafusion::arrow::datatypes::{DataType, SchemaRef, TimeUnit};
+use datafusion::arrow::record_batch::RecordBatch;
+use datafusion::error::{DataFusionError, Result as DataFusionResult};
+use datatypes::series::MetricName;
+use datatypes::series::raw::RawTimestamps;
+use datatypes::series::samples::{SampleRef, Samples, SeriesFields, SeriesSchema};
+
+#[cfg(test)]
+#[path = "native_tests.rs"]
+mod tests;
+
+use crate::extension_plan::instant_manipulate::InstantManipulate;
+use crate::instant_eval::{MAX_INSTANT_MANIPULATE_OUTPUT_POINTS, select_instants};
+
+#[derive(Clone)]
+enum Column {
+    Time,
+    Field(usize),
+    Tag(String),
+}
+
+/// Preflighted mapping from a complete typed series to the original query schema.
+pub(crate) struct NativeInstantMapping {
+    schema: SchemaRef,
+    series_schema: Arc<SeriesSchema>,
+    columns: Vec<Column>,
+    fields: Vec<usize>,
+    start: i64,
+    end: i64,
+    step: i64,
+    lookback: i64,
+}
+
+impl NativeInstantMapping {
+    /// Returns `None` when the node/schema is outside the bounded native subset.
+    pub(crate) fn try_new(
+        node: &InstantManipulate,
+        schema: SchemaRef,
+        series_schema: Arc<SeriesSchema>,
+        logical_tags: &[String],
+    ) -> DataFusionResult<Option<Self>> {
+        if node.offset != 0 || node.unfix.is_some() || node.start >= node.end {
+            return Ok(None);
+        }
+        if node.interval <= 0 {
+            return Err(DataFusionError::Execution(
+                "InstantManipulate interval must be positive".into(),
+            ));
+        }
+        if node.lookback_delta < 0 {
+            return Err(DataFusionError::Execution(
+                "InstantManipulate lookback must be non-negative".into(),
+            ));
+        }
+        if series_schema.eval_grid.is_some() {
+            return Ok(None);
+        }
+
+        let mut names = std::collections::HashSet::new();
+        for field in series_schema.fields.iter() {
+            if !names.insert(field.name.as_ref()) {
+                return Ok(None);
+            }
+        }
+        if logical_tags.iter().any(|name| !names.insert(name.as_str())) {
+            return Ok(None);
+        }
+
+        let stale_field_names = node.staleness_field_columns().collect::<Vec<_>>();
+        if stale_field_names.iter().any(|name| {
+            series_schema
+                .fields
+                .iter()
+                .all(|field| field.name.as_ref() != *name)
+        }) {
+            return Ok(None);
+        }
+
+        let mut columns = Vec::with_capacity(schema.fields().len());
+        let mut seen_fields = vec![false; series_schema.fields.len()];
+        let mut seen_tags = vec![false; logical_tags.len()];
+        let mut seen_time = false;
+        for field in schema.fields() {
+            let name = field.name();
+            let column = if name == node.time_index_column.as_str() {
+                if seen_time
+                    || field.data_type() != &DataType::Timestamp(TimeUnit::Millisecond, None)
+                {
+                    return Ok(None);
+                }
+                seen_time = true;
+                Column::Time
+            } else if let Some(index) = series_schema
+                .fields
+                .iter()
+                .position(|v| v.name.as_ref() == name)
+            {
+                if seen_fields[index] || field.data_type() != &DataType::Float64 {
+                    return Ok(None);
+                }
+                seen_fields[index] = true;
+                Column::Field(index)
+            } else if let Some(index) = logical_tags.iter().position(|tag| tag == name) {
+                if seen_tags[index] || !is_string_type(field.data_type()) {
+                    return Ok(None);
+                }
+                seen_tags[index] = true;
+                Column::Tag(name.clone())
+            } else {
+                return Ok(None);
+            };
+            columns.push(column);
+        }
+        if !seen_time || seen_fields.iter().any(|seen| !seen) || seen_tags.iter().any(|seen| !seen)
+        {
+            return Ok(None);
+        }
+
+        let fields = stale_field_names
+            .iter()
+            .filter_map(|name| {
+                series_schema
+                    .fields
+                    .iter()
+                    .position(|field| field.name.as_ref() == *name)
+            })
+            .collect();
+        Ok(Some(Self {
+            schema,
+            series_schema,
+            columns,
+            fields,
+            start: node.start,
+            end: node.end,
+            step: node.interval,
+            lookback: node.lookback_delta,
+        }))
+    }
+
+    /// The original output schema retained by this mapping.
+    pub(crate) fn schema(&self) -> SchemaRef {
+        self.schema.clone()
+    }
+
+    /// Estimates the clipped evaluation grid for complete source timestamp bounds.
+    pub(crate) fn estimated_points(&self, first: i64, last: i64) -> DataFusionResult<usize> {
+        let first = first as i128;
+        let last = last as i128;
+        let useful_last = if self.lookback == 0 {
+            last
+        } else {
+            last + self.lookback as i128 - 1
+        };
+        let query_start = self.start as i128;
+        let query_end = self.end as i128;
+        let max_start = first.max(query_start);
+        let min_end = useful_last.min(query_end);
+        let (aligned_start, aligned_end) = if max_start > min_end {
+            (1_i128, 0_i128)
+        } else {
+            (
+                query_start + (max_start - query_start) / self.step as i128 * self.step as i128,
+                query_end - (query_end - min_end) / self.step as i128 * self.step as i128,
+            )
+        };
+        let points = if aligned_end >= aligned_start {
+            (aligned_end - aligned_start) / self.step as i128 + 1
+        } else {
+            0
+        };
+        if points > MAX_INSTANT_MANIPULATE_OUTPUT_POINTS as i128 {
+            return Err(DataFusionError::Execution(format!(
+                "InstantManipulate output points exceed limit: {points} > {MAX_INSTANT_MANIPULATE_OUTPUT_POINTS}"
+            )));
+        }
+        usize::try_from(points).map_err(|_| {
+            DataFusionError::ResourcesExhausted("native instant point count overflow".into())
+        })
+    }
+
+    /// Selects instants against typed samples, then maps only selected values to Arrow.
+    pub(crate) fn evaluate(
+        &self,
+        metric: &MetricName,
+        time: &RawTimestamps,
+        fields: &SeriesFields,
+    ) -> DataFusionResult<RecordBatch> {
+        if fields.len() != self.series_schema.fields.len() {
+            return Err(DataFusionError::Execution(
+                "native instant field count mismatch".into(),
+            ));
+        }
+        let checked_fields = self
+            .fields
+            .iter()
+            .map(|i| fields.as_slice()[*i].clone())
+            .collect::<Vec<_>>();
+        let checked_fields = if checked_fields.is_empty() {
+            vec![Samples::from_float(Float64Array::new_null(time.len()))]
+        } else {
+            checked_fields
+        };
+        let checked = SeriesFields::try_from_vec(checked_fields)
+            .map_err(|e| DataFusionError::Execution(e.to_string()))?;
+        let (indices, timestamps) = select_instants(
+            time,
+            &checked,
+            self.start..=self.end,
+            self.step,
+            self.lookback,
+            0,
+        )?;
+        if indices.is_empty() {
+            return Ok(RecordBatch::new_empty(self.schema.clone()));
+        }
+
+        let arrays: DataFusionResult<Vec<ArrayRef>> = self
+            .columns
+            .iter()
+            .enumerate()
+            .map(|(output_index, column)| match column {
+                Column::Time => {
+                    Ok(Arc::new(TimestampMillisecondArray::from(timestamps.clone())) as ArrayRef)
+                }
+                Column::Field(index) => {
+                    let source = fields.get(*index).ok_or_else(|| {
+                        DataFusionError::Execution("native instant field missing".into())
+                    })?;
+                    let mut samples = source.iter();
+                    let mut cursor = 0usize;
+                    let mut selected_index = None;
+                    let mut selected_value = None;
+                    let mut values = Vec::with_capacity(indices.len());
+                    for index in &indices {
+                        let index = *index as usize;
+                        if selected_index != Some(index) {
+                            while cursor <= index {
+                                selected_value = Some(
+                                    match samples.next().ok_or_else(|| {
+                                        DataFusionError::Execution(
+                                        "native instant sample iterator ended before selected index"
+                                            .into(),
+                                    )
+                                    })? {
+                                        SampleRef::Float(value) => Some(value),
+                                        SampleRef::Missing => None,
+                                        SampleRef::Histogram(_) => {
+                                            return Err(DataFusionError::Execution(
+                                                "native instant expected Float64 sample".into(),
+                                            ));
+                                        }
+                                    },
+                                );
+                                selected_index = Some(cursor);
+                                cursor += 1;
+                            }
+                        }
+                        values.push(selected_value.ok_or_else(|| {
+                            DataFusionError::Execution(
+                                "native instant selected sample was not initialized".into(),
+                            )
+                        })?);
+                    }
+                    Ok(Arc::new(Float64Array::from(values)) as ArrayRef)
+                }
+                Column::Tag(name) => {
+                    let values = std::iter::repeat_with(|| metric.get(name))
+                        .take(indices.len())
+                        .collect::<StringArray>();
+                    let target = self.schema.field(output_index).data_type();
+                    let cast = datafusion::arrow::compute::cast(&values, target)?;
+                    Ok(cast)
+                }
+            })
+            .collect();
+        Ok(RecordBatch::try_new(self.schema.clone(), arrays?)?)
+    }
+}
+
+fn is_string_type(data_type: &DataType) -> bool {
+    matches!(
+        data_type,
+        DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View
+    ) || matches!(data_type, DataType::Dictionary(key, value) if key.as_ref() == &DataType::UInt32 && value.as_ref() == &DataType::Utf8)
+}

src/promql/src/extension_plan/instant_manipulate/native_exec.rs

@@ -0,0 +1,501 @@
+// 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::series_frame::{raw_series_frame_ordering, unpack_raw_series};
+use datafusion::arrow::array::Array;
+use datafusion::arrow::datatypes::SchemaRef;
+use datafusion::arrow::record_batch::RecordBatch;
+use datafusion::common::tree_node::TreeNodeRecursion;
+use datafusion::error::{DataFusionError, Result as DfResult};
+use datafusion::execution::context::TaskContext;
+use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation};
+use datafusion::physical_expr::EquivalenceProperties;
+use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion::physical_plan::{
+    DisplayAs, DisplayFormatType, ExecutionPlan, PhysicalExpr, PlanProperties, RecordBatchStream,
+    SendableRecordBatchStream,
+};
+use datatypes::series::SeriesOrdering;
+use datatypes::series::raw::{RawSeriesChunk, RawTimestamps};
+use datatypes::series::samples::{DenseSamplesRef, Samples, SeriesFields, SeriesSchema};
+use futures::{Stream, ready};
+use store_api::region_engine::{RawLabelSource, SendableRawSeriesStream};
+
+use crate::extension_plan::instant_manipulate::InstantManipulate;
+use crate::extension_plan::instant_manipulate::native::NativeInstantMapping;
+use crate::source::materialize_raw_labels;
+
+#[cfg(test)]
+#[path = "native_exec_tests.rs"]
+mod tests;
+
+impl InstantManipulate {
+    /// Returns whether the bounded native mapping accepts this output schema, series schema, and
+    /// complete logical tag set.
+    ///
+    /// This is the local half of [`Self::try_to_native_execution_plan`]: the same decision that
+    /// method makes before it inspects its input, so a caller can decline a plan without building
+    /// one. `Err` is the same error that method propagates for an invalid interval or lookback.
+    pub fn is_native_mapping_supported(
+        &self,
+        output_schema: SchemaRef,
+        series_schema: Arc<SeriesSchema>,
+        logical_tags: &[String],
+    ) -> DfResult<bool> {
+        Ok(
+            NativeInstantMapping::try_new(self, output_schema, series_schema, logical_tags)?
+                .is_some(),
+        )
+    }
+
+    /// Builds the bounded native execution path for a prepared raw-series frame.
+    pub fn try_to_native_execution_plan(
+        &self,
+        input: Arc<dyn ExecutionPlan>,
+        output_schema: SchemaRef,
+        label_source: RawLabelSource,
+        series_schema: Arc<SeriesSchema>,
+        logical_tags: &[String],
+    ) -> DfResult<Option<Arc<dyn ExecutionPlan>>> {
+        let Some(mapping) = NativeInstantMapping::try_new(
+            self,
+            output_schema,
+            series_schema.clone(),
+            logical_tags,
+        )?
+        else {
+            return Ok(None);
+        };
+        if input.properties().partitioning.partition_count() != 1 {
+            return Ok(None);
+        }
+        let declared_schema = input.schema();
+        let ordering = raw_series_frame_ordering(declared_schema.as_ref())?;
+        // Validate the declared raw-frame schema without executing the child.
+        unpack_raw_series(
+            Box::pin(RecordBatchStreamAdapter::new(
+                declared_schema.clone(),
+                futures::stream::empty(),
+            )),
+            series_schema.clone(),
+            label_source.clone(),
+            ordering,
+        )?;
+        Ok(Some(Arc::new(NativeInstantManipulateExec::new(
+            input,
+            Arc::new(mapping),
+            declared_schema,
+            series_schema,
+            label_source,
+            ordering,
+        ))))
+    }
+}
+
+struct NativeInstantManipulateExec {
+    input: Arc<dyn ExecutionPlan>,
+    mapping: Arc<NativeInstantMapping>,
+    raw_schema: SchemaRef,
+    series_schema: Arc<SeriesSchema>,
+    label_source: RawLabelSource,
+    ordering: SeriesOrdering,
+    properties: Arc<PlanProperties>,
+}
+impl NativeInstantManipulateExec {
+    fn new(
+        input: Arc<dyn ExecutionPlan>,
+        mapping: Arc<NativeInstantMapping>,
+        raw_schema: SchemaRef,
+        series_schema: Arc<SeriesSchema>,
+        label_source: RawLabelSource,
+        ordering: SeriesOrdering,
+    ) -> Self {
+        let properties = Arc::new(PlanProperties::new(
+            EquivalenceProperties::new(mapping.schema()),
+            datafusion::physical_plan::Partitioning::UnknownPartitioning(1),
+            EmissionType::Incremental,
+            Boundedness::Bounded,
+        ));
+        Self {
+            input,
+            mapping,
+            raw_schema,
+            series_schema,
+            label_source,
+            ordering,
+            properties,
+        }
+    }
+}
+impl std::fmt::Debug for NativeInstantManipulateExec {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("NativeInstantManipulateExec").finish()
+    }
+}
+impl DisplayAs for NativeInstantManipulateExec {
+    fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(f, "NativeInstantManipulateExec")
+    }
+}
+impl ExecutionPlan for NativeInstantManipulateExec {
+    fn name(&self) -> &str {
+        "NativeInstantManipulateExec"
+    }
+    fn schema(&self) -> SchemaRef {
+        self.mapping.schema()
+    }
+    fn apply_expressions(
+        &self,
+        _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> DfResult<TreeNodeRecursion>,
+    ) -> DfResult<TreeNodeRecursion> {
+        Ok(TreeNodeRecursion::Continue)
+    }
+    fn properties(&self) -> &Arc<PlanProperties> {
+        &self.properties
+    }
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+        vec![&self.input]
+    }
+    fn with_new_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+    ) -> DfResult<Arc<dyn ExecutionPlan>> {
+        if children.len() != 1
+            || children[0].schema() != self.raw_schema
+            || children[0].properties().partitioning.partition_count() != 1
+        {
+            return Err(DataFusionError::Plan(
+                "NativeInstantManipulateExec requires one unchanged one-partition raw child".into(),
+            ));
+        }
+        Ok(Arc::new(Self::new(
+            children[0].clone(),
+            self.mapping.clone(),
+            self.raw_schema.clone(),
+            self.series_schema.clone(),
+            self.label_source.clone(),
+            self.ordering,
+        )))
+    }
+    fn execute(
+        &self,
+        partition: usize,
+        context: Arc<TaskContext>,
+    ) -> DfResult<SendableRecordBatchStream> {
+        if partition != 0 {
+            return Err(DataFusionError::Execution(
+                "native instant partition must be zero".into(),
+            ));
+        }
+        let frame = self.input.execute(partition, context.clone())?;
+        let source = unpack_raw_series(
+            frame,
+            self.series_schema.clone(),
+            self.label_source.clone(),
+            self.ordering,
+        )?;
+        let source = materialize_raw_labels(source)
+            .map_err(|e| DataFusionError::Execution(e.to_string()))?;
+        let reservation =
+            MemoryConsumer::new("NativeInstantStream").register(&context.runtime_env().memory_pool);
+        Ok(Box::pin(NativeInstantStream {
+            source: Some(source),
+            mapping: self.mapping.clone(),
+            schema: self.mapping.schema(),
+            series_schema: self.series_schema.clone(),
+            pending: Vec::new(),
+            reservation,
+        }))
+    }
+}
+
+struct NativeInstantStream {
+    source: Option<SendableRawSeriesStream>,
+    mapping: Arc<NativeInstantMapping>,
+    schema: SchemaRef,
+    series_schema: Arc<SeriesSchema>,
+    pending: Vec<RawSeriesChunk>,
+    reservation: MemoryReservation,
+}
+impl NativeInstantStream {
+    fn fail(&mut self, error: DataFusionError) -> Poll<Option<DfResult<RecordBatch>>> {
+        self.source = None;
+        self.pending = Vec::new();
+        self.reservation.free();
+        Poll::Ready(Some(Err(error)))
+    }
+    fn finish_series(&mut self) -> DfResult<RecordBatch> {
+        let chunks = std::mem::take(&mut self.pending);
+        let rows = chunks
+            .iter()
+            .try_fold(0usize, |n, c| n.checked_add(c.time.len()))
+            .ok_or_else(|| {
+                DataFusionError::ResourcesExhausted("native instant row count overflow".into())
+            })?;
+        let first = chunks
+            .first()
+            .and_then(|c| c.time.timestamp(0))
+            .ok_or_else(|| {
+                DataFusionError::Execution("native instant series has no timestamps".into())
+            })?;
+        let last = chunks
+            .last()
+            .and_then(|c| c.time.timestamp(c.time.len().saturating_sub(1)))
+            .ok_or_else(|| {
+                DataFusionError::Execution("native instant series has no timestamps".into())
+            })?;
+        let points = self.mapping.estimated_points(first, last)?;
+        let label_bytes = chunks
+            .first()
+            .and_then(|c| c.metric.as_ref())
+            .map(|m| {
+                m.name().map_or(0, str::len)
+                    + m.tags()
+                        .iter()
+                        .map(|l| l.name.len() + l.value.len() + 16)
+                        .sum::<usize>()
+            })
+            .unwrap_or(0);
+        // Reserve timestamp concatenation by row count, but field descriptors by their
+        // actual chunk/segment counts: a dense chunk is still one field descriptor.
+        let field_count = self.series_schema.fields.len();
+        let chunk_count = chunks.len();
+        let segment_count = chunks
+            .iter()
+            .flat_map(|chunk| chunk.fields.as_slice())
+            .try_fold(0usize, |count, field| {
+                count.checked_add(field.segments().len())
+            })
+            .ok_or_else(|| {
+                DataFusionError::ResourcesExhausted("native instant segment count overflow".into())
+            })?;
+        let chunk_field_descriptors = chunk_count
+            .checked_mul(field_count)
+            .and_then(|n| n.checked_mul(std::mem::size_of::<Samples>()))
+            .ok_or_else(|| {
+                DataFusionError::ResourcesExhausted(
+                    "native instant descriptor estimate overflow".into(),
+                )
+            })?;
+        // Flattening and converting to Arc can transiently retain the target descriptors and
+        // two source/flattened vectors; three segment-counts conservatively cover those owners.
+        let flattened_descriptors = segment_count
+            .checked_mul(3)
+            .and_then(|n| n.checked_mul(std::mem::size_of::<Samples>()))
+            .ok_or_else(|| {
+                DataFusionError::ResourcesExhausted(
+                    "native instant segment estimate overflow".into(),
+                )
+            })?;
+        let scratch = rows
+            .checked_mul(16)
+            .and_then(|n| n.checked_add(chunk_field_descriptors))
+            .and_then(|n| n.checked_add(flattened_descriptors))
+            .and_then(|n| {
+                n.checked_add(field_count.checked_mul(std::mem::size_of::<Vec<Samples>>())?)
+            })
+            .and_then(|n| {
+                points
+                    .checked_mul(
+                        self.schema
+                            .fields()
+                            .len()
+                            .checked_mul(24)?
+                            .checked_add(label_bytes)?
+                            .checked_add(48)?,
+                    )?
+                    .checked_add(n)
+            })
+            .ok_or_else(|| {
+                DataFusionError::ResourcesExhausted(
+                    "native instant scratch estimate overflow".into(),
+                )
+            })?;
+        self.reservation
+            .try_grow(scratch)
+            .map_err(|e| DataFusionError::ResourcesExhausted(e.to_string()))?;
+        let metric = chunks
+            .first()
+            .and_then(|c| c.metric.clone())
+            .ok_or_else(|| {
+                DataFusionError::Execution("native instant series has no metric labels".into())
+            })?;
+        let mut times = Vec::new();
+        times
+            .try_reserve_exact(rows)
+            .map_err(|e| DataFusionError::ResourcesExhausted(e.to_string()))?;
+        let mut field_chunks = (0..field_count).map(|_| Vec::new()).collect::<Vec<_>>();
+        for field in &mut field_chunks {
+            field
+                .try_reserve_exact(chunk_count)
+                .map_err(|e| DataFusionError::ResourcesExhausted(e.to_string()))?;
+        }
+        for chunk in chunks {
+            for i in 0..chunk.time.len() {
+                times.push(chunk.time.timestamp(i).ok_or_else(|| {
+                    DataFusionError::Execution("native instant timestamp missing".into())
+                })?);
+            }
+            for (target, field) in field_chunks.iter_mut().zip(chunk.fields.as_slice()) {
+                target.push(field.clone());
+            }
+        }
+        let time = RawTimestamps::try_new(times.into())
+            .map_err(|e| DataFusionError::Execution(e.to_string()))?;
+        let fields = field_chunks
+            .into_iter()
+            .map(Samples::try_from_chunks)
+            .collect::<std::result::Result<Vec<_>, _>>()
+            .map_err(|e| DataFusionError::Execution(e.to_string()))?;
+        let fields = SeriesFields::try_from_vec(fields)
+            .map_err(|e| DataFusionError::Execution(e.to_string()))?;
+        self.mapping.evaluate(&metric, &time, &fields)
+    }
+}
+fn retained_chunk_bytes(chunk: &RawSeriesChunk) -> Option<usize> {
+    let timestamp_bytes = chunk.time.as_buffer().inner().capacity();
+    let labels = chunk
+        .metric
+        .as_ref()
+        .map(|metric| {
+            metric
+                .tags()
+                .iter()
+                .try_fold(metric.name().map_or(0, str::len), |total, label| {
+                    total
+                        .checked_add(label.name.len())?
+                        .checked_add(label.value.len())?
+                        .checked_add(std::mem::size_of::<datatypes::series::Label>())?
+                        .checked_add(2 * std::mem::size_of::<usize>())
+                })
+        })
+        .unwrap_or(Some(0))?;
+    let field_descriptors = std::mem::size_of::<RawSeriesChunk>()
+        .checked_add(std::mem::size_of::<SeriesFields>())?
+        .checked_add(
+            chunk
+                .fields
+                .as_slice()
+                .len()
+                .checked_mul(std::mem::size_of::<Samples>())?,
+        )?;
+    let (segment_descriptors, sample_buffers) = chunk
+        .fields
+        .as_slice()
+        .iter()
+        .flat_map(|field| field.segments())
+        .try_fold((0usize, 0usize), |(descriptors, buffers), segment| {
+            let bytes = match segment.dense()? {
+                DenseSamplesRef::Float(array) => array.get_buffer_memory_size(),
+                _ => return None,
+            };
+            Some((
+                descriptors.checked_add(std::mem::size_of::<Samples>())?,
+                buffers.checked_add(bytes)?,
+            ))
+        })?;
+    timestamp_bytes
+        .checked_add(chunk.primary_key.capacity())?
+        .checked_add(labels)?
+        .checked_add(field_descriptors)?
+        .checked_add(segment_descriptors)?
+        .checked_add(sample_buffers)
+}
+
+impl Stream for NativeInstantStream {
+    type Item = DfResult<RecordBatch>;
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        for _ in 0..64 {
+            let Some(source) = self.source.as_mut() else {
+                self.pending = Vec::new();
+                self.reservation.free();
+                return Poll::Ready(None);
+            };
+            match ready!(source.as_mut().poll_next(cx)) {
+                Some(Err(error)) => {
+                    return self.fail(DataFusionError::External(Box::new(error)));
+                }
+                None if !self.pending.is_empty() => {
+                    return self.fail(DataFusionError::Execution(
+                        "raw series ended without final chunk".into(),
+                    ));
+                }
+                None => {
+                    self.source = None;
+                    self.pending = Vec::new();
+                    self.reservation.free();
+                    return Poll::Ready(None);
+                }
+                Some(Ok(chunk)) => {
+                    let Some(bytes) = retained_chunk_bytes(&chunk) else {
+                        return self.fail(DataFusionError::ResourcesExhausted(
+                            "native instant retained-byte estimate overflow".into(),
+                        ));
+                    };
+                    if self.pending.len() == self.pending.capacity() {
+                        let old_capacity = self.pending.capacity();
+                        let new_capacity = old_capacity.saturating_mul(2).max(4);
+                        let slot_growth =
+                            new_capacity.checked_sub(old_capacity).and_then(|slots| {
+                                slots.checked_mul(std::mem::size_of::<RawSeriesChunk>())
+                            });
+                        let Some(slot_growth) = slot_growth else {
+                            return self.fail(DataFusionError::ResourcesExhausted(
+                                "native instant pending capacity overflow".into(),
+                            ));
+                        };
+                        let Some(growth) = bytes.checked_add(slot_growth) else {
+                            return self.fail(DataFusionError::ResourcesExhausted(
+                                "native instant chunk estimate overflow".into(),
+                            ));
+                        };
+                        if let Err(e) = self.reservation.try_grow(growth) {
+                            return self.fail(DataFusionError::ResourcesExhausted(e.to_string()));
+                        }
+                        let additional = new_capacity - self.pending.capacity();
+                        if let Err(e) = self.pending.try_reserve_exact(additional) {
+                            return self.fail(DataFusionError::ResourcesExhausted(e.to_string()));
+                        }
+                    } else if let Err(e) = self.reservation.try_grow(bytes) {
+                        return self.fail(DataFusionError::ResourcesExhausted(e.to_string()));
+                    }
+                    let last = chunk.is_last;
+                    self.pending.push(chunk);
+                    if last {
+                        match self.finish_series() {
+                            Ok(batch) => {
+                                self.pending = Vec::new();
+                                self.reservation.free();
+                                return Poll::Ready(Some(Ok(batch)));
+                            }
+                            Err(e) => return self.fail(e),
+                        }
+                    }
+                }
+            }
+        }
+        cx.waker().wake_by_ref();
+        Poll::Pending
+    }
+}
+impl RecordBatchStream for NativeInstantStream {
+    fn schema(&self) -> SchemaRef {
+        self.schema.clone()
+    }
+}

src/promql/src/extension_plan/instant_manipulate/native_exec_tests.rs

@@ -0,0 +1,427 @@
+// 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::VecDeque;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::task::{Context, Poll};
+
+use common_error::ext::{BoxedError, ErrorExt, PlainError};
+use common_error::status_code::StatusCode;
+use datafusion::arrow::array::{Array, Float64Array, TimestampMillisecondArray};
+use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit};
+use datafusion::common::DFSchema;
+use datafusion::error::DataFusionError;
+use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryConsumer, MemoryPool};
+use datafusion::logical_expr::{EmptyRelation, LogicalPlan};
+use datatypes::arrow::buffer::ScalarBuffer;
+use datatypes::series::raw::{RawSeriesChunk, RawTimestamps};
+use datatypes::series::samples::{FieldSpec, Samples, SeriesFields, SeriesSchema};
+use datatypes::series::{MetricName, SeriesOrdering};
+use futures::Stream;
+use store_api::region_engine::{RawLabelSource, RawSeriesStream, SendableRawSeriesStream};
+
+use super::*;
+
+fn raw_schema() -> Arc<SeriesSchema> {
+    Arc::new(SeriesSchema {
+        fields: Arc::from([FieldSpec {
+            name: Arc::from("value"),
+        }]),
+        eval_grid: None,
+    })
+}
+
+fn output_schema() -> Arc<Schema> {
+    Arc::new(Schema::new(vec![
+        Field::new("value", DataType::Float64, true),
+        Field::new(
+            "time",
+            DataType::Timestamp(TimeUnit::Millisecond, None),
+            false,
+        ),
+    ]))
+}
+
+fn mapping_for(
+    start: i64,
+    end: i64,
+    step: i64,
+    lookback: i64,
+) -> (Arc<NativeInstantMapping>, Arc<SeriesSchema>) {
+    let logical = LogicalPlan::EmptyRelation(EmptyRelation {
+        produce_one_row: false,
+        schema: Arc::new(DFSchema::empty()),
+    });
+    let node = InstantManipulate::new(
+        start,
+        end,
+        lookback,
+        step,
+        0,
+        "time".into(),
+        vec![],
+        Some("value".into()),
+        logical,
+    );
+    let series = raw_schema();
+    let mapping = NativeInstantMapping::try_new(&node, output_schema(), series.clone(), &[])
+        .unwrap()
+        .unwrap();
+    (Arc::new(mapping), series)
+}
+
+fn chunk(time: &[i64], values: &[f64], last: bool) -> RawSeriesChunk {
+    RawSeriesChunk::try_new(
+        &raw_schema(),
+        datatypes::arrow::buffer::Buffer::from(vec![1, 2, 3]),
+        Some(MetricName::try_new(None, vec![]).unwrap()),
+        RawTimestamps::try_new(ScalarBuffer::from(time.to_vec())).unwrap(),
+        SeriesFields::one(Samples::from_float(Float64Array::from(values.to_vec()))),
+        last,
+    )
+    .unwrap()
+}
+
+struct RawMock {
+    schema: Arc<SeriesSchema>,
+    labels: RawLabelSource,
+    items: VecDeque<Poll<Option<std::result::Result<RawSeriesChunk, BoxedError>>>>,
+    polls: Arc<AtomicUsize>,
+    drops: Arc<AtomicUsize>,
+}
+impl Stream for RawMock {
+    type Item = std::result::Result<RawSeriesChunk, BoxedError>;
+    fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        self.polls.fetch_add(1, Ordering::SeqCst);
+        self.items.pop_front().unwrap_or(Poll::Ready(None))
+    }
+}
+impl Drop for RawMock {
+    fn drop(&mut self) {
+        self.drops.fetch_add(1, Ordering::SeqCst);
+    }
+}
+impl RawSeriesStream for RawMock {
+    fn schema(&self) -> &Arc<SeriesSchema> {
+        &self.schema
+    }
+    fn label_source(&self) -> &RawLabelSource {
+        &self.labels
+    }
+    fn ordering(&self) -> SeriesOrdering {
+        SeriesOrdering::Contiguous
+    }
+}
+
+fn raw(
+    items: impl IntoIterator<Item = Poll<Option<std::result::Result<RawSeriesChunk, BoxedError>>>>,
+) -> (SendableRawSeriesStream, Arc<AtomicUsize>, Arc<AtomicUsize>) {
+    let polls = Arc::new(AtomicUsize::new(0));
+    let drops = Arc::new(AtomicUsize::new(0));
+    (
+        Box::pin(RawMock {
+            schema: raw_schema(),
+            labels: RawLabelSource::Explicit,
+            items: items.into_iter().collect(),
+            polls: polls.clone(),
+            drops: drops.clone(),
+        }),
+        polls,
+        drops,
+    )
+}
+
+fn stream(source: SendableRawSeriesStream, pool: &Arc<dyn MemoryPool>) -> NativeInstantStream {
+    stream_for(source, pool, 10, 30, 10, 10)
+}
+
+fn stream_for(
+    source: SendableRawSeriesStream,
+    pool: &Arc<dyn MemoryPool>,
+    start: i64,
+    end: i64,
+    step: i64,
+    lookback: i64,
+) -> NativeInstantStream {
+    let (mapping, schema) = mapping_for(start, end, step, lookback);
+    NativeInstantStream {
+        source: Some(source),
+        mapping: mapping.clone(),
+        schema: mapping.schema(),
+        series_schema: schema,
+        pending: Vec::new(),
+        reservation: MemoryConsumer::new("NativeInstantStreamTest").register(pool),
+    }
+}
+
+fn poll(stream: &mut NativeInstantStream) -> Poll<Option<DfResult<RecordBatch>>> {
+    let waker = futures::task::noop_waker();
+    let mut cx = Context::from_waker(&waker);
+    Pin::new(stream).poll_next(&mut cx)
+}
+
+#[test]
+fn emits_complete_series_before_polling_later_pending_input() {
+    let (input, polls, _) = raw([
+        Poll::Ready(Some(Ok(chunk(&[10], &[1.0], false)))),
+        Poll::Ready(Some(Ok(chunk(&[20], &[2.0], true)))),
+        Poll::Pending,
+    ]);
+    let pool = Arc::new(GreedyMemoryPool::new(1_000_000)) as Arc<dyn MemoryPool>;
+    let mut stream = stream(input, &pool);
+    let Poll::Ready(Some(Ok(batch))) = poll(&mut stream) else {
+        panic!("completed series should emit")
+    };
+    assert_eq!(polls.load(Ordering::SeqCst), 2);
+    assert_eq!(
+        batch
+            .column(0)
+            .as_any()
+            .downcast_ref::<Float64Array>()
+            .unwrap()
+            .values(),
+        &[1.0, 2.0]
+    );
+    assert_eq!(
+        batch
+            .column(1)
+            .as_any()
+            .downcast_ref::<TimestampMillisecondArray>()
+            .unwrap()
+            .values(),
+        &[10, 20]
+    );
+    assert!(matches!(poll(&mut stream), Poll::Pending));
+    assert_eq!(polls.load(Ordering::SeqCst), 3);
+}
+
+#[test]
+fn source_error_after_prefix_is_returned_once_and_clears_reservation_and_input() {
+    let failure = BoxedError::new(PlainError::new(
+        "mock source error".into(),
+        StatusCode::Internal,
+    ));
+    let (input, polls, drops) = raw([
+        Poll::Ready(Some(Ok(chunk(&[10], &[1.0], true)))),
+        Poll::Ready(Some(Err(failure))),
+    ]);
+    let pool = Arc::new(GreedyMemoryPool::new(1_000_000)) as Arc<dyn MemoryPool>;
+    let mut stream = stream(input, &pool);
+    assert!(matches!(poll(&mut stream), Poll::Ready(Some(Ok(_)))));
+    let Poll::Ready(Some(Err(error))) = poll(&mut stream) else {
+        panic!("source error must be returned once")
+    };
+    assert_eq!(
+        common_query::error::datafusion_status_code::<common_query::error::Error>(&error, None),
+        StatusCode::Internal
+    );
+    let DataFusionError::External(source) = error else {
+        panic!("source error must stay a typed DataFusionError::External payload")
+    };
+    let source = source
+        .downcast_ref::<BoxedError>()
+        .expect("source error must still be the original BoxedError");
+    assert_eq!(source.status_code(), StatusCode::Internal);
+    assert_eq!(source.to_string(), "mock source error");
+    assert!(matches!(poll(&mut stream), Poll::Ready(None)));
+    assert!(matches!(poll(&mut stream), Poll::Ready(None)));
+    assert_eq!(polls.load(Ordering::SeqCst), 2);
+    assert_eq!(drops.load(Ordering::SeqCst), 1);
+    assert_eq!(pool.reserved(), 0);
+}
+
+#[test]
+fn unfinished_series_eof_is_error_once_and_drops_partial_chunks() {
+    let (input, _, drops) = raw([Poll::Ready(Some(Ok(chunk(&[10], &[1.0], false))))]);
+    let pool = Arc::new(GreedyMemoryPool::new(1_000_000)) as Arc<dyn MemoryPool>;
+    let mut stream = stream(input, &pool);
+    assert!(matches!(poll(&mut stream), Poll::Ready(Some(Err(_)))));
+    assert_eq!(stream.pending.len(), 0);
+    assert_eq!(drops.load(Ordering::SeqCst), 1);
+    assert_eq!(pool.reserved(), 0);
+    assert!(matches!(poll(&mut stream), Poll::Ready(None)));
+}
+
+#[test]
+fn dropping_pending_stream_releases_source_chunks_and_memory() {
+    let array = Arc::new(Float64Array::from(vec![1.0]));
+    let samples = Samples::from_float(array.as_ref().clone());
+    drop(array);
+    let pending_chunk = RawSeriesChunk::try_new(
+        &raw_schema(),
+        datatypes::arrow::buffer::Buffer::from(vec![1]),
+        Some(MetricName::try_new(None, vec![]).unwrap()),
+        RawTimestamps::try_new(ScalarBuffer::from(vec![10])).unwrap(),
+        SeriesFields::one(samples),
+        false,
+    )
+    .unwrap();
+    let (input, _, drops) = raw([Poll::Ready(Some(Ok(pending_chunk))), Poll::Pending]);
+    let pool = Arc::new(GreedyMemoryPool::new(1_000_000)) as Arc<dyn MemoryPool>;
+    let mut stream = stream(input, &pool);
+    assert!(matches!(poll(&mut stream), Poll::Pending));
+    assert!(pool.reserved() > 0);
+    drop(stream);
+    assert_eq!(drops.load(Ordering::SeqCst), 1);
+    assert_eq!(pool.reserved(), 0);
+}
+
+#[test]
+fn reservation_failure_is_terminal_and_releases_input_and_partial_series() {
+    let (input, _, drops) = raw([Poll::Ready(Some(Ok(chunk(&[10], &[1.0], false))))]);
+    let pool = Arc::new(GreedyMemoryPool::new(1)) as Arc<dyn MemoryPool>;
+    let mut stream = stream(input, &pool);
+    assert!(matches!(
+        poll(&mut stream),
+        Poll::Ready(Some(Err(DataFusionError::ResourcesExhausted(_))))
+    ));
+    assert_eq!(drops.load(Ordering::SeqCst), 1);
+    assert_eq!(stream.pending.len(), 0);
+    assert_eq!(stream.pending.capacity(), 0);
+    assert_eq!(pool.reserved(), 0);
+    assert!(matches!(poll(&mut stream), Poll::Ready(None)));
+}
+
+#[test]
+fn source_error_drops_unfinished_series_without_emitting_it() {
+    let failure = BoxedError::new(PlainError::new(
+        "mock source error".into(),
+        StatusCode::Internal,
+    ));
+    let (input, polls, drops) = raw([
+        Poll::Ready(Some(Ok(chunk(&[10], &[1.0], false)))),
+        Poll::Ready(Some(Err(failure))),
+    ]);
+    let pool = Arc::new(GreedyMemoryPool::new(1_000_000)) as Arc<dyn MemoryPool>;
+    let mut stream = stream(input, &pool);
+    assert!(matches!(poll(&mut stream), Poll::Ready(Some(Err(_)))));
+    assert!(stream.pending.is_empty());
+    assert_eq!(stream.pending.capacity(), 0);
+    assert!(matches!(poll(&mut stream), Poll::Ready(None)));
+    assert_eq!(polls.load(Ordering::SeqCst), 2);
+    assert_eq!(drops.load(Ordering::SeqCst), 1);
+    assert_eq!(pool.reserved(), 0);
+}
+
+#[test]
+fn descriptor_heavy_series_fails_reservation_before_accumulating_unbounded_chunks() {
+    let items = (0..512)
+        .map(|i| Poll::Ready(Some(Ok(chunk(&[i], &[i as f64], false)))))
+        .chain([Poll::Pending]);
+    let (input, polls, drops) = raw(items);
+    let pool = Arc::new(GreedyMemoryPool::new(32 * 1024)) as Arc<dyn MemoryPool>;
+    let mut stream = stream(input, &pool);
+    assert!(matches!(
+        poll(&mut stream),
+        Poll::Ready(Some(Err(DataFusionError::ResourcesExhausted(_))))
+    ));
+    assert!(stream.pending.is_empty());
+    assert_eq!(stream.pending.capacity(), 0);
+    assert!(matches!(poll(&mut stream), Poll::Ready(None)));
+    assert!(polls.load(Ordering::SeqCst) < 512);
+    assert_eq!(drops.load(Ordering::SeqCst), 1);
+    assert_eq!(pool.reserved(), 0);
+}
+
+#[test]
+fn dense_large_chunk_uses_segment_count_not_row_count_for_descriptors() {
+    let count = 100_000usize;
+    let timestamps = (0..count).map(|i| i as i64).collect::<Vec<_>>();
+    let values = Float64Array::from(vec![7.0; count]);
+    let dense_chunk = RawSeriesChunk::try_new(
+        &raw_schema(),
+        datatypes::arrow::buffer::Buffer::from(vec![1]),
+        Some(MetricName::try_new(None, vec![]).unwrap()),
+        RawTimestamps::try_new(ScalarBuffer::from(timestamps)).unwrap(),
+        SeriesFields::one(Samples::from_float(values)),
+        true,
+    )
+    .unwrap();
+    let (input, _, _) = raw([Poll::Ready(Some(Ok(dense_chunk)))]);
+    // Input timestamp/field buffers and concatenation scratch fit; row-scaled Samples
+    // descriptors would exceed this pool despite there being only one dense segment.
+    let pool = Arc::new(GreedyMemoryPool::new(4 * 1024 * 1024)) as Arc<dyn MemoryPool>;
+    let mut stream = stream_for(input, &pool, 10, 11, 1, 0);
+    let Poll::Ready(Some(Ok(batch))) = poll(&mut stream) else {
+        panic!("dense chunk should fit its bounded pool")
+    };
+    assert_eq!(batch.num_rows(), 2);
+    assert_eq!(
+        batch
+            .column(0)
+            .as_any()
+            .downcast_ref::<Float64Array>()
+            .unwrap()
+            .values(),
+        &[7.0, 7.0]
+    );
+    assert_eq!(
+        batch
+            .column(1)
+            .as_any()
+            .downcast_ref::<TimestampMillisecondArray>()
+            .unwrap()
+            .values(),
+        &[10, 11]
+    );
+    assert_eq!(pool.reserved(), 0);
+}
+
+#[test]
+fn output_grid_scratch_is_reserved_before_large_evaluation() {
+    let (input, polls, drops) = raw([Poll::Ready(Some(Ok(chunk(&[0], &[7.0], true))))]);
+    let pool = Arc::new(GreedyMemoryPool::new(128 * 1024)) as Arc<dyn MemoryPool>;
+    // One source point with a long lookback produces a large grid despite the tiny source.
+    let mut stream = stream_for(input, &pool, 0, 900_000, 1, 900_001);
+    assert!(matches!(
+        poll(&mut stream),
+        Poll::Ready(Some(Err(DataFusionError::ResourcesExhausted(_))))
+    ));
+    assert_eq!(polls.load(Ordering::SeqCst), 1);
+    assert_eq!(drops.load(Ordering::SeqCst), 1);
+    assert!(stream.pending.is_empty());
+    assert_eq!(pool.reserved(), 0);
+    assert!(matches!(poll(&mut stream), Poll::Ready(None)));
+}
+
+#[test]
+fn larger_pool_completes_one_point_many_evaluation_grid() {
+    let (input, _, _) = raw([Poll::Ready(Some(Ok(chunk(&[0], &[7.0], true))))]);
+    let pool = Arc::new(GreedyMemoryPool::new(32 * 1024 * 1024)) as Arc<dyn MemoryPool>;
+    let mut stream = stream_for(input, &pool, 0, 4, 1, 5);
+    let Poll::Ready(Some(Ok(batch))) = poll(&mut stream) else {
+        panic!("sufficient reservation should evaluate")
+    };
+    assert_eq!(
+        batch
+            .column(0)
+            .as_any()
+            .downcast_ref::<Float64Array>()
+            .unwrap()
+            .values(),
+        &[7.0; 5]
+    );
+    assert_eq!(
+        batch
+            .column(1)
+            .as_any()
+            .downcast_ref::<TimestampMillisecondArray>()
+            .unwrap()
+            .values(),
+        &[0, 1, 2, 3, 4]
+    );
+    assert_eq!(pool.reserved(), 0);
+}

src/promql/src/extension_plan/instant_manipulate/native_tests.rs

@@ -0,0 +1,597 @@
+// 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::HashMap;
+use std::sync::Arc;
+
+use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS;
+use datafusion::arrow::array::{Array, Float64Array, StringArray, TimestampMillisecondArray};
+use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit};
+use datafusion::common::DFSchema;
+use datafusion::logical_expr::{EmptyRelation, LogicalPlan};
+use datatypes::arrow::buffer::ScalarBuffer;
+use datatypes::series::raw::RawTimestamps;
+use datatypes::series::samples::{FieldSpec, Samples, SeriesFields, SeriesSchema};
+use datatypes::series::{Label, MetricName};
+
+use super::*;
+
+fn node() -> InstantManipulate {
+    let input = LogicalPlan::EmptyRelation(EmptyRelation {
+        produce_one_row: false,
+        schema: Arc::new(DFSchema::empty()),
+    });
+    InstantManipulate::new(
+        10,
+        20,
+        10,
+        10,
+        0,
+        "time".into(),
+        vec![],
+        Some("value".into()),
+        input,
+    )
+}
+
+fn schema(tag_type: DataType) -> SchemaRef {
+    let mut metadata = HashMap::new();
+    metadata.insert("schema".to_string(), "kept".to_string());
+    Arc::new(Schema::new_with_metadata(
+        vec![
+            Field::new("value", DataType::Float64, true)
+                .with_metadata([("field-key".to_string(), "field-value".to_string())].into()),
+            Field::new(
+                "time",
+                DataType::Timestamp(TimeUnit::Millisecond, None),
+                false,
+            ),
+            Field::new("zone", tag_type, true),
+        ],
+        metadata,
+    ))
+}
+
+fn series_schema() -> Arc<SeriesSchema> {
+    Arc::new(SeriesSchema {
+        fields: Arc::from([FieldSpec {
+            name: Arc::from("value"),
+        }]),
+        eval_grid: None,
+    })
+}
+
+fn metric(zone: Option<&str>) -> MetricName {
+    let labels = zone
+        .map(|value| {
+            vec![Label {
+                name: Arc::from("zone"),
+                value: Arc::from(value),
+            }]
+        })
+        .unwrap_or_default();
+    MetricName::try_new(None, labels).unwrap()
+}
+
+/// The ordered, deliberately non-lexical field names of the multifloat regression.
+const MULTIFLOAT_FIELDS: [&str; 3] = ["value", "a_second", "m_third"];
+
+/// The strict millisecond timestamps of every multifloat case.
+const MULTIFLOAT_TIMESTAMPS: [i64; 5] = [1000, 2000, 3000, 4000, 5000];
+
+/// The kept source rows of each configured staleness field: the stale row of the configured
+/// field is suppressed, so the first field drops row 3 (4000) and the second field drops row 2
+/// (3000) without ever retrying an older sample.
+const MULTIFLOAT_KEPT: [[usize; 4]; 2] = [[0, 1, 2, 4], [0, 1, 3, 4]];
+
+/// An ordinary, non-stale NaN payload that must survive the mapping as exact bits.
+const ORDINARY_NAN_BITS: u64 = 0x7ff8_0000_0000_0001;
+
+/// The output schema of the regression: the declared field order, then the time index and the
+/// tag, with metadata that every case must retain.
+///
+/// Every declared field is nullable because every fixture lane keeps a NULL row in each of them:
+/// the first row of the second and third fields and the second row of the first field.
+fn multifloat_schema(field_names: &[&str], tag_type: DataType) -> SchemaRef {
+    let mut metadata = HashMap::new();
+    metadata.insert("schema".to_string(), "kept".to_string());
+    let mut fields = field_names
+        .iter()
+        .enumerate()
+        .map(|(index, name)| {
+            Field::new(*name, DataType::Float64, true)
+                .with_metadata([("field-key".to_string(), format!("field-value-{index}"))].into())
+        })
+        .collect::<Vec<_>>();
+    fields.push(Field::new(
+        "time",
+        DataType::Timestamp(TimeUnit::Millisecond, None),
+        false,
+    ));
+    fields.push(Field::new("zone", tag_type, true));
+    Arc::new(Schema::new_with_metadata(fields, metadata))
+}
+
+/// The series schema of the same declared field order.
+fn multifloat_series_schema(field_names: &[&str]) -> Arc<SeriesSchema> {
+    Arc::new(SeriesSchema {
+        fields: field_names
+            .iter()
+            .map(|name| FieldSpec {
+                name: Arc::from(*name),
+            })
+            .collect::<Vec<_>>()
+            .into(),
+        eval_grid: None,
+    })
+}
+
+/// The five source rows of each multifloat field.
+///
+/// Row `i` is the strict millisecond timestamp `MULTIFLOAT_TIMESTAMPS[i]`. The first field is
+/// stale at 4000 and carries an ordinary NaN payload at 5000; the unrelated second field is
+/// stale at 3000; the fields carry NULLs and a signed zero as well.
+fn multifloat_source(field_names: &[&str]) -> Vec<Float64Array> {
+    let stale = f64::from_bits(PROMETHEUS_STALE_NAN_BITS);
+    let nan = f64::from_bits(ORDINARY_NAN_BITS);
+    let rows: [&[Option<f64>]; 3] = [
+        &[Some(1.0), None, Some(-0.0), Some(stale), Some(nan)],
+        &[None, Some(-0.0), Some(stale), Some(2.0), Some(4.0)],
+        &[None, Some(3.0), Some(nan), Some(3.0), Some(5.0)],
+    ];
+    field_names
+        .iter()
+        .enumerate()
+        .map(|(index, _)| Float64Array::from(rows[index].to_vec()))
+        .collect()
+}
+
+/// The chunked samples of the multifloat input: the first field arrives in three chunks, two of
+/// which are slices with non-zero offsets, and every other field arrives as one dense array.
+fn multifloat_fields(source: &[Float64Array]) -> SeriesFields {
+    let fields = source
+        .iter()
+        .enumerate()
+        .map(|(index, values)| {
+            if index > 0 {
+                return Samples::from_float(values.clone());
+            }
+            Samples::try_from_chunks(vec![
+                Samples::from_float(values.slice(0, 1)),
+                Samples::from_float(values.slice(1, 3)),
+                Samples::from_float(values.slice(4, 1)),
+            ])
+            .unwrap()
+        })
+        .collect::<Vec<_>>();
+    SeriesFields::try_from_vec(fields).unwrap()
+}
+
+/// The bounded instant declaration of the regression: one output point per input millisecond
+/// with one lookback step, whose staleness field is exactly one declared field.
+fn multifloat_node(field_column: &str) -> InstantManipulate {
+    let mut query = node();
+    query.start = MULTIFLOAT_TIMESTAMPS[0];
+    query.end = MULTIFLOAT_TIMESTAMPS[4];
+    query.lookback_delta = 1000;
+    query.interval = 1000;
+    query.field_column = Some(field_column.to_string());
+    query
+}
+
+#[test]
+fn evaluates_sliced_chunked_samples_and_preserves_original_schema() {
+    let output_schema = schema(DataType::Dictionary(
+        Box::new(DataType::UInt32),
+        Box::new(DataType::Utf8),
+    ));
+    let mapping = NativeInstantMapping::try_new(
+        &node(),
+        output_schema.clone(),
+        series_schema(),
+        &["zone".into()],
+    )
+    .unwrap()
+    .unwrap();
+    let a = Float64Array::from(vec![Some(99.0), Some(1.0), None]);
+    let b = Float64Array::from(vec![Some(2.0), Some(3.0)]);
+    let samples = Samples::try_from_chunks(vec![
+        Samples::from_float(
+            a.slice(1, 2)
+                .as_any()
+                .downcast_ref::<Float64Array>()
+                .unwrap()
+                .clone(),
+        ),
+        Samples::from_float(b),
+    ])
+    .unwrap();
+    let fields = SeriesFields::one(samples);
+    let time = RawTimestamps::try_new(ScalarBuffer::from(vec![0, 10, 20, 30])).unwrap();
+    let batch = mapping.evaluate(&metric(Some("")), &time, &fields).unwrap();
+    assert_eq!(batch.num_rows(), 2);
+    assert_eq!(batch.schema(), output_schema);
+    let output_values = batch
+        .column(0)
+        .as_any()
+        .downcast_ref::<Float64Array>()
+        .unwrap();
+    assert!(output_values.is_null(0));
+    assert_eq!(output_values.value(1), 2.0);
+    assert_eq!(
+        batch
+            .column(1)
+            .as_any()
+            .downcast_ref::<TimestampMillisecondArray>()
+            .unwrap()
+            .values(),
+        &[10, 20]
+    );
+    let tag_values = datafusion::arrow::compute::cast(batch.column(2), &DataType::Utf8).unwrap();
+    let tag_values = tag_values.as_any().downcast_ref::<StringArray>().unwrap();
+    assert_eq!(tag_values.value(0), "");
+    assert_eq!(tag_values.value(1), "");
+}
+
+#[test]
+fn selected_repeated_indices_across_sliced_chunks_preserve_missing_and_nan() {
+    let mut query = node();
+    query.end = 30;
+    query.lookback_delta = 20;
+    let mapping = NativeInstantMapping::try_new(
+        &query,
+        schema(DataType::Utf8),
+        series_schema(),
+        &["zone".into()],
+    )
+    .unwrap()
+    .unwrap();
+    let first = Float64Array::from(vec![Some(99.0), None]).slice(1, 1);
+    let second = Float64Array::from(vec![Some(f64::NAN), Some(88.0)]).slice(0, 1);
+    let samples = Samples::try_from_chunks(vec![
+        Samples::from_float(
+            first
+                .as_any()
+                .downcast_ref::<Float64Array>()
+                .unwrap()
+                .clone(),
+        ),
+        Samples::from_float(
+            second
+                .as_any()
+                .downcast_ref::<Float64Array>()
+                .unwrap()
+                .clone(),
+        ),
+    ])
+    .unwrap();
+    let fields = SeriesFields::one(samples);
+    let time = RawTimestamps::try_new(ScalarBuffer::from(vec![10, 30])).unwrap();
+
+    let batch = mapping.evaluate(&metric(None), &time, &fields).unwrap();
+    let values = batch
+        .column(0)
+        .as_any()
+        .downcast_ref::<Float64Array>()
+        .unwrap();
+    assert_eq!(batch.num_rows(), 3);
+    assert!(values.is_null(0));
+    assert!(values.is_null(1));
+    assert!(!values.is_null(2));
+    assert!(values.value(2).is_nan());
+}
+
+#[test]
+fn maps_all_supported_string_types_and_distinguishes_missing_from_empty() {
+    let tag_types = [
+        DataType::Utf8,
+        DataType::LargeUtf8,
+        DataType::Utf8View,
+        DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)),
+    ];
+    let fields = SeriesFields::one(Samples::from_float(Float64Array::from(vec![
+        Some(1.0),
+        Some(2.0),
+    ])));
+    let time = RawTimestamps::try_new(ScalarBuffer::from(vec![10, 20])).unwrap();
+
+    for tag_type in tag_types {
+        let output_schema = schema(tag_type);
+        let mapping = NativeInstantMapping::try_new(
+            &node(),
+            output_schema.clone(),
+            series_schema(),
+            &["zone".into()],
+        )
+        .unwrap()
+        .unwrap();
+
+        for (label, expected) in [(Some(""), Some("")), (None, None)] {
+            let batch = mapping.evaluate(&metric(label), &time, &fields).unwrap();
+            assert_eq!(batch.schema(), output_schema);
+            let values =
+                datafusion::arrow::compute::cast(batch.column(2), &DataType::Utf8).unwrap();
+            let values = values.as_any().downcast_ref::<StringArray>().unwrap();
+            match expected {
+                Some(value) => {
+                    assert_eq!(values.value(0), value);
+                    assert_eq!(values.value(1), value);
+                    assert_eq!(values.null_count(), 0);
+                }
+                None => assert_eq!(values.null_count(), 2),
+            }
+        }
+    }
+}
+
+#[test]
+fn selected_field_staleness_and_no_field_semantics_match_existing_selector_contract() {
+    let stale = f64::from_bits(PROMETHEUS_STALE_NAN_BITS);
+    let times = RawTimestamps::try_new(ScalarBuffer::from(vec![10, 20])).unwrap();
+    let values = SeriesFields::one(Samples::from_float(Float64Array::from(vec![
+        Some(1.0),
+        Some(stale),
+    ])));
+    let mapping = NativeInstantMapping::try_new(
+        &node(),
+        schema(DataType::Utf8),
+        series_schema(),
+        &["zone".into()],
+    )
+    .unwrap()
+    .unwrap();
+    let batch = mapping.evaluate(&metric(None), &times, &values).unwrap();
+    assert_eq!(batch.num_rows(), 1);
+    assert_eq!(
+        batch
+            .column(1)
+            .as_any()
+            .downcast_ref::<TimestampMillisecondArray>()
+            .unwrap()
+            .value(0),
+        10
+    );
+
+    let mut no_staleness_node = node();
+    no_staleness_node.field_column = None;
+    let no_staleness_mapping = NativeInstantMapping::try_new(
+        &no_staleness_node,
+        schema(DataType::Utf8),
+        series_schema(),
+        &["zone".into()],
+    )
+    .unwrap()
+    .unwrap();
+    let batch = no_staleness_mapping
+        .evaluate(&metric(None), &times, &values)
+        .unwrap();
+    let output = batch
+        .column(0)
+        .as_any()
+        .downcast_ref::<Float64Array>()
+        .unwrap();
+    assert_eq!(output.len(), 2);
+    assert_eq!(output.value(1).to_bits(), PROMETHEUS_STALE_NAN_BITS);
+}
+
+#[test]
+fn rejects_unmapped_missing_duplicate_and_non_millisecond_schema_columns() {
+    let valid = schema(DataType::Utf8);
+    let mut cases = Vec::new();
+
+    cases.push(Arc::new(Schema::new(vec![
+        valid.field(0).clone(),
+        valid.field(1).clone(),
+        valid.field(2).clone(),
+        Field::new("helper", DataType::UInt64, true),
+    ])));
+    cases.push(Arc::new(Schema::new(vec![
+        valid.field(0).clone(),
+        valid.field(1).clone(),
+        valid.field(2).clone(),
+        Field::new("__tsid", DataType::UInt64, true),
+    ])));
+    cases.push(Arc::new(Schema::new(vec![
+        valid.field(0).clone(),
+        valid.field(1).clone(),
+    ])));
+    cases.push(Arc::new(Schema::new(vec![
+        valid.field(1).clone(),
+        valid.field(2).clone(),
+    ])));
+    cases.push(Arc::new(Schema::new(vec![
+        valid.field(0).clone(),
+        valid.field(1).clone(),
+        valid.field(2).clone(),
+        valid.field(0).clone(),
+    ])));
+    cases.push(Arc::new(Schema::new(vec![
+        valid.field(0).clone(),
+        Field::new(
+            "time",
+            DataType::Timestamp(TimeUnit::Microsecond, None),
+            false,
+        ),
+        valid.field(2).clone(),
+    ])));
+    cases.push(Arc::new(Schema::new(vec![
+        Field::new("value", DataType::Int64, true),
+        valid.field(1).clone(),
+        valid.field(2).clone(),
+    ])));
+
+    for output_schema in cases {
+        assert!(
+            NativeInstantMapping::try_new(
+                &node(),
+                output_schema,
+                series_schema(),
+                &["zone".into()],
+            )
+            .unwrap()
+            .is_none()
+        );
+    }
+}
+
+#[test]
+fn absent_tag_is_null_and_nonzero_offset_is_rejected() {
+    let output_schema = schema(DataType::Utf8);
+    let mapping =
+        NativeInstantMapping::try_new(&node(), output_schema, series_schema(), &["zone".into()])
+            .unwrap()
+            .unwrap();
+    let fields = SeriesFields::one(Samples::from_float(Float64Array::from(vec![
+        Some(1.0),
+        Some(2.0),
+    ])));
+    let time = RawTimestamps::try_new(ScalarBuffer::from(vec![10, 20])).unwrap();
+    let batch = mapping.evaluate(&metric(None), &time, &fields).unwrap();
+    assert_eq!(batch.column(2).null_count(), 2);
+
+    let mut unsupported = node();
+    unsupported.offset = 1;
+    assert!(
+        NativeInstantMapping::try_new(
+            &unsupported,
+            schema(DataType::Utf8),
+            series_schema(),
+            &["zone".into()],
+        )
+        .unwrap()
+        .is_none()
+    );
+}
+
+#[test]
+fn multifloat_native_mapping_keeps_ordered_fields_and_checks_only_the_configured_field() {
+    let time = RawTimestamps::try_new(ScalarBuffer::from(MULTIFLOAT_TIMESTAMPS.to_vec())).unwrap();
+
+    for field_names in [&MULTIFLOAT_FIELDS[..2], &MULTIFLOAT_FIELDS[..3]] {
+        for (configured, kept) in MULTIFLOAT_KEPT.iter().enumerate() {
+            let output_schema = multifloat_schema(field_names, DataType::Utf8);
+            let source = multifloat_source(field_names);
+            let fields = multifloat_fields(&source);
+            let mapping = NativeInstantMapping::try_new(
+                &multifloat_node(field_names[configured]),
+                output_schema.clone(),
+                multifloat_series_schema(field_names),
+                &["zone".into()],
+            )
+            .unwrap()
+            .expect("the complete ordered Float64 fields are mapped instead of declined");
+            let batch = mapping
+                .evaluate(&metric(Some("multifloat")), &time, &fields)
+                .unwrap();
+
+            // The complete output schema, including its metadata and nullability, is retained.
+            assert_eq!(batch.schema(), output_schema);
+            assert_eq!(batch.num_rows(), kept.len());
+            let kept_timestamps = batch
+                .column(field_names.len())
+                .as_any()
+                .downcast_ref::<TimestampMillisecondArray>()
+                .unwrap()
+                .values()
+                .to_vec();
+            assert_eq!(
+                kept_timestamps,
+                kept.iter()
+                    .map(|row| MULTIFLOAT_TIMESTAMPS[*row])
+                    .collect::<Vec<_>>()
+            );
+            let suppressed = if configured == 0 { 3 } else { 2 };
+            let unrelated = if configured == 0 { 2 } else { 3 };
+            assert!(
+                !kept_timestamps.contains(&MULTIFLOAT_TIMESTAMPS[suppressed]),
+                "the stale row of the configured field is suppressed instead of retried"
+            );
+            assert!(
+                kept_timestamps.contains(&MULTIFLOAT_TIMESTAMPS[unrelated]),
+                "the stale marker of another field does not suppress its row"
+            );
+
+            // Every output field keeps the nullness and the exact bits of its own source row.
+            for (index, values) in source.iter().enumerate() {
+                let output = batch
+                    .column(index)
+                    .as_any()
+                    .downcast_ref::<Float64Array>()
+                    .unwrap();
+                for (row, source_row) in kept.iter().enumerate() {
+                    if values.is_null(*source_row) {
+                        assert!(output.is_null(row));
+                    } else {
+                        assert!(!output.is_null(row));
+                        assert_eq!(
+                            output.value(row).to_bits(),
+                            values.value(*source_row).to_bits(),
+                            "field {index} row {row} must keep its exact source bits"
+                        );
+                    }
+                }
+            }
+
+            // The unchecked stale marker travels through as exact bits, so the row of another
+            // field is kept with the marker intact instead of being compared or re-encoded.
+            let (unchecked_field, unchecked_row) = if configured == 0 { (1, 2) } else { (0, 3) };
+            let unchecked_output_row = kept
+                .iter()
+                .position(|row| *row == unchecked_row)
+                .expect("the unchecked stale row is kept");
+            let marker = batch
+                .column(unchecked_field)
+                .as_any()
+                .downcast_ref::<Float64Array>()
+                .unwrap()
+                .value(unchecked_output_row);
+            assert_eq!(marker.to_bits(), PROMETHEUS_STALE_NAN_BITS);
+
+            // The ordinary NaN payload and the signed zero of the first two fields survive as
+            // exact bits as well.
+            let nan_row = kept
+                .iter()
+                .position(|row| *row == 4)
+                .expect("the ordinary NaN row is kept");
+            let nan = batch
+                .column(0)
+                .as_any()
+                .downcast_ref::<Float64Array>()
+                .unwrap()
+                .value(nan_row);
+            assert_eq!(nan.to_bits(), ORDINARY_NAN_BITS);
+            let signed_zero_row = kept
+                .iter()
+                .position(|row| *row == 1)
+                .expect("the signed zero row is kept");
+            let signed_zero = batch
+                .column(1)
+                .as_any()
+                .downcast_ref::<Float64Array>()
+                .unwrap()
+                .value(signed_zero_row);
+            assert_eq!(signed_zero.to_bits(), (-0.0_f64).to_bits());
+
+            // The series key of every kept row is the requested label value.
+            let tags = datafusion::arrow::compute::cast(
+                batch.column(field_names.len() + 1),
+                &DataType::Utf8,
+            )
+            .unwrap();
+            let tags = tags.as_any().downcast_ref::<StringArray>().unwrap();
+            assert_eq!(tags.null_count(), 0);
+            assert!(tags.iter().all(|value| value == Some("multifloat")));
+        }
+    }
+}