implement wait and fetch for the client

cpp/src/plasma/client.cc

@@ -320,6 +320,10 @@ Status PlasmaClient::PerformRelease(const ObjectID& object_id) {
 }
 
 Status PlasmaClient::Release(const ObjectID& object_id) {
+  // If the client is already disconnected, ignore release requests.
+  if (store_conn_ < 0) {
+    return Status::OK();
+  }
   // Add the new object to the release history.
   release_history_.push_front(object_id);
   // If there are too many bytes in use by the client or if there are too many
@@ -525,7 +529,11 @@ Status PlasmaClient::Disconnect() {
   // Close the connections to Plasma. The Plasma store will release the objects
   // that were in use by us when handling the SIGPIPE.
   close(store_conn_);
-  if (manager_conn_ >= 0) { close(manager_conn_); }
+  store_conn_ = -1;
+  if (manager_conn_ >= 0) {
+    close(manager_conn_);
+    manager_conn_ = -1;
+  }
   return Status::OK();
 }
 

cpp/src/plasma/common.cc

@@ -84,4 +84,7 @@ Status plasma_error_status(int plasma_error) {
   return Status::OK();
 }
 
+ARROW_EXPORT int ObjectStatusLocal = ObjectStatus_Local;
+ARROW_EXPORT int ObjectStatusRemote = ObjectStatus_Remote;
+
 } // namespace plasma

cpp/src/plasma/common.h

@@ -65,6 +65,36 @@ arrow::Status plasma_error_status(int plasma_error);
 /// Size of object hash digests.
 constexpr int64_t kDigestSize = sizeof(uint64_t);
 
+/// Object request data structure. Used for Wait.
+struct ObjectRequest {
+  /// The ID of the requested object. If ID_NIL request any object.
+  ObjectID object_id;
+  /// Request associated to the object. It can take one of the following values:
+  ///  - PLASMA_QUERY_LOCAL: return if or when the object is available in the
+  ///    local Plasma Store.
+  ///  - PLASMA_QUERY_ANYWHERE: return if or when the object is available in
+  ///    the system (i.e., either in the local or a remote Plasma Store).
+  int type;
+  /// Object status. Same as the status returned by plasma_status() function
+  /// call. This is filled in by plasma_wait_for_objects1():
+  ///  - ObjectStatus_Local: object is ready at the local Plasma Store.
+  ///  - ObjectStatus_Remote: object is ready at a remote Plasma Store.
+  ///  - ObjectStatus_Nonexistent: object does not exist in the system.
+  ///  - PLASMA_CLIENT_IN_TRANSFER, if the object is currently being scheduled
+  ///    for being transferred or it is transferring.
+  int status;
+};
+
+enum ObjectRequestType {
+  /// Query for object in the local plasma store.
+  PLASMA_QUERY_LOCAL = 1,
+  /// Query for object in the local plasma store or in a remote plasma store.
+  PLASMA_QUERY_ANYWHERE
+};
+
+extern int ObjectStatusLocal;
+extern int ObjectStatusRemote;
+
 } // namespace plasma
 
 #endif  // PLASMA_COMMON_H

cpp/src/plasma/plasma.h

@@ -58,27 +58,6 @@ namespace plasma {
 
 struct Client;
 
-/// Object request data structure. Used in the plasma_wait_for_objects()
-/// argument.
-struct ObjectRequest {
-  /// The ID of the requested object. If ID_NIL request any object.
-  ObjectID object_id;
-  /// Request associated to the object. It can take one of the following values:
-  ///  - PLASMA_QUERY_LOCAL: return if or when the object is available in the
-  ///    local Plasma Store.
-  ///  - PLASMA_QUERY_ANYWHERE: return if or when the object is available in
-  ///    the system (i.e., either in the local or a remote Plasma Store).
-  int type;
-  /// Object status. Same as the status returned by plasma_status() function
-  /// call. This is filled in by plasma_wait_for_objects1():
-  ///  - ObjectStatus_Local: object is ready at the local Plasma Store.
-  ///  - ObjectStatus_Remote: object is ready at a remote Plasma Store.
-  ///  - ObjectStatus_Nonexistent: object does not exist in the system.
-  ///  - PLASMA_CLIENT_IN_TRANSFER, if the object is currently being scheduled
-  ///    for being transferred or it is transferring.
-  int status;
-};
-
 /// Mapping from object IDs to type and status of the request.
 typedef std::unordered_map<ObjectID, ObjectRequest, UniqueIDHasher> ObjectRequestMap;
 
@@ -120,13 +99,6 @@ enum object_status {
   OBJECT_FOUND = 1
 };
 
-enum object_request_type {
-  /// Query for object in the local plasma store.
-  PLASMA_QUERY_LOCAL = 1,
-  /// Query for object in the local plasma store or in a remote plasma store.
-  PLASMA_QUERY_ANYWHERE
-};
-
 /// This type is used by the Plasma store. It is here because it is exposed to
 /// the eviction policy.
 struct ObjectTableEntry {

python/pyarrow/plasma.pyx

@@ -24,12 +24,16 @@ from libcpp.memory cimport shared_ptr, unique_ptr, make_shared
 from libcpp.string cimport string as c_string
 from libcpp.vector cimport vector as c_vector
 from libc.stdint cimport int64_t, uint8_t, uintptr_t
+from cpython.pycapsule cimport *
 
 from pyarrow.lib cimport Buffer, MutableBuffer, NativeFile, check_status
 from pyarrow.includes.libarrow cimport (CMutableBuffer, CBuffer,
                                         CFixedSizeBufferWrite, CStatus)
 
 
+PLASMA_WAIT_TIMEOUT = 2 ** 30
+
+
 cdef class FixedSizeBufferOutputStream(NativeFile):
 
     def __cinit__(self, Buffer buffer):
@@ -52,10 +56,21 @@ cdef extern from "plasma/common.h" nogil:
 
         c_string binary() const
 
+    cdef struct CObjectRequest" plasma::ObjectRequest":
+        CUniqueID object_id
+        int type
+        int status
+
 
 cdef extern from "plasma/common.h":
     cdef int64_t kDigestSize" plasma::kDigestSize"
 
+    cdef enum ObjectRequestType:
+        PLASMA_QUERY_LOCAL"plasma::PLASMA_QUERY_LOCAL",
+        PLASMA_QUERY_ANYWHERE"plasma::PLASMA_QUERY_ANYWHERE"
+
+    cdef int ObjectStatusLocal"plasma::ObjectStatusLocal";
+    cdef int ObjectStatusRemote"plasma::ObjectStatusRemote";
 
 cdef extern from "plasma/client.h" nogil:
 
@@ -90,6 +105,11 @@ cdef extern from "plasma/client.h" nogil:
 
         CStatus Disconnect()
 
+        CStatus Fetch(int num_object_ids, const CUniqueID* object_ids)
+
+        CStatus Wait(int64_t num_object_requests, CObjectRequest* object_requests,
+           int num_ready_objects, int64_t timeout_ms, int* num_objects_ready);
+
 
 cdef extern from "plasma/client.h" nogil:
 
@@ -117,12 +137,18 @@ cdef class ObjectID:
             raise ValueError("operation != 2 (only equality is supported)")
         return self.data == object_id.data
 
+    def __hash__(self):
+        return hash(self.data.binary())
+
     def __repr__(self):
         return "ObjectID(" + self.data.hex().decode() + ")"
 
     def __reduce__(self):
         return (make_object_id, (self.data.binary(),))
 
+    def binary(self):
+        return self.data.binary()
+
 
 cdef class PlasmaBuffer(Buffer):
     """
@@ -422,6 +448,69 @@ cdef class PlasmaClient:
             check_status(self.client.get().Evict(num_bytes, num_bytes_evicted))
         return num_bytes_evicted
 
+    def fetch(self, object_ids):
+        """
+        Fetch the objects with the given IDs from other plasma managers.
+
+        Parameters
+        ----------
+        object_ids : list
+            A list of strings used to identify the objects.
+        """
+        cdef c_vector[CUniqueID] ids
+        cdef ObjectID object_id
+        for object_id in object_ids:
+            ids.push_back(object_id.data)
+        with nogil:
+            check_status(self.client.get().Fetch(ids.size(), ids.data()))
+
+    def wait(self, object_ids, int64_t timeout=PLASMA_WAIT_TIMEOUT, int num_returns=1):
+        """
+        Wait until num_returns objects in object_ids are ready.
+        Currently, the object ID arguments to wait must be unique.
+
+        Parameters
+        ----------
+        object_ids : list
+            List of object IDs to wait for.
+        timeout :int
+            Return to the caller after timeout milliseconds.
+        num_returns : int
+            We are waiting for this number of objects to be ready.
+
+        Returns
+        -------
+        list
+            List of object IDs that are ready.
+        list
+            List of object IDs we might still wait on.
+        """
+        # Check that the object ID arguments are unique. The plasma manager
+        # currently crashes if given duplicate object IDs.
+        if len(object_ids) != len(set(object_ids)):
+            raise Exception("Wait requires a list of unique object IDs.")
+        cdef int64_t num_object_requests = len(object_ids)
+        cdef c_vector[CObjectRequest] object_requests = c_vector[CObjectRequest](num_object_requests)
+        cdef int num_objects_ready = 0
+        cdef ObjectID object_id
+        for i, object_id in enumerate(object_ids):
+            object_requests[i].object_id = object_id.data
+            object_requests[i].type = PLASMA_QUERY_ANYWHERE
+        with nogil:
+            check_status(self.client.get().Wait(num_object_requests, object_requests.data(), num_returns, timeout, &num_objects_ready))
+        cdef int num_to_return = min(num_objects_ready, num_returns);
+        ready_ids = []
+        waiting_ids = set(object_ids)
+        cdef int num_returned = 0
+        for i in range(len(object_ids)):
+            if num_returned == num_to_return:
+                break
+            if object_requests[i].status == ObjectStatusLocal or object_requests[i].status == ObjectStatusRemote:
+                ready_ids.append(ObjectID(object_requests[i].object_id.binary()))
+                waiting_ids.discard(ObjectID(object_requests[i].object_id.binary()))
+                num_returned += 1
+        return ready_ids, list(waiting_ids)
+
     def subscribe(self):
         """Subscribe to notifications about sealed objects."""
         with nogil:
@@ -450,6 +539,9 @@ cdef class PlasmaClient:
                                                            &metadata_size))
         return object_id, data_size, metadata_size
 
+    def to_capsule(self):
+        return PyCapsule_New(<void *>self.client.get(), "plasma", NULL)
+
     def disconnect(self):
         """
         Disconnect this client from the Plasma store.