Support ArrowDtype based pandas columns

pydiverse · Apr 20, 2024 · 637663c · 637663c
1 parent ea971b0
commit 637663c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/docs/source/changelog.md b/docs/source/changelog.md
@@ -1,5 +1,8 @@
 # Changelog
 
+## 0.1.5 (2024-04-20)
+- support ArrowDtype based columns in pandas dataframes
+
 ## 0.1.4 (2024-04-20)
 - better support for apache arrow backed pandas dataframes
 - fix handling of boolean literals

diff --git a/src/pydiverse/transform/eager/pandas_table.py b/src/pydiverse/transform/eager/pandas_table.py
@@ -68,6 +68,9 @@ def __init__(self, name: str, df: pd.DataFrame):
         super().__init__(name=name, columns=columns)
 
     def _dtype_from_pd(self, type_: Any) -> dtypes.DType:
+        if hasattr(type_, "numpy_dtype"):
+            # support for ArrowDtype based pandas columns
+            type_ = type_.numpy_dtype
         if pd.api.types.is_integer_dtype(type_):
             return dtypes.Int()
         if pd.api.types.is_float_dtype(type_):