ARROW-11290: [Rust][DataFusion] Address hash aggregate performance issue with high number of groups

Currently, we loop to the hashmap for every key.

However, as we receive a batch, if we a lot of groups in the group by expression (or receive sorted data, etc.) then we could create a lot of empty batches and call `update_batch` for each of the key already in the hashmap.

In the PR we keep track of which keys we received in the batch and only update the accumulators with the same keys instead of all accumulators.

On the db-benchmark h2oai/db-benchmark#182 this is the difference (mainly q3 and q5, others seem to be noise). It doesn't seem to completely solve the problem, but it reduces the problem already quite a bit.

This PR:
```
q1 took 340 ms
q2 took 1768 ms
q3 took 10975 ms
q4 took 337 ms
q5 took 13529 ms
```
Master:
```
q1 took 330 ms
q2 took 1648 ms
q3 took 16408 ms
q4 took 335 ms
q5 took 21074 ms
```

Closes #9234 from Dandandan/hash_agg_speed2

Authored-by: Heres, Daniel <[email protected]>
Signed-off-by: Andrew Lamb <[email protected]>

Author: Dandandan

rust/datafusion/src/physical_plan/hash_aggregate.rs

@@ -288,6 +288,9 @@ fn group_aggregate_batch(
     // Make sure we can create the accumulators or otherwise return an error
     create_accumulators(aggr_expr).map_err(DataFusionError::into_arrow_external_error)?;
 
+    // Keys received in this batch
+    let mut batch_keys = vec![];
+
     for row in 0..batch.num_rows() {
         // 1.1
         create_key(&group_values, row, &mut key)
@@ -297,11 +300,17 @@ fn group_aggregate_batch(
             .raw_entry_mut()
             .from_key(&key)
             // 1.3
-            .and_modify(|_, (_, _, v)| v.push(row as u32))
+            .and_modify(|_, (_, _, v)| {
+                if v.is_empty() {
+                    batch_keys.push(key.clone())
+                };
+                v.push(row as u32)
+            })
             // 1.2
             .or_insert_with(|| {
                 // We can safely unwrap here as we checked we can create an accumulator before
                 let accumulator_set = create_accumulators(aggr_expr).unwrap();
+                batch_keys.push(key.clone());
                 let _ = create_group_by_values(&group_values, row, &mut group_by_values);
                 (
                     key.clone(),
@@ -310,48 +319,48 @@ fn group_aggregate_batch(
             });
     }
 
-    // 2.1 for each key
+    // 2.1 for each key in this batch
     // 2.2 for each aggregation
     // 2.3 `take` from each of its arrays the keys' values
     // 2.4 update / merge the accumulator with the values
     // 2.5 clear indices
-    accumulators
-        .iter_mut()
-        .try_for_each(|(_, (_, accumulator_set, indices))| {
-            // 2.2
-            accumulator_set
-                .iter_mut()
-                .zip(&aggr_input_values)
-                .map(|(accumulator, aggr_array)| {
-                    (
-                        accumulator,
-                        aggr_array
-                            .iter()
-                            .map(|array| {
-                                // 2.3
-                                compute::take(
-                                    array.as_ref(),
-                                    &UInt32Array::from(indices.clone()),
-                                    None, // None: no index check
-                                )
-                                .unwrap()
-                            })
-                            .collect::<Vec<ArrayRef>>(),
-                    )
-                })
-                .try_for_each(|(accumulator, values)| match mode {
-                    AggregateMode::Partial => accumulator.update_batch(&values),
-                    AggregateMode::Final => {
-                        // note: the aggregation here is over states, not values, thus the merge
-                        accumulator.merge_batch(&values)
-                    }
-                })
-                // 2.5
-                .and({
-                    indices.clear();
-                    Ok(())
-                })
-        })?;
+    batch_keys.iter_mut().try_for_each(|key| {
+        let (_, accumulator_set, indices) = accumulators.get_mut(key).unwrap();
+        let primitive_indices = UInt32Array::from(indices.clone());
+        // 2.2
+        accumulator_set
+            .iter_mut()
+            .zip(&aggr_input_values)
+            .map(|(accumulator, aggr_array)| {
+                (
+                    accumulator,
+                    aggr_array
+                        .iter()
+                        .map(|array| {
+                            // 2.3
+                            compute::take(
+                                array.as_ref(),
+                                &primitive_indices,
+                                None, // None: no index check
+                            )
+                            .unwrap()
+                        })
+                        .collect::<Vec<ArrayRef>>(),
+                )
+            })
+            .try_for_each(|(accumulator, values)| match mode {
+                AggregateMode::Partial => accumulator.update_batch(&values),
+                AggregateMode::Final => {
+                    // note: the aggregation here is over states, not values, thus the merge
+                    accumulator.merge_batch(&values)
+                }
+            })
+            // 2.5
+            .and({
+                indices.clear();
+                Ok(())
+            })
+    })?;
     Ok(accumulators)
 }