diff --git a/leanframe/core/frame.py b/leanframe/core/frame.py index e9e08d1..546ab7b 100644 --- a/leanframe/core/frame.py +++ b/leanframe/core/frame.py @@ -30,7 +30,8 @@ class DataFrame: Session, instead. """ - def __init__(self, data: ibis_types.Table): + def __init__(self, session, data: ibis_types.Table): + self._session = session self._data = data @property @@ -57,7 +58,7 @@ def __getitem__(self, key: str): # TODO(tswast): Support filtering by a boolean Series if we get a Series # instead of a key? If so, the Series would have to be a column of the # current DataFrame, only. No joins by index key are available. - return leanframe.core.series.Series(self._data[key]) + return leanframe.core.series.Series(self._session, self._data[key]) def assign(self, **kwargs): """Assign new columns to a DataFrame. @@ -81,7 +82,7 @@ def assign(self, **kwargs): new_exprs[name] = expr named_exprs.update(new_exprs) - return DataFrame(self._data.select(**named_exprs)) + return DataFrame(self._session, self._data.select(**named_exprs)) def to_pandas(self) -> pd.DataFrame: """Convert the DataFrame to a pandas.DataFrame. diff --git a/leanframe/core/series.py b/leanframe/core/series.py index ce7521b..755f05b 100644 --- a/leanframe/core/series.py +++ b/leanframe/core/series.py @@ -29,7 +29,8 @@ class Series: Session, instead. """ - def __init__(self, data: ibis_types.Column): + def __init__(self, session, data: ibis_types.Column): + self._session = session self._data = data @property @@ -83,34 +84,34 @@ def empty(self) -> bool: return self.size == 0 def __add__(self, other) -> Series: - return Series(self._data + getattr(other, "_data", other)) + return Series(self._session, self._data + getattr(other, "_data", other)) def __radd__(self, other) -> Series: - return Series(getattr(other, "_data", other) + self._data) + return Series(self._session, getattr(other, "_data", other) + self._data) def __mul__(self, other) -> Series: - return Series(self._data * getattr(other, "_data", other)) + return Series(self._session, self._data * getattr(other, "_data", other)) def __rmul__(self, other) -> Series: - return Series(getattr(other, "_data", other) * self._data) + return Series(self._session, getattr(other, "_data", other) * self._data) def __lt__(self, other) -> Series: - return Series(self._data < getattr(other, "_data", other)) + return Series(self._session, self._data < getattr(other, "_data", other)) def __gt__(self, other) -> Series: - return Series(self._data > getattr(other, "_data", other)) + return Series(self._session, self._data > getattr(other, "_data", other)) def __le__(self, other) -> Series: - return Series(self._data <= getattr(other, "_data", other)) + return Series(self._session, self._data <= getattr(other, "_data", other)) def __ge__(self, other) -> Series: - return Series(self._data >= getattr(other, "_data", other)) + return Series(self._session, self._data >= getattr(other, "_data", other)) def __ne__(self, other) -> Series: # type: ignore[override] - return Series(self._data != getattr(other, "_data", other)) + return Series(self._session, self._data != getattr(other, "_data", other)) def __eq__(self, other) -> Series: # type: ignore[override] - return Series(self._data == getattr(other, "_data", other)) + return Series(self._session, self._data == getattr(other, "_data", other)) def lt(self, other) -> "Series": """Return a boolean Series showing whether each element in the Series is less than the other.""" @@ -137,11 +138,11 @@ def eq(self, other) -> "Series": return self == other def __round__(self, n) -> Series: - return Series(self._data.round(n)) + return Series(self._session, self._data.round(n)) def abs(self) -> "Series": """Return a Series with the absolute value of each element.""" - return Series(self._data.abs()) + return Series(self._session, self._data.abs()) def all(self) -> bool: """Return whether all elements are True.""" @@ -179,18 +180,56 @@ def count(self) -> int: """Return the number of non-null observations in the Series.""" return self._data.count().to_pyarrow().as_py() + def cummax(self) -> "Series": + """Return a Series with the cumulative maximum of each element.""" + return Series(self._session, self._data.cummax()) + + def cummin(self) -> "Series": + """Return a Series with the cumulative minimum of each element.""" + return Series(self._session, self._data.cummin()) + + def cumprod(self) -> "Series": + """Return a Series with the cumulative product of each element.""" + return Series(self._session, self._data.log().cumsum().exp().cast(self._data.type())) + + def cumsum(self) -> "Series": + """Return a Series with the cumulative sum of each element.""" + return Series(self._session, self._data.cumsum()) + + def describe(self) -> "Series": + """Return a Series with descriptive statistics.""" + stats = { + "count": self.count(), + "mean": self.mean(), + "std": self.std(), + "min": self.min(), + "25%": self._data.quantile(0.25).to_pyarrow().as_py(), + "50%": self._data.quantile(0.50).to_pyarrow().as_py(), + "75%": self._data.quantile(0.75).to_pyarrow().as_py(), + "max": self.max(), + } + + index = ["count", "mean", "std", "min", "25%", "50%", "75%", "max"] + pandas_series = pd.Series(stats, name=self.name, index=index) + + return self._session.DataFrame(pandas_series.to_frame())[self.name] + + def diff(self) -> "Series": + """Return a Series with the difference between each element and the previous element.""" + return Series(self._session, self._data - self._data.lag()) + def copy(self) -> Series: """Return a copy of the Series.""" - return Series(self._data) + return Series(self._session, 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)) + return Series(self._session, 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)) + return Series(self._session, self._data.cast(ibis_type)) def to_pandas(self) -> pd.Series: """Convert to a pandas Series.""" diff --git a/leanframe/core/session.py b/leanframe/core/session.py index 1061744..d7ad567 100644 --- a/leanframe/core/session.py +++ b/leanframe/core/session.py @@ -43,18 +43,18 @@ def read_sql_table(self, table_name: str): """Create a DataFrame pointing to the table called ``table_name``.""" import leanframe.core.frame - return leanframe.core.frame.DataFrame(self._backend.table(table_name)) + return leanframe.core.frame.DataFrame(self, self._backend.table(table_name)) def DataFrame(self, data: ibis_types.Table | pandas.DataFrame): """Construct a DataFrame.""" import leanframe.core.frame if isinstance(data, ibis_types.Table): - return leanframe.core.frame.DataFrame(data) + return leanframe.core.frame.DataFrame(self, data) elif isinstance(data, pandas.DataFrame): table_name = f"lf_{''.join(random.choices(_ALPHABET, k=10))}" table = self._backend.create_table(table_name, data, temp=True) - return leanframe.core.frame.DataFrame(table) + return leanframe.core.frame.DataFrame(self, table) else: raise NotImplementedError( f"DataFrame constructor doesn't support {type(data)} data yet." diff --git a/tests/unit/test_series.py b/tests/unit/test_series.py index 23995dd..096ff3b 100644 --- a/tests/unit/test_series.py +++ b/tests/unit/test_series.py @@ -310,6 +310,86 @@ def test_series_abs(session): ) +def test_series_cummax(session): + pandas_df = pd.DataFrame( + {"a": [1, 2, 3, 2, 1]}, + dtype=pd.ArrowDtype(pa.int64()), + ) + df = session.DataFrame(pandas_df) + series = df["a"] + result = series.cummax() + expected = pd.Series( + [1, 2, 3, 3, 3], + name="a", + dtype=pd.ArrowDtype(pa.int64()), + ) + pd.testing.assert_series_equal( + result.to_pandas(), + expected, + check_names=False, + ) + + +def test_series_cummin(session): + pandas_df = pd.DataFrame( + {"a": [3, 2, 1, 2, 3]}, + dtype=pd.ArrowDtype(pa.int64()), + ) + df = session.DataFrame(pandas_df) + series = df["a"] + result = series.cummin() + expected = pd.Series( + [3, 2, 1, 1, 1], + name="a", + dtype=pd.ArrowDtype(pa.int64()), + ) + pd.testing.assert_series_equal( + result.to_pandas(), + expected, + check_names=False, + ) + + +def test_series_cumprod(session): + pandas_df = pd.DataFrame( + {"a": [1, 2, 3, 4, 5]}, + dtype=pd.ArrowDtype(pa.int64()), + ) + df = session.DataFrame(pandas_df) + series = df["a"] + result = series.cumprod() + expected = pd.Series( + [1, 2, 6, 24, 120], + name="a", + dtype=pd.ArrowDtype(pa.int64()), + ) + pd.testing.assert_series_equal( + result.to_pandas(), + expected, + check_names=False, + ) + + +def test_series_cumsum(session): + pandas_df = pd.DataFrame( + {"a": [1, 2, 3, 4, 5]}, + dtype=pd.ArrowDtype(pa.int64()), + ) + df = session.DataFrame(pandas_df) + series = df["a"] + result = series.cumsum() + expected = pd.Series( + [1, 3, 6, 10, 15], + 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]}, @@ -330,6 +410,49 @@ def test_series_astype(session): ) +def test_series_describe(session): + pandas_df = pd.DataFrame( + {"a": [1, 2, 3, 4, 5]}, + dtype=pd.ArrowDtype(pa.int64()), + ) + df = session.DataFrame(pandas_df) + series = df["a"] + result = series.describe() + # Note: leanframe does not support indexes, so we don't check the index. + # The order of the values is guaranteed by the implementation of `describe`. + expected = pd.Series( + [5.0, 3.0, 1.5811388300841898, 1.0, 2.0, 3.0, 4.0, 5.0], + name="a", + dtype=pd.ArrowDtype(pa.float64()), + ) + pd.testing.assert_series_equal( + result.to_pandas(), + expected, + check_names=False, + check_index=False, + ) + + +def test_series_diff(session): + pandas_df = pd.DataFrame( + {"a": [1, 2, 3, 4, 5]}, + dtype=pd.ArrowDtype(pa.int64()), + ) + df = session.DataFrame(pandas_df) + series = df["a"] + result = series.diff() + expected = pd.Series( + [None, 1, 1, 1, 1], + name="a", + dtype=pd.ArrowDtype(pa.int64()), + ) + 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()