count group

Signed-off-by: Jay Zhan <[email protected]>
jayzhan211 · Nov 17, 2024 · a889299 · a889299
1 parent 3721e74
commit a889299
Showing 1 changed file with 5 additions and 9 deletions.
diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs
@@ -481,15 +481,11 @@ impl GroupsAccumulator for CountGroupsAccumulator {
         // Adds the counts with the partial counts
         self.counts.resize(total_num_groups, 0);
         match opt_filter {
-            Some(filter) => filter
-                .iter()
-                .zip(group_indices.iter())
-                .zip(partial_counts.iter())
-                .for_each(|((filter_value, &group_index), partial_count)| {
-                    if let Some(true) = filter_value {
-                        self.counts[group_index] += partial_count;
-                    }
-                }),
+            Some(filter) => {
+                filter.values().set_indices().for_each(|index| {
+                    self.counts[group_indices[index]] += partial_counts[index]
+                });
+            }
             None => group_indices.iter().zip(partial_counts.iter()).for_each(
                 |(&group_index, partial_count)| {
                     self.counts[group_index] += partial_count;