diff --git a/c/driver/postgresql/postgres_copy_reader.h b/c/driver/postgresql/postgres_copy_reader.h index f4faaaef56..bef0fcde8d 100644 --- a/c/driver/postgresql/postgres_copy_reader.h +++ b/c/driver/postgresql/postgres_copy_reader.h @@ -1237,6 +1237,7 @@ static inline ArrowErrorCode MakeCopyFieldWriter(const enum ArrowType arrow_type case NANOARROW_TYPE_BOOL: *out = new PostgresCopyBooleanFieldWriter(); return NANOARROW_OK; + case NANOARROW_TYPE_INT8: case NANOARROW_TYPE_INT16: *out = new PostgresCopyNetworkEndianFieldWriter(); return NANOARROW_OK; diff --git a/c/driver/postgresql/postgres_copy_reader_test.cc b/c/driver/postgresql/postgres_copy_reader_test.cc index f2db5b7698..76b56ef514 100644 --- a/c/driver/postgresql/postgres_copy_reader_test.cc +++ b/c/driver/postgresql/postgres_copy_reader_test.cc @@ -197,6 +197,30 @@ TEST(PostgresCopyUtilsTest, PostgresCopyReadSmallInt) { ASSERT_EQ(data_buffer[4], 0); } +TEST(PostgresCopyUtilsTest, PostgresCopyWriteInt8) { + adbc_validation::Handle schema; + adbc_validation::Handle array; + struct ArrowError na_error; + ASSERT_EQ(adbc_validation::MakeSchema(&schema.value, {{"col", NANOARROW_TYPE_INT8}}), + ADBC_STATUS_OK); + ASSERT_EQ(adbc_validation::MakeBatch(&schema.value, &array.value, &na_error, + {-123, -1, 1, 123, std::nullopt}), + ADBC_STATUS_OK); + + PostgresCopyStreamWriteTester tester; + ASSERT_EQ(tester.Init(&schema.value, &array.value), NANOARROW_OK); + ASSERT_EQ(tester.WriteAll(nullptr), ENODATA); + + const struct ArrowBuffer buf = tester.WriteBuffer(); + // The last 2 bytes of a message can be transmitted via PQputCopyData + // so no need to test those bytes from the Writer + constexpr size_t buf_size = sizeof(kTestPgCopySmallInt) - 2; + ASSERT_EQ(buf.size_bytes, buf_size); + for (size_t i = 0; i < buf_size; i++) { + ASSERT_EQ(buf.data[i], kTestPgCopySmallInt[i]); + } +} + TEST(PostgresCopyUtilsTest, PostgresCopyWriteInt16) { adbc_validation::Handle schema; adbc_validation::Handle array;