diff --git a/leanframe/core/dtypes.py b/leanframe/core/dtypes.py index d27c57d..768e2d9 100644 --- a/leanframe/core/dtypes.py +++ b/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) diff --git a/leanframe/core/series.py b/leanframe/core/series.py index b898010..fc9ea0e 100644 --- a/leanframe/core/series.py +++ b/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( diff --git a/tests/unit/test_series.py b/tests/unit/test_series.py index 9abe56b..0baf93d 100644 --- a/tests/unit/test_series.py +++ b/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)