feat: nulls first kernels

daft/dataframe/dataframe.py

@@ -1709,6 +1709,7 @@ def sort(
         self,
         by: Union[ColumnInputType, List[ColumnInputType]],
         desc: Union[bool, List[bool]] = False,
+        nulls_first: Optional[Union[bool, List[bool]]] = None,
     ) -> "DataFrame":
         """Sorts DataFrame globally.
 
@@ -1768,9 +1769,12 @@ def sort(
                 by,
             ]
 
+        if nulls_first is None:
+            nulls_first = desc
+
         sort_by = self.__column_input_to_expression(by)
 
-        builder = self._builder.sort(sort_by=sort_by, descending=desc, nulls_first=desc)
+        builder = self._builder.sort(sort_by=sort_by, descending=desc, nulls_first=nulls_first)
         return DataFrame(builder)
 
     @DataframePublicAPI

daft/execution/execution_step.py

@@ -935,12 +935,15 @@ class ReduceMergeAndSort(ReduceInstruction):
     sort_by: ExpressionsProjection
     descending: list[bool]
     bounds: MicroPartition
+    nulls_first: list[bool] | None = None
 
     def run(self, inputs: list[MicroPartition]) -> list[MicroPartition]:
         return self._reduce_merge_and_sort(inputs)
 
     def _reduce_merge_and_sort(self, inputs: list[MicroPartition]) -> list[MicroPartition]:
-        partition = MicroPartition.concat(inputs).sort(self.sort_by, descending=self.descending)
+        partition = MicroPartition.concat(inputs).sort(
+            self.sort_by, descending=self.descending, nulls_first=self.nulls_first
+        )
         return [partition]
 
     def run_partial_metadata(self, input_metadatas: list[PartialPartitionMetadata]) -> list[PartialPartitionMetadata]:
@@ -969,7 +972,6 @@ def _reduce_to_quantiles(self, inputs: list[MicroPartition]) -> list[MicroPartit
         merged = MicroPartition.concat(inputs)
 
         nulls_first = self.nulls_first if self.nulls_first is not None else self.descending
-
         # Skip evaluation of expressions by converting to Column Expression, since evaluation was done in Sample
         merged_sorted = merged.sort(
             self.sort_by.to_column_expressions(), descending=self.descending, nulls_first=nulls_first

daft/execution/physical_plan.py

@@ -1646,6 +1646,7 @@ def sort(
             execution_step.ReduceMergeAndSort(
                 sort_by=sort_by,
                 descending=descending,
+                nulls_first=nulls_first,
                 bounds=per_part_boundaries,
             )
             for per_part_boundaries in per_partition_bounds

src/daft-core/src/array/ops/arrow2/sort/primitive/common.rs

@@ -9,13 +9,14 @@ pub fn idx_sort<I, F>(
     cmp: F,
     length: usize,
     descending: bool,
+    nulls_first: bool,
 ) -> PrimitiveArray<I>
 where
     I: Index,
     F: Fn(&I, &I) -> std::cmp::Ordering,
 {
     let (mut indices, start_idx, end_idx) =
-        generate_initial_indices::<I>(validity, length, descending);
+        generate_initial_indices::<I>(validity, length, descending, nulls_first);
     let indices_slice = &mut indices.as_mut_slice()[start_idx..end_idx];
 
     if !descending {
@@ -33,13 +34,18 @@ pub fn multi_column_idx_sort<I, F>(
     others_cmp: &DynComparator,
     length: usize,
     first_col_desc: bool,
+    first_col_nulls_first: bool,
 ) -> PrimitiveArray<I>
 where
     I: Index,
     F: Fn(&I, &I) -> std::cmp::Ordering,
 {
-    let (mut indices, start_idx, end_idx) =
-        generate_initial_indices::<I>(first_col_validity, length, first_col_desc);
+    let (mut indices, start_idx, end_idx) = generate_initial_indices::<I>(
+        first_col_validity,
+        length,
+        first_col_desc,
+        first_col_nulls_first,
+    );
     let indices_slice = &mut indices.as_mut_slice()[start_idx..end_idx];
 
     indices_slice.sort_unstable_by(|a, b| overall_cmp(a, b));
@@ -60,6 +66,7 @@ fn generate_initial_indices<I>(
     validity: Option<&Bitmap>,
     length: usize,
     descending: bool,
+    nulls_first: bool,
 ) -> (Vec<I>, usize, usize)
 where
     I: Index,
@@ -72,16 +79,33 @@ where
         if descending {
             let mut nulls = 0;
             let mut valids = 0;
+            let last_valid_index = length.saturating_sub(validity.unset_bits());
+
             validity
                 .iter()
                 .zip(I::range(0, length).unwrap())
-                .for_each(|(is_valid, index)| {
-                    if is_valid {
-                        indices[validity.unset_bits() + valids] = index;
-                        valids += 1;
-                    } else {
-                        indices[nulls] = index;
-                        nulls += 1;
+                .for_each(|(is_not_null, index)| {
+                    match (is_not_null, nulls_first) {
+                        // value && nulls first
+                        (true, true) => {
+                            indices[validity.unset_bits() + valids] = index;
+                            valids += 1;
+                        }
+                        // value && nulls last
+                        (true, false) => {
+                            indices[valids] = index;
+                            valids += 1;
+                        }
+                        // null && nulls first
+                        (false, true) => {
+                            indices[nulls] = index;
+                            nulls += 1;
+                        }
+                        // null && nulls last
+                        (false, false) => {
+                            indices[last_valid_index + nulls] = index;
+                            nulls += 1;
+                        }
                     }
                 });
             start_idx = validity.unset_bits();
@@ -92,16 +116,35 @@ where
             validity
                 .iter()
                 .zip(I::range(0, length).unwrap())
-                .for_each(|(x, index)| {
-                    if x {
-                        indices[valids] = index;
-                        valids += 1;
-                    } else {
-                        indices[last_valid_index + nulls] = index;
-                        nulls += 1;
+                .for_each(|(is_not_null, index)| {
+                    match (is_not_null, nulls_first) {
+                        // value && nulls_first
+                        (true, true) => {
+                            indices[validity.unset_bits() + valids] = index;
+                            valids += 1;
+                        }
+                        // value && nulls last
+                        (true, false) => {
+                            indices[valids] = index;
+                            valids += 1;
+                        }
+                        // null && nulls first
+                        (false, true) => {
+                            indices[nulls] = index;
+                            nulls += 1;
+                        }
+                        // null && nulls last
+                        (false, false) => {
+                            indices[last_valid_index + nulls] = index;
+                            nulls += 1;
+                        }
                     }
                 });
-            end_idx = last_valid_index;
+            if nulls_first {
+                start_idx = validity.unset_bits();
+            } else {
+                end_idx = last_valid_index;
+            }
         }
         (indices, start_idx, end_idx)
     } else {

src/daft-core/src/array/ops/arrow2/sort/primitive/indices.rs

@@ -10,6 +10,7 @@ pub fn indices_sorted_unstable_by<I, T, F>(
     array: &PrimitiveArray<T>,
     cmp: F,
     descending: bool,
+    nulls_first: bool,
 ) -> PrimitiveArray<I>
 where
     I: Index,
@@ -29,6 +30,7 @@ where
             },
             array.len(),
             descending,
+            nulls_first,
         )
     }
 }