You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by li...@apache.org on 2023/05/04 17:10:53 UTC
[arrow-adbc] branch main updated: refactor(c/driver/postgresql): Factor out COPY reader and test it independently (#636)
This is an automated email from the ASF dual-hosted git repository.
lidavidm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-adbc.git
The following commit(s) were added to refs/heads/main by this push:
new 3af5665 refactor(c/driver/postgresql): Factor out COPY reader and test it independently (#636)
3af5665 is described below
commit 3af5665bf11fd33b45393a4fd99d00bdfe9d2791
Author: Dewey Dunnington <de...@voltrondata.com>
AuthorDate: Thu May 4 13:10:47 2023 -0400
refactor(c/driver/postgresql): Factor out COPY reader and test it independently (#636)
I think there are some holes here but I think it's ready for a round of
review!
I'm aware that this may be slower than the current approach (notably:
virtual method calls and more assertions about buffer sizes). At this
point I think that type coverage and safety are worth it, and at least
this approach allows one to write benchmarks to determine if this is, in
fact, a bottleneck.
A few possible follow-ups:
- Add support for more types. Date/time/timestamp are highest on my
personal list.
- Do a similar refactoring for generating COPY data. In particular, bulk
inserts are very slow with the current approach and we could probably
speed up the bulk ingest considerably using COPY there, too.
I think the 374 exit code on Windows is not related to this PR since it
also occurs in the sqlite driver.
---
c/driver/postgresql/CMakeLists.txt | 1 +
c/driver/postgresql/database.cc | 2 +-
c/driver/postgresql/postgres_copy_reader.h | 772 +++++++++++++++++++++++
c/driver/postgresql/postgres_copy_reader_test.cc | 501 +++++++++++++++
c/driver/postgresql/postgres_type.h | 23 +-
c/driver/postgresql/postgres_type_test.cc | 231 ++++---
c/driver/postgresql/postgres_util.h | 2 +
c/driver/postgresql/statement.cc | 305 +++------
c/driver/postgresql/statement.h | 24 +-
r/adbcpostgresql/bootstrap.R | 1 +
r/adbcpostgresql/src/.gitignore | 5 +-
11 files changed, 1523 insertions(+), 344 deletions(-)
diff --git a/c/driver/postgresql/CMakeLists.txt b/c/driver/postgresql/CMakeLists.txt
index 8cf6f28..337873f 100644
--- a/c/driver/postgresql/CMakeLists.txt
+++ b/c/driver/postgresql/CMakeLists.txt
@@ -73,6 +73,7 @@ if(ADBC_BUILD_TESTS)
driver-postgresql
SOURCES
postgres_type_test.cc
+ postgres_copy_reader_test.cc
postgresql_test.cc
../../validation/adbc_validation.cc
../../validation/adbc_validation_util.cc
diff --git a/c/driver/postgresql/database.cc b/c/driver/postgresql/database.cc
index 92accae..1d9bd41 100644
--- a/c/driver/postgresql/database.cc
+++ b/c/driver/postgresql/database.cc
@@ -251,7 +251,7 @@ static inline int32_t InsertPgTypeResult(
// If there's an array type and the insert succeeded, add that now too
if (result == NANOARROW_OK && typarray != 0) {
- std::string array_typname = "_" + *typname;
+ std::string array_typname = "_" + std::string(typname);
item.oid = typarray;
item.typname = array_typname.c_str();
item.typreceive = "array_recv";
diff --git a/c/driver/postgresql/postgres_copy_reader.h b/c/driver/postgresql/postgres_copy_reader.h
new file mode 100644
index 0000000..c92e8fa
--- /dev/null
+++ b/c/driver/postgresql/postgres_copy_reader.h
@@ -0,0 +1,772 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <cerrno>
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include <nanoarrow/nanoarrow.hpp>
+
+#include "postgres_type.h"
+#include "postgres_util.h"
+
+namespace adbcpq {
+
+static int8_t kPgCopyBinarySignature[] = {'P', 'G', 'C', 'O', 'P', 'Y',
+ '\n', '\377', '\r', '\n', '\0'};
+
+// Read a value from the buffer without checking the buffer size. Advances
+// the cursor of data and reduces its size by sizeof(T).
+template <typename T>
+inline T ReadUnsafe(ArrowBufferView* data) {
+ T out;
+ memcpy(&out, data->data.data, sizeof(T));
+ out = SwapNetworkToHost(out);
+ data->data.as_uint8 += sizeof(T);
+ data->size_bytes -= sizeof(T);
+ return out;
+}
+
+// Define some explicit specializations for types that don't have a SwapNetworkToHost
+// overload.
+template <>
+inline int8_t ReadUnsafe(ArrowBufferView* data) {
+ int8_t out = data->data.as_int8[0];
+ data->data.as_uint8 += sizeof(int8_t);
+ data->size_bytes -= sizeof(int8_t);
+ return out;
+}
+
+template <>
+inline int16_t ReadUnsafe(ArrowBufferView* data) {
+ return static_cast<int16_t>(ReadUnsafe<uint16_t>(data));
+}
+
+template <>
+inline int32_t ReadUnsafe(ArrowBufferView* data) {
+ return static_cast<int32_t>(ReadUnsafe<uint32_t>(data));
+}
+
+template <>
+inline int64_t ReadUnsafe(ArrowBufferView* data) {
+ return static_cast<int64_t>(ReadUnsafe<uint64_t>(data));
+}
+
+template <typename T>
+ArrowErrorCode ReadChecked(ArrowBufferView* data, T* out, ArrowError* error) {
+ if (data->size_bytes < static_cast<int64_t>(sizeof(T))) {
+ ArrowErrorSet(error, "Unexpected end of input (expected %d bytes but found %ld)",
+ static_cast<int>(sizeof(T)),
+ static_cast<long>(data->size_bytes)); // NOLINT(runtime/int)
+ return EINVAL;
+ }
+
+ *out = ReadUnsafe<T>(data);
+ return NANOARROW_OK;
+}
+
+class PostgresCopyFieldReader {
+ public:
+ PostgresCopyFieldReader() : validity_(nullptr), offsets_(nullptr), data_(nullptr) {
+ memset(&schema_view_, 0, sizeof(ArrowSchemaView));
+ }
+
+ virtual ~PostgresCopyFieldReader() {}
+
+ void Init(const PostgresType& pg_type) { pg_type_ = pg_type; }
+
+ const PostgresType& InputType() const { return pg_type_; }
+
+ virtual ArrowErrorCode InitSchema(ArrowSchema* schema) {
+ NANOARROW_RETURN_NOT_OK(ArrowSchemaViewInit(&schema_view_, schema, nullptr));
+ return NANOARROW_OK;
+ }
+
+ virtual ArrowErrorCode InitArray(ArrowArray* array) {
+ // Cache some buffer pointers
+ validity_ = ArrowArrayValidityBitmap(array);
+ for (int32_t i = 0; i < 3; i++) {
+ switch (schema_view_.layout.buffer_type[i]) {
+ case NANOARROW_BUFFER_TYPE_DATA_OFFSET:
+ if (schema_view_.layout.element_size_bits[i] == 32) {
+ offsets_ = ArrowArrayBuffer(array, i);
+ }
+ break;
+ case NANOARROW_BUFFER_TYPE_DATA:
+ data_ = ArrowArrayBuffer(array, i);
+ break;
+ default:
+ break;
+ }
+ }
+
+ return NANOARROW_OK;
+ }
+
+ virtual ArrowErrorCode Read(ArrowBufferView* data, int32_t field_size_bytes,
+ ArrowArray* array, ArrowError* error) {
+ return ENOTSUP;
+ }
+
+ virtual ArrowErrorCode FinishArray(ArrowArray* array, ArrowError* error) {
+ return NANOARROW_OK;
+ }
+
+ protected:
+ PostgresType pg_type_;
+ ArrowSchemaView schema_view_;
+ ArrowBitmap* validity_;
+ ArrowBuffer* offsets_;
+ ArrowBuffer* data_;
+ std::vector<std::unique_ptr<PostgresCopyFieldReader>> children_;
+
+ ArrowErrorCode AppendValid(ArrowArray* array) {
+ if (validity_->buffer.data != nullptr) {
+ NANOARROW_RETURN_NOT_OK(ArrowBitmapAppend(validity_, true, 1));
+ }
+
+ array->length++;
+ return NANOARROW_OK;
+ }
+};
+
+// Reader for a Postgres boolean (one byte -> bitmap)
+class PostgresCopyBooleanFieldReader : public PostgresCopyFieldReader {
+ public:
+ ArrowErrorCode Read(ArrowBufferView* data, int32_t field_size_bytes, ArrowArray* array,
+ ArrowError* error) override {
+ if (field_size_bytes <= 0) {
+ return ArrowArrayAppendNull(array, 1);
+ }
+
+ if (field_size_bytes != 1) {
+ ArrowErrorSet(error, "Expected field with one byte but found field with %d bytes",
+ static_cast<int>(field_size_bytes)); // NOLINT(runtime/int)
+ return EINVAL;
+ }
+
+ int64_t bytes_required = _ArrowBytesForBits(array->length + 1);
+ if (bytes_required > data_->size_bytes) {
+ NANOARROW_RETURN_NOT_OK(
+ ArrowBufferAppendFill(data_, 0, bytes_required - data_->size_bytes));
+ }
+
+ if (ReadUnsafe<int8_t>(data)) {
+ ArrowBitSet(data_->data, array->length);
+ } else {
+ ArrowBitClear(data_->data, array->length);
+ }
+
+ return AppendValid(array);
+ }
+};
+
+// Reader for Pg->Arrow conversions whose representations are identical minus
+// the bswap from network endian. This includes all integral and float types.
+template <typename T>
+class PostgresCopyNetworkEndianFieldReader : public PostgresCopyFieldReader {
+ public:
+ ArrowErrorCode Read(ArrowBufferView* data, int32_t field_size_bytes, ArrowArray* array,
+ ArrowError* error) override {
+ if (field_size_bytes <= 0) {
+ return ArrowArrayAppendNull(array, 1);
+ }
+
+ if (field_size_bytes != static_cast<int32_t>(sizeof(T))) {
+ ArrowErrorSet(error, "Expected field with %d bytes but found field with %d bytes",
+ static_cast<int>(sizeof(T)),
+ static_cast<int>(field_size_bytes)); // NOLINT(runtime/int)
+ return EINVAL;
+ }
+
+ T value = ReadUnsafe<T>(data);
+ NANOARROW_RETURN_NOT_OK(ArrowBufferAppend(data_, &value, sizeof(T)));
+ return AppendValid(array);
+ }
+};
+
+// Reader for Pg->Arrow conversions whose Arrow representation is simply the
+// bytes of the field representation. This can be used with binary and string
+// Arrow types and any Postgres type.
+class PostgresCopyBinaryFieldReader : public PostgresCopyFieldReader {
+ public:
+ ArrowErrorCode Read(ArrowBufferView* data, int32_t field_size_bytes, ArrowArray* array,
+ ArrowError* error) override {
+ // -1 for NULL (0 would be empty string)
+ if (field_size_bytes < 0) {
+ return ArrowArrayAppendNull(array, 1);
+ }
+
+ if (field_size_bytes > data->size_bytes) {
+ ArrowErrorSet(error, "Expected %d bytes of field data but got %d bytes of input",
+ static_cast<int>(field_size_bytes),
+ static_cast<int>(data->size_bytes)); // NOLINT(runtime/int)
+ return EINVAL;
+ }
+
+ NANOARROW_RETURN_NOT_OK(ArrowBufferAppend(data_, data->data.data, field_size_bytes));
+ data->data.as_uint8 += field_size_bytes;
+ data->size_bytes -= field_size_bytes;
+
+ int32_t* offsets = reinterpret_cast<int32_t*>(offsets_->data);
+ NANOARROW_RETURN_NOT_OK(
+ ArrowBufferAppendInt32(offsets_, offsets[array->length] + field_size_bytes));
+
+ return AppendValid(array);
+ }
+};
+
+class PostgresCopyArrayFieldReader : public PostgresCopyFieldReader {
+ public:
+ void InitChild(std::unique_ptr<PostgresCopyFieldReader> child) {
+ child_ = std::move(child);
+ child_->Init(pg_type_.child(0));
+ }
+
+ ArrowErrorCode InitSchema(ArrowSchema* schema) override {
+ NANOARROW_RETURN_NOT_OK(PostgresCopyFieldReader::InitSchema(schema));
+ NANOARROW_RETURN_NOT_OK(child_->InitSchema(schema->children[0]));
+ return NANOARROW_OK;
+ }
+
+ ArrowErrorCode InitArray(ArrowArray* array) override {
+ NANOARROW_RETURN_NOT_OK(PostgresCopyFieldReader::InitArray(array));
+ NANOARROW_RETURN_NOT_OK(child_->InitArray(array->children[0]));
+ return NANOARROW_OK;
+ }
+
+ ArrowErrorCode Read(ArrowBufferView* data, int32_t field_size_bytes, ArrowArray* array,
+ ArrowError* error) override {
+ if (field_size_bytes <= 0) {
+ return ArrowArrayAppendNull(array, 1);
+ }
+
+ // Keep the cursor where we start to parse the array so we can check
+ // the number of bytes read against the field size when finished
+ const uint8_t* data0 = data->data.as_uint8;
+
+ int32_t n_dim;
+ NANOARROW_RETURN_NOT_OK(ReadChecked<int32_t>(data, &n_dim, error));
+ int32_t flags;
+ NANOARROW_RETURN_NOT_OK(ReadChecked<int32_t>(data, &flags, error));
+ uint32_t element_type_oid;
+ NANOARROW_RETURN_NOT_OK(ReadChecked<uint32_t>(data, &element_type_oid, error));
+
+ // We could validate the OID here, but this is a poor fit for all cases
+ // (e.g. testing) since the OID can be specific to each database
+
+ if (n_dim < 0) {
+ ArrowErrorSet(error, "Expected array n_dim > 0 but got %d",
+ static_cast<int>(n_dim)); // NOLINT(runtime/int)
+ return EINVAL;
+ }
+
+ // This is apparently allowed
+ if (n_dim == 0) {
+ NANOARROW_RETURN_NOT_OK(ArrowArrayFinishElement(array));
+ return NANOARROW_OK;
+ }
+
+ int64_t n_items = 1;
+ for (int32_t i = 0; i < n_dim; i++) {
+ int32_t dim_size;
+ NANOARROW_RETURN_NOT_OK(ReadChecked<int32_t>(data, &dim_size, error));
+ n_items *= dim_size;
+
+ int32_t lower_bound;
+ NANOARROW_RETURN_NOT_OK(ReadChecked<int32_t>(data, &lower_bound, error));
+ if (lower_bound != 1) {
+ ArrowErrorSet(error, "Array value with lower bound != 1 is not supported");
+ return EINVAL;
+ }
+ }
+
+ for (int64_t i = 0; i < n_items; i++) {
+ int32_t child_field_size_bytes;
+ NANOARROW_RETURN_NOT_OK(ReadChecked<int32_t>(data, &child_field_size_bytes, error));
+ NANOARROW_RETURN_NOT_OK(
+ child_->Read(data, child_field_size_bytes, array->children[0], error));
+ }
+
+ int64_t bytes_read = data->data.as_uint8 - data0;
+ if (bytes_read != field_size_bytes) {
+ ArrowErrorSet(error, "Expected to read %d bytes from array field but read %d bytes",
+ static_cast<int>(field_size_bytes),
+ static_cast<int>(bytes_read)); // NOLINT(runtime/int)
+ return EINVAL;
+ }
+
+ NANOARROW_RETURN_NOT_OK(ArrowArrayFinishElement(array));
+ return NANOARROW_OK;
+ }
+
+ private:
+ std::unique_ptr<PostgresCopyFieldReader> child_;
+};
+
+class PostgresCopyRecordFieldReader : public PostgresCopyFieldReader {
+ public:
+ void AppendChild(std::unique_ptr<PostgresCopyFieldReader> child) {
+ int64_t child_i = static_cast<int64_t>(children_.size());
+ children_.push_back(std::move(child));
+ children_[child_i]->Init(pg_type_.child(child_i));
+ }
+
+ ArrowErrorCode InitSchema(ArrowSchema* schema) override {
+ NANOARROW_RETURN_NOT_OK(PostgresCopyFieldReader::InitSchema(schema));
+ for (int64_t i = 0; i < schema->n_children; i++) {
+ NANOARROW_RETURN_NOT_OK(children_[i]->InitSchema(schema->children[i]));
+ }
+
+ return NANOARROW_OK;
+ }
+
+ ArrowErrorCode InitArray(ArrowArray* array) override {
+ NANOARROW_RETURN_NOT_OK(PostgresCopyFieldReader::InitArray(array));
+ for (int64_t i = 0; i < array->n_children; i++) {
+ NANOARROW_RETURN_NOT_OK(children_[i]->InitArray(array->children[i]));
+ }
+
+ return NANOARROW_OK;
+ }
+
+ ArrowErrorCode Read(ArrowBufferView* data, int32_t field_size_bytes, ArrowArray* array,
+ ArrowError* error) override {
+ if (field_size_bytes < 0) {
+ return ArrowArrayAppendNull(array, 1);
+ }
+
+ // Keep the cursor where we start to parse the field so we can check
+ // the number of bytes read against the field size when finished
+ const uint8_t* data0 = data->data.as_uint8;
+
+ int32_t n_fields;
+ NANOARROW_RETURN_NOT_OK(ReadChecked<int32_t>(data, &n_fields, error));
+ if (n_fields != array->n_children) {
+ ArrowErrorSet(error, "Expected nested record type to have %ld fields but got %d",
+ static_cast<long>(array->n_children), // NOLINT(runtime/int)
+ static_cast<int>(n_fields)); // NOLINT(runtime/int)
+ return EINVAL;
+ }
+
+ for (int32_t i = 0; i < n_fields; i++) {
+ uint32_t child_oid;
+ NANOARROW_RETURN_NOT_OK(ReadChecked<uint32_t>(data, &child_oid, error));
+
+ int32_t child_field_size_bytes;
+ NANOARROW_RETURN_NOT_OK(ReadChecked<int32_t>(data, &child_field_size_bytes, error));
+ int result =
+ children_[i]->Read(data, child_field_size_bytes, array->children[i], error);
+
+ // On overflow, pretend all previous children for this struct were never
+ // appended to. This leaves array in a valid state in the specific case
+ // where EOVERFLOW was returned so that a higher level caller can attempt
+ // to try again after creating a new array.
+ if (result == EOVERFLOW) {
+ for (int16_t j = 0; j < i; j++) {
+ array->children[j]->length--;
+ }
+ }
+
+ if (result != NANOARROW_OK) {
+ return result;
+ }
+ }
+
+ // field size == -1 means don't check (e.g., for a top-level row tuple)
+ int64_t bytes_read = data->data.as_uint8 - data0;
+ if (field_size_bytes != -1 && bytes_read != field_size_bytes) {
+ ArrowErrorSet(error,
+ "Expected to read %d bytes from record field but read %d bytes",
+ static_cast<int>(field_size_bytes),
+ static_cast<int>(bytes_read)); // NOLINT(runtime/int)
+ return EINVAL;
+ }
+
+ array->length++;
+ return NANOARROW_OK;
+ }
+
+ private:
+ std::vector<std::unique_ptr<PostgresCopyFieldReader>> children_;
+};
+
+// Subtely different from a Record field item: field count is an int16_t
+// instead of an int32_t and each field is not prefixed by its OID.
+class PostgresCopyFieldTupleReader : public PostgresCopyFieldReader {
+ public:
+ void AppendChild(std::unique_ptr<PostgresCopyFieldReader> child) {
+ int64_t child_i = static_cast<int64_t>(children_.size());
+ children_.push_back(std::move(child));
+ children_[child_i]->Init(pg_type_.child(child_i));
+ }
+
+ ArrowErrorCode InitSchema(ArrowSchema* schema) override {
+ NANOARROW_RETURN_NOT_OK(PostgresCopyFieldReader::InitSchema(schema));
+ for (int64_t i = 0; i < schema->n_children; i++) {
+ NANOARROW_RETURN_NOT_OK(children_[i]->InitSchema(schema->children[i]));
+ }
+
+ return NANOARROW_OK;
+ }
+
+ ArrowErrorCode InitArray(ArrowArray* array) override {
+ NANOARROW_RETURN_NOT_OK(PostgresCopyFieldReader::InitArray(array));
+ for (int64_t i = 0; i < array->n_children; i++) {
+ NANOARROW_RETURN_NOT_OK(children_[i]->InitArray(array->children[i]));
+ }
+
+ return NANOARROW_OK;
+ }
+
+ ArrowErrorCode Read(ArrowBufferView* data, int32_t field_size_bytes, ArrowArray* array,
+ ArrowError* error) override {
+ int16_t n_fields;
+ NANOARROW_RETURN_NOT_OK(ReadChecked<int16_t>(data, &n_fields, error));
+ if (n_fields == -1) {
+ return ENODATA;
+ } else if (n_fields != array->n_children) {
+ ArrowErrorSet(error,
+ "Expected -1 for end-of-stream or number of fields in output array "
+ "(%ld) but got %d",
+ static_cast<long>(array->n_children), // NOLINT(runtime/int)
+ static_cast<int>(n_fields)); // NOLINT(runtime/int)
+ return EINVAL;
+ }
+
+ for (int16_t i = 0; i < n_fields; i++) {
+ int32_t child_field_size_bytes;
+ NANOARROW_RETURN_NOT_OK(ReadChecked<int32_t>(data, &child_field_size_bytes, error));
+ int result =
+ children_[i]->Read(data, child_field_size_bytes, array->children[i], error);
+
+ // On overflow, pretend all previous children for this struct were never
+ // appended to. This leaves array in a valid state in the specific case
+ // where EOVERFLOW was returned so that a higher level caller can attempt
+ // to try again after creating a new array.
+ if (result == EOVERFLOW) {
+ for (int16_t j = 0; j < i; j++) {
+ array->children[j]->length--;
+ }
+ }
+
+ if (result != NANOARROW_OK) {
+ return result;
+ }
+ }
+
+ array->length++;
+ return NANOARROW_OK;
+ }
+
+ private:
+ std::vector<std::unique_ptr<PostgresCopyFieldReader>> children_;
+};
+
+// Factory for a PostgresCopyFieldReader that instantiates the proper subclass
+// and gives a nice error for Postgres type -> Arrow type conversions that aren't
+// supported.
+static inline ArrowErrorCode ErrorCantConvert(ArrowError* error,
+ const PostgresType& pg_type,
+ const ArrowSchemaView& schema_view) {
+ ArrowErrorSet(error, "Can't convert Postgres type '%s' to Arrow type '%s'",
+ pg_type.typname().c_str(),
+ ArrowTypeString(schema_view.type)); // NOLINT(runtime/int)
+ return EINVAL;
+}
+
+static inline ArrowErrorCode MakeCopyFieldReader(const PostgresType& pg_type,
+ ArrowSchema* schema,
+ PostgresCopyFieldReader** out,
+ ArrowError* error) {
+ ArrowSchemaView schema_view;
+ NANOARROW_RETURN_NOT_OK(ArrowSchemaViewInit(&schema_view, schema, nullptr));
+
+ switch (schema_view.type) {
+ case NANOARROW_TYPE_BOOL:
+ switch (pg_type.type_id()) {
+ case PostgresTypeId::kBool:
+ *out = new PostgresCopyBooleanFieldReader();
+ return NANOARROW_OK;
+ default:
+ return ErrorCantConvert(error, pg_type, schema_view);
+ }
+
+ case NANOARROW_TYPE_INT16:
+ switch (pg_type.type_id()) {
+ case PostgresTypeId::kInt2:
+ *out = new PostgresCopyNetworkEndianFieldReader<int16_t>();
+ return NANOARROW_OK;
+ default:
+ return ErrorCantConvert(error, pg_type, schema_view);
+ }
+
+ case NANOARROW_TYPE_INT32:
+ switch (pg_type.type_id()) {
+ case PostgresTypeId::kInt4:
+ case PostgresTypeId::kOid:
+ case PostgresTypeId::kRegproc:
+ *out = new PostgresCopyNetworkEndianFieldReader<int32_t>();
+ return NANOARROW_OK;
+ default:
+ return ErrorCantConvert(error, pg_type, schema_view);
+ }
+
+ case NANOARROW_TYPE_INT64:
+ switch (pg_type.type_id()) {
+ case PostgresTypeId::kInt8:
+ *out = new PostgresCopyNetworkEndianFieldReader<int64_t>();
+ return NANOARROW_OK;
+ default:
+ return ErrorCantConvert(error, pg_type, schema_view);
+ }
+
+ case NANOARROW_TYPE_FLOAT:
+ switch (pg_type.type_id()) {
+ case PostgresTypeId::kFloat4:
+ *out = new PostgresCopyNetworkEndianFieldReader<uint32_t>();
+ return NANOARROW_OK;
+ default:
+ return ErrorCantConvert(error, pg_type, schema_view);
+ }
+
+ case NANOARROW_TYPE_DOUBLE:
+ switch (pg_type.type_id()) {
+ case PostgresTypeId::kFloat8:
+ *out = new PostgresCopyNetworkEndianFieldReader<uint64_t>();
+ return NANOARROW_OK;
+ default:
+ return ErrorCantConvert(error, pg_type, schema_view);
+ }
+
+ case NANOARROW_TYPE_STRING:
+ switch (pg_type.type_id()) {
+ case PostgresTypeId::kChar:
+ case PostgresTypeId::kVarchar:
+ case PostgresTypeId::kText:
+ case PostgresTypeId::kBpchar:
+ case PostgresTypeId::kName:
+ *out = new PostgresCopyBinaryFieldReader();
+ return NANOARROW_OK;
+ default:
+ return ErrorCantConvert(error, pg_type, schema_view);
+ }
+
+ case NANOARROW_TYPE_BINARY:
+ // No need to check pg_type here: we can return the bytes of any
+ // Postgres type as binary.
+ *out = new PostgresCopyBinaryFieldReader();
+ return NANOARROW_OK;
+
+ case NANOARROW_TYPE_LIST:
+ switch (pg_type.type_id()) {
+ case PostgresTypeId::kArray: {
+ if (pg_type.n_children() != 1) {
+ ArrowErrorSet(
+ error, "Expected Postgres array type to have one child but found %ld",
+ static_cast<long>(pg_type.n_children())); // NOLINT(runtime/int)
+ return EINVAL;
+ }
+
+ auto array_reader = std::unique_ptr<PostgresCopyArrayFieldReader>(
+ new PostgresCopyArrayFieldReader());
+ array_reader->Init(pg_type);
+
+ PostgresCopyFieldReader* child_reader;
+ NANOARROW_RETURN_NOT_OK(MakeCopyFieldReader(
+ pg_type.child(0), schema->children[0], &child_reader, error));
+ array_reader->InitChild(std::unique_ptr<PostgresCopyFieldReader>(child_reader));
+
+ *out = array_reader.release();
+ return NANOARROW_OK;
+ }
+ default:
+ return ErrorCantConvert(error, pg_type, schema_view);
+ }
+
+ case NANOARROW_TYPE_STRUCT:
+ switch (pg_type.type_id()) {
+ case PostgresTypeId::kRecord: {
+ if (pg_type.n_children() != schema->n_children) {
+ ArrowErrorSet(error,
+ "Can't convert Postgres record type with %ld chlidren to Arrow "
+ "struct type with %ld children",
+ static_cast<long>(pg_type.n_children()), // NOLINT(runtime/int)
+ static_cast<long>(schema->n_children)); // NOLINT(runtime/int)
+ return EINVAL;
+ }
+
+ auto record_reader = std::unique_ptr<PostgresCopyRecordFieldReader>(
+ new PostgresCopyRecordFieldReader());
+ record_reader->Init(pg_type);
+
+ for (int64_t i = 0; i < pg_type.n_children(); i++) {
+ PostgresCopyFieldReader* child_reader;
+ NANOARROW_RETURN_NOT_OK(MakeCopyFieldReader(
+ pg_type.child(i), schema->children[i], &child_reader, error));
+ record_reader->AppendChild(
+ std::unique_ptr<PostgresCopyFieldReader>(child_reader));
+ }
+
+ *out = record_reader.release();
+ return NANOARROW_OK;
+ }
+ default:
+ return ErrorCantConvert(error, pg_type, schema_view);
+ }
+ default:
+ return ErrorCantConvert(error, pg_type, schema_view);
+ }
+}
+
+class PostgresCopyStreamReader {
+ public:
+ ArrowErrorCode Init(const PostgresType& pg_type) {
+ if (pg_type.type_id() != PostgresTypeId::kRecord) {
+ return EINVAL;
+ }
+
+ root_reader_.Init(pg_type);
+ return NANOARROW_OK;
+ }
+
+ ArrowErrorCode SetOutputSchema(ArrowSchema* schema, ArrowError* error) {
+ if (std::string(schema_->format) != "+s") {
+ ArrowErrorSet(
+ error,
+ "Expected output schema of type struct but got output schema with format '%s'",
+ schema_->format); // NOLINT(runtime/int)
+ return EINVAL;
+ }
+
+ if (schema_->n_children != root_reader_.InputType().n_children()) {
+ ArrowErrorSet(error,
+ "Expected output schema with %ld columns to match Postgres input but "
+ "got schema with %ld columns",
+ static_cast<long>( // NOLINT(runtime/int)
+ root_reader_.InputType().n_children()),
+ static_cast<long>(schema->n_children)); // NOLINT(runtime/int)
+ return EINVAL;
+ }
+
+ schema_.reset(schema);
+ return NANOARROW_OK;
+ }
+
+ ArrowErrorCode InferOutputSchema(ArrowError* error) {
+ schema_.reset();
+ ArrowSchemaInit(schema_.get());
+ NANOARROW_RETURN_NOT_OK(root_reader_.InputType().SetSchema(schema_.get()));
+ return NANOARROW_OK;
+ }
+
+ ArrowErrorCode InitFieldReaders(ArrowError* error) {
+ if (schema_->release == nullptr) {
+ return EINVAL;
+ }
+
+ const PostgresType& root_type = root_reader_.InputType();
+
+ for (int64_t i = 0; i < root_type.n_children(); i++) {
+ const PostgresType& child_type = root_type.child(i);
+ PostgresCopyFieldReader* child_reader;
+ NANOARROW_RETURN_NOT_OK(
+ MakeCopyFieldReader(child_type, schema_->children[i], &child_reader, error));
+ root_reader_.AppendChild(std::unique_ptr<PostgresCopyFieldReader>(child_reader));
+ }
+
+ NANOARROW_RETURN_NOT_OK(root_reader_.InitSchema(schema_.get()));
+ return NANOARROW_OK;
+ }
+
+ ArrowErrorCode ReadHeader(ArrowBufferView* data, ArrowError* error) {
+ if (data->size_bytes < static_cast<int64_t>(sizeof(kPgCopyBinarySignature))) {
+ ArrowErrorSet(
+ error,
+ "Expected PGCOPY signature of %ld bytes at beginning of stream but "
+ "found %ld bytes of input",
+ static_cast<long>(sizeof(kPgCopyBinarySignature)), // NOLINT(runtime/int)
+ static_cast<long>(data->size_bytes)); // NOLINT(runtime/int)
+ return EINVAL;
+ }
+
+ if (memcmp(data->data.data, kPgCopyBinarySignature, sizeof(kPgCopyBinarySignature)) !=
+ 0) {
+ ArrowErrorSet(error, "Invalid PGCOPY signature at beginning of stream");
+ return EINVAL;
+ }
+
+ data->data.as_uint8 += sizeof(kPgCopyBinarySignature);
+ data->size_bytes -= sizeof(kPgCopyBinarySignature);
+
+ uint32_t flags;
+ NANOARROW_RETURN_NOT_OK(ReadChecked<uint32_t>(data, &flags, error));
+ uint32_t extension_length;
+ NANOARROW_RETURN_NOT_OK(ReadChecked<uint32_t>(data, &extension_length, error));
+
+ if (data->size_bytes < static_cast<int64_t>(extension_length)) {
+ ArrowErrorSet(error,
+ "Expected %ld bytes of extension metadata at start of stream but "
+ "found %ld bytes of input",
+ static_cast<long>(extension_length), // NOLINT(runtime/int)
+ static_cast<long>(data->size_bytes)); // NOLINT(runtime/int)
+ return EINVAL;
+ }
+
+ data->data.as_uint8 += extension_length;
+ data->size_bytes -= extension_length;
+ return NANOARROW_OK;
+ }
+
+ ArrowErrorCode ReadRecord(ArrowBufferView* data, ArrowError* error) {
+ if (array_->release == nullptr) {
+ NANOARROW_RETURN_NOT_OK(
+ ArrowArrayInitFromSchema(array_.get(), schema_.get(), error));
+ NANOARROW_RETURN_NOT_OK(ArrowArrayStartAppending(array_.get()));
+ NANOARROW_RETURN_NOT_OK(root_reader_.InitArray(array_.get()));
+ }
+
+ NANOARROW_RETURN_NOT_OK(root_reader_.Read(data, -1, array_.get(), error));
+ return NANOARROW_OK;
+ }
+
+ ArrowErrorCode GetSchema(ArrowSchema* out) {
+ return ArrowSchemaDeepCopy(schema_.get(), out);
+ }
+
+ ArrowErrorCode GetArray(ArrowArray* out, ArrowError* error) {
+ if (array_->release == nullptr) {
+ return EINVAL;
+ }
+
+ NANOARROW_RETURN_NOT_OK(ArrowArrayFinishBuildingDefault(array_.get(), error));
+ ArrowArrayMove(array_.get(), out);
+ return NANOARROW_OK;
+ }
+
+ private:
+ PostgresCopyFieldTupleReader root_reader_;
+ nanoarrow::UniqueSchema schema_;
+ nanoarrow::UniqueArray array_;
+};
+
+} // namespace adbcpq
diff --git a/c/driver/postgresql/postgres_copy_reader_test.cc b/c/driver/postgresql/postgres_copy_reader_test.cc
new file mode 100644
index 0000000..26a1ab3
--- /dev/null
+++ b/c/driver/postgresql/postgres_copy_reader_test.cc
@@ -0,0 +1,501 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gtest/gtest.h>
+#include <nanoarrow/nanoarrow.hpp>
+
+#include "postgres_copy_reader.h"
+
+namespace adbcpq {
+
+class PostgresCopyStreamTester {
+ public:
+ ArrowErrorCode Init(const PostgresType& root_type, ArrowError* error = nullptr) {
+ NANOARROW_RETURN_NOT_OK(reader_.Init(root_type));
+ NANOARROW_RETURN_NOT_OK(reader_.InferOutputSchema(error));
+ NANOARROW_RETURN_NOT_OK(reader_.InitFieldReaders(error));
+ return NANOARROW_OK;
+ }
+
+ ArrowErrorCode ReadAll(ArrowBufferView* data, ArrowError* error = nullptr) {
+ NANOARROW_RETURN_NOT_OK(reader_.ReadHeader(data, error));
+
+ int result;
+ do {
+ result = reader_.ReadRecord(data, error);
+ } while (result == NANOARROW_OK);
+
+ return result;
+ }
+
+ void GetSchema(ArrowSchema* out) { reader_.GetSchema(out); }
+
+ ArrowErrorCode GetArray(ArrowArray* out, ArrowError* error = nullptr) {
+ return reader_.GetArray(out, error);
+ }
+
+ private:
+ PostgresCopyStreamReader reader_;
+};
+
+// COPY (SELECT CAST("col" AS BOOLEAN) AS "col" FROM ( VALUES (TRUE), (FALSE), (NULL)) AS
+// drvd("col")) TO STDOUT;
+static uint8_t kTestPgCopyBoolean[] = {
+ 0x50, 0x47, 0x43, 0x4f, 0x50, 0x59, 0x0a, 0xff, 0x0d, 0x0a, 0x00, 0x00, 0x00, 0x00,
+ 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x01, 0x00, 0x01,
+ 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x01, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff};
+
+TEST(PostgresCopyUtilsTest, PostgresCopyReadBoolean) {
+ ArrowBufferView data;
+ data.data.as_uint8 = kTestPgCopyBoolean;
+ data.size_bytes = sizeof(kTestPgCopyBoolean);
+
+ auto col_type = PostgresType(PostgresTypeId::kBool);
+ PostgresType input_type(PostgresTypeId::kRecord);
+ input_type.AppendChild("col", col_type);
+
+ PostgresCopyStreamTester tester;
+ ASSERT_EQ(tester.Init(input_type), NANOARROW_OK);
+ ASSERT_EQ(tester.ReadAll(&data), ENODATA);
+
+ ASSERT_EQ(data.data.as_uint8 - kTestPgCopyBoolean, sizeof(kTestPgCopyBoolean));
+ ASSERT_EQ(data.size_bytes, 0);
+
+ nanoarrow::UniqueArray array;
+ ASSERT_EQ(tester.GetArray(array.get()), NANOARROW_OK);
+ ASSERT_EQ(array->length, 3);
+ ASSERT_EQ(array->n_children, 1);
+
+ const uint8_t* validity =
+ reinterpret_cast<const uint8_t*>(array->children[0]->buffers[0]);
+ const uint8_t* data_buffer =
+ reinterpret_cast<const uint8_t*>(array->children[0]->buffers[1]);
+ ASSERT_NE(validity, nullptr);
+ ASSERT_NE(data_buffer, nullptr);
+
+ ASSERT_TRUE(ArrowBitGet(validity, 0));
+ ASSERT_TRUE(ArrowBitGet(validity, 1));
+ ASSERT_FALSE(ArrowBitGet(validity, 2));
+
+ ASSERT_TRUE(ArrowBitGet(data_buffer, 0));
+ ASSERT_FALSE(ArrowBitGet(data_buffer, 1));
+ ASSERT_FALSE(ArrowBitGet(data_buffer, 2));
+}
+
+// COPY (SELECT CAST("col" AS SMALLINT) AS "col" FROM ( VALUES (-123), (-1), (1), (123),
+// (NULL)) AS drvd("col")) TO STDOUT WITH (FORMAT binary);
+static uint8_t kTestPgCopySmallInt[] = {
+ 0x50, 0x47, 0x43, 0x4f, 0x50, 0x59, 0x0a, 0xff, 0x0d, 0x0a, 0x00, 0x00,
+ 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
+ 0x02, 0xff, 0x85, 0x00, 0x01, 0x00, 0x00, 0x00, 0x02, 0xff, 0xff, 0x00,
+ 0x01, 0x00, 0x00, 0x00, 0x02, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00,
+ 0x02, 0x00, 0x7b, 0x00, 0x01, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff};
+
+TEST(PostgresCopyUtilsTest, PostgresCopyReadSmallInt) {
+ ArrowBufferView data;
+ data.data.as_uint8 = kTestPgCopySmallInt;
+ data.size_bytes = sizeof(kTestPgCopySmallInt);
+
+ auto col_type = PostgresType(PostgresTypeId::kInt2);
+ PostgresType input_type(PostgresTypeId::kRecord);
+ input_type.AppendChild("col", col_type);
+
+ PostgresCopyStreamTester tester;
+ ASSERT_EQ(tester.Init(input_type), NANOARROW_OK);
+ ASSERT_EQ(tester.ReadAll(&data), ENODATA);
+ ASSERT_EQ(data.data.as_uint8 - kTestPgCopySmallInt, sizeof(kTestPgCopySmallInt));
+ ASSERT_EQ(data.size_bytes, 0);
+
+ nanoarrow::UniqueArray array;
+ ASSERT_EQ(tester.GetArray(array.get()), NANOARROW_OK);
+ ASSERT_EQ(array->length, 5);
+ ASSERT_EQ(array->n_children, 1);
+
+ auto validity = reinterpret_cast<const uint8_t*>(array->children[0]->buffers[0]);
+ auto data_buffer = reinterpret_cast<const int16_t*>(array->children[0]->buffers[1]);
+ ASSERT_NE(validity, nullptr);
+ ASSERT_NE(data_buffer, nullptr);
+
+ ASSERT_TRUE(ArrowBitGet(validity, 0));
+ ASSERT_TRUE(ArrowBitGet(validity, 1));
+ ASSERT_TRUE(ArrowBitGet(validity, 2));
+ ASSERT_TRUE(ArrowBitGet(validity, 3));
+ ASSERT_FALSE(ArrowBitGet(validity, 4));
+
+ ASSERT_EQ(data_buffer[0], -123);
+ ASSERT_EQ(data_buffer[1], -1);
+ ASSERT_EQ(data_buffer[2], 1);
+ ASSERT_EQ(data_buffer[3], 123);
+ ASSERT_EQ(data_buffer[4], 0);
+}
+
+// COPY (SELECT CAST("col" AS INTEGER) AS "col" FROM ( VALUES (-123), (-1), (1), (123),
+// (NULL)) AS drvd("col")) TO STDOUT WITH (FORMAT binary);
+static uint8_t kTestPgCopyInteger[] = {
+ 0x50, 0x47, 0x43, 0x4f, 0x50, 0x59, 0x0a, 0xff, 0x0d, 0x0a, 0x00, 0x00, 0x00, 0x00,
+ 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x04, 0xff, 0xff, 0xff,
+ 0x85, 0x00, 0x01, 0x00, 0x00, 0x00, 0x04, 0xff, 0xff, 0xff, 0xff, 0x00, 0x01, 0x00,
+ 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x04, 0x00,
+ 0x00, 0x00, 0x7b, 0x00, 0x01, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff};
+
+TEST(PostgresCopyUtilsTest, PostgresCopyReadInteger) {
+ ArrowBufferView data;
+ data.data.as_uint8 = kTestPgCopyInteger;
+ data.size_bytes = sizeof(kTestPgCopyInteger);
+
+ auto col_type = PostgresType(PostgresTypeId::kInt4);
+ PostgresType input_type(PostgresTypeId::kRecord);
+ input_type.AppendChild("col", col_type);
+
+ PostgresCopyStreamTester tester;
+ ASSERT_EQ(tester.Init(input_type), NANOARROW_OK);
+ ASSERT_EQ(tester.ReadAll(&data), ENODATA);
+ ASSERT_EQ(data.data.as_uint8 - kTestPgCopyInteger, sizeof(kTestPgCopyInteger));
+ ASSERT_EQ(data.size_bytes, 0);
+
+ nanoarrow::UniqueArray array;
+ ASSERT_EQ(tester.GetArray(array.get()), NANOARROW_OK);
+ ASSERT_EQ(array->length, 5);
+ ASSERT_EQ(array->n_children, 1);
+
+ auto validity = reinterpret_cast<const uint8_t*>(array->children[0]->buffers[0]);
+ auto data_buffer = reinterpret_cast<const int32_t*>(array->children[0]->buffers[1]);
+ ASSERT_NE(validity, nullptr);
+ ASSERT_NE(data_buffer, nullptr);
+
+ ASSERT_TRUE(ArrowBitGet(validity, 0));
+ ASSERT_TRUE(ArrowBitGet(validity, 1));
+ ASSERT_TRUE(ArrowBitGet(validity, 2));
+ ASSERT_TRUE(ArrowBitGet(validity, 3));
+ ASSERT_FALSE(ArrowBitGet(validity, 4));
+
+ ASSERT_EQ(data_buffer[0], -123);
+ ASSERT_EQ(data_buffer[1], -1);
+ ASSERT_EQ(data_buffer[2], 1);
+ ASSERT_EQ(data_buffer[3], 123);
+ ASSERT_EQ(data_buffer[4], 0);
+}
+
+// COPY (SELECT CAST("col" AS BIGINT) AS "col" FROM ( VALUES (-123), (-1), (1), (123),
+// (NULL)) AS drvd("col")) TO STDOUT WITH (FORMAT binary);
+static uint8_t kTestPgCopyBigInt[] = {
+ 0x50, 0x47, 0x43, 0x4f, 0x50, 0x59, 0x0a, 0xff, 0x0d, 0x0a, 0x00, 0x00, 0x00, 0x00,
+ 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x08, 0xff, 0xff, 0xff,
+ 0xff, 0xff, 0xff, 0xff, 0x85, 0x00, 0x01, 0x00, 0x00, 0x00, 0x08, 0xff, 0xff, 0xff,
+ 0xff, 0xff, 0xff, 0xff, 0xff, 0x00, 0x01, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00,
+ 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00,
+ 0x00, 0x00, 0x00, 0x00, 0x7b, 0x00, 0x01, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff};
+
+TEST(PostgresCopyUtilsTest, PostgresCopyReadBigInt) {
+ ArrowBufferView data;
+ data.data.as_uint8 = kTestPgCopyBigInt;
+ data.size_bytes = sizeof(kTestPgCopyBigInt);
+
+ auto col_type = PostgresType(PostgresTypeId::kInt8);
+ PostgresType input_type(PostgresTypeId::kRecord);
+ input_type.AppendChild("col", col_type);
+
+ PostgresCopyStreamTester tester;
+ ASSERT_EQ(tester.Init(input_type), NANOARROW_OK);
+ ASSERT_EQ(tester.ReadAll(&data), ENODATA);
+ ASSERT_EQ(data.data.as_uint8 - kTestPgCopyBigInt, sizeof(kTestPgCopyBigInt));
+ ASSERT_EQ(data.size_bytes, 0);
+
+ nanoarrow::UniqueArray array;
+ ASSERT_EQ(tester.GetArray(array.get()), NANOARROW_OK);
+ ASSERT_EQ(array->length, 5);
+ ASSERT_EQ(array->n_children, 1);
+
+ auto validity = reinterpret_cast<const uint8_t*>(array->children[0]->buffers[0]);
+ auto data_buffer = reinterpret_cast<const int64_t*>(array->children[0]->buffers[1]);
+ ASSERT_NE(validity, nullptr);
+ ASSERT_NE(data_buffer, nullptr);
+
+ ASSERT_TRUE(ArrowBitGet(validity, 0));
+ ASSERT_TRUE(ArrowBitGet(validity, 1));
+ ASSERT_TRUE(ArrowBitGet(validity, 2));
+ ASSERT_TRUE(ArrowBitGet(validity, 3));
+ ASSERT_FALSE(ArrowBitGet(validity, 4));
+
+ ASSERT_EQ(data_buffer[0], -123);
+ ASSERT_EQ(data_buffer[1], -1);
+ ASSERT_EQ(data_buffer[2], 1);
+ ASSERT_EQ(data_buffer[3], 123);
+ ASSERT_EQ(data_buffer[4], 0);
+}
+
+// COPY (SELECT CAST("col" AS REAL) AS "col" FROM ( VALUES (-123.456), (-1), (1),
+// (123.456), (NULL)) AS drvd("col")) TO STDOUT WITH (FORMAT binary);
+static uint8_t kTestPgCopyReal[] = {
+ 0x50, 0x47, 0x43, 0x4f, 0x50, 0x59, 0x0a, 0xff, 0x0d, 0x0a, 0x00, 0x00, 0x00, 0x00,
+ 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x04, 0xc2, 0xf6, 0xe9,
+ 0x79, 0x00, 0x01, 0x00, 0x00, 0x00, 0x04, 0xbf, 0x80, 0x00, 0x00, 0x00, 0x01, 0x00,
+ 0x00, 0x00, 0x04, 0x3f, 0x80, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x04, 0x42,
+ 0xf6, 0xe9, 0x79, 0x00, 0x01, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff};
+
+TEST(PostgresCopyUtilsTest, PostgresCopyReadReal) {
+ ArrowBufferView data;
+ data.data.as_uint8 = kTestPgCopyReal;
+ data.size_bytes = sizeof(kTestPgCopyReal);
+
+ auto col_type = PostgresType(PostgresTypeId::kFloat4);
+ PostgresType input_type(PostgresTypeId::kRecord);
+ input_type.AppendChild("col", col_type);
+
+ PostgresCopyStreamTester tester;
+ ASSERT_EQ(tester.Init(input_type), NANOARROW_OK);
+ ASSERT_EQ(tester.ReadAll(&data), ENODATA);
+ ASSERT_EQ(data.data.as_uint8 - kTestPgCopyReal, sizeof(kTestPgCopyReal));
+ ASSERT_EQ(data.size_bytes, 0);
+
+ nanoarrow::UniqueArray array;
+ ASSERT_EQ(tester.GetArray(array.get()), NANOARROW_OK);
+ ASSERT_EQ(array->length, 5);
+ ASSERT_EQ(array->n_children, 1);
+
+ auto validity = reinterpret_cast<const uint8_t*>(array->children[0]->buffers[0]);
+ auto data_buffer = reinterpret_cast<const float*>(array->children[0]->buffers[1]);
+ ASSERT_NE(validity, nullptr);
+ ASSERT_NE(data_buffer, nullptr);
+
+ ASSERT_TRUE(ArrowBitGet(validity, 0));
+ ASSERT_TRUE(ArrowBitGet(validity, 1));
+ ASSERT_TRUE(ArrowBitGet(validity, 2));
+ ASSERT_TRUE(ArrowBitGet(validity, 3));
+ ASSERT_FALSE(ArrowBitGet(validity, 4));
+
+ ASSERT_FLOAT_EQ(data_buffer[0], -123.456);
+ ASSERT_EQ(data_buffer[1], -1);
+ ASSERT_EQ(data_buffer[2], 1);
+ ASSERT_FLOAT_EQ(data_buffer[3], 123.456);
+ ASSERT_EQ(data_buffer[4], 0);
+}
+
+// COPY (SELECT CAST("col" AS DOUBLE PRECISION) AS "col" FROM ( VALUES (-123.456), (-1),
+// (1), (123.456), (NULL)) AS drvd("col")) TO STDOUT WITH (FORMAT binary);
+static uint8_t kTestPgCopyDoublePrecision[] = {
+ 0x50, 0x47, 0x43, 0x4f, 0x50, 0x59, 0x0a, 0xff, 0x0d, 0x0a, 0x00, 0x00, 0x00, 0x00,
+ 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x08, 0xc0, 0x5e, 0xdd,
+ 0x2f, 0x1a, 0x9f, 0xbe, 0x77, 0x00, 0x01, 0x00, 0x00, 0x00, 0x08, 0xbf, 0xf0, 0x00,
+ 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x08, 0x3f, 0xf0, 0x00,
+ 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x08, 0x40, 0x5e, 0xdd,
+ 0x2f, 0x1a, 0x9f, 0xbe, 0x77, 0x00, 0x01, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff};
+
+TEST(PostgresCopyUtilsTest, PostgresCopyReadDoublePrecision) {
+ ArrowBufferView data;
+ data.data.as_uint8 = kTestPgCopyDoublePrecision;
+ data.size_bytes = sizeof(kTestPgCopyDoublePrecision);
+
+ auto col_type = PostgresType(PostgresTypeId::kFloat8);
+ PostgresType input_type(PostgresTypeId::kRecord);
+ input_type.AppendChild("col", col_type);
+
+ PostgresCopyStreamTester tester;
+ ASSERT_EQ(tester.Init(input_type), NANOARROW_OK);
+ ASSERT_EQ(tester.ReadAll(&data), ENODATA);
+ ASSERT_EQ(data.data.as_uint8 - kTestPgCopyDoublePrecision,
+ sizeof(kTestPgCopyDoublePrecision));
+ ASSERT_EQ(data.size_bytes, 0);
+
+ nanoarrow::UniqueArray array;
+ ASSERT_EQ(tester.GetArray(array.get()), NANOARROW_OK);
+ ASSERT_EQ(array->length, 5);
+ ASSERT_EQ(array->n_children, 1);
+
+ auto validity = reinterpret_cast<const uint8_t*>(array->children[0]->buffers[0]);
+ auto data_buffer = reinterpret_cast<const double*>(array->children[0]->buffers[1]);
+ ASSERT_NE(validity, nullptr);
+ ASSERT_NE(data_buffer, nullptr);
+
+ ASSERT_TRUE(ArrowBitGet(validity, 0));
+ ASSERT_TRUE(ArrowBitGet(validity, 1));
+ ASSERT_TRUE(ArrowBitGet(validity, 2));
+ ASSERT_TRUE(ArrowBitGet(validity, 3));
+ ASSERT_FALSE(ArrowBitGet(validity, 4));
+
+ ASSERT_DOUBLE_EQ(data_buffer[0], -123.456);
+ ASSERT_EQ(data_buffer[1], -1);
+ ASSERT_EQ(data_buffer[2], 1);
+ ASSERT_DOUBLE_EQ(data_buffer[3], 123.456);
+ ASSERT_EQ(data_buffer[4], 0);
+}
+
+// COPY (SELECT CAST("col" AS TEXT) AS "col" FROM ( VALUES ('abc'), ('1234'),
+// (NULL::text)) AS drvd("col")) TO STDOUT WITH (FORMAT binary);
+static uint8_t kTestPgCopyText[] = {
+ 0x50, 0x47, 0x43, 0x4f, 0x50, 0x59, 0x0a, 0xff, 0x0d, 0x0a, 0x00, 0x00,
+ 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
+ 0x03, 0x61, 0x62, 0x63, 0x00, 0x01, 0x00, 0x00, 0x00, 0x04, 0x31, 0x32,
+ 0x33, 0x34, 0x00, 0x01, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff};
+
+TEST(PostgresCopyUtilsTest, PostgresCopyReadText) {
+ ArrowBufferView data;
+ data.data.as_uint8 = kTestPgCopyText;
+ data.size_bytes = sizeof(kTestPgCopyText);
+
+ auto col_type = PostgresType(PostgresTypeId::kText);
+ PostgresType input_type(PostgresTypeId::kRecord);
+ input_type.AppendChild("col", col_type);
+
+ PostgresCopyStreamTester tester;
+ ASSERT_EQ(tester.Init(input_type), NANOARROW_OK);
+ ASSERT_EQ(tester.ReadAll(&data), ENODATA);
+ ASSERT_EQ(data.data.as_uint8 - kTestPgCopyText, sizeof(kTestPgCopyText));
+ ASSERT_EQ(data.size_bytes, 0);
+
+ nanoarrow::UniqueArray array;
+ ASSERT_EQ(tester.GetArray(array.get()), NANOARROW_OK);
+ ASSERT_EQ(array->length, 3);
+ ASSERT_EQ(array->n_children, 1);
+
+ auto validity = reinterpret_cast<const uint8_t*>(array->children[0]->buffers[0]);
+ auto offsets = reinterpret_cast<const int32_t*>(array->children[0]->buffers[1]);
+ auto data_buffer = reinterpret_cast<const char*>(array->children[0]->buffers[2]);
+ ASSERT_NE(validity, nullptr);
+ ASSERT_NE(data_buffer, nullptr);
+
+ ASSERT_TRUE(ArrowBitGet(validity, 0));
+ ASSERT_TRUE(ArrowBitGet(validity, 1));
+ ASSERT_FALSE(ArrowBitGet(validity, 2));
+
+ ASSERT_EQ(offsets[0], 0);
+ ASSERT_EQ(offsets[1], 3);
+ ASSERT_EQ(offsets[2], 7);
+ ASSERT_EQ(offsets[3], 7);
+
+ ASSERT_EQ(std::string(data_buffer + 0, 3), "abc");
+ ASSERT_EQ(std::string(data_buffer + 3, 4), "1234");
+}
+
+// COPY (SELECT CAST("col" AS INTEGER ARRAY) AS "col" FROM ( VALUES ('{-123, -1}'), ('{0,
+// 1, 123}'), (NULL)) AS drvd("col")) TO STDOUT WITH (FORMAT binary);
+static uint8_t kTestPgCopyIntegerArray[] = {
+ 0x50, 0x47, 0x43, 0x4f, 0x50, 0x59, 0x0a, 0xff, 0x0d, 0x0a, 0x00, 0x00, 0x00, 0x00,
+ 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x24, 0x00, 0x00, 0x00,
+ 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x17, 0x00, 0x00, 0x00, 0x02, 0x00,
+ 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x04, 0xff, 0xff, 0xff, 0x85, 0x00, 0x00, 0x00,
+ 0x04, 0xff, 0xff, 0xff, 0xff, 0x00, 0x01, 0x00, 0x00, 0x00, 0x2c, 0x00, 0x00, 0x00,
+ 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x17, 0x00, 0x00, 0x00, 0x03, 0x00,
+ 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+ 0x04, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x7b, 0x00,
+ 0x01, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff};
+
+TEST(PostgresCopyUtilsTest, PostgresCopyReadArray) {
+ ArrowBufferView data;
+ data.data.as_uint8 = kTestPgCopyIntegerArray;
+ data.size_bytes = sizeof(kTestPgCopyIntegerArray);
+
+ auto col_type = PostgresType(PostgresTypeId::kInt4).Array();
+ PostgresType input_type(PostgresTypeId::kRecord);
+ input_type.AppendChild("col", col_type);
+
+ PostgresCopyStreamTester tester;
+ ASSERT_EQ(tester.Init(input_type), NANOARROW_OK);
+ ASSERT_EQ(tester.ReadAll(&data), ENODATA);
+ ASSERT_EQ(data.data.as_uint8 - kTestPgCopyIntegerArray,
+ sizeof(kTestPgCopyIntegerArray));
+ ASSERT_EQ(data.size_bytes, 0);
+
+ nanoarrow::UniqueArray array;
+ ASSERT_EQ(tester.GetArray(array.get()), NANOARROW_OK);
+ ASSERT_EQ(array->length, 3);
+ ASSERT_EQ(array->n_children, 1);
+ ASSERT_EQ(array->children[0]->n_children, 1);
+ ASSERT_EQ(array->children[0]->children[0]->length, 5);
+
+ auto validity = reinterpret_cast<const uint8_t*>(array->children[0]->buffers[0]);
+ auto offsets = reinterpret_cast<const int32_t*>(array->children[0]->buffers[1]);
+ auto data_buffer =
+ reinterpret_cast<const int32_t*>(array->children[0]->children[0]->buffers[1]);
+ ASSERT_NE(validity, nullptr);
+ ASSERT_NE(data_buffer, nullptr);
+
+ ASSERT_TRUE(ArrowBitGet(validity, 0));
+ ASSERT_TRUE(ArrowBitGet(validity, 1));
+ ASSERT_FALSE(ArrowBitGet(validity, 2));
+
+ ASSERT_EQ(offsets[0], 0);
+ ASSERT_EQ(offsets[1], 2);
+ ASSERT_EQ(offsets[2], 5);
+ ASSERT_EQ(offsets[3], 5);
+
+ ASSERT_EQ(data_buffer[0], -123);
+ ASSERT_EQ(data_buffer[1], -1);
+ ASSERT_EQ(data_buffer[2], 0);
+ ASSERT_EQ(data_buffer[3], 1);
+ ASSERT_EQ(data_buffer[4], 123);
+}
+
+// CREATE TYPE custom_record AS (nested1 integer, nested2 double precision);
+// COPY (SELECT CAST("col" AS custom_record) AS "col" FROM ( VALUES ('(123, 456.789)'),
+// ('(12, 345.678)'), (NULL)) AS drvd("col")) TO STDOUT WITH (FORMAT binary);
+static uint8_t kTestPgCopyCustomRecord[] = {
+ 0x50, 0x47, 0x43, 0x4f, 0x50, 0x59, 0x0a, 0xff, 0x0d, 0x0a, 0x00, 0x00, 0x00,
+ 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x20, 0x00,
+ 0x00, 0x00, 0x02, 0x00, 0x00, 0x00, 0x17, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00,
+ 0x00, 0x7b, 0x00, 0x00, 0x02, 0xbd, 0x00, 0x00, 0x00, 0x08, 0x40, 0x7c, 0x8c,
+ 0x9f, 0xbe, 0x76, 0xc8, 0xb4, 0x00, 0x01, 0x00, 0x00, 0x00, 0x20, 0x00, 0x00,
+ 0x00, 0x02, 0x00, 0x00, 0x00, 0x17, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00,
+ 0x0c, 0x00, 0x00, 0x02, 0xbd, 0x00, 0x00, 0x00, 0x08, 0x40, 0x75, 0x9a, 0xd9,
+ 0x16, 0x87, 0x2b, 0x02, 0x00, 0x01, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff};
+
+TEST(PostgresCopyUtilsTest, PostgresCopyReadCustomRecord) {
+ ArrowBufferView data;
+ data.data.as_uint8 = kTestPgCopyCustomRecord;
+ data.size_bytes = sizeof(kTestPgCopyCustomRecord);
+
+ auto col_type = PostgresType(PostgresTypeId::kRecord);
+ col_type.AppendChild("nested1", PostgresType(PostgresTypeId::kInt4));
+ col_type.AppendChild("nested2", PostgresType(PostgresTypeId::kFloat8));
+ PostgresType input_type(PostgresTypeId::kRecord);
+ input_type.AppendChild("col", col_type);
+
+ PostgresCopyStreamTester tester;
+ ASSERT_EQ(tester.Init(input_type), NANOARROW_OK);
+ ASSERT_EQ(tester.ReadAll(&data), ENODATA);
+ ASSERT_EQ(data.data.as_uint8 - kTestPgCopyCustomRecord,
+ sizeof(kTestPgCopyCustomRecord));
+ ASSERT_EQ(data.size_bytes, 0);
+
+ nanoarrow::UniqueArray array;
+ ASSERT_EQ(tester.GetArray(array.get()), NANOARROW_OK);
+ ASSERT_EQ(array->length, 3);
+ ASSERT_EQ(array->n_children, 1);
+ ASSERT_EQ(array->children[0]->n_children, 2);
+ ASSERT_EQ(array->children[0]->children[0]->length, 3);
+ ASSERT_EQ(array->children[0]->children[1]->length, 3);
+
+ auto validity = reinterpret_cast<const uint8_t*>(array->children[0]->buffers[0]);
+ auto data_buffer1 =
+ reinterpret_cast<const int32_t*>(array->children[0]->children[0]->buffers[1]);
+ auto data_buffer2 =
+ reinterpret_cast<const double*>(array->children[0]->children[1]->buffers[1]);
+
+ ASSERT_TRUE(ArrowBitGet(validity, 0));
+ ASSERT_TRUE(ArrowBitGet(validity, 1));
+ ASSERT_FALSE(ArrowBitGet(validity, 2));
+
+ ASSERT_EQ(data_buffer1[0], 123);
+ ASSERT_EQ(data_buffer1[1], 12);
+ ASSERT_EQ(data_buffer1[2], 0);
+
+ ASSERT_DOUBLE_EQ(data_buffer2[0], 456.789);
+ ASSERT_DOUBLE_EQ(data_buffer2[1], 345.678);
+ ASSERT_DOUBLE_EQ(data_buffer2[2], 0);
+}
+
+} // namespace adbcpq
diff --git a/c/driver/postgresql/postgres_type.h b/c/driver/postgresql/postgres_type.h
index b327eb5..a2152a5 100644
--- a/c/driver/postgresql/postgres_type.h
+++ b/c/driver/postgresql/postgres_type.h
@@ -110,6 +110,7 @@ enum class PostgresTypeId {
kXid8,
kXid,
kXml,
+ kUserDefined
};
// Returns the receive function name as defined in the typrecieve column
@@ -198,6 +199,8 @@ class PostgresType {
NANOARROW_RETURN_NOT_OK(ArrowSchemaSetType(schema, NANOARROW_TYPE_INT16));
break;
case PostgresTypeId::kInt4:
+ case PostgresTypeId::kOid:
+ case PostgresTypeId::kRegproc:
NANOARROW_RETURN_NOT_OK(ArrowSchemaSetType(schema, NANOARROW_TYPE_INT32));
break;
case PostgresTypeId::kInt8:
@@ -213,6 +216,7 @@ class PostgresType {
case PostgresTypeId::kBpchar:
case PostgresTypeId::kVarchar:
case PostgresTypeId::kText:
+ case PostgresTypeId::kName:
NANOARROW_RETURN_NOT_OK(ArrowSchemaSetType(schema, NANOARROW_TYPE_STRING));
break;
case PostgresTypeId::kBytea:
@@ -230,9 +234,12 @@ class PostgresType {
NANOARROW_RETURN_NOT_OK(ArrowSchemaSetType(schema, NANOARROW_TYPE_LIST));
NANOARROW_RETURN_NOT_OK(children_[0].SetSchema(schema->children[0]));
break;
+
+ case PostgresTypeId::kUserDefined:
default: {
- // For any types we don't explicitly know how to deal with, we can still
- // return the bytes postgres gives us and attach the type name as metadata
+ // For user-defined types or types we don't explicitly know how to deal with, we
+ // can still return the bytes postgres gives us and attach the type name as
+ // metadata
NANOARROW_RETURN_NOT_OK(ArrowSchemaSetType(schema, NANOARROW_TYPE_BINARY));
nanoarrow::UniqueBuffer buffer;
ArrowMetadataBuilderInit(buffer.get(), nullptr);
@@ -326,13 +333,17 @@ class PostgresTypeResolver {
// class type using InsertClass().
ArrowErrorCode Insert(const Item& item, ArrowError* error) {
auto result = base_.find(item.typreceive);
+ PostgresType base;
+
if (result == base_.end()) {
- ArrowErrorSet(error, "Base type not found for type '%s' with receive function '%s'",
- item.typname, item.typreceive);
- return ENOTSUP;
+ // This occurs when a user-defined type has defined a custom receive function
+ // (e.g., PostGIS/geometry). The only way these types can be supported is
+ // by returning binary unless we hard-code support for some extensions.
+ base = PostgresType(PostgresTypeId::kUserDefined);
+ } else {
+ base = result->second;
}
- const PostgresType& base = result->second;
PostgresType type = base.WithPgTypeInfo(item.oid, item.typname);
switch (base.type_id()) {
diff --git a/c/driver/postgresql/postgres_type_test.cc b/c/driver/postgresql/postgres_type_test.cc
index 74fe53c..02e8da9 100644
--- a/c/driver/postgresql/postgres_type_test.cc
+++ b/c/driver/postgresql/postgres_type_test.cc
@@ -18,25 +18,23 @@
#include <utility>
#include <gtest/gtest.h>
-#include <nanoarrow/nanoarrow.h>
+#include <nanoarrow/nanoarrow.hpp>
#include "postgres_type.h"
-using adbcpq::PostgresType;
-using adbcpq::PostgresTypeId;
-using adbcpq::PostgresTypeResolver;
+namespace adbcpq {
class MockTypeResolver : public PostgresTypeResolver {
public:
ArrowErrorCode Init() {
- auto all_types = adbcpq::PostgresTypeIdAll(false);
+ auto all_types = PostgresTypeIdAll(false);
PostgresTypeResolver::Item item;
item.oid = 0;
// Insert all the base types
for (auto type_id : all_types) {
- std::string typreceive = adbcpq::PostgresTyprecv(type_id);
- std::string typname = adbcpq::PostgresTypname(type_id);
+ std::string typreceive = PostgresTyprecv(type_id);
+ std::string typname = PostgresTypname(type_id);
item.oid++;
item.typname = typname.c_str();
item.typreceive = typreceive.c_str();
@@ -125,168 +123,169 @@ TEST(PostgresTypeTest, PostgresTypeBasic) {
}
TEST(PostgresTypeTest, PostgresTypeSetSchema) {
- ArrowSchema schema;
-
- ArrowSchemaInit(&schema);
- EXPECT_EQ(PostgresType(PostgresTypeId::kBool).SetSchema(&schema), NANOARROW_OK);
- EXPECT_STREQ(schema.format, "b");
- schema.release(&schema);
-
- ArrowSchemaInit(&schema);
- EXPECT_EQ(PostgresType(PostgresTypeId::kInt2).SetSchema(&schema), NANOARROW_OK);
- EXPECT_STREQ(schema.format, "s");
- schema.release(&schema);
-
- ArrowSchemaInit(&schema);
- EXPECT_EQ(PostgresType(PostgresTypeId::kInt4).SetSchema(&schema), NANOARROW_OK);
- EXPECT_STREQ(schema.format, "i");
- schema.release(&schema);
-
- ArrowSchemaInit(&schema);
- EXPECT_EQ(PostgresType(PostgresTypeId::kInt8).SetSchema(&schema), NANOARROW_OK);
- EXPECT_STREQ(schema.format, "l");
- schema.release(&schema);
-
- ArrowSchemaInit(&schema);
- EXPECT_EQ(PostgresType(PostgresTypeId::kFloat4).SetSchema(&schema), NANOARROW_OK);
- EXPECT_STREQ(schema.format, "f");
- schema.release(&schema);
-
- ArrowSchemaInit(&schema);
- EXPECT_EQ(PostgresType(PostgresTypeId::kFloat8).SetSchema(&schema), NANOARROW_OK);
- EXPECT_STREQ(schema.format, "g");
- schema.release(&schema);
-
- ArrowSchemaInit(&schema);
- EXPECT_EQ(PostgresType(PostgresTypeId::kText).SetSchema(&schema), NANOARROW_OK);
- EXPECT_STREQ(schema.format, "u");
- schema.release(&schema);
-
- ArrowSchemaInit(&schema);
- EXPECT_EQ(PostgresType(PostgresTypeId::kBytea).SetSchema(&schema), NANOARROW_OK);
- EXPECT_STREQ(schema.format, "z");
- schema.release(&schema);
-
- ArrowSchemaInit(&schema);
- EXPECT_EQ(PostgresType(PostgresTypeId::kBool).Array().SetSchema(&schema), NANOARROW_OK);
- EXPECT_STREQ(schema.format, "+l");
- EXPECT_STREQ(schema.children[0]->format, "b");
- schema.release(&schema);
-
- ArrowSchemaInit(&schema);
+ nanoarrow::UniqueSchema schema;
+
+ ArrowSchemaInit(schema.get());
+ EXPECT_EQ(PostgresType(PostgresTypeId::kBool).SetSchema(schema.get()), NANOARROW_OK);
+ EXPECT_STREQ(schema->format, "b");
+ schema.reset();
+
+ ArrowSchemaInit(schema.get());
+ EXPECT_EQ(PostgresType(PostgresTypeId::kInt2).SetSchema(schema.get()), NANOARROW_OK);
+ EXPECT_STREQ(schema->format, "s");
+ schema.reset();
+
+ ArrowSchemaInit(schema.get());
+ EXPECT_EQ(PostgresType(PostgresTypeId::kInt4).SetSchema(schema.get()), NANOARROW_OK);
+ EXPECT_STREQ(schema->format, "i");
+ schema.reset();
+
+ ArrowSchemaInit(schema.get());
+ EXPECT_EQ(PostgresType(PostgresTypeId::kInt8).SetSchema(schema.get()), NANOARROW_OK);
+ EXPECT_STREQ(schema->format, "l");
+ schema.reset();
+
+ ArrowSchemaInit(schema.get());
+ EXPECT_EQ(PostgresType(PostgresTypeId::kFloat4).SetSchema(schema.get()), NANOARROW_OK);
+ EXPECT_STREQ(schema->format, "f");
+ schema.reset();
+
+ ArrowSchemaInit(schema.get());
+ EXPECT_EQ(PostgresType(PostgresTypeId::kFloat8).SetSchema(schema.get()), NANOARROW_OK);
+ EXPECT_STREQ(schema->format, "g");
+ schema.reset();
+
+ ArrowSchemaInit(schema.get());
+ EXPECT_EQ(PostgresType(PostgresTypeId::kText).SetSchema(schema.get()), NANOARROW_OK);
+ EXPECT_STREQ(schema->format, "u");
+ schema.reset();
+
+ ArrowSchemaInit(schema.get());
+ EXPECT_EQ(PostgresType(PostgresTypeId::kBytea).SetSchema(schema.get()), NANOARROW_OK);
+ EXPECT_STREQ(schema->format, "z");
+ schema.reset();
+
+ ArrowSchemaInit(schema.get());
+ EXPECT_EQ(PostgresType(PostgresTypeId::kBool).Array().SetSchema(schema.get()),
+ NANOARROW_OK);
+ EXPECT_STREQ(schema->format, "+l");
+ EXPECT_STREQ(schema->children[0]->format, "b");
+ schema.reset();
+
+ ArrowSchemaInit(schema.get());
PostgresType record(PostgresTypeId::kRecord);
record.AppendChild("col1", PostgresType(PostgresTypeId::kBool));
- EXPECT_EQ(record.SetSchema(&schema), NANOARROW_OK);
- EXPECT_STREQ(schema.format, "+s");
- EXPECT_STREQ(schema.children[0]->format, "b");
- schema.release(&schema);
+ EXPECT_EQ(record.SetSchema(schema.get()), NANOARROW_OK);
+ EXPECT_STREQ(schema->format, "+s");
+ EXPECT_STREQ(schema->children[0]->format, "b");
+ schema.reset();
- ArrowSchemaInit(&schema);
+ ArrowSchemaInit(schema.get());
PostgresType unknown(PostgresTypeId::kBrinMinmaxMultiSummary);
- EXPECT_EQ(unknown.WithPgTypeInfo(0, "some_name").SetSchema(&schema), NANOARROW_OK);
- EXPECT_STREQ(schema.format, "z");
+ EXPECT_EQ(unknown.WithPgTypeInfo(0, "some_name").SetSchema(schema.get()), NANOARROW_OK);
+ EXPECT_STREQ(schema->format, "z");
ArrowStringView value = ArrowCharView("<not found>");
- ArrowMetadataGetValue(schema.metadata, ArrowCharView("ADBC:postgresql:typname"),
+ ArrowMetadataGetValue(schema->metadata, ArrowCharView("ADBC:postgresql:typname"),
&value);
EXPECT_EQ(std::string(value.data, value.size_bytes), "some_name");
- schema.release(&schema);
+ schema.reset();
}
TEST(PostgresTypeTest, PostgresTypeFromSchema) {
- ArrowSchema schema;
+ nanoarrow::UniqueSchema schema;
PostgresType type;
MockTypeResolver resolver;
ASSERT_EQ(resolver.Init(), NANOARROW_OK);
- ASSERT_EQ(ArrowSchemaInitFromType(&schema, NANOARROW_TYPE_BOOL), NANOARROW_OK);
- EXPECT_EQ(adbcpq::PostgresType::FromSchema(resolver, &schema, &type, nullptr),
+ ASSERT_EQ(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_BOOL), NANOARROW_OK);
+ EXPECT_EQ(PostgresType::FromSchema(resolver, schema.get(), &type, nullptr),
NANOARROW_OK);
EXPECT_EQ(type.type_id(), PostgresTypeId::kBool);
- schema.release(&schema);
+ schema.reset();
- ASSERT_EQ(ArrowSchemaInitFromType(&schema, NANOARROW_TYPE_INT8), NANOARROW_OK);
- EXPECT_EQ(adbcpq::PostgresType::FromSchema(resolver, &schema, &type, nullptr),
+ ASSERT_EQ(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_INT8), NANOARROW_OK);
+ EXPECT_EQ(PostgresType::FromSchema(resolver, schema.get(), &type, nullptr),
NANOARROW_OK);
EXPECT_EQ(type.type_id(), PostgresTypeId::kInt2);
- schema.release(&schema);
+ schema.reset();
- ASSERT_EQ(ArrowSchemaInitFromType(&schema, NANOARROW_TYPE_UINT8), NANOARROW_OK);
- EXPECT_EQ(adbcpq::PostgresType::FromSchema(resolver, &schema, &type, nullptr),
+ ASSERT_EQ(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_UINT8), NANOARROW_OK);
+ EXPECT_EQ(PostgresType::FromSchema(resolver, schema.get(), &type, nullptr),
NANOARROW_OK);
EXPECT_EQ(type.type_id(), PostgresTypeId::kInt2);
- schema.release(&schema);
+ schema.reset();
- ASSERT_EQ(ArrowSchemaInitFromType(&schema, NANOARROW_TYPE_INT16), NANOARROW_OK);
- EXPECT_EQ(adbcpq::PostgresType::FromSchema(resolver, &schema, &type, nullptr),
+ ASSERT_EQ(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_INT16), NANOARROW_OK);
+ EXPECT_EQ(PostgresType::FromSchema(resolver, schema.get(), &type, nullptr),
NANOARROW_OK);
EXPECT_EQ(type.type_id(), PostgresTypeId::kInt2);
- schema.release(&schema);
+ schema.reset();
- ASSERT_EQ(ArrowSchemaInitFromType(&schema, NANOARROW_TYPE_UINT16), NANOARROW_OK);
- EXPECT_EQ(adbcpq::PostgresType::FromSchema(resolver, &schema, &type, nullptr),
+ ASSERT_EQ(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_UINT16), NANOARROW_OK);
+ EXPECT_EQ(PostgresType::FromSchema(resolver, schema.get(), &type, nullptr),
NANOARROW_OK);
EXPECT_EQ(type.type_id(), PostgresTypeId::kInt4);
- schema.release(&schema);
+ schema.reset();
- ASSERT_EQ(ArrowSchemaInitFromType(&schema, NANOARROW_TYPE_INT32), NANOARROW_OK);
- EXPECT_EQ(adbcpq::PostgresType::FromSchema(resolver, &schema, &type, nullptr),
+ ASSERT_EQ(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_INT32), NANOARROW_OK);
+ EXPECT_EQ(PostgresType::FromSchema(resolver, schema.get(), &type, nullptr),
NANOARROW_OK);
EXPECT_EQ(type.type_id(), PostgresTypeId::kInt4);
- schema.release(&schema);
+ schema.reset();
- ASSERT_EQ(ArrowSchemaInitFromType(&schema, NANOARROW_TYPE_UINT32), NANOARROW_OK);
- EXPECT_EQ(adbcpq::PostgresType::FromSchema(resolver, &schema, &type, nullptr),
+ ASSERT_EQ(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_UINT32), NANOARROW_OK);
+ EXPECT_EQ(PostgresType::FromSchema(resolver, schema.get(), &type, nullptr),
NANOARROW_OK);
EXPECT_EQ(type.type_id(), PostgresTypeId::kInt8);
- schema.release(&schema);
+ schema.reset();
- ASSERT_EQ(ArrowSchemaInitFromType(&schema, NANOARROW_TYPE_INT64), NANOARROW_OK);
- EXPECT_EQ(adbcpq::PostgresType::FromSchema(resolver, &schema, &type, nullptr),
+ ASSERT_EQ(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_INT64), NANOARROW_OK);
+ EXPECT_EQ(PostgresType::FromSchema(resolver, schema.get(), &type, nullptr),
NANOARROW_OK);
EXPECT_EQ(type.type_id(), PostgresTypeId::kInt8);
- schema.release(&schema);
+ schema.reset();
- ASSERT_EQ(ArrowSchemaInitFromType(&schema, NANOARROW_TYPE_FLOAT), NANOARROW_OK);
- EXPECT_EQ(adbcpq::PostgresType::FromSchema(resolver, &schema, &type, nullptr),
+ ASSERT_EQ(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_FLOAT), NANOARROW_OK);
+ EXPECT_EQ(PostgresType::FromSchema(resolver, schema.get(), &type, nullptr),
NANOARROW_OK);
EXPECT_EQ(type.type_id(), PostgresTypeId::kFloat4);
- schema.release(&schema);
+ schema.reset();
- ASSERT_EQ(ArrowSchemaInitFromType(&schema, NANOARROW_TYPE_DOUBLE), NANOARROW_OK);
- EXPECT_EQ(adbcpq::PostgresType::FromSchema(resolver, &schema, &type, nullptr),
+ ASSERT_EQ(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_DOUBLE), NANOARROW_OK);
+ EXPECT_EQ(PostgresType::FromSchema(resolver, schema.get(), &type, nullptr),
NANOARROW_OK);
EXPECT_EQ(type.type_id(), PostgresTypeId::kFloat8);
- schema.release(&schema);
+ schema.reset();
- ASSERT_EQ(ArrowSchemaInitFromType(&schema, NANOARROW_TYPE_BINARY), NANOARROW_OK);
- EXPECT_EQ(adbcpq::PostgresType::FromSchema(resolver, &schema, &type, nullptr),
+ ASSERT_EQ(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_BINARY), NANOARROW_OK);
+ EXPECT_EQ(PostgresType::FromSchema(resolver, schema.get(), &type, nullptr),
NANOARROW_OK);
EXPECT_EQ(type.type_id(), PostgresTypeId::kBytea);
- schema.release(&schema);
+ schema.reset();
- ASSERT_EQ(ArrowSchemaInitFromType(&schema, NANOARROW_TYPE_STRING), NANOARROW_OK);
- EXPECT_EQ(adbcpq::PostgresType::FromSchema(resolver, &schema, &type, nullptr),
+ ASSERT_EQ(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_STRING), NANOARROW_OK);
+ EXPECT_EQ(PostgresType::FromSchema(resolver, schema.get(), &type, nullptr),
NANOARROW_OK);
EXPECT_EQ(type.type_id(), PostgresTypeId::kText);
- schema.release(&schema);
+ schema.reset();
- ArrowSchemaInit(&schema);
- ASSERT_EQ(ArrowSchemaSetType(&schema, NANOARROW_TYPE_LIST), NANOARROW_OK);
- ASSERT_EQ(ArrowSchemaSetType(schema.children[0], NANOARROW_TYPE_BOOL), NANOARROW_OK);
- EXPECT_EQ(adbcpq::PostgresType::FromSchema(resolver, &schema, &type, nullptr),
+ ArrowSchemaInit(schema.get());
+ ASSERT_EQ(ArrowSchemaSetType(schema.get(), NANOARROW_TYPE_LIST), NANOARROW_OK);
+ ASSERT_EQ(ArrowSchemaSetType(schema->children[0], NANOARROW_TYPE_BOOL), NANOARROW_OK);
+ EXPECT_EQ(PostgresType::FromSchema(resolver, schema.get(), &type, nullptr),
NANOARROW_OK);
EXPECT_EQ(type.type_id(), PostgresTypeId::kArray);
EXPECT_EQ(type.child(0).type_id(), PostgresTypeId::kBool);
- schema.release(&schema);
+ schema.reset();
ArrowError error;
- ASSERT_EQ(ArrowSchemaInitFromType(&schema, NANOARROW_TYPE_INTERVAL_MONTH_DAY_NANO),
+ ASSERT_EQ(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_INTERVAL_MONTH_DAY_NANO),
NANOARROW_OK);
- EXPECT_EQ(adbcpq::PostgresType::FromSchema(resolver, &schema, &type, &error), ENOTSUP);
+ EXPECT_EQ(PostgresType::FromSchema(resolver, schema.get(), &type, &error), ENOTSUP);
EXPECT_STREQ(error.message,
"Can't map Arrow type 'interval_month_day_nano' to Postgres type");
- schema.release(&schema);
+ schema.reset();
}
TEST(PostgresTypeTest, PostgresTypeResolver) {
@@ -299,16 +298,8 @@ TEST(PostgresTypeTest, PostgresTypeResolver) {
EXPECT_EQ(resolver.Find(123, &type, &error), EINVAL);
EXPECT_STREQ(ArrowErrorMessage(&error), "Postgres type with oid 123 not found");
- // Check error for unsupported type_id name
- item.oid = 123;
- item.typname = "invalid";
- item.typreceive = "invalid_recv";
- EXPECT_EQ(resolver.Insert(item, &error), ENOTSUP);
- EXPECT_STREQ(
- ArrowErrorMessage(&error),
- "Base type not found for type 'invalid' with receive function 'invalid_recv'");
-
// Check error for Array with unknown child
+ item.oid = 123;
item.typname = "some_array";
item.typreceive = "array_recv";
item.child_oid = 1234;
@@ -403,3 +394,5 @@ TEST(PostgresTypeTest, PostgresTypeResolveRecord) {
EXPECT_EQ(type.child(1).field_name(), "text_col");
EXPECT_EQ(type.child(1).type_id(), PostgresTypeId::kText);
}
+
+} // namespace adbcpq
diff --git a/c/driver/postgresql/postgres_util.h b/c/driver/postgresql/postgres_util.h
index d03660b..df45bdd 100644
--- a/c/driver/postgresql/postgres_util.h
+++ b/c/driver/postgresql/postgres_util.h
@@ -17,6 +17,8 @@
#pragma once
+#include <cstring>
+
#ifdef _WIN32
#include <winsock2.h>
#else
diff --git a/c/driver/postgresql/statement.cc b/c/driver/postgresql/statement.cc
index 01d30d6..494c736 100644
--- a/c/driver/postgresql/statement.cc
+++ b/c/driver/postgresql/statement.cc
@@ -28,9 +28,10 @@
#include <adbc.h>
#include <libpq-fe.h>
-#include <nanoarrow/nanoarrow.h>
+#include <nanoarrow/nanoarrow.hpp>
#include "connection.h"
+#include "postgres_copy_reader.h"
#include "postgres_type.h"
#include "postgres_util.h"
#include "utils.h"
@@ -38,13 +39,8 @@
namespace adbcpq {
namespace {
-/// The header that comes at the start of a binary COPY stream
-constexpr std::array<char, 11> kPgCopyBinarySignature = {
- 'P', 'G', 'C', 'O', 'P', 'Y', '\n', '\377', '\r', '\n', '\0'};
/// The flag indicating to PostgreSQL that we want binary-format values.
constexpr int kPgBinaryFormat = 1;
-// A negative field length indicates a null.
-constexpr int32_t kNullFieldLength = -1;
/// One-value ArrowArrayStream used to unify the implementations of Bind
struct OneValueStream {
@@ -108,13 +104,14 @@ struct Handle {
Resource* operator->() { return &value; }
};
-/// Build an Arrow schema from a PostgreSQL result set
-AdbcStatusCode InferSchema(const PostgresTypeResolver& type_resolver, PGresult* result,
- struct ArrowSchema* out, struct AdbcError* error) {
+/// Build an PostgresType object from a PGresult*
+AdbcStatusCode ResolvePostgresType(const PostgresTypeResolver& type_resolver,
+ PGresult* result, PostgresType* out,
+ struct AdbcError* error) {
ArrowError na_error;
const int num_fields = PQnfields(result);
- ArrowSchemaInit(out);
- CHECK_NA(INTERNAL, ArrowSchemaSetTypeStruct(out, num_fields), error);
+ PostgresType root_type(PostgresTypeId::kRecord);
+
for (int i = 0; i < num_fields; i++) {
const Oid pg_oid = PQftype(result, i);
PostgresType pg_type;
@@ -124,10 +121,10 @@ AdbcStatusCode InferSchema(const PostgresTypeResolver& type_resolver, PGresult*
return ADBC_STATUS_NOT_IMPLEMENTED;
}
- CHECK_NA(INTERNAL,
- pg_type.WithFieldName(PQfname(result, i)).SetSchema(out->children[i]),
- error);
+ root_type.AppendChild(PQfname(result, i), pg_type);
}
+
+ *out = root_type;
return ADBC_STATUS_OK;
}
@@ -343,14 +340,17 @@ struct BindStream {
} // namespace
int TupleReader::GetSchema(struct ArrowSchema* out) {
- if (!schema_.release) {
- last_error_ = "[libpq] Result set was already consumed or freed";
+ int na_res = copy_reader_->GetSchema(out);
+ if (out->release == nullptr) {
+ StringBuilderAppend(&error_builder_,
+ "[libpq] Result set was already consumed or freed");
return EINVAL;
+ } else if (na_res != NANOARROW_OK) {
+ // e.g., Can't allocate memory
+ StringBuilderAppend(&error_builder_, "[libpq] Error copying schema");
}
- std::memset(out, 0, sizeof(*out));
- NANOARROW_RETURN_NOT_OK(ArrowSchemaDeepCopy(&schema_, out));
- return 0;
+ return na_res;
}
int TupleReader::GetNext(struct ArrowArray* out) {
@@ -363,101 +363,82 @@ int TupleReader::GetNext(struct ArrowArray* out) {
PQclear(result_);
result_ = nullptr;
+ // Clear the error builder
+ error_builder_.size = 0;
+
struct ArrowError error;
- // TODO: consistently release out on error (use another trampoline?)
- int na_res = ArrowArrayInitFromSchema(out, &schema_, &error);
- if (na_res != 0) {
- last_error_ = "[libpq] Failed to init output array: " + std::to_string(na_res) +
- std::strerror(na_res) + ": " + error.message;
- if (out->release) out->release(out);
- return na_res;
+ error.message[0] = '\0';
+ struct ArrowBufferView data;
+ data.data.data = nullptr;
+ data.size_bytes = 0;
+
+ // Fetch + parse the header
+ int get_copy_res = PQgetCopyData(conn_, &pgbuf_, /*async=*/0);
+ if (get_copy_res == -2) {
+ StringBuilderAppend(&error_builder_, "[libpq] Fetch header failed: %s",
+ PQerrorMessage(conn_));
+ return EIO;
}
- std::vector<ArrowSchemaView> fields(schema_.n_children);
- NANOARROW_RETURN_NOT_OK(ArrowArrayStartAppending(out));
- for (int col = 0; col < schema_.n_children; col++) {
- na_res = ArrowSchemaViewInit(&fields[col], schema_.children[col], &error);
- if (na_res != 0) {
- last_error_ = "[libpq] Failed to init schema view: " + std::to_string(na_res) +
- std::strerror(na_res) + ": " + error.message;
- if (out->release) out->release(out);
- return na_res;
- }
+ data.size_bytes = get_copy_res;
+ data.data.as_char = pgbuf_;
+ int na_res = copy_reader_->ReadHeader(&data, &error);
+ if (na_res != NANOARROW_OK) {
+ StringBuilderAppend(&error_builder_, "[libpq] ReadHeader failed: %s", error.message);
+ return na_res;
}
- // TODO: we need to always PQgetResult
-
- char* buf = nullptr;
- int buf_size = 0;
-
- // Get the header
- {
- constexpr size_t kPqHeaderLength =
- kPgCopyBinarySignature.size() + sizeof(uint32_t) + sizeof(uint32_t);
- // https://www.postgresql.org/docs/14/sql-copy.html#id-1.9.3.55.9.4.5
- const int size = PQgetCopyData(conn_, &pgbuf_, /*async=*/0);
- if (size < static_cast<int>(kPqHeaderLength)) {
- return EIO;
- } else if (std::strcmp(pgbuf_, kPgCopyBinarySignature.data()) != 0) {
- return EIO;
- }
- buf = pgbuf_ + kPgCopyBinarySignature.size();
-
- uint32_t flags = LoadNetworkUInt32(buf);
- buf += sizeof(uint32_t);
- if (flags != 0) {
- return EIO;
+ int64_t row_id = 0;
+ do {
+ // Parse the result (the header AND the first row are included in the first
+ // call to PQgetCopyData())
+ 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;
}
- // XXX: is this signed or unsigned? not stated by the docs
- uint32_t extension_length = LoadNetworkUInt32(buf);
- buf += sizeof(uint32_t) + extension_length;
+ row_id++;
- buf_size = size - (kPqHeaderLength + extension_length);
- }
-
- // Append each row
- int result_code = 0;
- int64_t num_rows = 0;
- last_error_.clear();
- do {
- result_code = AppendNext(fields.data(), buf, buf_size, &num_rows, out);
+ // Fetch + check
PQfreemem(pgbuf_);
- pgbuf_ = buf = nullptr;
- if (result_code != 0) break;
-
- buf_size = PQgetCopyData(conn_, &pgbuf_, /*async=*/0);
- if (buf_size < 0) {
- pgbuf_ = buf = nullptr;
+ pgbuf_ = nullptr;
+ get_copy_res = PQgetCopyData(conn_, &pgbuf_, /*async=*/0);
+ if (get_copy_res == -2) {
+ 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
break;
}
- buf = pgbuf_;
+
+ data.size_bytes = get_copy_res;
+ data.data.as_char = pgbuf_;
} while (true);
- // Finish the result array
- for (int col = 0; col < schema_.n_children; col++) {
- out->children[col]->length = num_rows;
- }
- out->length = num_rows;
- na_res = ArrowArrayFinishBuildingDefault(out, &error);
- if (na_res != 0) {
- result_code = na_res;
- if (!last_error_.empty()) last_error_ += '\n';
- last_error_ += "[libpq] Failed to build result array" + std::string(error.message);
+ na_res = copy_reader_->GetArray(out, &error);
+ if (na_res != NANOARROW_OK) {
+ StringBuilderAppend(&error_builder_, "[libpq] Failed to build result array: %s",
+ error.message);
+ return na_res;
}
// Check the server-side response
result_ = PQgetResult(conn_);
const int pq_status = PQresultStatus(result_);
if (pq_status != PGRES_COMMAND_OK) {
- if (!last_error_.empty()) last_error_ += '\n';
- last_error_ += "[libpq] Query failed: (" + std::to_string(pq_status) + ") " +
- PQresultErrorMessage(result_);
- result_code = EIO;
+ StringBuilderAppend(&error_builder_, "[libpq] Query failed [%d]: %s", pq_status,
+ PQresultErrorMessage(result_));
+ return EIO;
}
+
PQclear(result_);
result_ = nullptr;
- return result_code;
+ return NANOARROW_OK;
}
void TupleReader::Release() {
@@ -465,9 +446,6 @@ void TupleReader::Release() {
PQclear(result_);
result_ = nullptr;
}
- if (schema_.release) {
- schema_.release(&schema_);
- }
if (pgbuf_) {
PQfreemem(pgbuf_);
pgbuf_ = nullptr;
@@ -482,117 +460,6 @@ void TupleReader::ExportTo(struct ArrowArrayStream* stream) {
stream->private_data = this;
}
-int TupleReader::AppendNext(struct ArrowSchemaView* fields, const char* buf, int buf_size,
- int64_t* row_count, struct ArrowArray* out) {
- // https://www.postgresql.org/docs/14/sql-copy.html#id-1.9.3.55.9.4.6
- // TODO: DCHECK_GE(buf_size, 2) << "Buffer too short to contain field count";
-
- int16_t field_count = 0;
- std::memcpy(&field_count, buf, sizeof(int16_t));
- buf += sizeof(int16_t);
- field_count = ntohs(field_count);
-
- if (field_count == -1) {
- // end-of-stream
- return 0;
- } else if (field_count != schema_.n_children) {
- last_error_ = "[libpq] Expected " + std::to_string(schema_.n_children) +
- " fields but found " + std::to_string(field_count);
- return EIO;
- }
-
- for (int col = 0; col < schema_.n_children; col++) {
- int32_t field_length = LoadNetworkInt32(buf);
- buf += sizeof(int32_t);
-
- // TODO: set error message here
- if (field_length != kNullFieldLength) {
- NANOARROW_RETURN_NOT_OK(
- AppendValue(fields, buf, col, *row_count, field_length, out));
- buf += field_length;
- } else {
- NANOARROW_RETURN_NOT_OK(ArrowArrayAppendNull(out->children[col], 1));
- }
- }
- (*row_count)++;
- return 0;
-}
-
-int TupleReader::AppendValue(struct ArrowSchemaView* fields, const char* buf, int col,
- int64_t row_count, int32_t field_length,
- struct ArrowArray* out) {
- switch (fields[col].type) {
- case NANOARROW_TYPE_BOOL: {
- uint8_t raw_value = buf[0];
- buf += 1;
-
- if (raw_value != 0 && raw_value != 1) {
- last_error_ = "[libpq] Column #" + std::to_string(col + 1) + " (\"" +
- schema_.children[col]->name + "\"): invalid BOOL value " +
- std::to_string(raw_value);
- return EINVAL;
- }
- NANOARROW_RETURN_NOT_OK(ArrowArrayAppendInt(out->children[col], raw_value));
- break;
- }
- case NANOARROW_TYPE_DOUBLE: {
- static_assert(sizeof(double) == sizeof(uint64_t),
- "Float is not same size as uint64_t");
- uint64_t raw_value = LoadNetworkUInt64(buf);
- buf += sizeof(uint64_t);
- double value = 0.0;
- std::memcpy(&value, &raw_value, sizeof(double));
- NANOARROW_RETURN_NOT_OK(ArrowArrayAppendDouble(out->children[col], value));
- break;
- }
- case NANOARROW_TYPE_FLOAT: {
- static_assert(sizeof(float) == sizeof(uint32_t),
- "Float is not same size as uint32_t");
- uint32_t raw_value = LoadNetworkUInt32(buf);
- buf += sizeof(uint32_t);
- float value = 0.0;
- std::memcpy(&value, &raw_value, sizeof(float));
- NANOARROW_RETURN_NOT_OK(ArrowArrayAppendDouble(out->children[col], value));
- break;
- }
- case NANOARROW_TYPE_INT16: {
- int32_t value = LoadNetworkInt16(buf);
- buf += sizeof(int32_t);
- NANOARROW_RETURN_NOT_OK(ArrowArrayAppendInt(out->children[col], value));
- break;
- }
- case NANOARROW_TYPE_INT32: {
- int32_t value = LoadNetworkInt32(buf);
- buf += sizeof(int32_t);
- NANOARROW_RETURN_NOT_OK(ArrowArrayAppendInt(out->children[col], value));
- break;
- }
- case NANOARROW_TYPE_INT64: {
- int64_t value = LoadNetworkInt64(buf);
- buf += sizeof(int64_t);
- NANOARROW_RETURN_NOT_OK(ArrowArrayAppendInt(out->children[col], value));
- break;
- }
- case NANOARROW_TYPE_BINARY: {
- NANOARROW_RETURN_NOT_OK(
- ArrowArrayAppendBytes(out->children[col], {buf, field_length}));
- break;
- }
- case NANOARROW_TYPE_STRING: {
- // textsend() in varlena.c
- NANOARROW_RETURN_NOT_OK(
- ArrowArrayAppendString(out->children[col], {buf, field_length}));
- break;
- }
- default:
- last_error_ = "[libpq] Column #" + std::to_string(col + 1) + " (\"" +
- schema_.children[col]->name + "\") has unsupported type " +
- std::to_string(fields[col].type);
- return ENOTSUP;
- }
- return 0;
-}
-
int TupleReader::GetSchemaTrampoline(struct ArrowArrayStream* self,
struct ArrowSchema* out) {
if (!self || !self->private_data) return EINVAL;
@@ -797,9 +664,33 @@ AdbcStatusCode PostgresStatement::ExecuteQuery(struct ArrowArrayStream* stream,
PQclear(result);
return ADBC_STATUS_IO;
}
- AdbcStatusCode status = InferSchema(*type_resolver_, result, &reader_.schema_, error);
+
+ // Resolve the information from the PGresult into a PostgresType
+ PostgresType root_type;
+ AdbcStatusCode status =
+ ResolvePostgresType(*type_resolver_, result, &root_type, error);
PQclear(result);
if (status != ADBC_STATUS_OK) return status;
+
+ // Initialize the copy reader and infer the output schema (i.e., error for
+ // unsupported types before issuing the COPY query)
+ reader_.copy_reader_.reset(new PostgresCopyStreamReader());
+ reader_.copy_reader_->Init(root_type);
+ struct ArrowError na_error;
+ int na_res = reader_.copy_reader_->InferOutputSchema(&na_error);
+ if (na_res != NANOARROW_OK) {
+ SetError(error, "[libpq] Failed to infer output schema: %s", na_error.message);
+ return na_res;
+ }
+
+ // This resolves the reader specific to each PostgresType -> ArrowSchema
+ // conversion. It is unlikely that this will fail given that we have just
+ // inferred these conversions ourselves.
+ na_res = reader_.copy_reader_->InitFieldReaders(&na_error);
+ if (na_res != NANOARROW_OK) {
+ SetError(error, "[libpq] Failed to initialize field readers: %s", na_error.message);
+ return na_res;
+ }
}
// 2. Execute the query with COPY to get binary tuples
diff --git a/c/driver/postgresql/statement.h b/c/driver/postgresql/statement.h
index 5dfe04f..5629323 100644
--- a/c/driver/postgresql/statement.h
+++ b/c/driver/postgresql/statement.h
@@ -26,7 +26,9 @@
#include <libpq-fe.h>
#include <nanoarrow/nanoarrow.h>
+#include "postgres_copy_reader.h"
#include "postgres_type.h"
+#include "utils.h"
namespace adbcpq {
class PostgresConnection;
@@ -35,19 +37,21 @@ class PostgresStatement;
/// \brief An ArrowArrayStream that reads tuples from a PGresult.
class TupleReader final {
public:
- TupleReader(PGconn* conn) : conn_(conn), result_(nullptr), pgbuf_(nullptr) {
- std::memset(&schema_, 0, sizeof(schema_));
+ TupleReader(PGconn* conn)
+ : conn_(conn), result_(nullptr), pgbuf_(nullptr), copy_reader_(nullptr) {
+ StringBuilderInit(&error_builder_, 0);
}
int GetSchema(struct ArrowSchema* out);
int GetNext(struct ArrowArray* out);
- const char* last_error() const { return last_error_.c_str(); }
+ const char* last_error() const {
+ if (error_builder_.size > 0) {
+ return error_builder_.buffer;
+ } else {
+ return nullptr;
+ }
+ }
void Release();
-
- int AppendNext(struct ArrowSchemaView* fields, const char* buf, int buf_size,
- int64_t* row_count, struct ArrowArray* out);
- int AppendValue(struct ArrowSchemaView* fields, const char* buf, int col,
- int64_t row_count, int32_t field_length, struct ArrowArray* out);
void ExportTo(struct ArrowArrayStream* stream);
private:
@@ -61,8 +65,8 @@ class TupleReader final {
PGconn* conn_;
PGresult* result_;
char* pgbuf_;
- struct ArrowSchema schema_;
- std::string last_error_;
+ struct StringBuilder error_builder_;
+ std::unique_ptr<PostgresCopyStreamReader> copy_reader_;
};
class PostgresStatement {
diff --git a/r/adbcpostgresql/bootstrap.R b/r/adbcpostgresql/bootstrap.R
index fcfce07..64687a8 100644
--- a/r/adbcpostgresql/bootstrap.R
+++ b/r/adbcpostgresql/bootstrap.R
@@ -21,6 +21,7 @@ files_to_vendor <- c(
"../../adbc.h",
"../../c/driver/postgresql/postgres_util.h",
"../../c/driver/postgresql/postgres_type.h",
+ "../../c/driver/postgresql/postgres_copy_reader.h",
"../../c/driver/postgresql/statement.h",
"../../c/driver/postgresql/statement.cc",
"../../c/driver/postgresql/connection.h",
diff --git a/r/adbcpostgresql/src/.gitignore b/r/adbcpostgresql/src/.gitignore
index 1e652d7..ee58a5e 100644
--- a/r/adbcpostgresql/src/.gitignore
+++ b/r/adbcpostgresql/src/.gitignore
@@ -27,5 +27,8 @@ postgresql.cc
statement.h
statement.cc
postgres_type.h
-util.h
+postgres_copy_reader.h
+postgres_util.h
+utils.h
+utils.c
Makevars