Skip to content

Commit

Permalink
Add DictionaryArray Constructors (#3879) (#4068)
Browse files Browse the repository at this point in the history
  • Loading branch information
@tustvold
tustvold authored Apr 18, 2023
1 parent 472c977 commit 295ca86
Show file tree
Hide file tree
Showing 14 changed files with 218 additions and 214 deletions.
17 changes: 10 additions & 7 deletions arrow-arith/src/aggregate.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -759,6 +759,7 @@ mod tests {
use super::*;
use crate::arithmetic::add;
use arrow_array::types::*;
use std::sync::Arc;

#[test]
fn test_primitive_array_sum() {
Expand Down Expand Up @@ -1142,22 +1143,23 @@ mod tests {
#[test]
fn test_sum_dyn() {
let values = Int8Array::from_iter_values([10_i8, 11, 12, 13, 14, 15, 16, 17]);
let values = Arc::new(values) as ArrayRef;
let keys = Int8Array::from_iter_values([2_i8, 3, 4]);

let dict_array = DictionaryArray::try_new(&keys, &values).unwrap();
let dict_array = DictionaryArray::new(keys, values.clone());
let array = dict_array.downcast_dict::<Int8Array>().unwrap();
assert_eq!(39, sum_array::<Int8Type, _>(array).unwrap());

let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
assert_eq!(15, sum_array::<Int32Type, _>(&a).unwrap());

let keys = Int8Array::from(vec![Some(2_i8), None, Some(4)]);
let dict_array = DictionaryArray::try_new(&keys, &values).unwrap();
let dict_array = DictionaryArray::new(keys, values.clone());
let array = dict_array.downcast_dict::<Int8Array>().unwrap();
assert_eq!(26, sum_array::<Int8Type, _>(array).unwrap());

let keys = Int8Array::from(vec![None, None, None]);
let dict_array = DictionaryArray::try_new(&keys, &values).unwrap();
let dict_array = DictionaryArray::new(keys, values.clone());
let array = dict_array.downcast_dict::<Int8Array>().unwrap();
assert!(sum_array::<Int8Type, _>(array).is_none());
}
Expand All @@ -1166,8 +1168,9 @@ mod tests {
fn test_max_min_dyn() {
let values = Int8Array::from_iter_values([10_i8, 11, 12, 13, 14, 15, 16, 17]);
let keys = Int8Array::from_iter_values([2_i8, 3, 4]);
let values = Arc::new(values) as ArrayRef;

let dict_array = DictionaryArray::try_new(&keys, &values).unwrap();
let dict_array = DictionaryArray::new(keys, values.clone());
let array = dict_array.downcast_dict::<Int8Array>().unwrap();
assert_eq!(14, max_array::<Int8Type, _>(array).unwrap());

Expand All @@ -1179,14 +1182,14 @@ mod tests {
assert_eq!(1, min_array::<Int32Type, _>(&a).unwrap());

let keys = Int8Array::from(vec![Some(2_i8), None, Some(7)]);
let dict_array = DictionaryArray::try_new(&keys, &values).unwrap();
let dict_array = DictionaryArray::new(keys, values.clone());
let array = dict_array.downcast_dict::<Int8Array>().unwrap();
assert_eq!(17, max_array::<Int8Type, _>(array).unwrap());
let array = dict_array.downcast_dict::<Int8Array>().unwrap();
assert_eq!(12, min_array::<Int8Type, _>(array).unwrap());

let keys = Int8Array::from(vec![None, None, None]);
let dict_array = DictionaryArray::try_new(&keys, &values).unwrap();
let dict_array = DictionaryArray::new(keys, values.clone());
let array = dict_array.downcast_dict::<Int8Array>().unwrap();
assert!(max_array::<Int8Type, _>(array).is_none());
let array = dict_array.downcast_dict::<Int8Array>().unwrap();
Expand All @@ -1198,7 +1201,7 @@ mod tests {
let values = Float32Array::from(vec![5.0_f32, 2.0_f32, f32::NAN]);
let keys = Int8Array::from_iter_values([0_i8, 1, 2]);

let dict_array = DictionaryArray::try_new(&keys, &values).unwrap();
let dict_array = DictionaryArray::new(keys, Arc::new(values));
let array = dict_array.downcast_dict::<Float32Array>().unwrap();
assert!(max_array::<Float32Type, _>(array).unwrap().is_nan());

Expand Down
8 changes: 4 additions & 4 deletions arrow-arith/src/arithmetic.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3624,11 +3624,11 @@ mod tests {
fn test_dict_decimal() {
let values = Decimal128Array::from_iter_values([0, 1, 2, 3, 4, 5]);
let keys = Int8Array::from_iter_values([1_i8, 2, 5, 4, 3, 0]);
let array1 = DictionaryArray::try_new(&keys, &values).unwrap();
let array1 = DictionaryArray::new(keys, Arc::new(values));

let values = Decimal128Array::from_iter_values([7, -3, 4, 3, 5]);
let keys = Int8Array::from_iter_values([0_i8, 0, 1, 2, 3, 4]);
let array2 = DictionaryArray::try_new(&keys, &values).unwrap();
let array2 = DictionaryArray::new(keys, Arc::new(values));

let result = add_dyn(&array1, &array2).unwrap();
let expected =
Expand All @@ -3650,7 +3650,7 @@ mod tests {
]);
let keys =
Int8Array::from(vec![Some(1_i8), None, Some(5), Some(4), Some(3), None]);
let array1 = DictionaryArray::try_new(&keys, &values).unwrap();
let array1 = DictionaryArray::new(keys, Arc::new(values));

let values = Decimal256Array::from_iter_values([
i256::from_i128(7),
Expand All @@ -3661,7 +3661,7 @@ mod tests {
]);
let keys =
Int8Array::from(vec![Some(0_i8), Some(0), None, Some(2), Some(3), Some(4)]);
let array2 = DictionaryArray::try_new(&keys, &values).unwrap();
let array2 = DictionaryArray::new(keys, Arc::new(values));

let result = add_dyn(&array1, &array2).unwrap();
let expected = Arc::new(Decimal256Array::from(vec![
Expand Down
5 changes: 1 addition & 4 deletions arrow-arith/src/arity.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -603,10 +603,7 @@ mod tests {
fn test_unary_dict_mut() {
let values = Int32Array::from(vec![Some(10), Some(20), None]);
let keys = Int8Array::from_iter_values([0, 0, 1, 2]);
let dictionary = DictionaryArray::<Int8Type>::try_new(&keys, &values).unwrap();

drop(keys);
drop(values);
let dictionary = DictionaryArray::new(keys, Arc::new(values));

let updated = dictionary.unary_mut::<_, Int32Type>(|x| x + 1).unwrap();
let typed = updated.downcast_dict::<Int32Array>().unwrap();
Expand Down
108 changes: 43 additions & 65 deletions arrow-arith/src/temporal.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,10 @@

//! Defines temporal kernels for time and date related functions.

use chrono::{DateTime, Datelike, NaiveDateTime, NaiveTime, Offset, Timelike};
use std::sync::Arc;

use chrono::{DateTime, Datelike, NaiveDateTime, NaiveTime, Offset, Timelike};

use arrow_array::builder::*;
use arrow_array::iterator::ArrayIter;
use arrow_array::temporal_conversions::{
Expand Down Expand Up @@ -970,12 +971,14 @@ mod tests {
.with_timezone("+01:00".to_string());

let keys = Int8Array::from_iter_values([0_i8, 0, 1, 2, 1]);
let dict = DictionaryArray::try_new(&keys, &a).unwrap();
let dict = DictionaryArray::try_new(keys.clone(), Arc::new(a)).unwrap();

let b = hour_dyn(&dict).unwrap();

let expected_dict =
DictionaryArray::try_new(&keys, &Int32Array::from(vec![11, 21, 7])).unwrap();
let expected_dict = DictionaryArray::new(
keys.clone(),
Arc::new(Int32Array::from(vec![11, 21, 7])),
);
let expected = Arc::new(expected_dict) as ArrayRef;
assert_eq!(&expected, &b);

Expand All @@ -984,7 +987,7 @@ mod tests {
let b_old = minute_dyn(&dict).unwrap();

let expected_dict =
DictionaryArray::try_new(&keys, &Int32Array::from(vec![1, 2, 3])).unwrap();
DictionaryArray::new(keys.clone(), Arc::new(Int32Array::from(vec![1, 2, 3])));
let expected = Arc::new(expected_dict) as ArrayRef;
assert_eq!(&expected, &b);
assert_eq!(&expected, &b_old);
Expand All @@ -994,7 +997,7 @@ mod tests {
let b_old = second_dyn(&dict).unwrap();

let expected_dict =
DictionaryArray::try_new(&keys, &Int32Array::from(vec![1, 2, 3])).unwrap();
DictionaryArray::new(keys.clone(), Arc::new(Int32Array::from(vec![1, 2, 3])));
let expected = Arc::new(expected_dict) as ArrayRef;
assert_eq!(&expected, &b);
assert_eq!(&expected, &b_old);
Expand All @@ -1003,8 +1006,7 @@ mod tests {
time_fraction_dyn(&dict, "nanosecond", |t| t.nanosecond() as i32).unwrap();

let expected_dict =
DictionaryArray::try_new(&keys, &Int32Array::from(vec![0, 0, 0, 0, 0]))
.unwrap();
DictionaryArray::new(keys, Arc::new(Int32Array::from(vec![0, 0, 0, 0, 0])));
let expected = Arc::new(expected_dict) as ArrayRef;
assert_eq!(&expected, &b);
}
Expand All @@ -1015,15 +1017,14 @@ mod tests {
vec![Some(1514764800000), Some(1550636625000)].into();

let keys = Int8Array::from_iter_values([0_i8, 1, 1, 0]);
let dict = DictionaryArray::try_new(&keys, &a).unwrap();
let dict = DictionaryArray::new(keys.clone(), Arc::new(a));

let b = year_dyn(&dict).unwrap();

let expected_dict = DictionaryArray::try_new(
&keys,
&Int32Array::from(vec![2018, 2019, 2019, 2018]),
)
.unwrap();
let expected_dict = DictionaryArray::new(
keys,
Arc::new(Int32Array::from(vec![2018, 2019, 2019, 2018])),
);
let expected = Arc::new(expected_dict) as ArrayRef;
assert_eq!(&expected, &b);
}
Expand All @@ -1036,21 +1037,21 @@ mod tests {
vec![Some(1514764800000), Some(1566275025000)].into();

let keys = Int8Array::from_iter_values([0_i8, 1, 1, 0]);
let dict = DictionaryArray::try_new(&keys, &a).unwrap();
let dict = DictionaryArray::new(keys.clone(), Arc::new(a));

let b = quarter_dyn(&dict).unwrap();

let expected_dict =
DictionaryArray::try_new(&keys, &Int32Array::from(vec![1, 3, 3, 1])).unwrap();
let expected = Arc::new(expected_dict) as ArrayRef;
assert_eq!(&expected, &b);
let expected = DictionaryArray::new(
keys.clone(),
Arc::new(Int32Array::from(vec![1, 3, 3, 1])),
);
assert_eq!(b.as_ref(), &expected);

let b = month_dyn(&dict).unwrap();

let expected_dict =
DictionaryArray::try_new(&keys, &Int32Array::from(vec![1, 8, 8, 1])).unwrap();
let expected = Arc::new(expected_dict) as ArrayRef;
assert_eq!(&expected, &b);
let expected =
DictionaryArray::new(keys, Arc::new(Int32Array::from(vec![1, 8, 8, 1])));
assert_eq!(b.as_ref(), &expected);
}

#[test]
Expand All @@ -1061,57 +1062,37 @@ mod tests {
vec![Some(1514764800000), Some(1550636625000)].into();

let keys = Int8Array::from(vec![Some(0_i8), Some(1), Some(1), Some(0), None]);
let dict = DictionaryArray::try_new(&keys, &a).unwrap();
let dict = DictionaryArray::new(keys.clone(), Arc::new(a));

let b = num_days_from_monday_dyn(&dict).unwrap();

let expected_dict = DictionaryArray::try_new(
&keys,
&Int32Array::from(vec![Some(0), Some(2), Some(2), Some(0), None]),
)
.unwrap();
let expected = Arc::new(expected_dict) as ArrayRef;
assert_eq!(&expected, &b);
let a = Int32Array::from(vec![Some(0), Some(2), Some(2), Some(0), None]);
let expected = DictionaryArray::new(keys.clone(), Arc::new(a));
assert_eq!(b.as_ref(), &expected);

let b = num_days_from_sunday_dyn(&dict).unwrap();

let expected_dict = DictionaryArray::try_new(
&keys,
&Int32Array::from(vec![Some(1), Some(3), Some(3), Some(1), None]),
)
.unwrap();
let expected = Arc::new(expected_dict) as ArrayRef;
assert_eq!(&expected, &b);
let a = Int32Array::from(vec![Some(1), Some(3), Some(3), Some(1), None]);
let expected = DictionaryArray::new(keys.clone(), Arc::new(a));
assert_eq!(b.as_ref(), &expected);

let b = day_dyn(&dict).unwrap();

let expected_dict = DictionaryArray::try_new(
&keys,
&Int32Array::from(vec![Some(1), Some(20), Some(20), Some(1), None]),
)
.unwrap();
let expected = Arc::new(expected_dict) as ArrayRef;
assert_eq!(&expected, &b);
let a = Int32Array::from(vec![Some(1), Some(20), Some(20), Some(1), None]);
let expected = DictionaryArray::new(keys.clone(), Arc::new(a));
assert_eq!(b.as_ref(), &expected);

let b = doy_dyn(&dict).unwrap();

let expected_dict = DictionaryArray::try_new(
&keys,
&Int32Array::from(vec![Some(1), Some(51), Some(51), Some(1), None]),
)
.unwrap();
let expected = Arc::new(expected_dict) as ArrayRef;
assert_eq!(&expected, &b);
let a = Int32Array::from(vec![Some(1), Some(51), Some(51), Some(1), None]);
let expected = DictionaryArray::new(keys.clone(), Arc::new(a));
assert_eq!(b.as_ref(), &expected);

let b = week_dyn(&dict).unwrap();

let expected_dict = DictionaryArray::try_new(
&keys,
&Int32Array::from(vec![Some(1), Some(8), Some(8), Some(1), None]),
)
.unwrap();
let expected = Arc::new(expected_dict) as ArrayRef;
assert_eq!(&expected, &b);
let a = Int32Array::from(vec![Some(1), Some(8), Some(8), Some(1), None]);
let expected = DictionaryArray::new(keys, Arc::new(a));
assert_eq!(b.as_ref(), &expected);
}

#[test]
Expand All @@ -1129,14 +1110,11 @@ mod tests {
assert_eq!(453_000_000, b.value(1));

let keys = Int8Array::from(vec![Some(0_i8), Some(1), Some(1)]);
let dict = DictionaryArray::try_new(&keys, &a).unwrap();
let dict = DictionaryArray::new(keys.clone(), Arc::new(a));
let b = nanosecond_dyn(&dict).unwrap();

let expected_dict = DictionaryArray::try_new(
&keys,
&Int32Array::from(vec![None, Some(453_000_000)]),
)
.unwrap();
let a = Int32Array::from(vec![None, Some(453_000_000)]);
let expected_dict = DictionaryArray::new(keys, Arc::new(a));
let expected = Arc::new(expected_dict) as ArrayRef;
assert_eq!(&expected, &b);
}
Expand Down
Loading

0 comments on commit 295ca86

Please sign in to comment.