perf: Improve unique performance by adding RangedUniqueKernel for primitive arrays

crates/polars-compute/src/unique/boolean.rs

@@ -1,113 +1,128 @@
 use arrow::array::{Array, BooleanArray};
 use arrow::bitmap::MutableBitmap;
+use arrow::datatypes::ArrowDataType;
 
-use super::UniqueKernel;
-
-fn bool_unique_fold<'a>(
-    fst: &'a BooleanArray,
-    arrs: impl Iterator<Item = &'a BooleanArray>,
-) -> BooleanArray {
-    // can be None, Some(true), Some(false)
-    //
-    // We assign values to each value
-    // None        = 1
-    // Some(false) = 2
-    // Some(true)  = 3
-    //
-    // And keep track of 2 things
-    // - `found_set`: which values have already appeared
-    // - `order`:     in which order did the values appear
-
-    #[inline(always)]
-    fn append_arr(arr: &BooleanArray, found_set: &mut u32, order: &mut u32) {
-        for v in arr {
-            let value = v.map_or(1, |v| 2 + u32::from(v));
-            let nulled_value = if *found_set & (1 << value) != 0 {
-                0
-            } else {
-                value
-            };
-
-            *order |= nulled_value << (found_set.count_ones() * 2);
-            *found_set |= 1 << value;
-
-            if *found_set == 0b1110 {
-                break;
-            }
+use super::{GenericUniqueKernel, RangedUniqueKernel};
+
+pub struct BooleanUniqueKernelState {
+    seen: u32,
+    has_null: bool,
+    data_type: ArrowDataType,
+}
+
+const fn to_value(scalar: Option<bool>) -> u8 {
+    match scalar {
+        None => 0,
+        Some(false) => 1,
+        Some(true) => 2,
+    }
+}
+
+impl BooleanUniqueKernelState {
+    pub fn new(has_null: bool, data_type: ArrowDataType) -> Self {
+        Self {
+            seen: 0,
+            has_null,
+            data_type,
         }
     }
 
-    let mut found_set = 0u32;
-    let mut order = 0u32;
+    fn has_seen_null(&self) -> bool {
+        self.has_null && self.seen & (1 << to_value(None)) != 0
+    }
+}
 
-    append_arr(fst, &mut found_set, &mut order);
-    for arr in arrs {
-        append_arr(arr, &mut found_set, &mut order);
+impl RangedUniqueKernel for BooleanUniqueKernelState {
+    type Array = BooleanArray;
+
+    fn has_seen_all(&self) -> bool {
+        self.seen == 0b111
     }
 
-    let mut values = MutableBitmap::with_capacity(3);
-    let validity = if found_set & 0b10 != 0 {
-        let mut validity = MutableBitmap::with_capacity(3);
-        while order != 0 {
-            values.push(order & 0b11 > 2);
-            validity.push(order & 0b11 > 1);
-            order >>= 2;
+    fn append(&mut self, array: &Self::Array) {
+        if array.len() == 0 {
+            return;
         }
-        Some(validity.freeze())
-    } else {
-        while order != 0 {
-            values.push(order & 0b11 > 2);
-            order >>= 2;
+
+        let null_count = array.null_count();
+        let values = array.values();
+
+        if !self.has_null || null_count == 0 {
+            let set_bits = values.set_bits();
+            self.seen |= u32::from(set_bits != 0) << to_value(Some(true));
+            self.seen |= u32::from(set_bits != values.len()) << to_value(Some(false));
+
+            return;
         }
-        None
-    };
 
-    let values = values.freeze();
+        self.seen |= u32::from(null_count > 0) << to_value(None);
 
-    BooleanArray::new(fst.data_type().clone(), values, validity)
-}
+        if array.len() != null_count {
+            let validity = array.validity().unwrap();
 
-impl UniqueKernel for BooleanArray {
-    fn unique_fold<'a>(fst: &'a Self, others: impl Iterator<Item = &'a Self>) -> Self {
-        bool_unique_fold(fst, others)
+            let set_bits = values.num_intersections_with(validity);
+            self.seen |= u32::from(set_bits != 0) << to_value(Some(true));
+            self.seen |= u32::from(set_bits != values.len() - null_count) << to_value(Some(false));
+        }
     }
 
-    fn unique(&self) -> Self {
-        Self::unique_fold(self, [].iter())
-    }
+    fn finalize_unique(self) -> Self::Array {
+        let mut values = MutableBitmap::with_capacity(3);
+        let validity = if self.has_seen_null() {
+            let mut validity = MutableBitmap::with_capacity(3);
 
-    fn unique_sorted(&self) -> Self {
-        Self::unique_fold(self, [].iter())
-    }
+            for i in 0..3 {
+                if self.seen & (1 << i) != 0 {
+                    values.push(i > 1);
+                    validity.push(i > 0);
+                }
+            }
 
-    fn n_unique(&self) -> usize {
-        if self.len() == 0 {
-            return 0;
-        }
+            Some(validity.freeze())
+        } else {
+            for i in 1..3 {
+                if self.seen & (1 << i) != 0 {
+                    values.push(i > 1);
+                }
+            }
 
-        let null_count = self.null_count();
+            None
+        };
 
-        if self.len() == null_count {
-            return 1;
-        }
+        let values = values.freeze();
 
-        let values = self.values();
+        BooleanArray::new(self.data_type, values, validity)
+    }
 
-        if null_count == 0 {
-            let unset_bits = values.unset_bits();
-            let is_uniform = unset_bits == 0 || unset_bits == values.len();
-            return 2 - usize::from(is_uniform);
-        }
+    fn finalize_n_unique(self) -> usize {
+        self.seen.count_ones() as usize
+    }
+
+    fn finalize_n_unique_non_null(self) -> usize {
+        (self.seen & !1).count_ones() as usize
+    }
+}
 
-        let validity = self.validity().unwrap();
-        let set_bits = values.num_intersections_with(validity);
-        let is_uniform = set_bits == 0 || set_bits == validity.set_bits();
-        2 + usize::from(!is_uniform)
+impl GenericUniqueKernel for BooleanArray {
+    fn unique(&self) -> Self {
+        let mut state =
+            BooleanUniqueKernelState::new(self.null_count() > 0, self.data_type().clone());
+        state.append(self);
+        state.finalize_unique()
+    }
+
+    fn n_unique(&self) -> usize {
+        let mut state =
+            BooleanUniqueKernelState::new(self.null_count() > 0, self.data_type().clone());
+        state.append(self);
+        state.finalize_n_unique()
     }
 
-    #[inline]
     fn n_unique_non_null(&self) -> usize {
-        self.n_unique() - usize::from(self.null_count() > 0)
+        let mut state =
+            BooleanUniqueKernelState::new(self.null_count() > 0, self.data_type().clone());
+        state.append(self);
+        state.finalize_n_unique_non_null()
     }
 }
 

crates/polars-compute/src/unique/mod.rs

@@ -1,7 +1,7 @@
 use arrow::array::Array;
 
-/// Kernel to calculate the number of unique elements
-pub trait UniqueKernel: Array {
+/// Kernel to calculate the number of unique elements where the elements are already sorted.
+pub trait SortedUniqueKernel: Array {
     /// Calculate the set of unique elements in `fst` and `others` and fold the result into one
     /// array.
     fn unique_fold<'a>(fst: &'a Self, others: impl Iterator<Item = &'a Self>) -> Self;
@@ -10,9 +10,6 @@ pub trait UniqueKernel: Array {
     /// `self`.
     fn unique(&self) -> Self;
 
-    /// Calculate the set of unique elements in [`Self`] where `self` is sorted.
-    fn unique_sorted(&self) -> Self;
-
     /// Calculate the number of unique elements in [`Self`]
     ///
     /// A null is also considered a unique value
@@ -22,4 +19,46 @@ pub trait UniqueKernel: Array {
     fn n_unique_non_null(&self) -> usize;
 }
 
+/// Optimized kernel to calculate the unique elements of an array.
+///
+/// This kernel is a specialized for where all values are known to be in some small range of
+/// values. In this case, you can usually get by with a bitset and bit-arithmetic instead of using
+/// vectors and hashsets. Consequently, this kernel is usually called when further information is
+/// known about the underlying array.
+///
+/// This trait is not implemented directly on the `Array` as with many other kernels. Rather, it is
+/// implemented on a `State` struct to which `Array`s can be appended. This allows for sharing of
+/// `State` between many chunks and allows for different implementations for the same array (e.g. a
+/// maintain order and no maintain-order variant).
+pub trait RangedUniqueKernel {
+    type Array: Array;
+
+    /// Returns whether all the values in the whole range are in the state
+    fn has_seen_all(&self) -> bool;
+
+    /// Append an `Array`'s values to the `State`
+    fn append(&mut self, array: &Self::Array);
+
+    /// Consume the state to get the unique elements
+    fn finalize_unique(self) -> Self::Array;
+    /// Consume the state to get the number of unique elements including null
+    fn finalize_n_unique(self) -> usize;
+    /// Consume the state to get the number of unique elements excluding null
+    fn finalize_n_unique_non_null(self) -> usize;
+}
+
+/// A generic unique kernel that selects the generally applicable unique kernel for an `Array`.
+pub trait GenericUniqueKernel {
+    /// Calculate the set of unique elements
+    fn unique(&self) -> Self;
+    /// Calculate the number of unique elements including null
+    fn n_unique(&self) -> usize;
+    /// Calculate the number of unique elements excluding null
+    fn n_unique_non_null(&self) -> usize;
+}
+
 mod boolean;
+mod primitive;
+
+pub use boolean::BooleanUniqueKernelState;
+pub use primitive::PrimitiveRangedUniqueState;