Add simple tests

Author: jrbourbeau

conftest.py

@@ -26,6 +26,18 @@
 
 try:
     import pandas  # noqa: F401
+
+    # Temporary changes to look for pyarrow string failures
+    import dask
+    from dask.dataframe._compat import PANDAS_GT_130
+
+    try:
+        import pyarrow
+    except ImportError:
+        pyarrow = False
+
+    if PANDAS_GT_130 and pyarrow:
+        dask.config.set({"dataframe.object_as_pyarrow_string": True})
 except ImportError:
     collect_ignore_glob.append("dask/dataframe/*")
 
@@ -68,16 +80,3 @@ def pytest_runtest_setup(item):
 def shuffle_method(request):
     with dask.config.set(shuffle=request.param):
         yield request.param
-
-
-# Temporary changes to look for pyarrow string failures
-import dask
-from dask.dataframe._compat import PANDAS_GT_130
-
-try:
-    import pyarrow
-except ImportError:
-    pyarrow = False
-
-if PANDAS_GT_130 and pyarrow:
-    dask.config.set({"dataframe.object_as_pyarrow_string": True})

dask/dataframe/core.py

@@ -323,11 +323,40 @@ def _scalar_binary(op, self, other, inv=False):
         return Scalar(graph, name, meta)
 
 
-def _maybe_convert_dtype(dtype):
-    if is_object_dtype(dtype):
-        return pd.StringDtype("pyarrow")
-    else:
-        return dtype
+def _is_object_or_string_dtype(dtype):
+    """Determine if input dtype is `object` or `string[python]`"""
+    if is_object_dtype(dtype) or (
+        isinstance(dtype, pd.StringDtype) and dtype.storage == "python"
+    ):
+        return True
+    return False
+
+
+def to_pyarrow_string(df):
+    if not (is_dataframe_like(df) or is_series_like(df) or is_index_like(df)):
+        return df
+
+    dtypes = None
+    if is_dataframe_like(df):
+        dtypes = {
+            col: pd.StringDtype("pyarrow")
+            for col, s in df.items()
+            if _is_object_or_string_dtype(s.dtype)
+        }
+    elif _is_object_or_string_dtype(df.dtype):
+        dtypes = pd.StringDtype("pyarrow")
+
+    if dtypes is not None:
+        df = df.astype(dtypes)
+
+    # Convert DataFrame and Series index too
+    if (
+        (is_dataframe_like(df) or is_series_like(df))
+        and not isinstance(df.index, pd.MultiIndex)
+        and _is_object_or_string_dtype(df.index.dtype)
+    ):
+        df.index = df.index.astype(pd.StringDtype("pyarrow"))
+    return df
 
 
 class _Frame(DaskMethodsMixin, OperatorMethodMixin):
@@ -378,45 +407,26 @@ def __init__(self, dsk, name, meta, divisions):
                 )
 
             def _index_check(x):
-                # MultiIndex don't support non-object dtypes
                 return (
                     is_index_like(x)
-                    and is_object_dtype(x)
-                    and not isinstance(x, pd.MultiIndex)
+                    and _is_object_or_string_dtype(x)
+                    and not isinstance(
+                        x, pd.MultiIndex
+                    )  # MultiIndex don't support non-object dtypes
                 )
 
             def _series_check(x):
                 return is_series_like(x) and (
-                    is_object_dtype(x) or _index_check(x.index)
+                    _is_object_or_string_dtype(x) or _index_check(x.index)
                 )
 
-            def _df_check(x):
+            def _dataframe_check(x):
                 return is_dataframe_like(x) and (
                     any(_series_check(s) for _, s in x.items()) or _index_check(x.index)
                 )
 
-            if _df_check(meta) or _series_check(meta) or _index_check(meta):
-
-                def _object_to_pyarrow_string(df):
-                    if not (
-                        is_dataframe_like(df) or is_series_like(df) or is_index_like(df)
-                    ):
-                        return df
-                    if is_dataframe_like(df):
-                        dtypes = {
-                            col: _maybe_convert_dtype(df[col].dtype) for col in df
-                        }
-                    else:
-                        dtypes = _maybe_convert_dtype(df.dtype)
-                    df = df.astype(dtypes)
-
-                    if (is_dataframe_like(df) or is_series_like(df)) and not isinstance(
-                        df.index, pd.MultiIndex
-                    ):
-                        df.index = df.index.astype(_maybe_convert_dtype(df.index.dtype))
-                    return df
-
-                result = self.map_partitions(_object_to_pyarrow_string)
+            if _dataframe_check(meta) or _series_check(meta) or _index_check(meta):
+                result = self.map_partitions(to_pyarrow_string)
                 self.dask = result.dask
                 self._name = result._name
                 self._meta = result._meta

dask/dataframe/io/tests/test_csv.py

@@ -360,6 +360,21 @@ def test_read_csv(dd_read, pd_read, text, sep):
         assert_eq(result, pd_read(fn, sep=sep))
 
 
+def test_read_csv_object_as_pyarrow_string_config():
+    pytest.importorskip(
+        "pandas",
+        minversion="1.3.0",
+        reason="Requires support for pyarrow strings",
+    )
+    pytest.importorskip("pyarrow", reason="Requires pyarrow")
+    with filetext(csv_text) as fn:
+        df = pd.read_csv(fn)
+        with dask.config.set({"dataframe.object_as_pyarrow_string": True}):
+            ddf = dd.read_csv(fn)
+        df_pyarrow = df.astype({"name": "string[pyarrow]"})
+        assert_eq(df_pyarrow, ddf, check_index=False)
+
+
 @pytest.mark.parametrize(
     "dd_read,pd_read,text,skip",
     [

dask/dataframe/io/tests/test_io.py

@@ -4,6 +4,7 @@
 import pandas as pd
 import pytest
 
+import dask
 import dask.array as da
 import dask.dataframe as dd
 from dask import config
@@ -274,6 +275,45 @@ def test_from_pandas_npartitions_duplicates(index):
     assert ddf.divisions == ("A", "B", "C", "C")
 
 
+def test_from_pandas_object_as_pyarrow_string_config():
+    pytest.importorskip(
+        "pandas",
+        minversion="1.3.0",
+        reason="Requires support for pyarrow strings",
+    )
+    pytest.importorskip("pyarrow", reason="Requires pyarrow")
+
+    # `dataframe.object_as_pyarrow_string` defaults to `False`
+    s = pd.Series(["foo", "bar", "ricky", "bobby"], index=["a", "b", "c", "d"])
+    df = pd.DataFrame(
+        {
+            "x": [1, 2, 3, 4],
+            "y": [5.0, 6.0, 7.0, 8.0],
+            "z": ["foo", "bar", "ricky", "bobby"],
+        },
+        index=["a", "b", "c", "d"],
+    )
+
+    ds = dd.from_pandas(s, npartitions=2)
+    ddf = dd.from_pandas(df, npartitions=2)
+
+    assert_eq(s, ds)
+    assert_eq(df, ddf)
+
+    # When `dataframe.object_as_pyarrow_string = True`, dask should automatically
+    # cast `object`s to pyarrow strings
+    with dask.config.set({"dataframe.object_as_pyarrow_string": True}):
+        ds = dd.from_pandas(s, npartitions=2)
+        ddf = dd.from_pandas(df, npartitions=2)
+
+    s_pyarrow = s.astype("string[pyarrow]")
+    s_pyarrow.index = s_pyarrow.index.astype("string[pyarrow]")
+    df_pyarrow = df.astype({"z": "string[pyarrow]"})
+    df_pyarrow.index = df_pyarrow.index.astype("string[pyarrow]")
+    assert_eq(s_pyarrow, ds)
+    assert_eq(df_pyarrow, ddf)
+
+
 @pytest.mark.gpu
 def test_gpu_from_pandas_npartitions_duplicates():
     cudf = pytest.importorskip("cudf")

dask/dataframe/utils.py

@@ -12,7 +12,7 @@
 
 import numpy as np
 import pandas as pd
-from pandas.api.types import is_categorical_dtype, is_dtype_equal, is_object_dtype
+from pandas.api.types import is_categorical_dtype, is_dtype_equal
 
 from dask.base import get_scheduler, is_dask_collection
 from dask.core import get_deps
@@ -530,25 +530,6 @@ def _maybe_sort(a, check_index: bool):
     return a.sort_index() if check_index else a
 
 
-def _maybe_convert_to_pyarrow(a):
-    from dask.dataframe.core import _maybe_convert_dtype
-
-    if isinstance(a, pd.DataFrame):
-        dtypes = {col: _maybe_convert_dtype(a[col].dtype) for col in a}
-        a = a.astype(dtypes)
-    elif isinstance(a, pd.Series):
-        a = a.astype(_maybe_convert_dtype(a.dtype))
-    elif isinstance(a, pd.Index) and not isinstance(a, pd.MultiIndex):
-        a = a.astype(_maybe_convert_dtype(a.dtype))
-    if (
-        isinstance(a, (pd.DataFrame, pd.Series))
-        and not isinstance(a.index, pd.MultiIndex)
-        and is_object_dtype(a.index)
-    ):
-        a.index = a.index.astype(_maybe_convert_dtype(a.index))
-    return a
-
-
 def assert_eq(
     a,
     b,
@@ -564,8 +545,10 @@ def assert_eq(
     import dask
 
     if dask.config.get("dataframe.object_as_pyarrow_string"):
-        a = _maybe_convert_to_pyarrow(a)
-        b = _maybe_convert_to_pyarrow(b)
+        from dask.dataframe.core import to_pyarrow_string
+
+        a = to_pyarrow_string(a)
+        b = to_pyarrow_string(b)
 
     if check_divisions:
         assert_divisions(a, scheduler=scheduler)