Use a separate bool flag to track reader state and move reset() from …

…GetSchema to Release to fix the segfault
apache · Aug 7, 2023 · f158e89 · f158e89
1 parent a87481f
commit f158e89
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/c/driver/postgresql/statement.cc b/c/driver/postgresql/statement.cc
@@ -634,7 +634,7 @@ void TupleReader::ResetQuery() {
 }
 
 int TupleReader::GetNext(struct ArrowArray* out) {
-  if (!copy_reader_) {
+  if (is_finished_) {
     out->release = nullptr;
     return 0;
   }
@@ -668,9 +668,6 @@ int TupleReader::GetNext(struct ArrowArray* out) {
   struct ArrowArray tmp;
   NANOARROW_RETURN_NOT_OK(BuildOutput(&tmp, &error));
 
-  // Clear the copy reader to mark this reader as finished
-  copy_reader_.reset();
-
   // Check the server-side response
   result_ = PQgetResult(conn_);
   const int pq_status = PQresultStatus(result_);
@@ -702,6 +699,12 @@ void TupleReader::Release() {
     PQfreemem(pgbuf_);
     pgbuf_ = nullptr;
   }
+
+  // Clear the copy reader to mark this reader as finished
+  if (copy_reader_) {
+    copy_reader_.reset();
+    is_finished_ = true;
+  }
 }
 
 void TupleReader::ExportTo(struct ArrowArrayStream* stream) {

diff --git a/c/driver/postgresql/statement.h b/c/driver/postgresql/statement.h
@@ -50,6 +50,7 @@ class TupleReader final {
     StringBuilderInit(&error_builder_, 0);
     data_.data.as_char = nullptr;
     data_.size_bytes = 0;
+    is_finished_ = false;
   }
 
   int GetSchema(struct ArrowSchema* out);
@@ -85,6 +86,7 @@ class TupleReader final {
   std::unique_ptr<PostgresCopyStreamReader> copy_reader_;
   int64_t row_id_;
   int64_t batch_size_hint_bytes_;
+  bool is_finished_;
 };
 
 class PostgresStatement {