Add arrow schema extraction dispatch (#9169)

galipremsagar · web-flow · commit 6369cdb06813 · 2022-06-07T13:42:35.000-05:00
diff --git a/dask/dataframe/backends.py b/dask/dataframe/backends.py
@@ -30,6 +30,7 @@
     make_meta_dispatch,
     make_meta_obj,
     meta_nonempty,
+    pyarrow_schema_dispatch,
     tolist_dispatch,
     union_categoricals_dispatch,
 )
@@ -81,6 +82,13 @@ def _(x, index=None):
     pass
 
 
+@pyarrow_schema_dispatch.register((pd.DataFrame,))
+def get_pyarrow_schema_pandas(obj):
+    import pyarrow as pa
+
+    return pa.Schema.from_pandas(obj)
+
+
 @meta_nonempty.register(pd.DatetimeTZDtype)
 @make_meta_dispatch.register(pd.DatetimeTZDtype)
 def make_meta_pandas_datetime_tz(x, index=None):
diff --git a/dask/dataframe/dispatch.py b/dask/dataframe/dispatch.py
@@ -22,6 +22,7 @@
 is_categorical_dtype_dispatch = Dispatch("is_categorical_dtype")
 union_categoricals_dispatch = Dispatch("union_categoricals")
 grouper_dispatch = Dispatch("grouper")
+pyarrow_schema_dispatch = Dispatch("pyarrow_schema_dispatch")
 
 
 def concat(
diff --git a/dask/dataframe/io/parquet/arrow.py b/dask/dataframe/io/parquet/arrow.py
@@ -11,6 +11,7 @@
 
 from dask.base import tokenize
 from dask.core import flatten
+from dask.dataframe.backends import pyarrow_schema_dispatch
 from dask.dataframe.io.parquet.utils import (
     Engine,
     _get_aggregation_depth,
@@ -512,7 +513,7 @@ def initialize_write(
     ):
         if schema == "infer" or isinstance(schema, dict):
             # Start with schema from _meta_nonempty
-            inferred_schema = pa.Schema.from_pandas(
+            inferred_schema = pyarrow_schema_dispatch(
                 df._meta_nonempty.set_index(index_cols)
                 if index_cols
                 else df._meta_nonempty
diff --git a/dask/dataframe/io/tests/test_parquet.py b/dask/dataframe/io/tests/test_parquet.py
@@ -4177,3 +4177,26 @@ def test_deprecate_gather_statistics(tmp_path, engine):
             gather_statistics=True,
         )
     assert_eq(out, df)
+
+
+@pytest.mark.gpu
+def test_gpu_write_parquet_simple(tmpdir):
+    fn = str(tmpdir)
+    cudf = pytest.importorskip("cudf")
+    dask_cudf = pytest.importorskip("dask_cudf")
+    from dask.dataframe.dispatch import pyarrow_schema_dispatch
+
+    @pyarrow_schema_dispatch.register((cudf.DataFrame,))
+    def get_pyarrow_schema_cudf(obj):
+        return obj.to_arrow().schema
+
+    df = cudf.DataFrame(
+        {
+            "a": ["abc", "def"],
+            "b": ["a", "z"],
+        }
+    )
+    ddf = dask_cudf.from_cudf(df, 3)
+    ddf.to_parquet(fn)
+    got = dask_cudf.read_parquet(fn)
+    assert_eq(df, got)
