fix callbacks

NonVolatileComputing · Aug 18, 2017 · aeafd82 · aeafd82
1 parent 389bfc6
commit aeafd82
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 22 deletions.
diff --git a/cpp/src/arrow/python/arrow_to_python.cc b/cpp/src/arrow/python/arrow_to_python.cc
@@ -24,6 +24,11 @@
 #include "arrow/python/helpers.h"
 #include "arrow/python/numpy_convert.h"
 
+extern "C" {
+extern PyObject* pyarrow_serialize_callback;
+extern PyObject* pyarrow_deserialize_callback;
+}
+
 namespace arrow {
 namespace py {
 

diff --git a/cpp/src/arrow/python/arrow_to_python.h b/cpp/src/arrow/python/arrow_to_python.h
@@ -25,16 +25,9 @@
 
 #include <vector>
 
-extern "C" {
-extern PyObject* pyarrow_serialize_callback;
-extern PyObject* pyarrow_deserialize_callback;
-}
-
 namespace arrow {
 namespace py {
 
-Status CallCustomCallback(PyObject* callback, PyObject* elem, PyObject** result);
-
 Status ReadSerializedPythonSequence(std::shared_ptr<io::RandomAccessFile> src,
                                     std::shared_ptr<RecordBatch>* batch_out,
                                     std::vector<std::shared_ptr<Tensor>>* tensors_out);

diff --git a/cpp/src/arrow/python/python_to_arrow.cc b/cpp/src/arrow/python/python_to_arrow.cc
@@ -116,6 +116,12 @@ Status CallCustomCallback(PyObject* callback, PyObject* elem, PyObject** result)
   return Status::OK();
 }
 
+void set_serialization_callbacks(PyObject* serialize_callback,
+                                 PyObject* deserialize_callback) {
+  pyarrow_serialize_callback = serialize_callback;
+  pyarrow_deserialize_callback = deserialize_callback;
+}
+
 Status CallCustomSerializationCallback(PyObject* elem, PyObject** serialized_object) {
   RETURN_NOT_OK(CallCustomCallback(pyarrow_serialize_callback, elem, serialized_object));
   if (!PyDict_Check(*serialized_object)) {

diff --git a/cpp/src/arrow/python/python_to_arrow.h b/cpp/src/arrow/python/python_to_arrow.h
@@ -27,14 +27,12 @@
 
 #include <vector>
 
-extern "C" {
-extern PyObject* pyarrow_serialize_callback;
-extern PyObject* pyarrow_deserialize_callback;
-}
-
 namespace arrow {
 namespace py {
 
+void set_serialization_callbacks(PyObject* serialize_callback,
+                                 PyObject* deserialize_callback);
+
 Status SerializePythonSequence(PyObject* sequence,
                                std::shared_ptr<RecordBatch>* batch_out,
                                std::vector<std::shared_ptr<Tensor>>* tensors_out);

diff --git a/python/pyarrow/serialization.pxi b/python/pyarrow/serialization.pxi
@@ -41,18 +41,14 @@ cdef extern from "arrow/python/api.h" namespace 'arrow::py' nogil:
         shared_ptr[CRecordBatch] batch,
         c_vector[shared_ptr[CTensor]] tensors,
         OutputStream* dst)
-    
+
     cdef CStatus ReadSerializedPythonSequence(
         shared_ptr[RandomAccessFile] src,
         shared_ptr[CRecordBatch]* batch_out,
         c_vector[shared_ptr[CTensor]]* tensors_out)
 
-cdef extern from "arrow/python/python_to_arrow.h":
-
-    cdef extern PyObject *pyarrow_serialize_callback
-
-    cdef extern PyObject *pyarrow_deserialize_callback
-
+    void set_serialization_callbacks(PyObject* serialize_callback,
+                                     PyObject* deserialize_callback);
 
 def is_named_tuple(cls):
     """Return True if cls is a namedtuple and False otherwise."""
@@ -157,9 +153,8 @@ def _deserialization_callback(serialized_obj):
                 obj.__dict__.update(serialized_obj)
     return obj
 
-pyarrow_serialize_callback = <PyObject*> _serialization_callback
-
-pyarrow_deserialize_callback = <PyObject*> _deserialization_callback
+set_serialization_callbacks(<PyObject*> _serialization_callback,
+                            <PyObject*> _deserialization_callback)
 
 def serialize_sequence(object value, NativeFile sink):
     """Serialize a Python sequence to a file.