Repository: NVIDIA-Merlin/NVTabular

Issue:
[FEA] Write shared "_metadata" file for output parquet data in DaskWorkflow
<a href="https://github.com/rjzamora"><img src="https://avatars0.githubusercontent.com/u/20461013?v=4" align="left" width="96" height="96" hspace="10"></img></a> **Issue by [rjzamora](https://github.com/rjzamora)**
_Thursday May 28, 2020 at 16:43 GMT_
_Originally opened as https://github.com/rapidsai/recsys/issues/179_

----

**Is your feature request related to a problem? Please describe.**
The `DaskWorkflow` implemntation introduced in #136 does not currently write a shared parquet `"_metadata"` file when the `to_dataset` API is used.  This is because of the limitation discussed in [cudf#284](https://github.com/rapidsai/cudf/issues/5284).  Once that issue is resolved upstream, the `DaskWorkflow` code should be modified to write the shared file (having this file is a significant performance boost if/when dask is used to read the dataset back in).

**Describe the solution you'd like**
Once the upstream changes are made, metadata file creation will look very similar to [this existing cudf method](https://github.com/rapidsai/cudf/blob/56769d4755091e3fe892fcdd020bd1e2a9a19cfe/python/dask_cudf/dask_cudf/io/parquet.py#L130).

**Describe alternatives you've considered**
The alternative is to skip `"_metadata"` creation, or to create the shared file with pyarrow after the fact (requiring a followup python process to parse all the footer metadata after it is written).




Discussion:

Resolving patch:
diff --git a/nvtabular/io.py b/nvtabular/io.py
--- a/nvtabular/io.py
+++ b/nvtabular/io.py
@@ -39,7 +39,9 @@
 from dask.base import tokenize
 from dask.dataframe.core import new_dd_object
 from dask.dataframe.io.parquet.utils import _analyze_paths
+from dask.delayed import Delayed
 from dask.distributed import get_worker
+from dask.highlevelgraph import HighLevelGraph
 from dask.utils import natural_sort_key, parse_bytes
 from fsspec.core import get_fs_token_paths
 from fsspec.utils import stringify_path
@@ -612,17 +614,21 @@ def clean_pw_cache():
         return
 
 
+def close_cached_pw(fs):
+    md_dict = {}
+    with get_cache() as cache:
+        for path, (pw, bio) in cache.pq_writer_cache.items():
+            fn = bio.split(fs.sep)[-1]
+            md_dict[fn] = pw.close(metadata_file_path=fn)
+    return md_dict
+
+
 def guid():
     """ Simple utility function to get random hex string
     """
     return uuid4().hex
 
 
-def _write_metadata(meta_list):
-    # TODO: Write _metadata file here (need to collect metadata)
-    return meta_list
-
-
 @annotate("write_output_partition", color="green", domain="nvt_python")
 def _write_output_partition(gdf, processed_path, shuffle, out_files_per_proc, fs):
     gdf_size = len(gdf)
@@ -639,13 +645,42 @@ def _write_output_partition(gdf, processed_path, shuffle, out_files_per_proc, fs
             pw = cache.get_pq_writer(prefix, s, mem=(shuffle == "full"))
             pw.write_table(gdf.iloc[s * split_size : (s + 1) * split_size])
 
-    return gdf_size  # TODO: Make this metadata
+    return gdf_size
+
+
+def _to_parquet_dataset(ddf, fs, output_path, shuffle, out_files_per_proc):
+    name = "write-processed"
+    write_name = name + tokenize(ddf, shuffle, out_files_per_proc)
+    task_list = []
+    dsk = {}
+    for idx in range(ddf.npartitions):
+        key = (write_name, idx)
+        dsk[key] = (
+            _write_output_partition,
+            (ddf._name, idx),
+            output_path,
+            shuffle,
+            out_files_per_proc,
+            fs,
+        )
+        task_list.append(key)
+    dsk[name] = (lambda x: x, task_list)
+    graph = HighLevelGraph.from_collections(name, dsk, dependencies=[ddf])
+    return Delayed(name, graph)
+
+
+def _write_pq_metadata_file(md_list, fs, path):
+    if md_list:
+        metadata_path = fs.sep.join([path, "_metadata"])
+        _meta = cudf.io.merge_parquet_filemetadata(md_list) if len(md_list) > 1 else md_list[0]
+        with fs.open(metadata_path, "wb") as fil:
+            _meta.tofile(fil)
+    return
 
 
 @annotate("worker_shuffle", color="green", domain="nvt_python")
 def _worker_shuffle(processed_path, fs):
-    paths = []
-
+    metadata_dict = {}
     with get_cache() as cache:
         for path, (pw, bio) in cache.pq_writer_cache.items():
             pw.close()
@@ -656,9 +691,10 @@ def _worker_shuffle(processed_path, fs):
             gdf = _shuffle_gdf(gdf)
             rel_path = "shuffled.%s.parquet" % (guid())
             full_path = fs.sep.join([processed_path, rel_path])
-            gdf.to_parquet(full_path, compression=None, index=False)
-            paths.append(full_path)
-    return paths
+            metadata_dict[rel_path] = gdf.to_parquet(
+                full_path, compression=None, index=False, metadata_file_path=rel_path
+            )
+        return metadata_dict
 
 
 class Dataset:
diff --git a/nvtabular/workflow.py b/nvtabular/workflow.py
--- a/nvtabular/workflow.py
+++ b/nvtabular/workflow.py
@@ -22,9 +22,7 @@
 import dask
 import yaml
 from cudf._lib.nvtx import annotate
-from dask.base import tokenize
-from dask.delayed import Delayed
-from dask.highlevelgraph import HighLevelGraph
+from dask.utils import natural_sort_key
 from fsspec.core import get_fs_token_paths
 
 import nvtabular.io as nvt_io
@@ -833,48 +831,58 @@ def to_dataset(self, output_path, shuffle=None, out_files_per_proc=None):
         fs = get_fs_token_paths(output_path)[0]
         fs.mkdirs(output_path, exist_ok=True)
         if shuffle or out_files_per_proc:
-            name = "write-processed"
-            write_name = name + tokenize(ddf, shuffle, out_files_per_proc)
-            task_list = []
-            dsk = {}
-            for idx in range(ddf.npartitions):
-                key = (write_name, idx)
-                dsk[key] = (
-                    nvt_io._write_output_partition,
-                    (ddf._name, idx),
-                    output_path,
-                    shuffle,
-                    out_files_per_proc,
-                    fs,
-                )
-                task_list.append(key)
-            dsk[name] = (nvt_io._write_metadata, task_list)
-            graph = HighLevelGraph.from_collections(name, dsk, dependencies=[ddf])
-            out = Delayed(name, graph)
 
-            # Would also be nice to clean the categorical
+            # Construct graph for Dask-based dataset write
+            out = nvt_io._to_parquet_dataset(ddf, fs, output_path, shuffle, out_files_per_proc)
+
+            # Would be nice to clean the categorical
             # cache before the write (TODO)
 
-            # Trigger the Dask-based write and do a
-            # full (per-worker) shuffle if requested
+            # Trigger write execution
             if self.client:
                 self.client.cancel(self.ddf_base_dataset)
                 self.ddf_base_dataset = None
                 out = self.client.compute(out).result()
-                if shuffle == "full":
-                    self.client.cancel(self.ddf)
-                    self.ddf = None
-                    self.client.run(nvt_io._worker_shuffle, output_path, fs)
-                self.client.run(nvt_io.clean_pw_cache)
             else:
                 self.ddf_base_dataset = None
                 out = dask.compute(out, scheduler="synchronous")[0]
-                if shuffle == "full":
+
+            # Deal with "full" (per-worker) shuffle here.
+            if shuffle == "full":
+                if self.client:
+                    self.client.cancel(self.ddf)
                     self.ddf = None
-                    nvt_io._worker_shuffle(output_path, fs)
+                    worker_md = self.client.run(nvt_io._worker_shuffle, output_path, fs)
+                    worker_md = list(collections.ChainMap(*worker_md.values()).items())
+                else:
+                    self.ddf = None
+                    worker_md = nvt_io._worker_shuffle(output_path, fs)
+                    worker_md = list(worker_md.items())
+
+            else:
+                # Collect parquet metadata while closing
+                # ParquetWriter object(s)
+                if self.client:
+                    worker_md = self.client.run(nvt_io.close_cached_pw, fs)
+                    worker_md = list(collections.ChainMap(*worker_md.values()).items())
+                else:
+                    worker_md = nvt_io.close_cached_pw(fs)
+                    worker_md = list(worker_md.items())
+
+            # Sort metadata by file name and convert list of
+            # tuples to a list of metadata byte-blobs
+            md_list = [m[1] for m in sorted(worker_md, key=lambda x: natural_sort_key(x[0]))]
+
+            # Aggregate metadata and write _metadata file
+            nvt_io._write_pq_metadata_file(md_list, fs, output_path)
+
+            # Close ParquetWriter Objects
+            if self.client:
+                self.client.run(nvt_io.clean_pw_cache)
+            else:
                 nvt_io.clean_pw_cache()
 
-            return out
+            return
 
         # Default (shuffle=False): Just use dask_cudf.to_parquet
         fut = ddf.to_parquet(output_path, compression=None, write_index=False, compute=False)


Test patch:
diff --git a/tests/unit/test_workflow.py b/tests/unit/test_workflow.py
--- a/tests/unit/test_workflow.py
+++ b/tests/unit/test_workflow.py
@@ -266,24 +266,38 @@ def get_norms(tar: cudf.Series):
 
 
 @pytest.mark.parametrize("shuffle", ["full", "partial", None])
-def test_output_count(tmpdir, shuffle):
+@pytest.mark.parametrize("use_client", [True, False])
+def test_parquet_output(client, use_client, tmpdir, shuffle):
     out_files_per_proc = 2
-    out_path = os.path.join(tmpdir, "processed")
-    path = os.path.join(tmpdir, "simple.parquet")
+    n_workers = len(client.cluster.workers) if use_client else 1
+    out_path = str(tmpdir.mkdir("processed"))
+    path = str(tmpdir.join("simple.parquet"))
 
-    df = pd.DataFrame({"a": np.arange(25)})
-    df.to_parquet(path, row_group_size=5, engine="pyarrow")
+    size = 25
+    row_group_size = 5
+    df = pd.DataFrame({"a": np.arange(size)})
+    df.to_parquet(path, row_group_size=row_group_size, engine="pyarrow")
 
     columns = ["a"]
     dataset = nvt.Dataset(path, engine="parquet", row_groups_per_part=1)
-    processor = nvt.Workflow(cat_names=[], cont_names=columns, label_name=[])
-
+    processor = nvt.Workflow(
+        cat_names=[], cont_names=columns, label_name=[], client=client if use_client else None
+    )
     processor.add_preprocess(ops.Normalize())
     processor.finalize()
-
     processor.update_stats(
         dataset, output_path=out_path, shuffle=shuffle, out_files_per_proc=out_files_per_proc
     )
-    result = glob.glob(os.path.join(out_path, "*"))
 
-    assert len(result) == out_files_per_proc
+    # Check that the number of output files is correct
+    result = glob.glob(os.path.join(out_path, "*.parquet"))
+    assert len(result) == out_files_per_proc * n_workers
+
+    # Make sure _metadata exists
+    meta_path = os.path.join(out_path, "_metadata")
+    assert os.path.exists(meta_path)
+
+    # Make sure _metadata makes sense
+    _metadata = cudf.io.read_parquet_metadata(meta_path)
+    assert _metadata[0] == size
+    assert _metadata[2] == columns