diff --git a/arrow-ord/src/cmp.rs b/arrow-ord/src/cmp.rs index e4d8498276a2..54b3b505d0f6 100644 --- a/arrow-ord/src/cmp.rs +++ b/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 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 { @@ -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); diff --git a/arrow-ord/src/partition.rs b/arrow-ord/src/partition.rs index fa00edab69fd..ba4210fb76a6 100644 --- a/arrow-ord/src/partition.rs +++ b/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 { 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![