quiver_types 0.1.0

Core types for the quiver crate
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148

//! [`Timestamp`]: a logical type for columns of points in time.
//!
//! Each element is an `i64` counting time units since the Unix epoch
//! (1970-01-01 00:00:00), e.g. `TimestampNanosecond<Utc>` for nanoseconds in UTC.
//! Stored as the [`arrow::array::TimestampNanosecondArray`] family
//! ([`DataType::Timestamp`]).
//!
//! The time unit and the (optional) timezone are part of the type,
//! via the [`TimeUnitSpec`] and [`TimezoneSpec`] marker types.

use std::marker::PhantomData;

use arrow::array::{Array, ArrayRef};
use arrow::datatypes::DataType;

use crate::datatype::{
    ColumnError, Datatype, InfallibleBuild, PrimitiveDatatype, RefDatatype, downcast_array,
};

/// Marker for an arrow `Timestamp` column, e.g. `Timestamp<Nanosecond, Utc>`.
///
/// The values are raw `i64` ticks in the given [`TimeUnitSpec`],
/// counted from the unix epoch.
///
/// The timezone defaults to [`NoTimezone`]. Note that timezones are matched
/// *exactly*: a column declared `Timestamp<Nanosecond, Utc>` ("UTC") will not
/// accept an array with the timezone "+00:00".
///
/// This type is never instantiated — it only appears as a type parameter.
pub struct Timestamp<U, Z = NoTimezone> {
    _marker: PhantomData<fn() -> (U, Z)>,
}

/// A [`Timestamp`]/[`Duration`](crate::Duration) time unit:
/// [`Second`], [`Millisecond`], [`Microsecond`], or [`Nanosecond`].
pub trait TimeUnitSpec {
    /// The corresponding arrow timestamp type, e.g. `TimestampNanosecondType`.
    type TimestampType: arrow::datatypes::ArrowTimestampType;

    /// The corresponding arrow duration type, e.g. `DurationNanosecondType`.
    type DurationType: arrow::datatypes::ArrowPrimitiveType<Native = i64>;
}

pub struct Second;
pub struct Millisecond;
pub struct Microsecond;
pub struct Nanosecond;

impl TimeUnitSpec for Second {
    type TimestampType = arrow::datatypes::TimestampSecondType;
    type DurationType = arrow::datatypes::DurationSecondType;
}
impl TimeUnitSpec for Millisecond {
    type TimestampType = arrow::datatypes::TimestampMillisecondType;
    type DurationType = arrow::datatypes::DurationMillisecondType;
}
impl TimeUnitSpec for Microsecond {
    type TimestampType = arrow::datatypes::TimestampMicrosecondType;
    type DurationType = arrow::datatypes::DurationMicrosecondType;
}
impl TimeUnitSpec for Nanosecond {
    type TimestampType = arrow::datatypes::TimestampNanosecondType;
    type DurationType = arrow::datatypes::DurationNanosecondType;
}

/// The timezone of a [`Timestamp`]: [`NoTimezone`], [`Utc`], or your own marker type.
pub trait TimezoneSpec {
    /// E.g. `Some("UTC")`, `Some("+02:00")`, or `None` for timezone-naive timestamps.
    fn timezone() -> Option<std::sync::Arc<str>>;
}

/// Timezone-naive timestamps.
pub struct NoTimezone;

impl TimezoneSpec for NoTimezone {
    fn timezone() -> Option<std::sync::Arc<str>> {
        None
    }
}

/// The "UTC" timezone.
pub struct Utc;

impl TimezoneSpec for Utc {
    fn timezone() -> Option<std::sync::Arc<str>> {
        Some("UTC".into())
    }
}

impl<U: TimeUnitSpec + 'static, Z: TimezoneSpec + 'static> Datatype for Timestamp<U, Z> {
    type Typed = arrow::array::PrimitiveArray<U::TimestampType>;
    type Value<'a>
        = i64
    where
        Self: 'a;
    type Owned = i64;

    fn datatype() -> DataType {
        DataType::Timestamp(
            <U::TimestampType as arrow::datatypes::ArrowTimestampType>::UNIT,
            Z::timezone(),
        )
    }

    fn downcast(array: &dyn Array) -> Result<Self::Typed, ColumnError> {
        downcast_array::<Self::Typed>(array)
    }

    fn is_null(typed: &Self::Typed, index: usize) -> bool {
        typed.is_null(index)
    }

    fn value(typed: &Self::Typed, index: usize) -> Self::Value<'_> {
        typed.value(index)
    }

    fn build(values: impl Iterator<Item = Option<Self::Owned>>) -> Result<ArrayRef, ColumnError> {
        let array: arrow::array::PrimitiveArray<U::TimestampType> = values.collect();
        Ok(std::sync::Arc::new(array.with_timezone_opt(Z::timezone())))
    }

    fn to_owned_value(value: Self::Value<'_>) -> Self::Owned {
        value
    }
}

pub type TimestampSecond<Z = NoTimezone> = Timestamp<Second, Z>;
pub type TimestampMillisecond<Z = NoTimezone> = Timestamp<Millisecond, Z>;
pub type TimestampMicrosecond<Z = NoTimezone> = Timestamp<Microsecond, Z>;
pub type TimestampNanosecond<Z = NoTimezone> = Timestamp<Nanosecond, Z>;

impl<U: TimeUnitSpec + 'static, Z: TimezoneSpec + 'static> InfallibleBuild for Timestamp<U, Z> {}

impl<U: TimeUnitSpec + 'static, Z: TimezoneSpec + 'static> PrimitiveDatatype for Timestamp<U, Z> {
    type Native = i64;

    fn values(typed: &Self::Typed) -> &[i64] {
        typed.values()
    }
}

impl<U: TimeUnitSpec + 'static, Z: TimezoneSpec + 'static> RefDatatype for Timestamp<U, Z> {
    type Ref = i64;

    fn value_ref(typed: &Self::Typed, index: usize) -> &i64 {
        &typed.values()[index]
    }
}