Implement CardinalityAwareRowConverter while doing streaming merge

datafusion/core/src/physical_plan/aggregates/group_values/row.rs

@@ -16,8 +16,9 @@
 // under the License.
 
 use crate::physical_plan::aggregates::group_values::GroupValues;
+use crate::physical_plan::row_converter::CardinalityAwareRowConverter;
 use ahash::RandomState;
-use arrow::row::{RowConverter, Rows, SortField};
+use arrow::row::{Rows, SortField};
 use arrow_array::ArrayRef;
 use arrow_schema::SchemaRef;
 use datafusion_common::Result;
@@ -29,7 +30,7 @@ use hashbrown::raw::RawTable;
 /// A [`GroupValues`] making use of [`Rows`]
 pub struct GroupValuesRows {
     /// Converter for the group values
-    row_converter: RowConverter,
+    row_converter: CardinalityAwareRowConverter,
 
     /// Logically maps group values to a group_index in
     /// [`Self::group_values`] and in each accumulator
@@ -52,7 +53,7 @@ pub struct GroupValuesRows {
     /// important for multi-column group keys.
     ///
     /// [`Row`]: arrow::row::Row
-    group_values: Rows,
+    group_values: Option<Rows>,
 
     // buffer to be reused to store hashes
     hashes_buffer: Vec<u64>,
@@ -63,7 +64,7 @@ pub struct GroupValuesRows {
 
 impl GroupValuesRows {
     pub fn try_new(schema: SchemaRef) -> Result<Self> {
-        let row_converter = RowConverter::new(
+        let row_converter = CardinalityAwareRowConverter::new(
             schema
                 .fields()
                 .iter()
@@ -72,13 +73,12 @@ impl GroupValuesRows {
         )?;
 
         let map = RawTable::with_capacity(0);
-        let group_values = row_converter.empty_rows(0, 0);
 
         Ok(Self {
             row_converter,
             map,
             map_size: 0,
-            group_values,
+            group_values: None,
             hashes_buffer: Default::default(),
             random_state: Default::default(),
         })
@@ -92,6 +92,11 @@ impl GroupValues for GroupValuesRows {
         let group_rows = self.row_converter.convert_columns(cols)?;
         let n_rows = group_rows.num_rows();
 
+        let mut group_values = match self.group_values.take() {
+            Some(group_values) => group_values,
+            None => self.row_converter.empty_rows(0, 0)?,
+        };
+
         // tracks to which group each of the input rows belongs
         groups.clear();
 
@@ -106,7 +111,7 @@ impl GroupValues for GroupValuesRows {
                 // verify that a group that we are inserting with hash is
                 // actually the same key value as the group in
                 // existing_idx  (aka group_values @ row)
-                group_rows.row(row) == self.group_values.row(*group_idx)
+                group_rows.row(row) == group_values.row(*group_idx)
             });
 
             let group_idx = match entry {
@@ -115,8 +120,8 @@ impl GroupValues for GroupValuesRows {
                 //  1.2 Need to create new entry for the group
                 None => {
                     // Add new entry to aggr_state and save newly created index
-                    let group_idx = self.group_values.num_rows();
-                    self.group_values.push(group_rows.row(row));
+                    let group_idx = group_values.num_rows();
+                    group_values.push(group_rows.row(row));
 
                     // for hasher function, use precomputed hash value
                     self.map.insert_accounted(
@@ -130,12 +135,14 @@ impl GroupValues for GroupValuesRows {
             groups.push(group_idx);
         }
 
+        self.group_values = Some(group_values);
+
         Ok(())
     }
 
     fn size(&self) -> usize {
         self.row_converter.size()
-            + self.group_values.size()
+            + self.group_values.as_ref().unwrap().size()
             + self.map_size
             + self.hashes_buffer.allocated_size()
     }
@@ -145,25 +152,34 @@ impl GroupValues for GroupValuesRows {
     }
 
     fn len(&self) -> usize {
-        self.group_values.num_rows()
+        self.group_values
+            .as_ref()
+            .map(|group_values| group_values.num_rows())
+            .unwrap_or(0)
     }
 
     fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
-        Ok(match emit_to {
+        let mut group_values = self
+            .group_values
+            .take()
+            .expect("Can not emit from empty rows");
+
+        let output = match emit_to {
             EmitTo::All => {
-                // Eventually we may also want to clear the hash table here
-                self.row_converter.convert_rows(&self.group_values)?
+                let output = self.row_converter.convert_rows(&group_values)?;
+                group_values.clear();
+                output
             }
             EmitTo::First(n) => {
-                let groups_rows = self.group_values.iter().take(n);
+                let groups_rows = group_values.iter().take(n);
                 let output = self.row_converter.convert_rows(groups_rows)?;
                 // Clear out first n group keys by copying them to a new Rows.
                 // TODO file some ticket in arrow-rs to make this more efficent?
-                let mut new_group_values = self.row_converter.empty_rows(0, 0);
-                for row in self.group_values.iter().skip(n) {
+                let mut new_group_values = self.row_converter.empty_rows(0, 0)?;
+                for row in group_values.iter().skip(n) {
                     new_group_values.push(row);
                 }
-                std::mem::swap(&mut new_group_values, &mut self.group_values);
+                std::mem::swap(&mut new_group_values, &mut group_values);
 
                 // SAFETY: self.map outlives iterator and is not modified concurrently
                 unsafe {
@@ -179,6 +195,9 @@ impl GroupValues for GroupValuesRows {
                 }
                 output
             }
-        })
+        };
+
+        self.group_values = Some(group_values);
+        Ok(output)
     }
 }

datafusion/core/src/physical_plan/aggregates/order/partial.rs

@@ -15,8 +15,10 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::physical_expr::EmitTo;
-use arrow::row::{OwnedRow, RowConverter, Rows, SortField};
+use crate::{
+    physical_expr::EmitTo, physical_plan::row_converter::CardinalityAwareRowConverter,
+};
+use arrow::row::{OwnedRow, Rows, SortField};
 use arrow_array::ArrayRef;
 use arrow_schema::Schema;
 use datafusion_common::Result;
@@ -70,7 +72,7 @@ pub(crate) struct GroupOrderingPartial {
 
     /// Converter for the sort key (used on the group columns
     /// specified in `order_indexes`)
-    row_converter: RowConverter,
+    row_converter: CardinalityAwareRowConverter,
 }
 
 #[derive(Debug, Default)]
@@ -124,7 +126,7 @@ impl GroupOrderingPartial {
         Ok(Self {
             state: State::Start,
             order_indices: order_indices.to_vec(),
-            row_converter: RowConverter::new(fields)?,
+            row_converter: CardinalityAwareRowConverter::new(fields)?,
         })
     }
 
@@ -141,7 +143,7 @@ impl GroupOrderingPartial {
             .map(|&idx| group_values[idx].clone())
             .collect();
 
-        Ok(self.row_converter.convert_columns(&sort_values)?)
+        self.row_converter.convert_columns(&sort_values)
     }
 
     /// How many groups be emitted, or None if no data can be emitted

datafusion/core/src/physical_plan/mod.rs

@@ -387,6 +387,7 @@ pub mod memory;
 pub mod metrics;
 pub mod projection;
 pub mod repartition;
+mod row_converter;
 pub mod sorts;
 pub mod stream;
 pub mod streaming;