[cudf_streaming] Using rapidsmpf spill unspill partition

cpp/libcudf_streaming/include/cudf_streaming/integrations/partition.hpp

@@ -138,52 +138,4 @@ partition_and_pack(
   rapidsmpf::BufferResource* br,
   rapidsmpf::AllowOverbooking allow_overbooking = rapidsmpf::AllowOverbooking::YES);
 
-/**
- * @brief Spill partitions from device memory to host memory.
- *
- * Moves the buffer of each `PackedData` from device memory to host memory using
- * the provided buffer resource and the buffer's CUDA stream. Partitions that are
- * already in host memory are passed through unchanged.
- *
- * For device-resident partitions, a host memory reservation is made before moving
- * the buffer. If the reservation fails due to insufficient host memory, an exception
- * is thrown. Overbooking is not allowed.
- *
- * @param partitions The partitions to spill.
- * @param br Buffer resource used to reserve host memory and perform the move.
- *
- * @return A vector of `PackedData`, where each buffer resides in host memory.
- *
- * @throws rapidsmpf::reservation_error If host memory reservation fails.
- */
-std::vector<rapidsmpf::PackedData> spill_partitions(std::vector<rapidsmpf::PackedData>&& partitions,
-                                                    rapidsmpf::BufferResource* br);
-
-/**
- * @brief Move spilled partitions (i.e., packed tables in host memory) back to device
- * memory.
- *
- * Each partition is inspected to determine whether its buffer resides in device memory.
- * Buffers already in device memory are left untouched. Host-resident buffers are moved
- * to device memory using the provided buffer resource and the buffer's CUDA stream.
- *
- * If insufficient device memory is available, the buffer resource's spill manager is
- * invoked to free memory. If overbooking occurs and spilling fails to reclaim enough
- * memory, behavior depends on the `allow_overbooking` flag.
- *
- * @param partitions The partitions to unspill, potentially containing host-resident data.
- * @param br Buffer resource responsible for memory reservation and spills.
- * @param allow_overbooking If false, ensures enough memory is freed to satisfy the
- * reservation; otherwise, allows overbooking even if spilling was insufficient.
- *
- * @return A vector of `PackedData`, each with a buffer in device memory.
- *
- * @throws rapidsmpf::reservation_error If overbooking exceeds the amount spilled and
- *         `allow_overbooking` is false.
- */
-std::vector<rapidsmpf::PackedData> unspill_partitions(
-  std::vector<rapidsmpf::PackedData>&& partitions,
-  rapidsmpf::BufferResource* br,
-  rapidsmpf::AllowOverbooking allow_overbooking);
-
 }  // namespace cudf_streaming::integrations

cpp/libcudf_streaming/src/integrations/partition.cpp

@@ -183,44 +183,4 @@ std::unique_ptr<cudf::table> unpack_and_concat(std::vector<rapidsmpf::PackedData
   return cudf::concatenate(unpacked, stream, br->device_mr());
 }
 
-std::vector<rapidsmpf::PackedData> spill_partitions(std::vector<rapidsmpf::PackedData>&& partitions,
-                                                    rapidsmpf::BufferResource* br)
-{
-  // Sum the total size of all packed data in device memory.
-  std::size_t device_size{0};
-  for (auto& [_, data] : partitions) {
-    if (data->mem_type() == rapidsmpf::MemoryType::DEVICE) { device_size += data->size; }
-  }
-  // Spill each partition to host memory.
-  auto reservation = br->reserve_or_fail(device_size, rapidsmpf::SPILL_TARGET_MEMORY_TYPES);
-  std::vector<rapidsmpf::PackedData> ret;
-  ret.reserve(partitions.size());
-  for (auto& [metadata, data] : partitions) {
-    ret.emplace_back(std::move(metadata), br->move(std::move(data), reservation));
-  }
-  return ret;
-}
-
-std::vector<rapidsmpf::PackedData> unspill_partitions(
-  std::vector<rapidsmpf::PackedData>&& partitions,
-  rapidsmpf::BufferResource* br,
-  rapidsmpf::AllowOverbooking allow_overbooking)
-{
-  auto statistics = br->statistics();
-  // Sum the total size of all packed data not in device memory already.
-  std::size_t non_device_size{0};
-  for (auto& [_, data] : partitions) {
-    if (data->mem_type() != rapidsmpf::MemoryType::DEVICE) { non_device_size += data->size; }
-  }
-
-  // Unspill each partition.
-  auto reservation = br->reserve_device_memory_and_spill(non_device_size, allow_overbooking);
-  std::vector<rapidsmpf::PackedData> ret;
-  ret.reserve(partitions.size());
-  for (auto& [metadata, data] : partitions) {
-    ret.emplace_back(std::move(metadata), br->move(std::move(data), reservation));
-  }
-
-  return ret;
-}
 }  // namespace cudf_streaming::integrations

cpp/libcudf_streaming/src/streaming/partition.cpp

@@ -8,6 +8,7 @@
 #include <cudf_streaming/streaming/partition.hpp>
 #include <cudf_streaming/streaming/table_chunk.hpp>
 #include <rapidsmpf/cuda_stream.hpp>
+#include <rapidsmpf/memory/spill.hpp>
 #include <rapidsmpf/streaming/chunks/partition.hpp>
 #include <rapidsmpf/streaming/core/lineariser.hpp>
 
@@ -76,7 +77,7 @@ rapidsmpf::streaming::Actor unpack_and_concat(std::shared_ptr<rapidsmpf::streami
     auto stream = ctx->br()->stream_pool().get_stream();
 
     std::unique_ptr<cudf::table> ret = cudf_streaming::integrations::unpack_and_concat(
-      cudf_streaming::integrations::unspill_partitions(
+      rapidsmpf::unspill_partitions(
         std::move(data), ctx->br().get(), rapidsmpf::AllowOverbooking::NO),
       stream,
       ctx->br().get());

cpp/libcudf_streaming/tests/test_partition.cpp

@@ -91,42 +91,3 @@ TEST_P(NumOfPartitions, split_and_pack)
   // Compare the input table with the result.
   CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expect, *result);
 }
-
-class SpillingTest : public ::testing::Test {
- protected:
-  void SetUp() override
-  {
-    br     = rapidsmpf::BufferResource::create(cudf::get_current_device_resource_ref());
-    stream = cudf::get_default_stream();
-  }
-
-  std::shared_ptr<rapidsmpf::BufferResource> br;
-  rmm::cuda_stream_view stream;
-};
-
-TEST_F(SpillingTest, SpillUnspillRoundtripPreservesDataAndMetadata)
-{
-  std::vector<std::uint8_t> metadata{42, 99};
-  std::vector<std::uint8_t> payload{10, 20, 30};
-
-  // Create device input.
-  std::vector<rapidsmpf::PackedData> input;
-  input.push_back(create_packed_data(metadata, payload, stream, br.get()));
-
-  // Device -> Device (moves data)
-  auto on_gpu = unspill_partitions(std::move(input), br.get(), rapidsmpf::AllowOverbooking::YES);
-  ASSERT_EQ(on_gpu.size(), 1);
-  EXPECT_EQ(on_gpu[0].data->mem_type(), rapidsmpf::MemoryType::DEVICE);
-  EXPECT_EQ(*on_gpu[0].metadata, metadata);
-
-  // Device -> Host
-  auto back_on_host = spill_partitions(std::move(on_gpu), br.get());
-  ASSERT_EQ(back_on_host.size(), 1);
-  EXPECT_EQ(back_on_host[0].data->mem_type(), rapidsmpf::MemoryType::HOST);
-  EXPECT_EQ(*back_on_host[0].metadata, metadata);
-
-  // Check that contents match original
-  auto res    = br->reserve_or_fail(back_on_host[0].data->size, rapidsmpf::MemoryType::HOST);
-  auto actual = br->move_to_host_buffer(std::move(back_on_host[0].data), res);
-  EXPECT_EQ(actual->copy_to_uint8_vector(), payload);
-}

cpp/libcudf_streaming/tests/utils.hpp

@@ -152,21 +152,6 @@ template <std::integral T = std::int64_t>
   return sort_table(table->view(), column_indices);
 }
 
-/// @brief Create a PackedData object from a host buffer
-[[nodiscard]] inline rapidsmpf::PackedData create_packed_data(
-  std::span<std::uint8_t const> metadata,
-  std::span<std::uint8_t const> data,
-  rmm::cuda_stream_view stream,
-  rapidsmpf::BufferResource* br)
-{
-  auto metadata_ptr = std::make_unique<std::vector<std::uint8_t>>(metadata.begin(), metadata.end());
-
-  auto reservation =
-    br->reserve(rapidsmpf::MemoryType::DEVICE, data.size(), rapidsmpf::AllowOverbooking::YES);
-  auto data_ptr = std::make_unique<rmm::device_buffer>(data.data(), data.size(), stream);
-  return rapidsmpf::PackedData{std::move(metadata_ptr), br->move(std::move(data_ptr), stream)};
-}
-
 /**
  * @brief Generate a packed data object with the given number of elements and offset.
  *

python/cudf_streaming/cudf_streaming/integrations/partition.pxd

@@ -22,12 +22,6 @@ cpdef object split_and_pack(
     BufferResource br,
 )
 cpdef object unpack_and_concat(object partitions, Stream stream, BufferResource br)
-cpdef object spill_partitions(object partitions, BufferResource br)
-cpdef object unspill_partitions(
-    object partitions,
-    BufferResource br,
-    object allow_overbooking,
-)
 cpdef object packed_data_from_cudf_packed_columns(
     PackedColumns packed_columns,
     Stream stream,

python/cudf_streaming/cudf_streaming/integrations/partition.pyi

@@ -29,17 +29,6 @@ def unpack_and_concat(
     stream: Stream,
     br: BufferResource,
 ) -> Table: ...
-def spill_partitions(
-    partitions: Iterable[PackedData],
-    *,
-    br: BufferResource,
-) -> list[PackedData]: ...
-def unspill_partitions(
-    partitions: Iterable[PackedData],
-    *,
-    br: BufferResource,
-    allow_overbooking: bool,
-) -> list[PackedData]: ...
 def packed_data_from_cudf_packed_columns(
     packed_columns: PackedColumns,
     stream: Stream,

python/cudf_streaming/cudf_streaming/integrations/partition.pyx

@@ -19,11 +19,8 @@ from rmm.librmm.device_buffer cimport device_buffer
 from rmm.pylibrmm.stream cimport Stream
 
 from rapidsmpf._detail.exception_handling cimport ex_handler
-from rapidsmpf.memory.buffer_resource cimport (AllowOverbooking,
-                                               BufferResource,
-                                               cpp_BufferResource)
-from rapidsmpf.memory.packed_data cimport (PackedData, cpp_PackedData,
-                                           packed_data_vector_to_list)
+from rapidsmpf.memory.buffer_resource cimport BufferResource, cpp_BufferResource
+from rapidsmpf.memory.packed_data cimport PackedData, cpp_PackedData
 
 
 cdef extern from "<cudf_streaming/integrations/partition.hpp>" nogil:
@@ -256,68 +253,6 @@ cpdef object unpack_and_concat(
         )
     return Table.from_libcudf(move(_ret), stream, br._device_mr)
 
-
-cdef extern from "<cudf_streaming/integrations/partition.hpp>" nogil:
-    cdef vector[cpp_PackedData] cpp_spill_partitions \
-        "cudf_streaming::integrations::spill_partitions"(
-            vector[cpp_PackedData] partitions,
-            cpp_BufferResource* br,
-        ) except +ex_handler
-
-
-cpdef object spill_partitions(
-    object partitions,
-    BufferResource br,
-):
-    """
-    Spill partitions from device memory to host memory.
-
-    Moves the buffer of each ``PackedData`` from device memory to host memory using
-    the provided buffer resource and the buffer's CUDA stream. Partitions already
-    in host memory are returned unchanged.
-
-    For device-resident partitions, a host memory reservation is made before moving
-    the buffer. If the reservation fails due to insufficient host memory, an
-    exception is raised. Overbooking is not allowed.
-
-    The input partitions are released and are left empty on return.
-
-    Parameters
-    ----------
-    partitions
-        The partitions to spill.
-    br
-        Buffer resource used to reserve host memory and perform the move.
-
-    Returns
-    -------
-    A list of partitions whose buffers reside in host memory.
-
-    Raises
-    ------
-    ReservationError
-        If host memory reservation fails.
-    """
-    cdef cpp_BufferResource* _br = br.ptr()
-    cdef vector[cpp_PackedData] _partitions = _partitions_py_to_cpp(partitions)
-    cdef vector[cpp_PackedData] _ret
-    with nogil:
-        _ret = cpp_spill_partitions(
-            move(_partitions),
-            _br,
-        )
-    return packed_data_vector_to_list(move(_ret), br)
-
-
-cdef extern from "<cudf_streaming/integrations/partition.hpp>" nogil:
-    cdef vector[cpp_PackedData] cpp_unspill_partitions \
-        "cudf_streaming::integrations::unspill_partitions"(
-            vector[cpp_PackedData] partitions,
-            cpp_BufferResource* br,
-            AllowOverbooking allow_overbooking,
-        ) except +ex_handler
-
-
 cdef extern from *:
     """
     #include <rapidsmpf/error.hpp>
@@ -344,61 +279,6 @@ cdef extern from *:
     ) except +ex_handler nogil
 
 
-cpdef object unspill_partitions(
-    object partitions,
-    BufferResource br,
-    object allow_overbooking,
-):
-    """
-    Move spilled partitions back to device memory.
-
-    Each partition is inspected to determine whether its buffer resides in device
-    memory. Buffers already in device memory are left untouched. Host-resident buffers
-    are moved to device memory using the provided buffer resource and the buffer's CUDA
-    stream.
-
-    If insufficient device memory is available, the buffer resource's spill manager is
-    invoked to free memory. If overbooking occurs and spilling fails to reclaim enough
-    memory, behavior depends on ``allow_overbooking``.
-
-    The input partitions are released and are left empty on return.
-
-    Parameters
-    ----------
-    partitions
-        The partitions to unspill, potentially containing host-resident data.
-    br
-        Buffer resource responsible for memory reservation and spills.
-    allow_overbooking
-        If False, ensures enough memory is freed to satisfy the reservation;
-        otherwise, allows overbooking even if spilling was insufficient.
-
-    Returns
-    -------
-    A list of partitions whose buffers reside in device memory.
-
-    Raises
-    ------
-    ReservationError
-        If overbooking exceeds the amount spilled and ``allow_overbooking is False``.
-    """
-    if not isinstance(allow_overbooking, bool):
-        raise TypeError("allow_overbooking must be a bool")
-    cdef cpp_BufferResource* _br = br.ptr()
-    cdef vector[cpp_PackedData] _partitions = _partitions_py_to_cpp(partitions)
-    cdef vector[cpp_PackedData] _ret
-    cdef AllowOverbooking ab = (
-        AllowOverbooking.YES if allow_overbooking else AllowOverbooking.NO
-    )
-    with nogil:
-        _ret = cpp_unspill_partitions(
-            move(_partitions),
-            _br,
-            ab,
-        )
-    return packed_data_vector_to_list(move(_ret), br)
-
-
 cpdef object packed_data_from_cudf_packed_columns(
     PackedColumns packed_columns,
     Stream stream,