From 0663d9e31b0d5e5c989934580d93ac7340767989 Mon Sep 17 00:00:00 2001 From: Lawrence Qiu Date: Tue, 13 Oct 2026 03:50:24 +0000 Subject: [PATCH] perf(bigquery): [T1-9] replace LinkedBlockingDeque with ArrayBlockingQueue and batch-drain in ConnectionImpl --- .../cloud/bigquery/BigQueryResultImpl.java | 12 ++++++- .../google/cloud/bigquery/ConnectionImpl.java | 16 ++++----- .../bigquery/BigQueryResultImplTest.java | 33 +++++++++++++++++++ 3 files changed, 52 insertions(+), 9 deletions(-) diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryResultImpl.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryResultImpl.java index c84724083f63..ec086b9f91cf 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryResultImpl.java +++ b/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 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; } diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ConnectionImpl.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ConnectionImpl.java index b8832edb19de..8b1d28e676d9 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ConnectionImpl.java +++ b/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, Boolean>> pageCache = - new LinkedBlockingDeque<>( + new ArrayBlockingQueue<>( getPageCacheSize(connectionSettings.getNumBufferedRows(), schema)); // Keeps the raw RPC responses BlockingQueue> 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, Boolean>> pageCache = - new LinkedBlockingDeque<>( + new ArrayBlockingQueue<>( getPageCacheSize(connectionSettings.getNumBufferedRows(), schema)); BlockingQueue> 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 arrowNameToIndex = new HashMap<>(); // deserialize and populate the buffer async, so that the client isn't blocked processArrowStreamAsync( diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryResultImplTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryResultImplTest.java index ab5543e1b6e9..b5066c74a133 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryResultImplTest.java +++ b/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> 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> 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(); + } }