From f28427d61e858ff0bd4b803f9c5849434cfb22d8 Mon Sep 17 00:00:00 2001 From: Jorge Leitao Date: Tue, 7 Sep 2021 21:56:07 +0100 Subject: [PATCH] Moved compute tests to tests/ (#388) --- src/compute/boolean.rs | 343 ----------------- src/compute/concat.rs | 296 +-------------- src/compute/filter.rs | 182 --------- src/compute/hash.rs | 54 --- src/compute/if_then_else.rs | 44 --- src/compute/sort/mod.rs | 619 +------------------------------ tests/it/compute/boolean.rs | 338 +++++++++++++++++ tests/it/compute/concat.rs | 117 ++++++ tests/it/compute/filter.rs | 179 +++++++++ tests/it/compute/hash.rs | 49 +++ tests/it/compute/if_then_else.rs | 42 +++ tests/it/compute/mod.rs | 6 + tests/it/compute/sort.rs | 603 ++++++++++++++++++++++++++++++ 13 files changed, 1336 insertions(+), 1536 deletions(-) create mode 100644 tests/it/compute/boolean.rs create mode 100644 tests/it/compute/concat.rs create mode 100644 tests/it/compute/filter.rs create mode 100644 tests/it/compute/hash.rs create mode 100644 tests/it/compute/if_then_else.rs create mode 100644 tests/it/compute/sort.rs diff --git a/src/compute/boolean.rs b/src/compute/boolean.rs index 03a55f48f27..9d547a2e750 100644 --- a/src/compute/boolean.rs +++ b/src/compute/boolean.rs @@ -145,346 +145,3 @@ pub fn is_not_null(input: &dyn Array) -> BooleanArray { }; BooleanArray::from_data(DataType::Boolean, values, None) } - -#[cfg(test)] -mod tests { - use super::*; - - use crate::array::*; - - #[test] - fn test_bool_array_and() { - let a = BooleanArray::from_slice(vec![false, false, true, true]); - let b = BooleanArray::from_slice(vec![false, true, false, true]); - let c = and(&a, &b).unwrap(); - - let expected = BooleanArray::from_slice(vec![false, false, false, true]); - - assert_eq!(c, expected); - } - - #[test] - fn test_bool_array_or() { - let a = BooleanArray::from_slice(vec![false, false, true, true]); - let b = BooleanArray::from_slice(vec![false, true, false, true]); - let c = or(&a, &b).unwrap(); - - let expected = BooleanArray::from_slice(vec![false, true, true, true]); - - assert_eq!(c, expected); - } - - #[test] - fn test_bool_array_or_validity() { - let a = BooleanArray::from(vec![ - None, - None, - None, - Some(false), - Some(false), - Some(false), - Some(true), - Some(true), - Some(true), - ]); - let b = BooleanArray::from(vec![ - None, - Some(false), - Some(true), - None, - Some(false), - Some(true), - None, - Some(false), - Some(true), - ]); - let c = or(&a, &b).unwrap(); - - let expected = BooleanArray::from(vec![ - None, - None, - None, - None, - Some(false), - Some(true), - None, - Some(true), - Some(true), - ]); - - assert_eq!(c, expected); - } - - #[test] - fn test_bool_array_not() { - let a = BooleanArray::from_slice(vec![false, true]); - let c = not(&a); - - let expected = BooleanArray::from_slice(vec![true, false]); - - assert_eq!(c, expected); - } - - #[test] - fn test_bool_array_and_validity() { - let a = BooleanArray::from(vec![ - None, - None, - None, - Some(false), - Some(false), - Some(false), - Some(true), - Some(true), - Some(true), - ]); - let b = BooleanArray::from(vec![ - None, - Some(false), - Some(true), - None, - Some(false), - Some(true), - None, - Some(false), - Some(true), - ]); - let c = and(&a, &b).unwrap(); - - let expected = BooleanArray::from(vec![ - None, - None, - None, - None, - Some(false), - Some(false), - None, - Some(false), - Some(true), - ]); - - assert_eq!(c, expected); - } - - #[test] - fn test_bool_array_and_sliced_same_offset() { - let a = BooleanArray::from_slice(vec![ - false, false, false, false, false, false, false, false, false, false, true, true, - ]); - let b = BooleanArray::from_slice(vec![ - false, false, false, false, false, false, false, false, false, true, false, true, - ]); - - let a = a.slice(8, 4); - let b = b.slice(8, 4); - let c = and(&a, &b).unwrap(); - - let expected = BooleanArray::from_slice(vec![false, false, false, true]); - - assert_eq!(expected, c); - } - - #[test] - fn test_bool_array_and_sliced_same_offset_mod8() { - let a = BooleanArray::from_slice(vec![ - false, false, true, true, false, false, false, false, false, false, false, false, - ]); - let b = BooleanArray::from_slice(vec![ - false, false, false, false, false, false, false, false, false, true, false, true, - ]); - - let a = a.slice(0, 4); - let b = b.slice(8, 4); - - let c = and(&a, &b).unwrap(); - - let expected = BooleanArray::from_slice(vec![false, false, false, true]); - - assert_eq!(expected, c); - } - - #[test] - fn test_bool_array_and_sliced_offset1() { - let a = BooleanArray::from_slice(vec![ - false, false, false, false, false, false, false, false, false, false, true, true, - ]); - let b = BooleanArray::from_slice(vec![false, true, false, true]); - - let a = a.slice(8, 4); - - let c = and(&a, &b).unwrap(); - - let expected = BooleanArray::from_slice(vec![false, false, false, true]); - - assert_eq!(expected, c); - } - - #[test] - fn test_bool_array_and_sliced_offset2() { - let a = BooleanArray::from_slice(vec![false, false, true, true]); - let b = BooleanArray::from_slice(vec![ - false, false, false, false, false, false, false, false, false, true, false, true, - ]); - - let b = b.slice(8, 4); - - let c = and(&a, &b).unwrap(); - - let expected = BooleanArray::from_slice(vec![false, false, false, true]); - - assert_eq!(expected, c); - } - - #[test] - fn test_bool_array_and_validity_offset() { - let a = BooleanArray::from(vec![None, Some(false), Some(true), None, Some(true)]); - let a = a.slice(1, 4); - let a = a.as_any().downcast_ref::().unwrap(); - - let b = BooleanArray::from(vec![ - None, - None, - Some(true), - Some(false), - Some(true), - Some(true), - ]); - - let b = b.slice(2, 4); - let b = b.as_any().downcast_ref::().unwrap(); - - let c = and(a, b).unwrap(); - - let expected = BooleanArray::from(vec![Some(false), Some(false), None, Some(true)]); - - assert_eq!(expected, c); - } - - #[test] - fn test_nonnull_array_is_null() { - let a = Int32Array::from_slice(&[1, 2, 3, 4]); - - let res = is_null(&a); - - let expected = BooleanArray::from_slice(vec![false, false, false, false]); - - assert_eq!(expected, res); - } - - #[test] - fn test_nonnull_array_with_offset_is_null() { - let a = Int32Array::from_slice(vec![1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1]); - let a = a.slice(8, 4); - - let res = is_null(&a); - - let expected = BooleanArray::from_slice(vec![false, false, false, false]); - - assert_eq!(expected, res); - } - - #[test] - fn test_nonnull_array_is_not_null() { - let a = Int32Array::from_slice(&[1, 2, 3, 4]); - - let res = is_not_null(&a); - - let expected = BooleanArray::from_slice(vec![true, true, true, true]); - - assert_eq!(expected, res); - } - - #[test] - fn test_nonnull_array_with_offset_is_not_null() { - let a = Int32Array::from_slice(&[1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1]); - let a = a.slice(8, 4); - - let res = is_not_null(&a); - - let expected = BooleanArray::from_slice(&[true, true, true, true]); - - assert_eq!(expected, res); - } - - #[test] - fn test_nullable_array_is_null() { - let a = Int32Array::from(vec![Some(1), None, Some(3), None]); - - let res = is_null(&a); - - let expected = BooleanArray::from_slice(vec![false, true, false, true]); - - assert_eq!(expected, res); - } - - #[test] - fn test_nullable_array_with_offset_is_null() { - let a = Int32Array::from(vec![ - None, - None, - None, - None, - None, - None, - None, - None, - // offset 8, previous None values are skipped by the slice - Some(1), - None, - Some(2), - None, - Some(3), - Some(4), - None, - None, - ]); - let a = a.slice(8, 4); - - let res = is_null(&a); - - let expected = BooleanArray::from_slice(vec![false, true, false, true]); - - assert_eq!(expected, res); - } - - #[test] - fn test_nullable_array_is_not_null() { - let a = Int32Array::from(vec![Some(1), None, Some(3), None]); - - let res = is_not_null(&a); - - let expected = BooleanArray::from_slice(vec![true, false, true, false]); - - assert_eq!(expected, res); - } - - #[test] - fn test_nullable_array_with_offset_is_not_null() { - let a = Int32Array::from(vec![ - None, - None, - None, - None, - None, - None, - None, - None, - // offset 8, previous None values are skipped by the slice - Some(1), - None, - Some(2), - None, - Some(3), - Some(4), - None, - None, - ]); - let a = a.slice(8, 4); - - let res = is_not_null(&a); - - let expected = BooleanArray::from_slice(vec![true, false, true, false]); - - assert_eq!(expected, res); - } -} diff --git a/src/compute/concat.rs b/src/compute/concat.rs index 665818abc96..b6c410c04e6 100644 --- a/src/compute/concat.rs +++ b/src/compute/concat.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -//! Defines concat kernel for `ArrayRef` +//! Contains the concatenate kernel //! //! Example: //! @@ -61,297 +61,3 @@ pub fn concatenate(arrays: &[&dyn Array]) -> Result> { Ok(mutable.as_box()) } - -#[cfg(test)] -mod tests { - use super::*; - - use crate::array::*; - - #[test] - fn test_concat_empty_vec() { - let re = concatenate(&[]); - assert!(re.is_err()); - } - - #[test] - fn test_concat_incompatible_datatypes() { - let re = concatenate(&[ - &Int64Array::from(vec![Some(-1), Some(2), None]), - &Utf8Array::::from(&vec![Some("hello"), Some("bar"), Some("world")]), - ]); - assert!(re.is_err()); - } - - #[test] - fn test_concat_string_arrays() -> Result<()> { - let arr = concatenate(&[ - &Utf8Array::::from_slice(&vec!["hello", "world"]), - &Utf8Array::::from_slice(&vec!["2", "3", "4"]), - &Utf8Array::::from(&vec![Some("foo"), Some("bar"), None, Some("baz")]), - ])?; - - let expected_output = Utf8Array::::from(&vec![ - Some("hello"), - Some("world"), - Some("2"), - Some("3"), - Some("4"), - Some("foo"), - Some("bar"), - None, - Some("baz"), - ]); - - assert_eq!(expected_output, arr.as_ref()); - - Ok(()) - } - - #[test] - fn test_concat_primitive_arrays() -> Result<()> { - let arr = concatenate(&[ - &Int64Array::from(&[Some(-1), Some(-1), Some(2), None, None]), - &Int64Array::from(&[Some(101), Some(102), Some(103), None]), - &Int64Array::from(&[Some(256), Some(512), Some(1024)]), - ])?; - - let expected_output = Int64Array::from(vec![ - Some(-1), - Some(-1), - Some(2), - None, - None, - Some(101), - Some(102), - Some(103), - None, - Some(256), - Some(512), - Some(1024), - ]); - - assert_eq!(expected_output, arr.as_ref()); - - Ok(()) - } - - #[test] - fn test_concat_primitive_array_slices() -> Result<()> { - let input_1 = Int64Array::from(&[Some(-1), Some(-1), Some(2), None, None]).slice(1, 3); - - let input_2 = Int64Array::from(&[Some(101), Some(102), Some(103), None]).slice(1, 3); - let arr = concatenate(&[&input_1, &input_2])?; - - let expected_output = - Int64Array::from(&[Some(-1), Some(2), None, Some(102), Some(103), None]); - - assert_eq!(expected_output, arr.as_ref()); - - Ok(()) - } - - #[test] - fn test_concat_boolean_primitive_arrays() -> Result<()> { - let arr = concatenate(&[ - &BooleanArray::from(vec![ - Some(true), - Some(true), - Some(false), - None, - None, - Some(false), - ]), - &BooleanArray::from(vec![None, Some(false), Some(true), Some(false)]), - ])?; - - let expected_output = BooleanArray::from(vec![ - Some(true), - Some(true), - Some(false), - None, - None, - Some(false), - None, - Some(false), - Some(true), - Some(false), - ]); - - assert_eq!(expected_output, arr.as_ref()); - - Ok(()) - } - - // todo: migrate me - /* - #[test] - fn test_concat_primitive_list_arrays() -> Result<()> { - let list1 = vec![ - Some(vec![Some(-1), Some(-1), Some(2), None, None]), - Some(vec![]), - None, - Some(vec![Some(10)]), - ]; - let list1_array = - ListArray::from_iter_primitive::(list1.clone()); - - let list2 = vec![ - None, - Some(vec![Some(100), None, Some(101)]), - Some(vec![Some(102)]), - ]; - let list2_array = - ListArray::from_iter_primitive::(list2.clone()); - - let list3 = vec![Some(vec![Some(1000), Some(1001)])]; - let list3_array = - ListArray::from_iter_primitive::(list3.clone()); - - let array_result = concatenate(&[&list1_array, &list2_array, &list3_array])?; - - let expected = list1 - .into_iter() - .chain(list2.into_iter()) - .chain(list3.into_iter()); - let array_expected = ListArray::from_iter_primitive::(expected); - - assert_eq!(array_result.as_ref(), &array_expected as &dyn Array); - - Ok(()) - } - - #[test] - fn test_concat_struct_arrays() -> Result<()> { - let field = Field::new("field", DataType::Int64, true); - let input_primitive_1: ArrayRef = - Arc::new(PrimitiveArray::::from(vec![ - Some(-1), - Some(-1), - Some(2), - None, - None, - ])); - let input_struct_1 = StructArray::from(vec![(field.clone(), input_primitive_1)]); - - let input_primitive_2: ArrayRef = - Arc::new(PrimitiveArray::::from(vec![ - Some(101), - Some(102), - Some(103), - None, - ])); - let input_struct_2 = StructArray::from(vec![(field.clone(), input_primitive_2)]); - - let input_primitive_3: ArrayRef = - Arc::new(PrimitiveArray::::from(vec![ - Some(256), - Some(512), - Some(1024), - ])); - let input_struct_3 = StructArray::from(vec![(field, input_primitive_3)]); - - let arr = concatenate(&[&input_struct_1, &input_struct_2, &input_struct_3])?; - - let expected_primitive_output = Arc::new(PrimitiveArray::::from(vec![ - Some(-1), - Some(-1), - Some(2), - None, - None, - Some(101), - Some(102), - Some(103), - None, - Some(256), - Some(512), - Some(1024), - ])) as ArrayRef; - - let actual_primitive = arr - .as_any() - .downcast_ref::() - .unwrap() - .column(0); - assert_eq!(actual_primitive, &expected_primitive_output); - - Ok(()) - } - - #[test] - fn test_concat_struct_array_slices() -> Result<()> { - let field = Field::new("field", DataType::Int64, true); - let input_primitive_1: ArrayRef = - Arc::new(PrimitiveArray::::from(vec![ - Some(-1), - Some(-1), - Some(2), - None, - None, - ])); - let input_struct_1 = StructArray::from(vec![(field.clone(), input_primitive_1)]); - - let input_primitive_2: ArrayRef = - Arc::new(PrimitiveArray::::from(vec![ - Some(101), - Some(102), - Some(103), - None, - ])); - let input_struct_2 = StructArray::from(vec![(field, input_primitive_2)]); - - let arr = concatenate(&[ - input_struct_1.slice(1, 3).as_ref(), - input_struct_2.slice(1, 2).as_ref(), - ])?; - - let expected_primitive_output = Arc::new(PrimitiveArray::::from(vec![ - Some(-1), - Some(2), - None, - Some(102), - Some(103), - ])) as ArrayRef; - - let actual_primitive = arr - .as_any() - .downcast_ref::() - .unwrap() - .column(0); - assert_eq!(actual_primitive, &expected_primitive_output); - - Ok(()) - } - - #[test] - fn test_string_array_slices() -> Result<()> { - let input_1 = StringArray::from(vec!["hello", "A", "B", "C"]); - let input_2 = StringArray::from(vec!["world", "D", "E", "Z"]); - - let arr = concatenate(&[input_1.slice(1, 3).as_ref(), input_2.slice(1, 2).as_ref()])?; - - let expected_output = StringArray::from(vec!["A", "B", "C", "D", "E"]); - - let actual_output = arr.as_any().downcast_ref::().unwrap(); - assert_eq!(actual_output, &expected_output); - - Ok(()) - } - - #[test] - fn test_string_array_with_null_slices() -> Result<()> { - let input_1 = StringArray::from(vec![Some("hello"), None, Some("A"), Some("C")]); - let input_2 = StringArray::from(vec![None, Some("world"), Some("D"), None]); - - let arr = concatenate(&[input_1.slice(1, 3).as_ref(), input_2.slice(1, 2).as_ref()])?; - - let expected_output = - StringArray::from(vec![None, Some("A"), Some("C"), Some("world"), Some("D")]); - - let actual_output = arr.as_any().downcast_ref::().unwrap(); - assert_eq!(actual_output, &expected_output); - - Ok(()) - } - */ -} diff --git a/src/compute/filter.rs b/src/compute/filter.rs index 1ad87a8a998..12ac2a2ada0 100644 --- a/src/compute/filter.rs +++ b/src/compute/filter.rs @@ -182,185 +182,3 @@ pub fn filter_record_batch( }; RecordBatch::try_new(record_batch.schema().clone(), filtered_arrays) } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_filter_array_slice() { - let a = Int32Array::from_slice(&[5, 6, 7, 8, 9]).slice(1, 4); - let b = BooleanArray::from_slice(vec![true, true, false, false, true]).slice(1, 4); - let c = filter(&a, &b).unwrap(); - - let expected = Int32Array::from_slice(&[6, 9]); - - assert_eq!(expected, c.as_ref()); - } - - #[test] - fn test_filter_array_low_density() { - // this test exercises the all 0's branch of the filter algorithm - let mut data_values = (1..=65).collect::>(); - let mut filter_values = (1..=65).map(|i| matches!(i % 65, 0)).collect::>(); - // set up two more values after the batch - data_values.extend_from_slice(&[66, 67]); - filter_values.extend_from_slice(&[false, true]); - let a = Int32Array::from_slice(&data_values); - let b = BooleanArray::from_slice(filter_values); - let c = filter(&a, &b).unwrap(); - - let expected = Int32Array::from_slice(&[65, 67]); - - assert_eq!(expected, c.as_ref()); - } - - #[test] - fn test_filter_array_high_density() { - // this test exercises the all 1's branch of the filter algorithm - let mut data_values = (1..=65).map(Some).collect::>(); - let mut filter_values = (1..=65) - .map(|i| !matches!(i % 65, 0)) - .collect::>(); - // set second data value to null - data_values[1] = None; - // set up two more values after the batch - data_values.extend_from_slice(&[Some(66), None, Some(67), None]); - filter_values.extend_from_slice(&[false, true, true, true]); - let a = Int32Array::from(data_values); - let b = BooleanArray::from_slice(filter_values); - let c = filter(&a, &b).unwrap(); - let d = c.as_ref().as_any().downcast_ref::().unwrap(); - assert_eq!(67, d.len()); - assert_eq!(3, d.null_count()); - assert_eq!(1, d.value(0)); - assert!(d.is_null(1)); - assert_eq!(64, d.value(63)); - assert!(d.is_null(64)); - assert_eq!(67, d.value(65)); - } - - #[test] - fn test_filter_string_array_simple() { - let a = Utf8Array::::from_slice(&["hello", " ", "world", "!"]); - let b = BooleanArray::from_slice(&[true, false, true, false]); - let c = filter(&a, &b).unwrap(); - let d = c - .as_ref() - .as_any() - .downcast_ref::>() - .unwrap(); - assert_eq!(2, d.len()); - assert_eq!("hello", d.value(0)); - assert_eq!("world", d.value(1)); - } - - #[test] - fn test_filter_primative_array_with_null() { - let a = Int32Array::from(&[Some(5), None]); - let b = BooleanArray::from_slice(vec![false, true]); - let c = filter(&a, &b).unwrap(); - let d = c.as_ref().as_any().downcast_ref::().unwrap(); - assert_eq!(1, d.len()); - assert!(d.is_null(0)); - } - - #[test] - fn test_filter_string_array_with_null() { - let a = Utf8Array::::from(&vec![Some("hello"), None, Some("world"), None]); - let b = BooleanArray::from_slice(vec![true, false, false, true]); - let c = filter(&a, &b).unwrap(); - let d = c - .as_ref() - .as_any() - .downcast_ref::>() - .unwrap(); - assert_eq!(2, d.len()); - assert_eq!("hello", d.value(0)); - assert!(!d.is_null(0)); - assert!(d.is_null(1)); - } - - #[test] - fn test_filter_binary_array_with_null() { - let data: Vec> = vec![Some(b"hello"), None, Some(b"world"), None]; - let a = BinaryArray::::from(&data); - let b = BooleanArray::from_slice(vec![true, false, false, true]); - let c = filter(&a, &b).unwrap(); - let d = c - .as_ref() - .as_any() - .downcast_ref::>() - .unwrap(); - assert_eq!(2, d.len()); - assert_eq!(b"hello", d.value(0)); - assert!(!d.is_null(0)); - assert!(d.is_null(1)); - } - - /* - #[test] - fn test_filter_dictionary_array() { - let values = vec![Some("hello"), None, Some("world"), Some("!")]; - let a: Int8DictionaryArray = values.iter().copied().collect(); - let b = BooleanArray::from(vec![false, true, true, false]); - let c = filter(&a, &b).unwrap(); - let d = c - .as_ref() - .as_any() - .downcast_ref::() - .unwrap(); - let value_array = d.values(); - let values = value_array.as_any().downcast_ref::().unwrap(); - // values are cloned in the filtered dictionary array - assert_eq!(3, values.len()); - // but keys are filtered - assert_eq!(2, d.len()); - assert_eq!(true, d.is_null(0)); - assert_eq!("world", values.value(d.keys().value(1) as usize)); - } - - #[test] - fn test_filter_list_array() { - let value_data = ArrayData::builder(DataType::Int32) - .len(8) - .add_buffer(Buffer::from_slice_ref(&[0, 1, 2, 3, 4, 5, 6, 7])) - .build(); - - let value_offsets = Buffer::from_slice_ref(&[0i64, 3, 6, 8, 8]); - - let list_data_type = - DataType::LargeList(Box::new(Field::new("item", DataType::Int32, false))); - let list_data = ArrayData::builder(list_data_type) - .len(4) - .add_buffer(value_offsets) - .add_child_data(value_data) - .null_bit_buffer(Buffer::from([0b00000111])) - .build(); - - // a = [[0, 1, 2], [3, 4, 5], [6, 7], null] - let a = LargeListArray::from(list_data); - let b = BooleanArray::from(vec![false, true, false, true]); - let result = filter(&a, &b).unwrap(); - - // expected: [[3, 4, 5], null] - let value_data = ArrayData::builder(DataType::Int32) - .len(3) - .add_buffer(Buffer::from_slice_ref(&[3, 4, 5])) - .build(); - - let value_offsets = Buffer::from_slice_ref(&[0i64, 3, 3]); - - let list_data_type = - DataType::LargeList(Box::new(Field::new("item", DataType::Int32, false))); - let expected = ArrayData::builder(list_data_type) - .len(2) - .add_buffer(value_offsets) - .add_child_data(value_data) - .null_bit_buffer(Buffer::from([0b00000001])) - .build(); - - assert_eq!(&make_array(expected), &result); - } - */ -} diff --git a/src/compute/hash.rs b/src/compute/hash.rs index 3bf84f9c315..50133c7870f 100644 --- a/src/compute/hash.rs +++ b/src/compute/hash.rs @@ -162,57 +162,3 @@ pub fn can_hash(data_type: &DataType) -> bool { | DataType::LargeUtf8 ) } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn consistency() { - use crate::array::new_null_array; - use crate::datatypes::DataType::*; - use crate::datatypes::TimeUnit; - - let datatypes = vec![ - Null, - Boolean, - UInt8, - UInt16, - UInt32, - UInt64, - Int8, - Int16, - Int32, - Int64, - Float32, - Float64, - Timestamp(TimeUnit::Second, None), - Timestamp(TimeUnit::Millisecond, None), - Timestamp(TimeUnit::Microsecond, None), - Timestamp(TimeUnit::Nanosecond, None), - Time64(TimeUnit::Microsecond), - Time64(TimeUnit::Nanosecond), - Date32, - Time32(TimeUnit::Second), - Time32(TimeUnit::Millisecond), - Date64, - Utf8, - LargeUtf8, - Binary, - LargeBinary, - Duration(TimeUnit::Second), - Duration(TimeUnit::Millisecond), - Duration(TimeUnit::Microsecond), - Duration(TimeUnit::Nanosecond), - ]; - - datatypes.into_iter().for_each(|d1| { - let array = new_null_array(d1.clone(), 10); - if can_hash(&d1) { - assert!(hash(array.as_ref()).is_ok()); - } else { - assert!(hash(array.as_ref()).is_err()); - } - }); - } -} diff --git a/src/compute/if_then_else.rs b/src/compute/if_then_else.rs index 93e2e11d0aa..9cfb1cc6f8e 100644 --- a/src/compute/if_then_else.rs +++ b/src/compute/if_then_else.rs @@ -72,47 +72,3 @@ pub fn if_then_else( }; Ok(result) } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn basics() -> Result<()> { - let lhs = Int32Array::from_slice(&[1, 2, 3]); - let rhs = Int32Array::from_slice(&[4, 5, 6]); - let predicate = BooleanArray::from_slice(vec![true, false, true]); - let c = if_then_else(&predicate, &lhs, &rhs)?; - - let expected = Int32Array::from_slice(&[1, 5, 3]); - - assert_eq!(expected, c.as_ref()); - Ok(()) - } - - #[test] - fn basics_nulls() -> Result<()> { - let lhs = Int32Array::from(&[Some(1), None, None]); - let rhs = Int32Array::from(&[None, Some(5), Some(6)]); - let predicate = BooleanArray::from_slice(vec![true, false, true]); - let c = if_then_else(&predicate, &lhs, &rhs)?; - - let expected = Int32Array::from(&[Some(1), Some(5), None]); - - assert_eq!(expected, c.as_ref()); - Ok(()) - } - - #[test] - fn basics_nulls_pred() -> Result<()> { - let lhs = Int32Array::from_slice(&[1, 2, 3]); - let rhs = Int32Array::from_slice(&[4, 5, 6]); - let predicate = BooleanArray::from(&[Some(true), None, Some(false)]); - let result = if_then_else(&predicate, &lhs, &rhs)?; - - let expected = Int32Array::from(&[Some(1), None, Some(6)]); - - assert_eq!(expected, result.as_ref()); - Ok(()) - } -} diff --git a/src/compute/sort/mod.rs b/src/compute/sort/mod.rs index 036bc8b4e84..d7c0a08046d 100644 --- a/src/compute/sort/mod.rs +++ b/src/compute/sort/mod.rs @@ -1,6 +1,7 @@ use std::cmp::Ordering; use crate::array::ord; +use crate::buffer::MutableBuffer; use crate::compute::take; use crate::datatypes::*; use crate::error::{ArrowError, Result}; @@ -9,8 +10,6 @@ use crate::{ types::{Index, NativeType}, }; -use crate::buffer::MutableBuffer; - mod binary; mod boolean; mod common; @@ -398,619 +397,3 @@ fn cmp_array(a: &dyn Array, b: &dyn Array) -> Ordering { } Ordering::Equal } - -#[cfg(test)] -mod tests { - use super::*; - - fn test_sort_to_indices_boolean_arrays( - data: &[Option], - options: SortOptions, - expected_data: &[i32], - ) { - let output = BooleanArray::from(data); - let expected = Int32Array::from_slice(expected_data); - let output = sort_to_indices(&output, &options, None).unwrap(); - assert_eq!(output, expected) - } - - fn test_sort_primitive_arrays( - data: &[Option], - data_type: DataType, - options: SortOptions, - expected_data: &[Option], - ) where - T: NativeType, - { - let input = PrimitiveArray::::from(data).to(data_type.clone()); - let expected = PrimitiveArray::::from(expected_data).to(data_type); - let output = sort(&input, &options, None).unwrap(); - assert_eq!(expected, output.as_ref()) - } - - fn test_sort_to_indices_string_arrays( - data: &[Option<&str>], - options: SortOptions, - expected_data: &[i32], - ) { - let input = Utf8Array::::from(&data.to_vec()); - let expected = Int32Array::from_slice(expected_data); - let output = sort_to_indices(&input, &options, None).unwrap(); - assert_eq!(output, expected) - } - - fn test_sort_string_arrays( - data: &[Option<&str>], - options: SortOptions, - expected_data: &[Option<&str>], - ) { - let input = Utf8Array::::from(&data.to_vec()); - let expected = Utf8Array::::from(&expected_data.to_vec()); - let output = sort(&input, &options, None).unwrap(); - assert_eq!(expected, output.as_ref()) - } - - fn test_sort_string_dict_arrays( - data: &[Option<&str>], - options: SortOptions, - expected_data: &[Option<&str>], - ) { - let mut input = MutableDictionaryArray::>::new(); - input.try_extend(data.iter().copied()).unwrap(); - let input = input.into_arc(); - - let mut expected = MutableDictionaryArray::>::new(); - expected.try_extend(expected_data.iter().copied()).unwrap(); - let expected = expected.into_arc(); - - let output = sort(input.as_ref(), &options, None).unwrap(); - assert_eq!(expected.as_ref(), output.as_ref()) - } - - /* - fn test_sort_list_arrays( - data: Vec>>>, - options: Option, - expected_data: Vec>>>, - fixed_length: Option, - ) where - T: ArrowPrimitiveType, - PrimitiveArray: From>>, - { - // for FixedSizedList - if let Some(length) = fixed_length { - let input = Arc::new(build_fixed_size_list_nullable(data.clone(), length)); - let sorted = sort(&(input as ArrayRef), options).unwrap(); - let expected = Arc::new(build_fixed_size_list_nullable( - expected_data.clone(), - length, - )) as ArrayRef; - - assert_eq!(&sorted, &expected); - } - - // for List - let input = Arc::new(build_generic_list_nullable::(data.clone())); - let sorted = sort(&(input as ArrayRef), options).unwrap(); - let expected = - Arc::new(build_generic_list_nullable::(expected_data.clone())) - as ArrayRef; - - assert_eq!(&sorted, &expected); - - // for LargeList - let input = Arc::new(build_generic_list_nullable::(data)); - let sorted = sort(&(input as ArrayRef), options).unwrap(); - let expected = - Arc::new(build_generic_list_nullable::(expected_data)) as ArrayRef; - - assert_eq!(&sorted, &expected); - } - - fn test_lex_sort_arrays(input: Vec, expected_output: Vec) { - let sorted = lexsort(&input).unwrap(); - - for (result, expected) in sorted.iter().zip(expected_output.iter()) { - assert_eq!(result, expected); - } - } - */ - - #[test] - fn test_sort_boolean() { - // boolean - test_sort_to_indices_boolean_arrays( - &[None, Some(false), Some(true), Some(true), Some(false), None], - SortOptions { - descending: false, - nulls_first: true, - }, - &[0, 5, 1, 4, 2, 3], - ); - - // boolean, descending - test_sort_to_indices_boolean_arrays( - &[None, Some(false), Some(true), Some(true), Some(false), None], - SortOptions { - descending: true, - nulls_first: false, - }, - &[2, 3, 1, 4, 5, 0], - ); - - // boolean, descending, nulls first - test_sort_to_indices_boolean_arrays( - &[None, Some(false), Some(true), Some(true), Some(false), None], - SortOptions { - descending: true, - nulls_first: true, - }, - &[5, 0, 2, 3, 1, 4], - ); - } - - #[test] - #[ignore] // improve equality for NaN values. These are right but the equality fails - fn test_nans() { - test_sort_primitive_arrays::( - &[None, Some(0.0), Some(2.0), Some(-1.0), Some(f64::NAN), None], - DataType::Float64, - SortOptions { - descending: true, - nulls_first: true, - }, - &[None, None, Some(f64::NAN), Some(2.0), Some(0.0), Some(-1.0)], - ); - test_sort_primitive_arrays::( - &[Some(f64::NAN), Some(f64::NAN), Some(f64::NAN), Some(1.0)], - DataType::Float64, - SortOptions { - descending: true, - nulls_first: true, - }, - &[Some(f64::NAN), Some(f64::NAN), Some(f64::NAN), Some(1.0)], - ); - - test_sort_primitive_arrays::( - &[None, Some(0.0), Some(2.0), Some(-1.0), Some(f64::NAN), None], - DataType::Float64, - SortOptions { - descending: false, - nulls_first: true, - }, - &[None, None, Some(-1.0), Some(0.0), Some(2.0), Some(f64::NAN)], - ); - // nans - test_sort_primitive_arrays::( - &[Some(f64::NAN), Some(f64::NAN), Some(f64::NAN), Some(1.0)], - DataType::Float64, - SortOptions { - descending: false, - nulls_first: true, - }, - &[Some(1.0), Some(f64::NAN), Some(f64::NAN), Some(f64::NAN)], - ); - } - - #[test] - fn test_sort_to_indices_strings() { - test_sort_to_indices_string_arrays( - &[ - None, - Some("bad"), - Some("sad"), - None, - Some("glad"), - Some("-ad"), - ], - SortOptions { - descending: false, - nulls_first: true, - }, - // &[3, 0, 5, 1, 4, 2] is also valid - &[0, 3, 5, 1, 4, 2], - ); - - test_sort_to_indices_string_arrays( - &[ - None, - Some("bad"), - Some("sad"), - None, - Some("glad"), - Some("-ad"), - ], - SortOptions { - descending: true, - nulls_first: false, - }, - // &[2, 4, 1, 5, 3, 0] is also valid - &[2, 4, 1, 5, 0, 3], - ); - - test_sort_to_indices_string_arrays( - &[ - None, - Some("bad"), - Some("sad"), - None, - Some("glad"), - Some("-ad"), - ], - SortOptions { - descending: false, - nulls_first: true, - }, - // &[3, 0, 5, 1, 4, 2] is also valid - &[0, 3, 5, 1, 4, 2], - ); - - test_sort_to_indices_string_arrays( - &[ - None, - Some("bad"), - Some("sad"), - None, - Some("glad"), - Some("-ad"), - ], - SortOptions { - descending: true, - nulls_first: true, - }, - // &[3, 0, 2, 4, 1, 5] is also valid - &[0, 3, 2, 4, 1, 5], - ); - } - - #[test] - fn test_sort_strings() { - test_sort_string_arrays( - &[ - None, - Some("bad"), - Some("sad"), - None, - Some("glad"), - Some("-ad"), - ], - SortOptions { - descending: false, - nulls_first: true, - }, - &[ - None, - None, - Some("-ad"), - Some("bad"), - Some("glad"), - Some("sad"), - ], - ); - - test_sort_string_arrays( - &[ - None, - Some("bad"), - Some("sad"), - None, - Some("glad"), - Some("-ad"), - ], - SortOptions { - descending: true, - nulls_first: false, - }, - &[ - Some("sad"), - Some("glad"), - Some("bad"), - Some("-ad"), - None, - None, - ], - ); - - test_sort_string_arrays( - &[ - None, - Some("bad"), - Some("sad"), - None, - Some("glad"), - Some("-ad"), - ], - SortOptions { - descending: false, - nulls_first: true, - }, - &[ - None, - None, - Some("-ad"), - Some("bad"), - Some("glad"), - Some("sad"), - ], - ); - - test_sort_string_arrays( - &[ - None, - Some("bad"), - Some("sad"), - None, - Some("glad"), - Some("-ad"), - ], - SortOptions { - descending: true, - nulls_first: true, - }, - &[ - None, - None, - Some("sad"), - Some("glad"), - Some("bad"), - Some("-ad"), - ], - ); - } - - #[test] - fn test_sort_string_dicts() { - test_sort_string_dict_arrays::( - &[ - None, - Some("bad"), - Some("sad"), - None, - Some("glad"), - Some("-ad"), - ], - SortOptions { - descending: false, - nulls_first: true, - }, - &[ - None, - None, - Some("-ad"), - Some("bad"), - Some("glad"), - Some("sad"), - ], - ); - - test_sort_string_dict_arrays::( - &[ - None, - Some("bad"), - Some("sad"), - None, - Some("glad"), - Some("-ad"), - ], - SortOptions { - descending: true, - nulls_first: false, - }, - &[ - Some("sad"), - Some("glad"), - Some("bad"), - Some("-ad"), - None, - None, - ], - ); - - test_sort_string_dict_arrays::( - &[ - None, - Some("bad"), - Some("sad"), - None, - Some("glad"), - Some("-ad"), - ], - SortOptions { - descending: false, - nulls_first: true, - }, - &[ - None, - None, - Some("-ad"), - Some("bad"), - Some("glad"), - Some("sad"), - ], - ); - - test_sort_string_dict_arrays::( - &[ - None, - Some("bad"), - Some("sad"), - None, - Some("glad"), - Some("-ad"), - ], - SortOptions { - descending: true, - nulls_first: true, - }, - &[ - None, - None, - Some("sad"), - Some("glad"), - Some("bad"), - Some("-ad"), - ], - ); - } - - /* - #[test] - fn test_sort_list() { - test_sort_list_arrays::( - vec![ - Some(vec![Some(1)]), - Some(vec![Some(4)]), - Some(vec![Some(2)]), - Some(vec![Some(3)]), - ], - Some(SortOptions { - descending: false, - nulls_first: false, - }), - vec![ - Some(vec![Some(1)]), - Some(vec![Some(2)]), - Some(vec![Some(3)]), - Some(vec![Some(4)]), - ], - Some(1), - ); - - test_sort_list_arrays::( - vec![ - Some(vec![Some(1), Some(0)]), - Some(vec![Some(4), Some(3), Some(2), Some(1)]), - Some(vec![Some(2), Some(3), Some(4)]), - Some(vec![Some(3), Some(3), Some(3), Some(3)]), - Some(vec![Some(1), Some(1)]), - ], - Some(SortOptions { - descending: false, - nulls_first: false, - }), - vec![ - Some(vec![Some(1), Some(0)]), - Some(vec![Some(1), Some(1)]), - Some(vec![Some(2), Some(3), Some(4)]), - Some(vec![Some(3), Some(3), Some(3), Some(3)]), - Some(vec![Some(4), Some(3), Some(2), Some(1)]), - ], - None, - ); - - test_sort_list_arrays::( - vec![ - None, - Some(vec![Some(4), None, Some(2)]), - Some(vec![Some(2), Some(3), Some(4)]), - None, - Some(vec![Some(3), Some(3), None]), - ], - Some(SortOptions { - descending: false, - nulls_first: false, - }), - vec![ - Some(vec![Some(2), Some(3), Some(4)]), - Some(vec![Some(3), Some(3), None]), - Some(vec![Some(4), None, Some(2)]), - None, - None, - ], - Some(3), - ); - } - - #[test] - fn test_lex_sort_single_column() { - let input = vec![SortColumn { - values: Arc::new(PrimitiveArray::::from(vec![ - Some(17), - Some(2), - Some(-1), - Some(0), - ])) as ArrayRef, - options: None, - }]; - let expected = vec![Arc::new(PrimitiveArray::::from(vec![ - Some(-1), - Some(0), - Some(2), - Some(17), - ])) as ArrayRef]; - test_lex_sort_arrays(input, expected); - } - - #[test] - fn test_lex_sort_unaligned_rows() { - let input = vec![ - SortColumn { - values: Arc::new(PrimitiveArray::::from(vec![None, Some(-1)])) - as ArrayRef, - options: None, - }, - SortColumn { - values: Arc::new(StringArray::from(vec![Some("foo")])) as ArrayRef, - options: None, - }, - ]; - assert!( - lexsort(&input).is_err(), - "lexsort should reject columns with different row counts" - ); - } - */ - - #[test] - fn consistency() { - use crate::array::new_null_array; - use crate::datatypes::DataType::*; - use crate::datatypes::TimeUnit; - - let datatypes = vec![ - Null, - Boolean, - UInt8, - UInt16, - UInt32, - UInt64, - Int8, - Int16, - Int32, - Int64, - Float32, - Float64, - Timestamp(TimeUnit::Second, None), - Timestamp(TimeUnit::Millisecond, None), - Timestamp(TimeUnit::Microsecond, None), - Timestamp(TimeUnit::Nanosecond, None), - Time64(TimeUnit::Microsecond), - Time64(TimeUnit::Nanosecond), - Date32, - Time32(TimeUnit::Second), - Time32(TimeUnit::Millisecond), - Date64, - Utf8, - LargeUtf8, - Binary, - LargeBinary, - Duration(TimeUnit::Second), - Duration(TimeUnit::Millisecond), - Duration(TimeUnit::Microsecond), - Duration(TimeUnit::Nanosecond), - ]; - - datatypes.into_iter().for_each(|d1| { - let array = new_null_array(d1.clone(), 10); - let options = SortOptions { - descending: true, - nulls_first: true, - }; - if can_sort(&d1) { - assert!(sort(array.as_ref(), &options, None).is_ok()); - } else { - assert!(sort(array.as_ref(), &options, None).is_err()); - } - }); - } -} diff --git a/tests/it/compute/boolean.rs b/tests/it/compute/boolean.rs new file mode 100644 index 00000000000..5468909df10 --- /dev/null +++ b/tests/it/compute/boolean.rs @@ -0,0 +1,338 @@ +use arrow2::array::*; +use arrow2::compute::boolean::*; + +#[test] +fn array_and() { + let a = BooleanArray::from_slice(vec![false, false, true, true]); + let b = BooleanArray::from_slice(vec![false, true, false, true]); + let c = and(&a, &b).unwrap(); + + let expected = BooleanArray::from_slice(vec![false, false, false, true]); + + assert_eq!(c, expected); +} + +#[test] +fn array_or() { + let a = BooleanArray::from_slice(vec![false, false, true, true]); + let b = BooleanArray::from_slice(vec![false, true, false, true]); + let c = or(&a, &b).unwrap(); + + let expected = BooleanArray::from_slice(vec![false, true, true, true]); + + assert_eq!(c, expected); +} + +#[test] +fn array_or_validity() { + let a = BooleanArray::from(vec![ + None, + None, + None, + Some(false), + Some(false), + Some(false), + Some(true), + Some(true), + Some(true), + ]); + let b = BooleanArray::from(vec![ + None, + Some(false), + Some(true), + None, + Some(false), + Some(true), + None, + Some(false), + Some(true), + ]); + let c = or(&a, &b).unwrap(); + + let expected = BooleanArray::from(vec![ + None, + None, + None, + None, + Some(false), + Some(true), + None, + Some(true), + Some(true), + ]); + + assert_eq!(c, expected); +} + +#[test] +fn array_not() { + let a = BooleanArray::from_slice(vec![false, true]); + let c = not(&a); + + let expected = BooleanArray::from_slice(vec![true, false]); + + assert_eq!(c, expected); +} + +#[test] +fn array_and_validity() { + let a = BooleanArray::from(vec![ + None, + None, + None, + Some(false), + Some(false), + Some(false), + Some(true), + Some(true), + Some(true), + ]); + let b = BooleanArray::from(vec![ + None, + Some(false), + Some(true), + None, + Some(false), + Some(true), + None, + Some(false), + Some(true), + ]); + let c = and(&a, &b).unwrap(); + + let expected = BooleanArray::from(vec![ + None, + None, + None, + None, + Some(false), + Some(false), + None, + Some(false), + Some(true), + ]); + + assert_eq!(c, expected); +} + +#[test] +fn array_and_sliced_same_offset() { + let a = BooleanArray::from_slice(vec![ + false, false, false, false, false, false, false, false, false, false, true, true, + ]); + let b = BooleanArray::from_slice(vec![ + false, false, false, false, false, false, false, false, false, true, false, true, + ]); + + let a = a.slice(8, 4); + let b = b.slice(8, 4); + let c = and(&a, &b).unwrap(); + + let expected = BooleanArray::from_slice(vec![false, false, false, true]); + + assert_eq!(expected, c); +} + +#[test] +fn array_and_sliced_same_offset_mod8() { + let a = BooleanArray::from_slice(vec![ + false, false, true, true, false, false, false, false, false, false, false, false, + ]); + let b = BooleanArray::from_slice(vec![ + false, false, false, false, false, false, false, false, false, true, false, true, + ]); + + let a = a.slice(0, 4); + let b = b.slice(8, 4); + + let c = and(&a, &b).unwrap(); + + let expected = BooleanArray::from_slice(vec![false, false, false, true]); + + assert_eq!(expected, c); +} + +#[test] +fn array_and_sliced_offset1() { + let a = BooleanArray::from_slice(vec![ + false, false, false, false, false, false, false, false, false, false, true, true, + ]); + let b = BooleanArray::from_slice(vec![false, true, false, true]); + + let a = a.slice(8, 4); + + let c = and(&a, &b).unwrap(); + + let expected = BooleanArray::from_slice(vec![false, false, false, true]); + + assert_eq!(expected, c); +} + +#[test] +fn array_and_sliced_offset2() { + let a = BooleanArray::from_slice(vec![false, false, true, true]); + let b = BooleanArray::from_slice(vec![ + false, false, false, false, false, false, false, false, false, true, false, true, + ]); + + let b = b.slice(8, 4); + + let c = and(&a, &b).unwrap(); + + let expected = BooleanArray::from_slice(vec![false, false, false, true]); + + assert_eq!(expected, c); +} + +#[test] +fn array_and_validity_offset() { + let a = BooleanArray::from(vec![None, Some(false), Some(true), None, Some(true)]); + let a = a.slice(1, 4); + let a = a.as_any().downcast_ref::().unwrap(); + + let b = BooleanArray::from(vec![ + None, + None, + Some(true), + Some(false), + Some(true), + Some(true), + ]); + + let b = b.slice(2, 4); + let b = b.as_any().downcast_ref::().unwrap(); + + let c = and(a, b).unwrap(); + + let expected = BooleanArray::from(vec![Some(false), Some(false), None, Some(true)]); + + assert_eq!(expected, c); +} + +#[test] +fn test_nonnull_array_is_null() { + let a = Int32Array::from_slice(&[1, 2, 3, 4]); + + let res = is_null(&a); + + let expected = BooleanArray::from_slice(vec![false, false, false, false]); + + assert_eq!(expected, res); +} + +#[test] +fn test_nonnull_array_with_offset_is_null() { + let a = Int32Array::from_slice(vec![1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1]); + let a = a.slice(8, 4); + + let res = is_null(&a); + + let expected = BooleanArray::from_slice(vec![false, false, false, false]); + + assert_eq!(expected, res); +} + +#[test] +fn test_nonnull_array_is_not_null() { + let a = Int32Array::from_slice(&[1, 2, 3, 4]); + + let res = is_not_null(&a); + + let expected = BooleanArray::from_slice(vec![true, true, true, true]); + + assert_eq!(expected, res); +} + +#[test] +fn test_nonnull_array_with_offset_is_not_null() { + let a = Int32Array::from_slice(&[1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1]); + let a = a.slice(8, 4); + + let res = is_not_null(&a); + + let expected = BooleanArray::from_slice(&[true, true, true, true]); + + assert_eq!(expected, res); +} + +#[test] +fn test_nullable_array_is_null() { + let a = Int32Array::from(vec![Some(1), None, Some(3), None]); + + let res = is_null(&a); + + let expected = BooleanArray::from_slice(vec![false, true, false, true]); + + assert_eq!(expected, res); +} + +#[test] +fn test_nullable_array_with_offset_is_null() { + let a = Int32Array::from(vec![ + None, + None, + None, + None, + None, + None, + None, + None, + // offset 8, previous None values are skipped by the slice + Some(1), + None, + Some(2), + None, + Some(3), + Some(4), + None, + None, + ]); + let a = a.slice(8, 4); + + let res = is_null(&a); + + let expected = BooleanArray::from_slice(vec![false, true, false, true]); + + assert_eq!(expected, res); +} + +#[test] +fn test_nullable_array_is_not_null() { + let a = Int32Array::from(vec![Some(1), None, Some(3), None]); + + let res = is_not_null(&a); + + let expected = BooleanArray::from_slice(vec![true, false, true, false]); + + assert_eq!(expected, res); +} + +#[test] +fn test_nullable_array_with_offset_is_not_null() { + let a = Int32Array::from(vec![ + None, + None, + None, + None, + None, + None, + None, + None, + // offset 8, previous None values are skipped by the slice + Some(1), + None, + Some(2), + None, + Some(3), + Some(4), + None, + None, + ]); + let a = a.slice(8, 4); + + let res = is_not_null(&a); + + let expected = BooleanArray::from_slice(vec![true, false, true, false]); + + assert_eq!(expected, res); +} diff --git a/tests/it/compute/concat.rs b/tests/it/compute/concat.rs new file mode 100644 index 00000000000..80b3fa7c100 --- /dev/null +++ b/tests/it/compute/concat.rs @@ -0,0 +1,117 @@ +use arrow2::array::*; +use arrow2::compute::concat::concatenate; +use arrow2::error::Result; + +#[test] +fn empty_vec() { + let re = concatenate(&[]); + assert!(re.is_err()); +} + +#[test] +fn incompatible_datatypes() { + let re = concatenate(&[ + &Int64Array::from(vec![Some(-1), Some(2), None]), + &Utf8Array::::from(&vec![Some("hello"), Some("bar"), Some("world")]), + ]); + assert!(re.is_err()); +} + +#[test] +fn string_arrays() -> Result<()> { + let arr = concatenate(&[ + &Utf8Array::::from_slice(&vec!["hello", "world"]), + &Utf8Array::::from_slice(&vec!["2", "3", "4"]), + &Utf8Array::::from(&vec![Some("foo"), Some("bar"), None, Some("baz")]), + ])?; + + let expected_output = Utf8Array::::from(&vec![ + Some("hello"), + Some("world"), + Some("2"), + Some("3"), + Some("4"), + Some("foo"), + Some("bar"), + None, + Some("baz"), + ]); + + assert_eq!(expected_output, arr.as_ref()); + + Ok(()) +} + +#[test] +fn primitive_arrays() -> Result<()> { + let arr = concatenate(&[ + &Int64Array::from(&[Some(-1), Some(-1), Some(2), None, None]), + &Int64Array::from(&[Some(101), Some(102), Some(103), None]), + &Int64Array::from(&[Some(256), Some(512), Some(1024)]), + ])?; + + let expected_output = Int64Array::from(vec![ + Some(-1), + Some(-1), + Some(2), + None, + None, + Some(101), + Some(102), + Some(103), + None, + Some(256), + Some(512), + Some(1024), + ]); + + assert_eq!(expected_output, arr.as_ref()); + + Ok(()) +} + +#[test] +fn primitive_array_slices() -> Result<()> { + let input_1 = Int64Array::from(&[Some(-1), Some(-1), Some(2), None, None]).slice(1, 3); + + let input_2 = Int64Array::from(&[Some(101), Some(102), Some(103), None]).slice(1, 3); + let arr = concatenate(&[&input_1, &input_2])?; + + let expected_output = Int64Array::from(&[Some(-1), Some(2), None, Some(102), Some(103), None]); + + assert_eq!(expected_output, arr.as_ref()); + + Ok(()) +} + +#[test] +fn boolean_primitive_arrays() -> Result<()> { + let arr = concatenate(&[ + &BooleanArray::from(vec![ + Some(true), + Some(true), + Some(false), + None, + None, + Some(false), + ]), + &BooleanArray::from(vec![None, Some(false), Some(true), Some(false)]), + ])?; + + let expected_output = BooleanArray::from(vec![ + Some(true), + Some(true), + Some(false), + None, + None, + Some(false), + None, + Some(false), + Some(true), + Some(false), + ]); + + assert_eq!(expected_output, arr.as_ref()); + + Ok(()) +} diff --git a/tests/it/compute/filter.rs b/tests/it/compute/filter.rs new file mode 100644 index 00000000000..1fce99ab91c --- /dev/null +++ b/tests/it/compute/filter.rs @@ -0,0 +1,179 @@ +use arrow2::array::*; +use arrow2::compute::filter::*; + +#[test] +fn array_slice() { + let a = Int32Array::from_slice(&[5, 6, 7, 8, 9]).slice(1, 4); + let b = BooleanArray::from_slice(vec![true, true, false, false, true]).slice(1, 4); + let c = filter(&a, &b).unwrap(); + + let expected = Int32Array::from_slice(&[6, 9]); + + assert_eq!(expected, c.as_ref()); +} + +#[test] +fn array_low_density() { + // this test exercises the all 0's branch of the filter algorithm + let mut data_values = (1..=65).collect::>(); + let mut filter_values = (1..=65).map(|i| matches!(i % 65, 0)).collect::>(); + // set up two more values after the batch + data_values.extend_from_slice(&[66, 67]); + filter_values.extend_from_slice(&[false, true]); + let a = Int32Array::from_slice(&data_values); + let b = BooleanArray::from_slice(filter_values); + let c = filter(&a, &b).unwrap(); + + let expected = Int32Array::from_slice(&[65, 67]); + + assert_eq!(expected, c.as_ref()); +} + +#[test] +fn array_high_density() { + // this test exercises the all 1's branch of the filter algorithm + let mut data_values = (1..=65).map(Some).collect::>(); + let mut filter_values = (1..=65) + .map(|i| !matches!(i % 65, 0)) + .collect::>(); + // set second data value to null + data_values[1] = None; + // set up two more values after the batch + data_values.extend_from_slice(&[Some(66), None, Some(67), None]); + filter_values.extend_from_slice(&[false, true, true, true]); + let a = Int32Array::from(data_values); + let b = BooleanArray::from_slice(filter_values); + let c = filter(&a, &b).unwrap(); + let d = c.as_ref().as_any().downcast_ref::().unwrap(); + assert_eq!(67, d.len()); + assert_eq!(3, d.null_count()); + assert_eq!(1, d.value(0)); + assert!(d.is_null(1)); + assert_eq!(64, d.value(63)); + assert!(d.is_null(64)); + assert_eq!(67, d.value(65)); +} + +#[test] +fn string_array_simple() { + let a = Utf8Array::::from_slice(&["hello", " ", "world", "!"]); + let b = BooleanArray::from_slice(&[true, false, true, false]); + let c = filter(&a, &b).unwrap(); + let d = c + .as_ref() + .as_any() + .downcast_ref::>() + .unwrap(); + assert_eq!(2, d.len()); + assert_eq!("hello", d.value(0)); + assert_eq!("world", d.value(1)); +} + +#[test] +fn primative_array_with_null() { + let a = Int32Array::from(&[Some(5), None]); + let b = BooleanArray::from_slice(vec![false, true]); + let c = filter(&a, &b).unwrap(); + let d = c.as_ref().as_any().downcast_ref::().unwrap(); + assert_eq!(1, d.len()); + assert!(d.is_null(0)); +} + +#[test] +fn string_array_with_null() { + let a = Utf8Array::::from(&vec![Some("hello"), None, Some("world"), None]); + let b = BooleanArray::from_slice(vec![true, false, false, true]); + let c = filter(&a, &b).unwrap(); + let d = c + .as_ref() + .as_any() + .downcast_ref::>() + .unwrap(); + assert_eq!(2, d.len()); + assert_eq!("hello", d.value(0)); + assert!(!d.is_null(0)); + assert!(d.is_null(1)); +} + +#[test] +fn binary_array_with_null() { + let data: Vec> = vec![Some(b"hello"), None, Some(b"world"), None]; + let a = BinaryArray::::from(&data); + let b = BooleanArray::from_slice(vec![true, false, false, true]); + let c = filter(&a, &b).unwrap(); + let d = c + .as_ref() + .as_any() + .downcast_ref::>() + .unwrap(); + assert_eq!(2, d.len()); + assert_eq!(b"hello", d.value(0)); + assert!(!d.is_null(0)); + assert!(d.is_null(1)); +} + +/* +#[test] +fn dictionary_array() { + let values = vec![Some("hello"), None, Some("world"), Some("!")]; + let a: Int8DictionaryArray = values.iter().copied().collect(); + let b = BooleanArray::from(vec![false, true, true, false]); + let c = filter(&a, &b).unwrap(); + let d = c + .as_ref() + .as_any() + .downcast_ref::() + .unwrap(); + let value_array = d.values(); + let values = value_array.as_any().downcast_ref::().unwrap(); + // values are cloned in the filtered dictionary array + assert_eq!(3, values.len()); + // but keys are filtered + assert_eq!(2, d.len()); + assert_eq!(true, d.is_null(0)); + assert_eq!("world", values.value(d.keys().value(1) as usize)); +} + +#[test] +fn list_array() { + let value_data = ArrayData::builder(DataType::Int32) + .len(8) + .add_buffer(Buffer::from_slice_ref(&[0, 1, 2, 3, 4, 5, 6, 7])) + .build(); + + let value_offsets = Buffer::from_slice_ref(&[0i64, 3, 6, 8, 8]); + + let list_data_type = + DataType::LargeList(Box::new(Field::new("item", DataType::Int32, false))); + let list_data = ArrayData::builder(list_data_type) + .len(4) + .add_buffer(value_offsets) + .add_child_data(value_data) + .null_bit_buffer(Buffer::from([0b00000111])) + .build(); + + // a = [[0, 1, 2], [3, 4, 5], [6, 7], null] + let a = LargeListArray::from(list_data); + let b = BooleanArray::from(vec![false, true, false, true]); + let result = filter(&a, &b).unwrap(); + + // expected: [[3, 4, 5], null] + let value_data = ArrayData::builder(DataType::Int32) + .len(3) + .add_buffer(Buffer::from_slice_ref(&[3, 4, 5])) + .build(); + + let value_offsets = Buffer::from_slice_ref(&[0i64, 3, 3]); + + let list_data_type = + DataType::LargeList(Box::new(Field::new("item", DataType::Int32, false))); + let expected = ArrayData::builder(list_data_type) + .len(2) + .add_buffer(value_offsets) + .add_child_data(value_data) + .null_bit_buffer(Buffer::from([0b00000001])) + .build(); + + assert_eq!(&make_array(expected), &result); +} +*/ diff --git a/tests/it/compute/hash.rs b/tests/it/compute/hash.rs new file mode 100644 index 00000000000..92e263f235d --- /dev/null +++ b/tests/it/compute/hash.rs @@ -0,0 +1,49 @@ +use arrow2::array::new_null_array; +use arrow2::compute::hash::*; +use arrow2::datatypes::DataType::*; +use arrow2::datatypes::TimeUnit; + +#[test] +fn consistency() { + let datatypes = vec![ + Null, + Boolean, + UInt8, + UInt16, + UInt32, + UInt64, + Int8, + Int16, + Int32, + Int64, + Float32, + Float64, + Timestamp(TimeUnit::Second, None), + Timestamp(TimeUnit::Millisecond, None), + Timestamp(TimeUnit::Microsecond, None), + Timestamp(TimeUnit::Nanosecond, None), + Time64(TimeUnit::Microsecond), + Time64(TimeUnit::Nanosecond), + Date32, + Time32(TimeUnit::Second), + Time32(TimeUnit::Millisecond), + Date64, + Utf8, + LargeUtf8, + Binary, + LargeBinary, + Duration(TimeUnit::Second), + Duration(TimeUnit::Millisecond), + Duration(TimeUnit::Microsecond), + Duration(TimeUnit::Nanosecond), + ]; + + datatypes.into_iter().for_each(|d1| { + let array = new_null_array(d1.clone(), 10); + if can_hash(&d1) { + assert!(hash(array.as_ref()).is_ok()); + } else { + assert!(hash(array.as_ref()).is_err()); + } + }); +} diff --git a/tests/it/compute/if_then_else.rs b/tests/it/compute/if_then_else.rs new file mode 100644 index 00000000000..35e280eb4bf --- /dev/null +++ b/tests/it/compute/if_then_else.rs @@ -0,0 +1,42 @@ +use arrow2::array::*; +use arrow2::compute::if_then_else::if_then_else; +use arrow2::error::Result; + +#[test] +fn basics() -> Result<()> { + let lhs = Int32Array::from_slice(&[1, 2, 3]); + let rhs = Int32Array::from_slice(&[4, 5, 6]); + let predicate = BooleanArray::from_slice(vec![true, false, true]); + let c = if_then_else(&predicate, &lhs, &rhs)?; + + let expected = Int32Array::from_slice(&[1, 5, 3]); + + assert_eq!(expected, c.as_ref()); + Ok(()) +} + +#[test] +fn basics_nulls() -> Result<()> { + let lhs = Int32Array::from(&[Some(1), None, None]); + let rhs = Int32Array::from(&[None, Some(5), Some(6)]); + let predicate = BooleanArray::from_slice(vec![true, false, true]); + let c = if_then_else(&predicate, &lhs, &rhs)?; + + let expected = Int32Array::from(&[Some(1), Some(5), None]); + + assert_eq!(expected, c.as_ref()); + Ok(()) +} + +#[test] +fn basics_nulls_pred() -> Result<()> { + let lhs = Int32Array::from_slice(&[1, 2, 3]); + let rhs = Int32Array::from_slice(&[4, 5, 6]); + let predicate = BooleanArray::from(&[Some(true), None, Some(false)]); + let result = if_then_else(&predicate, &lhs, &rhs)?; + + let expected = Int32Array::from(&[Some(1), None, Some(6)]); + + assert_eq!(expected, result.as_ref()); + Ok(()) +} diff --git a/tests/it/compute/mod.rs b/tests/it/compute/mod.rs index f64549ca027..0924e72627f 100644 --- a/tests/it/compute/mod.rs +++ b/tests/it/compute/mod.rs @@ -1 +1,7 @@ +mod boolean; mod cast; +mod concat; +mod filter; +mod hash; +mod if_then_else; +mod sort; diff --git a/tests/it/compute/sort.rs b/tests/it/compute/sort.rs new file mode 100644 index 00000000000..1f2df489150 --- /dev/null +++ b/tests/it/compute/sort.rs @@ -0,0 +1,603 @@ +use arrow2::array::*; +use arrow2::compute::sort::*; +use arrow2::datatypes::*; +use arrow2::types::NativeType; + +fn to_indices_boolean_arrays(data: &[Option], options: SortOptions, expected_data: &[i32]) { + let output = BooleanArray::from(data); + let expected = Int32Array::from_slice(expected_data); + let output = sort_to_indices(&output, &options, None).unwrap(); + assert_eq!(output, expected) +} + +fn primitive_arrays( + data: &[Option], + data_type: DataType, + options: SortOptions, + expected_data: &[Option], +) where + T: NativeType, +{ + let input = PrimitiveArray::::from(data).to(data_type.clone()); + let expected = PrimitiveArray::::from(expected_data).to(data_type); + let output = sort(&input, &options, None).unwrap(); + assert_eq!(expected, output.as_ref()) +} + +fn to_indices_string_arrays(data: &[Option<&str>], options: SortOptions, expected_data: &[i32]) { + let input = Utf8Array::::from(&data.to_vec()); + let expected = Int32Array::from_slice(expected_data); + let output = sort_to_indices(&input, &options, None).unwrap(); + assert_eq!(output, expected) +} + +fn string_arrays(data: &[Option<&str>], options: SortOptions, expected_data: &[Option<&str>]) { + let input = Utf8Array::::from(&data.to_vec()); + let expected = Utf8Array::::from(&expected_data.to_vec()); + let output = sort(&input, &options, None).unwrap(); + assert_eq!(expected, output.as_ref()) +} + +fn string_dict_arrays( + data: &[Option<&str>], + options: SortOptions, + expected_data: &[Option<&str>], +) { + let mut input = MutableDictionaryArray::>::new(); + input.try_extend(data.iter().copied()).unwrap(); + let input = input.into_arc(); + + let mut expected = MutableDictionaryArray::>::new(); + expected.try_extend(expected_data.iter().copied()).unwrap(); + let expected = expected.into_arc(); + + let output = sort(input.as_ref(), &options, None).unwrap(); + assert_eq!(expected.as_ref(), output.as_ref()) +} + +/* +fn list_arrays( + data: Vec>>>, + options: Option, + expected_data: Vec>>>, + fixed_length: Option, +) where + T: ArrowPrimitiveType, + PrimitiveArray: From>>, +{ + // for FixedSizedList + if let Some(length) = fixed_length { + let input = Arc::new(build_fixed_size_list_nullable(data.clone(), length)); + let sorted = sort(&(input as ArrayRef), options).unwrap(); + let expected = Arc::new(build_fixed_size_list_nullable( + expected_data.clone(), + length, + )) as ArrayRef; + + assert_eq!(&sorted, &expected); + } + + // for List + let input = Arc::new(build_generic_list_nullable::(data.clone())); + let sorted = sort(&(input as ArrayRef), options).unwrap(); + let expected = + Arc::new(build_generic_list_nullable::(expected_data.clone())) + as ArrayRef; + + assert_eq!(&sorted, &expected); + + // for LargeList + let input = Arc::new(build_generic_list_nullable::(data)); + let sorted = sort(&(input as ArrayRef), options).unwrap(); + let expected = + Arc::new(build_generic_list_nullable::(expected_data)) as ArrayRef; + + assert_eq!(&sorted, &expected); +} + +fn test_lex_sort_arrays(input: Vec, expected_output: Vec) { + let sorted = lexsort(&input).unwrap(); + + for (result, expected) in sorted.iter().zip(expected_output.iter()) { + assert_eq!(result, expected); + } +} +*/ + +#[test] +fn boolean() { + // boolean + to_indices_boolean_arrays( + &[None, Some(false), Some(true), Some(true), Some(false), None], + SortOptions { + descending: false, + nulls_first: true, + }, + &[0, 5, 1, 4, 2, 3], + ); + + // boolean, descending + to_indices_boolean_arrays( + &[None, Some(false), Some(true), Some(true), Some(false), None], + SortOptions { + descending: true, + nulls_first: false, + }, + &[2, 3, 1, 4, 5, 0], + ); + + // boolean, descending, nulls first + to_indices_boolean_arrays( + &[None, Some(false), Some(true), Some(true), Some(false), None], + SortOptions { + descending: true, + nulls_first: true, + }, + &[5, 0, 2, 3, 1, 4], + ); +} + +#[test] +#[ignore] // improve equality for NaN values. These are right but the equality fails +fn test_nans() { + primitive_arrays::( + &[None, Some(0.0), Some(2.0), Some(-1.0), Some(f64::NAN), None], + DataType::Float64, + SortOptions { + descending: true, + nulls_first: true, + }, + &[None, None, Some(f64::NAN), Some(2.0), Some(0.0), Some(-1.0)], + ); + primitive_arrays::( + &[Some(f64::NAN), Some(f64::NAN), Some(f64::NAN), Some(1.0)], + DataType::Float64, + SortOptions { + descending: true, + nulls_first: true, + }, + &[Some(f64::NAN), Some(f64::NAN), Some(f64::NAN), Some(1.0)], + ); + + primitive_arrays::( + &[None, Some(0.0), Some(2.0), Some(-1.0), Some(f64::NAN), None], + DataType::Float64, + SortOptions { + descending: false, + nulls_first: true, + }, + &[None, None, Some(-1.0), Some(0.0), Some(2.0), Some(f64::NAN)], + ); + // nans + primitive_arrays::( + &[Some(f64::NAN), Some(f64::NAN), Some(f64::NAN), Some(1.0)], + DataType::Float64, + SortOptions { + descending: false, + nulls_first: true, + }, + &[Some(1.0), Some(f64::NAN), Some(f64::NAN), Some(f64::NAN)], + ); +} + +#[test] +fn to_indices_strings() { + to_indices_string_arrays( + &[ + None, + Some("bad"), + Some("sad"), + None, + Some("glad"), + Some("-ad"), + ], + SortOptions { + descending: false, + nulls_first: true, + }, + // &[3, 0, 5, 1, 4, 2] is also valid + &[0, 3, 5, 1, 4, 2], + ); + + to_indices_string_arrays( + &[ + None, + Some("bad"), + Some("sad"), + None, + Some("glad"), + Some("-ad"), + ], + SortOptions { + descending: true, + nulls_first: false, + }, + // &[2, 4, 1, 5, 3, 0] is also valid + &[2, 4, 1, 5, 0, 3], + ); + + to_indices_string_arrays( + &[ + None, + Some("bad"), + Some("sad"), + None, + Some("glad"), + Some("-ad"), + ], + SortOptions { + descending: false, + nulls_first: true, + }, + // &[3, 0, 5, 1, 4, 2] is also valid + &[0, 3, 5, 1, 4, 2], + ); + + to_indices_string_arrays( + &[ + None, + Some("bad"), + Some("sad"), + None, + Some("glad"), + Some("-ad"), + ], + SortOptions { + descending: true, + nulls_first: true, + }, + // &[3, 0, 2, 4, 1, 5] is also valid + &[0, 3, 2, 4, 1, 5], + ); +} + +#[test] +fn strings() { + string_arrays( + &[ + None, + Some("bad"), + Some("sad"), + None, + Some("glad"), + Some("-ad"), + ], + SortOptions { + descending: false, + nulls_first: true, + }, + &[ + None, + None, + Some("-ad"), + Some("bad"), + Some("glad"), + Some("sad"), + ], + ); + + string_arrays( + &[ + None, + Some("bad"), + Some("sad"), + None, + Some("glad"), + Some("-ad"), + ], + SortOptions { + descending: true, + nulls_first: false, + }, + &[ + Some("sad"), + Some("glad"), + Some("bad"), + Some("-ad"), + None, + None, + ], + ); + + string_arrays( + &[ + None, + Some("bad"), + Some("sad"), + None, + Some("glad"), + Some("-ad"), + ], + SortOptions { + descending: false, + nulls_first: true, + }, + &[ + None, + None, + Some("-ad"), + Some("bad"), + Some("glad"), + Some("sad"), + ], + ); + + string_arrays( + &[ + None, + Some("bad"), + Some("sad"), + None, + Some("glad"), + Some("-ad"), + ], + SortOptions { + descending: true, + nulls_first: true, + }, + &[ + None, + None, + Some("sad"), + Some("glad"), + Some("bad"), + Some("-ad"), + ], + ); +} + +#[test] +fn string_dicts() { + string_dict_arrays::( + &[ + None, + Some("bad"), + Some("sad"), + None, + Some("glad"), + Some("-ad"), + ], + SortOptions { + descending: false, + nulls_first: true, + }, + &[ + None, + None, + Some("-ad"), + Some("bad"), + Some("glad"), + Some("sad"), + ], + ); + + string_dict_arrays::( + &[ + None, + Some("bad"), + Some("sad"), + None, + Some("glad"), + Some("-ad"), + ], + SortOptions { + descending: true, + nulls_first: false, + }, + &[ + Some("sad"), + Some("glad"), + Some("bad"), + Some("-ad"), + None, + None, + ], + ); + + string_dict_arrays::( + &[ + None, + Some("bad"), + Some("sad"), + None, + Some("glad"), + Some("-ad"), + ], + SortOptions { + descending: false, + nulls_first: true, + }, + &[ + None, + None, + Some("-ad"), + Some("bad"), + Some("glad"), + Some("sad"), + ], + ); + + string_dict_arrays::( + &[ + None, + Some("bad"), + Some("sad"), + None, + Some("glad"), + Some("-ad"), + ], + SortOptions { + descending: true, + nulls_first: true, + }, + &[ + None, + None, + Some("sad"), + Some("glad"), + Some("bad"), + Some("-ad"), + ], + ); +} + +/* +#[test] +fn list() { + list_arrays::( + vec![ + Some(vec![Some(1)]), + Some(vec![Some(4)]), + Some(vec![Some(2)]), + Some(vec![Some(3)]), + ], + Some(SortOptions { + descending: false, + nulls_first: false, + }), + vec![ + Some(vec![Some(1)]), + Some(vec![Some(2)]), + Some(vec![Some(3)]), + Some(vec![Some(4)]), + ], + Some(1), + ); + + list_arrays::( + vec![ + Some(vec![Some(1), Some(0)]), + Some(vec![Some(4), Some(3), Some(2), Some(1)]), + Some(vec![Some(2), Some(3), Some(4)]), + Some(vec![Some(3), Some(3), Some(3), Some(3)]), + Some(vec![Some(1), Some(1)]), + ], + Some(SortOptions { + descending: false, + nulls_first: false, + }), + vec![ + Some(vec![Some(1), Some(0)]), + Some(vec![Some(1), Some(1)]), + Some(vec![Some(2), Some(3), Some(4)]), + Some(vec![Some(3), Some(3), Some(3), Some(3)]), + Some(vec![Some(4), Some(3), Some(2), Some(1)]), + ], + None, + ); + + list_arrays::( + vec![ + None, + Some(vec![Some(4), None, Some(2)]), + Some(vec![Some(2), Some(3), Some(4)]), + None, + Some(vec![Some(3), Some(3), None]), + ], + Some(SortOptions { + descending: false, + nulls_first: false, + }), + vec![ + Some(vec![Some(2), Some(3), Some(4)]), + Some(vec![Some(3), Some(3), None]), + Some(vec![Some(4), None, Some(2)]), + None, + None, + ], + Some(3), + ); +} + +#[test] +fn test_lex_sort_single_column() { + let input = vec![SortColumn { + values: Arc::new(PrimitiveArray::::from(vec![ + Some(17), + Some(2), + Some(-1), + Some(0), + ])) as ArrayRef, + options: None, + }]; + let expected = vec![Arc::new(PrimitiveArray::::from(vec![ + Some(-1), + Some(0), + Some(2), + Some(17), + ])) as ArrayRef]; + test_lex_sort_arrays(input, expected); +} + +#[test] +fn test_lex_sort_unaligned_rows() { + let input = vec![ + SortColumn { + values: Arc::new(PrimitiveArray::::from(vec![None, Some(-1)])) + as ArrayRef, + options: None, + }, + SortColumn { + values: Arc::new(StringArray::from(vec![Some("foo")])) as ArrayRef, + options: None, + }, + ]; + assert!( + lexsort(&input).is_err(), + "lexsort should reject columns with different row counts" + ); +} +*/ + +#[test] +fn consistency() { + use arrow2::array::new_null_array; + use arrow2::datatypes::DataType::*; + use arrow2::datatypes::TimeUnit; + + let datatypes = vec![ + Null, + Boolean, + UInt8, + UInt16, + UInt32, + UInt64, + Int8, + Int16, + Int32, + Int64, + Float32, + Float64, + Timestamp(TimeUnit::Second, None), + Timestamp(TimeUnit::Millisecond, None), + Timestamp(TimeUnit::Microsecond, None), + Timestamp(TimeUnit::Nanosecond, None), + Time64(TimeUnit::Microsecond), + Time64(TimeUnit::Nanosecond), + Date32, + Time32(TimeUnit::Second), + Time32(TimeUnit::Millisecond), + Date64, + Utf8, + LargeUtf8, + Binary, + LargeBinary, + Duration(TimeUnit::Second), + Duration(TimeUnit::Millisecond), + Duration(TimeUnit::Microsecond), + Duration(TimeUnit::Nanosecond), + ]; + + datatypes.into_iter().for_each(|d1| { + let array = new_null_array(d1.clone(), 10); + let options = SortOptions { + descending: true, + nulls_first: true, + }; + if can_sort(&d1) { + assert!(sort(array.as_ref(), &options, None).is_ok()); + } else { + assert!(sort(array.as_ref(), &options, None).is_err()); + } + }); +}