Rich-T-kid commented on code in PR #27272:
URL: https://github.com/apache/datafusion/pull/27272#discussion_r4551752892
##########
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:
Fixed. Retained null groups are shifted after partial emit.
##########
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:
Done. Partial emits drop emitted entries and shift retained cache ids.
##########
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:
Done. The selection map is now sized from the selected groups.
##########
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:
Done. Identity selections use the direct path, and other selections size the
map to the selection.
##########
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:
Restored. Null-free output uses the raw slot indices directly.
--
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]
