Move support_distinct next to distinct

Author: liamzwbao

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

@@ -21,9 +21,9 @@ use std::ops::Range;
 
 use arrow_array::{Array, ArrayRef};
 use arrow_buffer::BooleanBuffer;
-use arrow_schema::{ArrowError, DataType, SortOptions};
+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`]
@@ -152,23 +152,6 @@ pub fn partition(columns: &[ArrayRef]) -> Result<Partitions, ArrowError> {
     Ok(Partitions(Some(acc)))
 }
 
-/// 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.
-fn supports_distinct(dt: &DataType) -> bool {
-    let leaf = match dt {
-        DataType::Dictionary(_, v) => v.as_ref(),
-        dt => dt,
-    };
-    !leaf.is_nested()
-        && !matches!(
-            leaf,
-            DataType::Dictionary(_, _) | DataType::RunEndEncoded(_, _)
-        )
-}
-
 /// Returns a mask with bits set whenever the value or nullability changes
 fn find_boundaries(v: &dyn Array) -> Result<BooleanBuffer, ArrowError> {
     let slice_len = v.len() - 1;