rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-09-13 08:52:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: apply PromQL offsets without native timestamp overflow

Browse Source 
Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
This commit is contained in:
discord9
2026-09-09 16:23:48 +08:00
parent ee3fe9d35e
commit 8e89692a70
5 changed files with 687 additions and 135 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/promql/src/extension_plan.rs
+137 -1
View File
@@ -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::<SeriesNormalize>()
.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(&divide, "timestamp"));
}
}
src/promql/src/extension_plan/instant_manipulate.rs
+150 -5
View File
@@ -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::<Float64Array>()
.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::<Float64Array>()
.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,
src/promql/src/extension_plan/normalize.rs
+127 -105
View File
@@ -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<Millisecond> {
let index = self.input.schema().maybe_index_of_column(time_index)?;
let (qualifier, field) = self.input.schema().qualified_field(index);
(field.name() == &self.time_index_column_name
&& time_index == &Column::new(qualifier.cloned(), field.name().clone()))
.then_some(self.offset)
}
pub fn new<N: AsRef<str>>(
offset: Millisecond,
time_index_column_name: N,
@@ -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<RecordBatch> {
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<dyn Array> = match time_unit {
TimeUnit::Second => {
let column = input
.column(self.time_index)
.as_any()
.downcast_ref::<TimestampSecondArray>()
.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::<TimestampMillisecondArray>()
.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::<TimestampMicrosecondArray>()
.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::<TimestampNanosecondArray>()
.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::<DictionaryArray<Int64Type>>()
.unwrap()
.clone(),
)
.unwrap();
let values = values.get(0).unwrap();
let values = values
.as_any()
.downcast_ref::<datafusion::arrow::array::StructArray>()
.unwrap();
assert_eq!(
read_histogram(values, 0).unwrap().unwrap().start_timestamp,
Some(1_500)
);
assert_eq!(
read_histogram(values, 1).unwrap().unwrap().start_timestamp,
Some(0)
);
assert_eq!(
read_histogram(values, 2).unwrap().unwrap().start_timestamp,
None
);
let timestamps = RangeArray::try_new(
output[0]
.column(2)
.as_any()
.downcast_ref::<DictionaryArray<Int64Type>>()
.unwrap()
.clone(),
)
.unwrap();
assert_eq!(
timestamps.get(0).unwrap().to_data(),
TimestampMillisecondArray::from(vec![2_000; 3]).to_data()
);
}
}
src/promql/src/extension_plan/range_manipulate.rs
+238 -10
View File
@@ -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::<std::result::Result<Vec<_>, _>>()?;
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::<TimestampMillisecondArray>()
.unwrap()
.values()
.as_ref(),
&[1_000]
);
let values = RangeArray::try_new(
output
.column(1)
.as_any()
.downcast_ref::<DictionaryArray<Int64Type>>()
.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::<DictionaryArray<Int64Type>>()
.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::<DictionaryArray<Int64Type>>()
.unwrap()
.clone(),
)
.unwrap();
let payload = timestamps.get(0).unwrap();
let payload = payload
.as_any()
.downcast_ref::<TimestampMillisecondArray>()
.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,
src/query/src/promql/planner.rs
+35 -14
View File
@@ -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::<Vec<_>>();
// `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]\