Preallocate buffers for parsing timestamp into DateTimeParts

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

@@ -9,8 +9,10 @@ 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;
 pub struct TemporalParts {
@@ -39,25 +41,41 @@ 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);
-    }
+    let (days, seconds, subseconds) = match array.temporal_metadata().time_unit() {
+        TimeUnit::Nanoseconds => split_slice::<1_000_000_000>(timestamps.as_slice::<i64>()),
+        TimeUnit::Microseconds => split_slice::<1_000_000>(timestamps.as_slice::<i64>()),
+        TimeUnit::Milliseconds => split_slice::<1_000>(timestamps.as_slice::<i64>()),
+        TimeUnit::Seconds => split_slice::<1>(timestamps.as_slice::<i64>()),
+        TimeUnit::Days => vortex_bail!("Cannot handle day-level data"),
+    };
 
     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(),
     })
 }
 
+fn split_slice<const DIVISOR: i64>(
+    timestamps: &[i64],
+) -> (BufferMut<i64>, BufferMut<i64>, BufferMut<i64>) {
+    let mut days = BufferMut::zeroed(timestamps.len());
+    let mut seconds = BufferMut::zeroed(timestamps.len());
+    let mut subseconds = BufferMut::zeroed(timestamps.len());
+    let (days_out, seconds_out, subseconds_out) = (
+        days.as_mut_slice(),
+        seconds.as_mut_slice(),
+        subseconds.as_mut_slice(),
+    );
+    for (i, &ts) in timestamps.iter().enumerate() {
+        let ts_parts = timestamp::split_with_divisor::<DIVISOR>(ts);
+        days_out[i] = ts_parts.days;
+        seconds_out[i] = ts_parts.seconds;
+        subseconds_out[i] = ts_parts.subseconds;
+    }
+    (days, seconds, subseconds)
+}
+
 #[cfg(test)]
 mod tests {
     use rstest::rstest;

encodings/datetime-parts/src/compress.rs

@@ -31,22 +31,26 @@ 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
+pub(crate) fn split_with_divisor<const DIVISOR: i64>(timestamp: i64) -> TimestampParts {
+    let days = timestamp / (SECONDS_PER_DAY * DIVISOR);
+    let total_seconds = timestamp / DIVISOR;
+    TimestampParts {
+        days,
+        seconds: total_seconds - days * SECONDS_PER_DAY,
+        subseconds: timestamp - total_seconds * DIVISOR,
+    }
+}
+
 /// Combines timestamp parts back into a Unix timestamp.
 ///
 /// # Arguments
@@ -71,8 +75,39 @@ pub fn combine(ts_parts: TimestampParts, time_unit: TimeUnit) -> 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);
+            assert_eq!(parts.seconds, (ts % ticks_per_day) / divisor);
+            assert_eq!(parts.subseconds, (ts % ticks_per_day) % divisor);
+            assert_eq!(combine(parts, unit), ts);
+        }
+        Ok(())
+    }
+
     #[test]
     fn test_split_seconds() {
         // 1970-01-02 01:02:03 UTC