python arr.contains

pola-rs · Dec 9, 2021 · 6f72cdf · 6f72cdf
1 parent ac767f1
commit 6f72cdf
Show file tree

Hide file tree

Showing 6 changed files with 141 additions and 33 deletions.
diff --git a/polars/polars-core/src/chunked_array/ops/is_in.rs b/polars/polars-core/src/chunked_array/ops/is_in.rs
@@ -52,17 +52,32 @@ where
                     return left.is_in(&right);
                 }
 
-                let ca: BooleanChunked = self
-                    .into_iter()
-                    .zip(other.list()?.into_iter())
-                    .map(|(value, series)| match (value, series) {
-                        (val, Some(series)) => {
-                            let ca = series.unpack::<T>().unwrap();
-                            ca.into_iter().any(|a| a == val)
-                        }
-                        _ => false,
-                    })
-                    .collect_trusted();
+                let mut ca: BooleanChunked = if self.len() == 1 && other.len() != 1 {
+                    let value = self.get(0);
+                    other
+                        .list()?
+                        .amortized_iter()
+                        .map(|opt_s| {
+                            opt_s.map(|s| {
+                                let ca = s.as_ref().unpack::<T>().unwrap();
+                                ca.into_iter().any(|a| a == value)
+                            }) == Some(true)
+                        })
+                        .trust_my_length(other.len())
+                        .collect_trusted()
+                } else {
+                    self.into_iter()
+                        .zip(other.list()?.amortized_iter())
+                        .map(|(value, series)| match (value, series) {
+                            (val, Some(series)) => {
+                                let ca = series.as_ref().unpack::<T>().unwrap();
+                                ca.into_iter().any(|a| a == val)
+                            }
+                            _ => false,
+                        })
+                        .collect_trusted()
+                };
+                ca.rename(self.name());
                 Ok(ca)
             }
             _ => {
@@ -109,17 +124,32 @@ impl IsIn for Utf8Chunked {
     fn is_in(&self, other: &Series) -> Result<BooleanChunked> {
         match other.dtype() {
             DataType::List(dt) if self.dtype() == &**dt => {
-                let ca: BooleanChunked = self
-                    .into_iter()
-                    .zip(other.list()?.into_iter())
-                    .map(|(value, series)| match (value, series) {
-                        (val, Some(series)) => {
-                            let ca = series.unpack::<Utf8Type>().unwrap();
-                            ca.into_iter().any(|a| a == val)
-                        }
-                        _ => false,
-                    })
-                    .collect_trusted();
+                let mut ca: BooleanChunked = if self.len() == 1 && other.len() != 1 {
+                    let value = self.get(0);
+                    other
+                        .list()?
+                        .amortized_iter()
+                        .map(|opt_s| {
+                            opt_s.map(|s| {
+                                let ca = s.as_ref().unpack::<Utf8Type>().unwrap();
+                                ca.into_iter().any(|a| a == value)
+                            }) == Some(true)
+                        })
+                        .trust_my_length(other.len())
+                        .collect_trusted()
+                } else {
+                    self.into_iter()
+                        .zip(other.list()?.amortized_iter())
+                        .map(|(value, series)| match (value, series) {
+                            (val, Some(series)) => {
+                                let ca = series.as_ref().unpack::<Utf8Type>().unwrap();
+                                ca.into_iter().any(|a| a == val)
+                            }
+                            _ => false,
+                        })
+                        .collect_trusted()
+                };
+                ca.rename(self.name());
                 Ok(ca)
             }
             DataType::Utf8 => {
@@ -158,17 +188,32 @@ impl IsIn for BooleanChunked {
     fn is_in(&self, other: &Series) -> Result<BooleanChunked> {
         match other.dtype() {
             DataType::List(dt) if self.dtype() == &**dt => {
-                let ca: BooleanChunked = self
-                    .into_iter()
-                    .zip(other.list()?.into_iter())
-                    .map(|(value, series)| match (value, series) {
-                        (val, Some(series)) => {
-                            let ca = series.unpack::<BooleanType>().unwrap();
-                            ca.into_iter().any(|a| a == val)
-                        }
-                        _ => false,
-                    })
-                    .collect_trusted();
+                let mut ca: BooleanChunked = if self.len() == 1 && other.len() != 1 {
+                    let value = self.get(0);
+                    other
+                        .list()?
+                        .amortized_iter()
+                        .map(|opt_s| {
+                            opt_s.map(|s| {
+                                let ca = s.as_ref().unpack::<BooleanType>().unwrap();
+                                ca.into_iter().any(|a| a == value)
+                            }) == Some(true)
+                        })
+                        .trust_my_length(other.len())
+                        .collect_trusted()
+                } else {
+                    self.into_iter()
+                        .zip(other.list()?.amortized_iter())
+                        .map(|(value, series)| match (value, series) {
+                            (val, Some(series)) => {
+                                let ca = series.as_ref().unpack::<BooleanType>().unwrap();
+                                ca.into_iter().any(|a| a == val)
+                            }
+                            _ => false,
+                        })
+                        .collect_trusted()
+                };
+                ca.rename(self.name());
                 Ok(ca)
             }
             _ => Err(PolarsError::SchemaMisMatch(

diff --git a/py-polars/docs/source/reference/expression.rst b/py-polars/docs/source/reference/expression.rst
@@ -291,3 +291,4 @@ The following methods are available under the `expr.arr` attribute.
     ExprListNameSpace.get
     ExprListNameSpace.first
     ExprListNameSpace.last
+    ExprListNameSpace.contains
diff --git a/py-polars/docs/source/reference/series.rst b/py-polars/docs/source/reference/series.rst
@@ -251,3 +251,4 @@ The following methods are available under the `Series.arr` attribute.
     ListNameSpace.get
     ListNameSpace.first
     ListNameSpace.last
+    ListNameSpace.contains
diff --git a/py-polars/polars/internals/expr.py b/py-polars/polars/internals/expr.py
@@ -2161,6 +2161,21 @@ def last(self) -> "Expr":
         """
         return self.get(-1)
 
+    def contains(self, item: Union[float, str, bool, int, date, datetime]) -> "Expr":
+        """
+        Check if sublists contain the given item.
+
+        Parameters
+        ----------
+        item
+            Item that will be checked for membership
+
+        Returns
+        -------
+        Boolean mask
+        """
+        return wrap_expr(self._pyexpr).map(lambda s: s.arr.contains(item))
+
 
 class ExprStringNameSpace:
     """

diff --git a/py-polars/polars/internals/series.py b/py-polars/polars/internals/series.py
@@ -1,3 +1,4 @@
+import sys
 import typing as tp
 from datetime import date, datetime, timedelta
 from numbers import Number
@@ -61,6 +62,11 @@
 except ImportError:  # pragma: no cover
     _PANDAS_AVAILABLE = False
 
+if sys.version_info >= (3, 8):
+    from typing import Literal
+else:
+    from typing_extensions import Literal
+
 
 def match_dtype(value: Any, dtype: "Type[DataType]") -> Any:
     """
@@ -923,6 +929,18 @@ def alias(self, name: str) -> "Series":
         s._s.rename(name)
         return s
 
+    @tp.overload
+    def rename(self, name: str, in_place: Literal[False] = ...) -> "Series":
+        ...
+
+    @tp.overload
+    def rename(self, name: str, in_place: Literal[True]) -> None:
+        ...
+
+    @tp.overload
+    def rename(self, name: str, in_place: bool) -> Optional["Series"]:
+        ...
+
     def rename(self, name: str, in_place: bool = False) -> Optional["Series"]:
         """
         Rename this Series.
@@ -3414,6 +3432,24 @@ def last(self) -> "Series":
         """
         return self.get(-1)
 
+    def contains(self, item: Union[float, str, bool, int, date, datetime]) -> "Series":
+        """
+        Check if sublists contain the given item.
+
+        Parameters
+        ----------
+        item
+            Item that will be checked for membership
+
+        Returns
+        -------
+        Boolean mask
+        """
+        s = pli.Series("", [item])
+        s_list = wrap_s(self._s)
+        out = s.is_in(s_list)
+        return out.rename(s_list.name)
+
 
 class DateTimeNameSpace:
     """

diff --git a/py-polars/tests/test_lists.py b/py-polars/tests/test_lists.py
@@ -24,3 +24,13 @@ def test_list_arr_get() -> None:
     out = a.arr.get(-3)
     expected = pl.Series("a", [1, None, 7])
     testing.assert_series_equal(out, expected)
+
+
+def test_contains() -> None:
+    a = pl.Series("a", [[1, 2, 3], [2, 5], [6, 7, 8, 9]])
+    out = a.arr.contains(2)
+    expected = pl.Series("a", [True, True, False])
+    testing.assert_series_equal(out, expected)
+
+    out = pl.select(pl.lit(a).arr.contains(2)).to_series()
+    testing.assert_series_equal(out, expected)