Optimize partition_validity function used in sort kernels

arrow-ord/src/sort.rs

@@ -180,13 +180,36 @@ where
 
 // partition indices into valid and null indices
 fn partition_validity(array: &dyn Array) -> (Vec<u32>, Vec<u32>) {
-    match array.null_count() {
-        // faster path
-        0 => ((0..(array.len() as u32)).collect(), vec![]),
-        _ => {
-            let indices = 0..(array.len() as u32);
-            indices.partition(|index| array.is_valid(*index as usize))
+    let len = array.len();
+    match array.nulls() {
+        Some(nulls) if nulls.null_count() > 0 => {
+            let mut valid_indices = Vec::with_capacity(len - nulls.null_count());
+            let mut null_indices = Vec::with_capacity(nulls.null_count());
+
+            let valid_slice = valid_indices.spare_capacity_mut();
+            let null_slice = null_indices.spare_capacity_mut();
+            let mut valid_idx = 0;
+            let mut null_idx = 0;
+
+            nulls.into_iter().enumerate().for_each(|(i, v)| {
+                if v {
+                    valid_slice[valid_idx].write(i as u32);
+                    valid_idx += 1;
+                } else {
+                    null_slice[null_idx].write(i as u32);
+                    null_idx += 1;
+                }
+            });
+
+            unsafe {
+                valid_indices.set_len(valid_idx);
+                null_indices.set_len(null_idx);
+            }
+
+            (valid_indices, null_indices)
         }
+        // faster path
+        _ => ((0..(len as u32)).collect(), vec![]),
     }
 }