Update strings split APIs with stream parameters (#19909)

Contributes to #15163

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

Approvers:
  - Matthew Roeschke (https://github.com/mroeschke)

URL: #19909

Author: vyasr

python/pylibcudf/pylibcudf/libcudf/strings/split/partition.pxd

@@ -1,20 +1,23 @@
-# Copyright (c) 2020-2024, NVIDIA CORPORATION.
+# Copyright (c) 2020-2025, NVIDIA CORPORATION.
 from libcpp.memory cimport unique_ptr
 from libcpp.string cimport string
 from pylibcudf.exception_handler cimport libcudf_exception_handler
 from pylibcudf.libcudf.column.column cimport column
 from pylibcudf.libcudf.column.column_view cimport column_view
 from pylibcudf.libcudf.scalar.scalar cimport string_scalar
 from pylibcudf.libcudf.table.table cimport table
+from rmm.librmm.cuda_stream_view cimport cuda_stream_view
 
 
 cdef extern from "cudf/strings/split/partition.hpp" namespace \
         "cudf::strings" nogil:
 
     cdef unique_ptr[table] partition(
         column_view input,
-        string_scalar delimiter) except +libcudf_exception_handler
+        string_scalar delimiter,
+        cuda_stream_view stream) except +libcudf_exception_handler
 
     cdef unique_ptr[table] rpartition(
         column_view input,
-        string_scalar delimiter) except +libcudf_exception_handler
+        string_scalar delimiter,
+        cuda_stream_view stream) except +libcudf_exception_handler

python/pylibcudf/pylibcudf/libcudf/strings/split/split.pxd

@@ -1,4 +1,4 @@
-# Copyright (c) 2020-2024, NVIDIA CORPORATION.
+# Copyright (c) 2020-2025, NVIDIA CORPORATION.
 from libcpp.memory cimport unique_ptr
 from libcpp.string cimport string
 from pylibcudf.exception_handler cimport libcudf_exception_handler
@@ -8,6 +8,7 @@ from pylibcudf.libcudf.scalar.scalar cimport string_scalar
 from pylibcudf.libcudf.strings.regex_program cimport regex_program
 from pylibcudf.libcudf.table.table cimport table
 from pylibcudf.libcudf.types cimport size_type
+from rmm.librmm.cuda_stream_view cimport cuda_stream_view
 
 
 cdef extern from "cudf/strings/split/split.hpp" namespace \
@@ -16,22 +17,26 @@ cdef extern from "cudf/strings/split/split.hpp" namespace \
     cdef unique_ptr[table] split(
         column_view strings_column,
         string_scalar delimiter,
-        size_type maxsplit) except +libcudf_exception_handler
+        size_type maxsplit,
+        cuda_stream_view stream) except +libcudf_exception_handler
 
     cdef unique_ptr[table] rsplit(
         column_view strings_column,
         string_scalar delimiter,
-        size_type maxsplit) except +libcudf_exception_handler
+        size_type maxsplit,
+        cuda_stream_view stream) except +libcudf_exception_handler
 
     cdef unique_ptr[column] split_record(
         column_view strings,
         string_scalar delimiter,
-        size_type maxsplit) except +libcudf_exception_handler
+        size_type maxsplit,
+        cuda_stream_view stream) except +libcudf_exception_handler
 
     cdef unique_ptr[column] rsplit_record(
         column_view strings,
         string_scalar delimiter,
-        size_type maxsplit) except +libcudf_exception_handler
+        size_type maxsplit,
+        cuda_stream_view stream) except +libcudf_exception_handler
 
 
 cdef extern from "cudf/strings/split/split_re.hpp" namespace \
@@ -40,19 +45,23 @@ cdef extern from "cudf/strings/split/split_re.hpp" namespace \
     cdef unique_ptr[table] split_re(
         const column_view& input,
         regex_program prog,
-        size_type maxsplit) except +libcudf_exception_handler
+        size_type maxsplit,
+        cuda_stream_view stream) except +libcudf_exception_handler
 
     cdef unique_ptr[table] rsplit_re(
         const column_view& input,
         regex_program prog,
-        size_type maxsplit) except +libcudf_exception_handler
+        size_type maxsplit,
+        cuda_stream_view stream) except +libcudf_exception_handler
 
     cdef unique_ptr[column] split_record_re(
         const column_view& input,
         regex_program prog,
-        size_type maxsplit) except +libcudf_exception_handler
+        size_type maxsplit,
+        cuda_stream_view stream) except +libcudf_exception_handler
 
     cdef unique_ptr[column] rsplit_record_re(
         const column_view& input,
         regex_program prog,
-        size_type maxsplit) except +libcudf_exception_handler
+        size_type maxsplit,
+        cuda_stream_view stream) except +libcudf_exception_handler

python/pylibcudf/pylibcudf/strings/split/partition.pxd

@@ -1,10 +1,11 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.
+# Copyright (c) 2024-2025, NVIDIA CORPORATION.
 
 from pylibcudf.column cimport Column
 from pylibcudf.scalar cimport Scalar
 from pylibcudf.table cimport Table
+from rmm.pylibrmm.stream cimport Stream
 
 
-cpdef Table partition(Column input, Scalar delimiter=*)
+cpdef Table partition(Column input, Scalar delimiter=*, Stream stream=*)
 
-cpdef Table rpartition(Column input, Scalar delimiter=*)
+cpdef Table rpartition(Column input, Scalar delimiter=*, Stream stream=*)

python/pylibcudf/pylibcudf/strings/split/partition.pyi

@@ -1,8 +1,18 @@
 # Copyright (c) 2024, NVIDIA CORPORATION.
 
+from rmm.pylibrmm.stream import Stream
+
 from pylibcudf.column import Column
 from pylibcudf.scalar import Scalar
 from pylibcudf.table import Table
 
-def partition(input: Column, delimiter: Scalar | None = None) -> Table: ...
-def rpartition(input: Column, delimiter: Scalar | None = None) -> Table: ...
+def partition(
+    input: Column,
+    delimiter: Scalar | None = None,
+    stream: Stream | None = None,
+) -> Table: ...
+def rpartition(
+    input: Column,
+    delimiter: Scalar | None = None,
+    stream: Stream | None = None,
+) -> Table: ...

python/pylibcudf/pylibcudf/strings/split/partition.pyx

@@ -17,7 +17,7 @@ from cython.operator import dereference
 
 __all__ = ["partition", "rpartition"]
 
-cpdef Table partition(Column input, Scalar delimiter=None):
+cpdef Table partition(Column input, Scalar delimiter=None, Stream stream=None):
     """
     Returns a set of 3 columns by splitting each string using the
     specified delimiter.
@@ -41,23 +41,25 @@ cpdef Table partition(Column input, Scalar delimiter=None):
     cdef const string_scalar* c_delimiter = <const string_scalar*>(
         delimiter.c_obj.get()
     )
-    cdef Stream stream
+    cdef Stream stream_local
+
+    stream_local = _get_stream(stream)
 
     if delimiter is None:
-        stream = _get_stream(None)
         delimiter = Scalar.from_libcudf(
-            cpp_make_string_scalar("".encode(), stream.view())
+            cpp_make_string_scalar("".encode(), stream_local.view())
         )
 
     with nogil:
         c_result = cpp_partition.partition(
             input.view(),
-            dereference(c_delimiter)
+            dereference(c_delimiter),
+            stream_local.view()
         )
 
-    return Table.from_libcudf(move(c_result))
+    return Table.from_libcudf(move(c_result), stream_local)
 
-cpdef Table rpartition(Column input, Scalar delimiter=None):
+cpdef Table rpartition(Column input, Scalar delimiter=None, Stream stream=None):
     """
     Returns a set of 3 columns by splitting each string using the
     specified delimiter starting from the end of each string.
@@ -81,18 +83,20 @@ cpdef Table rpartition(Column input, Scalar delimiter=None):
     cdef const string_scalar* c_delimiter = <const string_scalar*>(
         delimiter.c_obj.get()
     )
-    cdef Stream stream
+    cdef Stream stream_local
+
+    stream_local = _get_stream(stream)
 
     if delimiter is None:
-        stream = _get_stream(None)
         delimiter = Scalar.from_libcudf(
-            cpp_make_string_scalar("".encode(), stream.view())
+            cpp_make_string_scalar("".encode(), stream_local.view())
         )
 
     with nogil:
         c_result = cpp_partition.rpartition(
             input.view(),
-            dereference(c_delimiter)
+            dereference(c_delimiter),
+            stream_local.view()
         )
 
-    return Table.from_libcudf(move(c_result))
+    return Table.from_libcudf(move(c_result), stream_local)

python/pylibcudf/pylibcudf/strings/split/split.pxd

@@ -1,24 +1,41 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.
+# Copyright (c) 2024-2025, NVIDIA CORPORATION.
 
 from pylibcudf.column cimport Column
 from pylibcudf.libcudf.types cimport size_type
 from pylibcudf.scalar cimport Scalar
 from pylibcudf.strings.regex_program cimport RegexProgram
 from pylibcudf.table cimport Table
+from rmm.pylibrmm.stream cimport Stream
 
 
-cpdef Table split(Column strings_column, Scalar delimiter, size_type maxsplit)
+cpdef Table split(
+    Column strings_column, Scalar delimiter, size_type maxsplit, Stream stream=*
+)
 
-cpdef Table rsplit(Column strings_column, Scalar delimiter, size_type maxsplit)
+cpdef Table rsplit(
+    Column strings_column, Scalar delimiter, size_type maxsplit, Stream stream=*
+)
 
-cpdef Column split_record(Column strings, Scalar delimiter, size_type maxsplit)
+cpdef Column split_record(
+    Column strings, Scalar delimiter, size_type maxsplit, Stream stream=*
+)
 
-cpdef Column rsplit_record(Column strings, Scalar delimiter, size_type maxsplit)
+cpdef Column rsplit_record(
+    Column strings, Scalar delimiter, size_type maxsplit, Stream stream=*
+)
 
-cpdef Table split_re(Column input, RegexProgram prog, size_type maxsplit)
+cpdef Table split_re(
+    Column input, RegexProgram prog, size_type maxsplit, Stream stream=*
+)
 
-cpdef Table rsplit_re(Column input, RegexProgram prog, size_type maxsplit)
+cpdef Table rsplit_re(
+    Column input, RegexProgram prog, size_type maxsplit, Stream stream=*
+)
 
-cpdef Column split_record_re(Column input, RegexProgram prog, size_type maxsplit)
+cpdef Column split_record_re(
+    Column input, RegexProgram prog, size_type maxsplit, Stream stream=*
+)
 
-cpdef Column rsplit_record_re(Column input, RegexProgram prog, size_type maxsplit)
+cpdef Column rsplit_record_re(
+    Column input, RegexProgram prog, size_type maxsplit, Stream stream=*
+)

python/pylibcudf/pylibcudf/strings/split/split.pyi

@@ -1,27 +1,57 @@
 # Copyright (c) 2024, NVIDIA CORPORATION.
 
+from rmm.pylibrmm.stream import Stream
+
 from pylibcudf.column import Column
 from pylibcudf.scalar import Scalar
 from pylibcudf.strings.regex_program import RegexProgram
 from pylibcudf.table import Table
 
 def split(
-    strings_column: Column, delimiter: Scalar, maxsplit: int
+    strings_column: Column,
+    delimiter: Scalar,
+    maxsplit: int,
+    stream: Stream | None = None,
 ) -> Table: ...
 def rsplit(
-    strings_column: Column, delimiter: Scalar, maxsplit: int
+    strings_column: Column,
+    delimiter: Scalar,
+    maxsplit: int,
+    stream: Stream | None = None,
 ) -> Table: ...
 def split_record(
-    strings: Column, delimiter: Scalar, maxsplit: int
+    strings: Column,
+    delimiter: Scalar,
+    maxsplit: int,
+    stream: Stream | None = None,
 ) -> Column: ...
 def rsplit_record(
-    strings: Column, delimiter: Scalar, maxsplit: int
+    strings: Column,
+    delimiter: Scalar,
+    maxsplit: int,
+    stream: Stream | None = None,
 ) -> Column: ...
-def split_re(input: Column, prog: RegexProgram, maxsplit: int) -> Table: ...
-def rsplit_re(input: Column, prog: RegexProgram, maxsplit: int) -> Table: ...
+def split_re(
+    input: Column,
+    prog: RegexProgram,
+    maxsplit: int,
+    stream: Stream | None = None,
+) -> Table: ...
+def rsplit_re(
+    input: Column,
+    prog: RegexProgram,
+    maxsplit: int,
+    stream: Stream | None = None,
+) -> Table: ...
 def split_record_re(
-    input: Column, prog: RegexProgram, maxsplit: int
+    input: Column,
+    prog: RegexProgram,
+    maxsplit: int,
+    stream: Stream | None = None,
 ) -> Column: ...
 def rsplit_record_re(
-    input: Column, prog: RegexProgram, maxsplit: int
+    input: Column,
+    prog: RegexProgram,
+    maxsplit: int,
+    stream: Stream | None = None,
 ) -> Column: ...