Hack to fix test_spilling_local_cuda_cluster[True]

rapidsai · Dec 16, 2020 · df1c0fd · df1c0fd
1 parent 78353fc
commit df1c0fd
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/dask_cuda/tests/test_proxy.py b/dask_cuda/tests/test_proxy.py
@@ -1,6 +1,7 @@
 import operator
 import pickle
 
+import pandas
 import pytest
 from pandas.testing import assert_frame_equal
 
@@ -222,7 +223,11 @@ def task(x):
             ddf = dask_cudf.from_cudf(df, npartitions=1)
             ddf = ddf.map_partitions(task, meta=df.head())
             got = ddf.compute()
-            assert_frame_equal(got.to_pandas(), df.to_pandas())
+            if not isinstance(got, pandas.Series):
+                # TODO: `got` should be of type cudf.core.dataframe.DataFrame
+                # but in CI, `got` is a `pandas.Series`
+                got = got.to_pandas()
+            assert_frame_equal(got, df.to_pandas())
 
 
 class _PxyObjTest(proxy_object.ProxyObject):