(draft) Try fixing the segfaults without the bool flag

apache · Aug 9, 2023 · c50cc3a · c50cc3a
1 parent 5e7289d
commit c50cc3a
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 19 deletions.
diff --git a/c/driver/postgresql/statement.cc b/c/driver/postgresql/statement.cc
@@ -611,18 +611,20 @@ int TupleReader::BuildOutput(struct ArrowArray* out, struct ArrowError* error) {
   return NANOARROW_OK;
 }
 
-void TupleReader::ResetQuery() {
-  // Clear result
+void TupleReader::FreePQ() {
   if (result_) {
     PQclear(result_);
     result_ = nullptr;
   }
 
-  // Reset result buffer
-  if (pgbuf_ != nullptr) {
+  if (pgbuf_) {
     PQfreemem(pgbuf_);
     pgbuf_ = nullptr;
   }
+}
+
+void TupleReader::ResetQuery() {
+  FreePQ();
 
   // Clear the error builder
   error_builder_.size = 0;
@@ -631,7 +633,7 @@ void TupleReader::ResetQuery() {
 }
 
 int TupleReader::GetNext(struct ArrowArray* out) {
-  if (is_finished_) {
+  if (!copy_reader_) {
     out->release = nullptr;
     return 0;
   }
@@ -679,25 +681,15 @@ int TupleReader::GetNext(struct ArrowArray* out) {
     return EIO;
   }
 
-  is_finished_ = true;
   ResetQuery();
   ArrowArrayMove(&tmp, out);
   return NANOARROW_OK;
 }
 
 void TupleReader::Release() {
+  FreePQ();
   StringBuilderReset(&error_builder_);
 
-  if (result_) {
-    PQclear(result_);
-    result_ = nullptr;
-  }
-
-  if (pgbuf_) {
-    PQfreemem(pgbuf_);
-    pgbuf_ = nullptr;
-  }
-
   if (copy_reader_) {
     copy_reader_.reset();
   }

diff --git a/c/driver/postgresql/statement.h b/c/driver/postgresql/statement.h
@@ -46,8 +46,7 @@ class TupleReader final {
         pgbuf_(nullptr),
         copy_reader_(nullptr),
         row_id_(-1),
-        batch_size_hint_bytes_(16777216),
-        is_finished_(false) {
+        batch_size_hint_bytes_(16777216) {
     StringBuilderInit(&error_builder_, 0);
     data_.data.as_char = nullptr;
     data_.size_bytes = 0;
@@ -71,6 +70,7 @@ class TupleReader final {
   int InitQueryAndFetchFirst(struct ArrowError* error);
   int AppendRowAndFetchNext(struct ArrowError* error);
   int BuildOutput(struct ArrowArray* out, struct ArrowError* error);
+  void FreePQ();
   void ResetQuery();
 
   static int GetSchemaTrampoline(struct ArrowArrayStream* self, struct ArrowSchema* out);
@@ -86,7 +86,6 @@ class TupleReader final {
   std::unique_ptr<PostgresCopyStreamReader> copy_reader_;
   int64_t row_id_;
   int64_t batch_size_hint_bytes_;
-  bool is_finished_;
 };
 
 class PostgresStatement {