Repository: NVIDIA-Merlin/NVTabular (language: python)

Issue:
[FEA] Need optimized column dtypes outputted from NVTabular workflow
**Is your feature request related to a problem? Please describe.**

I am trying to do the following with NVTabular via an LambdaOp, which does not work bcs of the error here: https://github.com/NVIDIA-Merlin/NVTabular/issues/1396.
```
def reduce_mem(col):
    if np.issubdtype(col.dtype, np.integer):
        if (col.min()>-32768)&(col.max()<32768):
            col = col.astype('int16')
        elif (col.min()>-2147483648)&(col.max()<2147483648):
            col = col.astype('int32')
    if np.issubdtype(col.dtype, np.float):
        col = col.astype('float32')
    return col

reduced_cols = cols >> nvt.ops.LambdaOp(reduce_mem)

dataset = nvt.Dataset(data_path, part_size='500MB')
workflow = nvt.Workflow(reduced_cols)
workflow.fit(dataset)
# Write to new "shuffled" and "processed" dataset
workflow.transform(dataset).to_parquet(
    OUTPUT_FOLDER,
    out_files_per_proc=1,
    shuffle=False
)

```

**Describe the solution you'd like**

Reducing the column dtypes is very common practice in DS pipeline, and it should be a feature of NVT to use optimized data types. We should not try to create a LambdaOP or a custom op to integrate that process in our pipeline. It'd be very useful if it comes naturally. 




Discussion / hints:

Fix patch:
diff --git a/nvtabular/ops/__init__.py b/nvtabular/ops/__init__.py
index eed2bde6019..ce7e08fa179 100644
--- a/nvtabular/ops/__init__.py
+++ b/nvtabular/ops/__init__.py
@@ -36,6 +36,7 @@
 from .logop import LogOp
 from .normalize import Normalize, NormalizeMinMax
 from .operator import ColumnSelector, Operator
+from .reduce_dtype_size import ReduceDtypeSize
 from .rename import Rename
 from .stat_operator import StatOperator
 from .target_encoding import TargetEncoding
diff --git a/nvtabular/ops/reduce_dtype_size.py b/nvtabular/ops/reduce_dtype_size.py
new file mode 100644
index 00000000000..f916e564a23
--- /dev/null
+++ b/nvtabular/ops/reduce_dtype_size.py
@@ -0,0 +1,80 @@
+#
+# Copyright (c) 2021, 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 dask.dataframe as dd
+import numpy as np
+
+from nvtabular.graph.schema import Schema
+
+from ..dispatch import DataFrameType, annotate
+from .operator import ColumnSelector
+from .stat_operator import StatOperator
+
+_INT_DTYPES = [np.int8, np.int16, np.int32, np.int64]
+
+
+class ReduceDtypeSize(StatOperator):
+    """
+    ReduceDtypeSize changes the dtypes of numeric columns. For integer columns
+    this will choose a dtype such that the minimum and maximum values in the
+    column will fit. For float columns this will cast to a float32.
+    """
+
+    def __init__(self, float_dtype=np.float32):
+        super().__init__()
+        self.float_dtype = float_dtype
+        self.ranges = {}
+        self.dtypes = {}
+
+    @annotate("reduce_dtype_size_fit", color="green", domain="nvt_python")
+    def fit(self, col_selector: ColumnSelector, ddf: dd.DataFrame):
+        return {col: (ddf[col].min(), ddf[col].max()) for col in col_selector.names}
+
+    def fit_finalize(self, dask_stats):
+        self.ranges = dask_stats
+
+    def clear(self):
+        self.dtypes = {}
+        self.ranges = {}
+
+    @annotate("reduce_dtype_size_transform", color="darkgreen", domain="nvt_python")
+    def transform(self, col_selector: ColumnSelector, df: DataFrameType) -> DataFrameType:
+        for col, dtype in self.dtypes.items():
+            df[col] = df[col].astype(dtype)
+        return df
+
+    def compute_output_schema(self, input_schema, selector, prev_output_schema=None):
+        if not self.ranges:
+            return input_schema
+
+        output_columns = []
+        for column, (min_value, max_value) in self.ranges.items():
+            column = input_schema[column]
+
+            dtype = column.dtype
+            if np.issubdtype(column.dtype, np.integer):
+                for possible_dtype in _INT_DTYPES:
+                    dtype_range = np.iinfo(possible_dtype)
+                    if min_value >= dtype_range.min and max_value <= dtype_range.max:
+                        dtype = possible_dtype
+                        break
+
+            elif np.issubdtype(column.dtype, np.float):
+                dtype = self.float_dtype
+
+            output_columns.append(column.with_dtype(dtype))
+
+        self.dtypes = {column.name: column.dtype for column in output_columns}
+        return Schema(output_columns)


Test patch:
diff --git a/tests/unit/ops/test_reduce_dtype_size.py b/tests/unit/ops/test_reduce_dtype_size.py
new file mode 100644
index 00000000000..59154b58e7d
--- /dev/null
+++ b/tests/unit/ops/test_reduce_dtype_size.py
@@ -0,0 +1,49 @@
+#
+# Copyright (c) 2021, 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 numpy as np
+import pandas as pd
+import pytest
+
+import nvtabular as nvt
+from tests.conftest import assert_eq
+
+try:
+    import cudf
+
+    _CPU = [True, False]
+except ImportError:
+    _CPU = [True]
+
+
+@pytest.mark.parametrize("cpu", _CPU)
+def test_reduce_size(tmpdir, cpu):
+    df = pd.DataFrame()
+    if not cpu:
+        df = cudf.DataFrame(df)
+
+    df["int16"] = np.array([2 ** 15 - 1, 0], dtype="int64")
+    df["int32"] = np.array([2 ** 30, -(2 ** 30)], dtype="int64")
+    df["int64"] = np.array([2 ** 60, -(2 ** 60)], dtype="int64")
+    df["float32"] = np.array([1.0, 2.0], dtype="float64")
+
+    workflow = nvt.Workflow(list(df.columns) >> nvt.ops.ReduceDtypeSize())
+    transformed = workflow.fit_transform(nvt.Dataset(df)).to_ddf().compute()
+
+    expected = df
+    for column in df:
+        expected[column] = expected[column].astype(column)
+
+    assert_eq(expected, transformed)