feat: implement DataFrame.dtypes and Series.dtype

leanframe/core/dtypes.py

@@ -0,0 +1,36 @@
+# Copyright 2024 Google LLC
+#
+# 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.
+
+"""Helpers to convert between ibis and pandas dtypes."""
+
+from __future__ import annotations
+
+import ibis.expr.datatypes as ibis_dtypes
+import pandas as pd
+
+
+def convert_ibis_to_pandas(
+    ibis_type: ibis_dtypes.DataType,
+) -> pd.ArrowDtype:
+    """
+    Convert an ibis type to a pandas ArrowDtype.
+
+    Args:
+        ibis_type: The ibis type to convert.
+
+    Returns:
+        The corresponding pandas ArrowDtype.
+    """
+    arrow_type = ibis_type.to_pyarrow()
+    return pd.ArrowDtype(arrow_type)

leanframe/core/frame.py

@@ -17,7 +17,9 @@
 from __future__ import annotations
 
 import ibis.expr.types as ibis_types
-import pandas
+import pandas as pd
+
+from leanframe.core.dtypes import convert_ibis_to_pandas
 
 
 class DataFrame:
@@ -31,9 +33,16 @@ def __init__(self, data: ibis_types.Table):
         self._data = data
 
     @property
-    def columns(self) -> pandas.Index:
+    def columns(self) -> pd.Index:
         """The column labels of the DataFrame."""
-        return pandas.Index(self._data.columns, dtype="object")
+        return pd.Index(self._data.columns, dtype="object")
+
+    @property
+    def dtypes(self) -> pd.Series:
+        """Return the dtypes in the DataFrame."""
+        names = self._data.columns
+        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):
         """Get a column.
@@ -49,12 +58,12 @@ def __getitem__(self, key: str):
         # current DataFrame, only. No joins by index key are available.
         return leanframe.core.series.Series(self._data[key])
 
-    def to_pandas(self) -> pandas.DataFrame:
+    def to_pandas(self) -> pd.DataFrame:
         """Convert the DataFrame to a pandas.DataFrame.
 
         Where possible, pandas.ArrowDtype is used to avoid lossy conversions
         from the database types to pandas.
         """
         return self._data.to_pyarrow().to_pandas(
-            types_mapper=lambda type_: pandas.ArrowDtype(type_)
+            types_mapper=lambda type_: pd.ArrowDtype(type_)
         )

leanframe/core/series.py

@@ -16,8 +16,10 @@
 
 from __future__ import annotations
 
-import pandas
 import ibis.expr.types as ibis_types
+import pandas as pd
+
+from leanframe.core.dtypes import convert_ibis_to_pandas
 
 
 class Series:
@@ -30,13 +32,18 @@ class Series:
     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()
 
-    def to_pandas(self) -> pandas.Series:
+    def to_pandas(self) -> pd.Series:
         """Convert to a pandas Series."""
         return self._data.to_pyarrow().to_pandas(
-            types_mapper=lambda type_: pandas.ArrowDtype(type_)
+            types_mapper=lambda type_: pd.ArrowDtype(type_)
         )

pyproject.toml

@@ -14,6 +14,8 @@ dependencies = [
 dev = [
     "ibis-framework[duckdb]>=10.6.0",
     "ipython>=8.37.0",
+    "mypy>=1.17.1",
+    "pandas-stubs>=2.3.0.250703",
     "pytest>=8.4.1",
     "ruff>=0.12.4",
 ]

specs/2025-08-04-dtypes.md

@@ -0,0 +1,65 @@
+# Implementing pandas-compatible dtypes
+
+The goal of this project is to implement a `dtypes` property on the `leanframe.core.frame.DataFrame` class and a `dtype` property on the `leanframe.core.series.Series` class. These properties should be compatible with the `pandas` library's `dtypes` and `dtype` properties.
+
+## Background
+
+The `leanframe` library is a lightweight wrapper around `ibis`. The `ibis` library uses its own type system. The `leanframe` `DataFrame` and `Series` objects are backed by `ibis` expressions. The `dtypes` and `dtype` properties should return `pandas` types, specifically `pandas.ArrowDtype` objects that wrap the underlying Arrow types corresponding to the `ibis` types.
+
+## Acceptance Criteria
+
+- A `DataFrame.dtypes` property exists and returns a `pandas.Series` object.
+  - The index of the series is the column names of the `DataFrame`.
+  - The values of the series are `pandas.ArrowDtype` objects.
+- A `Series.dtype` property exists and returns a `pandas.ArrowDtype` object.
+- The returned `ArrowDtype` objects correctly represent the underlying `ibis` types.
+- The implementation should be covered by unit tests.
+
+## Detailed Steps
+
+### 1. Create a type mapping module
+
+- [x] Create a new module `leanframe/core/dtypes.py` to house the type
+      conversion logic.
+
+### 2. Implement the type conversion function
+
+- [x] In `leanframe/core/dtypes.py`, create a function
+      `convert_ibis_to_pandas(ibis_type: ibis.expr.datatypes.DataType) ->
+      pandas.ArrowDtype`. This function will take an `ibis` type as input and return
+      the corresponding `pandas.ArrowDtype`.
+
+### 3. Implement `Series.dtype`
+
+- [x] In `leanframe/core/series.py`, implement the `dtype` property on the
+      `Series` class. This property should use the `convert_ibis_to_pandas` function
+      to convert the `ibis` type of the series to a `pandas.ArrowDtype`.
+
+### 4. Implement `DataFrame.dtypes`
+
+- [x] In `leanframe/core/frame.py`, implement the `dtypes` property on the
+      `DataFrame` class. This property should iterate over the columns of the `ibis`
+      table, convert each column's `ibis` type to a `pandas.ArrowDtype` using the
+      `convert_ibis_to_pandas` function, and return a `pandas.Series` with the column
+      names as the index and the `ArrowDtype` objects as the values.
+
+### 5. Write unit tests
+
+Create unit tests in `tests/unit/test_frame.py` and `tests/unit/test_series.py`
+to verify the correctness of the `dtypes` and `dtype` properties.
+
+- [ ] **`test_frame.py`**:
+  - [x] Create a test case for `DataFrame.dtypes`.
+    Assert that the `dtypes` property returns a `pandas.Series` with the expected `ArrowDtype` objects.
+  - [ ] Create a `test_dataframe_dtypes_nested` that includes complex types like `array` and `struct`.
+- [ ] **`test_series.py`**:
+  - [x] Create parameterized tests for `Series.dtype` that covers some `ibis` types.
+        For each `ibis` type, assert that the `dtype` property returns the expected `pandas.ArrowDtype`.
+  - [ ] Include complex "nested" types in the parameterized tests.
+
+## Verification
+
+- All new and existing unit tests should pass.
+- The `uv run mypy leanframe tests` static type checker should pass.
+- The `uv run ruff check` linter should pass.
+- Only add git commits. Do not change git history.

tests/unit/test_frame.py

@@ -14,17 +14,77 @@
 
 from __future__ import annotations
 
-import pandas
+import pandas as pd
+import pandas.testing as tm
+import pyarrow as pa
 
 import leanframe
 import leanframe.core.series
 
 
+def test_dataframe_dtypes(session: leanframe.Session):
+    df_pd = pd.DataFrame(
+        {
+            "col1": [1, 2, 3],
+            "col2": ["a", "b", "c"],
+            "col3": [1.1, 2.2, 3.3],
+        }
+    ).astype(
+        {
+            "col1": pd.ArrowDtype(pa.int64()),
+            "col2": pd.ArrowDtype(pa.string()),
+            "col3": pd.ArrowDtype(pa.float64()),
+        }
+    )
+    df_lf = session.DataFrame(df_pd)
+    result = df_lf.dtypes
+    expected = pd.Series(
+        [
+            pd.ArrowDtype(pa.int64()),
+            pd.ArrowDtype(pa.string()),
+            pd.ArrowDtype(pa.float64()),
+        ],
+        index=["col1", "col2", "col3"],
+        name="dtypes",
+    )
+    tm.assert_series_equal(result, expected)
+
+
+def test_dataframe_dtypes_complex(session: leanframe.Session):
+    pa_table = pa.Table.from_pydict(
+        {
+            "array_col": [[1, 2], [3, 4]],
+            "struct_col": [{"a": 1, "b": "c"}, {"a": 2, "b": "d"}],
+        },
+        schema=pa.schema(
+            [
+                pa.field("array_col", pa.list_(pa.int64())),
+                pa.field(
+                    "struct_col",
+                    pa.struct([("a", pa.int64()), ("b", pa.string())]),
+                ),
+            ]
+        ),
+    )
+    df_pd = pa_table.to_pandas(types_mapper=pd.ArrowDtype)
+    df_lf = session.DataFrame(df_pd)
+    result = df_lf.dtypes
+    expected = pd.Series(
+        [
+            pd.ArrowDtype(pa.list_(pa.int64())),
+            pd.ArrowDtype(pa.struct([("a", pa.int64()), ("b", pa.string())])),
+        ],
+        index=["array_col", "struct_col"],
+        name="dtypes",
+    )
+    tm.assert_series_equal(result, expected)
+
+
 def test_dataframe_getitem_with_column(session: leanframe.Session):
     """Read a table with simple scalar values."""
 
     df_lf = session.DataFrame(
-        pandas.DataFrame(
+        pd.DataFrame(
             {
                 "col1": [1, 2, 3],
                 "col2": ["a", "b", "c"],

tests/unit/test_series.py

@@ -14,31 +14,71 @@
 
 from __future__ import annotations
 
-import pandas
+import pandas as pd
 import pandas.testing
-import pyarrow
+import pyarrow as pa
 import pytest
 
 import leanframe
 
 
+@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(
-            pandas.Series([1, 2, 3], dtype=pandas.ArrowDtype(pyarrow.int64())),
+            pd.Series([1, 2, 3], dtype=pd.ArrowDtype(pa.int64())),
             id="int64",
         ),
         pytest.param(
-            pandas.Series(
-                [1.0, float("nan"), 3.0], dtype=pandas.ArrowDtype(pyarrow.float64())
-            ),
+            pd.Series([1.0, float("nan"), 3.0], dtype=pd.ArrowDtype(pa.float64())),
             id="float64",
         ),
     ),
 )
-def test_to_pandas(session: leanframe.Session, series_pd: pandas.Series):
-    df_pd = pandas.DataFrame(
+def test_to_pandas(session: leanframe.Session, series_pd: pd.Series):
+    df_pd = pd.DataFrame(
         {
             "my_col": series_pd,
         }
@@ -48,4 +88,4 @@ def test_to_pandas(session: leanframe.Session, series_pd: pandas.Series):
     result = df_lf["my_col"].to_pandas()
 
     # TODO(tswast): Allow input dtype != output dtype with an "expected_dtype" parameter.
-    pandas.testing.assert_series_equal(result, series_pd, check_names=False)
+    pd.testing.assert_series_equal(result, series_pd, check_names=False)