Handle exceptions on RowGroupWriter::Close better

cpp/src/arrow/parquet/utils.h

@@ -31,6 +31,11 @@ namespace parquet {
     (s);                        \
   } catch (const ::parquet::ParquetException& e) { return Status::Invalid(e.what()); }
 
+#define PARQUET_IGNORE_NOT_OK(s) \
+  try {                         \
+    (s);                        \
+  } catch (const ::parquet::ParquetException& e) { }
+
 }  // namespace parquet
 
 }  // namespace arrow

cpp/src/arrow/parquet/writer.cc

@@ -181,6 +181,7 @@ Status WriteFlatTable(const Table* table, MemoryPool* pool,
     std::shared_ptr<Array> array = table->column(i)->data()->chunk(0);
     auto primitive_array = std::dynamic_pointer_cast<PrimitiveArray>(array);
     if (!primitive_array) {
+      PARQUET_IGNORE_NOT_OK(writer.Close());
       return Status::NotImplemented("Table must consist of PrimitiveArray instances");
     }
     arrays[i] = primitive_array;
@@ -189,9 +190,9 @@ Status WriteFlatTable(const Table* table, MemoryPool* pool,
   for (int chunk = 0; chunk * chunk_size < table->num_rows(); chunk++) {
     int64_t offset = chunk * chunk_size;
     int64_t size = std::min(chunk_size, table->num_rows() - offset);
-    RETURN_NOT_OK(writer.NewRowGroup(size));
+    RETURN_NOT_OK_ELSE(writer.NewRowGroup(size), PARQUET_IGNORE_NOT_OK(writer.Close()));
     for (int i = 0; i < table->num_columns(); i++) {
-      RETURN_NOT_OK(writer.WriteFlatColumnChunk(arrays[i].get(), offset, size));
+      RETURN_NOT_OK_ELSE(writer.WriteFlatColumnChunk(arrays[i].get(), offset, size), PARQUET_IGNORE_NOT_OK(writer.Close()));
     }
   }
 

cpp/src/arrow/util/status.h

@@ -63,6 +63,15 @@ namespace arrow {
     if (!_s.ok()) { return _s; } \
   } while (0);
 
+#define RETURN_NOT_OK_ELSE(s, else_) \
+  do {                               \
+    Status _s = (s);                 \
+    if (!_s.ok()) {                  \
+      else_;                         \
+      return _s;                     \
+    }                                \
+  } while (0);
+
 enum class StatusCode : char {
   OK = 0,
   OutOfMemory = 1,

python/pyarrow/tests/test_parquet.py

@@ -21,12 +21,10 @@ import pyarrow.parquet
 
 A = arrow
 
-from shutil import rmtree
-from tempfile import mkdtemp
-
 import numpy as np
 import os.path
 import pandas as pd
+
 import pandas.util.testing as pdt