Fix concat_columns for DataFrames with list features (#165)

* Update instance checks in `concat_columns` to use cudf/pd classes.

Adds test for concat_columns with list columns

* Support CPU-only environment by checking cudf and pandas separately

* Update import formatting

merlin/core/dispatch.py

@@ -350,9 +350,13 @@ def concat_columns(args: list):
     """Dispatch function to concatenate DataFrames with axis=1"""
     if len(args) == 1:
         return args[0]
-    elif isinstance(args[0], DataFrameLike):
-        _lib = cudf if HAS_GPU and isinstance(args[0], cudf.DataFrame) else pd
-        return _lib.concat(
+    elif cudf is not None and isinstance(args[0], cudf.DataFrame):
+        return cudf.concat(
+            [a.reset_index(drop=True) for a in args],
+            axis=1,
+        )
+    elif isinstance(args[0], pd.DataFrame):
+        return pd.concat(
             [a.reset_index(drop=True) for a in args],
             axis=1,
         )
@@ -361,12 +365,6 @@ def concat_columns(args: list):
         for arg in args:
             result.update(arg)
         return result
-    else:
-        _lib = cudf if HAS_GPU and isinstance(args[0], cudf.DataFrame) else pd
-        return _lib.concat(
-            [a.reset_index(drop=True) for a in args],
-            axis=1,
-        )
     return None
 
 

tests/unit/core/test_dispatch.py

@@ -16,17 +16,17 @@
 import numpy as np
 import pytest
 
-from merlin.core.dispatch import HAS_GPU, is_list_dtype, list_val_dtype, make_df
+from merlin.core.dispatch import HAS_GPU, concat_columns, is_list_dtype, list_val_dtype, make_df
 
 if HAS_GPU:
-    _CPU = [True, False]
+    _DEVICES = ["cpu", "gpu"]
 else:
-    _CPU = [True]
+    _DEVICES = ["cpu"]
 
 
-@pytest.mark.parametrize("cpu", _CPU)
-def test_list_dtypes(tmpdir, cpu):
-    df = make_df(device="cpu" if cpu else "gpu")
+@pytest.mark.parametrize("device", _DEVICES)
+def test_list_dtypes(tmpdir, device):
+    df = make_df(device=device)
     df["vals"] = [
         [[0, 1, 2], [3, 4], [5]],
     ]
@@ -35,3 +35,12 @@ def test_list_dtypes(tmpdir, cpu):
 
     assert is_list_dtype(df["vals"])
     assert list_val_dtype(df["vals"]) == np.dtype(np.int64)
+
+
+@pytest.mark.parametrize("device", _DEVICES)
+def test_concat_columns(device):
+    df1 = make_df({"a": [1, 2], "b": [[3], [4, 5]]}, device=device)
+    df2 = make_df({"c": [3, 4, 5]}, device=device)
+    data_frames = [df1, df2]
+    res = concat_columns(data_frames)
+    assert res.columns.to_list() == ["a", "b", "c"]