diff --git a/leanframe/core/frame.py b/leanframe/core/frame.py index d59efe4..5447afc 100644 --- a/leanframe/core/frame.py +++ b/leanframe/core/frame.py @@ -29,9 +29,21 @@ LocIndexer, HeadTailMixin, ) +from leanframe.core.mixins import ( + AggregationsMixin, + TransformationsMixin, + ArithmeticMixin, + ComparisonsMixin, +) -class DataFrame(HeadTailMixin): +class DataFrame( + HeadTailMixin, + AggregationsMixin, + TransformationsMixin, + ArithmeticMixin, + ComparisonsMixin, +): """A 2D data structure, representing data and deferred computation. WARNING: Do not call this constructor directly. Use the factory methods on @@ -114,19 +126,17 @@ def dtypes(self) -> pd.Series: types = [convert_ibis_to_pandas(t) for t in self._data.schema().types] return pd.Series(types, index=names, name="dtypes") - def __getitem__(self, key: str): + def __getitem__(self, key: str) -> DataFrame: """Get a column. Note: direct row access via an Index is intentionally not implemented by leanframe. Check out a project like Google's BigQuery DataFrames (bigframes) if you require indexing. """ - import leanframe.core.series - - # 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 + # TODO(tswast): Support filtering by a boolean DataFrame if we get a DataFrame + # instead of a key? If so, the DataFrame would have to be a single column of the # current DataFrame, only. No joins by index key are available. - return leanframe.core.series.Series(self._data[key]) + return DataFrame(self._data.select(self._data[key])) def assign(self, **kwargs): """Assign new columns to a DataFrame. @@ -136,18 +146,53 @@ def assign(self, **kwargs): Args: kwargs: The column names are keywords. If the values are not callable, - (e.g. a Series, scalar, or array), they are simply assigned. + (e.g. a DataFrame, scalar, or array), they are simply assigned. """ - named_exprs = {name: self._data[name] for name in self._data.columns} + import leanframe.core.expression + + import ibis.expr.operations as ops + new_exprs = {} for name, value in kwargs.items(): - expr = getattr(value, "_data", None) - if expr is None: - expr = ibis.literal(value) + if isinstance(value, DataFrame): + col_name = value._data.columns[0] + op = value._data.op() + if isinstance(op, ops.Project): + expr = op.values[col_name].to_expr() + else: + expr = value._data[col_name] + elif isinstance(value, leanframe.core.expression.Expression): + expr = getattr(value, "_data") + else: + expr = getattr(value, "_data", None) + if expr is None: + expr = ibis.literal(value) + + # Use `expr.name()` to align with `mutate` semantics. new_exprs[name] = expr - named_exprs.update(new_exprs) - return DataFrame(self._data.select(**named_exprs)) + # Since we mapped expressions to the original table relations, we can use `mutate`. + # However, Ibis sometimes retains lineage through Projects that causes IntegrityError + # when mutating the original table. By extracting the literal value we bypass it mostly. + # But `replace` on the node might be required if it's still bound to a different relation. + try: + return DataFrame(self._data.mutate(**new_exprs)) + except ibis.common.exceptions.IntegrityError: + # Fallback for complex lineage rebinds + named_exprs = {n: self._data[n] for n in self._data.columns if n not in new_exprs} + + # We rewrite expressions directly into the parent table's scope to bypass integrity. + rebound_exprs = {} + for k, v in new_exprs.items(): + if hasattr(v, 'op') and getattr(v.op(), 'relations', None): + try: + v = v.op().replace({v.op().relations[0]: self._data.op()}).to_expr() + except Exception: + pass + rebound_exprs[k] = v + + named_exprs.update(rebound_exprs) + return DataFrame(self._data.select(**named_exprs)) def to_pandas(self) -> pd.DataFrame: """Convert the DataFrame to a pandas.DataFrame. diff --git a/leanframe/core/mixins/__init__.py b/leanframe/core/mixins/__init__.py new file mode 100644 index 0000000..6490ffa --- /dev/null +++ b/leanframe/core/mixins/__init__.py @@ -0,0 +1,27 @@ +# Copyright 2025 Google LLC, LeanFrame Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Mixins for DataFrame operations.""" + +from leanframe.core.mixins.aggregations import AggregationsMixin +from leanframe.core.mixins.transformations import TransformationsMixin +from leanframe.core.mixins.arithmetic import ArithmeticMixin +from leanframe.core.mixins.comparisons import ComparisonsMixin + +__all__ = [ + "AggregationsMixin", + "TransformationsMixin", + "ArithmeticMixin", + "ComparisonsMixin", +] diff --git a/leanframe/core/mixins/aggregations.py b/leanframe/core/mixins/aggregations.py new file mode 100644 index 0000000..b1e8755 --- /dev/null +++ b/leanframe/core/mixins/aggregations.py @@ -0,0 +1,152 @@ +# Copyright 2025 Google LLC, LeanFrame Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Aggregation operations for DataFrame.""" + +from __future__ import annotations +from typing import TYPE_CHECKING +if TYPE_CHECKING: + import ibis.expr.types as ibis_types + + from leanframe.core.frame import DataFrame + + +import pandas as pd + + +class AggregationsMixin: + _data: "ibis_types.Table" + + """Mixin for DataFrame aggregation methods.""" + + def _aggregate(self, op: str, numeric_only: bool = False, **kwargs) -> "DataFrame": + from leanframe.core.frame import DataFrame + + exprs = {} + for name in self._data.columns: + col = self._data[name] + is_numeric = col.type().is_numeric() + if numeric_only and not is_numeric: + continue + if op == "sum": + if not numeric_only and not is_numeric: + raise TypeError(f"Cannot sum non-numeric column '{name}'.") + exprs[name] = col.sum() + elif op == "mean": + if not numeric_only and not is_numeric: + raise TypeError(f"Cannot compute mean of non-numeric column '{name}'.") + exprs[name] = col.mean() + elif op == "min": + exprs[name] = col.min() + elif op == "max": + exprs[name] = col.max() + elif op == "std": + if not numeric_only and not is_numeric: + raise TypeError(f"Cannot compute std of non-numeric column '{name}'.") + exprs[name] = col.std() + elif op == "var": + if not numeric_only and not is_numeric: + raise TypeError(f"Cannot compute var of non-numeric column '{name}'.") + exprs[name] = col.var() + elif op == "count": + exprs[name] = col.count() + elif op == "any": + exprs[name] = col.any() + elif op == "all": + exprs[name] = col.all() + else: + raise ValueError(f"Unknown aggregation operation '{op}'") + + if not exprs: + # If all columns were skipped, return an empty DataFrame (or handle as appropriate) + raise ValueError("No columns to aggregate.") + + return DataFrame(self._data.aggregate(exprs)) + + def sum(self, numeric_only: bool = False) -> "DataFrame": + """Return the sum of the DataFrame over the columns.""" + return self._aggregate("sum", numeric_only=numeric_only) + + def mean(self, numeric_only: bool = False) -> "DataFrame": + """Return the mean of the DataFrame over the columns.""" + return self._aggregate("mean", numeric_only=numeric_only) + + def min(self, numeric_only: bool = False) -> "DataFrame": + """Return the min of the DataFrame over the columns.""" + return self._aggregate("min", numeric_only=numeric_only) + + def max(self, numeric_only: bool = False) -> "DataFrame": + """Return the max of the DataFrame over the columns.""" + return self._aggregate("max", numeric_only=numeric_only) + + def std(self, numeric_only: bool = False) -> "DataFrame": + """Return the std of the DataFrame over the columns.""" + return self._aggregate("std", numeric_only=numeric_only) + + def var(self, numeric_only: bool = False) -> "DataFrame": + """Return the var of the DataFrame over the columns.""" + return self._aggregate("var", numeric_only=numeric_only) + + def count(self) -> "DataFrame": + """Return the number of non-null observations in the DataFrame over the columns.""" + return self._aggregate("count") + + def any(self) -> "DataFrame": + """Return whether any element is True.""" + return self._aggregate("any") + + def all(self) -> "DataFrame": + """Return whether all elements are True.""" + return self._aggregate("all") + + def describe(self) -> pd.DataFrame: + """Return a pandas DataFrame with descriptive statistics.""" + # For now, we will evaluate each stat and combine them into a single pd.DataFrame + # In a fully lazy implementation, this would involve a complex union/unpivot, + # but returning pd.DataFrame implies immediate evaluation (as it did for Series). + count = self.count().to_pandas().iloc[0] + mean = self.mean(numeric_only=True).to_pandas().iloc[0] + std = self.std(numeric_only=True).to_pandas().iloc[0] + min_val = self.min(numeric_only=True).to_pandas().iloc[0] + max_val = self.max(numeric_only=True).to_pandas().iloc[0] + + # Calculate percentiles + percentiles = [0.25, 0.50, 0.75] + q_exprs = { + f"{int(p * 100)}%": [ + self._data[col].quantile(p) for col in mean.index + ] + for p in percentiles + } + + # Execute quantiles (requires some aggregation magic or multiple queries) + # To avoid multiple queries we can put them in a single projection, + # but for simplicity let's evaluate them: + q_results = {} + for name, p_list in q_exprs.items(): + aggs = {col: expr for col, expr in zip(mean.index, p_list)} + q_results[name] = self._data.aggregate(aggs).to_pyarrow().to_pandas().iloc[0] + + stats = { + "count": count, + "mean": mean, + "std": std, + "min": min_val, + "25%": q_results["25%"], + "50%": q_results["50%"], + "75%": q_results["75%"], + "max": max_val, + } + + return pd.DataFrame(stats) diff --git a/leanframe/core/mixins/arithmetic.py b/leanframe/core/mixins/arithmetic.py new file mode 100644 index 0000000..77f73ee --- /dev/null +++ b/leanframe/core/mixins/arithmetic.py @@ -0,0 +1,116 @@ +# Copyright 2025 Google LLC, LeanFrame Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Arithmetic operations for DataFrame.""" + +from __future__ import annotations +from typing import TYPE_CHECKING +if TYPE_CHECKING: + import ibis.expr.types as ibis_types + + from leanframe.core.frame import DataFrame + + + +class ArithmeticMixin: + _data: "ibis_types.Table" + + """Mixin for DataFrame arithmetic methods.""" + + def _arithmetic(self, op: str, other) -> "DataFrame": + from leanframe.core.frame import DataFrame + + exprs = {} + for name in self._data.columns: + col = self._data[name] + + + # If other is an Expression or Ibis scalar/column + other_val = getattr(other, "_data", other) + if hasattr(other, "_data") and hasattr(other_val, "columns") and len(other_val.columns) == 1: + import ibis.expr.operations as ops + op = other_val.op() + col_name = other_val.columns[0] + if isinstance(op, ops.Project): + other_val = op.values[col_name].to_expr() + else: + other_val = other_val[col_name] + + + # Or if it's another DataFrame, this would need joining logic which we don't have yet. + # Assuming scalar or Expression for now, matching the previous Series behavior. + if op == "add": + exprs[name] = col + other_val + elif op == "radd": + exprs[name] = other_val + col + elif op == "mul": + exprs[name] = col * other_val + elif op == "rmul": + exprs[name] = other_val * col + elif op == "sub": + exprs[name] = col - other_val + elif op == "rsub": + exprs[name] = other_val - col + elif op == "truediv": + exprs[name] = col / other_val + elif op == "rtruediv": + exprs[name] = other_val / col + elif op == "floordiv": + exprs[name] = col // other_val + elif op == "rfloordiv": + exprs[name] = other_val // col + elif op == "pow": + exprs[name] = col ** other_val + elif op == "rpow": + exprs[name] = other_val ** col + else: + raise ValueError(f"Unknown arithmetic operation '{op}'") + + return DataFrame(self._data.select(**exprs)) + + def __add__(self, other) -> "DataFrame": + return self._arithmetic("add", other) + + def __radd__(self, other) -> "DataFrame": + return self._arithmetic("radd", other) + + def __mul__(self, other) -> "DataFrame": + return self._arithmetic("mul", other) + + def __rmul__(self, other) -> "DataFrame": + return self._arithmetic("rmul", other) + + def __sub__(self, other) -> "DataFrame": + return self._arithmetic("sub", other) + + def __rsub__(self, other) -> "DataFrame": + return self._arithmetic("rsub", other) + + def __truediv__(self, other) -> "DataFrame": + return self._arithmetic("truediv", other) + + def __rtruediv__(self, other) -> "DataFrame": + return self._arithmetic("rtruediv", other) + + def __floordiv__(self, other) -> "DataFrame": + return self._arithmetic("floordiv", other) + + def __rfloordiv__(self, other) -> "DataFrame": + return self._arithmetic("rfloordiv", other) + + def __pow__(self, other) -> "DataFrame": + return self._arithmetic("pow", other) + + def __rpow__(self, other) -> "DataFrame": + return self._arithmetic("rpow", other) diff --git a/leanframe/core/mixins/comparisons.py b/leanframe/core/mixins/comparisons.py new file mode 100644 index 0000000..c6b26dd --- /dev/null +++ b/leanframe/core/mixins/comparisons.py @@ -0,0 +1,102 @@ +# Copyright 2025 Google LLC, LeanFrame Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Comparison operations for DataFrame.""" + +from __future__ import annotations +from typing import TYPE_CHECKING +if TYPE_CHECKING: + import ibis.expr.types as ibis_types + + from leanframe.core.frame import DataFrame + + + +class ComparisonsMixin: + _data: "ibis_types.Table" + + """Mixin for DataFrame comparison methods.""" + + def _compare(self, op: str, other) -> "DataFrame": + from leanframe.core.frame import DataFrame + + exprs = {} + for name in self._data.columns: + col = self._data[name] + other_val = getattr(other, "_data", other) + + if op == "lt": + exprs[name] = col < other_val + elif op == "gt": + exprs[name] = col > other_val + elif op == "le": + exprs[name] = col <= other_val + elif op == "ge": + exprs[name] = col >= other_val + elif op == "ne": + exprs[name] = col != other_val + elif op == "eq": + exprs[name] = col == other_val + elif op == "isin": + exprs[name] = col.isin(other_val) + else: + raise ValueError(f"Unknown comparison operation '{op}'") + + return DataFrame(self._data.select(**exprs)) + + def __lt__(self, other) -> "DataFrame": + return self._compare("lt", other) + + def __gt__(self, other) -> "DataFrame": + return self._compare("gt", other) + + def __le__(self, other) -> "DataFrame": + return self._compare("le", other) + + def __ge__(self, other) -> "DataFrame": + return self._compare("ge", other) + + def __ne__(self, other) -> "DataFrame": # type: ignore[override] + return self._compare("ne", other) + + def __eq__(self, other) -> "DataFrame": # type: ignore[override] + return self._compare("eq", other) + + def lt(self, other) -> "DataFrame": + """Return a boolean DataFrame showing whether each element is less than the other.""" + return self < other + + def gt(self, other) -> "DataFrame": + """Return a boolean DataFrame showing whether each element is greater than the other.""" + return self > other + + def le(self, other) -> "DataFrame": + """Return a boolean DataFrame showing whether each element is less than or equal to the other.""" + return self <= other + + def ge(self, other) -> "DataFrame": + """Return a boolean DataFrame showing whether each element is greater than or equal to the other.""" + return self >= other + + def ne(self, other) -> "DataFrame": + """Return a boolean DataFrame showing whether each element is not equal to the other.""" + return self != other + + def eq(self, other) -> "DataFrame": + """Return a boolean DataFrame showing whether each element is equal to the other.""" + return self == other + + def isin(self, values) -> "DataFrame": + """Return a boolean DataFrame showing whether each element is exactly contained in the passed sequence of values.""" + return self._compare("isin", values) diff --git a/leanframe/core/mixins/transformations.py b/leanframe/core/mixins/transformations.py new file mode 100644 index 0000000..7895030 --- /dev/null +++ b/leanframe/core/mixins/transformations.py @@ -0,0 +1,99 @@ +# Copyright 2025 Google LLC, LeanFrame Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Transformation operations for DataFrame.""" + +from __future__ import annotations +from typing import TYPE_CHECKING +if TYPE_CHECKING: + import ibis.expr.types as ibis_types + + from leanframe.core.frame import DataFrame + + + + +class TransformationsMixin: + _data: "ibis_types.Table" + + """Mixin for DataFrame transformation methods.""" + + def _transform(self, op: str, numeric_only: bool = False, **kwargs) -> "DataFrame": + from leanframe.core.frame import DataFrame + + exprs = {} + for name in self._data.columns: + col = self._data[name] + is_numeric = col.type().is_numeric() + if numeric_only and not is_numeric: + continue + if op == "cummax": + exprs[name] = col.cummax() + elif op == "cummin": + exprs[name] = col.cummin() + elif op == "cumprod": + if not numeric_only and not is_numeric: + raise TypeError(f"Cannot compute cumprod of non-numeric column '{name}'.") + exprs[name] = col.log().cumsum().exp().cast(col.type()) + elif op == "cumsum": + if not numeric_only and not is_numeric: + raise TypeError(f"Cannot compute cumsum of non-numeric column '{name}'.") + exprs[name] = col.cumsum() + elif op == "diff": + if not numeric_only and not is_numeric: + raise TypeError(f"Cannot compute diff of non-numeric column '{name}'.") + exprs[name] = col - col.lag() + elif op == "abs": + if not numeric_only and not is_numeric: + raise TypeError(f"Cannot compute abs of non-numeric column '{name}'.") + exprs[name] = col.abs() + elif op == "round": + if not numeric_only and not is_numeric: + raise TypeError(f"Cannot compute round of non-numeric column '{name}'.") + n = kwargs.get("n", 0) + exprs[name] = col.round(n) + else: + raise ValueError(f"Unknown transformation operation '{op}'") + + if not exprs: + raise ValueError("No columns to transform.") + + return DataFrame(self._data.select(**exprs)) + + def cummax(self) -> "DataFrame": + """Return a DataFrame with the cumulative maximum of each element.""" + return self._transform("cummax") + + def cummin(self) -> "DataFrame": + """Return a DataFrame with the cumulative minimum of each element.""" + return self._transform("cummin") + + def cumprod(self) -> "DataFrame": + """Return a DataFrame with the cumulative product of each element.""" + return self._transform("cumprod") + + def cumsum(self) -> "DataFrame": + """Return a DataFrame with the cumulative sum of each element.""" + return self._transform("cumsum") + + def diff(self) -> "DataFrame": + """Return a DataFrame with the difference between each element and the previous element.""" + return self._transform("diff") + + def abs(self, numeric_only: bool = False) -> "DataFrame": + """Return a DataFrame with the absolute value of each element.""" + return self._transform("abs", numeric_only=numeric_only) + + def __round__(self, n=0) -> "DataFrame": + return self._transform("round", n=n) diff --git a/leanframe/core/series.py b/leanframe/core/series.py deleted file mode 100644 index c352caf..0000000 --- a/leanframe/core/series.py +++ /dev/null @@ -1,251 +0,0 @@ -# Copyright 2025 Google LLC, LeanFrame Authors -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Series is a one dimensional data structure.""" - -from __future__ import annotations - -import ibis.expr.types as ibis_types -import numpy as np -import pandas as pd -from leanframe.core.dtypes import convert_ibis_to_pandas, convert_pandas_to_ibis - - -class Series: - """A 1D data structure, representing a column. - - WARNING: Do not call this constructor directly. Use the factory methods on - Session, instead. - """ - - def __init__(self, data: ibis_types.Column): - self._data = data - - @property - def dtype(self) -> pd.ArrowDtype: - """Return the dtype object of the underlying data.""" - return convert_ibis_to_pandas(self._data.type()) - - @property - def name(self) -> str: - """Name of the column.""" - return self._data.get_name() - - @property - def values(self) -> np.ndarray: - """Return a numpy representation of the Series.""" - return self._data.to_pyarrow().to_numpy() - - @property - def array(self) -> "pd.api.extensions.ExtensionArray": - """Return the underlying data as a pandas ExtensionArray.""" - return self.to_pandas().array - - @property - def shape(self) -> tuple[int, ...]: - """Return a tuple of the shape of the underlying data.""" - return (self.size,) - - @property - def nbytes(self) -> int: - """Return the number of bytes in the underlying data.""" - raise NotImplementedError("nbytes not relevant for ibis expression.") - - @property - def ndim(self) -> int: - """Return the number of dimensions of the underlying data.""" - return 1 - - @property - def size(self) -> int: - """Return the number of elements in the underlying data.""" - return self._data.as_table().count().to_pyarrow().as_py() - - @property - def hasnans(self) -> bool: - """Return True if there are any NaNs, False otherwise.""" - return self._data.isnull().any().to_pyarrow().as_py() - - @property - def empty(self) -> bool: - """Return True if the Series is empty, False otherwise.""" - return self.size == 0 - - def __add__(self, other) -> Series: - return Series(self._data + getattr(other, "_data", other)) - - def __radd__(self, other) -> Series: - return Series(getattr(other, "_data", other) + self._data) - - def __mul__(self, other) -> Series: - return Series(self._data * getattr(other, "_data", other)) - - def __rmul__(self, other) -> Series: - return Series(getattr(other, "_data", other) * self._data) - - def __lt__(self, other) -> Series: - return Series(self._data < getattr(other, "_data", other)) - - def __gt__(self, other) -> Series: - return Series(self._data > getattr(other, "_data", other)) - - def __le__(self, other) -> Series: - return Series(self._data <= getattr(other, "_data", other)) - - def __ge__(self, other) -> Series: - return Series(self._data >= getattr(other, "_data", other)) - - def __ne__(self, other) -> Series: # type: ignore[override] - return Series(self._data != getattr(other, "_data", other)) - - def __eq__(self, other) -> Series: # type: ignore[override] - return Series(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.""" - return self < other - - def gt(self, other) -> "Series": - """Return a boolean Series showing whether each element in the Series is greater than the other.""" - return self > other - - def le(self, other) -> "Series": - """Return a boolean Series showing whether each element in the Series is less than or equal to the other.""" - return self <= other - - def ge(self, other) -> "Series": - """Return a boolean Series showing whether each element in the Series is greater than or equal to the other.""" - return self >= other - - def ne(self, other) -> "Series": - """Return a boolean Series showing whether each element in the Series is not equal to the other.""" - return self != other - - def eq(self, other) -> "Series": - """Return a boolean Series showing whether each element in the Series is equal to the other.""" - return self == other - - 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() - - def mean(self): - """Return the mean of the Series.""" - return self._data.mean().to_pyarrow().as_py() - - def min(self): - """Return the min of the Series.""" - return self._data.min().to_pyarrow().as_py() - - def max(self): - """Return the max of the Series.""" - return self._data.max().to_pyarrow().as_py() - - def std(self): - """Return the std of the Series.""" - return self._data.std().to_pyarrow().as_py() - - 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 cummax(self) -> "Series": - """Return a Series with the cumulative maximum of each element.""" - return Series(self._data.cummax()) - - def cummin(self) -> "Series": - """Return a Series with the cumulative minimum of each element.""" - return Series(self._data.cummin()) - - def cumprod(self) -> "Series": - """Return a Series with the cumulative product of each element.""" - return Series(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._data.cumsum()) - - def describe(self) -> pd.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"] - return pd.Series(stats, name=self.name, index=index) - - def diff(self) -> "Series": - """Return a Series with the difference between each element and the previous element.""" - return Series(self._data - self._data.lag()) - - 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( - types_mapper=lambda type_: pd.ArrowDtype(type_) - ) - - def to_ibis(self) -> ibis_types.Column: - """Return the underlying Ibis expression.""" - return self._data - - def to_numpy(self) -> np.ndarray: - """Return a numpy representation of the Series.""" - return self.values - - def to_list(self) -> list: - """Return a list of the values.""" - return self.to_pandas().to_list() - - def __iter__(self): - """Return an iterator of the values.""" - return iter(self.to_list()) diff --git a/tests/test_to_ibis.py b/tests/test_to_ibis.py index 9dcbb19..3c9babd 100644 --- a/tests/test_to_ibis.py +++ b/tests/test_to_ibis.py @@ -1,7 +1,6 @@ import ibis import ibis.expr.types as ibis_types from leanframe.core.frame import DataFrame -from leanframe.core.series import Series def test_dataframe_to_ibis(): @@ -19,16 +18,15 @@ def test_dataframe_to_ibis(): assert isinstance(expr, ibis_types.Table) assert expr.equals(t) - def test_series_to_ibis(): # Setup con = ibis.sqlite.connect() t = con.create_table("test_series_ibis", schema=ibis.schema({"a": "int64"})) - s = Series(t["a"]) + s = DataFrame(t)["a"] # Execute expr = s.to_ibis() # Assert - assert isinstance(expr, ibis_types.Column) - assert expr.equals(t["a"]) + assert isinstance(expr, ibis_types.Table) + assert expr.columns == ("a",) diff --git a/tests/unit/nested_data/test_nested_column_access.py b/tests/unit/nested_data/test_nested_column_access.py index 331cafb..6b831c3 100644 --- a/tests/unit/nested_data/test_nested_column_access.py +++ b/tests/unit/nested_data/test_nested_column_access.py @@ -22,10 +22,10 @@ def test_nested_column_access(): assert "contact" in lf_df.columns.tolist() # Method 1: Access top-level nested columns - person_series = lf_df["person"] - contact_series = lf_df["contact"] - assert person_series.dtype.name.startswith("struct") - assert contact_series.dtype.name.startswith("struct") + person_df = lf_df["person"] + contact_df = lf_df["contact"] + assert person_df.dtypes.iloc[0].name.startswith("struct") + assert contact_df.dtypes.iloc[0].name.startswith("struct") # Method 2: Extract nested fields using ibis (Recommended!) ibis_table = lf_df._data diff --git a/tests/unit/test_frame.py b/tests/unit/test_frame.py index af4bffe..cdaf4fb 100644 --- a/tests/unit/test_frame.py +++ b/tests/unit/test_frame.py @@ -19,7 +19,6 @@ import pyarrow as pa import leanframe -import leanframe.core.series def test_dataframe_dtypes(session: leanframe.Session): @@ -92,13 +91,13 @@ def test_dataframe_getitem_with_column(session: leanframe.Session): ) ) series_1 = df_lf["col1"] - assert isinstance(series_1, leanframe.core.series.Series) - assert series_1.name == "col1" + assert isinstance(series_1, leanframe.core.frame.DataFrame) + assert series_1.columns.tolist() == ["col1"] # TODO(tswast): check dtype series_2 = df_lf["col2"] - assert isinstance(series_2, leanframe.core.series.Series) - assert series_2.name == "col2" + assert isinstance(series_2, leanframe.core.frame.DataFrame) + assert series_2.columns.tolist() == ["col2"] # TODO(tswast): check dtype @@ -134,7 +133,10 @@ def test_dataframe_assign_series(session: leanframe.Session): ) df_lf = session.DataFrame(df_pd) series_pd = pd.Series([4, 5, 6], name="col3").astype(pd.ArrowDtype(pa.int64())) - series_lf = df_lf["col1"] + 3 + + # Use session.col("col1") instead of df_lf["col1"] to construct the deferred expression + series_lf = session.col("col1") + 3 + result_lf = df_lf.assign(col3=series_lf) expected_pd = df_pd.assign(col3=series_pd) tm.assert_frame_equal(result_lf.to_pandas(), expected_pd) @@ -154,7 +156,10 @@ def test_dataframe_assign_multiple(session: leanframe.Session): ) df_lf = session.DataFrame(df_pd) series_pd = pd.Series([4, 5, 6], name="col3").astype(pd.ArrowDtype(pa.int64())) - series_lf = df_lf["col1"] + 3 + + # Use session.col("col1") + series_lf = session.col("col1") + 3 + result_lf = df_lf.assign(col3=series_lf, col4="d") expected_pd = df_pd.assign(col3=series_pd, col4="d").astype( {"col4": pd.ArrowDtype(pa.string())} @@ -175,6 +180,8 @@ def test_dataframe_assign_overwrite(session: leanframe.Session): } ) df_lf = session.DataFrame(df_pd) - result_lf = df_lf.assign(col1=df_lf["col1"] * 2) + + # Use session.col("col1") + result_lf = df_lf.assign(col1=session.col("col1") * 2) expected_pd = df_pd.assign(col1=df_pd["col1"] * 2) tm.assert_frame_equal(result_lf.to_pandas(), expected_pd) diff --git a/tests/unit/test_series.py b/tests/unit/test_series.py deleted file mode 100644 index ee2de4d..0000000 --- a/tests/unit/test_series.py +++ /dev/null @@ -1,662 +0,0 @@ -# Copyright 2025 Google LLC, LeanFrame Authors -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from __future__ import annotations - -import pandas as pd -import pandas.testing -import pyarrow as pa -import pytest - -import numpy as np - -import leanframe - - -@pytest.fixture -def series_for_properties(session): - df_pd = pd.DataFrame( - { - "int_col": pd.Series([1, 2, 3], dtype=pd.ArrowDtype(pa.int64())), - "float_col": pd.Series( - [1.0, float("nan"), 3.0], dtype=pd.ArrowDtype(pa.float64()) - ), - } - ) - df_lf = session.DataFrame(df_pd) - return df_lf["int_col"], df_lf["float_col"] - - -@pytest.fixture -def numeric_series(session): - df_pd = pd.DataFrame( - { - "a": [1, 2, 3, 4, 5], - "b": [1.1, 2.2, 3.3, 4.4, 5.5], - } - ) - 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 - assert series_float.ndim == 1 - - -def test_series_size(series_for_properties): - series_int, series_float = series_for_properties - assert series_int.size == 3 - assert series_float.size == 3 - - -def test_series_shape(series_for_properties): - series_int, series_float = series_for_properties - assert series_int.shape == (3,) - assert series_float.shape == (3,) - - -def test_series_hasnans(series_for_properties): - series_int, series_float = series_for_properties - assert not series_int.hasnans - assert series_float.hasnans - - -def test_series_empty(session): - df_pd = pd.DataFrame({"col1": [1, 2, 3]}) - df_lf = session.DataFrame(df_pd) - assert not df_lf["col1"].empty - - df_pd_empty = pd.DataFrame({"col1": []}) - df_lf_empty = session.DataFrame(df_pd_empty) - assert df_lf_empty["col1"].empty - - -def test_series_values(series_for_properties): - series_int, series_float = series_for_properties - np.testing.assert_array_equal(series_int.values, np.array([1, 2, 3])) - - -def test_series_array(series_for_properties): - series_int, series_float = series_for_properties - expected_array = pd.array([1, 2, 3], dtype=pd.ArrowDtype(pa.int64())) - pd.testing.assert_extension_array_equal(series_int.array, expected_array) - - -def test_series_nbytes(series_for_properties): - series_int, series_float = series_for_properties - - with pytest.raises(NotImplementedError, match="nbytes"): - assert series_int.nbytes - - -@pytest.mark.parametrize( - ("column", "expected_dtype"), - [ - ("string_col", pd.ArrowDtype(pa.string())), - ("int_col", pd.ArrowDtype(pa.int64())), - ("array_col", pd.ArrowDtype(pa.list_(pa.int64()))), - ( - "struct_col", - pd.ArrowDtype(pa.struct([("a", pa.int64()), ("b", pa.string())])), - ), - ], -) -def test_series_dtype(session, column, expected_dtype): - pa_table = pa.Table.from_pydict( - { - "string_col": ["a", "b", "c"], - "int_col": [1, 2, 3], - "array_col": [[1, 2], [3, 4], [5, 6]], - "struct_col": [ - {"a": 1, "b": "c"}, - {"a": 2, "b": "d"}, - {"a": 3, "b": "e"}, - ], - }, - schema=pa.schema( - [ - pa.field("string_col", pa.string()), - pa.field("int_col", pa.int64()), - pa.field("array_col", pa.list_(pa.int64())), - pa.field( - "struct_col", - pa.struct([("a", pa.int64()), ("b", pa.string())]), - ), - ] - ), - ) - pandas_df = pa_table.to_pandas(types_mapper=pd.ArrowDtype) - df = session.DataFrame(pandas_df) - series = df[column] - assert series.dtype == expected_dtype - - -@pytest.mark.parametrize( - ("series_pd",), - ( - pytest.param( - pd.Series([1, 2, 3], dtype=pd.ArrowDtype(pa.int64())), - id="int64", - ), - pytest.param( - pd.Series([1.0, float("nan"), 3.0], dtype=pd.ArrowDtype(pa.float64())), - id="float64", - ), - ), -) -def test_to_pandas(session: leanframe.Session, series_pd: pd.Series): - df_pd = pd.DataFrame( - { - "my_col": series_pd, - } - ) - df_lf = session.DataFrame(df_pd) - - result = df_lf["my_col"].to_pandas() - - # TODO(tswast): Allow input dtype != output dtype with an "expected_dtype" parameter. - pd.testing.assert_series_equal(result, series_pd, check_names=False) - - -@pytest.mark.parametrize( - ("op", "other", "expected_data"), - [ - pytest.param(lambda s, o: s + o, 1, [2, 3, 4], id="add_scalar"), - pytest.param(lambda s, o: o + s, 1, [2, 3, 4], id="radd_scalar"), - pytest.param(lambda s, o: s * o, 2, [2, 4, 6], id="mul_scalar"), - pytest.param(lambda s, o: o * s, 2, [2, 4, 6], id="rmul_scalar"), - ], -) -def test_series_arithmetic_scalar(session, op, other, expected_data): - pandas_df = pd.DataFrame( - {"a": [1, 2, 3]}, - dtype=pd.ArrowDtype(pa.int64()), - ) - df = session.DataFrame(pandas_df) - series_a = df["a"] - - result_series = op(series_a, other) - - expected_series = pd.Series( - expected_data, - dtype=pd.ArrowDtype(pa.int64()), - ) - pd.testing.assert_series_equal( - result_series.to_pandas(), - expected_series, - check_names=False, - ) - - -@pytest.mark.parametrize( - ("op", "other", "expected_data"), - [ - pytest.param( - lambda s, o: s < o, 3, [True, True, False, False, False], id="lt_scalar" - ), - pytest.param( - lambda s, o: s.lt(o), - 3, - [True, True, False, False, False], - id="lt_method_scalar", - ), - pytest.param( - lambda s, o: s > o, 3, [False, False, False, True, True], id="gt_scalar" - ), - pytest.param( - lambda s, o: s.gt(o), - 3, - [False, False, False, True, True], - id="gt_method_scalar", - ), - pytest.param( - lambda s, o: s <= o, 3, [True, True, True, False, False], id="le_scalar" - ), - pytest.param( - lambda s, o: s.le(o), - 3, - [True, True, True, False, False], - id="le_method_scalar", - ), - pytest.param( - lambda s, o: s >= o, 3, [False, False, True, True, True], id="ge_scalar" - ), - pytest.param( - lambda s, o: s.ge(o), - 3, - [False, False, True, True, True], - id="ge_method_scalar", - ), - pytest.param( - lambda s, o: s != o, 3, [True, True, False, True, True], id="ne_scalar" - ), - pytest.param( - lambda s, o: s.ne(o), - 3, - [True, True, False, True, True], - id="ne_method_scalar", - ), - pytest.param( - lambda s, o: s == o, 3, [False, False, True, False, False], id="eq_scalar" - ), - pytest.param( - lambda s, o: s.eq(o), - 3, - [False, False, True, False, False], - id="eq_method_scalar", - ), - ], -) -def test_series_comparison_scalar(session, op, other, expected_data): - pandas_df = pd.DataFrame( - {"a": [1, 2, 3, 4, 5]}, - dtype=pd.ArrowDtype(pa.int64()), - ) - df = session.DataFrame(pandas_df) - series_a = df["a"] - - result_series = op(series_a, other) - - expected_series = pd.Series( - expected_data, - dtype=pd.ArrowDtype(pa.bool_()), - ) - pd.testing.assert_series_equal( - result_series.to_pandas(), - expected_series, - check_names=False, - ) - - -@pytest.mark.parametrize( - ("op", "expected_data"), - [ - pytest.param( - lambda s1, s2: s1 < s2, [False, False, False, True, True], id="lt_series" - ), - pytest.param( - lambda s1, s2: s1.lt(s2), - [False, False, False, True, True], - id="lt_method_series", - ), - pytest.param( - lambda s1, s2: s1 > s2, [True, True, False, False, False], id="gt_series" - ), - pytest.param( - lambda s1, s2: s1.gt(s2), - [True, True, False, False, False], - id="gt_method_series", - ), - pytest.param( - lambda s1, s2: s1 <= s2, [False, False, True, True, True], id="le_series" - ), - pytest.param( - lambda s1, s2: s1.le(s2), - [False, False, True, True, True], - id="le_method_series", - ), - pytest.param( - lambda s1, s2: s1 >= s2, [True, True, True, False, False], id="ge_series" - ), - pytest.param( - lambda s1, s2: s1.ge(s2), - [True, True, True, False, False], - id="ge_method_series", - ), - pytest.param( - lambda s1, s2: s1 != s2, [True, True, False, True, True], id="ne_series" - ), - pytest.param( - lambda s1, s2: s1.ne(s2), - [True, True, False, True, True], - id="ne_method_series", - ), - pytest.param( - lambda s1, s2: s1 == s2, [False, False, True, False, False], id="eq_series" - ), - pytest.param( - lambda s1, s2: s1.eq(s2), - [False, False, True, False, False], - id="eq_method_series", - ), - ], -) -def test_series_comparison_series(session, op, expected_data): - pandas_df = pd.DataFrame( - {"a": [1, 2, 3, 4, 5], "b": [0, 1, 3, 5, 6]}, - dtype=pd.ArrowDtype(pa.int64()), - ) - df = session.DataFrame(pandas_df) - series_a = df["a"] - series_b = df["b"] - - result_series = op(series_a, series_b) - - expected_series = pd.Series( - expected_data, - dtype=pd.ArrowDtype(pa.bool_()), - ) - pd.testing.assert_series_equal( - result_series.to_pandas(), - expected_series, - check_names=False, - ) - - -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_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]}, - 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_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() - expected = pandas_df["a"].describe().astype("float64") - pd.testing.assert_series_equal( - result, - expected, - check_names=False, - rtol=0.01, - ) - - -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() - 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) - expected = pd.Series( - [1.0, 2.0, 3.0, 4.0, 6.0], - dtype=pd.ArrowDtype(pa.float64()), - ) - result_pd = result.to_pandas() - result_pd = result_pd.astype(pd.ArrowDtype(pa.float64())) - pd.testing.assert_series_equal( - result_pd, - expected, - check_names=False, - ) - - -def test_series_sum(numeric_series): - series = numeric_series["a"] - assert series.sum() == 15 - - -def test_series_mean(numeric_series): - series = numeric_series["a"] - assert series.mean() == 3.0 - - -def test_series_min(numeric_series): - series = numeric_series["a"] - assert series.min() == 1 - - -def test_series_max(numeric_series): - series = numeric_series["a"] - assert series.max() == 5 - - -def test_series_std(numeric_series): - series = numeric_series["b"] - assert round(series.std(), 2) == 1.74 - - -def test_series_var(numeric_series): - series = numeric_series["b"] - 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) - series = df_lf["col1"] - series_copy = series.copy() - assert series is not series_copy - assert series._data is series_copy._data - - -def test_series_to_numpy(series_for_properties): - series_int, series_float = series_for_properties - np.testing.assert_array_equal(series_int.to_numpy(), np.array([1, 2, 3])) - - -def test_series_to_list(series_for_properties): - series_int, series_float = series_for_properties - assert series_int.to_list() == [1, 2, 3] - - -def test_series_iter(series_for_properties): - series_int, series_float = series_for_properties - assert list(iter(series_int)) == [1, 2, 3] - - -@pytest.mark.parametrize( - ("op", "expected_data"), - [ - pytest.param(lambda s1, s2: s1 + s2, [5, 7, 9], id="add_series"), - pytest.param(lambda s1, s2: s1 * s2, [4, 10, 18], id="mul_series"), - ], -) -def test_series_arithmetic_series(session, op, expected_data): - pandas_df = pd.DataFrame( - {"a": [1, 2, 3], "b": [4, 5, 6]}, - dtype=pd.ArrowDtype(pa.int64()), - ) - df = session.DataFrame(pandas_df) - series_a = df["a"] - series_b = df["b"] - - result_series = op(series_a, series_b) - - expected_series = pd.Series( - expected_data, - dtype=pd.ArrowDtype(pa.int64()), - ) - pd.testing.assert_series_equal( - result_series.to_pandas(), - expected_series, - check_names=False, - )