diff --git a/leanframe/core/session.py b/leanframe/core/session.py index 401e28b..836e4b6 100644 --- a/leanframe/core/session.py +++ b/leanframe/core/session.py @@ -45,6 +45,12 @@ def read_sql_table(self, table_name: str): #TODO: will crash if self._backend is None. return leanframe.core.frame.DataFrame(self._backend.table(table_name)) # type: ignore + def read_ibis(self, table_expression: ibis_types.Table): + """Create a DataFrame from an Ibis table expression.""" + import leanframe.core.frame + + return leanframe.core.frame.DataFrame(table_expression) + def DataFrame(self, data: ibis_types.Table | pandas.DataFrame): """Construct a DataFrame.""" import leanframe.core.frame diff --git a/tests/unit/session/test_read_ibis.py b/tests/unit/session/test_read_ibis.py new file mode 100644 index 0000000..4454b21 --- /dev/null +++ b/tests/unit/session/test_read_ibis.py @@ -0,0 +1,41 @@ + +import ibis +import pandas as pd +import pytest +from unittest.mock import MagicMock + +from leanframe.core import session + +@pytest.fixture +def mock_backend(): + return MagicMock() + +def test_read_ibis_returns_dataframe(mock_backend): + """Test that read_ibis returns a leanframe DataFrame.""" + s = session.Session(mock_backend) + df = pd.DataFrame({'a': [1, 2, 3]}) + t = ibis.memtable(df) + + lf_df = s.read_ibis(t) + + # Check that it is a leanframe DataFrame + # Note: importing DataFrame inside the function to avoid circular imports if any, + # though it should be fine at module level for type checking + from leanframe.core.frame import DataFrame + assert isinstance(lf_df, DataFrame) + + # Check that the underlying data is correct + assert lf_df.to_ibis().equals(t) + +def test_read_ibis_with_complex_expression(mock_backend): + """Test read_ibis with a more complex ibis expression.""" + s = session.Session(mock_backend) + df = pd.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6]}) + t = ibis.memtable(df) + expr = t.mutate(c=t.a + t.b) + + lf_df = s.read_ibis(expr) + + from leanframe.core.frame import DataFrame + assert isinstance(lf_df, DataFrame) + assert lf_df.to_ibis().equals(expr)