Add transpose API to pylibcudf

docs/cudf/source/user_guide/api_docs/pylibcudf/index.rst

@@ -37,6 +37,7 @@ This page provides API documentation for pylibcudf.
     table
     traits
     transform
+    transpose
     types
     unary
 

docs/cudf/source/user_guide/api_docs/pylibcudf/transpose.rst

@@ -0,0 +1,6 @@
+=========
+transpose
+=========
+
+.. automodule:: pylibcudf.transpose
+   :members:

python/cudf/cudf/_lib/transpose.pyx

@@ -1,32 +1,15 @@
 # Copyright (c) 2020-2024, NVIDIA CORPORATION.
 
-from libcpp.memory cimport unique_ptr
-from libcpp.pair cimport pair
-from libcpp.utility cimport move
-
-from pylibcudf.libcudf.column.column cimport column
-from pylibcudf.libcudf.table.table_view cimport table_view
-from pylibcudf.libcudf.transpose cimport transpose as cpp_transpose
+import pylibcudf as plc
 
 from cudf._lib.column cimport Column
-from cudf._lib.utils cimport columns_from_table_view, table_view_from_columns
 
 
 def transpose(list source_columns):
     """Transpose m n-row columns into n m-row columns
     """
-    cdef pair[unique_ptr[column], table_view] c_result
-    cdef table_view c_input = table_view_from_columns(source_columns)
-
-    with nogil:
-        c_result = move(cpp_transpose(c_input))
-
-    # Notice, the data pointer of `result_owner` has been exposed
-    # through `c_result.second` at this point.
-    result_owner = Column.from_unique_ptr(
-        move(c_result.first), data_ptr_exposed=True
-    )
-    return columns_from_table_view(
-        c_result.second,
-        owners=[result_owner] * c_result.second.num_columns()
+    input_table = plc.table.Table(
+        [col.to_pylibcudf(mode="read") for col in source_columns]
     )
+    _, result_table = plc.transpose.transpose(input_table)
+    return [Column.from_pylibcudf(col) for col in result_table.columns()]

python/pylibcudf/pylibcudf/CMakeLists.txt

@@ -43,6 +43,7 @@ set(cython_sources
     table.pyx
     traits.pyx
     transform.pyx
+    transpose.pyx
     types.pyx
     unary.pyx
     utils.pyx

python/pylibcudf/pylibcudf/__init__.pxd

@@ -28,6 +28,7 @@ from . cimport (
     strings,
     traits,
     transform,
+    transpose,
     types,
     unary,
 )
@@ -71,6 +72,7 @@ __all__ = [
     "sorting",
     "traits",
     "transform",
+    "transpose",
     "types",
     "unary",
 ]

python/pylibcudf/pylibcudf/__init__.py

@@ -39,6 +39,7 @@
     strings,
     traits,
     transform,
+    transpose,
     types,
     unary,
 )
@@ -83,6 +84,7 @@
     "sorting",
     "traits",
     "transform",
+    "transpose",
     "types",
     "unary",
 ]

python/pylibcudf/pylibcudf/tests/test_transpose.py

@@ -0,0 +1,27 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.
+
+import pyarrow as pa
+import pylibcudf as plc
+import pytest
+
+
+@pytest.mark.parametrize(
+    "arr",
+    [
+        [],
+        [1, 2, 3],
+        [1, 2],
+        [1],
+    ],
+)
+def test_transpose(arr):
+    data = {"a": arr, "b": arr}
+    arrow_tbl = pa.table(data)
+    plc_tbl = plc.interop.from_arrow(arrow_tbl)
+    _, plc_result = plc.transpose.transpose(plc_tbl)
+    arrow_result = plc.interop.to_arrow(plc_result)
+    if len(arr) == 0:
+        expected = (0, 0)
+    else:
+        expected = (len(data), len(arr))
+    assert arrow_result.shape == expected

python/pylibcudf/pylibcudf/transpose.pxd

@@ -0,0 +1,5 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.
+from .table cimport Table
+
+
+cpdef tuple transpose(Table input_table)

python/pylibcudf/pylibcudf/transpose.pyx

@@ -0,0 +1,39 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.
+from libcpp.memory cimport unique_ptr
+from libcpp.pair cimport pair
+from libcpp.utility cimport move
+from pylibcudf.libcudf cimport transpose as cpp_transpose
+from pylibcudf.libcudf.column.column cimport column
+from pylibcudf.libcudf.table.table_view cimport table_view
+
+from .column cimport Column
+from .table cimport Table
+
+
+cpdef tuple transpose(Table input_table):
+    """Transpose a Table.
+
+    For details, see :cpp:func:`transpose`.
+
+    Parameters
+    ----------
+    input_table : Table
+        Table to transpose
+
+    Returns
+    -------
+    tuple[Column, Table]
+        Two-tuple of the owner column and transposed table.
+    """
+    cdef pair[unique_ptr[column], table_view] c_result
+
+    with nogil:
+        c_result = move(cpp_transpose.transpose(input_table.view()))
+
+    owner_column = Column.from_libcudf(move(c_result.first))
+    owner_table = Table([owner_column] * c_result.second.num_columns())
+
+    return (
+        owner_column,
+        Table.from_table_view(c_result.second, owner_table)
+    )