Repository navigation
Config option dataframe.dtype_backend: pyarrow doesn't seem to work #9840
Description
From https://docs.dask.org/en/stable/generated/dask.dataframe.read_parquet.html:
Use the dataframe.dtype_backend config option to select which dtype implementation to use.
dataframe.dtype_backend="pandas" (the default) will use pandas’ numpy-backed nullable dtypes (e.g. Int64, string[python], etc.) while dataframe.dtype_backend="pyarrow" will use pyarrow-backed extension dtypes (e.g. int64[pyarrow], string[pyarrow], etc.). dataframe.dtype_backend="pyarrow" requires pandas 1.5+.
coiled-runtime 0.2.1 + dask 2023.1.0:
>>> import dask.config
>>> import dask.dataframe as dd
>>> dask.config.set({"dataframe.dtype_backend": "pyarrow"})
>>> ddf = dd.read_parquet(
... "s3://coiled-datasets/prefect-dask/nyc-uber-lyft/processed_data.parquet",
... engine="pyarrow",
... filesystem="arrow",
... )
>>> ddf.dtypes['hvfhs_license_num']
string[python]Workaround
from distributed import WorkerPlugin
class SetPandasOptions(WorkerPlugin):
def setup(self, worker):
pd.set_option("string_storage", "pyarrow")
pd.set_option("string_storage", "pyarrow") # Set on the client
client.register_worker_plugin(SetPandasOptions()) # Set on the workers
Activity
jrbourbeau commentedon Jan 17, 2023
Thanks for the report @crusaderky. You also need to set
use_nullable_dtypes=True.use_nullable_dtypestells Dask whether or not to use extension dtypes and thedataframe.dtype_backendconfig option specifies which type of dtypes to usecrusaderky commentedon Jan 18, 2023
@jrbourbeau doesn't work:
I find the documentation confusing. From it, I can't understand what happens (or what's the benefit) from having
dataframe.dtype_backend="pyarrow"anduse_nullable_dtypes=False, or vice versa.In pandas all you need is a single global config option,
pd.set_option("string_storage", "pyarrow"), to switch to arrow strings; I don't see why it should be any different in dask?jrbourbeau commentedon Jan 20, 2023
I agree this is still erroring, but the dispatching to pyarrow dtypes appears to be working. The
Not all columns are categoricalsseems like a separate issue.We could add a note about
dataframe.dtype_backendonly having an effect whenuse_nullable_dtypes=True#9711 implements the
string_storageapproach in dask. Thedataframe.dtype_backendconfig option matches the upcoming behavior in pandas 2.0.crusaderky commentedon Sep 26, 2023
I believe this is fixed now.
Retested on dask 2023.9.2, pandas 2.1, pyarrow 13