From 1647def7af110f9c0def9a420275288b2a8a53ce Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Thu, 25 Jun 2026 12:02:25 +0100 Subject: [PATCH] fix Signed-off-by: Mikhail Kot --- Cargo.lock | 1 + Cargo.toml | 1 + encodings/datetime-parts/Cargo.toml | 1 + encodings/datetime-parts/src/array.rs | 11 +-- encodings/datetime-parts/src/canonical.rs | 114 +++++++++++++++------- 5 files changed, 85 insertions(+), 43 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 09ba6496805..7aa816ef9e3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9858,6 +9858,7 @@ dependencies = [ "num-traits", "prost 0.14.4", "rstest", + "seq-macro", "vortex-array", "vortex-buffer", "vortex-error", diff --git a/Cargo.toml b/Cargo.toml index acab68656aa..7f569518663 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -235,6 +235,7 @@ roaring = "0.11.0" rstest = "0.26.1" rstest_reuse = "0.7.0" rustc-hash = "2.1.1" +seq-macro = "0.3.6" serde = "1.0.220" serde_json = "1.0.138" serde_test = "1.0.176" diff --git a/encodings/datetime-parts/Cargo.toml b/encodings/datetime-parts/Cargo.toml index 52894fc7c07..7ff9f127385 100644 --- a/encodings/datetime-parts/Cargo.toml +++ b/encodings/datetime-parts/Cargo.toml @@ -19,6 +19,7 @@ workspace = true [dependencies] num-traits = { workspace = true } prost = { workspace = true } +seq-macro = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } diff --git a/encodings/datetime-parts/src/array.rs b/encodings/datetime-parts/src/array.rs index 42569da9728..539f94b1517 100644 --- a/encodings/datetime-parts/src/array.rs +++ b/encodings/datetime-parts/src/array.rs @@ -7,7 +7,7 @@ use std::fmt::Formatter; use std::hash::Hasher; use prost::Message; -use vortex_array::AnyCanonical; +use vortex_array::AnyColumnar; use vortex_array::Array; use vortex_array::ArrayEq; use vortex_array::ArrayHash; @@ -186,17 +186,16 @@ impl VTable for DateTimeParts { DateTimePartsSlots::NAMES[idx].to_string() } - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { let array = require_child!(array, array.days(), DateTimePartsSlots::DAYS => Primitive); let array = - require_child!(array, array.seconds(), DateTimePartsSlots::SECONDS => AnyCanonical); - let array = require_child!(array, array.subseconds(), DateTimePartsSlots::SUBSECONDS => AnyCanonical); - + require_child!(array, array.seconds(), DateTimePartsSlots::SECONDS => AnyColumnar); + let array = require_child!(array, array.subseconds(), DateTimePartsSlots::SUBSECONDS => AnyColumnar); let dtype = array.dtype().clone(); let parts = array.into_parts(); Ok(ExecutionResult::done( - decode_to_temporal(parts, &dtype, ctx)?.into_array(), + decode_to_temporal(parts, &dtype)?.into_array(), )) } diff --git a/encodings/datetime-parts/src/canonical.rs b/encodings/datetime-parts/src/canonical.rs index b08b76b148c..a4cf2dfe2c3 100644 --- a/encodings/datetime-parts/src/canonical.rs +++ b/encodings/datetime-parts/src/canonical.rs @@ -2,7 +2,10 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use num_traits::AsPrimitive; -use vortex_array::ExecutionCtx; +use vortex_array::AnyColumnar; +use vortex_array::CanonicalView; +use vortex_array::ColumnarView::Canonical; +use vortex_array::ColumnarView::Constant; use vortex_array::IntoArray; use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; @@ -20,11 +23,7 @@ use vortex_error::vortex_panic; use crate::array::DateTimePartsParts; /// Decode [`DateTimePartsParts`] into a [`TemporalArray`]. -pub fn decode_to_temporal( - parts: DateTimePartsParts, - dtype: &DType, - ctx: &mut ExecutionCtx, -) -> VortexResult { +pub fn decode_to_temporal(parts: DateTimePartsParts, dtype: &DType) -> VortexResult { let DType::Extension(ext) = dtype else { vortex_panic!(Compute: "expected dtype to be DType::Extension variant") }; @@ -45,45 +44,56 @@ pub fn decode_to_temporal( let days = parts.days.as_::(); let validity = days.validity()?; - let mut values: BufferMut = match_each_integer_ptype!(days.ptype(), |D| { - BufferMut::from_iter(days.as_slice::().iter().map(|d| { - let d: i64 = d.as_(); - d * 86_400 * divisor - })) - }); - - // Seconds/subseconds may be Constant — handle the fast path. - if let Some(seconds) = parts.seconds.as_constant() { + let seconds = parts.seconds.as_::(); + let constant_seconds = if let Constant(seconds) = seconds { let seconds = seconds + .scalar() + .value() + .vortex_expect("no value") .as_primitive() - .as_::() + .as_i64() .vortex_expect("non-nullable"); - let seconds = seconds * divisor; - for v in values.iter_mut() { - *v += seconds; - } + seconds * divisor } else { - let seconds_buf = parts.seconds.execute::(ctx)?; - match_each_integer_ptype!(seconds_buf.ptype(), |S| { - for (v, second) in values.iter_mut().zip(seconds_buf.as_slice::()) { + 0 + }; + + let subseconds = parts.subseconds.as_::(); + let constant_subseconds = if let Constant(subseconds) = subseconds { + subseconds + .scalar() + .value() + .vortex_expect("no value") + .as_primitive() + .as_i64() + .vortex_expect("non-nullable") + } else { + 0 + }; + + let mut values = decode_days( + &days, + 86_400i64 * divisor, + constant_seconds + constant_subseconds, + ); + + if let Canonical(seconds) = seconds { + let CanonicalView::Primitive(seconds) = seconds else { + vortex_panic!("not a primitive"); + }; + match_each_integer_ptype!(seconds.ptype(), |S| { + for (v, second) in values.iter_mut().zip(seconds.as_slice::()) { let second: i64 = second.as_(); *v += second * divisor; } }); } - - if let Some(subseconds) = parts.subseconds.as_constant() { - let subseconds = subseconds - .as_primitive() - .as_::() - .vortex_expect("non-nullable"); - for v in values.iter_mut() { - *v += subseconds; - } - } else { - let subseconds_buf = parts.subseconds.execute::(ctx)?; - match_each_integer_ptype!(subseconds_buf.ptype(), |S| { - for (v, subsecond) in values.iter_mut().zip(subseconds_buf.as_slice::()) { + if let Canonical(subseconds) = subseconds { + let CanonicalView::Primitive(subseconds) = subseconds else { + vortex_panic!("not a primitive"); + }; + match_each_integer_ptype!(subseconds.ptype(), |S| { + for (v, subsecond) in values.iter_mut().zip(subseconds.as_slice::()) { let subsecond: i64 = subsecond.as_(); *v += subsecond; } @@ -97,6 +107,36 @@ pub fn decode_to_temporal( )) } +// For constant seconds and subseconds, compute day * day_to_unit + const_offset +fn decode_days(days: &P, day_to_unit: i64, offset: i64) -> BufferMut { + /// If "days" are u16 or u32, LLVM doesn't auto-vectorize the code due to + /// widening to i64. If we process the code in fixed-size chunks and unroll + /// the chunks with seq_macro, vectorization happens. + const CHUNK: usize = 64; + let n = days.len(); + let mut values = BufferMut::::with_capacity(n); + match_each_integer_ptype!(days.ptype(), |D| { + let src = days.as_slice::(); + let (src_chunks, src_rem) = src.as_chunks::(); + let (dst_chunks, _) = values.spare_capacity_mut()[..n].as_chunks_mut::(); + for (s_chunk, d_chunk) in src_chunks.iter().zip(dst_chunks.iter_mut()) { + seq_macro::seq!(I in 0..64 { + let day: i64 = s_chunk[I].as_(); + d_chunk[I].write(day * day_to_unit + offset); + }); + } + let tail_start = src_chunks.len() * CHUNK; + let dst_tail = &mut values.spare_capacity_mut()[tail_start..n]; + for (s, d) in src_rem.iter().zip(dst_tail.iter_mut()) { + let day: i64 = s.as_(); + d.write(day * day_to_unit + offset); + } + }); + // SAFETY: every element in 0..n was written above. + unsafe { values.set_len(n) }; + values +} + #[cfg(test)] mod test { use std::sync::LazyLock; @@ -164,7 +204,7 @@ mod test { subseconds: date_times.subseconds().clone(), }; - let primitive_values = decode_to_temporal(parts, &dtype, &mut ctx)? + let primitive_values = decode_to_temporal(parts, &dtype)? .temporal_values() .clone() .execute::(&mut ctx)?;