Skip to content

Commit

Permalink
Replace complicated temporal macro with generic functions (#2850)
Browse files Browse the repository at this point in the history
* Remove complicated macro

* Fix clippy

* Add doc.
  • Loading branch information
viirya committed Oct 12, 2022
1 parent 3bb7f38 commit f8254e7
Show file tree
Hide file tree
Showing 2 changed files with 280 additions and 309 deletions.
117 changes: 100 additions & 17 deletions arrow/src/compute/kernels/cast.rs
Expand Up @@ -37,14 +37,13 @@

use chrono::format::strftime::StrftimeItems;
use chrono::format::{parse, Parsed};
use chrono::Timelike;
use chrono::{NaiveDateTime, Timelike};
use std::ops::{Div, Mul};
use std::str;
use std::sync::Arc;

use crate::buffer::MutableBuffer;
use crate::compute::kernels::cast_utils::string_to_timestamp_nanos;
use crate::compute::kernels::temporal::extract_component_from_array;
use crate::compute::kernels::temporal::return_compute_error_with;
use crate::compute::{divide_scalar, multiply_scalar};
use crate::compute::{try_unary, using_chrono_tz_and_utc_naive_date_time};
Expand Down Expand Up @@ -1638,6 +1637,98 @@ where
unsafe { PrimitiveArray::<R>::from_trusted_len_iter(iter) }
}

fn as_time_with_string_op<
A: ArrayAccessor<Item = T::Native>,
OffsetSize,
T: ArrowTemporalType,
F,
>(
iter: ArrayIter<A>,
mut builder: GenericStringBuilder<OffsetSize>,
op: F,
) -> ArrayRef
where
OffsetSize: OffsetSizeTrait,
F: Fn(NaiveDateTime) -> String,
i64: From<T::Native>,
{
iter.into_iter().for_each(|value| {
if let Some(value) = value {
match as_datetime::<T>(<i64 as From<_>>::from(value)) {
Some(dt) => builder.append_value(op(dt)),
None => builder.append_null(),
}
} else {
builder.append_null();
}
});

Arc::new(builder.finish())
}

fn extract_component_from_datatime_array<
A: ArrayAccessor<Item = T::Native>,
OffsetSize,
T: ArrowTemporalType,
F,
>(
iter: ArrayIter<A>,
mut builder: GenericStringBuilder<OffsetSize>,
tz: &str,
mut parsed: Parsed,
op: F,
) -> Result<ArrayRef>
where
OffsetSize: OffsetSizeTrait,
F: Fn(NaiveDateTime) -> String,
i64: From<T::Native>,
{
if (tz.starts_with('+') || tz.starts_with('-')) && !tz.contains(':') {
return_compute_error_with!(
"Invalid timezone",
"Expected format [+-]XX:XX".to_string()
)
} else {
let tz_parse_result = parse(&mut parsed, tz, StrftimeItems::new("%z"));
let fixed_offset_from_parsed = match tz_parse_result {
Ok(_) => match parsed.to_fixed_offset() {
Ok(fo) => Some(fo),
err => return_compute_error_with!("Invalid timezone", err),
},
_ => None,
};

for value in iter {
if let Some(value) = value {
match as_datetime::<T>(<i64 as From<_>>::from(value)) {
Some(utc) => {
let fixed_offset = match fixed_offset_from_parsed {
Some(fo) => fo,
None => {
match using_chrono_tz_and_utc_naive_date_time(tz, utc) {
Some(fo) => fo,
err => return_compute_error_with!(
"Unable to parse timezone",
err
),
}
}
};
builder.append_value(op(utc + fixed_offset));
}
err => return_compute_error_with!(
"Unable to read value as datetime",
err
),
}
} else {
builder.append_null();
}
}
}
Ok(Arc::new(builder.finish()))
}

/// Cast timestamp types to Utf8/LargeUtf8
fn cast_timestamp_to_string<T, OffsetSize>(
array: &ArrayRef,
Expand All @@ -1650,38 +1741,30 @@ where
{
let array = array.as_any().downcast_ref::<PrimitiveArray<T>>().unwrap();

let mut builder = GenericStringBuilder::<OffsetSize>::new();
let builder = GenericStringBuilder::<OffsetSize>::new();

if let Some(tz) = tz {
let mut scratch = Parsed::new();
let scratch = Parsed::new();
// The macro calls `as_datetime` on timestamp values of the array.
// After applying timezone offset on the datatime, calling `to_string` to get
// the strings.
let iter = ArrayIter::new(array);
extract_component_from_array!(
extract_component_from_datatime_array::<_, OffsetSize, T, _>(
iter,
builder,
to_string,
|value, tz| as_datetime::<T>(<i64 as From<_>>::from(value))
.map(|datetime| datetime + tz),
tz,
scratch,
|value| as_datetime::<T>(<i64 as From<_>>::from(value)),
|h| h
|t| t.to_string(),
)
} else {
// No timezone available. Calling `to_string` on the datatime value simply.
let iter = ArrayIter::new(array);
extract_component_from_array!(
Ok(as_time_with_string_op::<_, OffsetSize, T, _>(
iter,
builder,
to_string,
|value| as_datetime::<T>(<i64 as From<_>>::from(value)),
|h| h
)
|t| t.to_string(),
))
}

Ok(Arc::new(builder.finish()) as ArrayRef)
}

/// Cast date32 types to Utf8/LargeUtf8
Expand Down

0 comments on commit f8254e7

Please sign in to comment.