[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 · Feb 22, 2023 · fad8d66 · fad8d66
1 parent 886cb56
commit fad8d66
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/velox/exec/ArrowStream.cpp b/velox/exec/ArrowStream.cpp
@@ -34,9 +34,16 @@ ArrowStream::~ArrowStream() {
   close();
 }
 
+ArrowStream::~ArrowStream() {
+  if (!isFinished0()) {
+    close0();
+  }
+}
+
 RowVectorPtr ArrowStream::getOutput() {
   // Get Arrow array.
   struct ArrowArray arrowArray;
+
   if (arrowStream_->get_next(arrowStream_.get(), &arrowArray)) {
     if (arrowArray.release) {
       arrowArray.release(&arrowArray);
@@ -52,6 +59,7 @@ RowVectorPtr ArrowStream::getOutput() {
 
   // Get Arrow schema.
   struct ArrowSchema arrowSchema;
+
   if (arrowStream_->get_schema(arrowStream_.get(), &arrowSchema)) {
     if (arrowSchema.release) {
       arrowSchema.release(&arrowSchema);

diff --git a/velox/exec/ArrowStream.h b/velox/exec/ArrowStream.h
@@ -29,6 +29,8 @@ class ArrowStream : public SourceOperator {
 
   virtual ~ArrowStream();
 
+  virtual ~ArrowStream();
+
   RowVectorPtr getOutput() override;
 
   BlockingReason isBlocked(ContinueFuture* /* unused */) override {
@@ -37,6 +39,7 @@ class ArrowStream : public SourceOperator {
 
   bool isFinished() override;
 
+  bool isFinished() override;
   void close() override;
 
  private:
@@ -45,6 +48,10 @@ class ArrowStream : public SourceOperator {
 
   bool finished_ = false;
   std::shared_ptr<ArrowArrayStream> arrowStream_;
+
+  // For calls from destructor
+  bool isFinished0();
+  void close0();
 };
 
 } // namespace facebook::velox::exec