Repository: dask/dask
Pull request: Test round-tripping dataframe parquet I/O including pyspark
This is a proposed fix for #4096.  It adds a new `test_pyspark_compat` module to make sure that we can round-trip data from spark. I've opted to create a new CI workflow for this to prevent the additional weight and pain of including scala in our normal CI environments. It runs nightly on just the pyspark compatibility tests. There could be other ways to set this up, however, so I'm open to discussion on that point.

I'm trying to take the "Naively" from the linked issue seriously. The goal of these tests is to do as little trickery as possible, and look as close to user code as I can (with the caveat that I'm not very familiar with spark). So anything that involves any additional data transformation or non-default arguments is a failure (an there are a few here). However, I mostly *don't* care if the dataframe metadata are 100% identical after going through the round-trip process. So if we lose some information about the pandas index, or if the order of columns is different, that is probably okay. Instead, we should be testing that the data is faithfully round-tripped without loss or (excessive) coercion.

- [x] Closes #4096 
- [x] Tests added / passed
- [x] Passes `pre-commit run --all-files`

TODO
- [x] Figure out a workaround for timestamp-localization issues
- [ ] ~~Possibly trigger an issue if things fail (similar to the upstream builds)~~


Problem statement:
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

Hints from issue 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.

Commit message: Add test for round-tripping hive-partitioned datasets

Patch to review:
diff --git a/.github/workflows/spark.yml b/.github/workflows/spark.yml
new file mode 100644
index 00000000000..3bee3dde653
--- /dev/null
+++ b/.github/workflows/spark.yml
@@ -0,0 +1,40 @@
+name: PySpark compatibility tests
+
+on:
+  schedule:
+    - cron: "0 1 * * *"
+  push:
+  workflow_dispatch:
+
+jobs:
+  test:
+    runs-on: "ubuntu-latest"
+    strategy:
+      fail-fast: false
+    steps:
+      - name: Checkout source
+        uses: actions/checkout@v3
+      - name: Setup Java
+        uses: actions/setup-java@v3
+        with:
+          distribution: "zulu"
+          java-version: "11"
+      - name: Setup Conda Environment
+        uses: conda-incubator/setup-miniconda@v2
+        with:
+          miniforge-variant: Mambaforge
+          miniforge-version: latest
+          use-mamba: true
+          channel-priority: strict
+          python-version: 3.9 
+          environment-file: continuous_integration/environment-spark.yml
+          activate-environment: test-environment
+          auto-activate-base: false
+
+      - name: Install
+        shell: bash -l {0}
+        run: python -m pip install -e .[dataframe]
+
+      - name: Run tests
+        shell: bash -l {0}
+        run: pytest dask -m spark 
diff --git a/continuous_integration/environment-spark.yml b/continuous_integration/environment-spark.yml
new file mode 100644
index 00000000000..3f8031ce9a5
--- /dev/null
+++ b/continuous_integration/environment-spark.yml
@@ -0,0 +1,17 @@
+name: test-environment
+channels:
+  - conda-forge
+  - nodefaults 
+dependencies:
+  - pyspark
+  - python=3.9
+  - fastparquet
+  - numpy=1.21
+  - pandas=1.3
+  - pyarrow
+  - pytest
+  - pytest-rerunfailures
+  - pytest-timeout
+  - pip
+  - pip:
+    - git+https://github.com/dask/distributed
diff --git a/dask/tests/test_spark_compat.py b/dask/tests/test_spark_compat.py
new file mode 100644
index 00000000000..88c6ce05e15
--- /dev/null
+++ b/dask/tests/test_spark_compat.py
@@ -0,0 +1,83 @@
+import pytest
+
+from dask.datasets import timeseries
+
+dd = pytest.importorskip("dask.dataframe")
+pyspark = pytest.importorskip("pyspark")
+
+from dask.dataframe.utils import assert_eq
+
+pytestmark = pytest.mark.spark
+
+
+@pytest.fixture(scope="module")
+def spark_session():
+    spark = (
+        pyspark.sql.SparkSession.builder.master("local")
+        .appName("Dask Testing")
+        .getOrCreate()
+    )
+    yield spark
+
+    spark.stop()
+
+
+@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)
+    # TODO: pyspark auto-converts timezones -- round-tripping timestamps
+    # is tricky, so drop timeseries for now.
+    pdf = timeseries(freq="1H").compute().reset_index(drop=True)
+    sdf = spark_session.createDataFrame(pdf)
+    # not overwriting any data, but spark complains if the directory
+    # already exists and we don't set overwrite
+    sdf.repartition(npartitions).write.parquet(tmpdir, mode="overwrite")
+
+    # TODO: fastparquet requires the glob, pyarrow does not.
+    ddf = dd.read_parquet(tmpdir + "/**.parquet", engine=engine)
+    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)
+    # TODO: pyspark auto-converts timezones -- round-tripping timestamps
+    # is tricky, so drop timeseries for now.
+    pdf = timeseries(freq="1H").compute().reset_index(drop=True)
+    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")
+
+    # TODO: fastparquet requires the glob, pyarrow does not.
+    ddf = dd.read_parquet(tmpdir + "/**.parquet", engine=engine)
+
+    # 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, 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)
+    # TODO: pyspark auto-converts timezones -- round-tripping timestamps
+    # is tricky, so drop timeseries for now.
+    ddf = (
+        timeseries(freq="1H")
+        .repartition(npartitions=npartitions)
+        .reset_index(drop=True)
+    )
+    ddf.to_parquet(tmpdir, engine=engine, write_index=False)
+
+    sdf = spark_session.read.parquet(tmpdir)
+
+    assert_eq(sdf.toPandas(), ddf, check_index=False)
diff --git a/setup.cfg b/setup.cfg
index 071058da2e1..2a804faa082 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -46,6 +46,7 @@ markers:
   network: Test requires an internet connection
   slow: Skipped unless --runslow passed
   gpu: marks tests we want to run on GPUs
+  spark: marks tests for which we need a Spark environment
 addopts = -v -rsxfE --durations=10 --color=yes
 filterwarnings =
     # From Cython-1753


Code review comments:
diff_hunk: @@ -0,0 +1,83 @@
+import pytest
+
+from dask.datasets import timeseries
+
+dd = pytest.importorskip("dask.dataframe")
+pyspark = pytest.importorskip("pyspark")
+
+from dask.dataframe.utils import assert_eq
+
+pytestmark = pytest.mark.spark
+
+
+@pytest.fixture(scope="module")
+def spark_session():
+    spark = (
+        pyspark.sql.SparkSession.builder.master("local")
+        .appName("Dask Testing")
+        .getOrCreate()
+    )
+    yield spark
+
+    spark.stop()
+
+
+@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)
+    # TODO: pyspark auto-converts timezones -- round-tripping timestamps
+    # is tricky, so drop timeseries for now.
path: dask/tests/test_spark_compat.py
text: @author:
This is the biggest TODO in this PR from my perspective. Something in the pyspark stack is doing some timezone coercion that is resulting in deltas in datetime columns corresponding to one's UTC offset. This is a vector for real data loss.

@user1:
I think https://databricks.com/blog/2020/07/22/a-comprehensive-look-at-dates-and-timestamps-in-apache-spark-3-0.html covers the explanation.  

> When writing timestamp values out to non-text data sources like Parquet, the values are just instants (like timestamp in UTC) that have no time zone information

@author:
Thanks @user1, I'm not sure there's anything much to do here other than be aware of the gotcha. I resolved it here by ensuring that everything is in UTC