fix(promql): retain column indices in instant plan ordering

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

src/promql/src/extension_plan.rs

@@ -52,25 +52,25 @@ pub use union_distinct_on::{UnionDistinctOn, UnionDistinctOnExec, UnionDistinctO
 
 pub type Millisecond = <TimestampMillisecondType as ArrowPrimitiveType>::Native;
 
-/// Returns a timestamp value without reducing its Arrow storage precision.
-pub(crate) fn native_timestamp_values(array: &dyn Array) -> datafusion::error::Result<Vec<i64>> {
+/// Borrows timestamp values without reducing their Arrow storage precision.
+pub(crate) fn native_timestamp_values(array: &dyn Array) -> datafusion::error::Result<&[i64]> {
     let value = match array.data_type() {
         DataType::Timestamp(TimeUnit::Second, _) => array
             .as_any()
             .downcast_ref::<TimestampSecondArray>()
-            .map(|a| a.values().to_vec()),
+            .map(|a| a.values().as_ref()),
         DataType::Timestamp(TimeUnit::Millisecond, _) => array
             .as_any()
             .downcast_ref::<TimestampMillisecondArray>()
-            .map(|a| a.values().to_vec()),
+            .map(|a| a.values().as_ref()),
         DataType::Timestamp(TimeUnit::Microsecond, _) => array
             .as_any()
             .downcast_ref::<TimestampMicrosecondArray>()
-            .map(|a| a.values().to_vec()),
+            .map(|a| a.values().as_ref()),
         DataType::Timestamp(TimeUnit::Nanosecond, _) => array
             .as_any()
             .downcast_ref::<TimestampNanosecondArray>()
-            .map(|a| a.values().to_vec()),
+            .map(|a| a.values().as_ref()),
         _ => None,
     };
     value.ok_or_else(|| {
@@ -87,7 +87,7 @@ pub(crate) fn timestamp_unit(data_type: &DataType) -> datafusion::error::Result<
     }
 }
 
-pub(crate) fn native_per_nanosecond(unit: TimeUnit) -> i128 {
+pub(crate) fn nanoseconds_per_native_tick(unit: TimeUnit) -> i128 {
     match unit {
         TimeUnit::Second => 1_000_000_000,
         TimeUnit::Millisecond => 1_000_000,

src/promql/src/extension_plan/instant_manipulate.rs

@@ -46,7 +46,7 @@ use snafu::ResultExt;
 use crate::error::{DeserializeSnafu, Result};
 use crate::extension_plan::series_divide::SeriesDivide;
 use crate::extension_plan::{
-    METRIC_NUM_SERIES, Millisecond, is_prometheus_stale_sample, native_per_nanosecond,
+    METRIC_NUM_SERIES, Millisecond, is_prometheus_stale_sample, nanoseconds_per_native_tick,
     native_timestamp_values, prometheus_stale_sample_column, resolve_column_name,
     serialize_column_index, timestamp_unit,
 };
@@ -95,6 +95,7 @@ impl PartialOrd for InstantManipulate {
             &self.tag_columns,
             &self.field_column,
             &self.input,
+            &self.unfix,
         )
             .partial_cmp(&(
                 other.start,
@@ -105,6 +106,7 @@ impl PartialOrd for InstantManipulate {
                 &other.tag_columns,
                 &other.field_column,
                 &other.input,
+                &other.unfix,
             ))
     }
 }
@@ -645,7 +647,7 @@ impl InstantManipulateStream {
         if ts_column.is_empty() {
             return Ok(RecordBatch::new_empty(self.schema.clone()));
         }
-        let scale = native_per_nanosecond(self.time_unit);
+        let scale = nanoseconds_per_native_tick(self.time_unit);
         let stale_sample_columns = self.field_indices.map(|index| {
             index.and_then(|index| prometheus_stale_sample_column(input.column(index).as_ref()))
         });
@@ -958,6 +960,20 @@ mod test {
         }
     }
 
+    #[test]
+    fn deserialized_ordering_preserves_column_indices() {
+        let mut wire = pb::InstantManipulate::default();
+        let first = InstantManipulate::deserialize(&wire.encode_to_vec()).unwrap();
+        wire.time_index_idx = 1;
+        let second = InstantManipulate::deserialize(&wire.encode_to_vec()).unwrap();
+        assert_ne!(first, second);
+        assert_eq!(first.partial_cmp(&second), Some(std::cmp::Ordering::Less));
+        wire.field_index_idx = 2;
+        let third = InstantManipulate::deserialize(&wire.encode_to_vec()).unwrap();
+        assert_ne!(second, third);
+        assert_eq!(second.partial_cmp(&third), Some(std::cmp::Ordering::Less));
+    }
+
     #[test]
     fn pruning_should_keep_time_and_field_columns_for_exec() {
         let df_schema = prepare_test_data().schema().to_dfschema_ref().unwrap();

src/promql/src/extension_plan/range_manipulate.rs

@@ -46,7 +46,7 @@ use snafu::ResultExt;
 
 use crate::error::{DeserializeSnafu, Result};
 use crate::extension_plan::{
-    METRIC_NUM_SERIES, Millisecond, native_per_nanosecond, native_timestamp_values,
+    METRIC_NUM_SERIES, Millisecond, nanoseconds_per_native_tick, native_timestamp_values,
     resolve_column_name, serialize_column_index, timestamp_unit,
 };
 use crate::metrics::PROMQL_SERIES_COUNT;
@@ -714,7 +714,7 @@ impl RangeManipulateStream {
         input: &RecordBatch,
     ) -> DataFusionResult<(Vec<(u32, u32)>, (i64, i64))> {
         let ts_column = input.column(self.time_index);
-        let scale = native_per_nanosecond(self.time_unit);
+        let scale = nanoseconds_per_native_tick(self.time_unit);
         let timestamps = native_timestamp_values(ts_column.as_ref())?;
         let timestamp = |index| (timestamps[index] as i128) * scale;
         let len = timestamps.len();