Support ArrowDtype based pandas columns (#19)

* Support ArrowDtype based pandas columns
pydiverse · Apr 20, 2024 · 2421147 · 2421147
1 parent ea971b0
commit 2421147
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 1 deletion.
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/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pydiverse-transform"
-version = "0.1.4"
+version = "0.1.5"
 description = "Pipe based dataframe manipulation library that can also transform data on SQL databases"
 authors = [
   "QuantCo, Inc.",

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_):