Skip to content

Commit 49b1634

Browse files
authored
feat: teach TemporalMetadata about its relationship to DTypes (#5864)
I discovered this implicit mapping from TemporalMetadata to DType by trial-and-error and looking at arrays/datetime/mod.rs. AFAIK, this relationship between temporal type and bitwidth comes from [the arrow specification](https://github.com/apache/arrow/blob/7151dcd6faa6a507d99b33ecc797f3d6a74f45ae/format/Schema.fbs#L246-L277). For arrays, TemporalArray already addresses these issues [1], but for Vortex users who create or manipulate scalars, these methods are valuable. [1] Insofar as temporal array expects you to give it a valid PType (it panics if the width is wrong) but will construct the extension DType for you. --------- Signed-off-by: Daniel King <[email protected]>
1 parent 6a9f07a commit 49b1634

File tree

1 file changed

+57
-0
lines changed

1 file changed

+57
-0
lines changed

vortex-dtype/src/datetime/temporal.rs

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,12 @@ use vortex_error::vortex_bail;
1616
use vortex_error::vortex_err;
1717
use vortex_error::vortex_panic;
1818

19+
use crate::DType;
1920
use crate::ExtDType;
2021
use crate::ExtID;
2122
use crate::ExtMetadata;
23+
use crate::Nullability;
24+
use crate::PType;
2225
use crate::datetime::unit::TimeUnit;
2326

2427
/// ID for the Vortex time type.
@@ -120,6 +123,60 @@ impl TemporalMetadata {
120123
)),
121124
}
122125
}
126+
127+
/// The (only) allowed storage PType for this temporal type.
128+
///
129+
/// Each temporal extension type has exactly one allowed representation. This function returns
130+
/// that representation.
131+
pub fn storage_ptype(&self) -> VortexResult<PType> {
132+
let storage_ptype = match self {
133+
TemporalMetadata::Time(TimeUnit::Nanoseconds | TimeUnit::Microseconds) => PType::I64,
134+
TemporalMetadata::Time(TimeUnit::Milliseconds | TimeUnit::Seconds) => PType::I32,
135+
TemporalMetadata::Time(TimeUnit::Days) => {
136+
vortex_bail!("invalid TimeUnit days for vortex.time");
137+
}
138+
TemporalMetadata::Date(
139+
unit @ (TimeUnit::Nanoseconds | TimeUnit::Microseconds | TimeUnit::Seconds),
140+
) => vortex_bail!("invalid TimeUnit {unit} for vortex.date"),
141+
TemporalMetadata::Date(TimeUnit::Milliseconds) => PType::I64,
142+
TemporalMetadata::Date(TimeUnit::Days) => PType::I32,
143+
TemporalMetadata::Timestamp(
144+
TimeUnit::Nanoseconds
145+
| TimeUnit::Microseconds
146+
| TimeUnit::Milliseconds
147+
| TimeUnit::Seconds,
148+
..,
149+
) => PType::I64,
150+
TemporalMetadata::Timestamp(TimeUnit::Days, ..) => {
151+
vortex_bail!("invalid TimeUnit days for vortex.timestamp");
152+
}
153+
};
154+
155+
Ok(storage_ptype)
156+
}
157+
158+
/// The Extension (DType) ID for this temporal type.
159+
pub fn extension_id(&self) -> ExtID {
160+
match self {
161+
TemporalMetadata::Time(..) => TIME_ID.clone(),
162+
TemporalMetadata::Date(..) => DATE_ID.clone(),
163+
TemporalMetadata::Timestamp(..) => TIMESTAMP_ID.clone(),
164+
}
165+
}
166+
167+
/// The (only) allowed storage DType for this temporal type.
168+
///
169+
/// Each temporal extension type has exactly one allowed representation for a given
170+
/// nullability. This function returns that representation.
171+
pub fn dtype(self, nullability: Nullability) -> VortexResult<DType> {
172+
let extension_id = self.extension_id();
173+
let storage_ptype = self.storage_ptype()?;
174+
let metadata = ExtMetadata::from(self);
175+
let storage_dtype = DType::Primitive(storage_ptype, nullability);
176+
let ext_dtype = ExtDType::new(extension_id, Arc::new(storage_dtype), Some(metadata));
177+
178+
Ok(DType::Extension(Arc::new(ext_dtype)))
179+
}
123180
}
124181

125182
macro_rules! impl_temporal_metadata_try_from {

0 commit comments

Comments
 (0)