get all python tests in place

Author: pcmoritz

cpp/src/plasma/client.cc

@@ -388,22 +388,6 @@ static uint64_t compute_object_hash(const ObjectBuffer& obj_buffer) {
   return XXH64_digest(&hash_state);
 }
 
-bool plasma_compute_object_hash(
-    PlasmaClient* conn, ObjectID object_id, unsigned char* digest) {
-  // Get the plasma object data. We pass in a timeout of 0 to indicate that
-  // the operation should timeout immediately.
-  ObjectBuffer object_buffer;
-  ARROW_CHECK_OK(conn->Get(&object_id, 1, 0, &object_buffer));
-  // If the object was not retrieved, return false.
-  if (object_buffer.data_size == -1) { return false; }
-  // Compute the hash.
-  uint64_t hash = compute_object_hash(object_buffer);
-  memcpy(digest, &hash, sizeof(hash));
-  // Release the plasma object.
-  ARROW_CHECK_OK(conn->Release(object_id));
-  return true;
-}
-
 Status PlasmaClient::Seal(const ObjectID& object_id) {
   // Make sure this client has a reference to the object before sending the
   // request to Plasma.
@@ -415,7 +399,7 @@ Status PlasmaClient::Seal(const ObjectID& object_id) {
   object_entry->second->is_sealed = true;
   /// Send the seal request to Plasma.
   static unsigned char digest[kDigestSize];
-  ARROW_CHECK(plasma_compute_object_hash(this, object_id, &digest[0]));
+  RETURN_NOT_OK(Hash(object_id, &digest[0]));
   RETURN_NOT_OK(SendSealRequest(store_conn_, object_id, &digest[0]));
   // We call PlasmaClient::Release to decrement the number of instances of this
   // object
@@ -441,6 +425,22 @@ Status PlasmaClient::Evict(int64_t num_bytes, int64_t& num_bytes_evicted) {
   return ReadEvictReply(buffer.data(), num_bytes_evicted);
 }
 
+Status PlasmaClient::Hash(const ObjectID& object_id, uint8_t* digest) {
+  // Get the plasma object data. We pass in a timeout of 0 to indicate that
+  // the operation should timeout immediately.
+  ObjectBuffer object_buffer;
+  RETURN_NOT_OK(Get(&object_id, 1, 0, &object_buffer));
+  // If the object was not retrieved, return false.
+  if (object_buffer.data_size == -1) {
+    return Status::PlasmaObjectNonexistent("Object not found");
+  }
+  // Compute the hash.
+  uint64_t hash = compute_object_hash(object_buffer);
+  memcpy(digest, &hash, sizeof(hash));
+  // Release the plasma object.
+  return Release(object_id);
+}
+
 Status PlasmaClient::Subscribe(int* fd) {
   int sock[2];
   // Create a non-blocking socket pair. This will only be used to send
@@ -461,6 +461,25 @@ Status PlasmaClient::Subscribe(int* fd) {
   return Status::OK();
 }
 
+Status PlasmaClient::GetNotification(int fd, ObjectID *object_id, int64_t* data_size, int64_t* metadata_size) {
+  uint8_t* notification = read_message_async(fd);
+  if (notification == NULL) {
+    return Status::IOError("Failed to read object notification from Plasma socket");
+  }
+  auto object_info = flatbuffers::GetRoot<ObjectInfo>(notification);
+  ARROW_CHECK(object_info->object_id()->size() == sizeof(ObjectID));
+  memcpy(object_id, object_info->object_id()->data(), sizeof(ObjectID));
+  if (object_info->is_deletion()) {
+    *data_size = -1;
+    *metadata_size = -1;
+  } else {
+    *data_size = object_info->data_size();
+    *metadata_size = object_info->metadata_size();
+  }
+  delete[] notification;
+  return Status::OK();
+}
+
 Status PlasmaClient::Connect(const std::string& store_socket_name,
     const std::string& manager_socket_name, int release_delay) {
   store_conn_ = connect_ipc_sock_retry(store_socket_name, -1, -1);

cpp/src/plasma/client.h

@@ -179,18 +179,35 @@ class ARROW_EXPORT PlasmaClient {
   /// retrieved.
   /// @return The return status.
   Status Evict(int64_t num_bytes, int64_t& num_bytes_evicted);
+  
+  /// Compute the hash of an object in the object store.
+  ///
+  /// @param conn The object containing the connection state.
+  /// @param object_id The ID of the object we want to hash.
+  /// @param digest A pointer at which to return the hash digest of the object.
+  ///        The pointer must have at least kDigestSize bytes allocated.
+  /// @return The return status.
+  Status Hash(const ObjectID& object_id, uint8_t* digest);
 
   /// Subscribe to notifications when objects are sealed in the object store.
   /// Whenever an object is sealed, a message will be written to the client
-  /// socket
-  /// that is returned by this method.
+  /// socket that is returned by this method.
   ///
   /// @param fd Out parameter for the file descriptor the client should use to
   /// read notifications
   ///         from the object store about sealed objects.
   /// @return The return status.
   Status Subscribe(int* fd);
 
+  /// Receive next object notification for this client if Subscribe has been called.
+  ///
+  /// @param fd The file descriptor we are reading the notification from.
+  /// @param object_id Out parameter, the object_id of the object that was sealed.
+  /// @param data_size Out parameter, the data size of the object that was sealed.
+  /// @param metadata_size Out parameter, the metadata size of the object that was sealed.
+  /// @return The return status.
+  Status GetNotification(int fd, ObjectID* object_id, int64_t* data_size, int64_t* metadata_size);
+
   /// Disconnect from the local plasma instance, including the local store and
   /// manager.
   ///
@@ -333,14 +350,4 @@ class ARROW_EXPORT PlasmaClient {
   int64_t store_capacity_;
 };
 
-/// Compute the hash of an object in the object store.
-///
-/// @param conn The object containing the connection state.
-/// @param object_id The ID of the object we want to hash.
-/// @param digest A pointer at which to return the hash digest of the object.
-///        The pointer must have at least DIGEST_SIZE bytes allocated.
-/// @return A boolean representing whether the hash operation succeeded.
-bool plasma_compute_object_hash(
-    PlasmaClient* conn, ObjectID object_id, unsigned char* digest);
-
 #endif  // PLASMA_CLIENT_H

cpp/src/plasma/plasma.pyx

@@ -26,8 +26,13 @@ cdef extern from "plasma/common.h" nogil:
     @staticmethod
     CUniqueID from_binary(const c_string& binary)
 
+    c_bool operator==(const CUniqueID& rhs) const
+
     c_string hex()
 
+cdef extern from "plasma/plasma.h":
+    cdef int64_t kDigestSize
+
 cdef extern from "plasma/client.h" nogil:
 
   cdef cppclass CPlasmaClient" PlasmaClient":
@@ -43,8 +48,18 @@ cdef extern from "plasma/client.h" nogil:
 
     CStatus Seal(const CUniqueID& object_id)
 
+    CStatus Evict(int64_t num_bytes, int64_t& num_bytes_evicted)
+
+    CStatus Hash(const CUniqueID& object_id, uint8_t* digest)
+
     CStatus Release(const CUniqueID& object_id)
 
+    CStatus Contains(const CUniqueID& object_id, c_bool* has_object)
+
+    CStatus Subscribe(int* fd)
+
+    CStatus GetNotification(int fd, CUniqueID* object_id, int64_t* data_size, int64_t* metadata_size)
+
     CStatus Disconnect()
 
 cdef extern from "plasma/client.h" nogil:
@@ -63,6 +78,10 @@ cdef class ObjectID:
   def __cinit__(self, object_id):
     self.data = CUniqueID.from_binary(object_id)
 
+  def __richcmp__(ObjectID self, ObjectID object_id, operation):
+    assert operation == 2, "only equality implemented so far"
+    return self.data == object_id.data
+
   def __repr__(self):
     return "ObjectID(" + self.data.hex().decode() + ")"
 
@@ -83,9 +102,11 @@ cdef class PlasmaClient:
 
   cdef:
     shared_ptr[CPlasmaClient] client
+    int notification_fd
 
   def __cinit__(self):
     self.client.reset(new CPlasmaClient())
+    self.notification_fd = -1
 
   cdef _get_object_buffers(self, object_ids, int64_t timeout_ms, c_vector[CObjectBuffer]* result):
     cdef c_vector[CUniqueID] ids
@@ -135,5 +156,30 @@ cdef class PlasmaClient:
   def release(self, ObjectID object_id):
     check_status(self.client.get().Release(object_id.data))
 
+  def contains(self, ObjectID object_id):
+    cdef c_bool is_contained
+    check_status(self.client.get().Contains(object_id.data, &is_contained))
+    return is_contained
+
+  def hash(self, ObjectID object_id):
+    cdef c_vector[uint8_t] digest = c_vector[uint8_t](kDigestSize)
+    check_status(self.client.get().Hash(object_id.data, digest.data()))
+    return bytes(digest[:])
+
+  def evict(self, int64_t num_bytes):
+    cdef int64_t num_bytes_evicted
+    check_status(self.client.get().Evict(num_bytes, num_bytes_evicted))
+    return num_bytes_evicted
+
+  def subscribe(self):
+    check_status(self.client.get().Subscribe(&self.notification_fd))
+
+  def get_next_notification(self):
+    cdef ObjectID object_id = ObjectID(20 * b"\0")
+    cdef int64_t data_size
+    cdef int64_t metadata_size
+    check_status(self.client.get().GetNotification(self.notification_fd, &object_id.data, &data_size, &metadata_size))
+    return object_id, data_size, metadata_size
+
   def disconnect(self):
     check_status(self.client.get().Disconnect())