feat(series): implement abs, all, any, count, isin, and astype

leanframe/core/dtypes.py

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+import ibis
 import ibis.expr.datatypes as ibis_dtypes
 import pandas as pd
 
@@ -34,3 +35,19 @@ def convert_ibis_to_pandas(
     """
     arrow_type = ibis_type.to_pyarrow()
     return pd.ArrowDtype(arrow_type)
+
+
+def convert_pandas_to_ibis(
+    pandas_type: pd.ArrowDtype,
+) -> ibis_dtypes.DataType:
+    """
+    Convert a pandas ArrowDtype to an ibis type.
+
+    Args:
+        pandas_type: The pandas type to convert.
+
+    Returns:
+        The corresponding ibis type.
+    """
+    arrow_type = pandas_type.pyarrow_dtype
+    return ibis.dtype(arrow_type)

leanframe/core/series.py

@@ -19,7 +19,7 @@
 import ibis.expr.types as ibis_types
 import numpy as np
 import pandas as pd
-from leanframe.core.dtypes import convert_ibis_to_pandas
+from leanframe.core.dtypes import convert_ibis_to_pandas, convert_pandas_to_ibis
 
 
 class Series:
@@ -97,6 +97,18 @@ def __rmul__(self, other) -> Series:
     def __round__(self, n) -> Series:
         return Series(self._data.round(n))
 
+    def abs(self) -> "Series":
+        """Return a Series with the absolute value of each element."""
+        return Series(self._data.abs())
+
+    def all(self) -> bool:
+        """Return whether all elements are True."""
+        return self._data.all().to_pyarrow().as_py()
+
+    def any(self) -> bool:
+        """Return whether any element is True."""
+        return self._data.any().to_pyarrow().as_py()
+
     def sum(self):
         """Return the sum of the Series."""
         return self._data.sum().to_pyarrow().as_py()
@@ -121,10 +133,23 @@ def var(self):
         """Return the var of the Series."""
         return self._data.var().to_pyarrow().as_py()
 
+    def count(self) -> int:
+        """Return the number of non-null observations in the Series."""
+        return self._data.count().to_pyarrow().as_py()
+
     def copy(self) -> Series:
         """Return a copy of the Series."""
         return Series(self._data)
 
+    def isin(self, values) -> "Series":
+        """Return a boolean Series showing whether each element in the Series is exactly contained in the passed sequence of values."""
+        return Series(self._data.isin(values))
+
+    def astype(self, dtype: pd.ArrowDtype) -> "Series":
+        """Cast a Series to a specified dtype."""
+        ibis_type = convert_pandas_to_ibis(dtype)
+        return Series(self._data.cast(ibis_type))
+
     def to_pandas(self) -> pd.Series:
         """Convert to a pandas Series."""
         return self._data.to_pyarrow().to_pandas(

tests/unit/test_series.py

@@ -49,6 +49,19 @@ def numeric_series(session):
     return session.DataFrame(df_pd)
 
 
+@pytest.fixture
+def bool_series(session):
+    df_pd = pd.DataFrame(
+        {
+            "all_true": [True, True, True],
+            "some_true": [True, False, True],
+            "all_false": [False, False, False],
+        },
+        dtype=pd.ArrowDtype(pa.bool_()),
+    )
+    return session.DataFrame(df_pd)
+
+
 def test_series_ndim(series_for_properties):
     series_int, series_float = series_for_properties
     assert series_int.ndim == 1
@@ -200,6 +213,58 @@ def test_series_arithmetic_scalar(session, op, other, expected_data):
     )
 
 
+def test_series_abs(session):
+    pandas_df = pd.DataFrame(
+        {"a": [-1, 2, -3]},
+        dtype=pd.ArrowDtype(pa.int64()),
+    )
+    df = session.DataFrame(pandas_df)
+    series = df["a"]
+    result = series.abs()
+    expected = pd.Series(
+        [1, 2, 3],
+        name="a",
+        dtype=pd.ArrowDtype(pa.int64()),
+    )
+    pd.testing.assert_series_equal(
+        result.to_pandas(),
+        expected,
+        check_names=False,
+    )
+
+
+def test_series_astype(session):
+    pandas_df = pd.DataFrame(
+        {"a": [1, 2, 3]},
+        dtype=pd.ArrowDtype(pa.int64()),
+    )
+    df = session.DataFrame(pandas_df)
+    series = df["a"]
+    result = series.astype(pd.ArrowDtype(pa.float64()))
+    expected = pd.Series(
+        [1.0, 2.0, 3.0],
+        name="a",
+        dtype=pd.ArrowDtype(pa.float64()),
+    )
+    pd.testing.assert_series_equal(
+        result.to_pandas(),
+        expected,
+        check_names=False,
+    )
+
+
+def test_series_all(bool_series):
+    assert bool_series["all_true"].all()
+    assert not bool_series["some_true"].all()
+    assert not bool_series["all_false"].all()
+
+
+def test_series_any(bool_series):
+    assert bool_series["all_true"].any()
+    assert bool_series["some_true"].any()
+    assert not bool_series["all_false"].any()
+
+
 def test_series_round(numeric_series):
     series = numeric_series["b"]
     result = round(series, 0)
@@ -246,6 +311,32 @@ def test_series_var(numeric_series):
     assert round(series.var(), 2) == 3.03
 
 
+def test_series_count(series_for_properties):
+    series_int, series_float = series_for_properties
+    assert series_int.count() == 3
+    assert series_float.count() == 2
+
+
+def test_series_isin(session):
+    pandas_df = pd.DataFrame(
+        {"a": ["a", "b", "c"]},
+        dtype=pd.ArrowDtype(pa.string()),
+    )
+    df = session.DataFrame(pandas_df)
+    series = df["a"]
+    result = series.isin(["a", "c"])
+    expected = pd.Series(
+        [True, False, True],
+        name="a",
+        dtype=pd.ArrowDtype(pa.bool_()),
+    )
+    pd.testing.assert_series_equal(
+        result.to_pandas(),
+        expected,
+        check_names=False,
+    )
+
+
 def test_series_copy(session):
     df_pd = pd.DataFrame({"col1": [1, 2, 3]})
     df_lf = session.DataFrame(df_pd)