Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add DictionaryArray Constructors (#3879) #4068

Merged
merged 1 commit into from
Apr 18, 2023
Merged

Add DictionaryArray Constructors (#3879) #4068

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 10 additions & 7 deletions arrow-arith/src/aggregate.rs
View file
Open in desktop
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
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3426,11 +3426,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 @@ -3452,7 +3452,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 @@ -3463,7 +3463,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
Open in desktop
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
Open in desktop
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