kumarUjjawal commented on code in PR #27272:
URL: https://github.com/apache/datafusion/pull/27272#discussion_r4546038292
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,606 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF)
under one or more contributor license agreements
+
+use std::marker::PhantomData;
+use std::mem::size_of;
+use std::sync::{Arc, Weak};
+
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, Int64Array, PrimitiveArray,
+};
+use arrow::compute::take;
+use arrow::datatypes::{
+ ArrowDictionaryKeyType, ArrowNativeType, DataType, Field, Int8Type,
Int16Type,
+ Int32Type, Int64Type, Schema, UInt8Type, UInt16Type, UInt32Type,
UInt64Type,
+};
+use datafusion_common::{Result, exec_err};
+use datafusion_expr::{EmitTo, GroupSelection};
+use hashbrown::HashMap;
+
+use crate::aggregates::group_values::{
+ GroupValues, new_group_values as new_inner_group_values,
+};
+use crate::aggregates::order::GroupClustering;
+
+const UNRESOLVED: usize = usize::MAX;
+const PENDING: usize = usize::MAX - 1;
+// Sparse entries store both the dictionary index and group id plus hash-table
+// metadata, so prefer the faster dense cache until it would be over 4x larger.
+const SPARSE_CACHE_RATIO: usize = 4;
+// Arrow arrays are bounded by `isize::MAX`, leaving the high bit available to
+// tag unresolved row entries without allocating another row-sized buffer.
+const PENDING_VALUE_BIT: usize = 1 << (usize::BITS - 1);
+
+enum ValueToGroup {
+ Dense(Vec<usize>),
+ Sparse(HashMap<usize, usize>),
+}
+
+impl ValueToGroup {
+ fn new(num_values: usize, num_keys: usize) -> Self {
+ // A dense cache has a much cheaper hot-path lookup and less overhead
per
+ // entry. Use a sparse cache only when initializing the dense cache
would
+ // touch substantially more entries than this batch can reference.
+ if num_values <= num_keys.saturating_mul(SPARSE_CACHE_RATIO) {
+ Self::Dense(vec![UNRESOLVED; num_values])
+ } else {
+ Self::Sparse(HashMap::new())
+ }
+ }
+
+ fn maybe_promote(&mut self, num_values: usize) {
+ let Self::Sparse(values) = self else {
+ return;
+ };
+ if num_values > values.len().saturating_mul(SPARSE_CACHE_RATIO) {
+ return;
+ }
+
+ let mut dense = vec![UNRESOLVED; num_values];
+ for (&value_index, &group_index) in values.iter() {
+ dense[value_index] = group_index;
+ }
+ *self = Self::Dense(dense);
+ }
+ #[inline]
+ fn get(&self, value_index: usize) -> usize {
+ match self {
+ Self::Dense(values) => values[value_index],
+ Self::Sparse(values) => {
+ values.get(&value_index).copied().unwrap_or(UNRESOLVED)
+ }
+ }
+ }
+
+ #[inline]
+ fn insert(&mut self, value_index: usize, group_index: usize) {
+ match self {
+ Self::Dense(values) => values[value_index] = group_index,
+ Self::Sparse(values) => {
+ values.insert(value_index, group_index);
+ }
+ }
+ }
+
+ fn reset(&mut self, value_index: usize) {
+ match self {
+ Self::Dense(values) => values[value_index] = UNRESOLVED,
+ Self::Sparse(values) => {
+ values.remove(&value_index);
+ }
+ }
+ }
+
+ fn allocation_size(&self) -> usize {
+ match self {
+ Self::Dense(values) => values.capacity() * size_of::<usize>(),
+ Self::Sparse(values) => values.allocation_size(),
+ }
+ }
+}
+
+impl Default for ValueToGroup {
+ fn default() -> Self {
+ Self::Dense(vec![])
+ }
+}
+
+enum PendingValue {
+ Value(usize),
+ NullKey,
+}
+
+/// [`GroupValues`] for a single dictionary-encoded group key.
+struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send + Sync> {
+ /// Group values stored using the specialized single-column implementation
+ /// for the dictionary's value type.
+ values: Box<dyn GroupValues>,
+ /// Maps a value index in `cached_values` to its group id.
+ value_to_group: ValueToGroup,
+ /// Weak identity cache for the current dictionary values array.
+ cached_values: Option<Weak<dyn Array>>,
+ /// Group id for both null keys and keys that point to null values.
+ null_group: Option<usize>,
+ _phantom: PhantomData<K>,
+}
+
+impl<K: ArrowDictionaryKeyType + Send + Sync> GroupValuesDictionary<K> {
+ fn try_new(value_field: Field) -> Result<Self> {
+ let schema = Arc::new(Schema::new(vec![value_field]));
+ let values = new_inner_group_values(schema, &GroupClustering::None)?;
+ Ok(Self {
+ values,
+ value_to_group: ValueToGroup::default(),
+ cached_values: None,
+ null_group: None,
+ _phantom: PhantomData,
+ })
+ }
+
+ fn sync_cache(&mut self, values: &ArrayRef, num_keys: usize) {
+ if self
+ .cached_values
+ .as_ref()
+ .is_some_and(|cached| std::ptr::addr_eq(cached.as_ptr(),
Arc::as_ptr(values)))
+ {
+ return;
+ }
+
+ self.value_to_group = ValueToGroup::new(values.len(), num_keys);
+ self.cached_values = Some(Arc::downgrade(values));
+ }
+
+ fn clear_cache(&mut self) {
Review Comment:
`emit` calls `clear_cache`, which sets `null_group = None` even when
`EmitTo::First(n)` keeps the null group in `self.values`. `check_key_overflow`
then counts that null group as a non-null value and reports an overflow that
isn't real.**
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,606 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF)
under one or more contributor license agreements
+
+use std::marker::PhantomData;
+use std::mem::size_of;
+use std::sync::{Arc, Weak};
+
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, Int64Array, PrimitiveArray,
+};
+use arrow::compute::take;
+use arrow::datatypes::{
+ ArrowDictionaryKeyType, ArrowNativeType, DataType, Field, Int8Type,
Int16Type,
+ Int32Type, Int64Type, Schema, UInt8Type, UInt16Type, UInt32Type,
UInt64Type,
+};
+use datafusion_common::{Result, exec_err};
+use datafusion_expr::{EmitTo, GroupSelection};
+use hashbrown::HashMap;
+
+use crate::aggregates::group_values::{
+ GroupValues, new_group_values as new_inner_group_values,
+};
+use crate::aggregates::order::GroupClustering;
+
+const UNRESOLVED: usize = usize::MAX;
+const PENDING: usize = usize::MAX - 1;
+// Sparse entries store both the dictionary index and group id plus hash-table
+// metadata, so prefer the faster dense cache until it would be over 4x larger.
+const SPARSE_CACHE_RATIO: usize = 4;
+// Arrow arrays are bounded by `isize::MAX`, leaving the high bit available to
+// tag unresolved row entries without allocating another row-sized buffer.
+const PENDING_VALUE_BIT: usize = 1 << (usize::BITS - 1);
+
+enum ValueToGroup {
+ Dense(Vec<usize>),
+ Sparse(HashMap<usize, usize>),
+}
+
+impl ValueToGroup {
+ fn new(num_values: usize, num_keys: usize) -> Self {
+ // A dense cache has a much cheaper hot-path lookup and less overhead
per
+ // entry. Use a sparse cache only when initializing the dense cache
would
+ // touch substantially more entries than this batch can reference.
+ if num_values <= num_keys.saturating_mul(SPARSE_CACHE_RATIO) {
+ Self::Dense(vec![UNRESOLVED; num_values])
+ } else {
+ Self::Sparse(HashMap::new())
+ }
+ }
+
+ fn maybe_promote(&mut self, num_values: usize) {
+ let Self::Sparse(values) = self else {
+ return;
+ };
+ if num_values > values.len().saturating_mul(SPARSE_CACHE_RATIO) {
+ return;
+ }
+
+ let mut dense = vec![UNRESOLVED; num_values];
+ for (&value_index, &group_index) in values.iter() {
+ dense[value_index] = group_index;
+ }
+ *self = Self::Dense(dense);
+ }
+ #[inline]
+ fn get(&self, value_index: usize) -> usize {
+ match self {
+ Self::Dense(values) => values[value_index],
+ Self::Sparse(values) => {
+ values.get(&value_index).copied().unwrap_or(UNRESOLVED)
+ }
+ }
+ }
+
+ #[inline]
+ fn insert(&mut self, value_index: usize, group_index: usize) {
+ match self {
+ Self::Dense(values) => values[value_index] = group_index,
+ Self::Sparse(values) => {
+ values.insert(value_index, group_index);
+ }
+ }
+ }
+
+ fn reset(&mut self, value_index: usize) {
+ match self {
+ Self::Dense(values) => values[value_index] = UNRESOLVED,
+ Self::Sparse(values) => {
+ values.remove(&value_index);
+ }
+ }
+ }
+
+ fn allocation_size(&self) -> usize {
+ match self {
+ Self::Dense(values) => values.capacity() * size_of::<usize>(),
+ Self::Sparse(values) => values.allocation_size(),
+ }
+ }
+}
+
+impl Default for ValueToGroup {
+ fn default() -> Self {
+ Self::Dense(vec![])
+ }
+}
+
+enum PendingValue {
+ Value(usize),
+ NullKey,
+}
+
+/// [`GroupValues`] for a single dictionary-encoded group key.
+struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send + Sync> {
+ /// Group values stored using the specialized single-column implementation
+ /// for the dictionary's value type.
+ values: Box<dyn GroupValues>,
+ /// Maps a value index in `cached_values` to its group id.
+ value_to_group: ValueToGroup,
+ /// Weak identity cache for the current dictionary values array.
+ cached_values: Option<Weak<dyn Array>>,
+ /// Group id for both null keys and keys that point to null values.
+ null_group: Option<usize>,
+ _phantom: PhantomData<K>,
+}
+
+impl<K: ArrowDictionaryKeyType + Send + Sync> GroupValuesDictionary<K> {
+ fn try_new(value_field: Field) -> Result<Self> {
+ let schema = Arc::new(Schema::new(vec![value_field]));
+ let values = new_inner_group_values(schema, &GroupClustering::None)?;
Review Comment:
The inner group values are always built with `GroupClustering::None`,
ignoring the caller's clustering. Value types without a single-column
implementation (e.g. FixedSizeBinary, Struct, List) therefore get
`GroupValuesColumn::<false>` (vectorized intern) under streaming aggregation.
Before this PR, the outer `new_group_values` gave a single dictionary column
`GroupValuesColumn::<true>`.
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,606 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF)
under one or more contributor license agreements
+
+use std::marker::PhantomData;
+use std::mem::size_of;
+use std::sync::{Arc, Weak};
+
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, Int64Array, PrimitiveArray,
+};
+use arrow::compute::take;
+use arrow::datatypes::{
+ ArrowDictionaryKeyType, ArrowNativeType, DataType, Field, Int8Type,
Int16Type,
+ Int32Type, Int64Type, Schema, UInt8Type, UInt16Type, UInt32Type,
UInt64Type,
+};
+use datafusion_common::{Result, exec_err};
+use datafusion_expr::{EmitTo, GroupSelection};
+use hashbrown::HashMap;
+
+use crate::aggregates::group_values::{
+ GroupValues, new_group_values as new_inner_group_values,
+};
+use crate::aggregates::order::GroupClustering;
+
+const UNRESOLVED: usize = usize::MAX;
+const PENDING: usize = usize::MAX - 1;
+// Sparse entries store both the dictionary index and group id plus hash-table
+// metadata, so prefer the faster dense cache until it would be over 4x larger.
+const SPARSE_CACHE_RATIO: usize = 4;
+// Arrow arrays are bounded by `isize::MAX`, leaving the high bit available to
+// tag unresolved row entries without allocating another row-sized buffer.
+const PENDING_VALUE_BIT: usize = 1 << (usize::BITS - 1);
+
+enum ValueToGroup {
+ Dense(Vec<usize>),
+ Sparse(HashMap<usize, usize>),
+}
+
+impl ValueToGroup {
+ fn new(num_values: usize, num_keys: usize) -> Self {
+ // A dense cache has a much cheaper hot-path lookup and less overhead
per
+ // entry. Use a sparse cache only when initializing the dense cache
would
+ // touch substantially more entries than this batch can reference.
+ if num_values <= num_keys.saturating_mul(SPARSE_CACHE_RATIO) {
+ Self::Dense(vec![UNRESOLVED; num_values])
+ } else {
+ Self::Sparse(HashMap::new())
+ }
+ }
+
+ fn maybe_promote(&mut self, num_values: usize) {
+ let Self::Sparse(values) = self else {
+ return;
+ };
+ if num_values > values.len().saturating_mul(SPARSE_CACHE_RATIO) {
+ return;
+ }
+
+ let mut dense = vec![UNRESOLVED; num_values];
+ for (&value_index, &group_index) in values.iter() {
+ dense[value_index] = group_index;
+ }
+ *self = Self::Dense(dense);
+ }
+ #[inline]
+ fn get(&self, value_index: usize) -> usize {
+ match self {
+ Self::Dense(values) => values[value_index],
+ Self::Sparse(values) => {
+ values.get(&value_index).copied().unwrap_or(UNRESOLVED)
+ }
+ }
+ }
+
+ #[inline]
+ fn insert(&mut self, value_index: usize, group_index: usize) {
+ match self {
+ Self::Dense(values) => values[value_index] = group_index,
+ Self::Sparse(values) => {
+ values.insert(value_index, group_index);
+ }
+ }
+ }
+
+ fn reset(&mut self, value_index: usize) {
+ match self {
+ Self::Dense(values) => values[value_index] = UNRESOLVED,
+ Self::Sparse(values) => {
+ values.remove(&value_index);
+ }
+ }
+ }
+
+ fn allocation_size(&self) -> usize {
+ match self {
+ Self::Dense(values) => values.capacity() * size_of::<usize>(),
+ Self::Sparse(values) => values.allocation_size(),
+ }
+ }
+}
+
+impl Default for ValueToGroup {
+ fn default() -> Self {
+ Self::Dense(vec![])
+ }
+}
+
+enum PendingValue {
+ Value(usize),
+ NullKey,
+}
+
+/// [`GroupValues`] for a single dictionary-encoded group key.
+struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send + Sync> {
+ /// Group values stored using the specialized single-column implementation
+ /// for the dictionary's value type.
+ values: Box<dyn GroupValues>,
+ /// Maps a value index in `cached_values` to its group id.
+ value_to_group: ValueToGroup,
+ /// Weak identity cache for the current dictionary values array.
+ cached_values: Option<Weak<dyn Array>>,
+ /// Group id for both null keys and keys that point to null values.
+ null_group: Option<usize>,
+ _phantom: PhantomData<K>,
+}
+
+impl<K: ArrowDictionaryKeyType + Send + Sync> GroupValuesDictionary<K> {
+ fn try_new(value_field: Field) -> Result<Self> {
+ let schema = Arc::new(Schema::new(vec![value_field]));
+ let values = new_inner_group_values(schema, &GroupClustering::None)?;
+ Ok(Self {
+ values,
+ value_to_group: ValueToGroup::default(),
+ cached_values: None,
+ null_group: None,
+ _phantom: PhantomData,
+ })
+ }
+
+ fn sync_cache(&mut self, values: &ArrayRef, num_keys: usize) {
+ if self
+ .cached_values
+ .as_ref()
+ .is_some_and(|cached| std::ptr::addr_eq(cached.as_ptr(),
Arc::as_ptr(values)))
+ {
+ return;
+ }
+
+ self.value_to_group = ValueToGroup::new(values.len(), num_keys);
+ self.cached_values = Some(Arc::downgrade(values));
+ }
+
+ fn clear_cache(&mut self) {
+ self.value_to_group = ValueToGroup::default();
+ self.cached_values = None;
+ self.null_group = None;
+ }
+
+ fn reset_pending(&mut self, pending: &[PendingValue]) {
+ for value in pending {
+ if let PendingValue::Value(value_index) = value {
+ self.value_to_group.reset(*value_index);
+ }
+ }
+ }
+
+ #[inline]
+ fn group_or_queue(
+ &mut self,
+ value_index: usize,
+ pending: &mut Vec<PendingValue>,
+ ) -> usize {
+ match self.value_to_group.get(value_index) {
+ UNRESOLVED => {
+ self.value_to_group.insert(value_index, PENDING);
+ pending.push(PendingValue::Value(value_index));
+ PENDING_VALUE_BIT | value_index
+ }
+ PENDING => PENDING_VALUE_BIT | value_index,
+ group_index => group_index,
+ }
+ }
+
+ fn check_key_overflow(&self) -> Result<()> {
+ let non_null_count = self.values.len() - self.null_group.is_some() as
usize;
+ let max = match K::DATA_TYPE {
+ DataType::Int8 => i8::MAX as usize,
+ DataType::Int16 => i16::MAX as usize,
+ DataType::Int32 => i32::MAX as usize,
+ DataType::Int64 => i64::MAX as usize,
+ DataType::UInt8 => u8::MAX as usize,
+ DataType::UInt16 => u16::MAX as usize,
+ DataType::UInt32 => u32::MAX as usize,
+ DataType::UInt64 => usize::MAX,
+ _ => unreachable!("validated dictionary key type"),
+ };
+ if non_null_count != 0 && non_null_count - 1 > max {
+ return exec_err!(
+ "Dictionary key type {:?} cannot represent {} distinct values",
+ K::DATA_TYPE,
+ non_null_count
+ );
+ }
+ Ok(())
+ }
+
+ fn build_dictionary(values: ArrayRef, groups: &[usize]) -> ArrayRef {
+ let mut value_keys = Vec::with_capacity(values.len());
+ let mut non_null_indices = Vec::with_capacity(values.len());
+ for value_index in 0..values.len() {
+ if values.is_null(value_index) {
+ value_keys.push(None);
+ } else {
+
value_keys.push(Some(K::Native::usize_as(non_null_indices.len())));
+ non_null_indices.push(value_index as i64);
+ }
+ }
+
+ let values = if non_null_indices.len() == values.len() {
+ values
+ } else {
+ let indices = Int64Array::from(non_null_indices);
+ take(values.as_ref(), &indices, None).expect("compact dictionary
values")
+ };
+ let keys: PrimitiveArray<K> =
+ groups.iter().map(|&group| value_keys[group]).collect();
+ Arc::new(DictionaryArray::<K>::new(keys, values))
+ }
+}
+
+impl<K: ArrowDictionaryKeyType + Send + Sync> GroupValues for
GroupValuesDictionary<K> {
+ fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) ->
Result<()> {
+ assert_eq!(cols.len(), 1);
+ let dictionary = cols[0].as_dictionary::<K>();
+ let values = dictionary.values();
+ let keys = dictionary.keys();
+ self.sync_cache(values, keys.len() - keys.null_count());
+
+ groups.clear();
+ groups.resize(dictionary.len(), PENDING);
+ let mut pending = vec![];
+ let mut null_key_pending = false;
+ let raw_keys = keys.values();
+ if keys.null_count() == 0 {
+ for (group, key) in groups.iter_mut().zip(raw_keys) {
+ *group = self.group_or_queue(key.as_usize(), &mut pending);
+ }
+ } else {
+ let nulls = keys.nulls().unwrap();
+ for (row, group) in groups.iter_mut().enumerate() {
+ if nulls.is_null(row) {
+ if let Some(null_group) = self.null_group {
+ *group = null_group;
+ } else if !null_key_pending {
+ pending.push(PendingValue::NullKey);
+ null_key_pending = true;
+ }
+ } else {
+ *group = self.group_or_queue(raw_keys[row].as_usize(),
&mut pending);
+ }
+ }
+ }
+
+ if !pending.is_empty() {
+ let indices: Int64Array = pending
+ .iter()
+ .map(|pending| match pending {
+ PendingValue::Value(value_index) => Some(*value_index as
i64),
+ PendingValue::NullKey => None,
+ })
+ .collect();
+ let pending_values = match take(values.as_ref(), &indices, None) {
+ Ok(values) => values,
+ Err(error) => {
+ self.reset_pending(&pending);
+ return Err(error.into());
+ }
+ };
+ let mut pending_groups = Vec::with_capacity(pending.len());
+ if let Err(error) = self.values.intern(&[pending_values], &mut
pending_groups)
+ {
+ self.reset_pending(&pending);
+ return Err(error);
+ }
+ debug_assert_eq!(pending_groups.len(), pending.len());
+
+ for (pending, group_index) in
pending.into_iter().zip(pending_groups) {
+ match pending {
+ PendingValue::Value(value_index) => {
+ self.value_to_group.insert(value_index, group_index);
+ if values.is_null(value_index) {
+ self.null_group = Some(group_index);
+ }
+ }
+ PendingValue::NullKey => self.null_group =
Some(group_index),
+ }
+ }
+ self.value_to_group.maybe_promote(values.len());
+
+ for group in groups.iter_mut() {
+ if *group == PENDING {
+ *group = self.null_group.expect("null key was resolved");
+ } else if *group & PENDING_VALUE_BIT != 0 {
+ let value_index = *group & !PENDING_VALUE_BIT;
+ *group = self.value_to_group.get(value_index);
+ }
+ }
+ }
+
+ debug_assert!(
+ groups
+ .iter()
+ .all(|&group| group != PENDING && group & PENDING_VALUE_BIT ==
0)
+ );
+ self.check_key_overflow()
+ }
+
+ fn size(&self) -> usize {
+ self.values.size() + self.value_to_group.allocation_size() +
size_of::<Self>()
+ }
+
+ fn is_empty(&self) -> bool {
+ self.values.is_empty()
+ }
+
+ fn len(&self) -> usize {
+ self.values.len()
+ }
+
+ fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
+ let mut output = self.values.emit(emit_to)?;
+ debug_assert_eq!(output.len(), 1);
+ let values = output.pop().expect("one dictionary value column");
+ self.clear_cache();
+ let groups = (0..values.len()).collect::<Vec<_>>();
Review Comment:
`emit` builds an identity `groups` vec `(0..values.len())` only so that
`build_dictionary` can map `value_keys[group]` back to `value_keys`.
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,606 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF)
under one or more contributor license agreements
+
+use std::marker::PhantomData;
+use std::mem::size_of;
+use std::sync::{Arc, Weak};
+
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, Int64Array, PrimitiveArray,
+};
+use arrow::compute::take;
+use arrow::datatypes::{
+ ArrowDictionaryKeyType, ArrowNativeType, DataType, Field, Int8Type,
Int16Type,
+ Int32Type, Int64Type, Schema, UInt8Type, UInt16Type, UInt32Type,
UInt64Type,
+};
+use datafusion_common::{Result, exec_err};
+use datafusion_expr::{EmitTo, GroupSelection};
+use hashbrown::HashMap;
+
+use crate::aggregates::group_values::{
+ GroupValues, new_group_values as new_inner_group_values,
+};
+use crate::aggregates::order::GroupClustering;
+
+const UNRESOLVED: usize = usize::MAX;
+const PENDING: usize = usize::MAX - 1;
+// Sparse entries store both the dictionary index and group id plus hash-table
+// metadata, so prefer the faster dense cache until it would be over 4x larger.
+const SPARSE_CACHE_RATIO: usize = 4;
+// Arrow arrays are bounded by `isize::MAX`, leaving the high bit available to
+// tag unresolved row entries without allocating another row-sized buffer.
+const PENDING_VALUE_BIT: usize = 1 << (usize::BITS - 1);
+
+enum ValueToGroup {
+ Dense(Vec<usize>),
+ Sparse(HashMap<usize, usize>),
+}
+
+impl ValueToGroup {
+ fn new(num_values: usize, num_keys: usize) -> Self {
+ // A dense cache has a much cheaper hot-path lookup and less overhead
per
+ // entry. Use a sparse cache only when initializing the dense cache
would
+ // touch substantially more entries than this batch can reference.
+ if num_values <= num_keys.saturating_mul(SPARSE_CACHE_RATIO) {
+ Self::Dense(vec![UNRESOLVED; num_values])
+ } else {
+ Self::Sparse(HashMap::new())
+ }
+ }
+
+ fn maybe_promote(&mut self, num_values: usize) {
+ let Self::Sparse(values) = self else {
+ return;
+ };
+ if num_values > values.len().saturating_mul(SPARSE_CACHE_RATIO) {
+ return;
+ }
+
+ let mut dense = vec![UNRESOLVED; num_values];
+ for (&value_index, &group_index) in values.iter() {
+ dense[value_index] = group_index;
+ }
+ *self = Self::Dense(dense);
+ }
+ #[inline]
+ fn get(&self, value_index: usize) -> usize {
+ match self {
+ Self::Dense(values) => values[value_index],
+ Self::Sparse(values) => {
+ values.get(&value_index).copied().unwrap_or(UNRESOLVED)
+ }
+ }
+ }
+
+ #[inline]
+ fn insert(&mut self, value_index: usize, group_index: usize) {
+ match self {
+ Self::Dense(values) => values[value_index] = group_index,
+ Self::Sparse(values) => {
+ values.insert(value_index, group_index);
+ }
+ }
+ }
+
+ fn reset(&mut self, value_index: usize) {
+ match self {
+ Self::Dense(values) => values[value_index] = UNRESOLVED,
+ Self::Sparse(values) => {
+ values.remove(&value_index);
+ }
+ }
+ }
+
+ fn allocation_size(&self) -> usize {
+ match self {
+ Self::Dense(values) => values.capacity() * size_of::<usize>(),
+ Self::Sparse(values) => values.allocation_size(),
+ }
+ }
+}
+
+impl Default for ValueToGroup {
+ fn default() -> Self {
+ Self::Dense(vec![])
+ }
+}
+
+enum PendingValue {
+ Value(usize),
+ NullKey,
+}
+
+/// [`GroupValues`] for a single dictionary-encoded group key.
+struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send + Sync> {
+ /// Group values stored using the specialized single-column implementation
+ /// for the dictionary's value type.
+ values: Box<dyn GroupValues>,
+ /// Maps a value index in `cached_values` to its group id.
+ value_to_group: ValueToGroup,
+ /// Weak identity cache for the current dictionary values array.
+ cached_values: Option<Weak<dyn Array>>,
+ /// Group id for both null keys and keys that point to null values.
+ null_group: Option<usize>,
+ _phantom: PhantomData<K>,
+}
+
+impl<K: ArrowDictionaryKeyType + Send + Sync> GroupValuesDictionary<K> {
+ fn try_new(value_field: Field) -> Result<Self> {
+ let schema = Arc::new(Schema::new(vec![value_field]));
+ let values = new_inner_group_values(schema, &GroupClustering::None)?;
+ Ok(Self {
+ values,
+ value_to_group: ValueToGroup::default(),
+ cached_values: None,
+ null_group: None,
+ _phantom: PhantomData,
+ })
+ }
+
+ fn sync_cache(&mut self, values: &ArrayRef, num_keys: usize) {
+ if self
+ .cached_values
+ .as_ref()
+ .is_some_and(|cached| std::ptr::addr_eq(cached.as_ptr(),
Arc::as_ptr(values)))
+ {
+ return;
+ }
+
+ self.value_to_group = ValueToGroup::new(values.len(), num_keys);
+ self.cached_values = Some(Arc::downgrade(values));
+ }
+
+ fn clear_cache(&mut self) {
+ self.value_to_group = ValueToGroup::default();
+ self.cached_values = None;
+ self.null_group = None;
+ }
+
+ fn reset_pending(&mut self, pending: &[PendingValue]) {
+ for value in pending {
+ if let PendingValue::Value(value_index) = value {
+ self.value_to_group.reset(*value_index);
+ }
+ }
+ }
+
+ #[inline]
+ fn group_or_queue(
+ &mut self,
+ value_index: usize,
+ pending: &mut Vec<PendingValue>,
+ ) -> usize {
+ match self.value_to_group.get(value_index) {
+ UNRESOLVED => {
+ self.value_to_group.insert(value_index, PENDING);
+ pending.push(PendingValue::Value(value_index));
+ PENDING_VALUE_BIT | value_index
+ }
+ PENDING => PENDING_VALUE_BIT | value_index,
+ group_index => group_index,
+ }
+ }
+
+ fn check_key_overflow(&self) -> Result<()> {
Review Comment:
`check_key_overflow`, `build_dictionary` and the Weak-pointer identity check
in `sync_cache` repeat existing code in `multi_group_by/dictionary.rs`:
`DictionaryGroupValuesColumn::check_key_overflow`/`key_type_fits`, `into_dict`
and `is_cached`.
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,606 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF)
under one or more contributor license agreements
+
+use std::marker::PhantomData;
+use std::mem::size_of;
+use std::sync::{Arc, Weak};
+
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, Int64Array, PrimitiveArray,
+};
+use arrow::compute::take;
+use arrow::datatypes::{
+ ArrowDictionaryKeyType, ArrowNativeType, DataType, Field, Int8Type,
Int16Type,
+ Int32Type, Int64Type, Schema, UInt8Type, UInt16Type, UInt32Type,
UInt64Type,
+};
+use datafusion_common::{Result, exec_err};
+use datafusion_expr::{EmitTo, GroupSelection};
+use hashbrown::HashMap;
+
+use crate::aggregates::group_values::{
+ GroupValues, new_group_values as new_inner_group_values,
+};
+use crate::aggregates::order::GroupClustering;
+
+const UNRESOLVED: usize = usize::MAX;
+const PENDING: usize = usize::MAX - 1;
+// Sparse entries store both the dictionary index and group id plus hash-table
+// metadata, so prefer the faster dense cache until it would be over 4x larger.
+const SPARSE_CACHE_RATIO: usize = 4;
+// Arrow arrays are bounded by `isize::MAX`, leaving the high bit available to
+// tag unresolved row entries without allocating another row-sized buffer.
+const PENDING_VALUE_BIT: usize = 1 << (usize::BITS - 1);
+
+enum ValueToGroup {
+ Dense(Vec<usize>),
+ Sparse(HashMap<usize, usize>),
+}
+
+impl ValueToGroup {
+ fn new(num_values: usize, num_keys: usize) -> Self {
+ // A dense cache has a much cheaper hot-path lookup and less overhead
per
+ // entry. Use a sparse cache only when initializing the dense cache
would
+ // touch substantially more entries than this batch can reference.
+ if num_values <= num_keys.saturating_mul(SPARSE_CACHE_RATIO) {
+ Self::Dense(vec![UNRESOLVED; num_values])
+ } else {
+ Self::Sparse(HashMap::new())
+ }
+ }
+
+ fn maybe_promote(&mut self, num_values: usize) {
+ let Self::Sparse(values) = self else {
+ return;
+ };
+ if num_values > values.len().saturating_mul(SPARSE_CACHE_RATIO) {
+ return;
+ }
+
+ let mut dense = vec![UNRESOLVED; num_values];
+ for (&value_index, &group_index) in values.iter() {
+ dense[value_index] = group_index;
+ }
+ *self = Self::Dense(dense);
+ }
+ #[inline]
+ fn get(&self, value_index: usize) -> usize {
+ match self {
+ Self::Dense(values) => values[value_index],
+ Self::Sparse(values) => {
+ values.get(&value_index).copied().unwrap_or(UNRESOLVED)
+ }
+ }
+ }
+
+ #[inline]
+ fn insert(&mut self, value_index: usize, group_index: usize) {
+ match self {
+ Self::Dense(values) => values[value_index] = group_index,
+ Self::Sparse(values) => {
+ values.insert(value_index, group_index);
+ }
+ }
+ }
+
+ fn reset(&mut self, value_index: usize) {
+ match self {
+ Self::Dense(values) => values[value_index] = UNRESOLVED,
+ Self::Sparse(values) => {
+ values.remove(&value_index);
+ }
+ }
+ }
+
+ fn allocation_size(&self) -> usize {
+ match self {
+ Self::Dense(values) => values.capacity() * size_of::<usize>(),
+ Self::Sparse(values) => values.allocation_size(),
+ }
+ }
+}
+
+impl Default for ValueToGroup {
+ fn default() -> Self {
+ Self::Dense(vec![])
+ }
+}
+
+enum PendingValue {
+ Value(usize),
+ NullKey,
+}
+
+/// [`GroupValues`] for a single dictionary-encoded group key.
+struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send + Sync> {
+ /// Group values stored using the specialized single-column implementation
+ /// for the dictionary's value type.
+ values: Box<dyn GroupValues>,
+ /// Maps a value index in `cached_values` to its group id.
+ value_to_group: ValueToGroup,
+ /// Weak identity cache for the current dictionary values array.
+ cached_values: Option<Weak<dyn Array>>,
+ /// Group id for both null keys and keys that point to null values.
+ null_group: Option<usize>,
+ _phantom: PhantomData<K>,
+}
+
+impl<K: ArrowDictionaryKeyType + Send + Sync> GroupValuesDictionary<K> {
+ fn try_new(value_field: Field) -> Result<Self> {
+ let schema = Arc::new(Schema::new(vec![value_field]));
+ let values = new_inner_group_values(schema, &GroupClustering::None)?;
+ Ok(Self {
+ values,
+ value_to_group: ValueToGroup::default(),
+ cached_values: None,
+ null_group: None,
+ _phantom: PhantomData,
+ })
+ }
+
+ fn sync_cache(&mut self, values: &ArrayRef, num_keys: usize) {
+ if self
+ .cached_values
+ .as_ref()
+ .is_some_and(|cached| std::ptr::addr_eq(cached.as_ptr(),
Arc::as_ptr(values)))
+ {
+ return;
+ }
+
+ self.value_to_group = ValueToGroup::new(values.len(), num_keys);
+ self.cached_values = Some(Arc::downgrade(values));
+ }
+
+ fn clear_cache(&mut self) {
+ self.value_to_group = ValueToGroup::default();
+ self.cached_values = None;
+ self.null_group = None;
+ }
+
+ fn reset_pending(&mut self, pending: &[PendingValue]) {
+ for value in pending {
+ if let PendingValue::Value(value_index) = value {
+ self.value_to_group.reset(*value_index);
+ }
+ }
+ }
+
+ #[inline]
+ fn group_or_queue(
+ &mut self,
+ value_index: usize,
+ pending: &mut Vec<PendingValue>,
+ ) -> usize {
+ match self.value_to_group.get(value_index) {
+ UNRESOLVED => {
+ self.value_to_group.insert(value_index, PENDING);
+ pending.push(PendingValue::Value(value_index));
+ PENDING_VALUE_BIT | value_index
+ }
+ PENDING => PENDING_VALUE_BIT | value_index,
+ group_index => group_index,
+ }
+ }
+
+ fn check_key_overflow(&self) -> Result<()> {
+ let non_null_count = self.values.len() - self.null_group.is_some() as
usize;
+ let max = match K::DATA_TYPE {
+ DataType::Int8 => i8::MAX as usize,
+ DataType::Int16 => i16::MAX as usize,
+ DataType::Int32 => i32::MAX as usize,
+ DataType::Int64 => i64::MAX as usize,
+ DataType::UInt8 => u8::MAX as usize,
+ DataType::UInt16 => u16::MAX as usize,
+ DataType::UInt32 => u32::MAX as usize,
+ DataType::UInt64 => usize::MAX,
+ _ => unreachable!("validated dictionary key type"),
+ };
+ if non_null_count != 0 && non_null_count - 1 > max {
+ return exec_err!(
+ "Dictionary key type {:?} cannot represent {} distinct values",
+ K::DATA_TYPE,
+ non_null_count
+ );
+ }
+ Ok(())
+ }
+
+ fn build_dictionary(values: ArrayRef, groups: &[usize]) -> ArrayRef {
+ let mut value_keys = Vec::with_capacity(values.len());
+ let mut non_null_indices = Vec::with_capacity(values.len());
+ for value_index in 0..values.len() {
+ if values.is_null(value_index) {
+ value_keys.push(None);
+ } else {
+
value_keys.push(Some(K::Native::usize_as(non_null_indices.len())));
+ non_null_indices.push(value_index as i64);
+ }
+ }
+
+ let values = if non_null_indices.len() == values.len() {
+ values
+ } else {
+ let indices = Int64Array::from(non_null_indices);
+ take(values.as_ref(), &indices, None).expect("compact dictionary
values")
+ };
+ let keys: PrimitiveArray<K> =
+ groups.iter().map(|&group| value_keys[group]).collect();
+ Arc::new(DictionaryArray::<K>::new(keys, values))
+ }
+}
+
+impl<K: ArrowDictionaryKeyType + Send + Sync> GroupValues for
GroupValuesDictionary<K> {
+ fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) ->
Result<()> {
+ assert_eq!(cols.len(), 1);
+ let dictionary = cols[0].as_dictionary::<K>();
+ let values = dictionary.values();
+ let keys = dictionary.keys();
+ self.sync_cache(values, keys.len() - keys.null_count());
+
+ groups.clear();
+ groups.resize(dictionary.len(), PENDING);
+ let mut pending = vec![];
+ let mut null_key_pending = false;
+ let raw_keys = keys.values();
+ if keys.null_count() == 0 {
+ for (group, key) in groups.iter_mut().zip(raw_keys) {
+ *group = self.group_or_queue(key.as_usize(), &mut pending);
+ }
+ } else {
+ let nulls = keys.nulls().unwrap();
+ for (row, group) in groups.iter_mut().enumerate() {
+ if nulls.is_null(row) {
+ if let Some(null_group) = self.null_group {
+ *group = null_group;
+ } else if !null_key_pending {
+ pending.push(PendingValue::NullKey);
+ null_key_pending = true;
+ }
+ } else {
+ *group = self.group_or_queue(raw_keys[row].as_usize(),
&mut pending);
+ }
+ }
+ }
+
+ if !pending.is_empty() {
+ let indices: Int64Array = pending
+ .iter()
+ .map(|pending| match pending {
+ PendingValue::Value(value_index) => Some(*value_index as
i64),
+ PendingValue::NullKey => None,
+ })
+ .collect();
+ let pending_values = match take(values.as_ref(), &indices, None) {
+ Ok(values) => values,
+ Err(error) => {
+ self.reset_pending(&pending);
+ return Err(error.into());
+ }
+ };
+ let mut pending_groups = Vec::with_capacity(pending.len());
+ if let Err(error) = self.values.intern(&[pending_values], &mut
pending_groups)
+ {
+ self.reset_pending(&pending);
+ return Err(error);
+ }
+ debug_assert_eq!(pending_groups.len(), pending.len());
+
+ for (pending, group_index) in
pending.into_iter().zip(pending_groups) {
+ match pending {
+ PendingValue::Value(value_index) => {
+ self.value_to_group.insert(value_index, group_index);
+ if values.is_null(value_index) {
+ self.null_group = Some(group_index);
+ }
+ }
+ PendingValue::NullKey => self.null_group =
Some(group_index),
+ }
+ }
+ self.value_to_group.maybe_promote(values.len());
+
+ for group in groups.iter_mut() {
+ if *group == PENDING {
+ *group = self.null_group.expect("null key was resolved");
+ } else if *group & PENDING_VALUE_BIT != 0 {
+ let value_index = *group & !PENDING_VALUE_BIT;
+ *group = self.value_to_group.get(value_index);
+ }
+ }
+ }
+
+ debug_assert!(
+ groups
+ .iter()
+ .all(|&group| group != PENDING && group & PENDING_VALUE_BIT ==
0)
+ );
+ self.check_key_overflow()
+ }
+
+ fn size(&self) -> usize {
+ self.values.size() + self.value_to_group.allocation_size() +
size_of::<Self>()
+ }
+
+ fn is_empty(&self) -> bool {
+ self.values.is_empty()
+ }
+
+ fn len(&self) -> usize {
+ self.values.len()
+ }
+
+ fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
+ let mut output = self.values.emit(emit_to)?;
+ debug_assert_eq!(output.len(), 1);
+ let values = output.pop().expect("one dictionary value column");
+ self.clear_cache();
Review Comment:
Every `emit`, including `EmitTo::First(n)`, throws away the whole
value-index→group cache. It could instead shift the cached ids down by n and
drop those below n.
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,606 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF)
under one or more contributor license agreements
+
+use std::marker::PhantomData;
+use std::mem::size_of;
+use std::sync::{Arc, Weak};
+
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, Int64Array, PrimitiveArray,
+};
+use arrow::compute::take;
+use arrow::datatypes::{
+ ArrowDictionaryKeyType, ArrowNativeType, DataType, Field, Int8Type,
Int16Type,
+ Int32Type, Int64Type, Schema, UInt8Type, UInt16Type, UInt32Type,
UInt64Type,
+};
+use datafusion_common::{Result, exec_err};
+use datafusion_expr::{EmitTo, GroupSelection};
+use hashbrown::HashMap;
+
+use crate::aggregates::group_values::{
+ GroupValues, new_group_values as new_inner_group_values,
+};
+use crate::aggregates::order::GroupClustering;
+
+const UNRESOLVED: usize = usize::MAX;
+const PENDING: usize = usize::MAX - 1;
+// Sparse entries store both the dictionary index and group id plus hash-table
+// metadata, so prefer the faster dense cache until it would be over 4x larger.
+const SPARSE_CACHE_RATIO: usize = 4;
+// Arrow arrays are bounded by `isize::MAX`, leaving the high bit available to
+// tag unresolved row entries without allocating another row-sized buffer.
+const PENDING_VALUE_BIT: usize = 1 << (usize::BITS - 1);
+
+enum ValueToGroup {
+ Dense(Vec<usize>),
+ Sparse(HashMap<usize, usize>),
+}
+
+impl ValueToGroup {
+ fn new(num_values: usize, num_keys: usize) -> Self {
+ // A dense cache has a much cheaper hot-path lookup and less overhead
per
+ // entry. Use a sparse cache only when initializing the dense cache
would
+ // touch substantially more entries than this batch can reference.
+ if num_values <= num_keys.saturating_mul(SPARSE_CACHE_RATIO) {
+ Self::Dense(vec![UNRESOLVED; num_values])
+ } else {
+ Self::Sparse(HashMap::new())
+ }
+ }
+
+ fn maybe_promote(&mut self, num_values: usize) {
+ let Self::Sparse(values) = self else {
+ return;
+ };
+ if num_values > values.len().saturating_mul(SPARSE_CACHE_RATIO) {
+ return;
+ }
+
+ let mut dense = vec![UNRESOLVED; num_values];
+ for (&value_index, &group_index) in values.iter() {
+ dense[value_index] = group_index;
+ }
+ *self = Self::Dense(dense);
+ }
+ #[inline]
+ fn get(&self, value_index: usize) -> usize {
+ match self {
+ Self::Dense(values) => values[value_index],
+ Self::Sparse(values) => {
+ values.get(&value_index).copied().unwrap_or(UNRESOLVED)
+ }
+ }
+ }
+
+ #[inline]
+ fn insert(&mut self, value_index: usize, group_index: usize) {
+ match self {
+ Self::Dense(values) => values[value_index] = group_index,
+ Self::Sparse(values) => {
+ values.insert(value_index, group_index);
+ }
+ }
+ }
+
+ fn reset(&mut self, value_index: usize) {
+ match self {
+ Self::Dense(values) => values[value_index] = UNRESOLVED,
+ Self::Sparse(values) => {
+ values.remove(&value_index);
+ }
+ }
+ }
+
+ fn allocation_size(&self) -> usize {
+ match self {
+ Self::Dense(values) => values.capacity() * size_of::<usize>(),
+ Self::Sparse(values) => values.allocation_size(),
+ }
+ }
+}
+
+impl Default for ValueToGroup {
+ fn default() -> Self {
+ Self::Dense(vec![])
+ }
+}
+
+enum PendingValue {
+ Value(usize),
+ NullKey,
+}
+
+/// [`GroupValues`] for a single dictionary-encoded group key.
+struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send + Sync> {
+ /// Group values stored using the specialized single-column implementation
+ /// for the dictionary's value type.
+ values: Box<dyn GroupValues>,
+ /// Maps a value index in `cached_values` to its group id.
+ value_to_group: ValueToGroup,
+ /// Weak identity cache for the current dictionary values array.
+ cached_values: Option<Weak<dyn Array>>,
+ /// Group id for both null keys and keys that point to null values.
+ null_group: Option<usize>,
+ _phantom: PhantomData<K>,
+}
+
+impl<K: ArrowDictionaryKeyType + Send + Sync> GroupValuesDictionary<K> {
+ fn try_new(value_field: Field) -> Result<Self> {
+ let schema = Arc::new(Schema::new(vec![value_field]));
+ let values = new_inner_group_values(schema, &GroupClustering::None)?;
+ Ok(Self {
+ values,
+ value_to_group: ValueToGroup::default(),
+ cached_values: None,
+ null_group: None,
+ _phantom: PhantomData,
+ })
+ }
+
+ fn sync_cache(&mut self, values: &ArrayRef, num_keys: usize) {
+ if self
+ .cached_values
+ .as_ref()
+ .is_some_and(|cached| std::ptr::addr_eq(cached.as_ptr(),
Arc::as_ptr(values)))
+ {
+ return;
+ }
+
+ self.value_to_group = ValueToGroup::new(values.len(), num_keys);
+ self.cached_values = Some(Arc::downgrade(values));
+ }
+
+ fn clear_cache(&mut self) {
+ self.value_to_group = ValueToGroup::default();
+ self.cached_values = None;
+ self.null_group = None;
+ }
+
+ fn reset_pending(&mut self, pending: &[PendingValue]) {
+ for value in pending {
+ if let PendingValue::Value(value_index) = value {
+ self.value_to_group.reset(*value_index);
+ }
+ }
+ }
+
+ #[inline]
+ fn group_or_queue(
+ &mut self,
+ value_index: usize,
+ pending: &mut Vec<PendingValue>,
+ ) -> usize {
+ match self.value_to_group.get(value_index) {
+ UNRESOLVED => {
+ self.value_to_group.insert(value_index, PENDING);
+ pending.push(PendingValue::Value(value_index));
+ PENDING_VALUE_BIT | value_index
+ }
+ PENDING => PENDING_VALUE_BIT | value_index,
+ group_index => group_index,
+ }
+ }
+
+ fn check_key_overflow(&self) -> Result<()> {
+ let non_null_count = self.values.len() - self.null_group.is_some() as
usize;
+ let max = match K::DATA_TYPE {
+ DataType::Int8 => i8::MAX as usize,
+ DataType::Int16 => i16::MAX as usize,
+ DataType::Int32 => i32::MAX as usize,
+ DataType::Int64 => i64::MAX as usize,
+ DataType::UInt8 => u8::MAX as usize,
+ DataType::UInt16 => u16::MAX as usize,
+ DataType::UInt32 => u32::MAX as usize,
+ DataType::UInt64 => usize::MAX,
+ _ => unreachable!("validated dictionary key type"),
+ };
+ if non_null_count != 0 && non_null_count - 1 > max {
+ return exec_err!(
+ "Dictionary key type {:?} cannot represent {} distinct values",
+ K::DATA_TYPE,
+ non_null_count
+ );
+ }
+ Ok(())
+ }
+
+ fn build_dictionary(values: ArrayRef, groups: &[usize]) -> ArrayRef {
+ let mut value_keys = Vec::with_capacity(values.len());
+ let mut non_null_indices = Vec::with_capacity(values.len());
+ for value_index in 0..values.len() {
Review Comment:
Restore a null-free dictionary construction fast path
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,606 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF)
under one or more contributor license agreements
+
+use std::marker::PhantomData;
+use std::mem::size_of;
+use std::sync::{Arc, Weak};
+
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, Int64Array, PrimitiveArray,
+};
+use arrow::compute::take;
+use arrow::datatypes::{
+ ArrowDictionaryKeyType, ArrowNativeType, DataType, Field, Int8Type,
Int16Type,
+ Int32Type, Int64Type, Schema, UInt8Type, UInt16Type, UInt32Type,
UInt64Type,
+};
+use datafusion_common::{Result, exec_err};
+use datafusion_expr::{EmitTo, GroupSelection};
+use hashbrown::HashMap;
+
+use crate::aggregates::group_values::{
+ GroupValues, new_group_values as new_inner_group_values,
+};
+use crate::aggregates::order::GroupClustering;
+
+const UNRESOLVED: usize = usize::MAX;
+const PENDING: usize = usize::MAX - 1;
+// Sparse entries store both the dictionary index and group id plus hash-table
+// metadata, so prefer the faster dense cache until it would be over 4x larger.
+const SPARSE_CACHE_RATIO: usize = 4;
+// Arrow arrays are bounded by `isize::MAX`, leaving the high bit available to
+// tag unresolved row entries without allocating another row-sized buffer.
+const PENDING_VALUE_BIT: usize = 1 << (usize::BITS - 1);
+
+enum ValueToGroup {
+ Dense(Vec<usize>),
+ Sparse(HashMap<usize, usize>),
+}
+
+impl ValueToGroup {
+ fn new(num_values: usize, num_keys: usize) -> Self {
+ // A dense cache has a much cheaper hot-path lookup and less overhead
per
+ // entry. Use a sparse cache only when initializing the dense cache
would
+ // touch substantially more entries than this batch can reference.
+ if num_values <= num_keys.saturating_mul(SPARSE_CACHE_RATIO) {
+ Self::Dense(vec![UNRESOLVED; num_values])
+ } else {
+ Self::Sparse(HashMap::new())
+ }
+ }
+
+ fn maybe_promote(&mut self, num_values: usize) {
+ let Self::Sparse(values) = self else {
+ return;
+ };
+ if num_values > values.len().saturating_mul(SPARSE_CACHE_RATIO) {
+ return;
+ }
+
+ let mut dense = vec![UNRESOLVED; num_values];
+ for (&value_index, &group_index) in values.iter() {
+ dense[value_index] = group_index;
+ }
+ *self = Self::Dense(dense);
+ }
+ #[inline]
+ fn get(&self, value_index: usize) -> usize {
+ match self {
+ Self::Dense(values) => values[value_index],
+ Self::Sparse(values) => {
+ values.get(&value_index).copied().unwrap_or(UNRESOLVED)
+ }
+ }
+ }
+
+ #[inline]
+ fn insert(&mut self, value_index: usize, group_index: usize) {
+ match self {
+ Self::Dense(values) => values[value_index] = group_index,
+ Self::Sparse(values) => {
+ values.insert(value_index, group_index);
+ }
+ }
+ }
+
+ fn reset(&mut self, value_index: usize) {
+ match self {
+ Self::Dense(values) => values[value_index] = UNRESOLVED,
+ Self::Sparse(values) => {
+ values.remove(&value_index);
+ }
+ }
+ }
+
+ fn allocation_size(&self) -> usize {
+ match self {
+ Self::Dense(values) => values.capacity() * size_of::<usize>(),
+ Self::Sparse(values) => values.allocation_size(),
+ }
+ }
+}
+
+impl Default for ValueToGroup {
+ fn default() -> Self {
+ Self::Dense(vec![])
+ }
+}
+
+enum PendingValue {
+ Value(usize),
+ NullKey,
+}
+
+/// [`GroupValues`] for a single dictionary-encoded group key.
+struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send + Sync> {
+ /// Group values stored using the specialized single-column implementation
+ /// for the dictionary's value type.
+ values: Box<dyn GroupValues>,
+ /// Maps a value index in `cached_values` to its group id.
+ value_to_group: ValueToGroup,
+ /// Weak identity cache for the current dictionary values array.
+ cached_values: Option<Weak<dyn Array>>,
+ /// Group id for both null keys and keys that point to null values.
+ null_group: Option<usize>,
+ _phantom: PhantomData<K>,
+}
+
+impl<K: ArrowDictionaryKeyType + Send + Sync> GroupValuesDictionary<K> {
+ fn try_new(value_field: Field) -> Result<Self> {
+ let schema = Arc::new(Schema::new(vec![value_field]));
+ let values = new_inner_group_values(schema, &GroupClustering::None)?;
+ Ok(Self {
+ values,
+ value_to_group: ValueToGroup::default(),
+ cached_values: None,
+ null_group: None,
+ _phantom: PhantomData,
+ })
+ }
+
+ fn sync_cache(&mut self, values: &ArrayRef, num_keys: usize) {
+ if self
+ .cached_values
+ .as_ref()
+ .is_some_and(|cached| std::ptr::addr_eq(cached.as_ptr(),
Arc::as_ptr(values)))
+ {
+ return;
+ }
+
+ self.value_to_group = ValueToGroup::new(values.len(), num_keys);
+ self.cached_values = Some(Arc::downgrade(values));
+ }
+
+ fn clear_cache(&mut self) {
+ self.value_to_group = ValueToGroup::default();
+ self.cached_values = None;
+ self.null_group = None;
+ }
+
+ fn reset_pending(&mut self, pending: &[PendingValue]) {
+ for value in pending {
+ if let PendingValue::Value(value_index) = value {
+ self.value_to_group.reset(*value_index);
+ }
+ }
+ }
+
+ #[inline]
+ fn group_or_queue(
+ &mut self,
+ value_index: usize,
+ pending: &mut Vec<PendingValue>,
+ ) -> usize {
+ match self.value_to_group.get(value_index) {
+ UNRESOLVED => {
+ self.value_to_group.insert(value_index, PENDING);
+ pending.push(PendingValue::Value(value_index));
+ PENDING_VALUE_BIT | value_index
+ }
+ PENDING => PENDING_VALUE_BIT | value_index,
+ group_index => group_index,
+ }
+ }
+
+ fn check_key_overflow(&self) -> Result<()> {
+ let non_null_count = self.values.len() - self.null_group.is_some() as
usize;
+ let max = match K::DATA_TYPE {
+ DataType::Int8 => i8::MAX as usize,
+ DataType::Int16 => i16::MAX as usize,
+ DataType::Int32 => i32::MAX as usize,
+ DataType::Int64 => i64::MAX as usize,
+ DataType::UInt8 => u8::MAX as usize,
+ DataType::UInt16 => u16::MAX as usize,
+ DataType::UInt32 => u32::MAX as usize,
+ DataType::UInt64 => usize::MAX,
+ _ => unreachable!("validated dictionary key type"),
+ };
+ if non_null_count != 0 && non_null_count - 1 > max {
+ return exec_err!(
+ "Dictionary key type {:?} cannot represent {} distinct values",
+ K::DATA_TYPE,
+ non_null_count
+ );
+ }
+ Ok(())
+ }
+
+ fn build_dictionary(values: ArrayRef, groups: &[usize]) -> ArrayRef {
+ let mut value_keys = Vec::with_capacity(values.len());
+ let mut non_null_indices = Vec::with_capacity(values.len());
+ for value_index in 0..values.len() {
+ if values.is_null(value_index) {
+ value_keys.push(None);
+ } else {
+
value_keys.push(Some(K::Native::usize_as(non_null_indices.len())));
+ non_null_indices.push(value_index as i64);
+ }
+ }
+
+ let values = if non_null_indices.len() == values.len() {
+ values
+ } else {
+ let indices = Int64Array::from(non_null_indices);
+ take(values.as_ref(), &indices, None).expect("compact dictionary
values")
+ };
+ let keys: PrimitiveArray<K> =
+ groups.iter().map(|&group| value_keys[group]).collect();
+ Arc::new(DictionaryArray::<K>::new(keys, values))
+ }
+}
+
+impl<K: ArrowDictionaryKeyType + Send + Sync> GroupValues for
GroupValuesDictionary<K> {
+ fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) ->
Result<()> {
+ assert_eq!(cols.len(), 1);
+ let dictionary = cols[0].as_dictionary::<K>();
+ let values = dictionary.values();
+ let keys = dictionary.keys();
+ self.sync_cache(values, keys.len() - keys.null_count());
+
+ groups.clear();
+ groups.resize(dictionary.len(), PENDING);
+ let mut pending = vec![];
+ let mut null_key_pending = false;
+ let raw_keys = keys.values();
+ if keys.null_count() == 0 {
+ for (group, key) in groups.iter_mut().zip(raw_keys) {
+ *group = self.group_or_queue(key.as_usize(), &mut pending);
+ }
+ } else {
+ let nulls = keys.nulls().unwrap();
+ for (row, group) in groups.iter_mut().enumerate() {
+ if nulls.is_null(row) {
+ if let Some(null_group) = self.null_group {
+ *group = null_group;
+ } else if !null_key_pending {
+ pending.push(PendingValue::NullKey);
+ null_key_pending = true;
+ }
+ } else {
+ *group = self.group_or_queue(raw_keys[row].as_usize(),
&mut pending);
+ }
+ }
+ }
+
+ if !pending.is_empty() {
+ let indices: Int64Array = pending
+ .iter()
+ .map(|pending| match pending {
+ PendingValue::Value(value_index) => Some(*value_index as
i64),
+ PendingValue::NullKey => None,
+ })
+ .collect();
+ let pending_values = match take(values.as_ref(), &indices, None) {
+ Ok(values) => values,
+ Err(error) => {
+ self.reset_pending(&pending);
+ return Err(error.into());
+ }
+ };
+ let mut pending_groups = Vec::with_capacity(pending.len());
+ if let Err(error) = self.values.intern(&[pending_values], &mut
pending_groups)
+ {
+ self.reset_pending(&pending);
+ return Err(error);
+ }
+ debug_assert_eq!(pending_groups.len(), pending.len());
+
+ for (pending, group_index) in
pending.into_iter().zip(pending_groups) {
+ match pending {
+ PendingValue::Value(value_index) => {
+ self.value_to_group.insert(value_index, group_index);
+ if values.is_null(value_index) {
+ self.null_group = Some(group_index);
+ }
+ }
+ PendingValue::NullKey => self.null_group =
Some(group_index),
+ }
+ }
+ self.value_to_group.maybe_promote(values.len());
+
+ for group in groups.iter_mut() {
+ if *group == PENDING {
+ *group = self.null_group.expect("null key was resolved");
+ } else if *group & PENDING_VALUE_BIT != 0 {
+ let value_index = *group & !PENDING_VALUE_BIT;
+ *group = self.value_to_group.get(value_index);
+ }
+ }
+ }
+
+ debug_assert!(
+ groups
+ .iter()
+ .all(|&group| group != PENDING && group & PENDING_VALUE_BIT ==
0)
+ );
+ self.check_key_overflow()
+ }
+
+ fn size(&self) -> usize {
+ self.values.size() + self.value_to_group.allocation_size() +
size_of::<Self>()
+ }
+
+ fn is_empty(&self) -> bool {
+ self.values.is_empty()
+ }
+
+ fn len(&self) -> usize {
+ self.values.len()
+ }
+
+ fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
+ let mut output = self.values.emit(emit_to)?;
+ debug_assert_eq!(output.len(), 1);
+ let values = output.pop().expect("one dictionary value column");
+ self.clear_cache();
+ let groups = (0..values.len()).collect::<Vec<_>>();
+ Ok(vec![Self::build_dictionary(values, &groups)])
+ }
+
+ fn values_preserving(
+ &mut self,
+ selection: GroupSelection<'_>,
+ ) -> Result<Vec<ArrayRef>> {
+ selection.validate_num_groups(self.len())?;
+
+ let mut group_to_unique = vec![UNRESOLVED; self.len()];
Review Comment:
`values_preserving` always allocates a `group_to_unique` vec sized to the
total group count and rebuilds an index selection, even for
`GroupSelection::all`, where every group is already unique.
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,606 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF)
under one or more contributor license agreements
+
+use std::marker::PhantomData;
+use std::mem::size_of;
+use std::sync::{Arc, Weak};
+
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, Int64Array, PrimitiveArray,
+};
+use arrow::compute::take;
+use arrow::datatypes::{
+ ArrowDictionaryKeyType, ArrowNativeType, DataType, Field, Int8Type,
Int16Type,
+ Int32Type, Int64Type, Schema, UInt8Type, UInt16Type, UInt32Type,
UInt64Type,
+};
+use datafusion_common::{Result, exec_err};
+use datafusion_expr::{EmitTo, GroupSelection};
+use hashbrown::HashMap;
+
+use crate::aggregates::group_values::{
+ GroupValues, new_group_values as new_inner_group_values,
+};
+use crate::aggregates::order::GroupClustering;
+
+const UNRESOLVED: usize = usize::MAX;
+const PENDING: usize = usize::MAX - 1;
+// Sparse entries store both the dictionary index and group id plus hash-table
+// metadata, so prefer the faster dense cache until it would be over 4x larger.
+const SPARSE_CACHE_RATIO: usize = 4;
+// Arrow arrays are bounded by `isize::MAX`, leaving the high bit available to
+// tag unresolved row entries without allocating another row-sized buffer.
+const PENDING_VALUE_BIT: usize = 1 << (usize::BITS - 1);
+
+enum ValueToGroup {
+ Dense(Vec<usize>),
+ Sparse(HashMap<usize, usize>),
+}
+
+impl ValueToGroup {
+ fn new(num_values: usize, num_keys: usize) -> Self {
+ // A dense cache has a much cheaper hot-path lookup and less overhead
per
+ // entry. Use a sparse cache only when initializing the dense cache
would
+ // touch substantially more entries than this batch can reference.
+ if num_values <= num_keys.saturating_mul(SPARSE_CACHE_RATIO) {
+ Self::Dense(vec![UNRESOLVED; num_values])
+ } else {
+ Self::Sparse(HashMap::new())
+ }
+ }
+
+ fn maybe_promote(&mut self, num_values: usize) {
+ let Self::Sparse(values) = self else {
+ return;
+ };
+ if num_values > values.len().saturating_mul(SPARSE_CACHE_RATIO) {
+ return;
+ }
+
+ let mut dense = vec![UNRESOLVED; num_values];
+ for (&value_index, &group_index) in values.iter() {
+ dense[value_index] = group_index;
+ }
+ *self = Self::Dense(dense);
+ }
+ #[inline]
+ fn get(&self, value_index: usize) -> usize {
+ match self {
+ Self::Dense(values) => values[value_index],
+ Self::Sparse(values) => {
+ values.get(&value_index).copied().unwrap_or(UNRESOLVED)
+ }
+ }
+ }
+
+ #[inline]
+ fn insert(&mut self, value_index: usize, group_index: usize) {
+ match self {
+ Self::Dense(values) => values[value_index] = group_index,
+ Self::Sparse(values) => {
+ values.insert(value_index, group_index);
+ }
+ }
+ }
+
+ fn reset(&mut self, value_index: usize) {
+ match self {
+ Self::Dense(values) => values[value_index] = UNRESOLVED,
+ Self::Sparse(values) => {
+ values.remove(&value_index);
+ }
+ }
+ }
+
+ fn allocation_size(&self) -> usize {
+ match self {
+ Self::Dense(values) => values.capacity() * size_of::<usize>(),
+ Self::Sparse(values) => values.allocation_size(),
+ }
+ }
+}
+
+impl Default for ValueToGroup {
+ fn default() -> Self {
+ Self::Dense(vec![])
+ }
+}
+
+enum PendingValue {
+ Value(usize),
+ NullKey,
+}
+
+/// [`GroupValues`] for a single dictionary-encoded group key.
+struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send + Sync> {
+ /// Group values stored using the specialized single-column implementation
+ /// for the dictionary's value type.
+ values: Box<dyn GroupValues>,
+ /// Maps a value index in `cached_values` to its group id.
+ value_to_group: ValueToGroup,
+ /// Weak identity cache for the current dictionary values array.
+ cached_values: Option<Weak<dyn Array>>,
+ /// Group id for both null keys and keys that point to null values.
+ null_group: Option<usize>,
+ _phantom: PhantomData<K>,
+}
+
+impl<K: ArrowDictionaryKeyType + Send + Sync> GroupValuesDictionary<K> {
+ fn try_new(value_field: Field) -> Result<Self> {
+ let schema = Arc::new(Schema::new(vec![value_field]));
+ let values = new_inner_group_values(schema, &GroupClustering::None)?;
+ Ok(Self {
+ values,
+ value_to_group: ValueToGroup::default(),
+ cached_values: None,
+ null_group: None,
+ _phantom: PhantomData,
+ })
+ }
+
+ fn sync_cache(&mut self, values: &ArrayRef, num_keys: usize) {
+ if self
+ .cached_values
+ .as_ref()
+ .is_some_and(|cached| std::ptr::addr_eq(cached.as_ptr(),
Arc::as_ptr(values)))
+ {
+ return;
+ }
+
+ self.value_to_group = ValueToGroup::new(values.len(), num_keys);
+ self.cached_values = Some(Arc::downgrade(values));
+ }
+
+ fn clear_cache(&mut self) {
+ self.value_to_group = ValueToGroup::default();
+ self.cached_values = None;
+ self.null_group = None;
+ }
+
+ fn reset_pending(&mut self, pending: &[PendingValue]) {
+ for value in pending {
+ if let PendingValue::Value(value_index) = value {
+ self.value_to_group.reset(*value_index);
+ }
+ }
+ }
+
+ #[inline]
+ fn group_or_queue(
+ &mut self,
+ value_index: usize,
+ pending: &mut Vec<PendingValue>,
+ ) -> usize {
+ match self.value_to_group.get(value_index) {
+ UNRESOLVED => {
+ self.value_to_group.insert(value_index, PENDING);
+ pending.push(PendingValue::Value(value_index));
+ PENDING_VALUE_BIT | value_index
+ }
+ PENDING => PENDING_VALUE_BIT | value_index,
+ group_index => group_index,
+ }
+ }
+
+ fn check_key_overflow(&self) -> Result<()> {
+ let non_null_count = self.values.len() - self.null_group.is_some() as
usize;
+ let max = match K::DATA_TYPE {
+ DataType::Int8 => i8::MAX as usize,
+ DataType::Int16 => i16::MAX as usize,
+ DataType::Int32 => i32::MAX as usize,
+ DataType::Int64 => i64::MAX as usize,
+ DataType::UInt8 => u8::MAX as usize,
+ DataType::UInt16 => u16::MAX as usize,
+ DataType::UInt32 => u32::MAX as usize,
+ DataType::UInt64 => usize::MAX,
+ _ => unreachable!("validated dictionary key type"),
+ };
+ if non_null_count != 0 && non_null_count - 1 > max {
+ return exec_err!(
+ "Dictionary key type {:?} cannot represent {} distinct values",
+ K::DATA_TYPE,
+ non_null_count
+ );
+ }
+ Ok(())
+ }
+
+ fn build_dictionary(values: ArrayRef, groups: &[usize]) -> ArrayRef {
+ let mut value_keys = Vec::with_capacity(values.len());
+ let mut non_null_indices = Vec::with_capacity(values.len());
+ for value_index in 0..values.len() {
+ if values.is_null(value_index) {
+ value_keys.push(None);
+ } else {
+
value_keys.push(Some(K::Native::usize_as(non_null_indices.len())));
+ non_null_indices.push(value_index as i64);
+ }
+ }
+
+ let values = if non_null_indices.len() == values.len() {
+ values
+ } else {
+ let indices = Int64Array::from(non_null_indices);
+ take(values.as_ref(), &indices, None).expect("compact dictionary
values")
+ };
+ let keys: PrimitiveArray<K> =
+ groups.iter().map(|&group| value_keys[group]).collect();
+ Arc::new(DictionaryArray::<K>::new(keys, values))
+ }
+}
+
+impl<K: ArrowDictionaryKeyType + Send + Sync> GroupValues for
GroupValuesDictionary<K> {
+ fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) ->
Result<()> {
+ assert_eq!(cols.len(), 1);
+ let dictionary = cols[0].as_dictionary::<K>();
+ let values = dictionary.values();
+ let keys = dictionary.keys();
+ self.sync_cache(values, keys.len() - keys.null_count());
+
+ groups.clear();
+ groups.resize(dictionary.len(), PENDING);
+ let mut pending = vec![];
+ let mut null_key_pending = false;
+ let raw_keys = keys.values();
+ if keys.null_count() == 0 {
+ for (group, key) in groups.iter_mut().zip(raw_keys) {
+ *group = self.group_or_queue(key.as_usize(), &mut pending);
+ }
+ } else {
+ let nulls = keys.nulls().unwrap();
+ for (row, group) in groups.iter_mut().enumerate() {
+ if nulls.is_null(row) {
+ if let Some(null_group) = self.null_group {
+ *group = null_group;
+ } else if !null_key_pending {
+ pending.push(PendingValue::NullKey);
+ null_key_pending = true;
+ }
+ } else {
+ *group = self.group_or_queue(raw_keys[row].as_usize(),
&mut pending);
+ }
+ }
+ }
+
+ if !pending.is_empty() {
+ let indices: Int64Array = pending
+ .iter()
+ .map(|pending| match pending {
+ PendingValue::Value(value_index) => Some(*value_index as
i64),
+ PendingValue::NullKey => None,
+ })
+ .collect();
+ let pending_values = match take(values.as_ref(), &indices, None) {
+ Ok(values) => values,
+ Err(error) => {
+ self.reset_pending(&pending);
+ return Err(error.into());
+ }
+ };
+ let mut pending_groups = Vec::with_capacity(pending.len());
+ if let Err(error) = self.values.intern(&[pending_values], &mut
pending_groups)
+ {
+ self.reset_pending(&pending);
+ return Err(error);
+ }
+ debug_assert_eq!(pending_groups.len(), pending.len());
+
+ for (pending, group_index) in
pending.into_iter().zip(pending_groups) {
+ match pending {
+ PendingValue::Value(value_index) => {
+ self.value_to_group.insert(value_index, group_index);
+ if values.is_null(value_index) {
+ self.null_group = Some(group_index);
+ }
+ }
+ PendingValue::NullKey => self.null_group =
Some(group_index),
+ }
+ }
+ self.value_to_group.maybe_promote(values.len());
+
+ for group in groups.iter_mut() {
+ if *group == PENDING {
+ *group = self.null_group.expect("null key was resolved");
+ } else if *group & PENDING_VALUE_BIT != 0 {
+ let value_index = *group & !PENDING_VALUE_BIT;
+ *group = self.value_to_group.get(value_index);
+ }
+ }
+ }
+
+ debug_assert!(
+ groups
+ .iter()
+ .all(|&group| group != PENDING && group & PENDING_VALUE_BIT ==
0)
+ );
+ self.check_key_overflow()
+ }
+
+ fn size(&self) -> usize {
+ self.values.size() + self.value_to_group.allocation_size() +
size_of::<Self>()
+ }
+
+ fn is_empty(&self) -> bool {
+ self.values.is_empty()
+ }
+
+ fn len(&self) -> usize {
+ self.values.len()
+ }
+
+ fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
+ let mut output = self.values.emit(emit_to)?;
+ debug_assert_eq!(output.len(), 1);
+ let values = output.pop().expect("one dictionary value column");
+ self.clear_cache();
Review Comment:
Preserve retained null-group state after partial emission
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,606 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF)
under one or more contributor license agreements
+
+use std::marker::PhantomData;
+use std::mem::size_of;
+use std::sync::{Arc, Weak};
+
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, Int64Array, PrimitiveArray,
+};
+use arrow::compute::take;
+use arrow::datatypes::{
+ ArrowDictionaryKeyType, ArrowNativeType, DataType, Field, Int8Type,
Int16Type,
+ Int32Type, Int64Type, Schema, UInt8Type, UInt16Type, UInt32Type,
UInt64Type,
+};
+use datafusion_common::{Result, exec_err};
+use datafusion_expr::{EmitTo, GroupSelection};
+use hashbrown::HashMap;
+
+use crate::aggregates::group_values::{
+ GroupValues, new_group_values as new_inner_group_values,
+};
+use crate::aggregates::order::GroupClustering;
+
+const UNRESOLVED: usize = usize::MAX;
+const PENDING: usize = usize::MAX - 1;
+// Sparse entries store both the dictionary index and group id plus hash-table
+// metadata, so prefer the faster dense cache until it would be over 4x larger.
+const SPARSE_CACHE_RATIO: usize = 4;
+// Arrow arrays are bounded by `isize::MAX`, leaving the high bit available to
+// tag unresolved row entries without allocating another row-sized buffer.
+const PENDING_VALUE_BIT: usize = 1 << (usize::BITS - 1);
+
+enum ValueToGroup {
+ Dense(Vec<usize>),
+ Sparse(HashMap<usize, usize>),
+}
+
+impl ValueToGroup {
+ fn new(num_values: usize, num_keys: usize) -> Self {
+ // A dense cache has a much cheaper hot-path lookup and less overhead
per
+ // entry. Use a sparse cache only when initializing the dense cache
would
+ // touch substantially more entries than this batch can reference.
+ if num_values <= num_keys.saturating_mul(SPARSE_CACHE_RATIO) {
+ Self::Dense(vec![UNRESOLVED; num_values])
+ } else {
+ Self::Sparse(HashMap::new())
+ }
+ }
+
+ fn maybe_promote(&mut self, num_values: usize) {
+ let Self::Sparse(values) = self else {
+ return;
+ };
+ if num_values > values.len().saturating_mul(SPARSE_CACHE_RATIO) {
+ return;
+ }
+
+ let mut dense = vec![UNRESOLVED; num_values];
+ for (&value_index, &group_index) in values.iter() {
+ dense[value_index] = group_index;
+ }
+ *self = Self::Dense(dense);
+ }
+ #[inline]
+ fn get(&self, value_index: usize) -> usize {
+ match self {
+ Self::Dense(values) => values[value_index],
+ Self::Sparse(values) => {
+ values.get(&value_index).copied().unwrap_or(UNRESOLVED)
+ }
+ }
+ }
+
+ #[inline]
+ fn insert(&mut self, value_index: usize, group_index: usize) {
+ match self {
+ Self::Dense(values) => values[value_index] = group_index,
+ Self::Sparse(values) => {
+ values.insert(value_index, group_index);
+ }
+ }
+ }
+
+ fn reset(&mut self, value_index: usize) {
+ match self {
+ Self::Dense(values) => values[value_index] = UNRESOLVED,
+ Self::Sparse(values) => {
+ values.remove(&value_index);
+ }
+ }
+ }
+
+ fn allocation_size(&self) -> usize {
+ match self {
+ Self::Dense(values) => values.capacity() * size_of::<usize>(),
+ Self::Sparse(values) => values.allocation_size(),
+ }
+ }
+}
+
+impl Default for ValueToGroup {
+ fn default() -> Self {
+ Self::Dense(vec![])
+ }
+}
+
+enum PendingValue {
+ Value(usize),
+ NullKey,
+}
+
+/// [`GroupValues`] for a single dictionary-encoded group key.
+struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send + Sync> {
+ /// Group values stored using the specialized single-column implementation
+ /// for the dictionary's value type.
+ values: Box<dyn GroupValues>,
+ /// Maps a value index in `cached_values` to its group id.
+ value_to_group: ValueToGroup,
+ /// Weak identity cache for the current dictionary values array.
+ cached_values: Option<Weak<dyn Array>>,
+ /// Group id for both null keys and keys that point to null values.
+ null_group: Option<usize>,
+ _phantom: PhantomData<K>,
+}
+
+impl<K: ArrowDictionaryKeyType + Send + Sync> GroupValuesDictionary<K> {
+ fn try_new(value_field: Field) -> Result<Self> {
+ let schema = Arc::new(Schema::new(vec![value_field]));
+ let values = new_inner_group_values(schema, &GroupClustering::None)?;
+ Ok(Self {
+ values,
+ value_to_group: ValueToGroup::default(),
+ cached_values: None,
+ null_group: None,
+ _phantom: PhantomData,
+ })
+ }
+
+ fn sync_cache(&mut self, values: &ArrayRef, num_keys: usize) {
+ if self
+ .cached_values
+ .as_ref()
+ .is_some_and(|cached| std::ptr::addr_eq(cached.as_ptr(),
Arc::as_ptr(values)))
+ {
+ return;
+ }
+
+ self.value_to_group = ValueToGroup::new(values.len(), num_keys);
+ self.cached_values = Some(Arc::downgrade(values));
+ }
+
+ fn clear_cache(&mut self) {
+ self.value_to_group = ValueToGroup::default();
+ self.cached_values = None;
+ self.null_group = None;
+ }
+
+ fn reset_pending(&mut self, pending: &[PendingValue]) {
+ for value in pending {
+ if let PendingValue::Value(value_index) = value {
+ self.value_to_group.reset(*value_index);
+ }
+ }
+ }
+
+ #[inline]
+ fn group_or_queue(
+ &mut self,
+ value_index: usize,
+ pending: &mut Vec<PendingValue>,
+ ) -> usize {
+ match self.value_to_group.get(value_index) {
+ UNRESOLVED => {
+ self.value_to_group.insert(value_index, PENDING);
+ pending.push(PendingValue::Value(value_index));
+ PENDING_VALUE_BIT | value_index
+ }
+ PENDING => PENDING_VALUE_BIT | value_index,
+ group_index => group_index,
+ }
+ }
+
+ fn check_key_overflow(&self) -> Result<()> {
+ let non_null_count = self.values.len() - self.null_group.is_some() as
usize;
+ let max = match K::DATA_TYPE {
+ DataType::Int8 => i8::MAX as usize,
+ DataType::Int16 => i16::MAX as usize,
+ DataType::Int32 => i32::MAX as usize,
+ DataType::Int64 => i64::MAX as usize,
+ DataType::UInt8 => u8::MAX as usize,
+ DataType::UInt16 => u16::MAX as usize,
+ DataType::UInt32 => u32::MAX as usize,
+ DataType::UInt64 => usize::MAX,
+ _ => unreachable!("validated dictionary key type"),
+ };
+ if non_null_count != 0 && non_null_count - 1 > max {
+ return exec_err!(
+ "Dictionary key type {:?} cannot represent {} distinct values",
+ K::DATA_TYPE,
+ non_null_count
+ );
+ }
+ Ok(())
+ }
+
+ fn build_dictionary(values: ArrayRef, groups: &[usize]) -> ArrayRef {
+ let mut value_keys = Vec::with_capacity(values.len());
+ let mut non_null_indices = Vec::with_capacity(values.len());
+ for value_index in 0..values.len() {
+ if values.is_null(value_index) {
+ value_keys.push(None);
+ } else {
+
value_keys.push(Some(K::Native::usize_as(non_null_indices.len())));
+ non_null_indices.push(value_index as i64);
+ }
+ }
+
+ let values = if non_null_indices.len() == values.len() {
+ values
+ } else {
+ let indices = Int64Array::from(non_null_indices);
+ take(values.as_ref(), &indices, None).expect("compact dictionary
values")
+ };
+ let keys: PrimitiveArray<K> =
+ groups.iter().map(|&group| value_keys[group]).collect();
+ Arc::new(DictionaryArray::<K>::new(keys, values))
+ }
+}
+
+impl<K: ArrowDictionaryKeyType + Send + Sync> GroupValues for
GroupValuesDictionary<K> {
+ fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) ->
Result<()> {
+ assert_eq!(cols.len(), 1);
+ let dictionary = cols[0].as_dictionary::<K>();
+ let values = dictionary.values();
+ let keys = dictionary.keys();
+ self.sync_cache(values, keys.len() - keys.null_count());
+
+ groups.clear();
+ groups.resize(dictionary.len(), PENDING);
+ let mut pending = vec![];
+ let mut null_key_pending = false;
+ let raw_keys = keys.values();
+ if keys.null_count() == 0 {
+ for (group, key) in groups.iter_mut().zip(raw_keys) {
+ *group = self.group_or_queue(key.as_usize(), &mut pending);
+ }
+ } else {
+ let nulls = keys.nulls().unwrap();
+ for (row, group) in groups.iter_mut().enumerate() {
+ if nulls.is_null(row) {
+ if let Some(null_group) = self.null_group {
+ *group = null_group;
+ } else if !null_key_pending {
+ pending.push(PendingValue::NullKey);
+ null_key_pending = true;
+ }
+ } else {
+ *group = self.group_or_queue(raw_keys[row].as_usize(),
&mut pending);
+ }
+ }
+ }
+
+ if !pending.is_empty() {
+ let indices: Int64Array = pending
+ .iter()
+ .map(|pending| match pending {
+ PendingValue::Value(value_index) => Some(*value_index as
i64),
+ PendingValue::NullKey => None,
+ })
+ .collect();
+ let pending_values = match take(values.as_ref(), &indices, None) {
+ Ok(values) => values,
+ Err(error) => {
+ self.reset_pending(&pending);
+ return Err(error.into());
+ }
+ };
+ let mut pending_groups = Vec::with_capacity(pending.len());
+ if let Err(error) = self.values.intern(&[pending_values], &mut
pending_groups)
+ {
+ self.reset_pending(&pending);
+ return Err(error);
+ }
+ debug_assert_eq!(pending_groups.len(), pending.len());
+
+ for (pending, group_index) in
pending.into_iter().zip(pending_groups) {
+ match pending {
+ PendingValue::Value(value_index) => {
+ self.value_to_group.insert(value_index, group_index);
+ if values.is_null(value_index) {
+ self.null_group = Some(group_index);
+ }
+ }
+ PendingValue::NullKey => self.null_group =
Some(group_index),
+ }
+ }
+ self.value_to_group.maybe_promote(values.len());
+
+ for group in groups.iter_mut() {
+ if *group == PENDING {
+ *group = self.null_group.expect("null key was resolved");
+ } else if *group & PENDING_VALUE_BIT != 0 {
+ let value_index = *group & !PENDING_VALUE_BIT;
+ *group = self.value_to_group.get(value_index);
+ }
+ }
+ }
+
+ debug_assert!(
+ groups
+ .iter()
+ .all(|&group| group != PENDING && group & PENDING_VALUE_BIT ==
0)
+ );
+ self.check_key_overflow()
+ }
+
+ fn size(&self) -> usize {
+ self.values.size() + self.value_to_group.allocation_size() +
size_of::<Self>()
+ }
+
+ fn is_empty(&self) -> bool {
+ self.values.is_empty()
+ }
+
+ fn len(&self) -> usize {
+ self.values.len()
+ }
+
+ fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
+ let mut output = self.values.emit(emit_to)?;
+ debug_assert_eq!(output.len(), 1);
+ let values = output.pop().expect("one dictionary value column");
+ self.clear_cache();
+ let groups = (0..values.len()).collect::<Vec<_>>();
+ Ok(vec![Self::build_dictionary(values, &groups)])
+ }
+
+ fn values_preserving(
+ &mut self,
+ selection: GroupSelection<'_>,
+ ) -> Result<Vec<ArrayRef>> {
+ selection.validate_num_groups(self.len())?;
+
+ let mut group_to_unique = vec![UNRESOLVED; self.len()];
+ let mut unique_groups = vec![];
+ let mut selected_groups = Vec::with_capacity(selection.len());
Review Comment:
Size selection scratch by the number of selected groups
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
