Thread missing streams in column/table view creation to char size calculation (#20351)

All of the APIs for creating pylibcudf Columns from arbitrary other owners and a libcudf `column_view` of type string will eventually call into `strings_column_view::chars_size`, which is stream-ordered and needs an input stream. This PR pipes the necessary streams through.

Authors:
  - Vyas Ramasubramani (https://github.com/vyasr)
  - Matthew Murray (https://github.com/Matt711)

Approvers:
  - Bradley Dice (https://github.com/bdice)
  - Matthew Roeschke (https://github.com/mroeschke)
  - Lawrence Mitchell (https://github.com/wence-)

URL: #20351

Author: vyasr

python/pylibcudf/pylibcudf/column.pxd

@@ -26,7 +26,7 @@ cdef class OwnerWithCAI:
     cdef dict cai
 
     @staticmethod
-    cdef create(column_view cv, object owner)
+    cdef create(column_view cv, object owner, Stream stream)
 
 
 cdef class OwnerMaskWithCAI:
@@ -68,7 +68,11 @@ cdef class Column:
     cdef Column from_column_view(const column_view& cv, Column owner)
 
     @staticmethod
-    cdef Column from_column_view_of_arbitrary(const column_view& cv, object owner)
+    cdef Column from_column_view_of_arbitrary(
+        const column_view& cv,
+        object owner,
+        Stream stream,
+    )
 
     @staticmethod
     cdef Column _wrap_nested_list_column(

python/pylibcudf/pylibcudf/column.pyx

@@ -93,7 +93,7 @@ cdef class _ArrowColumnHolder:
 cdef class OwnerWithCAI:
     """An interface for column view's data with gpumemoryview via CAI."""
     @staticmethod
-    cdef create(column_view cv, object owner):
+    cdef create(column_view cv, object owner, Stream stream):
         obj = OwnerWithCAI()
         obj.owner = owner
         # The default size of 0 will be applied for any type that stores data in the
@@ -105,8 +105,7 @@ cdef class OwnerWithCAI:
             # Cast to Python integers before multiplying to avoid overflow.
             size = int(cv.size()) * int(cpp_size_of(cv.type()))
         elif cv.type().id() == type_id.STRING:
-            # TODO: stream-ordered
-            size = strings_column_view(cv).chars_size(_get_stream().view())
+            size = strings_column_view(cv).chars_size(stream.view())
 
         obj.cai = {
             "shape": (size,),
@@ -453,7 +452,11 @@ cdef class Column:
                 )
             result.col.swap(c_result)
 
-            return Column.from_column_view_of_arbitrary(result.col.get().view(), result)
+            return Column.from_column_view_of_arbitrary(
+                result.col.get().view(),
+                result,
+                stream,
+            )
         elif hasattr(obj, "__arrow_c_array__"):
             schema, h_array = obj.__arrow_c_array__()
             c_schema = <ArrowSchema*>PyCapsule_GetPointer(schema, "arrow_schema")
@@ -470,7 +473,11 @@ cdef class Column:
                 )
             result.col.swap(c_result)
 
-            return Column.from_column_view_of_arbitrary(result.col.get().view(), result)
+            return Column.from_column_view_of_arbitrary(
+                result.col.get().view(),
+                result,
+                stream,
+            )
         elif hasattr(obj, "__arrow_c_stream__"):
             arrow_stream = obj.__arrow_c_stream__()
             c_arrow_stream = (
@@ -490,7 +497,11 @@ cdef class Column:
                 )
             result.col.swap(c_result)
 
-            return Column.from_column_view_of_arbitrary(result.col.get().view(), result)
+            return Column.from_column_view_of_arbitrary(
+                result.col.get().view(),
+                result,
+                stream,
+            )
         elif hasattr(obj, "__arrow_c_device_stream__"):
             # TODO: When we add support for this case, it should be moved above
             # the __arrow_c_array__ case since we should prioritize device
@@ -716,7 +727,11 @@ cdef class Column:
     # from_column_view, but this does not work due to
     # https://github.com/cython/cython/issues/6740
     @staticmethod
-    cdef Column from_column_view_of_arbitrary(const column_view& cv, object owner):
+    cdef Column from_column_view_of_arbitrary(
+        const column_view& cv,
+        object owner,
+        Stream stream,
+    ):
         """Create a Column from a libcudf column_view into an arbitrary owner.
 
         This method accepts shared ownership of the underlying data from the owner.
@@ -730,16 +745,19 @@ cdef class Column:
         """
         # For efficiency, prohibit calling this overload with a Column owner.
         assert not isinstance(owner, Column)
+        stream = _get_stream(stream)
 
         children = []
         cdef size_type i
         if cv.num_children() != 0:
             for i in range(cv.num_children()):
                 children.append(
-                    Column.from_column_view_of_arbitrary(cv.child(i), owner)
+                    Column.from_column_view_of_arbitrary(cv.child(i), owner, stream)
                 )
 
-        cdef gpumemoryview owning_data = gpumemoryview(OwnerWithCAI.create(cv, owner))
+        cdef gpumemoryview owning_data = gpumemoryview(
+            OwnerWithCAI.create(cv, owner, stream)
+        )
         cdef gpumemoryview owning_mask = None
         if cv.null_count() > 0:
             owning_mask = gpumemoryview(OwnerMaskWithCAI.create(cv, owner))

python/pylibcudf/pylibcudf/contiguous_split.pxd

@@ -54,9 +54,10 @@ cdef class ChunkedPack:
 
 cpdef PackedColumns pack(Table input)
 
-cpdef Table unpack(PackedColumns input)
+cpdef Table unpack(PackedColumns input, Stream stream=*)
 
 cpdef Table unpack_from_memoryviews(
     memoryview metadata,
     gpumemoryview gpu_data,
+    Stream stream=*,
 )

python/pylibcudf/pylibcudf/contiguous_split.pyx

@@ -350,7 +350,7 @@ cpdef PackedColumns pack(Table input, Stream stream=None, DeviceMemoryResource m
     return PackedColumns.from_libcudf(move(pack), stream, mr)
 
 
-cpdef Table unpack(PackedColumns input):
+cpdef Table unpack(PackedColumns input, Stream stream=None):
     """Deserialize the result of `pack`.
 
     Copies the result of a serialized table into a table.
@@ -361,21 +361,25 @@ cpdef Table unpack(PackedColumns input):
     ----------
     input : PackedColumns
         The packed columns to unpack.
+    stream : Stream | None
+        CUDA stream on which to perform the operation.
 
     Returns
     -------
     Table
         Copy of the packed columns.
     """
     cdef table_view v
+    stream = _get_stream(stream)
     with nogil:
         v = cpp_unpack(dereference(input.c_obj))
-    return Table.from_table_view_of_arbitrary(v, input)
+    return Table.from_table_view_of_arbitrary(v, input, stream)
 
 
 cpdef Table unpack_from_memoryviews(
     memoryview metadata,
     gpumemoryview gpu_data,
+    Stream stream=None,
 ):
     """Deserialize the result of `pack`.
 
@@ -389,20 +393,23 @@ cpdef Table unpack_from_memoryviews(
         The packed metadata to unpack.
     gpu_data : gpumemoryview
         The packed gpu_data to unpack.
+    stream : Stream | None
+        CUDA stream on which to perform the operation.
 
     Returns
     -------
     Table
         Copy of the packed columns.
     """
+    stream = _get_stream(stream)
     if metadata.nbytes == 0:
         if gpu_data.__cuda_array_interface__["data"][0] != 0:
             raise ValueError("Expected an empty gpu_data from unpacking an empty table")
-        # For an empty table we just attach the default stream and mr since neither will
-        # be used for any operations.
+        # For an empty table we just attach the default mr since it will not be
+        # used for any operations.
         return Table.from_libcudf(
             make_unique[table](table_view()),
-            _get_stream(),
+            stream,
             _get_memory_resource(),
         )
 
@@ -414,4 +421,4 @@ cpdef Table unpack_from_memoryviews(
     cdef table_view v
     with nogil:
         v = cpp_unpack(metadata_ptr, gpu_data_ptr)
-    return Table.from_table_view_of_arbitrary(v, gpu_data)
+    return Table.from_table_view_of_arbitrary(v, gpu_data, stream)

python/pylibcudf/pylibcudf/table.pxd

@@ -28,6 +28,10 @@ cdef class Table:
     cdef Table from_table_view(const table_view& tv, Table owner)
 
     @staticmethod
-    cdef Table from_table_view_of_arbitrary(const table_view& tv, object owner)
+    cdef Table from_table_view_of_arbitrary(
+        const table_view& tv,
+        object owner,
+        Stream stream,
+    )
 
     cpdef list columns(self)

python/pylibcudf/pylibcudf/table.pyx

@@ -175,7 +175,11 @@ cdef class Table:
                 )
             result.tbl.swap(c_result)
 
-            return Table.from_table_view_of_arbitrary(result.tbl.get().view(), result)
+            return Table.from_table_view_of_arbitrary(
+                result.tbl.get().view(),
+                result,
+                stream,
+            )
         elif hasattr(obj, "__arrow_c_stream__"):
             arrow_stream = obj.__arrow_c_stream__()
             c_stream = (
@@ -194,7 +198,11 @@ cdef class Table:
                 )
             result.tbl.swap(c_result)
 
-            return Table.from_table_view_of_arbitrary(result.tbl.get().view(), result)
+            return Table.from_table_view_of_arbitrary(
+                result.tbl.get().view(),
+                result,
+                stream,
+            )
         elif hasattr(obj, "__arrow_c_device_stream__"):
             # TODO: When we add support for this case, it should be moved above
             # the __arrow_c_stream__ case since we should prioritize device
@@ -262,7 +270,11 @@ cdef class Table:
     # from_table_view, but this does not work due to
     # https://github.com/cython/cython/issues/6740
     @staticmethod
-    cdef Table from_table_view_of_arbitrary(const table_view& tv, object owner):
+    cdef Table from_table_view_of_arbitrary(
+        const table_view& tv,
+        object owner,
+        Stream stream,
+    ):
         """Create a Table from a libcudf table_view into an arbitrary owner.
 
         This method accepts shared ownership of the underlying data from the owner.
@@ -279,7 +291,7 @@ cdef class Table:
         assert not isinstance(owner, Table)
         cdef int i
         return Table([
-            Column.from_column_view_of_arbitrary(tv.column(i), owner)
+            Column.from_column_view_of_arbitrary(tv.column(i), owner, stream)
             for i in range(tv.num_columns())
         ])