feat(expr): add timestamp timezone replacement

Signed-off-by: Daniel King <dan@spiraldb.com>

Author: danking

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;
+}

vortex-array/src/expr/exprs.rs

@@ -52,6 +52,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;
@@ -1282,6 +1284,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
@@ -1330,6 +1355,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;

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;

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,
+}

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;