Simplify

Author: Dandandan

datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs

@@ -59,9 +59,7 @@ pub struct NullState {
     /// If `seen_values[i]` is false, have not seen any values that
     /// pass the filter yet for group `i`
     seen_values: BooleanBufferBuilder,
-    /// If true, all groups seen so far have seen at least one non-null value
-    /// and no filters have been applied.
-    all_seen: bool,
+    /// Tracks number of total groups
     seen_values_size: usize,
 }
 
@@ -75,7 +73,6 @@ impl NullState {
     pub fn new() -> Self {
         Self {
             seen_values: BooleanBufferBuilder::new(0),
-            all_seen: true,
             seen_values_size: 0,
         }
     }
@@ -113,13 +110,12 @@ impl NullState {
         T: ArrowPrimitiveType + Send,
         F: FnMut(usize, T::Native) + Send,
     {
-        if self.all_seen && opt_filter.is_none() && values.null_count() == 0 {
+        if self.seen_values.capacity() == 0 && opt_filter.is_none() && values.null_count() == 0 {
             accumulate(group_indices, values, None, value_fn);
             self.seen_values_size = total_num_groups;
         } else {
-            let prev_seen: bool = self.all_seen && self.seen_values_size > 0;
+            let prev_seen: bool = self.all_seen() && self.seen_values_size > 0;
 
-            self.all_seen = false;
             let seen_values =
                 initialize_builder(&mut self.seen_values, total_num_groups, false);
 
@@ -159,7 +155,7 @@ impl NullState {
         let data = values.values();
         assert_eq!(data.len(), group_indices.len());
 
-        if self.all_seen && opt_filter.is_none() && values.null_count() == 0 {
+        if self.seen_values.capacity() == 0 && opt_filter.is_none() && values.null_count() == 0 {
             group_indices
                 .iter()
                 .zip(data.iter())
@@ -168,9 +164,7 @@ impl NullState {
 
             return;
         }
-        let prev_seen: bool = self.all_seen && self.seen_values_size > 0;
-
-        self.all_seen = false;
+        let prev_seen: bool = self.seen_values.capacity() == 0 && self.seen_values_size > 0;
         // ensure the seen_values is big enough (start everything at
         // "not seen" valid)
         let seen_values =
@@ -253,15 +247,15 @@ impl NullState {
         let nulls = match emit_to {
             EmitTo::All => {
                 self.seen_values_size = 0;
-                if self.all_seen {
+                if self.all_seen() {
                     // all groups have seen at least one non null value
                     return None;
                 } else {
                     nulls
                 }
             }
             EmitTo::First(n) => {
-                if self.all_seen {
+                if self.all_seen() {
                     self.seen_values_size -= n;
                     return None;
                 }
@@ -275,6 +269,11 @@ impl NullState {
         };
         Some(NullBuffer::new(nulls))
     }
+
+    /// Returns true if all groups have seen at least one non null
+    fn all_seen(&self) -> bool {
+        self.seen_values.capacity() == 0
+    }
 }
 
 /// Invokes `value_fn(group_index, value)` for each non null, non
@@ -883,7 +882,7 @@ mod test {
                 "\n\naccumulated_values:{accumulated_values:#?}\n\nexpected_values:{expected_values:#?}"
             );
             let seen_values = null_state.seen_values.finish_cloned();
-            if null_state.all_seen {
+            if null_state.all_seen() {
                 assert_eq!(null_state.seen_values_size, total_num_groups);
             } else {
                 mock.validate_seen_values(&seen_values);
@@ -1016,7 +1015,7 @@ mod test {
 
             let null_buffer = null_state.build(EmitTo::All);
 
-            if !null_state.all_seen {
+            if !null_state.all_seen() {
                 assert_eq!(null_buffer, expected_null_buffer);
             }
         }