refactor(c/driver/postgresql): Refactor COPY Writer NULL handling (#1175

)
apache · Oct 6, 2023 · 13fc595 · 13fc595
1 parent 27903bd
commit 13fc595
Showing 1 changed file with 11 additions and 30 deletions.
diff --git a/c/driver/postgresql/postgres_copy_reader.h b/c/driver/postgresql/postgres_copy_reader.h
@@ -1125,7 +1125,13 @@ class PostgresCopyFieldTupleWriter : public PostgresCopyFieldWriter {
     NANOARROW_RETURN_NOT_OK(WriteChecked<int16_t>(buffer, n_fields, error));
 
     for (int16_t i = 0; i < n_fields; i++) {
-      children_[i]->Write(buffer, index, error);
+      const int8_t is_null = ArrowArrayViewIsNull(array_view_->children[i], index);
+      if (is_null) {
+        constexpr int32_t field_size_bytes = -1;
+        NANOARROW_RETURN_NOT_OK(WriteChecked<int32_t>(buffer, field_size_bytes, error));
+      } else {
+        children_[i]->Write(buffer, index, error);
+      }
     }
 
     return NANOARROW_OK;
@@ -1138,13 +1144,8 @@ class PostgresCopyFieldTupleWriter : public PostgresCopyFieldWriter {
 class PostgresCopyBooleanFieldWriter : public PostgresCopyFieldWriter {
  public:
   ArrowErrorCode Write(ArrowBuffer* buffer, int64_t index, ArrowError* error) override {
-    const int8_t is_null = ArrowArrayViewIsNull(array_view_, index);
-    const int32_t field_size_bytes = is_null ? -1 : 1;
+    constexpr int32_t field_size_bytes = 1;
     NANOARROW_RETURN_NOT_OK(WriteChecked<int32_t>(buffer, field_size_bytes, error));
-    if (is_null) {
-      return ADBC_STATUS_OK;
-    }
-
     const int8_t value =
         static_cast<int8_t>(ArrowArrayViewGetIntUnsafe(array_view_, index));
     NANOARROW_RETURN_NOT_OK(WriteChecked<int8_t>(buffer, value, error));
@@ -1157,13 +1158,8 @@ template <typename T, T kOffset = 0>
 class PostgresCopyNetworkEndianFieldWriter : public PostgresCopyFieldWriter {
  public:
   ArrowErrorCode Write(ArrowBuffer* buffer, int64_t index, ArrowError* error) override {
-    const int8_t is_null = ArrowArrayViewIsNull(array_view_, index);
-    const int32_t field_size_bytes = is_null ? -1 : sizeof(T);
+    constexpr int32_t field_size_bytes = sizeof(T);
     NANOARROW_RETURN_NOT_OK(WriteChecked<int32_t>(buffer, field_size_bytes, error));
-    if (is_null) {
-      return ADBC_STATUS_OK;
-    }
-
     const T value =
         static_cast<T>(ArrowArrayViewGetIntUnsafe(array_view_, index)) - kOffset;
     NANOARROW_RETURN_NOT_OK(WriteChecked<T>(buffer, value, error));
@@ -1175,12 +1171,8 @@ class PostgresCopyNetworkEndianFieldWriter : public PostgresCopyFieldWriter {
 class PostgresCopyFloatFieldWriter : public PostgresCopyFieldWriter {
  public:
   ArrowErrorCode Write(ArrowBuffer* buffer, int64_t index, ArrowError* error) override {
-    const int8_t is_null = ArrowArrayViewIsNull(array_view_, index);
-    const int32_t field_size_bytes = is_null ? -1 : sizeof(uint32_t);
+    constexpr int32_t field_size_bytes = sizeof(uint32_t);
     NANOARROW_RETURN_NOT_OK(WriteChecked<int32_t>(buffer, field_size_bytes, error));
-    if (is_null) {
-      return ADBC_STATUS_OK;
-    }
 
     uint32_t value;
     float raw_value = ArrowArrayViewGetDoubleUnsafe(array_view_, index);
@@ -1194,12 +1186,8 @@ class PostgresCopyFloatFieldWriter : public PostgresCopyFieldWriter {
 class PostgresCopyDoubleFieldWriter : public PostgresCopyFieldWriter {
  public:
   ArrowErrorCode Write(ArrowBuffer* buffer, int64_t index, ArrowError* error) override {
-    const int8_t is_null = ArrowArrayViewIsNull(array_view_, index);
-    const int32_t field_size_bytes = is_null ? -1 : sizeof(uint64_t);
+    constexpr int32_t field_size_bytes = sizeof(uint64_t);
     NANOARROW_RETURN_NOT_OK(WriteChecked<int32_t>(buffer, field_size_bytes, error));
-    if (is_null) {
-      return ADBC_STATUS_OK;
-    }
 
     uint64_t value;
     double raw_value = ArrowArrayViewGetDoubleUnsafe(array_view_, index);
@@ -1213,13 +1201,6 @@ class PostgresCopyDoubleFieldWriter : public PostgresCopyFieldWriter {
 class PostgresCopyBinaryFieldWriter : public PostgresCopyFieldWriter {
  public:
   ArrowErrorCode Write(ArrowBuffer* buffer, int64_t index, ArrowError* error) override {
-    const int8_t is_null = ArrowArrayViewIsNull(array_view_, index);
-    if (is_null) {
-      constexpr int32_t field_size_bytes = -1;
-      NANOARROW_RETURN_NOT_OK(WriteChecked<int32_t>(buffer, field_size_bytes, error));
-      return ADBC_STATUS_OK;
-    }
-
     struct ArrowStringView string_view =
         ArrowArrayViewGetStringUnsafe(array_view_, index);
     NANOARROW_RETURN_NOT_OK(WriteChecked<int32_t>(buffer, string_view.size_bytes, error));