paleolimbot commented on code in PR #870:
URL: https://github.com/apache/arrow-adbc/pull/870#discussion_r1254489446
##########
c/driver/postgresql/statement.cc:
##########
@@ -438,91 +438,145 @@ int TupleReader::GetSchema(struct ArrowSchema* out) {
return na_res;
}
-int TupleReader::GetNext(struct ArrowArray* out) {
- if (!result_) {
- out->release = nullptr;
- return 0;
+int TupleReader::InitQuery(struct ArrowError* error) {
+ ResetQuery();
+
+ // Fetch + parse the header
+ int get_copy_res = PQgetCopyData(conn_, &pgbuf_, /*async=*/0);
+ data_.size_bytes = get_copy_res;
+ data_.data.as_char = pgbuf_;
+
+ if (get_copy_res == -2) {
+ StringBuilderAppend(&error_builder_, "[libpq] Fetch header failed: %s",
+ PQerrorMessage(conn_));
+ return EIO;
}
- // Clear the result, since the data is actually read from the connection
- PQclear(result_);
- result_ = nullptr;
+ int na_res = copy_reader_->ReadHeader(&data_, error);
+ if (na_res != NANOARROW_OK) {
+ StringBuilderAppend(&error_builder_, "[libpq] ReadHeader failed: %s",
error->message);
+ return EIO;
+ }
- // Clear the error builder
- error_builder_.size = 0;
+ return NANOARROW_OK;
+}
- struct ArrowError error;
- error.message[0] = '\0';
- struct ArrowBufferView data;
- data.data.data = nullptr;
- data.size_bytes = 0;
+int TupleReader::AppendRow(struct ArrowError* error) {
+ // Parse the result (the header AND the first row are included in the first
+ // call to PQgetCopyData())
+ int na_res = copy_reader_->ReadRecord(&data_, error);
+ if (na_res != NANOARROW_OK && na_res != ENODATA) {
+ StringBuilderAppend(&error_builder_, "[libpq] ReadRecord failed at row
%ld: %s",
+ static_cast<long>(row_id_), // NOLINT(runtime/int)
+ error->message);
+ return na_res;
+ }
- // Fetch + parse the header
+ row_id_++;
+
+ // Fetch + check
+ PQfreemem(pgbuf_);
+ pgbuf_ = nullptr;
int get_copy_res = PQgetCopyData(conn_, &pgbuf_, /*async=*/0);
+ data_.size_bytes = get_copy_res;
+ data_.data.as_char = pgbuf_;
+
if (get_copy_res == -2) {
- StringBuilderAppend(&error_builder_, "[libpq] Fetch header failed: %s",
+ StringBuilderAppend(&error_builder_, "[libpq] Fetch row %ld failed: %s",
+ static_cast<long>(row_id_), // NOLINT(runtime/int)
PQerrorMessage(conn_));
return EIO;
+ } else if (get_copy_res == -1) {
+ // Returned when COPY has finished successfully
+ return ENODATA;
+ } else if ((copy_reader_->array_size_approx_bytes() + get_copy_res) >=
+ array_size_hint_bytes_) {
+ // Appending the next row will result in an array larger than requested.
+ // Return EOVERFLOW to force GetNext() to build the current result and
return.
+ return EOVERFLOW;
Review Comment:
Ok, this is now tested!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
