[OPPRO-135] ArrowStream::getOutput swallows errors, and destructor do…

…esn't release ArrowArrayStream (facebookincubator#16) * ArrowStream::getOutput swallows errors * Also, call release hook after ArrowArrayStream is used
rui-mo · Oct 20, 2022 · ebf45e1 · ebf45e1
1 parent 1c4c0f4
commit ebf45e1
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 5 deletions.
diff --git a/velox/exec/ArrowStream.cpp b/velox/exec/ArrowStream.cpp
@@ -30,27 +30,53 @@ ArrowStream::ArrowStream(
   arrowStream_ = arrowStream->arrowStream();
 }
 
+ArrowStream::~ArrowStream() {
+  if (!isFinished0()) {
+    close0();
+  }
+}
+
 RowVectorPtr ArrowStream::getOutput() {
   struct ArrowArray arrowArray;
-  arrowStream_->get_next(&(*arrowStream_), &arrowArray);
+  if (arrowStream_->get_next(&(*arrowStream_), &arrowArray)) {
+    VELOX_FAIL(
+        "Failed to call get_next on ArrowStream: " + std::string(GetError()))
+  }
   if (arrowArray.release == NULL) {
     // End of Stream.
     closed_ = true;
     return nullptr;
   }
   struct ArrowSchema arrowSchema;
-  arrowStream_->get_schema(&(*arrowStream_), &arrowSchema);
+  if (arrowStream_->get_schema(&(*arrowStream_), &arrowSchema)) {
+    VELOX_FAIL(
+        "Failed to call get_schema on ArrowStream: " + std::string(GetError()))
+  }
   // Convert Arrow data into RowVector.
   rowVector_ = std::dynamic_pointer_cast<RowVector>(
       facebook::velox::importFromArrowAsViewer(arrowSchema, arrowArray));
   return rowVector_;
 }
 
+const char* ArrowStream::GetError() {
+  return arrowStream_->get_last_error(arrowStream_.get());
+}
+
 void ArrowStream::close() {
-  closed_ = true;
+  close0();
+  SourceOperator::close();
 }
 
 bool ArrowStream::isFinished() {
+  return isFinished0();
+}
+
+void ArrowStream::close0() {
+  arrowStream_->release(arrowStream_.get());
+  closed_ = true;
+}
+
+bool ArrowStream::isFinished0() {
   return closed_;
 }
 

diff --git a/velox/exec/ArrowStream.h b/velox/exec/ArrowStream.h
@@ -27,25 +27,31 @@ class ArrowStream : public SourceOperator {
       DriverCtx* driverCtx,
       std::shared_ptr<const core::ArrowStreamNode> arrowStream);
 
+  virtual ~ArrowStream();
+
   RowVectorPtr getOutput() override;
 
   BlockingReason isBlocked(ContinueFuture* /* unused */) override {
     return BlockingReason::kNotBlocked;
   }
-
   void noMoreInput() override {
     Operator::noMoreInput();
     close();
   }
 
-  bool isFinished() override;
+  const char* GetError();
 
+  bool isFinished() override;
   void close() override;
 
  private:
   bool closed_ = false;
   RowVectorPtr rowVector_;
   std::shared_ptr<ArrowArrayStream> arrowStream_;
+
+  // For calls from destructor
+  bool isFinished0();
+  void close0();
 };
 
 } // namespace facebook::velox::exec