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

Code before:
    /// assert_eq!(df["Carbon"], sv[0]);
    /// assert_eq!(df["Hydrogen"], sv[1]);
    /// # Ok::<(), PolarsError>(())
    /// ```
    pub fn select_to_vec(
        &self,
        selection: impl IntoIterator<Item = impl AsRef<str>>,
    ) -> PolarsResult<Vec<Column>> {
        AmortizedColumnSelector::new(self).select_multiple(selection)
    }

    /// Take the [`DataFrame`] rows by a boolean mask.
    ///
    /// # Example
    ///
    /// ```
    /// # use polars_core::prelude::*;
    /// fn example(df: &DataFrame) -> PolarsResult<DataFrame> {
    ///     let mask = df.column("sepal_width")?.is_not_null();
    ///     df.filter(&mask)
    /// }
    /// ```
    pub fn filter(&self, mask: &BooleanChunked) -> PolarsResult<Self> {
        if self.width() == 0 {
            filter_zero_width(self.height(), mask)
        } else if mask.len() == 1 && mask.null_count() == 0 && self.len() >= 1 {
            if mask.all() {
                Ok(self.clone())
            } else {
                Ok(self.clear())
            }
        } else {
            let new_columns: Vec<Column> = self.try_apply_columns_par(|s| s.filter(mask))?;
            let out = unsafe {
                DataFrame::new_unchecked(new_columns[0].len(), new_columns).with_schema_from(self)
            };

            Ok(out)
        }
    }

    /// Same as `filter` but does not parallelize.
    pub fn filter_seq(&self, mask: &BooleanChunked) -> PolarsResult<Self> {
        if self.width() == 0 {
            filter_zero_width(self.height(), mask)
        } else if mask.len() == 1 && mask.null_count() == 0 && self.len() >= 1 {
            if mask.all() {
                Ok(self.clone())
            } else {
                Ok(self.clear())
            }

Review comment (suggestion):
I think we can also do nulls - in filters they behave as false

```rust
if mask.len() == 1 {
if mask.first() == Some(true) {
  self.clone()
} else {
  self.clear()
}
}
```


Code after:
    /// assert_eq!(df["Carbon"], sv[0]);
    /// assert_eq!(df["Hydrogen"], sv[1]);
    /// # Ok::<(), PolarsError>(())
    /// ```
    pub fn select_to_vec(
        &self,
        selection: impl IntoIterator<Item = impl AsRef<str>>,
    ) -> PolarsResult<Vec<Column>> {
        AmortizedColumnSelector::new(self).select_multiple(selection)
    }

    /// Take the [`DataFrame`] rows by a boolean mask.
    ///
    /// # Example
    ///
    /// ```
    /// # use polars_core::prelude::*;
    /// fn example(df: &DataFrame) -> PolarsResult<DataFrame> {
    ///     let mask = df.column("sepal_width")?.is_not_null();
    ///     df.filter(&mask)
    /// }
    /// ```
    pub fn filter(&self, mask: &BooleanChunked) -> PolarsResult<Self> {
        if self.width() == 0 {
            filter_zero_width(self.height(), mask)
        } else if mask.len() == 1 && self.len() >= 1 {
            if mask.all() && mask.null_count() == 0 {
                Ok(self.clone())
            } else {
                Ok(self.clear())
            }
        } else {
            let new_columns: Vec<Column> = self.try_apply_columns_par(|s| s.filter(mask))?;
            let out = unsafe {
                DataFrame::new_unchecked(new_columns[0].len(), new_columns).with_schema_from(self)
            };

            Ok(out)
        }
    }

    /// Same as `filter` but does not parallelize.
    pub fn filter_seq(&self, mask: &BooleanChunked) -> PolarsResult<Self> {
        if self.width() == 0 {
            filter_zero_width(self.height(), mask)
        } else if mask.len() == 1 && mask.null_count() == 0 && self.len() >= 1 {
            if mask.all() && mask.null_count() == 0 {
                Ok(self.clone())
            } else {
                Ok(self.clear())
            }