diff --git a/datafusion/core/src/physical_plan/aggregates/mod.rs b/datafusion/core/src/physical_plan/aggregates/mod.rs index 88eda63edeef0..b6a0320c12f6a 100644 --- a/datafusion/core/src/physical_plan/aggregates/mod.rs +++ b/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; diff --git a/datafusion/core/src/physical_plan/aggregates/row_hash.rs b/datafusion/core/src/physical_plan/aggregates/row_hash.rs index d1c61cd287f9d..aa0503a1bce9b 100644 --- a/datafusion/core/src/physical_plan/aggregates/row_hash.rs +++ b/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, 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> = create_group_rows(group_values, group_schema); + let mut group_rows: Vec> = 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, } -/// 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, } +impl AggregationState { + fn new(agg_schema: Arc) -> 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], schema: &Schema, row_type: RowType) -> Vec; u8::MAX as usize], Arc)), + Hash((RawTable<(u64, usize)>, RandomState)), +} + +/// chooses group by map based on the key length. +fn choose_group_by_map(schema: Arc) -> 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) -> 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, + hash: u64, + group_states: &mut Vec, + ) -> Option { + 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", + } + } +} diff --git a/datafusion/core/src/physical_plan/hash_utils.rs b/datafusion/core/src/physical_plan/hash_utils.rs index 61e72b07c551f..c6a650cfdaf35 100644 --- a/datafusion/core/src/physical_plan/hash_utils.rs +++ b/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, random_state: &RandomState) -> u64 { + random_state.hash_one(row) +} + +#[cfg(feature = "force_hash_collisions")] +pub fn create_row_hash(row: &Vec, random_state: &RandomState) -> u64 { + 0 +} /// Creates hash values for every row, based on the values in the /// columns. /// diff --git a/datafusion/core/tests/sql/group_by.rs b/datafusion/core/tests/sql/group_by.rs index 2e1007be81c9e..d1bde7e9b63aa 100644 --- a/datafusion/core/tests/sql/group_by.rs +++ b/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(()) +} diff --git a/datafusion/row/src/accessor.rs b/datafusion/row/src/accessor.rs index f8e34578dbdac..f114207c6546e 100644 --- a/datafusion/row/src/accessor.rs +++ b/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] diff --git a/datafusion/row/src/layout.rs b/datafusion/row/src/layout.rs index 1518df9bf55a2..a14ed738bc3fa 100644 --- a/datafusion/row/src/layout.rs +++ b/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!(