ENH: Implement arrow string option for various I/O methods

pandas/_config/__init__.py

@@ -38,3 +38,8 @@ def using_copy_on_write() -> bool:
 def using_nullable_dtypes() -> bool:
     _mode_options = _global_config["mode"]
     return _mode_options["nullable_dtypes"]
+
+
+def using_pyarrow_string_dtype() -> bool:
+    _mode_options = _global_config["future"]
+    return _mode_options["infer_string"]

pandas/_libs/lib.pyx

@@ -1299,6 +1299,7 @@ cdef class Seen:
         bint datetimetz_      # seen_datetimetz
         bint period_          # seen_period
         bint interval_        # seen_interval
+        bint str_             # seen_str
 
     def __cinit__(self, bint coerce_numeric=False):
         """
@@ -1325,6 +1326,7 @@ cdef class Seen:
         self.datetimetz_ = False
         self.period_ = False
         self.interval_ = False
+        self.str_ = False
         self.coerce_numeric = coerce_numeric
 
     cdef bint check_uint64_conflict(self) except -1:
@@ -2615,6 +2617,13 @@ def maybe_convert_objects(ndarray[object] objects,
             else:
                 seen.object_ = True
                 break
+        elif isinstance(val, str):
+            if convert_non_numeric:
+                seen.str_ = True
+                break
+            else:
+                seen.object_ = True
+                break
         else:
             seen.object_ = True
             break
@@ -2669,6 +2678,20 @@ def maybe_convert_objects(ndarray[object] objects,
             return pi._data
         seen.object_ = True
 
+    elif seen.str_:
+        if is_string_array(objects):
+            from pandas._config import get_option
+            opt = get_option("future.infer_string")
+            if opt is True:
+                import pyarrow as pa
+
+                from pandas.core.dtypes.dtypes import ArrowDtype
+
+                obj = pa.array(objects)
+                dtype = ArrowDtype(obj.type)
+                return dtype.construct_array_type()(obj)
+
+        seen.object_ = True
     elif seen.interval_:
         if is_interval_array(objects):
             from pandas import IntervalIndex

pandas/core/config_init.py

@@ -889,3 +889,14 @@ def register_converter_cb(key) -> None:
         styler_environment,
         validator=is_instance_factory([type(None), str]),
     )
+
+
+with cf.config_prefix("future"):
+    cf.register_option(
+        "infer_string",
+        False,
+        "Whether to infer sequence of str objects as pyarrow string "
+        "dtype, which will be the default in pandas 3.0 "
+        "(at which point this option will be deprecated).",
+        validator=is_one_of_factory([True, False]),
+    )

pandas/core/dtypes/cast.py

@@ -18,6 +18,8 @@
 
 import numpy as np
 
+from pandas._config import get_option
+
 from pandas._libs import lib
 from pandas._libs.missing import (
     NA,
@@ -796,6 +798,12 @@ def infer_dtype_from_scalar(val) -> tuple[DtypeObj, Any]:
         # coming out as np.str_!
 
         dtype = _dtype_obj
+        opt = get_option("future.infer_string")
+        if opt is True:
+            import pyarrow as pa
+
+            pa_dtype = pa.string()
+            dtype = ArrowDtype(pa_dtype)
 
     elif isinstance(val, (np.datetime64, dt.datetime)):
         try:

pandas/io/_util.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing import Callable
+
 from pandas.compat._optional import import_optional_dependency
 
 import pandas as pd
@@ -21,3 +23,9 @@ def _arrow_dtype_mapping() -> dict:
         pa.float32(): pd.Float32Dtype(),
         pa.float64(): pd.Float64Dtype(),
     }
+
+
+def arrow_string_types_mapper() -> Callable:
+    pa = import_optional_dependency("pyarrow")
+
+    return {pa.string(): pd.ArrowDtype(pa.string())}.get

pandas/io/feather_format.py

@@ -6,6 +6,8 @@
     Any,
 )
 
+from pandas._config import using_pyarrow_string_dtype
+
 from pandas._libs import lib
 from pandas.compat._optional import import_optional_dependency
 from pandas.util._decorators import doc
@@ -15,6 +17,7 @@
 from pandas.core.api import DataFrame
 from pandas.core.shared_docs import _shared_docs
 
+from pandas.io._util import arrow_string_types_mapper
 from pandas.io.common import get_handle
 
 if TYPE_CHECKING:
@@ -119,7 +122,7 @@ def read_feather(
     with get_handle(
         path, "rb", storage_options=storage_options, is_text=False
     ) as handles:
-        if dtype_backend is lib.no_default:
+        if dtype_backend is lib.no_default and not using_pyarrow_string_dtype():
             return feather.read_feather(
                 handles.handle, columns=columns, use_threads=bool(use_threads)
             )
@@ -135,3 +138,8 @@ def read_feather(
 
         elif dtype_backend == "pyarrow":
             return pa_table.to_pandas(types_mapper=pd.ArrowDtype)
+
+        elif using_pyarrow_string_dtype():
+            return pa_table.to_pandas(types_mapper=arrow_string_types_mapper())
+        else:
+            raise NotImplementedError

pandas/io/orc.py

@@ -9,6 +9,8 @@
     Literal,
 )
 
+from pandas._config import using_pyarrow_string_dtype
+
 from pandas._libs import lib
 from pandas.compat import pa_version_under8p0
 from pandas.compat._optional import import_optional_dependency
@@ -24,6 +26,7 @@
 import pandas as pd
 from pandas.core.indexes.api import default_index
 
+from pandas.io._util import arrow_string_types_mapper
 from pandas.io.common import (
     get_handle,
     is_fsspec_url,
@@ -132,7 +135,12 @@ def read_orc(
             df = pa_table.to_pandas(types_mapper=mapping.get)
         return df
     else:
-        return pa_table.to_pandas()
+        print("Ts")
+        if using_pyarrow_string_dtype():
+            types_mapper = arrow_string_types_mapper()
+        else:
+            types_mapper = None
+        return pa_table.to_pandas(types_mapper=types_mapper)
 
 
 def to_orc(

pandas/io/parquet.py

@@ -12,6 +12,8 @@
 import warnings
 from warnings import catch_warnings
 
+from pandas._config import using_pyarrow_string_dtype
+
 from pandas._libs import lib
 from pandas.compat._optional import import_optional_dependency
 from pandas.errors import AbstractMethodError
@@ -26,6 +28,7 @@
 )
 from pandas.core.shared_docs import _shared_docs
 
+from pandas.io._util import arrow_string_types_mapper
 from pandas.io.common import (
     IOHandles,
     get_handle,
@@ -252,6 +255,8 @@ def read(
             to_pandas_kwargs["types_mapper"] = mapping.get
         elif dtype_backend == "pyarrow":
             to_pandas_kwargs["types_mapper"] = pd.ArrowDtype  # type: ignore[assignment]  # noqa: E501
+        elif using_pyarrow_string_dtype():
+            to_pandas_kwargs["types_mapper"] = arrow_string_types_mapper()
 
         manager = get_option("mode.data_manager")
         if manager == "array":

pandas/io/parsers/arrow_parser_wrapper.py

@@ -2,6 +2,8 @@
 
 from typing import TYPE_CHECKING
 
+from pandas._config import using_pyarrow_string_dtype
+
 from pandas._libs import lib
 from pandas.compat._optional import import_optional_dependency
 
@@ -10,7 +12,10 @@
 import pandas as pd
 from pandas import DataFrame
 
-from pandas.io._util import _arrow_dtype_mapping
+from pandas.io._util import (
+    _arrow_dtype_mapping,
+    arrow_string_types_mapper,
+)
 from pandas.io.parsers.base_parser import ParserBase
 
 if TYPE_CHECKING:
@@ -215,6 +220,8 @@ def read(self) -> DataFrame:
             dtype_mapping = _arrow_dtype_mapping()
             dtype_mapping[pa.null()] = pd.Int64Dtype()
             frame = table.to_pandas(types_mapper=dtype_mapping.get)
+        elif using_pyarrow_string_dtype():
+            frame = table.to_pandas(types_mapper=arrow_string_types_mapper())
         else:
             frame = table.to_pandas()
         return self._finalize_pandas_output(frame)

pandas/io/pytables.py

@@ -30,6 +30,7 @@
 from pandas._config import (
     config,
     get_option,
+    using_pyarrow_string_dtype,
 )
 
 from pandas._libs import (
@@ -66,6 +67,7 @@
 )
 from pandas.core.dtypes.missing import array_equivalent
 
+import pandas as pd
 from pandas import (
     DataFrame,
     DatetimeIndex,
@@ -3219,7 +3221,12 @@ def read(
         self.validate_read(columns, where)
         index = self.read_index("index", start=start, stop=stop)
         values = self.read_array("values", start=start, stop=stop)
-        return Series(values, index=index, name=self.name, copy=False)
+        result = Series(values, index=index, name=self.name, copy=False)
+        if result.dtype.kind == "O" and using_pyarrow_string_dtype():
+            import pyarrow as pa
+
+            result = result.astype(pd.ArrowDtype(pa.string()))
+        return result
 
     # error: Signature of "write" incompatible with supertype "Fixed"
     def write(self, obj, **kwargs) -> None:  # type: ignore[override]
@@ -3287,6 +3294,10 @@ def read(
 
             columns = items[items.get_indexer(blk_items)]
             df = DataFrame(values.T, columns=columns, index=axes[1], copy=False)
+            if values.dtype.kind == "O" and using_pyarrow_string_dtype():
+                import pyarrow as pa
+
+                df = df.astype(pd.ArrowDtype(pa.string()))
             dfs.append(df)
 
         if len(dfs) > 0:
@@ -4669,6 +4680,10 @@ def read(
                 # Categorical
                 df = DataFrame._from_arrays([values], columns=cols_, index=index_)
             assert (df.dtypes == values.dtype).all(), (df.dtypes, values.dtype)
+            if values.dtype.kind == "O" and using_pyarrow_string_dtype():
+                import pyarrow as pa
+
+                df = df.astype(pd.ArrowDtype(pa.string()))
             frames.append(df)
 
         if len(frames) == 1:

pandas/tests/frame/test_constructors.py

@@ -2689,6 +2689,41 @@ def test_construct_with_strings_and_none(self):
         expected = DataFrame({"a": ["1", "2", None]}, dtype="str")
         tm.assert_frame_equal(df, expected)
 
+    def test_frame_string_inference(self):
+        # GH#54430
+        pa = pytest.importorskip("pyarrow")
+        dtype = pd.ArrowDtype(pa.string())
+        expected = DataFrame(
+            {"a": ["a", "b"]}, dtype=dtype, columns=Index(["a"], dtype=dtype)
+        )
+        with pd.option_context("future.infer_string", True):
+            df = DataFrame({"a": ["a", "b"]})
+        tm.assert_frame_equal(df, expected)
+
+        expected = DataFrame(
+            {"a": ["a", "b"]},
+            dtype=dtype,
+            columns=Index(["a"], dtype=dtype),
+            index=Index(["x", "y"], dtype=dtype),
+        )
+        with pd.option_context("future.infer_string", True):
+            df = DataFrame({"a": ["a", "b"]}, index=["x", "y"])
+        tm.assert_frame_equal(df, expected)
+
+        expected = DataFrame(
+            {"a": ["a", 1]}, dtype="object", columns=Index(["a"], dtype=dtype)
+        )
+        with pd.option_context("future.infer_string", True):
+            df = DataFrame({"a": ["a", 1]})
+        tm.assert_frame_equal(df, expected)
+
+        expected = DataFrame(
+            {"a": ["a", "b"]}, dtype="object", columns=Index(["a"], dtype=dtype)
+        )
+        with pd.option_context("future.infer_string", True):
+            df = DataFrame({"a": ["a", "b"]}, dtype="object")
+        tm.assert_frame_equal(df, expected)
+
 
 class TestDataFrameConstructorIndexInference:
     def test_frame_from_dict_of_series_overlapping_monthly_period_indexes(self):

pandas/tests/indexes/base_class/test_constructors.py

@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
 
+import pandas as pd
 from pandas import (
     Index,
     MultiIndex,
@@ -42,3 +43,17 @@ def test_construct_empty_tuples(self, tuple_list):
         expected = MultiIndex.from_tuples(tuple_list)
 
         tm.assert_index_equal(result, expected)
+
+    def test_index_string_inference(self):
+        # GH#54430
+        pa = pytest.importorskip("pyarrow")
+        dtype = pd.ArrowDtype(pa.string())
+        expected = Index(["a", "b"], dtype=dtype)
+        with pd.option_context("future.infer_string", True):
+            ser = Index(["a", "b"])
+        tm.assert_index_equal(ser, expected)
+
+        expected = Index(["a", 1], dtype="object")
+        with pd.option_context("future.infer_string", True):
+            ser = Index(["a", 1])
+        tm.assert_index_equal(ser, expected)

pandas/tests/io/parser/dtypes/test_dtypes_basic.py

@@ -538,3 +538,22 @@ def test_ea_int_avoid_overflow(all_parsers):
         }
     )
     tm.assert_frame_equal(result, expected)
+
+
+def test_string_inference(all_parsers):
+    # GH#54430
+    pa = pytest.importorskip("pyarrow")
+    dtype = pd.ArrowDtype(pa.string())
+
+    data = """a,b
+x,1
+y,2"""
+    parser = all_parsers
+    with pd.option_context("future.infer_string", True):
+        result = parser.read_csv(StringIO(data))
+
+    expected = DataFrame(
+        {"a": pd.Series(["x", "y"], dtype=dtype), "b": [1, 2]},
+        columns=pd.Index(["a", "b"], dtype=dtype),
+    )
+    tm.assert_frame_equal(result, expected)

pandas/tests/io/test_feather.py

@@ -219,3 +219,17 @@ def test_invalid_dtype_backend(self):
             df.to_feather(path)
             with pytest.raises(ValueError, match=msg):
                 read_feather(path, dtype_backend="numpy")
+
+    def test_string_inference(self, tmp_path):
+        # GH#54431
+        import pyarrow as pa
+
+        path = tmp_path / "test_string_inference.p"
+        df = pd.DataFrame(data={"a": ["x", "y"]})
+        df.to_feather(path)
+        with pd.option_context("future.infer_string", True):
+            result = read_feather(path)
+        expected = pd.DataFrame(
+            data={"a": ["x", "y"]}, dtype=pd.ArrowDtype(pa.string())
+        )
+        tm.assert_frame_equal(result, expected)