Skip to content

Commit

Permalink
Test values and formatting
Browse files Browse the repository at this point in the history
  • Loading branch information
@avantgardnerio
avantgardnerio committed Jul 8, 2022
1 parent 97cb277 commit ef88801
Show file tree
Hide file tree
Showing 2 changed files with 171 additions and 63 deletions.
193 changes: 144 additions & 49 deletions arrow/src/compute/kernels/arithmetic.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,10 @@ use crate::compute::kernels::arity::unary;
use crate::compute::unary_dyn;
use crate::compute::util::combine_option_bitmap;
use crate::datatypes;
use crate::datatypes::{ArrowNumericType, DataType, Date32Type, IntervalDayTimeType, IntervalMonthDayNanoType, IntervalUnit, IntervalYearMonthType};
use crate::datatypes::{
ArrowNumericType, DataType, Date32Type, Date64Type, IntervalDayTimeType,
IntervalMonthDayNanoType, IntervalUnit, IntervalYearMonthType,
};
use crate::datatypes::{
Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type, UInt16Type,
UInt32Type, UInt64Type, UInt8Type,
Expand Down Expand Up @@ -776,45 +779,103 @@ pub fn add_dyn(left: &dyn Array, right: &dyn Array) -> Result<ArrayRef> {
typed_dict_math_op!(left, right, |a, b| a + b, math_op_dict)
}
DataType::Date32 => {
let l = left.as_any()
let l = left
.as_any()
.downcast_ref::<PrimitiveArray<Date32Type>>()
.ok_or_else(|| {ArrowError::CastError(format!(
"Left array cannot be cast to Date32Type",
))})?;
.ok_or_else(|| {
ArrowError::CastError(format!(
"Left array cannot be cast to Date32Type",
))
})?;
match right.data_type() {
DataType::Interval(IntervalUnit::YearMonth) => {
let r = right.as_any()
.downcast_ref::<PrimitiveArray<IntervalYearMonthType>>()
.ok_or_else(|| {ArrowError::CastError(format!(
"Right array cannot be cast to IntervalYearMonthType",
))})?;
let res = math_op(l, r, |a, b| {
Date32Type::add_year_months(a, b)
})?;
let r = right
.as_any()
.downcast_ref::<PrimitiveArray<IntervalYearMonthType>>()
.ok_or_else(|| {
ArrowError::CastError(format!(
"Right array cannot be cast to IntervalYearMonthType",
))
})?;
let res = math_op(l, r, |a, b| Date32Type::add_year_months(a, b))?;
return Ok(Arc::new(res));
},
}
DataType::Interval(IntervalUnit::DayTime) => {
let r = right.as_any()
let r = right
.as_any()
.downcast_ref::<PrimitiveArray<IntervalDayTimeType>>()
.ok_or_else(|| {ArrowError::CastError(format!(
"Right array cannot be cast to IntervalDayTimeType",
))})?;
let res = math_op(l, r, |a, b| {
Date32Type::add_day_time(a, b)
})?;
.ok_or_else(|| {
ArrowError::CastError(format!(
"Right array cannot be cast to IntervalDayTimeType",
))
})?;
let res = math_op(l, r, |a, b| Date32Type::add_day_time(a, b))?;
return Ok(Arc::new(res));
},
}
DataType::Interval(IntervalUnit::MonthDayNano) => {
let r = right.as_any()
let r = right
.as_any()
.downcast_ref::<PrimitiveArray<IntervalMonthDayNanoType>>()
.ok_or_else(|| {ArrowError::CastError(format!(
"Right array cannot be cast to IntervalMonthDayNanoType",
))})?;
let res = math_op(l, r, |a, b| {
Date32Type::add_month_day_nano(a, b)
})?;
.ok_or_else(|| {
ArrowError::CastError(format!(
"Right array cannot be cast to IntervalMonthDayNanoType",
))
})?;
let res = math_op(l, r, |a, b| Date32Type::add_month_day_nano(a, b))?;
return Ok(Arc::new(res));
},
}
t => Err(ArrowError::CastError(format!(
"Cannot perform arithmetic operation on arrays of type {}",
t
))),
}
}
DataType::Date64 => {
let l = left
.as_any()
.downcast_ref::<PrimitiveArray<Date64Type>>()
.ok_or_else(|| {
ArrowError::CastError(format!(
"Left array cannot be cast to Date64Type",
))
})?;
match right.data_type() {
DataType::Interval(IntervalUnit::YearMonth) => {
let r = right
.as_any()
.downcast_ref::<PrimitiveArray<IntervalYearMonthType>>()
.ok_or_else(|| {
ArrowError::CastError(format!(
"Right array cannot be cast to IntervalYearMonthType",
))
})?;
let res = math_op(l, r, |a, b| Date64Type::add_year_months(a, b))?;
return Ok(Arc::new(res));
}
DataType::Interval(IntervalUnit::DayTime) => {
let r = right
.as_any()
.downcast_ref::<PrimitiveArray<IntervalDayTimeType>>()
.ok_or_else(|| {
ArrowError::CastError(format!(
"Right array cannot be cast to IntervalDayTimeType",
))
})?;
let res = math_op(l, r, |a, b| Date64Type::add_day_time(a, b))?;
return Ok(Arc::new(res));
}
DataType::Interval(IntervalUnit::MonthDayNano) => {
let r = right
.as_any()
.downcast_ref::<PrimitiveArray<IntervalMonthDayNanoType>>()
.ok_or_else(|| {
ArrowError::CastError(format!(
"Right array cannot be cast to IntervalMonthDayNanoType",
))
})?;
let res = math_op(l, r, |a, b| Date64Type::add_month_day_nano(a, b))?;
return Ok(Arc::new(res));
}
t => Err(ArrowError::CastError(format!(
"Cannot perform arithmetic operation on arrays of type {}",
t
Expand Down Expand Up @@ -1100,10 +1161,10 @@ where

#[cfg(test)]
mod tests {
use chrono::NaiveDate;
use super::*;
use crate::array::Int32Array;
use crate::datatypes::Date64Type;
use chrono::NaiveDate;

#[test]
fn test_primitive_array_add() {
Expand All @@ -1119,56 +1180,90 @@ mod tests {

#[test]
fn test_date32_month_add() {
let a = Date32Array::from(vec![Date32Type::from_naive_date(NaiveDate::from_ymd(2000, 01, 01))]);
let b = IntervalYearMonthArray::from(vec![IntervalYearMonthType::from(1, 1)]);
let a = Date32Array::from(vec![Date32Type::from_naive_date(
NaiveDate::from_ymd(2000, 01, 01),
)]);
let b = IntervalYearMonthArray::from(vec![IntervalYearMonthType::from(1, 2)]);
let c = add_dyn(&a, &b).unwrap();
let c = c.as_any().downcast_ref::<Date32Array>().unwrap();
assert_eq!(c.value(0), Date32Type::from_naive_date(NaiveDate::from_ymd(2001, 02, 01)));
assert_eq!(
c.value(0),
Date32Type::from_naive_date(NaiveDate::from_ymd(2001, 03, 01))
);
}

#[test]
fn test_date32_day_time_add() {
let a = Date32Array::from(vec![Date32Type::from_naive_date(NaiveDate::from_ymd(2000, 01, 01))]);
let b = IntervalDayTimeArray::from(vec![IntervalDayTimeType::from(1, 1)]);
let a = Date32Array::from(vec![Date32Type::from_naive_date(
NaiveDate::from_ymd(2000, 01, 01),
)]);
let b = IntervalDayTimeArray::from(vec![IntervalDayTimeType::from(1, 2)]);
let c = add_dyn(&a, &b).unwrap();
let c = c.as_any().downcast_ref::<Date32Array>().unwrap();
assert_eq!(c.value(0), Date32Type::from_naive_date(NaiveDate::from_ymd(2001, 01, 02)));
assert_eq!(
c.value(0),
Date32Type::from_naive_date(NaiveDate::from_ymd(2000, 01, 02))
);
}

#[test]
fn test_date32_month_day_nano_add() {
let a = Date32Array::from(vec![Date32Type::from_naive_date(NaiveDate::from_ymd(2000, 01, 01))]);
let b = IntervalMonthDayNanoArray::from(vec![IntervalMonthDayNanoType::from(1, 1, 1)]);
let a = Date32Array::from(vec![Date32Type::from_naive_date(
NaiveDate::from_ymd(2000, 01, 01),
)]);
let b = IntervalMonthDayNanoArray::from(vec![IntervalMonthDayNanoType::from(
1, 2, 3,
)]);
let c = add_dyn(&a, &b).unwrap();
let c = c.as_any().downcast_ref::<Date32Array>().unwrap();
assert_eq!(c.value(0), Date32Type::from_naive_date(NaiveDate::from_ymd(2001, 02, 01)));
assert_eq!(
c.value(0),
Date32Type::from_naive_date(NaiveDate::from_ymd(2000, 02, 03))
);
}

#[test]
fn test_date64_month_add() {
let a = Date64Array::from(vec![Date64Type::from_naive_date(NaiveDate::from_ymd(2000, 01, 01))]);
let b = IntervalYearMonthArray::from(vec![IntervalYearMonthType::from(1, 1)]);
let a = Date64Array::from(vec![Date64Type::from_naive_date(
NaiveDate::from_ymd(2000, 01, 01),
)]);
let b = IntervalYearMonthArray::from(vec![IntervalYearMonthType::from(1, 2)]);
let c = add_dyn(&a, &b).unwrap();
let c = c.as_any().downcast_ref::<Date64Array>().unwrap();
assert_eq!(c.value(0), Date64Type::from_naive_date(NaiveDate::from_ymd(2001, 02, 01)));
assert_eq!(
c.value(0),
Date64Type::from_naive_date(NaiveDate::from_ymd(2001, 03, 01))
);
}

#[test]
fn test_date64_day_time_add() {
let a = Date64Array::from(vec![Date64Type::from_naive_date(NaiveDate::from_ymd(2000, 01, 01))]);
let b = IntervalDayTimeArray::from(vec![IntervalDayTimeType::from(1, 1)]);
let a = Date64Array::from(vec![Date64Type::from_naive_date(
NaiveDate::from_ymd(2000, 01, 01),
)]);
let b = IntervalDayTimeArray::from(vec![IntervalDayTimeType::from(1, 2)]);
let c = add_dyn(&a, &b).unwrap();
let c = c.as_any().downcast_ref::<Date64Array>().unwrap();
assert_eq!(c.value(0), Date64Type::from_naive_date(NaiveDate::from_ymd(2001, 01, 02)));
assert_eq!(
c.value(0),
Date64Type::from_naive_date(NaiveDate::from_ymd(2000, 01, 02))
);
}

#[test]
fn test_date64_month_day_nano_add() {
let a = Date64Array::from(vec![Date64Type::from_naive_date(NaiveDate::from_ymd(2000, 01, 01))]);
let b = IntervalMonthDayNanoArray::from(vec![IntervalMonthDayNanoType::from(1, 1, 1)]);
let a = Date64Array::from(vec![Date64Type::from_naive_date(
NaiveDate::from_ymd(2000, 01, 01),
)]);
let b = IntervalMonthDayNanoArray::from(vec![IntervalMonthDayNanoType::from(
1, 2, 3,
)]);
let c = add_dyn(&a, &b).unwrap();
let c = c.as_any().downcast_ref::<Date64Array>().unwrap();
assert_eq!(c.value(0), Date64Type::from_naive_date(NaiveDate::from_ymd(2001, 02, 01)));
assert_eq!(
c.value(0),
Date64Type::from_naive_date(NaiveDate::from_ymd(2000, 02, 03))
);
}

#[test]
Expand Down
41 changes: 27 additions & 14 deletions arrow/src/datatypes/types.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,11 @@
// specific language governing permissions and limitations
// under the License.

use std::ops::{Add, Sub};
use super::{ArrowPrimitiveType, DataType, IntervalUnit, TimeUnit};
use chrono::{Duration, NaiveDate};
use chronoutil::shift_months;
use super::{ArrowPrimitiveType, DataType, IntervalUnit, TimeUnit};
use half::f16;
use std::ops::{Add, Sub};

// BooleanType is special: its bit-width is not the size of the primitive type, and its `index`
// operation assumes bit-packing.
Expand Down Expand Up @@ -196,7 +196,10 @@ impl ArrowTimestampType for TimestampNanosecondType {
}

impl IntervalYearMonthType {
pub fn from(years: i32, months: i32) -> <IntervalYearMonthType as ArrowPrimitiveType>::Native {
pub fn from(
years: i32,
months: i32,
) -> <IntervalYearMonthType as ArrowPrimitiveType>::Native {
years * 12 + months
}

Expand All @@ -206,36 +209,46 @@ impl IntervalYearMonthType {
}

impl IntervalDayTimeType {
pub fn from(days: i32, millis: i32) -> <IntervalDayTimeType as ArrowPrimitiveType>::Native {
pub fn from(
days: i32,
millis: i32,
) -> <IntervalDayTimeType as ArrowPrimitiveType>::Native {
let m = millis as u64 & u32::MAX as u64;
let d = (days as u64 & u32::MAX as u64) << 32;
(m | d) as <IntervalDayTimeType as ArrowPrimitiveType>::Native
}

pub fn to_parts(i: <IntervalDayTimeType as ArrowPrimitiveType>::Native) -> (i32, i32) {
pub fn to_parts(
i: <IntervalDayTimeType as ArrowPrimitiveType>::Native,
) -> (i32, i32) {
let days = (i >> 32) as i32;
let ms = i as i32;
(days, ms)
}
}

impl IntervalMonthDayNanoType {
pub fn from(months: i32, days: i32, nanos: i64) -> <IntervalMonthDayNanoType as ArrowPrimitiveType>::Native {
pub fn from(
months: i32,
days: i32,
nanos: i64,
) -> <IntervalMonthDayNanoType as ArrowPrimitiveType>::Native {
let m = months as u128 & u32::MAX as u128;
let d = (days as u128 & u32::MAX as u128) << 32;
let n = (nanos as u128) << 64;
(m | d | n) as <IntervalMonthDayNanoType as ArrowPrimitiveType>::Native
}

pub fn to_parts(i: <IntervalMonthDayNanoType as ArrowPrimitiveType>::Native) -> (i32, i32, i64) {
pub fn to_parts(
i: <IntervalMonthDayNanoType as ArrowPrimitiveType>::Native,
) -> (i32, i32, i64) {
let nanos = (i >> 64) as i64;
let days = (i >> 32) as i32;
let months = i as i32;
(months, days, nanos)
}
}


impl Date32Type {
pub fn to_naive_date(i: <Date32Type as ArrowPrimitiveType>::Native) -> NaiveDate {
let epoch = NaiveDate::from_ymd(1970, 1, 1);
Expand All @@ -249,7 +262,7 @@ impl Date32Type {

pub fn add_year_months(
date: <Date32Type as ArrowPrimitiveType>::Native,
delta: <IntervalYearMonthType as ArrowPrimitiveType>::Native
delta: <IntervalYearMonthType as ArrowPrimitiveType>::Native,
) -> <Date32Type as ArrowPrimitiveType>::Native {
let prior = Date32Type::to_naive_date(date);
let months = IntervalYearMonthType::to_months(delta);
Expand All @@ -259,7 +272,7 @@ impl Date32Type {

pub fn add_day_time(
date: <Date32Type as ArrowPrimitiveType>::Native,
delta: <IntervalDayTimeType as ArrowPrimitiveType>::Native
delta: <IntervalDayTimeType as ArrowPrimitiveType>::Native,
) -> <Date32Type as ArrowPrimitiveType>::Native {
let (days, ms) = IntervalDayTimeType::to_parts(delta);
let res = Date32Type::to_naive_date(date);
Expand All @@ -270,7 +283,7 @@ impl Date32Type {

pub fn add_month_day_nano(
date: <Date32Type as ArrowPrimitiveType>::Native,
delta: <IntervalMonthDayNanoType as ArrowPrimitiveType>::Native
delta: <IntervalMonthDayNanoType as ArrowPrimitiveType>::Native,
) -> <Date32Type as ArrowPrimitiveType>::Native {
let (months, days, nanos) = IntervalMonthDayNanoType::to_parts(delta);
let res = Date32Type::to_naive_date(date);
Expand All @@ -294,7 +307,7 @@ impl Date64Type {

pub fn add_year_months(
date: <Date64Type as ArrowPrimitiveType>::Native,
delta: <IntervalYearMonthType as ArrowPrimitiveType>::Native
delta: <IntervalYearMonthType as ArrowPrimitiveType>::Native,
) -> <Date64Type as ArrowPrimitiveType>::Native {
let prior = Date64Type::to_naive_date(date);
let months = IntervalYearMonthType::to_months(delta);
Expand All @@ -304,7 +317,7 @@ impl Date64Type {

pub fn add_day_time(
date: <Date64Type as ArrowPrimitiveType>::Native,
delta: <IntervalDayTimeType as ArrowPrimitiveType>::Native
delta: <IntervalDayTimeType as ArrowPrimitiveType>::Native,
) -> <Date64Type as ArrowPrimitiveType>::Native {
let (days, ms) = IntervalDayTimeType::to_parts(delta);
let res = Date64Type::to_naive_date(date);
Expand All @@ -315,7 +328,7 @@ impl Date64Type {

pub fn add_month_day_nano(
date: <Date64Type as ArrowPrimitiveType>::Native,
delta: <IntervalMonthDayNanoType as ArrowPrimitiveType>::Native
delta: <IntervalMonthDayNanoType as ArrowPrimitiveType>::Native,
) -> <Date64Type as ArrowPrimitiveType>::Native {
let (months, days, nanos) = IntervalMonthDayNanoType::to_parts(delta);
let res = Date64Type::to_naive_date(date);
Expand Down

0 comments on commit ef88801

Please sign in to comment.