Repository: NVIDIA-Merlin/NVTabular

Issue:
[BUG] Unique parquet files are not accurately generated after Categorify() op
**Describe the bug**
After Categorify() op when I check out the unique parquet files, I see that original id column, and count column in the parquet file do not match.  Somehow an existing element in the categorical column is encoded to `NA` and the last elements valu_count also appears as `NA` meaning that first column is somehow shifted 1 row and then this creates mismatch between two columns.   See screenshot below.



![image](https://user-images.githubusercontent.com/16246900/133690005-81400ea4-579c-432a-a267-f27b12bf642d.png)


**Steps/Code to reproduce bug**
Run the code below to reproduce the issue:
```
NUM_ROWS = 100000
long_tailed_item_distribution = np.clip(np.random.lognormal(3., 1., NUM_ROWS).astype(np.int32), 1, 50000)

# generate random item interaction features 
df = pd.DataFrame(np.random.randint(1, 20000, NUM_ROWS), columns=['session_id'])
df['item_id'] = long_tailed_item_distribution

# generate category mapping for each item-id
df['category'] = pd.cut(df['item_id'], bins=334, labels=np.arange(1, 335)).astype(np.int32)
df['timestamp/age_days'] = np.random.uniform(0, 1, NUM_ROWS)
df['timestamp/weekday/sin']= np.random.uniform(0, 1, NUM_ROWS)

# generate day mapping for each session 
map_day = dict(zip(df.session_id.unique(), np.random.randint(1, 10, size=(df.session_id.nunique()))))
df['day'] =  df.session_id.map(map_day)

categ_feats = ['session_id', 'item_id', 'category'] >> nvt.ops.Categorify()
workflow = nvt.Workflow(categ_feats)
dataset = nvt.Dataset(df, cpu=False)
workflow.fit(dataset)
gdf = workflow.transform(dataset).to_ddf().compute()

unique_sessionid=cudf.read_parquet('./categories/unique.session_id.parquet')
print(unique_sessionid)
```

**Expected behavior**
Correct encoding is expected.

**Environment details (please complete the following information):**
 - Environment location: [Bare-metal, Docker, Cloud(specify cloud provider)]
 - Method of NVTabular install: [conda, Docker, or from source]: NVT 0.6 Pytorch Docker latest main branch is pulled.
   - If method of install is [Docker], provide `docker pull` & `docker run` commands used
 



Discussion:
Same reproduction - reducing the number of rows and getting rid of the unneeded columns:

```
import numpy as np
import pandas as pd
import nvtabular as nvt
import cudf

NUM_ROWS = 50

# generate random item interaction features 
df = pd.DataFrame(np.random.randint(1, NUM_ROWS//4, NUM_ROWS), columns=['session_id'])

categ_feats = ['session_id'] >> nvt.ops.Categorify()
workflow = nvt.Workflow(categ_feats)
dataset = nvt.Dataset(df, cpu=False)
workflow.fit(dataset)
gdf = workflow.transform(dataset).to_ddf().compute()

unique_sessionid=cudf.read_parquet('./categories/unique.session_id.parquet')
print("got")
print(unique_sessionid)

print("expected")
print(df["session_id"].value_counts())
```

This prints out:

```python
got
   session_id session_id_count
0        <NA>                7
1           4                6
2           2                6
3           3                6
4           7                5
5           6                5
6           9                4
7           5                4
8           8                4
9          10                3
10         11             <NA>
expected
4     7
7     6
2     6
3     6
9     5
6     5
8     4
5     4
10    4
11    3
Name: session_id, dtype: int64
```

Showing that session '11' should have value counts of 3, but is instead getting NA

Resolving patch:
diff --git a/nvtabular/dispatch.py b/nvtabular/dispatch.py
--- a/nvtabular/dispatch.py
+++ b/nvtabular/dispatch.py
@@ -379,9 +379,8 @@ def _make_df(_like_df=None, device=None):
         return pd.DataFrame(_like_df)
     elif isinstance(_like_df, (cudf.DataFrame, cudf.Series)):
         return cudf.DataFrame(_like_df)
-    elif isinstance(_like_df, dict) and len(_like_df) > 0:
+    elif device is None and isinstance(_like_df, dict) and len(_like_df) > 0:
         is_pandas = all(isinstance(v, pd.Series) for v in _like_df.values())
-
         return pd.DataFrame(_like_df) if is_pandas else cudf.DataFrame(_like_df)
     if device == "cpu":
         return pd.DataFrame(_like_df)
diff --git a/nvtabular/ops/categorify.py b/nvtabular/ops/categorify.py
--- a/nvtabular/ops/categorify.py
+++ b/nvtabular/ops/categorify.py
@@ -981,6 +981,12 @@ def _write_uniques(dfs, base_path, col_selector: ColumnSelector, options: FitOpt
                     [_nullable_series([None], df, df[col].dtype), df[col]],
                     ignore_index=True,
                 )
+                if name_count in df:
+                    new_cols[name_count] = _concat(
+                        [_nullable_series([0], df, df[name_count].dtype), df[name_count]],
+                        ignore_index=True,
+                    )
+
             else:
                 # ensure None aka "unknown" stays at index 0
                 if name_count in df:
@@ -988,8 +994,9 @@ def _write_uniques(dfs, base_path, col_selector: ColumnSelector, options: FitOpt
                     df_1 = df.iloc[1:].sort_values(name_count, 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_count in df:
+                    new_cols[name_count] = df[name_count].copy(deep=False)
         if nulls_missing:
             df = type(df)(new_cols)
         df.to_parquet(path, index=False, compression=None)
@@ -998,6 +1005,7 @@ def _write_uniques(dfs, base_path, col_selector: ColumnSelector, options: FitOpt
         for c in col_selector.names:
             df_null[c] = df_null[c].astype(df[c].dtype)
         df_null.to_parquet(path, index=False, compression=None)
+
     del df
     return path
 


Test patch:
diff --git a/tests/unit/ops/test_ops.py b/tests/unit/ops/test_ops.py
--- a/tests/unit/ops/test_ops.py
+++ b/tests/unit/ops/test_ops.py
@@ -15,6 +15,8 @@
 #
 import copy
 import math
+import os
+import random
 import string
 
 import dask.dataframe as dd
@@ -469,6 +471,50 @@ def test_lambdaop_misalign(cpu):
     )
 
 
+@pytest.mark.parametrize("cpu", _CPU)
+@pytest.mark.parametrize("include_nulls", [True, False])
+def test_categorify_counts(tmpdir, cpu, include_nulls):
+    num_rows = 50
+    num_distinct = 10
+
+    possible_session_ids = list(range(num_distinct))
+    if include_nulls:
+        possible_session_ids.append(None)
+
+    df = dispatch._make_df(
+        {"session_id": [random.choice(possible_session_ids) for _ in range(num_rows)]},
+        device="cpu" if cpu else None,
+    )
+
+    cat_features = ["session_id"] >> nvt.ops.Categorify(out_path=str(tmpdir))
+    workflow = nvt.Workflow(cat_features)
+    workflow.fit_transform(nvt.Dataset(df, cpu=cpu)).to_ddf().compute()
+
+    vals = df["session_id"].value_counts()
+    vocab = dispatch._read_dispatch(cpu=cpu)(
+        os.path.join(tmpdir, "categories", "unique.session_id.parquet")
+    )
+
+    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
+        }
+    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
+            )
+            if count
+        }
+
+    assert computed == expected
+
+
 @pytest.mark.parametrize("freq_threshold", [0, 1, 2])
 @pytest.mark.parametrize("cpu", _CPU)
 @pytest.mark.parametrize("dtype", [None, np.int32, np.int64])