vortex_array/scalar_fn/fns/
replace_time_zone.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Reinterpret timestamp wall times in a different timezone.
5
6use std::fmt;
7use std::sync::Arc;
8
9use jiff::Timestamp as JiffTimestamp;
10use jiff::tz::AmbiguousOffset;
11use jiff::tz::TimeZone;
12use prost::Message;
13use vortex_buffer::BitBufferMut;
14use vortex_buffer::BufferMut;
15use vortex_error::VortexResult;
16use vortex_error::vortex_bail;
17use vortex_error::vortex_err;
18use vortex_session::VortexSession;
19use vortex_session::registry::CachedId;
20
21use crate::ArrayRef;
22use crate::ExecutionCtx;
23use crate::IntoArray;
24use crate::arrays::ConstantArray;
25use crate::arrays::ExtensionArray;
26use crate::arrays::PrimitiveArray;
27use crate::arrays::VarBinViewArray;
28use crate::arrays::extension::ExtensionArrayExt;
29use crate::dtype::DType;
30use crate::dtype::Nullability;
31use crate::dtype::PType;
32use crate::expr::BoundExpression;
33use crate::expr::bound;
34use crate::extension::datetime::TimeUnit;
35use crate::extension::datetime::Timestamp;
36use crate::extension::datetime::TimestampOptions;
37use crate::extension::datetime::resolve_time_zone as resolve_zone;
38use crate::proto::expr as pb;
39use crate::scalar::Scalar;
40use crate::scalar_fn::Arity;
41use crate::scalar_fn::ChildName;
42use crate::scalar_fn::ExecutionArgs;
43use crate::scalar_fn::ScalarFnId;
44use crate::scalar_fn::ScalarFnVTable;
45use crate::scalar_fn::fns::literal::Literal;
46use crate::validity::Validity;
47
48/// Timezone replacement options. The ambiguity policy is a UTF-8 expression child.
49#[derive(Clone, Debug, PartialEq, Eq, Hash)]
50pub struct ReplaceTimeZoneOptions {
51    /// Destination timezone, or `None` to remove the timezone while preserving wall time.
52    pub time_zone: Option<Arc<str>>,
53    /// Return null for nonexistent wall times instead of raising an error.
54    pub null_on_non_existent: bool,
55}
56
57impl fmt::Display for ReplaceTimeZoneOptions {
58    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
59        match &self.time_zone {
60            Some(tz) => write!(f, "tz={tz}")?,
61            None => write!(f, "tz=none")?,
62        }
63        write!(f, ", null_on_non_existent={}", self.null_on_non_existent)
64    }
65}
66
67/// Replace a timestamp's timezone, preserving its local wall time.
68///
69/// The second child selects `raise`, `earliest`, `latest`, or `null` for ambiguous times.
70#[derive(Clone)]
71pub struct ReplaceTimeZone;
72
73impl ScalarFnVTable for ReplaceTimeZone {
74    type Options = ReplaceTimeZoneOptions;
75
76    fn id(&self) -> ScalarFnId {
77        static ID: CachedId = CachedId::new("vortex.replace_time_zone");
78        *ID
79    }
80
81    fn serialize(&self, options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
82        Ok(Some(
83            pb::ReplaceTimeZoneOpts {
84                time_zone: options.time_zone.as_ref().map(ToString::to_string),
85                null_on_non_existent: options.null_on_non_existent,
86            }
87            .encode_to_vec(),
88        ))
89    }
90
91    fn deserialize(
92        &self,
93        metadata: &[u8],
94        _session: &VortexSession,
95    ) -> VortexResult<Self::Options> {
96        let options = pb::ReplaceTimeZoneOpts::decode(metadata)?;
97        Ok(ReplaceTimeZoneOptions {
98            time_zone: options.time_zone.map(Arc::from),
99            null_on_non_existent: options.null_on_non_existent,
100        })
101    }
102
103    fn arity(&self, _options: &Self::Options) -> Arity {
104        Arity::Exact(2)
105    }
106
107    fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName {
108        match child_idx {
109            0 => ChildName::from("input"),
110            1 => ChildName::from("ambiguous"),
111            _ => unreachable!("Invalid child index {child_idx} for replace_time_zone()"),
112        }
113    }
114
115    fn simplify(
116        &self,
117        options: &Self::Options,
118        expr: &BoundExpression,
119    ) -> VortexResult<Option<BoundExpression>> {
120        let (Some(value), Some(policy)) = (
121            expr.child(0).as_opt::<Literal>(),
122            expr.child(1).as_opt::<Literal>(),
123        ) else {
124            return Ok(None);
125        };
126        // Fold to a literal so stats pruning sees a bare literal operand. A failing replacement
127        // (e.g. a raised ambiguity) is left in place so the error surfaces at execution time.
128        Ok(replace_scalar(options, value, policy).ok().map(bound::lit))
129    }
130
131    fn return_dtype(&self, options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult<DType> {
132        let input = timestamp_options(&arg_dtypes[0])?;
133        if input.unit == TimeUnit::Days {
134            vortex_bail!("Timestamp cannot use day units");
135        }
136        if !matches!(arg_dtypes[1], DType::Utf8(_) | DType::Null) {
137            vortex_bail!(
138                "replace_time_zone() requires a UTF-8 ambiguity policy, got {}",
139                arg_dtypes[1]
140            );
141        }
142        resolve_zone(options.time_zone.as_deref())?;
143        // A row's ambiguity policy can request null even for non-nullable children.
144        Ok(DType::Extension(
145            Timestamp::new_with_tz(input.unit, options.time_zone.clone(), Nullability::Nullable)
146                .erased(),
147        ))
148    }
149
150    fn execute(
151        &self,
152        options: &Self::Options,
153        args: &dyn ExecutionArgs,
154        ctx: &mut ExecutionCtx,
155    ) -> VortexResult<ArrayRef> {
156        let input = args.get(0)?;
157        let ambiguous = args.get(1)?;
158        let dtype =
159            self.return_dtype(options, &[input.dtype().clone(), ambiguous.dtype().clone()])?;
160        let metadata = timestamp_options(input.dtype())?.clone();
161        let source = resolve_zone(metadata.tz.as_deref())?;
162        let target = resolve_zone(options.time_zone.as_deref())?;
163        let convert = |value, policy: &str| {
164            replace(
165                value,
166                metadata.unit,
167                &source,
168                &target,
169                policy,
170                options.null_on_non_existent,
171            )
172        };
173
174        if let (Some(value), Some(policy)) = (input.as_constant(), ambiguous.as_constant()) {
175            let scalar = replace_scalar(options, &value, &policy)?;
176            return Ok(ConstantArray::new(scalar, args.row_count()).into_array());
177        }
178
179        if matches!(ambiguous.dtype(), DType::Null) {
180            return Ok(ConstantArray::new(Scalar::null(dtype), args.row_count()).into_array());
181        }
182        let input = input.execute::<ExtensionArray>(ctx)?;
183        let storage = input
184            .storage_array()
185            .clone()
186            .execute::<PrimitiveArray>(ctx)?;
187        let values = storage.to_buffer::<i64>();
188        let valid = storage.validity()?.execute_mask(values.len(), ctx)?;
189        let policies = ambiguous.execute::<VarBinViewArray>(ctx)?;
190        let policies_valid = policies.validity()?.execute_mask(values.len(), ctx)?;
191        let mut output = BufferMut::<i64>::with_capacity(values.len());
192        let mut output_valid = BitBufferMut::with_capacity(values.len());
193        for (i, ((value, valid), policy_valid)) in values
194            .iter()
195            .zip(valid.iter())
196            .zip(policies_valid.iter())
197            .enumerate()
198        {
199            let replaced = if valid && policy_valid {
200                let bytes = policies.bytes_at(i);
201                // SAFETY: policies has Utf8 dtype and this row is valid.
202                let policy = unsafe { std::str::from_utf8_unchecked(&bytes) };
203                convert(*value, policy)?
204            } else {
205                None
206            };
207            output.push(replaced.unwrap_or_default());
208            output_valid.append(replaced.is_some());
209        }
210        let DType::Extension(ext) = dtype else {
211            unreachable!()
212        };
213        let storage = PrimitiveArray::new(
214            output.freeze(),
215            Validity::from_bit_buffer(output_valid.freeze(), Nullability::Nullable),
216        )
217        .into_array();
218        Ok(ExtensionArray::new(ext, storage).into_array())
219    }
220}
221
222/// Replace the timezone of a single timestamp scalar using a scalar ambiguity policy.
223fn replace_scalar(
224    options: &ReplaceTimeZoneOptions,
225    value: &Scalar,
226    policy: &Scalar,
227) -> VortexResult<Scalar> {
228    let DType::Extension(ext) =
229        ReplaceTimeZone.return_dtype(options, &[value.dtype().clone(), policy.dtype().clone()])?
230    else {
231        unreachable!("replace_time_zone() returns a timestamp")
232    };
233    let metadata = timestamp_options(value.dtype())?;
234    let result = if value.is_null() || policy.is_null() {
235        None
236    } else {
237        let storage = value.as_extension().to_storage_scalar();
238        let policy = policy
239            .as_utf8()
240            .value()
241            .ok_or_else(|| vortex_err!("Missing ambiguity policy"))?;
242        replace(
243            i64::try_from(&storage)?,
244            metadata.unit,
245            &resolve_zone(metadata.tz.as_deref())?,
246            &resolve_zone(options.time_zone.as_deref())?,
247            policy.as_str(),
248            options.null_on_non_existent,
249        )?
250    };
251    let storage = result.map_or_else(
252        || Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable)),
253        |v| Scalar::primitive(v, Nullability::Nullable),
254    );
255    Ok(Scalar::extension_ref(ext, storage))
256}
257
258fn timestamp_options(dtype: &DType) -> VortexResult<&TimestampOptions> {
259    if let DType::Extension(ext) = dtype
260        && let Some(options) = ext.metadata_opt::<Timestamp>()
261    {
262        return Ok(options);
263    }
264    vortex_bail!("replace_time_zone() requires Timestamp, got {dtype}")
265}
266
267fn replace(
268    value: i64,
269    unit: TimeUnit,
270    source: &TimeZone,
271    target: &TimeZone,
272    policy: &str,
273    null_on_non_existent: bool,
274) -> VortexResult<Option<i64>> {
275    if !matches!(policy, "raise" | "earliest" | "latest" | "null") {
276        vortex_bail!(
277            "Invalid ambiguity policy {policy:?}: expected raise, earliest, latest, or null"
278        );
279    }
280    // Keeping the same timezone does not introduce a new ambiguity to reject.
281    if source == target && policy == "raise" {
282        return Ok(Some(value));
283    }
284    let timestamp = match unit {
285        TimeUnit::Seconds => JiffTimestamp::from_second(value)?,
286        TimeUnit::Milliseconds => JiffTimestamp::from_millisecond(value)?,
287        TimeUnit::Microseconds => JiffTimestamp::from_microsecond(value)?,
288        TimeUnit::Nanoseconds => JiffTimestamp::from_nanosecond(i128::from(value))?,
289        TimeUnit::Days => vortex_bail!("Timestamp cannot use day units"),
290    };
291    let wall_time = source.to_datetime(timestamp);
292    let ambiguous = target.to_ambiguous_timestamp(wall_time);
293    let timestamp = match ambiguous.offset() {
294        AmbiguousOffset::Unambiguous { .. } => ambiguous.unambiguous()?,
295        AmbiguousOffset::Gap { .. } if null_on_non_existent => return Ok(None),
296        AmbiguousOffset::Gap { .. } => ambiguous.unambiguous()?,
297        AmbiguousOffset::Fold { .. } => match policy {
298            "earliest" => ambiguous.earlier()?,
299            "latest" => ambiguous.later()?,
300            "null" => return Ok(None),
301            _ => ambiguous.unambiguous()?,
302        },
303    };
304    Ok(Some(match unit {
305        TimeUnit::Seconds => timestamp.as_second(),
306        TimeUnit::Milliseconds => timestamp.as_millisecond(),
307        TimeUnit::Microseconds => timestamp.as_microsecond(),
308        TimeUnit::Nanoseconds => i64::try_from(timestamp.as_nanosecond())
309            .map_err(|e| vortex_err!("Timestamp overflow: {e}"))?,
310        TimeUnit::Days => unreachable!(),
311    }))
312}
313
314#[cfg(test)]
315mod tests {
316    use rstest::rstest;
317    use vortex_error::VortexResult;
318    use vortex_error::vortex_err;
319    use vortex_session::VortexSession;
320
321    use super::ReplaceTimeZone;
322    use super::ReplaceTimeZoneOptions;
323    use super::replace;
324    use super::resolve_zone;
325    use crate::IntoArray;
326    use crate::VortexSessionExecute;
327    use crate::array_session;
328    use crate::arrays::ExtensionArray;
329    use crate::arrays::PrimitiveArray;
330    use crate::arrays::StructArray;
331    use crate::arrays::VarBinViewArray;
332    use crate::assert_arrays_eq;
333    use crate::dtype::DType;
334    use crate::dtype::Nullability;
335    use crate::dtype::StructFields;
336    use crate::expr::get_item;
337    use crate::expr::lit;
338    use crate::expr::replace_time_zone;
339    use crate::expr::root;
340    use crate::extension::datetime::TimeUnit;
341    use crate::extension::datetime::Timestamp;
342    use crate::scalar::Scalar;
343    use crate::scalar_fn::ScalarFnVTable;
344    use crate::scalar_fn::fns::literal::Literal;
345
346    #[rstest]
347    #[case(TimeUnit::Seconds, 1)]
348    #[case(TimeUnit::Milliseconds, 1_000)]
349    #[case(TimeUnit::Microseconds, 1_000_000)]
350    #[case(TimeUnit::Nanoseconds, 1_000_000_000)]
351    fn preserves_units_and_fractional_values(
352        #[case] unit: TimeUnit,
353        #[case] scale: i64,
354    ) -> VortexResult<()> {
355        let utc = resolve_zone(None)?;
356        let new_york = resolve_zone(Some("America/New_York"))?;
357        let value = 1_705_320_000 * scale + scale / 2;
358        assert_eq!(
359            replace(value, unit, &utc, &new_york, "raise", false)?,
360            Some(value + 18_000 * scale)
361        );
362        assert_eq!(
363            replace(
364                value + 18_000 * scale,
365                unit,
366                &new_york,
367                &utc,
368                "raise",
369                false
370            )?,
371            Some(value)
372        );
373        Ok(())
374    }
375
376    #[test]
377    fn resolves_folds_and_gaps_separately() -> VortexResult<()> {
378        let utc = resolve_zone(None)?;
379        let new_york = resolve_zone(Some("America/New_York"))?;
380        let fold = "2024-11-03T01:30:00Z"
381            .parse::<jiff::Timestamp>()?
382            .as_second();
383        let gap = "2024-03-10T02:30:00Z"
384            .parse::<jiff::Timestamp>()?
385            .as_second();
386        let convert = |value, policy, null_on_non_existent| {
387            replace(
388                value,
389                TimeUnit::Seconds,
390                &utc,
391                &new_york,
392                policy,
393                null_on_non_existent,
394            )
395        };
396        assert_eq!(convert(fold, "earliest", false)?, Some(fold + 14_400));
397        assert_eq!(convert(fold, "latest", false)?, Some(fold + 18_000));
398        assert_eq!(convert(fold, "null", false)?, None);
399        assert!(convert(fold, "raise", false).is_err());
400        assert_eq!(convert(gap, "raise", true)?, None);
401        // Ambiguity policies must not resolve nonexistent wall times.
402        assert!(convert(gap, "earliest", false).is_err());
403        Ok(())
404    }
405
406    #[test]
407    fn rejects_invalid_policy_and_timestamp_overflow() -> VortexResult<()> {
408        let utc = resolve_zone(None)?;
409        assert!(replace(0, TimeUnit::Seconds, &utc, &utc, "invalid", false).is_err());
410        let new_york = resolve_zone(Some("America/New_York"))?;
411        assert!(
412            replace(
413                i64::MAX,
414                TimeUnit::Nanoseconds,
415                &utc,
416                &new_york,
417                "raise",
418                false
419            )
420            .is_err()
421        );
422        assert!(resolve_zone(Some("Not/A/Timezone")).is_err());
423        assert!(resolve_zone(Some("+99:00")).is_err());
424        assert!(resolve_zone(Some("-05:60")).is_err());
425        Ok(())
426    }
427
428    #[rstest]
429    #[case(None, false)]
430    #[case(Some("America/New_York"), true)]
431    fn options_round_trip(
432        #[case] time_zone: Option<&str>,
433        #[case] null_on_non_existent: bool,
434    ) -> VortexResult<()> {
435        let options = ReplaceTimeZoneOptions {
436            time_zone: time_zone.map(Into::into),
437            null_on_non_existent,
438        };
439        let metadata = ReplaceTimeZone.serialize(&options)?.unwrap();
440        assert_eq!(
441            ReplaceTimeZone.deserialize(&metadata, &VortexSession::empty())?,
442            options
443        );
444        Ok(())
445    }
446
447    #[rstest]
448    #[case(None, Some("UTC"), 0)]
449    #[case(Some("+01:00"), None, 3_600)]
450    #[case(Some("-05:30"), None, -19_800)]
451    #[case(None, Some("+05:30"), -19_800)]
452    #[case(Some("+01:00"), Some("-05:30"), 23_400)]
453    fn simplify_folds_literal_input(
454        #[case] source: Option<&str>,
455        #[case] target: Option<&str>,
456        #[case] offset_seconds: i64,
457    ) -> VortexResult<()> {
458        let timestamp = |tz: Option<&str>| {
459            Timestamp::new_with_tz(
460                TimeUnit::Microseconds,
461                tz.map(Into::into),
462                Nullability::Nullable,
463            )
464            .erased()
465        };
466        let micros = Scalar::primitive(1_704_153_600_000_000i64, Nullability::Nullable);
467        let expr = replace_time_zone(
468            lit(Scalar::extension_ref(timestamp(source), micros)),
469            lit("earliest"),
470            ReplaceTimeZoneOptions {
471                time_zone: target.map(Into::into),
472                null_on_non_existent: false,
473            },
474        );
475        let scope = DType::Struct(StructFields::empty(), Nullability::NonNullable);
476        let optimized = expr.bind(&scope)?.optimize()?;
477
478        // Stats pruning only matches bare literal operands, so the replacement must fold.
479        let scalar = optimized
480            .as_opt::<Literal>()
481            .ok_or_else(|| vortex_err!("expected a bare literal, got {optimized}"))?;
482        assert_eq!(
483            scalar,
484            &Scalar::extension_ref(
485                timestamp(target),
486                Scalar::primitive(
487                    1_704_153_600_000_000i64 + offset_seconds * 1_000_000,
488                    Nullability::Nullable,
489                ),
490            )
491        );
492        assert_eq!(scalar.dtype(), &DType::Extension(timestamp(target)));
493        Ok(())
494    }
495
496    #[test]
497    fn applies_per_row_policies_and_nulls() -> VortexResult<()> {
498        let timestamp = |tz: Option<&str>| {
499            Timestamp::new_with_tz(
500                TimeUnit::Microseconds,
501                tz.map(Into::into),
502                Nullability::Nullable,
503            )
504            .erased()
505        };
506        // 2024-11-03 01:30 occurs twice in New York.
507        let fold = 1_730_597_400_000_000i64;
508        let hour = 3_600_000_000i64;
509        let input = ExtensionArray::new(
510            timestamp(None),
511            PrimitiveArray::from_option_iter([
512                Some(fold),
513                Some(fold),
514                Some(fold),
515                Some(fold),
516                None,
517            ])
518            .into_array(),
519        );
520        let policies = VarBinViewArray::from_iter_nullable_str([
521            Some("earliest"),
522            Some("latest"),
523            Some("null"),
524            None,
525            Some("raise"),
526        ]);
527        let array =
528            StructArray::from_fields(&[("t", input.into_array()), ("p", policies.into_array())])?
529                .into_array();
530        let result = array.apply(&replace_time_zone(
531            get_item("t", root()),
532            get_item("p", root()),
533            ReplaceTimeZoneOptions {
534                time_zone: Some("America/New_York".into()),
535                null_on_non_existent: false,
536            },
537        ))?;
538
539        let expected = ExtensionArray::new(
540            timestamp(Some("America/New_York")),
541            PrimitiveArray::from_option_iter([
542                Some(fold + 4 * hour),
543                Some(fold + 5 * hour),
544                None,
545                None,
546                None,
547            ])
548            .into_array(),
549        );
550        let mut ctx = array_session().create_execution_ctx();
551        assert_arrays_eq!(result, expected, &mut ctx);
552        Ok(())
553    }
554}
vortex_array/scalar_fn/fns/replace_time_zone.rs

vortex_array/scalar_fn/fns/
replace_time_zone.rs
