arrow_schema/extension/canonical/
timestamp_with_offset.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! Timestamp with an offset in minutes
6//!
7//! <https://arrow.apache.org/docs/format/CanonicalExtensions.html#timestamp-with-offset>
8
9use crate::{ArrowError, DataType, extension::ExtensionType};
10
11/// The extension type for `TimestampWithOffset`.
12///
13/// Extension name: `arrow.timestamp_with_offset`.
14///
15/// This type represents a timestamp column that stores potentially different timezone offsets per
16/// value. The timestamp is stored in UTC alongside the original timezone offset in minutes. This
17/// extension type is intended to be compatible with ANSI SQL's `TIMESTAMP WITH TIME ZONE`, which
18/// is supported by multiple database engines.
19///
20/// The storage type of the extension is a `Struct` with 2 fields, in order: - `timestamp`: a
21/// non-nullable `Timestamp(time_unit, "UTC")`, where `time_unit` is any Arrow `TimeUnit` (s, ms,
22/// us or ns). - `offset_minutes`: a non-nullable signed 16-bit integer (`Int16`) representing the
23/// offset in minutes from the UTC timezone. Negative offsets represent time zones west of UTC,
24/// while positive offsets represent east. Offsets normally range from -779 (-12:59) to +780
25/// (+13:00).
26///
27/// This type has no type parameters.
28///
29/// Metadata is either empty or an empty string.
30///
31/// It is also *permissible* for the `offset_minutes` field to be dictionary-encoded with a
32/// preferred (*but not required*) index type of `int8`, or run-end-encoded with a preferred (*but
33/// not required*) runs type of `int8`.
34///
35/// It's worth noting that the data source needs to resolve timezone strings such as `UTC` or
36/// `Americas/Los_Angeles` into an offset in minutes in order to construct a `TimestampWithOffset`.
37/// This makes `TimestampWithOffset` type "lossy" in the sense that any original "unresolved"
38/// timezone string gets lost in this conversion. It's a tradeoff for optimizing the row
39/// representation and simplifying the client code, which does not need to know how to convert from
40/// timezone string to its corresponding offset in minutes.
41///
42/// <https://arrow.apache.org/docs/format/CanonicalExtensions.html#timestamp-with-offset>
43#[derive(Debug, Default, Clone, Copy, PartialEq)]
44pub struct TimestampWithOffset;
45
46const TIMESTAMP_FIELD_NAME: &str = "timestamp";
47const OFFSET_FIELD_NAME: &str = "offset_minutes";
48
49impl ExtensionType for TimestampWithOffset {
50    const NAME: &'static str = "arrow.timestamp_with_offset";
51
52    type Metadata = ();
53
54    fn metadata(&self) -> &Self::Metadata {
55        &()
56    }
57
58    fn serialize_metadata(&self) -> Option<String> {
59        None
60    }
61
62    fn deserialize_metadata(metadata: Option<&str>) -> Result<Self::Metadata, ArrowError> {
63        metadata.map_or_else(
64            || Ok(()),
65            |v| {
66                if !v.is_empty() {
67                    Err(ArrowError::InvalidArgumentError(
68                        "TimestampWithOffset extension type expects no metadata".to_owned(),
69                    ))
70                } else {
71                    Ok(())
72                }
73            },
74        )
75    }
76
77    fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
78        let ok = match data_type {
79            DataType::Struct(fields) => match fields.len() {
80                2 => {
81                    let maybe_timestamp = fields.first().unwrap();
82                    let maybe_offset = fields.get(1).unwrap();
83
84                    let timestamp_type_ok = matches!(maybe_timestamp.data_type(), DataType::Timestamp(_, tz) if {
85                        match tz {
86                            Some(tz) => {
87                                tz.as_ref() == "UTC"
88                            },
89                            None => false
90                        }
91                    });
92
93                    let offset_type_ok = match maybe_offset.data_type() {
94                        DataType::Int16 => true,
95                        DataType::Dictionary(key_type, value_type) => {
96                            key_type.is_dictionary_key_type()
97                                && matches!(value_type.as_ref(), DataType::Int16)
98                        }
99                        DataType::RunEndEncoded(run_ends, values) => {
100                            run_ends.data_type().is_run_ends_type()
101                                && matches!(values.data_type(), DataType::Int16)
102                        }
103                        _ => false,
104                    };
105
106                    maybe_timestamp.name() == TIMESTAMP_FIELD_NAME
107                        && timestamp_type_ok
108                        && !maybe_timestamp.is_nullable()
109                        && maybe_offset.name() == OFFSET_FIELD_NAME
110                        && offset_type_ok
111                        && !maybe_offset.is_nullable()
112                }
113                _ => false,
114            },
115            _ => false,
116        };
117
118        match ok {
119            true => Ok(()),
120            false => Err(ArrowError::InvalidArgumentError(format!(
121                "TimestampWithOffset data type mismatch, expected Struct(\"timestamp\": Timestamp(_, Some(\"UTC\")), \"offset_minutes\": Int16), found {data_type}"
122            ))),
123        }
124    }
125
126    fn try_new(data_type: &DataType, _metadata: Self::Metadata) -> Result<Self, ArrowError> {
127        Self.supports_data_type(data_type).map(|()| Self)
128    }
129
130    fn validate(data_type: &DataType, _metadata: Self::Metadata) -> Result<(), ArrowError> {
131        Self.supports_data_type(data_type)
132    }
133}
134
135#[cfg(test)]
136mod tests {
137    use std::sync::Arc;
138
139    #[cfg(feature = "canonical_extension_types")]
140    use crate::extension::CanonicalExtensionType;
141    use crate::{
142        Field, Fields, TimeUnit,
143        extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY},
144    };
145
146    use super::*;
147
148    fn make_valid_field_primitive(time_unit: TimeUnit) -> Field {
149        Field::new(
150            "",
151            DataType::Struct(Fields::from_iter([
152                Field::new(
153                    TIMESTAMP_FIELD_NAME,
154                    DataType::Timestamp(time_unit, Some("UTC".into())),
155                    false,
156                ),
157                Field::new(OFFSET_FIELD_NAME, DataType::Int16, false),
158            ])),
159            false,
160        )
161    }
162
163    fn make_valid_field_dict_encoded(time_unit: TimeUnit, key_type: DataType) -> Field {
164        assert!(key_type.is_dictionary_key_type());
165
166        Field::new(
167            "",
168            DataType::Struct(Fields::from_iter([
169                Field::new(
170                    TIMESTAMP_FIELD_NAME,
171                    DataType::Timestamp(time_unit, Some("UTC".into())),
172                    false,
173                ),
174                Field::new(
175                    OFFSET_FIELD_NAME,
176                    DataType::Dictionary(Box::new(key_type), Box::new(DataType::Int16)),
177                    false,
178                ),
179            ])),
180            false,
181        )
182    }
183
184    fn make_valid_field_run_end_encoded(time_unit: TimeUnit, run_ends_type: DataType) -> Field {
185        assert!(run_ends_type.is_run_ends_type());
186        Field::new(
187            "",
188            DataType::Struct(Fields::from_iter([
189                Field::new(
190                    TIMESTAMP_FIELD_NAME,
191                    DataType::Timestamp(time_unit, Some("UTC".into())),
192                    false,
193                ),
194                Field::new(
195                    OFFSET_FIELD_NAME,
196                    DataType::RunEndEncoded(
197                        Arc::new(Field::new(
198                            Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
199                            run_ends_type,
200                            false,
201                        )),
202                        Arc::new(Field::new(
203                            Field::REE_VALUES_FIELD_DEFAULT_NAME,
204                            DataType::Int16,
205                            false,
206                        )),
207                    ),
208                    false,
209                ),
210            ])),
211            false,
212        )
213    }
214
215    #[test]
216    fn valid_primitive_offsets() -> Result<(), ArrowError> {
217        let time_units = [
218            TimeUnit::Second,
219            TimeUnit::Millisecond,
220            TimeUnit::Microsecond,
221            TimeUnit::Nanosecond,
222        ];
223
224        for time_unit in time_units {
225            let mut field = make_valid_field_primitive(time_unit);
226            field.try_with_extension_type(TimestampWithOffset)?;
227            field.try_extension_type::<TimestampWithOffset>()?;
228            #[cfg(feature = "canonical_extension_types")]
229            assert_eq!(
230                field.try_canonical_extension_type()?,
231                CanonicalExtensionType::TimestampWithOffset(TimestampWithOffset)
232            );
233        }
234
235        Ok(())
236    }
237
238    #[test]
239    fn valid_dict_encoded_offsets() -> Result<(), ArrowError> {
240        let time_units = [
241            TimeUnit::Second,
242            TimeUnit::Millisecond,
243            TimeUnit::Microsecond,
244            TimeUnit::Nanosecond,
245        ];
246
247        let key_types = [
248            DataType::UInt8,
249            DataType::UInt16,
250            DataType::UInt32,
251            DataType::UInt64,
252            DataType::Int8,
253            DataType::Int16,
254            DataType::Int32,
255            DataType::Int64,
256        ];
257
258        for time_unit in time_units {
259            for key_type in &key_types {
260                let mut field = make_valid_field_dict_encoded(time_unit, key_type.clone());
261                field.try_with_extension_type(TimestampWithOffset)?;
262                field.try_extension_type::<TimestampWithOffset>()?;
263                #[cfg(feature = "canonical_extension_types")]
264                assert_eq!(
265                    field.try_canonical_extension_type()?,
266                    CanonicalExtensionType::TimestampWithOffset(TimestampWithOffset)
267                );
268            }
269        }
270
271        Ok(())
272    }
273
274    #[test]
275    fn valid_run_end_encoded_offsets() -> Result<(), ArrowError> {
276        let time_units = [
277            TimeUnit::Second,
278            TimeUnit::Millisecond,
279            TimeUnit::Microsecond,
280            TimeUnit::Nanosecond,
281        ];
282
283        let run_ends_types = [DataType::Int16, DataType::Int32, DataType::Int64];
284
285        for time_unit in time_units {
286            for run_ends_type in &run_ends_types {
287                let mut field = make_valid_field_run_end_encoded(time_unit, run_ends_type.clone());
288                field.try_with_extension_type(TimestampWithOffset)?;
289                field.try_extension_type::<TimestampWithOffset>()?;
290                #[cfg(feature = "canonical_extension_types")]
291                assert_eq!(
292                    field.try_canonical_extension_type()?,
293                    CanonicalExtensionType::TimestampWithOffset(TimestampWithOffset)
294                );
295            }
296        }
297
298        Ok(())
299    }
300
301    #[test]
302    #[should_panic(expected = "Extension type name missing")]
303    fn missing_name() {
304        let field = make_valid_field_primitive(TimeUnit::Second)
305            .with_metadata([(EXTENSION_TYPE_METADATA_KEY, "")]);
306        field.extension_type::<TimestampWithOffset>();
307    }
308
309    #[test]
310    #[should_panic(
311        expected = "expected Struct(\"timestamp\": Timestamp(_, Some(\"UTC\")), \"offset_minutes\": Int16), found Boolean"
312    )]
313    fn invalid_type_top_level() {
314        Field::new("", DataType::Boolean, false).with_extension_type(TimestampWithOffset);
315    }
316
317    #[test]
318    #[should_panic(
319        expected = "expected Struct(\"timestamp\": Timestamp(_, Some(\"UTC\")), \"offset_minutes\": Int16), found Struct"
320    )]
321    fn invalid_type_struct_field_count() {
322        let data_type =
323            DataType::Struct(Fields::from_iter([Field::new("", DataType::Int16, false)]));
324        Field::new("", data_type, false).with_extension_type(TimestampWithOffset);
325    }
326
327    #[test]
328    #[should_panic(
329        expected = "expected Struct(\"timestamp\": Timestamp(_, Some(\"UTC\")), \"offset_minutes\": Int16), found Struct"
330    )]
331    fn invalid_type_wrong_timestamp_type() {
332        let data_type = DataType::Struct(Fields::from_iter([
333            Field::new(TIMESTAMP_FIELD_NAME, DataType::Int16, false),
334            Field::new(OFFSET_FIELD_NAME, DataType::Int16, false),
335        ]));
336        Field::new("", data_type, false).with_extension_type(TimestampWithOffset);
337    }
338
339    #[test]
340    #[should_panic(
341        expected = "expected Struct(\"timestamp\": Timestamp(_, Some(\"UTC\")), \"offset_minutes\": Int16), found Struct"
342    )]
343    fn invalid_type_wrong_offset_type() {
344        let data_type = DataType::Struct(Fields::from_iter([
345            Field::new(
346                TIMESTAMP_FIELD_NAME,
347                DataType::Timestamp(TimeUnit::Second, Some("UTC".into())),
348                false,
349            ),
350            Field::new(OFFSET_FIELD_NAME, DataType::UInt64, false),
351        ]));
352        Field::new("", data_type, false).with_extension_type(TimestampWithOffset);
353    }
354
355    #[test]
356    #[should_panic(
357        expected = "expected Struct(\"timestamp\": Timestamp(_, Some(\"UTC\")), \"offset_minutes\": Int16), found Struct"
358    )]
359    fn invalid_type_wrong_offset_key_dict_encoded() {
360        let data_type = DataType::Struct(Fields::from_iter([
361            Field::new(
362                TIMESTAMP_FIELD_NAME,
363                DataType::Timestamp(TimeUnit::Second, Some("UTC".into())),
364                false,
365            ),
366            Field::new(
367                OFFSET_FIELD_NAME,
368                DataType::Dictionary(Box::new(DataType::Boolean), Box::new(DataType::Int16)),
369                false,
370            ),
371        ]));
372        Field::new("", data_type, false).with_extension_type(TimestampWithOffset);
373    }
374
375    #[test]
376    #[should_panic(
377        expected = "expected Struct(\"timestamp\": Timestamp(_, Some(\"UTC\")), \"offset_minutes\": Int16), found Struct"
378    )]
379    fn invalid_type_wrong_offset_value_dict_encoded() {
380        let data_type = DataType::Struct(Fields::from_iter([
381            Field::new(
382                TIMESTAMP_FIELD_NAME,
383                DataType::Timestamp(TimeUnit::Second, Some("UTC".into())),
384                false,
385            ),
386            Field::new(
387                OFFSET_FIELD_NAME,
388                DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Int32)),
389                false,
390            ),
391        ]));
392        Field::new("", data_type, false).with_extension_type(TimestampWithOffset);
393    }
394
395    #[test]
396    #[should_panic(
397        expected = "expected Struct(\"timestamp\": Timestamp(_, Some(\"UTC\")), \"offset_minutes\": Int16), found Struct"
398    )]
399    fn invalid_type_wrong_run_ends_run_end_encoded() {
400        let data_type = DataType::Struct(Fields::from_iter([
401            Field::new(
402                TIMESTAMP_FIELD_NAME,
403                DataType::Timestamp(TimeUnit::Second, Some("UTC".into())),
404                false,
405            ),
406            Field::new(
407                OFFSET_FIELD_NAME,
408                DataType::RunEndEncoded(
409                    Arc::new(Field::new(
410                        Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
411                        DataType::Boolean,
412                        false,
413                    )),
414                    Arc::new(Field::new(
415                        Field::REE_VALUES_FIELD_DEFAULT_NAME,
416                        DataType::Int16,
417                        false,
418                    )),
419                ),
420                false,
421            ),
422        ]));
423        Field::new("", data_type, false).with_extension_type(TimestampWithOffset);
424    }
425
426    #[test]
427    #[should_panic(
428        expected = "expected Struct(\"timestamp\": Timestamp(_, Some(\"UTC\")), \"offset_minutes\": Int16), found Struct"
429    )]
430    fn invalid_type_wrong_values_run_end_encoded() {
431        let data_type = DataType::Struct(Fields::from_iter([
432            Field::new(
433                TIMESTAMP_FIELD_NAME,
434                DataType::Timestamp(TimeUnit::Second, Some("UTC".into())),
435                false,
436            ),
437            Field::new(
438                OFFSET_FIELD_NAME,
439                DataType::RunEndEncoded(
440                    Arc::new(Field::new(
441                        Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
442                        DataType::UInt16,
443                        false,
444                    )),
445                    Arc::new(Field::new(
446                        Field::REE_VALUES_FIELD_DEFAULT_NAME,
447                        DataType::Int32,
448                        false,
449                    )),
450                ),
451                false,
452            ),
453        ]));
454        Field::new("", data_type, false).with_extension_type(TimestampWithOffset);
455    }
456
457    #[test]
458    #[should_panic(
459        expected = "expected Struct(\"timestamp\": Timestamp(_, Some(\"UTC\")), \"offset_minutes\": Int16), found Struct"
460    )]
461    fn invalid_type_nullable_timestamp() {
462        let data_type = DataType::Struct(Fields::from_iter([
463            Field::new(
464                TIMESTAMP_FIELD_NAME,
465                DataType::Timestamp(TimeUnit::Second, Some("UTC".into())),
466                true,
467            ),
468            Field::new(OFFSET_FIELD_NAME, DataType::Int16, false),
469        ]));
470        Field::new("", data_type, false).with_extension_type(TimestampWithOffset);
471    }
472
473    #[test]
474    #[should_panic(
475        expected = "expected Struct(\"timestamp\": Timestamp(_, Some(\"UTC\")), \"offset_minutes\": Int16), found Struct"
476    )]
477    fn invalid_type_nullable_offset() {
478        let data_type = DataType::Struct(Fields::from_iter([
479            Field::new(
480                TIMESTAMP_FIELD_NAME,
481                DataType::Timestamp(TimeUnit::Second, Some("UTC".into())),
482                false,
483            ),
484            Field::new(OFFSET_FIELD_NAME, DataType::Int16, true),
485        ]));
486        Field::new("", data_type, false).with_extension_type(TimestampWithOffset);
487    }
488
489    #[test]
490    #[should_panic(
491        expected = "expected Struct(\"timestamp\": Timestamp(_, Some(\"UTC\")), \"offset_minutes\": Int16), found Struct"
492    )]
493    fn invalid_type_no_timezone() {
494        let data_type = DataType::Struct(Fields::from_iter([
495            Field::new(
496                TIMESTAMP_FIELD_NAME,
497                DataType::Timestamp(TimeUnit::Second, None),
498                false,
499            ),
500            Field::new(OFFSET_FIELD_NAME, DataType::Int16, false),
501        ]));
502        Field::new("", data_type, false).with_extension_type(TimestampWithOffset);
503    }
504
505    #[test]
506    #[should_panic(
507        expected = "expected Struct(\"timestamp\": Timestamp(_, Some(\"UTC\")), \"offset_minutes\": Int16), found Struct"
508    )]
509    fn invalid_type_wrong_timezone() {
510        let data_type = DataType::Struct(Fields::from_iter([
511            Field::new(
512                TIMESTAMP_FIELD_NAME,
513                DataType::Timestamp(TimeUnit::Second, Some("Americas/Sao_Paulo".into())),
514                false,
515            ),
516            Field::new(OFFSET_FIELD_NAME, DataType::Int16, false),
517        ]));
518        Field::new("", data_type, false).with_extension_type(TimestampWithOffset);
519    }
520
521    #[test]
522    fn no_metadata() {
523        let field = make_valid_field_primitive(TimeUnit::Second)
524            .with_metadata([(EXTENSION_TYPE_NAME_KEY, TimestampWithOffset::NAME)]);
525        field.extension_type::<TimestampWithOffset>();
526    }
527
528    #[test]
529    fn empty_metadata() {
530        let field = make_valid_field_primitive(TimeUnit::Second).with_metadata([
531            (EXTENSION_TYPE_NAME_KEY, TimestampWithOffset::NAME),
532            (EXTENSION_TYPE_METADATA_KEY, ""),
533        ]);
534        field.extension_type::<TimestampWithOffset>();
535    }
536}
arrow_schema/extension/canonical/timestamp_with_offset.rs

arrow_schema/extension/canonical/
timestamp_with_offset.rs
