move put and get into PlasmaClient

python/pyarrow/plasma.pyx

@@ -303,7 +303,7 @@ cdef class PlasmaClient:
                                                   metadata.size(), &data))
         return self._make_mutable_plasma_buffer(object_id, data, data_size)
 
-    def get(self, object_ids, timeout_ms=-1):
+    def get_buffer(self, object_ids, timeout_ms=-1):
         """
         Returns data buffer from the PlasmaStore based on object ID.
 
@@ -370,6 +370,65 @@ cdef class PlasmaClient:
                                          object_buffers[i].metadata_size))
         return result
 
+    def put(self, list value, ObjectID object_id=None):
+        """
+        Store a Python value into the object store.
+
+        Parameters
+        ----------
+        value : list
+            A Python object to store. Currently only lists are supported.
+        object_id : ObjectID, default None
+            If this is provided, the specified object ID will be used to refer
+            to the object.
+
+        Returns
+        -------
+        The object ID associated to the Python object.
+        """
+        cdef ObjectID target_id = object_id if object_id else ObjectID.from_random()
+        serialized = pyarrow.serialize(value)
+        buffer = self.create(target_id, serialized.total_bytes)
+        stream = pyarrow.FixedSizeBufferOutputStream(buffer)
+        stream.set_memcopy_threads(4)
+        serialized.write_to(stream)
+        self.seal(target_id)
+        return target_id
+
+    def get(self, object_ids, int timeout_ms=-1):
+        """
+        Get one or more Python values from the object store.
+
+        Parameters
+        ----------
+        object_ids : list or ObjectID
+            Object ID or list of object IDs associated to the values 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
+        -------
+        list or object
+            Python value or list of Python values for the data associated with
+            the object_ids and ObjectNotAvailable if the object was not available.
+        """
+        if isinstance(object_ids, collections.Sequence):
+            results = []
+            buffers = self.get_buffer(object_ids, timeout_ms)
+            for i in range(len(object_ids)):
+                # buffers[i] is None if this object was not available within the
+                # timeout
+                if buffers[i]:
+                    results.append(pyarrow.deserialize(buffers[i]))
+                else:
+                    results.append(ObjectNotAvailable)
+            return results
+        else:
+            return self.get([object_ids], timeout_ms)[0]
+
     def seal(self, ObjectID object_id):
         """
         Seal the buffer in the PlasmaStore for a particular object ID.
@@ -606,66 +665,3 @@ def connect(store_socket_name, manager_socket_name, int release_delay,
                               result.manager_socket_name,
                               release_delay, num_retries))
     return result
-
-def put(PlasmaClient client, list value, ObjectID object_id=None):
-    """
-    Store a Python value into the object store.
-
-    Parameters
-    ----------
-    client : PlasmaClient
-        The client connected to the object store we put the value in.
-    value : list
-        A Python object to store. Currently only lists are supported.
-    object_id : ObjectID, default None
-        If this is provided, the specified object ID will be used to refer
-        to the object.
-
-    Returns
-    -------
-    The object ID associated to the Python object.
-    """
-    cdef ObjectID id = object_id if object_id else ObjectID.from_random()
-    serialized = pyarrow.serialize(value)
-    buffer = client.create(id, serialized.total_bytes)
-    stream = pyarrow.FixedSizeBufferOutputStream(buffer)
-    stream.set_memcopy_threads(4)
-    serialized.write_to(stream)
-    client.seal(id)
-    return id
-
-def get(PlasmaClient client, object_ids, int timeout_ms=-1):
-    """
-    Get one or more Python values from the object store.
-
-    Parameters
-    ----------
-    client : PlasmaClient
-        The client connected to the object store we get objects from.
-    object_ids : list or ObjectID
-        Object ID or list of object IDs associated to the values 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
-    -------
-    list or object
-        Python value or list of Python values for the data associated with
-        the object_ids and ObjectNotAvailable if the object was not available.
-    """
-    if isinstance(object_ids, collections.Sequence):
-        results = []
-        buffers = client.get(object_ids, timeout_ms)
-        for i in range(len(object_ids)):
-            # buffers[i] is None if this object was not available within the
-            # timeout
-            if buffers[i]:
-                results.append(pyarrow.deserialize(buffers[i]))
-            else:
-                results.append(ObjectNotAvailable)
-        return results
-    else:
-        return get(client, [object_ids], timeout_ms)[0]

python/pyarrow/tests/test_plasma.py

@@ -82,8 +82,8 @@ def create_object(client, data_size, metadata_size, seal=True):
 def assert_get_object_equal(unit_test, client1, client2, object_id,
                             memory_buffer=None, metadata=None):
     import pyarrow.plasma as plasma
-    client1_buff = client1.get([object_id])[0]
-    client2_buff = client2.get([object_id])[0]
+    client1_buff = client1.get_buffer([object_id])[0]
+    client2_buff = client2.get_buffer([object_id])[0]
     client1_metadata = client1.get_metadata([object_id])[0]
     client2_metadata = client2.get_metadata([object_id])[0]
     assert len(client1_buff) == len(client2_buff)
@@ -187,7 +187,7 @@ class TestPlasmaClient(object):
         # Seal the object.
         self.plasma_client.seal(object_id)
         # Get the object.
-        memory_buffer = np.frombuffer(self.plasma_client.get([object_id])[0],
+        memory_buffer = np.frombuffer(self.plasma_client.get_buffer([object_id])[0],
                                       dtype="uint8")
         for i in range(length):
             assert memory_buffer[i] == i % 256
@@ -209,7 +209,7 @@ class TestPlasmaClient(object):
             self.plasma_client.seal(object_id)
             # Get the object.
             memory_buffer = np.frombuffer(
-                self.plasma_client.get([object_id])[0], dtype="uint8")
+                self.plasma_client.get_buffer([object_id])[0], dtype="uint8")
             for i in range(length):
                 assert memory_buffer[i] == i % 256
             # Get the metadata.
@@ -241,7 +241,7 @@ class TestPlasmaClient(object):
         # Test timing out of get with various timeouts.
         for timeout in [0, 10, 100, 1000]:
             object_ids = [random_object_id() for _ in range(num_object_ids)]
-            results = self.plasma_client.get(object_ids, timeout_ms=timeout)
+            results = self.plasma_client.get_buffer(object_ids, timeout_ms=timeout)
             assert results == num_object_ids * [None]
 
         data_buffers = []
@@ -256,8 +256,8 @@ class TestPlasmaClient(object):
         # Test timing out from some but not all get calls with various
         # timeouts.
         for timeout in [0, 10, 100, 1000]:
-            data_results = self.plasma_client.get(object_ids,
-                                                  timeout_ms=timeout)
+            data_results = self.plasma_client.get_buffer(object_ids,
+                                                         timeout_ms=timeout)
             # metadata_results = self.plasma_client.get_metadata(
             #     object_ids, timeout_ms=timeout)
             for i in range(num_object_ids):
@@ -275,15 +275,15 @@ class TestPlasmaClient(object):
 
     def test_put_and_get(self):
         value = ["hello", "world", 3, 1.0]
-        object_id = pa.plasma.put(self.plasma_client, value)
-        [result] = pa.plasma.get(self.plasma_client, [object_id])
+        object_id = self.plasma_client.put(value)
+        [result] = self.plasma_client.get([object_id])
         assert result == value
 
-        result = pa.plasma.get(self.plasma_client, object_id)
+        result = self.plasma_client.get(object_id)
         assert result == value
 
         object_id = pa.plasma.ObjectID.from_random()
-        [result] = pa.plasma.get(self.plasma_client, [object_id], timeout_ms=0)
+        [result] = self.plasma_client.get([object_id], timeout_ms=0)
         assert result == pa.plasma.ObjectNotAvailable
 
     def test_store_arrow_objects(self):
@@ -297,7 +297,7 @@ class TestPlasmaClient(object):
         pa.write_tensor(tensor, stream)
         self.plasma_client.seal(object_id)
         # Read the arrow object.
-        [tensor] = self.plasma_client.get([object_id])
+        [tensor] = self.plasma_client.get_buffer([object_id])
         reader = pa.BufferReader(tensor)
         array = pa.read_tensor(reader).to_numpy()
         # Assert that they are equal.
@@ -326,7 +326,7 @@ class TestPlasmaClient(object):
         self.plasma_client.seal(object_id)
 
         # Read the DataFrame.
-        [data] = self.plasma_client.get([object_id])
+        [data] = self.plasma_client.get_buffer([object_id])
         reader = pa.RecordBatchStreamReader(pa.BufferReader(data))
         result = reader.get_next_batch().to_pandas()
 
@@ -564,7 +564,7 @@ class TestPlasmaClient(object):
         # with pytest.raises(Exception):
         # illegal_assignment()
         # Get the object.
-        memory_buffer = self.plasma_client.get([object_id])[0]
+        memory_buffer = self.plasma_client.get_buffer([object_id])[0]
 
         # Make sure the object is read only.
         def illegal_assignment():