diff --git a/python/cudf/cudf/_lib/transpose.pyx b/python/cudf/cudf/_lib/transpose.pyx index b33a3cefba7..931a2702612 100644 --- a/python/cudf/cudf/_lib/transpose.pyx +++ b/python/cudf/cudf/_lib/transpose.pyx @@ -28,20 +28,19 @@ def transpose(source): return source cats = None - dtype = source._columns[0].dtype + columns = source._columns + dtype = columns[0].dtype if is_categorical_dtype(dtype): - if any(not is_categorical_dtype(c.dtype) for c in source._columns): + if any(not is_categorical_dtype(c.dtype) for c in columns): raise ValueError('Columns must all have the same dtype') - cats = list(c.categories for c in source._columns) + cats = list(c.categories for c in columns) cats = cudf.core.column.concat_columns(cats).unique() source = cudf.core.frame.Frame(index=source._index, data=[ (name, col._set_categories(cats, is_unique=True).codes) for name, col in source._data.items() ]) - elif dtype.kind in 'OU': - raise NotImplementedError('Cannot transpose string columns') - elif any(c.dtype != dtype for c in source._columns): + elif any(c.dtype != dtype for c in columns): raise ValueError('Columns must all have the same dtype') cdef pair[unique_ptr[column], table_view] c_result diff --git a/python/cudf/cudf/tests/test_dataframe.py b/python/cudf/cudf/tests/test_dataframe.py index 33c993cc56a..f42920b7c50 100644 --- a/python/cudf/cudf/tests/test_dataframe.py +++ b/python/cudf/cudf/tests/test_dataframe.py @@ -1769,7 +1769,7 @@ def test_dataframe_shape_empty(): @pytest.mark.parametrize("num_cols", [1, 2, 10]) @pytest.mark.parametrize("num_rows", [1, 2, 20]) -@pytest.mark.parametrize("dtype", dtypes) +@pytest.mark.parametrize("dtype", dtypes + ["object"]) @pytest.mark.parametrize("nulls", ["none", "some", "all"]) def test_dataframe_transpose(nulls, num_cols, num_rows, dtype): # In case of `bool` dtype: pandas <= 1.2.5 type-casts