GH-36103: [C++] Initial device sync API

cpp/src/arrow/acero/aggregate_internal.h

@@ -52,8 +52,8 @@
 // segment-keys is used to refine the partitioning. However, segment-keys are different in
 // that they partition only consecutive rows into a single group. Such a partition of
 // consecutive rows is called a segment group. For example, consider a column X with
-// values [A, A, B, A] at row-indices [0, 1, 2, 3]. A regular group-by aggregation with keys
-// [X] yields a row-index partitioning [[0, 1, 3], [2]] whereas a segmented-group-by
+// values [A, A, B, A] at row-indices [0, 1, 2, 3]. A regular group-by aggregation with
+// keys [X] yields a row-index partitioning [[0, 1, 3], [2]] whereas a segmented-group-by
 // aggregation with segment-keys [X] yields [[0, 1], [2], [3]].
 //
 // The implementation first segments the input using the segment-keys, then groups by the

cpp/src/arrow/buffer.h

@@ -346,6 +346,8 @@ class ARROW_EXPORT Buffer {
   static Result<std::shared_ptr<Buffer>> ViewOrCopy(
       std::shared_ptr<Buffer> source, const std::shared_ptr<MemoryManager>& to);
 
+  virtual std::shared_ptr<DeviceSync> get_device_sync() { return nullptr; }
+
  protected:
   bool is_mutable_;
   bool is_cpu_;

cpp/src/arrow/c/bridge.cc

@@ -521,8 +521,7 @@ struct ExportedArrayPrivateData : PoolAllocationMixin<ExportedArrayPrivateData>
   SmallVector<struct ArrowArray*, 4> child_pointers_;
 
   std::shared_ptr<ArrayData> data_;
-
-  RawSyncEvent sync_event_;
+  std::shared_ptr<DeviceSync> sync_;
 
   ExportedArrayPrivateData() = default;
   ARROW_DEFAULT_MOVE_AND_ASSIGN(ExportedArrayPrivateData);
@@ -547,10 +546,10 @@ void ReleaseExportedArray(struct ArrowArray* array) {
   }
   DCHECK_NE(array->private_data, nullptr);
   auto* pdata = reinterpret_cast<ExportedArrayPrivateData*>(array->private_data);
-  if (pdata->sync_event_.sync_event != nullptr &&
-      pdata->sync_event_.release_func != nullptr) {
-    pdata->sync_event_.release_func(pdata->sync_event_.sync_event);
+  if (pdata->sync_) {
+    pdata->sync_->clear_event();
   }
+
   delete pdata;
 
   ArrowArrayMarkReleased(array);
@@ -591,7 +590,7 @@ struct ArrayExporter {
     // Store owning pointer to ArrayData
     export_.data_ = data;
 
-    export_.sync_event_ = RawSyncEvent();
+    export_.sync_ = nullptr;
     return Status::OK();
   }
 
@@ -714,11 +713,11 @@ Result<std::pair<std::optional<DeviceAllocationType>, int64_t>> ValidateDeviceIn
   return std::make_pair(device_type, device_id);
 }
 
-Status ExportDeviceArray(const Array& array, RawSyncEvent sync_event,
+Status ExportDeviceArray(const Array& array, std::shared_ptr<DeviceSync>& sync,
                          struct ArrowDeviceArray* out, struct ArrowSchema* out_schema) {
-  if (sync_event.sync_event != nullptr && sync_event.release_func) {
-    return Status::Invalid(
-        "Must provide a release event function if providing a non-null event");
+  void* sync_event{nullptr};
+  if (sync) {
+    ARROW_ASSIGN_OR_RAISE(sync_event, sync->get_event());
   }
 
   SchemaExportGuard guard(out_schema);
@@ -739,19 +738,20 @@ Status ExportDeviceArray(const Array& array, RawSyncEvent sync_event,
   exporter.Finish(&out->array);
 
   auto* pdata = reinterpret_cast<ExportedArrayPrivateData*>(out->array.private_data);
-  pdata->sync_event_ = sync_event;
-  out->sync_event = sync_event.sync_event;
+  pdata->sync_ = sync;
+  out->sync_event = sync_event;
 
   guard.Detach();
   return Status::OK();
 }
 
-Status ExportDeviceRecordBatch(const RecordBatch& batch, RawSyncEvent sync_event,
+Status ExportDeviceRecordBatch(const RecordBatch& batch,
+                               std::shared_ptr<DeviceSync>& sync,
                                struct ArrowDeviceArray* out,
                                struct ArrowSchema* out_schema) {
-  if (sync_event.sync_event != nullptr && sync_event.release_func == nullptr) {
-    return Status::Invalid(
-        "Must provide a release event function if providing a non-null event");
+  void* sync_event{nullptr};
+  if (sync) {
+    ARROW_ASSIGN_OR_RAISE(sync_event, sync->get_event());
   }
 
   // XXX perhaps bypass ToStructArray for speed?
@@ -776,8 +776,8 @@ Status ExportDeviceRecordBatch(const RecordBatch& batch, RawSyncEvent sync_event
   exporter.Finish(&out->array);
 
   auto* pdata = reinterpret_cast<ExportedArrayPrivateData*>(out->array.private_data);
-  pdata->sync_event_ = sync_event;
-  out->sync_event = sync_event.sync_event;
+  pdata->sync_ = sync;
+  out->sync_event = sync_event;
 
   guard.Detach();
   return Status::OK();
@@ -1362,7 +1362,7 @@ namespace {
 // The ArrowArray is released on destruction.
 struct ImportedArrayData {
   struct ArrowArray array_;
-  void* sync_event_;
+  std::shared_ptr<DeviceSync> device_sync;
 
   ImportedArrayData() {
     ArrowArrayMarkReleased(&array_);  // Initially released
@@ -1395,6 +1395,8 @@ class ImportedBuffer : public Buffer {
 
   ~ImportedBuffer() override {}
 
+  std::shared_ptr<DeviceSync> get_device_sync() override { return import_->device_sync; }
+
  protected:
   std::shared_ptr<ImportedArrayData> import_;
 };
@@ -1409,7 +1411,8 @@ struct ArrayImporter {
     ARROW_ASSIGN_OR_RAISE(memory_mgr_, mapper(src->device_type, src->device_id));
     device_type_ = static_cast<DeviceAllocationType>(src->device_type);
     RETURN_NOT_OK(Import(&src->array));
-    import_->sync_event_ = src->sync_event;
+    ARROW_ASSIGN_OR_RAISE(import_->device_sync,
+                          memory_mgr_->MakeDeviceSync(src->sync_event));
     // reset internal state before next import
     memory_mgr_.reset();
     device_type_ = DeviceAllocationType::kCPU;

cpp/src/arrow/c/bridge.h

@@ -22,6 +22,7 @@
 #include <string>
 
 #include "arrow/c/abi.h"
+#include "arrow/device.h"
 #include "arrow/result.h"
 #include "arrow/status.h"
 #include "arrow/type_fwd.h"
@@ -172,33 +173,22 @@ Result<std::shared_ptr<RecordBatch>> ImportRecordBatch(struct ArrowArray* array,
 ///
 /// @{
 
-/// \brief EXPERIMENTAL: Type for freeing a sync event
-///
-/// If synchronization is necessary for accessing the data on a device,
-/// a pointer to an event needs to be passed when exporting the device
-/// array. It's the responsibility of the release function for the array
-/// to release the event. Both can be null if no sync'ing is necessary.
-struct RawSyncEvent {
-  void* sync_event = NULL;
-  std::function<void(void*)> release_func;
-};
-
 /// \brief EXPERIMENTAL: Export C++ Array as an ArrowDeviceArray.
 ///
 /// The resulting ArrowDeviceArray struct keeps the array data and buffers alive
 /// until its release callback is called by the consumer. All buffers in
 /// the provided array MUST have the same device_type, otherwise an error
 /// will be returned.
 ///
-/// If a non-null sync_event is provided, then the sync_release func must also be
-/// non-null. If the sync_event is null, then the sync_release parameter is not called.
+/// If sync is non-null, get_event will be called on it in order to
+/// potentially provide an event for consumers to synchronize on.
 ///
 /// \param[in] array Array object to export
-/// \param[in] sync_event A struct containing what is needed for syncing if necessary
+/// \param[in] sync shared_ptr to object derived from DeviceSync or null
 /// \param[out] out C struct to export the array to
 /// \param[out] out_schema optional C struct to export the array type to
 ARROW_EXPORT
-Status ExportDeviceArray(const Array& array, RawSyncEvent sync_event,
+Status ExportDeviceArray(const Array& array, std::shared_ptr<DeviceSync>& sync,
                          struct ArrowDeviceArray* out,
                          struct ArrowSchema* out_schema = NULLPTR);
 
@@ -212,15 +202,16 @@ Status ExportDeviceArray(const Array& array, RawSyncEvent sync_event,
 /// otherwise an error will be returned. If columns are on different devices,
 /// they should be exported using different ArrowDeviceArray instances.
 ///
-/// If a non-null sync_event is provided, then the sync_release func must also be
-/// non-null. If the sync_event is null, then the sync_release parameter is ignored.
+/// If sync is non-null, get_event will be called on it in order to
+/// potentially provide an event for consumers to synchronize on.
 ///
 /// \param[in] batch Record batch to export
-/// \param[in] sync_event A struct containing what is needed for syncing if necessary
+/// \param[in] sync shared_ptr to object derived from DeviceSync or null
 /// \param[out] out C struct where to export the record batch
 /// \param[out] out_schema optional C struct where to export the record batch schema
 ARROW_EXPORT
-Status ExportDeviceRecordBatch(const RecordBatch& batch, RawSyncEvent sync_event,
+Status ExportDeviceRecordBatch(const RecordBatch& batch,
+                               std::shared_ptr<DeviceSync>& sync,
                                struct ArrowDeviceArray* out,
                                struct ArrowSchema* out_schema = NULLPTR);