Use direct indexing for small sized keys to improve group by aggregation performance

datafusion/core/src/physical_plan/aggregates/mod.rs

@@ -16,7 +16,6 @@
 // under the License.
 
 //! Aggregates functionalities
-
 use crate::execution::context::TaskContext;
 use crate::physical_plan::aggregates::hash::GroupedHashAggregateStream;
 use crate::physical_plan::aggregates::no_grouping::AggregateStream;

datafusion/core/src/physical_plan/aggregates/row_hash.rs

@@ -16,7 +16,6 @@
 // under the License.
 
 //! Hash aggregation through row format
-
 use std::sync::Arc;
 use std::task::{Context, Poll};
 use std::vec;
@@ -32,10 +31,11 @@ use crate::physical_plan::aggregates::{
     evaluate_group_by, evaluate_many, group_schema, AccumulatorItemV2, AggregateMode,
     PhysicalGroupBy,
 };
-use crate::physical_plan::hash_utils::create_row_hashes;
+use crate::physical_plan::hash_utils::{create_row_hash, create_row_hashes};
 use crate::physical_plan::metrics::{BaselineMetrics, RecordOutput};
 use crate::physical_plan::{aggregates, AggregateExpr, PhysicalExpr};
 use crate::physical_plan::{RecordBatchStream, SendableRecordBatchStream};
+use arrow::datatypes::DataType;
 
 use arrow::compute::cast;
 use arrow::datatypes::Schema;
@@ -84,7 +84,6 @@ pub(crate) struct GroupedHashAggregateStreamV2 {
     aggr_layout: Arc<RowLayout>,
 
     baseline_metrics: BaselineMetrics,
-    random_state: RandomState,
     finished: bool,
 }
 
@@ -129,12 +128,11 @@ impl GroupedHashAggregateStreamV2 {
             group_by,
             accumulators,
             group_schema,
-            aggr_schema,
+            aggr_schema: aggr_schema.clone(),
             aggr_layout,
             baseline_metrics,
             aggregate_expressions,
-            aggr_state: Default::default(),
-            random_state: Default::default(),
+            aggr_state: AggregationState::new(aggr_schema),
             finished: false,
         })
     }
@@ -160,7 +158,6 @@ impl Stream for GroupedHashAggregateStreamV2 {
                     let timer = elapsed_compute.timer();
                     let result = group_aggregate_batch(
                         &this.mode,
-                        &this.random_state,
                         &this.group_by,
                         &mut this.accumulators,
                         &this.group_schema,
@@ -212,7 +209,6 @@ impl RecordBatchStream for GroupedHashAggregateStreamV2 {
 #[allow(clippy::too_many_arguments)]
 fn group_aggregate_batch(
     mode: &AggregateMode,
-    random_state: &RandomState,
     grouping_set: &PhysicalGroupBy,
     accumulators: &mut [AccumulatorItemV2],
     group_schema: &Schema,
@@ -225,7 +221,7 @@ fn group_aggregate_batch(
     let grouping_by_values = evaluate_group_by(grouping_set, &batch)?;
 
     for group_values in grouping_by_values {
-        let group_rows: Vec<Vec<u8>> = create_group_rows(group_values, group_schema);
+        let mut group_rows: Vec<Vec<u8>> = create_group_rows(group_values, group_schema);
 
         // evaluate the aggregation expressions.
         // We could evaluate them after the `take`, but since we need to evaluate all
@@ -239,47 +235,35 @@ fn group_aggregate_batch(
         // track which entries in `aggr_state` have rows in this batch to aggregate
         let mut groups_with_rows = vec![];
 
-        // 1.1 Calculate the group keys for the group values
-        let mut batch_hashes = vec![0; batch.num_rows()];
-        create_row_hashes(&group_rows, random_state, &mut batch_hashes)?;
-
-        for (row, hash) in batch_hashes.into_iter().enumerate() {
+        for (row_idx, row) in group_rows.iter_mut().enumerate() {
             let AggregationState { map, group_states } = aggr_state;
 
-            let entry = map.get_mut(hash, |(_hash, group_idx)| {
-                // verify that a group that we are inserting with hash is
-                // actually the same key value as the group in
-                // existing_idx  (aka group_values @ row)
-                let group_state = &group_states[*group_idx];
-                group_rows[row] == group_state.group_by_values
-            });
-
-            match entry {
-                // Existing entry for this group value
-                Some((_hash, group_idx)) => {
-                    let group_state = &mut group_states[*group_idx];
+            // 1.1 Calculate the group keys for the group values
+            let hash = map.generate_hash(row);
+            match map.get_group_idx(row, hash, group_states) {
+                Some(group_idx) => {
+                    let group_state = &mut group_states[group_idx];
                     // 1.3
                     if group_state.indices.is_empty() {
-                        groups_with_rows.push(*group_idx);
+                        groups_with_rows.push(group_idx);
                     };
-                    group_state.indices.push(row as u32); // remember this row
+                    group_state.indices.push(row_idx as u32); // remember this row
                 }
-                //  1.2 Need to create new entry
                 None => {
-                    // Add new entry to group_states and save newly created index
+                    // 1.2 Add new entry to group_states and save newly created index
                     let group_state = RowGroupState {
-                        group_by_values: group_rows[row].clone(),
+                        group_by_values: row.clone(),
                         aggregation_buffer: vec![0; state_layout.fixed_part_width()],
-                        indices: vec![row as u32], // 1.3
+                        indices: vec![row_idx as u32], // 1.3
                     };
                     let group_idx = group_states.len();
                     group_states.push(group_state);
                     groups_with_rows.push(group_idx);
 
-                    // for hasher function, use precomputed hash value
-                    map.insert(hash, (hash, group_idx), |(hash, _group_idx)| *hash);
+                    // store the index using computed hash.
+                    map.store_group_idx(hash, group_idx);
                 }
-            };
+            }
         }
 
         // Collect all indices + offsets based on keys in this vec
@@ -379,26 +363,27 @@ struct RowGroupState {
     indices: Vec<u32>,
 }
 
-/// The state of all the groups
-#[derive(Default)]
 struct AggregationState {
     /// Logically maps group values to an index in `group_states`
-    ///
-    /// Uses the raw API of hashbrown to avoid actually storing the
-    /// keys in the table
-    ///
-    /// keys: u64 hashes of the GroupValue
-    /// values: (hash, index into `group_states`)
-    map: RawTable<(u64, usize)>,
+    map: GroupByMap,
 
     /// State for each group
     group_states: Vec<RowGroupState>,
 }
 
+impl AggregationState {
+    fn new(agg_schema: Arc<Schema>) -> AggregationState {
+        AggregationState {
+            map: choose_group_by_map(agg_schema),
+            group_states: Default::default(),
+        }
+    }
+}
+
 impl std::fmt::Debug for AggregationState {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
         // hashes are not store inline, so could only get values
-        let map_string = "RawTable";
+        let map_string = self.map.strategy_type();
         f.debug_struct("AggregationState")
             .field("map", &map_string)
             .field("group_states", &self.group_states)
@@ -486,3 +471,81 @@ fn read_as_batch(rows: &[Vec<u8>], schema: &Schema, row_type: RowType) -> Vec<Ar
 
     output.output_as_columns()
 }
+
+/// GroupByMap is used to store index of `group_state`
+///
+/// different strategy will be used to store index based on the key size.
+enum GroupByMap {
+    DirectIndexing(([Option<u8>; u8::MAX as usize], Arc<RowLayout>)),
+    Hash((RawTable<(u64, usize)>, RandomState)),
+}
+
+/// chooses group by map based on the key length.  
+fn choose_group_by_map(schema: Arc<Schema>) -> GroupByMap {
+    if schema.fields().len() == 1
+        && matches!(
+            schema.field(0).data_type(),
+            DataType::Boolean | DataType::UInt8
+        )
+    {
+        return GroupByMap::DirectIndexing((
+            [None; u8::MAX as usize],
+            Arc::new(RowLayout::new(&schema, RowType::Compact)),
+        ));
+    }
+    GroupByMap::Hash(Default::default())
+}
+
+impl GroupByMap {
+    /// generates hash based on the given row. it'll be used to store or retrive
+    /// group index.
+    fn generate_hash(&mut self, row: &mut Vec<u8>) -> u64 {
+        match self {
+            GroupByMap::Hash((_, random_state)) => create_row_hash(row, random_state),
+            GroupByMap::DirectIndexing((_, layout)) => {
+                let mut accessor = RowAccessor::new_from_layout(layout.clone());
+                accessor.point_to(0, row.as_mut_slice());
+                accessor.get_u8(0) as u64
+            }
+        }
+    }
+
+    /// retrive group index for the given hash and row.
+    fn get_group_idx(
+        &mut self,
+        row: &Vec<u8>,
+        hash: u64,
+        group_states: &mut Vec<RowGroupState>,
+    ) -> Option<usize> {
+        match self {
+            GroupByMap::DirectIndexing((map, _)) => {
+                map[hash as usize].map(|idx| idx as usize)
+            }
+            GroupByMap::Hash((map, _)) => {
+                let entry = map.get_mut(hash, |(_hash, group_idx)| {
+                    *row == group_states[*group_idx].group_by_values
+                });
+                entry.map(|(_hash, group_idx)| *group_idx)
+            }
+        }
+    }
+
+    /// store the group index for the given hash.
+    fn store_group_idx(&mut self, hash: u64, group_idx: usize) {
+        match self {
+            GroupByMap::DirectIndexing((map, _)) => {
+                map[hash as usize] = Some(group_idx as u8)
+            }
+            GroupByMap::Hash((map, _)) => {
+                map.insert(hash, (hash, group_idx), |(hash, _group_idx)| *hash);
+            }
+        }
+    }
+
+    fn strategy_type(&self) -> &str {
+        match self {
+            GroupByMap::DirectIndexing(..) => "DirectIndexing",
+            GroupByMap::Hash(..) => "RawTable",
+        }
+    }
+}

datafusion/core/src/physical_plan/hash_utils.rs

@@ -305,6 +305,15 @@ pub fn create_row_hashes<'a>(
     Ok(hashes_buffer)
 }
 
+#[cfg(not(feature = "force_hash_collisions"))]
+pub fn create_row_hash(row: &Vec<u8>, random_state: &RandomState) -> u64 {
+    random_state.hash_one(row)
+}
+
+#[cfg(feature = "force_hash_collisions")]
+pub fn create_row_hash(row: &Vec<u8>, random_state: &RandomState) -> u64 {
+    0
+}
 /// Creates hash values for every row, based on the values in the
 /// columns.
 ///

datafusion/core/tests/sql/group_by.rs

@@ -753,3 +753,38 @@ async fn csv_query_group_by_order_by_avg_group_by_substr() -> Result<()> {
     assert_batches_sorted_eq!(expected, &actual);
     Ok(())
 }
+
+#[tokio::test]
+async fn test_u8_columns() -> Result<()> {
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("a", DataType::UInt8, true),
+        Field::new("b", DataType::UInt8, false),
+    ]));
+    let col_a = UInt8Array::from(vec![1, 1, 2, 3, 4, 5]);
+    let col_b = UInt8Array::from(vec![1, 1, 1, 1, 1, 1]);
+
+    let batch =
+        RecordBatch::try_new(schema.clone(), vec![Arc::new(col_a), Arc::new(col_b)])
+            .unwrap();
+
+    let provider = MemTable::try_new(schema, vec![vec![batch]]).unwrap();
+    let ctx = SessionContext::new();
+    ctx.register_table("t", Arc::new(provider)).unwrap();
+    let results = plan_and_collect(&ctx, "SELECT a, count(b) FROM t GROUP BY a")
+        .await
+        .expect("ran plan correctly");
+
+    let expected = vec![
+        "+---+------------+",
+        "| a | COUNT(t.b) |",
+        "+---+------------+",
+        "| 1 | 2          |",
+        "| 2 | 1          |",
+        "| 3 | 1          |",
+        "| 4 | 1          |",
+        "| 5 | 1          |",
+        "+---+------------+",
+    ];
+    assert_batches_sorted_eq!(expected, &results);
+    Ok(())
+}

datafusion/row/src/accessor.rs

@@ -178,7 +178,7 @@ impl<'a> RowAccessor<'a> {
         value[0] != 0
     }
 
-    fn get_u8(&self, idx: usize) -> u8 {
+    pub fn get_u8(&self, idx: usize) -> u8 {
         self.assert_index_valid(idx);
         let offset = self.field_offsets()[idx];
         self.data[self.base_offset + offset]

datafusion/row/src/layout.rs

@@ -204,7 +204,6 @@ pub fn row_supported(schema: &Schema, row_type: RowType) -> bool {
 
 fn supported_type(dt: &DataType, row_type: RowType) -> bool {
     use DataType::*;
-
     match row_type {
         RowType::Compact => {
             matches!(