perf(bigquery): [T1-9] replace LinkedBlockingDeque with ArrayBlockingQueue and batch-drain in ConnectionImpl

java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryResultImpl.java

@@ -26,6 +26,7 @@
 import java.sql.Timestamp;
 import java.time.LocalDateTime;
 import java.time.LocalTime;
+import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -114,6 +115,8 @@ public ResultSet getResultSet() {
   }
 
   private class BigQueryResultSet extends AbstractJdbcResultSet {
+    private static final int DRAIN_BATCH_SIZE = 256;
+    private final ArrayDeque<T> localBatch = new ArrayDeque<>(DRAIN_BATCH_SIZE);
     private boolean hasReachedEnd =
         false; // flag which will be set to true when we have encountered a EndOfStream or when
     // curTup.isLast(). Ref: https://github.com/googleapis/java-bigquery/issues/2033
@@ -200,15 +203,22 @@ public boolean next() throws SQLException {
         return false;
       }
       try {
-        cursor = buffer.take(); // advance the cursor,Potentially blocking operation
+        cursor = localBatch.pollFirst();
+        if (cursor == null) {
+          cursor = buffer.take(); // advance the cursor,Potentially blocking operation
+          // Drain available queued rows into a consumer-local batch to amortize lock acquisitions.
+          buffer.drainTo(localBatch, DRAIN_BATCH_SIZE);
+        }
         if (isEndOfStream(cursor)) { // check for end of stream
           cursor = null;
+          localBatch.clear();
           hasReachedEnd = true;
           return false;
         } else if (cursor instanceof Row) {
           Row curTup = (Row) cursor;
           if (curTup.isLast()) { // last Tuple
             cursor = null;
+            localBatch.clear();
             hasReachedEnd = true;
             return false;
           }

java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ConnectionImpl.java

@@ -64,10 +64,10 @@
 import java.util.Map;
 import java.util.Queue;
 import java.util.UUID;
+import java.util.concurrent.ArrayBlockingQueue;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
-import java.util.concurrent.LinkedBlockingDeque;
 import java.util.concurrent.TimeUnit;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -544,16 +544,16 @@ BigQueryResult tableDataList(GetQueryResultsResponse firstPage, JobId jobId) {
     BigQueryResultStats bigQueryResultStats = getBigQueryResultSetStats(jobId);
 
     // Keeps the deserialized records at the row level, which is consumed by BigQueryResult
-    bufferFvl = new LinkedBlockingDeque<>(getBufferSize());
+    bufferFvl = new ArrayBlockingQueue<>(getBufferSize());
 
     // Keeps the parsed FieldValueLists
     BlockingQueue<Tuple<Iterable<FieldValueList>, Boolean>> pageCache =
-        new LinkedBlockingDeque<>(
+        new ArrayBlockingQueue<>(
             getPageCacheSize(connectionSettings.getNumBufferedRows(), schema));
 
     // Keeps the raw RPC responses
     BlockingQueue<Tuple<TableDataList, Boolean>> rpcResponseQueue =
-        new LinkedBlockingDeque<>(
+        new ArrayBlockingQueue<>(
             getPageCacheSize(connectionSettings.getNumBufferedRows(), schema));
 
     runNextPageTaskAsync(firstPage.getPageToken(), getDestinationTable(jobId), rpcResponseQueue);
@@ -598,12 +598,12 @@ BigQueryResult processQueryResponseResults(
     BigQueryResultStats bigQueryResultStats =
         new BigQueryResultStatsImpl(queryStatistics, sessionInfo);
 
-    bufferFvl = new LinkedBlockingDeque<>(getBufferSize());
+    bufferFvl = new ArrayBlockingQueue<>(getBufferSize());
     BlockingQueue<Tuple<Iterable<FieldValueList>, Boolean>> pageCache =
-        new LinkedBlockingDeque<>(
+        new ArrayBlockingQueue<>(
             getPageCacheSize(connectionSettings.getNumBufferedRows(), schema));
     BlockingQueue<Tuple<TableDataList, Boolean>> rpcResponseQueue =
-        new LinkedBlockingDeque<>(
+        new ArrayBlockingQueue<>(
             getPageCacheSize(connectionSettings.getNumBufferedRows(), schema));
 
     JobId jobId = JobId.fromPb(results.getJobReference());
@@ -1027,7 +1027,7 @@ BigQueryResult highThroughPutRead(
           // DO a regex check using order by and use multiple streams
           ;
       ReadSession readSession = bqReadClient.createReadSession(builder.build());
-      bufferRow = new LinkedBlockingDeque<>(getBufferSize());
+      bufferRow = new ArrayBlockingQueue<>(getBufferSize());
       Map<String, Integer> arrowNameToIndex = new HashMap<>();
       // deserialize and populate the buffer async, so that the client isn't blocked
       processArrowStreamAsync(

java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryResultImplTest.java

@@ -316,4 +316,37 @@ void testResultSetReadApiGetLongByColumnIndexPreservesInt64Value()
     assertThat(resultSet.wasNull()).isFalse();
     assertThat(resultSet.next()).isFalse();
   }
+
+  @Test
+  void testResultSetBatchDrainingPreservesFifoOrderAndEof()
+      throws InterruptedException, SQLException {
+    int numRows = 600;
+    BlockingQueue<AbstractList<FieldValue>> buffer =
+        new java.util.concurrent.ArrayBlockingQueue<>(numRows + 2);
+    Schema singleColSchema =
+        Schema.of(
+            Field.newBuilder("id", StandardSQLTypeName.INT64)
+                .setMode(Field.Mode.NULLABLE)
+                .build());
+    FieldList singleColFields = FieldList.of(Field.of("id", LegacySQLTypeName.INTEGER));
+
+    for (int i = 0; i < numRows; i++) {
+      buffer.put(
+          FieldValueList.of(
+              ImmutableList.of(FieldValue.of(Attribute.PRIMITIVE, Integer.toString(i))),
+              singleColFields));
+    }
+    buffer.put(new EndOfFieldValueList());
+
+    BigQueryResultImpl<AbstractList<FieldValue>> bigQueryResult =
+        new BigQueryResultImpl<>(singleColSchema, numRows, buffer, null);
+    ResultSet resultSet = bigQueryResult.getResultSet();
+
+    for (int i = 0; i < numRows; i++) {
+      assertThat(resultSet.next()).isTrue();
+      assertThat(resultSet.getInt("id")).isEqualTo(i);
+    }
+    assertThat(resultSet.next()).isFalse();
+    assertThat(resultSet.next()).isFalse();
+  }
 }