Remove Series class in favor of single-column DataFrame

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.

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",
+]

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)