Repository: dask-contrib/dask-expr
Pull request #911 (2024-02-29 13:49:00)

Wrap fragments
This is not my preferred way of doing things but this effectively bypasses https://github.com/apache/arrow/issues/40279 and delays the initialization cost of the filesystem to the workers. This is only necessarily because we no longer have the "don't deserialize on scheduler" mechanics (at least not properly). The most important part is the delay of the filesystem deserialization. However, storing this as a bytestring comes at a cost since pickle can't seem to deduplicate this any longer such that the overall size of the serialized graph is larger. 

I still have to run a few more tests to see if this is all really worth it.


Patch:
diff --git a/dask_expr/_backends.py b/dask_expr/_backends.py
index 2a9b93e..9913f37 100644
--- a/dask_expr/_backends.py
+++ b/dask_expr/_backends.py
@@ -1,11 +1,28 @@
 from __future__ import annotations
 
+import numpy as np
 import pandas as pd
 from dask.backends import CreationDispatch
 from dask.dataframe.backends import DataFrameBackendEntrypoint
 
 from dask_expr._dispatch import get_collection_type
 
+try:
+    import sparse
+
+    sparse_installed = True
+except ImportError:
+    sparse_installed = False
+
+
+try:
+    import scipy.sparse as sp
+
+    scipy_installed = True
+except ImportError:
+    scipy_installed = False
+
+
 dataframe_creation_dispatch = CreationDispatch(
     module_name="dataframe",
     default="pandas",
@@ -60,6 +77,38 @@ def get_collection_type_index(_):
     return Index
 
 
+def create_array_collection(expr):
+    # This is hacky and an abstraction leak, but utilizing get_collection_type
+    # to infer that we want to create an array is the only way that is guaranteed
+    # to be a general solution.
+    # We can get rid of this when we have an Array expression
+    from dask.dataframe.core import new_dd_object
+
+    result = expr.optimize()
+    return new_dd_object(
+        result.__dask_graph__(), result._name, result._meta, result.divisions
+    )
+
+
+@get_collection_type.register(np.ndarray)
+def get_collection_type_array(_):
+    return create_array_collection
+
+
+if sparse_installed:
+
+    @get_collection_type.register(sparse.COO)
+    def get_collection_type_array(_):
+        return create_array_collection
+
+
+if scipy_installed:
+
+    @get_collection_type.register(sp.csr_matrix)
+    def get_collection_type_array(_):
+        return create_array_collection
+
+
 @get_collection_type.register(object)
 def get_collection_type_object(_):
     from dask_expr._collection import Scalar
diff --git a/dask_expr/_collection.py b/dask_expr/_collection.py
index 1b5807a..0e8a201 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 (
@@ -3841,6 +3842,9 @@ class Series(FrameBase):
                 split_out = 1
             else:
                 split_out = True
+        if split_out == 1 and split_out is not True and sort is None:
+            sort = True
+
         length = None
         if (split_out > 1 or split_out is True) and normalize:
             frame = self if not dropna else self.dropna()
@@ -4426,6 +4430,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 +5002,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 +5017,7 @@ def from_map(
             )
         )
     else:
-        return new_collection(
+        result = new_collection(
             FromMap(
                 func,
                 iterables,
@@ -5024,6 +5029,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 +5067,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..fe0c3b6 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"}
@@ -2645,8 +2653,10 @@ class Partitions(Expr):
         return (self.frame._name, self.partitions[index])
 
     def _simplify_down(self):
+        from dask_expr import SetIndexBlockwise
+
         if isinstance(self.frame, Blockwise) and not isinstance(
-            self.frame, (BlockwiseIO, Fused)
+            self.frame, (BlockwiseIO, Fused, SetIndexBlockwise)
         ):
             operands = [
                 Partitions(op, self.partitions)
diff --git a/dask_expr/_groupby.py b/dask_expr/_groupby.py
index 5754542..76b5452 100644
--- a/dask_expr/_groupby.py
+++ b/dask_expr/_groupby.py
@@ -671,6 +671,34 @@ class GroupByReduction(Reduction, GroupByBase):
         meta = meta_nonempty(self.frame._meta)
         return self.chunk(meta, *self._by_meta, **self.chunk_kwargs)
 
+    def _divisions(self):
+        if self.sort:
+            return (None, None)
+        split_out = self.split_out
+        if split_out is True:
+            split_out = self.frame.npartitions
+        return (None,) * (split_out + 1)
+
+    def _simplify_up(self, parent, dependents):
+        return groupby_projection(self, parent, dependents)
+
+    @functools.cached_property
+    def combine_kwargs(self):
+        return {"levels": self.levels, "observed": self.observed, "dropna": self.dropna}
+
+    @functools.cached_property
+    def chunk_kwargs(self):
+        return {"observed": self.observed, "dropna": self.dropna}
+
+    @functools.cached_property
+    def aggregate_kwargs(self):
+        return {
+            "levels": self.levels,
+            "sort": self.sort,
+            "observed": self.observed,
+            "dropna": self.dropna,
+        }
+
 
 def _var_combine(g, levels, sort=False, observed=False, dropna=True):
     return g.groupby(level=levels, sort=sort, observed=observed, dropna=dropna).sum()
@@ -696,68 +724,24 @@ class Var(GroupByReduction):
         "split_every": None,
         "shuffle_method": None,
     }
-    reduction_aggregate = _var_agg
-    reduction_combine = _var_combine
-
-    @staticmethod
-    def chunk(frame, *by, **kwargs):
-        return _var_chunk(frame, *by, **kwargs)
+    reduction_aggregate = staticmethod(_var_agg)
+    reduction_combine = staticmethod(_var_combine)
+    chunk = staticmethod(_var_chunk)
 
     @functools.cached_property
     def aggregate_kwargs(self):
         return {
             "ddof": self.ddof,
-            "levels": self.levels,
             "numeric_only": self.numeric_only,
-            "sort": self.sort,
-            "observed": self.observed,
-            "dropna": self.dropna,
+            **super().aggregate_kwargs,
         }
 
     @functools.cached_property
     def chunk_kwargs(self):
-        return {
-            "numeric_only": self.numeric_only,
-            "observed": self.observed,
-            "dropna": self.dropna,
-        }
-
-    @functools.cached_property
-    def combine_kwargs(self):
-        return {"levels": self.levels, "observed": self.observed, "dropna": self.dropna}
-
-    def _divisions(self):
-        if self.sort:
-            return (None, None)
-        return (None,) * (self.split_out + 1)
-
-    def _simplify_up(self, parent, dependents):
-        return groupby_projection(self, parent, dependents)
+        return {"numeric_only": self.numeric_only, **super().chunk_kwargs}
 
 
-class Std(SingleAggregation):
-    _parameters = [
-        "frame",
-        "ddof",
-        "numeric_only",
-        "split_out",
-        "split_every",
-        "sort",
-        "dropna",
-        "observed",
-        "shuffle_method",
-    ]
-    _defaults = {
-        "split_out": 1,
-        "sort": None,
-        "split_every": None,
-        "shuffle_method": None,
-    }
-
-    @functools.cached_property
-    def _meta(self):
-        return self._lower()._meta
-
+class Std(Var):
     def _lower(self):
         v = Var(*self.operands)
         return MapPartitions(
@@ -780,8 +764,8 @@ def _mean_chunk(df, *by, observed=None, dropna=None):
     return concat([x, n], axis=1)
 
 
-def _mean_combine(g, levels, sort=False):
-    return g.groupby(level=levels, sort=sort).sum()
+def _mean_combine(g, levels, sort=False, observed=None, dropna=None):
+    return g.groupby(level=levels, sort=sort, observed=observed, dropna=dropna).sum()
 
 
 def _mean_agg(g, levels, sort=False, observed=False, dropna=True):
@@ -799,34 +783,6 @@ class Mean(GroupByReduction):
     reduction_combine = staticmethod(_mean_combine)
     chunk = staticmethod(_mean_chunk)
 
-    @functools.cached_property
-    def aggregate_kwargs(self):
-        return {
-            "levels": self.levels,
-            "sort": self.sort,
-            "observed": self.observed,
-            "dropna": self.dropna,
-        }
-
-    @functools.cached_property
-    def chunk_kwargs(self):
-        return {"observed": self.observed, "dropna": self.dropna}
-
-    @functools.cached_property
-    def combine_kwargs(self):
-        return {"levels": self.levels}
-
-    def _divisions(self):
-        if self.sort:
-            return (None, None)
-        split_out = self.split_out
-        if split_out is True:
-            split_out = self.frame.npartitions
-        return (None,) * (split_out + 1)
-
-    def _simplify_up(self, parent, dependents):
-        return groupby_projection(self, parent, dependents)
-
 
 def nunique_df_combine(dfs, *args, **kwargs):
     return _nunique_df_combine(concat(dfs), *args, **kwargs)
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/_shuffle.py b/dask_expr/_shuffle.py
index 710adc3..b99037d 100644
--- a/dask_expr/_shuffle.py
+++ b/dask_expr/_shuffle.py
@@ -545,6 +545,10 @@ class DiskShuffle(SimpleShuffle):
 class P2PShuffle(SimpleShuffle):
     """P2P worker-based shuffle implementation"""
 
+    @functools.cached_property
+    def _meta(self):
+        return self.frame._meta.drop(columns=self.partitioning_index)
+
     def _layer(self):
         from distributed.shuffle._shuffle import (
             ShuffleId,
@@ -574,6 +578,7 @@ class P2PShuffle(SimpleShuffle):
                 self.frame._meta,
                 set(parts_out),
                 True,
+                True,
             )
 
         dsk[_barrier_key] = (shuffle_barrier, token, transfer_keys)
diff --git a/dask_expr/_version.py b/dask_expr/_version.py
index f7b60f2..d8fe571 100644
--- a/dask_expr/_version.py
+++ b/dask_expr/_version.py
@@ -50,8 +50,8 @@ def get_config() -> VersioneerConfig:
     cfg = VersioneerConfig()
     cfg.VCS = "git"
     cfg.style = "pep440"
-    cfg.tag_prefix = ""
-    cfg.parentdir_prefix = "dask-"
+    cfg.tag_prefix = "v"
+    cfg.parentdir_prefix = "dask_expr-"
     cfg.versionfile_source = "dask_expr/_version.py"
     cfg.verbose = False
     return cfg
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
 
diff --git a/dask_expr/io/parquet.py b/dask_expr/io/parquet.py
index 2845bad..d176548 100644
--- a/dask_expr/io/parquet.py
+++ b/dask_expr/io/parquet.py
@@ -3,7 +3,9 @@ from __future__ import annotations
 import contextlib
 import itertools
 import operator
+import pickle
 import warnings
+import weakref
 from abc import abstractmethod
 from collections import defaultdict
 from functools import cached_property
@@ -28,7 +30,7 @@ from dask.dataframe.io.parquet.core import (
 from dask.dataframe.io.parquet.utils import _split_user_options
 from dask.dataframe.io.utils import _is_local_fs
 from dask.delayed import delayed
-from dask.utils import apply, funcname, natural_sort_key, typename
+from dask.utils import apply, funcname, natural_sort_key, parse_bytes, typename
 from fsspec.utils import stringify_path
 from toolz import identity
 
@@ -86,6 +88,66 @@ _CACHED_PLAN_SIZE = 10
 _cached_plan = {}
 
 
+class FragmentWrapper:
+    _filesystems = weakref.WeakValueDictionary()
+
+    def __init__(self, fragment=None, file_size=None, fragment_packed=None) -> None:
+        """Wrap a pyarrow Fragment to only deserialize when needed."""
+        # https://github.com/apache/arrow/issues/40279
+        self._fragment = fragment
+        self._fragment_packed = fragment_packed
+        self._file_size = file_size
+        self._fs = None
+
+    def pack(self):
+        if self._fragment_packed is None:
+            self._fragment_packed = (
+                self._fragment.format,
+                (
+                    self._fragment.path
+                    if self._fragment.buffer is None
+                    else self._fragment.buffer
+                ),
+                pickle.dumps(self._fragment.filesystem),
+                self._fragment.partition_expression,
+                self._file_size,
+            )
+        self._fs = self._fragment = None
+
+    def unpack(self):
+        if self._fragment is None:
+            (
+                pqformat,
+                path_or_buffer,
+                fs_raw,
+                partition_expression,
+                file_size,
+            ) = self._fragment_packed
+            fs = FragmentWrapper._filesystems.get(fs_raw)
+            if fs is None:
+                fs = pickle.loads(fs_raw)
+                FragmentWrapper._filesystems[fs_raw] = fs
+            # arrow doesn't keep the python object alive so if we want to reuse
+            # we need to keep a reference
+            self._fs = fs
+            self._fragment = pqformat.make_fragment(
+                path_or_buffer,
+                filesystem=fs,
+                partition_expression=partition_expression,
+                file_size=file_size,
+            )
+        self._fragment_packed = None
+
+    @property
+    def fragment(self):
+        self.unpack()
+        return self._fragment
+
+    def __reduce__(self):
+        self.pack()
+        return FragmentWrapper, (None, None, self._fragment_packed)
+
+
 def _control_cached_plan(key):
     if len(_cached_plan) > _CACHED_PLAN_SIZE and key not in _cached_plan:
         key_to_pop = list(_cached_plan.keys())[0]
@@ -675,9 +737,11 @@ class ReadParquetPyarrowFS(ReadParquet):
                     filesystem=self.fs,
                 )
                 dataset_info["using_metadata_file"] = True
-                dataset_info["fragments"] = dataset.get_fragments()
+                dataset_info["fragments"] = _frags = dataset.get_fragments()
+                dataset_info["file_sizes"] = [None for fi in _frags]
         if checksum is None:
             checksum = tokenize(all_files)
+            dataset_info["file_sizes"] = [fi.size for fi in all_files]
         dataset_info["checksum"] = checksum
         if dataset is None:
             import pyarrow.parquet as pq
@@ -704,7 +768,7 @@ class ReadParquetPyarrowFS(ReadParquet):
     def _divisions(self):
         return tuple([None] * (len(self.fragments) + 1))
 
-    @property
+    @cached_property
     def fragments(self):
         if self.filters is not None:
             if self._dataset_info["using_metadata_file"]:
@@ -714,54 +778,51 @@ class ReadParquetPyarrowFS(ReadParquet):
             return list(ds.get_fragments(filter=pq.filters_to_expression(self.filters)))
         return self._dataset_info["fragments"]
 
-    @staticmethod
-    def _fragment_to_pandas(fragment, columns, filters, schema):
-        from dask.utils import parse_bytes
-
-        if isinstance(filters, list):
-            filters = pq.filters_to_expression(filters)
-        # TODO: There should be a way for users to define the type mapper
-        table = fragment.to_table(
-            schema=schema,
-            columns=columns,
-            filter=filters,
-            # Batch size determines how many rows are read at once and will
-            # cause the underlying array to be split into chunks of this size
-            # (max). We'd like to avoid fragmentation as much as possible and
-            # and to set this to something like inf but we have to set a finite,
-            # positive number.
-            # In the presence of row groups, the underlying array will still be
-            # chunked per rowgroup
-            batch_size=10_000_000,
-            # batch_readahead=16,
-            # fragment_readahead=4,
-            fragment_scan_options=pa.dataset.ParquetFragmentScanOptions(
-                pre_buffer=True,
-                cache_options=pa.CacheOptions(
-                    hole_size_limit=parse_bytes("4 MiB"),
-                    range_size_limit=parse_bytes("32.00 MiB"),
-                ),
-            ),
-            # TODO: Reconsider this. The OMP_NUM_THREAD variable makes it harmful to enable this
-            use_threads=True,
-        )
-        df = table.to_pandas(
-            types_mapper=_determine_type_mapper(),
-            use_threads=False,
-            self_destruct=True,
-        )
-        return df
-
     def _filtered_task(self, index: int):
         return (
-            ReadParquetPyarrowFS._fragment_to_pandas,
-            self.fragments[index],
+            _fragment_to_pandas,
+            FragmentWrapper(self.fragments[index]),
             self.columns,
             self.filters,
-            self._dataset_info["schema"],
+            self._dataset_info["schema"].remove_metadata(),
         )
 
 
+def _fragment_to_pandas(fragment_wrapper, columns, filters, schema):
+    fragment = fragment_wrapper.fragment
+    if isinstance(filters, list):
+        filters = pq.filters_to_expression(filters)
+    # TODO: There should be a way for users to define the type mapper
+    table = fragment.to_table(
+        schema=schema,
+        columns=columns,
+        filter=filters,
+        # Batch size determines how many rows are read at once and will
+        # cause the underlying array to be split into chunks of this size
+        # (max). We'd like to avoid fragmentation as much as possible and
+        # and to set this to something like inf but we have to set a finite,
+        # positive number.
+        # In the presence of row groups, the underlying array will still be
+        # chunked per rowgroup
+        batch_size=10_000_000,
+        fragment_scan_options=pa.dataset.ParquetFragmentScanOptions(
+            pre_buffer=True,
+            cache_options=pa.CacheOptions(
+                hole_size_limit=parse_bytes("4 MiB"),
+                range_size_limit=parse_bytes("32.00 MiB"),
+            ),
+        ),
+        # TODO: Reconsider this. The OMP_NUM_THREAD variable makes it harmful to enable this
+        use_threads=True,
+    )
+    df = table.to_pandas(
+        types_mapper=_determine_type_mapper(),
+        use_threads=False,
+        self_destruct=True,
+    )
+    return df
+
+
 class ReadParquetFSSpec(ReadParquet):
     """Read a parquet dataset"""
 
diff --git a/pyproject.toml b/pyproject.toml
index 4fd8bac..b6ac865 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -61,7 +61,7 @@ style = "pep440"
 versionfile_source = "dask_expr/_version.py"
 versionfile_build = "dask_expr/_version.py"
 tag_prefix = "v"
-parentdir_prefix = "dask-"
+parentdir_prefix = "dask_expr-"
 
 
 [tool.pytest.ini_options]


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/io/tests/test_parquet.py b/dask_expr/io/tests/test_parquet.py
index 9ba0aaf..479302d 100644
--- a/dask_expr/io/tests/test_parquet.py
+++ b/dask_expr/io/tests/test_parquet.py
@@ -1,11 +1,13 @@
 import os
+import pickle
 
 import pandas as pd
 import pytest
 from dask.dataframe.utils import assert_eq
+from dask.utils import key_split
 from pyarrow import fs
 
-from dask_expr import from_pandas, read_parquet
+from dask_expr import from_graph, from_pandas, read_parquet
 from dask_expr._expr import Filter, Lengths, Literal
 from dask_expr._reductions import Len
 from dask_expr.io import FusedIO, ReadParquet
@@ -82,6 +84,22 @@ def test_pyarrow_filesystem(parquet_file):
     assert assert_eq(df, df_pa)
 
 
+def test_pyarrow_filesystem_serialize(parquet_file):
+    filesystem = fs.LocalFileSystem()
+
+    df_pa = read_parquet(parquet_file, filesystem=filesystem)
+
+    roundtripped = pickle.loads(pickle.dumps(df_pa.optimize().dask))
+    roundtripped_df = from_graph(
+        roundtripped,
+        df_pa._meta,
+        df_pa.divisions,
+        df_pa.__dask_keys__(),
+        key_split(df_pa._name),
+    )
+    assert assert_eq(df_pa, roundtripped_df)
+
+
 def test_pyarrow_filesystem_filters(parquet_file):
     filesystem = fs.LocalFileSystem()
 
diff --git a/dask_expr/tests/test_collection.py b/dask_expr/tests/test_collection.py
index 4e317b5..a0e9440 100644
--- a/dask_expr/tests/test_collection.py
+++ b/dask_expr/tests/test_collection.py
@@ -129,12 +129,18 @@ def test_setitem(pdf, df):
     assert_eq(df, pdf)
 
 
+def test_set_index_blockwise_culling(df, pdf):
+    result = df.set_index("x").partitions[0]
+    assert_eq(result, pdf.set_index("x").head(10))
+
+
 @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 +1468,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 +1483,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 +2245,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()
 
 
@@ -2425,7 +2432,15 @@ def test_scalar_repr(df):
     assert result == "<dask_expr.expr.Scalar: expr=df.size(), dtype=int64>"
 
 
-def test_reset_index_filter_pushdown(df):
+def test_array_conversion(df, pdf):
+    import dask.array as da
+
+    result = df.map_partitions(lambda x: np.asarray(x))
+    expected = np.asarray(pdf)
+    da.utils.assert_eq(result, expected)
+
+
+def test_reset_index_filter_pushdown(df, pdf):
     q = df.reset_index()
     result = q[q.x > 5]
     expected = df[df.x > 5].reset_index()
@@ -2433,8 +2448,9 @@ def test_reset_index_filter_pushdown(df):
 
     # We are accessing the index, so don't do anything
     result = q[q["index"] > 5]
-    expected = df[df.index > 5].reset_index()
-    assert result.simplify()._name == expected._name
+    expected = pdf.reset_index()
+    expected = expected[expected["index"] > 5]
+    assert_eq(result, expected, check_index=False)
 
     q = df.x.reset_index(drop=True)
     result = q[q > 5]
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_reductions.py b/dask_expr/tests/test_reductions.py
index f56e0a5..77b987d 100644
--- a/dask_expr/tests/test_reductions.py
+++ b/dask_expr/tests/test_reductions.py
@@ -105,6 +105,15 @@ def test_value_counts(pdf, df, split_every, split_out):
     )
 
 
+def test_value_counts_sort(pdf):
+    pdf = pdf.iloc[4:]
+    df = from_pandas(pdf, npartitions=10)
+    pd.testing.assert_series_equal(
+        df.y.value_counts(split_out=1).compute(),
+        pdf.y.value_counts(),
+    )
+
+
 @pytest.mark.parametrize("split_every", [None, 5])
 @pytest.mark.parametrize("split_out", [1, True])
 def test_unique(pdf, df, split_every, split_out):
diff --git a/dask_expr/tests/test_shuffle.py b/dask_expr/tests/test_shuffle.py
index 3627762..8c627ee 100644
--- a/dask_expr/tests/test_shuffle.py
+++ b/dask_expr/tests/test_shuffle.py
@@ -587,10 +587,10 @@ def test_set_index_predicate_pushdown(df, pdf):
     expected_query = df[df.y > 5].set_index("x").optimize()
     assert expected_query._name == result.optimize()._name
 
-    result = query[query.index > 5]
+    result = query[query.index.to_series() > 5]
     assert_eq(result, pdf[pdf.index > 5])
 
-    result = query[(query.index > 5) & (query.y > -1)]
+    result = query[(query.index.to_series() > 5) & (query.y > -1)]
     assert_eq(result, pdf[(pdf.index > 5) & (pdf.y > -1)])
 
 
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