pydiverse-common 0.3.8__tar.gz → 0.3.9__tar.gz

{pydiverse_common-0.3.8 → pydiverse_common-0.3.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydiverse-common
-Version: 0.3.8
+Version: 0.3.9
 Summary: Common functionality shared between pydiverse libraries
 Author: QuantCo, Inc.
 Author-email: Martin Trautmann <windiana@users.sf.net>, Finn Rudolph <finn.rudolph@t-online.de>

{pydiverse_common-0.3.8 → pydiverse_common-0.3.9}/docs/source/changelog.md

@@ -1,5 +1,8 @@
 # Changelog
 
+## 0.3.9 (2025-08-21)
+- fix enum pyarrow dtype
+
 ## 0.3.8 (2025-08-19)
 - fixed util.hashing.hash_polars_dataframe for simple dataframe
 

{pydiverse_common-0.3.8 → pydiverse_common-0.3.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "pydiverse-common"
-version = "0.3.8"
+version = "0.3.9"
 description = "Common functionality shared between pydiverse libraries"
 authors = [
   { name = "QuantCo, Inc." },

{pydiverse_common-0.3.8 → pydiverse_common-0.3.9}/src/pydiverse/common/dtypes.py

@@ -201,6 +201,11 @@ class Dtype:
             return NullType()
         if pa.types.is_list(arrow_type):
             return List(Dtype.from_arrow(arrow_type.value_type))
+        if pa.types.is_dictionary(arrow_type):
+            raise RuntimeError(
+                "Most likely this is an Enum type. But metadata about categories is "
+                "only in the pyarrow field and not in the pyarrow dtype"
+            )
         raise TypeError
 
     @staticmethod
@@ -268,7 +273,7 @@ class Dtype:
         if backend == PandasBackend.NUMPY:
             return self.to_pandas_nullable(backend)
         if backend == PandasBackend.ARROW:
-            if self == String():
+            if self == String() or isinstance(self, Enum):
                 return pd.StringDtype(storage="pyarrow")
             return pd.ArrowDtype(self.to_arrow())
 
@@ -355,6 +360,12 @@ class Dtype:
             NullType(): pa.null(),
         }[self]
 
+    def to_arrow_field(self, name: str, nullable: bool = True):
+        """Convert this Dtype to a PyArrow Field."""
+        import pyarrow as pa
+
+        return pa.field(name, self.to_arrow(), nullable=nullable)
+
     def to_polars(self: "Dtype"):
         """Convert this Dtype to a Polars type."""
         import polars as pl
@@ -506,6 +517,19 @@ class Enum(String):
     def to_arrow(self):
         import pyarrow as pa
 
-        # There is also pa.dictionary(), which seems to be kind of similar to an enum.
-        # Maybe it is better to convert to this.
+        # enum categories can only be maintained in pyarrow field (see to_arrow_field)
         return pa.string()
+
+    def to_arrow_field(self, name: str, nullable: bool = True):
+        """Convert this Dtype to a PyArrow Field."""
+        import pyarrow as pa
+
+        # try to mimic what polars does
+        return pa.field(
+            name,
+            pa.dictionary(pa.uint32(), pa.large_string()),
+            nullable=nullable,
+            metadata={
+                "_PL_ENUM_VALUES": "".join([f"{len(c)};{c}" for c in self.categories])
+            },
+        )

{pydiverse_common-0.3.8 → pydiverse_common-0.3.9}/tests/dtypes/test_dtype_pyarrow.py

@@ -95,6 +95,20 @@ def test_dtype_to_pyarrow():
     assert_conversion(Datetime(), pa.timestamp("us"))
 
 
+@pytest.mark.skipif(pa is None, reason="requires pyarrow")
+def test_dtype_to_pyarrow_enum():
+    import polars as pl
+
+    def assert_conversion(type_: Dtype, expected_dtype):
+        df = pl.DataFrame(dict(x=["a"]), schema=dict(x=expected_dtype))
+        expected = df.to_arrow().schema
+        actual = pa.schema([type_.to_arrow_field("x", nullable=True)])
+        assert actual == expected
+        assert actual.field(0).metadata == expected.field(0).metadata
+
+    assert_conversion(Enum("a", "b;c"), pl.Enum(["a", "b;c"]))
+
+
 @pytest.mark.skipif(pa is None, reason="requires pyarrow")
 @pytest.mark.parametrize(
     "type_",