bjchambers commented on code in PR #1499:
URL: https://github.com/apache/arrow-rs/pull/1499#discussion_r843282687
##########
arrow/src/compute/kernels/boolean.rs:
##########
@@ -1173,13 +1412,284 @@ mod tests {
]);
let comp = comp.slice(2, 3); // Some(false), None, Some(true)
let comp = comp.as_any().downcast_ref::<BooleanArray>().unwrap();
- let res = nullif(a, comp).unwrap();
+ let res = nullif_alternative(&a, comp).unwrap();
- let expected = Int32Array::from(vec![
+ let expected: ArrayRef = Arc::new(Int32Array::from(vec![
Some(15), // False => keep it
Some(8), // None => keep it
None, // true => None
- ]);
+ ]));
assert_eq!(&expected, &res)
}
+
+ #[test]
+ fn test_nullif_dict() {
+ let array = DictionaryArray::<Int32Type>::from_iter(["a", "b", "c"]);
+ let array_ref = Arc::new(array) as ArrayRef;
+ let condition = BooleanArray::from(vec![true, false, true]);
+
+ let result = nullif_alternative(&array_ref, &condition).unwrap();
+ let result = result
+ .as_any()
+ .downcast_ref::<DictionaryArray<Int32Type>>()
+ .unwrap();
+ let expected_keys = &[None, Some(1),
None].iter().collect::<Int32Array>();
+ let actual_keys = result.keys();
+
+ assert_eq!(expected_keys, actual_keys);
+ }
+
+ #[test]
+ fn test_nullif_dict_sliced() {
+ let array =
+ DictionaryArray::<Int32Type>::from_iter(["a", "b", "c", "b", "c",
"b"]);
+ let array_ref = Arc::new(array) as ArrayRef;
+ let array_ref = array_ref.slice(3, 3);
+ let condition = BooleanArray::from(vec![true, true, false]);
+
+ let result = nullif_alternative(&array_ref, &condition).unwrap();
+ let result = result
+ .as_any()
+ .downcast_ref::<DictionaryArray<Int32Type>>()
+ .unwrap();
+ let expected_keys = &[None, None, Some(1)]
+ .iter()
+ .collect::<PrimitiveArray<Int32Type>>();
+ let actual_keys = result.keys();
+
+ assert_eq!(expected_keys, actual_keys);
+ }
+
+ #[test]
+ fn test_nullif_dict_cond_sliced() {
+ let array = DictionaryArray::<Int32Type>::from_iter(["a", "b", "c"]);
+ let array_ref = Arc::new(array) as ArrayRef;
+ let condition = BooleanArray::from(vec![true, false, true, true, true,
false]);
+ let condition = condition.slice(3, 3);
+ let condition =
condition.as_any().downcast_ref::<BooleanArray>().unwrap();
+
+ let result = nullif_alternative(&array_ref, condition).unwrap();
+ let result = result
+ .as_any()
+ .downcast_ref::<DictionaryArray<Int32Type>>()
+ .unwrap();
+ let expected_keys = &[None, None, Some(2)]
+ .iter()
+ .collect::<PrimitiveArray<Int32Type>>();
+ let actual_keys = result.keys();
+
+ assert_eq!(expected_keys, actual_keys);
+ }
+
+ #[test]
+ fn test_nullif_list() {
+ let array = ListArray::from_iter_primitive::<Float64Type, _, _>(vec![
+ Some(vec![Some(1.0), Some(2.0), Some(3.0)]),
+ None,
+ Some(vec![Some(4.0), None, Some(6.0)]),
+ Some(vec![]),
+ Some(vec![None, Some(8.0), None]),
+ ]);
+ let array_ref = Arc::new(array) as ArrayRef;
+ let condition = BooleanArray::from(vec![true, false, false, false,
true]);
+
+ let result = nullif_alternative(&array_ref, &condition).unwrap();
+ let result = result.as_any().downcast_ref::<ListArray>().unwrap();
+
+ let expected = ListArray::from_iter_primitive::<Float64Type, _,
_>(vec![
+ None,
+ None,
+ Some(vec![Some(4.0), None, Some(6.0)]),
+ Some(vec![]),
+ None,
+ ]);
+ let expected = &expected;
+
+ // simple assert_eq on ListArrays does not work if their offsets are
different
+ // see https://github.com/apache/arrow-rs/issues/626
+ assert_eq!(expected.len(), result.len());
+ for i in 0..expected.len() {
+ assert_eq!(expected.is_valid(i), result.is_valid(i));
+ if expected.is_valid(i) {
+ assert_eq!(&expected.value(i), &result.value(i));
+ }
+ }
+ }
+
+ #[test]
+ fn test_nullif_list_sliced() {
+ let array = ListArray::from_iter_primitive::<Float64Type, _, _>(vec![
+ Some(vec![Some(1.0), Some(2.0), Some(3.0)]),
+ None,
+ Some(vec![Some(4.0), None, Some(6.0)]),
+ Some(vec![]),
+ Some(vec![None, Some(8.0), None]),
+ ]);
+ let array_ref = array.slice(1, 4);
+ let condition = BooleanArray::from(vec![false, false, false, true]);
+
+ let result = nullif_alternative(&array_ref, &condition).unwrap();
+ let result = result.as_any().downcast_ref::<ListArray>().unwrap();
+
+ let expected = ListArray::from_iter_primitive::<Float64Type, _,
_>(vec![
+ None,
+ Some(vec![Some(4.0), None, Some(6.0)]),
+ Some(vec![]),
+ None,
+ ]);
+ let expected = &expected;
+
+ // simple assert_eq on ListArrays does not work if their offsets are
different
+ // see https://github.com/apache/arrow-rs/issues/626
+ assert_eq!(expected.len(), result.len());
+ for i in 0..expected.len() {
+ assert_eq!(expected.is_valid(i), result.is_valid(i));
+ if expected.is_valid(i) {
+ assert_eq!(&expected.value(i), &result.value(i));
+ }
+ }
+ }
+
+ #[test]
+ fn test_nullif_strings() {
+ let array = StringArray::from(vec!["abc", "", "def", "gh", "ijkl"]);
+ let array_ref = Arc::new(array) as ArrayRef;
+ let condition = BooleanArray::from(vec![true, false, false, false,
true]);
+
+ let result = nullif_alternative(&array_ref, &condition).unwrap();
+ let result = result.as_any().downcast_ref::<StringArray>().unwrap();
+
+ let expected =
+ StringArray::from(vec![None, Some(""), Some("def"), Some("gh"),
None]);
+
+ assert_eq!(result, &expected);
+ }
+
+ #[test]
+ fn test_nullif_strings_sliced() {
+ let array = StringArray::from(vec!["abc", "", "def", "gh", "ijkl"]);
+ let array_ref = array.slice(1, 4);
+ let condition = BooleanArray::from(vec![false, false, false, true]);
+
+ let result = nullif_alternative(&array_ref, &condition).unwrap();
+ let result = result.as_any().downcast_ref::<StringArray>().unwrap();
+
+ let expected = StringArray::from(vec![Some(""), Some("def"),
Some("gh"), None]);
+
+ assert_eq!(result, &expected);
+ }
+
+ #[test]
+ fn test_nullif_fixed_size_binary_sliced() {
+ let array = FixedSizeBinaryArray::try_from_iter(
+ vec![b"abc", b"def", b"ghi", b"jkl"].into_iter(),
+ )
+ .unwrap();
+ let array_ref = array.slice(1, 2);
+ let condition = BooleanArray::from(vec![false, false]);
+
+ let result = nullif_alternative(&array_ref, &condition).unwrap();
+ let result = result
+ .as_any()
+ .downcast_ref::<FixedSizeBinaryArray>()
+ .unwrap();
+
+ let expected = FixedSizeBinaryArray::try_from_sparse_iter(
+ vec![Some(b"def"), Some(b"ghi")].into_iter(),
+ )
+ .unwrap();
+
+ assert_eq!(result, &expected);
+ }
+
+ #[test]
+ fn test_nullif_fixed_size_list_sliced() {
+ let primitives =
+ Int16Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,
14, 15]);
+ let fixed_size_array = make_array(
+ ArrayData::try_new(
+ DataType::FixedSizeList(
+ Box::new(Field::new("item",
primitives.data_type().clone(), true)),
+ 3,
+ ),
+ 5,
+ None,
+ None,
+ 0,
+ vec![],
+ vec![primitives.data().clone()],
+ )
+ .unwrap(),
+ );
+ let array_ref = fixed_size_array.slice(1, 3);
+ let condition = BooleanArray::from(vec![false, true, false]);
+
+ let result = nullif_alternative(&array_ref, &condition).unwrap();
+ let result = result
+ .as_any()
+ .downcast_ref::<FixedSizeListArray>()
+ .unwrap();
+
+ let expected = make_array(
+ ArrayData::try_new(
+ DataType::FixedSizeList(
+ Box::new(Field::new("item",
primitives.data_type().clone(), true)),
+ 3,
+ ),
+ 3,
+ None,
+ Some(MutableBuffer::from_iter(vec![true, false, true]).into()),
+ 0,
+ vec![],
+ vec![primitives.slice(3, 9).data().clone()],
+ )
+ .unwrap(),
+ );
+ let expected = expected
+ .as_any()
+ .downcast_ref::<FixedSizeListArray>()
+ .unwrap();
+
+ assert_eq!(result, expected);
+ }
+
+ #[test]
+ fn test_nullif_struct_sliced() {
Review Comment:
For many of these -- it looks like we test for the case of the data array
being sliced, but not the boolean array being sliced. Should we add some
additional tests for that?
Also note -- I've seen problems with specific values of slicing. Not sure if
it's worth considering something like proptests for this?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
