Skip to content

Improved conversion between pyarrow and pandas in P2P shuffling #7896

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jrbourbeau merged 8 commits into from
Jun 15, 2023
Merged

Improved conversion between pyarrow and pandas in P2P shuffling #7896

Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
87fef7c
Add types mapper
hendrikmakait Jun 9, 2023
3f46581
large_string
hendrikmakait Jun 12, 2023
21f129d
Minor
hendrikmakait Jun 12, 2023
684d988
[skip-caching]
hendrikmakait Jun 12, 2023
38898f7
Copy only if necessary
hendrikmakait Jun 12, 2023
c27be14
Merge branch 'main' into dont-convert-strings
hendrikmakait Jun 13, 2023
f793bcc
[skip-caching]
hendrikmakait Jun 13, 2023
8e93913
[skip-caching]
hendrikmakait Jun 14, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 14 additions & 1 deletion distributed/shuffle/_arrow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ def check_minimal_arrow_version() -> None:


def convert_partition(data: bytes, meta: pd.DataFrame) -> pd.DataFrame:
import pandas as pd
import pyarrow as pa

file = BytesIO(data)
Expand All @@ -56,7 +57,19 @@ def convert_partition(data: bytes, meta: pd.DataFrame) -> pd.DataFrame:
shards.append(sr.read_all())
table = pa.concat_tables(shards)
df = table.to_pandas(self_destruct=True)
return df.astype(meta.dtypes)

def default_types_mapper(pyarrow_dtype: pa.DataType) -> object:
# Avoid converting strings from `string[pyarrow]` to `string[python]`
# if we have *some* `string[pyarrow]`
if (
pyarrow_dtype in {pa.large_string(), pa.string()}
and pd.StringDtype("pyarrow") in meta.dtypes.values
Copy link
Member

@jrbourbeau jrbourbeau Jun 12, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are actually two implementations of pyarrow-strings in pandas. The one you have here and also pd.ArrowDtype(pa.string()). What you have here is fine for now, especially since it's just a performance optimization. Over in dask/dask we're also using pd.StringDtype("pyarrow") as it, historically, has been more feature complete than pd.ArrowDtype(pa.string()). That said, I think the situation has changed in pandas=2, so we may switch to pd.ArrowDtype(pa.string()) at some point in the future. This is mostly just an FYI in case we need to circle back to here in the future.

):
return pd.StringDtype("pyarrow")
return None

df = table.to_pandas(self_destruct=True, types_mapper=default_types_mapper)
return df.astype(meta.dtypes, copy=False)
Copy link
Member Author

@hendrikmakait hendrikmakait Jun 12, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

copy=False to make this a no-op for columns with matching dtypes.



def list_of_buffers_to_table(data: list[bytes]) -> pa.Table:
Expand Down