Repository: NVIDIA-Merlin/NVTabular

Issue:
[BUG] Recent regressions in Criteo-benchmark performance
**Describe the bug**
It looks like a few PRs have introduced significant performance regressions in the DGX-1 dask-nvtabular-criteo-benchmark.py run time over the previous month. If I go back to the commit before  #1136 , I get a total run time of ~125s. However, on the latest main I am getting ~434s.

The biggest performance hit comes from the addition of a HugeCTR keyset output file write in #1136 (which is apparently very slow). However, #1162 also adds a new `len(ddf)` call to `Categorify`. For some reason, that new `len(ddf)` call was running relatively fast when it was first introduced, but it became considerably slower after #1194. 

When I manually remove the keyset write and the `len(ddf)` call from `main`, the benchmark goes back down to ~125s.

**Steps/Code to reproduce bug**:

Running on a DGX-1 machine with 8x32GB-V100:

```
python NVTabular/bench/examples/dask-nvtabular-criteo-benchmark.py --part-mem-frac 0.125 --device-limit-frac 0.7 --device-pool-frac 0.8  --data-path /raid/.../criteo_parquet--out-path /raid/.../scratch/ --out-files-per-proc 3 --cat-cache-high device -d 0,1,2,3,4,5,6,7
```

**Output before #1136**:
```
Dask-NVTabular DLRM/Criteo benchmark
--------------------------------------
partition size     | 4261216256
protocol           | tcp
device(s)          | 0,1,2,3,4,5,6,7
rmm-pool-frac      | 0.8
out-files-per-proc | 3
num_io_threads     | 0
shuffle            | PER_PARTITION
cats-on-device     | False
======================================
Runtime[s]         | 125.49687790870667
======================================
```

**Output on `main`**:

```
Dask-NVTabular DLRM/Criteo benchmark
--------------------------------------
partition size     | 4261216256
protocol           | tcp
device(s)          | 0,1,2,3,4,5,6,7
rmm-pool-frac      | 0.8
out-files-per-proc | 3
num_io_threads     | 0
shuffle            | PER_PARTITION
cats-on-device     | False
======================================
Runtime[s]         | 434.0827877521515
======================================
```

**Environment details (please complete the following information):**
 - Environment location: Bare-metal
 - Method of NVTabular install: conda (cudf-21.08 conda package)




Discussion:

Resolving patch:
diff --git a/nvtabular/io/dataset.py b/nvtabular/io/dataset.py
--- a/nvtabular/io/dataset.py
+++ b/nvtabular/io/dataset.py
@@ -668,6 +668,7 @@ def to_parquet(
         suffix=".parquet",
         partition_on=None,
         method="subgraph",
+        write_hugectr_keyset=False,
     ):
         """Writes out to a parquet dataset
 
@@ -742,6 +743,10 @@ def to_parquet(
             a single large task). In some cases, it may be more ideal to prioritize
             concurrency. In that case, a worker-based approach can be used by
             specifying `method="worker"`.
+        write_hugectr_keyset : bool, optional
+            Whether to write a HugeCTR keyset output file ("_hugectr.keyset").
+            Writing this file can be very slow, and should only be done if you
+            are planning to ingest the output data with HugeCTR. Default is False.
         """
 
         if partition_on:
@@ -905,7 +910,7 @@ def to_parquet(
             self.cpu,
             suffix=suffix,
             partition_on=partition_on,
-            schema=self.schema,
+            schema=self.schema if write_hugectr_keyset else None,
         )
 
     def to_hugectr(
diff --git a/nvtabular/ops/categorify.py b/nvtabular/ops/categorify.py
--- a/nvtabular/ops/categorify.py
+++ b/nvtabular/ops/categorify.py
@@ -685,7 +685,6 @@ class FitOptions:
     max_size: Optional[Union[int, dict]] = None
     num_buckets: Optional[Union[int, dict]] = None
     start_index: int = 0
-    dataset_size: int = 0
 
     def __post_init__(self):
         if not isinstance(self.col_groups, ColumnSelector):
@@ -742,7 +741,16 @@ def _top_level_groupby(df, options: FitOptions):
             df_gb = df[combined_col_selector.names].copy(deep=False)
 
         agg_dict = {}
-        agg_dict[cat_col_selector.names[0]] = ["count"]
+        base_aggs = []
+        if "size" in options.agg_list:
+            # This is either for a Categorify operation,
+            # or "size" is in the list of aggregations
+            base_aggs.append("size")
+        if set(options.agg_list).difference({"size", "min", "max"}):
+            # This is a groupby aggregation that may
+            # require "count" statistics
+            base_aggs.append("count")
+        agg_dict[cat_col_selector.names[0]] = base_aggs
         if isinstance(options.agg_cols, list):
             options.agg_cols = ColumnSelector(options.agg_cols)
         for col in options.agg_cols.names:
@@ -807,13 +815,15 @@ def _mid_level_groupby(dfs, col_selector: ColumnSelector, freq_limit_val, option
     gb.reset_index(drop=False, inplace=True)
 
     name_count = _make_name(*(col_selector.names + ["count"]), sep=options.name_sep)
+    name_size = _make_name(*(col_selector.names + ["size"]), sep=options.name_sep)
     if options.freq_limit and not options.max_size:
-        gb = gb[gb[name_count] >= freq_limit_val]
+        gb = gb[gb[name_size] >= freq_limit_val]
 
     required = col_selector.names.copy()
     if "count" in options.agg_list:
         required.append(name_count)
-
+    if "size" in options.agg_list:
+        required.append(name_size)
     ddof = 1
     if isinstance(options.agg_cols, list):
         options.agg_cols = ColumnSelector(options.agg_cols)
@@ -959,10 +969,10 @@ def _write_uniques(dfs, base_path, col_selector: ColumnSelector, options: FitOpt
         new_cols = {}
         nulls_missing = False
         for col in col_selector.names:
-            name_count = col + "_count"
+            name_size = col + "_size"
             null_size = 0
-            if name_count in df:
-                null_size = options.dataset_size - df[name_count].sum()
+            if name_size in df:
+                null_size = df[name_size].iloc[0]
             if options.max_size:
                 max_emb_size = options.max_size
                 if isinstance(options.max_size, dict):
@@ -979,33 +989,32 @@ def _write_uniques(dfs, base_path, col_selector: ColumnSelector, options: FitOpt
                     raise ValueError("`nlargest` cannot be 0 or negative")
 
                 if nlargest < len(df):
-                    df = df.nlargest(n=nlargest, columns=name_count)
+                    df = df.nlargest(n=nlargest, columns=name_size)
             if not dispatch._series_has_nulls(df[col]):
-                if name_count in df:
-                    df = df.sort_values(name_count, ascending=False, ignore_index=True)
+                if name_size in df:
+                    df = df.sort_values(name_size, ascending=False, ignore_index=True)
 
                 nulls_missing = True
                 new_cols[col] = _concat(
                     [_nullable_series([None], df, df[col].dtype), df[col]],
                     ignore_index=True,
                 )
-                if name_count in df:
-                    new_cols[name_count] = _concat(
-                        [_nullable_series([null_size], df, df[name_count].dtype), df[name_count]],
+                if name_size in df:
+                    new_cols[name_size] = _concat(
+                        [_nullable_series([null_size], df, df[name_size].dtype), df[name_size]],
                         ignore_index=True,
                     )
 
             else:
                 # ensure None aka "unknown" stays at index 0
-                if name_count in df:
+                if name_size in df:
                     df_0 = df.iloc[0:1]
-                    df_0[name_count] = null_size
-                    df_1 = df.iloc[1:].sort_values(name_count, ascending=False, ignore_index=True)
+                    df_1 = df.iloc[1:].sort_values(name_size, ascending=False, ignore_index=True)
                     df = _concat([df_0, df_1])
                 new_cols[col] = df[col].copy(deep=False)
 
-                if name_count in df:
-                    new_cols[name_count] = df[name_count].copy(deep=False)
+                if name_size in df:
+                    new_cols[name_size] = df[name_size].copy(deep=False)
         if nulls_missing:
             df = type(df)(new_cols)
         df.to_parquet(path, index=False, compression=None)
@@ -1028,8 +1037,10 @@ def _groupby_to_disk(ddf, write_func, options: FitOptions):
         return {}
 
     if options.concat_groups:
-        if options.agg_list and options.agg_list != ["count"]:
-            raise ValueError("Cannot use concat_groups=True with aggregations other than count")
+        if options.agg_list and not set(options.agg_list).issubset({"count", "size"}):
+            raise ValueError(
+                "Cannot use concat_groups=True with aggregations other than count and size"
+            )
         if options.agg_cols:
             raise ValueError("Cannot aggregate continuous-column stats with concat_groups=True")
 
@@ -1068,7 +1079,6 @@ def _groupby_to_disk(ddf, write_func, options: FitOptions):
     level_2_name = "level_2-" + token
     level_3_name = "level_3-" + token
     finalize_labels_name = options.stat_name + "-" + token
-    options.dataset_size = len(ddf)
     for p in range(ddf.npartitions):
         dsk[(level_1_name, p)] = (_top_level_groupby, (ddf._name, p), options)
         k = 0
@@ -1120,7 +1130,7 @@ def _groupby_to_disk(ddf, write_func, options: FitOptions):
 def _category_stats(ddf, options: FitOptions):
     # Check if we only need categories
     if options.agg_cols == [] and options.agg_list == []:
-        options.agg_list = ["count"]
+        options.agg_list = ["size"]
         return _groupby_to_disk(ddf, _write_uniques, options)
 
     # Otherwise, getting category-statistics


Test patch:
diff --git a/tests/unit/ops/test_categorify.py b/tests/unit/ops/test_categorify.py
--- a/tests/unit/ops/test_categorify.py
+++ b/tests/unit/ops/test_categorify.py
@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-import math
 import os
 import random
 
@@ -37,7 +36,7 @@
 
 @pytest.mark.parametrize("cpu", _CPU)
 @pytest.mark.parametrize("include_nulls", [True, False])
-def test_categorify_counts(tmpdir, cpu, include_nulls):
+def test_categorify_size(tmpdir, cpu, include_nulls):
     num_rows = 50
     num_distinct = 10
 
@@ -62,21 +61,21 @@ def test_categorify_counts(tmpdir, cpu, include_nulls):
     if cpu:
         expected = dict(zip(vals.index, vals))
         computed = {
-            session: count
-            for session, count in zip(vocab["session_id"], vocab["session_id_count"])
-            if count
+            session: size
+            for session, size in zip(vocab["session_id"], vocab["session_id_size"])
+            if size
         }
     else:
         expected = dict(zip(vals.index.values_host, vals.values_host))
         computed = {
-            session: count
-            for session, count in zip(
-                vocab["session_id"].values_host, vocab["session_id_count"].values_host
+            session: size
+            for session, size in zip(
+                vocab["session_id"].values_host, vocab["session_id_size"].values_host
             )
-            if count
+            if size
         }
     first_key = list(computed.keys())[0]
-    if math.isnan(first_key):
+    if pd.isna(first_key):
         computed.pop(first_key)
     assert computed == expected
 
@@ -104,8 +103,8 @@ def test_na_value_count(tmpdir):
     second_cat = dispatch._read_dispatch("./categories/unique.productID.parquet")(
         "./categories/unique.productID.parquet"
     )
-    assert single_cat["brand_count"][0] == 5
-    assert second_cat["productID_count"][0] == 3
+    assert single_cat["brand_size"][0] == 5
+    assert second_cat["productID_size"][0] == 3
 
 
 @pytest.mark.parametrize("freq_threshold", [0, 1, 2])
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
@@ -346,6 +346,7 @@ def test_hugectr(
             output_path=outdir,
             out_files_per_proc=nfiles,
             num_threads=num_io_threads,
+            write_hugectr_keyset=True,
         )
 
     # Check for _file_list.txt