Feat: add dictionaries as a supported group column type - #23187
Open
Rich-T-kid wants to merge 8 commits into
Open
Feat: add dictionaries as a supported group column type#23187Rich-T-kid wants to merge 8 commits into
Rich-T-kid wants to merge 8 commits into
Conversation
Rich-T-kid
force-pushed
the
rich-T-kid/dictionary-groupValuesColumn-impl
branch
2 times, most recently
from
July 1, 2026 04:25
3f7ff57 to
e6b6dce
Compare
Contributor
Author
|
@kumarUjjawal could you run the dictionary benchmarks on this PR? Thx |
Rich-T-kid
commented
Jul 1, 2026
| } | ||
| } | ||
| DataType::Dictionary(key_dt, value_dt) => { | ||
| let new_field = Field::new("", *value_dt.clone(), true); |
Contributor
Author
There was a problem hiding this comment.
Since this field is never read again it may be fine to ignore the name field.
should be weary of similar issues to #21765 (comment)
Contributor
There was a problem hiding this comment.
Kind of annoying that make_group_column takes a field instead of a DataType. Maybe we can change that in a follow up PR?
Rich-T-kid
commented
Jul 1, 2026
Rich-T-kid
commented
Jul 1, 2026
Contributor
Author
|
@kumarUjjawal wanted to bump this 😄 |
Contributor
|
run benchmark dictionary_group_values |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/dictionary-groupValuesColumn-impl (eb41915) to 01bf68c (merge-base) diff using: dictionary_group_values File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagedictionary_group_values — base (merge-base)
dictionary_group_values — branch
File an issue against this benchmark runner |
Rich-T-kid
force-pushed
the
rich-T-kid/dictionary-groupValuesColumn-impl
branch
2 times, most recently
from
July 1, 2026 15:32
770abfe to
243a557
Compare
Rich-T-kid
changed the title
Contributor
Author
|
@codex review |
Contributor
@Rich-T-kid Thank you! I have been sick so I won't be available for review. I will probably get back next week. |
Contributor
Author
@kumarUjjawal Sorry to hear that. I hope you feel better! no rush on the review! |
Rich-T-kid
commented
Jul 2, 2026
Rich-T-kid
force-pushed
the
rich-T-kid/dictionary-groupValuesColumn-impl
branch
from
July 2, 2026 19:18
4ee52da to
7af7080
Compare
Rich-T-kid
commented
Jul 4, 2026
Rich-T-kid
force-pushed
the
rich-T-kid/dictionary-groupValuesColumn-impl
branch
from
July 5, 2026 03:59
152c1f0 to
f3387c5
Compare
Contributor
Author
Latest PR speeds up the low cardinality case
|
Contributor
Author
|
@geoffreyclaude could you run the benchmarks command again when you get a chance. Thanks 🚀 |
Contributor
|
run benchmark dictionary_group_values |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/dictionary-groupValuesColumn-impl (3d1e1c9) to 01bf68c (merge-base) diff using: dictionary_group_values File an issue against this benchmark runner |
Rich-T-kid
commented
Jul 6, 2026
Rich-T-kid
left a comment
Rich-T-kid
left a comment
Contributor
Author
There was a problem hiding this comment.
Revision #4 Last revision before making this open for review
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagedictionary_group_values — base (merge-base)
dictionary_group_values — branch
File an issue against this benchmark runner |
Contributor
Author
|
🤔 benchmarks show good improvement in every case but it should be much larger. Currently every |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/dictionary-groupValuesColumn-impl (c98c47f) to 30ae8bf (merge-base) diff Run configurationrun benchmark dictionary_group_valuesResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing rich-T-kid/dictionary-groupValuesColumn-impl (c98c47f) to 30ae8bf (merge-base) diff Run configurationrun benchmark dictionary_group_valuesCPU Details (lscpu)Details
Resource Usagedictionary_group_values — base (merge-base)
dictionary_group_values — branch
File an issue against this benchmark runner |
zhuqi-lucas
approved these changes
Jul 31, 2026
| } | ||
|
|
||
| fn take_n(&mut self, n: usize) -> ArrayRef { | ||
| let old_inner_len = self.inner.len(); |
Contributor
There was a problem hiding this comment.
Every partial emit still drains, hashes, and rebuilds all surviving values, retaining the O(G² / batch_size) cost and high peak memory. Also, Arrow take retains full backing storage for Utf8View, BinaryView, and nested dictionaries.
Contributor
Author
There was a problem hiding this comment.
Agreed, this isn't the most performant way to implement take_n() addressing them properly would require either:
- Lazy compaction with a base-offset scheme (defer the rebuild until the surviving set actually shrinks meaningfully), or
- Tracking which inner slots are still live without a full drain-and-rehash cycle
Both touch internals in a very meaningful way, and we currently have no benchmark that targets partial-emit patterns, so there's no good baseline to validate against. I'd rather land this with correct behaviour and open a follow-up issue to tackle the performance side with proper benchmarks in place. If that sounds good I can create the issue or add onto what @kosiew wrote in #23993
add schema support for dictionarys introduce high level GroupValuesColumn test introduce groupColumn trait test git issues introduce edge case/ regression section inital impl working implementation of dictionary for groupValuesCOlumns benchmarks show perf boost over groupvaluerows, TODO:dedupe items before inner append fix clippy errors & inline final builder add cache for arc ptr trim down test trim test LOC again trim PR revision 3 speed up low cardinlaity case working version introduce inter-batch caching break complex types into seperate parts fixed breaking test, re-allocate hashtable on each intern() call re-introduce cache remove mutex add cache to concat pointers to avoid un-needed allocations remove ptr caches and concat call reduce LOC revised PR comments add regression test to align with GroupValueRows & re-order overflow check add test to assert de-duplicated output dictionary tmp low card speed up optimize low-card case wip re-use allocations across calls add test final clean up
Rich-T-kid
force-pushed
the
rich-T-kid/dictionary-groupValuesColumn-impl
branch
from
August 4, 2026 03:45
c98c47f to
41192e5
Compare
Rich-T-kid
force-pushed
the
rich-T-kid/dictionary-groupValuesColumn-impl
branch
from
August 4, 2026 03:49
41192e5 to
44ad7d1
Compare
Contributor
Author
|
commit 44ad7d1 should address all of your comments except for the |
Contributor
|
run benchmark dictionary_group_values |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/dictionary-groupValuesColumn-impl (63d4021) to 47ca490 (merge-base) diff Run configurationrun benchmark dictionary_group_valuesResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing rich-T-kid/dictionary-groupValuesColumn-impl (63d4021) to 47ca490 (merge-base) diff Run configurationrun benchmark dictionary_group_valuesCPU Details (lscpu)Details
Resource Usagedictionary_group_values — base (merge-base)
dictionary_group_values — branch
File an issue against this benchmark runner |
kumarUjjawal
approved these changes
Aug 4, 2026
kumarUjjawal
left a comment
kumarUjjawal
left a comment
Contributor
There was a problem hiding this comment.
Thank you @Rich-T-kid
This looks good. Please create a separate issue for the optimization as you mentioned.
Contributor
Author
|
created a follow up issue here #24089 thank you for the reviews @kumarUjjawal @zhuqi-lucas @jayshrivastava 🚀 |
Contributor
|
I will merge this in 1-2 days in no one else has any objections. |
alamb
reviewed
Aug 5, 2026
alamb
left a comment
alamb
left a comment
Contributor
There was a problem hiding this comment.
Thank you for this @Rich-T-kid and @kumarUjjawal
I think it would be really great to show that this works end to end (and improves performance) -- I suspect it does but it would be good to verify
One way we could do this is run TPCH Q1 https://github.com/apache/datafusion/blob/main/benchmarks/queries/q1.sql
You could test this locally like
Make data
tpcgen-cli tpch parquet --scale-factor=10 --tables=lineitem
lineitem [==================] (100%)
Then convert the relevant columns to dictioanry
> COPY (select arrow_cast(l_returnflag, 'Dictionary(Int32, Utf8)') as l_returnflag, arrow_cast(l_linestatus, 'Dictionary(Int32, Utf8)') as l_linestatus, l_quantity, l_extendedprice, l_discount, l_shipdate, l_tax from 'lineitem.parquet') to 'lineitem_dict.parquet';
+----------+
| count |
+----------+
| 59986052 |
+----------+
1 row(s) fetched.
Elapsed 1.767 seconds.
Then run query
> select
l_returnflag,
l_linestatus,
sum(l_quantity) as sum_qty,
sum(l_extendedprice) as sum_base_price,
sum(l_extendedprice * (1 - l_discount)) as sum_disc_price,
sum(l_extendedprice * (1 - l_discount) * (1 + l_tax)) as sum_charge,
avg(l_quantity) as avg_qty,
avg(l_extendedprice) as avg_price,
avg(l_discount) as avg_disc,
count(*) as count_order
from
'lineitem_dict.parquet'
where
l_shipdate <= date '1998-09-02'
group by
l_returnflag,
l_linestatus
order by
l_returnflag,
l_linestatus;
+--------------+--------------+--------------+------------------+--------------------+----------------------+-----------+--------------+----------+-------------+
| l_returnflag | l_linestatus | sum_qty | sum_base_price | sum_disc_price | sum_charge | avg_qty | avg_price | avg_disc | count_order |
+--------------+--------------+--------------+------------------+--------------------+----------------------+-----------+--------------+----------+-------------+
| A | F | 377518399.00 | 566065727797.25 | 537759104278.0656 | 559276670892.116819 | 25.500975 | 38237.151008 | 0.050006 | 14804077 |
| N | F | 9851614.00 | 14767438399.17 | 14028805792.2114 | 14590490998.366737 | 25.522448 | 38257.810660 | 0.049973 | 385998 |
| N | O | 743124873.00 | 1114302286901.88 | 1058580922144.9638 | 1100937000170.591854 | 25.498075 | 38233.902923 | 0.050000 | 29144351 |
| R | F | 377732830.00 | 566431054976.00 | 538110922664.7677 | 559634780885.086257 | 25.508384 | 38251.219273 | 0.049996 | 14808183 |
+--------------+--------------+--------------+------------------+--------------------+----------------------+-----------+--------------+----------+-------------+
4 row(s) fetched.
Elapsed 0.370 seconds.| } | ||
| }; | ||
|
|
||
| // TODO: Materialize dictionaries in group keys (#7647) |
Contributor
|
Another thing that might be valuable is to make a variant of the tpch benchmarks, where you tell the parquet reader that the low cardinality columns (that would be good for dictionary encoding) should be read as Dictionary... Arrays -- and then run the queries and see how much improvement this PR gives I actually think that would be helpful in general if your goal is to improve the overall performance of the dictionaries in DataFusion. I will file a ticket with more detail |
Contributor
Author
|
@alamb I think the second approach is easier to work with and less error-prone. But from my understanding, this isn't currently possible, or at least isn't exposed in DataFusion. The only config that currently exists is for writing: https://docs.rs/datafusion/latest/datafusion/config/struct.ParquetOptions.html#structfield.dictionary_enabled
I think it'd be nice to expose this as a session prop so users can configure it with SET notation in SQL: The only caveat that comes to mind is this may cause the Parquet reader to read all string columns into dictionary arrays, which isn't what we want for precision. the concept your referring to is actually quite similar to what I raised in this issue: #22891. Ideally we would be able to determine which columns are low-cardinality and change their type dictionary at physical planning time. with #23187 this should provide a nice perf boost. the only issue would be determining what threshold counts as "low cardinality". cc @asolimando |
Contributor
Author
|
So I think some nice follow ups would be
by hand we can experiment with transforming low-card string group keys into |
Contributor
Author
|
I ran TPCH Q4 against main and these were my results main cast the columns to dictionary using this branch |
This was referenced Aug 5, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
8 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Rationale for this change
This PR introduces a specialized
GroupColumnimplementation for dictionary-typed columns insideGroupValuesColumn, allowing dictionary columns to participate in the columnar, vectorized aggregation path instead of the row-based fallback.The Implementation is only about 175+ lines of code. the remaining LOC is adding extensive test at the
GroupColumntrait level as well as testing theGroupValuesColumnGroupValues trait and how it inter-opts with multi-dictionary group by's.What changes are included in this PR?
DictionaryGroupValueBuilderstruct implementing theGroupColumntrait forDictionary-typed group-by columns, supporting a configurable subset of value typesGroupValuesColumn::try_new(thematches!block) to acceptDictionary(_, value_type)wherevalue_typeis already supported.emitAre these changes tested?
yes. a majority of this PR is test
Are there any user-facing changes?
no. this is a pure perf boost for users.