You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "Rich-T-kid (via GitHub)" <gi...@apache.org> on 2026/04/21 14:28:58 UTC
[PR] Optimize Dictionary groupings [datafusion]
Rich-T-kid opened a new pull request, #21765:
URL: https://github.com/apache/datafusion/pull/21765
## Which issue does this PR close?
This PR make an effort towards https://github.com/apache/datafusion/issues/7000 & closing https://github.com/apache/datafusion/issues/7647 + https://github.com/apache/datafusion/issues/21466
This PR aim to close half of https://github.com/apache/datafusion/issues/7647
**A separate follow up PR aims to close the multi-column + dictionary column case**
<!--
We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
-->
- Closes #7647.
## Rationale for this change
<!--
Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.
-->
Issue #7647 ([Materialize Dictionaries in Group Keys](https://github.com/apache/datafusion/issues/7647)) identified that DataFusion was not taking advantage of dictionary encoding during hash aggregation, instead of operating on the compact dictionary representation, it was deserializing dictionary arrays into a generic row-based format, throwing away all the encoding benefits. An initial attempt was made to fix this but it caused regressions and was ultimately rolled back.
### This PR
This PR takes a different approach. Rather than materializing the dictionary into a generic representation, it introduces GroupValuesDictionary, a specialized implementation that operates directly on the dictionary's structure. The key insight is that dictionary encoded columns have inherently low cardinality by design, meaning the same values repeat frequently across rows. Instead of hashing and comparing every row independently, we maintain a mapping of unique value hashes to group IDs so that repeated values are resolved in O(1) after the first encounter. Null values are similarly cached so that subsequent null rows never pay the lookup cost again. For a more detailed explanation of the implementation see[ this design doc](https://docs.google.com/document/d/13mEwd08jKzgw3Dcu4-B0ARSpFsZKPUEwrSBMm1VlRzo/edit?tab=t.ulp2jcmea65z).
## What changes are included in this PR?
Update the match statement in new_group_values to include a custom dictionary encoding branch that works for single fields that are of type Dictionary array
<!--
There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
-->
## Are these changes tested?
Yes, about half of the code in this PR is about testing edge cases.
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example, are they covered by existing tests)?
-->
## Are there any user-facing changes?
No
<!--
If there are user-facing changes then we may require documentation to be updated before approving the PR.
-->
<!--
If there are any breaking changes to public APIs, please add the `api change` label.
-->
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4726293654
🤖 Benchmark completed (GKE) | [trigger](https://github.com/apache/datafusion/pull/21765#issuecomment-4726105593)
**Instance:** `c4a-highmem-16` (12 vCPU / 65 GiB)
<details><summary>CPU Details (lscpu)</summary>
```
Architecture: aarch64
CPU op-mode(s): 64-bit
Byte Order: Little Endian
CPU(s): 16
On-line CPU(s) list: 0-15
Vendor ID: ARM
Model name: Neoverse-V2
Model: 1
Thread(s) per core: 1
Core(s) per cluster: 16
Socket(s): -
Cluster(s): 1
Stepping: r0p1
BogoMIPS: 2000.00
Flags: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache: 1 MiB (16 instances)
L1i cache: 1 MiB (16 instances)
L2 cache: 32 MiB (16 instances)
L3 cache: 80 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; __user pointer sanitization
Vulnerability Spectre v2: Mitigation; CSV2, BHB
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected
```
</details>
<details><summary>Details</summary>
<p>
```
group HEAD rich-t-kid_optimize-dictionary-grouping
----- ---- ---------------------------------------
dict_intern_emit/intern_emit/size_65536_card_1000_null_0.00 5.04 875.0±5.77µs 71.4 MElem/sec 1.00 173.7±1.01µs 359.8 MElem/sec
dict_intern_emit/intern_emit/size_65536_card_20_null_0.00 5.78 788.0±5.09µs 79.3 MElem/sec 1.00 136.3±0.67µs 458.5 MElem/sec
dict_intern_emit/intern_emit/size_65536_card_300_null_0.00 5.54 817.0±5.09µs 76.5 MElem/sec 1.00 147.4±0.45µs 424.0 MElem/sec
dict_intern_emit/intern_emit/size_65536_card_65536_null_0.00 1.54 6.3±0.01ms 9.9 MElem/sec 1.00 4.1±0.03ms 15.3 MElem/sec
dict_intern_emit/intern_emit/size_65536_card_75_null_0.00 5.83 806.4±6.12µs 77.5 MElem/sec 1.00 138.4±0.50µs 451.5 MElem/sec
dict_intern_emit/intern_emit/size_8192_card_1000_null_0.00 2.91 161.6±1.17µs 48.3 MElem/sec 1.00 55.6±0.39µs 140.6 MElem/sec
dict_intern_emit/intern_emit/size_8192_card_20_null_0.00 5.59 105.0±1.02µs 74.4 MElem/sec 1.00 18.8±0.14µs 416.3 MElem/sec
dict_intern_emit/intern_emit/size_8192_card_300_null_0.00 4.14 122.7±1.44µs 63.7 MElem/sec 1.00 29.7±0.22µs 263.3 MElem/sec
dict_intern_emit/intern_emit/size_8192_card_75_null_0.00 5.32 111.4±1.35µs 70.1 MElem/sec 1.00 20.9±0.21µs 373.1 MElem/sec
dict_intern_emit/intern_emit/size_8192_card_8192_null_0.00 1.89 667.5±2.09µs 11.7 MElem/sec 1.00 353.7±4.95µs 22.1 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_65536_card_1000_null_0.10 5.22 4.3±0.03ms 57.6 MElem/sec 1.00 832.5±3.83µs 300.3 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_65536_card_20_null_0.10 5.67 4.1±0.03ms 60.3 MElem/sec 1.00 730.8±4.59µs 342.1 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_65536_card_300_null_0.10 5.55 4.2±0.03ms 59.3 MElem/sec 1.00 760.2±4.13µs 328.9 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_65536_card_65536_null_0.10 1.13 17.2±0.05ms 14.5 MElem/sec 1.00 15.2±0.08ms 16.5 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_65536_card_75_null_0.10 5.70 4.2±0.02ms 59.3 MElem/sec 1.00 739.7±4.35µs 338.0 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_8192_card_1000_null_0.10 3.23 605.4±3.21µs 51.6 MElem/sec 1.00 187.4±1.25µs 166.7 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_8192_card_20_null_0.10 5.42 504.4±3.93µs 62.0 MElem/sec 1.00 93.0±0.34µs 336.1 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_8192_card_300_null_0.10 4.61 536.1±3.02µs 58.3 MElem/sec 1.00 116.4±0.70µs 268.5 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_8192_card_75_null_0.10 5.49 516.1±3.25µs 60.5 MElem/sec 1.00 94.0±0.44µs 332.6 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_8192_card_8192_null_0.10 1.67 1472.0±7.21µs 21.2 MElem/sec 1.00 881.4±13.47µs 35.5 MElem/sec
```
</p>
</details>
<details><summary>Resource Usage</summary>
**dictionary_group_values — base (merge-base)**
| Metric | Value |
|--------|-------|
| Wall time | 450.1s |
| Peak memory | 598.7 MiB |
| Avg memory | 83.5 MiB |
| CPU user | 230.5s |
| CPU sys | 15.1s |
| Peak spill | 0 B |
**dictionary_group_values — branch**
| Metric | Value |
|--------|-------|
| Wall time | 430.1s |
| Peak memory | 856.6 MiB |
| Avg memory | 50.5 MiB |
| CPU user | 214.0s |
| CPU sys | 19.4s |
| Peak spill | 0 B |
</details>
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4731088959
re-ran the E2E benchmarks as well in `benchmarks/src/dict.rs`
<img width="795" height="147" alt="Image 6-17-26 at 10 00 AM" src="https://github.com/user-attachments/assets/d2faf734-bbaf-47d1-a0f3-4bf4855c2d52" />
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4673219877
> This provides a performance increase in every case where dictionary arrays are expected to be used — low to medium cardinality. If the data is extremely high cardinality, dictionary arrays are the wrong data type. I don't think it's possible to get the best of both worlds in this case.
Is there any way to reduce the overhead so the difference is not as much?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "kumarUjjawal (via GitHub)" <gi...@apache.org>.
kumarUjjawal commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4408385862
run benchmarks
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4408256796
cc @Dandandan @alamb
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4422445146
actually it seems that the @/adriangbot list of supported benchmark commands is hardcoded. ill just attach a screenshot. thank you again @kumarUjjawal 🎆
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "kumarUjjawal (via GitHub)" <gi...@apache.org>.
kumarUjjawal commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4726105593
run benchmark dictionary_group_values
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4649570634
@kumarUjjawal for the E2E sql test there are already a couple that exist for dictionary so I just added one more that uses StringViews as a value type. As for streaming, I assume you were referring to the partial emit that datafusion uses for spills to disk? I added a [test](https://github.com/apache/datafusion/pull/21765/changes#diff-4515126fc8522ddcc05a5673023bf83367b0c21e31f2933e792d42da8f2d6a63R838) that validates the original group ID is correct across partial emits. let me know if you were referring to something else & I can add a test for that as well!
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on code in PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#discussion_r3454836051
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,1003 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, LargeStringArray, LargeStringBuilder,
+ ListArray, ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ offset: usize,
+ len: usize,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ /// Packed byte storage for all group values.
+ row_buffer: Vec<u8>,
+ /// `row_offsets[g]` = start of group `g` in `row_buffer`;
+ row_offsets: Vec<usize>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ // key to group vector scratch space, used to avoid re-allocating a new vector on each call to intern
+ key_to_group: Vec<Option<usize>>,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ row_buffer: Vec::new(),
+ row_offsets: Vec::new(),
+ unique_dict_value_mapping: HashTable::new(),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ key_to_group: Vec::new(),
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ let row_buffer = &self.row_buffer;
+ if let Some(e) = self.unique_dict_value_mapping.find(hash, |e| {
+ e.hash == hash && &row_buffer[e.offset..e.offset + e.len] == raw
+ }) {
+ return e.group_id;
+ }
+ let new_group_id = self.row_offsets.len();
+ let offset = self.row_buffer.len();
+ self.row_offsets.push(offset);
+ self.row_buffer.extend_from_slice(raw);
+ self.unique_dict_value_mapping.insert_unique(
+ hash,
+ DictEntry {
+ hash,
+ group_id: new_group_id,
+ offset,
+ len: raw.len(),
+ },
+ |e| e.hash,
+ );
+ new_group_id
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.row_offsets.len();
+ self.row_offsets.push(self.row_buffer.len()); // null: empty span
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+
+ fn transform_into_array(
+ &self,
+ n: usize,
+ null_group_id: Option<usize>,
+ ) -> Result<ArrayRef> {
+ let data_capacity = n * AVG_BYTES_PER_DICT_VALUE;
+ let raw = (0..n).map(|i| {
+ if Some(i) == null_group_id {
+ return None;
+ }
+ let start = self.row_offsets[i];
+ // last group has no i+1 entry; its end is the buffer tail.
+ let end = self
+ .row_offsets
+ .get(i + 1)
+ .copied()
+ .unwrap_or(self.row_buffer.len());
+ Some(&self.row_buffer[start..end])
+ });
+ match &self.value_dt {
Review Comment:
thank you for the review @alamb
I agree on avoiding ballooning LOC, adding support for every data type becomes more trouble than the compute savings justify. That said, I think restricting support to `Dictionary<_, Utf8/Utf8View>` is the right call. It significantly simplifies the code paths while still delivering a large performance boost for the most common production data shapes (low-cardinality group bys), and similar to other single-column group by specializations, it leaves a viable fallback that loses nothing when the fast path isn't hit.
With that in mind, @zhuqi-lucas's work looks very interesting -- I had a similar idea with #22891 and #21878. The issue was that multi-column group-bys on dictionaries fell through to `GroupValuesRows`, which was slow. I'm actually working on a multi-column dictionary group by in #22983, but I'm running into similar issues around handling multiple distinct types. For the multi-column case it makes more sense to build on @zhuqi-lucas's work, since the number of types that need to be supported explodes and there are diminishing returns to exploiting dictionary properties across multiple columns and intern calls ([key space cache](https://github.com/apache/datafusion/pull/22983/changes#diff-84ceab5d992f30b340d52094b28685b5ecfd6f68f27867641ca1a9453ad279baR209), [arc ptr cache](https://github.com/apache/datafusion/pull/22983/changes#diff-84ceab5d992f30b340d52094b28685b5ecfd6f68f27867641ca1a9453ad279baR193), [compute keys once and reuse](https://github.com/apache/datafusion/pull/22983/c
hanges#diff-84ceab5d992f30b340d52094b28685b5ecfd6f68f27867641ca1a9453ad279baR342)).
As for this PR specifically, I do think it adds significant value. Would restricting the supported value types be a viable path forward? It would keep the type combinations from exploding while still letting the community benefit from the speedup.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#discussion_r3454366652
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,1003 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, LargeStringArray, LargeStringBuilder,
+ ListArray, ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ offset: usize,
+ len: usize,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ /// Packed byte storage for all group values.
+ row_buffer: Vec<u8>,
+ /// `row_offsets[g]` = start of group `g` in `row_buffer`;
+ row_offsets: Vec<usize>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ // key to group vector scratch space, used to avoid re-allocating a new vector on each call to intern
+ key_to_group: Vec<Option<usize>>,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ row_buffer: Vec::new(),
+ row_offsets: Vec::new(),
+ unique_dict_value_mapping: HashTable::new(),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ key_to_group: Vec::new(),
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ let row_buffer = &self.row_buffer;
+ if let Some(e) = self.unique_dict_value_mapping.find(hash, |e| {
+ e.hash == hash && &row_buffer[e.offset..e.offset + e.len] == raw
+ }) {
+ return e.group_id;
+ }
+ let new_group_id = self.row_offsets.len();
+ let offset = self.row_buffer.len();
+ self.row_offsets.push(offset);
+ self.row_buffer.extend_from_slice(raw);
+ self.unique_dict_value_mapping.insert_unique(
+ hash,
+ DictEntry {
+ hash,
+ group_id: new_group_id,
+ offset,
+ len: raw.len(),
+ },
+ |e| e.hash,
+ );
+ new_group_id
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.row_offsets.len();
+ self.row_offsets.push(self.row_buffer.len()); // null: empty span
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+
+ fn transform_into_array(
+ &self,
+ n: usize,
+ null_group_id: Option<usize>,
+ ) -> Result<ArrayRef> {
+ let data_capacity = n * AVG_BYTES_PER_DICT_VALUE;
+ let raw = (0..n).map(|i| {
+ if Some(i) == null_group_id {
+ return None;
+ }
+ let start = self.row_offsets[i];
+ // last group has no i+1 entry; its end is the buffer tail.
+ let end = self
+ .row_offsets
+ .get(i + 1)
+ .copied()
+ .unwrap_or(self.row_buffer.len());
+ Some(&self.row_buffer[start..end])
+ });
+ match &self.value_dt {
Review Comment:
In general I am worried about needing special code for different element types in the dictionary
In general, this seems like:
1. It limits the types of DictionaryArray that can be supported (e.g. this doesn't support Dicts with struct elements)
2. It will have substantial amounts of code (b/c each type of dictionary now gets its own branch in this match statement)
What I think @zhuqi-lucas is trying to do as part of
- https://github.com/apache/datafusion/issues/22682
And some related PRs is to make a generic way to handle nested types that doesn't require generating code for all possible type combinations.
##########
datafusion/sqllogictest/test_files/dictionary.slt:
##########
@@ -515,3 +515,15 @@ DROP TABLE dict_hash_10;
statement ok
DROP TABLE dict_hash_src;
+
+query TI
+SELECT
Review Comment:
there is a bunch of code to handle many more types of dictionaries such as Lists, etc. -- can we please add more tests to cover those cases?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on code in PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#discussion_r3454840874
##########
datafusion/sqllogictest/test_files/dictionary.slt:
##########
@@ -515,3 +515,15 @@ DROP TABLE dict_hash_10;
statement ok
DROP TABLE dict_hash_src;
+
+query TI
+SELECT
Review Comment:
this PR is already at 1k LOC, I'll open a separate PR
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4723446208
@alamb could you re-run the micro-benchmarks?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4291560321
<html><head></head><body><p><strong>TPC-H In-Memory SF10 Benchmark Summary</strong></p>
Metric | Value
-- | --
Total Time (main) | 40.17ms
Total Time (rich-t-kid_optimize-dictionary-grouping) | 35.87ms
Average Time (main) | 1.83ms
Average Time (rich-t-kid_optimize-dictionary-grouping) | 1.63ms
Queries Faster | 20
Queries Slower | 1
Queries with No Change | 1
Queries with Failure | 0
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4290861738
Do we have any performance benchmarks for this PR?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4676981749
When arrow-rs reads RLE-encoded strings from a Parquet file, it decodes them into dictionary arrays. These flow through operators like any other array. For operators like `AggregateExec`, the dictionary gets expanded during grouping: the `GroupValuesRows` grouping implementation uses `RowConverter`, which expands dictionaries when converting columns to their row representation. ([source](https://github.com/apache/arrow-rs/blob/301eb26bb92362e531bd6c39980292ebcae8e8aa/arrow-row/src/lib.rs#L2060))
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "kumarUjjawal (via GitHub)" <gi...@apache.org>.
kumarUjjawal commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4742641276
@alamb This looks like a good improvement to me, what do you think?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on code in PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#discussion_r3455196070
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,1003 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, LargeStringArray, LargeStringBuilder,
+ ListArray, ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ offset: usize,
+ len: usize,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ /// Packed byte storage for all group values.
+ row_buffer: Vec<u8>,
+ /// `row_offsets[g]` = start of group `g` in `row_buffer`;
+ row_offsets: Vec<usize>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ // key to group vector scratch space, used to avoid re-allocating a new vector on each call to intern
+ key_to_group: Vec<Option<usize>>,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ row_buffer: Vec::new(),
+ row_offsets: Vec::new(),
+ unique_dict_value_mapping: HashTable::new(),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ key_to_group: Vec::new(),
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ let row_buffer = &self.row_buffer;
+ if let Some(e) = self.unique_dict_value_mapping.find(hash, |e| {
+ e.hash == hash && &row_buffer[e.offset..e.offset + e.len] == raw
+ }) {
+ return e.group_id;
+ }
+ let new_group_id = self.row_offsets.len();
+ let offset = self.row_buffer.len();
+ self.row_offsets.push(offset);
+ self.row_buffer.extend_from_slice(raw);
+ self.unique_dict_value_mapping.insert_unique(
+ hash,
+ DictEntry {
+ hash,
+ group_id: new_group_id,
+ offset,
+ len: raw.len(),
+ },
+ |e| e.hash,
+ );
+ new_group_id
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.row_offsets.len();
+ self.row_offsets.push(self.row_buffer.len()); // null: empty span
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+
+ fn transform_into_array(
+ &self,
+ n: usize,
+ null_group_id: Option<usize>,
+ ) -> Result<ArrayRef> {
+ let data_capacity = n * AVG_BYTES_PER_DICT_VALUE;
+ let raw = (0..n).map(|i| {
+ if Some(i) == null_group_id {
+ return None;
+ }
+ let start = self.row_offsets[i];
+ // last group has no i+1 entry; its end is the buffer tail.
+ let end = self
+ .row_offsets
+ .get(i + 1)
+ .copied()
+ .unwrap_or(self.row_buffer.len());
+ Some(&self.row_buffer[start..end])
+ });
+ match &self.value_dt {
Review Comment:
I touched on this in the above comment but there are properties of single-column dictionary group bys that multi-column group bys can't practically take advantage of:
**null handling**: in a single-column group by, all nulls map to the same group id. This avoids the combinatorial null space that emerges with multiple columns -- (null, _, _), (_, null, _), (_, _, null), etc. In the singular case we can store the null group id directly and look it up with no hashing, a meaningful win for data with a high null percentage.
**arc pointer identity cache**: consecutive batches of the same dictionary column often share the same underlying Arc buffer; a pointer equality check lets us skip re-hashing the values buffer entirely across batches.
**group id precomputation**: group ids can be computed once per unique dictionary key and stored in a flat lookup table, avoiding per-row hashing. In the multi-column case the number of possible key combinations grows multiplicatively -- two columns with cardinalities 100 and 200 produce up to 20,000 combinations, making this impractical.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on code in PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#discussion_r3373379676
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,878 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, LargeStringArray, LargeStringBuilder,
+ ListArray, ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
Review Comment:
ideally the 'cheap signal' would be the first dictionary arrays values, assuming that the values are unique. the tables capacity is pre-allocated at creation time not when `intern()` is first called so I think it makes sense to just remove the pre-allocation all together and let it grow.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4662502779
@kumarUjjawal that is expected. In the case where every value is unique we are adding overhead. a regression occurs when the values & keys are 1:1
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4662528621
This provides a performance increase in every case where dictionary arrays are expected to be used — low to medium cardinality. If the data is high cardinality, dictionary arrays are the wrong data type. I don't think it'd be possible to get the best of both worlds in this case.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4473626561
Shrunk the PR by 200 lines so this should be easier to review 🚀
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "kumarUjjawal (via GitHub)" <gi...@apache.org>.
kumarUjjawal commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4412614782
run benchmarks
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4422336556
🤖 Criterion benchmark completed (GKE) | [trigger](https://github.com/apache/datafusion/pull/21765#issuecomment-4422189550)
**Instance:** `c4a-highmem-16` (12 vCPU / 65 GiB)
<details><summary>CPU Details (lscpu)</summary>
```
Architecture: aarch64
CPU op-mode(s): 64-bit
Byte Order: Little Endian
CPU(s): 16
On-line CPU(s) list: 0-15
Vendor ID: ARM
Model name: Neoverse-V2
Model: 1
Thread(s) per core: 1
Core(s) per cluster: 16
Socket(s): -
Cluster(s): 1
Stepping: r0p1
BogoMIPS: 2000.00
Flags: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache: 1 MiB (16 instances)
L1i cache: 1 MiB (16 instances)
L2 cache: 32 MiB (16 instances)
L3 cache: 80 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; __user pointer sanitization
Vulnerability Spectre v2: Mitigation; CSV2, BHB
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected
```
</details>
<details><summary>Details</summary>
<p>
```
group main rich-t-kid_optimize-dictionary-grouping
----- ---- ---------------------------------------
dict_intern_emit/intern_emit/size_65536_card_1000_null_0.00 2.97 897.3±9.09µs 69.7 MElem/sec 1.00 301.8±14.24µs 207.1 MElem/sec
dict_intern_emit/intern_emit/size_65536_card_20_null_0.00 4.73 786.7±9.43µs 79.4 MElem/sec 1.00 166.2±3.56µs 375.9 MElem/sec
dict_intern_emit/intern_emit/size_65536_card_300_null_0.00 3.93 821.0±6.91µs 76.1 MElem/sec 1.00 208.7±3.69µs 299.5 MElem/sec
dict_intern_emit/intern_emit/size_65536_card_65536_null_0.00 1.00 6.3±0.01ms 9.8 MElem/sec 2.50 15.8±0.22ms 3.9 MElem/sec
dict_intern_emit/intern_emit/size_65536_card_75_null_0.00 4.54 804.7±8.12µs 77.7 MElem/sec 1.00 177.1±3.73µs 352.9 MElem/sec
dict_intern_emit/intern_emit/size_8192_card_1000_null_0.00 1.29 180.7±1.17µs 43.2 MElem/sec 1.00 140.4±6.50µs 55.6 MElem/sec
dict_intern_emit/intern_emit/size_8192_card_20_null_0.00 4.44 105.7±4.54µs 73.9 MElem/sec 1.00 23.8±0.22µs 328.3 MElem/sec
dict_intern_emit/intern_emit/size_8192_card_300_null_0.00 2.19 130.2±27.22µs 60.0 MElem/sec 1.00 59.3±1.67µs 131.7 MElem/sec
dict_intern_emit/intern_emit/size_8192_card_75_null_0.00 3.52 111.1±1.13µs 70.3 MElem/sec 1.00 31.6±0.52µs 247.5 MElem/sec
dict_intern_emit/intern_emit/size_8192_card_8192_null_0.00 1.00 688.6±2.06µs 11.3 MElem/sec 1.50 1032.1±55.93µs 7.6 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_65536_card_1000_null_0.10 2.83 4.4±0.02ms 57.1 MElem/sec 1.00 1545.6±12.39µs 161.8 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_65536_card_20_null_0.10 3.16 4.1±0.02ms 60.3 MElem/sec 1.00 1311.1±3.19µs 190.7 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_65536_card_300_null_0.10 3.07 4.2±0.02ms 59.0 MElem/sec 1.00 1380.9±5.30µs 181.0 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_65536_card_65536_null_0.10 1.00 18.2±0.06ms 13.7 MElem/sec 1.62 29.5±0.73ms 8.5 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_65536_card_75_null_0.10 3.17 4.2±0.02ms 59.4 MElem/sec 1.00 1328.8±3.13µs 188.1 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_8192_card_1000_null_0.10 1.82 621.0±3.51µs 50.3 MElem/sec 1.00 340.4±8.44µs 91.8 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_8192_card_20_null_0.10 3.09 495.5±2.37µs 63.1 MElem/sec 1.00 160.1±0.55µs 195.2 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_8192_card_300_null_0.10 2.51 534.6±2.83µs 58.5 MElem/sec 1.00 213.1±2.83µs 146.6 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_8192_card_75_null_0.10 2.99 510.1±2.46µs 61.3 MElem/sec 1.00 170.4±0.73µs 183.3 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_8192_card_8192_null_0.10 1.00 1504.0±8.86µs 20.8 MElem/sec 1.16 1741.3±17.55µs 17.9 MElem/sec
```
</p>
</details>
<details><summary>Resource Usage</summary>
**base (merge-base)**
| Metric | Value |
|--------|-------|
| Wall time | 200.0s |
| Peak memory | 5.3 GiB |
| Avg memory | 4.9 GiB |
| CPU user | 233.2s |
| CPU sys | 15.9s |
| Peak spill | 0 B |
**branch**
| Metric | Value |
|--------|-------|
| Wall time | 210.0s |
| Peak memory | 9.1 GiB |
| Avg memory | 8.9 GiB |
| CPU user | 255.5s |
| CPU sys | 1.8s |
| Peak spill | 0 B |
</details>
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4418375523
🤖 Benchmark running (GKE) | [trigger](https://github.com/apache/datafusion/pull/21765#issuecomment-4418355614)
**Instance:** `c4a-highmem-16` (12 vCPU / 65 GiB) | `Linux bench-c4418355614-12-pkghb 6.12.68+ #1 SMP Wed Apr 1 02:23:28 UTC 2026 aarch64 GNU/Linux`
<details><summary>CPU Details (lscpu)</summary>
```
Architecture: aarch64
CPU op-mode(s): 64-bit
Byte Order: Little Endian
CPU(s): 16
On-line CPU(s) list: 0-15
Vendor ID: ARM
Model name: Neoverse-V2
Model: 1
Thread(s) per core: 1
Core(s) per cluster: 16
Socket(s): -
Cluster(s): 1
Stepping: r0p1
BogoMIPS: 2000.00
Flags: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache: 1 MiB (16 instances)
L1i cache: 1 MiB (16 instances)
L2 cache: 32 MiB (16 instances)
L3 cache: 80 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; __user pointer sanitization
Vulnerability Spectre v2: Mitigation; CSV2, BHB
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected
```
</details>
Comparing rich-t-kid/optimize-dictionary-grouping (577baab1e8f652768893653cd1b02e782fb76f63) to aca4d13 (merge-base) [diff](https://github.com/apache/datafusion/compare/aca4d1377bf9785e6fa9a154c579cc761ca7b54b..577baab1e8f652768893653cd1b02e782fb76f63) using: tpcds
Results will be posted here when complete
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4293363131
Removing any extra allocations like **.to_vec()** and repeated hash look ups into map
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4299950312
### tpch_sft1
<img width="1052" height="529" alt="Image 4-22-26 at 5 02 PM" src="https://github.com/user-attachments/assets/67a2a519-ac21-47d9-80e0-af7e783f72f0" />
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4294548121
@Rich-T-kid are you sure tpch-1 looks at the data at all in your benchmarks? It looks suspicially fast, I think it might only plan/setup the queries but not look at the data.
Perhaps the path is wrong or using it from the wrong directory?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4418471662
🤖 Benchmark completed (GKE) | [trigger](https://github.com/apache/datafusion/pull/21765#issuecomment-4418355614)
**Instance:** `c4a-highmem-16` (12 vCPU / 65 GiB)
<details><summary>CPU Details (lscpu)</summary>
```
Architecture: aarch64
CPU op-mode(s): 64-bit
Byte Order: Little Endian
CPU(s): 16
On-line CPU(s) list: 0-15
Vendor ID: ARM
Model name: Neoverse-V2
Model: 1
Thread(s) per core: 1
Core(s) per cluster: 16
Socket(s): -
Cluster(s): 1
Stepping: r0p1
BogoMIPS: 2000.00
Flags: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache: 1 MiB (16 instances)
L1i cache: 1 MiB (16 instances)
L2 cache: 32 MiB (16 instances)
L3 cache: 80 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; __user pointer sanitization
Vulnerability Spectre v2: Mitigation; CSV2, BHB
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected
```
</details>
<details><summary>Details</summary>
<p>
```
Comparing HEAD and rich-t-kid_optimize-dictionary-grouping
--------------------
Benchmark clickbench_partitioned.json
--------------------
┏━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━┓
┃ Query ┃ HEAD ┃ rich-t-kid_optimize-dictionary-grouping ┃ Change ┃
┡━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━┩
│ QQuery 0 │ 1.21 / 4.63 ±6.77 / 18.18 ms │ 1.17 / 4.56 ±6.71 / 17.98 ms │ no change │
│ QQuery 1 │ 12.92 / 13.17 ±0.20 / 13.47 ms │ 12.44 / 12.55 ±0.08 / 12.64 ms │ no change │
│ QQuery 2 │ 36.01 / 36.42 ±0.37 / 36.98 ms │ 35.77 / 35.94 ±0.13 / 36.09 ms │ no change │
│ QQuery 3 │ 30.91 / 31.40 ±0.51 / 32.32 ms │ 30.22 / 31.88 ±2.84 / 37.54 ms │ no change │
│ QQuery 4 │ 228.69 / 233.10 ±2.95 / 236.94 ms │ 233.07 / 235.74 ±1.51 / 236.98 ms │ no change │
│ QQuery 5 │ 275.95 / 277.96 ±1.76 / 280.31 ms │ 274.40 / 277.57 ±3.62 / 284.63 ms │ no change │
│ QQuery 6 │ 6.69 / 7.16 ±0.32 / 7.61 ms │ 6.47 / 7.11 ±0.38 / 7.64 ms │ no change │
│ QQuery 7 │ 13.93 / 14.04 ±0.09 / 14.17 ms │ 13.49 / 13.60 ±0.06 / 13.68 ms │ no change │
│ QQuery 8 │ 309.13 / 314.12 ±3.66 / 319.92 ms │ 310.07 / 311.60 ±1.34 / 314.03 ms │ no change │
│ QQuery 9 │ 447.16 / 452.56 ±5.58 / 462.43 ms │ 443.39 / 446.68 ±2.40 / 449.72 ms │ no change │
│ QQuery 10 │ 68.91 / 70.08 ±0.68 / 70.98 ms │ 69.31 / 69.99 ±1.08 / 72.15 ms │ no change │
│ QQuery 11 │ 79.50 / 80.20 ±0.58 / 81.24 ms │ 80.44 / 80.91 ±0.55 / 81.93 ms │ no change │
│ QQuery 12 │ 267.53 / 271.81 ±3.57 / 278.07 ms │ 269.30 / 274.47 ±3.41 / 278.51 ms │ no change │
│ QQuery 13 │ 377.74 / 389.83 ±8.41 / 399.72 ms │ 380.61 / 385.91 ±7.53 / 400.77 ms │ no change │
│ QQuery 14 │ 273.34 / 279.19 ±6.50 / 291.70 ms │ 278.47 / 281.60 ±2.72 / 286.65 ms │ no change │
│ QQuery 15 │ 274.74 / 279.48 ±4.29 / 286.21 ms │ 275.09 / 286.51 ±13.16 / 312.21 ms │ no change │
│ QQuery 16 │ 598.18 / 604.84 ±5.53 / 613.86 ms │ 593.52 / 607.96 ±13.86 / 631.30 ms │ no change │
│ QQuery 17 │ 596.88 / 601.49 ±3.74 / 607.84 ms │ 598.32 / 609.37 ±8.28 / 618.27 ms │ no change │
│ QQuery 18 │ 1175.52 / 1207.11 ±22.50 / 1232.74 ms │ 1186.36 / 1199.01 ±9.38 / 1211.06 ms │ no change │
│ QQuery 19 │ 27.98 / 29.70 ±2.90 / 35.48 ms │ 27.69 / 36.62 ±13.53 / 62.81 ms │ 1.23x slower │
│ QQuery 20 │ 517.19 / 523.54 ±8.01 / 539.20 ms │ 515.26 / 520.57 ±5.50 / 531.18 ms │ no change │
│ QQuery 21 │ 593.52 / 603.08 ±7.33 / 613.64 ms │ 595.54 / 605.09 ±6.22 / 612.27 ms │ no change │
│ QQuery 22 │ 1042.89 / 1059.09 ±10.56 / 1074.32 ms │ 1058.76 / 1065.13 ±3.96 / 1070.63 ms │ no change │
│ QQuery 23 │ 3133.18 / 3159.01 ±15.57 / 3175.81 ms │ 3162.18 / 3189.13 ±16.04 / 3209.91 ms │ no change │
│ QQuery 24 │ 41.45 / 44.00 ±4.41 / 52.81 ms │ 41.27 / 42.18 ±0.96 / 43.96 ms │ no change │
│ QQuery 25 │ 109.76 / 111.52 ±1.61 / 113.52 ms │ 110.83 / 114.60 ±4.34 / 122.77 ms │ no change │
│ QQuery 26 │ 41.83 / 42.28 ±0.35 / 42.84 ms │ 41.58 / 42.50 ±0.59 / 43.18 ms │ no change │
│ QQuery 27 │ 669.72 / 681.00 ±7.11 / 691.24 ms │ 667.79 / 671.71 ±5.62 / 682.46 ms │ no change │
│ QQuery 28 │ 2998.21 / 3014.34 ±10.30 / 3029.28 ms │ 3021.73 / 3035.90 ±11.30 / 3056.43 ms │ no change │
│ QQuery 29 │ 41.73 / 48.96 ±14.21 / 77.37 ms │ 41.28 / 48.22 ±12.74 / 73.68 ms │ no change │
│ QQuery 30 │ 295.78 / 300.40 ±3.57 / 305.03 ms │ 297.28 / 300.64 ±2.53 / 303.31 ms │ no change │
│ QQuery 31 │ 291.19 / 296.63 ±6.61 / 309.55 ms │ 287.32 / 292.50 ±6.18 / 304.15 ms │ no change │
│ QQuery 32 │ 905.12 / 916.14 ±10.41 / 932.77 ms │ 892.73 / 911.21 ±14.60 / 936.91 ms │ no change │
│ QQuery 33 │ 1413.14 / 1425.68 ±7.76 / 1432.41 ms │ 1414.95 / 1429.33 ±12.75 / 1447.64 ms │ no change │
│ QQuery 34 │ 1423.33 / 1436.06 ±7.32 / 1443.59 ms │ 1426.73 / 1440.40 ±16.65 / 1471.95 ms │ no change │
│ QQuery 35 │ 282.60 / 295.33 ±9.45 / 309.29 ms │ 279.76 / 299.19 ±13.08 / 317.18 ms │ no change │
│ QQuery 36 │ 62.47 / 66.98 ±3.07 / 71.27 ms │ 63.37 / 65.81 ±2.95 / 71.37 ms │ no change │
│ QQuery 37 │ 34.83 / 36.18 ±1.32 / 38.63 ms │ 35.07 / 39.58 ±2.82 / 42.52 ms │ 1.09x slower │
│ QQuery 38 │ 40.05 / 42.37 ±3.03 / 48.22 ms │ 40.30 / 43.14 ±2.00 / 46.34 ms │ no change │
│ QQuery 39 │ 123.16 / 127.29 ±3.61 / 131.92 ms │ 122.53 / 130.37 ±4.43 / 135.02 ms │ no change │
│ QQuery 40 │ 14.24 / 15.77 ±1.96 / 19.39 ms │ 14.18 / 15.43 ±1.50 / 18.14 ms │ no change │
│ QQuery 41 │ 13.65 / 13.78 ±0.10 / 13.95 ms │ 13.54 / 16.38 ±4.61 / 25.48 ms │ 1.19x slower │
│ QQuery 42 │ 13.20 / 13.44 ±0.25 / 13.91 ms │ 12.78 / 14.01 ±1.74 / 17.45 ms │ no change │
└───────────┴───────────────────────────────────────┴─────────────────────────────────────────┴──────────────┘
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┓
┃ Benchmark Summary ┃ ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━┩
│ Total Time (HEAD) │ 19471.16ms │
│ Total Time (rich-t-kid_optimize-dictionary-grouping) │ 19542.58ms │
│ Average Time (HEAD) │ 452.82ms │
│ Average Time (rich-t-kid_optimize-dictionary-grouping) │ 454.48ms │
│ Queries Faster │ 0 │
│ Queries Slower │ 3 │
│ Queries with No Change │ 40 │
│ Queries with Failure │ 0 │
└────────────────────────────────────────────────────────┴────────────┘
```
</p>
</details>
<details><summary>Resource Usage</summary>
**clickbench_partitioned — base (merge-base)**
| Metric | Value |
|--------|-------|
| Wall time | 100.0s |
| Peak memory | 31.1 GiB |
| Avg memory | 23.1 GiB |
| CPU user | 1025.0s |
| CPU sys | 63.2s |
| Peak spill | 0 B |
**clickbench_partitioned — branch**
| Metric | Value |
|--------|-------|
| Wall time | 100.0s |
| Peak memory | 30.3 GiB |
| Avg memory | 22.9 GiB |
| CPU user | 1028.9s |
| CPU sys | 63.5s |
| Peak spill | 0 B |
</details>
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "kumarUjjawal (via GitHub)" <gi...@apache.org>.
kumarUjjawal commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4422189550
run benchmark dictionary_group_values
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4704021677
> > This provides a performance increase in every case where dictionary arrays are expected to be used — low to medium cardinality. If the data is extremely high cardinality, dictionary arrays are the wrong data type. I don't think it's possible to get the best of both worlds in this case.
>
> Is there any way to reduce the overhead so the difference is not as much?
There are spots that can be further optimized for example, storing a single byte buffer instead of a vector of vectors. This avoids a layer of indirection and the n-element shifts that occur [here](https://github.com/apache/datafusion/pull/21765/changes#diff-4515126fc8522ddcc05a5673023bf83367b0c21e31f2933e792d42da8f2d6a63R402). I've tracked this in a separate [issue](https://github.com/apache/datafusion/issues/22078), which provides a bit more detail. That said, this is an optimization on top of the current approach and won't change the fundamental characteristic of being slower with high-cardinality data. Do you have ideas on how to further close the gap with `GroupValueRows`?
@alamb
cc @kumarUjjawal
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "LiaCastaneda (via GitHub)" <gi...@apache.org>.
LiaCastaneda commented on code in PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#discussion_r3356276108
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,783 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, DictionaryArray, LargeStringArray, LargeStringBuilder, ListArray,
+ ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+/// Entry stored inside `unique_dict_value_mapping`. Caching `hash` on the entry
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ bytes: Vec<u8>,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ // stores the order new unique elements are seen for self.emit()
+ seen_elements: Vec<Option<Vec<u8>>>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ intern_called: bool,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ seen_elements: Vec::new(),
+ unique_dict_value_mapping: HashTable::with_capacity(INITIAL_PRE_ALLOCATION),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ intern_called: false,
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ /// Returns the existing `group_id` for the value with this hash and bytes,
+ /// or inserts a new entry and returns the freshly assigned group_id.
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ match self
+ .unique_dict_value_mapping
+ .entry(hash, |e| e.bytes == raw, |e| e.hash)
+ {
+ HashTableEntry::Occupied(o) => o.get().group_id,
+ HashTableEntry::Vacant(v) => {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(Some(raw.to_vec())); // replace this with raw buffer to avoid this double copy #TODO see https://github.com/apache/datafusion/issues/22078
+ v.insert(DictEntry {
+ hash,
+ group_id: new_group_id,
+ bytes: raw.to_vec(),
+ });
+ new_group_id
+ }
+ }
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(None);
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+ fn transform_into_array(&self, raw: &[Option<Vec<u8>>]) -> Result<ArrayRef> {
+ let item_capacity = raw.len();
+ let data_capacity = item_capacity * AVG_BYTES_PER_DICT_VALUE;
+ match &self.value_dt {
+ DataType::Utf8 => decode_scalar_string!(
+ raw,
+ StringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::LargeUtf8 => decode_scalar_string!(
+ raw,
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::Utf8View => decode_scalar_string!(
+ raw,
+ StringViewBuilder::with_capacity(item_capacity)
+ ),
+ DataType::List(field) => match field.data_type() {
+ DataType::Utf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::LargeUtf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::Utf8View => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringViewBuilder::with_capacity(item_capacity),
+ item_capacity,
+ )
+ ),
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for List<{other:?}>"
+ ))),
+ },
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for {other:?}"
+ ))),
+ }
+ }
+ #[inline]
+ fn keys_to_usize(key_array: &PrimitiveArray<K>) -> Vec<Option<usize>> {
+ (0..key_array.len())
+ .map(|i| {
+ if key_array.is_null(i) {
+ None
+ } else {
+ Some(key_array.value(i).to_usize().unwrap())
+ }
+ })
+ .collect()
+ }
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K> {
+ fn size(&self) -> usize {
+ let seen_elements_size = self.seen_elements.capacity()
+ * size_of::<Option<Vec<u8>>>()
+ + self
+ .seen_elements
+ .iter()
+ .filter_map(|opt| opt.as_ref())
+ .map(|inner| inner.capacity())
+ .sum::<usize>();
+
+ let unique_mapping_size = self.unique_dict_value_mapping.capacity()
+ * size_of::<DictEntry>()
+ + self
+ .unique_dict_value_mapping
+ .iter()
+ .map(|e| e.bytes.capacity())
+ .sum::<usize>();
+
+ let values_cache_size = self.values_cache.1.capacity() * size_of::<u64>()
+ + self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|a| a.get_array_memory_size())
+ .unwrap_or(0);
+
+ size_of::<Self>() + seen_elements_size + unique_mapping_size + values_cache_size
+ }
+ fn len(&self) -> usize {
+ self.seen_elements.len()
+ }
+ fn is_empty(&self) -> bool {
+ self.seen_elements.is_empty()
+ }
+ fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
+ assert_eq!(
+ cols.len(),
+ 1,
+ "GroupValuesDictionary only supports a single column"
+ );
+ let array = Arc::clone(&cols[0]);
+ groups.clear(); // zero out buffer
+ let dict_array = array
+ .as_any()
+ .downcast_ref::<DictionaryArray<K>>()
+ .ok_or_else(|| {
+ Internal(format!(
+ "GroupValuesDictionary expected DictionaryArray but got {:?}",
+ array.data_type()
+ ))
+ })?;
+
+ let values = dict_array.values();
+ let key_array = dict_array.keys();
+ if key_array.is_empty() {
+ return Ok(());
+ }
+ let keys_as_usize: Vec<Option<usize>> = Self::keys_to_usize(key_array);
Review Comment:
maybe this allocation can be removed and we can iterate the keys on the fly?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on code in PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#discussion_r3365098692
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,783 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, DictionaryArray, LargeStringArray, LargeStringBuilder, ListArray,
+ ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+/// Entry stored inside `unique_dict_value_mapping`. Caching `hash` on the entry
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ bytes: Vec<u8>,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ // stores the order new unique elements are seen for self.emit()
+ seen_elements: Vec<Option<Vec<u8>>>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ intern_called: bool,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ seen_elements: Vec::new(),
+ unique_dict_value_mapping: HashTable::with_capacity(INITIAL_PRE_ALLOCATION),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ intern_called: false,
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ /// Returns the existing `group_id` for the value with this hash and bytes,
+ /// or inserts a new entry and returns the freshly assigned group_id.
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ match self
+ .unique_dict_value_mapping
+ .entry(hash, |e| e.bytes == raw, |e| e.hash)
+ {
+ HashTableEntry::Occupied(o) => o.get().group_id,
+ HashTableEntry::Vacant(v) => {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(Some(raw.to_vec())); // replace this with raw buffer to avoid this double copy #TODO see https://github.com/apache/datafusion/issues/22078
+ v.insert(DictEntry {
+ hash,
+ group_id: new_group_id,
+ bytes: raw.to_vec(),
+ });
+ new_group_id
+ }
+ }
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(None);
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+ fn transform_into_array(&self, raw: &[Option<Vec<u8>>]) -> Result<ArrayRef> {
+ let item_capacity = raw.len();
+ let data_capacity = item_capacity * AVG_BYTES_PER_DICT_VALUE;
+ match &self.value_dt {
+ DataType::Utf8 => decode_scalar_string!(
+ raw,
+ StringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::LargeUtf8 => decode_scalar_string!(
+ raw,
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::Utf8View => decode_scalar_string!(
+ raw,
+ StringViewBuilder::with_capacity(item_capacity)
+ ),
+ DataType::List(field) => match field.data_type() {
+ DataType::Utf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::LargeUtf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::Utf8View => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringViewBuilder::with_capacity(item_capacity),
+ item_capacity,
+ )
+ ),
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for List<{other:?}>"
+ ))),
+ },
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for {other:?}"
+ ))),
+ }
+ }
+ #[inline]
+ fn keys_to_usize(key_array: &PrimitiveArray<K>) -> Vec<Option<usize>> {
+ (0..key_array.len())
+ .map(|i| {
+ if key_array.is_null(i) {
+ None
+ } else {
+ Some(key_array.value(i).to_usize().unwrap())
+ }
+ })
+ .collect()
+ }
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K> {
+ fn size(&self) -> usize {
+ let seen_elements_size = self.seen_elements.capacity()
+ * size_of::<Option<Vec<u8>>>()
+ + self
+ .seen_elements
+ .iter()
+ .filter_map(|opt| opt.as_ref())
+ .map(|inner| inner.capacity())
+ .sum::<usize>();
+
+ let unique_mapping_size = self.unique_dict_value_mapping.capacity()
+ * size_of::<DictEntry>()
+ + self
+ .unique_dict_value_mapping
+ .iter()
+ .map(|e| e.bytes.capacity())
+ .sum::<usize>();
+
+ let values_cache_size = self.values_cache.1.capacity() * size_of::<u64>()
+ + self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|a| a.get_array_memory_size())
+ .unwrap_or(0);
+
+ size_of::<Self>() + seen_elements_size + unique_mapping_size + values_cache_size
+ }
+ fn len(&self) -> usize {
+ self.seen_elements.len()
+ }
+ fn is_empty(&self) -> bool {
+ self.seen_elements.is_empty()
+ }
+ fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
+ assert_eq!(
+ cols.len(),
+ 1,
+ "GroupValuesDictionary only supports a single column"
+ );
+ let array = Arc::clone(&cols[0]);
+ groups.clear(); // zero out buffer
+ let dict_array = array
+ .as_any()
+ .downcast_ref::<DictionaryArray<K>>()
+ .ok_or_else(|| {
+ Internal(format!(
+ "GroupValuesDictionary expected DictionaryArray but got {:?}",
+ array.data_type()
+ ))
+ })?;
+
+ let values = dict_array.values();
+ let key_array = dict_array.keys();
+ if key_array.is_empty() {
+ return Ok(());
+ }
+ let keys_as_usize: Vec<Option<usize>> = Self::keys_to_usize(key_array);
+
+ let cache_hit = self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|cached| Arc::ptr_eq(cached, values))
+ .unwrap_or(false);
+
+ if !cache_hit {
+ // first time seeing this values array to compute hashes, and cache the pointer, cheap operation
+ self.compute_value_hashes(values)?;
+ self.values_cache.0 = Some(Arc::clone(values));
+ }
+ let mut key_to_group: Vec<Option<usize>> = vec![None; values.len()];
+ if self.intern_called {
Review Comment:
for multiple calls to intern this is a big performance boost.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4676296593
This looks very cool! I have a background question: how do these dictionary arrays enter the aggregation operator?
Do we have a way to read dictionary arrays directly from Parquet, or to preserve them through `FilterExec`? I might have missed some recent progress here, so I’d appreciate any pointers.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "LiaCastaneda (via GitHub)" <gi...@apache.org>.
LiaCastaneda commented on code in PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#discussion_r3356276108
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,783 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, DictionaryArray, LargeStringArray, LargeStringBuilder, ListArray,
+ ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+/// Entry stored inside `unique_dict_value_mapping`. Caching `hash` on the entry
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ bytes: Vec<u8>,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ // stores the order new unique elements are seen for self.emit()
+ seen_elements: Vec<Option<Vec<u8>>>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ intern_called: bool,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ seen_elements: Vec::new(),
+ unique_dict_value_mapping: HashTable::with_capacity(INITIAL_PRE_ALLOCATION),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ intern_called: false,
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ /// Returns the existing `group_id` for the value with this hash and bytes,
+ /// or inserts a new entry and returns the freshly assigned group_id.
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ match self
+ .unique_dict_value_mapping
+ .entry(hash, |e| e.bytes == raw, |e| e.hash)
+ {
+ HashTableEntry::Occupied(o) => o.get().group_id,
+ HashTableEntry::Vacant(v) => {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(Some(raw.to_vec())); // replace this with raw buffer to avoid this double copy #TODO see https://github.com/apache/datafusion/issues/22078
+ v.insert(DictEntry {
+ hash,
+ group_id: new_group_id,
+ bytes: raw.to_vec(),
+ });
+ new_group_id
+ }
+ }
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(None);
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+ fn transform_into_array(&self, raw: &[Option<Vec<u8>>]) -> Result<ArrayRef> {
+ let item_capacity = raw.len();
+ let data_capacity = item_capacity * AVG_BYTES_PER_DICT_VALUE;
+ match &self.value_dt {
+ DataType::Utf8 => decode_scalar_string!(
+ raw,
+ StringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::LargeUtf8 => decode_scalar_string!(
+ raw,
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::Utf8View => decode_scalar_string!(
+ raw,
+ StringViewBuilder::with_capacity(item_capacity)
+ ),
+ DataType::List(field) => match field.data_type() {
+ DataType::Utf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::LargeUtf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::Utf8View => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringViewBuilder::with_capacity(item_capacity),
+ item_capacity,
+ )
+ ),
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for List<{other:?}>"
+ ))),
+ },
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for {other:?}"
+ ))),
+ }
+ }
+ #[inline]
+ fn keys_to_usize(key_array: &PrimitiveArray<K>) -> Vec<Option<usize>> {
+ (0..key_array.len())
+ .map(|i| {
+ if key_array.is_null(i) {
+ None
+ } else {
+ Some(key_array.value(i).to_usize().unwrap())
+ }
+ })
+ .collect()
+ }
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K> {
+ fn size(&self) -> usize {
+ let seen_elements_size = self.seen_elements.capacity()
+ * size_of::<Option<Vec<u8>>>()
+ + self
+ .seen_elements
+ .iter()
+ .filter_map(|opt| opt.as_ref())
+ .map(|inner| inner.capacity())
+ .sum::<usize>();
+
+ let unique_mapping_size = self.unique_dict_value_mapping.capacity()
+ * size_of::<DictEntry>()
+ + self
+ .unique_dict_value_mapping
+ .iter()
+ .map(|e| e.bytes.capacity())
+ .sum::<usize>();
+
+ let values_cache_size = self.values_cache.1.capacity() * size_of::<u64>()
+ + self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|a| a.get_array_memory_size())
+ .unwrap_or(0);
+
+ size_of::<Self>() + seen_elements_size + unique_mapping_size + values_cache_size
+ }
+ fn len(&self) -> usize {
+ self.seen_elements.len()
+ }
+ fn is_empty(&self) -> bool {
+ self.seen_elements.is_empty()
+ }
+ fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
+ assert_eq!(
+ cols.len(),
+ 1,
+ "GroupValuesDictionary only supports a single column"
+ );
+ let array = Arc::clone(&cols[0]);
+ groups.clear(); // zero out buffer
+ let dict_array = array
+ .as_any()
+ .downcast_ref::<DictionaryArray<K>>()
+ .ok_or_else(|| {
+ Internal(format!(
+ "GroupValuesDictionary expected DictionaryArray but got {:?}",
+ array.data_type()
+ ))
+ })?;
+
+ let values = dict_array.values();
+ let key_array = dict_array.keys();
+ if key_array.is_empty() {
+ return Ok(());
+ }
+ let keys_as_usize: Vec<Option<usize>> = Self::keys_to_usize(key_array);
Review Comment:
maybe this allocation can be removed and we can iterate and convert the keys on the fly?
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,783 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, DictionaryArray, LargeStringArray, LargeStringBuilder, ListArray,
+ ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+/// Entry stored inside `unique_dict_value_mapping`. Caching `hash` on the entry
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ bytes: Vec<u8>,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ // stores the order new unique elements are seen for self.emit()
+ seen_elements: Vec<Option<Vec<u8>>>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ intern_called: bool,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ seen_elements: Vec::new(),
+ unique_dict_value_mapping: HashTable::with_capacity(INITIAL_PRE_ALLOCATION),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ intern_called: false,
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ /// Returns the existing `group_id` for the value with this hash and bytes,
+ /// or inserts a new entry and returns the freshly assigned group_id.
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ match self
+ .unique_dict_value_mapping
+ .entry(hash, |e| e.bytes == raw, |e| e.hash)
+ {
+ HashTableEntry::Occupied(o) => o.get().group_id,
+ HashTableEntry::Vacant(v) => {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(Some(raw.to_vec())); // replace this with raw buffer to avoid this double copy #TODO see https://github.com/apache/datafusion/issues/22078
+ v.insert(DictEntry {
+ hash,
+ group_id: new_group_id,
+ bytes: raw.to_vec(),
+ });
+ new_group_id
+ }
+ }
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(None);
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+ fn transform_into_array(&self, raw: &[Option<Vec<u8>>]) -> Result<ArrayRef> {
+ let item_capacity = raw.len();
+ let data_capacity = item_capacity * AVG_BYTES_PER_DICT_VALUE;
+ match &self.value_dt {
+ DataType::Utf8 => decode_scalar_string!(
+ raw,
+ StringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::LargeUtf8 => decode_scalar_string!(
+ raw,
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::Utf8View => decode_scalar_string!(
+ raw,
+ StringViewBuilder::with_capacity(item_capacity)
+ ),
+ DataType::List(field) => match field.data_type() {
+ DataType::Utf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::LargeUtf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::Utf8View => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringViewBuilder::with_capacity(item_capacity),
+ item_capacity,
+ )
+ ),
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for List<{other:?}>"
+ ))),
+ },
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for {other:?}"
+ ))),
+ }
+ }
+ #[inline]
+ fn keys_to_usize(key_array: &PrimitiveArray<K>) -> Vec<Option<usize>> {
+ (0..key_array.len())
+ .map(|i| {
+ if key_array.is_null(i) {
+ None
+ } else {
+ Some(key_array.value(i).to_usize().unwrap())
+ }
+ })
+ .collect()
+ }
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K> {
+ fn size(&self) -> usize {
+ let seen_elements_size = self.seen_elements.capacity()
+ * size_of::<Option<Vec<u8>>>()
+ + self
+ .seen_elements
+ .iter()
+ .filter_map(|opt| opt.as_ref())
+ .map(|inner| inner.capacity())
+ .sum::<usize>();
+
+ let unique_mapping_size = self.unique_dict_value_mapping.capacity()
+ * size_of::<DictEntry>()
+ + self
+ .unique_dict_value_mapping
+ .iter()
+ .map(|e| e.bytes.capacity())
+ .sum::<usize>();
+
+ let values_cache_size = self.values_cache.1.capacity() * size_of::<u64>()
+ + self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|a| a.get_array_memory_size())
+ .unwrap_or(0);
+
+ size_of::<Self>() + seen_elements_size + unique_mapping_size + values_cache_size
+ }
+ fn len(&self) -> usize {
+ self.seen_elements.len()
+ }
+ fn is_empty(&self) -> bool {
+ self.seen_elements.is_empty()
+ }
+ fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
+ assert_eq!(
+ cols.len(),
+ 1,
+ "GroupValuesDictionary only supports a single column"
+ );
+ let array = Arc::clone(&cols[0]);
+ groups.clear(); // zero out buffer
+ let dict_array = array
+ .as_any()
+ .downcast_ref::<DictionaryArray<K>>()
+ .ok_or_else(|| {
+ Internal(format!(
+ "GroupValuesDictionary expected DictionaryArray but got {:?}",
+ array.data_type()
+ ))
+ })?;
Review Comment:
```suggestion
array.as_dictionary();
```
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,783 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, DictionaryArray, LargeStringArray, LargeStringBuilder, ListArray,
+ ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+/// Entry stored inside `unique_dict_value_mapping`. Caching `hash` on the entry
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ bytes: Vec<u8>,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ // stores the order new unique elements are seen for self.emit()
+ seen_elements: Vec<Option<Vec<u8>>>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ intern_called: bool,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ seen_elements: Vec::new(),
+ unique_dict_value_mapping: HashTable::with_capacity(INITIAL_PRE_ALLOCATION),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ intern_called: false,
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ /// Returns the existing `group_id` for the value with this hash and bytes,
+ /// or inserts a new entry and returns the freshly assigned group_id.
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ match self
+ .unique_dict_value_mapping
+ .entry(hash, |e| e.bytes == raw, |e| e.hash)
+ {
+ HashTableEntry::Occupied(o) => o.get().group_id,
+ HashTableEntry::Vacant(v) => {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(Some(raw.to_vec())); // replace this with raw buffer to avoid this double copy #TODO see https://github.com/apache/datafusion/issues/22078
+ v.insert(DictEntry {
+ hash,
+ group_id: new_group_id,
+ bytes: raw.to_vec(),
+ });
+ new_group_id
+ }
+ }
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(None);
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+ fn transform_into_array(&self, raw: &[Option<Vec<u8>>]) -> Result<ArrayRef> {
+ let item_capacity = raw.len();
+ let data_capacity = item_capacity * AVG_BYTES_PER_DICT_VALUE;
+ match &self.value_dt {
+ DataType::Utf8 => decode_scalar_string!(
+ raw,
+ StringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::LargeUtf8 => decode_scalar_string!(
+ raw,
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::Utf8View => decode_scalar_string!(
+ raw,
+ StringViewBuilder::with_capacity(item_capacity)
+ ),
+ DataType::List(field) => match field.data_type() {
+ DataType::Utf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::LargeUtf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::Utf8View => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringViewBuilder::with_capacity(item_capacity),
+ item_capacity,
+ )
+ ),
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for List<{other:?}>"
+ ))),
+ },
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for {other:?}"
+ ))),
+ }
+ }
+ #[inline]
+ fn keys_to_usize(key_array: &PrimitiveArray<K>) -> Vec<Option<usize>> {
+ (0..key_array.len())
+ .map(|i| {
+ if key_array.is_null(i) {
+ None
+ } else {
+ Some(key_array.value(i).to_usize().unwrap())
+ }
+ })
+ .collect()
+ }
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K> {
+ fn size(&self) -> usize {
+ let seen_elements_size = self.seen_elements.capacity()
+ * size_of::<Option<Vec<u8>>>()
+ + self
+ .seen_elements
+ .iter()
+ .filter_map(|opt| opt.as_ref())
+ .map(|inner| inner.capacity())
+ .sum::<usize>();
+
+ let unique_mapping_size = self.unique_dict_value_mapping.capacity()
+ * size_of::<DictEntry>()
+ + self
+ .unique_dict_value_mapping
+ .iter()
+ .map(|e| e.bytes.capacity())
+ .sum::<usize>();
+
+ let values_cache_size = self.values_cache.1.capacity() * size_of::<u64>()
+ + self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|a| a.get_array_memory_size())
+ .unwrap_or(0);
+
+ size_of::<Self>() + seen_elements_size + unique_mapping_size + values_cache_size
+ }
+ fn len(&self) -> usize {
+ self.seen_elements.len()
+ }
+ fn is_empty(&self) -> bool {
+ self.seen_elements.is_empty()
+ }
+ fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
+ assert_eq!(
+ cols.len(),
+ 1,
+ "GroupValuesDictionary only supports a single column"
+ );
+ let array = Arc::clone(&cols[0]);
+ groups.clear(); // zero out buffer
+ let dict_array = array
+ .as_any()
+ .downcast_ref::<DictionaryArray<K>>()
+ .ok_or_else(|| {
+ Internal(format!(
+ "GroupValuesDictionary expected DictionaryArray but got {:?}",
+ array.data_type()
+ ))
+ })?;
+
+ let values = dict_array.values();
+ let key_array = dict_array.keys();
+ if key_array.is_empty() {
+ return Ok(());
+ }
+ let keys_as_usize: Vec<Option<usize>> = Self::keys_to_usize(key_array);
+
+ let cache_hit = self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|cached| Arc::ptr_eq(cached, values))
+ .unwrap_or(false);
+
+ if !cache_hit {
+ // first time seeing this values array to compute hashes, and cache the pointer, cheap operation
Review Comment:
wdyt about rephrasing this to something like
```suggestion
// values array changed since last batch - recompute hashes and update cached pointer
```
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,783 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, DictionaryArray, LargeStringArray, LargeStringBuilder, ListArray,
+ ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+/// Entry stored inside `unique_dict_value_mapping`. Caching `hash` on the entry
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ bytes: Vec<u8>,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ // stores the order new unique elements are seen for self.emit()
+ seen_elements: Vec<Option<Vec<u8>>>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ intern_called: bool,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ seen_elements: Vec::new(),
+ unique_dict_value_mapping: HashTable::with_capacity(INITIAL_PRE_ALLOCATION),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ intern_called: false,
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ /// Returns the existing `group_id` for the value with this hash and bytes,
+ /// or inserts a new entry and returns the freshly assigned group_id.
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ match self
+ .unique_dict_value_mapping
+ .entry(hash, |e| e.bytes == raw, |e| e.hash)
+ {
+ HashTableEntry::Occupied(o) => o.get().group_id,
+ HashTableEntry::Vacant(v) => {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(Some(raw.to_vec())); // replace this with raw buffer to avoid this double copy #TODO see https://github.com/apache/datafusion/issues/22078
+ v.insert(DictEntry {
+ hash,
+ group_id: new_group_id,
+ bytes: raw.to_vec(),
+ });
+ new_group_id
+ }
+ }
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(None);
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+ fn transform_into_array(&self, raw: &[Option<Vec<u8>>]) -> Result<ArrayRef> {
+ let item_capacity = raw.len();
+ let data_capacity = item_capacity * AVG_BYTES_PER_DICT_VALUE;
+ match &self.value_dt {
+ DataType::Utf8 => decode_scalar_string!(
+ raw,
+ StringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::LargeUtf8 => decode_scalar_string!(
+ raw,
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::Utf8View => decode_scalar_string!(
+ raw,
+ StringViewBuilder::with_capacity(item_capacity)
+ ),
+ DataType::List(field) => match field.data_type() {
+ DataType::Utf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::LargeUtf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::Utf8View => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringViewBuilder::with_capacity(item_capacity),
+ item_capacity,
+ )
+ ),
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for List<{other:?}>"
+ ))),
+ },
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for {other:?}"
+ ))),
+ }
+ }
+ #[inline]
+ fn keys_to_usize(key_array: &PrimitiveArray<K>) -> Vec<Option<usize>> {
+ (0..key_array.len())
+ .map(|i| {
+ if key_array.is_null(i) {
+ None
+ } else {
+ Some(key_array.value(i).to_usize().unwrap())
+ }
+ })
+ .collect()
+ }
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K> {
+ fn size(&self) -> usize {
+ let seen_elements_size = self.seen_elements.capacity()
+ * size_of::<Option<Vec<u8>>>()
+ + self
+ .seen_elements
+ .iter()
+ .filter_map(|opt| opt.as_ref())
+ .map(|inner| inner.capacity())
+ .sum::<usize>();
+
+ let unique_mapping_size = self.unique_dict_value_mapping.capacity()
+ * size_of::<DictEntry>()
+ + self
+ .unique_dict_value_mapping
+ .iter()
+ .map(|e| e.bytes.capacity())
+ .sum::<usize>();
+
+ let values_cache_size = self.values_cache.1.capacity() * size_of::<u64>()
+ + self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|a| a.get_array_memory_size())
+ .unwrap_or(0);
+
+ size_of::<Self>() + seen_elements_size + unique_mapping_size + values_cache_size
+ }
+ fn len(&self) -> usize {
+ self.seen_elements.len()
+ }
+ fn is_empty(&self) -> bool {
+ self.seen_elements.is_empty()
+ }
+ fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
+ assert_eq!(
+ cols.len(),
+ 1,
+ "GroupValuesDictionary only supports a single column"
+ );
+ let array = Arc::clone(&cols[0]);
+ groups.clear(); // zero out buffer
+ let dict_array = array
+ .as_any()
+ .downcast_ref::<DictionaryArray<K>>()
+ .ok_or_else(|| {
+ Internal(format!(
+ "GroupValuesDictionary expected DictionaryArray but got {:?}",
+ array.data_type()
+ ))
+ })?;
+
+ let values = dict_array.values();
+ let key_array = dict_array.keys();
+ if key_array.is_empty() {
+ return Ok(());
+ }
+ let keys_as_usize: Vec<Option<usize>> = Self::keys_to_usize(key_array);
+
+ let cache_hit = self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|cached| Arc::ptr_eq(cached, values))
+ .unwrap_or(false);
+
+ if !cache_hit {
+ // first time seeing this values array to compute hashes, and cache the pointer, cheap operation
+ self.compute_value_hashes(values)?;
+ self.values_cache.0 = Some(Arc::clone(values));
+ }
+ let mut key_to_group: Vec<Option<usize>> = vec![None; values.len()];
+ if self.intern_called {
+ for key_opt in keys_as_usize.iter() {
+ let Some(original_key) = *key_opt else {
+ continue;
+ };
+ if values.is_null(original_key) {
+ continue;
+ }
+ if key_to_group[original_key].is_some() {
+ continue;
+ }
+
+ let hash = self.values_cache.1[original_key];
+ let raw = Self::get_raw_bytes(values, original_key);
+ if let Some(entry) = self
+ .unique_dict_value_mapping
+ .find(hash, |e| e.bytes == raw.as_ref())
+ {
+ key_to_group[original_key] = Some(entry.group_id);
+ continue;
+ }
+ }
+ }
+ // iterate keys array (n iterations)
+ // only d insertions at most, repeated work is cached
+ for key_opt in keys_as_usize.iter() {
Review Comment:
This code is quite algorithmic, especially for people who won't read the PR description. I think it would be useful to add some documentation since there are too many nested branches without explanation -- wdyt about refactoring this loop with some inline comments?
```
for i in 0..key_array.len() {
// null key -> null row
if key_array.is_null(i) {
groups.push(self.get_null_group_id());
continue;
}
let key = key_array.value(i).to_usize().unwrap();
// already resolved this dictionary index earlier in this batch
if let Some(group_id) = key_to_group[key] {
groups.push(group_id);
continue;
}
// non-null key pointing to a null value slot — treat as null
if values.is_null(key) {
let gid = self.get_null_group_id();
key_to_group[key] = Some(gid);
groups.push(gid);
continue;
}
// new value — insert into hash table and assign a new group id
let hash = self.values_cache.1[key];
let raw = Self::get_raw_bytes(values, key);
let gid = self.lookup_or_insert_in_table(hash, raw.as_ref());
key_to_group[key] = Some(gid);
groups.push(gid);
}
```
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,783 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, DictionaryArray, LargeStringArray, LargeStringBuilder, ListArray,
+ ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+/// Entry stored inside `unique_dict_value_mapping`. Caching `hash` on the entry
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ bytes: Vec<u8>,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ // stores the order new unique elements are seen for self.emit()
+ seen_elements: Vec<Option<Vec<u8>>>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ intern_called: bool,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ seen_elements: Vec::new(),
+ unique_dict_value_mapping: HashTable::with_capacity(INITIAL_PRE_ALLOCATION),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ intern_called: false,
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ /// Returns the existing `group_id` for the value with this hash and bytes,
+ /// or inserts a new entry and returns the freshly assigned group_id.
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ match self
+ .unique_dict_value_mapping
+ .entry(hash, |e| e.bytes == raw, |e| e.hash)
+ {
+ HashTableEntry::Occupied(o) => o.get().group_id,
+ HashTableEntry::Vacant(v) => {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(Some(raw.to_vec())); // replace this with raw buffer to avoid this double copy #TODO see https://github.com/apache/datafusion/issues/22078
+ v.insert(DictEntry {
+ hash,
+ group_id: new_group_id,
+ bytes: raw.to_vec(),
+ });
+ new_group_id
+ }
+ }
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(None);
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+ fn transform_into_array(&self, raw: &[Option<Vec<u8>>]) -> Result<ArrayRef> {
+ let item_capacity = raw.len();
+ let data_capacity = item_capacity * AVG_BYTES_PER_DICT_VALUE;
+ match &self.value_dt {
+ DataType::Utf8 => decode_scalar_string!(
+ raw,
+ StringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::LargeUtf8 => decode_scalar_string!(
+ raw,
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::Utf8View => decode_scalar_string!(
+ raw,
+ StringViewBuilder::with_capacity(item_capacity)
+ ),
+ DataType::List(field) => match field.data_type() {
+ DataType::Utf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::LargeUtf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::Utf8View => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringViewBuilder::with_capacity(item_capacity),
+ item_capacity,
+ )
+ ),
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for List<{other:?}>"
+ ))),
+ },
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for {other:?}"
+ ))),
+ }
+ }
+ #[inline]
+ fn keys_to_usize(key_array: &PrimitiveArray<K>) -> Vec<Option<usize>> {
+ (0..key_array.len())
+ .map(|i| {
+ if key_array.is_null(i) {
+ None
+ } else {
+ Some(key_array.value(i).to_usize().unwrap())
+ }
+ })
+ .collect()
+ }
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K> {
+ fn size(&self) -> usize {
+ let seen_elements_size = self.seen_elements.capacity()
+ * size_of::<Option<Vec<u8>>>()
+ + self
+ .seen_elements
+ .iter()
+ .filter_map(|opt| opt.as_ref())
+ .map(|inner| inner.capacity())
+ .sum::<usize>();
+
+ let unique_mapping_size = self.unique_dict_value_mapping.capacity()
+ * size_of::<DictEntry>()
+ + self
+ .unique_dict_value_mapping
+ .iter()
+ .map(|e| e.bytes.capacity())
+ .sum::<usize>();
+
+ let values_cache_size = self.values_cache.1.capacity() * size_of::<u64>()
+ + self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|a| a.get_array_memory_size())
Review Comment:
`get_array_memory_size()` over counts if it has underlying buffers that are shared with another array ( the `RecordBatch` that owns this values array may already be tracked by the memory pool). This could cause the memory pool to falsely exceeded limits when it hasn't actually been exceeded. Although in practice I don't think this is very probable. Maybe it's worth using a more acurate method like `a.to_data().get_array_memory_size()` ([example](https://github.com/apache/datafusion/blob/b7761dc59155f04cd7d4b45eb7c5d4a1c3587d32/datafusion/functions-aggregate/src/array_agg.rs#L472))
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,783 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, DictionaryArray, LargeStringArray, LargeStringBuilder, ListArray,
+ ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+/// Entry stored inside `unique_dict_value_mapping`. Caching `hash` on the entry
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ bytes: Vec<u8>,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ // stores the order new unique elements are seen for self.emit()
+ seen_elements: Vec<Option<Vec<u8>>>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ intern_called: bool,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ seen_elements: Vec::new(),
+ unique_dict_value_mapping: HashTable::with_capacity(INITIAL_PRE_ALLOCATION),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ intern_called: false,
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ /// Returns the existing `group_id` for the value with this hash and bytes,
+ /// or inserts a new entry and returns the freshly assigned group_id.
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ match self
+ .unique_dict_value_mapping
+ .entry(hash, |e| e.bytes == raw, |e| e.hash)
+ {
+ HashTableEntry::Occupied(o) => o.get().group_id,
+ HashTableEntry::Vacant(v) => {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(Some(raw.to_vec())); // replace this with raw buffer to avoid this double copy #TODO see https://github.com/apache/datafusion/issues/22078
+ v.insert(DictEntry {
+ hash,
+ group_id: new_group_id,
+ bytes: raw.to_vec(),
+ });
+ new_group_id
+ }
+ }
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(None);
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+ fn transform_into_array(&self, raw: &[Option<Vec<u8>>]) -> Result<ArrayRef> {
+ let item_capacity = raw.len();
+ let data_capacity = item_capacity * AVG_BYTES_PER_DICT_VALUE;
+ match &self.value_dt {
+ DataType::Utf8 => decode_scalar_string!(
+ raw,
+ StringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::LargeUtf8 => decode_scalar_string!(
+ raw,
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::Utf8View => decode_scalar_string!(
+ raw,
+ StringViewBuilder::with_capacity(item_capacity)
+ ),
+ DataType::List(field) => match field.data_type() {
+ DataType::Utf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::LargeUtf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::Utf8View => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringViewBuilder::with_capacity(item_capacity),
+ item_capacity,
+ )
+ ),
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for List<{other:?}>"
+ ))),
+ },
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for {other:?}"
+ ))),
+ }
+ }
+ #[inline]
+ fn keys_to_usize(key_array: &PrimitiveArray<K>) -> Vec<Option<usize>> {
+ (0..key_array.len())
+ .map(|i| {
+ if key_array.is_null(i) {
+ None
+ } else {
+ Some(key_array.value(i).to_usize().unwrap())
+ }
+ })
+ .collect()
+ }
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K> {
+ fn size(&self) -> usize {
+ let seen_elements_size = self.seen_elements.capacity()
+ * size_of::<Option<Vec<u8>>>()
+ + self
+ .seen_elements
+ .iter()
+ .filter_map(|opt| opt.as_ref())
+ .map(|inner| inner.capacity())
+ .sum::<usize>();
+
+ let unique_mapping_size = self.unique_dict_value_mapping.capacity()
+ * size_of::<DictEntry>()
+ + self
+ .unique_dict_value_mapping
+ .iter()
+ .map(|e| e.bytes.capacity())
+ .sum::<usize>();
+
+ let values_cache_size = self.values_cache.1.capacity() * size_of::<u64>()
+ + self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|a| a.get_array_memory_size())
+ .unwrap_or(0);
+
+ size_of::<Self>() + seen_elements_size + unique_mapping_size + values_cache_size
+ }
+ fn len(&self) -> usize {
+ self.seen_elements.len()
+ }
+ fn is_empty(&self) -> bool {
+ self.seen_elements.is_empty()
+ }
+ fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
+ assert_eq!(
+ cols.len(),
+ 1,
+ "GroupValuesDictionary only supports a single column"
+ );
+ let array = Arc::clone(&cols[0]);
+ groups.clear(); // zero out buffer
+ let dict_array = array
+ .as_any()
+ .downcast_ref::<DictionaryArray<K>>()
+ .ok_or_else(|| {
+ Internal(format!(
+ "GroupValuesDictionary expected DictionaryArray but got {:?}",
+ array.data_type()
+ ))
+ })?;
+
+ let values = dict_array.values();
+ let key_array = dict_array.keys();
+ if key_array.is_empty() {
+ return Ok(());
+ }
+ let keys_as_usize: Vec<Option<usize>> = Self::keys_to_usize(key_array);
+
+ let cache_hit = self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|cached| Arc::ptr_eq(cached, values))
+ .unwrap_or(false);
+
+ if !cache_hit {
+ // first time seeing this values array to compute hashes, and cache the pointer, cheap operation
+ self.compute_value_hashes(values)?;
+ self.values_cache.0 = Some(Arc::clone(values));
+ }
+ let mut key_to_group: Vec<Option<usize>> = vec![None; values.len()];
+ if self.intern_called {
Review Comment:
Can we add a small comment on why we have this branch?
```
// For values already seen in prior batches, pre fill `key_to_group`
// so the main loop below can skip hash table lookups for those keys.
// This pre scan is skipped on the first batch since the hash table is empty.
```
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "LiaCastaneda (via GitHub)" <gi...@apache.org>.
LiaCastaneda commented on code in PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#discussion_r3371441053
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,878 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, LargeStringArray, LargeStringBuilder,
+ ListArray, ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+/// Entry stored inside `unique_dict_value_mapping`. Caching `hash` on the entry
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ bytes: Vec<u8>,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ // stores the order new unique elements are seen for self.emit()
+ seen_elements: Vec<Option<Vec<u8>>>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ // key to group vector scratch space, used to avoid re-allocating a new vector on each call to intern
+ key_to_group: Vec<Option<usize>>,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ seen_elements: Vec::new(),
+ unique_dict_value_mapping: HashTable::with_capacity(INITIAL_PRE_ALLOCATION),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ key_to_group: Vec::new(),
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ /// Returns the existing `group_id` for the value with this hash and bytes,
+ /// or inserts a new entry and returns the freshly assigned group_id.
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ match self
+ .unique_dict_value_mapping
+ .entry(hash, |e| e.bytes == raw, |e| e.hash)
+ {
+ HashTableEntry::Occupied(o) => o.get().group_id,
+ HashTableEntry::Vacant(v) => {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(Some(raw.to_vec())); // replace this with raw buffer to avoid this double copy #TODO see https://github.com/apache/datafusion/issues/22078
+ v.insert(DictEntry {
+ hash,
+ group_id: new_group_id,
+ bytes: raw.to_vec(),
+ });
+ new_group_id
+ }
+ }
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(None);
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+ fn transform_into_array(&self, raw: &[Option<Vec<u8>>]) -> Result<ArrayRef> {
+ let item_capacity = raw.len();
+ let data_capacity = item_capacity * AVG_BYTES_PER_DICT_VALUE;
+ match &self.value_dt {
+ DataType::Utf8 => decode_scalar_string!(
+ raw,
+ StringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::LargeUtf8 => decode_scalar_string!(
+ raw,
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::Utf8View => decode_scalar_string!(
+ raw,
+ StringViewBuilder::with_capacity(item_capacity)
+ ),
+ DataType::List(field) => match field.data_type() {
+ DataType::Utf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ .with_field(Arc::clone(field))
+ ),
+ DataType::LargeUtf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ .with_field(Arc::clone(field))
+ ),
+ DataType::Utf8View => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringViewBuilder::with_capacity(item_capacity),
+ item_capacity,
+ )
+ .with_field(Arc::clone(field))
+ ),
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for List<{other:?}>"
+ ))),
+ },
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for {other:?}"
+ ))),
+ }
+ }
+}
+
+fn valid_bounds<K: ArrowDictionaryKeyType>(n: usize) -> bool {
+ let max: usize = 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,
+ _ => return false,
+ };
+ n <= max
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K> {
+ fn size(&self) -> usize {
+ let seen_elements_size = self.seen_elements.capacity()
+ * size_of::<Option<Vec<u8>>>()
+ + self
+ .seen_elements
+ .iter()
+ .filter_map(|opt| opt.as_ref())
+ .map(|inner| inner.capacity())
+ .sum::<usize>();
+
+ let unique_mapping_size = self.unique_dict_value_mapping.capacity()
+ * size_of::<DictEntry>()
+ + self
+ .unique_dict_value_mapping
+ .iter()
+ .map(|e| e.bytes.capacity())
+ .sum::<usize>();
+
+ let values_cache_size = self.values_cache.1.capacity() * size_of::<u64>()
+ + self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|a| a.to_data().get_array_memory_size())
Review Comment:
@Rich-T-kid there was a typo in the comment I left before, the function is not `get_array_memory_size()`, it is `get_slice_memory_size()`
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4301680166
> ### tpch_sf10
> <img alt="Image 4-22-26 at 5 02 PM" width="1468" height="551" src="https://private-user-images.githubusercontent.com/137434454/582243706-7e321840-21a8-4000-95d5-e34e8cdc3d82.png?jwt=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3NzY5MTc2MzgsIm5iZiI6MTc3NjkxNzMzOCwicGF0aCI6Ii8xMzc0MzQ0NTQvNTgyMjQzNzA2LTdlMzIxODQwLTIxYTgtNDAwMC05NWQ1LWUzNGU4Y2RjM2Q4Mi5wbmc_WC1BbXotQWxnb3JpdGhtPUFXUzQtSE1BQy1TSEEyNTYmWC1BbXotQ3JlZGVudGlhbD1BS0lBVkNPRFlMU0E1M1BRSzRaQSUyRjIwMjYwNDIzJTJGdXMtZWFzdC0xJTJGczMlMkZhd3M0X3JlcXVlc3QmWC1BbXotRGF0ZT0yMDI2MDQyM1QwNDA4NThaJlgtQW16LUV4cGlyZXM9MzAwJlgtQW16LVNpZ25hdHVyZT0wNDA0MGVmZDAzZjY3MzdmZTNjNjE3Mzc3YWFlOTdiMzNkNmZjODQ1ZjU1NDdjYWZkNTUyNDAzOWJiODc0MTZlJlgtQW16LVNpZ25lZEhlYWRlcnM9aG9zdCZyZXNwb25zZS1jb250ZW50LXR5cGU9aW1hZ2UlMkZwbmcifQ.tQ6Em8f8lB1MNw297pQSBjbb2gKaGjwnxJZSbaGHg-A">
Still does not look correct
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4293114070
### tpch_mem_sf1
<img width="1188" height="670" alt="Image 4-21-26 at 10 18 PM" src="https://github.com/user-attachments/assets/a72394ed-e4e3-4bf3-b65f-6ab6dd30d1fa" />
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4293364792
### SF1
<img width="1325" height="667" alt="Image 4-21-26 at 11 43 PM" src="https://github.com/user-attachments/assets/e496b3f6-d46a-4a3a-8803-54c095d29851" />
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4500478254
I am losing track of all the open PRs. I am going to experiment adding a `review:waiting` label and see if that makes it easier to find PRs waiting for review going forward
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4723347832
> When arrow-rs reads RLE-encoded strings from a Parquet file, it decodes them into dictionary arrays.
I think you need to configure the parquet reader with a schema that specifies Dictionary as the data type when reading such a column from parquet files (we use this at InfluxData). However, I don't think there is way to do this via SQL today in DataFusion, though you can do it programmatically (by providing a schema to the table provider)
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on code in PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#discussion_r3454836051
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,1003 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, LargeStringArray, LargeStringBuilder,
+ ListArray, ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ offset: usize,
+ len: usize,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ /// Packed byte storage for all group values.
+ row_buffer: Vec<u8>,
+ /// `row_offsets[g]` = start of group `g` in `row_buffer`;
+ row_offsets: Vec<usize>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ // key to group vector scratch space, used to avoid re-allocating a new vector on each call to intern
+ key_to_group: Vec<Option<usize>>,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ row_buffer: Vec::new(),
+ row_offsets: Vec::new(),
+ unique_dict_value_mapping: HashTable::new(),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ key_to_group: Vec::new(),
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ let row_buffer = &self.row_buffer;
+ if let Some(e) = self.unique_dict_value_mapping.find(hash, |e| {
+ e.hash == hash && &row_buffer[e.offset..e.offset + e.len] == raw
+ }) {
+ return e.group_id;
+ }
+ let new_group_id = self.row_offsets.len();
+ let offset = self.row_buffer.len();
+ self.row_offsets.push(offset);
+ self.row_buffer.extend_from_slice(raw);
+ self.unique_dict_value_mapping.insert_unique(
+ hash,
+ DictEntry {
+ hash,
+ group_id: new_group_id,
+ offset,
+ len: raw.len(),
+ },
+ |e| e.hash,
+ );
+ new_group_id
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.row_offsets.len();
+ self.row_offsets.push(self.row_buffer.len()); // null: empty span
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+
+ fn transform_into_array(
+ &self,
+ n: usize,
+ null_group_id: Option<usize>,
+ ) -> Result<ArrayRef> {
+ let data_capacity = n * AVG_BYTES_PER_DICT_VALUE;
+ let raw = (0..n).map(|i| {
+ if Some(i) == null_group_id {
+ return None;
+ }
+ let start = self.row_offsets[i];
+ // last group has no i+1 entry; its end is the buffer tail.
+ let end = self
+ .row_offsets
+ .get(i + 1)
+ .copied()
+ .unwrap_or(self.row_buffer.len());
+ Some(&self.row_buffer[start..end])
+ });
+ match &self.value_dt {
Review Comment:
thank you for the review @alamb
I agree on avoiding ballooning LOC, adding support for every data type becomes more trouble than the compute savings justify. That said, I think restricting support to `Dictionary<_, Utf8/Utf8View>` is the right call. It significantly simplifies the code paths while still delivering a large performance boost for the most common production data shapes (low-cardinality group bys), and similar to other single-column group by specializations, it leaves a viable fallback that loses nothing when the fast path isn't hit.
With that in mind, @zhuqi-lucas's work looks very interesting -- I had a similar idea with #22891 and #21878. The issue was that multi-column group-bys on dictionaries fell through to `GroupValuesRows`, which was slower than we'd like. I'm actually working on a multi-column dictionary group by in #22983, but I'm running into similar issues around handling multiple distinct types. For the multi-column case it makes more sense to build on @zhuqi-lucas's work, since the number of types that need to be supported explodes and there are diminishing returns to exploiting dictionary properties across multiple columns and intern calls ([key space cache](https://github.com/apache/datafusion/pull/22983/changes#diff-84ceab5d992f30b340d52094b28685b5ecfd6f68f27867641ca1a9453ad279baR209), [arc ptr cache](https://github.com/apache/datafusion/pull/22983/changes#diff-84ceab5d992f30b340d52094b28685b5ecfd6f68f27867641ca1a9453ad279baR193), [compute keys once and reuse](https://github.com/apache/datafu
sion/pull/22983/changes#diff-84ceab5d992f30b340d52094b28685b5ecfd6f68f27867641ca1a9453ad279baR342)).
As for this PR specifically, I do think it adds significant value. Would restricting the supported value types be a viable path forward? It would keep the type combinations from exploding while still letting the community benefit from the speedup.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4662226826
🤖 Benchmark running (GKE) | [trigger](https://github.com/apache/datafusion/pull/21765#issuecomment-4662199962)
**Instance:** `c4a-highmem-16` (12 vCPU / 65 GiB) | `Linux bench-c4662199962-505-5wtgp 6.12.68+ #1 SMP Sat May 2 07:49:07 UTC 2026 aarch64 GNU/Linux`
<details><summary>CPU Details (lscpu)</summary>
```
Architecture: aarch64
CPU op-mode(s): 64-bit
Byte Order: Little Endian
CPU(s): 16
On-line CPU(s) list: 0-15
Vendor ID: ARM
Model name: Neoverse-V2
Model: 1
Thread(s) per core: 1
Core(s) per cluster: 16
Socket(s): -
Cluster(s): 1
Stepping: r0p1
BogoMIPS: 2000.00
Flags: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache: 1 MiB (16 instances)
L1i cache: 1 MiB (16 instances)
L2 cache: 32 MiB (16 instances)
L3 cache: 80 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; __user pointer sanitization
Vulnerability Spectre v2: Mitigation; CSV2, BHB
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected
```
</details>
Comparing rich-t-kid/optimize-dictionary-grouping (080daa9230eb8c27939b28815fd3364fd989ca05) to 9b81ff8 (merge-base) [diff](https://github.com/apache/datafusion/compare/9b81ff86d1adf35d953329eb05b4e6d3e0904ea8..080daa9230eb8c27939b28815fd3364fd989ca05) using: dictionary_group_values
Results will be posted here when complete
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on code in PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#discussion_r3507233746
##########
datafusion/sqllogictest/test_files/dictionary.slt:
##########
@@ -515,3 +515,15 @@ DROP TABLE dict_hash_10;
statement ok
DROP TABLE dict_hash_src;
+
+query TI
+SELECT
Review Comment:
opened #23280
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4420844268
the tpch_sp1/sp10 and click house benchmarks dont include dictionary encoded columns, so they aren't benchmarking anything thats included in this PR. ( as @Dandandan mentioned [here](https://github.com/apache/datafusion/pull/21765#issuecomment-4301675954))
benchmarks that include dictionary columns exist in `benchmarks/src/dict` which is run with `/benchmarks/bench.sh run dict` & micro benchmarks that exist in `datafusion/physical-plan/benches/dictionary_group_values.rs`. Im not sure what the command is to run these specific benchmarks through @/adriangbot.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "LiaCastaneda (via GitHub)" <gi...@apache.org>.
LiaCastaneda commented on code in PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#discussion_r3371441053
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,878 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, LargeStringArray, LargeStringBuilder,
+ ListArray, ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+/// Entry stored inside `unique_dict_value_mapping`. Caching `hash` on the entry
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ bytes: Vec<u8>,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ // stores the order new unique elements are seen for self.emit()
+ seen_elements: Vec<Option<Vec<u8>>>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ // key to group vector scratch space, used to avoid re-allocating a new vector on each call to intern
+ key_to_group: Vec<Option<usize>>,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ seen_elements: Vec::new(),
+ unique_dict_value_mapping: HashTable::with_capacity(INITIAL_PRE_ALLOCATION),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ key_to_group: Vec::new(),
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ /// Returns the existing `group_id` for the value with this hash and bytes,
+ /// or inserts a new entry and returns the freshly assigned group_id.
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ match self
+ .unique_dict_value_mapping
+ .entry(hash, |e| e.bytes == raw, |e| e.hash)
+ {
+ HashTableEntry::Occupied(o) => o.get().group_id,
+ HashTableEntry::Vacant(v) => {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(Some(raw.to_vec())); // replace this with raw buffer to avoid this double copy #TODO see https://github.com/apache/datafusion/issues/22078
+ v.insert(DictEntry {
+ hash,
+ group_id: new_group_id,
+ bytes: raw.to_vec(),
+ });
+ new_group_id
+ }
+ }
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(None);
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+ fn transform_into_array(&self, raw: &[Option<Vec<u8>>]) -> Result<ArrayRef> {
+ let item_capacity = raw.len();
+ let data_capacity = item_capacity * AVG_BYTES_PER_DICT_VALUE;
+ match &self.value_dt {
+ DataType::Utf8 => decode_scalar_string!(
+ raw,
+ StringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::LargeUtf8 => decode_scalar_string!(
+ raw,
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::Utf8View => decode_scalar_string!(
+ raw,
+ StringViewBuilder::with_capacity(item_capacity)
+ ),
+ DataType::List(field) => match field.data_type() {
+ DataType::Utf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ .with_field(Arc::clone(field))
+ ),
+ DataType::LargeUtf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ .with_field(Arc::clone(field))
+ ),
+ DataType::Utf8View => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringViewBuilder::with_capacity(item_capacity),
+ item_capacity,
+ )
+ .with_field(Arc::clone(field))
+ ),
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for List<{other:?}>"
+ ))),
+ },
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for {other:?}"
+ ))),
+ }
+ }
+}
+
+fn valid_bounds<K: ArrowDictionaryKeyType>(n: usize) -> bool {
+ let max: usize = 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,
+ _ => return false,
+ };
+ n <= max
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K> {
+ fn size(&self) -> usize {
+ let seen_elements_size = self.seen_elements.capacity()
+ * size_of::<Option<Vec<u8>>>()
+ + self
+ .seen_elements
+ .iter()
+ .filter_map(|opt| opt.as_ref())
+ .map(|inner| inner.capacity())
+ .sum::<usize>();
+
+ let unique_mapping_size = self.unique_dict_value_mapping.capacity()
+ * size_of::<DictEntry>()
+ + self
+ .unique_dict_value_mapping
+ .iter()
+ .map(|e| e.bytes.capacity())
+ .sum::<usize>();
+
+ let values_cache_size = self.values_cache.1.capacity() * size_of::<u64>()
+ + self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|a| a.to_data().get_array_memory_size())
Review Comment:
@Rich-T-kid there was a typo in the [comment](https://github.com/apache/datafusion/pull/21765#discussion_r3357174074) I left before, the function is not `get_array_memory_size()`, it is `get_slice_memory_size()`
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "adriangb (via GitHub)" <gi...@apache.org>.
adriangb commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4416549396
run benchmarks
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4401926750
Bench marks
<img width="941" height="1230" alt="Image 5-7-26 at 7 20 PM" src="https://github.com/user-attachments/assets/50125d91-f674-41e0-a030-8979170ad79c" />
<img width="1131" height="1308" alt="Image 5-7-26 at 7 20 PM (1)" src="https://github.com/user-attachments/assets/1596ff90-ac19-454a-91ba-b00f4642173f" />
<img width="955" height="565" alt="Image 5-7-26 at 7 20 PM (2)" src="https://github.com/user-attachments/assets/6eff0054-68c2-46bb-b24d-e692521caf59" />
the benchmarks in `physical-plan/benches/dictionary_group_values.rs` as well as the `datafusion/benchmarks/dict.rs` show a meaningful improvement. But I think there are still some improvements that can be made to make it even more efficent, One idea I have is to store intermediate bytes in one buffer as opposed to a vector of bytes, this removes the double memory allocation that is currently happening in intern. another improvement is to add caching to the value hashes that are computed.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4301676571
run benchmarks
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4301679144
Benchmark for [this request](https://github.com/apache/datafusion/pull/21765#issuecomment-4301676571) failed.
Last 20 lines of output:
<details><summary>Click to expand</summary>
```
* [new tag] 48.0.0-rc1 -> 48.0.0-rc1
* [new tag] 48.0.0-rc2 -> 48.0.0-rc2
* [new tag] 5.0.0 -> 5.0.0
* [new tag] 5.0.0-rc1 -> 5.0.0-rc1
* [new tag] 5.0.0-rc3 -> 5.0.0-rc3
* [new tag] 6.0.0 -> 6.0.0
* [new tag] 6.0.0-rc0 -> 6.0.0-rc0
* [new tag] 7.0.0 -> 7.0.0
* [new tag] 7.0.0-rc2 -> 7.0.0-rc2
* [new tag] 8.0.0 -> 8.0.0
* [new tag] 8.0.0-rc1 -> 8.0.0-rc1
* [new tag] 8.0.0-rc2 -> 8.0.0-rc2
* [new tag] 9.0.0 -> 9.0.0
* [new tag] 9.0.0-rc1 -> 9.0.0-rc1
* [new tag] ballista-0.5.0 -> ballista-0.5.0
* [new tag] ballista-0.6.0 -> ballista-0.6.0
* [new tag] ballista-0.7.0 -> ballista-0.7.0
* [new tag] python-0.3.0 -> python-0.3.0
* [new tag] python-0.4.0 -> python-0.4.0
Switched to branch 'rich-t-kid/optimize-dictionary-grouping'
```
</details>
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4291557213
<html><head></head><body><p><strong>TPC-H SF1 Benchmark Summary</strong></p>
Metric | Value
-- | --
Total Time (main) | 37.43ms
Total Time (rich-t-kid_optimize-dictionary-grouping) | 38.70ms
Average Time (main) | 1.70ms
Average Time (rich-t-kid_optimize-dictionary-grouping) | 1.76ms
Queries Faster | 2
Queries Slower | 13
Queries with No Change | 7
Queries with Failure | 0
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4291628227
There were also micro benchmarks linked in the[ previous PR](https://github.com/apache/datafusion/pull/21589#issuecomment-4253764562)
**Faster** : 25
**slower** : 17
**no-change** : 24
I didn't expect to see the tpc-h sf1 have this many regressions. I think I have an idea as to what is causing this. currently working on a fix
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4293109913
### tpch_sf1
<img width="1232" height="665" alt="Image 4-21-26 at 10 17 PM" src="https://github.com/user-attachments/assets/9ee80d1c-bad3-4b08-b09a-07fce6360621" />
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4292826785
### TPC-H SF10
<img width="1082" height="663" alt="Image 4-21-26 at 8 56 PM" src="https://github.com/user-attachments/assets/39803c21-d4df-4373-87e0-c2c42d093207" />
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4299954030
### tpch_sf10
<img width="1468" height="551" alt="Image 4-22-26 at 5 02 PM" src="https://github.com/user-attachments/assets/7e321840-21a8-4000-95d5-e34e8cdc3d82" />
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "LiaCastaneda (via GitHub)" <gi...@apache.org>.
LiaCastaneda commented on code in PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#discussion_r3360985739
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,783 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, DictionaryArray, LargeStringArray, LargeStringBuilder, ListArray,
+ ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+/// Entry stored inside `unique_dict_value_mapping`. Caching `hash` on the entry
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ bytes: Vec<u8>,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ // stores the order new unique elements are seen for self.emit()
+ seen_elements: Vec<Option<Vec<u8>>>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ intern_called: bool,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ seen_elements: Vec::new(),
+ unique_dict_value_mapping: HashTable::with_capacity(INITIAL_PRE_ALLOCATION),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ intern_called: false,
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ /// Returns the existing `group_id` for the value with this hash and bytes,
+ /// or inserts a new entry and returns the freshly assigned group_id.
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ match self
+ .unique_dict_value_mapping
+ .entry(hash, |e| e.bytes == raw, |e| e.hash)
+ {
+ HashTableEntry::Occupied(o) => o.get().group_id,
+ HashTableEntry::Vacant(v) => {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(Some(raw.to_vec())); // replace this with raw buffer to avoid this double copy #TODO see https://github.com/apache/datafusion/issues/22078
+ v.insert(DictEntry {
+ hash,
+ group_id: new_group_id,
+ bytes: raw.to_vec(),
+ });
+ new_group_id
+ }
+ }
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(None);
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+ fn transform_into_array(&self, raw: &[Option<Vec<u8>>]) -> Result<ArrayRef> {
+ let item_capacity = raw.len();
+ let data_capacity = item_capacity * AVG_BYTES_PER_DICT_VALUE;
+ match &self.value_dt {
+ DataType::Utf8 => decode_scalar_string!(
+ raw,
+ StringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::LargeUtf8 => decode_scalar_string!(
+ raw,
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::Utf8View => decode_scalar_string!(
+ raw,
+ StringViewBuilder::with_capacity(item_capacity)
+ ),
+ DataType::List(field) => match field.data_type() {
+ DataType::Utf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::LargeUtf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::Utf8View => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringViewBuilder::with_capacity(item_capacity),
+ item_capacity,
+ )
+ ),
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for List<{other:?}>"
+ ))),
+ },
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for {other:?}"
+ ))),
+ }
+ }
+ #[inline]
+ fn keys_to_usize(key_array: &PrimitiveArray<K>) -> Vec<Option<usize>> {
+ (0..key_array.len())
+ .map(|i| {
+ if key_array.is_null(i) {
+ None
+ } else {
+ Some(key_array.value(i).to_usize().unwrap())
+ }
+ })
+ .collect()
+ }
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K> {
+ fn size(&self) -> usize {
+ let seen_elements_size = self.seen_elements.capacity()
+ * size_of::<Option<Vec<u8>>>()
+ + self
+ .seen_elements
+ .iter()
+ .filter_map(|opt| opt.as_ref())
+ .map(|inner| inner.capacity())
+ .sum::<usize>();
+
+ let unique_mapping_size = self.unique_dict_value_mapping.capacity()
+ * size_of::<DictEntry>()
+ + self
+ .unique_dict_value_mapping
+ .iter()
+ .map(|e| e.bytes.capacity())
+ .sum::<usize>();
+
+ let values_cache_size = self.values_cache.1.capacity() * size_of::<u64>()
+ + self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|a| a.get_array_memory_size())
Review Comment:
@Rich-T-kid there was a typo in what I wrote, the function is not `get_array_memory_size()`, it is `get_slice_memory_size()`
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "kumarUjjawal (via GitHub)" <gi...@apache.org>.
kumarUjjawal commented on code in PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#discussion_r3361326084
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,768 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, LargeStringArray, LargeStringBuilder,
+ ListArray, ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+/// Entry stored inside `unique_dict_value_mapping`. Caching `hash` on the entry
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ bytes: Vec<u8>,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ // stores the order new unique elements are seen for self.emit()
+ seen_elements: Vec<Option<Vec<u8>>>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ intern_called: bool,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ seen_elements: Vec::new(),
+ unique_dict_value_mapping: HashTable::with_capacity(INITIAL_PRE_ALLOCATION),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ intern_called: false,
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ /// Returns the existing `group_id` for the value with this hash and bytes,
+ /// or inserts a new entry and returns the freshly assigned group_id.
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ match self
+ .unique_dict_value_mapping
+ .entry(hash, |e| e.bytes == raw, |e| e.hash)
+ {
+ HashTableEntry::Occupied(o) => o.get().group_id,
+ HashTableEntry::Vacant(v) => {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(Some(raw.to_vec())); // replace this with raw buffer to avoid this double copy #TODO see https://github.com/apache/datafusion/issues/22078
+ v.insert(DictEntry {
+ hash,
+ group_id: new_group_id,
+ bytes: raw.to_vec(),
+ });
+ new_group_id
+ }
+ }
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(None);
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+ fn transform_into_array(&self, raw: &[Option<Vec<u8>>]) -> Result<ArrayRef> {
+ let item_capacity = raw.len();
+ let data_capacity = item_capacity * AVG_BYTES_PER_DICT_VALUE;
+ match &self.value_dt {
+ DataType::Utf8 => decode_scalar_string!(
+ raw,
+ StringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::LargeUtf8 => decode_scalar_string!(
+ raw,
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::Utf8View => decode_scalar_string!(
+ raw,
+ StringViewBuilder::with_capacity(item_capacity)
+ ),
+ DataType::List(field) => match field.data_type() {
+ DataType::Utf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::LargeUtf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::Utf8View => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringViewBuilder::with_capacity(item_capacity),
+ item_capacity,
+ )
+ ),
+ other => Err(NotImplemented(format!(
Review Comment:
When the results are built back into a list column, the inner part of the list gets a hardcoded name ("item"). But Parquet files name it "element". DataFusion checks that names match exactly when assembling the output, so the query blows up at runtime. Also proved with a test.
when building the output, reuse the exact column description from the input table instead of letting the builder invent one.
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,768 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, LargeStringArray, LargeStringBuilder,
+ ListArray, ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+/// Entry stored inside `unique_dict_value_mapping`. Caching `hash` on the entry
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ bytes: Vec<u8>,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ // stores the order new unique elements are seen for self.emit()
+ seen_elements: Vec<Option<Vec<u8>>>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ intern_called: bool,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ seen_elements: Vec::new(),
+ unique_dict_value_mapping: HashTable::with_capacity(INITIAL_PRE_ALLOCATION),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ intern_called: false,
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ /// Returns the existing `group_id` for the value with this hash and bytes,
+ /// or inserts a new entry and returns the freshly assigned group_id.
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ match self
+ .unique_dict_value_mapping
+ .entry(hash, |e| e.bytes == raw, |e| e.hash)
+ {
+ HashTableEntry::Occupied(o) => o.get().group_id,
+ HashTableEntry::Vacant(v) => {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(Some(raw.to_vec())); // replace this with raw buffer to avoid this double copy #TODO see https://github.com/apache/datafusion/issues/22078
+ v.insert(DictEntry {
+ hash,
+ group_id: new_group_id,
+ bytes: raw.to_vec(),
+ });
+ new_group_id
+ }
+ }
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(None);
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+ fn transform_into_array(&self, raw: &[Option<Vec<u8>>]) -> Result<ArrayRef> {
+ let item_capacity = raw.len();
+ let data_capacity = item_capacity * AVG_BYTES_PER_DICT_VALUE;
+ match &self.value_dt {
+ DataType::Utf8 => decode_scalar_string!(
+ raw,
+ StringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::LargeUtf8 => decode_scalar_string!(
+ raw,
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::Utf8View => decode_scalar_string!(
+ raw,
+ StringViewBuilder::with_capacity(item_capacity)
+ ),
+ DataType::List(field) => match field.data_type() {
+ DataType::Utf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::LargeUtf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::Utf8View => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringViewBuilder::with_capacity(item_capacity),
+ item_capacity,
+ )
+ ),
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for List<{other:?}>"
+ ))),
+ },
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for {other:?}"
+ ))),
+ }
+ }
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K> {
+ fn size(&self) -> usize {
+ let seen_elements_size = self.seen_elements.capacity()
+ * size_of::<Option<Vec<u8>>>()
+ + self
+ .seen_elements
+ .iter()
+ .filter_map(|opt| opt.as_ref())
+ .map(|inner| inner.capacity())
+ .sum::<usize>();
+
+ let unique_mapping_size = self.unique_dict_value_mapping.capacity()
+ * size_of::<DictEntry>()
+ + self
+ .unique_dict_value_mapping
+ .iter()
+ .map(|e| e.bytes.capacity())
+ .sum::<usize>();
+
+ let values_cache_size = self.values_cache.1.capacity() * size_of::<u64>()
+ + self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|a| a.to_data().get_array_memory_size())
+ .unwrap_or(0);
+
+ size_of::<Self>() + seen_elements_size + unique_mapping_size + values_cache_size
+ }
+ fn len(&self) -> usize {
+ self.seen_elements.len()
+ }
+ fn is_empty(&self) -> bool {
+ self.seen_elements.is_empty()
+ }
+ fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
+ assert_eq!(
+ cols.len(),
+ 1,
+ "GroupValuesDictionary only supports a single column"
+ );
+ let array = Arc::clone(&cols[0]);
+ groups.clear(); // zero out buffer
+ let dict_array = array
+ .as_dictionary_opt()
+ .map_or_else(|| Err(Internal("Expected dictionary array".into())), Ok)?;
+
+ let values = dict_array.values();
+ let key_array: &PrimitiveArray<K> = dict_array.keys();
+ if key_array.is_empty() {
+ return Ok(());
+ }
+
+ let cache_hit = self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|cached| Arc::ptr_eq(cached, values))
+ .unwrap_or(false);
+
+ if !cache_hit {
+ // values array changed since last batch - recompute hashes and update cached pointer
+ self.compute_value_hashes(values)?;
+ self.values_cache.0 = Some(Arc::clone(values));
+ }
+ let mut key_to_group: Vec<Option<usize>> = vec![None; values.len()];
Review Comment:
Each time a chunk of data arrives, we builds a brand new list with one slot per dictionary entry, uses it, and throws it away. When the dictionary is large allocation is repeated over and over.
We should keep one scratch list on the tracker and reuse it between chunks
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,783 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, DictionaryArray, LargeStringArray, LargeStringBuilder, ListArray,
+ ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+/// Entry stored inside `unique_dict_value_mapping`. Caching `hash` on the entry
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ bytes: Vec<u8>,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ // stores the order new unique elements are seen for self.emit()
+ seen_elements: Vec<Option<Vec<u8>>>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ intern_called: bool,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ seen_elements: Vec::new(),
+ unique_dict_value_mapping: HashTable::with_capacity(INITIAL_PRE_ALLOCATION),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ intern_called: false,
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ /// Returns the existing `group_id` for the value with this hash and bytes,
+ /// or inserts a new entry and returns the freshly assigned group_id.
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ match self
+ .unique_dict_value_mapping
+ .entry(hash, |e| e.bytes == raw, |e| e.hash)
+ {
+ HashTableEntry::Occupied(o) => o.get().group_id,
+ HashTableEntry::Vacant(v) => {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(Some(raw.to_vec())); // replace this with raw buffer to avoid this double copy #TODO see https://github.com/apache/datafusion/issues/22078
+ v.insert(DictEntry {
+ hash,
+ group_id: new_group_id,
+ bytes: raw.to_vec(),
+ });
+ new_group_id
+ }
+ }
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(None);
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+ fn transform_into_array(&self, raw: &[Option<Vec<u8>>]) -> Result<ArrayRef> {
+ let item_capacity = raw.len();
+ let data_capacity = item_capacity * AVG_BYTES_PER_DICT_VALUE;
+ match &self.value_dt {
+ DataType::Utf8 => decode_scalar_string!(
+ raw,
+ StringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::LargeUtf8 => decode_scalar_string!(
+ raw,
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::Utf8View => decode_scalar_string!(
+ raw,
+ StringViewBuilder::with_capacity(item_capacity)
+ ),
+ DataType::List(field) => match field.data_type() {
+ DataType::Utf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::LargeUtf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::Utf8View => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringViewBuilder::with_capacity(item_capacity),
+ item_capacity,
+ )
+ ),
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for List<{other:?}>"
+ ))),
+ },
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for {other:?}"
+ ))),
+ }
+ }
+ #[inline]
+ fn keys_to_usize(key_array: &PrimitiveArray<K>) -> Vec<Option<usize>> {
+ (0..key_array.len())
+ .map(|i| {
+ if key_array.is_null(i) {
+ None
+ } else {
+ Some(key_array.value(i).to_usize().unwrap())
+ }
+ })
+ .collect()
+ }
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K> {
+ fn size(&self) -> usize {
+ let seen_elements_size = self.seen_elements.capacity()
+ * size_of::<Option<Vec<u8>>>()
+ + self
+ .seen_elements
+ .iter()
+ .filter_map(|opt| opt.as_ref())
+ .map(|inner| inner.capacity())
+ .sum::<usize>();
+
+ let unique_mapping_size = self.unique_dict_value_mapping.capacity()
+ * size_of::<DictEntry>()
+ + self
+ .unique_dict_value_mapping
+ .iter()
+ .map(|e| e.bytes.capacity())
+ .sum::<usize>();
+
+ let values_cache_size = self.values_cache.1.capacity() * size_of::<u64>()
+ + self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|a| a.get_array_memory_size())
+ .unwrap_or(0);
+
+ size_of::<Self>() + seen_elements_size + unique_mapping_size + values_cache_size
+ }
+ fn len(&self) -> usize {
+ self.seen_elements.len()
+ }
+ fn is_empty(&self) -> bool {
+ self.seen_elements.is_empty()
+ }
+ fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
+ assert_eq!(
+ cols.len(),
+ 1,
+ "GroupValuesDictionary only supports a single column"
+ );
+ let array = Arc::clone(&cols[0]);
+ groups.clear(); // zero out buffer
+ let dict_array = array
+ .as_any()
+ .downcast_ref::<DictionaryArray<K>>()
+ .ok_or_else(|| {
+ Internal(format!(
+ "GroupValuesDictionary expected DictionaryArray but got {:?}",
+ array.data_type()
+ ))
+ })?;
+
+ let values = dict_array.values();
+ let key_array = dict_array.keys();
+ if key_array.is_empty() {
+ return Ok(());
+ }
+ let keys_as_usize: Vec<Option<usize>> = Self::keys_to_usize(key_array);
+
+ let cache_hit = self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|cached| Arc::ptr_eq(cached, values))
+ .unwrap_or(false);
+
+ if !cache_hit {
+ // first time seeing this values array to compute hashes, and cache the pointer, cheap operation
+ self.compute_value_hashes(values)?;
+ self.values_cache.0 = Some(Arc::clone(values));
+ }
+ let mut key_to_group: Vec<Option<usize>> = vec![None; values.len()];
+ if self.intern_called {
Review Comment:
I think this whole pre-scan block can be removed. The main loop below already does at most one table lookup per unique dictionary entry and caches the result so the pre-scan doesn't save any lookups, it just moves them earlier and adds a second full pass over the keys.
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,768 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, LargeStringArray, LargeStringBuilder,
+ ListArray, ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+/// Entry stored inside `unique_dict_value_mapping`. Caching `hash` on the entry
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ bytes: Vec<u8>,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ // stores the order new unique elements are seen for self.emit()
+ seen_elements: Vec<Option<Vec<u8>>>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ intern_called: bool,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ seen_elements: Vec::new(),
+ unique_dict_value_mapping: HashTable::with_capacity(INITIAL_PRE_ALLOCATION),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ intern_called: false,
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ /// Returns the existing `group_id` for the value with this hash and bytes,
+ /// or inserts a new entry and returns the freshly assigned group_id.
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ match self
+ .unique_dict_value_mapping
+ .entry(hash, |e| e.bytes == raw, |e| e.hash)
+ {
+ HashTableEntry::Occupied(o) => o.get().group_id,
+ HashTableEntry::Vacant(v) => {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(Some(raw.to_vec())); // replace this with raw buffer to avoid this double copy #TODO see https://github.com/apache/datafusion/issues/22078
+ v.insert(DictEntry {
+ hash,
+ group_id: new_group_id,
+ bytes: raw.to_vec(),
+ });
+ new_group_id
+ }
+ }
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(None);
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+ fn transform_into_array(&self, raw: &[Option<Vec<u8>>]) -> Result<ArrayRef> {
+ let item_capacity = raw.len();
+ let data_capacity = item_capacity * AVG_BYTES_PER_DICT_VALUE;
+ match &self.value_dt {
+ DataType::Utf8 => decode_scalar_string!(
+ raw,
+ StringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::LargeUtf8 => decode_scalar_string!(
+ raw,
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::Utf8View => decode_scalar_string!(
+ raw,
+ StringViewBuilder::with_capacity(item_capacity)
+ ),
+ DataType::List(field) => match field.data_type() {
+ DataType::Utf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::LargeUtf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::Utf8View => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringViewBuilder::with_capacity(item_capacity),
+ item_capacity,
+ )
+ ),
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for List<{other:?}>"
+ ))),
+ },
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for {other:?}"
+ ))),
+ }
+ }
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K> {
+ fn size(&self) -> usize {
+ let seen_elements_size = self.seen_elements.capacity()
+ * size_of::<Option<Vec<u8>>>()
+ + self
+ .seen_elements
+ .iter()
+ .filter_map(|opt| opt.as_ref())
+ .map(|inner| inner.capacity())
+ .sum::<usize>();
+
+ let unique_mapping_size = self.unique_dict_value_mapping.capacity()
+ * size_of::<DictEntry>()
+ + self
+ .unique_dict_value_mapping
+ .iter()
+ .map(|e| e.bytes.capacity())
+ .sum::<usize>();
+
+ let values_cache_size = self.values_cache.1.capacity() * size_of::<u64>()
+ + self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|a| a.to_data().get_array_memory_size())
+ .unwrap_or(0);
+
+ size_of::<Self>() + seen_elements_size + unique_mapping_size + values_cache_size
+ }
+ fn len(&self) -> usize {
+ self.seen_elements.len()
+ }
+ fn is_empty(&self) -> bool {
+ self.seen_elements.is_empty()
+ }
+ fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
+ assert_eq!(
+ cols.len(),
+ 1,
+ "GroupValuesDictionary only supports a single column"
+ );
+ let array = Arc::clone(&cols[0]);
+ groups.clear(); // zero out buffer
+ let dict_array = array
+ .as_dictionary_opt()
+ .map_or_else(|| Err(Internal("Expected dictionary array".into())), Ok)?;
+
+ let values = dict_array.values();
+ let key_array: &PrimitiveArray<K> = dict_array.keys();
+ if key_array.is_empty() {
+ return Ok(());
+ }
+
+ let cache_hit = self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|cached| Arc::ptr_eq(cached, values))
+ .unwrap_or(false);
+
+ if !cache_hit {
+ // values array changed since last batch - recompute hashes and update cached pointer
+ self.compute_value_hashes(values)?;
+ self.values_cache.0 = Some(Arc::clone(values));
+ }
+ let mut key_to_group: Vec<Option<usize>> = vec![None; values.len()];
+ // For values already seen in prior batches, pre fill `key_to_group`
+ // so the main loop below can skip hash table lookups for those keys.
+ // This pre scan is skipped on the first batch since the hash table is empty.
+ if self.intern_called {
+ for i in 0..key_array.len() {
+ if key_array.is_null(i) {
+ continue;
+ }
+ let original_key = key_array.value(i).to_usize().unwrap();
+ if values.is_null(original_key) {
+ continue;
+ }
+ if key_to_group[original_key].is_some() {
+ continue;
+ }
+
+ let hash = self.values_cache.1[original_key];
+ let raw = Self::get_raw_bytes(values, original_key);
+ if let Some(entry) = self
+ .unique_dict_value_mapping
+ .find(hash, |e| e.bytes == raw.as_ref())
+ {
+ key_to_group[original_key] = Some(entry.group_id);
+ }
+ }
+ }
+ // iterate keys array (n iterations)
+ // only d insertions at most, repeated work is cached
+ for i in 0..key_array.len() {
+ let group_id = if key_array.is_null(i) {
+ self.get_null_group_id()
+ } else {
+ let key = key_array.value(i).to_usize().unwrap();
+ if let Some(group_id) = key_to_group[key] {
+ group_id
+ } else if values.is_null(key) {
+ let gid = self.get_null_group_id();
+ key_to_group[key] = Some(gid);
+ gid
+ } else {
+ let hash = self.values_cache.1[key];
+ let raw = Self::get_raw_bytes(values, key);
+ let gid = self.lookup_or_insert_in_table(hash, raw.as_ref());
+ key_to_group[key] = Some(gid);
+ gid
+ }
+ };
+ groups.push(group_id);
+ }
+ // set flag to true to repeated calls to intern will pre-fill key_to_group cache
+ self.intern_called = true;
+ Ok(())
+ }
+ fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
+ let (elements_to_emit, null_id) = match emit_to {
+ EmitTo::All => {
+ let original_null_id = self.null_group_id;
+ self.null_group_id = None;
+ self.unique_dict_value_mapping.clear();
+ (std::mem::take(&mut self.seen_elements), original_null_id)
+ }
+ EmitTo::First(n) => {
+ let n = n.min(self.seen_elements.len());
+ let first_n = self.seen_elements.drain(..n).collect::<Vec<_>>(); // again read from raw buffer directly
+ let original_null_id = self.null_group_id.filter(|&id| id < n);
+ // update null_group_id if the null group was in the first n
+ if let Some(null_id) = self.null_group_id {
+ if null_id < n {
+ self.null_group_id = None;
+ } else {
+ self.null_group_id = Some(null_id - n);
+ }
+ }
+ // shift all remaining group indices down by n in the table
+ self.unique_dict_value_mapping.retain(|e| {
+ if e.group_id < n {
+ false
+ } else {
+ e.group_id -= n;
+ true
+ }
+ });
+ (first_n, original_null_id)
+ }
+ };
+
+ let n = elements_to_emit.len();
+ let values_array = self.transform_into_array(&elements_to_emit)?; // ingest a buffer + n then produce an array. should also update the len+offset of other entries #TODO see https://github.com/apache/datafusion/issues/22078
+
+ let mut keys_builder = PrimitiveBuilder::<K>::with_capacity(n);
+ if let Some(null_id) = null_id {
+ for i in 0..n {
+ if i == null_id {
+ keys_builder.append_null();
+ } else {
+ keys_builder.append_value(K::Native::usize_as(i));
+ }
+ }
+ } else {
+ for i in 0..n {
+ keys_builder.append_value(K::Native::usize_as(i));
+ }
Review Comment:
Each chunk of data that flows in can carry its own lookup list of values. Small number types can only count so high, one of them maxes out at 256. But across many chunks, the total number of distinct groups can grow past that limit. When it does, the counter wraps around and starts over from zero, so a group near the end gets labeled with some other group's value.
Before producing results, we should check whether the group count still fits in the number type if not, stop with a clear error. Even nicer would be automatically switching to a bigger number type?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4723014865
local results of benchmarks for the 1:1 cardinality case (ran against main)
<img width="1354" height="254" alt="Image 6-16-26 at 4 06 PM" src="https://github.com/user-attachments/assets/d03c2318-5c1e-4f0b-ab4e-8a098cdfe9c4" />
@kumarUjjawal could you run the benchmarks again?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4418453084
🤖 Benchmark completed (GKE) | [trigger](https://github.com/apache/datafusion/pull/21765#issuecomment-4418355614)
**Instance:** `c4a-highmem-16` (12 vCPU / 65 GiB)
<details><summary>CPU Details (lscpu)</summary>
```
Architecture: aarch64
CPU op-mode(s): 64-bit
Byte Order: Little Endian
CPU(s): 16
On-line CPU(s) list: 0-15
Vendor ID: ARM
Model name: Neoverse-V2
Model: 1
Thread(s) per core: 1
Core(s) per cluster: 16
Socket(s): -
Cluster(s): 1
Stepping: r0p1
BogoMIPS: 2000.00
Flags: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache: 1 MiB (16 instances)
L1i cache: 1 MiB (16 instances)
L2 cache: 32 MiB (16 instances)
L3 cache: 80 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; __user pointer sanitization
Vulnerability Spectre v2: Mitigation; CSV2, BHB
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected
```
</details>
<details><summary>Details</summary>
<p>
```
Comparing HEAD and rich-t-kid_optimize-dictionary-grouping
--------------------
Benchmark tpcds_sf1.json
--------------------
┏━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━┓
┃ Query ┃ HEAD ┃ rich-t-kid_optimize-dictionary-grouping ┃ Change ┃
┡━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━┩
│ QQuery 1 │ 6.25 / 6.71 ±0.87 / 8.44 ms │ 6.28 / 6.75 ±0.80 / 8.34 ms │ no change │
│ QQuery 2 │ 82.12 / 82.40 ±0.29 / 82.83 ms │ 81.42 / 81.64 ±0.14 / 81.80 ms │ no change │
│ QQuery 3 │ 28.83 / 29.13 ±0.24 / 29.53 ms │ 29.13 / 29.30 ±0.19 / 29.65 ms │ no change │
│ QQuery 4 │ 510.31 / 516.82 ±5.56 / 526.65 ms │ 510.66 / 514.75 ±3.68 / 519.18 ms │ no change │
│ QQuery 5 │ 52.37 / 52.97 ±0.43 / 53.54 ms │ 52.55 / 53.20 ±0.50 / 54.04 ms │ no change │
│ QQuery 6 │ 35.10 / 35.75 ±0.43 / 36.34 ms │ 35.39 / 35.78 ±0.31 / 36.28 ms │ no change │
│ QQuery 7 │ 109.75 / 111.60 ±3.04 / 117.67 ms │ 110.44 / 111.91 ±1.16 / 113.45 ms │ no change │
│ QQuery 8 │ 39.02 / 39.20 ±0.12 / 39.37 ms │ 39.72 / 40.35 ±0.34 / 40.73 ms │ no change │
│ QQuery 9 │ 53.67 / 54.87 ±0.77 / 55.85 ms │ 54.55 / 55.45 ±0.72 / 56.15 ms │ no change │
│ QQuery 10 │ 80.94 / 81.90 ±1.43 / 84.75 ms │ 83.26 / 84.58 ±1.97 / 88.51 ms │ no change │
│ QQuery 11 │ 312.54 / 317.97 ±4.87 / 325.49 ms │ 340.65 / 344.59 ±2.88 / 349.21 ms │ 1.08x slower │
│ QQuery 12 │ 28.28 / 28.71 ±0.33 / 29.16 ms │ 29.44 / 29.74 ±0.34 / 30.37 ms │ no change │
│ QQuery 13 │ 128.66 / 129.09 ±0.34 / 129.49 ms │ 130.66 / 131.46 ±0.46 / 132.00 ms │ no change │
│ QQuery 14 │ 513.64 / 515.58 ±1.84 / 518.96 ms │ 521.92 / 522.12 ±0.27 / 522.66 ms │ no change │
│ QQuery 15 │ 60.55 / 61.63 ±0.88 / 62.69 ms │ 61.78 / 62.60 ±0.67 / 63.65 ms │ no change │
│ QQuery 16 │ 6.95 / 7.08 ±0.18 / 7.44 ms │ 6.92 / 7.16 ±0.18 / 7.48 ms │ no change │
│ QQuery 17 │ 83.52 / 84.43 ±0.64 / 85.20 ms │ 83.90 / 84.87 ±1.10 / 86.90 ms │ no change │
│ QQuery 18 │ 155.56 / 156.49 ±0.83 / 157.87 ms │ 156.52 / 156.93 ±0.51 / 157.89 ms │ no change │
│ QQuery 19 │ 41.14 / 41.54 ±0.36 / 42.20 ms │ 41.39 / 41.59 ±0.25 / 42.06 ms │ no change │
│ QQuery 20 │ 35.26 / 35.75 ±0.40 / 36.25 ms │ 35.24 / 35.88 ±0.35 / 36.22 ms │ no change │
│ QQuery 21 │ 17.93 / 18.31 ±0.24 / 18.53 ms │ 17.89 / 18.16 ±0.26 / 18.60 ms │ no change │
│ QQuery 22 │ 62.14 / 63.12 ±0.94 / 64.85 ms │ 62.67 / 62.93 ±0.25 / 63.28 ms │ no change │
│ QQuery 23 │ 476.87 / 479.49 ±1.79 / 482.36 ms │ 477.82 / 484.83 ±4.28 / 489.64 ms │ no change │
│ QQuery 24 │ 241.57 / 245.58 ±2.05 / 247.01 ms │ 243.24 / 246.92 ±3.47 / 252.83 ms │ no change │
│ QQuery 25 │ 116.94 / 118.67 ±1.77 / 121.54 ms │ 117.11 / 119.58 ±1.74 / 122.32 ms │ no change │
│ QQuery 26 │ 70.73 / 71.30 ±0.70 / 72.66 ms │ 71.95 / 72.57 ±0.68 / 73.90 ms │ no change │
│ QQuery 27 │ 6.89 / 7.03 ±0.11 / 7.24 ms │ 6.88 / 7.01 ±0.16 / 7.31 ms │ no change │
│ QQuery 28 │ 62.14 / 62.69 ±0.35 / 63.15 ms │ 62.42 / 62.90 ±0.26 / 63.14 ms │ no change │
│ QQuery 29 │ 101.94 / 102.20 ±0.21 / 102.48 ms │ 101.99 / 102.52 ±0.40 / 103.20 ms │ no change │
│ QQuery 30 │ FAIL │ FAIL │ incomparable │
│ QQuery 31 │ 111.97 / 113.23 ±1.15 / 115.19 ms │ 111.70 / 112.81 ±1.18 / 114.92 ms │ no change │
│ QQuery 32 │ 20.69 / 20.80 ±0.10 / 20.99 ms │ 20.19 / 20.55 ±0.38 / 21.23 ms │ no change │
│ QQuery 33 │ 39.02 / 39.25 ±0.21 / 39.60 ms │ 39.26 / 39.48 ±0.11 / 39.59 ms │ no change │
│ QQuery 34 │ 9.70 / 10.21 ±0.58 / 11.28 ms │ 10.11 / 10.29 ±0.16 / 10.50 ms │ no change │
│ QQuery 35 │ 80.24 / 80.73 ±0.48 / 81.60 ms │ 81.05 / 82.65 ±1.97 / 86.43 ms │ no change │
│ QQuery 36 │ 6.32 / 6.44 ±0.21 / 6.87 ms │ 6.45 / 6.58 ±0.18 / 6.94 ms │ no change │
│ QQuery 37 │ 7.00 / 7.17 ±0.13 / 7.39 ms │ 7.03 / 7.19 ±0.09 / 7.29 ms │ no change │
│ QQuery 38 │ 67.85 / 68.93 ±0.92 / 70.23 ms │ 67.75 / 68.81 ±0.60 / 69.29 ms │ no change │
│ QQuery 39 │ 100.36 / 101.21 ±0.52 / 101.91 ms │ 100.26 / 100.94 ±0.61 / 101.79 ms │ no change │
│ QQuery 40 │ 23.69 / 24.02 ±0.29 / 24.48 ms │ 23.47 / 24.15 ±0.48 / 24.91 ms │ no change │
│ QQuery 41 │ 14.16 / 14.59 ±0.58 / 15.70 ms │ 14.08 / 14.21 ±0.14 / 14.47 ms │ no change │
│ QQuery 42 │ 24.37 / 24.66 ±0.24 / 24.98 ms │ 24.01 / 24.15 ±0.14 / 24.40 ms │ no change │
│ QQuery 43 │ 5.19 / 5.31 ±0.17 / 5.64 ms │ 5.40 / 5.48 ±0.11 / 5.69 ms │ no change │
│ QQuery 44 │ 10.60 / 10.72 ±0.14 / 11.00 ms │ 10.87 / 11.08 ±0.14 / 11.30 ms │ no change │
│ QQuery 45 │ 40.67 / 42.26 ±0.81 / 42.91 ms │ 40.35 / 41.35 ±0.54 / 41.94 ms │ no change │
│ QQuery 46 │ 13.31 / 13.68 ±0.22 / 13.92 ms │ 13.71 / 13.85 ±0.11 / 14.01 ms │ no change │
│ QQuery 47 │ 230.97 / 236.45 ±4.70 / 243.89 ms │ 228.62 / 232.46 ±3.16 / 238.22 ms │ no change │
│ QQuery 48 │ 104.48 / 105.52 ±1.13 / 107.65 ms │ 104.20 / 105.49 ±0.68 / 106.17 ms │ no change │
│ QQuery 49 │ 82.01 / 82.80 ±0.91 / 84.43 ms │ 81.60 / 82.31 ±0.45 / 82.97 ms │ no change │
│ QQuery 50 │ 61.94 / 63.98 ±1.36 / 65.73 ms │ 62.63 / 64.41 ±1.74 / 66.67 ms │ no change │
│ QQuery 51 │ 94.04 / 95.55 ±1.08 / 96.89 ms │ 92.88 / 93.99 ±0.96 / 95.42 ms │ no change │
│ QQuery 52 │ 24.05 / 24.59 ±0.45 / 25.30 ms │ 24.44 / 24.73 ±0.20 / 24.96 ms │ no change │
│ QQuery 53 │ 30.19 / 30.52 ±0.17 / 30.69 ms │ 30.59 / 31.19 ±0.94 / 33.05 ms │ no change │
│ QQuery 54 │ 54.14 / 55.16 ±0.53 / 55.60 ms │ 54.54 / 55.72 ±1.80 / 59.29 ms │ no change │
│ QQuery 55 │ 23.59 / 25.27 ±1.88 / 28.84 ms │ 23.77 / 24.13 ±0.40 / 24.92 ms │ no change │
│ QQuery 56 │ 39.71 / 40.29 ±0.64 / 41.50 ms │ 39.36 / 39.88 ±0.33 / 40.34 ms │ no change │
│ QQuery 57 │ 177.15 / 181.16 ±2.44 / 184.50 ms │ 179.47 / 180.36 ±0.68 / 181.05 ms │ no change │
│ QQuery 58 │ 117.21 / 118.05 ±0.65 / 119.05 ms │ 118.28 / 118.77 ±0.46 / 119.35 ms │ no change │
│ QQuery 59 │ 120.02 / 121.34 ±1.49 / 124.14 ms │ 119.18 / 119.88 ±0.41 / 120.31 ms │ no change │
│ QQuery 60 │ 39.86 / 40.22 ±0.28 / 40.66 ms │ 40.67 / 41.49 ±0.42 / 41.85 ms │ no change │
│ QQuery 61 │ 13.79 / 13.88 ±0.12 / 14.12 ms │ 14.41 / 15.02 ±0.68 / 16.30 ms │ 1.08x slower │
│ QQuery 62 │ 46.67 / 47.62 ±1.51 / 50.62 ms │ 47.40 / 48.13 ±0.92 / 49.83 ms │ no change │
│ QQuery 63 │ 30.48 / 30.63 ±0.12 / 30.81 ms │ 31.17 / 31.66 ±0.38 / 32.19 ms │ no change │
│ QQuery 64 │ 477.06 / 484.49 ±7.03 / 495.36 ms │ 485.59 / 489.30 ±1.88 / 490.61 ms │ no change │
│ QQuery 65 │ 148.40 / 150.37 ±1.35 / 151.75 ms │ 149.29 / 154.14 ±2.67 / 156.38 ms │ no change │
│ QQuery 66 │ 83.05 / 83.32 ±0.44 / 84.19 ms │ 82.86 / 83.87 ±1.01 / 85.76 ms │ no change │
│ QQuery 67 │ 248.44 / 253.24 ±3.23 / 257.05 ms │ 242.10 / 246.20 ±2.81 / 249.34 ms │ no change │
│ QQuery 68 │ 13.56 / 13.72 ±0.22 / 14.14 ms │ 13.74 / 13.96 ±0.20 / 14.34 ms │ no change │
│ QQuery 69 │ 77.04 / 78.75 ±3.13 / 85.00 ms │ 77.87 / 78.05 ±0.11 / 78.14 ms │ no change │
│ QQuery 70 │ 105.50 / 108.46 ±2.54 / 111.75 ms │ 104.98 / 109.00 ±2.28 / 111.16 ms │ no change │
│ QQuery 71 │ 35.57 / 35.98 ±0.29 / 36.39 ms │ 35.33 / 35.98 ±0.48 / 36.67 ms │ no change │
│ QQuery 72 │ 2228.10 / 2300.40 ±59.68 / 2376.47 ms │ 2164.51 / 2221.19 ±42.45 / 2278.49 ms │ no change │
│ QQuery 73 │ 9.68 / 10.08 ±0.28 / 10.45 ms │ 9.57 / 9.79 ±0.25 / 10.27 ms │ no change │
│ QQuery 74 │ 181.27 / 183.48 ±1.75 / 185.60 ms │ 177.22 / 181.63 ±2.64 / 184.67 ms │ no change │
│ QQuery 75 │ 148.94 / 151.38 ±3.14 / 157.56 ms │ 150.82 / 151.44 ±0.49 / 152.16 ms │ no change │
│ QQuery 76 │ 35.33 / 35.76 ±0.33 / 36.30 ms │ 35.86 / 37.29 ±2.01 / 41.28 ms │ no change │
│ QQuery 77 │ 61.53 / 62.00 ±0.40 / 62.70 ms │ 61.85 / 62.44 ±0.35 / 62.79 ms │ no change │
│ QQuery 78 │ 195.26 / 198.14 ±1.78 / 199.81 ms │ 196.56 / 198.91 ±2.14 / 201.77 ms │ no change │
│ QQuery 79 │ 67.25 / 68.91 ±2.60 / 74.00 ms │ 67.15 / 68.30 ±0.98 / 70.13 ms │ no change │
│ QQuery 80 │ 103.08 / 104.48 ±0.74 / 105.03 ms │ 104.64 / 106.08 ±1.33 / 108.44 ms │ no change │
│ QQuery 81 │ 24.69 / 24.90 ±0.22 / 25.26 ms │ 24.67 / 24.93 ±0.17 / 25.20 ms │ no change │
│ QQuery 82 │ 16.59 / 16.76 ±0.15 / 17.01 ms │ 16.76 / 16.93 ±0.16 / 17.19 ms │ no change │
│ QQuery 83 │ 37.81 / 39.08 ±2.07 / 43.21 ms │ 37.01 / 37.63 ±0.38 / 38.19 ms │ no change │
│ QQuery 84 │ 43.02 / 44.62 ±1.83 / 48.19 ms │ 43.31 / 44.22 ±1.12 / 46.39 ms │ no change │
│ QQuery 85 │ 136.54 / 137.33 ±0.58 / 138.00 ms │ 138.04 / 138.82 ±0.90 / 140.55 ms │ no change │
│ QQuery 86 │ 24.88 / 25.13 ±0.22 / 25.54 ms │ 25.07 / 25.34 ±0.25 / 25.75 ms │ no change │
│ QQuery 87 │ 68.82 / 69.66 ±0.87 / 70.98 ms │ 68.79 / 69.69 ±0.81 / 71.05 ms │ no change │
│ QQuery 88 │ 64.59 / 65.21 ±0.66 / 66.45 ms │ 64.55 / 65.02 ±0.44 / 65.64 ms │ no change │
│ QQuery 89 │ 36.09 / 36.32 ±0.14 / 36.47 ms │ 36.01 / 36.20 ±0.18 / 36.47 ms │ no change │
│ QQuery 90 │ 17.32 / 17.45 ±0.13 / 17.68 ms │ 17.53 / 17.75 ±0.21 / 18.04 ms │ no change │
│ QQuery 91 │ 52.29 / 53.26 ±1.62 / 56.48 ms │ 51.73 / 52.50 ±0.55 / 53.43 ms │ no change │
│ QQuery 92 │ 29.47 / 29.68 ±0.20 / 30.05 ms │ 29.59 / 30.55 ±1.28 / 33.02 ms │ no change │
│ QQuery 93 │ 53.84 / 54.60 ±0.87 / 56.31 ms │ 53.17 / 54.52 ±1.41 / 56.49 ms │ no change │
│ QQuery 94 │ 37.99 / 38.75 ±0.65 / 39.96 ms │ 37.81 / 38.15 ±0.18 / 38.31 ms │ no change │
│ QQuery 95 │ 89.91 / 90.62 ±0.45 / 91.02 ms │ 90.22 / 90.66 ±0.32 / 91.15 ms │ no change │
│ QQuery 96 │ 24.31 / 24.44 ±0.21 / 24.86 ms │ 24.42 / 24.60 ±0.19 / 24.96 ms │ no change │
│ QQuery 97 │ 46.94 / 47.40 ±0.26 / 47.73 ms │ 46.73 / 47.19 ±0.30 / 47.56 ms │ no change │
│ QQuery 98 │ 42.43 / 42.98 ±0.38 / 43.57 ms │ 42.44 / 43.18 ±0.65 / 44.27 ms │ no change │
│ QQuery 99 │ 69.96 / 70.77 ±1.30 / 73.36 ms │ 69.71 / 70.29 ±0.44 / 71.06 ms │ no change │
└───────────┴───────────────────────────────────────┴─────────────────────────────────────────┴──────────────┘
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┓
┃ Benchmark Summary ┃ ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━┩
│ Total Time (HEAD) │ 10706.22ms │
│ Total Time (rich-t-kid_optimize-dictionary-grouping) │ 10676.97ms │
│ Average Time (HEAD) │ 109.25ms │
│ Average Time (rich-t-kid_optimize-dictionary-grouping) │ 108.95ms │
│ Queries Faster │ 0 │
│ Queries Slower │ 2 │
│ Queries with No Change │ 96 │
│ Queries with Failure │ 1 │
└────────────────────────────────────────────────────────┴────────────┘
```
</p>
</details>
<details><summary>Resource Usage</summary>
**tpcds — base (merge-base)**
| Metric | Value |
|--------|-------|
| Wall time | 55.0s |
| Peak memory | 6.8 GiB |
| Avg memory | 6.1 GiB |
| CPU user | 249.5s |
| CPU sys | 6.5s |
| Peak spill | 0 B |
**tpcds — branch**
| Metric | Value |
|--------|-------|
| Wall time | 55.0s |
| Peak memory | 6.9 GiB |
| Avg memory | 6.2 GiB |
| CPU user | 245.5s |
| CPU sys | 6.4s |
| Peak spill | 0 B |
</details>
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4418375138
🤖 Benchmark running (GKE) | [trigger](https://github.com/apache/datafusion/pull/21765#issuecomment-4418355614)
**Instance:** `c4a-highmem-16` (12 vCPU / 65 GiB) | `Linux bench-c4418355614-11-j82mv 6.12.68+ #1 SMP Wed Apr 1 02:23:28 UTC 2026 aarch64 GNU/Linux`
<details><summary>CPU Details (lscpu)</summary>
```
Architecture: aarch64
CPU op-mode(s): 64-bit
Byte Order: Little Endian
CPU(s): 16
On-line CPU(s) list: 0-15
Vendor ID: ARM
Model name: Neoverse-V2
Model: 1
Thread(s) per core: 1
Core(s) per cluster: 16
Socket(s): -
Cluster(s): 1
Stepping: r0p1
BogoMIPS: 2000.00
Flags: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache: 1 MiB (16 instances)
L1i cache: 1 MiB (16 instances)
L2 cache: 32 MiB (16 instances)
L3 cache: 80 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; __user pointer sanitization
Vulnerability Spectre v2: Mitigation; CSV2, BHB
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected
```
</details>
Comparing rich-t-kid/optimize-dictionary-grouping (577baab1e8f652768893653cd1b02e782fb76f63) to aca4d13 (merge-base) [diff](https://github.com/apache/datafusion/compare/aca4d1377bf9785e6fa9a154c579cc761ca7b54b..577baab1e8f652768893653cd1b02e782fb76f63) using: clickbench_partitioned
Results will be posted here when complete
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4418382686
🤖 Benchmark running (GKE) | [trigger](https://github.com/apache/datafusion/pull/21765#issuecomment-4418355614)
**Instance:** `c4a-highmem-16` (12 vCPU / 65 GiB) | `Linux bench-c4418355614-13-l6rjm 6.12.68+ #1 SMP Wed Apr 1 02:23:28 UTC 2026 aarch64 GNU/Linux`
<details><summary>CPU Details (lscpu)</summary>
```
Architecture: aarch64
CPU op-mode(s): 64-bit
Byte Order: Little Endian
CPU(s): 16
On-line CPU(s) list: 0-15
Vendor ID: ARM
Model name: Neoverse-V2
Model: 1
Thread(s) per core: 1
Core(s) per cluster: 16
Socket(s): -
Cluster(s): 1
Stepping: r0p1
BogoMIPS: 2000.00
Flags: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache: 1 MiB (16 instances)
L1i cache: 1 MiB (16 instances)
L2 cache: 32 MiB (16 instances)
L3 cache: 80 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; __user pointer sanitization
Vulnerability Spectre v2: Mitigation; CSV2, BHB
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected
```
</details>
Comparing rich-t-kid/optimize-dictionary-grouping (577baab1e8f652768893653cd1b02e782fb76f63) to aca4d13 (merge-base) [diff](https://github.com/apache/datafusion/compare/aca4d1377bf9785e6fa9a154c579cc761ca7b54b..577baab1e8f652768893653cd1b02e782fb76f63) using: tpch
Results will be posted here when complete
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4422408709
🥳 the results look good. thank you @kumarUjjawal !! could you also run the end to end query benchmarks that was introduced in https://github.com/apache/datafusion/pull/21860
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4408765861
Does it normally take this long to for @adriangbot ?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4407934429
@alamb the PR is ready, Can you run the benchmarks?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on code in PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#discussion_r3365277937
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,768 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, LargeStringArray, LargeStringBuilder,
+ ListArray, ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+/// Entry stored inside `unique_dict_value_mapping`. Caching `hash` on the entry
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ bytes: Vec<u8>,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ // stores the order new unique elements are seen for self.emit()
+ seen_elements: Vec<Option<Vec<u8>>>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ intern_called: bool,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ seen_elements: Vec::new(),
+ unique_dict_value_mapping: HashTable::with_capacity(INITIAL_PRE_ALLOCATION),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ intern_called: false,
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ /// Returns the existing `group_id` for the value with this hash and bytes,
+ /// or inserts a new entry and returns the freshly assigned group_id.
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ match self
+ .unique_dict_value_mapping
+ .entry(hash, |e| e.bytes == raw, |e| e.hash)
+ {
+ HashTableEntry::Occupied(o) => o.get().group_id,
+ HashTableEntry::Vacant(v) => {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(Some(raw.to_vec())); // replace this with raw buffer to avoid this double copy #TODO see https://github.com/apache/datafusion/issues/22078
+ v.insert(DictEntry {
+ hash,
+ group_id: new_group_id,
+ bytes: raw.to_vec(),
+ });
+ new_group_id
+ }
+ }
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(None);
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+ fn transform_into_array(&self, raw: &[Option<Vec<u8>>]) -> Result<ArrayRef> {
+ let item_capacity = raw.len();
+ let data_capacity = item_capacity * AVG_BYTES_PER_DICT_VALUE;
+ match &self.value_dt {
+ DataType::Utf8 => decode_scalar_string!(
+ raw,
+ StringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::LargeUtf8 => decode_scalar_string!(
+ raw,
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::Utf8View => decode_scalar_string!(
+ raw,
+ StringViewBuilder::with_capacity(item_capacity)
+ ),
+ DataType::List(field) => match field.data_type() {
+ DataType::Utf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::LargeUtf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::Utf8View => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringViewBuilder::with_capacity(item_capacity),
+ item_capacity,
+ )
+ ),
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for List<{other:?}>"
+ ))),
+ },
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for {other:?}"
+ ))),
+ }
+ }
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K> {
+ fn size(&self) -> usize {
+ let seen_elements_size = self.seen_elements.capacity()
+ * size_of::<Option<Vec<u8>>>()
+ + self
+ .seen_elements
+ .iter()
+ .filter_map(|opt| opt.as_ref())
+ .map(|inner| inner.capacity())
+ .sum::<usize>();
+
+ let unique_mapping_size = self.unique_dict_value_mapping.capacity()
+ * size_of::<DictEntry>()
+ + self
+ .unique_dict_value_mapping
+ .iter()
+ .map(|e| e.bytes.capacity())
+ .sum::<usize>();
+
+ let values_cache_size = self.values_cache.1.capacity() * size_of::<u64>()
+ + self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|a| a.to_data().get_array_memory_size())
+ .unwrap_or(0);
+
+ size_of::<Self>() + seen_elements_size + unique_mapping_size + values_cache_size
+ }
+ fn len(&self) -> usize {
+ self.seen_elements.len()
+ }
+ fn is_empty(&self) -> bool {
+ self.seen_elements.is_empty()
+ }
+ fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
+ assert_eq!(
+ cols.len(),
+ 1,
+ "GroupValuesDictionary only supports a single column"
+ );
+ let array = Arc::clone(&cols[0]);
+ groups.clear(); // zero out buffer
+ let dict_array = array
+ .as_dictionary_opt()
+ .map_or_else(|| Err(Internal("Expected dictionary array".into())), Ok)?;
+
+ let values = dict_array.values();
+ let key_array: &PrimitiveArray<K> = dict_array.keys();
+ if key_array.is_empty() {
+ return Ok(());
+ }
+
+ let cache_hit = self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|cached| Arc::ptr_eq(cached, values))
+ .unwrap_or(false);
+
+ if !cache_hit {
+ // values array changed since last batch - recompute hashes and update cached pointer
+ self.compute_value_hashes(values)?;
+ self.values_cache.0 = Some(Arc::clone(values));
+ }
+ let mut key_to_group: Vec<Option<usize>> = vec![None; values.len()];
+ // For values already seen in prior batches, pre fill `key_to_group`
+ // so the main loop below can skip hash table lookups for those keys.
+ // This pre scan is skipped on the first batch since the hash table is empty.
+ if self.intern_called {
+ for i in 0..key_array.len() {
+ if key_array.is_null(i) {
+ continue;
+ }
+ let original_key = key_array.value(i).to_usize().unwrap();
+ if values.is_null(original_key) {
+ continue;
+ }
+ if key_to_group[original_key].is_some() {
+ continue;
+ }
+
+ let hash = self.values_cache.1[original_key];
+ let raw = Self::get_raw_bytes(values, original_key);
+ if let Some(entry) = self
+ .unique_dict_value_mapping
+ .find(hash, |e| e.bytes == raw.as_ref())
+ {
+ key_to_group[original_key] = Some(entry.group_id);
+ }
+ }
+ }
+ // iterate keys array (n iterations)
+ // only d insertions at most, repeated work is cached
+ for i in 0..key_array.len() {
+ let group_id = if key_array.is_null(i) {
+ self.get_null_group_id()
+ } else {
+ let key = key_array.value(i).to_usize().unwrap();
+ if let Some(group_id) = key_to_group[key] {
+ group_id
+ } else if values.is_null(key) {
+ let gid = self.get_null_group_id();
+ key_to_group[key] = Some(gid);
+ gid
+ } else {
+ let hash = self.values_cache.1[key];
+ let raw = Self::get_raw_bytes(values, key);
+ let gid = self.lookup_or_insert_in_table(hash, raw.as_ref());
+ key_to_group[key] = Some(gid);
+ gid
+ }
+ };
+ groups.push(group_id);
+ }
+ // set flag to true to repeated calls to intern will pre-fill key_to_group cache
+ self.intern_called = true;
+ Ok(())
+ }
+ fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
+ let (elements_to_emit, null_id) = match emit_to {
+ EmitTo::All => {
+ let original_null_id = self.null_group_id;
+ self.null_group_id = None;
+ self.unique_dict_value_mapping.clear();
+ (std::mem::take(&mut self.seen_elements), original_null_id)
+ }
+ EmitTo::First(n) => {
+ let n = n.min(self.seen_elements.len());
+ let first_n = self.seen_elements.drain(..n).collect::<Vec<_>>(); // again read from raw buffer directly
+ let original_null_id = self.null_group_id.filter(|&id| id < n);
+ // update null_group_id if the null group was in the first n
+ if let Some(null_id) = self.null_group_id {
+ if null_id < n {
+ self.null_group_id = None;
+ } else {
+ self.null_group_id = Some(null_id - n);
+ }
+ }
+ // shift all remaining group indices down by n in the table
+ self.unique_dict_value_mapping.retain(|e| {
+ if e.group_id < n {
+ false
+ } else {
+ e.group_id -= n;
+ true
+ }
+ });
+ (first_n, original_null_id)
+ }
+ };
+
+ let n = elements_to_emit.len();
+ let values_array = self.transform_into_array(&elements_to_emit)?; // ingest a buffer + n then produce an array. should also update the len+offset of other entries #TODO see https://github.com/apache/datafusion/issues/22078
+
+ let mut keys_builder = PrimitiveBuilder::<K>::with_capacity(n);
+ if let Some(null_id) = null_id {
+ for i in 0..n {
+ if i == null_id {
+ keys_builder.append_null();
+ } else {
+ keys_builder.append_value(K::Native::usize_as(i));
+ }
+ }
+ } else {
+ for i in 0..n {
+ keys_builder.append_value(K::Native::usize_as(i));
+ }
Review Comment:
I believe the appropriate thing to do is to return an error. from my understand altering a dictionaries type from `Dictionary<int8,utf8>` to `Dictionary<int16,utf8>` can have [compatibility problems for upstream operators](https://github.com/apache/datafusion/issues/7647). I added test that assert this behavoir.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4645679872
Benchmark for [this request](https://github.com/apache/datafusion/pull/21765#issuecomment-4645670621) failed.
Last 20 lines of output:
<details><summary>Click to expand</summary>
```
* [new tag] 48.0.0-rc1 -> 48.0.0-rc1
* [new tag] 48.0.0-rc2 -> 48.0.0-rc2
* [new tag] 5.0.0 -> 5.0.0
* [new tag] 5.0.0-rc1 -> 5.0.0-rc1
* [new tag] 5.0.0-rc3 -> 5.0.0-rc3
* [new tag] 6.0.0 -> 6.0.0
* [new tag] 6.0.0-rc0 -> 6.0.0-rc0
* [new tag] 7.0.0 -> 7.0.0
* [new tag] 7.0.0-rc2 -> 7.0.0-rc2
* [new tag] 8.0.0 -> 8.0.0
* [new tag] 8.0.0-rc1 -> 8.0.0-rc1
* [new tag] 8.0.0-rc2 -> 8.0.0-rc2
* [new tag] 9.0.0 -> 9.0.0
* [new tag] 9.0.0-rc1 -> 9.0.0-rc1
* [new tag] ballista-0.5.0 -> ballista-0.5.0
* [new tag] ballista-0.6.0 -> ballista-0.6.0
* [new tag] ballista-0.7.0 -> ballista-0.7.0
* [new tag] python-0.3.0 -> python-0.3.0
* [new tag] python-0.4.0 -> python-0.4.0
Switched to branch 'rich-t-kid/optimize-dictionary-grouping'
```
</details>
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4299955474
### tpch_mem_sf1
<img width="1458" height="548" alt="Image 4-22-26 at 5 03 PM" src="https://github.com/user-attachments/assets/ba583426-efeb-42dd-9660-1d91931d08f7" />
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4293112110
### tpch_sf10
<img width="1208" height="666" alt="Image 4-21-26 at 10 18 PM" src="https://github.com/user-attachments/assets/9d55335b-99c5-462d-8836-dc5621a7aae7" />
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4299958040
### tpch_mem_sf10
<img width="1425" height="565" alt="Image 4-22-26 at 5 03 PM" src="https://github.com/user-attachments/assets/94d8ac7e-3fc9-47a8-a4a9-e7e39d5f3e2c" />
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4299947460
reran the benchmarks after generating data.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4291566798
<html><head></head><body><p><strong>TPC-H In-Memory SF1 Benchmark Summary</strong></p>
Metric | Value
-- | --
Total Time (main) | 353.30ms
Total Time (rich-t-kid_optimize-dictionary-grouping) | 362.53ms
Average Time (main) | 16.06ms
Average Time (rich-t-kid_optimize-dictionary-grouping) | 16.48ms
Queries Faster | 3
Queries Slower | 3
Queries with No Change | 16
Queries with Failure | 0
</body></html>
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4289346457
If your interested in some of the discussions made leading up to this PR please view https://github.com/apache/datafusion/pull/21589
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4290869431
(basically it is hard to justify an optimization without benchmark results, even if they need to be run manually at first)
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4292829830
### TPC-H In-Memory SF1
<img width="1344" height="663" alt="Image 4-21-26 at 8 57 PM" src="https://github.com/user-attachments/assets/25e6fd52-bafc-4d37-9f2c-68e0f856291a" />
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4292823218
### TPC-H SF1
<img width="1201" height="667" alt="Image 4-21-26 at 8 55 PM" src="https://github.com/user-attachments/assets/a5c8d773-7197-42e2-936f-e4fc2d90217b" />
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on code in PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#discussion_r3364375245
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,783 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, DictionaryArray, LargeStringArray, LargeStringBuilder, ListArray,
+ ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+/// Entry stored inside `unique_dict_value_mapping`. Caching `hash` on the entry
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ bytes: Vec<u8>,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ // stores the order new unique elements are seen for self.emit()
+ seen_elements: Vec<Option<Vec<u8>>>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ intern_called: bool,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ seen_elements: Vec::new(),
+ unique_dict_value_mapping: HashTable::with_capacity(INITIAL_PRE_ALLOCATION),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ intern_called: false,
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ /// Returns the existing `group_id` for the value with this hash and bytes,
+ /// or inserts a new entry and returns the freshly assigned group_id.
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ match self
+ .unique_dict_value_mapping
+ .entry(hash, |e| e.bytes == raw, |e| e.hash)
+ {
+ HashTableEntry::Occupied(o) => o.get().group_id,
+ HashTableEntry::Vacant(v) => {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(Some(raw.to_vec())); // replace this with raw buffer to avoid this double copy #TODO see https://github.com/apache/datafusion/issues/22078
+ v.insert(DictEntry {
+ hash,
+ group_id: new_group_id,
+ bytes: raw.to_vec(),
+ });
+ new_group_id
+ }
+ }
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(None);
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+ fn transform_into_array(&self, raw: &[Option<Vec<u8>>]) -> Result<ArrayRef> {
+ let item_capacity = raw.len();
+ let data_capacity = item_capacity * AVG_BYTES_PER_DICT_VALUE;
+ match &self.value_dt {
+ DataType::Utf8 => decode_scalar_string!(
+ raw,
+ StringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::LargeUtf8 => decode_scalar_string!(
+ raw,
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::Utf8View => decode_scalar_string!(
+ raw,
+ StringViewBuilder::with_capacity(item_capacity)
+ ),
+ DataType::List(field) => match field.data_type() {
+ DataType::Utf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::LargeUtf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::Utf8View => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringViewBuilder::with_capacity(item_capacity),
+ item_capacity,
+ )
+ ),
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for List<{other:?}>"
+ ))),
+ },
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for {other:?}"
+ ))),
+ }
+ }
+ #[inline]
+ fn keys_to_usize(key_array: &PrimitiveArray<K>) -> Vec<Option<usize>> {
+ (0..key_array.len())
+ .map(|i| {
+ if key_array.is_null(i) {
+ None
+ } else {
+ Some(key_array.value(i).to_usize().unwrap())
+ }
+ })
+ .collect()
+ }
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K> {
+ fn size(&self) -> usize {
+ let seen_elements_size = self.seen_elements.capacity()
+ * size_of::<Option<Vec<u8>>>()
+ + self
+ .seen_elements
+ .iter()
+ .filter_map(|opt| opt.as_ref())
+ .map(|inner| inner.capacity())
+ .sum::<usize>();
+
+ let unique_mapping_size = self.unique_dict_value_mapping.capacity()
+ * size_of::<DictEntry>()
+ + self
+ .unique_dict_value_mapping
+ .iter()
+ .map(|e| e.bytes.capacity())
+ .sum::<usize>();
+
+ let values_cache_size = self.values_cache.1.capacity() * size_of::<u64>()
+ + self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|a| a.get_array_memory_size())
+ .unwrap_or(0);
+
+ size_of::<Self>() + seen_elements_size + unique_mapping_size + values_cache_size
+ }
+ fn len(&self) -> usize {
+ self.seen_elements.len()
+ }
+ fn is_empty(&self) -> bool {
+ self.seen_elements.is_empty()
+ }
+ fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
+ assert_eq!(
+ cols.len(),
+ 1,
+ "GroupValuesDictionary only supports a single column"
+ );
+ let array = Arc::clone(&cols[0]);
+ groups.clear(); // zero out buffer
+ let dict_array = array
+ .as_any()
+ .downcast_ref::<DictionaryArray<K>>()
+ .ok_or_else(|| {
+ Internal(format!(
+ "GroupValuesDictionary expected DictionaryArray but got {:?}",
+ array.data_type()
+ ))
+ })?;
+
+ let values = dict_array.values();
+ let key_array = dict_array.keys();
+ if key_array.is_empty() {
+ return Ok(());
+ }
+ let keys_as_usize: Vec<Option<usize>> = Self::keys_to_usize(key_array);
+
+ let cache_hit = self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|cached| Arc::ptr_eq(cached, values))
+ .unwrap_or(false);
+
+ if !cache_hit {
+ // first time seeing this values array to compute hashes, and cache the pointer, cheap operation
+ self.compute_value_hashes(values)?;
+ self.values_cache.0 = Some(Arc::clone(values));
+ }
+ let mut key_to_group: Vec<Option<usize>> = vec![None; values.len()];
+ if self.intern_called {
Review Comment:
I think that makes sense. just to be safe ill benchmark the difference & share the results if performance doesn't improve like we expect it to.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on code in PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#discussion_r3373423040
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,878 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, LargeStringArray, LargeStringBuilder,
+ ListArray, ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+/// Entry stored inside `unique_dict_value_mapping`. Caching `hash` on the entry
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ bytes: Vec<u8>,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ // stores the order new unique elements are seen for self.emit()
+ seen_elements: Vec<Option<Vec<u8>>>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ // key to group vector scratch space, used to avoid re-allocating a new vector on each call to intern
+ key_to_group: Vec<Option<usize>>,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ seen_elements: Vec::new(),
+ unique_dict_value_mapping: HashTable::with_capacity(INITIAL_PRE_ALLOCATION),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ key_to_group: Vec::new(),
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ /// Returns the existing `group_id` for the value with this hash and bytes,
+ /// or inserts a new entry and returns the freshly assigned group_id.
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ match self
+ .unique_dict_value_mapping
+ .entry(hash, |e| e.bytes == raw, |e| e.hash)
+ {
+ HashTableEntry::Occupied(o) => o.get().group_id,
+ HashTableEntry::Vacant(v) => {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(Some(raw.to_vec())); // replace this with raw buffer to avoid this double copy #TODO see https://github.com/apache/datafusion/issues/22078
+ v.insert(DictEntry {
+ hash,
+ group_id: new_group_id,
+ bytes: raw.to_vec(),
+ });
+ new_group_id
+ }
+ }
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(None);
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+ fn transform_into_array(&self, raw: &[Option<Vec<u8>>]) -> Result<ArrayRef> {
+ let item_capacity = raw.len();
+ let data_capacity = item_capacity * AVG_BYTES_PER_DICT_VALUE;
+ match &self.value_dt {
+ DataType::Utf8 => decode_scalar_string!(
+ raw,
+ StringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::LargeUtf8 => decode_scalar_string!(
+ raw,
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::Utf8View => decode_scalar_string!(
+ raw,
+ StringViewBuilder::with_capacity(item_capacity)
+ ),
+ DataType::List(field) => match field.data_type() {
+ DataType::Utf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ .with_field(Arc::clone(field))
+ ),
+ DataType::LargeUtf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ .with_field(Arc::clone(field))
+ ),
+ DataType::Utf8View => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringViewBuilder::with_capacity(item_capacity),
+ item_capacity,
+ )
+ .with_field(Arc::clone(field))
+ ),
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for List<{other:?}>"
+ ))),
+ },
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for {other:?}"
+ ))),
+ }
+ }
+}
+
+fn valid_bounds<K: ArrowDictionaryKeyType>(n: usize) -> bool {
+ let max: usize = 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,
+ _ => return false,
+ };
+ n <= max
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K> {
+ fn size(&self) -> usize {
+ let seen_elements_size = self.seen_elements.capacity()
+ * size_of::<Option<Vec<u8>>>()
+ + self
+ .seen_elements
+ .iter()
+ .filter_map(|opt| opt.as_ref())
+ .map(|inner| inner.capacity())
+ .sum::<usize>();
+
+ let unique_mapping_size = self.unique_dict_value_mapping.capacity()
+ * size_of::<DictEntry>()
+ + self
+ .unique_dict_value_mapping
+ .iter()
+ .map(|e| e.bytes.capacity())
+ .sum::<usize>();
+
+ let values_cache_size = self.values_cache.1.capacity() * size_of::<u64>()
+ + self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|a| a.to_data().get_array_memory_size())
Review Comment:
I've updated the function!
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "kumarUjjawal (via GitHub)" <gi...@apache.org>.
kumarUjjawal commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4662438041
@Rich-T-kid we are seeing regression in the high cardinality but low cardinality is excellent
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4662380765
🤖 Benchmark completed (GKE) | [trigger](https://github.com/apache/datafusion/pull/21765#issuecomment-4662199962)
**Instance:** `c4a-highmem-16` (12 vCPU / 65 GiB)
<details><summary>CPU Details (lscpu)</summary>
```
Architecture: aarch64
CPU op-mode(s): 64-bit
Byte Order: Little Endian
CPU(s): 16
On-line CPU(s) list: 0-15
Vendor ID: ARM
Model name: Neoverse-V2
Model: 1
Thread(s) per core: 1
Core(s) per cluster: 16
Socket(s): -
Cluster(s): 1
Stepping: r0p1
BogoMIPS: 2000.00
Flags: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache: 1 MiB (16 instances)
L1i cache: 1 MiB (16 instances)
L2 cache: 32 MiB (16 instances)
L3 cache: 80 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; __user pointer sanitization
Vulnerability Spectre v2: Mitigation; CSV2, BHB
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected
```
</details>
<details><summary>Details</summary>
<p>
```
group HEAD rich-t-kid_optimize-dictionary-grouping
----- ---- ---------------------------------------
dict_intern_emit/intern_emit/size_65536_card_1000_null_0.00 3.39 885.3±8.93µs 70.6 MElem/sec 1.00 261.3±2.07µs 239.2 MElem/sec
dict_intern_emit/intern_emit/size_65536_card_20_null_0.00 5.81 793.8±6.22µs 78.7 MElem/sec 1.00 136.6±1.03µs 457.6 MElem/sec
dict_intern_emit/intern_emit/size_65536_card_300_null_0.00 4.83 827.4±8.79µs 75.5 MElem/sec 1.00 171.2±0.82µs 365.0 MElem/sec
dict_intern_emit/intern_emit/size_65536_card_65536_null_0.00 1.00 6.3±0.01ms 10.0 MElem/sec 2.50 15.6±0.12ms 4.0 MElem/sec
dict_intern_emit/intern_emit/size_65536_card_75_null_0.00 5.70 817.0±26.55µs 76.5 MElem/sec 1.00 143.4±0.41µs 435.9 MElem/sec
dict_intern_emit/intern_emit/size_8192_card_1000_null_0.00 1.12 162.2±1.16µs 48.2 MElem/sec 1.00 144.7±2.12µs 54.0 MElem/sec
dict_intern_emit/intern_emit/size_8192_card_20_null_0.00 5.72 110.4±47.55µs 70.8 MElem/sec 1.00 19.3±0.17µs 404.7 MElem/sec
dict_intern_emit/intern_emit/size_8192_card_300_null_0.00 2.25 122.8±0.98µs 63.6 MElem/sec 1.00 54.6±0.51µs 143.1 MElem/sec
dict_intern_emit/intern_emit/size_8192_card_75_null_0.00 4.17 111.3±1.12µs 70.2 MElem/sec 1.00 26.7±0.16µs 292.6 MElem/sec
dict_intern_emit/intern_emit/size_8192_card_8192_null_0.00 1.00 671.0±2.56µs 11.6 MElem/sec 1.66 1112.2±18.96µs 7.0 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_65536_card_1000_null_0.10 4.54 4.4±0.03ms 56.9 MElem/sec 1.00 967.0±4.56µs 258.5 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_65536_card_20_null_0.10 5.52 4.2±0.03ms 60.1 MElem/sec 1.00 754.1±2.08µs 331.5 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_65536_card_300_null_0.10 5.25 4.2±0.03ms 59.0 MElem/sec 1.00 806.3±1.24µs 310.0 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_65536_card_65536_null_0.10 1.00 17.3±0.07ms 14.5 MElem/sec 1.63 28.2±0.26ms 8.9 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_65536_card_75_null_0.10 5.53 4.2±0.03ms 59.2 MElem/sec 1.00 764.0±1.13µs 327.2 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_8192_card_1000_null_0.10 2.19 611.9±3.09µs 51.1 MElem/sec 1.00 279.8±1.90µs 111.7 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_8192_card_20_null_0.10 5.54 503.4±4.10µs 62.1 MElem/sec 1.00 90.9±0.22µs 343.8 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_8192_card_300_null_0.10 3.75 540.5±5.84µs 57.8 MElem/sec 1.00 144.2±1.07µs 216.7 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_8192_card_75_null_0.10 5.23 518.9±6.50µs 60.2 MElem/sec 1.00 99.2±0.35µs 315.1 MElem/sec
dict_repeated_intern_emit/repeated_intern_emit/size_8192_card_8192_null_0.10 1.00 1478.4±6.62µs 21.1 MElem/sec 1.22 1806.9±20.28µs 17.3 MElem/sec
```
</p>
</details>
<details><summary>Resource Usage</summary>
**dictionary_group_values — base (merge-base)**
| Metric | Value |
|--------|-------|
| Wall time | 535.1s |
| Peak memory | 617.0 MiB |
| Avg memory | 74.9 MiB |
| CPU user | 231.1s |
| CPU sys | 15.4s |
| Peak spill | 0 B |
**dictionary_group_values — branch**
| Metric | Value |
|--------|-------|
| Wall time | 535.1s |
| Peak memory | 895.4 MiB |
| Avg memory | 176.9 MiB |
| CPU user | 255.4s |
| CPU sys | 3.9s |
| Peak spill | 0 B |
</details>
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4673419896
I've experimented with switching between `GroupValuesRows` and `GroupValuesDictionary` at runtime depending on the size of the dictionary values array, but this isn't a good approach. DataFusion is a streaming engine, so the first batch of data we see isn't necessarily an indicator of future batches. If we decide to use `GroupValuesRows` due to high cardinality in the first incoming RecordBatch and each subsequent batch turns out to be low cardinality, we lose out on the benefits `GroupValuesDictionary` provides (a possible 5x speedup). The inverse is also true.
Switching between the two approaches mid-aggregation wouldn't be effective either. Having to `Emit(EmitTo::All)` from one approach and then `intern()` into another would cause more overhead than any savings it could produce.
I think the best case to trust users to use the correct data types to represent their data, with `GroupValuesDictionary` it give them another option besides `GroupValuesRows` in case they do have low-medium cardinality.
I'm open to hearing other possible approaches @alamb
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4418448238
🤖 Benchmark completed (GKE) | [trigger](https://github.com/apache/datafusion/pull/21765#issuecomment-4418355614)
**Instance:** `c4a-highmem-16` (12 vCPU / 65 GiB)
<details><summary>CPU Details (lscpu)</summary>
```
Architecture: aarch64
CPU op-mode(s): 64-bit
Byte Order: Little Endian
CPU(s): 16
On-line CPU(s) list: 0-15
Vendor ID: ARM
Model name: Neoverse-V2
Model: 1
Thread(s) per core: 1
Core(s) per cluster: 16
Socket(s): -
Cluster(s): 1
Stepping: r0p1
BogoMIPS: 2000.00
Flags: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache: 1 MiB (16 instances)
L1i cache: 1 MiB (16 instances)
L2 cache: 32 MiB (16 instances)
L3 cache: 80 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; __user pointer sanitization
Vulnerability Spectre v2: Mitigation; CSV2, BHB
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected
```
</details>
<details><summary>Details</summary>
<p>
```
Comparing HEAD and rich-t-kid_optimize-dictionary-grouping
--------------------
Benchmark tpch_sf1.json
--------------------
┏━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Query ┃ HEAD ┃ rich-t-kid_optimize-dictionary-grouping ┃ Change ┃
┡━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ QQuery 1 │ 38.88 / 39.93 ±0.96 / 41.46 ms │ 38.44 / 39.39 ±1.10 / 41.04 ms │ no change │
│ QQuery 2 │ 20.30 / 20.53 ±0.19 / 20.89 ms │ 20.09 / 20.60 ±0.45 / 21.35 ms │ no change │
│ QQuery 3 │ 34.94 / 36.67 ±1.28 / 38.27 ms │ 35.50 / 37.17 ±0.94 / 38.39 ms │ no change │
│ QQuery 4 │ 17.42 / 17.97 ±0.70 / 19.35 ms │ 17.37 / 17.51 ±0.11 / 17.65 ms │ no change │
│ QQuery 5 │ 43.09 / 44.47 ±0.99 / 45.89 ms │ 42.73 / 44.07 ±1.03 / 45.19 ms │ no change │
│ QQuery 6 │ 16.35 / 17.23 ±0.83 / 18.65 ms │ 16.17 / 16.21 ±0.06 / 16.33 ms │ +1.06x faster │
│ QQuery 7 │ 48.82 / 50.48 ±1.39 / 52.63 ms │ 48.00 / 49.19 ±0.94 / 50.18 ms │ no change │
│ QQuery 8 │ 45.44 / 46.02 ±0.69 / 47.37 ms │ 45.37 / 46.63 ±1.03 / 47.80 ms │ no change │
│ QQuery 9 │ 50.15 / 50.97 ±0.80 / 52.44 ms │ 49.01 / 50.66 ±0.94 / 51.79 ms │ no change │
│ QQuery 10 │ 63.84 / 64.12 ±0.21 / 64.40 ms │ 63.29 / 63.41 ±0.13 / 63.65 ms │ no change │
│ QQuery 11 │ 13.43 / 13.77 ±0.35 / 14.43 ms │ 13.31 / 13.62 ±0.18 / 13.80 ms │ no change │
│ QQuery 12 │ 25.36 / 25.56 ±0.24 / 26.04 ms │ 24.80 / 25.70 ±0.90 / 27.41 ms │ no change │
│ QQuery 13 │ 35.82 / 36.39 ±0.41 / 37.06 ms │ 35.78 / 36.06 ±0.34 / 36.69 ms │ no change │
│ QQuery 14 │ 25.73 / 25.84 ±0.13 / 26.09 ms │ 25.63 / 25.70 ±0.04 / 25.76 ms │ no change │
│ QQuery 15 │ 31.50 / 32.30 ±0.73 / 33.23 ms │ 31.31 / 31.75 ±0.58 / 32.89 ms │ no change │
│ QQuery 16 │ 14.75 / 14.95 ±0.12 / 15.09 ms │ 14.81 / 15.36 ±0.63 / 16.50 ms │ no change │
│ QQuery 17 │ 75.26 / 77.00 ±2.37 / 81.57 ms │ 74.29 / 76.19 ±2.21 / 80.40 ms │ no change │
│ QQuery 18 │ 68.63 / 70.56 ±1.20 / 72.01 ms │ 68.00 / 68.83 ±0.73 / 69.74 ms │ no change │
│ QQuery 19 │ 35.25 / 35.50 ±0.19 / 35.72 ms │ 34.87 / 35.18 ±0.21 / 35.43 ms │ no change │
│ QQuery 20 │ 37.60 / 38.23 ±0.46 / 38.77 ms │ 37.70 / 38.02 ±0.23 / 38.34 ms │ no change │
│ QQuery 21 │ 58.03 / 59.55 ±1.80 / 62.79 ms │ 58.15 / 60.10 ±2.12 / 64.24 ms │ no change │
│ QQuery 22 │ 23.09 / 23.47 ±0.29 / 23.87 ms │ 23.38 / 23.53 ±0.12 / 23.68 ms │ no change │
└───────────┴────────────────────────────────┴─────────────────────────────────────────┴───────────────┘
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━┓
┃ Benchmark Summary ┃ ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━┩
│ Total Time (HEAD) │ 841.50ms │
│ Total Time (rich-t-kid_optimize-dictionary-grouping) │ 834.88ms │
│ Average Time (HEAD) │ 38.25ms │
│ Average Time (rich-t-kid_optimize-dictionary-grouping) │ 37.95ms │
│ Queries Faster │ 1 │
│ Queries Slower │ 0 │
│ Queries with No Change │ 21 │
│ Queries with Failure │ 0 │
└────────────────────────────────────────────────────────┴──────────┘
```
</p>
</details>
<details><summary>Resource Usage</summary>
**tpch — base (merge-base)**
| Metric | Value |
|--------|-------|
| Wall time | 5.0s |
| Peak memory | 5.5 GiB |
| Avg memory | 5.0 GiB |
| CPU user | 31.2s |
| CPU sys | 2.1s |
| Peak spill | 0 B |
**tpch — branch**
| Metric | Value |
|--------|-------|
| Wall time | 5.0s |
| Peak memory | 5.5 GiB |
| Avg memory | 5.0 GiB |
| CPU user | 31.0s |
| CPU sys | 2.1s |
| Peak spill | 0 B |
</details>
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid closed pull request #21765: Optimize Dictionary groupings
URL: https://github.com/apache/datafusion/pull/21765
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4293366367
### sf10
<img width="1332" height="667" alt="Image 4-21-26 at 11 44 PM" src="https://github.com/user-attachments/assets/5479a4c6-1a35-4c75-a9e0-fc72426585e9" />
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4292837556
**faster**: 19
**slower**: 8
**no-change**: 39
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4293119744
**Faster** : 21
**Slower** : 9
**No-diff** : 36
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4293402012
**speed up** : 24
**Slower** : 6
**no change** : 36
last optimize that comes to mind is updating **normalize_dict_hash()** by eliminating per-insert allocations.
Currently, **normalize_dict_hash()** creates a **HashMap<Vec<u8>, usize>** where each unique key's raw bytes are heap allocated via **.to_vec()** before insertion. This allocation occurs once per unique value but is unnecessary since the underlying Arrow buffer already owns the bytes.
The plan is to pre-compute a **Vec<Option<Cow<[u8]>>>** of raw byte slices for all accessed values upfront, allowing the hashmap to store **&[u8]** references instead of owned **Vec<u8>** keys, eliminating the per-insert allocations entirely.
In a perfect world I would benchmark this to see if eliminating the **.to_vec()** allocations outweighs the cost of the upfront pass over the accessed value indices to build the slice cache. I'll implement both approaches, benchmark them, and push the more optimized version.
with that being said I think these results show a generally a large improvement over the current approach to dealing with dictionary encoded columns in data-fusion. @alamb
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4301681290
Benchmark for [this request](https://github.com/apache/datafusion/pull/21765#issuecomment-4301676571) failed.
Last 20 lines of output:
<details><summary>Click to expand</summary>
```
* [new tag] 48.0.0-rc1 -> 48.0.0-rc1
* [new tag] 48.0.0-rc2 -> 48.0.0-rc2
* [new tag] 5.0.0 -> 5.0.0
* [new tag] 5.0.0-rc1 -> 5.0.0-rc1
* [new tag] 5.0.0-rc3 -> 5.0.0-rc3
* [new tag] 6.0.0 -> 6.0.0
* [new tag] 6.0.0-rc0 -> 6.0.0-rc0
* [new tag] 7.0.0 -> 7.0.0
* [new tag] 7.0.0-rc2 -> 7.0.0-rc2
* [new tag] 8.0.0 -> 8.0.0
* [new tag] 8.0.0-rc1 -> 8.0.0-rc1
* [new tag] 8.0.0-rc2 -> 8.0.0-rc2
* [new tag] 9.0.0 -> 9.0.0
* [new tag] 9.0.0-rc1 -> 9.0.0-rc1
* [new tag] ballista-0.5.0 -> ballista-0.5.0
* [new tag] ballista-0.6.0 -> ballista-0.6.0
* [new tag] ballista-0.7.0 -> ballista-0.7.0
* [new tag] python-0.3.0 -> python-0.3.0
* [new tag] python-0.4.0 -> python-0.4.0
Switched to branch 'rich-t-kid/optimize-dictionary-grouping'
```
</details>
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4301675954
Are these queries even using dictionary data types?
I think they are string views and primitives
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "adriangb (via GitHub)" <gi...@apache.org>.
adriangb commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4416559411
run benchmarks
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4422217852
🤖 Criterion benchmark running (GKE) | [trigger](https://github.com/apache/datafusion/pull/21765#issuecomment-4422189550)
**Instance:** `c4a-highmem-16` (12 vCPU / 65 GiB) | `Linux bench-c4422189550-14-zhklw 6.12.68+ #1 SMP Wed Apr 1 02:23:28 UTC 2026 aarch64 GNU/Linux`
<details><summary>CPU Details (lscpu)</summary>
```
Architecture: aarch64
CPU op-mode(s): 64-bit
Byte Order: Little Endian
CPU(s): 16
On-line CPU(s) list: 0-15
Vendor ID: ARM
Model name: Neoverse-V2
Model: 1
Thread(s) per core: 1
Core(s) per cluster: 16
Socket(s): -
Cluster(s): 1
Stepping: r0p1
BogoMIPS: 2000.00
Flags: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache: 1 MiB (16 instances)
L1i cache: 1 MiB (16 instances)
L2 cache: 32 MiB (16 instances)
L3 cache: 80 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; __user pointer sanitization
Vulnerability Spectre v2: Mitigation; CSV2, BHB
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected
```
</details>
Comparing rich-t-kid/optimize-dictionary-grouping (577baab1e8f652768893653cd1b02e782fb76f63) to aca4d13 (merge-base) [diff](https://github.com/apache/datafusion/compare/aca4d1377bf9785e6fa9a154c579cc761ca7b54b..577baab1e8f652768893653cd1b02e782fb76f63)
BENCH_NAME=dictionary_group_values
BENCH_COMMAND=cargo bench --features=parquet --bench dictionary_group_values
BENCH_FILTER=
Results will be posted here when complete
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "kumarUjjawal (via GitHub)" <gi...@apache.org>.
kumarUjjawal commented on code in PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#discussion_r3371046924
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,878 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, LargeStringArray, LargeStringBuilder,
+ ListArray, ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
Review Comment:
The lookup table is always pre-sized to ~8,000 entries on creation. so most instances will hold a handful of values yet still pay for thousands of empty slots. We can start small and let it grow, or size it from a cheap signal (such as the number of distinct values already present), rather than a fixed large constant.
##########
datafusion/physical-plan/src/aggregates/group_values/mod.rs:
##########
@@ -196,6 +202,39 @@ pub fn new_group_values(
DataType::Boolean => {
return Ok(Box::new(GroupValuesBoolean::new()));
}
+ DataType::Dictionary(key_type, value_type)
+ if supported_single_dictionary_value(value_type) =>
+ {
+ return match key_type.as_ref() {
+ DataType::Int8 => {
+ Ok(Box::new(GroupValuesDictionary::<Int8Type>::new(value_type)))
+ }
+ DataType::Int16 => Ok(Box::new(
+ GroupValuesDictionary::<Int16Type>::new(value_type),
+ )),
+ DataType::Int32 => Ok(Box::new(
+ GroupValuesDictionary::<Int32Type>::new(value_type),
+ )),
+ DataType::Int64 => Ok(Box::new(
+ GroupValuesDictionary::<Int64Type>::new(value_type),
+ )),
+ DataType::UInt8 => Ok(Box::new(
+ GroupValuesDictionary::<UInt8Type>::new(value_type),
+ )),
+ DataType::UInt16 => Ok(Box::new(
+ GroupValuesDictionary::<UInt16Type>::new(value_type),
+ )),
+ DataType::UInt32 => Ok(Box::new(
+ GroupValuesDictionary::<UInt32Type>::new(value_type),
+ )),
+ DataType::UInt64 => Ok(Box::new(
+ GroupValuesDictionary::<UInt64Type>::new(value_type),
+ )),
+ _ => Err(datafusion_common::DataFusionError::NotImplemented(
Review Comment:
If a dictionary's value type is supported but its key type somehow isn't, the code returns an error rather than handing the work to the generic path. Every real dictionary key type is covered, so this branch is unreachable today but a graceful fallback would be more robust than a hard error if Arrow ever adds a key type.
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,878 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, LargeStringArray, LargeStringBuilder,
+ ListArray, ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+/// Entry stored inside `unique_dict_value_mapping`. Caching `hash` on the entry
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ bytes: Vec<u8>,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ // stores the order new unique elements are seen for self.emit()
+ seen_elements: Vec<Option<Vec<u8>>>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ // key to group vector scratch space, used to avoid re-allocating a new vector on each call to intern
+ key_to_group: Vec<Option<usize>>,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ seen_elements: Vec::new(),
+ unique_dict_value_mapping: HashTable::with_capacity(INITIAL_PRE_ALLOCATION),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ key_to_group: Vec::new(),
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ /// Returns the existing `group_id` for the value with this hash and bytes,
+ /// or inserts a new entry and returns the freshly assigned group_id.
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ match self
+ .unique_dict_value_mapping
+ .entry(hash, |e| e.bytes == raw, |e| e.hash)
+ {
+ HashTableEntry::Occupied(o) => o.get().group_id,
+ HashTableEntry::Vacant(v) => {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(Some(raw.to_vec())); // replace this with raw buffer to avoid this double copy #TODO see https://github.com/apache/datafusion/issues/22078
+ v.insert(DictEntry {
+ hash,
+ group_id: new_group_id,
+ bytes: raw.to_vec(),
+ });
+ new_group_id
+ }
+ }
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(None);
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+ fn transform_into_array(&self, raw: &[Option<Vec<u8>>]) -> Result<ArrayRef> {
+ let item_capacity = raw.len();
+ let data_capacity = item_capacity * AVG_BYTES_PER_DICT_VALUE;
+ match &self.value_dt {
+ DataType::Utf8 => decode_scalar_string!(
+ raw,
+ StringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::LargeUtf8 => decode_scalar_string!(
+ raw,
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::Utf8View => decode_scalar_string!(
+ raw,
+ StringViewBuilder::with_capacity(item_capacity)
+ ),
+ DataType::List(field) => match field.data_type() {
+ DataType::Utf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ .with_field(Arc::clone(field))
+ ),
+ DataType::LargeUtf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ .with_field(Arc::clone(field))
+ ),
+ DataType::Utf8View => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringViewBuilder::with_capacity(item_capacity),
+ item_capacity,
+ )
+ .with_field(Arc::clone(field))
+ ),
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for List<{other:?}>"
+ ))),
+ },
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for {other:?}"
+ ))),
+ }
+ }
+}
+
+fn valid_bounds<K: ArrowDictionaryKeyType>(n: usize) -> bool {
+ let max: usize = 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,
+ _ => return false,
+ };
+ n <= max
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K> {
+ fn size(&self) -> usize {
+ let seen_elements_size = self.seen_elements.capacity()
+ * size_of::<Option<Vec<u8>>>()
+ + self
+ .seen_elements
+ .iter()
+ .filter_map(|opt| opt.as_ref())
+ .map(|inner| inner.capacity())
+ .sum::<usize>();
+
+ let unique_mapping_size = self.unique_dict_value_mapping.capacity()
+ * size_of::<DictEntry>()
+ + self
+ .unique_dict_value_mapping
+ .iter()
+ .map(|e| e.bytes.capacity())
+ .sum::<usize>();
+
+ let values_cache_size = self.values_cache.1.capacity() * size_of::<u64>()
+ + self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|a| a.to_data().get_array_memory_size())
+ .unwrap_or(0);
+
+ size_of::<Self>() + seen_elements_size + unique_mapping_size + values_cache_size
+ }
+ fn len(&self) -> usize {
+ self.seen_elements.len()
+ }
+ fn is_empty(&self) -> bool {
+ self.seen_elements.is_empty()
+ }
+ fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
+ assert_eq!(
+ cols.len(),
+ 1,
+ "GroupValuesDictionary only supports a single column"
+ );
+ let array = Arc::clone(&cols[0]);
+ groups.clear(); // zero out buffer
+ let dict_array = array
+ .as_dictionary_opt()
+ .map_or_else(|| Err(Internal("Expected dictionary array".into())), Ok)?;
+
+ let values = dict_array.values();
+ let key_array: &PrimitiveArray<K> = dict_array.keys();
+ if key_array.is_empty() {
+ return Ok(());
+ }
+
+ let cache_hit = self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|cached| Arc::ptr_eq(cached, values))
+ .unwrap_or(false);
+
+ if !cache_hit {
+ // values array changed since last batch - recompute hashes and update cached pointer
+ self.compute_value_hashes(values)?;
+ self.values_cache.0 = Some(Arc::clone(values));
+ }
+ // avoid re-allocating the key_to_group vector on each call to intern by re-using it as scratch space and only updating self.key_to_group at the end of the function.
+ let mut key_to_group = std::mem::take(&mut self.key_to_group);
+ key_to_group.clear();
+ key_to_group.resize(values.len(), None);
+
+ // iterate keys array (n iterations)
+ // only d insertions at most, repeated work is cached
+ for i in 0..key_array.len() {
+ let group_id = if key_array.is_null(i) {
+ self.get_null_group_id()
+ } else {
+ let key = key_array.value(i).to_usize().unwrap();
+ if let Some(group_id) = key_to_group[key] {
+ group_id
+ } else if values.is_null(key) {
+ let gid = self.get_null_group_id();
+ key_to_group[key] = Some(gid);
+ gid
+ } else {
+ let hash = self.values_cache.1[key];
+ let raw = Self::get_raw_bytes(values, key);
+ let gid = self.lookup_or_insert_in_table(hash, raw.as_ref());
+ key_to_group[key] = Some(gid);
+ gid
+ }
+ };
+ groups.push(group_id);
+ }
+ self.key_to_group = key_to_group;
+ Ok(())
+ }
+ fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
+ let (elements_to_emit, null_id) = match emit_to {
+ EmitTo::All => {
+ let original_null_id = self.null_group_id;
+ self.null_group_id = None;
+ self.unique_dict_value_mapping.clear();
+ (std::mem::take(&mut self.seen_elements), original_null_id)
+ }
+ EmitTo::First(n) => {
+ let n = n.min(self.seen_elements.len());
+ let first_n = self.seen_elements.drain(..n).collect::<Vec<_>>(); // again read from raw buffer directly
+ let original_null_id = self.null_group_id.filter(|&id| id < n);
+ // update null_group_id if the null group was in the first n
+ if let Some(null_id) = self.null_group_id {
+ if null_id < n {
+ self.null_group_id = None;
+ } else {
+ self.null_group_id = Some(null_id - n);
+ }
+ }
+ // shift all remaining group indices down by n in the table
+ self.unique_dict_value_mapping.retain(|e| {
+ if e.group_id < n {
+ false
+ } else {
+ e.group_id -= n;
+ true
+ }
+ });
+ (first_n, original_null_id)
+ }
+ };
+
+ let n = elements_to_emit.len();
+ if !valid_bounds::<K>(n) {
Review Comment:
guard is off by one,compare against one more than the largest index, i.e. allow the count when the largest index it will produce still fits the key type
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "kumarUjjawal (via GitHub)" <gi...@apache.org>.
kumarUjjawal commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4662199962
run benchmark dictionary_group_values
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on code in PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#discussion_r3365180358
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,768 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, LargeStringArray, LargeStringBuilder,
+ ListArray, ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+/// Entry stored inside `unique_dict_value_mapping`. Caching `hash` on the entry
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ bytes: Vec<u8>,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ // stores the order new unique elements are seen for self.emit()
+ seen_elements: Vec<Option<Vec<u8>>>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ intern_called: bool,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ seen_elements: Vec::new(),
+ unique_dict_value_mapping: HashTable::with_capacity(INITIAL_PRE_ALLOCATION),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ intern_called: false,
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ /// Returns the existing `group_id` for the value with this hash and bytes,
+ /// or inserts a new entry and returns the freshly assigned group_id.
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ match self
+ .unique_dict_value_mapping
+ .entry(hash, |e| e.bytes == raw, |e| e.hash)
+ {
+ HashTableEntry::Occupied(o) => o.get().group_id,
+ HashTableEntry::Vacant(v) => {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(Some(raw.to_vec())); // replace this with raw buffer to avoid this double copy #TODO see https://github.com/apache/datafusion/issues/22078
+ v.insert(DictEntry {
+ hash,
+ group_id: new_group_id,
+ bytes: raw.to_vec(),
+ });
+ new_group_id
+ }
+ }
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(None);
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+ fn transform_into_array(&self, raw: &[Option<Vec<u8>>]) -> Result<ArrayRef> {
+ let item_capacity = raw.len();
+ let data_capacity = item_capacity * AVG_BYTES_PER_DICT_VALUE;
+ match &self.value_dt {
+ DataType::Utf8 => decode_scalar_string!(
+ raw,
+ StringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::LargeUtf8 => decode_scalar_string!(
+ raw,
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::Utf8View => decode_scalar_string!(
+ raw,
+ StringViewBuilder::with_capacity(item_capacity)
+ ),
+ DataType::List(field) => match field.data_type() {
+ DataType::Utf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::LargeUtf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::Utf8View => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringViewBuilder::with_capacity(item_capacity),
+ item_capacity,
+ )
+ ),
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for List<{other:?}>"
+ ))),
+ },
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for {other:?}"
+ ))),
+ }
+ }
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K> {
+ fn size(&self) -> usize {
+ let seen_elements_size = self.seen_elements.capacity()
+ * size_of::<Option<Vec<u8>>>()
+ + self
+ .seen_elements
+ .iter()
+ .filter_map(|opt| opt.as_ref())
+ .map(|inner| inner.capacity())
+ .sum::<usize>();
+
+ let unique_mapping_size = self.unique_dict_value_mapping.capacity()
+ * size_of::<DictEntry>()
+ + self
+ .unique_dict_value_mapping
+ .iter()
+ .map(|e| e.bytes.capacity())
+ .sum::<usize>();
+
+ let values_cache_size = self.values_cache.1.capacity() * size_of::<u64>()
+ + self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|a| a.to_data().get_array_memory_size())
+ .unwrap_or(0);
+
+ size_of::<Self>() + seen_elements_size + unique_mapping_size + values_cache_size
+ }
+ fn len(&self) -> usize {
+ self.seen_elements.len()
+ }
+ fn is_empty(&self) -> bool {
+ self.seen_elements.is_empty()
+ }
+ fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
+ assert_eq!(
+ cols.len(),
+ 1,
+ "GroupValuesDictionary only supports a single column"
+ );
+ let array = Arc::clone(&cols[0]);
+ groups.clear(); // zero out buffer
+ let dict_array = array
+ .as_dictionary_opt()
+ .map_or_else(|| Err(Internal("Expected dictionary array".into())), Ok)?;
+
+ let values = dict_array.values();
+ let key_array: &PrimitiveArray<K> = dict_array.keys();
+ if key_array.is_empty() {
+ return Ok(());
+ }
+
+ let cache_hit = self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|cached| Arc::ptr_eq(cached, values))
+ .unwrap_or(false);
+
+ if !cache_hit {
+ // values array changed since last batch - recompute hashes and update cached pointer
+ self.compute_value_hashes(values)?;
+ self.values_cache.0 = Some(Arc::clone(values));
+ }
+ let mut key_to_group: Vec<Option<usize>> = vec![None; values.len()];
Review Comment:
make sense.
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,768 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, LargeStringArray, LargeStringBuilder,
+ ListArray, ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+/// Entry stored inside `unique_dict_value_mapping`. Caching `hash` on the entry
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ bytes: Vec<u8>,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ // stores the order new unique elements are seen for self.emit()
+ seen_elements: Vec<Option<Vec<u8>>>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ intern_called: bool,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ seen_elements: Vec::new(),
+ unique_dict_value_mapping: HashTable::with_capacity(INITIAL_PRE_ALLOCATION),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ intern_called: false,
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ /// Returns the existing `group_id` for the value with this hash and bytes,
+ /// or inserts a new entry and returns the freshly assigned group_id.
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ match self
+ .unique_dict_value_mapping
+ .entry(hash, |e| e.bytes == raw, |e| e.hash)
+ {
+ HashTableEntry::Occupied(o) => o.get().group_id,
+ HashTableEntry::Vacant(v) => {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(Some(raw.to_vec())); // replace this with raw buffer to avoid this double copy #TODO see https://github.com/apache/datafusion/issues/22078
+ v.insert(DictEntry {
+ hash,
+ group_id: new_group_id,
+ bytes: raw.to_vec(),
+ });
+ new_group_id
+ }
+ }
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(None);
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+ fn transform_into_array(&self, raw: &[Option<Vec<u8>>]) -> Result<ArrayRef> {
+ let item_capacity = raw.len();
+ let data_capacity = item_capacity * AVG_BYTES_PER_DICT_VALUE;
+ match &self.value_dt {
+ DataType::Utf8 => decode_scalar_string!(
+ raw,
+ StringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::LargeUtf8 => decode_scalar_string!(
+ raw,
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::Utf8View => decode_scalar_string!(
+ raw,
+ StringViewBuilder::with_capacity(item_capacity)
+ ),
+ DataType::List(field) => match field.data_type() {
+ DataType::Utf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::LargeUtf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::Utf8View => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringViewBuilder::with_capacity(item_capacity),
+ item_capacity,
+ )
+ ),
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for List<{other:?}>"
+ ))),
+ },
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for {other:?}"
+ ))),
+ }
+ }
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K> {
+ fn size(&self) -> usize {
+ let seen_elements_size = self.seen_elements.capacity()
+ * size_of::<Option<Vec<u8>>>()
+ + self
+ .seen_elements
+ .iter()
+ .filter_map(|opt| opt.as_ref())
+ .map(|inner| inner.capacity())
+ .sum::<usize>();
+
+ let unique_mapping_size = self.unique_dict_value_mapping.capacity()
+ * size_of::<DictEntry>()
+ + self
+ .unique_dict_value_mapping
+ .iter()
+ .map(|e| e.bytes.capacity())
+ .sum::<usize>();
+
+ let values_cache_size = self.values_cache.1.capacity() * size_of::<u64>()
+ + self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|a| a.to_data().get_array_memory_size())
+ .unwrap_or(0);
+
+ size_of::<Self>() + seen_elements_size + unique_mapping_size + values_cache_size
+ }
+ fn len(&self) -> usize {
+ self.seen_elements.len()
+ }
+ fn is_empty(&self) -> bool {
+ self.seen_elements.is_empty()
+ }
+ fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
+ assert_eq!(
+ cols.len(),
+ 1,
+ "GroupValuesDictionary only supports a single column"
+ );
+ let array = Arc::clone(&cols[0]);
+ groups.clear(); // zero out buffer
+ let dict_array = array
+ .as_dictionary_opt()
+ .map_or_else(|| Err(Internal("Expected dictionary array".into())), Ok)?;
+
+ let values = dict_array.values();
+ let key_array: &PrimitiveArray<K> = dict_array.keys();
+ if key_array.is_empty() {
+ return Ok(());
+ }
+
+ let cache_hit = self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|cached| Arc::ptr_eq(cached, values))
+ .unwrap_or(false);
+
+ if !cache_hit {
+ // values array changed since last batch - recompute hashes and update cached pointer
+ self.compute_value_hashes(values)?;
+ self.values_cache.0 = Some(Arc::clone(values));
+ }
+ let mut key_to_group: Vec<Option<usize>> = vec![None; values.len()];
Review Comment:
make sense
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on code in PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#discussion_r3358358656
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,783 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, DictionaryArray, LargeStringArray, LargeStringBuilder, ListArray,
+ ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+/// Entry stored inside `unique_dict_value_mapping`. Caching `hash` on the entry
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ bytes: Vec<u8>,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ // stores the order new unique elements are seen for self.emit()
+ seen_elements: Vec<Option<Vec<u8>>>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ intern_called: bool,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ seen_elements: Vec::new(),
+ unique_dict_value_mapping: HashTable::with_capacity(INITIAL_PRE_ALLOCATION),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ intern_called: false,
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ /// Returns the existing `group_id` for the value with this hash and bytes,
+ /// or inserts a new entry and returns the freshly assigned group_id.
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ match self
+ .unique_dict_value_mapping
+ .entry(hash, |e| e.bytes == raw, |e| e.hash)
+ {
+ HashTableEntry::Occupied(o) => o.get().group_id,
+ HashTableEntry::Vacant(v) => {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(Some(raw.to_vec())); // replace this with raw buffer to avoid this double copy #TODO see https://github.com/apache/datafusion/issues/22078
+ v.insert(DictEntry {
+ hash,
+ group_id: new_group_id,
+ bytes: raw.to_vec(),
+ });
+ new_group_id
+ }
+ }
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(None);
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+ fn transform_into_array(&self, raw: &[Option<Vec<u8>>]) -> Result<ArrayRef> {
+ let item_capacity = raw.len();
+ let data_capacity = item_capacity * AVG_BYTES_PER_DICT_VALUE;
+ match &self.value_dt {
+ DataType::Utf8 => decode_scalar_string!(
+ raw,
+ StringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::LargeUtf8 => decode_scalar_string!(
+ raw,
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::Utf8View => decode_scalar_string!(
+ raw,
+ StringViewBuilder::with_capacity(item_capacity)
+ ),
+ DataType::List(field) => match field.data_type() {
+ DataType::Utf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::LargeUtf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::Utf8View => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringViewBuilder::with_capacity(item_capacity),
+ item_capacity,
+ )
+ ),
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for List<{other:?}>"
+ ))),
+ },
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for {other:?}"
+ ))),
+ }
+ }
+ #[inline]
+ fn keys_to_usize(key_array: &PrimitiveArray<K>) -> Vec<Option<usize>> {
+ (0..key_array.len())
+ .map(|i| {
+ if key_array.is_null(i) {
+ None
+ } else {
+ Some(key_array.value(i).to_usize().unwrap())
+ }
+ })
+ .collect()
+ }
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K> {
+ fn size(&self) -> usize {
+ let seen_elements_size = self.seen_elements.capacity()
+ * size_of::<Option<Vec<u8>>>()
+ + self
+ .seen_elements
+ .iter()
+ .filter_map(|opt| opt.as_ref())
+ .map(|inner| inner.capacity())
+ .sum::<usize>();
+
+ let unique_mapping_size = self.unique_dict_value_mapping.capacity()
+ * size_of::<DictEntry>()
+ + self
+ .unique_dict_value_mapping
+ .iter()
+ .map(|e| e.bytes.capacity())
+ .sum::<usize>();
+
+ let values_cache_size = self.values_cache.1.capacity() * size_of::<u64>()
+ + self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|a| a.get_array_memory_size())
+ .unwrap_or(0);
+
+ size_of::<Self>() + seen_elements_size + unique_mapping_size + values_cache_size
+ }
+ fn len(&self) -> usize {
+ self.seen_elements.len()
+ }
+ fn is_empty(&self) -> bool {
+ self.seen_elements.is_empty()
+ }
+ fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
+ assert_eq!(
+ cols.len(),
+ 1,
+ "GroupValuesDictionary only supports a single column"
+ );
+ let array = Arc::clone(&cols[0]);
+ groups.clear(); // zero out buffer
+ let dict_array = array
+ .as_any()
+ .downcast_ref::<DictionaryArray<K>>()
+ .ok_or_else(|| {
+ Internal(format!(
+ "GroupValuesDictionary expected DictionaryArray but got {:?}",
+ array.data_type()
+ ))
+ })?;
+
+ let values = dict_array.values();
+ let key_array = dict_array.keys();
+ if key_array.is_empty() {
+ return Ok(());
+ }
+ let keys_as_usize: Vec<Option<usize>> = Self::keys_to_usize(key_array);
+
+ let cache_hit = self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|cached| Arc::ptr_eq(cached, values))
+ .unwrap_or(false);
+
+ if !cache_hit {
+ // first time seeing this values array to compute hashes, and cache the pointer, cheap operation
+ self.compute_value_hashes(values)?;
+ self.values_cache.0 = Some(Arc::clone(values));
+ }
+ let mut key_to_group: Vec<Option<usize>> = vec![None; values.len()];
+ if self.intern_called {
+ for key_opt in keys_as_usize.iter() {
+ let Some(original_key) = *key_opt else {
+ continue;
+ };
+ if values.is_null(original_key) {
+ continue;
+ }
+ if key_to_group[original_key].is_some() {
+ continue;
+ }
+
+ let hash = self.values_cache.1[original_key];
+ let raw = Self::get_raw_bytes(values, original_key);
+ if let Some(entry) = self
+ .unique_dict_value_mapping
+ .find(hash, |e| e.bytes == raw.as_ref())
+ {
+ key_to_group[original_key] = Some(entry.group_id);
+ continue;
+ }
+ }
+ }
+ // iterate keys array (n iterations)
+ // only d insertions at most, repeated work is cached
+ for key_opt in keys_as_usize.iter() {
Review Comment:
I agree, added comments that explain whats happening
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4635555334
Thank you for taking a look @kumarUjjawal, pushed up some revisions. Could you take another look? 🚀
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on code in PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#discussion_r3365051991
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,768 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, LargeStringArray, LargeStringBuilder,
+ ListArray, ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+/// Entry stored inside `unique_dict_value_mapping`. Caching `hash` on the entry
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ bytes: Vec<u8>,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ // stores the order new unique elements are seen for self.emit()
+ seen_elements: Vec<Option<Vec<u8>>>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ intern_called: bool,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ seen_elements: Vec::new(),
+ unique_dict_value_mapping: HashTable::with_capacity(INITIAL_PRE_ALLOCATION),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ intern_called: false,
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ /// Returns the existing `group_id` for the value with this hash and bytes,
+ /// or inserts a new entry and returns the freshly assigned group_id.
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ match self
+ .unique_dict_value_mapping
+ .entry(hash, |e| e.bytes == raw, |e| e.hash)
+ {
+ HashTableEntry::Occupied(o) => o.get().group_id,
+ HashTableEntry::Vacant(v) => {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(Some(raw.to_vec())); // replace this with raw buffer to avoid this double copy #TODO see https://github.com/apache/datafusion/issues/22078
+ v.insert(DictEntry {
+ hash,
+ group_id: new_group_id,
+ bytes: raw.to_vec(),
+ });
+ new_group_id
+ }
+ }
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(None);
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+ fn transform_into_array(&self, raw: &[Option<Vec<u8>>]) -> Result<ArrayRef> {
+ let item_capacity = raw.len();
+ let data_capacity = item_capacity * AVG_BYTES_PER_DICT_VALUE;
+ match &self.value_dt {
+ DataType::Utf8 => decode_scalar_string!(
+ raw,
+ StringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::LargeUtf8 => decode_scalar_string!(
+ raw,
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::Utf8View => decode_scalar_string!(
+ raw,
+ StringViewBuilder::with_capacity(item_capacity)
+ ),
+ DataType::List(field) => match field.data_type() {
+ DataType::Utf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::LargeUtf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::Utf8View => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringViewBuilder::with_capacity(item_capacity),
+ item_capacity,
+ )
+ ),
+ other => Err(NotImplemented(format!(
Review Comment:
good catch! updated the `transform_array` logic & added a regression test
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on code in PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#discussion_r3454584104
##########
datafusion/sqllogictest/test_files/dictionary.slt:
##########
@@ -515,3 +515,15 @@ DROP TABLE dict_hash_10;
statement ok
DROP TABLE dict_hash_src;
+
+query TI
+SELECT
Review Comment:
will do!
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on code in PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#discussion_r3454836051
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,1003 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, AsArray, DictionaryArray, LargeStringArray, LargeStringBuilder,
+ ListArray, ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ offset: usize,
+ len: usize,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ /// Packed byte storage for all group values.
+ row_buffer: Vec<u8>,
+ /// `row_offsets[g]` = start of group `g` in `row_buffer`;
+ row_offsets: Vec<usize>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ // key to group vector scratch space, used to avoid re-allocating a new vector on each call to intern
+ key_to_group: Vec<Option<usize>>,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ row_buffer: Vec::new(),
+ row_offsets: Vec::new(),
+ unique_dict_value_mapping: HashTable::new(),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ key_to_group: Vec::new(),
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ let row_buffer = &self.row_buffer;
+ if let Some(e) = self.unique_dict_value_mapping.find(hash, |e| {
+ e.hash == hash && &row_buffer[e.offset..e.offset + e.len] == raw
+ }) {
+ return e.group_id;
+ }
+ let new_group_id = self.row_offsets.len();
+ let offset = self.row_buffer.len();
+ self.row_offsets.push(offset);
+ self.row_buffer.extend_from_slice(raw);
+ self.unique_dict_value_mapping.insert_unique(
+ hash,
+ DictEntry {
+ hash,
+ group_id: new_group_id,
+ offset,
+ len: raw.len(),
+ },
+ |e| e.hash,
+ );
+ new_group_id
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.row_offsets.len();
+ self.row_offsets.push(self.row_buffer.len()); // null: empty span
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+
+ fn transform_into_array(
+ &self,
+ n: usize,
+ null_group_id: Option<usize>,
+ ) -> Result<ArrayRef> {
+ let data_capacity = n * AVG_BYTES_PER_DICT_VALUE;
+ let raw = (0..n).map(|i| {
+ if Some(i) == null_group_id {
+ return None;
+ }
+ let start = self.row_offsets[i];
+ // last group has no i+1 entry; its end is the buffer tail.
+ let end = self
+ .row_offsets
+ .get(i + 1)
+ .copied()
+ .unwrap_or(self.row_buffer.len());
+ Some(&self.row_buffer[start..end])
+ });
+ match &self.value_dt {
Review Comment:
thank you for the review @alamb
I think it makes sense to only support `Dictionary<_, Utf8/Utf8View>`, this would significantly simplify the code paths while delivering a large performance boost for production-like data shapes (low-cardinality group bys). Similar to other single-column group by specializations, it leaves a viable fallback that loses nothing when the fast path isn't hit.
With that in mind, @zhuqi-lucas's work looks very interesting -- I had a similar idea with #22891 and #21878. The issue was that multi-column group-bys on dictionaries fell through to `GroupValuesRows`, which was slow. I'm actually working on a multi-column dictionary group by in #22983, but I'm running into similar issues around handling multiple distinct types. For the multi-column case it makes more sense to build on @zhuqi-lucas's work, since the number of types that need to be supported explodes and there are diminishing returns to exploiting dictionary properties across multiple columns and intern calls ([key space cache](https://github.com/apache/datafusion/pull/22983/changes#diff-84ceab5d992f30b340d52094b28685b5ecfd6f68f27867641ca1a9453ad279baR209), [arc ptr cache](https://github.com/apache/datafusion/pull/22983/changes#diff-84ceab5d992f30b340d52094b28685b5ecfd6f68f27867641ca1a9453ad279baR193), [compute keys once and reuse](https://github.com/apache/datafusion/pull/22983/c
hanges#diff-84ceab5d992f30b340d52094b28685b5ecfd6f68f27867641ca1a9453ad279baR342)).
As for this PR specifically, I do think it adds significant value. Would restricting the supported value types be a viable path forward? It would keep the type combinations from exploding while still letting the community benefit from the speedup.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4724015717
@alamb @adriangbot failed
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4723878833
Benchmark for [this request](https://github.com/apache/datafusion/pull/21765#issuecomment-4723859187) failed.
Last 20 lines of output:
<details><summary>Click to expand</summary>
```
* [new tag] 48.0.0-rc1 -> 48.0.0-rc1
* [new tag] 48.0.0-rc2 -> 48.0.0-rc2
* [new tag] 5.0.0 -> 5.0.0
* [new tag] 5.0.0-rc1 -> 5.0.0-rc1
* [new tag] 5.0.0-rc3 -> 5.0.0-rc3
* [new tag] 6.0.0 -> 6.0.0
* [new tag] 6.0.0-rc0 -> 6.0.0-rc0
* [new tag] 7.0.0 -> 7.0.0
* [new tag] 7.0.0-rc2 -> 7.0.0-rc2
* [new tag] 8.0.0 -> 8.0.0
* [new tag] 8.0.0-rc1 -> 8.0.0-rc1
* [new tag] 8.0.0-rc2 -> 8.0.0-rc2
* [new tag] 9.0.0 -> 9.0.0
* [new tag] 9.0.0-rc1 -> 9.0.0-rc1
* [new tag] ballista-0.5.0 -> ballista-0.5.0
* [new tag] ballista-0.6.0 -> ballista-0.6.0
* [new tag] ballista-0.7.0 -> ballista-0.7.0
* [new tag] python-0.3.0 -> python-0.3.0
* [new tag] python-0.4.0 -> python-0.4.0
Switched to branch 'rich-t-kid/optimize-dictionary-grouping'
```
</details>
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4723859187
run benchmark dictionary_group_values
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4730296213
@kumarUjjawal @alamb with #22078's performance gains, there are now no regressions even in the worst case 🚀
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4299983695
@adriangb these benchmarks reflect what my design doc mentions the more rows/greater scale the better it should perform. these were all being run on a mac-book maybe running it with the ```run benchmarks``` command may yield slightly different results
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4293367939
### In memory s1
<img width="1287" height="668" alt="Image 4-21-26 at 11 45 PM" src="https://github.com/user-attachments/assets/b057ab83-198b-4f6f-b5f3-ecc79a0a3b84" />
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4290601724
@alamb this PR should resolve the initial regression that was caused [(described here)](https://github.com/apache/datafusion/issues/8738) as well as provide performance boost. I would love to hear your thoughts this approach
> I think we need to create a benchmark that does aggregation queries on dictionary encoded string columns (I know the existing end to end TPCH, and ClickBench benchmarks do not do this). I will file a ticket shortly about this
[source](https://github.com/apache/datafusion/issues/7647#issuecomment-1880919163)
seeing as how this PR is already large, I think this would be a nice follow up PR since as far as i know this work was never finished since #9017 was closed due to inactivity
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4293541622
Pre-allocating the values buffer causes a regression
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4301684225
Benchmark for [this request](https://github.com/apache/datafusion/pull/21765#issuecomment-4301676571) failed.
Last 20 lines of output:
<details><summary>Click to expand</summary>
```
* [new tag] 48.0.0-rc1 -> 48.0.0-rc1
* [new tag] 48.0.0-rc2 -> 48.0.0-rc2
* [new tag] 5.0.0 -> 5.0.0
* [new tag] 5.0.0-rc1 -> 5.0.0-rc1
* [new tag] 5.0.0-rc3 -> 5.0.0-rc3
* [new tag] 6.0.0 -> 6.0.0
* [new tag] 6.0.0-rc0 -> 6.0.0-rc0
* [new tag] 7.0.0 -> 7.0.0
* [new tag] 7.0.0-rc2 -> 7.0.0-rc2
* [new tag] 8.0.0 -> 8.0.0
* [new tag] 8.0.0-rc1 -> 8.0.0-rc1
* [new tag] 8.0.0-rc2 -> 8.0.0-rc2
* [new tag] 9.0.0 -> 9.0.0
* [new tag] 9.0.0-rc1 -> 9.0.0-rc1
* [new tag] ballista-0.5.0 -> ballista-0.5.0
* [new tag] ballista-0.6.0 -> ballista-0.6.0
* [new tag] ballista-0.7.0 -> ballista-0.7.0
* [new tag] python-0.3.0 -> python-0.3.0
* [new tag] python-0.4.0 -> python-0.4.0
Switched to branch 'rich-t-kid/optimize-dictionary-grouping'
```
</details>
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4322753440
Yea it seems that no benchmarks currently perform a group by using dictionary encodeded columns. I created this PR to address that. @Dandandan could you please take a look? https://github.com/apache/datafusion/pull/21860
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4490852547
@alamb @Dandandan just wanted to bump this PR.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4771233445
FYI @zhuqi-lucas -- perhaps is this part of your work to complete type support of Dictionaries?
- https://github.com/apache/datafusion/issues/22715
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on code in PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#discussion_r3358263849
##########
datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs:
##########
@@ -0,0 +1,783 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::GroupValues;
+use crate::hash_utils::RandomState;
+use arrow::array::{
+ Array, ArrayRef, DictionaryArray, LargeStringArray, LargeStringBuilder, ListArray,
+ ListBuilder, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StringViewArray, StringViewBuilder,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use datafusion_common::DataFusionError::{Internal, NotImplemented};
+use datafusion_common::Result;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_expr::EmitTo;
+use hashbrown::HashTable;
+use hashbrown::hash_table::Entry as HashTableEntry;
+use std::borrow::Cow;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// Heuristic for sizing the values buffer of string builders during emit:
+/// dictionary-encoded values are short by design (categorical strings, short
+/// identifiers), so 16 B/item avoids the realloc-doubling chain in the common
+/// case while keeping over-allocation cheap when values are smaller.
+const AVG_BYTES_PER_DICT_VALUE: usize = 16;
+const INITIAL_PRE_ALLOCATION: usize = 8 * 1024; // avoid re-allocation`s for small-medium groups
+
+macro_rules! decode_list {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ None => builder.append_null(),
+ Some(raw_vector) => {
+ let mut offset = 0;
+ while offset < raw_vector.len() {
+ let len = i64::from_ne_bytes(
+ raw_vector[offset..offset + 8]
+ .try_into()
+ .expect("slice of length 8"),
+ );
+ offset += 8;
+ if len == -1 {
+ builder.values().append_null();
+ } else {
+ let s = unsafe {
+ std::str::from_utf8_unchecked(
+ &raw_vector[offset..offset + len as usize],
+ )
+ };
+ builder.values().append_value(s);
+ offset += len as usize;
+ }
+ }
+ builder.append(true);
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+macro_rules! decode_scalar_string {
+ ($raw:expr, $builder:expr) => {{
+ let mut builder = $builder;
+ for raw_bytes in $raw {
+ match raw_bytes {
+ Some(raw_vector) => {
+ let s = unsafe { std::str::from_utf8_unchecked(raw_vector) };
+ builder.append_value(s);
+ }
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+/// Entry stored inside `unique_dict_value_mapping`. Caching `hash` on the entry
+struct DictEntry {
+ hash: u64,
+ group_id: usize,
+ bytes: Vec<u8>,
+}
+
+pub struct GroupValuesDictionary<K: ArrowDictionaryKeyType + Send> {
+ // stores the order new unique elements are seen for self.emit()
+ seen_elements: Vec<Option<Vec<u8>>>,
+ value_dt: DataType,
+ _phantom: PhantomData<K>,
+ // keeps track of which values weve already seen, keyed by raw value hash.
+ unique_dict_value_mapping: HashTable<DictEntry>,
+
+ random_state: RandomState,
+
+ // cache the group id for nulls since they all map to the same group
+ null_group_id: Option<usize>,
+ intern_called: bool,
+ // 0. cache pointer of arrays, this avoids having to re-compute hashing for arrays weve already seen on past iterations
+ // 1. avoid re-allocating buffer inbetween calls, instead of allocating a new vector each time re-use inbetween calls
+ values_cache: (Option<ArrayRef>, Vec<u64>),
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
+ pub fn new(data_type: &DataType) -> Self {
+ Self {
+ seen_elements: Vec::new(),
+ unique_dict_value_mapping: HashTable::with_capacity(INITIAL_PRE_ALLOCATION),
+ value_dt: data_type.clone(),
+ _phantom: PhantomData,
+ random_state: RandomState::with_seed(0),
+ null_group_id: None,
+ intern_called: false,
+ values_cache: (None, Vec::new()),
+ }
+ }
+
+ /// Returns the existing `group_id` for the value with this hash and bytes,
+ /// or inserts a new entry and returns the freshly assigned group_id.
+ fn lookup_or_insert_in_table(&mut self, hash: u64, raw: &[u8]) -> usize {
+ match self
+ .unique_dict_value_mapping
+ .entry(hash, |e| e.bytes == raw, |e| e.hash)
+ {
+ HashTableEntry::Occupied(o) => o.get().group_id,
+ HashTableEntry::Vacant(v) => {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(Some(raw.to_vec())); // replace this with raw buffer to avoid this double copy #TODO see https://github.com/apache/datafusion/issues/22078
+ v.insert(DictEntry {
+ hash,
+ group_id: new_group_id,
+ bytes: raw.to_vec(),
+ });
+ new_group_id
+ }
+ }
+ }
+ fn compute_value_hashes(&mut self, values: &ArrayRef) -> Result<()> {
+ self.values_cache.1.clear();
+ self.values_cache.1.resize(values.len(), 0);
+ create_hashes(
+ [Arc::clone(values)],
+ &self.random_state,
+ &mut self.values_cache.1,
+ )?;
+ Ok(())
+ //Ok(hashes)
+ }
+
+ fn get_raw_bytes(values: &ArrayRef, index: usize) -> Cow<'_, [u8]> {
+ match values.data_type() {
+ DataType::Utf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("Expected StringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::LargeUtf8 => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .expect("Expected LargeStringArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::Utf8View => Cow::Borrowed(
+ values
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .expect("Expected StringViewArray")
+ .value(index)
+ .as_bytes(),
+ ),
+ DataType::List(_) => {
+ let list_array = values
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("Expected ListArray");
+
+ debug_assert!(!list_array.is_null(index));
+
+ let start = list_array.value_offsets()[index] as usize;
+ let end = list_array.value_offsets()[index + 1] as usize;
+ let child = list_array.values();
+
+ let mut bytes = Vec::new();
+ for i in start..end {
+ if child.is_null(i) {
+ // acts as a marker for transform_into_array to write a null
+ bytes.extend_from_slice(&(-1i64).to_ne_bytes());
+ } else {
+ let raw = Self::get_raw_bytes(child, i);
+ bytes.extend_from_slice(&(raw.len() as i64).to_ne_bytes());
+ bytes.extend_from_slice(&raw);
+ }
+ }
+ Cow::Owned(bytes)
+ }
+ other => unimplemented!("get_raw_bytes not implemented for {other:?}"),
+ }
+ }
+
+ #[inline]
+ fn get_null_group_id(&mut self) -> usize {
+ if let Some(group_id) = self.null_group_id {
+ group_id
+ } else {
+ let new_group_id = self.seen_elements.len();
+ self.seen_elements.push(None);
+ self.null_group_id = Some(new_group_id);
+ new_group_id
+ }
+ }
+ fn transform_into_array(&self, raw: &[Option<Vec<u8>>]) -> Result<ArrayRef> {
+ let item_capacity = raw.len();
+ let data_capacity = item_capacity * AVG_BYTES_PER_DICT_VALUE;
+ match &self.value_dt {
+ DataType::Utf8 => decode_scalar_string!(
+ raw,
+ StringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::LargeUtf8 => decode_scalar_string!(
+ raw,
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity)
+ ),
+ DataType::Utf8View => decode_scalar_string!(
+ raw,
+ StringViewBuilder::with_capacity(item_capacity)
+ ),
+ DataType::List(field) => match field.data_type() {
+ DataType::Utf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::LargeUtf8 => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ LargeStringBuilder::with_capacity(item_capacity, data_capacity),
+ item_capacity,
+ )
+ ),
+ DataType::Utf8View => decode_list!(
+ raw,
+ ListBuilder::with_capacity(
+ StringViewBuilder::with_capacity(item_capacity),
+ item_capacity,
+ )
+ ),
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for List<{other:?}>"
+ ))),
+ },
+ other => Err(NotImplemented(format!(
+ "transform_into_array not implemented for {other:?}"
+ ))),
+ }
+ }
+ #[inline]
+ fn keys_to_usize(key_array: &PrimitiveArray<K>) -> Vec<Option<usize>> {
+ (0..key_array.len())
+ .map(|i| {
+ if key_array.is_null(i) {
+ None
+ } else {
+ Some(key_array.value(i).to_usize().unwrap())
+ }
+ })
+ .collect()
+ }
+}
+
+impl<K: ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K> {
+ fn size(&self) -> usize {
+ let seen_elements_size = self.seen_elements.capacity()
+ * size_of::<Option<Vec<u8>>>()
+ + self
+ .seen_elements
+ .iter()
+ .filter_map(|opt| opt.as_ref())
+ .map(|inner| inner.capacity())
+ .sum::<usize>();
+
+ let unique_mapping_size = self.unique_dict_value_mapping.capacity()
+ * size_of::<DictEntry>()
+ + self
+ .unique_dict_value_mapping
+ .iter()
+ .map(|e| e.bytes.capacity())
+ .sum::<usize>();
+
+ let values_cache_size = self.values_cache.1.capacity() * size_of::<u64>()
+ + self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|a| a.get_array_memory_size())
+ .unwrap_or(0);
+
+ size_of::<Self>() + seen_elements_size + unique_mapping_size + values_cache_size
+ }
+ fn len(&self) -> usize {
+ self.seen_elements.len()
+ }
+ fn is_empty(&self) -> bool {
+ self.seen_elements.is_empty()
+ }
+ fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
+ assert_eq!(
+ cols.len(),
+ 1,
+ "GroupValuesDictionary only supports a single column"
+ );
+ let array = Arc::clone(&cols[0]);
+ groups.clear(); // zero out buffer
+ let dict_array = array
+ .as_any()
+ .downcast_ref::<DictionaryArray<K>>()
+ .ok_or_else(|| {
+ Internal(format!(
+ "GroupValuesDictionary expected DictionaryArray but got {:?}",
+ array.data_type()
+ ))
+ })?;
+
+ let values = dict_array.values();
+ let key_array = dict_array.keys();
+ if key_array.is_empty() {
+ return Ok(());
+ }
+ let keys_as_usize: Vec<Option<usize>> = Self::keys_to_usize(key_array);
+
+ let cache_hit = self
+ .values_cache
+ .0
+ .as_ref()
+ .map(|cached| Arc::ptr_eq(cached, values))
+ .unwrap_or(false);
+
+ if !cache_hit {
+ // first time seeing this values array to compute hashes, and cache the pointer, cheap operation
Review Comment:
sounds good!
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "kumarUjjawal (via GitHub)" <gi...@apache.org>.
kumarUjjawal commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4646029425
@Rich-T-kid Left some comments, I was also thinking we should add some tests like
- test for de-duplicating the same value across two batches with different value buffers
- coverage of the string-view value type
- end-to-end SQL test
- test for the streaming/ordered emit path
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "kumarUjjawal (via GitHub)" <gi...@apache.org>.
kumarUjjawal commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4645670621
run benchmark dictionary_group_values
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "kumarUjjawal (via GitHub)" <gi...@apache.org>.
kumarUjjawal commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4411196800
It shouldn't, I have no idea what's wrong 🤷🏽♂️
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4411562626
Could you try running the command again?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "kumarUjjawal (via GitHub)" <gi...@apache.org>.
kumarUjjawal commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4418355614
run benchmarks
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4423145724
<img width="908" height="367" alt="Image 5-11-26 at 1 24 PM" src="https://github.com/user-attachments/assets/9db4d248-293b-4795-ae68-14258b5d3ab9" />
Results of `/benchmarks/bench.sh run dict` command compared between my branch and main
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-4726181248
🤖 Benchmark running (GKE) | [trigger](https://github.com/apache/datafusion/pull/21765#issuecomment-4726105593)
**Instance:** `c4a-highmem-16` (12 vCPU / 65 GiB) | `Linux bench-c4726105593-579-2xdsl 6.12.68+ #1 SMP Sat May 2 07:49:07 UTC 2026 aarch64 GNU/Linux`
<details><summary>CPU Details (lscpu)</summary>
```
Architecture: aarch64
CPU op-mode(s): 64-bit
Byte Order: Little Endian
CPU(s): 16
On-line CPU(s) list: 0-15
Vendor ID: ARM
Model name: Neoverse-V2
Model: 1
Thread(s) per core: 1
Core(s) per cluster: 16
Socket(s): -
Cluster(s): 1
Stepping: r0p1
BogoMIPS: 2000.00
Flags: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache: 1 MiB (16 instances)
L1i cache: 1 MiB (16 instances)
L2 cache: 32 MiB (16 instances)
L3 cache: 80 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; __user pointer sanitization
Vulnerability Spectre v2: Mitigation; CSV2, BHB
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected
```
</details>
Comparing rich-t-kid/optimize-dictionary-grouping (d440824b6c605cd0e28308d2ac0442d81f115922) to 96a6096 (merge-base) [diff](https://github.com/apache/datafusion/compare/96a6096c6f4b924e8cab4bc1629759a948e12939..d440824b6c605cd0e28308d2ac0442d81f115922) using: dictionary_group_values
Results will be posted here when complete
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Optimize Dictionary groupings [datafusion]
Posted by "Rich-T-kid (via GitHub)" <gi...@apache.org>.
Rich-T-kid commented on PR #21765:
URL: https://github.com/apache/datafusion/pull/21765#issuecomment-5307936800
#23187 closes this
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org