fix

NonVolatileComputing · Aug 26, 2017 · 3518c71 · 3518c71
1 parent e1924a4
commit 3518c71
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/python/pyarrow/__init__.py b/python/pyarrow/__init__.py
@@ -85,12 +85,14 @@
                          ArrowIOError,
                          ArrowMemoryError,
                          ArrowNotImplementedError,
-                         ArrowTypeError)
+                         ArrowTypeError,
+                         PlasmaObjectExists)
 
 # Serialization
 from pyarrow.lib import (deserialize_from, deserialize,
                          serialize, serialize_to, read_serialized,
                          SerializedPyObject,
+                         SerializationException, DeserializationException,
                          # This is temporary
                          register_type, type_to_type_id, whitelisted_types,
                          types_to_pickle, custom_serializers, custom_deserializers)

diff --git a/python/pyarrow/plasma.pyx b/python/pyarrow/plasma.pyx
@@ -591,12 +591,14 @@ def connect(store_socket_name, manager_socket_name, int release_delay,
                               release_delay, num_retries))
     return result
 
-def put(PlasmaClient client, value):
-    cdef ObjectID object_id = ObjectID.from_random()
+def put(PlasmaClient client, value, object_id=None):
+    cdef ObjectID id = object_id if object_id else ObjectID.from_random()
     cdef SerializedPyObject serialized = pyarrow.serialize(value)
-    buffer = client.create(object_id, serialized.total_bytes)
-    serialized.write_to(buffer)
-    return object_id
+    buffer = client.create(id, serialized.total_bytes)
+    stream = pyarrow.FixedSizeBufferOutputStream(buffer)
+    stream.set_memcopy_threads(4)
+    serialized.write_to(stream)
+    return id
 
 def get(PlasmaClient client, object_ids, timeout_ms=-1):
     results = []