diff --git a/src/promql/src/extension_plan.rs b/src/promql/src/extension_plan.rs index 2ae9230f72..a8fcd10ca4 100644 --- a/src/promql/src/extension_plan.rs +++ b/src/promql/src/extension_plan.rs @@ -35,8 +35,9 @@ use datafusion::arrow::array::{ use datafusion::arrow::datatypes::{ ArrowPrimitiveType, DataType, TimeUnit, TimestampMillisecondType, }; -use datafusion::common::DFSchemaRef; +use datafusion::common::{Column, DFSchemaRef}; use datafusion::error::{DataFusionError, Result as DataFusionResult}; +use datafusion::logical_expr::{Expr, Extension, LogicalPlan}; use datatypes::data_type::DataType as _; pub use empty_metric::{EmptyMetric, EmptyMetricExec, EmptyMetricStream, build_special_time_expr}; pub use histogram_fold::{ @@ -96,6 +97,57 @@ pub(crate) fn nanoseconds_per_native_tick(unit: TimeUnit) -> i128 { } } +/// Returns the offset of an immediately underlying normalize node when the +/// requested time index retains its logical identity through projections. +pub(crate) fn local_offset(plan: &LogicalPlan, time_index: &str) -> Millisecond { + let Some(index) = plan.schema().index_of_column_by_name(None, time_index) else { + return 0; + }; + let (qualifier, field) = plan.schema().qualified_field(index); + let mut time_index = Column::new(qualifier.cloned(), field.name().clone()); + let mut plan = plan; + + loop { + match plan { + LogicalPlan::Extension(Extension { node }) => { + return node + .as_any() + .downcast_ref::() + .and_then(|normalize| normalize.offset_for_time_index(&time_index)) + .unwrap_or_default(); + } + LogicalPlan::Projection(projection) => { + let Some(output_index) = projection.schema.maybe_index_of_column(&time_index) + else { + return 0; + }; + let expr = &projection.expr[output_index]; + let source = match expr { + Expr::Column(column) => column, + Expr::Alias(alias) => { + let Expr::Column(column) = alias.expr.as_ref() else { + return 0; + }; + if alias.name != column.name { + return 0; + } + column + } + _ => return 0, + }; + let Some(input_index) = projection.input.schema().maybe_index_of_column(source) + else { + return 0; + }; + let (qualifier, field) = projection.input.schema().qualified_field(input_index); + time_index = Column::new(qualifier.cloned(), field.name().clone()); + plan = projection.input.as_ref(); + } + _ => return 0, + } + } +} + const METRIC_NUM_SERIES: &str = "num_series"; fn prometheus_stale_sample_column(column: &dyn Array) -> Option<(&dyn Array, &Float64Array)> { @@ -158,3 +210,87 @@ pub fn resolve_column_names( .map(|idx| resolve_column_name(*idx, schema, context, column_type)) .collect() } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit}; + use datafusion::common::ToDFSchema; + use datafusion::logical_expr::{EmptyRelation, Extension, LogicalPlan, Projection}; + use datafusion_expr::col; + + use super::*; + + fn input() -> LogicalPlan { + LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: Arc::new(Schema::new(vec![ + Field::new( + "timestamp", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new( + "other_ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("value", DataType::Float64, true), + ])) + .to_dfschema_ref() + .unwrap(), + }) + } + + fn normalized() -> LogicalPlan { + LogicalPlan::Extension(Extension { + node: Arc::new(SeriesNormalize::new( + 1_000, + "timestamp", + false, + Vec::new(), + input(), + )), + }) + } + + #[test] + fn local_offset_tracks_identity_preserving_projections() { + let projection = + Projection::try_new(vec![col("timestamp"), col("value")], Arc::new(normalized())) + .unwrap(); + let projection = Projection::try_new( + vec![col("timestamp").alias("timestamp"), col("value")], + Arc::new(LogicalPlan::Projection(projection)), + ) + .unwrap(); + + assert_eq!( + 1_000, + local_offset(&LogicalPlan::Projection(projection), "timestamp") + ); + } + + #[test] + fn local_offset_rejects_a_different_timestamp_or_manipulator() { + let renamed = Projection::try_new( + vec![col("other_ts").alias("timestamp"), col("value")], + Arc::new(normalized()), + ) + .unwrap(); + assert_eq!( + 0, + local_offset(&LogicalPlan::Projection(renamed), "timestamp") + ); + + let divide = LogicalPlan::Extension(Extension { + node: Arc::new(SeriesDivide::new( + Vec::new(), + "timestamp".to_string(), + normalized(), + )), + }); + assert_eq!(0, local_offset(÷, "timestamp")); + } +} diff --git a/src/promql/src/extension_plan/instant_manipulate.rs b/src/promql/src/extension_plan/instant_manipulate.rs index b0b50578e6..0c58afb598 100644 --- a/src/promql/src/extension_plan/instant_manipulate.rs +++ b/src/promql/src/extension_plan/instant_manipulate.rs @@ -46,9 +46,9 @@ 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, nanoseconds_per_native_tick, - native_timestamp_values, prometheus_stale_sample_column, resolve_column_name, - serialize_column_index, timestamp_unit, + METRIC_NUM_SERIES, Millisecond, is_prometheus_stale_sample, local_offset, + nanoseconds_per_native_tick, native_timestamp_values, prometheus_stale_sample_column, + resolve_column_name, serialize_column_index, timestamp_unit, }; use crate::metrics::PROMQL_SERIES_COUNT; @@ -354,6 +354,7 @@ impl InstantManipulate { input_properties.boundedness, )); Arc::new(InstantManipulateExec { + offset: local_offset(&self.input, &self.time_index_column), start: self.start, end: self.end, lookback_delta: self.lookback_delta, @@ -420,6 +421,7 @@ impl InstantManipulate { #[derive(Debug)] pub struct InstantManipulateExec { + offset: Millisecond, start: Millisecond, end: Millisecond, lookback_delta: Millisecond, @@ -474,6 +476,7 @@ impl ExecutionPlan for InstantManipulateExec { input_properties.boundedness, )); Ok(Arc::new(Self { + offset: self.offset, start: self.start, end: self.end, lookback_delta: self.lookback_delta, @@ -517,6 +520,7 @@ impl ExecutionPlan for InstantManipulateExec { .filter(|(_, field)| field.data_type() == &DataType::UInt64) .map(|(index, _)| index); Ok(Box::pin(InstantManipulateStream { + offset: self.offset, start: self.start, end: self.end, lookback_delta: self.lookback_delta, @@ -584,6 +588,7 @@ impl DisplayAs for InstantManipulateExec { } pub struct InstantManipulateStream { + offset: Millisecond, start: Millisecond, end: Millisecond, lookback_delta: Millisecond, @@ -654,7 +659,8 @@ impl InstantManipulateStream { }; let timestamps = native_timestamp_values(ts_column.as_ref())?; let len = timestamps.len(); - let to_nanoseconds = |timestamp: i64| (timestamp as i128) * scale; + let to_nanoseconds = + |timestamp: i64| (timestamp as i128) * scale + (self.offset as i128) * 1_000_000; let first_ns = to_nanoseconds(timestamps[0]); let last_ns = to_nanoseconds(timestamps[len - 1]); // An exact sample remains useful with zero lookback. Otherwise the lower @@ -780,7 +786,9 @@ mod test { use datafusion::common::ToDFSchema; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; - use datafusion::logical_expr::{EmptyRelation, LogicalPlan}; + use datafusion::logical_expr::{ + EmptyRelation, Extension, LogicalPlan, UserDefinedLogicalNodeCore, + }; use datafusion::prelude::SessionContext; use super::*; @@ -802,6 +810,7 @@ mod test { Arc::new(prepare_test_data()) }; let normalize_exec = Arc::new(InstantManipulateExec { + offset: 0, start, end, lookback_delta, @@ -955,6 +964,74 @@ mod test { } } + #[tokio::test] + async fn logical_normalize_offset_survives_rebuild_and_executes() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("value", DataType::Float64, true), + ])); + let input = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: schema.clone().to_dfschema_ref().unwrap(), + }); + let normalize = crate::extension_plan::SeriesNormalize::new( + 1_000, + TIME_INDEX_COLUMN, + false, + Vec::new(), + input.clone(), + ); + let normalize = crate::extension_plan::SeriesNormalize::deserialize(&normalize.serialize()) + .unwrap() + .with_exprs_and_inputs(vec![], vec![input]) + .unwrap(); + let normalized = LogicalPlan::Extension(Extension { + node: Arc::new(normalize), + }); + let plan = InstantManipulate::new( + 1_000, + 1_000, + 0, + 1, + TIME_INDEX_COLUMN.to_string(), + Vec::new(), + Some("value".to_string()), + normalized.clone(), + ); + let rebuilt = InstantManipulate::deserialize(&plan.serialize()) + .unwrap() + .with_exprs_and_inputs(vec![], vec![normalized]) + .unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![0])), + Arc::new(Float64Array::from(vec![7.0])), + ], + ) + .unwrap(); + let exec_input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), + ))); + let exec = rebuilt.to_execution_plan(exec_input); + let output = datafusion::physical_plan::collect(exec, SessionContext::default().task_ctx()) + .await + .unwrap(); + assert_eq!( + output[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 7.0 + ); + } + #[test] fn deserialized_ordering_preserves_column_indices() { let mut wire = pb::InstantManipulate::default(); @@ -1129,6 +1206,7 @@ mod test { MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), ))); let normalize_exec = Arc::new(InstantManipulateExec { + offset: 0, start: 0, end: 1_500, lookback_delta: 1_000, @@ -1192,6 +1270,7 @@ mod test { MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), ))); let normalize_exec = Arc::new(InstantManipulateExec { + offset: 0, start: 0, end: 1_500, lookback_delta: 1_000, @@ -1248,6 +1327,7 @@ mod test { ))); let too_many_points = MAX_INSTANT_MANIPULATE_OUTPUT_POINTS as Millisecond + 1; let normalize_exec = Arc::new(InstantManipulateExec { + offset: 0, start: 0, end: too_many_points, lookback_delta: too_many_points + 1, @@ -1567,6 +1647,7 @@ mod test { ) .unwrap(); let stream = InstantManipulateStream { + offset: 0, start: 1_000, end: 1_050, lookback_delta: 100, @@ -1615,6 +1696,7 @@ mod test { )])); let input = RecordBatch::new_empty(input_schema.clone()); let stream = InstantManipulateStream { + offset: 0, start: 0, end: 0, lookback_delta: 0, @@ -1664,6 +1746,7 @@ mod test { ) .unwrap(); let stream = InstantManipulateStream { + offset: 0, start: i64::MIN + 1, end: i64::MAX, lookback_delta: 0, @@ -1697,6 +1780,64 @@ mod test { assert_eq!(values.values(), &[7.0]); } + #[test] + fn native_nanosecond_offset_uses_wide_shifted_timeline() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + ), + Field::new("value", DataType::Float64, true), + ])); + let raw = 9_223_112_837_000_000_000_i64; + let input = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampNanosecondArray::from(vec![raw])), + Arc::new(Float64Array::from(vec![7.0])), + ], + ) + .unwrap(); + let stream = InstantManipulateStream { + offset: 3 * 24 * 60 * 60 * 1_000, + start: 9_223_372_037_000, + end: 9_223_372_037_000, + lookback_delta: 300_000, + interval: 1, + time_index: 0, + time_unit: TimeUnit::Nanosecond, + field_indices: [Some(1), None], + tsid_index: None, + reuse_tsid_column: false, + schema: Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("value", DataType::Float64, true), + ])), + input: Box::pin( + datafusion::physical_plan::memory::MemoryStream::try_new(vec![], schema, None) + .unwrap(), + ), + metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + num_series: Count::new(), + }; + let output = stream.manipulate(input).unwrap(); + assert_eq!(output.num_rows(), 1); + assert_eq!( + output + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 7.0 + ); + } + #[tokio::test] async fn ordinary_nan_is_selected_for_exact_and_lookback() { let schema = Arc::new(Schema::new(vec![ @@ -1719,6 +1860,7 @@ mod test { MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), ))); let exec = Arc::new(InstantManipulateExec { + offset: 0, start: 1_000, end: 1_500, lookback_delta: 1_000, @@ -1800,6 +1942,7 @@ mod test { MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), ))); let exec = Arc::new(InstantManipulateExec { + offset: 0, start: 750, end: 1_500, lookback_delta: 1_001, @@ -1865,6 +2008,7 @@ mod test { MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), ))); let exec = Arc::new(InstantManipulateExec { + offset: 0, start: 1_000, end: 1_500, lookback_delta: 1_001, @@ -1914,6 +2058,7 @@ mod test { MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), ))); let exec = Arc::new(InstantManipulateExec { + offset: 0, start: 1_000, end: 1_500, lookback_delta: 1_000, diff --git a/src/promql/src/extension_plan/normalize.rs b/src/promql/src/extension_plan/normalize.rs index fc499209ad..728e5d1a4a 100644 --- a/src/promql/src/extension_plan/normalize.rs +++ b/src/promql/src/extension_plan/normalize.rs @@ -20,7 +20,7 @@ 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::{DFSchema, DFSchemaRef, Result as DataFusionResult, Statistics}; +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}; @@ -33,14 +33,8 @@ use datafusion::physical_plan::{ SendableRecordBatchStream, }; use datafusion_expr::col; -use datatypes::arrow::array::{ - TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, - TimestampSecondArray, -}; -use datatypes::arrow::datatypes::{ - SchemaRef, TimeUnit, TimestampMicrosecondType, TimestampMillisecondType, - TimestampNanosecondType, TimestampSecondType, -}; +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; @@ -50,7 +44,7 @@ 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, timestamp_unit, + resolve_column_name, serialize_column_index, }; use crate::metrics::PROMQL_SERIES_COUNT; @@ -58,7 +52,7 @@ use crate::metrics::PROMQL_SERIES_COUNT; /// the input batch only contains sample points from one time series. /// /// Roughly speaking, this method does these things: -/// - bias sample and native histogram start timestamps by offset +/// - retain raw native sample timestamps while biasing native histogram start timestamps by offset /// - sort the record batch based on timestamp column /// - remove Prometheus stale markers (optional) #[derive(Debug, PartialEq, Eq, Hash, PartialOrd)] @@ -185,6 +179,14 @@ impl UserDefinedLogicalNodeCore for SeriesNormalize { } impl SeriesNormalize { + pub(crate) fn offset_for_time_index(&self, time_index: &Column) -> Option { + 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>( offset: Millisecond, time_index_column_name: N, @@ -335,13 +337,8 @@ impl ExecutionPlan for SeriesNormalizeExec { let input = self.input.execute(partition, context)?; let schema = input.schema(); - let time_index = schema - .column_with_name(&self.time_index_column_name) - .expect("time index column not found") - .0; Ok(Box::pin(SeriesNormalizeStream { offset: self.offset, - time_index, filter_stale_markers: self.filter_stale_markers, schema, input, @@ -381,8 +378,6 @@ impl DisplayAs for SeriesNormalizeExec { pub struct SeriesNormalizeStream { offset: Millisecond, - // Column index of TIME INDEX column's position in schema - time_index: usize, filter_stale_markers: bool, schema: SchemaRef, @@ -394,87 +389,12 @@ pub struct SeriesNormalizeStream { impl SeriesNormalizeStream { pub fn normalize(&self, input: RecordBatch) -> DataFusionResult { - let time_unit = timestamp_unit(input.column(self.time_index).data_type())?; - let offset = match time_unit { - TimeUnit::Second => (self.offset % 1_000 == 0).then_some(self.offset / 1_000), - TimeUnit::Millisecond => Some(self.offset), - TimeUnit::Microsecond => self.offset.checked_mul(1_000), - TimeUnit::Nanosecond => self.offset.checked_mul(1_000_000), - } - .ok_or_else(|| { - DataFusionError::Execution("SeriesNormalize: timestamp offset overflow".into()) - })?; - let bias_timestamp = |timestamp: i64| { - timestamp.checked_add(offset).ok_or_else(|| { - DataFusionError::Execution("SeriesNormalize: timestamp offset overflow".into()) - }) - }; - - // Bias timestamps in their native Arrow unit; histogram start timestamps - // intentionally remain millisecond payloads below. - let ts_column_biased: Arc = match time_unit { - TimeUnit::Second => { - let column = input - .column(self.time_index) - .as_any() - .downcast_ref::() - .ok_or_else(|| { - DataFusionError::Execution("Time index column downcast failed".into()) - })?; - if offset == 0 { - Arc::new(column.clone()) - } else { - Arc::new(column.try_unary::<_, TimestampSecondType, _>(&bias_timestamp)?) - } - } - TimeUnit::Millisecond => { - let column = input - .column(self.time_index) - .as_any() - .downcast_ref::() - .ok_or_else(|| { - DataFusionError::Execution("Time index column downcast failed".into()) - })?; - if offset == 0 { - Arc::new(column.clone()) - } else { - Arc::new(column.try_unary::<_, TimestampMillisecondType, _>(&bias_timestamp)?) - } - } - TimeUnit::Microsecond => { - let column = input - .column(self.time_index) - .as_any() - .downcast_ref::() - .ok_or_else(|| { - DataFusionError::Execution("Time index column downcast failed".into()) - })?; - if offset == 0 { - Arc::new(column.clone()) - } else { - Arc::new(column.try_unary::<_, TimestampMicrosecondType, _>(&bias_timestamp)?) - } - } - TimeUnit::Nanosecond => { - let column = input - .column(self.time_index) - .as_any() - .downcast_ref::() - .ok_or_else(|| { - DataFusionError::Execution("Time index column downcast failed".into()) - })?; - if offset == 0 { - Arc::new(column.clone()) - } else { - Arc::new(column.try_unary::<_, TimestampNanosecondType, _>(&bias_timestamp)?) - } - } - }; + // 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(); - columns[self.time_index] = ts_column_biased; - // Offset selectors move samples into the evaluation timeline. Keep native histogram - // start timestamps on the same timeline for rate and reset calculations. + // 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 { @@ -575,10 +495,12 @@ impl Stream for SeriesNormalizeStream { mod test { use common_query::native_histogram::{build_histogram_array, read_histogram}; use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; - use datafusion::arrow::array::Float64Array; + use datafusion::arrow::array::{ + DictionaryArray, Float64Array, TimestampMicrosecondArray, TimestampNanosecondArray, + }; use datafusion::arrow::buffer::NullBuffer; use datafusion::arrow::datatypes::{ - ArrowPrimitiveType, DataType, Field, Schema, TimestampMillisecondType, + ArrowPrimitiveType, DataType, Field, Int64Type, Schema, TimeUnit, TimestampMillisecondType, }; use datafusion::common::ToDFSchema; use datafusion::datasource::memory::MemorySourceConfig; @@ -589,7 +511,9 @@ mod test { 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"; @@ -689,11 +613,11 @@ mod test { "+---------------------+--------+------+\ \n| timestamp | value | path |\ \n+---------------------+--------+------+\ - \n| 1970-01-01T00:01:01 | 0.0 | foo |\ - \n| 1970-01-01T00:02:01 | 1.0 | foo |\ - \n| 1970-01-01T00:00:01 | 10.0 | foo |\ - \n| 1970-01-01T00:00:31 | 100.0 | foo |\ - \n| 1970-01-01T00:01:31 | 1000.0 | foo |\ + \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+---------------------+--------+------+", ); @@ -839,7 +763,7 @@ mod test { let expected_timestamps = timestamp_array( [1_000, 3_000, 4_000, 5_000] .into_iter() - .map(|timestamp| (timestamp + offset) * ticks_per_ms) + .map(|timestamp| timestamp * ticks_per_ms) .collect(), ); assert_eq!( @@ -869,5 +793,103 @@ mod test { 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, + 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::>() + .unwrap() + .clone(), + ) + .unwrap(); + let values = values.get(0).unwrap(); + let values = values + .as_any() + .downcast_ref::() + .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::>() + .unwrap() + .clone(), + ) + .unwrap(); + assert_eq!( + timestamps.get(0).unwrap().to_data(), + TimestampMillisecondArray::from(vec![2_000; 3]).to_data() + ); } } diff --git a/src/promql/src/extension_plan/range_manipulate.rs b/src/promql/src/extension_plan/range_manipulate.rs index f4f6795270..d18b9b271a 100644 --- a/src/promql/src/extension_plan/range_manipulate.rs +++ b/src/promql/src/extension_plan/range_manipulate.rs @@ -46,8 +46,8 @@ use snafu::ResultExt; use crate::error::{DeserializeSnafu, Result}; use crate::extension_plan::{ - METRIC_NUM_SERIES, Millisecond, nanoseconds_per_native_tick, native_timestamp_values, - resolve_column_name, serialize_column_index, timestamp_unit, + METRIC_NUM_SERIES, Millisecond, local_offset, nanoseconds_per_native_tick, + native_timestamp_values, resolve_column_name, serialize_column_index, timestamp_unit, }; use crate::metrics::PROMQL_SERIES_COUNT; use crate::range_array::RangeArray; @@ -190,6 +190,7 @@ impl RangeManipulate { properties.boundedness, )); Arc::new(RangeManipulateExec { + offset: local_offset(&self.input, &self.time_index), start: self.start, end: self.end, interval: self.interval, @@ -423,6 +424,7 @@ impl UserDefinedLogicalNodeCore for RangeManipulate { #[derive(Debug)] pub struct RangeManipulateExec { + offset: Millisecond, start: Millisecond, end: Millisecond, interval: Millisecond, @@ -483,6 +485,7 @@ impl ExecutionPlan for RangeManipulateExec { properties.boundedness, )); Ok(Arc::new(Self { + offset: self.offset, start: self.start, end: self.end, interval: self.interval, @@ -532,6 +535,7 @@ impl ExecutionPlan for RangeManipulateExec { let aligned_ts_array = RangeManipulateStream::build_aligned_ts_array(self.start, self.end, self.interval); Ok(Box::pin(RangeManipulateStream { + offset: self.offset, start: self.start, end: self.end, interval: self.interval, @@ -594,6 +598,7 @@ impl DisplayAs for RangeManipulateExec { } pub struct RangeManipulateStream { + offset: Millisecond, start: Millisecond, end: Millisecond, interval: Millisecond, @@ -671,12 +676,29 @@ impl RangeManipulateStream { new_columns[*index] = new_column; } - // push timestamp range column - let timestamp_values = compute::cast( - input.column(self.time_index), - &DataType::Timestamp(TimeUnit::Millisecond, None), - )?; - let ts_range_column = RangeArray::from_ranges(timestamp_values, ranges.clone()) + // The timestamp range payload is always millisecond ABI. Shift in wide + // native precision before truncating toward zero, preserving null validity. + let scale = nanoseconds_per_native_tick(self.time_unit); + let timestamps = native_timestamp_values(input.column(self.time_index).as_ref())?; + let timestamp_values = timestamps + .iter() + .enumerate() + .map(|(index, timestamp)| { + if !input.column(self.time_index).is_valid(index) { + return Ok(None); + } + let shifted_ns = (*timestamp as i128) * scale + (self.offset as i128) * 1_000_000; + i64::try_from(shifted_ns / 1_000_000) + .map(Some) + .map_err(|_| { + ArrowError::ComputeError( + "RangeManipulate timestamp payload overflow".into(), + ) + }) + }) + .collect::, _>>()?; + let timestamp_values = TimestampMillisecondArray::from(timestamp_values); + let ts_range_column = RangeArray::from_ranges(Arc::new(timestamp_values), ranges.clone()) .map_err(|e| ArrowError::InvalidArgumentError(e.to_string()))? .into_dict(); new_columns.push(Arc::new(ts_range_column)); @@ -716,7 +738,8 @@ impl RangeManipulateStream { let ts_column = input.column(self.time_index); 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 timestamp = + |index| (timestamps[index] as i128) * scale + (self.offset as i128) * 1_000_000; let len = timestamps.len(); if len == 0 { return Ok((vec![], (self.start, self.end))); @@ -785,13 +808,16 @@ mod test { ArrayRef, DictionaryArray, Float64Array, StringArray, TimestampMicrosecondArray, TimestampNanosecondArray, }; + use datafusion::arrow::buffer::NullBuffer; use datafusion::arrow::datatypes::{ ArrowPrimitiveType, DataType, Field, Int64Type, Schema, TimestampMillisecondType, }; use datafusion::common::ToDFSchema; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; - use datafusion::logical_expr::{EmptyRelation, LogicalPlan}; + use datafusion::logical_expr::{ + EmptyRelation, Extension, LogicalPlan, UserDefinedLogicalNodeCore, + }; use datafusion::physical_expr::Partitioning; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::memory::MemoryStream; @@ -873,6 +899,7 @@ mod test { Boundedness::Bounded, )); let normalize_exec = Arc::new(RangeManipulateExec { + offset: 0, start, end, interval, @@ -1040,6 +1067,205 @@ mod test { } } + #[tokio::test] + async fn logical_normalize_offset_survives_rebuild_and_executes() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + TimestampMillisecondType::DATA_TYPE, + false, + ), + Field::new("value", DataType::Float64, true), + ])); + let input = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: schema.clone().to_dfschema_ref().unwrap(), + }); + let normalize = crate::extension_plan::SeriesNormalize::new( + 1_000, + TIME_INDEX_COLUMN, + false, + Vec::new(), + input.clone(), + ); + let normalize = crate::extension_plan::SeriesNormalize::deserialize(&normalize.serialize()) + .unwrap() + .with_exprs_and_inputs(vec![], vec![input]) + .unwrap(); + let normalized = LogicalPlan::Extension(Extension { + node: Arc::new(normalize), + }); + let plan = RangeManipulate::new( + 1_000, + 1_000, + 1, + 1_000, + TIME_INDEX_COLUMN.to_string(), + vec!["value".to_string()], + normalized.clone(), + ) + .unwrap(); + let rebuilt = RangeManipulate::deserialize(&plan.serialize()) + .unwrap() + .with_exprs_and_inputs(vec![], vec![normalized]) + .unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![0])), + Arc::new(Float64Array::from(vec![7.0])), + ], + ) + .unwrap(); + let exec_input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), + ))); + let output = datafusion::physical_plan::collect( + rebuilt.to_execution_plan(exec_input), + SessionContext::default().task_ctx(), + ) + .await + .unwrap(); + assert_eq!(output.len(), 1); + let output = &output[0]; + assert_eq!(output.num_rows(), 1); + assert_eq!( + output + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_ref(), + &[1_000] + ); + let values = RangeArray::try_new( + output + .column(1) + .as_any() + .downcast_ref::>() + .unwrap() + .clone(), + ) + .unwrap(); + assert_eq!(values.get_offset_length(0), Some((0, 1))); + assert_eq!( + values.get(0).unwrap().to_data(), + Float64Array::from(vec![7.0]).to_data() + ); + let timestamps = RangeArray::try_new( + output + .column(2) + .as_any() + .downcast_ref::>() + .unwrap() + .clone(), + ) + .unwrap(); + assert_eq!(timestamps.get_offset_length(0), Some((0, 1))); + assert_eq!( + timestamps.get(0).unwrap().to_data(), + TimestampMillisecondArray::from(vec![1_000]).to_data() + ); + } + + #[tokio::test] + async fn range_payload_preserves_null_timestamp_and_rejects_offset_overflow() { + let schema = Arc::new(Schema::new(vec![ + Field::new(TIME_INDEX_COLUMN, TimestampMillisecondType::DATA_TYPE, true), + Field::new("value", DataType::Float64, true), + ])); + let null_timestamp = + TimestampMillisecondArray::new(vec![1_000].into(), Some(NullBuffer::from(vec![false]))); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(null_timestamp), + Arc::new(Float64Array::from(vec![7.0])), + ], + ) + .unwrap(); + let input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], schema.clone(), None).unwrap(), + ))); + let plan = RangeManipulate::new( + 1_000, + 1_000, + 1, + 1, + TIME_INDEX_COLUMN.to_string(), + vec!["value".to_string()], + LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: schema.clone().to_dfschema_ref().unwrap(), + }), + ) + .unwrap(); + let output = datafusion::physical_plan::collect( + plan.to_execution_plan(input), + SessionContext::default().task_ctx(), + ) + .await + .unwrap(); + let timestamps = RangeArray::try_new( + output[0] + .column(2) + .as_any() + .downcast_ref::>() + .unwrap() + .clone(), + ) + .unwrap(); + let payload = timestamps.get(0).unwrap(); + let payload = payload + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(payload.len(), 1); + assert!(!payload.is_valid(0)); + + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![0, i64::MAX])), + Arc::new(Float64Array::from(vec![7.0, 8.0])), + ], + ) + .unwrap(); + let input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], schema.clone(), None).unwrap(), + ))); + let normalized = crate::extension_plan::SeriesNormalize::new( + 1, + TIME_INDEX_COLUMN, + false, + Vec::new(), + LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: schema.to_dfschema_ref().unwrap(), + }), + ); + let plan = RangeManipulate::new( + 1, + 1, + 1, + 1, + TIME_INDEX_COLUMN.to_string(), + vec!["value".to_string()], + LogicalPlan::Extension(Extension { + node: Arc::new(normalized), + }), + ) + .unwrap(); + let error = datafusion::physical_plan::collect( + plan.to_execution_plan(input), + SessionContext::default().task_ctx(), + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("timestamp payload overflow")); + } + #[tokio::test] async fn pruning_should_keep_time_and_value_columns_for_exec() { let schema = Arc::new(Schema::new(vec![ @@ -1194,6 +1420,7 @@ mod test { let empty_stream = MemoryStream::try_new(vec![], schema.clone(), None).unwrap(); let stream = RangeManipulateStream { + offset: 0, start: 1758093274000, // ends in 4000 end: 1758093334000, // ends in 4000 interval: 30000, // 30s step @@ -1259,6 +1486,7 @@ mod test { )])); let empty_stream = MemoryStream::try_new(vec![], schema.clone(), None).unwrap(); let stream = RangeManipulateStream { + offset: 0, start: query_start, end: query_end, interval, diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 9eedf7c067..34d0e20840 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -1943,6 +1943,11 @@ impl PromPlanner { if let Some(empty_plan) = self.setup_context().await? { return Ok(empty_plan); } + let offset_ms = match offset { + Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond, + Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond), + None => 0, + }; let normalize = self .selector_to_series_normalize_plan(offset, matchers, false) .await?; @@ -1974,26 +1979,42 @@ impl PromPlanner { DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) }) .collect::>(); - // `timestamp()` carries the sample timestamp as a value. Cast native - // input here because the helper consumes millisecond ticks. - let sample_time = col(&time_index_column); - let sample_time = if sample_time + // `timestamp()` preserves the shifted selector timeline even though + // SeriesNormalize now retains raw native timestamp storage. Decimal + // arithmetic shifts before truncating to milliseconds. + let unit_factor = match col(&time_index_column) .get_type(normalize.schema()) .context(DataFusionPlanningSnafu)? - == ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None) { - sample_time - } else { - DfExpr::Cast(Cast { - expr: Box::new(sample_time), - data_type: ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), - }) + ArrowDataType::Timestamp(ArrowTimeUnit::Second, _) => (1_000_i128, 4, 0), + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, _) => (1, 1, 0), + ArrowDataType::Timestamp(ArrowTimeUnit::Microsecond, _) => (1, 4, 3), + ArrowDataType::Timestamp(ArrowTimeUnit::Nanosecond, _) => (1, 7, 6), + _ => unreachable!("time index is a timestamp"), }; - let sample_time = sample_time + let sample_time = col(&time_index_column) .cast_to(&ArrowDataType::Int64, normalize.schema()) .context(DataFusionPlanningSnafu)? - .cast_to(&ArrowDataType::Float64, normalize.schema()) + .cast_to(&ArrowDataType::Decimal128(19, 0), normalize.schema()) .context(DataFusionPlanningSnafu)?; + let sample_time = DfExpr::BinaryExpr(BinaryExpr { + left: Box::new(sample_time), + op: Operator::Multiply, + right: Box::new(lit(ScalarValue::Decimal128( + Some(unit_factor.0), + unit_factor.1, + unit_factor.2, + ))), + }); + let sample_time = DfExpr::BinaryExpr(BinaryExpr { + left: Box::new(sample_time), + op: Operator::Plus, + right: Box::new(lit(ScalarValue::Decimal128(Some(offset_ms as i128), 19, 0))), + }) + .cast_to(&ArrowDataType::Int64, normalize.schema()) + .context(DataFusionPlanningSnafu)? + .cast_to(&ArrowDataType::Float64, normalize.schema()) + .context(DataFusionPlanningSnafu)?; let sample_time = DfExpr::BinaryExpr(BinaryExpr { left: Box::new(sample_time), op: Operator::Divide, @@ -8877,7 +8898,7 @@ mod test { \n Projection: some_metric.timestamp, value AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\ \n Projection: some_metric.timestamp, __promql_timestamp_value_ AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\ \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\ - \n Projection: some_metric.tag_0, some_metric.timestamp, some_metric.field_0, CAST(CAST(some_metric.timestamp AS Int64) AS Float64) / Float64(1000) AS __promql_timestamp_value_ [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\ + \n Projection: some_metric.tag_0, some_metric.timestamp, some_metric.field_0, CAST(CAST(CAST(CAST(some_metric.timestamp AS Int64) AS Decimal128(19, 0)) * Decimal128(Some(1),1,0) + Decimal128(Some(0),19,0) AS Int64) AS Float64) / Float64(1000) AS __promql_timestamp_value_ [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\ \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ \n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\