polars_ops/frame/
gather.rs

1use polars_core::prelude::*;
2
3use crate::series::convert_and_bound_index;
4
5pub trait GatherDf {
6    /// Selects from this DataFrame the rows at the indices in the given column.
7    fn gather_with_column(&self, idxs: &Column, null_on_oob: bool) -> PolarsResult<DataFrame>;
8
9    /// Selects from this DataFrame the rows at the indices in the given series.
10    fn gather_with_series(&self, idxs: &Series, null_on_oob: bool) -> PolarsResult<DataFrame>;
11}
12
13impl GatherDf for DataFrame {
14    fn gather_with_column(&self, idxs: &Column, null_on_oob: bool) -> PolarsResult<DataFrame> {
15        match idxs {
16            Column::Series(s) => self.gather_with_series(s, null_on_oob),
17            Column::Scalar(idx_c) => {
18                if idx_c.is_empty() {
19                    return Ok(self.clear());
20                }
21
22                let idx_s = idx_c.as_single_value_series();
23                let idx_ca = convert_and_bound_index(&idx_s, self.height(), null_on_oob)?;
24                match idx_ca.get(0) {
25                    Some(idx) => Ok(self.new_from_index(idx as usize, idx_c.len())),
26                    None => Ok(DataFrame::full_null(self.schema(), idx_c.len())),
27                }
28            },
29        }
30    }
31
32    fn gather_with_series(&self, idxs: &Series, null_on_oob: bool) -> PolarsResult<DataFrame> {
33        if idxs.is_empty() {
34            return Ok(self.clear());
35        }
36
37        let idx_ca = convert_and_bound_index(idxs, self.height(), null_on_oob)?;
38        Ok(unsafe { self.take_unchecked_impl(&idx_ca, false) })
39    }
40}
polars_ops/frame/gather.rs

polars_ops/frame/
gather.rs
