Repository: dask/dask

Issue:
Switch default pyarrow-based read_parquet engine to "pyarrow-dataset" 
When the support for `pyarrow.dataset` was added to `read_parquet` (`ArrowDatasetEngine`, #6174, #6534), we didn't yet turn it on by default (when using pyarrow as engine). 

In the upcoming pyarrow 5.0.0, I am adding more deprecation warnings to the use of several `pyarrow.parquet.ParquetDataset` functionalities (eg the `pieces` and `partitions` attributes). So I think it would be good if dask could start to actually use "pyarrow-dataset" as the default (meaning that `engine="pyarrow"` would map to "pyarrow-dataset" instead of "pyarrow-legacy").  cc @rjzamora 

Are there still known issues with "pyarrow-dataset" compared to "pyarrow-legacy"?


Discussion:
Thanks for making a formal issue about this @jorisvandenbossche !

>Are there still known issues with "pyarrow-dataset" compared to "pyarrow-legacy"?

I will need to re-run some benchmarks.  There may be some cases where the legacy engine is still a bit faster, but I wouldn't want to delay the move to `ArrowDatasetEngine` unless the difference is significant.  I do happen to  know that Dask-cuDF is still on the legacy engine, but I can submit a PR to change that in the next week or so.

Resolving patch:
diff --git a/dask/dataframe/io/parquet/arrow.py b/dask/dataframe/io/parquet/arrow.py
--- a/dask/dataframe/io/parquet/arrow.py
+++ b/dask/dataframe/io/parquet/arrow.py
@@ -1703,6 +1703,32 @@ def _parquet_piece_as_arrow(
         )
         return arrow_table
 
+    @classmethod
+    def collect_file_metadata(cls, path, fs, file_path):
+        with fs.open(path, "rb") as f:
+            meta = pq.ParquetFile(f).metadata
+        if file_path:
+            meta.set_file_path(file_path)
+        return meta
+
+    @classmethod
+    def aggregate_metadata(cls, meta_list, fs, out_path):
+        meta = None
+        for _meta in meta_list:
+            if meta:
+                _append_row_groups(meta, _meta)
+            else:
+                meta = _meta
+        if out_path:
+            metadata_path = fs.sep.join([out_path, "_metadata"])
+            with fs.open(metadata_path, "wb") as fil:
+                if not meta:
+                    raise ValueError("Cannot write empty metdata!")
+                meta.write_metadata_file(fil)
+            return None
+        else:
+            return meta
+
 
 #
 #  PyArrow Legacy API [PyArrow<1.0.0]
@@ -2135,32 +2161,6 @@ def _read_table(
             **kwargs,
         )
 
-    @classmethod
-    def collect_file_metadata(cls, path, fs, file_path):
-        with fs.open(path, "rb") as f:
-            meta = pq.ParquetFile(f).metadata
-        if file_path:
-            meta.set_file_path(file_path)
-        return meta
-
-    @classmethod
-    def aggregate_metadata(cls, meta_list, fs, out_path):
-        meta = None
-        for _meta in meta_list:
-            if meta:
-                _append_row_groups(meta, _meta)
-            else:
-                meta = _meta
-        if out_path:
-            metadata_path = fs.sep.join([out_path, "_metadata"])
-            with fs.open(metadata_path, "wb") as fil:
-                if not meta:
-                    raise ValueError("Cannot write empty metdata!")
-                meta.write_metadata_file(fil)
-            return None
-        else:
-            return meta
-
     @classmethod
     def multi_support(cls):
         return cls == ArrowLegacyEngine
diff --git a/dask/dataframe/io/parquet/core.py b/dask/dataframe/io/parquet/core.py
--- a/dask/dataframe/io/parquet/core.py
+++ b/dask/dataframe/io/parquet/core.py
@@ -5,6 +5,7 @@
 from fsspec.core import get_fs_token_paths
 from fsspec.implementations.local import LocalFileSystem
 from fsspec.utils import stringify_path
+from packaging.version import parse as parse_version
 
 from ....base import tokenize
 from ....delayed import Delayed
@@ -164,15 +165,12 @@ def read_parquet(
         Parquet reader library to use. Options include: 'auto', 'fastparquet',
         'pyarrow', 'pyarrow-dataset', and 'pyarrow-legacy'. Defaults to 'auto',
         which selects the FastParquetEngine if fastparquet is installed (and
-        ArrowLegacyEngine otherwise).  If 'pyarrow-dataset' is specified, the
-        ArrowDatasetEngine (which leverages the pyarrow.dataset API) will be used
-        for newer PyArrow versions (>=1.0.0). If 'pyarrow' or 'pyarrow-legacy' are
-        specified, the ArrowLegacyEngine will be used (which leverages the
-        pyarrow.parquet.ParquetDataset API).
-        NOTE: 'pyarrow-dataset' enables row-wise filtering, but requires
-        pyarrow>=1.0. The behavior of 'pyarrow' will most likely change to
-        ArrowDatasetEngine in a future release, and the 'pyarrow-legacy'
-        option will be deprecated once the ParquetDataset API is deprecated.
+        ArrowDatasetEngine otherwise).  If 'pyarrow' or 'pyarrow-dataset' is
+        specified, the ArrowDatasetEngine (which leverages the pyarrow.dataset
+        API) will be used. If 'pyarrow-legacy' is specified, ArrowLegacyEngine
+        will be used (which leverages the pyarrow.parquet.ParquetDataset API).
+        NOTE: The 'pyarrow-legacy' option (ArrowLegacyEngine) is deprecated
+        for pyarrow>=5.
     gather_statistics : bool, default None
         Gather the statistics for each dataset partition. By default,
         this will only be done if the _metadata file is available. Otherwise,
@@ -911,6 +909,19 @@ def get_engine(engine):
         return eng
 
     elif engine in ("pyarrow", "arrow", "pyarrow-legacy", "pyarrow-dataset"):
+
+        pa = import_required("pyarrow", "`pyarrow` not installed")
+        pa_version = parse_version(pa.__version__)
+
+        if engine in ("pyarrow", "arrow"):
+            engine = "pyarrow-dataset"
+        elif pa_version.major >= 5 and engine == "pyarrow-legacy":
+            warnings.warn(
+                "`ArrowLegacyEngine` ('pyarrow-legacy') is deprecated "
+                "for pyarrow>=5. Please use `engine='pyarrow'`.",
+                FutureWarning,
+            )
+
         if engine == "pyarrow-dataset":
             from .arrow import ArrowDatasetEngine
 


Test patch:
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
@@ -59,8 +59,16 @@
     SKIP_PYARROW_REASON = "pyarrow not found"
 PYARROW_MARK = pytest.mark.skipif(SKIP_PYARROW, reason=SKIP_PYARROW_REASON)
 
+# "Legacy" and "Dataset"-specific MARK definitions
+SKIP_PYARROW_LE = SKIP_PYARROW
+SKIP_PYARROW_LE_REASON = "pyarrow not found"
 SKIP_PYARROW_DS = SKIP_PYARROW
 SKIP_PYARROW_DS_REASON = "pyarrow not found"
+if pa and pa_version.major >= 5:
+    SKIP_PYARROW_LE = True
+    SKIP_PYARROW_LE_REASON = "pyarrow < 5.0.0 required for pyarrow legacy API"
+PYARROW_LE_MARK = pytest.mark.skipif(SKIP_PYARROW_LE, reason=SKIP_PYARROW_LE_REASON)
+
 PYARROW_DS_MARK = pytest.mark.skipif(SKIP_PYARROW_DS, reason=SKIP_PYARROW_DS_REASON)
 
 ANY_ENGINE_MARK = pytest.mark.skipif(
@@ -85,7 +93,7 @@
 @pytest.fixture(
     params=[
         pytest.param("fastparquet", marks=FASTPARQUET_MARK),
-        pytest.param("pyarrow-legacy", marks=PYARROW_MARK),
+        pytest.param("pyarrow-legacy", marks=PYARROW_LE_MARK),
         pytest.param("pyarrow-dataset", marks=PYARROW_DS_MARK),
     ]
 )
@@ -103,7 +111,7 @@ def write_read_engines(**kwargs):
     # Skip if uninstalled
     skip_marks = {
         "fastparquet": FASTPARQUET_MARK,
-        "pyarrow-legacy": PYARROW_MARK,
+        "pyarrow-legacy": PYARROW_LE_MARK,
         "pyarrow-dataset": PYARROW_DS_MARK,
     }
     marks = {(w, r): [skip_marks[w], skip_marks[r]] for w in backends for r in backends}
@@ -160,6 +168,21 @@ def write_read_engines(**kwargs):
     fp_pandas_xfail = write_read_engines()
 
 
+@PYARROW_MARK
+def test_pyarrow_getengine():
+    from dask.dataframe.io.parquet.arrow import ArrowDatasetEngine
+    from dask.dataframe.io.parquet.core import get_engine
+
+    # Check that the default engine for "pyarrow"/"arrow"
+    # is the `pyarrow.dataset`-based engine
+    assert get_engine("pyarrow") == ArrowDatasetEngine
+    assert get_engine("arrow") == ArrowDatasetEngine
+
+    if SKIP_PYARROW_LE:
+        with pytest.raises(FutureWarning):
+            get_engine("pyarrow-legacy")
+
+
 @write_read_engines()
 def test_local(tmpdir, write_engine, read_engine):
     tmp = str(tmpdir)
@@ -1044,7 +1067,7 @@ def test_to_parquet_default_writes_nulls(tmpdir):
     assert table[1].null_count == 2
 
 
-@PYARROW_MARK
+@PYARROW_LE_MARK
 def test_to_parquet_pyarrow_w_inconsistent_schema_by_partition_fails_by_default(tmpdir):
     df = pd.DataFrame(
         {"partition_column": [0, 0, 1, 1], "strings": ["a", "b", None, None]}
@@ -2285,27 +2308,11 @@ def test_timeseries_nulls_in_schema(tmpdir, engine, schema):
 
     # Note: `append_row_groups` will fail with pyarrow>0.17.1 for _metadata write
     dataset = {"validate_schema": False}
-    engine_use = engine
-    if schema != "infer" and engine.startswith("pyarrow"):
-        engine_use = "pyarrow-legacy"
-    ddf2.to_parquet(
-        tmp_path, engine=engine_use, write_metadata_file=False, schema=schema
-    )
-    ddf_read = dd.read_parquet(tmp_path, engine=engine_use, dataset=dataset)
+    ddf2.to_parquet(tmp_path, engine=engine, write_metadata_file=False, schema=schema)
+    ddf_read = dd.read_parquet(tmp_path, engine=engine, dataset=dataset)
 
     assert_eq(ddf_read, ddf2, check_divisions=False, check_index=False)
 
-    # Can force schema validation on each partition in pyarrow
-    if engine.startswith("pyarrow") and schema is None:
-        dataset = {"validate_schema": True}
-        engine_use = engine
-        if schema != "infer":
-            engine_use = "pyarrow-legacy"
-        # The schema mismatch should raise an error if the
-        # dataset was written with `schema=None` (no inference)
-        with pytest.raises(ValueError):
-            ddf_read = dd.read_parquet(tmp_path, dataset=dataset, engine=engine_use)
-
 
 @PYARROW_MARK
 @pytest.mark.parametrize("numerical", [True, False])
@@ -2964,10 +2971,13 @@ def test_pandas_timestamp_overflow_pyarrow(tmpdir):
 
     # This will raise by default due to overflow
     with pytest.raises(pa.lib.ArrowInvalid) as e:
-        dd.read_parquet(str(tmpdir), engine="pyarrow-legacy").compute()
+        dd.read_parquet(str(tmpdir), engine="pyarrow").compute()
     assert "out of bounds" in str(e.value)
 
-    from dask.dataframe.io.parquet.arrow import ArrowEngine
+    if pa_version >= parse_version("5.0.0"):
+        from dask.dataframe.io.parquet.arrow import ArrowDatasetEngine as ArrowEngine
+    else:
+        from dask.dataframe.io.parquet.arrow import ArrowEngine
 
     class ArrowEngineWithTimestampClamp(ArrowEngine):
         @classmethod
@@ -3164,8 +3174,8 @@ def test_pyarrow_dataset_simple(tmpdir, engine):
     df.set_index("a", inplace=True, drop=True)
     ddf = dd.from_pandas(df, npartitions=2)
     ddf.to_parquet(fn, engine=engine)
-    read_df = dd.read_parquet(fn, engine="pyarrow-legacy")
-    read_df.compute(scheduler="synchronous")
+    read_df = dd.read_parquet(fn, engine="pyarrow-dataset")
+    read_df.compute()
     assert_eq(ddf, read_df)