Skip to content

Commit b768151

Browse files
aseyboldtaseyboldt
committed
feat: support datetime coordinates
1 parent 6e84e32 commit b768151

File tree

9 files changed

+156
-3
lines changed

9 files changed

+156
-3
lines changed

‎nuts-storable/src/lib.rs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,13 @@
11
use std::collections::HashMap;
22

3+
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4+
pub enum DateTimeUnit {
5+
Seconds,
6+
Milliseconds,
7+
Microseconds,
8+
Nanoseconds,
9+
}
10+
311
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
412
pub enum ItemType {
513
U64,
@@ -8,6 +16,8 @@ pub enum ItemType {
816
F32,
917
Bool,
1018
String,
19+
DateTime64(DateTimeUnit),
20+
TimeDelta64(DateTimeUnit),
1121
}
1222

1323
#[derive(Debug, Clone, PartialEq)]
@@ -18,6 +28,8 @@ pub enum Value {
1828
F32(Vec<f32>),
1929
Bool(Vec<bool>),
2030
ScalarString(String),
31+
DateTime64(DateTimeUnit, Vec<i64>),
32+
TimeDelta64(DateTimeUnit, Vec<i64>),
2133
ScalarU64(u64),
2234
ScalarI64(i64),
2335
ScalarF64(f64),

‎src/lib.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,7 @@ mod transform_adapt_strategy;
117117
mod transformed_hamiltonian;
118118

119119
pub use nuts_derive::Storable;
120-
pub use nuts_storable::{HasDims, ItemType, Storable, Value};
120+
pub use nuts_storable::{DateTimeUnit, HasDims, ItemType, Storable, Value};
121121

122122
pub use adapt_strategy::EuclideanAdaptOptions;
123123
pub use chain::Chain;

‎src/storage/arrow.rs‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,12 @@ impl ArrowBuilder {
3232
ItemType::I64 => Box::new(Int64Builder::with_capacity(capacity)),
3333
ItemType::U64 => Box::new(UInt64Builder::with_capacity(capacity)),
3434
ItemType::String => Box::new(StringBuilder::with_capacity(capacity, capacity)),
35+
ItemType::DateTime64(_) => {
36+
panic!("DateTime values not supported as values in arrow storage")
37+
}
38+
ItemType::TimeDelta64(_) => {
39+
panic!("TimeDelta values not supported as values in arrow storage")
40+
}
3541
};
3642

3743
if shape.is_empty() {
@@ -100,6 +106,12 @@ impl ArrowBuilder {
100106
string_builder.append_value(&item);
101107
}
102108
}
109+
Value::DateTime64(_, _) => {
110+
panic!("DateTime64 scalar values not supported in arrow storage")
111+
}
112+
Value::TimeDelta64(_, _) => {
113+
panic!("TimeDelta64 scalar values not supported in arrow storage")
114+
}
103115
},
104116
ArrowBuilder::Tensor(list_builder) => {
105117
match value {
@@ -154,6 +166,12 @@ impl ArrowBuilder {
154166
downcast_builder!(list_builder.values(), BooleanBuilder, ScalarBool)?
155167
.append_value(val);
156168
}
169+
Value::DateTime64(_, _) => {
170+
panic!("DateTime64 scalar values not supported in arrow storage")
171+
}
172+
Value::TimeDelta64(_, _) => {
173+
panic!("TimeDelta64 scalar values not supported in arrow storage")
174+
}
157175
}
158176
list_builder.append(true);
159177
}
@@ -211,6 +229,12 @@ fn item_type_to_arrow_type(item_type: ItemType) -> DataType {
211229
ItemType::I64 => DataType::Int64,
212230
ItemType::Bool => DataType::Boolean,
213231
ItemType::String => DataType::Utf8,
232+
ItemType::DateTime64(_) => {
233+
panic!("DateTime64 scalar values not supported in arrow storage")
234+
}
235+
ItemType::TimeDelta64(_) => {
236+
panic!("TimeDelta64 scalar values not supported in arrow storage")
237+
}
214238
}
215239
}
216240

‎src/storage/csv.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -220,6 +220,8 @@ impl CsvChainStorage {
220220
vec[0].clone()
221221
}
222222
}
223+
Value::DateTime64(_, _) => panic!("DateTime64 not supported in CSV output"),
224+
Value::TimeDelta64(_, _) => panic!("TimeDelta64 not supported in CSV output"),
223225
}
224226
}
225227

‎src/storage/hashmap.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ impl HashMapValue {
2626
ItemType::I64 => HashMapValue::I64(Vec::new()),
2727
ItemType::U64 => HashMapValue::U64(Vec::new()),
2828
ItemType::String => HashMapValue::String(Vec::new()),
29+
ItemType::DateTime64(_) | ItemType::TimeDelta64(_) => HashMapValue::I64(Vec::new()),
2930
}
3031
}
3132

@@ -45,6 +46,11 @@ impl HashMapValue {
4546
(HashMapValue::Bool(vec), Value::Bool(v)) => vec.extend(v),
4647
(HashMapValue::I64(vec), Value::I64(v)) => vec.extend(v),
4748

49+
(HashMapValue::String(vec), Value::Strings(v)) => vec.extend(v),
50+
(HashMapValue::String(vec), Value::ScalarString(v)) => vec.push(v),
51+
(HashMapValue::I64(vec), Value::DateTime64(_, v)) => vec.extend(v),
52+
(HashMapValue::I64(vec), Value::TimeDelta64(_, v)) => vec.extend(v),
53+
4854
_ => panic!("Mismatched item type"),
4955
}
5056
}

‎src/storage/ndarray.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,9 @@ impl NdarrayValue {
3030
ItemType::String => {
3131
NdarrayValue::String(ArrayD::from_elem(IxDyn(shape), String::new()))
3232
}
33+
ItemType::DateTime64(_) | ItemType::TimeDelta64(_) => {
34+
NdarrayValue::I64(ArrayD::zeros(IxDyn(shape)))
35+
}
3336
}
3437
}
3538

‎src/storage/zarr/async_impl.rs‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
use std::collections::HashMap;
22
use std::iter::once;
3+
use std::num::NonZero;
34
use std::sync::Arc;
45
use tokio::task::JoinHandle;
56

@@ -8,6 +9,7 @@ use nuts_storable::{ItemType, Value};
89
use zarrs::array::{ArrayBuilder, DataType, FillValue};
910
use zarrs::array_subset::ArraySubset;
1011
use zarrs::group::GroupBuilder;
12+
use zarrs::metadata_ext::data_type::NumpyTimeUnit;
1113
use zarrs::storage::{
1214
AsyncReadableWritableListableStorage, AsyncReadableWritableListableStorageTraits,
1315
};
@@ -140,6 +142,38 @@ async fn store_coords(
140142
&Value::I64(ref v) => (DataType::Int64, v.len(), FillValue::from(0i64)),
141143
&Value::Bool(ref v) => (DataType::Bool, v.len(), FillValue::from(false)),
142144
&Value::Strings(ref v) => (DataType::String, v.len(), FillValue::from("")),
145+
&Value::DateTime64(unit, ref v) => {
146+
let unit = match unit {
147+
nuts_storable::DateTimeUnit::Seconds => NumpyTimeUnit::Second,
148+
nuts_storable::DateTimeUnit::Milliseconds => NumpyTimeUnit::Millisecond,
149+
nuts_storable::DateTimeUnit::Microseconds => NumpyTimeUnit::Microsecond,
150+
nuts_storable::DateTimeUnit::Nanoseconds => NumpyTimeUnit::Nanosecond,
151+
};
152+
(
153+
DataType::NumpyDateTime64 {
154+
unit,
155+
scale_factor: NonZero::new(1).unwrap(),
156+
},
157+
v.len(),
158+
FillValue::from(0i64),
159+
)
160+
}
161+
&Value::TimeDelta64(unit, ref v) => {
162+
let unit = match unit {
163+
nuts_storable::DateTimeUnit::Seconds => NumpyTimeUnit::Second,
164+
nuts_storable::DateTimeUnit::Milliseconds => NumpyTimeUnit::Millisecond,
165+
nuts_storable::DateTimeUnit::Microseconds => NumpyTimeUnit::Microsecond,
166+
nuts_storable::DateTimeUnit::Nanoseconds => NumpyTimeUnit::Nanosecond,
167+
};
168+
(
169+
DataType::NumpyTimeDelta64 {
170+
unit,
171+
scale_factor: NonZero::new(1).unwrap(),
172+
},
173+
v.len(),
174+
FillValue::from(0i64),
175+
)
176+
}
143177
_ => panic!("Unsupported coordinate type for {}", name),
144178
};
145179
let name: &String = name;
@@ -179,6 +213,16 @@ async fn store_coords(
179213
.async_store_chunk_elements::<String>(&subset, v)
180214
.await?
181215
}
216+
&Value::DateTime64(_, ref data) => {
217+
coord_array
218+
.async_store_chunk_elements::<i64>(&subset, data)
219+
.await?
220+
}
221+
&Value::TimeDelta64(_, ref data) => {
222+
coord_array
223+
.async_store_chunk_elements::<i64>(&subset, data)
224+
.await?
225+
}
182226
_ => unreachable!(),
183227
}
184228
coord_array.async_store_metadata().await?;

‎src/storage/zarr/common.rs‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
1-
use std::collections::HashMap;
21
use std::mem::replace;
32
use std::sync::Arc;
3+
use std::{collections::HashMap, num::NonZero};
44

55
use anyhow::Result;
66
use nuts_storable::{ItemType, Value};
77
use zarrs::array::{Array, ArrayBuilder, DataType, FillValue};
8+
use zarrs::metadata_ext::data_type::NumpyTimeUnit;
89

910
/// Container for different types of sample values
1011
#[derive(Clone, Debug)]
@@ -51,6 +52,8 @@ impl SampleBuffer {
5152
ItemType::Bool => SampleBufferValue::Bool(Vec::with_capacity(chunk_size)),
5253
ItemType::I64 => SampleBufferValue::I64(Vec::with_capacity(chunk_size)),
5354
ItemType::String => panic!("String type not supported in SampleBuffer"),
55+
ItemType::DateTime64(_) => panic!("DateTime64 type not supported in SampleBuffer"),
56+
ItemType::TimeDelta64(_) => panic!("TimeDelta64 type not supported in SampleBuffer"),
5457
};
5558
Self {
5659
items: inner,
@@ -196,6 +199,24 @@ pub fn create_arrays<TStorage: ?Sized>(
196199
ItemType::I64 => DataType::Int64,
197200
ItemType::Bool => DataType::Bool,
198201
ItemType::String => DataType::String,
202+
ItemType::DateTime64(unit) => DataType::NumpyDateTime64 {
203+
unit: match unit {
204+
nuts_storable::DateTimeUnit::Seconds => NumpyTimeUnit::Second,
205+
nuts_storable::DateTimeUnit::Milliseconds => NumpyTimeUnit::Millisecond,
206+
nuts_storable::DateTimeUnit::Microseconds => NumpyTimeUnit::Microsecond,
207+
nuts_storable::DateTimeUnit::Nanoseconds => NumpyTimeUnit::Nanosecond,
208+
},
209+
scale_factor: NonZero::new(1).unwrap(),
210+
},
211+
ItemType::TimeDelta64(unit) => DataType::NumpyTimeDelta64 {
212+
unit: match unit {
213+
nuts_storable::DateTimeUnit::Seconds => NumpyTimeUnit::Second,
214+
nuts_storable::DateTimeUnit::Milliseconds => NumpyTimeUnit::Millisecond,
215+
nuts_storable::DateTimeUnit::Microseconds => NumpyTimeUnit::Microsecond,
216+
nuts_storable::DateTimeUnit::Nanoseconds => NumpyTimeUnit::Nanosecond,
217+
},
218+
scale_factor: NonZero::new(1).unwrap(),
219+
},
199220
};
200221
let fill_value = match item_type {
201222
ItemType::F64 => FillValue::from(f64::NAN),
@@ -204,6 +225,8 @@ pub fn create_arrays<TStorage: ?Sized>(
204225
ItemType::I64 => FillValue::from(0i64),
205226
ItemType::Bool => FillValue::from(false),
206227
ItemType::String => FillValue::from(""),
228+
ItemType::DateTime64(_) => FillValue::new_null(),
229+
ItemType::TimeDelta64(_) => FillValue::new_null(),
207230
};
208231
let grid: Vec<u64> = std::iter::once(1)
209232
.chain(std::iter::once(draw_chunk_size))

‎src/storage/zarr/sync_impl.rs‎

Lines changed: 40 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,14 @@
11
use std::collections::HashMap;
22
use std::iter::once;
3+
use std::num::NonZero;
34
use std::sync::Arc;
45

56
use anyhow::{Context, Result};
67
use nuts_storable::{ItemType, Value};
78
use zarrs::array::{ArrayBuilder, DataType, FillValue};
89
use zarrs::array_subset::ArraySubset;
910
use zarrs::group::GroupBuilder;
11+
use zarrs::metadata_ext::data_type::NumpyTimeUnit;
1012
use zarrs::storage::{ReadableWritableListableStorage, ReadableWritableListableStorageTraits};
1113

1214
use super::common::{Chunk, SampleBuffer, SampleBufferValue};
@@ -37,7 +39,38 @@ pub fn store_coords(
3739
&Value::I64(ref v) => (DataType::Int64, v.len(), FillValue::from(0i64)),
3840
&Value::Bool(ref v) => (DataType::Bool, v.len(), FillValue::from(false)),
3941
&Value::Strings(ref v) => (DataType::String, v.len(), FillValue::from("")),
40-
_ => panic!("Unsupported coordinate type for {}", name),
42+
&Value::DateTime64(ref unit, ref data) => (
43+
DataType::NumpyDateTime64 {
44+
unit: match unit {
45+
nuts_storable::DateTimeUnit::Seconds => NumpyTimeUnit::Second,
46+
nuts_storable::DateTimeUnit::Milliseconds => NumpyTimeUnit::Millisecond,
47+
nuts_storable::DateTimeUnit::Microseconds => NumpyTimeUnit::Microsecond,
48+
nuts_storable::DateTimeUnit::Nanoseconds => NumpyTimeUnit::Nanosecond,
49+
},
50+
scale_factor: NonZero::new(1).unwrap(),
51+
},
52+
data.len(),
53+
FillValue::new_null(),
54+
),
55+
&Value::TimeDelta64(ref unit, ref data) => (
56+
DataType::NumpyTimeDelta64 {
57+
unit: match unit {
58+
nuts_storable::DateTimeUnit::Seconds => NumpyTimeUnit::Second,
59+
nuts_storable::DateTimeUnit::Milliseconds => NumpyTimeUnit::Millisecond,
60+
nuts_storable::DateTimeUnit::Microseconds => NumpyTimeUnit::Microsecond,
61+
nuts_storable::DateTimeUnit::Nanoseconds => NumpyTimeUnit::Nanosecond,
62+
},
63+
scale_factor: NonZero::new(1).unwrap(),
64+
},
65+
data.len(),
66+
FillValue::new_null(),
67+
),
68+
_ => {
69+
return Err(anyhow::anyhow!(
70+
"Unsupported coordinate type for coordinate {}",
71+
name
72+
));
73+
}
4174
};
4275
let name: &String = name;
4376

@@ -53,6 +86,12 @@ pub fn store_coords(
5386
&Value::I64(ref v) => coord_array.store_chunk_elements::<i64>(&subset, v)?,
5487
&Value::Bool(ref v) => coord_array.store_chunk_elements::<bool>(&subset, v)?,
5588
&Value::Strings(ref v) => coord_array.store_chunk_elements::<String>(&subset, v)?,
89+
&Value::DateTime64(_, ref data) => {
90+
coord_array.store_chunk_elements::<i64>(&subset, data)?
91+
}
92+
&Value::TimeDelta64(_, ref data) => {
93+
coord_array.store_chunk_elements::<i64>(&subset, data)?
94+
}
5695
_ => unreachable!(),
5796
}
5897
coord_array.store_metadata()?;

0 commit comments

Comments
 (0)