Repository: NVIDIA-Merlin/NVTabular

Issue:
[BUG] Output-file count is incorrect for shuffle="partial"
**Describe the bug**
As originally discovered by @rnyak, there seems to be an extra output file written for `shuffle="partial"`.   That is, when using a single process with `out_files_per_proc=2`, there are **3** output files written.

**Steps/Code to reproduce bug**
```python
import nvtabular as nvt
import pandas as pd
import numpy as np

!rm -rf test.parquet processed
!mkdir processed
path = "test.parquet"
out_path = "processed/"

out_files_per_proc = 2

size = 25
df = pd.DataFrame({"a": np.arange(size)})
df.to_parquet(path, row_group_size=5, 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.add_preprocess(nvt.ops.Normalize())
processor.finalize()

processor.update_stats(
    dataset,
    output_path=out_path,
    shuffle="partial", 
    out_files_per_proc=out_files_per_proc,
)
!ls processed
```
**Output has 3 files (should be 2)**:
```
split.0.459d88bc3f52405bae6fd572570fd452.parquet
split.1.ff7fd6b1f7d44504b24351afd8b43d4c.parquet
split.2.f6cfed77c5a848e998c60b2d2613a4dd.parquet
```

**Expected behavior**
The `out_files_per_proc` parameter should result in the expected number of output files.



Discussion:

Resolving patch:
diff --git a/nvtabular/categorify.py b/nvtabular/categorify.py
--- a/nvtabular/categorify.py
+++ b/nvtabular/categorify.py
@@ -191,7 +191,9 @@ def _write_uniques(dfs, base_path, col, on_host):
         # Make sure first category is Null
         df = df.sort_values(col, na_position="first")
         if not df[col]._column.has_nulls:
-            df = cudf.DataFrame({col: _concat([cudf.Series([None]), df[col]], ignore_index)})
+            df = cudf.DataFrame(
+                {col: _concat([cudf.Series([None], dtype=df[col].dtype), df[col]], ignore_index)}
+            )
         df.to_parquet(path, write_index=False, compression=None)
     else:
         df_null = cudf.DataFrame({col: [None]})
diff --git a/nvtabular/io.py b/nvtabular/io.py
--- a/nvtabular/io.py
+++ b/nvtabular/io.py
@@ -17,6 +17,7 @@
 import io
 import json
 import logging
+import math
 import os
 import queue
 import threading
@@ -37,7 +38,6 @@
 from dask.base import tokenize
 from dask.dataframe.core import new_dd_object
 from dask.dataframe.io.parquet.utils import _analyze_paths
-from dask.dataframe.utils import group_split_dispatch
 from dask.distributed import get_worker
 from dask.utils import natural_sort_key, parse_bytes
 from fsspec.core import get_fs_token_paths
@@ -617,30 +617,18 @@ def _write_metadata(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)
-    if shuffle == "full":
-        # Dont need a real sort if we are doing in memory later
-        typ = np.min_scalar_type(out_files_per_proc * 2)
-        ind = cp.random.choice(cp.arange(out_files_per_proc, dtype=typ), gdf_size)
-        result = group_split_dispatch(gdf, ind, out_files_per_proc, ignore_index=True)
-        del ind
-        del gdf
-        # Write each split to a separate file
-        for s, df in result.items():
-            prefix = fs.sep.join([processed_path, "split." + str(s)])
-            pw = get_cache().get_pq_writer(prefix, s, mem=True)
-            pw.write_table(df)
-    else:
+    out_files_per_proc = out_files_per_proc or 1
+    if shuffle and shuffle != "full":
         # We should do a real sort here
-        if shuffle == "partial":
-            gdf = _shuffle_gdf(gdf, gdf_size=gdf_size)
-        splits = list(range(0, gdf_size, int(gdf_size / out_files_per_proc)))
-        if splits[-1] < gdf_size:
-            splits.append(gdf_size)
-        # Write each split to a separate file
-        for s in range(0, len(splits) - 1):
-            prefix = fs.sep.join([processed_path, "split." + str(s)])
-            pw = get_cache().get_pq_writer(prefix, s, mem=False)
-            pw.write_table(gdf.iloc[splits[s] : splits[s + 1]])
+        gdf = _shuffle_gdf(gdf, gdf_size=gdf_size)
+    split_size = math.ceil(gdf_size / out_files_per_proc)
+    ind = cp.ones(gdf_size).cumsum() // split_size
+    del ind
+    for s in range(out_files_per_proc):
+        prefix = fs.sep.join([processed_path, "split." + str(s)])
+        pw = get_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
 
 
diff --git a/nvtabular/workflow.py b/nvtabular/workflow.py
--- a/nvtabular/workflow.py
+++ b/nvtabular/workflow.py
@@ -705,10 +705,17 @@ def apply(
         record_stats : boolean
             record the stats in file or not. Only available
             for apply_offline=True
-        shuffle : boolean
-            shuffles the data or not
+        shuffle : {"full", "partial", None}
+            Whether to shuffle output dataset. "partial" means
+            each worker will randomly shuffle data into a number
+            (`out_files_per_proc`) of different output files as the data
+            is processed. The output files are distinctly mapped to
+            each worker process. "full" means the workers will perform the
+            "partial" shuffle into BytesIO files, and then perform
+            a full shuffle of each in-memory file before writing to
+            disk. A "true" full shuffle is not yet implemented.
         output_path : string
-            path to export stats
+            path to output data
         out_files_per_proc : integer
             number of files to create (per process) after
             shuffling the data
@@ -748,7 +755,7 @@ def apply(
             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=num_out_files)
+                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":
@@ -823,11 +830,9 @@ def update_stats(
 
     def to_dataset(self, output_path, shuffle=None, out_files_per_proc=None):
         ddf = self.get_ddf()
-        out_files_per_proc = out_files_per_proc or 1
         fs = get_fs_token_paths(output_path)[0]
         fs.mkdirs(output_path, exist_ok=True)
-
-        if shuffle:
+        if shuffle or out_files_per_proc:
             name = "write-processed"
             write_name = name + tokenize(ddf, shuffle, out_files_per_proc)
             task_list = []


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
@@ -16,9 +16,11 @@
 
 import glob
 import math
+import os
 
 import cudf
 import numpy as np
+import pandas as pd
 import pytest
 from cudf.tests.utils import assert_eq
 from pandas.api.types import is_integer_dtype
@@ -90,7 +92,7 @@ def get_norms(tar: cudf.Series):
     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,
+        glob.glob(str(tmpdir) + "/ds_part.*.parquet"), part_mem_fraction=gpu_memory_frac
     )
 
     df_pp = cudf.concat(list(dataset_2.to_iter()), axis=0)
@@ -166,7 +168,7 @@ def get_norms(tar: cudf.Series):
     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,
+        glob.glob(str(tmpdir) + "/ds_part.*.parquet"), part_mem_fraction=gpu_memory_frac
     )
 
     df_pp = cudf.concat(list(dataset_2.to_iter()), axis=0)
@@ -250,7 +252,7 @@ def get_norms(tar: cudf.Series):
     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,
+        glob.glob(str(tmpdir) + "/ds_part.*.parquet"), part_mem_fraction=gpu_memory_frac
     )
 
     df_pp = cudf.concat(list(dataset_2.to_iter()), axis=0)
@@ -261,3 +263,27 @@ def get_norms(tar: cudf.Series):
 
     num_rows, num_row_groups, col_names = cudf.io.read_parquet_metadata(str(tmpdir) + "/_metadata")
     assert num_rows == len(df_pp)
+
+
+@pytest.mark.parametrize("shuffle", ["full", "partial", None])
+def test_output_count(tmpdir, shuffle):
+    out_files_per_proc = 2
+    out_path = os.path.join(tmpdir, "processed")
+    path = os.path.join(tmpdir, "simple.parquet")
+
+    df = pd.DataFrame({"a": np.arange(25)})
+    df.to_parquet(path, row_group_size=5, 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.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