Add ree_distinct function that compares Run-End Encoded arrays directly without expanding to logical form.

Author: Rich-T-kid

arrow-arith/src/aggregate.rs

@@ -17,7 +17,7 @@
 
 //! Defines aggregations over Arrow arrays.
 
-use arrow_array::cast::{*};
+use arrow_array::cast::*;
 use arrow_array::iterator::ArrayIter;
 use arrow_array::*;
 use arrow_buffer::{ArrowNativeType, NullBuffer};
@@ -653,7 +653,7 @@ where
                 DataType::Int16 => AnyRunArray::new(&array, DataType::Int16),
                 _ => return Ok(None),
             };
-            
+
             if let Some(ree) = ree {
                 let mut sum = T::default_value();
 
@@ -666,15 +666,14 @@ where
                     let end = ree.run_ends_value(i);
                     let run_length = end - prev_end;
                     let run_length_native = T::Native::from_usize(run_length).unwrap();
-                    sum = sum.add_checked(values_data[i].mul_checked(run_length_native)?)?; 
+                    sum = sum.add_checked(values_data[i].mul_checked(run_length_native)?)?;
                     prev_end = end;
                 }
 
                 Ok(Some(sum))
             } else {
                 Ok(None)
             }
-            
         }
         _ => sum_checked::<T>(as_primitive_array(&array)),
     }
@@ -712,9 +711,7 @@ where
 {
     match array.data_type() {
         DataType::Dictionary(_, _) => min_max_helper::<T::Native, _, _>(array, cmp),
-        DataType::RunEndEncoded(_, _) => {
-            min_max_helper::<T::Native, _, _>(array, cmp)
-        }
+        DataType::RunEndEncoded(_, _) => min_max_helper::<T::Native, _, _>(array, cmp),
         _ => m(as_primitive_array(&array)),
     }
 }
@@ -1773,8 +1770,8 @@ mod tests {
     }
     mod ree_aggregation {
         use super::*;
-        use arrow_array::{RunArray, Int32Array, Int64Array, Float64Array};
-        use arrow_array::types::{Int32Type, Int64Type, Float64Type};
+        use arrow_array::types::{Float64Type, Int32Type, Int64Type};
+        use arrow_array::{Float64Array, Int32Array, Int64Array, RunArray};
 
         #[test]
         fn test_ree_sum_array_basic() {
@@ -1783,7 +1780,6 @@ mod tests {
             let values = Int32Array::from(vec![10, 20, 30]);
             let run_array = RunArray::<Int32Type>::try_new(&run_ends, &values).unwrap();
 
-            
             let typed_array = run_array.downcast::<Int32Array>().unwrap();
 
             let result = sum_array::<Int32Type, _>(typed_array);
@@ -1794,14 +1790,26 @@ mod tests {
         fn test_ree_sum_array_with_nulls() {
             // REE array with nulls: [10, NULL, 20, NULL, 30]
             let run_ends = Int32Array::from(vec![1, 2, 3, 4, 5]);
-            let values = Int32Array::from(vec![10, 0, 20, 0, 30]); // 0 represents null
+            let values = Int32Array::from(vec![Some(10), None, Some(20), None, Some(30)]);
             let run_array = RunArray::<Int32Type>::try_new(&run_ends, &values).unwrap();
 
             let typed_array = run_array.downcast::<Int32Array>().unwrap();
             let result = sum_array::<Int32Type, _>(typed_array);
             assert_eq!(result, Some(60)); // 10+20+30 = 60 (nulls ignored)
         }
 
+        #[test]
+        fn test_ree_sum_array_with_only_nulls() {
+            // REE array: [None, None, None, None, None] (logical length 5)
+            let run_ends = Int32Array::from(vec![1, 2, 3, 4, 5]);
+            let values = Int32Array::from(vec![None, None, None, None, None]);
+            let run_array = RunArray::<Int32Type>::try_new(&run_ends, &values).unwrap();
+
+            let typed_array = run_array.downcast::<Int32Array>().unwrap();
+            let result = sum_array::<Int32Type, _>(typed_array);
+            assert_eq!(result, Some(0)); // 0
+        }
+
         #[test]
         fn test_ree_sum_array_large_values() {
             // REE array with large values: [1000, 1000, 2000, 3000, 3000]

arrow-array/src/array/mod.rs

@@ -68,9 +68,6 @@ mod run_array;
 
 pub use run_array::*;
 
-// Re-export the unwrap_ree_array function for public use
-pub use run_array::unwrap_ree_array;
-
 mod byte_view_array;
 
 pub use byte_view_array::*;

arrow-array/src/array/run_array.rs

@@ -18,14 +18,24 @@
 use std::any::Any;
 use std::sync::Arc;
 
+use crate::{
+    builder::StringRunBuilder,
+    cast::AsArray,
+    make_array,
+    run_iterator::RunArrayIter,
+    types::{
+        Date32Type, Date64Type, Decimal128Type, Decimal256Type, DurationNanosecondType,
+        Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type, IntervalDayTimeType,
+        RunEndIndexType, Time32MillisecondType, Time64NanosecondType, TimestampMicrosecondType,
+        UInt16Type, UInt32Type, UInt64Type, UInt8Type,
+    },
+    Array, ArrayAccessor, ArrayRef, ArrowPrimitiveType, BooleanArray, Int16Array, Int32Array,
+    Int64Array, PrimitiveArray, StringArray,
+};
 use arrow_buffer::{ArrowNativeType, BooleanBufferBuilder, NullBuffer, RunEndBuffer};
 use arrow_data::{ArrayData, ArrayDataBuilder};
 use arrow_schema::{ArrowError, DataType, Field};
 
-use crate::{
-    builder::StringRunBuilder, cast::AsArray, make_array, run_iterator::RunArrayIter, types::{Int16Type, Int32Type, Int64Type, RunEndIndexType}, Array, ArrayAccessor, ArrayRef, ArrowPrimitiveType, PrimitiveArray, Int16Array, Int32Array, Int64Array
-};
-
 /// An array of [run-end encoded values](https://arrow.apache.org/docs/format/Columnar.html#run-end-encoded-layout)
 ///
 /// This encoding is variation on [run-length encoding (RLE)](https://en.wikipedia.org/wiki/Run-length_encoding)
@@ -247,14 +257,16 @@ impl<R: RunEndIndexType> RunArray<R> {
             values: self.values.clone(),
         }
     }
+
     /// Expands the REE array to its logical form
-    pub fn expand_to_logical<T: ArrowPrimitiveType>(&self) -> Result<Box<dyn Array>, ArrowError> 
+    pub fn expand_to_logical<T: ArrowPrimitiveType>(&self) -> Result<Box<dyn Array>, ArrowError>
     where
         T::Native: Default,
     {
-        let typed_ree = self.downcast::<PrimitiveArray<T>>()
-            .ok_or_else(|| ArrowError::InvalidArgumentError("Failed to downcast to typed REE".to_string()))?;
-        
+        let typed_ree = self.downcast::<PrimitiveArray<T>>().ok_or_else(|| {
+            ArrowError::InvalidArgumentError("Failed to downcast to typed REE".to_string())
+        })?;
+
         let mut builder = PrimitiveArray::<T>::builder(typed_ree.len());
         for i in 0..typed_ree.len() {
             if typed_ree.is_null(i) {
@@ -265,29 +277,6 @@ impl<R: RunEndIndexType> RunArray<R> {
         }
         Ok(Box::new(builder.finish()))
     }
-    /// Unwraps a REE array into a logical array
-    pub fn unwrap_ree_array(array: &dyn Array) -> Option<Box<dyn Array>> {
-    match array.data_type() {
-        arrow_schema::DataType::RunEndEncoded(run_ends_field, _) => {
-            match run_ends_field.data_type() {
-                arrow_schema::DataType::Int16 => {
-                    array.as_run_opt::<Int16Type>()
-                    .and_then(|ree| ree.expand_to_logical::<Int16Type>().ok())
-                }
-                arrow_schema::DataType::Int32 => {
-                    array.as_run_opt::<Int32Type>()
-                    .and_then(|ree| ree.expand_to_logical::<Int32Type>().ok())
-                }
-                arrow_schema::DataType::Int64 => {
-                    array.as_run_opt::<Int64Type>()
-                    .and_then(|ree| ree.expand_to_logical::<Int64Type>().ok())
-                }
-                _ => None,
-            }
-        }
-        _ => None,
-    }
-}
 }
 
 impl<R: RunEndIndexType> From<ArrayData> for RunArray<R> {
@@ -566,28 +555,140 @@ pub struct TypedRunArray<'a, R: RunEndIndexType, V> {
 }
 
 /// Unwraps a REE array into a logical array
-pub fn unwrap_ree_array(array: &dyn Array) -> Option<Box<dyn Array>> {
+pub fn ree_to_expanded_array(array: &dyn Array) -> Option<Box<dyn Array>> {
     match array.data_type() {
         arrow_schema::DataType::RunEndEncoded(run_ends_field, _) => {
             match run_ends_field.data_type() {
-                arrow_schema::DataType::Int16 => {
-                    array.as_run_opt::<Int16Type>()
-                    .and_then(|ree| ree.expand_to_logical::<Int16Type>().ok())
-                }
-                arrow_schema::DataType::Int32 => {
-                    array.as_run_opt::<Int32Type>()
-                    .and_then(|ree| ree.expand_to_logical::<Int32Type>().ok())
-                }
-                arrow_schema::DataType::Int64 => {
-                    array.as_run_opt::<Int64Type>()
-                    .and_then(|ree| ree.expand_to_logical::<Int64Type>().ok())
-                }
+                arrow_schema::DataType::Int16 => array
+                    .as_run_opt::<Int16Type>()
+                    .and_then(|ree| ree.expand_to_logical::<Int16Type>().ok()),
+                arrow_schema::DataType::Int32 => array
+                    .as_run_opt::<Int32Type>()
+                    .and_then(|ree| ree.expand_to_logical::<Int32Type>().ok()),
+                arrow_schema::DataType::Int64 => array
+                    .as_run_opt::<Int64Type>()
+                    .and_then(|ree| ree.expand_to_logical::<Int64Type>().ok()),
                 _ => None,
             }
         }
         _ => None,
     }
 }
+
+/// Generate a boolean array that indicates if two run arrays are equal
+pub fn ree_distinct(
+    lhs: &AnyRunArray,
+    rhs: &AnyRunArray,
+    size: usize,
+    flag: bool,
+) -> Option<BooleanArray> {
+    // Iterate through both run arrays and compare their logical indices
+    // we know that the run arrays of the exact same size.
+    let lhs_vals = lhs.values();
+    let rhs_vals = rhs.values();
+    if lhs_vals.data_type() != rhs_vals.data_type() {
+        return None;
+    }
+    match lhs_vals.data_type() {
+        // Integer types
+        DataType::Int8 => ree_distinct_primitive::<Int8Type>(lhs, rhs, size, flag),
+        DataType::Int16 => ree_distinct_primitive::<Int16Type>(lhs, rhs, size, flag),
+        DataType::Int32 => ree_distinct_primitive::<Int32Type>(lhs, rhs, size, flag),
+        DataType::Int64 => ree_distinct_primitive::<Int64Type>(lhs, rhs, size, flag),
+        DataType::UInt8 => ree_distinct_primitive::<UInt8Type>(lhs, rhs, size, flag),
+        DataType::UInt16 => ree_distinct_primitive::<UInt16Type>(lhs, rhs, size, flag),
+        DataType::UInt32 => ree_distinct_primitive::<UInt32Type>(lhs, rhs, size, flag),
+        DataType::UInt64 => ree_distinct_primitive::<UInt64Type>(lhs, rhs, size, flag),
+
+        // Floating point
+        DataType::Float32 => ree_distinct_primitive::<Float32Type>(lhs, rhs, size, flag),
+        DataType::Float64 => ree_distinct_primitive::<Float64Type>(lhs, rhs, size, flag),
+
+        // Temporal
+        DataType::Date32 => ree_distinct_primitive::<Date32Type>(lhs, rhs, size, flag),
+        DataType::Date64 => ree_distinct_primitive::<Date64Type>(lhs, rhs, size, flag),
+        DataType::Timestamp(_, _) => {
+            ree_distinct_primitive::<TimestampMicrosecondType>(lhs, rhs, size, flag)
+        }
+        DataType::Time32(_) => {
+            ree_distinct_primitive::<Time32MillisecondType>(lhs, rhs, size, flag)
+        }
+        DataType::Time64(_) => ree_distinct_primitive::<Time64NanosecondType>(lhs, rhs, size, flag),
+        DataType::Duration(_) => {
+            ree_distinct_primitive::<DurationNanosecondType>(lhs, rhs, size, flag)
+        }
+        DataType::Interval(_) => {
+            ree_distinct_primitive::<IntervalDayTimeType>(lhs, rhs, size, flag)
+        }
+
+        // Decimals
+        DataType::Decimal128(_, _) => {
+            ree_distinct_primitive::<Decimal128Type>(lhs, rhs, size, flag)
+        }
+        DataType::Decimal256(_, _) => {
+            ree_distinct_primitive::<Decimal256Type>(lhs, rhs, size, flag)
+        }
+        // Strings arent a primitive type, so we need to handle them separately
+        DataType::Utf8 => ree_distinct_string(lhs, rhs, size, flag),
+
+        // Not yet supported or complex
+        _ => None,
+    }
+}
+
+fn ree_distinct_primitive<T: ArrowPrimitiveType>(
+    lhs: &AnyRunArray,
+    rhs: &AnyRunArray,
+    size: usize,
+    flag: bool,
+) -> Option<BooleanArray> {
+    let lhs_vals = lhs.values().as_any().downcast_ref::<PrimitiveArray<T>>()?;
+    let rhs_vals = rhs.values().as_any().downcast_ref::<PrimitiveArray<T>>()?;
+    let mut builder = BooleanBufferBuilder::new(size);
+    for i in 0..size {
+        let li = lhs.get_physical_index(i);
+        let ri = rhs.get_physical_index(i);
+
+        let mut is_same = match (lhs_vals.is_null(li), rhs_vals.is_null(ri)) {
+            (true, true) => true,
+            (true, false) | (false, true) => false, // If one is null, result depends on flag
+            (false, false) => lhs_vals.value(li) == rhs_vals.value(ri),
+        };
+        if flag {
+            is_same = !is_same;
+        }
+        builder.append(is_same);
+    }
+    Some(BooleanArray::from(builder.finish()))
+}
+
+fn ree_distinct_string(
+    lhs: &AnyRunArray,
+    rhs: &AnyRunArray,
+    size: usize,
+    flag: bool,
+) -> Option<BooleanArray> {
+    let lhs_vals = lhs.values().as_any().downcast_ref::<StringArray>()?;
+    let rhs_vals = rhs.values().as_any().downcast_ref::<StringArray>()?;
+
+    let mut builder = BooleanBufferBuilder::new(size);
+    for i in 0..size {
+        let li = lhs.get_physical_index(i);
+        let ri = rhs.get_physical_index(i);
+
+        let mut is_same = match (lhs_vals.is_null(li), rhs_vals.is_null(ri)) {
+            (true, true) => true,
+            (true, false) | (false, true) => false,
+            (false, false) => lhs_vals.value(li) == rhs_vals.value(ri),
+        };
+        if flag {
+            is_same = !is_same;
+        }
+        builder.append(is_same);
+    }
+    Some(BooleanArray::from(builder.finish()))
+}
+
 // Manually implement `Clone` to avoid `V: Clone` type constraint
 impl<R: RunEndIndexType, V> Clone for TypedRunArray<'_, R, V> {
     fn clone(&self) -> Self {
@@ -720,7 +821,6 @@ where
     }
 }
 
-
 /// An AnyRunArray is a wrapper around a RunArray that can be used to aggregate over a RunEndEncodedArray
 /// This is used to avoid the need to downcast the RunEndEncodedArray to a specific type
 pub enum AnyRunArray<'a> {
@@ -777,7 +877,7 @@ impl<'a> AnyRunArray<'a> {
             AnyRunArray::Int16(array) => array.run_ends().values()[i].as_usize(),
         }
     }
-    
+
     /// Returns the length of run ends array
     pub fn run_ends_len(&self) -> usize {
         match self {
@@ -786,9 +886,16 @@ impl<'a> AnyRunArray<'a> {
             AnyRunArray::Int16(array) => array.values().len(),
         }
     }
-    
-}
 
+    /// Returns the physical index for the given logical index
+    pub fn get_physical_index(&self, logical_index: usize) -> usize {
+        match self {
+            AnyRunArray::Int64(array) => array.get_physical_index(logical_index),
+            AnyRunArray::Int32(array) => array.get_physical_index(logical_index),
+            AnyRunArray::Int16(array) => array.get_physical_index(logical_index),
+        }
+    }
+}
 
 #[cfg(test)]
 mod tests {