Add session.col() method for deferred series expressions

leanframe/core/expression.py

@@ -0,0 +1,56 @@
+# 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.
+
+"""Expression represents a deferred computation."""
+
+from __future__ import annotations
+
+import ibis.expr.types as ibis_types
+
+
+class Expression:
+    """A wrapper around an Ibis expression for deferred computation."""
+
+    def __init__(self, data: ibis_types.Value):
+        self._data = data
+
+    def __add__(self, other) -> Expression:
+        return Expression(self._data + getattr(other, "_data", other))
+
+    def __radd__(self, other) -> Expression:
+        return Expression(getattr(other, "_data", other) + self._data)
+
+    def __mul__(self, other) -> Expression:
+        return Expression(self._data * getattr(other, "_data", other))
+
+    def __rmul__(self, other) -> Expression:
+        return Expression(getattr(other, "_data", other) * self._data)
+
+    def __lt__(self, other) -> Expression:
+        return Expression(self._data < getattr(other, "_data", other))
+
+    def __gt__(self, other) -> Expression:
+        return Expression(self._data > getattr(other, "_data", other))
+
+    def __le__(self, other) -> Expression:
+        return Expression(self._data <= getattr(other, "_data", other))
+
+    def __ge__(self, other) -> Expression:
+        return Expression(self._data >= getattr(other, "_data", other))
+
+    def __ne__(self, other) -> Expression:  # type: ignore[override]
+        return Expression(self._data != getattr(other, "_data", other))
+
+    def __eq__(self, other) -> Expression:  # type: ignore[override]
+        return Expression(self._data == getattr(other, "_data", other))

leanframe/core/session.py

@@ -66,3 +66,9 @@ def DataFrame(self, data: ibis_types.Table | pandas.DataFrame):
             raise NotImplementedError(
                 f"DataFrame constructor doesn't support {type(data)} data yet."
             )
+
+    def col(self, name: str):
+        """Return a new expression object which is a deferred series."""
+        import leanframe.core.expression
+
+        return leanframe.core.expression.Expression(ibis.deferred[name])

tests/unit/session/test_col.py

@@ -0,0 +1,72 @@
+# 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.
+
+"""Tests for Session.col."""
+
+import ibis
+import pandas as pd
+
+import leanframe.core.expression
+from leanframe.core.session import Session
+
+
+def test_col_assign_with_memtable():
+    """Verify col() works with assign() on a memtable based DataFrame."""
+    # Use sqlite backend as dummy
+    backend = ibis.sqlite.connect()
+    session = Session(backend)
+
+    # Create data
+    data = pd.DataFrame({'a': [1, 2, 3]})
+    t = ibis.memtable(data)
+    df = session.DataFrame(t)
+
+    # Use col to create a new column based on existing one
+    deferred_col = session.col('a')
+
+    # Verify it returns an Expression
+    assert isinstance(deferred_col, leanframe.core.expression.Expression)
+
+    # Perform arithmetic (should return another deferred Expression)
+    expr = deferred_col + 1
+    assert isinstance(expr, leanframe.core.expression.Expression)
+
+    # Use in assign
+    df_new = df.assign(b=expr)
+
+    result = df_new.to_pandas()
+    expected = pd.DataFrame({'a': [1, 2, 3], 'b': [2, 3, 4]})
+    pd.testing.assert_frame_equal(result, expected, check_dtype=False)
+
+
+def test_col_arithmetic_chain():
+    """Verify complex arithmetic chains with col()."""
+    backend = ibis.sqlite.connect()
+    session = Session(backend)
+
+    col_a = session.col('a')
+    col_b = session.col('b')
+
+    # (a + b) * 2
+    expr = (col_a + col_b) * 2
+
+    data = pd.DataFrame({'a': [10], 'b': [5]})
+    t = ibis.memtable(data)
+    df = session.DataFrame(t)
+
+    df_new = df.assign(c=expr)
+
+    result = df_new.to_pandas()
+    expected = pd.DataFrame({'a': [10], 'b': [5], 'c': [30]})
+    pd.testing.assert_frame_equal(result, expected, check_dtype=False)