accept regex in filter (#3666)

pola-rs · Jun 11, 2022 · c627af4 · c627af4
1 parent c48823a
commit c627af4
Show file tree

Hide file tree

Showing 9 changed files with 47 additions and 79 deletions.
diff --git a/polars/polars-core/src/chunked_array/bitwise.rs b/polars/polars-core/src/chunked_array/bitwise.rs
@@ -253,35 +253,6 @@ impl BitAnd for BooleanChunked {
     }
 }
 
-macro_rules! impl_floats {
-    ($_type:ty) => {
-        impl BitXor for &$_type {
-            type Output = $_type;
-
-            fn bitxor(self, _rhs: Self) -> Self::Output {
-                unimplemented!()
-            }
-        }
-        impl BitAnd for &$_type {
-            type Output = $_type;
-
-            fn bitand(self, _rhs: Self) -> Self::Output {
-                unimplemented!()
-            }
-        }
-        impl BitOr for &$_type {
-            type Output = $_type;
-
-            fn bitor(self, _rhs: Self) -> Self::Output {
-                unimplemented!()
-            }
-        }
-    };
-}
-
-impl_floats!(Float64Chunked);
-impl_floats!(Float32Chunked);
-
 #[cfg(test)]
 mod test {
     use super::*;

diff --git a/polars/polars-core/src/series/implementations/floats.rs b/polars/polars-core/src/series/implementations/floats.rs
@@ -22,7 +22,6 @@ use ahash::RandomState;
 use arrow::array::ArrayRef;
 use polars_arrow::prelude::QuantileInterpolOptions;
 use std::borrow::Cow;
-use std::ops::{BitAnd, BitOr, BitXor};
 
 macro_rules! impl_dyn_series {
     ($ca: ident) => {
@@ -180,36 +179,6 @@ macro_rules! impl_dyn_series {
                 self.0.interpolate().into_series()
             }
 
-            fn bitand(&self, other: &Series) -> Result<Series> {
-                let other = if other.len() == 1 {
-                    Cow::Owned(other.cast(self.dtype())?)
-                } else {
-                    Cow::Borrowed(other)
-                };
-                let other = self.0.unpack_series_matching_type(&other)?;
-                Ok(self.0.bitand(&other).into_series())
-            }
-
-            fn bitor(&self, other: &Series) -> Result<Series> {
-                let other = if other.len() == 1 {
-                    Cow::Owned(other.cast(self.dtype())?)
-                } else {
-                    Cow::Borrowed(other)
-                };
-                let other = self.0.unpack_series_matching_type(&other)?;
-                Ok(self.0.bitor(&other).into_series())
-            }
-
-            fn bitxor(&self, other: &Series) -> Result<Series> {
-                let other = if other.len() == 1 {
-                    Cow::Owned(other.cast(self.dtype())?)
-                } else {
-                    Cow::Borrowed(other)
-                };
-                let other = self.0.unpack_series_matching_type(&other)?;
-                Ok(self.0.bitxor(&other).into_series())
-            }
-
             fn rename(&mut self, name: &str) {
                 self.0.rename(name);
             }

diff --git a/polars/polars-core/src/series/series_trait.rs b/polars/polars-core/src/series/series_trait.rs
@@ -203,24 +203,33 @@ pub trait SeriesTrait:
     fn rename(&mut self, name: &str);
 
     fn bitand(&self, _other: &Series) -> Result<Series> {
-        panic!(
-            "bitwise and operation not supported for dtype {:?}",
-            self.dtype()
-        )
+        Err(PolarsError::InvalidOperation(
+            format!(
+                "bitwise 'AND' operation not supported for dtype {:?}",
+                self.dtype()
+            )
+            .into(),
+        ))
     }
 
     fn bitor(&self, _other: &Series) -> Result<Series> {
-        panic!(
-            "bitwise or operation not fit supported for dtype {:?}",
-            self.dtype()
-        )
+        Err(PolarsError::InvalidOperation(
+            format!(
+                "bitwise 'OR' operation not supported for dtype {:?}",
+                self.dtype()
+            )
+            .into(),
+        ))
     }
 
     fn bitxor(&self, _other: &Series) -> Result<Series> {
-        panic!(
-            "bitwise xor operation not fit supported for dtype {:?}",
-            self.dtype()
-        )
+        Err(PolarsError::InvalidOperation(
+            format!(
+                "bitwise 'XOR' operation not supported for dtype {:?}",
+                self.dtype()
+            )
+            .into(),
+        ))
     }
 
     /// Get the lengths of the underlying chunks

diff --git a/polars/polars-lazy/src/logical_plan/builder.rs b/polars/polars-lazy/src/logical_plan/builder.rs
@@ -302,11 +302,10 @@ impl LogicalPlanBuilder {
 
     /// Apply a filter
     pub fn filter(self, predicate: Expr) -> Self {
-        let predicate = if has_expr(&predicate, |e| {
-            matches!(
-                e,
-                Expr::Wildcard | Expr::RenameAlias { .. } | Expr::Columns(_)
-            )
+        let predicate = if has_expr(&predicate, |e| match e {
+            Expr::Column(name) => name.starts_with('^') && name.ends_with('$'),
+            Expr::Wildcard | Expr::RenameAlias { .. } | Expr::Columns(_) => true,
+            _ => false,
         }) {
             let rewritten = rewrite_projections(vec![predicate], self.0.schema(), &[]);
             combine_predicates_expr(rewritten.into_iter())

diff --git a/py-polars/polars/internals/__init__.py b/py-polars/polars/internals/__init__.py
@@ -9,6 +9,7 @@
     _scan_ipc_fsspec,
     _scan_parquet_fsspec,
 )
+from .datatypes import IntoExpr
 from .expr import Expr, expr_to_lit_or_expr, selection_to_pyexpr_list, wrap_expr
 from .frame import DataFrame, LazyFrame, wrap_df, wrap_ldf
 from .functions import concat, date_range  # DataFrame.describe() & DataFrame.upsample()

diff --git a/py-polars/polars/internals/datatypes.py b/py-polars/polars/internals/datatypes.py
@@ -0,0 +1,5 @@
+from typing import Union
+
+from polars import internals as pli
+
+IntoExpr = Union[int, float, str, "pli.Expr", "pli.Series"]
diff --git a/py-polars/polars/internals/lazy_functions.py b/py-polars/polars/internals/lazy_functions.py
@@ -855,7 +855,7 @@ def map_binary(
 
 
 def fold(
-    acc: "pli.Expr",
+    acc: "pli.IntoExpr",
     f: Callable[["pli.Series", "pli.Series"], "pli.Series"],
     exprs: Union[Sequence[Union["pli.Expr", str]], "pli.Expr"],
 ) -> "pli.Expr":

diff --git a/py-polars/tests/test_df.py b/py-polars/tests/test_df.py
@@ -1587,7 +1587,7 @@ def test_filter_with_all_expansion() -> None:
             "a": [None, None, None],
         }
     )
-    out = df.filter(~pl.fold(True, lambda acc, s: acc & s.is_null(), pl.all()))  # type: ignore
+    out = df.filter(~pl.fold(True, lambda acc, s: acc & s.is_null(), pl.all()))
     assert out.shape == (2, 3)
 
 

diff --git a/py-polars/tests/test_exprs.py b/py-polars/tests/test_exprs.py
@@ -248,3 +248,17 @@ def test_expression_appends() -> None:
 
     assert out.n_chunks() == 1
     assert out.to_series().to_list() == [None, None, None, 1, 1, 2]
+
+
+def test_regex_in_filter() -> None:
+    df = pl.DataFrame(
+        {
+            "nrs": [1, 2, 3, None, 5],
+            "names": ["foo", "ham", "spam", "egg", None],
+            "flt": [1.0, None, 3.0, 1.0, None],
+        }
+    )
+
+    assert df.filter(
+        pl.fold(acc=False, f=lambda acc, s: acc | s, exprs=(pl.col("^nrs|flt*$") < 3))
+    ).row(0) == (1, "foo", 1.0)