refactor: improve run boundary computation in RunArray casting

Author: Weijun-H

arrow-cast/src/cast/run_array.rs

@@ -15,10 +15,21 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::vec;
-
 use crate::cast::*;
-use arrow_array::Array;
+use arrow_array::cast::AsArray;
+use arrow_array::types::{
+    ArrowDictionaryKeyType, ArrowPrimitiveType, Date32Type, Date64Type, Decimal128Type,
+    Decimal256Type, DurationMicrosecondType, DurationMillisecondType, DurationNanosecondType,
+    DurationSecondType, Float16Type, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type,
+    Int64Type, IntervalDayTimeType, IntervalMonthDayNanoType, IntervalYearMonthType,
+    Time32MillisecondType, Time32SecondType, Time64MicrosecondType, Time64NanosecondType,
+    TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType,
+    TimestampSecondType, UInt8Type, UInt16Type, UInt32Type, UInt64Type,
+};
+use arrow_array::{
+    Array, ArrayRef, BinaryViewArray, BooleanArray, DictionaryArray, FixedSizeBinaryArray,
+    GenericBinaryArray, GenericStringArray, OffsetSizeTrait, PrimitiveArray, StringViewArray,
+};
 
 /// Attempts to cast a `RunArray` with index type K into
 /// `to_type` for supported types.
@@ -137,18 +148,7 @@ pub(crate) fn cast_to_run_end_encoded<K: RunEndIndexType>(
     }
 
     // Identify run boundaries by comparing consecutive values
-    let mut run_ends = Vec::new();
-    let mut values_indexes = vec![0usize]; // Always include the first index
-    let mut current_data = cast_array.slice(0, 1).to_data();
-    for idx in 1..cast_array.len() {
-        let next_data = cast_array.slice(idx, 1).to_data();
-        if current_data != next_data {
-            run_ends.push(idx);
-            values_indexes.push(idx);
-            current_data = next_data;
-        }
-    }
-    run_ends.push(cast_array.len());
+    let (run_ends, values_indexes) = compute_run_boundaries(cast_array);
 
     // Build the run_ends array
     for run_end in run_ends {
@@ -167,3 +167,222 @@ pub(crate) fn cast_to_run_end_encoded<K: RunEndIndexType>(
     let run_array = RunArray::<K>::try_new(&run_ends_array, values_array.as_ref())?;
     Ok(Arc::new(run_array))
 }
+
+// Compute run boundaries for the given array.
+fn compute_run_boundaries(array: &ArrayRef) -> (Vec<usize>, Vec<usize>) {
+    use arrow_schema::{DataType::*, IntervalUnit, TimeUnit};
+    match array.data_type() {
+        Null => runs_for_null(array.len()),
+        Boolean => runs_for_boolean(array.as_boolean()),
+        Int8 => runs_for_primitive(array.as_primitive::<Int8Type>()),
+        Int16 => runs_for_primitive(array.as_primitive::<Int16Type>()),
+        Int32 => runs_for_primitive(array.as_primitive::<Int32Type>()),
+        Int64 => runs_for_primitive(array.as_primitive::<Int64Type>()),
+        UInt8 => runs_for_primitive(array.as_primitive::<UInt8Type>()),
+        UInt16 => runs_for_primitive(array.as_primitive::<UInt16Type>()),
+        UInt32 => runs_for_primitive(array.as_primitive::<UInt32Type>()),
+        UInt64 => runs_for_primitive(array.as_primitive::<UInt64Type>()),
+        Float16 => runs_for_primitive(array.as_primitive::<Float16Type>()),
+        Float32 => runs_for_primitive(array.as_primitive::<Float32Type>()),
+        Float64 => runs_for_primitive(array.as_primitive::<Float64Type>()),
+        Date32 => runs_for_primitive(array.as_primitive::<Date32Type>()),
+        Date64 => runs_for_primitive(array.as_primitive::<Date64Type>()),
+        Time32(TimeUnit::Second) => runs_for_primitive(array.as_primitive::<Time32SecondType>()),
+        Time32(TimeUnit::Millisecond) => {
+            runs_for_primitive(array.as_primitive::<Time32MillisecondType>())
+        }
+        Time64(TimeUnit::Microsecond) => {
+            runs_for_primitive(array.as_primitive::<Time64MicrosecondType>())
+        }
+        Time64(TimeUnit::Nanosecond) => {
+            runs_for_primitive(array.as_primitive::<Time64NanosecondType>())
+        }
+        Duration(TimeUnit::Second) => {
+            runs_for_primitive(array.as_primitive::<DurationSecondType>())
+        }
+        Duration(TimeUnit::Millisecond) => {
+            runs_for_primitive(array.as_primitive::<DurationMillisecondType>())
+        }
+        Duration(TimeUnit::Microsecond) => {
+            runs_for_primitive(array.as_primitive::<DurationMicrosecondType>())
+        }
+        Duration(TimeUnit::Nanosecond) => {
+            runs_for_primitive(array.as_primitive::<DurationNanosecondType>())
+        }
+        Timestamp(TimeUnit::Second, _) => {
+            runs_for_primitive(array.as_primitive::<TimestampSecondType>())
+        }
+        Timestamp(TimeUnit::Millisecond, _) => {
+            runs_for_primitive(array.as_primitive::<TimestampMillisecondType>())
+        }
+        Timestamp(TimeUnit::Microsecond, _) => {
+            runs_for_primitive(array.as_primitive::<TimestampMicrosecondType>())
+        }
+        Timestamp(TimeUnit::Nanosecond, _) => {
+            runs_for_primitive(array.as_primitive::<TimestampNanosecondType>())
+        }
+        Interval(IntervalUnit::YearMonth) => {
+            runs_for_primitive(array.as_primitive::<IntervalYearMonthType>())
+        }
+        Interval(IntervalUnit::DayTime) => {
+            runs_for_primitive(array.as_primitive::<IntervalDayTimeType>())
+        }
+        Interval(IntervalUnit::MonthDayNano) => {
+            runs_for_primitive(array.as_primitive::<IntervalMonthDayNanoType>())
+        }
+        Decimal128(_, _) => runs_for_primitive(array.as_primitive::<Decimal128Type>()),
+        Decimal256(_, _) => runs_for_primitive(array.as_primitive::<Decimal256Type>()),
+        Utf8 => runs_for_string(array.as_string::<i32>()),
+        LargeUtf8 => runs_for_string(array.as_string::<i64>()),
+        Utf8View => runs_for_string_view(array.as_string_view()),
+        Binary => runs_for_binary(array.as_binary::<i32>()),
+        LargeBinary => runs_for_binary(array.as_binary::<i64>()),
+        BinaryView => runs_for_binary_view(array.as_binary_view()),
+        FixedSizeBinary(_) => runs_for_fixed_size_binary(array.as_fixed_size_binary()),
+        Dictionary(key_type, _) => match key_type.as_ref() {
+            Int8 => runs_for_dictionary::<Int8Type>(array.as_dictionary()),
+            Int16 => runs_for_dictionary::<Int16Type>(array.as_dictionary()),
+            Int32 => runs_for_dictionary::<Int32Type>(array.as_dictionary()),
+            Int64 => runs_for_dictionary::<Int64Type>(array.as_dictionary()),
+            UInt8 => runs_for_dictionary::<UInt8Type>(array.as_dictionary()),
+            UInt16 => runs_for_dictionary::<UInt16Type>(array.as_dictionary()),
+            UInt32 => runs_for_dictionary::<UInt32Type>(array.as_dictionary()),
+            UInt64 => runs_for_dictionary::<UInt64Type>(array.as_dictionary()),
+            _ => runs_generic(array),
+        },
+        _ => runs_generic(array),
+    }
+}
+
+fn runs_for_null(len: usize) -> (Vec<usize>, Vec<usize>) {
+    (vec![len], vec![0])
+}
+
+fn runs_for_boolean(array: &BooleanArray) -> (Vec<usize>, Vec<usize>) {
+    build_runs_with_nulls(
+        array.len(),
+        array.null_count() > 0,
+        |idx| array.is_valid(idx),
+        |idx| array.value(idx),
+    )
+}
+
+fn runs_for_primitive<T: ArrowPrimitiveType>(
+    array: &PrimitiveArray<T>,
+) -> (Vec<usize>, Vec<usize>) {
+    build_runs_with_nulls(
+        array.len(),
+        array.null_count() > 0,
+        |idx| array.is_valid(idx),
+        |idx| array.value(idx),
+    )
+}
+
+fn runs_for_binary<O: OffsetSizeTrait>(array: &GenericBinaryArray<O>) -> (Vec<usize>, Vec<usize>) {
+    build_runs_with_nulls(
+        array.len(),
+        array.null_count() > 0,
+        |idx| array.is_valid(idx),
+        |idx| array.value(idx),
+    )
+}
+
+fn runs_for_string<O: OffsetSizeTrait>(array: &GenericStringArray<O>) -> (Vec<usize>, Vec<usize>) {
+    build_runs_with_nulls(
+        array.len(),
+        array.null_count() > 0,
+        |idx| array.is_valid(idx),
+        |idx| array.value(idx),
+    )
+}
+
+fn runs_for_binary_view(array: &BinaryViewArray) -> (Vec<usize>, Vec<usize>) {
+    build_runs_with_nulls(
+        array.len(),
+        array.null_count() > 0,
+        |idx| array.is_valid(idx),
+        |idx| array.value(idx),
+    )
+}
+
+fn runs_for_string_view(array: &StringViewArray) -> (Vec<usize>, Vec<usize>) {
+    build_runs_with_nulls(
+        array.len(),
+        array.null_count() > 0,
+        |idx| array.is_valid(idx),
+        |idx| array.value(idx),
+    )
+}
+
+fn runs_for_fixed_size_binary(array: &FixedSizeBinaryArray) -> (Vec<usize>, Vec<usize>) {
+    build_runs_with_nulls(
+        array.len(),
+        array.null_count() > 0,
+        |idx| array.is_valid(idx),
+        |idx| array.value(idx),
+    )
+}
+
+fn runs_for_dictionary<K: ArrowDictionaryKeyType>(
+    array: &DictionaryArray<K>,
+) -> (Vec<usize>, Vec<usize>) {
+    runs_for_primitive(array.keys())
+}
+
+fn build_runs(len: usize, mut equal: impl FnMut(usize, usize) -> bool) -> (Vec<usize>, Vec<usize>) {
+    let mut run_ends = Vec::new();
+    let mut values_indexes = vec![0usize];
+    if len == 0 {
+        return (run_ends, values_indexes);
+    }
+    let mut current = 0usize;
+    for idx in 1..len {
+        if !equal(current, idx) {
+            run_ends.push(idx);
+            values_indexes.push(idx);
+            current = idx;
+        }
+    }
+    run_ends.push(len);
+    (run_ends, values_indexes)
+}
+
+fn build_runs_with_nulls<T: PartialEq>(
+    len: usize,
+    has_nulls: bool,
+    mut is_valid: impl FnMut(usize) -> bool,
+    mut value_at: impl FnMut(usize) -> T,
+) -> (Vec<usize>, Vec<usize>) {
+    build_runs(len, |lhs, rhs| {
+        if has_nulls {
+            let lhs_valid = is_valid(lhs);
+            let rhs_valid = is_valid(rhs);
+            if lhs_valid && rhs_valid {
+                value_at(lhs) == value_at(rhs)
+            } else {
+                lhs_valid == rhs_valid
+            }
+        } else {
+            value_at(lhs) == value_at(rhs)
+        }
+    })
+}
+
+fn runs_generic(array: &ArrayRef) -> (Vec<usize>, Vec<usize>) {
+    let mut run_ends = Vec::new();
+    let mut values_indexes = vec![0usize];
+    if array.len() == 0 {
+        return (run_ends, values_indexes);
+    }
+    let mut current_data = array.slice(0, 1).to_data();
+    for idx in 1..array.len() {
+        let next_data = array.slice(idx, 1).to_data();
+        if current_data != next_data {
+            run_ends.push(idx);
+            values_indexes.push(idx);
+            current_data = next_data;
+        }
+    }
+    run_ends.push(array.len());
+    (run_ends, values_indexes)
+}