ghostiee-11 commented on issue #7647:
URL: https://github.com/apache/datafusion/issues/7647#issuecomment-4883024528
> Filed with AI assistance per the project's policy. The repro was written
and run against datafusion 54.0.0 (pyarrow 23.0.0, arrow-rs 58.3.0); the source
analysis below is mine and was reviewed before posting.
Related to this issue: on datafusion 54 a dictionary group key whose
combined cardinality across partitions exceeds its key width fails hard, rather
than just being suboptimal.
## Repro (pure datafusion + pyarrow, no network)
```python
import numpy as np, pyarrow as pa
from datafusion import SessionContext
PER_PARTITION, N_PARTITIONS = 100, 50 # combined 5000 >> Int8 max (127)
parts = []
for p in range(N_PARTITIONS):
values = pa.array(np.arange(p*PER_PARTITION,
p*PER_PARTITION+PER_PARTITION, dtype=np.int64))
keys = pa.array(np.arange(PER_PARTITION, dtype=np.int8))
k = pa.DictionaryArray.from_arrays(keys, values) # valid
Dictionary(Int8, Int64) per batch
parts.append(pa.record_batch({"k": k, "v":
pa.array(np.ones(PER_PARTITION))}))
ctx = SessionContext()
ctx.register_record_batches("t", [[b] for b in parts])
ctx.sql("SELECT k, SUM(v) FROM t GROUP BY k").to_arrow_table()
# -> Arrow error: Dictionary key bigger than the key type
```
Each partition is a valid `Dictionary(Int8, Int64)` (<=128 distinct), but
the partitions hold disjoint values, so the combined group-key set is 5000.
## Mechanism (datafusion 54)
- `Aggregate::group_fields` (physical-plan/src/aggregates/mod.rs) keeps the
dictionary type for the group key.
- `GroupValuesRows::emit` decodes the keys via the RowConverter, then
`dictionary_encode_if_necessary` (group_values/row.rs) re-encodes them back
into the narrow dictionary because the group schema says so.
- Across `RepartitionExec: Hash` / `FinalPartitioned`, the combined set
exceeds Int8, so the re-encode overflows. arrow-rs is behaving correctly here
(`concat` is type-preserving, so it errors in `merge_dictionary_values` rather
than silently widening); the decision to keep a narrow key for the combined
group keys is the datafusion side.
`#8291` materialized dict group keys to the value type and avoided this;
that path is not present in 54. This issue and `#21765` are moving toward
keeping dictionaries and optimizing on them, which is great for low cardinality
but needs to handle the case where the combined cardinality outgrows the key
width.
## Question before I send a PR
Two ways to fix it, and I would rather match your direction:
- **A. Decode** dict group keys to their value type at plan time (restores
`#8291` semantics). Simple, output type changes to the value type.
- **B. Widen** the group key to a safe key width (e.g. Int32/Int64) and keep
a dictionary output (aligns with `#21765`). Preserves "group by dict returns
dict".
I have both a candidate patch and a sqllogictest ready. Which direction do
you prefer for the aggregate group-key path? Happy to open the PR against
whichever you pick.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
