Support nested REE in arrow-ord partition function

arrow-ord/src/cmp.rs

@@ -31,7 +31,7 @@ use arrow_array::{
 };
 use arrow_buffer::bit_util::ceil;
 use arrow_buffer::{BooleanBuffer, NullBuffer};
-use arrow_schema::ArrowError;
+use arrow_schema::{ArrowError, DataType};
 use arrow_select::take::take;
 use std::cmp::Ordering;
 use std::ops::Not;
@@ -201,6 +201,20 @@ pub fn not_distinct(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<BooleanArray, Ar
     compare_op(Op::NotDistinct, lhs, rhs)
 }
 
+/// Returns true if `distinct` (via `compare_op`) can handle this data type.
+///
+/// `compare_op` unwraps at most one level of dictionary, then dispatches on
+/// the leaf type. Anything else (REE, nested dictionary, nested/complex types)
+/// must go through `make_comparator` instead.
+pub(crate) fn supports_distinct(dt: &DataType) -> bool {
+    use arrow_schema::DataType::*;
+    let leaf = match dt {
+        Dictionary(_, v) => v.as_ref(),
+        dt => dt,
+    };
+    !leaf.is_nested() && !matches!(leaf, Dictionary(_, _) | RunEndEncoded(_, _))
+}
+
 /// Perform `op` on the provided `Datum`
 #[inline(never)]
 fn compare_op(op: Op, lhs: &dyn Datum, rhs: &dyn Datum) -> Result<BooleanArray, ArrowError> {
@@ -832,6 +846,38 @@ mod tests {
         assert_eq!(not_distinct(&b, &a).unwrap(), expected);
     }
 
+    #[test]
+    fn test_supports_distinct() {
+        use arrow_schema::{DataType::*, Field};
+
+        assert!(supports_distinct(&Int32));
+        assert!(supports_distinct(&Float64));
+        assert!(supports_distinct(&Utf8));
+        assert!(supports_distinct(&Boolean));
+
+        // One level of dictionary unwrap is supported.
+        assert!(supports_distinct(&Dictionary(
+            Box::new(Int16),
+            Box::new(Utf8),
+        )));
+
+        // REE, nested dictionary, and complex types are not supported.
+        assert!(!supports_distinct(&RunEndEncoded(
+            Arc::new(Field::new("run_ends", Int32, false)),
+            Arc::new(Field::new("values", Int32, true)),
+        )));
+        assert!(!supports_distinct(&Dictionary(
+            Box::new(Int16),
+            Box::new(Dictionary(Box::new(Int8), Box::new(Utf8))),
+        )));
+        assert!(!supports_distinct(&List(Arc::new(Field::new(
+            "item", Int32, true,
+        )))));
+        assert!(!supports_distinct(&Struct(
+            vec![Field::new("a", Int32, true)].into()
+        )));
+    }
+
     #[test]
     fn test_scalar_negation() {
         let a = Int32Array::new_scalar(54);

arrow-ord/src/partition.rs

@@ -23,7 +23,7 @@ use arrow_array::{Array, ArrayRef};
 use arrow_buffer::BooleanBuffer;
 use arrow_schema::{ArrowError, SortOptions};
 
-use crate::cmp::distinct;
+use crate::cmp::{distinct, supports_distinct};
 use crate::ord::make_comparator;
 
 /// A computed set of partitions, see [`partition`]
@@ -158,7 +158,7 @@ fn find_boundaries(v: &dyn Array) -> Result<BooleanBuffer, ArrowError> {
     let v1 = v.slice(0, slice_len);
     let v2 = v.slice(1, slice_len);
 
-    if !v.data_type().is_nested() {
+    if supports_distinct(v.data_type()) {
         return Ok(distinct(&v1, &v2)?.values().clone());
     }
     // Given that we're only comparing values, null ordering in the input or
@@ -306,6 +306,88 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_partition_run_end_encoded() {
+        let run_ends = Int32Array::from(vec![2, 3, 5]);
+        let values = StringArray::from(vec!["x", "y", "x"]);
+        let ree = RunArray::try_new(&run_ends, &values).unwrap();
+        // logical: ["x", "x", "y", "x", "x"]
+        let input = vec![Arc::new(ree) as _];
+        assert_eq!(partition(&input).unwrap().ranges(), vec![0..2, 2..3, 3..5],);
+    }
+
+    #[test]
+    fn test_partition_nested_run_end_encoded() {
+        // Inner REE (values of the outer): run_ends [1, 2, 3], values ["x", "y", "x"]
+        // logical length 3: ["x", "y", "x"]
+        let inner_run_ends = Int32Array::from(vec![1, 2, 3]);
+        let inner_values = StringArray::from(vec!["x", "y", "x"]);
+        let inner_ree = RunArray::try_new(&inner_run_ends, &inner_values).unwrap();
+
+        // Outer REE: run_ends [2, 3, 5], values = inner_ree (length 3)
+        // logical: rows 0,1 → inner[0]="x", row 2 → inner[1]="y", rows 3,4 → inner[2]="x"
+        // = ["x", "x", "y", "x", "x"]
+        let outer_run_ends = Int32Array::from(vec![2, 3, 5]);
+        let outer_ree = RunArray::try_new(&outer_run_ends, &inner_ree).unwrap();
+
+        let input = vec![Arc::new(outer_ree) as ArrayRef];
+        assert_eq!(partition(&input).unwrap().ranges(), vec![0..2, 2..3, 3..5]);
+    }
+
+    #[test]
+    fn test_partition_ree_with_dictionary_values() {
+        // Dictionary values: keys [0, 1, 0], dict ["x", "y"] → logical ["x", "y", "x"]
+        let dict_values = StringArray::from(vec!["x", "y"]);
+        let keys = Int32Array::from(vec![0, 1, 0]);
+        let dict = DictionaryArray::try_new(keys, Arc::new(dict_values)).unwrap();
+
+        // REE wrapping dict: run_ends [2, 3, 5] → logical [dict[0], dict[0], dict[1], dict[2], dict[2]]
+        // = ["x", "x", "y", "x", "x"]
+        let run_ends = Int32Array::from(vec![2, 3, 5]);
+        let ree = RunArray::try_new(&run_ends, &dict).unwrap();
+        let input = vec![Arc::new(ree) as ArrayRef];
+        assert_eq!(partition(&input).unwrap().ranges(), vec![0..2, 2..3, 3..5],);
+    }
+
+    #[test]
+    fn test_partition_dictionary() {
+        let values = StringArray::from(vec!["x", "y"]);
+        let keys = Int32Array::from(vec![0, 0, 1, 0, 0]);
+        let dict = DictionaryArray::try_new(keys, Arc::new(values)).unwrap();
+        // logical: ["x", "x", "y", "x", "x"]
+        let input = vec![Arc::new(dict) as _];
+        assert_eq!(partition(&input).unwrap().ranges(), vec![0..2, 2..3, 3..5],);
+    }
+
+    #[test]
+    fn test_partition_nested_dictionary() {
+        let inner_values = StringArray::from(vec!["x", "y"]);
+        let inner_keys = Int32Array::from(vec![0, 1, 0]);
+        let inner_dict = DictionaryArray::try_new(inner_keys, Arc::new(inner_values)).unwrap();
+
+        // Outer dict keys index into inner dict's logical values: ["x", "y", "x"]
+        // keys [0, 0, 1, 2, 2] → logical ["x", "x", "y", "x", "x"]
+        let outer_keys = Int32Array::from(vec![0, 0, 1, 2, 2]);
+        let outer_dict = DictionaryArray::try_new(outer_keys, Arc::new(inner_dict)).unwrap();
+        let input = vec![Arc::new(outer_dict) as ArrayRef];
+        assert_eq!(partition(&input).unwrap().ranges(), vec![0..2, 2..3, 3..5],);
+    }
+
+    #[test]
+    fn test_partition_dictionary_with_ree_values() {
+        // REE values: run_ends [2, 3], values ["x", "y"] → logical ["x", "x", "y"]
+        let run_ends = Int32Array::from(vec![2, 3]);
+        let str_values = StringArray::from(vec!["x", "y"]);
+        let ree = RunArray::try_new(&run_ends, &str_values).unwrap();
+
+        // Dictionary keys index into the REE's logical values
+        // keys [0, 0, 2, 0, 0] → logical ["x", "x", "y", "x", "x"]
+        let keys = Int32Array::from(vec![0, 0, 2, 0, 0]);
+        let dict = DictionaryArray::try_new(keys, Arc::new(ree)).unwrap();
+        let input = vec![Arc::new(dict) as ArrayRef];
+        assert_eq!(partition(&input).unwrap().ranges(), vec![0..2, 2..3, 3..5],);
+    }
+
     #[test]
     fn test_partition_nested() {
         let input = vec![