Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Commit

Permalink
feat: add duration type to json writer (#1522)
Browse files Browse the repository at this point in the history
  • Loading branch information
ritchie46 authored Jul 27, 2023
1 parent f175c1c commit 15c5ec1
Show file tree
Hide file tree
Showing 2 changed files with 65 additions and 4 deletions.
43 changes: 40 additions & 3 deletions src/io/json/write/serialize.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use chrono::{NaiveDate, NaiveDateTime};
use chrono::{Duration, NaiveDate, NaiveDateTime};
use lexical_core::ToLexical;
use std::io::Write;
use streaming_iterator::StreamingIterator;
Expand All @@ -8,8 +8,9 @@ use crate::datatypes::{IntegerType, TimeUnit};
use crate::io::iterator::BufStreamingIterator;
use crate::offset::Offset;
use crate::temporal_conversions::{
date32_to_date, date64_to_date, timestamp_ms_to_datetime, timestamp_ns_to_datetime,
timestamp_s_to_datetime, timestamp_us_to_datetime,
date32_to_date, date64_to_date, duration_ms_to_duration, duration_ns_to_duration,
duration_s_to_duration, duration_us_to_duration, timestamp_ms_to_datetime,
timestamp_ns_to_datetime, timestamp_s_to_datetime, timestamp_us_to_datetime,
};
use crate::util::lexical_to_bytes_mut;
use crate::{array::*, datatypes::DataType, types::NativeType};
Expand Down Expand Up @@ -266,6 +267,28 @@ where
materialize_serializer(f, array.iter(), offset, take)
}

fn duration_serializer<'a, T, F>(
array: &'a PrimitiveArray<T>,
convert: F,
offset: usize,
take: usize,
) -> Box<dyn StreamingIterator<Item = [u8]> + 'a + Send + Sync>
where
T: NativeType,
F: Fn(T) -> Duration + 'static + Send + Sync,
{
let f = move |x: Option<&T>, buf: &mut Vec<u8>| {
if let Some(x) = x {
let duration = convert(*x);
write!(buf, "\"{duration}\"").unwrap();
} else {
buf.extend_from_slice(b"null")
}
};

materialize_serializer(f, array.iter(), offset, take)
}

fn timestamp_serializer<'a, F>(
array: &'a PrimitiveArray<i64>,
convert: F,
Expand Down Expand Up @@ -385,6 +408,20 @@ pub(crate) fn new_serializer<'a>(
)
}
}
DataType::Duration(tu) => {
let convert = match tu {
TimeUnit::Nanosecond => duration_ns_to_duration,
TimeUnit::Microsecond => duration_us_to_duration,
TimeUnit::Millisecond => duration_ms_to_duration,
TimeUnit::Second => duration_s_to_duration,
};
duration_serializer(
array.as_any().downcast_ref().unwrap(),
convert,
offset,
take,
)
}
DataType::Null => null_serializer(array.len(), offset, take),
other => todo!("Writing {:?} to JSON", other),
}
Expand Down
26 changes: 25 additions & 1 deletion src/temporal_conversions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

use chrono::{
format::{parse, Parsed, StrftimeItems},
Datelike, FixedOffset, NaiveDate, NaiveDateTime, NaiveTime,
Datelike, Duration, FixedOffset, NaiveDate, NaiveDateTime, NaiveTime,
};

use crate::error::Result;
Expand Down Expand Up @@ -66,6 +66,30 @@ pub fn time32s_to_time(v: i32) -> NaiveTime {
NaiveTime::from_num_seconds_from_midnight_opt(v as u32, 0).expect("invalid time")
}

/// converts a `i64` representing a `duration(s)` to [`Duration`]
#[inline]
pub fn duration_s_to_duration(v: i64) -> Duration {
Duration::seconds(v)
}

/// converts a `i64` representing a `duration(ms)` to [`Duration`]
#[inline]
pub fn duration_ms_to_duration(v: i64) -> Duration {
Duration::milliseconds(v)
}

/// converts a `i64` representing a `duration(us)` to [`Duration`]
#[inline]
pub fn duration_us_to_duration(v: i64) -> Duration {
Duration::microseconds(v)
}

/// converts a `i64` representing a `duration(ns)` to [`Duration`]
#[inline]
pub fn duration_ns_to_duration(v: i64) -> Duration {
Duration::nanoseconds(v)
}

/// converts a `i32` representing a `time32(ms)` to [`NaiveTime`]
#[inline]
pub fn time32ms_to_time(v: i32) -> NaiveTime {
Expand Down

0 comments on commit 15c5ec1

Please sign in to comment.