kosiew commented on code in PR #23695:
URL: https://github.com/apache/datafusion/pull/23695#discussion_r3663622885
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -2549,6 +2571,119 @@ impl AggregateExec {
}
}
+/// The type used internally to intern and emit a group-by key of the given
+/// input type.
+///
+/// String and binary group keys are accumulated with 64-bit offsets
+/// (`LargeUtf8` / `LargeBinary`) so that the single array holding all
+/// distinct group keys is not limited to `i32::MAX` total bytes. This is an
+/// implementation detail of the aggregation streams: output batches are
+/// narrowed back to the declared output types after being sliced to
+/// `batch_size` rows, where the narrow representation always fits.
+pub(crate) fn internal_group_key_type(input_type: &DataType) -> DataType {
+ match input_type {
+ DataType::Utf8 => DataType::LargeUtf8,
+ DataType::Binary => DataType::LargeBinary,
+ other => other.clone(),
+ }
+}
+
+/// Returns `schema` with any string/binary fields among the leading
+/// `num_group_columns` group key fields widened to the internal group key
+/// representation (see [`internal_group_key_type`]).
+///
+/// Returns the original schema when nothing needs widening.
+pub(crate) fn widen_group_key_schema(
+ schema: &SchemaRef,
+ num_group_columns: usize,
+) -> SchemaRef {
+ let mut changed = false;
+ let fields = schema
+ .fields()
+ .iter()
+ .enumerate()
+ .map(|(idx, field)| {
+ if idx < num_group_columns {
+ let wide = internal_group_key_type(field.data_type());
+ if wide != *field.data_type() {
+ changed = true;
+ return
Arc::new(field.as_ref().clone().with_data_type(wide));
+ }
+ }
+ Arc::clone(field)
+ })
+ .collect::<Vec<_>>();
+
+ if changed {
+ Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone()))
+ } else {
+ Arc::clone(schema)
+ }
+}
+
+/// Widens string/binary group key arrays to the internal representation
+/// interned by `group_values` (see `PhysicalGroupBy::group_schema`).
+pub(crate) fn widen_group_key_arrays(arrays: Vec<ArrayRef>) ->
Result<Vec<ArrayRef>> {
Review Comment:
One performance thought. `ByteGroupValueBuilder` now accepts both narrow and
large input arrays, so I wonder if the per input batch `arrow::compute::cast`
in `widen_group_key_arrays` is still necessary.
If the hashing, ordering, spilling and merge paths do not actually require
the grouping arrays to be widened ahead of time, it may be possible to remove
this cast entirely. That would avoid some extra offset and value buffer work on
the hot aggregation path.
##########
datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs:
##########
@@ -253,7 +274,13 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
}
};
- let batch = output.next_batch(batch_size);
+ // Narrow internally-widened group key columns back to the output
+ // schema types. Each slice holds at most `batch_size` rows, so unlike
+ // the full materialized batch, the narrow representation fits.
+ let batch = output
+ .next_batch(batch_size)
Review Comment:
I think there is still one important edge case here. The output is split by
row count before the `LargeUtf8` / `LargeBinary` group keys are narrowed back
to the declared `Utf8` / `Binary` schema.
The Arrow limit for narrow string and binary arrays is based on the total
value buffer size, not the number of rows. That means a `batch_size` slice can
still exceed `i32::MAX` bytes if it contains many large distinct keys
accumulated across input batches.
In that case this change moves the overflow from group value accumulation to
the later `narrow_group_key_columns` cast, so the query can still fail.
I noticed the same assumption in
`datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs:827` and
`datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs:432`.
Could we instead split emitted output based on the narrow offset byte limit,
or otherwise guarantee that each emitted slice is below that limit? It would
also be great to add a regression that exercises this boundary.
--
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]
