bjchambers commented on a change in pull request #521:
URL: https://github.com/apache/arrow-rs/pull/521#discussion_r670114367
##########
File path: arrow/src/compute/kernels/boolean.rs
##########
@@ -1148,12 +1222,215 @@ 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 = res.as_any().downcast_ref::<Int32Array>().unwrap();
let expected = Int32Array::from(vec![
Some(15), // False => keep it
Some(8), // None => keep it
None, // true => None
]);
- assert_eq!(&expected, &res)
+ assert_eq!(&expected, res)
+ }
+
+ #[test]
+ fn test_nullif_int_large_right_offset() {
+ let a: ArrayRef = Arc::new(Int32Array::from(vec![
+ None, // 0
+ Some(15), // 1
+ Some(8),
+ Some(1),
+ Some(9),
+ ]));
+ let a = a.slice(1, 3); // Some(15), Some(8), Some(1)
+
+ let comp = BooleanArray::from(vec![
+ Some(false), // 0
+ Some(false),
+ Some(false),
+ Some(false),
+ Some(false),
+ Some(false),
+ Some(false),
+ Some(false),
+ Some(false), // 8
+ Some(false),
+ Some(false),
+ Some(false),
+ Some(false),
+ Some(false),
+ Some(false),
+ Some(false),
+ Some(false), // 16
+ Some(false), // 17
+ Some(false), // 18
+ None,
+ Some(true),
+ Some(false),
+ None,
+ ]);
+ let comp = comp.slice(18, 3); // Some(false), None, Some(true)
+ let comp = comp.as_any().downcast_ref::<BooleanArray>().unwrap();
+ let res = nullif(&a, &comp).unwrap();
+ let res = res.as_any().downcast_ref::<Int32Array>().unwrap();
+
+ let expected = Int32Array::from(vec![
+ Some(15), // False => keep it
+ Some(8), // None => keep it
+ None, // true => None
+ ]);
+ assert_eq!(&expected, res)
+ }
+
+ #[test]
+ fn test_nullif_boolean_offset() {
+ let a: ArrayRef = Arc::new(BooleanArray::from(vec![
+ None, // 0
+ Some(true), // 1
+ Some(false),
+ Some(true),
+ Some(true),
+ ]));
+ let a = a.slice(1, 3); // Some(true), Some(false), Some(true)
+
+ let comp = BooleanArray::from(vec![
+ Some(false), // 0
+ Some(false), // 1
+ Some(false), // 2
+ None,
+ Some(true),
+ Some(false),
+ None,
+ ]);
+ 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 = res.as_any().downcast_ref::<BooleanArray>().unwrap();
+
+ let expected = BooleanArray::from(vec![
+ Some(true), // False => keep it
+ Some(false), // None => keep it
+ None, // true => None
+ ]);
+ assert_eq!(&expected, res)
+ }
+
+ #[test]
+ fn test_nullif_passthrough_all_nonnull() {
+ // DO NOT SUBMIT: Write this test
+ }
+
+ #[test]
+ fn test_nullif_passthrough_all_null_or_true() {
+ // DO NOT SUBMIT: Write this test
+ }
+
+ struct Foo {
+ a: Option<i32>,
+ b: Option<bool>,
+ /// Whether the entry should be valid.
+ is_valid: bool,
+ }
+
+ impl Foo {
+ fn new_valid(a: i32, b: bool) -> Foo {
+ Self {
+ a: Some(a),
+ b: Some(b),
+ is_valid: true,
+ }
+ }
+
+ fn new_null() -> Foo {
+ Self {
+ a: None,
+ b: None,
+ is_valid: false,
+ }
+ }
+ }
+
+ /// Struct Array equality is a bit weird -- we need to have the *child
values*
+ /// correct even if the enclosing struct indicates it is null. But we
+ /// also need the top level is_valid bits to be correct.
+ fn create_foo_struct(values: Vec<Foo>) -> StructArray {
+ let mut struct_array = StructBuilder::new(
+ vec![
+ Field::new("a", DataType::Int32, true),
+ Field::new("b", DataType::Boolean, true),
+ ],
+ vec![
+ Box::new(Int32Builder::new(values.len())),
+ Box::new(BooleanBuilder::new(values.len())),
+ ],
+ );
+
+ for value in values {
+ struct_array
+ .field_builder::<Int32Builder>(0)
+ .unwrap()
+ .append_option(value.a)
+ .unwrap();
+ struct_array
+ .field_builder::<BooleanBuilder>(1)
+ .unwrap()
+ .append_option(value.b)
+ .unwrap();
+ struct_array.append(value.is_valid).unwrap();
+ }
+
+ struct_array.finish()
+ }
+
+ #[test]
+ fn test_nullif_struct() {
+ let struct_array = create_foo_struct(vec![
+ Foo::new_valid(7, true),
+ Foo::new_valid(15, false),
+ Foo::new_valid(8, true),
+ Foo::new_valid(12, false),
+ Foo::new_null(),
+ Foo::new_null(),
+ Foo::new_valid(42, true),
+ ]);
+
+ let struct_array: ArrayRef = Arc::new(struct_array);
+ // Some({a: 15, b: false}), Some({a: 8, b: true}), Some({a: 12, b:
false}),
+ // None, None
+ let struct_array = struct_array.slice(1, 5);
+ let comp = BooleanArray::from(vec![
+ Some(false), // 0
+ Some(false), // 1
+ Some(false), // 2
+ None,
+ Some(true),
+ Some(false),
+ None,
+ ]);
+ let comp = comp.slice(2, 5); // Some(false), None, Some(true),
Some(false), None
+ let comp = comp.as_any().downcast_ref::<BooleanArray>().unwrap();
+ let res = nullif(&struct_array, &comp).unwrap();
+ let res = res.as_any().downcast_ref::<StructArray>().unwrap();
+
+ let expected = create_foo_struct(vec![
+ // Some(false) -> keep
+ Foo::new_valid(15, false),
+ // None -> keep
+ Foo::new_valid(8, true),
+ // Some(true) -> null out. But child values are still there.
+ Foo {
+ a: Some(12),
+ b: Some(false),
+ is_valid: false,
+ },
+ // Some(false) -> keep, but was null
+ Foo::new_null(),
+ // None -> keep, but was null
+ Foo::new_null(),
+ ]);
+
+ // This should pass. And in fact the arrays will *print* as equal.
+ // But, due to how the comparison is implemented, the slicing of the
outer
+ // struct array isn't properly propagated.
+ assert_eq!(&expected, res);
Review comment:
#514 was fixed by #389, so this assertion passes after having rebased.
--
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]
