polars_core/frame/
top_k.rs

1use super::*;
2use crate::prelude::sort::arg_bottom_k::_arg_bottom_k;
3
4impl DataFrame {
5    pub(crate) fn bottom_k_impl(
6        &self,
7        k: usize,
8        by_column: Vec<Column>,
9        mut sort_options: SortMultipleOptions,
10    ) -> PolarsResult<DataFrame> {
11        let first_descending = sort_options.descending[0];
12        let first_by_column = by_column[0].name().to_string();
13
14        let idx = _arg_bottom_k(k, &by_column, &mut sort_options)?;
15
16        let mut df = unsafe { self.take_unchecked(&idx.into_inner()) };
17
18        // Mark the first sort column as sorted
19        // if the column did not exists it is ok, because we sorted by an expression
20        // not present in the dataframe
21        let _ = df.apply(&first_by_column, |s| {
22            let mut s = s.clone();
23            if first_descending {
24                s.set_sorted_flag(IsSorted::Descending)
25            } else {
26                s.set_sorted_flag(IsSorted::Ascending)
27            }
28            s
29        });
30        Ok(df)
31    }
32}