Revert changes to dispatch.py. Now part of NVIDIA-Merlin#244

Author: oliverholworthy

merlin/core/dispatch.py

@@ -24,16 +24,16 @@
 import pyarrow as pa
 import pyarrow.parquet as pq
 
-# unused HAS_GPU import is here for backwards compatibility
-from merlin.core.compat import HAS_GPU  # pylint: disable=unused-import # noqa: F401
-from merlin.core.compat import cupy as cp
+from merlin.core.compat import HAS_GPU
 from merlin.core.protocols import DataFrameLike, DictLike, SeriesLike
 
 cudf = None
+cp = None
 rmm = None
 
-if cudf:
+if HAS_GPU:
     try:
+        import cudf  # type: ignore[no-redef]
         import dask_cudf
         import rmm  # type: ignore[no-redef]
         from cudf.core.column import as_column, build_column
@@ -48,6 +48,10 @@
             from cudf.utils.dtypes import is_string_dtype as cudf_is_string_dtype
     except ImportError:
         pass
+    try:
+        import cupy as cp  # type: ignore[no-redef]
+    except ImportError:
+        pass
 
 try:
     # Dask >= 2021.5.1
@@ -73,7 +77,7 @@ def inner2(*args, **kwargs):
         return inner1
 
 
-if cudf:
+if HAS_GPU and cudf:
     DataFrameType = Union[pd.DataFrame, cudf.DataFrame]  # type: ignore
     SeriesType = Union[pd.Series, cudf.Series]  # type: ignore
 else:
@@ -120,7 +124,7 @@ def create_merlin_dataset(df):
 
 def read_parquet_metadata(path):
     """Read parquet metadata from path"""
-    if cudf:
+    if HAS_GPU:
         return cudf.io.read_parquet_metadata(path)
     full_meta = pq.read_metadata(path)
     pf = pq.ParquetFile(path)
@@ -129,7 +133,7 @@ def read_parquet_metadata(path):
 
 def get_lib():
     """Dispatch to the appropriate library (cudf or pandas) for the current environment"""
-    return cudf or pd
+    return cudf if HAS_GPU else pd
 
 
 def reinitialize(managed_memory=False):
@@ -140,15 +144,15 @@ def reinitialize(managed_memory=False):
 
 def random_uniform(size):
     """Dispatch for numpy.random.RandomState"""
-    if cp:
+    if HAS_GPU:
         return cp.random.uniform(size=size)
     else:
         return np.random.uniform(size=size)
 
 
 def coo_matrix(data, row, col):
     """Dispatch for scipy.sparse.coo_matrix"""
-    if cp:
+    if HAS_GPU:
         return cp.sparse.coo_matrix((data, row, col))
     else:
         import scipy
@@ -159,9 +163,9 @@ def coo_matrix(data, row, col):
 def is_dataframe_object(x):
     # Simple check if object is a cudf or pandas
     # DataFrame object
-    if cudf:
-        return isinstance(x, (cudf.DataFrame, pd.DataFrame))
-    return isinstance(x, pd.DataFrame)
+    if not HAS_GPU:
+        return isinstance(x, pd.DataFrame)
+    return isinstance(x, (cudf.DataFrame, pd.DataFrame))
 
 
 def nullable_series(data, like_df, dtype):
@@ -178,9 +182,9 @@ def nullable_series(data, like_df, dtype):
 def is_series_object(x):
     # Simple check if object is a cudf or pandas
     # Series object
-    if cudf:
-        return isinstance(x, (cudf.Series, pd.Series))
-    return isinstance(x, pd.Series)
+    if not HAS_GPU:
+        return isinstance(x, pd.Series)
+    return isinstance(x, (cudf.Series, pd.Series))
 
 
 def is_cpu_object(x):
@@ -213,78 +217,50 @@ def pd_convert_hex(x):
 
 def random_state(seed, like_df=None):
     """Dispatch for numpy.random.RandomState"""
-    if isinstance(like_df, (pd.DataFrame, pd.Series)):
+    if not HAS_GPU or isinstance(like_df, (pd.DataFrame, pd.Series)):
         return np.random.RandomState(seed)
-    elif cudf and isinstance(like_df, (cudf.DataFrame, cudf.Series)):
-        return cp.random.RandomState(seed)
     else:
-        return ValueError(
-            "Unsupported dataframe type: "
-            f"{type(like_df)}"
-            " Expected either a pandas or cudf DataFrame or Series."
-        )
+        return cp.random.RandomState(seed)
 
 
 def arange(size, like_df=None, dtype=None):
     """Dispatch for numpy.arange"""
-    if isinstance(like_df, (np.ndarray, pd.DataFrame, pd.Series)):
+    if not HAS_GPU or isinstance(like_df, (np.ndarray, pd.DataFrame, pd.Series)):
         return np.arange(size, dtype=dtype)
-    elif cudf and isinstance(like_df, (cp.ndarray, cudf.DataFrame, cudf.Series)):
-        return cp.arange(size, dtype=dtype)
     else:
-        return ValueError(
-            "Unsupported dataframe type: "
-            f"{type(like_df)}"
-            " Expected either a pandas or cudf DataFrame or Series."
-        )
+        return cp.arange(size, dtype=dtype)
 
 
 def array(x, like_df=None, dtype=None):
     """Dispatch for numpy.array"""
-    if isinstance(like_df, (np.ndarray, pd.DataFrame, pd.Series)):
+    if not HAS_GPU or isinstance(like_df, (np.ndarray, pd.DataFrame, pd.Series)):
         return np.array(x, dtype=dtype)
-    elif cudf and isinstance(like_df, (cp.ndarray, cudf.DataFrame, cudf.Series)):
-        return cp.array(x, dtype=dtype)
     else:
-        return ValueError(
-            "Unsupported dataframe type: "
-            f"{type(like_df)}"
-            " Expected either a pandas or cudf DataFrame or Series."
-        )
+        return cp.array(x, dtype=dtype)
 
 
 def zeros(size, like_df=None, dtype=None):
     """Dispatch for numpy.array"""
-    if isinstance(like_df, (np.ndarray, pd.DataFrame, pd.Series)):
+    if not HAS_GPU or isinstance(like_df, (np.ndarray, pd.DataFrame, pd.Series)):
         return np.zeros(size, dtype=dtype)
-    elif cudf and isinstance(like_df, (cp.ndarray, cudf.DataFrame, cudf.Series)):
-        return cp.zeros(size, dtype=dtype)
     else:
-        return ValueError(
-            "Unsupported dataframe type: "
-            f"{type(like_df)}"
-            " Expected either a pandas or cudf DataFrame or Series."
-        )
+        return cp.zeros(size, dtype=dtype)
 
 
 def hash_series(ser):
     """Row-wise Series hash"""
-    if isinstance(ser, pd.Series):
+    if not HAS_GPU or isinstance(ser, pd.Series):
         # Using pandas hashing, which does not produce the
         # same result as cudf.Series.hash_values().  Do not
         # expect hash-based data transformations to be the
         # same on CPU and CPU.  TODO: Fix this (maybe use
         # murmurhash3 manually on CPU).
         return hash_object_dispatch(ser).values
-    elif cudf and isinstance(ser, cudf.Series):
+    else:
         if is_list_dtype(ser):
             return ser.list.leaves.hash_values()
         else:
             return ser.hash_values()
-    else:
-        return ValueError(
-            "Unsupported series type: " f"{type(ser)}" " Expected either a pandas or cudf Series."
-        )
 
 
 def series_has_nulls(s):
@@ -333,11 +309,12 @@ def is_list_dtype(ser):
         if not len(ser):  # pylint: disable=len-as-condition
             return False
         return pd.api.types.is_list_like(ser.values[0])
-    elif cudf and isinstance(ser, cudf.Series):
-        return cudf_is_list_dtype(ser)
-    elif isinstance(ser, np.ndarray):
-        return pd.api.types.is_list_like(ser[0])
-    return pd.api.types.is_list_like(ser)
+    elif not HAS_GPU:
+        # either np.ndarray or a dtype
+        if isinstance(ser, np.ndarray):
+            ser = ser[0]
+        return pd.api.types.is_list_like(ser)
+    return cudf_is_list_dtype(ser)
 
 
 def is_string_dtype(dtype: np.dtype) -> bool:
@@ -353,21 +330,18 @@ def is_string_dtype(dtype: np.dtype) -> bool:
     bool
         `True` if the dtype of `obj` is a string type
     """
-    if cudf:
+    if not HAS_GPU:
+        return pd.api.types.is_string_dtype(dtype)
+    else:
         return cudf_is_string_dtype(dtype)
-    return pd.api.types.is_string_dtype(dtype)
 
 
 def flatten_list_column_values(s):
     """returns a flattened list from a list column"""
-    if isinstance(s, pd.Series):
+    if isinstance(s, pd.Series) or not cudf:
         return pd.Series(itertools.chain(*s))
-    elif cudf and isinstance(s, cudf.Series):
-        return s.list.leaves
     else:
-        return ValueError(
-            "Unsupported series type: " f"{type(s)}" " Expected either a pandas or cudf Series."
-        )
+        return s.list.leaves
 
 
 def flatten_list_column(s):
@@ -410,7 +384,7 @@ def read_dispatch(df: DataFrameLike = None, cpu=None, collection=False, fmt="par
     """Return the necessary read_parquet function to generate
     data of a specified type.
     """
-    if cpu or isinstance(df, pd.DataFrame):
+    if cpu or isinstance(df, pd.DataFrame) or not HAS_GPU:
         _mod = dd if collection else pd
     else:
         if collection:
@@ -511,16 +485,10 @@ def concat(objs, **kwargs):
     """dispatch function for concat"""
     if isinstance(objs[0], dd.DataFrame):
         return dd.multi.concat(objs)
-    elif isinstance(objs[0], (pd.DataFrame, pd.Series)):
+    elif isinstance(objs[0], (pd.DataFrame, pd.Series)) or not HAS_GPU:
         return pd.concat(objs, **kwargs)
-    elif cudf and isinstance(objs[0], (cudf.DataFrame, cudf.Series)):
-        return cudf.core.reshape.concat(objs, **kwargs)
     else:
-        return ValueError(
-            "Unsupported dataframe type: "
-            f"{type(objs[0])}"
-            " Expected a pandas, cudf, or dask DataFrame."
-        )
+        return cudf.core.reshape.concat(objs, **kwargs)
 
 
 def make_df(_like_df=None, device=None):
@@ -616,7 +584,7 @@ def convert_data(x, cpu=True, to_collection=None, npartitions=1):
             _x = x if isinstance(x, pd.DataFrame) else x.to_pandas()
             # Output a collection if `to_collection=True`
             return dd.from_pandas(_x, sort=False, npartitions=npartitions) if to_collection else _x
-    elif cudf and dask_cudf:
+    else:
         if isinstance(x, dd.DataFrame):
             # If input is a Dask collection, convert to dask_cudf
             if isinstance(x, dask_cudf.DataFrame):
@@ -639,30 +607,26 @@ def convert_data(x, cpu=True, to_collection=None, npartitions=1):
                 if to_collection
                 else _x
             )
-    else:
-        raise RuntimeError(
-            "Unable to move data to GPU. "
-            "cudf and dask_cudf are not available. "
-            "Make sure these packages are installed and can be imported in this environment. "
-        )
 
 
 def to_host(x):
     """Move cudf.DataFrame to host memory for caching.
 
     All other data will pass through unchanged.
     """
-    if cudf and isinstance(x, cudf.DataFrame):
+    if not HAS_GPU or isinstance(x, (pd.DataFrame, dd.DataFrame)):
+        return x
+    else:
         return x.to_arrow()
-    return x
 
 
 def from_host(x):
-    if isinstance(x, pd.DataFrame):
-        return cudf.DataFrame.from_pandas(x)
-    elif isinstance(x, pa.Table):
+    if not HAS_GPU:
+        return x
+    elif isinstance(x, cudf.DataFrame):
+        return x
+    else:
         return cudf.DataFrame.from_arrow(x)
-    return x
 
 
 def build_cudf_list_column(new_elements, new_offsets):
@@ -681,14 +645,14 @@ def build_cudf_list_column(new_elements, new_offsets):
     cudf.Series
         The list column with corresponding elements and row_lengths as a series.
     """
-    if cudf:
-        return build_column(
-            None,
-            dtype=cudf.core.dtypes.ListDtype(new_elements.dtype),
-            size=new_offsets.size - 1,
-            children=(as_column(new_offsets), as_column(new_elements)),
-        )
-    return []
+    if not HAS_GPU:
+        return []
+    return build_column(
+        None,
+        dtype=cudf.core.dtypes.ListDtype(new_elements.dtype),
+        size=new_offsets.size - 1,
+        children=(as_column(new_offsets), as_column(new_elements)),
+    )
 
 
 def build_pandas_list_column(elements, row_lengths):