Refactor PrimitiveGroupValueBuilder to use MaybeNullBufferBuilder

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

@@ -32,11 +32,11 @@ use arrow::datatypes::GenericBinaryType;
 use arrow::datatypes::GenericStringType;
 use datafusion_common::utils::proxy::VecAllocExt;
 
+use crate::aggregates::group_values::null_builder::MaybeNullBufferBuilder;
+use datafusion_physical_expr_common::binary_map::{OutputType, INITIAL_BUFFER_CAPACITY};
 use std::sync::Arc;
 use std::vec;
 
-use datafusion_physical_expr_common::binary_map::{OutputType, INITIAL_BUFFER_CAPACITY};
-
 /// Trait for group values column-wise row comparison
 ///
 /// Implementations of this trait store a in-progress collection of group values
@@ -46,6 +46,8 @@ use datafusion_physical_expr_common::binary_map::{OutputType, INITIAL_BUFFER_CAP
 pub trait GroupColumn: Send + Sync {
     /// Returns equal if the row stored in this builder at `lhs_row` is equal to
     /// the row in `array` at `rhs_row`
+    ///
+    /// Note that this comparison returns true if both elements are NULL
     fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool;
     /// Appends the row at `row` in `array` to this builder
     fn append_val(&mut self, array: &ArrayRef, row: usize);
@@ -62,57 +64,60 @@ pub trait GroupColumn: Send + Sync {
 
 pub struct PrimitiveGroupValueBuilder<T: ArrowPrimitiveType> {
     group_values: Vec<T::Native>,
-    nulls: Vec<bool>,
-    // whether the array contains at least one null, for fast non-null path
-    has_null: bool,
-    nullable: bool,
+    /// Null state (when None, input is guaranteed not to have nulls)
+    nulls: Option<MaybeNullBufferBuilder>,
 }
 
 impl<T> PrimitiveGroupValueBuilder<T>
 where
     T: ArrowPrimitiveType,
 {
+    /// Create a new [`PrimitiveGroupValueBuilder`]
+    ///
+    /// If `nullable` is false, it means the input will never have nulls
     pub fn new(nullable: bool) -> Self {
+        let nulls = if nullable {
+            Some(MaybeNullBufferBuilder::new())
+        } else {
+            None
+        };
+
         Self {
             group_values: vec![],
-            nulls: vec![],
-            has_null: false,
-            nullable,
+            nulls,
         }
     }
 }
 
 impl<T: ArrowPrimitiveType> GroupColumn for PrimitiveGroupValueBuilder<T> {
     fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool {
-        // non-null fast path
-        // both non-null
-        if !self.nullable {
-            return self.group_values[lhs_row]
-                == array.as_primitive::<T>().value(rhs_row);
-        }
-
-        // lhs is non-null
-        if self.nulls[lhs_row] {
-            if array.is_null(rhs_row) {
-                return false;
+        // fast path when input has no nulls
+        match self.nulls.as_ref() {
+            None => {
+                self.group_values[lhs_row] == array.as_primitive::<T>().value(rhs_row)
+            }
+            Some(nulls) => {
+                // slower path if the input could have nulls
+                nulls.is_null(lhs_row) == array.is_null(rhs_row)
+                    && self.group_values[lhs_row]
+                        == array.as_primitive::<T>().value(rhs_row)
             }
-
-            return self.group_values[lhs_row]
-                == array.as_primitive::<T>().value(rhs_row);
         }
-
-        array.is_null(rhs_row)
     }
 
     fn append_val(&mut self, array: &ArrayRef, row: usize) {
-        if self.nullable && array.is_null(row) {
-            self.group_values.push(T::default_value());
-            self.nulls.push(false);
-            self.has_null = true;
-        } else {
-            let elem = array.as_primitive::<T>().value(row);
-            self.group_values.push(elem);
-            self.nulls.push(true);
+        match self.nulls.as_mut() {
+            // input can't possibly have nulls, so don't worry about them
+            None => self.group_values.push(array.as_primitive::<T>().value(row)),
+            Some(nulls) => {
+                if array.is_null(row) {
+                    nulls.append(true);
+                    self.group_values.push(T::default_value());
+                } else {
+                    nulls.append(false);
+                    self.group_values.push(array.as_primitive::<T>().value(row));
+                }
+            }
         }
     }
 
@@ -121,36 +126,37 @@ impl<T: ArrowPrimitiveType> GroupColumn for PrimitiveGroupValueBuilder<T> {
     }
 
     fn size(&self) -> usize {
-        self.group_values.allocated_size() + self.nulls.allocated_size()
+        let nulls_size = self
+            .nulls
+            .as_ref()
+            .map(|nulls| nulls.allocated_size())
+            .unwrap_or(0);
+
+        self.group_values.allocated_size() + nulls_size
     }
 
     fn build(self: Box<Self>) -> ArrayRef {
-        if self.has_null {
-            Arc::new(PrimitiveArray::<T>::new(
-                ScalarBuffer::from(self.group_values),
-                Some(NullBuffer::from(self.nulls)),
-            ))
-        } else {
-            Arc::new(PrimitiveArray::<T>::new(
-                ScalarBuffer::from(self.group_values),
-                None,
-            ))
-        }
+        let Self {
+            group_values,
+            nulls,
+        } = *self;
+
+        let nulls = nulls.and_then(|nulls| nulls.build());
+
+        Arc::new(PrimitiveArray::<T>::new(
+            ScalarBuffer::from(group_values),
+            nulls,
+        ))
     }
 
     fn take_n(&mut self, n: usize) -> ArrayRef {
-        if self.has_null {
-            let first_n = self.group_values.drain(0..n).collect::<Vec<_>>();
-            let first_n_nulls = self.nulls.drain(0..n).collect::<Vec<_>>();
-            Arc::new(PrimitiveArray::<T>::new(
-                ScalarBuffer::from(first_n),
-                Some(NullBuffer::from(first_n_nulls)),
-            ))
-        } else {
-            let first_n = self.group_values.drain(0..n).collect::<Vec<_>>();
-            self.nulls.truncate(self.nulls.len() - n);
-            Arc::new(PrimitiveArray::<T>::new(ScalarBuffer::from(first_n), None))
-        }
+        let first_n = self.group_values.drain(0..n).collect::<Vec<_>>();
+        let first_n_nulls = self.nulls.as_mut().and_then(|nulls| nulls.take_n(n));
+
+        Arc::new(PrimitiveArray::<T>::new(
+            ScalarBuffer::from(first_n),
+            first_n_nulls,
+        ))
     }
 }
 

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

@@ -36,6 +36,7 @@ use bytes::GroupValuesByes;
 use datafusion_physical_expr::binary_map::OutputType;
 
 mod group_column;
+mod null_builder;
 
 /// An interning store for group keys
 pub trait GroupValues: Send {

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

@@ -0,0 +1,115 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow_buffer::{BooleanBufferBuilder, NullBuffer};
+
+/// Builder for an (optional) null mask
+///
+/// Optimized for avoid creating the bitmask when all values are non-null
+#[derive(Debug)]
+pub(crate) enum MaybeNullBufferBuilder {
+    ///  seen `row_count` rows but no nulls yet
+    NoNulls { row_count: usize },
+    /// have at least one null value
+    ///
+    /// Note this is an Arrow *VALIDITY* buffer (so it is false for nulls, true
+    /// for non-nulls)
+    Nulls(BooleanBufferBuilder),
+}
+
+impl MaybeNullBufferBuilder {
+    /// Create a new builder
+    pub fn new() -> Self {
+        Self::NoNulls { row_count: 0 }
+    }
+
+    /// Return true if the row at index `row` is null
+    pub fn is_null(&self, row: usize) -> bool {
+        match self {
+            Self::NoNulls { .. } => false,
+            // validity mask means a unset bit is NULL
+            Self::Nulls(builder) => !builder.get_bit(row),
+        }
+    }
+
+    /// Set the nullness of the next row to `is_null`
+    ///
+    /// num_values is the current length of the rows being tracked
+    ///
+    /// If `value` is true, the row is null.
+    /// If `value` is false, the row is non null
+    pub fn append(&mut self, 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(false);
+                *self = Self::Nulls(nulls);
+            }
+            Self::NoNulls { row_count } => {
+                *row_count += 1;
+            }
+            Self::Nulls(builder) => builder.append(!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 {
+            Self::NoNulls { .. } => 0,
+            // BooleanBufferBuilder builder::capacity returns capacity in bits (not bytes)
+            Self::Nulls(builder) => builder.capacity() / 8,
+        }
+    }
+
+    /// Return a NullBuffer representing the accumulated nulls so far
+    pub fn build(self) -> Option<NullBuffer> {
+        match self {
+            Self::NoNulls { .. } => None,
+            Self::Nulls(mut builder) => Some(NullBuffer::from(builder.finish())),
+        }
+    }
+
+    /// Returns a NullBuffer representing the first `n` rows accumulated so far
+    /// shifting any remaining down by `n`
+    pub fn take_n(&mut self, n: usize) -> Option<NullBuffer> {
+        match self {
+            Self::NoNulls { row_count } => {
+                *row_count -= n;
+                None
+            }
+            Self::Nulls(builder) => {
+                // Copy over the values at  n..len-1 values to the start of a
+                // new builder and leave it in self
+                //
+                // TODO: it would be great to use something like `set_bits` from arrow here.
+                let mut new_builder = BooleanBufferBuilder::new(builder.len());
+                for i in n..builder.len() {
+                    new_builder.append(builder.get_bit(i));
+                }
+                std::mem::swap(&mut new_builder, builder);
+
+                // take only first n values from the original builder
+                new_builder.truncate(n);
+                Some(NullBuffer::from(new_builder.finish()))
+            }
+        }
+    }
+}