Add plasma queue python interface

Author: songqing

cpp/cmake_modules/CompilerInfo.cmake

@@ -78,6 +78,11 @@ elseif("${COMPILER_VERSION_FULL}" MATCHES ".*clang-9")
   set(COMPILER_FAMILY "clang")
   set(COMPILER_VERSION "4.0.0svn")
 
+# clang on Mac OS X, XCode 9.
+elseif("${COMPILER_VERSION_FULL}" MATCHES ".*clang-10")
+  set(COMPILER_FAMILY "clang")
+  set(COMPILER_VERSION "4.1.0svn")
+
 # gcc
 elseif("${COMPILER_VERSION_FULL_LOWER}" MATCHES ".*gcc[ -]version.*")
   set(COMPILER_FAMILY "gcc")

python/pyarrow/_plasma.pyx

@@ -23,7 +23,7 @@ from libcpp cimport bool as c_bool, nullptr
 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 libc.stdint cimport int64_t, uint8_t, uintptr_t, uint32_t, uint64_t
 from cpython.pycapsule cimport *
 
 import collections
@@ -113,6 +113,22 @@ cdef extern from "plasma/client.h" nogil:
         CStatus Transfer(const char* addr, int port,
                          const CUniqueID& object_id)
 
+        # Interfaces that are related to plasma queue.
+        CStatus CreateQueue(const CUniqueID& object_id, int64_t data_size,
+                            const shared_ptr[CBuffer]* data)
+
+        CStatus GetQueue(const CUniqueID& object_id, int64_t timeout_ms,
+                         int* fd)
+
+        CStatus CreateQueueItem(const CUniqueID& object_id, uint32_t data_size,
+                                shared_ptr[CBuffer]* data,
+                                uint64_t& seq_id)
+
+        CStatus SealQueueItem(const CUniqueID& object_id, uint64_t seq_id,
+                              shared_ptr[CBuffer] data)
+
+        CStatus GetQueueItem(const CUniqueID& object_id,
+                             CObjectBuffer* object_buffer, uint64_t& seq_id)
 
 cdef extern from "plasma/client.h" nogil:
 
@@ -396,6 +412,138 @@ cdef class PlasmaClient:
         self.seal(target_id)
         return target_id
 
+    def create_queue(self, ObjectID queue_id=None, total_bytes=1024000):
+        """
+        Store a plasma queue into the object store.
+
+        Parameters
+        ----------
+        queue_id : ObjectID, default None
+            If this is provided, the specified object ID will be used to refer
+            to the object.
+        -------
+        """
+        cdef ObjectID target_id = (queue_id if queue_id
+                                   else ObjectID.from_random())
+        """
+        Create a new buffer in the PlasmaStore for a particular object ID.
+        The returned buffer is unmutable as seal is called.
+
+        Raises
+        ------
+        PlasmaObjectExists
+            This exception is raised if the object could not be created because
+            there already is an object with the same ID in the plasma store.
+
+        PlasmaStoreFull: This exception is raised if the object could
+                not be created because the plasma store is unable to evict
+                enough objects to create room for it.
+        """
+        # Apply a buffer from plasma store for the plasma queue
+        cdef shared_ptr[CBuffer] data
+        check_status(self.client.get().CreateQueue(target_id.data, 
+                                                   total_bytes, &data))
+
+    def push_queue(self, object value, ObjectID queue_id=None,
+            int memcopy_threads=6, serialization_context=None):
+        """
+        Push a Python value into the plasma queue.
+
+        Parameters
+        ----------
+        value : object
+            A Python object to store.
+        queue_id : ObjectID, default None
+            If this is provided, the specified object ID will be used to refer
+            to the plasma queue.
+        memcopy_threads : int, default 6
+            The number of threads to use to write the serialized object into
+            the object store for large objects.
+        serialization_context : pyarrow.SerializationContext, default None
+            Custom serialization and deserialization context.
+
+        """
+        cdef ObjectID target_id = (queue_id if queue_id
+                                   else ObjectID.from_random())
+        serialized = pyarrow.serialize(value, serialization_context)
+
+        # Return buffer from the plasma queue to write
+        cdef uint64_t seq_id = 0
+        cdef shared_ptr[CBuffer] data
+        check_status(self.client.get().CreateQueueItem(target_id.data,
+            serialized.total_bytes, &data, seq_id))
+        buffer = self._make_mutable_plasma_buffer(target_id,
+                                                data.get().mutable_data(),
+                                                serialized.total_bytes)
+
+        stream = pyarrow.FixedSizeBufferWriter(buffer)
+        stream.set_memcopy_threads(memcopy_threads)
+        serialized.write_to(stream)
+        with nogil:
+            check_status(self.client.get().SealQueueItem(target_id.data,
+                                                        seq_id, data))
+        
+    def get_queue(self, queue_id, timeout_ms=-1):
+        """
+        Subscribe the corresponding plasma queue.
+
+        Parameters
+        ----------
+        queue_id : ObjectID
+            Object ID associated to the plasma queue we get
+            from the store.
+        timeout_ms : int, default -1
+            The number of milliseconds that the get call should block before
+            timing out and returning. Pass -1 if the call should block and 0
+            if the call should return immediately.
+
+        Returns
+        -------
+        """
+
+        cdef ObjectID target_id = (queue_id if queue_id
+                                   else ObjectID.from_random())
+        cdef int nofify_fd
+        check_status(self.client.get().GetQueue(target_id.data,
+                                                   timeout_ms, &nofify_fd))
+
+    def read_queue(self, queue_id, index, int timeout_ms=-1,
+            serialization_context=None):
+        """
+        Get one Python value from the plasma queue.
+
+        Parameters
+        ----------
+        queue_id : ObjectID
+            Object ID associated to the plasma queue we get
+            from the store.
+        timeout_ms : int, default -1
+            The number of milliseconds that the get call should block before
+            timing out and returning. Pass -1 if the call should block and 0
+            if the call should return immediately.
+        serialization_context : pyarrow.SerializationContext, default None
+            Custom serialization and deserialization context.
+
+        Returns
+        -------
+        object
+            Python value associated with the current index in the plasma queue,
+            and ObjectNotAvailable if the object was not available.
+        """
+
+        cdef ObjectID target_id = (queue_id if queue_id
+                                   else ObjectID.from_random())
+        cdef CObjectBuffer object_buffer
+        check_status(self.client.get().GetQueueItem(target_id.data,
+                                                    &object_buffer, index))
+        buffer = pyarrow_wrap_buffer(object_buffer.data)
+        # buffer is None if this object was not available within the timeout
+        if buffer:
+            value = pyarrow.deserialize(buffer, serialization_context)
+            return value
+        else:
+            return ObjectNotAvailable
+
     def get(self, object_ids, int timeout_ms=-1, serialization_context=None):
         """
         Get one or more Python values from the object store.