maxburke commented on code in PR #23695:
URL: https://github.com/apache/datafusion/pull/23695#discussion_r4539922598
##########
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:
Done!
--
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]
