Merge pull request #136 from hafenkran/fix_appendrows_with_over_10mb

Soft-limit AppendRows requests and fix post-create scans for BigQuery extension

Author: hafenkran

src/bigquery_arrow_scan.cpp

@@ -285,10 +285,17 @@ void BigqueryArrowScanFunction::BigqueryArrowScanExecute(ClientContext &ctx,
         NumericCast<idx_t>(state.chunk->arrow_array.length) - state.chunk_offset);
     data.lines_read += output_size;
 
+    auto ensure_column_id_coverage = [&](DataChunk &chunk) {
+        if (!state.column_ids.empty() && state.column_ids.size() != chunk.ColumnCount()) {
+            state.column_ids.clear();
+        }
+    };
+
     if (gstate.CanRemoveFilterColumns()) {
         state.all_columns.Reset();
         state.all_columns.SetCardinality(output_size);
 
+        ensure_column_id_coverage(state.all_columns);
         ArrowTableFunction::ArrowToDuckDB(state,
                                           data.arrow_table.GetColumns(),
                                           state.all_columns,
@@ -358,13 +365,15 @@ void BigqueryArrowScanFunction::BigqueryArrowScanExecute(ClientContext &ctx,
         bool do_cast = data.requires_cast || geometry_cast_needed;
         if (!do_cast) {
             // Direct write to output
+            ensure_column_id_coverage(output);
             ArrowTableFunction::ArrowToDuckDB(state,
                                               data.arrow_table.GetColumns(),
                                               output,
                                               data.lines_read - output_size);
         } else {
             state.all_columns.Reset();
             state.all_columns.SetCardinality(output_size);
+            ensure_column_id_coverage(state.all_columns);
             ArrowTableFunction::ArrowToDuckDB(state,
                                               data.arrow_table.GetColumns(),
                                               state.all_columns,

src/bigquery_proto_writer.cpp

@@ -238,11 +238,6 @@ void BigqueryProtoWriter::WriteChunk(DataChunk &chunk, const std::map<std::strin
         throw BinderException("Cannot get message prototype from message descriptor");
     }
 
-    // Create the append request
-    google::cloud::bigquery::storage::v1::AppendRowsRequest request;
-    request.set_write_stream(write_stream.name());
-    msg_descriptor->CopyTo(request.mutable_proto_rows()->mutable_writer_schema()->mutable_proto_descriptor());
-
     vector<idx_t> column_indexes;
     if (column_idxs.empty()) {
         column_indexes.resize(chunk.ColumnCount());
@@ -255,7 +250,18 @@ void BigqueryProtoWriter::WriteChunk(DataChunk &chunk, const std::map<std::strin
         }
     }
 
+    auto create_request = [this]() {
+        google::cloud::bigquery::storage::v1::AppendRowsRequest new_request;
+        new_request.set_write_stream(write_stream.name());
+        msg_descriptor->CopyTo(new_request.mutable_proto_rows()->mutable_writer_schema()->mutable_proto_descriptor());
+        return new_request;
+    };
+
+    auto request = create_request();
     auto *rows = request.mutable_proto_rows()->mutable_rows();
+    idx_t rows_in_batch = 0;
+    size_t current_request_bytes = request.ByteSizeLong();
+
     for (idx_t i = 0; i < chunk.size(); i++) {
         google::protobuf::Message *msg = msg_prototype->New();
         const google::protobuf::Reflection *reflection = msg->GetReflection();
@@ -290,10 +296,40 @@ void BigqueryProtoWriter::WriteChunk(DataChunk &chunk, const std::map<std::strin
         if (!msg->SerializeToString(&serialized_msg)) {
             throw std::runtime_error("Failed to serialize message");
         }
+        auto estimated_size_increase = serialized_msg.size() + APPEND_ROWS_ROW_OVERHEAD;
+
+        if (rows_in_batch > 0 && current_request_bytes + estimated_size_increase > DEFAULT_APPEND_ROWS_SOFT_LIMIT) {
+            SendAppendRequest(request);
+            request = create_request();
+            rows = request.mutable_proto_rows()->mutable_rows();
+            rows_in_batch = 0;
+            current_request_bytes = request.ByteSizeLong();
+        }
+
         rows->add_serialized_rows(serialized_msg);
+        rows_in_batch++;
+        current_request_bytes += estimated_size_increase;
+
+        if (current_request_bytes >= DEFAULT_APPEND_ROWS_SOFT_LIMIT) {
+            SendAppendRequest(request);
+            request = create_request();
+            rows = request.mutable_proto_rows()->mutable_rows();
+            rows_in_batch = 0;
+            current_request_bytes = request.ByteSizeLong();
+        }
         delete msg;
     }
 
+    if (rows_in_batch > 0) {
+        SendAppendRequest(request);
+    }
+}
+
+void BigqueryProtoWriter::SendAppendRequest(const google::cloud::bigquery::storage::v1::AppendRowsRequest &request) {
+    if (!request.has_proto_rows() || request.proto_rows().rows().serialized_rows_size() == 0) {
+        return;
+    }
+
     int max_retries = 100;
     for (int attempt = 0; attempt < max_retries; attempt++) {
         auto handle_broken_stream = [this](char const *where) {
@@ -322,7 +358,6 @@ void BigqueryProtoWriter::WriteChunk(DataChunk &chunk, const std::map<std::strin
             }
         }
 
-        // GET THE RESPONSE AND ERROR HANDLING
         auto response = grpc_stream->Read().get();
         if (!response) {
             if (attempt < max_retries - 1) {

src/include/bigquery_proto_writer.hpp

@@ -39,6 +39,11 @@ class BigqueryProtoWriter {
     void Finalize();
 
 private:
+    static constexpr idx_t DEFAULT_APPEND_ROWS_SOFT_LIMIT = 9 * 1024 * 1024;
+    static constexpr idx_t APPEND_ROWS_ROW_OVERHEAD = 32;
+
+    void SendAppendRequest(const google::cloud::bigquery::storage::v1::AppendRowsRequest &request);
+
     string table_string;
 
     google::protobuf::DescriptorPool pool;
@@ -51,7 +56,6 @@ class BigqueryProtoWriter {
     std::unique_ptr<google::cloud::AsyncStreamingReadWriteRpc<google::cloud::bigquery::storage::v1::AppendRowsRequest,
                                                               google::cloud::bigquery::storage::v1::AppendRowsResponse>>
         grpc_stream;
-    // google::cloud::bigquery::storage::v1::AppendRowsRequest append_request;
 };
 
 } // namespace bigquery

test/sql/storage/attach_large_insert.test

@@ -0,0 +1,41 @@
+# name: test/sql/storage/attach_large_insert.test
+# description: Ensure large string inserts (single chunk > 10 MiB) succeed by batching AppendRows requests
+# group: [storage]
+
+require bigquery
+
+require-env BQ_TEST_PROJECT
+
+require-env BQ_TEST_DATASET
+
+statement ok
+ATTACH 'project=${BQ_TEST_PROJECT} dataset=${BQ_TEST_DATASET}' AS bq (TYPE bigquery);
+
+statement ok
+CREATE OR REPLACE TABLE bq.${BQ_TEST_DATASET}.large_json_insert (
+    id INTEGER,
+    payload_a STRING,
+    payload_b STRING,
+    payload_c STRING,
+    payload_d STRING,
+    payload_e STRING
+);
+
+statement ok
+INSERT INTO bq.${BQ_TEST_DATASET}.large_json_insert
+SELECT
+    i,
+    '{"payload_a":"' || repeat('A', 100000) || '"}',
+    '{"payload_b":"' || repeat('B', 100000) || '"}',
+    '{"payload_c":"' || repeat('C', 100000) || '"}',
+    '{"payload_d":"' || repeat('D', 100000) || '"}',
+    '{"payload_e":"' || repeat('E', 100000) || '"}'
+FROM range(100) tbl(i);
+
+query I
+SELECT COUNT(*) FROM bq.${BQ_TEST_DATASET}.large_json_insert;
+----
+100
+
+statement ok
+DROP TABLE bq.${BQ_TEST_DATASET}.large_json_insert;