From 4f08feeffc563da852375a99fd9c1ce024c11603 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:58:40 +0800 Subject: [PATCH] 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 | 14 ++++++------- .../src/extension_plan/instant_manipulate.rs | 20 +++++++++++++++++-- .../src/extension_plan/range_manipulate.rs | 4 ++-- 3 files changed, 27 insertions(+), 11 deletions(-) diff --git a/src/promql/src/extension_plan.rs b/src/promql/src/extension_plan.rs index 2a94763766..2ae9230f72 100644 --- a/src/promql/src/extension_plan.rs +++ b/src/promql/src/extension_plan.rs @@ -52,25 +52,25 @@ pub use union_distinct_on::{UnionDistinctOn, UnionDistinctOnExec, UnionDistinctO pub type Millisecond = ::Native; -/// Returns a timestamp value without reducing its Arrow storage precision. -pub(crate) fn native_timestamp_values(array: &dyn Array) -> datafusion::error::Result> { +/// 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::() - .map(|a| a.values().to_vec()), + .map(|a| a.values().as_ref()), DataType::Timestamp(TimeUnit::Millisecond, _) => array .as_any() .downcast_ref::() - .map(|a| a.values().to_vec()), + .map(|a| a.values().as_ref()), DataType::Timestamp(TimeUnit::Microsecond, _) => array .as_any() .downcast_ref::() - .map(|a| a.values().to_vec()), + .map(|a| a.values().as_ref()), DataType::Timestamp(TimeUnit::Nanosecond, _) => array .as_any() .downcast_ref::() - .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, diff --git a/src/promql/src/extension_plan/instant_manipulate.rs b/src/promql/src/extension_plan/instant_manipulate.rs index 7a76ab83ba..0f3d91b7de 100644 --- a/src/promql/src/extension_plan/instant_manipulate.rs +++ b/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(); diff --git a/src/promql/src/extension_plan/range_manipulate.rs b/src/promql/src/extension_plan/range_manipulate.rs index dab3596395..f4f6795270 100644 --- a/src/promql/src/extension_plan/range_manipulate.rs +++ b/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();