Feat: add dictionaries as a supported group column type#23187
Open
Rich-T-kid wants to merge 27 commits into
Open
Feat: add dictionaries as a supported group column type#23187Rich-T-kid wants to merge 27 commits into
Rich-T-kid wants to merge 27 commits into
Conversation
Rich-T-kid
force-pushed
the
rich-T-kid/dictionary-groupValuesColumn-impl
branch
from
June 30, 2026 20:23
2b60132 to
1ffc7f0
Compare
Rich-T-kid
force-pushed
the
rich-T-kid/dictionary-groupValuesColumn-impl
branch
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
from
July 1, 2026 15:30
eb41915 to
770abfe
Compare
Rich-T-kid
force-pushed
the
rich-T-kid/dictionary-groupValuesColumn-impl
branch
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
|
I believe this PR is ready for a maintainer to take a look 👍 |
Contributor
|
run benchmark dictionary_group_values |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/dictionary-groupValuesColumn-impl (fb7856f) 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 |
| } | ||
| } | ||
|
|
||
| fn into_dict(values: ArrayRef) -> ArrayRef { |
Contributor
There was a problem hiding this comment.
The emitted DictionaryArray uses identity keys (key[i] = i), so output is fully expanded with no value sharing. Could you add a short doc-comment explaining this?
Contributor
There was a problem hiding this comment.
valid_bounds is guarded with assert! which panics. A Dictionary(Int8, Utf8) column with >127 distinct groups would cause a panic in production rather than a graceful error. All other error paths here use Result. converting into_dict to -> Result and returning internal_err! on overflow, or checking bounds at intern time so build / take_n stay infallible.
Contributor
Author
There was a problem hiding this comment.
valid_bounds is guarded with assert! which panics. A Dictionary(Int8, Utf8) column with >127 distinct groups would cause a panic in production rather than a graceful error.
This is already an issue with GroupValueRows, I created this issue to track it #23127
a Dictionary(Int8, Utf8) with > 127 groups cannot be represented using the int8, this means we would be forced to bump the key type to the next support range being int16. The issue with bumping the key type is the output schema of the dictionary must match the input schema. Changing this will cause panics for downstream operators.
references : #7647
I can include a check in append_val()/vectorized_append_val() to return an error if this operation would overflow the key type.
Contributor
Author
There was a problem hiding this comment.
I'll also add the comment to the code so that readers understand why these bounds check are running. when #23127 is closed we can remove these checks
| &value_indices, | ||
| equal_to_results, | ||
| ); | ||
| } else if let Some(sentinel) = Self::null_sentinel_index(dict_values) { |
Contributor
There was a problem hiding this comment.
a null key in rhs_rows is mapped to sentinel_idx inside dict_values (the batch's values array), and then passed to self.inner.vectorized_equal_to . But the null that was originally appended went through self.inner.append_val(&self.null_array, 0), a null from a different array than this batch's dict_values .
Can you confirm that all inner builder implementations (bytes, primitive, etc.) short-circuit on null equality before reading the array content so comparing "null from null_array " vs "null from dict_values [sentinel]" always returns true ?
Contributor
Author
There was a problem hiding this comment.
Yes each of the traits use nulls_equal_to() to short-circuit null array checks.
bytes.rs/180
let exist_null = self.nulls.is_null(lhs_row);
let input_null = array.is_null(rhs_row);
if let Some(result) = nulls_equal_to(exist_null, input_null) {
return result;
}bytes_view.rs/263
// Check if nulls equal firstly
if HAS_NULLS {
let exist_null = self.nulls.is_null(lhs_row);
let input_null = array.is_null(rhs_row);
if let Some(result) = nulls_equal_to(exist_null, input_null) {
return result;
}
}primitive.rs/154
// Perf: skip null check (by short circuit) if input is not nullable
if NULLABLE {
let exist_null = self.nulls.is_null(lhs_row);
let input_null = array.is_null(rhs_row);
if let Some(result) = nulls_equal_to(exist_null, input_null) {
return result;
}
// Otherwise, we need to check their values
}boolean.rs/51
if NULLABLE {
let exist_null = self.nulls.is_null(lhs_row);
let input_null = array.is_null(rhs_row);
if let Some(result) = nulls_equal_to(exist_null, input_null) {
return result;
}
}mod.rs/116
pub fn nulls_equal_to(lhs_null: bool, rhs_null: bool) -> Option<bool> {
match (lhs_null, rhs_null) {
(true, true) => Some(true),
(false, true) | (true, false) => Some(false),
_ => None,
}
}since we know that rhs_null will always be true this function will always return Some(true/false) meaning further comparisons never run.
Contributor
Author
|
thanks for the review @kumarUjjawal I pushed up revisions to the PR, could you take another look 🚀 |
kosiew
pushed a commit
to kosiew/datafusion
that referenced
this pull request
Jul 11, 2026
…23280) ## Which issue does this PR close? <!-- 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 apache#123` indicates that this PR will close issue apache#123. --> works towards apache#22682. ## Rationale for this change There is a lack of testing for multi-dictionary group bys. It make sense to introduce these test before the implementation of `Dict<K,V>` in apache#23187 <!-- 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. --> ## What changes are included in this PR? introduces a couple test - grouping by `Dict<_,largeutf8>` - mixing grouping by dictionarys and non dictionary columns - a 3-way group by where each column is `dict<_,_>` - the test also have nulls sprinkled in to verify null handling <!-- 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? the changes are test. <!-- 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. -->
Contributor
|
run benchmark dictionary_group_values |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/dictionary-groupValuesColumn-impl (dc097bd) to 01bf68c (merge-base) diff using: dictionary_group_values File an issue against this benchmark runner |
| /// Returns an error if adding `additional` groups would overflow the key type's range. | ||
| // https://github.com/apache/datafusion/issues/23127 | ||
| fn check_key_overflow(current_len: usize, additional: usize) -> Result<()> { | ||
| if !Self::valid_bounds::<K>(current_len + additional) { |
Contributor
There was a problem hiding this comment.
check_key_overflow currently checks the dictionary key capacity against the total number of group rows. In a multi-column group-by, there can be more than 128 groups while an Int8 dictionary still has only a few distinct values. Told row path will successfully emits 129 groups with two dictionary values, while this path return will be overflow error. Could we preserve dictionary encoding during emit, or keep narrow key types on the fallback path? Please also add this as a regression test.
Contributor
Author
There was a problem hiding this comment.
In a multi-column group-by, there can be more than 128 groups while an Int8 dictionary still has only a few distinct values. Told row path will successfully emits 129 groups with two dictionary values while this path return will be overflow error.
I created this test off of the main branch to test GroupValueRows
fn test_dictionary_emit_logic() {
let schema = Arc::new(Schema::new(vec![Field::new(
"categories",
DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
false,
)]));
let keys: Vec<u8> = (0..1000).map(|i| (i % 255) as u8).collect();
let values = (0..255)
.map(|i| format!("category_{}", i))
.collect::<Vec<_>>();
let values_vec: Vec<Option<&str>> =
values.iter().map(|s| Some(s.as_str())).collect();
let values_arr = Arc::new(StringArray::from(values_vec));
let key_array = PrimitiveArray::<UInt8Type>::from(keys);
let dict_array = Arc::new(
DictionaryArray::<UInt8Type>::try_new(key_array, values_arr).unwrap(),
) as ArrayRef;
let mut group_values_rows =
new_group_values(schema, &GroupOrdering::None).unwrap();
let mut groups = Vec::new();
group_values_rows
.intern(&[dict_array], &mut groups)
.unwrap();
let size = group_values_rows.len();
// we are now at the max valid size to represent a uint8 key
assert_eq!(size, 255);
// Add two more groups and verify count rises to 257
let additional_keys: Vec<u8> = vec![0, 1];
let additional_values = vec![Some("extra_0"), Some("extra_1")];
let key_array = PrimitiveArray::<UInt8Type>::from(additional_keys);
let values_arr = Arc::new(StringArray::from(additional_values));
let additional_dict_array = Arc::new(
DictionaryArray::<UInt8Type>::try_new(key_array, values_arr).unwrap(),
) as ArrayRef;
group_values_rows
.intern(&[additional_dict_array], &mut groups)
.unwrap();
let new_size = group_values_rows.len();
assert_eq!(new_size, 257);
let output = group_values_rows.emit(super::EmitTo::All).unwrap();
assert_eq!(output.len(), 1);
let output_array = &output[0];
assert_eq!(
output_array.data_type(),
&DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8))
);
}it panics on the .emit(All).unwrap call with this error
thread 'aggregates::group_values::row::test::test_dictionary_emit_logic' (51728094) panicked at datafusion/physical-plan/src/aggregates/group_values/row.rs:390:65:
called `Result::unwrap()` on an `Err` value: ArrowError(DictionaryKeyOverflowError, Some(""))
This behavior already exist in GroupValuesRows this is because of the cast the happens in the emit path. if the strings RowConverter produces cannot be cast to the input schemas dictionary key type it will just fail.
Could we preserve dictionary encoding during emit, or keep narrow key types on the fallback path? Please also add this as a regression test.
If emit::all/build is called, we need to produce every distinct value we've stored. If the number of stored values exceeds what the key type can represent, we can only emit the first key-type-sized number of outputs — e.g., if we've stored 300 distinct values but the input key type is uint8, we can only emit 255, which is logically incorrect.
This matches existing behavior, so it won't introduce a regression, but it's worth fixing properly down the line, I've opened #23127 to track it, and I'll add a regression test to document the current behavior.
Contributor
Author
There was a problem hiding this comment.
Actually I think I understand what your refering to now. the check for key-overflow needs to be placed after we store new values in the internal GroupColumns.
so instead of
fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()> {
let dict = array.as_dictionary::<K>();
.... write to inner groupColumnwe'd do
... write to inner groupColumn
Self::check_key_overflow(self.inner.len(), 0)?;now if the actually groups has exceed what we can store we'd error out here.
|
🤖 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
from
July 17, 2026 14:47
3dbb88d to
aaad2c0
Compare
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #23187 +/- ##
=======================================
Coverage ? 80.69%
=======================================
Files ? 1087
Lines ? 367086
Branches ? 367086
=======================================
Hits ? 296205
Misses ? 53244
Partials ? 17637 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Contributor
|
run benchmark clickbench |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/dictionary-groupValuesColumn-impl (aaad2c0) to 01bf68c (merge-base) diff using: clickbench File an issue against this benchmark runner |
|
Benchmark for this request failed. Last 20 lines of output: Click to expandFile an issue against this benchmark runner |
Contributor
|
run benchmarks |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/dictionary-groupValuesColumn-impl (aaad2c0) to 01bf68c (merge-base) diff using: tpcds File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/dictionary-groupValuesColumn-impl (aaad2c0) to 01bf68c (merge-base) diff using: clickbench_partitioned File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/dictionary-groupValuesColumn-impl (aaad2c0) to 01bf68c (merge-base) diff using: tpch File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagetpch — base (merge-base)
tpch — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagetpcds — base (merge-base)
tpcds — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usageclickbench_partitioned — base (merge-base)
clickbench_partitioned — branch
File an issue against this benchmark runner |
Contributor
Author
|
@zhuqi-lucas clickbench and tpch benchmarks dont contain dictionary arrays |
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
7 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.