From b3c756a7dd05cd7e1b6685489714c4a5795a33df Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 9 Oct 2026 10:45:15 +0000 Subject: [PATCH] perf(promql): reduce timestamp payload allocations and use narrow cursor scans (#9417) * perf(promql): zero-copy timestamp payload for millisecond zero-offset ranges RangeManipulate rebuilt the timestamp range payload per sample through wide-precision i128 conversion even for millisecond input with zero offset, where the input is already millisecond and nothing changes. That turned a previously zero-copy path into a per-sample loop plus a Vec allocation, regressing range-function queries. Reuse the input timestamp array directly when the time index is already millisecond and the offset is zero. Other units and non-zero offsets keep the wide-precision conversion (required so ns/us data near the native type limits stays correct under large offsets). Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> Signed-off-by: discord9 * perf(promql): vectorize timestamp payload with plain i64 RangeManipulate rebuilt the timestamp range payload per sample through wide-precision i128 conversion, turning a hot path into a scalar loop with a Vec allocation. Realistic timestamps and offsets stay far inside i64, so the wide path is unnecessary. - millisecond + zero offset reuses the input array with zero copy; - other units / offsets shift in one vectorized i64 pass (no i128, no per-sample Option, no overflow branch). Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> Signed-off-by: discord9 * perf(promql): check timestamp payload overflow via arrow try_unary Replace the silent wrapping shift with try_unary + checked_mul/checked_add. checked arithmetic lowers to a per-lane overflow flag so the single pass still vectorizes over the dense i64 buffer; any overflow short-circuits to an error instead of silently wrapping to a wrong timestamp. No i128, no per-sample Option. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> Signed-off-by: discord9 * perf(promql): vectorize cursor scans with i64 fast path RangeManipulate::calculate_range and InstantManipulate::manipulate compared every sample through i128 wide-precision arithmetic. For realistic timestamps and offsets the shifted values stay far inside i64, so the wide path was pure overhead. Split both scans into an i64 fast path (checked arithmetic precomputed once per batch) and an i128 exact fallback. Any input that would overflow i64 falls back to the i128 scan, keeping window membership semantics identical. Verified: 47/47 range_manipulate + instant_manipulate tests, 275/275 promql crate tests, 1612 randomized cases byte-identical to i128 scan. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> Signed-off-by: discord9 * refactor(promql): generalize cursor scans and share endpoint precheck Address review comments on PR #9417: - Generalize scan_ranges/scan_instants over Nanoseconds (i64/i128), eliminating the duplicated scan bodies. - Extract shared narrow_scan precheck used by both RangeManipulate and InstantManipulate, with identical checked endpoint logic. - Trim doc comments to one line; fix factual errors about fallback trigger (2262 for ms/s, not native type limits). Semantics unchanged: same cursor algorithm, same overflow fallback, same i64-vs-i128 selection, same stale/null/tie handling. Verified: 47/47 range_manipulate + instant_manipulate tests, clippy clean. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> Signed-off-by: discord9 * perf(promql): vectorize payload with try_unary and add scan carrier bench Replace the per-sample i128 payload loop with a single arrow try_unary pass: i128 intermediate holds any native timestamp shift, only the final millisecond result must fit i64. This is semantically equivalent to the original i128 implementation (no silent wrap, no per-sample branching). Add cursor_scan_carriers A/B benchmark: i64 fast path vs i128 exact fallback on four workloads. Add five deterministic differential tests pinning scan_ranges:: and scan_ranges:: to be pointwise equivalent on boundary shapes (sub-ms samples, fallback triggers, nanosecond extremes, seeded precision matrix). Bench (N=4096, 15s cadence): - payload wall_time 20x15s: 74.94us -> 48.98us (-35%) - payload wall_time 240x15s: 102.44us - scan i64 fast path: 10.97us vs i128 19.84us (1.81x) Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * perf(promql): restore zero-copy payload for ms+offset=0 and fix comment Restore the zero-copy fast path for millisecond timestamps with zero offset, guarded by the input column's exact data type (Timestamp (Millisecond, None)) rather than self.time_unit. The old time_unit guard would accept a tz-carrying millisecond column and produce a dictionary value type disagreeing with the output schema; the datatype guard is the correct check. Also fix the comment: arrow's try_unary is a fallible kernel that LLVM cannot auto-vectorize; the win over the hand-written loop comes from eliminating the per-sample Option/Vec and branches, not SIMD. Add a regression test with a timezone-carrying millisecond index to prove it falls back to the try_unary path. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * perf(promql): replace try_unary with unary + Cell overflow flag Arrow's try_unary is a fallible kernel that LLVM cannot auto-vectorize. Replace it with the infallible unary kernel: the closure is a plain i64 -> i64 mapping that sets a Cell on overflow instead of returning Err per sample. The flag is checked once after the loop. This removes the per-lane Result construction and lets LLVM vectorize the shift+truncate sequence. Bench (N=4096, 15s cadence): - window4x15s: 42.12us -> 28.77us (-32%) - window20x15s: 48.98us -> 27.68us (-43%) - window240x15s: 102.44us -> 27.61us (-73%) - regular_sparse: 90.04us -> 29.20us (-68%) All cases now run at ~27-29us regardless of window size; the zero-copy path dominates for ms+offset=0. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * perf(promql): dispatch payload conversion on null_count; drop wall-time bench Null-free batches (dominant PromQL case) take arrow's vectorizable unary kernel with a Cell overflow flag; nullable batches take try_unary so a NULL slot's arbitrary underlying value cannot raise a false batch-wide overflow error. Add a regression test exercising both paths. Remove bench_range_manipulate_wall_time (all-ms zero-offset cases only exercised the zero-copy path, measuring nothing about the conversion) and its orphaned helpers; restore the cursor-scan bench helpers it shared. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): keep nullable timestamp cursor scans on i128 NULL slots may contain unordered arbitrary values, so the endpoint-only narrow-scan proof cannot bound their interior timestamp arithmetic. Both range and instant scans now attempt i64 only for null-free timestamp arrays and otherwise retain the existing exact i128 path. Add focused interior-NULL regressions through both active entrypoints, checking concrete selections and output against the existing wide scan. The null-free fast path and timestamp payload conversion are unchanged. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> Signed-off-by: discord9 --- src/promql/benches/bench_range_fn.rs | 468 ++++--- .../src/extension_plan/instant_manipulate.rs | 156 ++- .../src/extension_plan/range_manipulate.rs | 1203 ++++++++++++++++- .../promql_long_range_rate/case.toml | 119 ++ 4 files changed, 1686 insertions(+), 260 deletions(-) create mode 100644 tests/perf/query_cases/promql_long_range_rate/case.toml diff --git a/src/promql/benches/bench_range_fn.rs b/src/promql/benches/bench_range_fn.rs index b4e026bae3d..51e0aefdc79 100644 --- a/src/promql/benches/bench_range_fn.rs +++ b/src/promql/benches/bench_range_fn.rs @@ -17,25 +17,14 @@ use std::sync::Arc; use criterion::{BenchmarkId, Criterion, criterion_group}; -use datafusion::arrow::array::{ - Array, ArrayRef, DictionaryArray, Float64Array, TimestampMillisecondArray, -}; -use datafusion::arrow::datatypes::{ - DataType as ArrowDataType, Field as ArrowField, Int64Type, Schema, -}; -use datafusion::arrow::record_batch::RecordBatch; -use datafusion::common::ToDFSchema; -use datafusion::datasource::memory::MemorySourceConfig; -use datafusion::datasource::source::DataSourceExec; -use datafusion::execution::context::TaskContext; -use datafusion::logical_expr::{EmptyRelation, LogicalPlan}; -use datafusion::physical_plan::{ColumnarValue, ExecutionPlan}; +use datafusion::arrow::array::{Array, Float64Array, TimestampMillisecondArray}; +use datafusion::arrow::datatypes::TimeUnit; +use datafusion::physical_plan::ColumnarValue; use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; use datafusion_expr::ScalarFunctionArgs; use datatypes::arrow::datatypes::{DataType, Field}; -use futures::StreamExt; -use promql::extension_plan::RangeManipulate; +use promql::extension_plan::RangeManipulateStream; use promql::functions::{ AbsentOverTime, Changes, CountOverTime, Delta, DoubleExponentialSmoothing, IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, PredictLinear, PresentOverTime, QuantileOverTime, Rate, @@ -964,223 +953,294 @@ fn bench_edge_count_functions(c: &mut Criterion) { group.finish(); } -const RANGE_MANIPULATE_CADENCE_MS: i64 = 15_000; +/// One cursor-scan workload: the timeline, the window and the aligned instants to visit. +#[derive(Clone, Debug)] +struct CursorScanWorkload { + time_unit: TimeUnit, + offset_ms: i64, + timestamps: Vec, + window_ms: i64, + interval_ms: i64, + start_ms: i64, + end_ms: i64, +} -fn make_range_manipulate_batch(timestamps: Vec, field_count: usize) -> RecordBatch { - let mut fields = Vec::with_capacity(field_count + 1); - let mut columns: Vec = Vec::with_capacity(field_count + 1); - fields.push(ArrowField::new( - "timestamp", - ArrowDataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None), - false, - )); - columns.push(Arc::new(TimestampMillisecondArray::from(timestamps)) as _); - - for field_index in 0..field_count { - fields.push(ArrowField::new( - format!("value_{field_index}"), - ArrowDataType::Float64, - false, - )); - let values = (0..columns[0].len()) - .map(|row| row as f64 + field_index as f64 * 0.01) - .collect::>(); - columns.push(Arc::new(Float64Array::from(values)) as _); +impl CursorScanWorkload { + /// The workload through the narrow `i64` carrier, or `None` when production would fall + /// back to `i128` for it. + /// + /// For the fallback cases the fast path is timed on a twin workload: the same shape + /// translated onto a timeline the narrow carrier holds, which keeps the number of sample + /// and window comparisons identical to the fallback run. + fn bench_scan_narrow(&self) -> Option> { + RangeManipulateStream::bench_scan_narrow( + self.time_unit, + self.offset_ms, + self.window_ms, + self.interval_ms, + &self.timestamps, + self.start_ms, + self.end_ms, + ) } - RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap() -} + /// The workload through the exact `i128` carrier, the fallback production takes + /// whenever [`Self::bench_scan_narrow`] returns `None`. + fn bench_scan_wide(&self) -> Vec<(u32, u32)> { + RangeManipulateStream::bench_scan_wide( + self.time_unit, + self.offset_ms, + self.window_ms, + self.interval_ms, + &self.timestamps, + self.start_ms, + self.end_ms, + ) + } -async fn execute_and_drain( - plan: Arc, - task_ctx: Arc, -) -> Vec { - let mut stream = plan.execute(0, task_ctx).unwrap(); - let mut output = Vec::new(); - while let Some(batch) = stream.next().await { - output.push(batch.unwrap()); + fn evaluations(&self) -> usize { + ((self.end_ms - self.start_ms) / self.interval_ms) as usize + 1 + } + + fn label(&self) -> String { + format!( + "N{}_eval{}_window{}x{}ms", + self.timestamps.len(), + self.evaluations(), + self.window_ms / self.interval_ms, + self.interval_ms + ) } - output } -fn range_keys(batch: &RecordBatch, index: usize) -> Vec<(u32, u32)> { - let ranges = batch - .column(index) - .as_any() - .downcast_ref::>() - .unwrap() - .clone(); - RangeArray::try_new(ranges) - .unwrap() - .ranges() - .map(Option::unwrap) - .collect() +/// Cadence of the cursor-scan benchmarks: PromQL's usual 15s step. +const CURSOR_SCAN_CADENCE_MS: i64 = 15_000; + +/// Milliseconds of a one-hour window, the widest window benchmarked here. +const CURSOR_SCAN_HOUR_MS: i64 = 3_600_000; + +/// Builds a workload from a first timestamp and a sample count. +/// +/// `sample_spacing_native` is the distance between samples in native ticks: 15s worth, so +/// every native-precision workload has the same shape (one sample per cadence, 20 samples +/// per 300s window). The aligned instants start at the first shifted sample, so every +/// window overlaps the input the same way in both arms of the A/B comparison. +fn cursor_scan_workload( + time_unit: TimeUnit, + offset_ms: i64, + first_timestamp: i64, + samples: usize, + window_ms: i64, + sample_spacing_native: i64, +) -> CursorScanWorkload { + let interval_ms = CURSOR_SCAN_CADENCE_MS; + let timestamps = (0..samples as i64) + .map(|index| first_timestamp + index * sample_spacing_native) + .collect::>(); + // `Millisecond` instants of the first shifted sample, flooring like `calculate_range`. + let start_ms = shifted_floor_millis(time_unit, offset_ms, first_timestamp); + let evaluations = samples - 1; + CursorScanWorkload { + time_unit, + offset_ms, + timestamps, + window_ms, + interval_ms, + start_ms, + end_ms: start_ms + evaluations as i64 * interval_ms, + } } -fn brute_force_range_keys( - timestamps: &[i64], - evaluations: usize, - window_points: usize, -) -> Vec<(u32, u32)> { - let range = window_points as i64 * RANGE_MANIPULATE_CADENCE_MS; - (0..evaluations) - .map(|evaluation| { - let current = evaluation as i64 * RANGE_MANIPULATE_CADENCE_MS; - let mut offset = None; - let mut length = 0; - for (index, timestamp) in timestamps.iter().enumerate() { - if *timestamp > current - range && *timestamp <= current { - offset.get_or_insert(index as u32); - length += 1; - } - } - (offset.unwrap_or(0), length) - }) - .collect() +/// Native ticks of one benchmark cadence, i.e. 15s in the given precision. +fn cadence_in_native_ticks(time_unit: TimeUnit) -> i64 { + match time_unit { + TimeUnit::Second => CURSOR_SCAN_CADENCE_MS / 1_000, + TimeUnit::Millisecond => CURSOR_SCAN_CADENCE_MS, + TimeUnit::Microsecond => CURSOR_SCAN_CADENCE_MS * 1_000, + TimeUnit::Nanosecond => CURSOR_SCAN_CADENCE_MS * 1_000_000, + } } -fn validate_range_manipulate_output( - output: &[RecordBatch], - field_count: usize, - expected_range_keys: Option<&[(u32, u32)]>, -) { - let Some(expected_range_keys) = expected_range_keys else { - assert!(output.is_empty()); - return; +/// Floor of a shifted sample in milliseconds, in exact `i128` native precision. +fn shifted_floor_millis(time_unit: TimeUnit, offset_ms: i64, timestamp: i64) -> i64 { + let scale = match time_unit { + TimeUnit::Second => 1_000_000_000i128, + TimeUnit::Millisecond => 1_000_000i128, + TimeUnit::Microsecond => 1_000i128, + TimeUnit::Nanosecond => 1i128, + }; + i64::try_from( + ((timestamp as i128) * scale + (offset_ms as i128) * 1_000_000).div_euclid(1_000_000), + ) + .expect("the shifted millisecond must stay legal") +} + +/// A/B benchmark of the cursor-scan carriers behind `RangeManipulate::scan_ranges`. +/// +/// `RangeManipulate` scans the aligned instants in `i64` nanoseconds whenever +/// `narrow_scan` proves the batch and the bounds fit, and in exact `i128` nanoseconds +/// otherwise. Both arms below drive the production scan through the carrier the operator +/// would pick, so the numbers isolate the carrier and not the surrounding operator work: +/// +/// - `ms_zero_offset`: millisecond batch with no offset, the `i64` fast path. The same +/// workload is also timed through `i128` to price the fallback on this input. +/// - `ms_zero_offset_long_window`: same, with a one-hour window, where the scan advances +/// both cursors much further per evaluation. +/// - `ms_large_offset`: an offset whose nanosecond product overflows `i64`, so production +/// falls back to `i128`. The `i64` arm times the identical workload shape translated onto +/// a narrow timeline, which keeps the selections and thus the work identical. +/// - `ns_native_floor`: nanosecond batch at the `i64::MIN` floor, where the shifted start +/// bound underflows `i64` nanoseconds, so production falls back to `i128`. The `i64` arm +/// again times the translated narrow shape. +/// +/// Every arm asserts that both carriers selected the same ranges, so the comparison cannot +/// be won by doing less work. The fallback triggers themselves are pinned by the unit tests +/// in `range_manipulate.rs` (`scan_falls_back_to_i128_when_only_the_sample_tail_overflows` +/// and `scan_carriers_split_at_the_nanosecond_boundaries`). +fn bench_cursor_scan_carriers(c: &mut Criterion) { + const SAMPLES: usize = 4_096; + const SHORT_WINDOW_MS: i64 = 20 * CURSOR_SCAN_CADENCE_MS; + // `i64::MAX` milliseconds is ~9.22e18 ms; one microsecond past it in the nanosecond + // product makes `narrow_scan` reject the offset, while the shifted millisecond stays legal. + const LARGE_OFFSET_MS: i64 = 9_223_372_036_855; + const LARGE_OFFSET_BASE_MS: i64 = 1_000_000_000_000_000; + const NANOSECOND_FLOOR_SHIFT_NS: i64 = 9_223_372_036_000_000_000; + const NANOSECOND_FLOOR_SAMPLES: usize = 4_096; + + let ms_zero_offset = cursor_scan_workload( + TimeUnit::Millisecond, + 0, + 0, + SAMPLES, + SHORT_WINDOW_MS, + cadence_in_native_ticks(TimeUnit::Millisecond), + ); + let ms_zero_offset_long_window = cursor_scan_workload( + TimeUnit::Millisecond, + 0, + 0, + SAMPLES, + CURSOR_SCAN_HOUR_MS, + cadence_in_native_ticks(TimeUnit::Millisecond), + ); + // The large offset is shifted inside one millisecond of the nanosecond ceiling, so the + // payload stays legal: the fallback is driven by the product, not by an illegal instant. + let ms_large_offset = cursor_scan_workload( + TimeUnit::Millisecond, + LARGE_OFFSET_MS, + LARGE_OFFSET_BASE_MS, + SAMPLES, + SHORT_WINDOW_MS, + cadence_in_native_ticks(TimeUnit::Millisecond), + ); + // The narrow twin of `ms_large_offset`: the same batch shape without the offset. + let ms_large_offset_narrow = cursor_scan_workload( + TimeUnit::Millisecond, + 0, + 0, + SAMPLES, + SHORT_WINDOW_MS, + cadence_in_native_ticks(TimeUnit::Millisecond), + ); + + // Nanosecond batch on the native floor: samples at `i64::MIN + i * 15s`, whose shifted + // start bound underflows `i64` nanoseconds. + let nanosecond_floor_first = i64::MIN; + let nanosecond_floor = cursor_scan_workload( + TimeUnit::Nanosecond, + 0, + nanosecond_floor_first, + NANOSECOND_FLOOR_SAMPLES, + SHORT_WINDOW_MS, + cadence_in_native_ticks(TimeUnit::Nanosecond), + ); + // The narrow twin: the same shape translated by a multiple of one millisecond, so both + // the samples and the aligned instants keep their relative geometry. + let nanosecond_floor_narrow = CursorScanWorkload { + time_unit: TimeUnit::Nanosecond, + offset_ms: 0, + timestamps: nanosecond_floor + .timestamps + .iter() + .map(|timestamp| timestamp + NANOSECOND_FLOOR_SHIFT_NS) + .collect(), + window_ms: nanosecond_floor.window_ms, + interval_ms: nanosecond_floor.interval_ms, + start_ms: nanosecond_floor.start_ms + NANOSECOND_FLOOR_SHIFT_NS / 1_000_000, + end_ms: nanosecond_floor.end_ms + NANOSECOND_FLOOR_SHIFT_NS / 1_000_000, }; - assert_eq!(output.len(), 1); - let batch = &output[0]; - assert_eq!(batch.num_rows(), expected_range_keys.len()); - - let timestamp_range_keys = range_keys(batch, field_count + 1); - assert_eq!(timestamp_range_keys, expected_range_keys); - - for field_index in 0..field_count { - assert_eq!(range_keys(batch, field_index + 1), expected_range_keys); - } -} - -fn bench_range_manipulate_wall_time(c: &mut Criterion) { - let mut group = c.benchmark_group("range_manipulate_wall_time"); - let primary_points = 4_096; - let sparse_evaluations = primary_points - 1; - let mut cases = vec![ + let cases = [ ( - "one_field", - (0..primary_points) - .map(|point| point as i64 * RANGE_MANIPULATE_CADENCE_MS) - .collect::>(), - primary_points, - 4, - 1, + "ms_zero_offset", + ms_zero_offset.clone(), + ms_zero_offset, false, ), ( - "one_field", - (0..primary_points) - .map(|point| point as i64 * RANGE_MANIPULATE_CADENCE_MS) - .collect::>(), - primary_points, - 20, - 1, + "ms_zero_offset_long_window", + ms_zero_offset_long_window.clone(), + ms_zero_offset_long_window, false, ), ( - "one_field", - (0..primary_points) - .map(|point| point as i64 * RANGE_MANIPULATE_CADENCE_MS) - .collect::>(), - primary_points, - 240, - 1, - false, - ), - ( - "regular_sparse", - (0..sparse_evaluations) - .step_by(2) - .map(|point| point as i64 * RANGE_MANIPULATE_CADENCE_MS) - .collect::>(), - sparse_evaluations, - 20, - 1, - false, - ), - ( - "all_empty_pathological", - (0..primary_points) - .map(|point| { - (primary_points as i64 + 21 + point as i64) * RANGE_MANIPULATE_CADENCE_MS - }) - .collect::>(), - primary_points, - 20, - 1, + "ms_large_offset", + ms_large_offset_narrow, + ms_large_offset, true, ), ( - "multi_field_control", - (0..primary_points) - .map(|point| point as i64 * RANGE_MANIPULATE_CADENCE_MS) - .collect::>(), - primary_points, - 20, - 4, - false, + "ns_native_floor", + nanosecond_floor_narrow, + nanosecond_floor, + true, ), ]; - for (case_name, timestamps, evaluations, window_points, field_count, expect_empty) in - cases.drain(..) - { - let expected_range_keys = (!expect_empty) - .then(|| brute_force_range_keys(×tamps, evaluations, window_points)); - let input_batch = make_range_manipulate_batch(timestamps, field_count); - let input_rows = input_batch.num_rows(); - let input_schema = input_batch.schema(); - let logical_input = LogicalPlan::EmptyRelation(EmptyRelation { - produce_one_row: false, - schema: input_schema.clone().to_dfschema_ref().unwrap(), - }); - let field_columns = (0..field_count) - .map(|field_index| format!("value_{field_index}")) - .collect::>(); - let logical_plan = RangeManipulate::new( - 0, - (evaluations as i64 - 1) * RANGE_MANIPULATE_CADENCE_MS, - RANGE_MANIPULATE_CADENCE_MS, - 0, - window_points as i64 * RANGE_MANIPULATE_CADENCE_MS, - "timestamp".to_string(), - field_columns, - logical_input, - ) - .unwrap(); - let physical_input: Arc = Arc::new(DataSourceExec::new(Arc::new( - MemorySourceConfig::try_new(&[vec![input_batch]], input_schema, None).unwrap(), - ))); - let execution_plan = logical_plan.to_execution_plan(physical_input); - let runtime = tokio::runtime::Runtime::new().unwrap(); - let task_ctx = datafusion::prelude::SessionContext::new().task_ctx(); + let mut group = c.benchmark_group("cursor_scan_carriers"); + group.warm_up_time(std::time::Duration::from_millis(500)); + group.measurement_time(std::time::Duration::from_secs(2)); + group.sample_size(50); - // Validate the public operator output and RangeArray keys before timing. - let output = runtime.block_on(execute_and_drain(execution_plan.clone(), task_ctx.clone())); - validate_range_manipulate_output(&output, field_count, expected_range_keys.as_deref()); + for (case_name, narrow, wide, wide_falls_back) in cases { + // Both arms must be timed on work the operator would really do: the narrow arm in + // `i64`, and the wide arm in `i128` exactly when the narrow baselines overflow. + let narrow_expected = narrow + .bench_scan_narrow() + .unwrap_or_else(|| panic!("{case_name}: the i64 arm must fit the narrow carrier")); + assert_eq!( + wide.bench_scan_narrow().is_none(), + wide_falls_back, + "{case_name}: the i128 arm's fallback trigger" + ); + // Fairness check before timing: both carriers have to select the same ranges, which + // makes the sample and window comparisons of the two arms identical. + assert_eq!( + wide.bench_scan_wide(), + narrow_expected, + "{case_name}: the two carriers must select the same ranges" + ); + assert_eq!( + narrow_expected.len(), + narrow.evaluations(), + "{case_name}: one range per aligned evaluation instant" + ); - // This measures public RangeManipulate execution wall time, including stream draining. + group.throughput(criterion::Throughput::Elements( + narrow.timestamps.len() as u64 + )); group.bench_with_input( - BenchmarkId::new( - case_name, - format!("N{input_rows}_eval{evaluations}_window{window_points}x15s"), - ), - &(), - |b, _| { - b.iter(|| { - let output = runtime - .block_on(execute_and_drain(execution_plan.clone(), task_ctx.clone())); - std::hint::black_box(output); - }) + BenchmarkId::new(format!("{case_name}/i64"), narrow.label()), + &narrow, + |b, workload| { + b.iter(|| std::hint::black_box(workload.bench_scan_narrow())); + }, + ); + group.bench_with_input( + BenchmarkId::new(format!("{case_name}/i128"), wide.label()), + &wide, + |b, workload| { + b.iter(|| std::hint::black_box(workload.bench_scan_wide())); }, ); } @@ -1196,5 +1256,5 @@ criterion_group!( bench_rate_window_steps, bench_edge_count_functions, bench_extrema_functions, - bench_range_manipulate_wall_time + bench_cursor_scan_carriers ); diff --git a/src/promql/src/extension_plan/instant_manipulate.rs b/src/promql/src/extension_plan/instant_manipulate.rs index 4085d382c37..4a9a3fe5fb8 100644 --- a/src/promql/src/extension_plan/instant_manipulate.rs +++ b/src/promql/src/extension_plan/instant_manipulate.rs @@ -46,6 +46,7 @@ use prost::Message; use snafu::ResultExt; use crate::error::{DeserializeSnafu, Result}; +use crate::extension_plan::range_manipulate::{Nanoseconds, ScanNanoseconds, narrow_scan}; use crate::extension_plan::series_divide::SeriesDivide; use crate::extension_plan::{ METRIC_NUM_SERIES, Millisecond, is_prometheus_stale_sample, local_offset, @@ -691,6 +692,7 @@ impl InstantManipulateStream { let (timestamps, _) = timestamp_array_to_primitive(ts_column).ok_or_else(|| { DataFusionError::Execution("Time index column is not a timestamp".into()) })?; + let timestamp_null_count = timestamps.null_count(); let timestamps = timestamps.values(); let len = timestamps.len(); // Shift the native-tick timeline in i128 before comparing samples. Doing @@ -735,28 +737,77 @@ impl InstantManipulateStream { let estimated_points = estimated_points as usize; let aligned_start = aligned_start as i64; let aligned_end = aligned_end as i64; - let mut take_indices = Vec::with_capacity(estimated_points); - let mut aligned_ts = Vec::with_capacity(estimated_points); - let mut cursor = 0; - for expected_ms in (aligned_start..=aligned_end).step_by(self.interval as usize) { - let expected = (expected_ms as i128) * 1_000_000; + // Narrow `i64` nanoseconds whenever the batch fits them, exact `i128` otherwise. + let narrow = (timestamp_null_count == 0) + .then(|| { + narrow_scan( + self.time_unit, + self.offset, + self.lookback_delta, + timestamps, + aligned_start, + aligned_end, + ) + }) + .flatten(); + if let Some(scan) = &narrow { + debug_assert_eq!(i64::try_from(first_ns), Ok(scan.shift(timestamps[0]))); + debug_assert_eq!(i64::try_from(last_ns), Ok(scan.shift(timestamps[len - 1]))); + } + let (take_indices, aligned_ts) = match narrow { + Some(scan) => self.scan_instants::( + timestamps, + aligned_start, + aligned_end, + &scan, + estimated_points, + &is_stale, + ), + None => self.scan_instants::( + timestamps, + aligned_start, + aligned_end, + &ScanNanoseconds::::wide(self.time_unit, self.offset, self.lookback_delta), + estimated_points, + &is_stale, + ), + }; + self.take_record_batch_optional(input, take_indices, aligned_ts) + } + + /// Chooses the sample of every aligned evaluation instant, in `Ns` nanoseconds. + /// + /// Keeps the first row among exact timestamp ties and the latest preceding row + /// otherwise; zero lookback admits only exact samples. Staleness is tested after + /// choosing: a selected stale marker suppresses this evaluation rather than + /// falling back to an older finite value. + fn scan_instants( + &self, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, + scan: &ScanNanoseconds, + capacity: usize, + is_stale: impl Fn(usize) -> bool, + ) -> (Vec, Vec) { + let len = timestamps.len(); + let mut cursor = 0usize; + let mut take_indices = Vec::with_capacity(capacity); + let mut aligned_ts = Vec::with_capacity(capacity); + for expected_ms in (start..=end).step_by(self.interval as usize) { + let expected = scan.instant(expected_ms); let mut exact_candidate = None; - while cursor < len && to_nanoseconds(timestamps[cursor]) <= expected { - if to_nanoseconds(timestamps[cursor]) == expected && exact_candidate.is_none() { + while cursor < len && scan.shift(timestamps[cursor]) <= expected { + if scan.shift(timestamps[cursor]) == expected && exact_candidate.is_none() { exact_candidate = Some(cursor); } cursor += 1; } - // Keep the first row among exact timestamp ties; otherwise use the - // latest preceding row. Zero lookback admits only exact samples. - // Test staleness after choosing: a selected stale marker suppresses - // this evaluation rather than falling back to an older finite value. let Some(candidate) = exact_candidate.or_else(|| cursor.checked_sub(1)) else { continue; }; - let candidate_ts = to_nanoseconds(timestamps[candidate]); - let lower = expected - (self.lookback_delta as i128) * 1_000_000; - if (candidate_ts == expected || candidate_ts > lower) + let candidate_ts = scan.shift(timestamps[candidate]); + if (candidate_ts == expected || candidate_ts > scan.window_start(expected)) && candidate_ts <= expected && !is_stale(candidate) { @@ -764,7 +815,7 @@ impl InstantManipulateStream { aligned_ts.push(expected_ms); } } - self.take_record_batch_optional(input, take_indices, aligned_ts) + (take_indices, aligned_ts) } /// Helper function to apply "take" on record batch. @@ -1775,6 +1826,81 @@ mod test { .await; } + #[test] + fn nullable_interior_timestamp_matches_wide_scan() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Second, None), + true, + ), + Field::new("value", DataType::Float64, false), + ])); + let timestamps = [0, i64::MAX, 1]; + let array = TimestampSecondArray::new( + timestamps.to_vec().into(), + Some(NullBuffer::from(vec![false, false, true])), + ); + let input = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(array), + Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])), + ], + ) + .unwrap(); + let stream = InstantManipulateStream { + offset: 0, + start: 0, + end: 1_000, + lookback_delta: 1_000, + interval: 1_000, + time_index: 0, + time_unit: TimeUnit::Second, + 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, false), + ])), + 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(); + let output_timestamps = output + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let output_values = output + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let wide = stream.scan_instants::( + ×tamps, + 0, + 1_000, + &ScanNanoseconds::::wide(TimeUnit::Second, 0, 1_000), + 2, + |_| false, + ); + assert_eq!(output_timestamps.values(), &[0]); + assert_eq!(output_values.values(), &[1.0]); + assert_eq!(wide, (vec![0], vec![0])); + assert_eq!(wide.1, output_timestamps.values().to_vec()); + } + #[test] fn exact_ties_select_first_and_lookback_uses_latest() { for (values, expected_timestamp) in [ diff --git a/src/promql/src/extension_plan/range_manipulate.rs b/src/promql/src/extension_plan/range_manipulate.rs index 8ea7ba6b191..38ecb7aa4db 100644 --- a/src/promql/src/extension_plan/range_manipulate.rs +++ b/src/promql/src/extension_plan/range_manipulate.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::HashSet; +use std::ops::{Add, Mul, Sub}; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; @@ -20,7 +21,7 @@ use std::task::{Context, Poll}; use common_telemetry::{debug, warn}; use datafusion::arrow::array::{Array, ArrayRef, Int64Array, TimestampMillisecondArray}; use datafusion::arrow::compute; -use datafusion::arrow::datatypes::{DataType, Field, SchemaRef, TimeUnit}; +use datafusion::arrow::datatypes::{DataType, Field, Int64Type, SchemaRef, TimeUnit}; use datafusion::arrow::error::ArrowError; use datafusion::arrow::record_batch::RecordBatch; use datafusion::common::stats::Precision; @@ -702,33 +703,68 @@ impl RangeManipulateStream { new_columns[*index] = new_column; } - // 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, _) = timestamp_array_to_primitive(input.column(self.time_index)) - .ok_or_else(|| { + // The timestamp range payload is always millisecond ABI. A column is reusable as-is + // only when its own data type already is that ABI: the output schema fixes the + // dictionary value type, so a millisecond column carrying a timezone must be + // converted like any other native unit. + let ts_column = input.column(self.time_index); + let timestamp_values: ArrayRef = if self.offset == 0 + && matches!( + ts_column.data_type(), + DataType::Timestamp(TimeUnit::Millisecond, None) + ) { + ts_column.clone() + } else { + // Shift in wide native precision before truncating toward zero, preserving + // null validity. + let scale = nanoseconds_per_native_tick(self.time_unit); + let (timestamps, _) = timestamp_array_to_primitive(ts_column).ok_or_else(|| { DataFusionError::Execution("Time index column is not a timestamp".into()) })?; - let timestamp_values = timestamps - .values() - .iter() - .enumerate() - .map(|(index, timestamp)| { - if !input.column(self.time_index).is_valid(index) { - return Ok(None); + // Single pass over the timestamps: an i128 intermediate holds any native + // timestamp shift; only the final millisecond result must fit i64. Null-free + // batches (the dominant PromQL case) take arrow's infallible `unary` kernel, + // which LLVM can vectorize, and report overflow through a shared cell. But + // `unary` visits NULL slots' underlying values, so a NULL slot holding an + // arbitrary overflowing value would raise a false batch-wide error; nullable + // batches therefore take `try_unary`, which applies the closure to valid rows + // only. + let offset_ns = (self.offset as i128) * 1_000_000; + let timestamp_values = if timestamps.null_count() == 0 { + let overflow = std::cell::Cell::new(false); + let shifted = compute::unary::(×tamps, |v| { + let shifted_ns = (v as i128) * scale + offset_ns; + let shifted_ms = shifted_ns / 1_000_000; + match i64::try_from(shifted_ms) { + Ok(ms) => ms, + Err(_) => { + overflow.set(true); + 0 + } + } + }); + if overflow.get() { + return Err(DataFusionError::Execution( + "RangeManipulate timestamp payload overflow".into(), + )); } - 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(|_| { + shifted + } else { + compute::try_unary::(×tamps, |v| { + let shifted_ns = (v as i128) * scale + offset_ns; + i64::try_from(shifted_ns / 1_000_000).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()) + })? + }; + Arc::new(TimestampMillisecondArray::new( + timestamp_values.values().clone(), + timestamp_values.nulls().cloned(), + )) + }; + let ts_range_column = RangeArray::from_ranges(timestamp_values, ranges.clone()) .map_err(|e| ArrowError::InvalidArgumentError(e.to_string()))? .into_dict(); new_columns.push(Arc::new(ts_range_column)); @@ -770,6 +806,7 @@ impl RangeManipulateStream { let (timestamps, _) = timestamp_array_to_primitive(ts_column).ok_or_else(|| { DataFusionError::Execution("Time index column is not a timestamp".into()) })?; + let timestamp_null_count = timestamps.null_count(); let timestamps = timestamps.values(); let timestamp = |index| (timestamps[index] as i128) * scale + (self.offset as i128) * 1_000_000; @@ -799,36 +836,193 @@ impl RangeManipulateStream { // The intersection is within the declared i64 query bounds. let start = start as i64; let end = end as i64; - let mut ranges = Vec::new(); // Range membership is decided on shifted native ticks, before the // timestamp-range payload is converted to its millisecond ABI. This // keeps sub-millisecond samples distinct in a range; equal millisecond // payload values are not a reason to deduplicate input samples. // - // Calculate for every aligned timestamp (`curr_ts`), assuming ordered timestamps. - let mut left = 0usize; - let mut right = 0usize; - for curr_ts in (start..=end).step_by(self.interval as _) { - let start_ts = (curr_ts as i128) * 1_000_000 - (self.range as i128) * 1_000_000; - - while left < len && timestamp(left) <= start_ts { - left += 1; + // Calculate for every aligned timestamp (`curr_ts`), assuming ordered + // timestamps: the scan compares shifted samples in `i64` nanoseconds + // whenever they fit the batch, and in exact `i128` nanoseconds otherwise. + let narrow = (timestamp_null_count == 0) + .then(|| { + narrow_scan( + self.time_unit, + self.offset, + self.range, + timestamps, + start, + end, + ) + }) + .flatten(); + let ranges = match narrow { + Some(scan) => scan_ranges::(self.interval, timestamps, start, end, &scan), + None => { + let scan = ScanNanoseconds::::wide(self.time_unit, self.offset, self.range); + scan_ranges::(self.interval, timestamps, start, end, &scan) } - right = right.max(left); - while right < len && timestamp(right) <= (curr_ts as i128) * 1_000_000 { - right += 1; - } - - if left == right { - ranges.push((0, 0)); - } else { - ranges.push((left as _, (right - left) as _)); - } - } + }; Ok((ranges, (start, end))) } + + /// Cursor scan of `timestamps` through the narrow `i64` carrier, for benchmarks. + /// + /// `bench_range_fn` times the fast path directly, next to the exact `i128` fallback in + /// [`RangeManipulateStream::bench_scan_wide`]; production selects the carrier through + /// [`narrow_scan`] inside `calculate_range`. Returns `None` when that carrier is + /// rejected, which is exactly the condition that makes the operator fall back to `i128`. + #[doc(hidden)] + pub fn bench_scan_narrow( + time_unit: TimeUnit, + offset: Millisecond, + window: Millisecond, + interval: Millisecond, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, + ) -> Option> { + let scan = narrow_scan(time_unit, offset, window, timestamps, start, end)?; + Some(scan_ranges::(interval, timestamps, start, end, &scan)) + } + + /// Cursor scan of `timestamps` through the exact `i128` carrier, for benchmarks. + /// + /// The carrier is exact for every input, so this times the fallback that + /// [`narrow_scan`] selects whenever the narrow baselines would overflow. + #[doc(hidden)] + pub fn bench_scan_wide( + time_unit: TimeUnit, + offset: Millisecond, + window: Millisecond, + interval: Millisecond, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, + ) -> Vec<(u32, u32)> { + scan_ranges::( + interval, + timestamps, + start, + end, + &ScanNanoseconds::::wide(time_unit, offset, window), + ) + } +} + +/// Cursor scan over the aligned evaluation points, in `Ns` nanoseconds. +fn scan_ranges( + interval: Millisecond, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, + scan: &ScanNanoseconds, +) -> Vec<(u32, u32)> { + let len = timestamps.len(); + let mut ranges = Vec::new(); + let mut left = 0usize; + let mut right = 0usize; + for curr_ts in (start..=end).step_by(interval as _) { + let curr_ts_ns = scan.instant(curr_ts); + let start_ts = scan.window_start(curr_ts_ns); + + while left < len && scan.shift(timestamps[left]) <= start_ts { + left += 1; + } + right = right.max(left); + while right < len && scan.shift(timestamps[right]) <= curr_ts_ns { + right += 1; + } + + if left == right { + ranges.push((0, 0)); + } else { + ranges.push((left as _, (right - left) as _)); + } + } + ranges +} + +/// Cursor scan baselines in nanoseconds, shared with `InstantManipulate`. +/// +/// The narrow `i64` baselines come from [`narrow_scan`], the exact ones from +/// [`ScanNanoseconds::wide`]. +#[derive(Debug, Clone, Copy)] +pub(super) struct ScanNanoseconds { + /// Nanoseconds per timestamp tick. + scale: Ns, + /// The offset shift in nanoseconds. + offset_ns: Ns, + /// The range or lookback window in nanoseconds. + window_ns: Ns, +} + +/// Nanoseconds carrier of the cursor scans: `i64` where [`narrow_scan`] proved that +/// the batch fits, `i128` for the exact fallback. +pub(super) trait Nanoseconds: + Copy + PartialOrd + From + Add + Sub + Mul +{ +} + +impl Nanoseconds for T where + T: Copy + PartialOrd + From + Add + Sub + Mul +{ +} + +impl ScanNanoseconds { + /// The `i128` baselines, exact for every input. + pub(super) fn wide(time_unit: TimeUnit, offset: Millisecond, window: Millisecond) -> Self { + Self { + scale: nanoseconds_per_native_tick(time_unit), + offset_ns: i128::from(offset) * 1_000_000, + window_ns: i128::from(window) * 1_000_000, + } + } +} + +impl ScanNanoseconds { + /// Shifted nanoseconds of an input timestamp. + pub(super) fn shift(&self, timestamp: i64) -> Ns { + Ns::from(timestamp) * self.scale + self.offset_ns + } + + /// Nanoseconds of an evaluation instant given in milliseconds. + pub(super) fn instant(&self, millis: Millisecond) -> Ns { + Ns::from(millis) * Ns::from(1_000_000) + } + + /// The exclusive lower boundary of the window ending at `instant`. + pub(super) fn window_start(&self, instant: Ns) -> Ns { + instant - self.window_ns + } +} + +/// The narrow `i64` baselines of a sorted batch, or `None` when they would overflow. +pub(super) fn narrow_scan( + time_unit: TimeUnit, + offset: Millisecond, + window: Millisecond, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, +) -> Option> { + let scale = i64::try_from(nanoseconds_per_native_tick(time_unit)).ok()?; + let offset_ns = offset.checked_mul(1_000_000)?; + let window_ns = window.checked_mul(1_000_000)?; + // Sorted samples, so the checked endpoints bound every interior shift and product. + let shift = |timestamp: i64| timestamp.checked_mul(scale)?.checked_add(offset_ns); + shift(*timestamps.first()?)?; + shift(*timestamps.last()?)?; + // The checked `[start, end]` bounds every instant and every `instant - window`. + start.checked_mul(1_000_000)?.checked_sub(window_ns)?; + end.checked_mul(1_000_000)?.checked_sub(window_ns)?; + Some(ScanNanoseconds { + scale, + offset_ns, + window_ns, + }) } #[cfg(test)] @@ -859,6 +1053,9 @@ mod test { const TIME_INDEX_COLUMN: &str = "timestamp"; + /// The `(offset, length)` selections of one cursor scan. + type RangeSelections = Vec<(u32, u32)>; + fn project_batch(batch: &RecordBatch, indices: &[usize]) -> RecordBatch { let fields = indices .iter() @@ -1711,6 +1908,36 @@ mod test { stream.calculate_range(&batch).unwrap() } + #[test] + fn calculate_range_nullable_interior_timestamp_matches_wide_scan() { + let schema = Arc::new(Schema::new(vec![Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Second, None), + true, + )])); + let stream = scan_stream_for_test(TimeUnit::Second, 0, 0, 1_000, 1_000, 1_000); + let timestamps = [0, i64::MAX, 1]; + let array = TimestampSecondArray::new( + timestamps.to_vec().into(), + Some(NullBuffer::from(vec![false, false, true])), + ); + let batch = RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap(); + + let (ranges, bounds) = stream.calculate_range(&batch).unwrap(); + let wide = RangeManipulateStream::bench_scan_wide( + TimeUnit::Second, + 0, + 1_000, + 1_000, + ×tamps, + bounds.0, + bounds.1, + ); + assert_eq!(bounds, (0, 1_000)); + assert_eq!(ranges, vec![(0, 1), (0, 0)]); + assert_eq!(ranges, wide); + } + #[test] fn calculate_range_keeps_query_aligned_tail() { let (ranges, bounds) = calculate_range_for_test(4, 94, 30, 15, &[20, 50, 80]); @@ -1954,4 +2181,898 @@ mod test { ); } } + + /// Nanoseconds per native tick, as kept by the narrow carrier. + fn native_scale(time_unit: TimeUnit) -> i64 { + i64::try_from(nanoseconds_per_native_tick(time_unit)).unwrap() + } + + /// Shifted nanoseconds of a sample, computed in the exact `i128` carrier. + fn shifted_nanoseconds(time_unit: TimeUnit, offset: Millisecond, timestamp: i64) -> i128 { + (timestamp as i128) * nanoseconds_per_native_tick(time_unit) + (offset as i128) * 1_000_000 + } + + /// Shifted milliseconds of the instant grid of a sample, flooring like `calculate_range`. + fn shifted_floor_millis( + time_unit: TimeUnit, + offset: Millisecond, + timestamp: i64, + ) -> Millisecond { + Millisecond::try_from( + shifted_nanoseconds(time_unit, offset, timestamp).div_euclid(1_000_000), + ) + .unwrap() + } + + /// Shifted milliseconds kept by the range payload of a sample. + /// + /// The payload conversion truncates toward zero, so samples below the epoch may keep + /// the millisecond above their own native tick. + fn shifted_millis(time_unit: TimeUnit, offset: Millisecond, timestamp: i64) -> Millisecond { + Millisecond::try_from(shifted_nanoseconds(time_unit, offset, timestamp) / 1_000_000) + .unwrap() + } + + /// Stream whose cursor scan the differential tests below drive directly. + fn scan_stream_for_test( + time_unit: TimeUnit, + offset: Millisecond, + start: Millisecond, + end: Millisecond, + interval: Millisecond, + range: Millisecond, + ) -> RangeManipulateStream { + let schema = Arc::new(Schema::new(vec![Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(time_unit, None), + false, + )])); + let empty_stream = MemoryStream::try_new(vec![], schema.clone(), None).unwrap(); + RangeManipulateStream { + offset, + start, + end, + interval, + range, + time_index: 0, + time_unit, + field_columns: vec![], + aligned_ts_array: Arc::new(TimestampMillisecondArray::from(vec![0i64; 0])), + output_schema: schema, + input: Box::pin(empty_stream), + metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + num_series: Count::new(), + } + } + + /// Exact expectations of the cursor scan, formulated outside the production loops. + /// + /// Every evaluation instant keeps the samples of `(instant - range, instant]` in + /// shifted nanoseconds: the lower boundary is exclusive, the upper one inclusive, and + /// native ticks below a millisecond stay distinct. + fn scan_oracle( + time_unit: TimeUnit, + offset: Millisecond, + range: Millisecond, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, + interval: Millisecond, + ) -> Vec<(u32, u32)> { + let window_ns = (range as i128) * 1_000_000; + let shifted = timestamps + .iter() + .map(|timestamp| shifted_nanoseconds(time_unit, offset, *timestamp)) + .collect::>(); + (start..=end) + .step_by(interval as usize) + .map(|instant| { + let instant_ns = (instant as i128) * 1_000_000; + let mut offset = None; + let mut length = 0u32; + for (index, timestamp) in shifted.iter().enumerate() { + if *timestamp > instant_ns - window_ns && *timestamp <= instant_ns { + offset.get_or_insert(index as u32); + length += 1; + } + } + (offset.unwrap_or(0), length) + }) + .collect() + } + + /// Runs the cursor scan of one input through both carriers. + fn scan_both_carriers( + stream: &RangeManipulateStream, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, + ) -> (Option, RangeSelections) { + let exact = scan_ranges::( + stream.interval, + timestamps, + start, + end, + &ScanNanoseconds::::wide(stream.time_unit, stream.offset, stream.range), + ); + let narrow = narrow_scan( + stream.time_unit, + stream.offset, + stream.range, + timestamps, + start, + end, + ) + .map(|scan| scan_ranges::(stream.interval, timestamps, start, end, &scan)); + // The benchmark hooks must measure the very scans the operator runs. + assert_eq!( + RangeManipulateStream::bench_scan_wide( + stream.time_unit, + stream.offset, + stream.range, + stream.interval, + timestamps, + start, + end, + ), + exact, + "the i128 benchmark hook diverged from the cursor scan" + ); + if let Some(narrow) = &narrow { + assert_eq!( + RangeManipulateStream::bench_scan_narrow( + stream.time_unit, + stream.offset, + stream.range, + stream.interval, + timestamps, + start, + end, + ), + Some(narrow.clone()), + "the i64 benchmark hook diverged from the cursor scan" + ); + } else { + assert_eq!( + RangeManipulateStream::bench_scan_narrow( + stream.time_unit, + stream.offset, + stream.range, + stream.interval, + timestamps, + start, + end, + ), + None, + "the i64 benchmark hook must report the fallback" + ); + } + (narrow, exact) + } + + /// Asserts the narrow carrier was taken and matches the exact scan evaluation by evaluation. + fn assert_narrow_carrier_agrees( + narrow: Option>, + exact: &[(u32, u32)], + context: &str, + ) { + let Some(narrow) = narrow else { + panic!("the i64 carrier was expected, but narrow_scan rejected the input: {context}"); + }; + assert_eq!( + narrow.len(), + exact.len(), + "carriers disagree on the number of evaluations: {context}" + ); + for (evaluation, (narrow_range, exact_range)) in narrow.iter().zip(exact).enumerate() { + assert_eq!( + narrow_range, exact_range, + "carriers disagree on evaluation {evaluation}: {context}" + ); + } + } + + /// Replays the cursor scan in checked `i64` arithmetic. + /// + /// Reports whether the narrow carrier overflows on an operation the scan evaluates, + /// which is what makes the `i128` fallback unavoidable for an input. + fn narrow_carrier_overflows( + time_unit: TimeUnit, + offset: Millisecond, + range: Millisecond, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, + interval: Millisecond, + ) -> bool { + let Ok(scale) = i64::try_from(nanoseconds_per_native_tick(time_unit)) else { + return true; + }; + let Some(offset_ns) = offset.checked_mul(1_000_000) else { + return true; + }; + let Some(window_ns) = range.checked_mul(1_000_000) else { + return true; + }; + let shift = |timestamp: i64| timestamp.checked_mul(scale)?.checked_add(offset_ns); + let mut left = 0usize; + let mut right = 0usize; + for instant in (start..=end).step_by(interval as usize) { + let Some(instant_ns) = instant.checked_mul(1_000_000) else { + return true; + }; + let Some(window_start) = instant_ns.checked_sub(window_ns) else { + return true; + }; + while left < timestamps.len() { + let Some(shifted) = shift(timestamps[left]) else { + return true; + }; + if shifted <= window_start { + left += 1; + } else { + break; + } + } + right = right.max(left); + while right < timestamps.len() { + let Some(shifted) = shift(timestamps[right]) else { + return true; + }; + if shifted <= instant_ns { + right += 1; + } else { + break; + } + } + } + false + } + + /// The checked `i64` operation that makes [`narrow_scan`] reject an input, if any. + /// + /// Mirrors the rejection criteria of [`narrow_scan`] so the deterministic cases below + /// can pin which boundary drives the `i128` fallback. + fn narrow_carrier_rejection( + time_unit: TimeUnit, + offset: Millisecond, + range: Millisecond, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, + ) -> Option<&'static str> { + let Ok(scale) = i64::try_from(nanoseconds_per_native_tick(time_unit)) else { + return Some("nanoseconds per native tick do not fit i64"); + }; + let Some(offset_ns) = offset.checked_mul(1_000_000) else { + return Some("offset in nanoseconds does not fit i64"); + }; + let Some(window_ns) = range.checked_mul(1_000_000) else { + return Some("range in nanoseconds does not fit i64"); + }; + let shift = |timestamp: i64| timestamp.checked_mul(scale)?.checked_add(offset_ns); + let (Some(first), Some(last)) = (timestamps.first(), timestamps.last()) else { + return Some("empty batch has no endpoints"); + }; + if shift(*first).is_none() || shift(*last).is_none() { + return Some("sample nanoseconds do not fit i64"); + } + if start + .checked_mul(1_000_000) + .and_then(|nanoseconds| nanoseconds.checked_sub(window_ns)) + .is_none() + { + return Some("start bound in nanoseconds does not fit i64"); + } + if end + .checked_mul(1_000_000) + .and_then(|nanoseconds| nanoseconds.checked_sub(window_ns)) + .is_none() + { + return Some("end bound in nanoseconds do not fit i64"); + } + None + } + + #[test] + fn scan_carriers_agree_on_range_boundaries_within_one_millisecond() { + // Membership is `(lower, upper]` on native ticks: the sample exactly on the lower + // boundary is dropped, the sample exactly on the upper boundary is kept, and the + // three samples sharing one millisecond payload stay distinct and survive together. + let cases = [ + ( + TimeUnit::Microsecond, + vec![1_000_000, 1_000_001, 1_000_999, 1_001_000, 1_001_001], + ), + ( + TimeUnit::Nanosecond, + vec![ + 1_000_000_000, + 1_000_000_001, + 1_000_000_999, + 1_001_000_000, + 1_001_000_001, + ], + ), + ]; + + for (time_unit, timestamps) in cases { + let (start, end, interval, range) = (1_001, 1_001, 1, 1); + let context = format!("unit={time_unit:?}, timestamps={timestamps:?}"); + let stream = scan_stream_for_test(time_unit, 0, start, end, interval, range); + + let oracle = scan_oracle(time_unit, 0, range, ×tamps, start, end, interval); + assert_eq!( + oracle, + vec![(1, 3)], + "{context}: (lower, upper] expectations" + ); + + let (narrow, exact) = scan_both_carriers(&stream, ×tamps, start, end); + assert_eq!(exact, oracle, "{context}: i128 carrier must stay exact"); + assert_narrow_carrier_agrees(narrow, &exact, &context); + + // Two distinct native ticks below the upper boundary share the millisecond + // payload `1_000`; they must not be deduplicated by going through that ABI. + let payloads = timestamps[1..4] + .iter() + .map(|timestamp| shifted_millis(time_unit, 0, *timestamp)) + .collect::>(); + assert_eq!( + payloads, + vec![1_000, 1_000, 1_001], + "{context}: kept samples in one millisecond keep both native ticks" + ); + assert_ne!( + timestamps[1], timestamps[2], + "{context}: the two ticks sharing a payload must stay distinct" + ); + } + } + + #[test] + fn scan_falls_back_to_i128_when_only_the_sample_tail_overflows() { + // `narrow_scan` inspects the sorted endpoints, so a batch whose prefix fits the + // `i64` carrier but whose tail does not has to take the exact carrier. The tail + // sample stays outside every window, so the selections keep the prefix ones. + struct Case { + time_unit: TimeUnit, + narrow_safe: Vec, + overflowing: i64, + start: Millisecond, + end: Millisecond, + interval: Millisecond, + range: Millisecond, + } + + let cases = [ + Case { + time_unit: TimeUnit::Second, + narrow_safe: vec![9_223_372_034, 9_223_372_035, 9_223_372_036], + overflowing: 9_223_372_037, + start: 9_223_372_034_000, + end: 9_223_372_036_000, + interval: 1_000, + range: 1_000, + }, + Case { + time_unit: TimeUnit::Millisecond, + narrow_safe: vec![9_223_372_036_853, 9_223_372_036_854], + overflowing: 9_223_372_036_855, + start: 9_223_372_036_853, + end: 9_223_372_036_854, + interval: 1, + range: 1, + }, + Case { + time_unit: TimeUnit::Microsecond, + narrow_safe: vec![9_223_372_036_853_999, 9_223_372_036_854_000], + overflowing: 9_223_372_036_854_776, + start: 9_223_372_036_854, + end: 9_223_372_036_854, + interval: 1, + range: 1, + }, + ]; + + for case in cases { + let Case { + time_unit, + narrow_safe, + overflowing, + start, + end, + interval, + range, + } = case; + let context = format!("unit={time_unit:?}, timestamps={narrow_safe:?} + {overflowing}"); + let scale = native_scale(time_unit); + // The prefix fits the narrow carrier, the tail sample does not... + assert!( + narrow_safe + .iter() + .all(|timestamp| timestamp.checked_mul(scale).is_some()), + "{context}: prefix samples must fit i64 nanoseconds" + ); + assert!( + overflowing.checked_mul(scale).is_none(), + "{context}: the tail sample must overflow i64 nanoseconds" + ); + // ...even though the millisecond payload of the tail stays legal. + let payload = shifted_millis(time_unit, 0, overflowing); + assert!( + payload.abs() < i64::MAX / 2, + "{context}: the shifted payload {payload} must stay legal" + ); + + let stream = scan_stream_for_test(time_unit, 0, start, end, interval, range); + + let (narrow, prefix_exact) = scan_both_carriers(&stream, &narrow_safe, start, end); + let prefix_oracle = + scan_oracle(time_unit, 0, range, &narrow_safe, start, end, interval); + assert_eq!( + prefix_exact, prefix_oracle, + "{context}: prefix i128 carrier" + ); + assert_narrow_carrier_agrees(narrow, &prefix_exact, &format!("{context} prefix")); + + let mut timestamps = narrow_safe.clone(); + timestamps.push(overflowing); + assert_eq!( + narrow_carrier_rejection(time_unit, 0, range, ×tamps, start, end), + Some("sample nanoseconds do not fit i64"), + "{context}: the overflowing endpoint must reject the narrow carrier" + ); + assert!( + narrow_carrier_overflows(time_unit, 0, range, ×tamps, start, end, interval), + "{context}: the scan evaluates the overflowing tail" + ); + + let (narrow, exact) = scan_both_carriers(&stream, ×tamps, start, end); + assert!( + narrow.is_none(), + "{context}: the i128 carrier must be taken" + ); + let oracle = scan_oracle(time_unit, 0, range, ×tamps, start, end, interval); + assert_eq!(exact, oracle, "{context}: exact carrier on the full batch"); + assert_eq!( + exact, prefix_exact, + "{context}: the out-of-window tail keeps the prefix selections" + ); + } + } + + #[test] + fn scan_carriers_split_at_the_nanosecond_boundaries() { + // Nanoseconds are not scaled, so the `i64` carrier is rejected by the query bounds: + // below the millisecond of `i64::MIN` the start bound underflows, while a batch near + // `i64::MAX` still fits. Both keep a legal millisecond payload. + let low = i64::MIN; + let low_floor = shifted_floor_millis(TimeUnit::Nanosecond, 0, low); + assert_eq!(low_floor, -9_223_372_036_855); + assert!(low_floor.checked_mul(1_000_000).is_none()); + assert_eq!( + shifted_millis(TimeUnit::Nanosecond, 0, low), + -9_223_372_036_854 + ); + + let timestamps = vec![low, low + 1_000_000]; + let (start, end, interval, range) = (low_floor, low_floor + 1, 1, 1); + let context = format!("nanosecond floor, timestamps={timestamps:?}"); + let stream = scan_stream_for_test(TimeUnit::Nanosecond, 0, start, end, interval, range); + assert_eq!( + narrow_carrier_rejection(TimeUnit::Nanosecond, 0, range, ×tamps, start, end), + Some("start bound in nanoseconds does not fit i64"), + "{context}" + ); + let (narrow, exact) = scan_both_carriers(&stream, ×tamps, start, end); + assert!( + narrow.is_none(), + "{context}: the i128 carrier must be taken" + ); + let oracle = scan_oracle( + TimeUnit::Nanosecond, + 0, + range, + ×tamps, + start, + end, + interval, + ); + assert_eq!(exact, oracle, "{context}: i128 carrier must stay exact"); + + // The ceiling side still fits the narrow carrier, so both carriers have to agree. + let high = i64::MAX - 2_000_000; + let high_floor = shifted_floor_millis(TimeUnit::Nanosecond, 0, high); + assert!(high_floor.checked_mul(1_000_000).is_some()); + let timestamps = vec![high, high + 1_000_000, high + 2_000_000]; + let (start, end, interval, range) = (high_floor, high_floor + 2, 1, 1); + let context = format!("nanosecond ceiling, timestamps={timestamps:?}"); + let stream = scan_stream_for_test(TimeUnit::Nanosecond, 0, start, end, interval, range); + assert_eq!( + narrow_carrier_rejection(TimeUnit::Nanosecond, 0, range, ×tamps, start, end), + None, + "{context}: the ceiling batch still fits the narrow carrier" + ); + let (narrow, exact) = scan_both_carriers(&stream, ×tamps, start, end); + let oracle = scan_oracle( + TimeUnit::Nanosecond, + 0, + range, + ×tamps, + start, + end, + interval, + ); + assert_eq!(exact, oracle, "{context}: i128 carrier must stay exact"); + assert_narrow_carrier_agrees(narrow, &exact, &context); + } + + #[test] + fn scan_carriers_agree_for_seeded_native_precision_matrix() { + // Seeded matrix over all four native precisions: positive and negative offsets, + // negative times, sub-millisecond samples sharing a millisecond payload, and native + // extremes whose intermediate product overflows while the shifted millisecond stays + // legal. + let units = [ + TimeUnit::Second, + TimeUnit::Millisecond, + TimeUnit::Microsecond, + TimeUnit::Nanosecond, + ]; + // `narrow_scan` rejects the two large offsets: `offset * 1ms` overflows `i64`. + let offsets = [0, 1, -1, 250, -250, 9_223_372_036_855, -9_223_372_036_855]; + let extremes = |time_unit: TimeUnit| match time_unit { + TimeUnit::Second => [-9_223_372_037, 9_223_372_037], + TimeUnit::Millisecond => [-9_223_372_036_855, 9_223_372_036_855], + TimeUnit::Microsecond => [-9_223_372_036_854_776, 9_223_372_036_854_776], + TimeUnit::Nanosecond => [i64::MIN, i64::MAX], + }; + let mut prng = TinyPrng(0x5ca1_ab1e_d1ff_0001); + let mut narrow_cases = 0usize; + let mut wide_cases = 0usize; + let mut skipped = 0usize; + + for case in 0..256usize { + let time_unit = units[case % units.len()]; + let offset = offsets[(case / units.len()) % offsets.len()]; + let [low, high] = extremes(time_unit); + let anchor = match prng.next_i64(0, 3) { + 0 => prng.next_i64(-2_000, 2_000), + 1 => -prng.next_i64(1, 2_000), + 2 => low, + _ => high, + }; + // Samples stay within a few milliseconds, with a native granularity small + // enough that sub-millisecond samples share a millisecond payload. + let step = match time_unit { + TimeUnit::Second => 1, + TimeUnit::Millisecond => 1, + TimeUnit::Microsecond => 1 + prng.next_i64(0, 9), + TimeUnit::Nanosecond => 1 + prng.next_i64(0, 999_999), + }; + let count = 1 + prng.next_i64(0, 4); + let ascending = anchor != high; + let mut timestamps = (0..count) + .map(|index| { + let delta = index * step; + if ascending { + anchor + delta + } else { + anchor - delta + } + }) + .collect::>(); + if !ascending { + timestamps.reverse(); + } + // The review shape keeps a legal millisecond payload; skip the extremes that the + // large offsets push out of the millisecond space. + let nanoseconds = timestamps + .iter() + .map(|timestamp| shifted_nanoseconds(time_unit, offset, *timestamp)) + .collect::>(); + if nanoseconds + .iter() + .any(|value| i64::try_from(value / 1_000_000).is_err()) + { + skipped += 1; + continue; + } + + let first_millis = Millisecond::try_from(nanoseconds[0].div_euclid(1_000_000)).unwrap(); + let interval = prng.next_i64(1, 4); + let points = 1 + prng.next_i64(0, 4); + let start = first_millis - interval; + let end = start + interval * (points - 1); + let range = prng.next_i64(0, 3) * interval; + let context = format!( + "case={case}, unit={time_unit:?}, offset={offset}, start={start}, end={end}, \ + interval={interval}, range={range}, timestamps={timestamps:?}" + ); + + let stream = scan_stream_for_test(time_unit, offset, start, end, interval, range); + let oracle = scan_oracle(time_unit, offset, range, ×tamps, start, end, interval); + let (narrow, exact) = scan_both_carriers(&stream, ×tamps, start, end); + assert_eq!(exact, oracle, "{context}: i128 carrier must stay exact"); + + match ( + &narrow, + narrow_carrier_rejection(time_unit, offset, range, ×tamps, start, end), + ) { + (Some(_), None) => { + narrow_cases += 1; + assert!( + !narrow_carrier_overflows( + time_unit, + offset, + range, + ×tamps, + start, + end, + interval + ), + "{context}: every narrow operation must fit i64" + ); + assert_narrow_carrier_agrees(narrow, &exact, &context); + } + (None, Some(_reason)) => { + wide_cases += 1; + } + (narrow, rejection) => panic!( + "{context}: narrow_scan and its mirror disagree, narrow={}, rejection={rejection:?}", + narrow.is_some() + ), + } + } + + assert!( + narrow_cases > 0 && wide_cases > 0, + "both carriers must be covered: {narrow_cases} narrow, {wide_cases} wide \ + ({skipped} skipped)" + ); + assert!( + skipped < 64, + "the generator skipped too many cases: {skipped}" + ); + } + + /// Builds the single-column timestamp batch the cursor scans consume. + fn timestamp_batch(time_unit: TimeUnit, timestamps: &[i64]) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(time_unit, None), + false, + )])); + let array: ArrayRef = match time_unit { + TimeUnit::Second => Arc::new(TimestampSecondArray::from(timestamps.to_vec())), + TimeUnit::Millisecond => Arc::new(TimestampMillisecondArray::from(timestamps.to_vec())), + TimeUnit::Microsecond => Arc::new(TimestampMicrosecondArray::from(timestamps.to_vec())), + TimeUnit::Nanosecond => Arc::new(TimestampNanosecondArray::from(timestamps.to_vec())), + }; + RecordBatch::try_new(schema, vec![array]).unwrap() + } + + #[test] + fn calculate_range_dispatches_carriers_and_agrees_pointwise() { + // Production dispatch, not the hooks: `calculate_range` must scan through the narrow + // carrier while the batch fits and through the exact one once the shifted tail does + // not, and both must select the same ranges for the bounds it derives itself. + let narrow_timestamps = [0i64, 15_000, 30_000]; + let wide_timestamps = [9_223_372_036_853, 9_223_372_036_854, 9_223_372_036_855]; + let cases: [(&str, &[i64], Millisecond, Millisecond, bool); 2] = [ + ("narrow carrier", &narrow_timestamps, 0, 30_000, true), + ( + "i128 fallback", + &wide_timestamps, + wide_timestamps[0], + wide_timestamps[1], + false, + ), + ]; + + for (name, timestamps, query_start, query_end, expect_narrow) in cases { + let (interval, range) = if expect_narrow { + (15_000, 15_000) + } else { + (1, 1) + }; + // `scan_stream_for_test` panics on an empty batch when `narrow_scan` is asked for + // the endpoints, so the batch below is never empty. + let stream = scan_stream_for_test( + TimeUnit::Millisecond, + 0, + query_start, + query_end, + interval, + range, + ); + let batch = timestamp_batch(TimeUnit::Millisecond, timestamps); + let (ranges, (start, end)) = stream.calculate_range(&batch).unwrap(); + + let oracle = scan_oracle( + TimeUnit::Millisecond, + 0, + range, + timestamps, + start, + end, + interval, + ); + assert_eq!(ranges, oracle, "{name}: calculate_range pointwise"); + + let rejection = + narrow_carrier_rejection(TimeUnit::Millisecond, 0, range, timestamps, start, end); + assert_eq!( + rejection.is_none(), + expect_narrow, + "{name}: carrier choice by narrow_scan, rejection={rejection:?}" + ); + let (narrow, exact) = scan_both_carriers(&stream, timestamps, start, end); + assert_eq!(exact, ranges, "{name}: i128 carrier is the oracle"); + assert_eq!(narrow.is_some(), expect_narrow, "{name}: narrow carrier"); + assert_eq!( + RangeManipulateStream::bench_scan_wide( + stream.time_unit, + stream.offset, + stream.range, + interval, + timestamps, + start, + end, + ), + ranges, + "{name}: the benchmark hook matches calculate_range" + ); + } + } + + /// Stream whose single millisecond timestamp column is folded into the range payload. + fn payload_stream( + offset: Millisecond, + start: Millisecond, + end: Millisecond, + interval: Millisecond, + range: Millisecond, + ) -> RangeManipulateStream { + let payload_type = DataType::Timestamp(TimeUnit::Millisecond, None); + let schema = Arc::new(Schema::new(vec![ + Field::new(TIME_INDEX_COLUMN, payload_type.clone(), false), + Field::new( + RangeManipulate::build_timestamp_range_name(TIME_INDEX_COLUMN), + RangeArray::convert_data_type(payload_type), + false, + ), + ])); + let mut stream = + scan_stream_for_test(TimeUnit::Millisecond, offset, start, end, interval, range); + stream.output_schema = schema; + stream + } + + /// Builds the single-column millisecond timestamp batch the payload tests consume. + fn timestamp_batch_with_timezone(timezone: Option<&str>, timestamps: &[i64]) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Millisecond, timezone.map(Arc::from)), + false, + )])); + let array: ArrayRef = Arc::new( + TimestampMillisecondArray::from(timestamps.to_vec()).with_timezone_opt(timezone), + ); + RecordBatch::try_new(schema, vec![array]).unwrap() + } + + /// Millisecond payload values of the range timestamp column of a manipulated batch. + fn payload_millis(batch: &RecordBatch) -> Vec { + let dict = batch + .column(batch.num_columns() - 1) + .as_any() + .downcast_ref::>() + .expect("the timestamp range payload must be a folded dictionary") + .clone(); + let payload = RangeArray::try_new(dict).unwrap(); + assert_eq!( + payload.values().data_type(), + &DataType::Timestamp(TimeUnit::Millisecond, None), + "the payload value type must agree with the output schema" + ); + payload + .values() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + } + + #[test] + fn manipulate_timezone_time_index_skips_zero_copy_payload() { + // Oracle-verified regression: `Timestamp(Millisecond, Some(tz))` passes a + // `time_unit == Millisecond` check but is not the payload ABI. Cloning such a column + // would stamp a dictionary value type disagreeing with the output schema and + // `RecordBatch::try_new` would fail, so the reuse must be guarded on the data type. + let timestamps = [0i64, 1_000, 2_000, 3_000]; + let stream = payload_stream(0, 0, 3_000, 1_000, 1_000); + + // With offset 0 the data type is the only reason to convert, and the `try_unary` + // path is the identity for a millisecond index. + let zoned = stream + .manipulate(timestamp_batch_with_timezone(Some("+08:00"), ×tamps)) + .expect("a timezone-carrying index must still build the output batch") + .expect("the ranges are not all empty"); + assert_eq!( + payload_millis(&zoned), + timestamps, + "the converted payload must keep the try_unary values" + ); + + // The same logical samples without a timezone do take the zero-copy path; both paths + // must agree on the payload. + let plain = stream + .manipulate(timestamp_batch(TimeUnit::Millisecond, ×tamps)) + .unwrap() + .unwrap(); + assert_eq!( + payload_millis(&plain), + payload_millis(&zoned), + "the zero-copy and the converted path must agree" + ); + } + + /// A NULL slot's underlying value must never trigger the payload overflow + /// error: a nullable batch is converted with arrow's `try_unary`, which applies + /// the closure to valid rows only. The same value as a *valid* sample must still + /// error, through the null-free `unary` + overflow-cell path. + #[test] + fn manipulate_null_slot_underlying_value_does_not_overflow() { + // Second-slot underlying value i64::MAX seconds overflows on conversion, + // but the slot is NULL so the conversion must not see it. offset=1ms + // keeps the batch off the zero-copy path. The query range must be bounded + // and contain the first sample: with `end == 0` no range is selected and + // the conversion (and so the regression) is skipped, while an unbounded + // `end` makes the range scan visit `end - start` evaluation points. + let timestamps = [0i64, i64::MAX]; + let validity = NullBuffer::from(vec![true, false]); + let array: ArrayRef = Arc::new(TimestampSecondArray::new( + timestamps.to_vec().into(), + Some(validity), + )); + let schema = Arc::new(Schema::new(vec![Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Second, None), + true, + )])); + let batch = RecordBatch::try_new(schema, vec![array]).unwrap(); + + let mut stream = payload_stream(1, 0, 1_000, 1, 1_000); + stream.time_unit = TimeUnit::Second; + let output = stream + .manipulate(batch) + .expect("a NULL slot's underlying value must not overflow the payload") + .expect("the first sample (0s + 1ms) is inside the query range"); + assert_eq!( + payload_millis(&output)[0], + 1, + "the valid sample must be converted to its shifted millisecond payload" + ); + + // Same underlying value, but now valid: the conversion must reject it. + let array: ArrayRef = Arc::new(TimestampSecondArray::from(timestamps.to_vec())); + let schema = Arc::new(Schema::new(vec![Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Second, None), + false, + )])); + let batch = RecordBatch::try_new(schema, vec![array]).unwrap(); + let mut stream = payload_stream(1, 0, 1_000, 1, 1_000); + stream.time_unit = TimeUnit::Second; + let error = stream + .manipulate(batch) + .expect_err("a valid overflowing timestamp must still error"); + assert!( + error.to_string().contains("timestamp payload overflow"), + "unexpected error: {error}" + ); + } } diff --git a/tests/perf/query_cases/promql_long_range_rate/case.toml b/tests/perf/query_cases/promql_long_range_rate/case.toml new file mode 100644 index 00000000000..51dfe1f1a86 --- /dev/null +++ b/tests/perf/query_cases/promql_long_range_rate/case.toml @@ -0,0 +1,119 @@ +# PromQL long-window rate benchmark for issue #9070. +# +# 2048 series × 1440 samples at 15-second cadence = 2,949,120 rows spanning the +# six hours from 2024-01-01T00:00:00Z to 2024-01-01T06:00:00Z. Evaluation covers +# the last two hours of that data (04:00:00Z–06:00:00Z at a 15-second step, 481 +# evaluations), so every evaluation holds a full 240-sample one-hour window. +# `rate_1h` is the primary regressed shape: it drives the RangeManipulate +# timestamp payload and the cursor scan across a long, heavily overlapping +# window set. `rate_5m` is the small-window control, `offset 1h` exercises the +# shifted payload path, and the bare selector controls unrelated-path variance. +# +# The generator derives each tag value as `series % cardinality`, so the number +# of distinct (host, instance) series is the least common multiple of the two +# cardinalities. `instance` declares 2048 values so the host(16) × instance tag +# shape still yields 2048 distinct series; a 128-value instance tag would cap +# this fixture at 128 distinct series. +# +# The deterministic wave is not a true monotonic counter, so the rate queries +# measure latency on the regular rate path rather than a numeric counter-rate +# oracle. + +[case] +name = "promql_long_range_rate" +description = "PromQL long-window rate latency regression for #9070 over six hours of 15-second samples" +issue = "https://github.com/GreptimeTeam/greptimedb/issues/9070" + +[scenario] +kind = "direct_readable_sst" +seed = 9070 + +[[scenario.tables]] +database = "public" +name = "promql_long_range_rate" +engine = "mito" +append_mode = true +sst_format = "flat" +primary_key = ["host", "instance"] +time_index = "ts" + +[[scenario.tables.columns]] +name = "host" +type = "STRING" +semantic = "tag" +distribution = { kind = "cardinality", values = 16, prefix = "host" } + +[[scenario.tables.columns]] +name = "instance" +type = "STRING" +semantic = "tag" +distribution = { kind = "cardinality", values = 2048, prefix = "instance" } + +[[scenario.tables.columns]] +name = "value" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 1000.0 } + +[[scenario.tables.columns]] +name = "ts" +type = "TIMESTAMP(9)" +semantic = "timestamp" + +# Eight 45-minute SSTs (180 timestamps each) tile the six-hour range. +# A (03:00:00Z, 06:00:00Z] scan prunes to four of them; the offset query scans +# the five SSTs covering [02:00:00Z, 05:00:00Z). +[scenario.layout] +regions = 1 +sst_count = 8 +rows_per_sst = 368640 +row_group_size = 36864 +series_count = 2048 +start_unix_nanos = 1704067200000000000 +step_nanos = 15000000000 +time_range_layout = "non_overlapping_per_sst" +series_layout = "timestamp_major" + +# 240 samples per full window: the primary regressed shape. +[[scenario.queries]] +name = "rate_1h" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704081600, 1704088800, '15s') rate(promql_long_range_rate{host=~'host.*'}[1h])" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 20 + +# Twenty samples per full window: small-window control. +[[scenario.queries]] +name = "rate_5m" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704081600, 1704088800, '15s') rate(promql_long_range_rate{host=~'host.*'}[5m])" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 20 + +# The same long window shifted back an hour: non-zero-offset payload path. +[[scenario.queries]] +name = "rate_1h_offset_1h" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704081600, 1704088800, '15s') rate(promql_long_range_rate{host=~'host.*'}[1h] offset 1h)" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 20 + +# Bare selector over the same evaluation range: unrelated-path variance control. +[[scenario.queries]] +name = "plain_selector_control" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704081600, 1704088800, '15s') promql_long_range_rate{host=~'host.*'}" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 20