add unit tests for primitive group values.

Author: Rachelint

datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs

@@ -408,3 +408,86 @@ where
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::collections::{BTreeMap, HashMap};
+    use std::sync::Arc;
+
+    use crate::aggregates::group_values::single_group_by::primitive::GroupValuesPrimitive;
+    use crate::aggregates::group_values::GroupValues;
+    use arrow::array::{ArrayRef, AsArray, Int64Array, NullBufferBuilder, UInt32Array};
+    use arrow::datatypes::{DataType, Int64Type, UInt32Type};
+    use datafusion_expr::EmitTo;
+
+    #[test]
+    fn test_flat_primitive_group_values() {
+        // Will cover such insert cases:
+        //   1.1 Non-null row + distinct
+        //   1.2 Null row + distinct
+        //   1.3 Non-null row + non-distinct
+        //   1.4 Null row + non-distinct
+        //
+        // Will cover such emit cases:
+        //   2.1 Emit first n
+        //   2.2 Emit all
+        //   2.3 Insert again + emit
+        let mut group_values = GroupValuesPrimitive::<UInt32Type>::new(DataType::UInt32);
+        let mut group_indices = vec![];
+
+        let data1 = Arc::new(UInt32Array::from(vec![
+            Some(1),
+            None,
+            Some(1),
+            None,
+            Some(2),
+            Some(3),
+        ]));
+        let data2 = Arc::new(UInt32Array::from(vec![Some(3), None, Some(4), Some(5)]));
+
+        // Insert case 1.1, 1.3, 1.4 + Emit case 2.1
+        group_values
+            .intern(&vec![data1.clone() as _], &mut group_indices)
+            .unwrap();
+
+        let mut expected = BTreeMap::new();
+        for (&group_index, value) in group_indices.iter().zip(data1.iter()) {
+            expected.insert(group_index, value);
+        }
+
+        let emit_result = group_values.emit(EmitTo::First(3)).unwrap();
+        for (group_idx, actual_val) in emit_result[0]
+            .as_primitive::<UInt32Type>()
+            .iter()
+            .enumerate()
+        {
+            assert!(group_idx < expected.len() - 1);
+            let expected_val = *expected.get(&group_idx).unwrap();
+            assert_eq!(actual_val, expected_val);
+        }
+
+        // Insert case 1.1~1.3 + Emit case 2.2~2.3
+        group_values
+            .intern(&vec![data2.clone() as _], &mut group_indices)
+            .unwrap();
+
+        let (_, remain_val) = expected.pop_last().unwrap();
+        let mut expected = BTreeMap::new();
+        for (&group_index, value) in group_indices.iter().zip(data2.iter()) {
+            if group_index == 0 {
+                assert_eq!(remain_val, value);
+            }
+            expected.insert(group_index, value);
+        }
+
+        let emit_result = group_values.emit(EmitTo::All).unwrap();
+        for (group_idx, actual_val) in emit_result[0]
+            .as_primitive::<UInt32Type>()
+            .iter()
+            .enumerate()
+        {
+            let expected_val = *expected.get(&group_idx).unwrap();
+            assert_eq!(actual_val, expected_val);
+        }
+    }
+}