From 0cc328e8baac7b4e588ec50adab136f1fbd90afe Mon Sep 17 00:00:00 2001
From: David Li <li.davidm96@gmail.com>
Date: Fri, 18 Aug 2023 15:26:41 -0400
Subject: [PATCH] Add async.h

---
 python/pyarrow/includes/common.pxd       |  6 ++-
 python/pyarrow/src/arrow/python/async.h  | 50 ++++++++++++++++++++++++
 python/pyarrow/src/arrow/python/common.h | 18 ---------
 3 files changed, 54 insertions(+), 20 deletions(-)
 create mode 100644 python/pyarrow/src/arrow/python/async.h

diff --git a/python/pyarrow/includes/common.pxd b/python/pyarrow/includes/common.pxd
index 76b8c256da91b..77fc606a83a82 100644
--- a/python/pyarrow/includes/common.pxd
+++ b/python/pyarrow/includes/common.pxd
@@ -157,12 +157,14 @@ cdef extern from "arrow/util/future.h" namespace "arrow" nogil:
 ctypedef object PyWrapper(void*)
 
 
+cdef extern from "arrow/python/async.h" namespace "arrow::py" nogil:
+    void BindFuture[T](CFuture[T], object cb, PyWrapper wrapper)
+
+
 cdef extern from "arrow/python/common.h" namespace "arrow::py" nogil:
     T GetResultValue[T](CResult[T]) except *
     cdef function[F] BindFunction[F](void* unbound, object bound, ...)
 
-    void BindFuture[T](CFuture[T], object cb, PyWrapper wrapper)
-
 
 cdef inline object PyObject_to_object(PyObject* o):
     # Cast to "object" increments reference count
diff --git a/python/pyarrow/src/arrow/python/async.h b/python/pyarrow/src/arrow/python/async.h
new file mode 100644
index 0000000000000..da9402d0bdfae
--- /dev/null
+++ b/python/pyarrow/src/arrow/python/async.h
@@ -0,0 +1,50 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <utility>
+
+#include "arrow/python/common.h"
+#include "arrow/status.h"
+#include "arrow/util/future.h"
+
+namespace arrow::py {
+/// \brief Bind a Python callback to an arrow::Future.
+///
+/// \param future The future to bind to.
+/// \param py_cb The Python callback function. Will be passed the result of
+///   py_wrapper, or a Python exception if the future failed or one was
+///   raised by py_wrapper.
+/// \param py_wrapper A function (likely defined in Cython) to convert the C++
+///   result of the future to a Python object.
+template <typename T, typename Wrapper = PyObject* (*)(void*)>
+void BindFuture(Future<T> future, PyObject* py_cb, Wrapper py_wrapper) {
+  Py_INCREF(py_cb);
+  OwnedRefNoGIL cb_ref(py_cb);
+
+  auto future_cb = [cb_ref = std::move(cb_ref), py_wrapper](Result<T> result) {
+    SafeCallIntoPythonVoid([&]() {
+      OwnedRef py_value_or_exc{WrapResult(std::move(result), std::move(py_wrapper))};
+      Py_XDECREF(
+          PyObject_CallFunctionObjArgs(cb_ref.obj(), py_value_or_exc.obj(), NULLPTR));
+      ARROW_WARN_NOT_OK(CheckPyError(), "Internal error in async call");
+    });
+  };
+  future.AddCallback(std::move(future_cb));
+}
+}  // namespace arrow::py
diff --git a/python/pyarrow/src/arrow/python/common.h b/python/pyarrow/src/arrow/python/common.h
index 13db6dee1e90c..c8157f0f46383 100644
--- a/python/pyarrow/src/arrow/python/common.h
+++ b/python/pyarrow/src/arrow/python/common.h
@@ -25,7 +25,6 @@
 #include "arrow/python/pyarrow.h"
 #include "arrow/python/visibility.h"
 #include "arrow/result.h"
-#include "arrow/util/future.h"
 #include "arrow/util/macros.h"
 
 namespace arrow {
@@ -288,23 +287,6 @@ std::function<OutFn> BindFunction(Return (*unbound)(PyObject*, Args...),
       [bound_fn](Args... args) { return bound_fn->Invoke(std::forward<Args>(args)...); };
 }
 
-// XXX Put this in arrow/python/async.h to avoid adding more random stuff here?
-template <typename T, typename Wrapper = PyObject* (*)(void*)>
-void BindFuture(Future<T> future, PyObject* py_cb, Wrapper py_wrapper) {
-  Py_INCREF(py_cb);
-  OwnedRefNoGIL cb_ref(py_cb);
-
-  auto future_cb = [cb_ref = std::move(cb_ref), py_wrapper](Result<T> result) {
-    SafeCallIntoPythonVoid([&]() {
-      OwnedRef py_value_or_exc{WrapResult(std::move(result), std::move(py_wrapper))};
-      Py_XDECREF(
-          PyObject_CallFunctionObjArgs(cb_ref.obj(), py_value_or_exc.obj(), NULLPTR));
-      ARROW_WARN_NOT_OK(CheckPyError(), "Internal error in async call");
-    });
-  };
-  future.AddCallback(std::move(future_cb));
-}
-
 // A temporary conversion of a Python object to a bytes area.
 struct PyBytesView {
   const char* bytes;