From 7385cfd3217a1998fe62cbf0c42f2d0a2ce0d73d Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Tue, 8 Sep 2026 15:10:27 +0100 Subject: [PATCH 1/2] Change DateTimeParts to i32 Signed-off-by: Mikhail Kot --- Cargo.lock | 2 + encodings/datetime-parts/Cargo.toml | 6 ++ .../datetime-parts/benches/split_temporal.rs | 58 ++++++++++ encodings/datetime-parts/src/compress.rs | 101 ++++++++++++++++-- .../datetime-parts/src/compute/compare.rs | 2 +- encodings/datetime-parts/src/compute/rules.rs | 2 +- encodings/datetime-parts/src/ops.rs | 18 ++-- encodings/datetime-parts/src/timestamp.rs | 75 ++++++++++--- 8 files changed, 226 insertions(+), 38 deletions(-) create mode 100644 encodings/datetime-parts/benches/split_temporal.rs diff --git a/Cargo.lock b/Cargo.lock index f7d3103c9d5..a4994ea8006 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10907,8 +10907,10 @@ dependencies = [ name = "vortex-datetime-parts" version = "0.1.0" dependencies = [ + "codspeed-divan-compat", "num-traits", "prost 0.14.4", + "rand 0.10.2", "rstest", "vortex-array", "vortex-buffer", diff --git a/encodings/datetime-parts/Cargo.toml b/encodings/datetime-parts/Cargo.toml index 52894fc7c07..3bbdde465dd 100644 --- a/encodings/datetime-parts/Cargo.toml +++ b/encodings/datetime-parts/Cargo.toml @@ -26,6 +26,12 @@ vortex-mask = { workspace = true } vortex-session = { workspace = true } [dev-dependencies] +divan = { workspace = true } +rand = { workspace = true } rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } vortex-error = { workspace = true } + +[[bench]] +name = "split_temporal" +harness = false diff --git a/encodings/datetime-parts/benches/split_temporal.rs b/encodings/datetime-parts/benches/split_temporal.rs new file mode 100644 index 00000000000..213192899fd --- /dev/null +++ b/encodings/datetime-parts/benches/split_temporal.rs @@ -0,0 +1,58 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use rand::RngExt; +use rand::SeedableRng as _; +use rand::rngs::StdRng; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::TemporalArray; +use vortex_array::extension::datetime::TimeUnit; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_datetime_parts::split_temporal; +use vortex_session::VortexSession; + +fn main() { + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + +const BENCH_ARGS: &[(usize, TimeUnit)] = &[ + (65_536, TimeUnit::Seconds), + (65_536, TimeUnit::Milliseconds), + (65_536, TimeUnit::Microseconds), + (65_536, TimeUnit::Nanoseconds), +]; + +#[divan::bench(args = BENCH_ARGS)] +fn split(bencher: Bencher, args: (usize, TimeUnit)) { + let (n, unit) = args; + let divisor: i64 = match unit { + TimeUnit::Seconds => 1, + TimeUnit::Milliseconds => 1_000, + TimeUnit::Microseconds => 1_000_000, + TimeUnit::Nanoseconds => 1_000_000_000, + TimeUnit::Days => unreachable!(), + }; + let mut rng = StdRng::seed_from_u64(0); + let timestamps = Buffer::from_iter((0..n).map(|_| { + rng.random_range(1_500_000_000i64..1_800_000_000) * divisor + rng.random_range(0..divisor) + })); + let array = TemporalArray::new_timestamp( + PrimitiveArray::new(timestamps, Validity::NonNullable).into_array(), + unit, + Some("UTC".into()), + ); + + bencher + .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) + .bench_values(|(array, mut ctx)| split_temporal(array, &mut ctx).unwrap()) +} diff --git a/encodings/datetime-parts/src/compress.rs b/encodings/datetime-parts/src/compress.rs index a3ab3e6f23b..da019edcfbc 100644 --- a/encodings/datetime-parts/src/compress.rs +++ b/encodings/datetime-parts/src/compress.rs @@ -1,18 +1,26 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::mem::MaybeUninit; + use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::TemporalArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::PType; +use vortex_array::extension::datetime::TimeUnit; use vortex_buffer::BufferMut; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use crate::timestamp; +/// All parts are stored as i32. seconds and subseconds always fit in i32. +/// For days we have a upper day limitation due to "jiff" validation so days +/// also fit in i32. pub struct TemporalParts { pub days: ArrayRef, pub seconds: ArrayRef, @@ -39,25 +47,98 @@ pub fn split_temporal(array: TemporalArray, ctx: &mut ExecutionCtx) -> VortexRes ))? .execute::(ctx)?; + let time_unit = array.temporal_metadata().time_unit(); + if matches!(time_unit, TimeUnit::Days) { + vortex_bail!("Cannot handle day-level data"); + } + let length = timestamps.len(); - let mut days = BufferMut::with_capacity(length); - let mut seconds = BufferMut::with_capacity(length); - let mut subseconds = BufferMut::with_capacity(length); - - for &ts in timestamps.as_slice::() { - let ts_parts = timestamp::split(ts, array.temporal_metadata().time_unit())?; - days.push(ts_parts.days); - seconds.push(ts_parts.seconds); - subseconds.push(ts_parts.subseconds); + + // If we don't [..length], compiler can't infer all 3 or 4 slices are the + // same length, and zip() in split_slice_* checks iterator boundaries. + let timestamps = ×tamps.as_slice::()[..length]; + + let mut days: BufferMut = BufferMut::with_capacity(timestamps.len()); + let mut seconds: BufferMut = BufferMut::with_capacity(timestamps.len()); + let days_ptr = &mut days.spare_capacity_mut()[..length]; + let seconds_ptr = &mut seconds.spare_capacity_mut()[..length]; + + if matches!(time_unit, TimeUnit::Seconds) { + split_slice_seconds(days_ptr, seconds_ptr, timestamps); + + // SAFETY: all items in [0; length) are filled in split_slice_seconds + unsafe { + days.set_len(length); + seconds.set_len(length); + } + + return Ok(TemporalParts { + days: PrimitiveArray::new(days.freeze(), temporal_values.validity()?).into_array(), + seconds: seconds.into_array(), + subseconds: ConstantArray::new(0, length).into_array(), + }); + } + + let mut subseconds: BufferMut = BufferMut::with_capacity(timestamps.len()); + let subseconds_ptr = &mut subseconds.spare_capacity_mut()[..length]; + + match time_unit { + TimeUnit::Nanoseconds => { + split_slice::<1_000_000_000>(days_ptr, seconds_ptr, subseconds_ptr, timestamps) + } + TimeUnit::Microseconds => { + split_slice::<1_000_000>(days_ptr, seconds_ptr, subseconds_ptr, timestamps) + } + TimeUnit::Milliseconds => { + split_slice::<1_000>(days_ptr, seconds_ptr, subseconds_ptr, timestamps) + } + _ => unreachable!("Handled before"), + }; + + // SAFETY: all items in [0; length) are filled in split_slice + unsafe { + days.set_len(length); + seconds.set_len(length); + subseconds.set_len(length); } Ok(TemporalParts { - days: PrimitiveArray::new(days, temporal_values.validity()?).into_array(), + days: PrimitiveArray::new(days.freeze(), temporal_values.validity()?).into_array(), seconds: seconds.into_array(), subseconds: subseconds.into_array(), }) } +#[inline] +fn split_slice( + days: &mut [MaybeUninit], + seconds: &mut [MaybeUninit], + subseconds: &mut [MaybeUninit], + timestamps: &[i64], +) { + for (((day, second), subseconds), ts) in + days.iter_mut().zip(seconds).zip(subseconds).zip(timestamps) + { + let parts = timestamp::split_with_divisor::(*ts); + day.write(parts.days); + second.write(parts.seconds); + subseconds.write(parts.subseconds); + } +} + +#[inline] +fn split_slice_seconds( + days: &mut [MaybeUninit], + seconds: &mut [MaybeUninit], + timestamps: &[i64], +) { + for ((day, second), ts) in days.iter_mut().zip(seconds).zip(timestamps) { + let parts = timestamp::split_with_divisor::<1>(*ts); + day.write(parts.days); + second.write(parts.seconds); + } +} + #[cfg(test)] mod tests { use rstest::rstest; diff --git a/encodings/datetime-parts/src/compute/compare.rs b/encodings/datetime-parts/src/compute/compare.rs index 75b6c066773..a4336e0d0fd 100644 --- a/encodings/datetime-parts/src/compute/compare.rs +++ b/encodings/datetime-parts/src/compute/compare.rs @@ -176,7 +176,7 @@ fn compare_gt( fn compare_dtp( lhs: &ArrayRef, - rhs: i64, + rhs: i32, operator: CompareOperator, nullability: Nullability, ) -> VortexResult { diff --git a/encodings/datetime-parts/src/compute/rules.rs b/encodings/datetime-parts/src/compute/rules.rs index aac6d101101..12546c23fd8 100644 --- a/encodings/datetime-parts/src/compute/rules.rs +++ b/encodings/datetime-parts/src/compute/rules.rs @@ -143,7 +143,7 @@ impl ArrayParentReduceRule for DTPComparisonPushDownRule { /// Try to extract the days value from a constant timestamp. /// Returns None if the constant is not a timestamp or has non-zero seconds/subseconds. -fn try_extract_days_constant(array: &ArrayRef) -> Option { +fn try_extract_days_constant(array: &ArrayRef) -> Option { let constant = array.as_constant()?; // Extract the timestamp value diff --git a/encodings/datetime-parts/src/ops.rs b/encodings/datetime-parts/src/ops.rs index 57a1f60173f..d99e55c7542 100644 --- a/encodings/datetime-parts/src/ops.rs +++ b/encodings/datetime-parts/src/ops.rs @@ -37,24 +37,24 @@ impl OperationsVTable for DateTimeParts { return Ok(Scalar::null(DType::Extension(ext))); } - let days: i64 = array + let days: i32 = array .days() .execute_scalar(index, ctx)? .as_primitive() - .as_::() - .vortex_expect("days fits in i64"); - let seconds: i64 = array + .as_::() + .vortex_expect("days fits in i32"); + let seconds: i32 = array .seconds() .execute_scalar(index, ctx)? .as_primitive() - .as_::() - .vortex_expect("seconds fits in i64"); - let subseconds: i64 = array + .as_::() + .vortex_expect("seconds fits in i32"); + let subseconds: i32 = array .subseconds() .execute_scalar(index, ctx)? .as_primitive() - .as_::() - .vortex_expect("subseconds fits in i64"); + .as_::() + .vortex_expect("subseconds fits in i32"); let ts = timestamp::combine( TimestampParts { diff --git a/encodings/datetime-parts/src/timestamp.rs b/encodings/datetime-parts/src/timestamp.rs index 6cafeb530d5..dee08cfeb2d 100644 --- a/encodings/datetime-parts/src/timestamp.rs +++ b/encodings/datetime-parts/src/timestamp.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use num_traits::AsPrimitive; use vortex_array::extension::datetime::TimeUnit; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -11,13 +12,13 @@ pub const SECONDS_PER_DAY: i64 = 86_400; // 24 * 60 * 60 /// Parts of a Unix timestamp (time since 1970-01-01 00:00:00 UTC). /// /// Broken down into: -/// * Days since epoch +/// * Days since epoch (max range is "jiff"'s range) /// * Seconds within day (0-86399) /// * Subseconds (range depends on `TimeUnit`) pub struct TimestampParts { - pub days: i64, - pub seconds: i64, - pub subseconds: i64, + pub days: i32, + pub seconds: i32, + pub subseconds: i32, } /// Splits a Unix timestamp into its component parts. @@ -31,22 +32,29 @@ pub struct TimestampParts { /// /// Returns an error if `time_unit` is days, which cannot be split. pub fn split(timestamp: i64, time_unit: TimeUnit) -> VortexResult { - let divisor = match time_unit { - TimeUnit::Nanoseconds => 1_000_000_000, - TimeUnit::Microseconds => 1_000_000, - TimeUnit::Milliseconds => 1_000, - TimeUnit::Seconds => 1, + Ok(match time_unit { + TimeUnit::Nanoseconds => split_with_divisor::<1_000_000_000>(timestamp), + TimeUnit::Microseconds => split_with_divisor::<1_000_000>(timestamp), + TimeUnit::Milliseconds => split_with_divisor::<1_000>(timestamp), + TimeUnit::Seconds => split_with_divisor::<1>(timestamp), TimeUnit::Days => vortex_bail!("Cannot handle day-level data"), - }; - - let ticks_per_day = SECONDS_PER_DAY * divisor; - Ok(TimestampParts { - days: timestamp / ticks_per_day, - seconds: (timestamp % ticks_per_day) / divisor, - subseconds: (timestamp % ticks_per_day) % divisor, }) } +/// Split a Unix timestamp into parts using a constant DIVISOR. +#[inline] +pub(crate) fn split_with_divisor(timestamp: i64) -> TimestampParts { + let days = timestamp / (SECONDS_PER_DAY * DIVISOR); + let total_seconds = timestamp / DIVISOR; + let seconds = total_seconds - days * SECONDS_PER_DAY; + let subseconds = timestamp - total_seconds * DIVISOR; + TimestampParts { + days: days.as_(), + seconds: seconds.as_(), + subseconds: subseconds.as_(), + } +} + /// Combines timestamp parts back into a Unix timestamp. /// /// # Arguments @@ -66,13 +74,46 @@ pub fn combine(ts_parts: TimestampParts, time_unit: TimeUnit) -> i64 { TimeUnit::Days => vortex_panic!("Cannot handle day-level data"), }; - ts_parts.days * SECONDS_PER_DAY * divisor + ts_parts.seconds * divisor + ts_parts.subseconds + ts_parts.days as i64 * SECONDS_PER_DAY * divisor + + ts_parts.seconds as i64 * divisor + + ts_parts.subseconds as i64 } #[cfg(test)] mod tests { + use rstest::rstest; + use super::*; + #[rstest] + #[case(TimeUnit::Seconds, 1)] + #[case(TimeUnit::Milliseconds, 1_000)] + #[case(TimeUnit::Microseconds, 1_000_000)] + #[case(TimeUnit::Nanoseconds, 1_000_000_000)] + fn test_split_divmod(#[case] unit: TimeUnit, #[case] divisor: i64) -> VortexResult<()> { + let ticks_per_day = SECONDS_PER_DAY * divisor; + for ts in [ + 0, + 1, + -1, + divisor - 1, + 3723 * divisor + divisor / 2, + -(3723 * divisor + divisor / 2), + ticks_per_day - 1, + ticks_per_day, + -ticks_per_day, + i64::MAX, + i64::MIN, + ] { + let parts = split(ts, unit)?; + assert_eq!(parts.days, (ts / ticks_per_day) as i32); + assert_eq!(parts.seconds, ((ts % ticks_per_day) / divisor) as i32); + assert_eq!(parts.subseconds, ((ts % ticks_per_day) % divisor) as i32); + assert_eq!(combine(parts, unit), ts); + } + Ok(()) + } + #[test] fn test_split_seconds() { // 1970-01-02 01:02:03 UTC From f2d36c138e88300a9feb9565debfd0c8c3e99a9b Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Wed, 9 Sep 2026 16:32:37 +0100 Subject: [PATCH 2/2] chunks of 4 Signed-off-by: Mikhail Kot --- encodings/datetime-parts/src/compress.rs | 49 ++++++++++++++++++----- encodings/datetime-parts/src/timestamp.rs | 12 +++--- 2 files changed, 46 insertions(+), 15 deletions(-) diff --git a/encodings/datetime-parts/src/compress.rs b/encodings/datetime-parts/src/compress.rs index da019edcfbc..1198257ebae 100644 --- a/encodings/datetime-parts/src/compress.rs +++ b/encodings/datetime-parts/src/compress.rs @@ -32,6 +32,11 @@ pub struct TemporalParts { /// Splitting the components by granularity creates more small values, which enables better /// cascading compression. pub fn split_temporal(array: TemporalArray, ctx: &mut ExecutionCtx) -> VortexResult { + let time_unit = array.temporal_metadata().time_unit(); + if matches!(time_unit, TimeUnit::Days) { + vortex_bail!("Cannot handle day-level data"); + } + let temporal_values = array .temporal_values() .clone() @@ -47,11 +52,6 @@ pub fn split_temporal(array: TemporalArray, ctx: &mut ExecutionCtx) -> VortexRes ))? .execute::(ctx)?; - let time_unit = array.temporal_metadata().time_unit(); - if matches!(time_unit, TimeUnit::Days) { - vortex_bail!("Cannot handle day-level data"); - } - let length = timestamps.len(); // If we don't [..length], compiler can't infer all 3 or 4 slices are the @@ -116,13 +116,42 @@ fn split_slice( subseconds: &mut [MaybeUninit], timestamps: &[i64], ) { - for (((day, second), subseconds), ts) in - days.iter_mut().zip(seconds).zip(subseconds).zip(timestamps) - { - let parts = timestamp::split_with_divisor::(*ts); + // Computing chunks of 4 elements lets LLVM optimize stores into + // a 16-byte vector store per chunk. + let length = timestamps.len(); + let (timestamps, timestamps_rem) = timestamps.as_chunks::<4>(); + let (days, days_rem) = days[..length].as_chunks_mut::<4>(); + let (seconds, seconds_rem) = seconds[..length].as_chunks_mut::<4>(); + let (subseconds, subseconds_rem) = subseconds[..length].as_chunks_mut::<4>(); + + let chunks = timestamps.iter().zip(days).zip(seconds).zip(subseconds); + for (((timestamp, day), second), subsecond) in chunks { + let mut day_buf = [0i32; 4]; + let mut second_buf = [0i32; 4]; + let mut subsecond_buf = [0i32; 4]; + for k in 0..4 { + let parts = timestamp::split_with_divisor::(timestamp[k]); + day_buf[k] = parts.days; + second_buf[k] = parts.seconds; + subsecond_buf[k] = parts.subseconds; + } + for k in 0..4 { + day[k].write(day_buf[k]); + second[k].write(second_buf[k]); + subsecond[k].write(subsecond_buf[k]); + } + } + + let remainder = timestamps_rem + .iter() + .zip(days_rem) + .zip(seconds_rem) + .zip(subseconds_rem); + for (((&ts, day), second), subsecond) in remainder { + let parts = timestamp::split_with_divisor::(ts); day.write(parts.days); second.write(parts.seconds); - subseconds.write(parts.subseconds); + subsecond.write(parts.subseconds); } } diff --git a/encodings/datetime-parts/src/timestamp.rs b/encodings/datetime-parts/src/timestamp.rs index dee08cfeb2d..2aa10d4fac6 100644 --- a/encodings/datetime-parts/src/timestamp.rs +++ b/encodings/datetime-parts/src/timestamp.rs @@ -102,13 +102,15 @@ mod tests { ticks_per_day - 1, ticks_per_day, -ticks_per_day, - i64::MAX, - i64::MIN, ] { let parts = split(ts, unit)?; - assert_eq!(parts.days, (ts / ticks_per_day) as i32); - assert_eq!(parts.seconds, ((ts % ticks_per_day) / divisor) as i32); - assert_eq!(parts.subseconds, ((ts % ticks_per_day) % divisor) as i32); + + let days: i32 = (ts / ticks_per_day).as_(); + let seconds: i32 = ((ts % ticks_per_day) / divisor).as_(); + let subseconds: i32 = ((ts % ticks_per_day) % divisor).as_(); + assert_eq!(parts.days, days); + assert_eq!(parts.seconds, seconds); + assert_eq!(parts.subseconds, subseconds); assert_eq!(combine(parts, unit), ts); } Ok(())