Merge pull request #263 from data-engineering-collective/schema_compat_str_dtype

Author: fjetter

CHANGES.rst

@@ -2,9 +2,12 @@
 Changelog
 =========
 
-Next release
-============
+Plateau 4.6.0 (2025-08-12)
+==========================
 
+* Schema normalization for pandas 3.x `str` dtype. String fields are considered
+  compatible if they are using the same NA value but pandas storage backend is
+  being ignored.
 * Support for pyarrow 21.0.0
 * Drop support for pyarrow 15.0.2, 16.1.0 and 17.0.0
 

plateau/core/_compat.py

@@ -1,14 +1,22 @@
 import pandas as pd
+import pyarrow as pa
 import simplejson
 from packaging.version import parse as parse_version
+from pandas.errors import OptionError
 
 PANDAS_3 = parse_version(pd.__version__).major >= 3
 
+ARROW_GE_20 = parse_version(pa.__version__).major >= 20
 
-def pandas_infer_string():
-    return (
-        pd.get_option("future.infer_string") or parse_version(pd.__version__).major >= 3
-    )
+
+def pandas_infer_string() -> bool:
+    if parse_version(pd.__version__).major >= 3:
+        # In pandas 3, infer_string is always True
+        return True
+    try:
+        return pd.get_option("future.infer_string")
+    except OptionError:
+        return False
 
 
 def load_json(buf, **kwargs):

plateau/core/common_metadata.py

@@ -78,6 +78,33 @@ def _schema_compat(self):
                     if index_level_ix >= 0:
                         schema = schema.remove(index_level_ix)
 
+            for cmd in pandas_metadata["columns"]:
+                name = cmd.get("name")
+                if name is None:
+                    continue
+
+                field_name = cmd["field_name"]
+                field_idx = schema.get_field_index(field_name)
+                if field_idx < 0:
+                    continue
+                field = schema[field_idx]
+                if (
+                    pa.types.is_string(field.type)
+                    and cmd["pandas_type"] == "unicode"
+                    and cmd["numpy_type"] == "object"
+                ):
+                    schema = schema.remove(field_idx)
+                    new_field = pa.field(
+                        field.name,
+                        pa.large_string(),
+                        field.nullable,
+                        field.metadata,
+                    )
+                    schema = schema.insert(field_idx, new_field)
+                    cmd["pandas_type"] = "object"
+                    cmd["numpy_type"] = "str"
+                    cmd["metadata"] = None
+
             schema = schema.remove_metadata()
             md = {b"pandas": _dict_to_binary(pandas_metadata)}
             schema = schema.with_metadata(md)
@@ -319,8 +346,32 @@ def normalize_type(
         )
         return pa.list_(t_pa2), f"list[{t_pd2}]", "object", None
     elif pa.types.is_dictionary(t_pa):
-        # downcast to dictionary content, `t_pd` is useless in that case
-        return normalize_type(t_pa.value_type, t_np, t_np, None)
+        return normalize_type(t_pa.value_type, t_pd, t_np, None)
+    elif pa.types.is_string(t_pa) or pa.types.is_large_string(t_pa):
+        # Pyarrow only supports reading back
+        #
+        # pyarrow + np.nan
+        # pa.large_string(), "object", "str", None
+        # or
+        # python + pd.NA
+        # pa.string(), "unicode", "string", None
+        #
+        # unintuitively, the numpy type identifier `t_np` corresponds
+        # to the pandas dtypes `str` and `string`
+
+        # pandas also supports mixed types but those are rare and must be
+        # constructed explicitly
+        if t_pd == "categorical":
+            # We loose the information of the nullable type since the t_np type
+            # is set to the dtype of the codes but not the categories.
+            return pa.large_string(), "object", "str", None
+        elif t_np == "str":
+            return pa.large_string(), "object", "str", None
+        elif t_np == "string":
+            return pa.string(), "unicode", "string", None
+        else:
+            # This should be the ordinary object dtype
+            return t_pa, t_pd, t_np, metadata
     else:
         return t_pa, t_pd, t_np, metadata
 

plateau/core/index.py

@@ -11,6 +11,7 @@
 
 import plateau.core._time
 from plateau.core import naming
+from plateau.core._compat import ARROW_GE_20
 from plateau.core._mixins import CopyMixin
 from plateau.core.common_metadata import normalize_type
 from plateau.core.docs import default_docs
@@ -138,7 +139,7 @@ def observed_values(
     ) -> np.ndarray:
         """Return an array of all observed values."""
         keys = np.array(list(self.index_dct.keys()))
-        labeled_array = pa.array(keys, type=self.dtype)
+        labeled_array = _safe_paarray(keys, self.dtype)
 
         _coerce = {"coerce_temporal_nanoseconds": coerce_temporal_nanoseconds}
         return np.array(
@@ -918,10 +919,17 @@ def _index_dct_to_table(index_dct: IndexDictType, column: str, dtype: pa.DataTyp
         # the np.array dtype will be double which arrow cannot convert to the target type, so use an empty list instead
         labeled_array = pa.array([], type=dtype)
     else:
-        labeled_array = pa.array(keys, type=dtype)
+        labeled_array = _safe_paarray(keys, dtype)
 
     partition_array = pa.array(list(index_dct.values()), type=pa.list_(pa.string()))
 
     return pa.Table.from_arrays(
         [labeled_array, partition_array], names=[column, _PARTITION_COLUMN_NAME]
     )
+
+
+def _safe_paarray(arr: np.ndarray, dtype: pa.DataType) -> pa.Array:
+    if dtype is not None and pa.types.is_large_string(dtype) and not ARROW_GE_20:
+        return pa.array(iter(arr), type=dtype)
+    else:
+        return pa.array(arr, type=dtype)

plateau/io/testing/read.py

@@ -36,7 +36,7 @@
 import pytest
 from minimalkv import get_store_from_url
 
-from plateau.core._compat import pandas_infer_string
+from plateau.core._compat import PANDAS_3, pandas_infer_string
 from plateau.io.eager import store_dataframes_as_dataset
 from plateau.io.iter import store_dataframes_as_dataset__iter
 from plateau.io_components.metapartition import SINGLE_TABLE, MetaPartition
@@ -648,8 +648,14 @@ def test_binary_column_metadata(store_factory, bound_load_dataframes):
     assert set(df.columns.map(type)) == {str}
 
 
-def test_extensiondtype_roundtrip(store_factory, bound_load_dataframes):
-    df = pd.DataFrame({"str": pd.Series(["a", "b"], dtype="string")})
+def test_string_type_roundtrip(store_factory, bound_load_dataframes):
+    # Note: we're not actually roundtripping the string type since the loading
+    # type depends on the pandas version. Keeping the loading type aligned with
+    # what is typically initialized by pandas by default is likely the best
+    # option
+    df = pd.DataFrame(
+        {"str": pd.Series(["a", "b"], dtype="str" if PANDAS_3 else "string")}
+    )
 
     store_dataframes_as_dataset(
         dfs=[df], store=store_factory, dataset_uuid="dataset_uuid"

plateau/io/testing/update.py

@@ -6,13 +6,19 @@
 
 import numpy as np
 import pandas as pd
+import pandas.testing as pdt
 import pytest
 
 from plateau.api.dataset import read_dataset_as_ddf
+from plateau.core._compat import PANDAS_3
 from plateau.core.dataset import DatasetMetadata
 from plateau.core.naming import DEFAULT_METADATA_VERSION
 from plateau.core.testing import TIME_TO_FREEZE_ISO
-from plateau.io.eager import read_dataset_as_metapartitions, store_dataframes_as_dataset
+from plateau.io.eager import (
+    read_dataset_as_metapartitions,
+    read_table,
+    store_dataframes_as_dataset,
+)
 from plateau.io.iter import read_dataset_as_dataframes__iterator
 
 
@@ -665,3 +671,198 @@ def test_update_of_dataset_with_non_default_table_name(
     )
     df_expected = pd.concat([df_create, df_update]).reset_index(drop=True)
     pd.testing.assert_frame_equal(df_read, df_expected)
+
+
+def _dtype_from_storage_nan_value(storage_backend, na_value):
+    if PANDAS_3:
+        dtype = pd.StringDtype(storage=storage_backend, na_value=na_value)
+    else:
+        if storage_backend == "pyarrow" and na_value is pd.NA:
+            dtype = "string[pyarrow]"
+        elif storage_backend == "pyarrow" and na_value is np.nan:
+            dtype = "string[pyarrow_numpy]"
+        elif storage_backend == "python" and na_value is np.nan:
+            return None
+        elif storage_backend == "python" and na_value is pd.NA:
+            dtype = "string"
+        else:
+            raise ValueError(f"Unsupported storage backend: {storage_backend}")
+    return dtype
+
+
+@pytest.mark.parametrize("storage_backend", ["pyarrow", "python"])
+@pytest.mark.parametrize("na_value", [np.nan, pd.NA])
+def test_update_after_empty_partition_string_dtypes(
+    store_factory, bound_update_dataset, storage_backend, na_value, backend_identifier
+):
+    import pandas as pd
+
+    with pd.option_context("future.infer_string", True):
+        other_nan_value = {np.nan, pd.NA}
+        other_nan_value.remove(na_value)
+        other_nan_value = other_nan_value.pop()
+        dtype = _dtype_from_storage_nan_value(storage_backend, na_value)
+        if dtype is None:
+            pytest.skip()
+        df = pd.DataFrame({"str": pd.Series(["a", "b", None], dtype=dtype)})
+
+        dataset_uuid = "dataset_uuid"
+        bound_update_dataset(
+            [df.iloc[0:0]],  # empty partition
+            store=store_factory,
+            dataset_uuid=dataset_uuid,
+        )
+        # Schema verification should not fail
+        bound_update_dataset(
+            [df],
+            store=store_factory,
+            dataset_uuid=dataset_uuid,
+        )
+        if na_value is pd.NA:
+            expected_dtype = _dtype_from_storage_nan_value("python", pd.NA)
+        else:
+            expected_dtype = _dtype_from_storage_nan_value("pyarrow", np.nan)
+        # We have to cast to the expected dtype since pyarrow is only reading
+        # the above two data types in. They are ignoring the written storage
+        # backend and are defaulting to python for pd.NA and to pyarrow for
+        # np.nan
+        df["str"] = df["str"].astype(expected_dtype)
+
+        pdt.assert_frame_equal(read_table(dataset_uuid, store_factory()), df)
+        if backend_identifier == "dask.dataframe":
+            # FIXME: dask.dataframe triggers the schema validation error but somehow
+            # the exception is not properly forwarded and the test always fails
+            return
+        for storage in ["pyarrow", "python"]:
+            df = pd.DataFrame(
+                {
+                    "str": pd.Series(
+                        ["c", "d"],
+                        dtype=_dtype_from_storage_nan_value(storage, other_nan_value),
+                    )
+                }
+            )
+            # Should be a ValueError but dask sometimes raises a different exception
+            # type
+            with pytest.raises(ValueError, match="Schemas.*are not compatible.*"):
+                bound_update_dataset(
+                    [df],
+                    store=store_factory,
+                    dataset_uuid=dataset_uuid,
+                )
+
+
+@pytest.mark.parametrize("storage_backend", ["pyarrow", "python"])
+@pytest.mark.parametrize("na_value", [np.nan, pd.NA])
+def test_update_after_empty_partition_string_dtypes_categoricals(
+    store_factory, bound_update_dataset, storage_backend, na_value
+):
+    import pandas as pd
+
+    with pd.option_context("future.infer_string", True):
+        other_nan_value = {np.nan, pd.NA}
+        other_nan_value.remove(na_value)
+        other_nan_value = other_nan_value.pop()
+        dtype = _dtype_from_storage_nan_value(storage_backend, na_value)
+        if dtype is None:
+            pytest.skip()
+        df = pd.DataFrame(
+            {"str": pd.Series(["a", "b", None], dtype=dtype).astype("category")}
+        )
+
+        dataset_uuid = "dataset_uuid"
+        bound_update_dataset(
+            [df.iloc[0:0]],  # empty partition
+            store=store_factory,
+            dataset_uuid=dataset_uuid,
+        )
+        # Schema verification should not fail
+        bound_update_dataset(
+            [df],
+            store=store_factory,
+            dataset_uuid=dataset_uuid,
+        )
+        expected_dtype = _dtype_from_storage_nan_value("pyarrow", np.nan)
+        # We have to cast to the expected dtype since pyarrow is only reading
+        # categoricals with the pyarrow_numpy data type.
+        df["str"] = df["str"].astype(expected_dtype)
+
+        pdt.assert_frame_equal(read_table(dataset_uuid, store_factory()), df)
+        for storage in ["pyarrow", "python"]:
+            df = pd.DataFrame(
+                {
+                    "str": pd.Series(
+                        ["c", "d"],
+                        dtype=_dtype_from_storage_nan_value(storage, other_nan_value),
+                    ).astype("category")
+                }
+            )
+            bound_update_dataset(
+                [df],
+                store=store_factory,
+                dataset_uuid=dataset_uuid,
+            )
+    after_update = read_table(dataset_uuid, store_factory())
+
+    if not PANDAS_3:
+        expected_dtype = "object"
+
+    expected_after_update = pd.DataFrame(
+        {"str": pd.Series(["a", "b", None, "c", "d", "c", "d"], dtype=expected_dtype)}
+    )
+    pdt.assert_frame_equal(after_update, expected_after_update)
+
+    # Storage of categorical dtypes will only happen with np.nan If we try the other na_value we'll get a validation error
+
+    for storage in ["pyarrow", "python"]:
+        df = pd.DataFrame(
+            {
+                "str": pd.Series(
+                    ["e", "f", None],
+                    dtype=_dtype_from_storage_nan_value(storage, pd.NA),
+                )
+            }
+        )
+        with pytest.raises(ValueError, match="Schemas.*are not compatible.*"):
+            bound_update_dataset(
+                [df],
+                store=store_factory,
+                dataset_uuid=dataset_uuid,
+            )
+
+    # With np.nan works fine?
+    skipped = False
+    for storage in ["pyarrow", "python"]:
+        dtype = _dtype_from_storage_nan_value(storage, np.nan)
+        if dtype is None:
+            skipped = True
+            continue
+        df = pd.DataFrame(
+            {
+                "str": pd.Series(
+                    ["e", "f", None],
+                    dtype=dtype,
+                )
+            }
+        )
+        bound_update_dataset(
+            [df],
+            store=store_factory,
+            dataset_uuid=dataset_uuid,
+        )
+
+    after_update_as_cats = read_table(
+        dataset_uuid, store_factory(), categoricals=["str"]
+    )
+    values = ["a", "b", None, "c", "d", "c", "d", "e", "f", None, "e", "f", None]
+    if skipped:
+        values = values[:-3]
+    expected = pd.DataFrame(
+        {
+            "str": pd.Series(
+                values,
+                dtype=expected_dtype,
+            ).astype("category")
+        }
+    )
+    pdt.assert_frame_equal(after_update_as_cats, expected)