Repository: dask-contrib/dask-expr
Pull request #912 (2024-02-29 16:31:04)

Add support for convert string option
xref #691

this will probably break a bunch of tests in dask/dask

dask/dask pr: https://github.com/dask/dask/pull/10968

phofl: tests here: https://github.com/dask-contrib/dask-expr/pull/916
phofl: yep ci is happy on the other pr, merging here

Patch:
diff --git a/dask_expr/_collection.py b/dask_expr/_collection.py
index 1b5807a..c63f8e1 100644
--- a/dask_expr/_collection.py
+++ b/dask_expr/_collection.py
@@ -13,7 +13,7 @@ import dask.dataframe.methods as methods
 import numpy as np
 import pandas as pd
 import pyarrow as pa
-from dask import compute, delayed
+from dask import compute
 from dask.array import Array
 from dask.base import DaskMethodsMixin, is_dask_collection, named_schedulers
 from dask.core import flatten
@@ -42,6 +42,7 @@ from dask.dataframe.utils import (
     insert_meta_param_description,
     meta_frame_constructor,
     meta_series_constructor,
+    pyarrow_strings_enabled,
 )
 from dask.delayed import delayed
 from dask.utils import (
@@ -4426,6 +4427,7 @@ def from_pandas(data, npartitions=None, sort=True, chunksize=None):
             npartitions=npartitions,
             sort=sort,
             chunksize=chunksize,
+            pyarrow_strings_enabled=pyarrow_strings_enabled(),
         )
     )
 
@@ -4997,7 +4999,7 @@ def from_map(
     kwargs = {} if kwargs is None else kwargs
     if allow_projection:
         columns = kwargs.pop("columns", None)
-        return new_collection(
+        result = new_collection(
             FromMapProjectable(
                 func,
                 iterables,
@@ -5012,7 +5014,7 @@ def from_map(
             )
         )
     else:
-        return new_collection(
+        result = new_collection(
             FromMap(
                 func,
                 iterables,
@@ -5024,6 +5026,9 @@ def from_map(
                 label,
             )
         )
+    if pyarrow_strings_enabled():
+        return new_collection(expr.ArrowStringConversion(result))
+    return result
 
 
 def repartition(df, divisions, force=False):
@@ -5059,7 +5064,11 @@ def repartition(df, divisions, force=False):
         return df.repartition(divisions=divisions, force=force)
     elif is_dataframe_like(df) or is_series_like(df):
         return new_collection(
-            FromPandasDivisions(_BackendData(df), divisions=divisions)
+            FromPandasDivisions(
+                _BackendData(df),
+                divisions=divisions,
+                pyarrow_strings_enabled=pyarrow_strings_enabled(),
+            )
         )
     else:
         raise NotImplementedError(f"repartition is not implemented for {type(df)}.")
diff --git a/dask_expr/_expr.py b/dask_expr/_expr.py
index cc529d1..291178b 100644
--- a/dask_expr/_expr.py
+++ b/dask_expr/_expr.py
@@ -15,6 +15,7 @@ from dask.array import Array
 from dask.base import normalize_token
 from dask.core import flatten
 from dask.dataframe import methods
+from dask.dataframe._pyarrow import to_pyarrow_string
 from dask.dataframe.core import (
     _concat,
     _get_divisions_map_partitions,
@@ -1270,6 +1271,13 @@ class Clip(Elemwise):
             return plain_column_projection(self, parent, dependents)
 
 
+class ArrowStringConversion(Elemwise):
+    _projection_passthrough = True
+    _filter_passthrough = True
+    _parameters = ["frame"]
+    operation = staticmethod(to_pyarrow_string)
+
+
 class Between(Elemwise):
     _parameters = ["frame", "left", "right", "inclusive"]
     _defaults = {"inclusive": "both"}
diff --git a/dask_expr/_merge_asof.py b/dask_expr/_merge_asof.py
index 7efcc13..03d7249 100644
--- a/dask_expr/_merge_asof.py
+++ b/dask_expr/_merge_asof.py
@@ -4,6 +4,7 @@ import pandas as pd
 from dask.dataframe import methods
 from dask.dataframe.dispatch import make_meta, meta_nonempty
 from dask.dataframe.multi import merge_asof_padded, pair_partitions
+from dask.dataframe.utils import pyarrow_strings_enabled
 from dask.utils import apply
 
 from dask_expr import SetIndexBlockwise, new_collection
@@ -113,7 +114,11 @@ class MergeAsof(Merge):
         left_index, right_index = True, True
 
         if all(map(pd.isnull, left.divisions)):
-            return FromPandas(_BackendData(self._meta), npartitions=left.npartitions)
+            return FromPandas(
+                _BackendData(self._meta),
+                npartitions=left.npartitions,
+                pyarrow_strings_enabled=pyarrow_strings_enabled(),
+            )
 
         if all(map(pd.isnull, right.divisions)):
             return MapPartitions(
diff --git a/dask_expr/datasets.py b/dask_expr/datasets.py
index 3634b3b..3d728e9 100644
--- a/dask_expr/datasets.py
+++ b/dask_expr/datasets.py
@@ -3,8 +3,10 @@ import operator
 
 import numpy as np
 import pandas as pd
+from dask.dataframe.utils import pyarrow_strings_enabled
 
 from dask_expr._collection import new_collection
+from dask_expr._expr import ArrowStringConversion
 from dask_expr.io import BlockwiseIO, PartitionsFiltered
 
 __all__ = ["timeseries"]
@@ -238,4 +240,6 @@ def timeseries(
         kwargs,
         columns=list(dtypes.keys()),
     )
+    if pyarrow_strings_enabled():
+        return new_collection(ArrowStringConversion(expr))
     return new_collection(expr)
diff --git a/dask_expr/io/io.py b/dask_expr/io/io.py
index 1b6b34f..b2611ba 100644
--- a/dask_expr/io/io.py
+++ b/dask_expr/io/io.py
@@ -6,6 +6,7 @@ import operator
 
 import numpy as np
 from dask.dataframe import methods
+from dask.dataframe._pyarrow import to_pyarrow_string
 from dask.dataframe.core import apply_and_enforce, is_dataframe_like, make_meta
 from dask.dataframe.io.io import _meta_from_array, sorted_division_locations
 from dask.utils import apply, funcname, is_series_like
@@ -319,6 +320,7 @@ class FromPandas(PartitionsFiltered, BlockwiseIO):
         "sort",
         "chunksize",
         "columns",
+        "pyarrow_strings_enabled",
         "_partitions",
         "_series",
     ]
@@ -343,7 +345,11 @@ class FromPandas(PartitionsFiltered, BlockwiseIO):
 
     @functools.cached_property
     def _meta(self):
-        meta = self.frame.head(0)
+        if self.pyarrow_strings_enabled:
+            meta = make_meta(to_pyarrow_string(self.frame.head(1)))
+        else:
+            meta = self.frame.head(0)
+
         if self.operand("columns") is not None:
             return meta[self.columns[0]] if self._series else meta[self.columns]
         return meta
@@ -427,6 +433,8 @@ class FromPandas(PartitionsFiltered, BlockwiseIO):
     def _filtered_task(self, index: int):
         start, stop = self._locations()[index : index + 2]
         part = self.frame.iloc[start:stop]
+        if self.pyarrow_strings_enabled:
+            part = to_pyarrow_string(part)
         if self.operand("columns") is not None:
             return part[self.columns[0]] if self._series else part[self.columns]
         return part
@@ -442,7 +450,14 @@ class FromPandas(PartitionsFiltered, BlockwiseIO):
 
 
 class FromPandasDivisions(FromPandas):
-    _parameters = ["frame", "divisions", "columns", "_partitions", "_series"]
+    _parameters = [
+        "frame",
+        "divisions",
+        "columns",
+        "pyarrow_strings_enabled",
+        "_partitions",
+        "_series",
+    ]
     _defaults = {"columns": None, "_partitions": None, "_series": False}
     sort = True
 


Test patch:
diff --git a/dask_expr/io/tests/test_from_pandas.py b/dask_expr/io/tests/test_from_pandas.py
index f6c5962..2d430f6 100644
--- a/dask_expr/io/tests/test_from_pandas.py
+++ b/dask_expr/io/tests/test_from_pandas.py
@@ -1,3 +1,4 @@
+import dask
 import pytest
 from dask.dataframe.utils import assert_eq
 
@@ -108,3 +109,21 @@ def test_from_pandas_chunksize(pdf, chunksize, npartitions, sort):
 def test_from_pandas_npartitions_and_chunksize(pdf):
     with pytest.raises(ValueError, match="npartitions and chunksize"):
         from_pandas(pdf, npartitions=2, chunksize=3)
+
+
+def test_from_pandas_string_option():
+    pdf = pd.DataFrame({"x": [1, 2, 3], "y": "a"}, index=["a", "b", "c"])
+    df = from_pandas(pdf, npartitions=2)
+    assert df.dtypes["y"] == "string"
+    assert df.index.dtype == "string"
+    assert df.compute().dtypes["y"] == "string"
+    assert df.compute().index.dtype == "string"
+    assert_eq(df, pdf)
+
+    with dask.config.set({"dataframe.convert-string": False}):
+        df = from_pandas(pdf, npartitions=2)
+        assert df.dtypes["y"] == "object"
+        assert df.index.dtype == "object"
+        assert df.compute().dtypes["y"] == "object"
+        assert df.compute().index.dtype == "object"
+        assert_eq(df, pdf)
diff --git a/dask_expr/io/tests/test_io.py b/dask_expr/io/tests/test_io.py
index a76bf13..d3905b6 100644
--- a/dask_expr/io/tests/test_io.py
+++ b/dask_expr/io/tests/test_io.py
@@ -312,12 +312,12 @@ def test_from_map(tmpdir, meta, label, allow_projection, enforce_metadata):
     assert_eq(df[["a", "b"]], pdf[["a", "b"]], check_index=False)
 
     if label:
-        assert df.expr._name.startswith(label)
+        assert df.expr.frame._name.startswith(label)
 
     if allow_projection:
         got = df[["a", "b"]].optimize(fuse=False)
-        assert isinstance(got.expr, FromMap)
-        assert got.expr.operand("columns") == ["a", "b"]
+        assert isinstance(got.expr.frame, FromMap)
+        assert got.expr.frame.operand("columns") == ["a", "b"]
 
     # Check that we can always pass columns up front
     if meta:
@@ -359,6 +359,16 @@ def test_from_map_columns_required():
     assert actual._name == expected._name
 
 
+def func_object(path):
+    return pd.DataFrame({"a": "x", "b": range(10)})
+
+
+def test_from_map_string_conversion():
+    result = from_map(func_object, ["foo"])
+    assert result.a.dtype == "string"
+    assert result.a.compute().dtype == "string"
+
+
 def test_from_array():
     arr = np.random.randint(1, 100, (100,))
     assert_eq(from_array(arr, chunksize=5), pd.Series(arr))
diff --git a/dask_expr/tests/test_collection.py b/dask_expr/tests/test_collection.py
index 02a705a..a1b9f1d 100644
--- a/dask_expr/tests/test_collection.py
+++ b/dask_expr/tests/test_collection.py
@@ -131,10 +131,11 @@ def test_setitem(pdf, df):
 
 @xfail_gpu("https://github.com/rapidsai/cudf/issues/10271")
 def test_explode():
-    pdf = pd.DataFrame({"a": [[1, 2], [3, 4]]})
-    df = from_pandas(pdf)
-    assert_eq(pdf.explode(column="a"), df.explode(column="a"))
-    assert_eq(pdf.a.explode(), df.a.explode())
+    with dask.config.set({"dataframe.convert-string": False}):
+        pdf = pd.DataFrame({"a": [[1, 2], [3, 4]]})
+        df = from_pandas(pdf)
+        assert_eq(pdf.explode(column="a"), df.explode(column="a"))
+        assert_eq(pdf.a.explode(), df.a.explode())
 
 
 @xfail_gpu("https://github.com/rapidsai/cudf/issues/10271")
@@ -1462,8 +1463,9 @@ def test_tree_repr(fuse):
     assert "Sum:" in s
     assert "Add:" in s
     assert "Mean:" in s
+    assert "ArrowStringConversion" in s
     assert "AlignPartitions:" not in s
-    assert str(df.seed) in s.lower()
+    assert str(df.frame.seed) in s.lower()
 
     # Check result after optimization
     optimized = expr.optimize(fuse=fuse)
@@ -1476,7 +1478,7 @@ def test_tree_repr(fuse):
     assert "True" not in s
     assert "None" not in s
     assert "skipna=False" in s
-    assert str(df.seed) in s.lower()
+    assert str(df.frame.seed) in s.lower()
     if fuse:
         assert "Fused" in s
         assert "|" in s
@@ -2238,7 +2240,7 @@ def test_quantile(df):
         df.x.quantile(q=[]).compute()
 
     ser = from_pandas(pd.Series(["a", "b", "c"]), npartitions=2)
-    with pytest.raises(ValueError, match="object series"):
+    with pytest.raises(ValueError, match="not supported with string series"):
         ser.quantile()
 
 
diff --git a/dask_expr/tests/test_datasets.py b/dask_expr/tests/test_datasets.py
index 0d8534d..83422c6 100644
--- a/dask_expr/tests/test_datasets.py
+++ b/dask_expr/tests/test_datasets.py
@@ -18,12 +18,23 @@ def test_timeseries():
 def test_optimization():
     df = timeseries(dtypes={"x": int, "y": float}, seed=123)
     expected = timeseries(dtypes={"x": int}, seed=123)
-    result = df[["x"]].optimize()
-    assert result.expr.operand("columns") == expected.expr.operand("columns")
+    result = df[["x"]].optimize(fuse=False)
+    assert result.expr.frame.operand("columns") == expected.expr.frame.operand(
+        "columns"
+    )
 
     expected = timeseries(dtypes={"x": int}, seed=123)["x"].simplify()
     result = df["x"].optimize(fuse=False)
-    assert expected.expr.operand("columns") == result.expr.operand("columns")
+    assert expected.expr.frame.operand("columns") == result.expr.frame.operand(
+        "columns"
+    )
+
+
+def test_arrow_string_option():
+    df = timeseries(dtypes={"x": object, "y": float}, seed=123)
+    result = df.optimize(fuse=False)
+    assert result.x.dtype == "string"
+    assert result.x.compute().dtype == "string"
 
 
 def test_column_projection_deterministic():
diff --git a/dask_expr/tests/test_distributed.py b/dask_expr/tests/test_distributed.py
index af9b321..cabb150 100644
--- a/dask_expr/tests/test_distributed.py
+++ b/dask_expr/tests/test_distributed.py
@@ -362,7 +362,7 @@ async def test_merge_indicator(c, s, a, b):
         "id": ["101-a", "102-a", "103-a"],
         "test": ["val101a", "val102a", "val103a"],
     }
-    pdf = pd.DataFrame(data)
+    pdf = pd.DataFrame(data, dtype="string[pyarrow]")
     df = from_pandas(pdf, npartitions=2)
     result = df.merge(df, on="id", how="outer", indicator=True)
     x = c.compute(result)
diff --git a/dask_expr/tests/test_format.py b/dask_expr/tests/test_format.py
index c542181..926d315 100644
--- a/dask_expr/tests/test_format.py
+++ b/dask_expr/tests/test_format.py
@@ -26,7 +26,7 @@ def test_to_string():
         """\
                            A       B                C
         npartitions=3                                
-        0              int64  object  category[known]
+        0              int64  string  category[known]
         3                ...     ...              ...
         6                ...     ...              ...
         7                ...     ...              ..."""
@@ -52,7 +52,7 @@ def test_to_string():
     <tr>
       <th>0</th>
       <td>int64</td>
-      <td>object</td>
+      <td>string</td>
       <td>category[known]</td>
     </tr>
     <tr>
diff --git a/dask_expr/tests/test_string_accessor.py b/dask_expr/tests/test_string_accessor.py
index ad6a3be..f1901f8 100644
--- a/dask_expr/tests/test_string_accessor.py
+++ b/dask_expr/tests/test_string_accessor.py
@@ -1,4 +1,3 @@
-import numpy as np
 import pytest
 from dask.dataframe._compat import PANDAS_GE_200
 
@@ -15,7 +14,9 @@ def ser():
 
 @pytest.fixture()
 def dser(ser):
-    return from_pandas(ser, npartitions=3)
+    import dask.dataframe as dd
+
+    return dd.from_pandas(ser, npartitions=3)
 
 
 @pytest.mark.parametrize(
@@ -79,8 +80,28 @@ def dser(ser):
     ],
 )
 def test_string_accessor(ser, dser, func, kwargs):
+    ser = ser.astype("string[pyarrow]")
+
     assert_eq(getattr(ser.str, func)(**kwargs), getattr(dser.str, func)(**kwargs))
 
+    if func in (
+        "contains",
+        "endswith",
+        "fullmatch",
+        "isalnum",
+        "isalpha",
+        "isdecimal",
+        "isdigit",
+        "islower",
+        "isspace",
+        "istitle",
+        "isupper",
+        "startswith",
+        "match",
+    ):
+        # This returns arrays and doesn't work in dask/dask either
+        return
+
     ser.index = ser.values
     ser = ser.sort_index()
     dser = from_pandas(ser, npartitions=3)
@@ -89,9 +110,6 @@ def test_string_accessor(ser, dser, func, kwargs):
     if func == "cat" and len(kwargs) > 0:
         # Doesn't work with others on Index
         return
-
-    if isinstance(pdf_result, np.ndarray):
-        pdf_result = pd.Index(pdf_result)
     if isinstance(pdf_result, pd.DataFrame):
         assert_eq(
             getattr(dser.index.str, func)(**kwargs), pdf_result, check_index=False