Add memory resources to I/O modules (#20136)

Contributes to #15170

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

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

URL: #20136

Author: vyasr

python/pylibcudf/pylibcudf/io/avro.pxd

@@ -1,5 +1,6 @@
 # Copyright (c) 2024-2025, NVIDIA CORPORATION.
 from rmm.pylibrmm.stream cimport Stream
+from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource
 
 from pylibcudf.io.types cimport SourceInfo, TableWithMetadata
 
@@ -25,4 +26,6 @@ cdef class AvroReaderOptionsBuilder:
     cpdef AvroReaderOptionsBuilder num_rows(self, size_type num_rows)
     cpdef AvroReaderOptions build(self)
 
-cpdef TableWithMetadata read_avro(AvroReaderOptions options, Stream stream = *)
+cpdef TableWithMetadata read_avro(
+    AvroReaderOptions options, Stream stream = *, DeviceMemoryResource mr=*
+)

python/pylibcudf/pylibcudf/io/avro.pyi

@@ -1,4 +1,5 @@
 # Copyright (c) 2024, NVIDIA CORPORATION.
+from rmm.pylibrmm.memory_resource import DeviceMemoryResource
 from rmm.pylibrmm.stream import Stream
 
 from pylibcudf.io.types import SourceInfo, TableWithMetadata
@@ -16,5 +17,7 @@ class AvroReaderOptionsBuilder:
     def build(self) -> AvroReaderOptions: ...
 
 def read_avro(
-    options: AvroReaderOptions, stream: Stream = None
+    options: AvroReaderOptions,
+    stream: Stream = None,
+    mr: DeviceMemoryResource = None,
 ) -> TableWithMetadata: ...

python/pylibcudf/pylibcudf/io/avro.pyx

@@ -5,6 +5,7 @@ from libcpp.utility cimport move
 from libcpp.vector cimport vector
 
 from rmm.pylibrmm.stream cimport Stream
+from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource
 
 from pylibcudf.io.types cimport SourceInfo, TableWithMetadata
 
@@ -15,7 +16,7 @@ from pylibcudf.libcudf.io.avro cimport (
 
 from pylibcudf.libcudf.types cimport size_type
 
-from pylibcudf.utils cimport _get_stream
+from pylibcudf.utils cimport _get_stream, _get_memory_resource
 
 
 __all__ = ["read_avro", "AvroReaderOptions", "AvroReaderOptionsBuilder"]
@@ -136,6 +137,7 @@ cdef class AvroReaderOptionsBuilder:
 cpdef TableWithMetadata read_avro(
     AvroReaderOptions options,
     Stream stream = None,
+    DeviceMemoryResource mr=None,
 ):
     """
     Read from Avro format.
@@ -151,9 +153,12 @@ cpdef TableWithMetadata read_avro(
         Settings for controlling reading behavior
     stream : Stream | None
         CUDA stream used for device memory operations and kernel launches
+    mr : DeviceMemoryResource, optional
+        Device memory resource used to allocate the returned table's device memory.
     """
     cdef Stream s = _get_stream(stream)
+    mr = _get_memory_resource(mr)
     with nogil:
-        c_result = move(cpp_read_avro(options.c_obj, s.view()))
+        c_result = move(cpp_read_avro(options.c_obj, s.view(), mr.get_mr()))
 
     return TableWithMetadata.from_libcudf(c_result, s)

python/pylibcudf/pylibcudf/io/csv.pxd

@@ -4,6 +4,7 @@ from libcpp.string cimport string
 from libcpp.vector cimport vector
 
 from rmm.pylibrmm.stream cimport Stream
+from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource
 
 from pylibcudf.io.types cimport SinkInfo, SourceInfo, TableWithMetadata
 
@@ -70,7 +71,9 @@ cdef class CsvReaderOptionsBuilder:
     cpdef CsvReaderOptionsBuilder delimiter(self, str delimiter)
     cpdef CsvReaderOptions build(self)
 
-cpdef TableWithMetadata read_csv(CsvReaderOptions options, Stream stream = *)
+cpdef TableWithMetadata read_csv(
+    CsvReaderOptions options, Stream stream = *, DeviceMemoryResource mr=*
+)
 
 cdef class CsvWriterOptions:
     cdef csv_writer_options c_obj

python/pylibcudf/pylibcudf/io/csv.pyi

@@ -2,6 +2,7 @@
 
 from typing_extensions import Self
 
+from rmm.pylibrmm.memory_resource import DeviceMemoryResource
 from rmm.pylibrmm.stream import Stream
 
 from pylibcudf.io.types import (
@@ -59,8 +60,9 @@ class CsvReaderOptionsBuilder:
 def read_csv(
     options: CsvReaderOptions,
     stream: Stream = None,
+    mr: DeviceMemoryResource = None,
 ) -> TableWithMetadata: ...
-def write_csv(options: CsvWriterOptionsBuilder, stream: Stream = None): ...
+def write_csv(options: CsvWriterOptions, stream: Stream = None): ...
 
 class CsvWriterOptions:
     def __init__(self): ...

python/pylibcudf/pylibcudf/io/csv.pyx

@@ -7,6 +7,7 @@ from libcpp.utility cimport move
 from libcpp.vector cimport vector
 
 from rmm.pylibrmm.stream cimport Stream
+from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource
 
 from pylibcudf.io.types cimport SourceInfo, SinkInfo, TableWithMetadata
 
@@ -30,7 +31,7 @@ from pylibcudf.table cimport Table
 
 from pylibcudf.types cimport DataType
 
-from pylibcudf.utils cimport _get_stream
+from pylibcudf.utils cimport _get_stream, _get_memory_resource
 
 __all__ = [
     "read_csv",
@@ -656,6 +657,7 @@ cdef class CsvReaderOptionsBuilder:
 cpdef TableWithMetadata read_csv(
     CsvReaderOptions options,
     Stream stream = None,
+    DeviceMemoryResource mr=None,
 ):
     """
     Read from CSV format.
@@ -671,11 +673,14 @@ cpdef TableWithMetadata read_csv(
         Settings for controlling reading behavior
     stream : Stream | None
         CUDA stream used for device memory operations and kernel launches
+    mr : DeviceMemoryResource, optional
+        Device memory resource used to allocate the returned table's device memory.
     """
     cdef table_with_metadata c_result
     cdef Stream s = _get_stream(stream)
+    mr = _get_memory_resource(mr)
     with nogil:
-        c_result = move(cpp_read_csv(options.c_obj, s.view()))
+        c_result = move(cpp_read_csv(options.c_obj, s.view(), mr.get_mr()))
 
     cdef TableWithMetadata tbl_meta = TableWithMetadata.from_libcudf(c_result, s)
     return tbl_meta

python/pylibcudf/pylibcudf/io/json.pxd

@@ -79,7 +79,9 @@ cdef class JsonReaderOptionsBuilder:
     cpdef JsonReaderOptionsBuilder unquoted_control_chars(self, bool val)
     cpdef build(self)
 
-cpdef TableWithMetadata read_json(JsonReaderOptions options, Stream stream = *)
+cpdef TableWithMetadata read_json(
+    JsonReaderOptions options, Stream stream = *, DeviceMemoryResource mr = *
+)
 
 cpdef TableWithMetadata read_json_from_string_column(
     Column input,
@@ -120,4 +122,5 @@ cpdef tuple chunked_read_json(
     JsonReaderOptions options,
     int chunk_size= *,
     Stream stream = *,
+    DeviceMemoryResource mr = *,
 )

python/pylibcudf/pylibcudf/io/json.pyi

@@ -4,6 +4,7 @@ from typing import TypeAlias
 
 from typing_extensions import Self
 
+from rmm.pylibrmm.memory_resource import DeviceMemoryResource
 from rmm.pylibrmm.stream import Stream
 
 from pylibcudf.column import Column
@@ -71,7 +72,9 @@ class JsonReaderOptionsBuilder:
     def build(self) -> JsonReaderOptions: ...
 
 def read_json(
-    options: JsonReaderOptions, stream: Stream = None
+    options: JsonReaderOptions,
+    stream: Stream = None,
+    mr: DeviceMemoryResource = None,
 ) -> TableWithMetadata: ...
 def read_json_from_string_column(
     input: Column,
@@ -81,6 +84,7 @@ def read_json_from_string_column(
     compression: CompressionType,
     recovery_mode: JSONRecoveryMode,
     stream: Stream = None,
+    mr: DeviceMemoryResource = None,
 ) -> TableWithMetadata: ...
 
 class JsonWriterOptions:
@@ -104,5 +108,6 @@ def chunked_read_json(
     options: JsonReaderOptions,
     chunk_size: int = 100_000_000,
     stream: Stream = None,
+    mr: DeviceMemoryResource = None,
 ) -> tuple[list[Column], list[str], ChildNameToTypeMap]: ...
 def is_supported_write_json(type: DataType) -> bool: ...

python/pylibcudf/pylibcudf/io/json.pyx

@@ -689,6 +689,7 @@ cpdef tuple chunked_read_json(
     JsonReaderOptions options,
     int chunk_size=100_000_000,
     Stream stream = None,
+    DeviceMemoryResource mr = None,
 ):
     """
     Reads chunks of a JSON file into a :py:class:`~.types.TableWithMetadata`.
@@ -718,14 +719,15 @@ cpdef tuple chunked_read_json(
     child_names = None
     i = 0
     cdef Stream s = _get_stream(stream)
+    mr = _get_memory_resource(mr)
     while True:
         options.enable_lines(True)
         options.set_byte_range_offset(c_range_size * i)
         options.set_byte_range_size(c_range_size)
 
         try:
             with nogil:
-                c_result = move(cpp_read_json(options.c_obj, s.view()))
+                c_result = move(cpp_read_json(options.c_obj, s.view(), mr.get_mr()))
         except (ValueError, OverflowError):
             break
         if meta_names is None:
@@ -754,7 +756,8 @@ cpdef tuple chunked_read_json(
 
 cpdef TableWithMetadata read_json(
     JsonReaderOptions options,
-    Stream stream = None
+    Stream stream = None,
+    DeviceMemoryResource mr = None
 ):
     """
     Read from JSON format.
@@ -778,8 +781,9 @@ cpdef TableWithMetadata read_json(
     """
     cdef table_with_metadata c_result
     cdef Stream s = _get_stream(stream)
+    mr = _get_memory_resource(mr)
     with nogil:
-        c_result = move(cpp_read_json(options.c_obj, s.view()))
+        c_result = move(cpp_read_json(options.c_obj, s.view(), mr.get_mr()))
 
     return TableWithMetadata.from_libcudf(c_result, s)
 

python/pylibcudf/pylibcudf/io/orc.pxd

@@ -9,6 +9,7 @@ from libcpp.string cimport string
 from libcpp.vector cimport vector
 
 from rmm.pylibrmm.stream cimport Stream
+from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource
 
 from pylibcudf.io.types cimport (
     SourceInfo,
@@ -61,7 +62,9 @@ cdef class OrcReaderOptionsBuilder:
     cpdef OrcReaderOptionsBuilder use_index(self, bool use)
     cpdef OrcReaderOptions build(self)
 
-cpdef TableWithMetadata read_orc(OrcReaderOptions options)
+cpdef TableWithMetadata read_orc(
+    OrcReaderOptions options, Stream stream = *, DeviceMemoryResource mr=*
+)
 
 cdef class OrcColumnStatistics:
     cdef optional[uint64_t] number_of_values_c