Create standard Merlin dtypes in the merlin.dtypes module

merlin/dag/dictarray.py

@@ -17,6 +17,7 @@
 
 import numpy as np
 
+import merlin.dtypes as md
 from merlin.core.protocols import SeriesLike
 
 
@@ -29,7 +30,7 @@ def __init__(self, values, dtype=None):
         super().__init__()
 
         self.values = values
-        self.dtype = dtype or values.dtype
+        self.dtype = md.dtype(dtype or values.dtype)
 
     def __getitem__(self, index):
         return self.values[index]

merlin/dag/executors.py

@@ -19,6 +19,7 @@
 import pandas as pd
 from dask.core import flatten
 
+import merlin.dtypes as md
 from merlin.core.dispatch import concat_columns, is_list_dtype, list_val_dtype
 from merlin.core.utils import (
     ensure_optimize_dataframe_graph,
@@ -177,17 +178,27 @@ def _transform_data(self, node, input_data, capture_dtypes=False):
 
                 if is_list:
                     col_dtype = list_val_dtype(col_series)
-                if hasattr(col_dtype, "as_numpy_dtype"):
-                    col_dtype = col_dtype.as_numpy_dtype()
-                elif hasattr(col_series, "numpy"):
+
+                # TODO: Add a utility that condenses the known methods of fetching dtypes
+                # from series/arrays into a single function, so that Tensorflow specific
+                # code doesn't leak into the executors
+                if not hasattr(col_dtype, "as_numpy_dtype") and hasattr(col_series, "numpy"):
                     col_dtype = col_series[0].cpu().numpy().dtype
 
                 output_data_schema = output_col_schema.with_dtype(col_dtype, is_list=is_list)
 
                 if capture_dtypes:
                     node.output_schema.column_schemas[col_name] = output_data_schema
                 elif len(output_data):
-                    if output_col_schema.dtype != output_data_schema.dtype:
+                    # Validate that the dtypes match but only if they both exist
+                    # (since schemas may not have all dtypes specified, especially
+                    # in the tests)
+                    if (
+                        output_col_schema.dtype
+                        and output_data_schema.dtype
+                        and output_col_schema.dtype != md.string
+                        and output_col_schema.dtype != output_data_schema.dtype
+                    ):
                         raise TypeError(
                             f"Dtype discrepancy detected for column {col_name}: "
                             f"operator {node.op.label} reported dtype "
@@ -269,11 +280,16 @@ def transform(
         columns = list(flatten(wfn.output_columns.names for wfn in nodes))
         columns += additional_columns if additional_columns else []
 
+        if isinstance(output_dtypes, dict):
+            for col_name, col_dtype in output_dtypes.items():
+                if col_dtype:
+                    output_dtypes[col_name] = md.dtype(col_dtype).to_numpy
+
         if isinstance(output_dtypes, dict) and isinstance(ddf._meta, pd.DataFrame):
             dtypes = output_dtypes
             output_dtypes = type(ddf._meta)({k: [] for k in columns})
-            for column, dtype in dtypes.items():
-                output_dtypes[column] = output_dtypes[column].astype(dtype)
+            for col_name, col_dtype in dtypes.items():
+                output_dtypes[col_name] = output_dtypes[col_name].astype(col_dtype)
 
         elif not output_dtypes:
             # TODO: constructing meta like this loses dtype information on the ddf

merlin/dtypes/__init__.py

@@ -0,0 +1,61 @@
+#
+# Copyright (c) 2022, 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.
+#
+
+# flake8: noqa
+from merlin.dtypes import mappings
+from merlin.dtypes.aliases import *
+from merlin.dtypes.base import DType
+from merlin.dtypes.registry import _dtype_registry
+
+# Convenience alias for registering dtypes
+register = _dtype_registry.register
+
+
+def dtype(external_dtype):
+    # If the supplied dtype is None, then there's not a default dtype we can
+    # universally translate to across frameworks, so raise an error and help
+    # the downstream developer figure out how to handle that case explicitly
+    if external_dtype is None:
+        raise TypeError(
+            "Merlin doesn't provide a default dtype mapping for `None`. "
+            "This differs from the Numpy behavior you may be expecting, "
+            "which treats `None` as an alias for `np.float64`. If you're "
+            "expecting this dtype to be non-`None`, there may be an issue "
+            "in upstream code. If you'd like to allow this dtype to be `None`, "
+            "you can use a `try/except` to catch this error."
+        )
+
+    # If the supplied dtype is already a Merlin dtype, then there's
+    # nothing for us to do and we can exit early
+    if isinstance(external_dtype, DType):
+        return external_dtype
+
+    # If not, attempt to apply all the registered Merlin dtype mappings.
+    # If we don't find a match with those, fall back on converting to
+    # a numpy dtype and trying to match that instead.
+    try:
+        return _dtype_registry.to_merlin(external_dtype)
+    except TypeError as base_exc:
+
+        try:
+            return _dtype_registry.to_merlin_via_numpy(external_dtype)
+        except TypeError as exc:
+            # If we fail to find a match even after we try converting to
+            # numpy, re-raise the original exception because it has more
+            # information about the original external dtype that's causing
+            # the problem. (We want to highlight that dtype, not whatever
+            # numpy dtype it was converted to in the interim.)
+            raise base_exc from exc

merlin/dtypes/aliases.py

@@ -0,0 +1,52 @@
+#
+# Copyright (c) 2022, 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.
+#
+
+from merlin.dtypes.base import DType, ElementType, ElementUnit
+
+# Unsigned Integer
+uint8 = DType("uint8", ElementType.UInt, 8)
+uint16 = DType("uint16", ElementType.UInt, 16)
+uint32 = DType("uint32", ElementType.UInt, 32)
+uint64 = DType("uint64", ElementType.UInt, 64)
+
+# Signed Integer
+int8 = DType("int8", ElementType.Int, 8, signed=True)
+int16 = DType("int16", ElementType.Int, 16, signed=True)
+int32 = DType("int32", ElementType.Int, 32, signed=True)
+int64 = DType("int64", ElementType.Int, 64, signed=True)
+
+# Float
+float16 = DType("float16", ElementType.Float, 16, signed=True)
+float32 = DType("float32", ElementType.Float, 32, signed=True)
+float64 = DType("float64", ElementType.Float, 64, signed=True)
+
+# Date/Time
+datetime64 = DType("datetime64", ElementType.DateTime, 64)
+datetime64Y = DType("datetime64[Y]", ElementType.DateTime, 64, ElementUnit.Year)
+datetime64M = DType("datetime64[M]", ElementType.DateTime, 64, ElementUnit.Month)
+datetime64D = DType("datetime64[D]", ElementType.DateTime, 64, ElementUnit.Day)
+datetime64h = DType("datetime64[h]", ElementType.DateTime, 64, ElementUnit.Hour)
+datetime64m = DType("datetime64[m]", ElementType.DateTime, 64, ElementUnit.Minute)
+datetime64s = DType("datetime64[s]", ElementType.DateTime, 64, ElementUnit.Second)
+datetime64ms = DType("datetime64[ms]", ElementType.DateTime, 64, ElementUnit.Millisecond)
+datetime64us = DType("datetime64[us]", ElementType.DateTime, 64, ElementUnit.Microsecond)
+datetime64ns = DType("datetime64[ns]", ElementType.DateTime, 64, ElementUnit.Nanosecond)
+
+# Miscellaneous
+string = DType("str", ElementType.String)
+boolean = DType("bool", ElementType.Bool)
+object_ = DType("object", ElementType.Object)
+unknown = DType("unknown", ElementType.Unknown)

merlin/dtypes/base.py

@@ -0,0 +1,127 @@
+#
+# Copyright (c) 2022, 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.
+#
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional
+
+from merlin.dtypes.registry import _dtype_registry
+
+
+class ElementType(Enum):
+    """
+    Merlin DType base types
+
+    Since a Merlin DType may describe a list, these are the either the types of
+    scalars or the types of list elements.
+    """
+
+    Bool = "bool"
+    Int = "int"
+    UInt = "uint"
+    Float = "float"
+    String = "string"
+    DateTime = "datetime"
+    Object = "object"
+    Unknown = "unknown"
+
+
+class ElementUnit(Enum):
+    """
+    Dtype units, used only for datetime types
+
+    Since a Merlin DType may describe a list, these are the either the units of
+    scalars or the units of list elements.
+    """
+
+    Year = "year"
+    Month = "month"
+    Day = "day"
+    Hour = "hour"
+    Minute = "minute"
+    Second = "second"
+    Millisecond = "millisecond"
+    Microsecond = "microsecond"
+    Nanosecond = "nanosecond"
+
+
+@dataclass(eq=True, frozen=True)
+class DType:
+    """
+    Merlin dtypes are objects of this dataclass
+    """
+
+    name: str
+    element_type: ElementType
+    element_size: Optional[int] = None
+    element_unit: Optional[ElementUnit] = None
+    signed: Optional[bool] = None
+
+    def to(self, mapping_name: str):
+        """
+        Convert this Merlin dtype to another framework's dtypes
+
+        Parameters
+        ----------
+        mapping_name : str
+            Name of the framework dtype mapping to apply
+
+        Returns
+        -------
+        Any
+            An external framework dtype object
+
+        Raises
+        ------
+        ValueError
+            If there is no registered mapping for the given framework name
+        ValueError
+            The registered mapping for the given framework name doesn't map
+            this Merlin dtype to a framework dtype
+        """
+        try:
+            mapping = _dtype_registry.mappings[mapping_name]
+        except KeyError as exc:
+            raise ValueError(
+                f"Merlin doesn't have a registered dtype mapping for '{mapping_name}'. "
+                "If you'd like to register a new dtype mapping, use `merlin.dtype.register()`. "
+                "If you're expecting this mapping to already exist, has the library or package "
+                "that defines the mapping been imported successfully?"
+            ) from exc
+
+        try:
+            return mapping.from_merlin(self)
+        except KeyError as exc:
+            raise ValueError(
+                f"The registered dtype mapping for {mapping_name} doesn't contain type {self.name}."
+            ) from exc
+
+    @property
+    def to_numpy(self):
+        return self.to("numpy")
+
+    @property
+    def to_python(self):
+        return self.to("python")
+
+    # These properties refer to a single scalar (potentially a list element)
+    @property
+    def is_integer(self):
+        return self.element_type.value == "int"
+
+    @property
+    def is_float(self):
+        return self.element_type.value == "float"