Enable basic p2p shuffle for dask-cudf

[rjzamora]

These changes are meant to enable basic "p2p" functionality for dask_cudf.DataFrame collections by addressing some minor inconsistencies between pandas and cudf for DataFrame<->pa.Table conversion. This PR does not attempt to optimize the the p2p shuffle algorithm for GPU-backed data.

TODO (in follow-up work):

• Support null-index round-tripping?
• Clean up type hints?

[rjzamora]

@wence- @quasiben - I know we want to implement a more-efficient version of p2p for GPU-backed data, but I'm thinking basic compatibility is worth pursuing in the short-to-medium term. Let me know what you guys think.

[github-actions]

Unit Test Results

See test report for an extended history of previous test failures. This is useful for diagnosing flaky tests.

       21 files  ±0         21 suites  ±0   11h 25m 49s ⏱️ -52s
  3 771 tests +1    3 655 ✔️  - 6     108 💤 +1  8 ❌ +6 
36 465 runs  +8  34 648 ✔️  - 5  1 809 💤 +7  8 ❌ +6 

For more details on these failures, see this check.

Results for commit 50beefc. ± Comparison against base commit 3082d27.

♻️ This comment has been updated with latest results.

[fjetter]

I think most of the CI failures are not related and should've been fixed on main already.

[fjetter]

Do you want us to have a closer look? Is this intended to be merged at some point or is this just for show and tell?

[rjzamora]

Do you want us to have a closer look?

I don't think this is ready for a detailed review. However, I would welcome thoughts at any level. The immediate goal of this PR is to figure out a simple way to get the "p2p" shuffle working for cudf-backed data. That is, performance can be a secondary concern for now. Therefore, it would be good to know your thoughts on ways to relax the pandas-specific logic currently used in "p2p".

Is this intended to be merged at some point or is this just for show and tell?

This started off as an experiment, but it would be extremely useful on the RAPIDs side to not have to worry about adding workarounds to avoid the "p2p" shuffle default.

[quasiben]

Would probably be good to add a test for this

[rjzamora]

The p2p code is good about using type hints. Unfortunately, it uses pd.DataFrame in many places where we want something like pd.DataFrame | cudf.DataFrame. The obvious problem is that I cannot just add the | cudf.DataFrame, because there is no guarantee that cudf is installed (and we don't want to eagerly import it, even if it is installed).

What is the established type-hinting convention for cases like this? Do we need something like a run-time checkable BackendDataFrame protocol? cc @jrbourbeau

[wence-]

You could do this:

from typing import TYPE_CHECKING, Union
from typing_extensions import TypeAlias
import pandas as pd

DataFrameT: TypeAlias = Union[pd.DataFrame, "cudf.DataFrame"]

if TYPE_CHECKING:
    import cudf

I think?

[wence-]

Ah, but I guess mypy still complains because it can't find the import at type-checking time.

[wence-]

I suppose you could also make this generic and say:

T = TypeVar("T")
def _arrow_to_df(table: pa.Table, like: T) -> T:
   ...

[wence-]

Do we need something like a run-time checkable BackendDataFrame protocol?

The static-typing approach to this in python is PEP 544

[rjzamora]

The static-typing approach to this in python is PEP 544

Right exactly. This is why I asked if we should add a BackendDataFrame protocol. Perhaps it makes sense to (eventually) replace is_dataframe_like with a DataFrameLike protocol in dask/dask?

[jorisvandenbossche]

FWIW, this would also be useful for dask-geopandas, to allow letting us have control over how a geopandas.GeoDataFrame gets converted to a pyarrow table and back (xref geopandas/dask-geopandas#256)

[rjzamora]

FWIW, this would also be useful for dask-geopandas, to allow letting us have control over how a geopandas.GeoDataFrame gets converted to a pyarrow table and back (xref geopandas/dask-geopandas#256)

Dask main now offers to/from arrow-Table dispatch functions. For example, here is the to_pyarrow_table_dispatch implementation in dask_cudf. With that said, I'd still expect the supported kwargs to change a bit.

[hendrikmakait]

Thanks, @rjzamora!

[rjzamora]

Thanks for the review @hendrikmakait !

Let me know if there is anything left to do to get this merged (I don't think any of the CI failures are related to the changes in any way).

[hendrikmakait]

Ready to merge from my perspective, when I reviewed it, this still had a [WIP] flag.

[rjzamora]

Ready to merge from my perspective, when I reviewed it, this still had a [WIP] flag.

Sorry, forgot about the "[WIP]" label - makes sense. Note that I don't have write permissions in distributed :)

[hendrikmakait]

Note that I don't have write permissions in distributed

I was not aware of that, thanks for pointing it out!