Skip to content

Commit 9535bbb

Browse files
gatesnAdamGS
gatesn
authored and
AdamGS
committed
More ergonomic TimestampOptions (#6245)
Fixes TemporalMetadata such that matching on unit, tz is easier Signed-off-by: Nicholas Gates <nick@nickgates.com>
1 parent 84dadd3 commit 9535bbb

8 files changed

Lines changed: 32 additions & 40 deletions

File tree

‎vortex-array/src/arrays/datetime/test.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -161,7 +161,7 @@ fn test_timestamp() {
161161
);
162162
assert_eq!(
163163
temporal_array.temporal_metadata(),
164-
TemporalMetadata::Timestamp(&TimestampOptions { unit, tz })
164+
TemporalMetadata::Timestamp((&unit, &tz))
165165
);
166166
}
167167
}

‎vortex-array/src/arrow/executor/temporal.rs‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@ use vortex_dtype::NativePType;
2424
use vortex_dtype::datetime::AnyTemporal;
2525
use vortex_dtype::datetime::TemporalMetadata;
2626
use vortex_dtype::datetime::TimeUnit;
27-
use vortex_dtype::datetime::TimestampOptions;
2827
use vortex_error::VortexResult;
2928
use vortex_error::vortex_bail;
3029
use vortex_error::vortex_ensure;
@@ -77,10 +76,7 @@ pub(super) fn to_arrow_temporal(
7776
DataType::Time64(ArrowTimeUnit::Nanosecond),
7877
) => to_temporal::<Time64NanosecondType>(array, ctx),
7978

80-
(
81-
TemporalMetadata::Timestamp(TimestampOptions { unit, tz }),
82-
DataType::Timestamp(arrow_unit, arrow_tz),
83-
) => {
79+
(TemporalMetadata::Timestamp((unit, tz)), DataType::Timestamp(arrow_unit, arrow_tz)) => {
8480
vortex_ensure!(
8581
tz == arrow_tz,
8682
"Cannot convert {} array to Arrow type {} due to timezone mismatch",

‎vortex-datafusion/src/convert/scalars.rs‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@ use vortex::dtype::arrow::FromArrowType;
1212
use vortex::dtype::datetime::AnyTemporal;
1313
use vortex::dtype::datetime::TemporalMetadata;
1414
use vortex::dtype::datetime::TimeUnit;
15-
use vortex::dtype::datetime::TimestampOptions;
1615
use vortex::dtype::half::f16;
1716
use vortex::error::VortexResult;
1817
use vortex::error::vortex_bail;
@@ -126,7 +125,7 @@ impl TryToDataFusion<ScalarValue> for Scalar {
126125
// temporal physical types.
127126
let pv = storage_scalar.as_primitive();
128127
match temporal {
129-
TemporalMetadata::Timestamp(TimestampOptions { unit, tz }) => match unit {
128+
TemporalMetadata::Timestamp((unit, tz)) => match unit {
130129
TimeUnit::Nanoseconds => {
131130
ScalarValue::TimestampNanosecond(pv.as_::<i64>(), tz.clone())
132131
}

‎vortex-dtype/src/arrow.rs‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,6 @@ use crate::datetime::TemporalMetadata;
4242
use crate::datetime::Time;
4343
use crate::datetime::TimeUnit;
4444
use crate::datetime::Timestamp;
45-
use crate::datetime::TimestampOptions;
4645

4746
/// Trait for converting Arrow types to Vortex types.
4847
pub trait FromArrowType<T>: Sized {
@@ -302,7 +301,7 @@ impl DType {
302301
// Try and match against the known extension DTypes.
303302
if let Some(temporal) = ext_dtype.metadata_opt::<AnyTemporal>() {
304303
return Ok(match temporal {
305-
TemporalMetadata::Timestamp(TimestampOptions { unit, tz }) => {
304+
TemporalMetadata::Timestamp((unit, tz)) => {
306305
DataType::Timestamp(ArrowTimeUnit::try_from(*unit)?, tz.clone())
307306
}
308307
TemporalMetadata::Date(unit) => match unit {

‎vortex-dtype/src/datetime/matcher.rs‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,14 @@
22
// SPDX-FileCopyrightText: Copyright the Vortex contributors
33

44
use std::fmt::Display;
5+
use std::sync::Arc;
56

67
use vortex_error::VortexResult;
78

89
use crate::datetime::Date;
910
use crate::datetime::Time;
11+
use crate::datetime::TimeUnit;
1012
use crate::datetime::Timestamp;
11-
use crate::datetime::TimestampOptions;
1213
use crate::extension::ExtDTypeRef;
1314
use crate::extension::ExtDTypeVTable;
1415
use crate::extension::Matcher;
@@ -21,7 +22,7 @@ impl Matcher for AnyTemporal {
2122

2223
fn try_match<'a>(item: &'a ExtDTypeRef) -> Option<Self::Match<'a>> {
2324
if let Some(opts) = item.metadata_opt::<Timestamp>() {
24-
return Some(TemporalMetadata::Timestamp(opts));
25+
return Some(TemporalMetadata::Timestamp((&opts.unit, &opts.tz)));
2526
}
2627
if let Some(opts) = item.metadata_opt::<Date>() {
2728
return Some(TemporalMetadata::Date(opts));
@@ -36,8 +37,8 @@ impl Matcher for AnyTemporal {
3637
/// Metadata for temporal extension data types.
3738
#[derive(Debug, PartialEq, Eq)]
3839
pub enum TemporalMetadata<'a> {
39-
/// Metadata for Timestamp dtypes
40-
Timestamp(&'a <Timestamp as ExtDTypeVTable>::Metadata),
40+
/// Metadata for Timestamp dtypes, a tuple of time unit and optional timezone.
41+
Timestamp((&'a TimeUnit, &'a Option<Arc<str>>)),
4142
/// Metadata for Date dtypes
4243
Date(&'a <Date as ExtDTypeVTable>::Metadata),
4344
/// Metadata for Time dtypes
@@ -48,11 +49,11 @@ pub enum TemporalMetadata<'a> {
4849
// Currently this is used largely to implement scalar display hacks.
4950
impl TemporalMetadata<'_> {
5051
/// Get the time unit of the temporal dtype.
51-
pub fn time_unit(&self) -> crate::datetime::TimeUnit {
52+
pub fn time_unit(&self) -> TimeUnit {
5253
match self {
5354
TemporalMetadata::Time(unit) => **unit,
5455
TemporalMetadata::Date(unit) => **unit,
55-
TemporalMetadata::Timestamp(opts) => opts.unit,
56+
TemporalMetadata::Timestamp((unit, _tz)) => **unit,
5657
}
5758
}
5859

@@ -65,7 +66,7 @@ impl TemporalMetadata<'_> {
6566
TemporalMetadata::Date(unit) => Ok(TemporalJiff::Date(
6667
jiff::civil::Date::new(1970, 1, 1)?.checked_add(unit.to_jiff_span(v)?)?,
6768
)),
68-
TemporalMetadata::Timestamp(TimestampOptions { unit, tz }) => match tz {
69+
TemporalMetadata::Timestamp((unit, tz)) => match tz {
6970
None => Ok(TemporalJiff::Unzoned(
7071
jiff::civil::DateTime::new(1970, 1, 1, 0, 0, 0, 0)?
7172
.checked_add(unit.to_jiff_span(v)?)?,

‎vortex-duckdb/src/convert/dtype.rs‎

Lines changed: 18 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -235,27 +235,25 @@ impl TryFrom<&DType> for LogicalType {
235235
};
236236

237237
match temporal {
238-
TemporalMetadata::Timestamp(ts) => match &ts.tz {
239-
None => match ts.unit {
240-
TimeUnit::Nanoseconds => DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP_NS,
241-
TimeUnit::Microseconds => DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP,
242-
TimeUnit::Milliseconds => DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP_MS,
243-
TimeUnit::Seconds => DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP_S,
244-
_ => vortex_bail!("Invalid TimeUnit {} for timestamp", ts.unit),
245-
},
246-
Some(tz) => {
247-
if tz.as_ref() != "UTC" {
248-
vortex_bail!("Invalid timezone for timestamp_tz {tz}, must be UTC");
249-
}
250-
if ts.unit != TimeUnit::Microseconds {
251-
vortex_bail!(
252-
"Invalid TimeUnit {} for timestamp_tz, must be Microseconds",
253-
ts.unit
254-
);
255-
}
256-
DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP_TZ
257-
}
238+
TemporalMetadata::Timestamp((unit, None)) => match unit {
239+
TimeUnit::Nanoseconds => DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP_NS,
240+
TimeUnit::Microseconds => DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP,
241+
TimeUnit::Milliseconds => DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP_MS,
242+
TimeUnit::Seconds => DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP_S,
243+
_ => vortex_bail!("Invalid TimeUnit {} for timestamp", unit),
258244
},
245+
TemporalMetadata::Timestamp((unit, Some(tz))) => {
246+
if tz.as_ref() != "UTC" {
247+
vortex_bail!("Invalid timezone for timestamp_tz {tz}, must be UTC");
248+
}
249+
if unit != &TimeUnit::Microseconds {
250+
vortex_bail!(
251+
"Invalid TimeUnit {} for timestamp_tz, must be Microseconds",
252+
unit
253+
);
254+
}
255+
DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP_TZ
256+
}
259257
TemporalMetadata::Date(unit) => match unit {
260258
TimeUnit::Days => DUCKDB_TYPE::DUCKDB_TYPE_DATE,
261259
_ => vortex_bail!("Invalid TimeUnit {} for date", unit),

‎vortex-duckdb/src/convert/scalar.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -182,7 +182,7 @@ impl ToDuckDBScalar for ExtScalar<'_> {
182182
};
183183

184184
Ok(match temporal {
185-
TemporalMetadata::Timestamp(TimestampOptions { unit, tz }) => {
185+
TemporalMetadata::Timestamp((unit, tz)) => {
186186
if let Some(tz) = tz.as_ref() {
187187
if tz.as_ref() != "UTC" {
188188
// TODO(ngates): we should convert into UTC as DuckDB does internally.

‎vortex-scalar/src/arrow/mod.rs‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ use vortex_dtype::PType;
1010
use vortex_dtype::datetime::AnyTemporal;
1111
use vortex_dtype::datetime::TemporalMetadata;
1212
use vortex_dtype::datetime::TimeUnit;
13-
use vortex_dtype::datetime::TimestampOptions;
1413
use vortex_error::VortexError;
1514
use vortex_error::vortex_bail;
1615
use vortex_error::vortex_err;
@@ -134,7 +133,7 @@ impl TryFrom<&Scalar> for Arc<dyn Datum> {
134133
.ok_or_else(|| vortex_err!("Expected primitive scalar"))?;
135134

136135
match temporal {
137-
TemporalMetadata::Timestamp(TimestampOptions { unit, tz }) => {
136+
TemporalMetadata::Timestamp((unit, tz)) => {
138137
let value = primitive.as_::<i64>();
139138
match unit {
140139
TimeUnit::Nanoseconds => {

0 commit comments

Comments
 (0)