Repository: NVIDIA-Merlin/NVTabular

Issue:
[Task] Write global parquet metadata file 
<a href="https://github.com/benfred"><img src="https://avatars2.githubusercontent.com/u/69536?v=4" align="left" width="96" height="96" hspace="10"></img></a> **Issue by [benfred](https://github.com/benfred)**
_Friday May 22, 2020 at 19:52 GMT_
_Originally opened as https://github.com/rapidsai/recsys/issues/168_

----

We should write out a global metadata file when exporting a parquet dataset. 



Discussion:
<a href="https://github.com/benfred"><img src="https://avatars2.githubusercontent.com/u/69536?v=4" align="left" width="48" height="48" hspace="10"></img></a> **Comment by [benfred](https://github.com/benfred)**
_Wednesday May 27, 2020 at 00:04 GMT_

----

See also https://github.com/rapidsai/cudf/issues/5284


Resolving patch:
diff --git a/nvtabular/ds_writer.py b/nvtabular/ds_writer.py
deleted file mode 100644
--- a/nvtabular/ds_writer.py
+++ /dev/null
@@ -1,120 +0,0 @@
-#
-# Copyright (c) 2020, NVIDIA CORPORATION.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-
-import glob
-import os
-
-import cudf
-import cupy as cp
-import numpy as np
-import pyarrow.parquet as pq
-
-
-class FileIterator:
-    def __init__(self, path, nfiles, shuffle=True, **kwargs):
-        self.path = path
-        self.nfiles = nfiles
-        self.shuffle = shuffle
-        self.ind = 0
-        self.inds = np.arange(self.nfiles)
-        if self.shuffle:
-            np.random.shuffle(self.inds)
-
-    def __iter__(self):
-        self.ind = 0
-        self.inds = np.arange(self.nfiles)
-        if self.shuffle:
-            np.random.shuffle(self.inds)
-        return self
-
-    def __next__(self):
-        if self.ind >= self.nfiles:
-            raise StopIteration
-        self.ind += 1
-        # if self.name, return that naming convention.
-        return "%s/ds_part.%d.parquet" % (self.path, self.ind - 1)
-
-
-class DatasetWriter:
-    def __init__(self, path, nfiles=1, **kwargs):
-        self.path = path
-        self.nfiles = nfiles
-        self.writers = {fn: None for fn in FileIterator(path, nfiles)}
-        self.shared_meta_path = str(path) + "/_metadata"
-        self.metadata = None
-        self.new_metadata = {fn: [] for fn in FileIterator(path, nfiles)}
-
-        # Check for _metadata
-        metafile = glob.glob(self.shared_meta_path)
-        if metafile:
-            self.metadata = pq.ParquetDataset(metafile[0]).metadata
-
-    def write(self, gdf, shuffle=True):
-
-        # Shuffle the dataframe
-        gdf_size = len(gdf)
-        if shuffle:
-            sort_key = "__sort_index__"
-            arr = cp.arange(gdf_size)
-            cp.random.shuffle(arr)
-            gdf[sort_key] = cudf.Series(arr)
-            gdf = gdf.sort_values(sort_key).drop(columns=[sort_key])
-
-        # Write to
-        chunk_size = int(gdf_size / self.nfiles)
-        for i, fn in enumerate(FileIterator(self.path, self.nfiles)):
-            s1 = i * chunk_size
-            s2 = (i + 1) * chunk_size
-            if i == (self.nfiles - 1):
-                s2 = gdf_size
-            chunk = gdf[s1:s2]
-            pa_table = chunk.to_arrow()
-            if self.writers[fn] is None:
-                self.writers[fn] = pq.ParquetWriter(
-                    fn, pa_table.schema, metadata_collector=self.new_metadata[fn],
-                )
-            self.writers[fn].write_table(pa_table)
-
-    def write_metadata(self):
-        self.close_writers()  # Writers must be closed to get metadata
-        fns = [fn for fn in FileIterator(self.path, self.nfiles, shuffle=False)]
-        if self.metadata is not None:
-            _meta = self.metadata
-            i_start = 0
-        else:
-            _meta = self.new_metadata[fns[0]]
-            if _meta:
-                _meta = _meta[0]
-            i_start = 1
-        for i in range(i_start, len(fns)):
-            _meta_new = self.new_metadata[fns[i]]
-            if _meta_new:
-                _meta.append_row_groups(_meta_new[0])
-        with open(self.shared_meta_path, "wb") as fil:
-            _meta.write_metadata_file(fil)
-        self.metadata = _meta
-        return
-
-    def close_writers(self):
-        for fn, writer in self.writers.items():
-            if writer is not None:
-                writer.close()
-                # Set row-group file paths
-                self.new_metadata[fn][0].set_file_path(os.path.basename(fn))
-                writer = None
-
-    def __del__(self):
-        self.close_writers()
diff --git a/nvtabular/io.py b/nvtabular/io.py
--- a/nvtabular/io.py
+++ b/nvtabular/io.py
@@ -14,6 +14,7 @@
 # limitations under the License.
 #
 
+import collections
 import contextlib
 import functools
 import io
@@ -55,6 +56,80 @@
 
 LOG = logging.getLogger("nvtabular")
 
+
+#
+# Cache-Specific Code
+#
+
+
+class WriterCache:
+    def __init__(self):
+        self.pq_writer_cache = {}
+
+    def __del__(self):
+        for path, (pw, fpath) in self.pq_writer_cache.items():
+            pw.close()
+
+    def get_pq_writer(self, prefix, s, mem):
+        pw, fil = self.pq_writer_cache.get(prefix, (None, None))
+        if pw is None:
+            if mem:
+                fil = BytesIO()
+                pw = pwriter(fil, compression=None)
+                self.pq_writer_cache[prefix] = (pw, fil)
+            else:
+                outfile_id = guid() + ".parquet"
+                full_path = ".".join([prefix, outfile_id])
+                pw = pwriter(full_path, compression=None)
+                self.pq_writer_cache[prefix] = (pw, full_path)
+        return pw
+
+
+@contextlib.contextmanager
+def get_cache():
+    with _DEFAULT_CACHE_LOCK:
+        yield _get_cache()
+
+
+def _get_cache():
+    try:
+        worker = get_worker()
+    except ValueError:
+        # There is no dask.distributed worker.
+        # Assume client/worker are same process
+        global _DEFAULT_CACHE
+        if _DEFAULT_CACHE is None:
+            _DEFAULT_CACHE = WriterCache()
+        return _DEFAULT_CACHE
+    if not hasattr(worker, "pw_cache"):
+        worker.pw_cache = WriterCache()
+    return worker.pw_cache
+
+
+def clean_pw_cache():
+    with _DEFAULT_CACHE_LOCK:
+        try:
+            worker = get_worker()
+        except ValueError:
+            global _DEFAULT_CACHE
+            if _DEFAULT_CACHE is not None:
+                del _DEFAULT_CACHE
+                _DEFAULT_CACHE = None
+            return
+        if hasattr(worker, "pw_cache"):
+            del worker.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
+
+
 #
 # Helper Function definitions
 #
@@ -82,6 +157,15 @@ def _get_read_engine(engine, file_path, **kwargs):
         raise ValueError("Unrecognized read engine.")
 
 
+def _shuffle_gdf(gdf, gdf_size=None):
+    """ Shuffles a cudf dataframe, returning a new dataframe with randomly
+    ordered rows """
+    gdf_size = gdf_size or len(gdf)
+    arr = cp.arange(gdf_size)
+    cp.random.shuffle(arr)
+    return gdf.iloc[arr]
+
+
 #
 # GPUFileReader Base Class
 #
@@ -315,13 +399,19 @@ def _set_dtypes(self, chunk):
                 chunk[col] = chunk[col].astype(dtype)
 
 
-def _shuffle_gdf(gdf, gdf_size=None):
-    """ Shuffles a cudf dataframe, returning a new dataframe with randomly
-    ordered rows """
-    gdf_size = gdf_size or len(gdf)
-    arr = cp.arange(gdf_size)
-    cp.random.shuffle(arr)
-    return gdf.iloc[arr]
+def writer_factory(output_format, output_path, out_files_per_proc, shuffle):
+    if output_format is None:
+        return None
+
+    if output_format == "parquet":
+        writer_cls = ParquetWriter
+    elif output_format == "hugectr":
+        writer_cls = HugeCTRWriter
+    else:
+        raise ValueError("Output format not yet supported.")
+
+    fs = get_fs_token_paths(output_path)[0]
+    return writer_cls(output_path, num_out_files=out_files_per_proc, shuffle=shuffle, fs=fs)
 
 
 class Writer:
@@ -331,32 +421,39 @@ def __init__(self):
     def add_data(self, gdf):
         raise NotImplementedError()
 
-    def close(self):
-        pass
-
+    def package_general_metadata(self):
+        raise NotImplementedError()
 
-class Shuffler(Writer):
-    def __init__(
-        self, out_dir, num_out_files=30, num_threads=4, cats=None, conts=None, labels=None
-    ):
-        self.writer = ParquetWriter(out_dir, num_out_files, num_threads, cats, conts, labels)
+    @classmethod
+    def write_general_metadata(cls, data, fs, out_dir):
+        raise NotImplementedError()
 
-    def add_data(self, gdf):
-        self.writer.add_data(_shuffle_gdf(gdf))
+    @classmethod
+    def write_special_metadata(cls, data, fs, out_dir):
+        raise NotImplementedError()
 
     def close(self):
-        self.writer.close()
+        pass
 
 
 class ThreadedWriter(Writer):
     def __init__(
-        self, out_dir, num_out_files=30, num_threads=4, cats=None, conts=None, labels=None
+        self,
+        out_dir,
+        num_out_files=30,
+        num_threads=4,
+        cats=None,
+        conts=None,
+        labels=None,
+        shuffle=None,
+        fs=None,
     ):
         # set variables
         self.out_dir = out_dir
         self.cats = cats
         self.conts = conts
         self.labels = labels
+        self.shuffle = shuffle
         self.column_names = None
         if labels and conts:
             self.column_names = labels + conts
@@ -367,7 +464,11 @@ def __init__(
         self.num_out_files = num_out_files
         self.num_samples = [0] * num_out_files
 
-        self.data_files = None
+        self.data_paths = None
+        self.need_cal_col_names = True
+
+        # Resolve file system
+        self.fs = fs or get_fs_token_paths(str(out_dir))[0]
 
         # create thread queue and locks
         self.queue = queue.Queue(num_threads)
@@ -392,6 +493,11 @@ def _write_thread(self):
 
     @annotate("add_data", color="orange", domain="nvt_python")
     def add_data(self, gdf):
+
+        # Shuffle if necessary
+        if self.shuffle:
+            gdf = _shuffle_gdf(gdf)
+
         # Populate columns idxs
         if not self.col_idx:
             for i, x in enumerate(gdf.columns.values):
@@ -413,17 +519,59 @@ def add_data(self, gdf):
         # wait for all writes to finish before exitting (so that we aren't using memory)
         self.queue.join()
 
-    def _write_metadata(self):
-        return
+    def package_general_metadata(self):
+        data = {}
+        if self.cats is None:
+            return data
+        data["data_paths"] = self.data_paths
+        data["file_stats"] = []
+        for i, path in enumerate(self.data_paths):
+            fn = path.split(self.fs.sep)[-1]
+            data["file_stats"].append({"file_name": fn, "num_rows": self.num_samples[i]})
+        # cats
+        data["cats"] = []
+        for c in self.cats:
+            data["cats"].append({"col_name": c, "index": self.col_idx[c]})
+        # conts
+        data["conts"] = []
+        for c in self.conts:
+            data["conts"].append({"col_name": c, "index": self.col_idx[c]})
+        # labels
+        data["labels"] = []
+        for c in self.labels:
+            data["labels"].append({"col_name": c, "index": self.col_idx[c]})
+
+        return data
+
+    @classmethod
+    def write_general_metadata(cls, data, fs, out_dir):
+        if not data:
+            return
+        data_paths = data.pop("data_paths", [])
+        num_out_files = len(data_paths)
 
-    def _write_filelist(self):
-        file_list_writer = open(os.path.join(self.out_dir, "file_list.txt"), "w")
-        file_list_writer.write(str(self.num_out_files) + "\n")
-        for f in self.data_files:
+        # Write file_list
+        file_list_writer = fs.open(fs.sep.join([out_dir, "file_list.txt"]), "w")
+        file_list_writer.write(str(num_out_files) + "\n")
+        for f in data_paths:
             file_list_writer.write(f + "\n")
         file_list_writer.close()
 
-    def close(self):
+        # Write metadata json
+        metadata_writer = fs.open(fs.sep.join([out_dir, "metadata.json"]), "w")
+        json.dump(data, metadata_writer)
+        metadata_writer.close()
+
+    @classmethod
+    def write_special_metadata(cls, data, fs, out_dir):
+        pass
+
+    def _close_writers(self):
+        for writer in self.data_writers:
+            writer.close()
+        return None
+
+    def close(self, write_metadata=True):
         # wake up all the worker threads and signal for them to exit
         for _ in range(self.num_threads):
             self.queue.put(self._eod)
@@ -431,21 +579,25 @@ def close(self):
         # wait for pending writes to finish
         self.queue.join()
 
-        self._write_filelist()
-        self._write_metadata()
-
-        # Close writers
-        for writer in self.data_writers:
-            writer.close()
+        # Close writers and collect various metadata
+        _gen_meta = self.package_general_metadata()
+        _special_meta = self._close_writers()
+        return _gen_meta, _special_meta
 
 
 class ParquetWriter(ThreadedWriter):
-    def __init__(
-        self, out_dir, num_out_files=30, num_threads=4, cats=None, conts=None, labels=None
-    ):
-        super().__init__(out_dir, num_out_files, num_threads, cats, conts, labels)
-        self.data_files = [os.path.join(out_dir, f"{i}.parquet") for i in range(num_out_files)]
-        self.data_writers = [pwriter(f, compression=None) for f in self.data_files]
+    def __init__(self, out_dir, use_guid=False, **kwargs):
+        super().__init__(out_dir, **kwargs)
+        self.data_paths = []
+        self.data_writers = []
+        for i in range(self.num_out_files):
+            if use_guid:
+                fn = f"{i}.{guid()}.parquet"
+            else:
+                fn = f"{i}.parquet"
+            path = os.path.join(out_dir, fn)
+            self.data_paths.append(path)
+            self.data_writers.append(pwriter(path, compression=None))
 
     def _write_thread(self):
         while True:
@@ -459,38 +611,28 @@ def _write_thread(self):
             finally:
                 self.queue.task_done()
 
-    def _write_metadata(self):
-        if self.cats is None:
-            return
-        metadata_writer = open(os.path.join(self.out_dir, "metadata.json"), "w")
-        data = {}
-        data["file_stats"] = []
-        for i in range(len(self.data_files)):
-            data["file_stats"].append({"file_name": f"{i}.data", "num_rows": self.num_samples[i]})
-        # cats
-        data["cats"] = []
-        for c in self.cats:
-            data["cats"].append({"col_name": c, "index": self.col_idx[c]})
-        # conts
-        data["conts"] = []
-        for c in self.conts:
-            data["conts"].append({"col_name": c, "index": self.col_idx[c]})
-        # labels
-        data["labels"] = []
-        for c in self.labels:
-            data["labels"].append({"col_name": c, "index": self.col_idx[c]})
+    @classmethod
+    def write_special_metadata(cls, md, fs, out_dir):
+        # 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(list(md.items()), key=lambda x: natural_sort_key(x[0]))]
 
-        json.dump(data, metadata_writer)
-        metadata_writer.close()
+        # Aggregate metadata and write _metadata file
+        _write_pq_metadata_file(md_list, fs, out_dir)
+
+    def _close_writers(self):
+        md_dict = {}
+        for writer, path in zip(self.data_writers, self.data_paths):
+            fn = path.split(self.fs.sep)[-1]
+            md_dict[fn] = writer.close(metadata_file_path=fn)
+        return md_dict
 
 
 class HugeCTRWriter(ThreadedWriter):
-    def __init__(
-        self, out_dir, num_out_files=30, num_threads=4, cats=None, conts=None, labels=None
-    ):
-        super().__init__(out_dir, num_out_files, num_threads, cats, conts, labels)
-        self.data_files = [os.path.join(out_dir, f"{i}.data") for i in range(num_out_files)]
-        self.data_writers = [open(f, "ab") for f in self.data_files]
+    def __init__(self, out_dir, **kwargs):
+        super().__init__(out_dir, **kwargs)
+        self.data_paths = [os.path.join(out_dir, f"{i}.data") for i in range(self.num_out_files)]
+        self.data_writers = [open(f, "ab") for f in self.data_paths]
 
     def _write_thread(self):
         while True:
@@ -509,32 +651,33 @@ def _write_thread(self):
             finally:
                 self.queue.task_done()
 
-    def _write_metadata(self):
-        if self.cats is None:
-            return
-        for i in range(len(self.data_writers)):
-            self.data_writers[i].seek(0)
-            # error_check (0: no error check; 1: check_num)
-            # num of samples in this file
-            # Dimension of the labels
-            # Dimension of the features
-            # slot_num for each embedding
-            # reserved for future use
-            header = np.array(
-                [
-                    0,
-                    self.num_samples[i],
-                    len(self.labels),
-                    len(self.conts),
-                    len(self.cats),
-                    0,
-                    0,
-                    0,
-                ],
-                dtype=np.longlong,
-            )
-
-            self.data_writers[i].write(header.tobytes())
+    def _close_writers(self):
+        for i, writer in enumerate(self.data_writers):
+            if self.cats:
+                # Write HugeCTR Metadata
+                writer.seek(0)
+                # error_check (0: no error check; 1: check_num)
+                # num of samples in this file
+                # Dimension of the labels
+                # Dimension of the features
+                # slot_num for each embedding
+                # reserved for future use
+                header = np.array(
+                    [
+                        0,
+                        self.num_samples[i],
+                        len(self.labels),
+                        len(self.conts),
+                        len(self.cats),
+                        0,
+                        0,
+                        0,
+                    ],
+                    dtype=np.longlong,
+                )
+                writer.write(header.tobytes())
+            writer.close()
+        return None
 
 
 def device_mem_size(kind="total"):
@@ -556,74 +699,6 @@ def device_mem_size(kind="total"):
     return size
 
 
-class WriterCache:
-    def __init__(self):
-        self.pq_writer_cache = {}
-
-    def __del__(self):
-        for path, (pw, fpath) in self.pq_writer_cache.items():
-            pw.close()
-
-    def get_pq_writer(self, prefix, s, mem):
-        pw, fil = self.pq_writer_cache.get(prefix, (None, None))
-        if pw is None:
-            if mem:
-                fil = BytesIO()
-                pw = pwriter(fil, compression=None)
-                self.pq_writer_cache[prefix] = (pw, fil)
-            else:
-                outfile_id = guid() + ".parquet"
-                full_path = ".".join([prefix, outfile_id])
-                pw = pwriter(full_path, compression=None)
-                self.pq_writer_cache[prefix] = (pw, full_path)
-        return pw
-
-
-@contextlib.contextmanager
-def get_cache():
-    with _DEFAULT_CACHE_LOCK:
-        yield _get_cache()
-
-
-def _get_cache():
-    try:
-        worker = get_worker()
-    except ValueError:
-        # There is no dask.distributed worker.
-        # Assume client/worker are same process
-        global _DEFAULT_CACHE
-        if _DEFAULT_CACHE is None:
-            _DEFAULT_CACHE = WriterCache()
-        return _DEFAULT_CACHE
-    if not hasattr(worker, "pw_cache"):
-        worker.pw_cache = WriterCache()
-    return worker.pw_cache
-
-
-def clean_pw_cache():
-    with _DEFAULT_CACHE_LOCK:
-        try:
-            worker = get_worker()
-        except ValueError:
-            global _DEFAULT_CACHE
-            if _DEFAULT_CACHE is not None:
-                del _DEFAULT_CACHE
-                _DEFAULT_CACHE = None
-            return
-        if hasattr(worker, "pw_cache"):
-            del worker.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
     """
@@ -649,7 +724,7 @@ def _write_output_partition(gdf, processed_path, shuffle, out_files_per_proc, fs
     return gdf_size
 
 
-def _to_parquet_dataset(ddf, fs, output_path, shuffle, out_files_per_proc):
+def _ddf_to_pq_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 = []
@@ -670,6 +745,42 @@ def _to_parquet_dataset(ddf, fs, output_path, shuffle, out_files_per_proc):
     return Delayed(name, graph)
 
 
+def _finish_pq_dataset(client, ddf, shuffle, output_path, fs):
+    # Deal with "full" (per-worker) shuffle here.
+    if shuffle == "full":
+        if client:
+            client.cancel(ddf)
+            ddf = None
+            worker_md = client.run(_worker_shuffle, output_path, fs)
+            worker_md = list(collections.ChainMap(*worker_md.values()).items())
+        else:
+            ddf = None
+            worker_md = _worker_shuffle(output_path, fs)
+            worker_md = list(worker_md.items())
+    else:
+        # Collect parquet metadata while closing
+        # ParquetWriter object(s)
+        if client:
+            worker_md = client.run(close_cached_pw, fs)
+            worker_md = list(collections.ChainMap(*worker_md.values()).items())
+        else:
+            worker_md = 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
+    _write_pq_metadata_file(md_list, fs, output_path)
+
+    # Close ParquetWriter Objects
+    if client:
+        client.run(clean_pw_cache)
+    else:
+        clean_pw_cache()
+
+
 def _write_pq_metadata_file(md_list, fs, path):
     if md_list:
         metadata_path = fs.sep.join([path, "_metadata"])
diff --git a/nvtabular/workflow.py b/nvtabular/workflow.py
--- a/nvtabular/workflow.py
+++ b/nvtabular/workflow.py
@@ -15,18 +15,14 @@
 #
 import collections
 import logging
-import os
 import time
 import warnings
 
 import dask
 import yaml
-from cudf._lib.nvtx import annotate
-from dask.utils import natural_sort_key
 from fsspec.core import get_fs_token_paths
 
 import nvtabular.io as nvt_io
-from nvtabular.ds_writer import DatasetWriter
 from nvtabular.ops import DFOperator, StatOperator, TransformOperator
 from nvtabular.worker import clean_worker_cache
 
@@ -232,22 +228,6 @@ def finalize(self):
         """
         self.load_config(self.config)
 
-    def write_to_dataset(self, path, dataset, apply_ops=False, nfiles=1, shuffle=True, **kwargs):
-        """ Write data to shuffled parquet dataset.
-        """
-        if isinstance(dataset, nvt_io.Dataset):
-            itr = dataset.to_iter()
-        else:
-            itr = dataset
-
-        writer = DatasetWriter(path, nfiles=nfiles)
-
-        for gdf in itr:
-            if apply_ops:
-                gdf = self.apply_ops(gdf)
-            writer.write(gdf, shuffle=shuffle)
-        writer.write_metadata()
-
     def load_config(self, config, pro=False):
         """
         This function extracts all the operators from the given phases and produces a
@@ -492,14 +472,7 @@ def _run_trans_ops_for_phase(self, gdf, tasks):
         return gdf
 
     def apply_ops(
-        self,
-        gdf,
-        start_phase=None,
-        end_phase=None,
-        shuffler=None,
-        output_path=None,
-        num_out_files=None,
-        huge_ctr=None,
+        self, gdf, start_phase=None, end_phase=None, writer=None, output_path=None, shuffle=None
     ):
         """
         gdf: cudf dataframe
@@ -515,33 +488,22 @@ def apply_ops(
             start = time.time()
             gdf = self._run_trans_ops_for_phase(gdf, self.phases[phase_index])
             self.timings["preproc_apply"] += time.time() - start
-            if phase_index == len(self.phases) - 1 and output_path:
-                self.write_df(gdf, output_path, shuffler=shuffler, num_out_files=num_out_files)
+            if phase_index == len(self.phases) - 1 and writer and output_path:
 
-            if huge_ctr and phase_index == len(self.phases) - 1:
-                if not self.cal_col_names:
+                if writer.need_cal_col_names:
                     cat_names = self.get_final_cols_names("categorical")
                     cont_names = self.get_final_cols_names("continuous")
                     label_names = self.get_final_cols_names("label")
-                    huge_ctr.set_col_names(labels=label_names, cats=cat_names, conts=cont_names)
-                    self.cal_col_names = True
-                huge_ctr.add_data(gdf)
+                    writer.set_col_names(labels=label_names, cats=cat_names, conts=cont_names)
+                    writer.need_cal_col_names = False
 
-        return gdf
+                start_write = time.time()
+                writer.add_data(gdf)
+                self.timings["shuffle_df"] += time.time() - start_write
 
-    @annotate("Write_df", color="red", domain="nvt_python")
-    def write_df(self, gdf, export_path, shuffler, num_out_files):
-        if shuffler:
-            start = time.time()
-            shuffler.add_data(gdf)
-            self.timings["shuffle_df"] += time.time() - start
-        else:
-            file_name = f"{self.current_file_num}.parquet"
-            path = os.path.join(export_path, file_name)
-            gdf.to_parquet(path, compression=None)
-            self.current_file_num += 1
+        return gdf
 
-    def _update_stats(self, stat_op):
+    def _update_statistics(self, stat_op):
         self.stats.update(stat_op.stats_collected())
 
     def save_stats(self, path):
@@ -667,16 +629,38 @@ def exec_phase(self, phase_index, record_stats=True):
                 for r in self.client.compute(stats):
                     computed_stats, op = r.result()
                     op.finalize(computed_stats)
-                    self._update_stats(op)
+                    self._update_statistics(op)
                     op.clear()
             else:
                 for r in dask.compute(stats, scheduler="synchronous")[0]:
                     computed_stats, op = r
                     op.finalize(computed_stats)
-                    self._update_stats(op)
+                    self._update_statistics(op)
                     op.clear()
             del stats
 
+    def reorder_tasks(self, end):
+        if end != 2:
+            # Opt only works for two phases (for now)
+            return
+        stat_tasks = []
+        trans_tasks = []
+        for idx, _ in enumerate(self.phases[:end]):
+            for task in self.phases[idx]:
+                deps = task[2]
+                if isinstance(task[0], StatOperator):
+                    if deps == ["base"]:
+                        stat_tasks.append(task)
+                    else:
+                        # This statistics depends on a transform
+                        # (Opt wont work)
+                        return
+                elif isinstance(task[0], TransformOperator):
+                    trans_tasks.append(task)
+
+        self.phases[0] = stat_tasks
+        self.phases[1] = trans_tasks
+
     def apply(
         self,
         dataset,
@@ -684,11 +668,8 @@ def apply(
         record_stats=True,
         shuffle=None,
         output_path="./ds_export",
+        output_format="parquet",
         out_files_per_proc=None,
-        hugectr_gen_output=False,
-        hugectr_output_path="./hugectr",
-        hugectr_num_out_files=None,
-        hugectr_output_format=None,
         **kwargs,
     ):
         """
@@ -714,6 +695,9 @@ def apply(
             disk. A "true" full shuffle is not yet implemented.
         output_path : string
             path to output data
+        output_format : {"parquet", "hugectr", None}
+            Output format for processed/shuffled dataset. If None,
+            no output dataset will be written.
         out_files_per_proc : integer
             number of files to create (per process) after
             shuffling the data
@@ -734,83 +718,87 @@ def apply(
         # If no tasks have been loaded then we need to load internal config
         if not self.phases:
             self.finalize()
+
+        # Gather statstics (if apply_offline), and/or transform
+        # and write out processed data
         if apply_offline:
-            if hugectr_gen_output:
-                raise ValueError(
-                    "TODO: Support HugeCTR output for offline processing with Dask."
-                    " This is part of the larger task of aligning online/offline API."
-                )
-            self.update_stats(
+            self.build_and_process_graph(
                 dataset,
                 output_path=output_path,
                 record_stats=record_stats,
                 shuffle=shuffle,
+                output_format=output_format,
                 out_files_per_proc=out_files_per_proc,
             )
         else:
-            shuffler = None
-            huge_ctr = None
-            if shuffle:
-                if isinstance(shuffle, str):
-                    raise ValueError("TODO: Align shuffling/writing API for online/offline.")
-                shuffler = nvt_io.Shuffler(output_path, num_out_files=out_files_per_proc)
-            if hugectr_gen_output:
-                self.cal_col_names = False
-                if hugectr_output_format == "binary":
-                    huge_ctr = nvt_io.HugeCTRWriter(
-                        hugectr_output_path, num_out_files=hugectr_num_out_files
-                    )
-                elif hugectr_output_format == "parquet":
-                    huge_ctr = nvt_io.ParquetWriter(
-                        hugectr_output_path, num_out_files=hugectr_num_out_files
-                    )
-            # "Online" apply currently requires manual iteration
+            self.iterate_online(
+                dataset,
+                output_path=output_path,
+                shuffle=shuffle,
+                output_format=output_format,
+                out_files_per_proc=out_files_per_proc,
+            )
+
+    def iterate_online(
+        self,
+        dataset,
+        end_phase=None,
+        output_path=None,
+        shuffle=None,
+        output_format=None,
+        out_files_per_proc=None,
+        apply_ops=True,
+    ):
+        """ Iterate through dataset and (optionally) apply/shuffle/write.
+        """
+        # Check if we have a (supported) writer
+        output_path = output_path or "./"
+        output_path = str(output_path)
+        writer = nvt_io.writer_factory(output_format, output_path, out_files_per_proc, shuffle)
+
+        # Iterate through dataset, apply ops, and write out processed data
+        if apply_ops:
             for gdf in dataset.to_iter():
-                self.apply_ops(
-                    gdf,
-                    output_path=output_path,
-                    shuffler=shuffler,
-                    num_out_files=out_files_per_proc,
-                    huge_ctr=huge_ctr,
-                )
-            if shuffler:
-                shuffler.close()
-            if huge_ctr:
-                huge_ctr.close()
+                self.apply_ops(gdf, output_path=output_path, writer=writer, shuffle=shuffle)
 
-    def reorder_tasks(self, end):
-        if end != 2:
-            # Opt only works for two phases (for now)
-            return
-        stat_tasks = []
-        trans_tasks = []
-        for idx, _ in enumerate(self.phases[:end]):
-            for task in self.phases[idx]:
-                deps = task[2]
-                if isinstance(task[0], StatOperator):
-                    if deps == ["base"]:
-                        stat_tasks.append(task)
-                    else:
-                        # This statistics depends on a transform
-                        # (Opt wont work)
-                        return
-                elif isinstance(task[0], TransformOperator):
-                    trans_tasks.append(task)
+        # Close writer and write general/specialized metadata
+        if writer:
+            general_md, special_md = writer.close()
 
-        self.phases[0] = stat_tasks
-        self.phases[1] = trans_tasks
+            # Note that we "could" have the special and general metadata
+            # written during `writer.close()` (just above) for the single-GPU case.
+            # Instead, the metadata logic is separated from the `Writer` object to
+            # simplify multi-GPU integration. When using Dask, we cannot assume
+            # that the "shared" metadata files can/will be written by the same
+            # process that writes the data.
+            type(writer).write_special_metadata(special_md, writer.fs, output_path)
+            type(writer).write_general_metadata(general_md, writer.fs, output_path)
+
+    def update_stats(self, dataset, end_phase=None):
+        """ Colllect statistics only.
+        """
+        self.build_and_process_graph(dataset, end_phase=end_phase, record_stats=True)
 
-    def update_stats(
+    def build_and_process_graph(
         self,
         dataset,
         end_phase=None,
         output_path=None,
         record_stats=True,
         shuffle=None,
+        output_format=None,
         out_files_per_proc=None,
+        apply_ops=True,
     ):
+        """ Build Dask-task graph for workflow.
+
+            Full graph is only executed if `output_format` is specified.
+        """
         end = end_phase if end_phase else len(self.phases)
 
+        if output_format not in ("parquet", None):
+            raise ValueError("Output format not yet supported with Dask.")
+
         # Reorder tasks for two-phase workflows
         self.reorder_tasks(end)
 
@@ -821,19 +809,66 @@ def update_stats(
             clean_worker_cache()
 
         self.set_ddf(dataset)
-        for idx, _ in enumerate(self.phases[:end]):
-            self.exec_phase(idx, record_stats=record_stats)
-        if output_path:
-            self.to_dataset(output_path, shuffle=shuffle, out_files_per_proc=out_files_per_proc)
+        if apply_ops:
+            for idx, _ in enumerate(self.phases[:end]):
+                self.exec_phase(idx, record_stats=record_stats)
+        if output_format:
+            output_path = output_path or "./"
+            output_path = str(output_path)
+            self.ddf_to_dataset(output_path, shuffle=shuffle, out_files_per_proc=out_files_per_proc)
+
+    def write_to_dataset(
+        self,
+        path,
+        dataset,
+        apply_ops=False,
+        nfiles=1,
+        shuffle=True,
+        output_format="parquet",
+        iterate=True,
+    ):
+        """ Write data to shuffled parquet dataset.
 
-    def to_dataset(self, output_path, shuffle=None, out_files_per_proc=None):
+            Assumes statistics are already gathered.
+        """
+        path = str(path)
+        if iterate:
+            self.iterate_online(
+                dataset,
+                output_path=path,
+                shuffle=shuffle,
+                output_format=output_format,
+                out_files_per_proc=nfiles,
+                apply_ops=apply_ops,
+            )
+        else:
+            self.build_and_process_graph(
+                dataset,
+                output_path=path,
+                record_stats=False,
+                shuffle=shuffle,
+                output_format=output_format,
+                out_files_per_proc=nfiles,
+                apply_ops=apply_ops,
+            )
+
+    def ddf_to_dataset(
+        self, output_path, shuffle=None, out_files_per_proc=None, output_format="parquet"
+    ):
+        """ Dask-based dataset output.
+
+            Currently supports parquet only.
+            TODO: Leverage `ThreadedWriter` implementations.
+        """
+        if output_format != "parquet":
+            raise ValueError("Only parquet output supported with Dask.")
         ddf = self.get_ddf()
         fs = get_fs_token_paths(output_path)[0]
         fs.mkdirs(output_path, exist_ok=True)
         if shuffle or out_files_per_proc:
 
             # Construct graph for Dask-based dataset write
-            out = nvt_io._to_parquet_dataset(ddf, fs, output_path, shuffle, out_files_per_proc)
+            out = nvt_io._ddf_to_pq_dataset(ddf, fs, output_path, shuffle, out_files_per_proc)
 
             # Would be nice to clean the categorical
             # cache before the write (TODO)
@@ -847,40 +882,8 @@ def to_dataset(self, output_path, shuffle=None, out_files_per_proc=None):
                 self.ddf_base_dataset = None
                 out = dask.compute(out, scheduler="synchronous")[0]
 
-            # Deal with "full" (per-worker) shuffle here.
-            if shuffle == "full":
-                if self.client:
-                    self.client.cancel(self.ddf)
-                    self.ddf = None
-                    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()
+            # Follow-up Shuffling and _metadata creation
+            nvt_io._finish_pq_dataset(self.client, self.ddf, shuffle, output_path, fs)
 
             return
 


Test patch:
diff --git a/tests/unit/test_io.py b/tests/unit/test_io.py
--- a/tests/unit/test_io.py
+++ b/tests/unit/test_io.py
@@ -25,7 +25,7 @@
 import nvtabular as nvt
 import nvtabular.io
 import nvtabular.ops as ops
-from nvtabular.io import Shuffler
+from nvtabular.io import ParquetWriter
 from tests.conftest import allcols_csv, mycols_csv, mycols_pq
 
 
@@ -37,9 +37,9 @@ def test_shuffle_gpu(tmpdir, datasets, engine):
         df1 = cudf.read_parquet(paths[0])[mycols_pq]
     else:
         df1 = cudf.read_csv(paths[0], header=False, names=allcols_csv)[mycols_csv]
-    shuf = Shuffler(tmpdir, num_files)
+    shuf = ParquetWriter(tmpdir, num_out_files=num_files, shuffle=True)
     shuf.add_data(df1)
-    writer_files = shuf.writer.data_files
+    writer_files = shuf.data_paths
     shuf.close()
     if engine == "parquet":
         df3 = cudf.read_parquet(writer_files[0])[mycols_pq]
@@ -88,7 +88,7 @@ def test_dask_dataset(datasets, engine, num_files):
     assert_eq(ddf0, result)
 
 
-@pytest.mark.parametrize("output_format", ["binary", "parquet"])
+@pytest.mark.parametrize("output_format", ["hugectr", "parquet"])
 @pytest.mark.parametrize("engine", ["parquet", "csv", "csv-no-header"])
 @pytest.mark.parametrize("op_columns", [["x"], None])
 def test_hugectr(tmpdir, df, dataset, output_format, engine, op_columns):
@@ -99,10 +99,8 @@ def test_hugectr(tmpdir, df, dataset, output_format, engine, op_columns):
     # set variables
     nfiles = 10
     ext = ""
-    outdir = tmpdir + "/dontcare"
-    h_outdir = tmpdir + "/hugectr"
+    outdir = tmpdir + "/hugectr"
     os.mkdir(outdir)
-    os.mkdir(h_outdir)
 
     # process data
     processor = nvt.Workflow(cat_names=cat_names, cont_names=cont_names, label_name=label_names)
@@ -120,22 +118,19 @@ def test_hugectr(tmpdir, df, dataset, output_format, engine, op_columns):
         apply_offline=False,
         record_stats=False,
         output_path=outdir,
+        out_files_per_proc=nfiles,
+        output_format=output_format,
         shuffle=False,
-        hugectr_gen_output=True,
-        hugectr_output_path=h_outdir,
-        hugectr_num_out_files=nfiles,
-        hugectr_output_format=output_format,
     )
 
     # Check files
     ext = ""
     if output_format == "parquet":
         ext = "parquet"
-        assert os.path.isfile(h_outdir + "/metadata.json")
-    elif output_format == "binary":
+        assert os.path.isfile(outdir + "/metadata.json")
+    elif output_format == "hugectr":
         ext = "data"
 
-    assert os.path.isfile(h_outdir + "/file_list.txt")
-
+    assert os.path.isfile(outdir + "/file_list.txt")
     for n in range(nfiles):
-        assert os.path.isfile(os.path.join(h_outdir, str(n) + "." + ext))
+        assert os.path.isfile(os.path.join(outdir, str(n) + "." + ext))
diff --git a/tests/unit/test_notebooks.py b/tests/unit/test_notebooks.py
--- a/tests/unit/test_notebooks.py
+++ b/tests/unit/test_notebooks.py
@@ -35,6 +35,7 @@ def test_optimize_criteo(tmpdir):
 
 
 def test_rossman_example(tmpdir):
+    pytest.importorskip("tensorflow")
     _get_random_rossmann_data(1000).to_csv(os.path.join(tmpdir, "train.csv"))
     _get_random_rossmann_data(1000).to_csv(os.path.join(tmpdir, "valid.csv"))
 
diff --git a/tests/unit/test_ops.py b/tests/unit/test_ops.py
--- a/tests/unit/test_ops.py
+++ b/tests/unit/test_ops.py
@@ -400,7 +400,7 @@ def test_lambdaop(tmpdir, df, dataset, gpu_memory_frac, engine, client):
     processor.write_to_dataset(tmpdir, dataset, nfiles=10, shuffle=True, apply_ops=True)
 
     dataset_2 = nvtabular.io.Dataset(
-        glob.glob(str(tmpdir) + "/ds_part.*.parquet"), part_mem_fraction=gpu_memory_frac
+        glob.glob(str(tmpdir) + "/*.parquet"), part_mem_fraction=gpu_memory_frac
     )
     df_pp = cudf.concat(list(dataset_2.to_iter()), axis=0)
     assert is_integer_dtype(df_pp["name-cat"].dtype)
@@ -420,7 +420,7 @@ def test_lambdaop(tmpdir, df, dataset, gpu_memory_frac, engine, client):
     processor.write_to_dataset(tmpdir, dataset, nfiles=10, shuffle=True, apply_ops=True)
 
     dataset_2 = nvtabular.io.Dataset(
-        glob.glob(str(tmpdir) + "/ds_part.*.parquet"), part_mem_fraction=gpu_memory_frac
+        glob.glob(str(tmpdir) + "/*.parquet"), part_mem_fraction=gpu_memory_frac
     )
     df_pp = cudf.concat(list(dataset_2.to_iter()), axis=0)
     assert is_integer_dtype(df_pp["name-cat"].dtype)
@@ -446,7 +446,7 @@ def test_lambdaop(tmpdir, df, dataset, gpu_memory_frac, engine, client):
     processor.update_stats(dataset)
     processor.write_to_dataset(tmpdir, dataset, nfiles=10, shuffle=True, apply_ops=True)
     dataset_2 = nvtabular.io.Dataset(
-        glob.glob(str(tmpdir) + "/ds_part.*.parquet"), part_mem_fraction=gpu_memory_frac
+        glob.glob(str(tmpdir) + "/*.parquet"), part_mem_fraction=gpu_memory_frac
     )
     df_pp = cudf.concat(list(dataset_2.to_iter()), axis=0)
 
@@ -470,7 +470,7 @@ def test_lambdaop(tmpdir, df, dataset, gpu_memory_frac, engine, client):
     processor.write_to_dataset(tmpdir, dataset, nfiles=10, shuffle=True, apply_ops=True)
 
     dataset_2 = nvtabular.io.Dataset(
-        glob.glob(str(tmpdir) + "/ds_part.*.parquet"), part_mem_fraction=gpu_memory_frac
+        glob.glob(str(tmpdir) + "/*.parquet"), part_mem_fraction=gpu_memory_frac
     )
     df_pp = cudf.concat(list(dataset_2.to_iter()), axis=0)
     assert is_integer_dtype(df_pp["name-cat_add100"].dtype)
@@ -497,7 +497,7 @@ def test_lambdaop(tmpdir, df, dataset, gpu_memory_frac, engine, client):
     processor.write_to_dataset(tmpdir, dataset, nfiles=10, shuffle=True, apply_ops=True)
 
     dataset_2 = nvtabular.io.Dataset(
-        glob.glob(str(tmpdir) + "/ds_part.*.parquet"), part_mem_fraction=gpu_memory_frac
+        glob.glob(str(tmpdir) + "/*.parquet"), part_mem_fraction=gpu_memory_frac
     )
     df_pp = cudf.concat(list(dataset_2.to_iter()), axis=0)
     assert np.sum(df_pp["x_mul0_add100"] < 100) == 0
diff --git a/tests/unit/test_tf_dataloader.py b/tests/unit/test_tf_dataloader.py
--- a/tests/unit/test_tf_dataloader.py
+++ b/tests/unit/test_tf_dataloader.py
@@ -54,7 +54,7 @@ def test_tf_gpu_dl(tmpdir, paths, use_paths, dataset, batch_size, gpu_memory_fra
         engine=engine,
         shuffle=False,
     )
-    processor.update_stats(dataset, record_stats=True)
+    processor.update_stats(dataset)
     data_itr.map(processor)
 
     rows = 0
diff --git a/tests/unit/test_torch_dataloader.py b/tests/unit/test_torch_dataloader.py
--- a/tests/unit/test_torch_dataloader.py
+++ b/tests/unit/test_torch_dataloader.py
@@ -148,7 +148,7 @@ def get_norms(tar: cudf.Series):
         torch_dataloader.FileItrDataset(
             x, use_row_groups=True, gpu_memory_frac=gpu_memory_frac, names=allcols_csv
         )
-        for x in glob.glob(str(tmpdir) + "/ds_part.*.parquet")
+        for x in glob.glob(str(tmpdir) + "/*.parquet")
     ]
 
     data_itr = torch.utils.data.ChainDataset(data_files)
@@ -160,9 +160,7 @@ def get_norms(tar: cudf.Series):
     for chunk in dl:
         len_df_pp += len(chunk[0][0])
 
-    dataset = Dataset(
-        glob.glob(str(tmpdir) + "/ds_part.*.parquet"), part_mem_fraction=gpu_memory_frac,
-    )
+    dataset = Dataset(glob.glob(str(tmpdir) + "/*.parquet"), part_mem_fraction=gpu_memory_frac)
     x = processor.ds_to_tensors(dataset.to_iter(), apply_ops=False)
 
     num_rows, num_row_groups, col_names = cudf.io.read_parquet_metadata(str(tmpdir) + "/_metadata")
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
@@ -91,9 +91,7 @@ def get_norms(tar: cudf.Series):
     # Write to new "shuffled" and "processed" dataset
     processor.write_to_dataset(tmpdir, dataset, nfiles=10, shuffle=True, apply_ops=True)
 
-    dataset_2 = Dataset(
-        glob.glob(str(tmpdir) + "/ds_part.*.parquet"), part_mem_fraction=gpu_memory_frac
-    )
+    dataset_2 = Dataset(glob.glob(str(tmpdir) + "/*.parquet"), part_mem_fraction=gpu_memory_frac)
 
     df_pp = cudf.concat(list(dataset_2.to_iter()), axis=0)
 
@@ -167,9 +165,7 @@ def get_norms(tar: cudf.Series):
     # Write to new "shuffled" and "processed" dataset
     processor.write_to_dataset(tmpdir, dataset, nfiles=10, shuffle=True, apply_ops=True)
 
-    dataset_2 = Dataset(
-        glob.glob(str(tmpdir) + "/ds_part.*.parquet"), part_mem_fraction=gpu_memory_frac
-    )
+    dataset_2 = Dataset(glob.glob(str(tmpdir) + "/*.parquet"), part_mem_fraction=gpu_memory_frac)
 
     df_pp = cudf.concat(list(dataset_2.to_iter()), axis=0)
 
@@ -250,9 +246,7 @@ def get_norms(tar: cudf.Series):
     # Write to new "shuffled" and "processed" dataset
     processor.write_to_dataset(tmpdir, dataset, nfiles=10, shuffle=True, apply_ops=True)
 
-    dataset_2 = Dataset(
-        glob.glob(str(tmpdir) + "/ds_part.*.parquet"), part_mem_fraction=gpu_memory_frac
-    )
+    dataset_2 = Dataset(glob.glob(str(tmpdir) + "/*.parquet"), part_mem_fraction=gpu_memory_frac)
 
     df_pp = cudf.concat(list(dataset_2.to_iter()), axis=0)
 
@@ -284,7 +278,7 @@ def test_parquet_output(client, use_client, tmpdir, shuffle):
     )
     processor.add_preprocess(ops.Normalize())
     processor.finalize()
-    processor.update_stats(
+    processor.apply(
         dataset, output_path=out_path, shuffle=shuffle, out_files_per_proc=out_files_per_proc
     )