Repository: dask/dask

Issue:
Naively roundtrip parquet data from Spark
Currently it is not easy to roundtrip data between Spark and Dask Dataframe with Parquet.  There are a variety of details one needs to know to do this well.  This would be a good case study to improve our usability.

cc @martindurant 


Discussion:
In particular I'm thinking that the following should probably just work for a variety of data types.

```python
spark_df.write.parquet(fn)
dask_df = dd.read_parquet(fn)

dask_df.to_parquet(fn2)
spark.read.parquet(fn2)
```
Fastparquet does have a number of tests specifically for spark round-tripping: https://github.com/dask/fastparquet/blob/master/fastparquet/test/test_aroundtrips.py
(these are without dask's part, of course)
Understood.  To be clear, I'm not speaking about fastparquet, I'm talking about dask.dataframe.  My attempt to recreate the workflow above failed with dask dataframe.  I think that this is a fairly important issue.
Yes, agree, just thought I'd point it out, the FP tests might serve as a starting point. Would be easy to write out a set of tests like this (not for CI, though, probably), and then fix for FP and Arrow, as needed.
If you have time in the near future I think that you would be the obviously
skilled person to fix this :)

On Sat, Oct 13, 2018 at 2:21 PM Martin Durant <notifications@github.com>
wrote:

> Yes, agree, just thought I'd point it out, the FP tests might serve as a
> starting point. Would be easy to write out a set of tests like this (not
> for CI, though, probably), and then fix for FP and Arrow, as needed.
>
> —
> You are receiving this because you authored the thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/4096#issuecomment-429564031>, or mute
> the thread
> <https://github.com/notifications/unsubscribe-auth/AASszE5U_Z8s4P7V5ps4JWVHePSdA1jJks5uki8agaJpZM4Xaqs5>
> .
>

It's on the list...
We also have a `flavor='spark'` argument in `pyarrow` which handles some of the limitations that the Spark Parquet implementation has. Sadly there are still some features of Parquet that `pyarrow` uses but Spark cannot handle.
I suspect that Martin feels your frustration :)

On Tue, Oct 16, 2018 at 5:23 AM Uwe L. Korn <notifications@github.com>
wrote:

> We also have a flavor='spark' argument in pyarrow which handles some of
> the limitations that the Spark Parquet implementation has. Sadly there are
> still some features of Parquet that pyarrow uses but Spark cannot handle.
>
> —
> You are receiving this because you authored the thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/4096#issuecomment-430166110>, or mute
> the thread
> <https://github.com/notifications/unsubscribe-auth/AASszLL-47f3N9EsCoLwfbBj7YcjzVK-ks5ulaWVgaJpZM4Xaqs5>
> .
>

Is this still of interest?
Could be contemplated in the context of #4336 , but interop with arrow more important now that spark per se (plus pyspark's inherent limitations) - so maybe the issue of rountripping to spark should be something that they care about more than we do. In any case, I don't think there are any plans to work on this. I am ambivalent on whether to close.
I care about this.  I think that it is important for the Dask project to be
able to easily move data back and forth.  I agree with @martindurant that
this might be partially solved by doing work in the Arrow codebase.  I'd
still like to keep track of it here.

On Tue, Apr 30, 2019 at 12:08 PM Martin Durant <notifications@github.com>
wrote:

> Could be contemplated in the context of #4336
> <https://github.com/dask/dask/pull/4336> , but interop with arrow more
> important now that spark per se (plus pyspark's inherent limitations) - so
> maybe the issue of rountripping to spark should be something that they care
> about more than we do. In any case, I don't think there are any plans to
> work on this. I am ambivalent on whether to close.
>
> —
> You are receiving this because you authored the thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/4096#issuecomment-488035056>, or mute
> the thread
> <https://github.com/notifications/unsubscribe-auth/AACKZTHLIZMSR5OOOIY6TOLPTB4HVANCNFSM4F3KVM4Q>
> .
>


Resolving patch:
diff --git a/dask/dataframe/io/parquet/fastparquet.py b/dask/dataframe/io/parquet/fastparquet.py
--- a/dask/dataframe/io/parquet/fastparquet.py
+++ b/dask/dataframe/io/parquet/fastparquet.py
@@ -429,6 +429,7 @@ def _collect_dataset_info(
                     paths = [
                         path for path in paths if path.endswith(parquet_file_extension)
                     ]
+                    fns = [fn for fn in fns if fn.endswith(parquet_file_extension)]
                     if len0 and paths == []:
                         raise ValueError(
                             "No files satisfy the `parquet_file_extension` criteria "
@@ -535,13 +536,11 @@ def _create_dd_meta(cls, dataset_info):
             ) = _parse_pandas_metadata(pandas_md)
             #  auto-ranges should not be created by fastparquet
             column_names.extend(pf.cats)
-
         else:
             index_names = []
             column_names = pf.columns + list(pf.cats)
             storage_name_mapping = {k: k for k in column_names}
             column_index_names = [None]
-
         if index is None and len(index_names) > 0:
             if len(index_names) == 1 and index_names[0] is not None:
                 index = index_names[0]


Test patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -35,7 +35,12 @@ jobs:
         uses: actions/checkout@v2
         with:
           fetch-depth: 0  # Needed by codecov.io
-
+      - name: Setup Java
+        uses: actions/setup-java@v3
+        if: ${{ matrix.os == 'ubuntu-latest' }}
+        with:
+          distribution: "zulu"
+          java-version: "11"
       - name: Setup Conda Environment
         uses: conda-incubator/setup-miniconda@v2
         with:
diff --git a/dask/dataframe/io/tests/test_parquet.py b/dask/dataframe/io/tests/test_parquet.py
--- a/dask/dataframe/io/tests/test_parquet.py
+++ b/dask/dataframe/io/tests/test_parquet.py
@@ -3832,18 +3832,27 @@ def test_metadata_task_size(tmpdir, engine, write_metadata_file, metadata_task_s
     assert_eq(ddf2b, ddf2c)
 
 
-def test_extra_file(tmpdir, engine):
+@pytest.mark.parametrize("partition_on", ("b", None))
+def test_extra_file(tmpdir, engine, partition_on):
     # Check that read_parquet can handle spark output
     # See: https://github.com/dask/dask/issues/8087
     tmpdir = str(tmpdir)
     df = pd.DataFrame({"a": range(100), "b": ["dog", "cat"] * 50})
+    df = df.assign(b=df.b.astype("category"))
     ddf = dd.from_pandas(df, npartitions=2)
-    ddf.to_parquet(tmpdir, engine=engine, write_metadata_file=True)
+    ddf.to_parquet(
+        tmpdir,
+        engine=engine,
+        write_metadata_file=True,
+        partition_on=partition_on,
+    )
     open(os.path.join(tmpdir, "_SUCCESS"), "w").close()
     open(os.path.join(tmpdir, "part.0.parquet.crc"), "w").close()
     os.remove(os.path.join(tmpdir, "_metadata"))
     out = dd.read_parquet(tmpdir, engine=engine, calculate_divisions=True)
-    assert_eq(out, df)
+    # Weird two-step since that we don't care if category ordering changes
+    assert_eq(out, df, check_categorical=False)
+    assert_eq(out.b, df.b, check_category_order=False)
 
     # For "fastparquet" and "pyarrow", we can pass the
     # expected file extension, or avoid checking file extensions
@@ -3865,7 +3874,9 @@ def _parquet_file_extension(val, legacy=False):
         **_parquet_file_extension(".parquet"),
         calculate_divisions=True,
     )
-    assert_eq(out, df)
+    # Weird two-step since that we don't care if category ordering changes
+    assert_eq(out, df, check_categorical=False)
+    assert_eq(out.b, df.b, check_category_order=False)
 
     # Should Work (with FutureWarning)
     with pytest.warns(FutureWarning, match="require_extension is deprecated"):
@@ -3875,7 +3886,6 @@ def _parquet_file_extension(val, legacy=False):
             **_parquet_file_extension(".parquet", legacy=True),
             calculate_divisions=True,
         )
-        assert_eq(out, df)
 
     # Should Fail (for not capturing the _SUCCESS and crc files)
     with pytest.raises((OSError, pa.lib.ArrowInvalid)):
diff --git a/dask/tests/test_spark_compat.py b/dask/tests/test_spark_compat.py
new file mode 100644
--- /dev/null
+++ b/dask/tests/test_spark_compat.py
@@ -0,0 +1,109 @@
+import signal
+import sys
+import threading
+
+import pytest
+
+from dask.datasets import timeseries
+
+dd = pytest.importorskip("dask.dataframe")
+pyspark = pytest.importorskip("pyspark")
+pytest.importorskip("pyarrow")
+pytest.importorskip("fastparquet")
+
+from dask.dataframe.utils import assert_eq
+
+if not sys.platform.startswith("linux"):
+    pytest.skip(
+        "Unnecessary, and hard to get spark working on non-linux platforms",
+        allow_module_level=True,
+    )
+
+# pyspark auto-converts timezones -- round-tripping timestamps is easier if
+# we set everything to UTC.
+pdf = timeseries(freq="1H").compute()
+pdf.index = pdf.index.tz_localize("UTC")
+pdf = pdf.reset_index()
+
+
+@pytest.fixture(scope="module")
+def spark_session():
+    # Spark registers a global signal handler that can cause problems elsewhere
+    # in the test suite. In particular, the handler fails if the spark session
+    # is stopped (a bug in pyspark).
+    prev = signal.getsignal(signal.SIGINT)
+    # Create a spark session. Note that we set the timezone to UTC to avoid
+    # conversion to local time when reading parquet files.
+    spark = (
+        pyspark.sql.SparkSession.builder.master("local")
+        .appName("Dask Testing")
+        .config("spark.sql.session.timeZone", "UTC")
+        .getOrCreate()
+    )
+    yield spark
+
+    spark.stop()
+    # Make sure we get rid of the signal once we leave stop the session.
+    if threading.current_thread() is threading.main_thread():
+        signal.signal(signal.SIGINT, prev)
+
+
+@pytest.mark.parametrize("npartitions", (1, 5, 10))
+@pytest.mark.parametrize("engine", ("pyarrow", "fastparquet"))
+def test_roundtrip_parquet_spark_to_dask(spark_session, npartitions, tmpdir, engine):
+    tmpdir = str(tmpdir)
+
+    sdf = spark_session.createDataFrame(pdf)
+    # We are not overwriting any data, but spark complains if the directory
+    # already exists (as tmpdir does) and we don't set overwrite
+    sdf.repartition(npartitions).write.parquet(tmpdir, mode="overwrite")
+
+    ddf = dd.read_parquet(tmpdir, engine=engine)
+    # Papercut: pandas TZ localization doesn't survive roundtrip
+    ddf = ddf.assign(timestamp=ddf.timestamp.dt.tz_localize("UTC"))
+    assert ddf.npartitions == npartitions
+
+    assert_eq(ddf, pdf, check_index=False)
+
+
+@pytest.mark.parametrize("engine", ("pyarrow", "fastparquet"))
+def test_roundtrip_hive_parquet_spark_to_dask(spark_session, tmpdir, engine):
+    tmpdir = str(tmpdir)
+
+    sdf = spark_session.createDataFrame(pdf)
+    # not overwriting any data, but spark complains if the directory
+    # already exists and we don't set overwrite
+    sdf.write.parquet(tmpdir, mode="overwrite", partitionBy="name")
+
+    ddf = dd.read_parquet(tmpdir, engine=engine)
+    # Papercut: pandas TZ localization doesn't survive roundtrip
+    ddf = ddf.assign(timestamp=ddf.timestamp.dt.tz_localize("UTC"))
+
+    # Partitioning can change the column order. This is mostly okay,
+    # but we sort them here to ease comparison
+    ddf = ddf.compute().sort_index(axis=1)
+    # Dask automatically converts hive-partitioned columns to categories.
+    # This is fine, but convert back to strings for comparison.
+    ddf = ddf.assign(name=ddf.name.astype("str"))
+
+    assert_eq(ddf, pdf.sort_index(axis=1), check_index=False)
+
+
+@pytest.mark.parametrize("npartitions", (1, 5, 10))
+@pytest.mark.parametrize("engine", ("pyarrow", "fastparquet"))
+def test_roundtrip_parquet_dask_to_spark(spark_session, npartitions, tmpdir, engine):
+    tmpdir = str(tmpdir)
+    ddf = dd.from_pandas(pdf, npartitions=npartitions)
+
+    # Papercut: https://github.com/dask/fastparquet/issues/646#issuecomment-885614324
+    kwargs = {"times": "int96"} if engine == "fastparquet" else {}
+
+    ddf.to_parquet(tmpdir, engine=engine, write_index=False, **kwargs)
+
+    sdf = spark_session.read.parquet(tmpdir)
+    sdf = sdf.toPandas()
+
+    # Papercut: pandas TZ localization doesn't survive roundtrip
+    sdf = sdf.assign(timestamp=sdf.timestamp.dt.tz_localize("UTC"))
+
+    assert_eq(sdf, ddf, check_index=False)