impl the new interface functions for all GroupValues impls.

Author: Rachelint

datafusion/physical-plan/src/aggregates/group_values/bytes.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::aggregates::group_values::GroupValues;
+use crate::aggregates::group_values::{GroupIdx, GroupValues};
 use arrow_array::{Array, ArrayRef, OffsetSizeTrait, RecordBatch};
 use datafusion_expr::EmitTo;
 use datafusion_physical_expr_common::binary_map::{ArrowBytesMap, OutputType};
@@ -44,7 +44,7 @@ impl<O: OffsetSizeTrait> GroupValues for GroupValuesByes<O> {
     fn intern(
         &mut self,
         cols: &[ArrayRef],
-        groups: &mut Vec<usize>,
+        groups: &mut Vec<GroupIdx>,
     ) -> datafusion_common::Result<()> {
         assert_eq!(cols.len(), 1);
 
@@ -63,7 +63,7 @@ impl<O: OffsetSizeTrait> GroupValues for GroupValuesByes<O> {
             },
             // called for each group
             |group_idx| {
-                groups.push(group_idx);
+                groups.push(GroupIdx::new(0, group_idx as u64));
             },
         );
 
@@ -84,7 +84,7 @@ impl<O: OffsetSizeTrait> GroupValues for GroupValuesByes<O> {
         self.num_groups
     }
 
-    fn emit(&mut self, emit_to: EmitTo) -> datafusion_common::Result<Vec<ArrayRef>> {
+    fn emit(&mut self, emit_to: EmitTo) -> datafusion_common::Result<Vec<Vec<ArrayRef>>> {
         // Reset the map to default, and convert it into a single array
         let map_contents = self.map.take().into_state();
 
@@ -111,13 +111,13 @@ impl<O: OffsetSizeTrait> GroupValues for GroupValuesByes<O> {
                 self.intern(&[remaining_group_values], &mut group_indexes)?;
 
                 // Verify that the group indexes were assigned in the correct order
-                assert_eq!(0, group_indexes[0]);
+                assert_eq!(0, group_indexes[0].block_offset());
 
                 emit_group_values
             }
         };
 
-        Ok(vec![group_values])
+        Ok(vec![vec![group_values]])
     }
 
     fn clear_shrink(&mut self, _batch: &RecordBatch) {

datafusion/physical-plan/src/aggregates/group_values/bytes_view.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::aggregates::group_values::GroupValues;
+use crate::aggregates::group_values::{GroupIdx, GroupValues};
 use arrow_array::{Array, ArrayRef, RecordBatch};
 use datafusion_expr::EmitTo;
 use datafusion_physical_expr::binary_map::OutputType;
@@ -45,7 +45,7 @@ impl GroupValues for GroupValuesBytesView {
     fn intern(
         &mut self,
         cols: &[ArrayRef],
-        groups: &mut Vec<usize>,
+        groups: &mut Vec<GroupIdx>,
     ) -> datafusion_common::Result<()> {
         assert_eq!(cols.len(), 1);
 
@@ -64,7 +64,7 @@ impl GroupValues for GroupValuesBytesView {
             },
             // called for each group
             |group_idx| {
-                groups.push(group_idx);
+                groups.push(GroupIdx::new(0, group_idx as u64));
             },
         );
 
@@ -85,7 +85,7 @@ impl GroupValues for GroupValuesBytesView {
         self.num_groups
     }
 
-    fn emit(&mut self, emit_to: EmitTo) -> datafusion_common::Result<Vec<ArrayRef>> {
+    fn emit(&mut self, emit_to: EmitTo) -> datafusion_common::Result<Vec<Vec<ArrayRef>>> {
         // Reset the map to default, and convert it into a single array
         let map_contents = self.map.take().into_state();
 
@@ -112,13 +112,13 @@ impl GroupValues for GroupValuesBytesView {
                 self.intern(&[remaining_group_values], &mut group_indexes)?;
 
                 // Verify that the group indexes were assigned in the correct order
-                assert_eq!(0, group_indexes[0]);
+                assert_eq!(0, group_indexes[0].block_offset());
 
                 emit_group_values
             }
         };
 
-        Ok(vec![group_values])
+        Ok(vec![vec![group_values]])
     }
 
     fn clear_shrink(&mut self, _batch: &RecordBatch) {

datafusion/physical-plan/src/aggregates/group_values/mod.rs

@@ -61,7 +61,7 @@ impl GroupIdx {
 /// An interning store for group keys
 pub trait GroupValues: Send {
     /// Calculates the `groups` for each input row of `cols`
-    fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()>;
+    fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<GroupIdx>) -> Result<()>;
 
     /// Returns the number of bytes used by this [`GroupValues`]
     fn size(&self) -> usize;

datafusion/physical-plan/src/aggregates/group_values/primitive.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::aggregates::group_values::GroupValues;
+use crate::aggregates::group_values::{GroupIdx, GroupValues};
 use ahash::RandomState;
 use arrow::array::BooleanBufferBuilder;
 use arrow::buffer::NullBuffer;
@@ -111,7 +111,7 @@ impl<T: ArrowPrimitiveType> GroupValues for GroupValuesPrimitive<T>
 where
     T::Native: HashValue,
 {
-    fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
+    fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<GroupIdx>) -> Result<()> {
         assert_eq!(cols.len(), 1);
         groups.clear();
 
@@ -145,7 +145,7 @@ where
                     }
                 }
             };
-            groups.push(group_id)
+            groups.push(GroupIdx::new(0, group_id as u64))
         }
         Ok(())
     }
@@ -162,7 +162,7 @@ where
         self.values.len()
     }
 
-    fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
+    fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<Vec<ArrayRef>>> {
         fn build_primitive<T: ArrowPrimitiveType>(
             values: Vec<T::Native>,
             null_idx: Option<usize>,
@@ -207,7 +207,7 @@ where
                 build_primitive(split, null_group)
             }
         };
-        Ok(vec![Arc::new(array.with_data_type(self.data_type.clone()))])
+        Ok(vec![vec![Arc::new(array.with_data_type(self.data_type.clone()))]])
     }
 
     fn clear_shrink(&mut self, batch: &RecordBatch) {

datafusion/physical-plan/src/aggregates/group_values/row.rs

@@ -18,7 +18,7 @@
 use std::collections::VecDeque;
 use std::mem;
 
-use crate::aggregates::group_values::{GroupBlock, GroupIdx, GroupValues};
+use crate::aggregates::group_values::{GroupIdx, GroupValues};
 use ahash::RandomState;
 use arrow::compute::cast;
 use arrow::record_batch::RecordBatch;
@@ -189,7 +189,7 @@ impl GroupValues for GroupValuesRows {
     }
 
     fn size(&self) -> usize {
-        let group_values_size = self.group_values_blocks.as_ref().map(|v| v.size()).unwrap_or(0);
+        let group_values_size = self.group_values_blocks.iter().map(|v| v.size()).sum::<usize>();
         self.row_converter.size()
             + group_values_size
             + self.map_size
@@ -203,9 +203,9 @@ impl GroupValues for GroupValuesRows {
 
     fn len(&self) -> usize {
         self.group_values_blocks
-            .as_ref()
+            .iter()
             .map(|group_values| group_values.num_rows())
-            .unwrap_or(0)
+            .sum::<usize>()
     }
 
     fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<Vec<ArrayRef>>> {
@@ -293,10 +293,8 @@ impl GroupValues for GroupValuesRows {
 
     fn clear_shrink(&mut self, batch: &RecordBatch) {
         let count = batch.num_rows();
-        self.group_values_blocks = self.group_values_blocks.take().map(|mut rows| {
-            rows.clear();
-            rows
-        });
+        let mut old_blocks = mem::replace(&mut self.group_values_blocks, VecDeque::new());
+        old_blocks.iter_mut().for_each(|block| block.clear());
         self.map.clear();
         self.map.shrink_to(count, |_| 0); // hasher does not matter since the map is cleared
         self.map_size = self.map.capacity() * std::mem::size_of::<(u64, usize)>();