add append_n in MaybeNullBufferBuilder.

apache · Oct 19, 2024 · a83c2ea · a83c2ea
1 parent 04ea2d2
commit a83c2ea
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 2 deletions.
diff --git a/datafusion/physical-plan/src/aggregates/group_values/group_column.rs b/datafusion/physical-plan/src/aggregates/group_values/group_column.rs
@@ -61,6 +61,8 @@ pub trait GroupColumn: Send + Sync {
 
     fn append_non_nullable_val(&mut self, array: &ArrayRef, row: usize);
 
+    fn append_batch(&mut self, array: &ArrayRef, rows: &[usize]);
+
     /// Returns the number of rows stored in this builder
     fn len(&self) -> usize;
     /// Returns the number of bytes used by this [`GroupColumn`]
@@ -120,8 +122,11 @@ impl<T: ArrowPrimitiveType, const NULLABLE: bool> GroupColumn
         self.group_values[lhs_row] == array.as_primitive::<T>().value(rhs_row)
     }
 
+    fn append_batch(&mut self, array: &ArrayRef, rows: &[usize]) {
+        todo!()
+    }    
+
     fn append_val(&mut self, array: &ArrayRef, row: usize) {
-        self.nullable_call += 1;
         // Perf: skip null check if input can't have nulls
         if NULLABLE {
             if array.is_null(row) {
@@ -137,7 +142,6 @@ impl<T: ArrowPrimitiveType, const NULLABLE: bool> GroupColumn
     }
 
     fn append_non_nullable_val(&mut self, array: &ArrayRef, row: usize) {
-        self.non_nullable_call += 1;
         if NULLABLE {
             self.nulls.append(false);
             self.group_values.push(array.as_primitive::<T>().value(row));

diff --git a/datafusion/physical-plan/src/aggregates/group_values/null_builder.rs b/datafusion/physical-plan/src/aggregates/group_values/null_builder.rs
@@ -70,6 +70,24 @@ impl MaybeNullBufferBuilder {
         }
     }
 
+    pub fn append_n(&mut self, n: usize, is_null: bool) {
+        match self {
+            Self::NoNulls { row_count } if is_null => {
+                // have seen no nulls so far, this is the  first null,
+                // need to create the nulls buffer for all currently valid values
+                // alloc 2x the need given we push a new but immediately
+                let mut nulls = BooleanBufferBuilder::new(*row_count * 2);
+                nulls.append_n(*row_count, true);
+                nulls.append_n(n, false);
+                *self = Self::Nulls(nulls);
+            }
+            Self::NoNulls { row_count } => {
+                *row_count += n;
+            }
+            Self::Nulls(builder) => builder.append_n(n, !is_null),
+        }
+    }
+
     /// return the number of heap allocated bytes used by this structure to store boolean values
     pub fn allocated_size(&self) -> usize {
         match self {