Fix equal_to in ByteGroupValueBuilder

datafusion/physical-plan/src/aggregates/group_values/group_column.rs

@@ -95,19 +95,12 @@ impl<T: ArrowPrimitiveType, const NULLABLE: bool> GroupColumn
         if NULLABLE {
             // In nullable path, we should check if both `exist row` and `input row`
             // are null/not null
-            let is_exist_null = self.nulls.is_null(lhs_row);
-            let null_match = is_exist_null == array.is_null(rhs_row);
-            if !null_match {
-                // If `is_null`s in `exist row` and `input row` don't match, return not equal to
-                return false;
-            } else if is_exist_null {
-                // If `is_null`s in `exist row` and `input row` match, and they are `null`s,
-                // return equal to
-                //
-                // NOTICE: we should not check their values when they are `null`s, because they are
-                // meaningless actually, and not ensured to be same
-                //
-                return true;
+            let exist_null = self.nulls.is_null(lhs_row);
+            let input_null = array.is_null(rhs_row);
+            match (exist_null, input_null) {
+                (true, true) => return true,
+                (false, true) | (true, false) => return false,
+                _ => {}
             }
             // Otherwise, we need to check their values
         }
@@ -224,9 +217,19 @@ where
     where
         B: ByteArrayType,
     {
-        let arr = array.as_bytes::<B>();
-        self.nulls.is_null(lhs_row) == arr.is_null(rhs_row)
-            && self.value(lhs_row) == (arr.value(rhs_row).as_ref() as &[u8])
+        let array = array.as_bytes::<B>();
+
+        // In nullable path, we should check if both `exist row` and `input row`
+        // are null/not null
+        let exist_null = self.nulls.is_null(lhs_row);
+        let input_null = array.is_null(rhs_row);
+        match (exist_null, input_null) {
+            (true, true) => return true,
+            (false, true) | (true, false) => return false,
+            _ => {}
+        }
+        // Otherwise, we need to check their values
+        self.value(lhs_row) == (array.value(rhs_row).as_ref() as &[u8])
     }
 
     /// return the current value of the specified row irrespective of null
@@ -468,13 +471,14 @@ mod tests {
         builder.append_val(&builder_array, 5);
 
         // Define input array
-        let (_, values, _) =
+        let (_nulls, values, _) =
             Int64Array::from(vec![Some(1), Some(2), None, None, Some(1), Some(3)])
                 .into_parts();
 
+        // explicitly build a boolean buffer where one of the null values also happens to match
         let mut boolean_buffer_builder = BooleanBufferBuilder::new(6);
         boolean_buffer_builder.append(true);
-        boolean_buffer_builder.append(false);
+        boolean_buffer_builder.append(false); // this sets Some(2) to null above
         boolean_buffer_builder.append(false);
         boolean_buffer_builder.append(false);
         boolean_buffer_builder.append(true);
@@ -511,4 +515,63 @@ mod tests {
         assert!(builder.equal_to(0, &input_array, 0));
         assert!(!builder.equal_to(1, &input_array, 1));
     }
+
+    #[test]
+    fn test_byte_array_equal_to() {
+        // Will cover such cases:
+        //   - exist null, input not null
+        //   - exist null, input null; values not equal
+        //   - exist null, input null; values equal
+        //   - exist not null, input null
+        //   - exist not null, input not null; values not equal
+        //   - exist not null, input not null; values equal
+
+        // Define PrimitiveGroupValueBuilder
+        let mut builder = ByteGroupValueBuilder::<i32>::new(OutputType::Utf8);
+        let builder_array = Arc::new(StringArray::from(vec![
+            None,
+            None,
+            None,
+            Some("foo"),
+            Some("bar"),
+            Some("baz"),
+        ])) as ArrayRef;
+        builder.append_val(&builder_array, 0);
+        builder.append_val(&builder_array, 1);
+        builder.append_val(&builder_array, 2);
+        builder.append_val(&builder_array, 3);
+        builder.append_val(&builder_array, 4);
+        builder.append_val(&builder_array, 5);
+
+        // Define input array
+        let (offsets, buffer, _nulls) = StringArray::from(vec![
+            Some("foo"),
+            Some("bar"),
+            None,
+            None,
+            Some("foo"),
+            Some("baz"),
+        ])
+        .into_parts();
+
+        // explicitly build a boolean buffer where one of the null values also happens to match
+        let mut boolean_buffer_builder = BooleanBufferBuilder::new(6);
+        boolean_buffer_builder.append(true);
+        boolean_buffer_builder.append(false); // this sets Some(2) to null above
+        boolean_buffer_builder.append(false);
+        boolean_buffer_builder.append(false);
+        boolean_buffer_builder.append(true);
+        boolean_buffer_builder.append(true);
+        let nulls = NullBuffer::new(boolean_buffer_builder.finish());
+        let input_array =
+            Arc::new(StringArray::new(offsets, buffer, Some(nulls))) as ArrayRef;
+
+        // Check
+        assert!(!builder.equal_to(0, &input_array, 0));
+        assert!(builder.equal_to(1, &input_array, 1));
+        assert!(builder.equal_to(2, &input_array, 2));
+        assert!(!builder.equal_to(3, &input_array, 3));
+        assert!(!builder.equal_to(4, &input_array, 4));
+        assert!(builder.equal_to(5, &input_array, 5));
+    }
 }