Change DateTimeParts to i32

Signed-off-by: Mikhail Kot <mikhail@spiraldb.com>

Author: myrrc

Cargo.lock

@@ -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

encodings/datetime-parts/Cargo.toml

@@ -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<VortexSession> = 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())
+}

encodings/datetime-parts/benches/split_temporal.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::<PrimitiveArray>(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::<i64>() {
-        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 = &timestamps.as_slice::<i64>()[..length];
+
+    let mut days: BufferMut<i32> = BufferMut::with_capacity(timestamps.len());
+    let mut seconds: BufferMut<i32> = 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<i32> = 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<const DIVISOR: i64>(
+    days: &mut [MaybeUninit<i32>],
+    seconds: &mut [MaybeUninit<i32>],
+    subseconds: &mut [MaybeUninit<i32>],
+    timestamps: &[i64],
+) {
+    for (((day, second), subseconds), ts) in
+        days.iter_mut().zip(seconds).zip(subseconds).zip(timestamps)
+    {
+        let parts = timestamp::split_with_divisor::<DIVISOR>(*ts);
+        day.write(parts.days);
+        second.write(parts.seconds);
+        subseconds.write(parts.subseconds);
+    }
+}
+
+#[inline]
+fn split_slice_seconds(
+    days: &mut [MaybeUninit<i32>],
+    seconds: &mut [MaybeUninit<i32>],
+    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;

encodings/datetime-parts/src/compress.rs

@@ -176,7 +176,7 @@ fn compare_gt(
 
 fn compare_dtp(
     lhs: &ArrayRef,
-    rhs: i64,
+    rhs: i32,
     operator: CompareOperator,
     nullability: Nullability,
 ) -> VortexResult<ArrayRef> {

encodings/datetime-parts/src/compute/compare.rs

@@ -143,7 +143,7 @@ impl ArrayParentReduceRule<DateTimeParts> 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<i64> {
+fn try_extract_days_constant(array: &ArrayRef) -> Option<i32> {
     let constant = array.as_constant()?;
 
     // Extract the timestamp value

encodings/datetime-parts/src/compute/rules.rs

@@ -37,24 +37,24 @@ impl OperationsVTable<DateTimeParts> 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_::<i64>()
-            .vortex_expect("days fits in i64");
-        let seconds: i64 = array
+            .as_::<i32>()
+            .vortex_expect("days fits in i32");
+        let seconds: i32 = array
             .seconds()
             .execute_scalar(index, ctx)?
             .as_primitive()
-            .as_::<i64>()
-            .vortex_expect("seconds fits in i64");
-        let subseconds: i64 = array
+            .as_::<i32>()
+            .vortex_expect("seconds fits in i32");
+        let subseconds: i32 = array
             .subseconds()
             .execute_scalar(index, ctx)?
             .as_primitive()
-            .as_::<i64>()
-            .vortex_expect("subseconds fits in i64");
+            .as_::<i32>()
+            .vortex_expect("subseconds fits in i32");
 
         let ts = timestamp::combine(
             TimestampParts {

encodings/datetime-parts/src/ops.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<TimestampParts> {
-    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<const DIVISOR: i64>(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