diff --git a/arrow2_convert/src/deserialize.rs b/arrow2_convert/src/deserialize.rs index d8ea29d..a9cd255 100644 --- a/arrow2_convert/src/deserialize.rs +++ b/arrow2_convert/src/deserialize.rs @@ -1,7 +1,7 @@ //! Implementation and traits for deserializing from Arrow. use arrow2::{array::*, buffer::Buffer, types::NativeType}; -use chrono::{NaiveDate, NaiveDateTime}; +use chrono::{NaiveDate, NaiveDateTime, NaiveTime}; use crate::field::*; @@ -174,6 +174,15 @@ impl ArrowDeserialize for NaiveDate { } } +impl ArrowDeserialize for NaiveTime { + type ArrayType = PrimitiveArray; + + #[inline] + fn arrow_deserialize(v: Option<&i32>) -> Option { + v.map(|t| arrow2::temporal_conversions::time32s_to_time(*t)) + } +} + /// Iterator for for [`BufferBinaryArray`] pub struct BufferBinaryArrayIter<'a> { index: usize, diff --git a/arrow2_convert/src/field.rs b/arrow2_convert/src/field.rs index c182233..5a0f8d3 100644 --- a/arrow2_convert/src/field.rs +++ b/arrow2_convert/src/field.rs @@ -2,17 +2,17 @@ use arrow2::{ buffer::Buffer, - datatypes::{DataType, Field}, + datatypes::{DataType, Field, TimeUnit}, types::NativeType, }; -use chrono::{NaiveDate, NaiveDateTime}; +use chrono::{NaiveDate, NaiveDateTime, NaiveTime}; /// Trait implemented by all types that can be used as an Arrow field. /// /// Implementations are provided for types already supported by the arrow2 crate: /// - numeric types: [`u8`], [`u16`], [`u32`], [`u64`], [`i8`], [`i16`], [`i32`], [`i128`], [`i64`], [`f32`], [`f64`], /// - other types: [`bool`], [`String`] -/// - temporal types: [`chrono::NaiveDate`], [`chrono::NaiveDateTime`] +/// - temporal types: [`chrono::NaiveDate`], [`chrono::NaiveTime`], [`chrono::NaiveDateTime`] /// /// Custom implementations can be provided for other types. /// @@ -174,6 +174,15 @@ impl ArrowField for NaiveDate { } } +impl ArrowField for NaiveTime { + type Type = Self; + + #[inline] + fn data_type() -> arrow2::datatypes::DataType { + arrow2::datatypes::DataType::Time32(TimeUnit::Second) + } +} + impl ArrowField for Buffer { type Type = Self; @@ -281,6 +290,7 @@ arrow_enable_vec_for_type!(LargeString); arrow_enable_vec_for_type!(bool); arrow_enable_vec_for_type!(NaiveDateTime); arrow_enable_vec_for_type!(NaiveDate); +arrow_enable_vec_for_type!(NaiveTime); arrow_enable_vec_for_type!(Vec); arrow_enable_vec_for_type!(Buffer); arrow_enable_vec_for_type!(LargeBinary); diff --git a/arrow2_convert/src/serialize.rs b/arrow2_convert/src/serialize.rs index 3460445..e68dd73 100644 --- a/arrow2_convert/src/serialize.rs +++ b/arrow2_convert/src/serialize.rs @@ -1,14 +1,13 @@ //! Implementation and traits for serializing to Arrow. +use crate::field::*; use arrow2::array::*; use arrow2::chunk::Chunk; use arrow2::types::NativeType; use arrow2::{array::Array, buffer::Buffer}; -use chrono::{NaiveDate, NaiveDateTime}; +use chrono::{NaiveDate, NaiveDateTime, NaiveTime}; use std::sync::Arc; -use crate::field::*; - /// Trait that is implemented by all types that are serializable to Arrow. /// /// Implementations are provided for all built-in arrow types as well as Vec, and Option @@ -180,6 +179,20 @@ impl ArrowSerialize for NaiveDate { } } +impl ArrowSerialize for NaiveTime { + type MutableArrayType = MutablePrimitiveArray; + + #[inline] + fn new_array() -> Self::MutableArrayType { + Self::MutableArrayType::from(::data_type()) + } + + #[inline] + fn arrow_serialize(v: &Self, array: &mut Self::MutableArrayType) -> arrow2::error::Result<()> { + array.try_push(Some(chrono::Timelike::num_seconds_from_midnight(v) as i32)) + } +} + impl ArrowSerialize for Buffer { type MutableArrayType = MutableBinaryArray; diff --git a/arrow2_convert/tests/complex_example.rs b/arrow2_convert/tests/complex_example.rs index b572c49..b76d382 100644 --- a/arrow2_convert/tests/complex_example.rs +++ b/arrow2_convert/tests/complex_example.rs @@ -18,10 +18,12 @@ pub struct Root { a3: Option>, // date32 a4: chrono::NaiveDate, + // time32 + a5: chrono::NaiveTime, // timestamp(ns, None) - a5: chrono::NaiveDateTime, + a6: chrono::NaiveDateTime, // timestamp(ns, None) - a6: Option, + a7: Option, // array of date times date_time_list: Vec, // optional list array of optional strings @@ -118,7 +120,7 @@ impl arrow2_convert::deserialize::ArrowDeserialize for CustomType { arrow2_convert::arrow_enable_vec_for_type!(CustomType); fn item1() -> Root { - use chrono::{NaiveDate, NaiveDateTime}; + use chrono::{NaiveDate, NaiveDateTime, NaiveTime}; Root { name: Some("a".to_string()), @@ -127,8 +129,9 @@ fn item1() -> Root { a2: 1, a3: Some(b"aa".to_vec()), a4: NaiveDate::from_ymd_opt(1970, 1, 2).unwrap(), - a5: NaiveDateTime::from_timestamp_opt(10000, 0).unwrap(), - a6: Some(NaiveDateTime::from_timestamp_opt(10001, 0)).unwrap(), + a5: NaiveTime::from_num_seconds_from_midnight_opt(86340, 0).unwrap(), + a6: NaiveDateTime::from_timestamp_opt(10000, 0).unwrap(), + a7: Some(NaiveDateTime::from_timestamp_opt(10001, 0)).unwrap(), date_time_list: vec![ NaiveDateTime::from_timestamp_opt(10000, 10).unwrap(), NaiveDateTime::from_timestamp_opt(10000, 11).unwrap(), @@ -164,7 +167,7 @@ fn item1() -> Root { } fn item2() -> Root { - use chrono::{NaiveDate, NaiveDateTime}; + use chrono::{NaiveDate, NaiveDateTime, NaiveTime}; Root { name: Some("b".to_string()), @@ -173,8 +176,9 @@ fn item2() -> Root { a2: 1, a3: Some(b"aa".to_vec()), a4: NaiveDate::from_ymd_opt(1970, 1, 2).unwrap(), - a5: NaiveDateTime::from_timestamp_opt(10000, 0).unwrap(), - a6: None, + a5: NaiveTime::from_num_seconds_from_midnight_opt(60, 0).unwrap(), + a6: NaiveDateTime::from_timestamp_opt(10000, 0).unwrap(), + a7: None, date_time_list: vec![ NaiveDateTime::from_timestamp_opt(10000, 10).unwrap(), NaiveDateTime::from_timestamp_opt(10000, 11).unwrap(), @@ -222,7 +226,7 @@ fn test_round_trip() -> arrow2::error::Result<()> { assert_eq!(struct_array.len(), 2); let values = struct_array.values(); - assert_eq!(values.len(), 21); + assert_eq!(values.len(), 22); assert_eq!(struct_array.len(), 2); // can iterate one struct at a time without collecting diff --git a/arrow2_convert/tests/test_schema.rs b/arrow2_convert/tests/test_schema.rs index 1e465df..75b18ef 100644 --- a/arrow2_convert/tests/test_schema.rs +++ b/arrow2_convert/tests/test_schema.rs @@ -14,13 +14,15 @@ fn test_schema_types() { a3: Option>, // date32 a4: chrono::NaiveDate, + // time32 + a5: chrono::NaiveTime, // timestamp(ns, None) - a5: chrono::NaiveDateTime, + a6: chrono::NaiveDateTime, // timestamp(ns, None) - a6: Option, + a7: Option, // i128(precision, scale) #[arrow_field(type = "arrow2_convert::field::I128<32, 32>")] - a7: i128, + a8: i128, // array of date times date_time_list: Vec, // optional list array of optional strings @@ -125,9 +127,10 @@ fn test_schema_types() { Field::new("a2", DataType::Int64, false), Field::new("a3", DataType::Binary, true), Field::new("a4", DataType::Date32, false), - Field::new("a5", DataType::Timestamp(TimeUnit::Nanosecond, None), false), - Field::new("a6", DataType::Timestamp(TimeUnit::Nanosecond, None), true), - Field::new("a7", DataType::Decimal(32, 32), false), + Field::new("a5", DataType::Time32(TimeUnit::Second), false), + Field::new("a6", DataType::Timestamp(TimeUnit::Nanosecond, None), false), + Field::new("a7", DataType::Timestamp(TimeUnit::Nanosecond, None), true), + Field::new("a8", DataType::Decimal(32, 32), false), Field::new( "date_time_list", DataType::List(Box::new(Field::new(