Change DateTimeParts to i32

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

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<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/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,
@@ -24,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<TemporalParts> {
+    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()
@@ -40,24 +53,121 @@ pub fn split_temporal(array: TemporalArray, ctx: &mut ExecutionCtx) -> VortexRes
         .execute::<PrimitiveArray>(ctx)?;
 
     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],
+) {
+    // 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::<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::<DIVISOR>(ts);
+        day.write(parts.days);
+        second.write(parts.seconds);
+        subsecond.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/compute/compare.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/rules.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/ops.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 {