refactor(rust): Remove descending from sort options.limit

pola-rs · Jan 15, 2025 · 83885f9 · 83885f9
1 parent 1cd72ff
commit 83885f9
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 35 deletions.
diff --git a/crates/polars-core/src/chunked_array/ops/sort/arg_sort.rs b/crates/polars-core/src/chunked_array/ops/sort/arg_sort.rs
@@ -105,8 +105,7 @@ where
     {
         len = options
             .limit
-            .map(|(limit, _)| std::cmp::min(limit as usize, len))
-            .unwrap_or(len);
+            .map_or(len, |limit| std::cmp::min(limit.try_into().unwrap(), len));
         return ChunkedArray::with_chunk(
             name,
             IdxArr::from_data_default(
@@ -136,17 +135,12 @@ where
         vals.extend(iter);
     }
 
-    let vals = if let Some((limit, desc)) = options.limit {
+    let vals = if let Some(limit) = options.limit {
         let limit = limit as usize;
         // Overwrite output len.
         len = limit;
         let out = if limit >= vals.len() {
             vals.as_mut_slice()
-        } else if desc {
-            let (lower, _el, _upper) = vals
-                .as_mut_slice()
-                .select_nth_unstable_by(limit, |a, b| b.1.tot_cmp(&a.1));
-            lower
         } else {
             let (lower, _el, _upper) = vals
                 .as_mut_slice()
@@ -205,8 +199,7 @@ where
     if is_sorted_flag != IsSorted::Not {
         let len_final = options
             .limit
-            .map(|(limit, _)| std::cmp::min(limit as usize, len))
-            .unwrap_or(len);
+            .map_or(len, |limit| std::cmp::min(limit.try_into().unwrap(), len));
         if (options.descending && is_sorted_flag == IsSorted::Descending)
             || (!options.descending && is_sorted_flag == IsSorted::Ascending)
         {
@@ -237,15 +230,10 @@ where
         }));
     }
 
-    let vals = if let Some((limit, desc)) = options.limit {
+    let vals = if let Some(limit) = options.limit {
         let limit = limit as usize;
         let out = if limit >= vals.len() {
             vals.as_mut_slice()
-        } else if desc {
-            let (lower, _el, _upper) = vals
-                .as_mut_slice()
-                .select_nth_unstable_by(limit, |a, b| b.1.tot_cmp(&a.1));
-            lower
         } else {
             let (lower, _el, _upper) = vals
                 .as_mut_slice()

diff --git a/crates/polars-core/src/chunked_array/ops/sort/options.rs b/crates/polars-core/src/chunked_array/ops/sort/options.rs
@@ -42,9 +42,7 @@ pub struct SortOptions {
     /// Default `false`.
     pub maintain_order: bool,
     /// Limit a sort output, this is for optimization purposes and might be ignored.
-    /// - Len
-    /// - Descending
-    pub limit: Option<(IdxSize, bool)>,
+    pub limit: Option<IdxSize>,
 }
 
 /// Sort options for multi-series sorting.
@@ -101,9 +99,7 @@ pub struct SortMultipleOptions {
     /// Whether maintain the order of equal elements. Default `false`.
     pub maintain_order: bool,
     /// Limit a sort output, this is for optimization purposes and might be ignored.
-    /// - Len
-    /// - Descending
-    pub limit: Option<(IdxSize, bool)>,
+    pub limit: Option<IdxSize>,
 }
 
 impl Default for SortOptions {

diff --git a/crates/polars-core/src/frame/column/mod.rs b/crates/polars-core/src/frame/column/mod.rs
@@ -1041,14 +1041,9 @@ impl Column {
 
         // @NOTE: This can theoretically be pushed into the previous operation but it is really
         // worth it... probably not...
-        if let Some((limit, limit_dsc)) = options.limit {
+        if let Some(limit) = options.limit {
             let limit = limit.min(length);
-
-            if limit_dsc {
-                values = values.drain((length - limit) as usize..).collect();
-            } else {
-                values.truncate(limit as usize);
-            }
+            values.truncate(limit as usize);
         }
 
         IdxCa::from_vec(self.name().clone(), values)

diff --git a/crates/polars-core/src/frame/mod.rs b/crates/polars-core/src/frame/mod.rs
@@ -2080,12 +2080,7 @@ impl DataFrame {
         }
         if let Some((0, k)) = slice {
             if k < self.len() {
-                let desc = if sort_options.descending.len() == 1 {
-                    sort_options.descending[0]
-                } else {
-                    false
-                };
-                sort_options.limit = Some((k as IdxSize, desc));
+                sort_options.limit = Some(k.try_into().unwrap());
                 return self.bottom_k_impl(k, by_column, sort_options);
             }
         }