Add hashing of REE arrays and add REE aggregation example

Author: brancz

datafusion-examples/examples/ree_agg.rs

@@ -0,0 +1,73 @@
+use arrow::array::{Array, Int32Array, RunArray, StringViewArray};
+use arrow::datatypes::{DataType, Field, Int32Type, Schema};
+use arrow::record_batch::RecordBatch;
+use datafusion::datasource::MemTable;
+use datafusion::prelude::*;
+use std::sync::Arc;
+
+#[tokio::main]
+async fn main() -> Result<(), datafusion::error::DataFusionError> {
+    // Create a new DataFusion context
+    let ctx = SessionContext::new();
+
+    // First, let's create our data
+    // We'll have temperature readings where multiple consecutive readings come from the same sensor
+
+    // Temperature values (not run-length encoded)
+    // This represents all temperature readings in sequence
+    let temperatures = Int32Array::from(vec![
+        22, 23, 24, 25, 22, 21, 20, 21, 22, 23, 24, 25, 26, 27, 28,
+    ]);
+
+    // Create the string values for sensor IDs
+    let sensor_id_values =
+        StringViewArray::from(vec!["sensor_A", "sensor_B", "sensor_C", "sensor_D"]);
+
+    // Create the run ends array (positions where each run ends)
+    let sensor_id_run_ends = Int32Array::from(vec![4, 7, 12, 15]);
+
+    // Create RunArray for sensor IDs with Int32Type as run end type
+    let sensor_id_ree =
+        RunArray::<Int32Type>::try_new(&sensor_id_run_ends, &sensor_id_values)
+            .expect("Failed to create sensor ID RunArray");
+
+    // Get the exact data type of the RunArray for the schema
+    let sensor_id_type = sensor_id_ree.data_type().clone();
+
+    // Create schema
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("sensor_id", sensor_id_type, false),
+        Field::new("temperature", DataType::Int32, false),
+    ]));
+
+    // Create record batch
+    let batch = RecordBatch::try_new(
+        schema.clone(),
+        vec![Arc::new(sensor_id_ree), Arc::new(temperatures)],
+    )?;
+
+    // Register as a table
+    let provider = MemTable::try_new(schema, vec![vec![batch]])?;
+    ctx.register_table("sensor_readings", Arc::new(provider))?;
+
+    // Run aggregation query
+    // Group by sensor ID and calculate statistics
+    let sql = "
+        SELECT
+            sensor_id,
+            AVG(temperature) AS avg_temp,
+            MIN(temperature) AS min_temp,
+            MAX(temperature) AS max_temp,
+            COUNT(temperature) AS reading_count
+        FROM sensor_readings
+        GROUP BY sensor_id
+        ORDER BY sensor_id
+    ";
+
+    let results = ctx.sql(sql).await?.collect().await?;
+    for batch in results {
+        println!("{:?}", batch);
+    }
+
+    Ok(())
+}

datafusion/common/src/cast.rs

@@ -27,15 +27,16 @@ use arrow::array::{
     Int16Array, Int8Array, LargeBinaryArray, LargeListViewArray, LargeStringArray,
     ListViewArray, StringViewArray, UInt16Array,
 };
+use arrow::datatypes::RunEndIndexType;
 use arrow::{
     array::{
         Array, BinaryArray, BooleanArray, Date32Array, Date64Array, Decimal128Array,
         Decimal256Array, DictionaryArray, FixedSizeBinaryArray, FixedSizeListArray,
         Float32Array, Float64Array, GenericBinaryArray, GenericListArray,
         GenericStringArray, Int32Array, Int64Array, IntervalDayTimeArray,
         IntervalMonthDayNanoArray, IntervalYearMonthArray, LargeListArray, ListArray,
-        MapArray, NullArray, OffsetSizeTrait, PrimitiveArray, StringArray, StructArray,
-        Time32MillisecondArray, Time32SecondArray, Time64MicrosecondArray,
+        MapArray, NullArray, OffsetSizeTrait, PrimitiveArray, RunArray, StringArray,
+        StructArray, Time32MillisecondArray, Time32SecondArray, Time64MicrosecondArray,
         Time64NanosecondArray, TimestampMicrosecondArray, TimestampMillisecondArray,
         TimestampNanosecondArray, TimestampSecondArray, UInt32Array, UInt64Array,
         UInt8Array, UnionArray,
@@ -334,3 +335,8 @@ pub fn as_list_view_array(array: &dyn Array) -> Result<&ListViewArray> {
 pub fn as_large_list_view_array(array: &dyn Array) -> Result<&LargeListViewArray> {
     Ok(downcast_value!(array, LargeListViewArray))
 }
+
+// Downcast ArrayRef to RunArray
+pub fn as_run_array<R: RunEndIndexType>(array: &dyn Array) -> Result<&RunArray<R>> {
+    Ok(downcast_value!(array, RunArray, R))
+}

datafusion/common/src/hash_utils.rs

@@ -28,11 +28,12 @@ use arrow::{downcast_dictionary_array, downcast_primitive_array};
 use crate::cast::{
     as_binary_view_array, as_boolean_array, as_fixed_size_list_array,
     as_generic_binary_array, as_large_list_array, as_list_array, as_map_array,
-    as_string_array, as_string_view_array, as_struct_array, as_union_array,
+    as_run_array, as_string_array, as_string_view_array, as_struct_array, as_union_array,
 };
 use crate::error::Result;
 use crate::error::{_internal_datafusion_err, _internal_err};
 use std::cell::RefCell;
+use std::sync::Arc;
 
 // Combines two hashes into one hash
 #[inline]
@@ -484,6 +485,50 @@ fn hash_fixed_list_array(
     Ok(())
 }
 
+#[cfg(not(feature = "force_hash_collisions"))]
+fn hash_run_array<R: RunEndIndexType>(
+    array: &RunArray<R>,
+    random_state: &RandomState,
+    hashes_buffer: &mut [u64],
+    rehash: bool,
+) -> Result<()> {
+    // Get the physical values and create a clone for hashing
+    let values = array.values();
+    let values_len = values.len();
+    let mut values_hashes = vec![0u64; values_len];
+    create_hashes(&[Arc::clone(values)], random_state, &mut values_hashes)?;
+
+    // Get run ends buffer
+    let run_ends = array.run_ends();
+
+    // Previous run end (starting at 0 for the first run)
+    let mut prev_run_end = 0;
+
+    // Process each run - the number of runs equals the number of physical values
+    for i in 0..values_len {
+        // Get the run end for this run by accessing the run_ends buffer at the corresponding index
+        let run_end = run_ends.values()[i].as_usize();
+        // Apply the hash of the value at index i to all positions in this run
+        let value_hash = values_hashes[i];
+
+        // Apply this hash to all positions in the current run
+        if rehash {
+            for pos in prev_run_end..run_end {
+                hashes_buffer[pos] = combine_hashes(value_hash, hashes_buffer[pos]);
+            }
+        } else {
+            for pos in prev_run_end..run_end {
+                hashes_buffer[pos] = value_hash;
+            }
+        }
+
+        // Update the previous run end
+        prev_run_end = run_end;
+    }
+
+    Ok(())
+}
+
 /// Internal helper function that hashes a single array and either initializes or combines
 /// the hash values in the buffer.
 #[cfg(not(feature = "force_hash_collisions"))]
@@ -535,6 +580,23 @@ fn hash_single_array(
             let array = as_union_array(array)?;
             hash_union_array(array, random_state, hashes_buffer)?;
         }
+        DataType::RunEndEncoded(run_ends_type, _) => {
+            match run_ends_type.data_type() {
+                DataType::Int16 => {
+                    let array = as_run_array::<Int16Type>(array)?;
+                    hash_run_array(array, random_state, hashes_buffer, rehash)?;
+                }
+                DataType::Int32 => {
+                    let array = as_run_array::<Int32Type>(array)?;
+                    hash_run_array(array, random_state, hashes_buffer, rehash)?;
+                }
+                DataType::Int64 => {
+                    let array = as_run_array::<Int64Type>(array)?;
+                    hash_run_array(array, random_state, hashes_buffer, rehash)?;
+                }
+                _ => unreachable!("RunEndEncoded must have Int16, Int32, or Int64 run ends")
+            }
+        }
         _ => {
             // This is internal because we should have caught this before.
             return _internal_err!(
@@ -803,6 +865,91 @@ mod tests {
     create_hash_string!(string_view_array, StringArray);
     create_hash_string!(dict_string_array, DictionaryArray<Int8Type>);
 
+    #[test]
+    #[cfg(not(feature = "force_hash_collisions"))]
+    fn create_hashes_for_run_array() -> Result<()> {
+        // Create values and run_ends for RunArray
+        let values = Arc::new(Int32Array::from(vec![10, 20, 30]));
+        let run_ends = Arc::new(Int32Array::from(vec![2, 5, 7]));
+
+        // Create the RunArray
+        let array = Arc::new(RunArray::try_new(&run_ends, values.as_ref()).unwrap());
+
+        // Create hashes
+        let random_state = RandomState::with_seeds(0, 0, 0, 0);
+        let hashes_buff = &mut vec![0; array.len()];
+        let hashes =
+            create_hashes(&[array.clone() as ArrayRef], &random_state, hashes_buff)?;
+
+        // The length should be 7 (last run end)
+        assert_eq!(hashes.len(), 7);
+
+        // Values at indices 0,1 should have the same hash (value 10)
+        assert_eq!(hashes[0], hashes[1]);
+
+        // Values at indices 2,3,4 should have the same hash (value 20)
+        assert_eq!(hashes[2], hashes[3]);
+        assert_eq!(hashes[3], hashes[4]);
+
+        // Values at indices 5,6 should have the same hash (value 30)
+        assert_eq!(hashes[5], hashes[6]);
+
+        // Hashes for different values should be different
+        assert_ne!(hashes[0], hashes[2]);
+        assert_ne!(hashes[2], hashes[5]);
+        assert_ne!(hashes[0], hashes[5]);
+
+        Ok(())
+    }
+
+    #[test]
+    #[cfg(not(feature = "force_hash_collisions"))]
+    fn create_multi_column_hash_with_run_array() -> Result<()> {
+        // Create a regular Int32Array
+        let int_array = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7]));
+
+        // Create a RunArray with same logical length
+        let values = Arc::new(StringArray::from(vec!["foo", "bar", "baz"]));
+        let run_ends = Arc::new(Int32Array::from(vec![2, 5, 7]));
+        let run_array = Arc::new(RunArray::try_new(&run_ends, values.as_ref()).unwrap());
+
+        // Create hashes for single column (int array)
+        let random_state = RandomState::with_seeds(0, 0, 0, 0);
+        let mut one_col_hashes = vec![0; int_array.len()];
+        create_hashes(
+            &[int_array.clone() as ArrayRef],
+            &random_state,
+            &mut one_col_hashes,
+        )?;
+
+        // Create hashes for both columns
+        let mut two_col_hashes = vec![0; int_array.len()];
+        create_hashes(
+            &[int_array.clone() as ArrayRef, run_array.clone() as ArrayRef],
+            &random_state,
+            &mut two_col_hashes,
+        )?;
+
+        // Verify lengths
+        assert_eq!(one_col_hashes.len(), 7);
+        assert_eq!(two_col_hashes.len(), 7);
+
+        // Hashes should be different when including the Run array
+        assert_ne!(one_col_hashes, two_col_hashes);
+
+        // Indices 0,1 should have the same Run value ("foo") so their rehash pattern should be consistent
+        let diff_0_vs_1_one_col = one_col_hashes[0] != one_col_hashes[1];
+        let diff_0_vs_1_two_col = two_col_hashes[0] != two_col_hashes[1];
+        assert_eq!(diff_0_vs_1_one_col, diff_0_vs_1_two_col);
+
+        // Similarly for indices with the same Run value ("bar")
+        let diff_2_vs_3_one_col = one_col_hashes[2] != one_col_hashes[3];
+        let diff_2_vs_3_two_col = two_col_hashes[2] != two_col_hashes[3];
+        assert_eq!(diff_2_vs_3_one_col, diff_2_vs_3_two_col);
+
+        Ok(())
+    }
+
     #[test]
     // Tests actual values of hashes, which are different if forcing collisions
     #[cfg(not(feature = "force_hash_collisions"))]