Aggregate byte types internally with i64 offsets

datafusion/physical-expr-common/src/binary_map.rs

@@ -9,7 +9,7 @@ use arrow::array::{
     Array, ArrayRef, BufferBuilder, GenericBinaryArray, GenericStringArray,
     NullBufferBuilder, OffsetSizeTrait,
     cast::AsArray,
-    types::{ByteArrayType, GenericBinaryType, GenericStringType},
+    types::{BinaryType, ByteArrayType, GenericBinaryType, GenericStringType, Utf8Type},
 };
 use arrow::buffer::{Buffer, NullBuffer, OffsetBuffer, ScalarBuffer};
 use arrow::datatypes::DataType;
@@ -348,7 +348,9 @@ where
     ///
     /// # Arguments:
     ///
-    /// `values`: array whose values are inserted
+    /// `values`: the input array. A map with 64-bit offsets also accepts
+    /// 32-bit input offsets. The map can store more than 2 GiB of values
+    /// without input conversion.
     ///
     /// `make_payload_fn`:  invoked for each value that is not already present
     /// to create the payload, in order of the values in `values`
@@ -374,29 +376,47 @@ where
         MP: FnMut(Option<&[u8]>) -> V,
         OP: FnMut(V),
     {
+        let reads_32_bit_offsets = O::IS_LARGE
+            && matches!(values.data_type(), DataType::Utf8 | DataType::Binary);
         // Sanity array type
         match self.output_type {
             OutputType::Binary => {
                 assert!(matches!(
                     values.data_type(),
                     DataType::Binary | DataType::LargeBinary
                 ));
-                self.insert_if_new_inner::<MP, OP, GenericBinaryType<O>>(
-                    values,
-                    make_payload_fn,
-                    observe_payload_fn,
-                )
+                if reads_32_bit_offsets {
+                    self.insert_if_new_inner::<MP, OP, BinaryType>(
+                        values,
+                        make_payload_fn,
+                        observe_payload_fn,
+                    )
+                } else {
+                    self.insert_if_new_inner::<MP, OP, GenericBinaryType<O>>(
+                        values,
+                        make_payload_fn,
+                        observe_payload_fn,
+                    )
+                }
             }
             OutputType::Utf8 => {
                 assert!(matches!(
                     values.data_type(),
                     DataType::Utf8 | DataType::LargeUtf8
                 ));
-                self.insert_if_new_inner::<MP, OP, GenericStringType<O>>(
-                    values,
-                    make_payload_fn,
-                    observe_payload_fn,
-                )
+                if reads_32_bit_offsets {
+                    self.insert_if_new_inner::<MP, OP, Utf8Type>(
+                        values,
+                        make_payload_fn,
+                        observe_payload_fn,
+                    )
+                } else {
+                    self.insert_if_new_inner::<MP, OP, GenericStringType<O>>(
+                        values,
+                        make_payload_fn,
+                        observe_payload_fn,
+                    )
+                }
             }
             _ => unreachable!("View types should use `ArrowBytesViewMap`"),
         }
@@ -736,7 +756,7 @@ where
 #[cfg(test)]
 mod tests {
     use super::*;
-    use arrow::array::{BinaryArray, LargeBinaryArray, StringArray};
+    use arrow::array::{BinaryArray, LargeBinaryArray, LargeStringArray, StringArray};
     use std::collections::HashMap;
 
     /// A lower bound on the bytes a hashbrown table holding `entries` entries
@@ -1201,6 +1221,52 @@ mod tests {
         assert_eq!(&set.into_state(), &expected);
     }
 
+    #[test]
+    fn large_offset_set_inserts_arrays_with_32_bit_offsets() {
+        let mut set = ArrowBytesSet::<i64>::new(OutputType::Utf8);
+        let values: ArrayRef = Arc::new(StringArray::from(vec![
+            Some("a"),
+            Some("CXCCCCCCCC"),
+            None,
+            Some("CXCCCCCCCC"),
+        ]));
+        set.insert(&values);
+        let values: ArrayRef = Arc::new(LargeStringArray::from(vec![
+            Some("CXCCCCCCCC"),
+            Some("b"),
+            None,
+        ]));
+        set.insert(&values);
+        let expected: ArrayRef = Arc::new(LargeStringArray::from(vec![
+            Some("a"),
+            Some("CXCCCCCCCC"),
+            None,
+            Some("b"),
+        ]));
+        assert_eq!(&set.into_state(), &expected);
+
+        let mut set = ArrowBytesSet::<i64>::new(OutputType::Binary);
+        let values: ArrayRef = Arc::new(BinaryArray::from_opt_vec(vec![
+            Some(b"a"),
+            Some(b"CXCCCCCCCC"),
+            None,
+            Some(b"CXCCCCCCCC"),
+        ]));
+        set.insert(&values);
+        let values: ArrayRef = Arc::new(LargeBinaryArray::from_opt_vec(vec![
+            Some(b"CXCCCCCCCC"),
+            Some(b"b"),
+        ]));
+        set.insert(&values);
+        let expected: ArrayRef = Arc::new(LargeBinaryArray::from_opt_vec(vec![
+            Some(b"a"),
+            Some(b"CXCCCCCCCC"),
+            None,
+            Some(b"b"),
+        ]));
+        assert_eq!(&set.into_state(), &expected);
+    }
+
     #[test]
     #[should_panic(
         expected = "matches!(values.data_type(), DataType::Utf8 | DataType::LargeUtf8)"

datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs

@@ -9,7 +9,7 @@ use std::sync::Arc;
 
 use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, new_null_array};
 use arrow::compute::prep_null_mask_filter;
-use arrow::datatypes::SchemaRef;
+use arrow::datatypes::{Schema, SchemaRef};
 use arrow::downcast_dictionary_array;
 use arrow::record_batch::RecordBatch;
 use datafusion_common::utils::compact_dictionary;
@@ -29,7 +29,8 @@ use crate::aggregates::group_values::{
 use crate::aggregates::order::GroupClustering;
 use crate::aggregates::{
     AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by,
-    group_id_array, max_duplicate_ordinal,
+    group_id_array, group_schema_with_large_offsets, max_duplicate_ordinal,
+    output_batch_rows, restore_group_key_types,
 };
 
 use super::AggregateTableMetrics;
@@ -309,7 +310,11 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
                         Ok::<_, datafusion_common::DataFusionError>(columns)
                     })?;
 
-                    let batch = RecordBatch::try_new(output_schema, columns)?;
+                    let emit_schema = group_schema_with_large_offsets(
+                        &output_schema,
+                        state.group_by.num_group_exprs(),
+                    );
+                    let batch = RecordBatch::try_new(emit_schema, columns)?;
                     debug_assert!(batch.num_rows() > 0);
                     MaterializedAggregateOutput::new(batch)
                 }
@@ -322,7 +327,11 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
                 }
             };
 
-        let batch = output.next_batch(batch_size)?;
+        let rows = output.next_batch_rows(&output_schema, batch_size);
+        let batch = output
+            .next_batch(rows)?
+            .map(|batch| restore_group_key_types(batch, &output_schema))
+            .transpose()?;
         if output.is_exhausted() {
             self.state = AggregateHashTableState::Done;
         } else {
@@ -362,6 +371,10 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
     /// Unlike normal single aggregation output, this materializes intermediate
     /// states rather than final values. The states can therefore be merged after
     /// spilling without finalizing the same group more than once.
+    ///
+    /// The keys use 64-bit offsets because this batch can contain more than
+    /// 2 GiB of key data.
+    /// Callers divide the batch, then change the keys to the output types.
     pub(in crate::aggregates) fn take_state_batch(
         &mut self,
     ) -> Result<Option<RecordBatch>> {
@@ -385,6 +398,10 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
             Ok::<_, datafusion_common::DataFusionError>(output)
         })?;
 
+        let state_schema = group_schema_with_large_offsets(
+            &state_schema,
+            state.group_by.num_group_exprs(),
+        );
         let batch = RecordBatch::try_new(state_schema, output)?;
         debug_assert!(batch.num_rows() > 0);
 
@@ -643,6 +660,15 @@ impl MaterializedAggregateOutput {
         Self { batch, offset: 0 }
     }
 
+    /// Returns a row count for [`Self::next_batch`] within the row and key byte limits.
+    pub(super) fn next_batch_rows(
+        &self,
+        output_schema: &Schema,
+        batch_size: usize,
+    ) -> usize {
+        output_batch_rows(&self.batch, output_schema, self.offset, batch_size)
+    }
+
     pub(super) fn next_batch(
         &mut self,
         batch_size: usize,
@@ -972,6 +998,7 @@ mod tests {
 
     use super::*;
     use crate::aggregates::group_values::aggregate_sub_metrics;
+    use crate::aggregates::tests::{GIB, zeroed_group_keys};
     use crate::metrics::ExecutionPlanMetricsSet;
 
     #[test]
@@ -1017,6 +1044,30 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn next_batch_rows_limits_key_bytes() -> Result<()> {
+        let output_schema =
+            Schema::new(vec![Field::new("group_col", DataType::Binary, false)]);
+        let stored_schema = Arc::new(Schema::new(vec![Field::new(
+            "group_col",
+            DataType::LargeBinary,
+            false,
+        )]));
+        // The first two keys hold one byte more than `i32::MAX`.
+        let keys = zeroed_group_keys(&[(DataType::LargeBinary, &[GIB, GIB, 1, 1, 1])]);
+        let batch = RecordBatch::try_new(stored_schema, keys)?;
+        let mut output = MaterializedAggregateOutput::new(batch);
+
+        let mut rows = vec![];
+        while !output.is_exhausted() {
+            let next_rows = output.next_batch_rows(&output_schema, 8192);
+            rows.push(output.next_batch(next_rows)?.unwrap().num_rows());
+        }
+        assert_eq!(rows, vec![1, 4]);
+
+        Ok(())
+    }
+
     fn dictionary_batch<K: ArrowDictionaryKeyType>(
         keys: PrimitiveArray<K>,
         num_values: usize,

datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_clustered.rs

@@ -22,7 +22,7 @@ use crate::aggregates::group_values::{
 use crate::aggregates::order::{GroupClustering, GroupClusteringMode};
 use crate::aggregates::{
     AggregateExec, AggregateMode, PhysicalGroupBy, aggregate_expressions,
-    evaluate_group_by,
+    evaluate_group_by, group_schema_with_large_offsets,
 };
 
 use super::AggregateTableMetrics;
@@ -320,6 +320,10 @@ impl<AggrMode> ClusteredAggregateTable<AggrMode> {
     /// active (incomplete) groups. Partial aggregation can pass those states to
     /// its final stage, while single and final aggregation sort and spill them
     /// before replay.
+    ///
+    /// The keys use 64-bit offsets because this batch can contain more than
+    /// 2 GiB of key data.
+    /// Callers divide the batch, then change the keys to the output types.
     pub(in crate::aggregates) fn take_state_batch(
         &mut self,
     ) -> Result<Option<RecordBatch>> {
@@ -340,7 +344,11 @@ impl<AggrMode> ClusteredAggregateTable<AggrMode> {
             Ok::<_, datafusion_common::DataFusionError>(output)
         })?;
 
-        let batch = RecordBatch::try_new(Arc::clone(&self.state_schema), output)?;
+        let state_schema = group_schema_with_large_offsets(
+            &self.state_schema,
+            self.buffer.group_by.num_group_exprs(),
+        );
+        let batch = RecordBatch::try_new(state_schema, output)?;
         debug_assert!(batch.num_rows() > 0);
 
         // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the
@@ -416,6 +424,8 @@ impl<AggrMode> ClusteredAggregateTable<AggrMode> {
 
     /// Removes the selected groups once and materializes their output columns.
     /// The caller chooses the completed prefix.
+    ///
+    /// The keys keep 64-bit offsets, as in [`Self::take_state_batch`].
     pub(super) fn materialize_groups(
         &mut self,
         emit_to: EmitTo,
@@ -441,7 +451,11 @@ impl<AggrMode> ClusteredAggregateTable<AggrMode> {
             Ok::<_, datafusion_common::DataFusionError>(output)
         })?;
 
-        let batch = RecordBatch::try_new(Arc::clone(&self.output_schema), output)?;
+        let output_schema = group_schema_with_large_offsets(
+            &self.output_schema,
+            self.buffer.group_by.num_group_exprs(),
+        );
+        let batch = RecordBatch::try_new(output_schema, output)?;
         debug_assert!(batch.num_rows() > 0);
 
         Ok(batch)

datafusion/physical-plan/src/aggregates/clustered_final_stream.rs

@@ -19,6 +19,7 @@ use super::aggregate_hash_table::{
     ClusteredAggregateTable, ClusteredAggregateTableMetrics, FinalMarker,
 };
 use super::order::GroupClusteringMode;
+use super::restore_group_key_stream;
 use super::spill::AggregateSpill;
 use crate::SendableRecordBatchStream;
 use crate::aggregates::AggregateMode;
@@ -272,6 +273,7 @@ impl ClusteredFinalAggregateStream {
             stage,
         } = self;
         let schema = Arc::clone(&context.schema);
+        let batch_size = context.batch_size;
         let metrics = context.baseline_metrics.clone();
         let stream = async_try_stream(|mut emitter| async move {
             let mut stage = Some(stage);
@@ -296,6 +298,7 @@ impl ClusteredFinalAggregateStream {
             Ok(())
         });
         let stream = Box::pin(RecordBatchStreamAdapter::new(schema, stream));
+        let stream = restore_group_key_stream(stream, batch_size);
         Box::pin(ObservedStream::new(stream, metrics, None))
     }
 }

datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs

@@ -15,6 +15,7 @@ use futures::stream::StreamExt;
 
 use super::AggregateExec;
 use super::aggregate_hash_table::{ClusteredAggregateTable, PartialMarker};
+use super::restore_group_key_stream;
 use crate::aggregates::AggregateMode;
 use crate::aggregates::order::{GroupClustering, GroupClusteringMode};
 use crate::metrics::{BaselineMetrics, MetricBuilder, SpillMetrics};
@@ -269,6 +270,7 @@ impl ClusteredPartialAggregateStream {
             stage,
         } = self;
         let schema = Arc::clone(&context.schema);
+        let batch_size = context.batch_size;
         let metrics = context.baseline_metrics.clone();
         let stream = async_try_stream(|mut emitter| async move {
             let mut stage = Some(stage);
@@ -287,6 +289,7 @@ impl ClusteredPartialAggregateStream {
             Ok(())
         });
         let stream = Box::pin(RecordBatchStreamAdapter::new(schema, stream));
+        let stream = restore_group_key_stream(stream, batch_size);
         Box::pin(ObservedStream::new(stream, metrics, None))
     }
 }

datafusion/physical-plan/src/aggregates/clustered_single_stream.rs

@@ -17,6 +17,7 @@ use futures::stream::{Stream, StreamExt};
 
 use super::aggregate_hash_table::{ClusteredAggregateTable, SingleMarker};
 use super::order::GroupClusteringMode;
+use super::restore_group_key_stream;
 use super::spill::AggregateSpill;
 use super::{AggregateExec, create_schema};
 use crate::aggregates::AggregateMode;
@@ -209,6 +210,11 @@ impl ClusteredSingleAggregateStream {
         })
     }
 
+    pub(super) fn into_stream(self) -> SendableRecordBatchStream {
+        let batch_size = self.batch_size;
+        restore_group_key_stream(Box::pin(self), batch_size)
+    }
+
     fn close_input(&mut self) {
         let input_schema = self.input.schema();
         self.input = Box::pin(EmptyRecordBatchStream::new(input_schema));