Fix PythonErrorDetail to store Python error state

(and restore it in check_status())

cpp/src/arrow/python/common.cc

@@ -23,11 +23,15 @@
 
 #include "arrow/memory_pool.h"
 #include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
 #include "arrow/util/logging.h"
 
 #include "arrow/python/helpers.h"
 
 namespace arrow {
+
+using internal::checked_cast;
+
 namespace py {
 
 static std::mutex memory_pool_mutex;
@@ -48,22 +52,128 @@ MemoryPool* get_memory_pool() {
 }
 
 // ----------------------------------------------------------------------
-// PythonErroDetail
+// PythonErrorDetail
+
 namespace {
 
-const char kErrorDetailTypeId[] = "existing exception on python static status.";
-// PythonErrorDetail indicates a python exception was raised and not
-// reset in C++ code (i.e. it should be propagated up through the python
-// stack).
+const char kErrorDetailTypeId[] = "arrow::py::PythonErrorDetail";
+
+// Try to match the Python exception type with an appropriate Status code
+StatusCode MapPyError(PyObject* exc_type) {
+  StatusCode code;
+
+  if (PyErr_GivenExceptionMatches(exc_type, PyExc_MemoryError)) {
+    code = StatusCode::OutOfMemory;
+  } else if (PyErr_GivenExceptionMatches(exc_type, PyExc_IndexError)) {
+    code = StatusCode::IndexError;
+  } else if (PyErr_GivenExceptionMatches(exc_type, PyExc_KeyError)) {
+    code = StatusCode::KeyError;
+  } else if (PyErr_GivenExceptionMatches(exc_type, PyExc_TypeError)) {
+    code = StatusCode::TypeError;
+  } else if (PyErr_GivenExceptionMatches(exc_type, PyExc_ValueError) ||
+             PyErr_GivenExceptionMatches(exc_type, PyExc_OverflowError)) {
+    code = StatusCode::Invalid;
+  } else if (PyErr_GivenExceptionMatches(exc_type, PyExc_EnvironmentError)) {
+    code = StatusCode::IOError;
+  } else if (PyErr_GivenExceptionMatches(exc_type, PyExc_NotImplementedError)) {
+    code = StatusCode::NotImplemented;
+  } else {
+    code = StatusCode::UnknownError;
+  }
+  return code;
+}
+
+// PythonErrorDetail indicates a Python exception was raised.
 class PythonErrorDetail : public StatusDetail {
  public:
-  PythonErrorDetail() = default;
   const char* type_id() const override { return kErrorDetailTypeId; }
-  std::string ToString() const override { return "Uncaught Python."; }
+
+  std::string ToString() const override {
+    // This is simple enough not to need the GIL
+    const auto ty = reinterpret_cast<const PyTypeObject*>(exc_type_.obj());
+    // XXX Should we also print traceback?
+    return std::string("Python exception: ") + ty->tp_name;
+  }
+
+  void RestorePyError() const {
+    Py_INCREF(exc_type_.obj());
+    Py_INCREF(exc_value_.obj());
+    Py_INCREF(exc_traceback_.obj());
+    PyErr_Restore(exc_type_.obj(), exc_value_.obj(), exc_traceback_.obj());
+  }
+
+  PyObject* exc_type() const { return exc_type_.obj(); }
+
+  PyObject* exc_value() const { return exc_value_.obj(); }
+
+  static std::shared_ptr<PythonErrorDetail> FromPyError() {
+    PyObject* exc_type = nullptr;
+    PyObject* exc_value = nullptr;
+    PyObject* exc_traceback = nullptr;
+
+    PyErr_Fetch(&exc_type, &exc_value, &exc_traceback);
+    PyErr_NormalizeException(&exc_type, &exc_value, &exc_traceback);
+    ARROW_CHECK(exc_type)
+        << "PythonErrorDetail::FromPyError called without a Python error set";
+    DCHECK(PyType_Check(exc_type));
+    DCHECK(exc_value);  // Ensured by PyErr_NormalizeException, double-check
+    if (exc_traceback == nullptr) {
+      // Needed by PyErr_Restore()
+      Py_INCREF(Py_None);
+      exc_traceback = Py_None;
+    }
+
+    std::shared_ptr<PythonErrorDetail> detail(new PythonErrorDetail);
+    detail->exc_type_.reset(exc_type);
+    detail->exc_value_.reset(exc_value);
+    detail->exc_traceback_.reset(exc_traceback);
+    return detail;
+  }
+
+ protected:
+  PythonErrorDetail() = default;
+
+  OwnedRefNoGIL exc_type_, exc_value_, exc_traceback_;
 };
 
 }  // namespace
 
+// ----------------------------------------------------------------------
+// Python exception <-> Status
+
+Status ConvertPyError(StatusCode code) {
+  auto detail = PythonErrorDetail::FromPyError();
+  if (code == StatusCode::UnknownError) {
+    code = MapPyError(detail->exc_type());
+  }
+
+  std::string message;
+  RETURN_NOT_OK(internal::PyObject_StdStringStr(detail->exc_value(), &message));
+  return Status(code, message, detail);
+}
+
+Status PassPyError() {
+  if (PyErr_Occurred()) {
+    return ConvertPyError();
+  }
+  return Status::OK();
+}
+
+bool IsPyError(const Status& status) {
+  if (status.ok()) {
+    return false;
+  }
+  auto detail = status.detail();
+  bool result = detail != nullptr && detail->type_id() == kErrorDetailTypeId;
+  return result;
+}
+
+void RestorePyError(const Status& status) {
+  ARROW_CHECK(IsPyError(status));
+  const auto& detail = checked_cast<const PythonErrorDetail&>(*status.detail());
+  detail.RestorePyError();
+}
+
 // ----------------------------------------------------------------------
 // PyBuffer
 
@@ -81,8 +191,7 @@ Status PyBuffer::Init(PyObject* obj) {
     }
     return Status::OK();
   } else {
-    auto detail = std::make_shared<PythonErrorDetail>();
-    return Status(StatusCode::Invalid, "", detail);
+    return ConvertPyError(StatusCode::Invalid);
   }
 }
 
@@ -101,68 +210,5 @@ PyBuffer::~PyBuffer() {
   }
 }
 
-// ----------------------------------------------------------------------
-// Python exception -> Status
-
-Status ConvertPyError(StatusCode code) {
-  PyObject* exc_type = nullptr;
-  PyObject* exc_value = nullptr;
-  PyObject* traceback = nullptr;
-
-  PyErr_Fetch(&exc_type, &exc_value, &traceback);
-  PyErr_NormalizeException(&exc_type, &exc_value, &traceback);
-
-  DCHECK_NE(exc_type, nullptr) << "ConvertPyError called without an exception set";
-
-  OwnedRef exc_type_ref(exc_type);
-  OwnedRef exc_value_ref(exc_value);
-  OwnedRef traceback_ref(traceback);
-
-  std::string message;
-  RETURN_NOT_OK(internal::PyObject_StdStringStr(exc_value, &message));
-
-  if (code == StatusCode::UnknownError) {
-    // Try to match the Python exception type with an appropriate Status code
-    if (PyErr_GivenExceptionMatches(exc_type, PyExc_MemoryError)) {
-      code = StatusCode::OutOfMemory;
-    } else if (PyErr_GivenExceptionMatches(exc_type, PyExc_IndexError)) {
-      code = StatusCode::IndexError;
-    } else if (PyErr_GivenExceptionMatches(exc_type, PyExc_KeyError)) {
-      code = StatusCode::KeyError;
-    } else if (PyErr_GivenExceptionMatches(exc_type, PyExc_TypeError)) {
-      code = StatusCode::TypeError;
-    } else if (PyErr_GivenExceptionMatches(exc_type, PyExc_ValueError) ||
-               PyErr_GivenExceptionMatches(exc_type, PyExc_OverflowError)) {
-      code = StatusCode::Invalid;
-    } else if (PyErr_GivenExceptionMatches(exc_type, PyExc_EnvironmentError)) {
-      code = StatusCode::IOError;
-    } else if (PyErr_GivenExceptionMatches(exc_type, PyExc_NotImplementedError)) {
-      code = StatusCode::NotImplemented;
-    }
-  }
-
-  // We consume the python error here so don't construct with a detail.
-  return Status(code, message);
-}
-
-Status PassPyError() {
-  if (PyErr_Occurred()) {
-    // Do not call PyErr_Clear, the assumption is that someone further
-    // up the call stack will want to deal with the Python error.
-    auto detail = std::make_shared<PythonErrorDetail>();
-    return Status(StatusCode::UnknownError, "", detail);
-  }
-  return Status::OK();
-}
-
-bool IsPythonError(const Status& status) {
-  if (status.ok()) {
-    return false;
-  }
-  auto detail = status.detail();
-  bool result = detail != nullptr && detail->type_id() == kErrorDetailTypeId;
-  return result;
-}
-
 }  // namespace py
 }  // namespace arrow

cpp/src/arrow/python/common.h

@@ -36,7 +36,15 @@ class Result;
 
 namespace py {
 
+// Convert current Python error to a Status.  The Python error state is cleared
+// and can be restored with RestorePyError().
 ARROW_PYTHON_EXPORT Status ConvertPyError(StatusCode code = StatusCode::UnknownError);
+// Same as ConvertPyError(), but returns Status::OK() if no Python error is set.
+ARROW_PYTHON_EXPORT Status PassPyError();
+// Query whether the given Status is a Python error (as wrapped by ConvertPyError()).
+ARROW_PYTHON_EXPORT bool IsPyError(const Status& status);
+// Restore a Python error wrapped in a Status.
+ARROW_PYTHON_EXPORT void RestorePyError(const Status& status);
 
 // Catch a pending Python exception and return the corresponding Status.
 // If no exception is pending, Status::OK() is returned.
@@ -48,10 +56,6 @@ inline Status CheckPyError(StatusCode code = StatusCode::UnknownError) {
   }
 }
 
-ARROW_PYTHON_EXPORT Status PassPyError();
-ARROW_PYTHON_EXPORT bool IsPythonError(const Status& status);
-
-// TODO(wesm): We can just let errors pass through. To be explored later
 #define RETURN_IF_PYERROR() ARROW_RETURN_NOT_OK(CheckPyError());
 
 #define PY_RETURN_IF_ERROR(CODE) ARROW_RETURN_NOT_OK(CheckPyError(CODE));
@@ -98,6 +102,18 @@ class ARROW_PYTHON_EXPORT PyAcquireGIL {
   ARROW_DISALLOW_COPY_AND_ASSIGN(PyAcquireGIL);
 };
 
+// A RAII-style helper that releases the GIL until the end of a lexical block
+class ARROW_PYTHON_EXPORT PyReleaseGIL {
+ public:
+  PyReleaseGIL() { saved_state_ = PyEval_SaveThread(); }
+
+  ~PyReleaseGIL() { PyEval_RestoreThread(saved_state_); }
+
+ private:
+  PyThreadState* saved_state_;
+  ARROW_DISALLOW_COPY_AND_ASSIGN(PyReleaseGIL);
+};
+
 // A helper to call safely into the Python interpreter from arbitrary C++ code.
 // The GIL is acquired, and the current thread's error status is preserved.
 template <typename Function>
@@ -110,7 +126,7 @@ Status SafeCallIntoPython(Function&& func) {
   Status st = std::forward<Function>(func)();
   // If the return Status is a "Python error", the current Python error status
   // describes the error and shouldn't be clobbered.
-  if (!IsPythonError(st) && exc_type != NULLPTR) {
+  if (!IsPyError(st) && exc_type != NULLPTR) {
     PyErr_Restore(exc_type, exc_value, exc_traceback);
   }
   return st;

cpp/src/arrow/python/python-test.cc

@@ -40,22 +40,12 @@ using internal::checked_cast;
 
 namespace py {
 
-TEST(PyBuffer, InvalidInputObject) {
-  std::shared_ptr<Buffer> res;
-  PyObject* input = Py_None;
-  auto old_refcnt = Py_REFCNT(input);
-  Status s = PyBuffer::FromPyObject(input, &res);
-  ASSERT_TRUE(IsPythonError(s)) << s.ToString();
-  PyErr_Clear();
-  ASSERT_EQ(old_refcnt, Py_REFCNT(input));
-}
-
 TEST(OwnedRef, TestMoves) {
-  PyAcquireGIL lock;
   std::vector<OwnedRef> vec;
   PyObject *u, *v;
   u = PyList_New(0);
   v = PyList_New(0);
+
   {
     OwnedRef ref(u);
     vec.push_back(std::move(ref));
@@ -67,31 +57,42 @@ TEST(OwnedRef, TestMoves) {
 }
 
 TEST(OwnedRefNoGIL, TestMoves) {
-  std::vector<OwnedRefNoGIL> vec;
-  PyObject *u, *v;
+  PyAcquireGIL lock;
+  lock.release();
+
   {
-    PyAcquireGIL lock;
-    u = PyList_New(0);
-    v = PyList_New(0);
+    std::vector<OwnedRef> vec;
+    PyObject *u, *v;
+    {
+      lock.acquire();
+      u = PyList_New(0);
+      v = PyList_New(0);
+      lock.release();
+    }
+    {
+      OwnedRefNoGIL ref(u);
+      vec.push_back(std::move(ref));
+      ASSERT_EQ(ref.obj(), nullptr);
+    }
+    vec.emplace_back(v);
+    ASSERT_EQ(Py_REFCNT(u), 1);
+    ASSERT_EQ(Py_REFCNT(v), 1);
   }
-  {
-    OwnedRefNoGIL ref(u);
-    vec.push_back(std::move(ref));
-    ASSERT_EQ(ref.obj(), nullptr);
-  }
-  vec.emplace_back(v);
-  ASSERT_EQ(Py_REFCNT(u), 1);
-  ASSERT_EQ(Py_REFCNT(v), 1);
 }
 
 TEST(CheckPyError, TestStatus) {
-  PyAcquireGIL lock;
   Status st;
 
-  auto check_error = [](Status& st, const char* expected_message = "some error") {
+  auto check_error = [](Status& st, const char* expected_message = "some error",
+                        const char* expected_detail = nullptr) {
     st = CheckPyError();
     ASSERT_EQ(st.message(), expected_message);
     ASSERT_FALSE(PyErr_Occurred());
+    if (expected_detail) {
+      auto detail = st.detail();
+      ASSERT_NE(detail, nullptr);
+      ASSERT_EQ(detail->ToString(), expected_detail);
+    }
   };
 
   for (PyObject* exc_type : {PyExc_Exception, PyExc_SyntaxError}) {
@@ -101,7 +102,7 @@ TEST(CheckPyError, TestStatus) {
   }
 
   PyErr_SetString(PyExc_TypeError, "some error");
-  check_error(st);
+  check_error(st, "some error", "Python exception: TypeError");
   ASSERT_TRUE(st.IsTypeError());
 
   PyErr_SetString(PyExc_ValueError, "some error");
@@ -119,7 +120,7 @@ TEST(CheckPyError, TestStatus) {
   }
 
   PyErr_SetString(PyExc_NotImplementedError, "some error");
-  check_error(st);
+  check_error(st, "some error", "Python exception: NotImplementedError");
   ASSERT_TRUE(st.IsNotImplemented());
 
   // No override if a specific status code is given
@@ -130,6 +131,52 @@ TEST(CheckPyError, TestStatus) {
   ASSERT_FALSE(PyErr_Occurred());
 }
 
+TEST(CheckPyError, TestStatusNoGIL) {
+  PyAcquireGIL lock;
+  {
+    Status st;
+    PyErr_SetString(PyExc_ZeroDivisionError, "zzzt");
+    st = ConvertPyError();
+    ASSERT_FALSE(PyErr_Occurred());
+    lock.release();
+    ASSERT_TRUE(st.IsUnknownError());
+    ASSERT_EQ(st.message(), "zzzt");
+    ASSERT_EQ(st.detail()->ToString(), "Python exception: ZeroDivisionError");
+  }
+}
+
+TEST(RestorePyError, Basics) {
+  PyErr_SetString(PyExc_ZeroDivisionError, "zzzt");
+  auto st = ConvertPyError();
+  ASSERT_FALSE(PyErr_Occurred());
+  ASSERT_TRUE(st.IsUnknownError());
+  ASSERT_EQ(st.message(), "zzzt");
+  ASSERT_EQ(st.detail()->ToString(), "Python exception: ZeroDivisionError");
+
+  RestorePyError(st);
+  ASSERT_TRUE(PyErr_Occurred());
+  PyObject* exc_type;
+  PyObject* exc_value;
+  PyObject* exc_traceback;
+  PyErr_Fetch(&exc_type, &exc_value, &exc_traceback);
+  ASSERT_TRUE(PyErr_GivenExceptionMatches(exc_type, PyExc_ZeroDivisionError));
+  std::string py_message;
+  ASSERT_OK(internal::PyObject_StdStringStr(exc_value, &py_message));
+  ASSERT_EQ(py_message, "zzzt");
+}
+
+TEST(PyBuffer, InvalidInputObject) {
+  std::shared_ptr<Buffer> res;
+  PyObject* input = Py_None;
+  auto old_refcnt = Py_REFCNT(input);
+  {
+    Status st = PyBuffer::FromPyObject(input, &res);
+    ASSERT_TRUE(IsPyError(st)) << st.ToString();
+    ASSERT_FALSE(PyErr_Occurred());
+  }
+  ASSERT_EQ(old_refcnt, Py_REFCNT(input));
+}
+
 class DecimalTest : public ::testing::Test {
  public:
   DecimalTest() : lock_(), decimal_constructor_() {
@@ -254,8 +301,6 @@ TEST(PandasConversionTest, TestObjectBlockWriteFails) {
 }
 
 TEST(BuiltinConversionTest, TestMixedTypeFails) {
-  PyAcquireGIL lock;
-
   OwnedRef list_ref(PyList_New(3));
   PyObject* list = list_ref.obj();
 
@@ -406,8 +451,6 @@ TEST_F(DecimalTest, TestMixedPrecisionAndScale) {
 }
 
 TEST_F(DecimalTest, TestMixedPrecisionAndScaleSequenceConvert) {
-  PyAcquireGIL lock;
-
   PyObject* value1 = this->CreatePythonDecimal("0.01").detach();
   ASSERT_NE(value1, nullptr);
 

cpp/src/arrow/python/serialize.cc

@@ -332,8 +332,8 @@ Status SequenceBuilder::AppendDict(PyObject* context, PyObject* dict,
 
 Status CallCustomCallback(PyObject* context, PyObject* method_name, PyObject* elem,
                           PyObject** result) {
-  *result = NULL;
   if (context == Py_None) {
+    *result = NULL;
     return Status::SerializationError("error while calling callback on ",
                                       internal::PyObject_StdStringRepr(elem),
                                       ": handler not registered");

python/pyarrow/error.pxi

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-from pyarrow.includes.libarrow cimport CStatus, IsPythonError
+from pyarrow.includes.libarrow cimport CStatus, IsPyError, RestorePyError
 from pyarrow.includes.common cimport c_string
 from pyarrow.compat import frombytes
 
@@ -60,17 +60,17 @@ class ArrowSerializationError(ArrowException):
     pass
 
 
+# This function could be written directly in C++ if we didn't
+# define Arrow-specific subclasses (ArrowInvalid etc.)
 cdef int check_status(const CStatus& status) nogil except -1:
     if status.ok():
         return 0
 
-    # PythonError indicates there is already an exception raised
-    # so even though the cython documentation says not to ever return
-    # -1 here, we do so anyways.
-    if IsPythonError(status):
-        return -1
-
     with gil:
+        if IsPyError(status):
+            RestorePyError(status)
+            return -1
+
         message = frombytes(status.message())
         if status.IsInvalid():
             raise ArrowInvalid(message)

python/pyarrow/includes/libarrow.pxd

@@ -1207,7 +1207,6 @@ cdef extern from "arrow/python/api.h" namespace "arrow::py" nogil:
     shared_ptr[CDataType] GetPrimitiveType(Type type)
     shared_ptr[CDataType] GetTimestampType(TimeUnit unit)
 
-    c_bool IsPythonError(const CStatus& status)
     object PyHalf_FromHalf(npy_half value)
 
     cdef cppclass PyConversionOptions:
@@ -1218,6 +1217,8 @@ cdef extern from "arrow/python/api.h" namespace "arrow::py" nogil:
         CMemoryPool* pool
         c_bool from_pandas
 
+    # TODO Some functions below are not actually "nogil"
+
     CStatus ConvertPySequence(object obj, object mask,
                               const PyConversionOptions& options,
                               shared_ptr[CChunkedArray]* out)
@@ -1343,6 +1344,11 @@ cdef extern from 'arrow/python/init.h':
     int arrow_init_numpy() except -1
 
 
+cdef extern from 'arrow/python/common.h' namespace "arrow::py":
+    c_bool IsPyError(const CStatus& status)
+    void RestorePyError(const CStatus& status)
+
+
 cdef extern from 'arrow/python/pyarrow.h' namespace 'arrow::py':
     int import_pyarrow() except -1
 

python/pyarrow/tests/test_array.py

@@ -1479,7 +1479,7 @@ def test_array_masked():
 
 def test_array_from_large_pyints():
     # ARROW-5430
-    with pytest.raises(pa.ArrowInvalid):
+    with pytest.raises(OverflowError):
         # too large for int64 so dtype must be explicitly provided
         pa.array([int(2 ** 63)])
 

python/pyarrow/tests/test_convert_builtin.py

@@ -191,7 +191,7 @@ def test_nested_lists(seq):
 @parametrize_with_iterable_types
 def test_list_with_non_list(seq):
     # List types don't accept non-sequences
-    with pytest.raises(pa.ArrowTypeError):
+    with pytest.raises(TypeError):
         pa.array(seq([[], [1, 2], 3]), type=pa.list_(pa.int64()))
 
 
@@ -308,7 +308,7 @@ def test_numpy_scalars_mixed_type():
 
 
 @pytest.mark.xfail(reason="Type inference for uint64 not implemented",
-                   raises=pa.ArrowException)
+                   raises=OverflowError)
 def test_uint64_max_convert():
     data = [0, np.iinfo(np.uint64).max]
 
@@ -323,20 +323,20 @@ def test_uint64_max_convert():
 @pytest.mark.parametrize("bits", [8, 16, 32, 64])
 def test_signed_integer_overflow(bits):
     ty = getattr(pa, "int%d" % bits)()
-    # XXX ideally would raise OverflowError
-    with pytest.raises((ValueError, pa.ArrowException)):
+    # XXX ideally would always raise OverflowError
+    with pytest.raises((OverflowError, pa.ArrowInvalid)):
         pa.array([2 ** (bits - 1)], ty)
-    with pytest.raises((ValueError, pa.ArrowException)):
+    with pytest.raises((OverflowError, pa.ArrowInvalid)):
         pa.array([-2 ** (bits - 1) - 1], ty)
 
 
 @pytest.mark.parametrize("bits", [8, 16, 32, 64])
 def test_unsigned_integer_overflow(bits):
     ty = getattr(pa, "uint%d" % bits)()
-    # XXX ideally would raise OverflowError
-    with pytest.raises((ValueError, pa.ArrowException)):
+    # XXX ideally would always raise OverflowError
+    with pytest.raises((OverflowError, pa.ArrowInvalid)):
         pa.array([2 ** bits], ty)
-    with pytest.raises((ValueError, pa.ArrowException)):
+    with pytest.raises((OverflowError, pa.ArrowInvalid)):
         pa.array([-1], ty)
 
 
@@ -831,7 +831,7 @@ def test_sequence_timestamp_from_int_with_unit():
     assert repr(arr_ns[0]) == "Timestamp('1970-01-01 00:00:00.000000001')"
     assert str(arr_ns[0]) == "1970-01-01 00:00:00.000000001"
 
-    with pytest.raises(pa.ArrowException):
+    with pytest.raises(TypeError):
         class CustomClass():
             pass
         pa.array([1, CustomClass()], type=ns)