Repository: pola-rs/polars
File: crates/polars-core/src/frame/mod.rs (Rust)

Code before:
    /// ```
    /// # use polars_core::prelude::*;
    /// fn example(df: &DataFrame) -> PolarsResult<DataFrame> {
    ///     let idx = IdxCa::new("idx".into(), [0, 1, 9]);
    ///     df.take(&idx)
    /// }
    /// ```
    pub fn take(&self, indices: &IdxCa) -> PolarsResult<Self> {
        let new_col = POOL.install(|| self.try_apply_columns_par(&|s| s.take(indices)))?;

        Ok(unsafe { DataFrame::new_no_checks(indices.len(), new_col) })
    }

    /// # Safety
    /// The indices must be in-bounds.
    pub unsafe fn take_unchecked(&self, idx: &IdxCa) -> Self {
        self.take_unchecked_impl(idx, true)
    }

    /// # Safety
    /// The indices must be in-bounds.
    pub unsafe fn take_unchecked_impl(&self, idx: &IdxCa, allow_threads: bool) -> Self {
        let cols = if allow_threads {
            POOL.install(|| self._apply_columns_par(&|c| c.take_unchecked(idx)))
        } else {
            self._apply_columns(&|s| s.take_unchecked(idx))
        };
        unsafe { DataFrame::new_no_checks(idx.len(), cols) }
    }

    /// # Safety
    /// The indices must be in-bounds.
    pub unsafe fn take_slice_unchecked(&self, idx: &[IdxSize]) -> Self {
        self.take_slice_unchecked_impl(idx, true)
    }

    /// # Safety
    /// The indices must be in-bounds.
    pub unsafe fn take_slice_unchecked_impl(&self, idx: &[IdxSize], allow_threads: bool) -> Self {
        let cols = if allow_threads {
            POOL.install(|| self._apply_columns_par(&|s| s.take_slice_unchecked(idx)))
        } else {
            self._apply_columns(&|s| s.take_slice_unchecked(idx))
        };
        unsafe { DataFrame::new_no_checks(idx.len(), cols) }
    }

    /// Rename a column in the [`DataFrame`].
    ///
    /// Should not be called in a loop as that can lead to quadratic behavior.
    ///

Review comment (suggestion):
Could we also avoid this branch if `self.len() / stride < 2`


Code after:
    /// ```
    /// # use polars_core::prelude::*;
    /// fn example(df: &DataFrame) -> PolarsResult<DataFrame> {
    ///     let idx = IdxCa::new("idx".into(), [0, 1, 9]);
    ///     df.take(&idx)
    /// }
    /// ```
    pub fn take(&self, indices: &IdxCa) -> PolarsResult<Self> {
        let new_col = POOL.install(|| self.try_apply_columns_par(&|s| s.take(indices)))?;

        Ok(unsafe { DataFrame::new_no_checks(indices.len(), new_col) })
    }

    /// # Safety
    /// The indices must be in-bounds.
    pub unsafe fn take_unchecked(&self, idx: &IdxCa) -> Self {
        self.take_unchecked_impl(idx, true)
    }

    /// # Safety
    /// The indices must be in-bounds.
    pub unsafe fn take_unchecked_impl(&self, idx: &IdxCa, allow_threads: bool) -> Self {
        let cols = if allow_threads && POOL.current_num_threads() > 1 {
            POOL.install(|| {
                if POOL.current_num_threads() > self.width() {
                    let stride = usize::max(idx.len().div_ceil(POOL.current_num_threads()), 256);
                    self._apply_columns_par(&|c| {
                        // Nested types initiate a rechunk in their take_unchecked implementation.
                        // If we do not rechunk, it will result in rechunk storms downstream.
                        let c = if c.dtype().is_nested() {
                            &c.rechunk()
                        } else {
                            c
                        };

                        (0..idx.len().div_ceil(stride))
                            .into_par_iter()
                            .map(|i| c.take_unchecked(&idx.slice((i * stride) as i64, stride)))
                            .reduce(
                                || Column::new_empty(c.name().clone(), c.dtype()),
                                |mut a, b| {
                                    a.append_owned(b).unwrap();
                                    a
                                },
                            )
                    })
                } else {
                    self._apply_columns_par(&|c| c.take_unchecked(idx))
                }
            })
        } else {