From 1f1227692c00f1bfc6dafcd74b8a9353a3850f58 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 1 Oct 2026 12:04:32 -0400 Subject: [PATCH 1/7] feat(expr): add timestamp timezone replacement Signed-off-by: Daniel King --- vortex-array/proto/expr.proto | 6 + vortex-array/src/expr/exprs.rs | 26 ++ vortex-array/src/expr/mod.rs | 1 + .../src/proto/generated/vortex.expr.rs | 8 + vortex-array/src/scalar_fn/fns/mod.rs | 1 + .../src/scalar_fn/fns/replace_time_zone.rs | 381 ++++++++++++++++++ vortex-array/src/scalar_fn/session.rs | 2 + vortex-python/python/vortex/_lib/expr.pyi | 8 + vortex-python/python/vortex/expr.py | 2 + vortex-python/src/expr/mod.rs | 40 ++ vortex-python/test/test_expr.py | 10 + 11 files changed, 485 insertions(+) create mode 100644 vortex-array/src/scalar_fn/fns/replace_time_zone.rs diff --git a/vortex-array/proto/expr.proto b/vortex-array/proto/expr.proto index a53a75ad50c..508f954ccfb 100644 --- a/vortex-array/proto/expr.proto +++ b/vortex-array/proto/expr.proto @@ -129,3 +129,9 @@ message SelectOpts { message CaseWhenOpts { uint32 num_children = 1; } + +// Options for `vortex.replace_time_zone`. Ambiguity policy is an expression child. +message ReplaceTimeZoneOpts { + optional string time_zone = 1; + bool null_on_non_existent = 2; +} diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index 1b2a32df4c1..46e5b6c03a2 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -51,6 +51,8 @@ use crate::scalar_fn::fns::operators::CompareOperator; use crate::scalar_fn::fns::operators::Operator; use crate::scalar_fn::fns::pack::Pack; use crate::scalar_fn::fns::pack::PackOptions; +use crate::scalar_fn::fns::replace_time_zone::ReplaceTimeZone; +use crate::scalar_fn::fns::replace_time_zone::ReplaceTimeZoneOptions; use crate::scalar_fn::fns::select::FieldSelection; use crate::scalar_fn::fns::select::Select; use crate::scalar_fn::fns::variant_get::VariantGet; @@ -1265,6 +1267,29 @@ pub fn bound_list_sum_opts( .vortex_expect("list-sum expressions require a numeric list child") } +/// Reinterpret timestamp wall times in `options.time_zone`, preserving the time unit. +/// +/// `ambiguous` is a UTF-8 expression containing `raise`, `earliest`, `latest`, or `null`. +/// Removing the timezone preserves wall time in a timezone-naive timestamp. +pub fn replace_time_zone( + input: Expression, + ambiguous: Expression, + options: ReplaceTimeZoneOptions, +) -> Expression { + ReplaceTimeZone.new_expr(options, [input, ambiguous]) +} + +/// Create a bound timezone replacement expression. +pub fn bound_replace_time_zone( + input: BoundExpression, + ambiguous: BoundExpression, + options: ReplaceTimeZoneOptions, +) -> BoundExpression { + ReplaceTimeZone + .try_new_bound_expr(options, [input, ambiguous]) + .vortex_expect("timezone replacement requires a timestamp and a UTF-8 ambiguity policy") +} + /// Constructors for expressions whose children have already been bound and type-checked. /// /// These mirror the constructors in [`crate::expr`] and panic when the supplied children do not @@ -1313,6 +1338,7 @@ pub mod bound { pub use super::bound_or as or; pub use super::bound_or_collect as or_collect; pub use super::bound_pack as pack; + pub use super::bound_replace_time_zone as replace_time_zone; pub use super::bound_root as root; pub use super::bound_select as select; pub use super::bound_select_exclude as select_exclude; diff --git a/vortex-array/src/expr/mod.rs b/vortex-array/src/expr/mod.rs index 31c99f3f110..1f78d8d0911 100644 --- a/vortex-array/src/expr/mod.rs +++ b/vortex-array/src/expr/mod.rs @@ -104,6 +104,7 @@ pub use exprs::not_like; pub use exprs::or; pub use exprs::or_collect; pub use exprs::pack; +pub use exprs::replace_time_zone; pub use exprs::root; pub use exprs::select; pub use exprs::select_exclude; diff --git a/vortex-array/src/proto/generated/vortex.expr.rs b/vortex-array/src/proto/generated/vortex.expr.rs index b08822055a7..ceaa08c0324 100644 --- a/vortex-array/src/proto/generated/vortex.expr.rs +++ b/vortex-array/src/proto/generated/vortex.expr.rs @@ -217,3 +217,11 @@ pub struct CaseWhenOpts { #[prost(uint32, tag = "1")] pub num_children: u32, } +/// Options for `vortex.replace_time_zone`. Ambiguity policy is an expression child. +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct ReplaceTimeZoneOpts { + #[prost(string, optional, tag = "1")] + pub time_zone: ::core::option::Option<::prost::alloc::string::String>, + #[prost(bool, tag = "2")] + pub null_on_non_existent: bool, +} diff --git a/vortex-array/src/scalar_fn/fns/mod.rs b/vortex-array/src/scalar_fn/fns/mod.rs index ece71396faf..3cb30209dd9 100644 --- a/vortex-array/src/scalar_fn/fns/mod.rs +++ b/vortex-array/src/scalar_fn/fns/mod.rs @@ -23,6 +23,7 @@ pub mod merge; pub mod not; pub mod operators; pub mod pack; +pub mod replace_time_zone; pub mod select; pub mod stat; pub mod variant_get; diff --git a/vortex-array/src/scalar_fn/fns/replace_time_zone.rs b/vortex-array/src/scalar_fn/fns/replace_time_zone.rs new file mode 100644 index 00000000000..762fb7ca594 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/replace_time_zone.rs @@ -0,0 +1,381 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Reinterpret timestamp wall times in a different timezone. + +use std::fmt; +use std::sync::Arc; + +use jiff::Timestamp as JiffTimestamp; +use jiff::tz::AmbiguousOffset; +use jiff::tz::TimeZone; +use prost::Message; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::ConstantArray; +use crate::arrays::ExtensionArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::VarBinViewArray; +use crate::arrays::extension::ExtensionArrayExt; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::extension::datetime::TimeUnit; +use crate::extension::datetime::Timestamp; +use crate::extension::datetime::TimestampOptions; +use crate::proto::expr as pb; +use crate::scalar::Scalar; +use crate::scalar_fn::Arity; +use crate::scalar_fn::ChildName; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; + +/// Timezone replacement options. The ambiguity policy is a UTF-8 expression child. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct ReplaceTimeZoneOptions { + /// Destination timezone, or `None` to remove the timezone while preserving wall time. + pub time_zone: Option>, + /// Return null for nonexistent wall times instead of raising an error. + pub null_on_non_existent: bool, +} + +impl fmt::Display for ReplaceTimeZoneOptions { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "tz={:?}, null_on_non_existent={}", + self.time_zone, self.null_on_non_existent + ) + } +} + +/// Replace a timestamp's timezone, preserving its local wall time. +/// +/// The second child selects `raise`, `earliest`, `latest`, or `null` for ambiguous times. +#[derive(Clone)] +pub struct ReplaceTimeZone; + +impl ScalarFnVTable for ReplaceTimeZone { + type Options = ReplaceTimeZoneOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.replace_time_zone"); + *ID + } + + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + Ok(Some( + pb::ReplaceTimeZoneOpts { + time_zone: options.time_zone.as_ref().map(ToString::to_string), + null_on_non_existent: options.null_on_non_existent, + } + .encode_to_vec(), + )) + } + + fn deserialize( + &self, + metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + let options = pb::ReplaceTimeZoneOpts::decode(metadata)?; + Ok(ReplaceTimeZoneOptions { + time_zone: options.time_zone.map(Arc::from), + null_on_non_existent: options.null_on_non_existent, + }) + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(2) + } + + fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { + match child_idx { + 0 => ChildName::from("input"), + 1 => ChildName::from("ambiguous"), + _ => unreachable!("Invalid child index {child_idx} for replace_time_zone()"), + } + } + + fn return_dtype( + &self, + options: &Self::Options, + arg_dtypes: &[DType], + ) -> VortexResult { + let input = timestamp_options(&arg_dtypes[0])?; + if input.unit == TimeUnit::Days { + vortex_bail!("Timestamp cannot use day units"); + } + if !matches!(arg_dtypes[1], DType::Utf8(_) | DType::Null) { + vortex_bail!( + "replace_time_zone() requires a UTF-8 ambiguity policy, got {}", + arg_dtypes[1] + ); + } + resolve_zone(options.time_zone.as_deref())?; + // A row's ambiguity policy can request null even for non-nullable children. + Ok(DType::Extension( + Timestamp::new_with_tz( + input.unit, + options.time_zone.clone(), + Nullability::Nullable, + ) + .erased(), + )) + } + + fn execute( + &self, + options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let input = args.get(0)?; + let ambiguous = args.get(1)?; + let dtype = self.return_dtype( + options, + &[input.dtype().clone(), ambiguous.dtype().clone()], + )?; + let metadata = timestamp_options(input.dtype())?.clone(); + let source = resolve_zone(metadata.tz.as_deref())?; + let target = resolve_zone(options.time_zone.as_deref())?; + let convert = |value, policy: &str| { + replace( + value, + metadata.unit, + &source, + &target, + policy, + options.null_on_non_existent, + ) + }; + + if let (Some(value), Some(policy)) = (input.as_constant(), ambiguous.as_constant()) { + let result = if value.is_null() || policy.is_null() { + None + } else { + let storage = value.as_extension().to_storage_scalar(); + let policy = policy + .as_utf8() + .value() + .ok_or_else(|| vortex_err!("Missing ambiguity policy"))?; + convert(i64::try_from(&storage)?, policy.as_str())? + }; + let DType::Extension(ext) = dtype else { + unreachable!() + }; + let storage = result.map_or_else( + || Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable)), + |v| Scalar::primitive(v, Nullability::Nullable), + ); + return Ok( + ConstantArray::new(Scalar::extension_ref(ext, storage), args.row_count()) + .into_array(), + ); + } + + if matches!(ambiguous.dtype(), DType::Null) { + return Ok( + ConstantArray::new(Scalar::null(dtype), args.row_count()).into_array(), + ); + } + let input = input.execute::(ctx)?; + let storage = input + .storage_array() + .clone() + .execute::(ctx)?; + let values = storage.to_buffer::(); + let valid = storage.validity()?.execute_mask(values.len(), ctx)?; + let policies = ambiguous.execute::(ctx)?; + let policies_valid = policies + .validity()? + .execute_mask(values.len(), ctx)?; + let mut output = Vec::with_capacity(values.len()); + for (i, ((value, valid), policy_valid)) in values + .iter() + .zip(valid.iter()) + .zip(policies_valid.iter()) + .enumerate() + { + output.push(if valid && policy_valid { + let bytes = policies.bytes_at(i); + let policy = std::str::from_utf8(&bytes) + .map_err(|e| vortex_err!("Invalid ambiguity policy: {e}"))?; + convert(*value, policy)? + } else { + None + }); + } + let DType::Extension(ext) = dtype else { + unreachable!() + }; + let storage = PrimitiveArray::from_option_iter(output) + .into_array() + .cast(DType::Primitive(PType::I64, Nullability::Nullable))?; + Ok(ExtensionArray::new(ext, storage).into_array()) + } +} + +fn timestamp_options(dtype: &DType) -> VortexResult<&TimestampOptions> { + if let DType::Extension(ext) = dtype { + if let Some(options) = ext.metadata_opt::() { + return Ok(options); + } + } + vortex_bail!("replace_time_zone() requires Timestamp, got {dtype}") +} + +fn resolve_zone(name: Option<&str>) -> VortexResult { + Ok(match name { + Some(name) => TimeZone::get(name)?, + None => TimeZone::UTC, + }) +} + +fn replace( + value: i64, + unit: TimeUnit, + source: &TimeZone, + target: &TimeZone, + policy: &str, + null_on_non_existent: bool, +) -> VortexResult> { + if !matches!(policy, "raise" | "earliest" | "latest" | "null") { + vortex_bail!( + "Invalid ambiguity policy {policy:?}: expected raise, earliest, latest, or null" + ); + } + // Keeping the same timezone does not introduce a new ambiguity to reject. + if source == target && policy == "raise" { + return Ok(Some(value)); + } + let timestamp = match unit { + TimeUnit::Seconds => JiffTimestamp::from_second(value)?, + TimeUnit::Milliseconds => JiffTimestamp::from_millisecond(value)?, + TimeUnit::Microseconds => JiffTimestamp::from_microsecond(value)?, + TimeUnit::Nanoseconds => JiffTimestamp::from_nanosecond(i128::from(value))?, + TimeUnit::Days => vortex_bail!("Timestamp cannot use day units"), + }; + let wall_time = source.to_datetime(timestamp); + let ambiguous = target.to_ambiguous_timestamp(wall_time); + let timestamp = match ambiguous.offset() { + AmbiguousOffset::Unambiguous { .. } => ambiguous.unambiguous()?, + AmbiguousOffset::Gap { .. } if null_on_non_existent => return Ok(None), + AmbiguousOffset::Gap { .. } => ambiguous.unambiguous()?, + AmbiguousOffset::Fold { .. } => match policy { + "earliest" => ambiguous.earlier()?, + "latest" => ambiguous.later()?, + "null" => return Ok(None), + _ => ambiguous.unambiguous()?, + }, + }; + Ok(Some(match unit { + TimeUnit::Seconds => timestamp.as_second(), + TimeUnit::Milliseconds => timestamp.as_millisecond(), + TimeUnit::Microseconds => timestamp.as_microsecond(), + TimeUnit::Nanoseconds => i64::try_from(timestamp.as_nanosecond()) + .map_err(|e| vortex_err!("Timestamp overflow: {e}"))?, + TimeUnit::Days => unreachable!(), + })) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use super::ReplaceTimeZone; + use super::ReplaceTimeZoneOptions; + use super::replace; + use super::resolve_zone; + use crate::extension::datetime::TimeUnit; + use crate::scalar_fn::ScalarFnVTable; + + #[rstest] + #[case(TimeUnit::Seconds, 1)] + #[case(TimeUnit::Milliseconds, 1_000)] + #[case(TimeUnit::Microseconds, 1_000_000)] + #[case(TimeUnit::Nanoseconds, 1_000_000_000)] + fn preserves_units_and_fractional_values( + #[case] unit: TimeUnit, + #[case] scale: i64, + ) -> VortexResult<()> { + let utc = resolve_zone(None)?; + let new_york = resolve_zone(Some("America/New_York"))?; + let value = 1_705_320_000 * scale + scale / 2; + assert_eq!( + replace(value, unit, &utc, &new_york, "raise", false)?, + Some(value + 18_000 * scale) + ); + assert_eq!( + replace(value + 18_000 * scale, unit, &new_york, &utc, "raise", false)?, + Some(value) + ); + Ok(()) + } + + #[test] + fn resolves_folds_and_gaps_separately() -> VortexResult<()> { + let utc = resolve_zone(None)?; + let new_york = resolve_zone(Some("America/New_York"))?; + let fold = "2024-11-03T01:30:00Z" + .parse::()? + .as_second(); + let gap = "2024-03-10T02:30:00Z" + .parse::()? + .as_second(); + let convert = |value, policy, null_on_non_existent| { + replace(value, TimeUnit::Seconds, &utc, &new_york, policy, null_on_non_existent) + }; + assert_eq!(convert(fold, "earliest", false)?, Some(fold + 14_400)); + assert_eq!(convert(fold, "latest", false)?, Some(fold + 18_000)); + assert_eq!(convert(fold, "null", false)?, None); + assert!(convert(fold, "raise", false).is_err()); + assert_eq!(convert(gap, "raise", true)?, None); + // Ambiguity policies must not resolve nonexistent wall times. + assert!(convert(gap, "earliest", false).is_err()); + Ok(()) + } + + #[test] + fn rejects_invalid_policy_and_timestamp_overflow() -> VortexResult<()> { + let utc = resolve_zone(None)?; + assert!(replace(0, TimeUnit::Seconds, &utc, &utc, "invalid", false).is_err()); + let new_york = resolve_zone(Some("America/New_York"))?; + assert!( + replace(i64::MAX, TimeUnit::Nanoseconds, &utc, &new_york, "raise", false) + .is_err() + ); + assert!(resolve_zone(Some("Not/A/Timezone")).is_err()); + Ok(()) + } + + #[rstest] + #[case(None, false)] + #[case(Some("America/New_York"), true)] + fn options_round_trip( + #[case] time_zone: Option<&str>, + #[case] null_on_non_existent: bool, + ) -> VortexResult<()> { + let options = ReplaceTimeZoneOptions { + time_zone: time_zone.map(Into::into), + null_on_non_existent, + }; + let metadata = ReplaceTimeZone.serialize(&options)?.unwrap(); + assert_eq!( + ReplaceTimeZone.deserialize(&metadata, &VortexSession::empty())?, + options + ); + Ok(()) + } +} diff --git a/vortex-array/src/scalar_fn/session.rs b/vortex-array/src/scalar_fn/session.rs index 4e85b4d078c..3031ad0b444 100644 --- a/vortex-array/src/scalar_fn/session.rs +++ b/vortex-array/src/scalar_fn/session.rs @@ -31,6 +31,7 @@ use crate::scalar_fn::fns::mask::Mask; use crate::scalar_fn::fns::merge::Merge; use crate::scalar_fn::fns::not::Not; use crate::scalar_fn::fns::pack::Pack; +use crate::scalar_fn::fns::replace_time_zone::ReplaceTimeZone; use crate::scalar_fn::fns::select::Select; use crate::scalar_fn::fns::stat::StatFn; use crate::scalar_fn::fns::variant_get::VariantGet; @@ -83,6 +84,7 @@ impl Default for ScalarFnSession { this.register(Merge); this.register(Not); this.register(Pack); + this.register(ReplaceTimeZone); this.register(Select); this.register(StatFn); this.register(VariantGet); diff --git a/vortex-python/python/vortex/_lib/expr.pyi b/vortex-python/python/vortex/_lib/expr.pyi index 6cd2b58bfd5..347e3bdc4c0 100644 --- a/vortex-python/python/vortex/_lib/expr.pyi +++ b/vortex-python/python/vortex/_lib/expr.pyi @@ -120,3 +120,11 @@ def variant_get(child: IntoExpr, path: VariantPath, dtype: DType | None = None) # Serialization def deserialize(data: bytes) -> Expr: ... + +def replace_time_zone( + child: IntoExpr, + time_zone: str | None, + *, + ambiguous: IntoExpr = "raise", + non_existent: Literal["raise", "null"] = "raise", +) -> Expr: ... diff --git a/vortex-python/python/vortex/expr.py b/vortex-python/python/vortex/expr.py index d02ed5ebab6..d810e7f2657 100644 --- a/vortex-python/python/vortex/expr.py +++ b/vortex-python/python/vortex/expr.py @@ -40,6 +40,7 @@ or_, or_collect, pack, + replace_time_zone, root, select, select_exclude, @@ -86,6 +87,7 @@ "or_", "or_collect", "pack", + "replace_time_zone", "root", "select", "select_exclude", diff --git a/vortex-python/src/expr/mod.rs b/vortex-python/src/expr/mod.rs index 8d76d20b051..5600f25a8e8 100644 --- a/vortex-python/src/expr/mod.rs +++ b/vortex-python/src/expr/mod.rs @@ -25,6 +25,7 @@ use vortex::scalar_fn::fns::between::StrictComparison; use vortex::scalar_fn::fns::binary::Binary; use vortex::scalar_fn::fns::merge::DuplicateHandling; use vortex::scalar_fn::fns::operators::Operator; +use vortex::scalar_fn::fns::replace_time_zone::ReplaceTimeZoneOptions; use vortex::scalar_fn::fns::variant_get::VariantPath; use vortex::scalar_fn::fns::variant_get::VariantPathElement; @@ -88,6 +89,9 @@ pub(crate) fn init(py: Python, parent: &Bound) -> PyResult<()> { m.add_function(wrap_pyfunction!(list_length, &m)?)?; m.add_function(wrap_pyfunction!(list_sum, &m)?)?; + // Temporal + m.add_function(wrap_pyfunction!(replace_time_zone, &m)?)?; + // Conditionals and misc m.add_function(wrap_pyfunction!(case_when, &m)?)?; m.add_function(wrap_pyfunction!(zip_, &m)?)?; @@ -1361,3 +1365,39 @@ fn variant_path(path: &Bound<'_, PyAny>) -> PyResult { .collect::>>()?; Ok(VariantPath::new(elements)) } + +/// Replace a timestamp's timezone while preserving its local wall time and time unit. +/// +/// Pass ``None`` to remove the timezone. ``ambiguous`` accepts an expression or one of +/// ``"raise"``, ``"earliest"``, ``"latest"``, and ``"null"``. ``non_existent`` accepts +/// ``"raise"`` or ``"null"``. Null inputs produce null outputs. +#[pyfunction] +#[pyo3(signature = ( + child, + time_zone, + *, + ambiguous = PyIntoExpr(lit("raise")), + non_existent = "raise" +))] +pub fn replace_time_zone( + child: PyIntoExpr, + time_zone: Option, + ambiguous: PyIntoExpr, + non_existent: &str, +) -> PyResult { + let null_on_non_existent = match non_existent { + "raise" => false, + "null" => true, + _ => return Err(PyValueError::new_err("non_existent must be 'raise' or 'null'")), + }; + Ok(PyExpr { + inner: expr::replace_time_zone( + child.into_inner(), + ambiguous.into_inner(), + ReplaceTimeZoneOptions { + time_zone: time_zone.map(Into::into), + null_on_non_existent, + }, + ), + }) +} diff --git a/vortex-python/test/test_expr.py b/vortex-python/test/test_expr.py index eebaa9be130..f2f1db20ccd 100644 --- a/vortex-python/test/test_expr.py +++ b/vortex-python/test/test_expr.py @@ -260,3 +260,13 @@ def test_deserialize_rejects_garbage() -> None: def test_serialize_is_stable() -> None: expr = ve.column("age") > 21 assert expr.serialize() == expr.serialize() + +@pytest.mark.parametrize("value", [1_705_320_000_000_000_123, None]) +def test_replace_time_zone_constant(value): + dtype = vx.timestamp("ns", nullable=value is None) + expression = ve.replace_time_zone(ve.literal(dtype, value), "America/New_York") + expected_value = None if value is None else value + 18_000_000_000_000 + expected = pa.array([expected_value] * 3, type=pa.timestamp("ns", tz="America/New_York")) + actual = vx.array(pa.array([0, 1, 2])).apply(expression).to_arrow_array() + assert actual.equals(expected) + From 993b177654ac4c32a13b5fa37975611c468cca0b Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 8 Oct 2026 17:56:05 +0000 Subject: [PATCH 2/7] style: satisfy current clippy and formatting after rebase Avoid Debug formatting in the ReplaceTimeZoneOptions Display impl, collapse the nested timestamp-options check, and apply rustfmt and ruff formatting. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01E92wUoS3LaqWVJmmS7vKVq --- .../src/scalar_fn/fns/replace_time_zone.rs | 75 ++++++++++--------- vortex-python/python/vortex/_lib/expr.pyi | 1 - vortex-python/src/expr/mod.rs | 6 +- vortex-python/test/test_expr.py | 2 +- 4 files changed, 47 insertions(+), 37 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/replace_time_zone.rs b/vortex-array/src/scalar_fn/fns/replace_time_zone.rs index 762fb7ca594..ed0665d1472 100644 --- a/vortex-array/src/scalar_fn/fns/replace_time_zone.rs +++ b/vortex-array/src/scalar_fn/fns/replace_time_zone.rs @@ -50,11 +50,11 @@ pub struct ReplaceTimeZoneOptions { impl fmt::Display for ReplaceTimeZoneOptions { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!( - f, - "tz={:?}, null_on_non_existent={}", - self.time_zone, self.null_on_non_existent - ) + match &self.time_zone { + Some(tz) => write!(f, "tz={tz}")?, + None => write!(f, "tz=none")?, + } + write!(f, ", null_on_non_existent={}", self.null_on_non_existent) } } @@ -106,11 +106,7 @@ impl ScalarFnVTable for ReplaceTimeZone { } } - fn return_dtype( - &self, - options: &Self::Options, - arg_dtypes: &[DType], - ) -> VortexResult { + fn return_dtype(&self, options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { let input = timestamp_options(&arg_dtypes[0])?; if input.unit == TimeUnit::Days { vortex_bail!("Timestamp cannot use day units"); @@ -124,12 +120,8 @@ impl ScalarFnVTable for ReplaceTimeZone { resolve_zone(options.time_zone.as_deref())?; // A row's ambiguity policy can request null even for non-nullable children. Ok(DType::Extension( - Timestamp::new_with_tz( - input.unit, - options.time_zone.clone(), - Nullability::Nullable, - ) - .erased(), + Timestamp::new_with_tz(input.unit, options.time_zone.clone(), Nullability::Nullable) + .erased(), )) } @@ -141,10 +133,8 @@ impl ScalarFnVTable for ReplaceTimeZone { ) -> VortexResult { let input = args.get(0)?; let ambiguous = args.get(1)?; - let dtype = self.return_dtype( - options, - &[input.dtype().clone(), ambiguous.dtype().clone()], - )?; + let dtype = + self.return_dtype(options, &[input.dtype().clone(), ambiguous.dtype().clone()])?; let metadata = timestamp_options(input.dtype())?.clone(); let source = resolve_zone(metadata.tz.as_deref())?; let target = resolve_zone(options.time_zone.as_deref())?; @@ -184,9 +174,7 @@ impl ScalarFnVTable for ReplaceTimeZone { } if matches!(ambiguous.dtype(), DType::Null) { - return Ok( - ConstantArray::new(Scalar::null(dtype), args.row_count()).into_array(), - ); + return Ok(ConstantArray::new(Scalar::null(dtype), args.row_count()).into_array()); } let input = input.execute::(ctx)?; let storage = input @@ -196,9 +184,7 @@ impl ScalarFnVTable for ReplaceTimeZone { let values = storage.to_buffer::(); let valid = storage.validity()?.execute_mask(values.len(), ctx)?; let policies = ambiguous.execute::(ctx)?; - let policies_valid = policies - .validity()? - .execute_mask(values.len(), ctx)?; + let policies_valid = policies.validity()?.execute_mask(values.len(), ctx)?; let mut output = Vec::with_capacity(values.len()); for (i, ((value, valid), policy_valid)) in values .iter() @@ -226,10 +212,10 @@ impl ScalarFnVTable for ReplaceTimeZone { } fn timestamp_options(dtype: &DType) -> VortexResult<&TimestampOptions> { - if let DType::Extension(ext) = dtype { - if let Some(options) = ext.metadata_opt::() { - return Ok(options); - } + if let DType::Extension(ext) = dtype + && let Some(options) = ext.metadata_opt::() + { + return Ok(options); } vortex_bail!("replace_time_zone() requires Timestamp, got {dtype}") } @@ -318,7 +304,14 @@ mod tests { Some(value + 18_000 * scale) ); assert_eq!( - replace(value + 18_000 * scale, unit, &new_york, &utc, "raise", false)?, + replace( + value + 18_000 * scale, + unit, + &new_york, + &utc, + "raise", + false + )?, Some(value) ); Ok(()) @@ -335,7 +328,14 @@ mod tests { .parse::()? .as_second(); let convert = |value, policy, null_on_non_existent| { - replace(value, TimeUnit::Seconds, &utc, &new_york, policy, null_on_non_existent) + replace( + value, + TimeUnit::Seconds, + &utc, + &new_york, + policy, + null_on_non_existent, + ) }; assert_eq!(convert(fold, "earliest", false)?, Some(fold + 14_400)); assert_eq!(convert(fold, "latest", false)?, Some(fold + 18_000)); @@ -353,8 +353,15 @@ mod tests { assert!(replace(0, TimeUnit::Seconds, &utc, &utc, "invalid", false).is_err()); let new_york = resolve_zone(Some("America/New_York"))?; assert!( - replace(i64::MAX, TimeUnit::Nanoseconds, &utc, &new_york, "raise", false) - .is_err() + replace( + i64::MAX, + TimeUnit::Nanoseconds, + &utc, + &new_york, + "raise", + false + ) + .is_err() ); assert!(resolve_zone(Some("Not/A/Timezone")).is_err()); Ok(()) diff --git a/vortex-python/python/vortex/_lib/expr.pyi b/vortex-python/python/vortex/_lib/expr.pyi index 347e3bdc4c0..6a377ac9914 100644 --- a/vortex-python/python/vortex/_lib/expr.pyi +++ b/vortex-python/python/vortex/_lib/expr.pyi @@ -120,7 +120,6 @@ def variant_get(child: IntoExpr, path: VariantPath, dtype: DType | None = None) # Serialization def deserialize(data: bytes) -> Expr: ... - def replace_time_zone( child: IntoExpr, time_zone: str | None, diff --git a/vortex-python/src/expr/mod.rs b/vortex-python/src/expr/mod.rs index 5600f25a8e8..27d3c663b0c 100644 --- a/vortex-python/src/expr/mod.rs +++ b/vortex-python/src/expr/mod.rs @@ -1388,7 +1388,11 @@ pub fn replace_time_zone( let null_on_non_existent = match non_existent { "raise" => false, "null" => true, - _ => return Err(PyValueError::new_err("non_existent must be 'raise' or 'null'")), + _ => { + return Err(PyValueError::new_err( + "non_existent must be 'raise' or 'null'", + )); + } }; Ok(PyExpr { inner: expr::replace_time_zone( diff --git a/vortex-python/test/test_expr.py b/vortex-python/test/test_expr.py index f2f1db20ccd..5a41566bcbb 100644 --- a/vortex-python/test/test_expr.py +++ b/vortex-python/test/test_expr.py @@ -261,6 +261,7 @@ def test_serialize_is_stable() -> None: expr = ve.column("age") > 21 assert expr.serialize() == expr.serialize() + @pytest.mark.parametrize("value", [1_705_320_000_000_000_123, None]) def test_replace_time_zone_constant(value): dtype = vx.timestamp("ns", nullable=value is None) @@ -269,4 +270,3 @@ def test_replace_time_zone_constant(value): expected = pa.array([expected_value] * 3, type=pa.timestamp("ns", tz="America/New_York")) actual = vx.array(pa.array([0, 1, 2])).apply(expression).to_arrow_array() assert actual.equals(expected) - From d7839030d0313a28b8f73500437c32423f620a88 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 9 Oct 2026 08:50:38 +0000 Subject: [PATCH 3/7] fix(expr): fold replace_time_zone over literals so stats pruning applies Polars serializes every timezone-aware datetime literal as replace_time_zone(), and stats pruning only matches bare literal operands. Without folding, `col >= replace_time_zone(lit)` falsified to `max(col) < min(replace_time_zone(lit))`, which zone maps cannot bind, so no zones were pruned. Fold literal inputs in simplify, mirroring Cast, and share the scalar path with execute. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01E92wUoS3LaqWVJmmS7vKVq --- .../src/scalar_fn/fns/replace_time_zone.rs | 120 +++++++++++++++--- 1 file changed, 99 insertions(+), 21 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/replace_time_zone.rs b/vortex-array/src/scalar_fn/fns/replace_time_zone.rs index ed0665d1472..1afb9383d39 100644 --- a/vortex-array/src/scalar_fn/fns/replace_time_zone.rs +++ b/vortex-array/src/scalar_fn/fns/replace_time_zone.rs @@ -28,6 +28,8 @@ use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; +use crate::expr::BoundExpression; +use crate::expr::bound; use crate::extension::datetime::TimeUnit; use crate::extension::datetime::Timestamp; use crate::extension::datetime::TimestampOptions; @@ -38,6 +40,7 @@ use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::literal::Literal; /// Timezone replacement options. The ambiguity policy is a UTF-8 expression child. #[derive(Clone, Debug, PartialEq, Eq, Hash)] @@ -106,6 +109,22 @@ impl ScalarFnVTable for ReplaceTimeZone { } } + fn simplify( + &self, + options: &Self::Options, + expr: &BoundExpression, + ) -> VortexResult> { + let (Some(value), Some(policy)) = ( + expr.child(0).as_opt::(), + expr.child(1).as_opt::(), + ) else { + return Ok(None); + }; + // Fold to a literal so stats pruning sees a bare literal operand. A failing replacement + // (e.g. a raised ambiguity) is left in place so the error surfaces at execution time. + Ok(replace_scalar(options, value, policy).ok().map(bound::lit)) + } + fn return_dtype(&self, options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { let input = timestamp_options(&arg_dtypes[0])?; if input.unit == TimeUnit::Days { @@ -150,27 +169,8 @@ impl ScalarFnVTable for ReplaceTimeZone { }; if let (Some(value), Some(policy)) = (input.as_constant(), ambiguous.as_constant()) { - let result = if value.is_null() || policy.is_null() { - None - } else { - let storage = value.as_extension().to_storage_scalar(); - let policy = policy - .as_utf8() - .value() - .ok_or_else(|| vortex_err!("Missing ambiguity policy"))?; - convert(i64::try_from(&storage)?, policy.as_str())? - }; - let DType::Extension(ext) = dtype else { - unreachable!() - }; - let storage = result.map_or_else( - || Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable)), - |v| Scalar::primitive(v, Nullability::Nullable), - ); - return Ok( - ConstantArray::new(Scalar::extension_ref(ext, storage), args.row_count()) - .into_array(), - ); + let scalar = replace_scalar(options, &value, &policy)?; + return Ok(ConstantArray::new(scalar, args.row_count()).into_array()); } if matches!(ambiguous.dtype(), DType::Null) { @@ -211,6 +211,42 @@ impl ScalarFnVTable for ReplaceTimeZone { } } +/// Replace the timezone of a single timestamp scalar using a scalar ambiguity policy. +fn replace_scalar( + options: &ReplaceTimeZoneOptions, + value: &Scalar, + policy: &Scalar, +) -> VortexResult { + let DType::Extension(ext) = + ReplaceTimeZone.return_dtype(options, &[value.dtype().clone(), policy.dtype().clone()])? + else { + unreachable!("replace_time_zone() returns a timestamp") + }; + let metadata = timestamp_options(value.dtype())?; + let result = if value.is_null() || policy.is_null() { + None + } else { + let storage = value.as_extension().to_storage_scalar(); + let policy = policy + .as_utf8() + .value() + .ok_or_else(|| vortex_err!("Missing ambiguity policy"))?; + replace( + i64::try_from(&storage)?, + metadata.unit, + &resolve_zone(metadata.tz.as_deref())?, + &resolve_zone(options.time_zone.as_deref())?, + policy.as_str(), + options.null_on_non_existent, + )? + }; + let storage = result.map_or_else( + || Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable)), + |v| Scalar::primitive(v, Nullability::Nullable), + ); + Ok(Scalar::extension_ref(ext, storage)) +} + fn timestamp_options(dtype: &DType) -> VortexResult<&TimestampOptions> { if let DType::Extension(ext) = dtype && let Some(options) = ext.metadata_opt::() @@ -278,14 +314,23 @@ fn replace( mod tests { use rstest::rstest; use vortex_error::VortexResult; + use vortex_error::vortex_err; use vortex_session::VortexSession; use super::ReplaceTimeZone; use super::ReplaceTimeZoneOptions; use super::replace; use super::resolve_zone; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::StructFields; + use crate::expr::lit; + use crate::expr::replace_time_zone; use crate::extension::datetime::TimeUnit; + use crate::extension::datetime::Timestamp; + use crate::scalar::Scalar; use crate::scalar_fn::ScalarFnVTable; + use crate::scalar_fn::fns::literal::Literal; #[rstest] #[case(TimeUnit::Seconds, 1)] @@ -385,4 +430,37 @@ mod tests { ); Ok(()) } + + #[test] + fn simplify_folds_literal_input() -> VortexResult<()> { + let timestamp = |tz: Option<&str>| { + Timestamp::new_with_tz( + TimeUnit::Microseconds, + tz.map(Into::into), + Nullability::Nullable, + ) + .erased() + }; + let micros = Scalar::primitive(1_704_153_600_000_000i64, Nullability::Nullable); + let expr = replace_time_zone( + lit(Scalar::extension_ref(timestamp(None), micros.clone())), + lit("earliest"), + ReplaceTimeZoneOptions { + time_zone: Some("UTC".into()), + null_on_non_existent: false, + }, + ); + let scope = DType::Struct(StructFields::empty(), Nullability::NonNullable); + let optimized = expr.bind(&scope)?.optimize()?; + + // Stats pruning only matches bare literal operands, so the replacement must fold. + let scalar = optimized + .as_opt::() + .ok_or_else(|| vortex_err!("expected a bare literal, got {optimized}"))?; + assert_eq!( + scalar, + &Scalar::extension_ref(timestamp(Some("UTC")), micros) + ); + Ok(()) + } } From 937500fadd6e548286b961a7d8df9dcf5f2c0176 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 9 Oct 2026 13:21:19 +0100 Subject: [PATCH 4/7] fix(expr): support fixed-offset timestamp timezones Signed-off-by: Robert Kruszewski --- .../src/extension/datetime/timestamp.rs | 37 +++++++++++++++++-- .../src/scalar_fn/fns/replace_time_zone.rs | 35 ++++++++++++------ vortex-python/test/test_expr.py | 36 ++++++++++++++++-- 3 files changed, 89 insertions(+), 19 deletions(-) diff --git a/vortex-array/src/extension/datetime/timestamp.rs b/vortex-array/src/extension/datetime/timestamp.rs index d07f68aa684..7c0cdaeb1bb 100644 --- a/vortex-array/src/extension/datetime/timestamp.rs +++ b/vortex-array/src/extension/datetime/timestamp.rs @@ -7,6 +7,8 @@ use std::fmt; use std::sync::Arc; use jiff::Span; +use jiff::fmt::temporal::DateTimeParser; +use jiff::tz::TimeZone; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -75,6 +77,17 @@ impl fmt::Display for TimestampOptions { } } +/// Resolve an Arrow timestamp timezone, accepting IANA names and fixed offsets. +pub(crate) fn resolve_time_zone(name: Option<&str>) -> VortexResult { + Ok(match name { + Some(name) if name.starts_with('+') || name.starts_with('-') => { + DateTimeParser::new().parse_time_zone(name)? + } + Some(name) => TimeZone::get(name)?, + None => TimeZone::UTC, + }) +} + /// Unpacked value of a [`Timestamp`] extension scalar. /// /// Each variant carries the raw storage value and an optional timezone. @@ -102,7 +115,8 @@ impl fmt::Display for TimestampValue<'_> { match tz { None => write!(f, "{ts}"), Some(tz) => { - let adjusted_ts = ts.in_tz(tz.as_ref()).vortex_expect("unknown timezone"); + let zone = resolve_time_zone(Some(tz)).vortex_expect("unknown timezone"); + let adjusted_ts = ts.to_zoned(zone); write!(f, "{adjusted_ts}",) } } @@ -218,12 +232,12 @@ impl ExtVTable for Timestamp { }; // Validate the storage value is within the valid range for Timestamp. - let ts = jiff::Timestamp::UNIX_EPOCH + jiff::Timestamp::UNIX_EPOCH .checked_add(span) .map_err(|e| vortex_err!("Invalid timestamp scalar: {}", e))?; if let Some(tz) = tz { - ts.in_tz(tz.as_ref()) + resolve_time_zone(Some(tz)) .map_err(|e| vortex_err!("Invalid timezone for timestamp scalar: {}", e))?; } @@ -235,6 +249,7 @@ impl ExtVTable for Timestamp { mod tests { use std::sync::Arc; + use rstest::rstest; use vortex_error::VortexResult; use crate::dtype::DType; @@ -253,6 +268,22 @@ mod tests { Ok(()) } + #[cfg_attr(miri, ignore)] + #[rstest] + #[case("+01:00", "1970-01-01T01:00:00+01:00[+01:00]")] + #[case("-05:30", "1969-12-31T18:30:00-05:30[-05:30]")] + fn display_fixed_offset_timestamp( + #[case] timezone: &str, + #[case] expected: &str, + ) -> VortexResult<()> { + let dtype = DType::Extension( + Timestamp::new_with_tz(TimeUnit::Seconds, Some(timezone.into()), Nullable).erased(), + ); + let scalar = Scalar::try_new(dtype, Some(ScalarValue::Primitive(PValue::I64(0))))?; + assert_eq!(format!("{}", scalar.as_extension()), expected); + Ok(()) + } + #[cfg_attr(miri, ignore)] #[test] fn reject_timestamp_with_invalid_timezone() { diff --git a/vortex-array/src/scalar_fn/fns/replace_time_zone.rs b/vortex-array/src/scalar_fn/fns/replace_time_zone.rs index 1afb9383d39..9c65f376056 100644 --- a/vortex-array/src/scalar_fn/fns/replace_time_zone.rs +++ b/vortex-array/src/scalar_fn/fns/replace_time_zone.rs @@ -33,6 +33,7 @@ use crate::expr::bound; use crate::extension::datetime::TimeUnit; use crate::extension::datetime::Timestamp; use crate::extension::datetime::TimestampOptions; +use crate::extension::datetime::resolve_time_zone as resolve_zone; use crate::proto::expr as pb; use crate::scalar::Scalar; use crate::scalar_fn::Arity; @@ -256,13 +257,6 @@ fn timestamp_options(dtype: &DType) -> VortexResult<&TimestampOptions> { vortex_bail!("replace_time_zone() requires Timestamp, got {dtype}") } -fn resolve_zone(name: Option<&str>) -> VortexResult { - Ok(match name { - Some(name) => TimeZone::get(name)?, - None => TimeZone::UTC, - }) -} - fn replace( value: i64, unit: TimeUnit, @@ -409,6 +403,8 @@ mod tests { .is_err() ); assert!(resolve_zone(Some("Not/A/Timezone")).is_err()); + assert!(resolve_zone(Some("+99:00")).is_err()); + assert!(resolve_zone(Some("-05:60")).is_err()); Ok(()) } @@ -431,8 +427,17 @@ mod tests { Ok(()) } - #[test] - fn simplify_folds_literal_input() -> VortexResult<()> { + #[rstest] + #[case(None, Some("UTC"), 0)] + #[case(Some("+01:00"), None, 3_600)] + #[case(Some("-05:30"), None, -19_800)] + #[case(None, Some("+05:30"), -19_800)] + #[case(Some("+01:00"), Some("-05:30"), 23_400)] + fn simplify_folds_literal_input( + #[case] source: Option<&str>, + #[case] target: Option<&str>, + #[case] offset_seconds: i64, + ) -> VortexResult<()> { let timestamp = |tz: Option<&str>| { Timestamp::new_with_tz( TimeUnit::Microseconds, @@ -443,10 +448,10 @@ mod tests { }; let micros = Scalar::primitive(1_704_153_600_000_000i64, Nullability::Nullable); let expr = replace_time_zone( - lit(Scalar::extension_ref(timestamp(None), micros.clone())), + lit(Scalar::extension_ref(timestamp(source), micros)), lit("earliest"), ReplaceTimeZoneOptions { - time_zone: Some("UTC".into()), + time_zone: target.map(Into::into), null_on_non_existent: false, }, ); @@ -459,7 +464,13 @@ mod tests { .ok_or_else(|| vortex_err!("expected a bare literal, got {optimized}"))?; assert_eq!( scalar, - &Scalar::extension_ref(timestamp(Some("UTC")), micros) + &Scalar::extension_ref( + timestamp(target), + Scalar::primitive( + 1_704_153_600_000_000i64 + offset_seconds * 1_000_000, + Nullability::Nullable, + ), + ) ); Ok(()) } diff --git a/vortex-python/test/test_expr.py b/vortex-python/test/test_expr.py index 5a41566bcbb..7b404ca2807 100644 --- a/vortex-python/test/test_expr.py +++ b/vortex-python/test/test_expr.py @@ -263,10 +263,38 @@ def test_serialize_is_stable() -> None: @pytest.mark.parametrize("value", [1_705_320_000_000_000_123, None]) -def test_replace_time_zone_constant(value): +@pytest.mark.parametrize( + ("target", "offset_seconds"), [("America/New_York", 18_000), ("+05:30", -19_800)] +) +def test_replace_time_zone_constant(value, target, offset_seconds): dtype = vx.timestamp("ns", nullable=value is None) - expression = ve.replace_time_zone(ve.literal(dtype, value), "America/New_York") - expected_value = None if value is None else value + 18_000_000_000_000 - expected = pa.array([expected_value] * 3, type=pa.timestamp("ns", tz="America/New_York")) + expression = ve.replace_time_zone(ve.literal(dtype, value), target) + expected_value = None if value is None else value + offset_seconds * 1_000_000_000 + expected = pa.array([expected_value] * 3, type=pa.timestamp("ns", tz=target)) actual = vx.array(pa.array([0, 1, 2])).apply(expression).to_arrow_array() assert actual.equals(expected) + + +@pytest.mark.parametrize( + ("source", "target", "offset_seconds"), + [ + ("+01:00", None, 3_600), + ("-05:30", None, -19_800), + (None, "+05:30", -19_800), + ("+01:00", "-05:30", 23_400), + ("+01:00", "America/New_York", 21_600), + ("America/New_York", "+01:00", -21_600), + ("+00:00", None, 0), + ], +) +@pytest.mark.parametrize("unit", ["s", "ms", "us", "ns"]) +def test_replace_time_zone_fixed_offsets(source, target, offset_seconds, unit): + scale = {"s": 1, "ms": 1_000, "us": 1_000_000, "ns": 1_000_000_000}[unit] + values = [0, scale // 2, -scale - scale // 2, None] + array = vx.array(pa.array(values, type=pa.timestamp(unit, tz=source))) + actual = array.apply(ve.replace_time_zone(ve.root(), target)).to_arrow_array() + expected = pa.array( + [None if value is None else value + offset_seconds * scale for value in values], + type=pa.timestamp(unit, tz=target), + ) + assert actual.equals(expected) From 8441e38f067a4f3b123daba4c10c4f35c44d04e1 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 9 Oct 2026 13:28:27 +0100 Subject: [PATCH 5/7] style(python): format timezone replacement tests Signed-off-by: Robert Kruszewski --- vortex-python/test/test_expr.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/vortex-python/test/test_expr.py b/vortex-python/test/test_expr.py index 7b404ca2807..8471cfb5ab0 100644 --- a/vortex-python/test/test_expr.py +++ b/vortex-python/test/test_expr.py @@ -263,9 +263,7 @@ def test_serialize_is_stable() -> None: @pytest.mark.parametrize("value", [1_705_320_000_000_000_123, None]) -@pytest.mark.parametrize( - ("target", "offset_seconds"), [("America/New_York", 18_000), ("+05:30", -19_800)] -) +@pytest.mark.parametrize(("target", "offset_seconds"), [("America/New_York", 18_000), ("+05:30", -19_800)]) def test_replace_time_zone_constant(value, target, offset_seconds): dtype = vx.timestamp("ns", nullable=value is None) expression = ve.replace_time_zone(ve.literal(dtype, value), target) From 61d3d7d01ca843bb9eea7ebfe9a8cd8313e5299c Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 9 Oct 2026 13:47:38 +0100 Subject: [PATCH 6/7] fix(expr): trust UTF-8 validity of ambiguity policies Signed-off-by: Robert Kruszewski --- vortex-array/src/scalar_fn/fns/replace_time_zone.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/replace_time_zone.rs b/vortex-array/src/scalar_fn/fns/replace_time_zone.rs index 9c65f376056..5e1a27a0904 100644 --- a/vortex-array/src/scalar_fn/fns/replace_time_zone.rs +++ b/vortex-array/src/scalar_fn/fns/replace_time_zone.rs @@ -195,8 +195,8 @@ impl ScalarFnVTable for ReplaceTimeZone { { output.push(if valid && policy_valid { let bytes = policies.bytes_at(i); - let policy = std::str::from_utf8(&bytes) - .map_err(|e| vortex_err!("Invalid ambiguity policy: {e}"))?; + // SAFETY: policies has Utf8 dtype and this row is valid. + let policy = unsafe { std::str::from_utf8_unchecked(&bytes) }; convert(*value, policy)? } else { None From 92ddab240122ff0cba351ea50fe1d03826d2f5b8 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 9 Oct 2026 13:03:25 +0000 Subject: [PATCH 7/7] perf(expr): build replace_time_zone output buffers directly Write values into a BufferMut and validity into a BitBufferMut instead of collecting Vec>, rebuilding it with from_option_iter, and casting it to the nullable i64 it already was. Add a test that runs the array path with per-row policies and null inputs and policies. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01E92wUoS3LaqWVJmmS7vKVq --- .../src/scalar_fn/fns/replace_time_zone.rs | 90 +++++++++++++++++-- 1 file changed, 83 insertions(+), 7 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/replace_time_zone.rs b/vortex-array/src/scalar_fn/fns/replace_time_zone.rs index 5e1a27a0904..f146fdfdeaf 100644 --- a/vortex-array/src/scalar_fn/fns/replace_time_zone.rs +++ b/vortex-array/src/scalar_fn/fns/replace_time_zone.rs @@ -10,6 +10,8 @@ use jiff::Timestamp as JiffTimestamp; use jiff::tz::AmbiguousOffset; use jiff::tz::TimeZone; use prost::Message; +use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; @@ -24,7 +26,6 @@ use crate::arrays::ExtensionArray; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinViewArray; use crate::arrays::extension::ExtensionArrayExt; -use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -42,6 +43,7 @@ use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::literal::Literal; +use crate::validity::Validity; /// Timezone replacement options. The ambiguity policy is a UTF-8 expression child. #[derive(Clone, Debug, PartialEq, Eq, Hash)] @@ -186,28 +188,33 @@ impl ScalarFnVTable for ReplaceTimeZone { let valid = storage.validity()?.execute_mask(values.len(), ctx)?; let policies = ambiguous.execute::(ctx)?; let policies_valid = policies.validity()?.execute_mask(values.len(), ctx)?; - let mut output = Vec::with_capacity(values.len()); + let mut output = BufferMut::::with_capacity(values.len()); + let mut output_valid = BitBufferMut::with_capacity(values.len()); for (i, ((value, valid), policy_valid)) in values .iter() .zip(valid.iter()) .zip(policies_valid.iter()) .enumerate() { - output.push(if valid && policy_valid { + let replaced = if valid && policy_valid { let bytes = policies.bytes_at(i); // SAFETY: policies has Utf8 dtype and this row is valid. let policy = unsafe { std::str::from_utf8_unchecked(&bytes) }; convert(*value, policy)? } else { None - }); + }; + output.push(replaced.unwrap_or_default()); + output_valid.append(replaced.is_some()); } let DType::Extension(ext) = dtype else { unreachable!() }; - let storage = PrimitiveArray::from_option_iter(output) - .into_array() - .cast(DType::Primitive(PType::I64, Nullability::Nullable))?; + let storage = PrimitiveArray::new( + output.freeze(), + Validity::from_bit_buffer(output_valid.freeze(), Nullability::Nullable), + ) + .into_array(); Ok(ExtensionArray::new(ext, storage).into_array()) } } @@ -315,11 +322,21 @@ mod tests { use super::ReplaceTimeZoneOptions; use super::replace; use super::resolve_zone; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::ExtensionArray; + use crate::arrays::PrimitiveArray; + use crate::arrays::StructArray; + use crate::arrays::VarBinViewArray; + use crate::assert_arrays_eq; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::StructFields; + use crate::expr::get_item; use crate::expr::lit; use crate::expr::replace_time_zone; + use crate::expr::root; use crate::extension::datetime::TimeUnit; use crate::extension::datetime::Timestamp; use crate::scalar::Scalar; @@ -474,4 +491,63 @@ mod tests { ); Ok(()) } + + #[test] + fn applies_per_row_policies_and_nulls() -> VortexResult<()> { + let timestamp = |tz: Option<&str>| { + Timestamp::new_with_tz( + TimeUnit::Microseconds, + tz.map(Into::into), + Nullability::Nullable, + ) + .erased() + }; + // 2024-11-03 01:30 occurs twice in New York. + let fold = 1_730_597_400_000_000i64; + let hour = 3_600_000_000i64; + let input = ExtensionArray::new( + timestamp(None), + PrimitiveArray::from_option_iter([ + Some(fold), + Some(fold), + Some(fold), + Some(fold), + None, + ]) + .into_array(), + ); + let policies = VarBinViewArray::from_iter_nullable_str([ + Some("earliest"), + Some("latest"), + Some("null"), + None, + Some("raise"), + ]); + let array = + StructArray::from_fields(&[("t", input.into_array()), ("p", policies.into_array())])? + .into_array(); + let result = array.apply(&replace_time_zone( + get_item("t", root()), + get_item("p", root()), + ReplaceTimeZoneOptions { + time_zone: Some("America/New_York".into()), + null_on_non_existent: false, + }, + ))?; + + let expected = ExtensionArray::new( + timestamp(Some("America/New_York")), + PrimitiveArray::from_option_iter([ + Some(fold + 4 * hour), + Some(fold + 5 * hour), + None, + None, + None, + ]) + .into_array(), + ); + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } }