fix arbitrary output file number bug, shrink number of files and warn… (#1301)

* fix arbitrary output file number bug, shrink number of files and warn if not enough ddf partitions

* add test to verify behavior support

* made changes to ci script to clean pip packages before install to make it fresh

* adding user warning

* add correct if condition for warning

* Update nvtabular/io/dataset.py

Co-authored-by: Richard (Rick) Zamora <[email protected]>

* Update tests/unit/test_io.py

Co-authored-by: Richard (Rick) Zamora <[email protected]>

* fix issues in formatting of code

Co-authored-by: Ben Frederickson <[email protected]>
Co-authored-by: Alberto Alvarez <[email protected]>
Co-authored-by: Richard (Rick) Zamora <[email protected]>

Author: 4 people

nvtabular/io/dataset.py

@@ -861,7 +861,14 @@ def to_parquet(
                     fns = output_files[i : i + files_per_task]
                     start = i * split
                     stop = min(start + split * len(fns), ddf.npartitions)
-                    new[tuple(fns)] = np.arange(start, stop)
+                    if start < stop:
+                        new[tuple(fns)] = np.arange(start, stop)
+                # let user know they will not have expected number of output files.
+                if len(new.keys()) < len(output_files):
+                    warnings.warn(
+                        f"Only created {len(new.keys())} files did not have enough\n"
+                        f"partitions to create {len(output_files)} files."
+                    )
                 output_files = new
                 suffix = ""  # Don't add a suffix later - Names already include it
             if not isinstance(output_files, dict):

tests/unit/test_io.py

@@ -158,6 +158,36 @@ def test_dask_dataset_itr(tmpdir, datasets, engine, gpu_memory_frac):
     assert len(my_iter) == size
 
 
+def test_io_partitions_push(tmpdir):
+    os.makedirs(os.path.join(tmpdir, "csv"))
+
+    # Generate random csv files
+    files = [os.path.join(tmpdir, f"csv/day_{i}") for i in range(23)]
+    for file in files:
+        with open(file, "w") as f:
+            f.write("0,1,2,3,a,b,c\n" * 1000)
+
+    # Load csv files
+    label_columns = ["label"]
+    cont_columns = ["I1", "I2", "I3"]
+    cat_columns = ["C1", "C2", "C3"]
+    columns = label_columns + cont_columns + cat_columns
+    dataset = nvt.Dataset(files, engine="csv", names=columns)
+    print("npartitions of dataset:", dataset.npartitions)
+
+    for x in range(20):
+        dataset.to_parquet(
+            output_files=x,
+            output_path=os.path.join(tmpdir, f"parquet{x}"),
+            cats=cat_columns,
+            conts=cont_columns,
+            labels=label_columns,
+        )
+
+        df_lib = dispatch.get_lib()
+        df_lib.read_parquet(os.path.join(tmpdir, f"parquet{x}/part_0.parquet"))
+
+
 @pytest.mark.parametrize("engine", ["csv", "parquet", "csv-no-header"])
 @pytest.mark.parametrize("num_files", [1, 2])
 @pytest.mark.parametrize("cpu", [None, True])