From 7ce41072aec71b7cabed19b70f0f3739a5b7580d Mon Sep 17 00:00:00 2001 From: kirl Date: Tue, 20 Oct 2026 23:36:06 +0000 Subject: [PATCH 1/5] fix(bigquery-jdbc): bound Storage Read API Arrow batch queue to 32 batches --- .../bigquery/jdbc/BigQueryStatement.java | 7 ++++- .../bigquery/jdbc/BigQueryStatementTest.java | 29 +++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java index af50bf8b8df3..0654d6036f0f 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java @@ -104,6 +104,7 @@ public class BigQueryStatement extends BigQueryNoOpsStatement { private final BigQueryJdbcCustomLogger LOG = new BigQueryJdbcCustomLogger(this.toString()); public static final int DEFAULT_BUFFER_SIZE = BigQuerySettings.DEFAULT_NUM_BUFFERED_ROWS * 2; + static final int DEFAULT_ARROW_BUFFER_SIZE = 32; private static final String DEFAULT_DATASET_NAME = "_google_jdbc"; private static final String DEFAULT_TABLE_NAME = "temp_table_"; private static final String JDBC_JOB_PREFIX = "google-jdbc-"; @@ -931,7 +932,7 @@ ResultSet processArrowResultSet(TableResult results, Job job) throws SQLExceptio .setMaxStreamCount(1); ReadSession readSession = getReadSession(builder.build()); - this.arrowBatchWrapperBlockingQueue = new LinkedBlockingDeque<>(getBufferSize()); + this.arrowBatchWrapperBlockingQueue = new LinkedBlockingDeque<>(getArrowBufferSize()); // deserialize and populate the buffer async, so that the client isn't blocked populateBufferWorker = populateArrowBufferedQueue( @@ -1512,6 +1513,10 @@ private int getBufferSize() { : Math.min(this.querySettings.getNumBufferedRows() * 2, 100000)); } + int getArrowBufferSize() { + return DEFAULT_ARROW_BUFFER_SIZE; + } + /** Returns the destinationTable from jobId by calling `jobs.get` API */ TableId getDestinationTable(JobId jobId) { Job job = this.bigQuery.getJob(jobId); diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java index 9245f627c147..a442efd61ed4 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java @@ -1318,4 +1318,33 @@ public void testJoblessQueryPopulatesQueryStatistics() throws Exception { verify(bigquery, Mockito.never()).getJob(any(JobId.class)); } + + @Test + public void testProcessArrowResultSetUsesBoundedQueueCapacity() throws SQLException { + BigQueryStatement defaultStmtSpy = Mockito.spy(bigQueryStatement); + BigQueryReadClient bigQueryReadClient = mock(BigQueryReadClient.class); + Schema schema = Schema.of(fieldList); + ReadSession readSession = ReadSession.getDefaultInstance(); + doReturn(bigQueryReadClient).when(defaultStmtSpy).getBigQueryReadClient(); + doReturn(readSession).when(defaultStmtSpy).getReadSession(any(CreateReadSessionRequest.class)); + Future mockWorker = mock(Future.class); + @SuppressWarnings("unchecked") + ArgumentCaptor> defaultQueueCaptor = + ArgumentCaptor.forClass(BlockingQueue.class); + doReturn(mockWorker) + .when(defaultStmtSpy) + .populateArrowBufferedQueue( + any(ReadSession.class), defaultQueueCaptor.capture(), any(BigQueryReadClient.class)); + doReturn(arrowSchema).when(defaultStmtSpy).getArrowSchema(any(ReadSession.class)); + + TableResult result = mock(TableResult.class); + doReturn(schema).when(result).getSchema(); + doReturn(10L).when(result).getTotalRows(); + doReturn(TABLE_ID).when(defaultStmtSpy).getDestinationTable(any()); + + defaultStmtSpy.processArrowResultSet(result, null); + assertEquals( + BigQueryStatement.DEFAULT_ARROW_BUFFER_SIZE, + defaultQueueCaptor.getValue().remainingCapacity()); + } } From 02f9592015eb88e11f5de8ae5788c1955bf0a497 Mon Sep 17 00:00:00 2001 From: kirl Date: Tue, 20 Oct 2026 23:38:50 +0000 Subject: [PATCH 2/5] fix(bigquery-jdbc): stop producer and release resources on ResultSet.close() --- .../bigquery/jdbc/BigQueryArrowResultSet.java | 19 ++- .../bigquery/jdbc/BigQueryJsonResultSet.java | 7 + .../bigquery/jdbc/BigQueryStatement.java | 125 +++++++++++++++--- .../jdbc/BigQueryArrowResultSetTest.java | 30 +++++ .../bigquery/jdbc/BigQueryStatementTest.java | 99 ++++++++++++++ 5 files changed, 258 insertions(+), 22 deletions(-) diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java index b6cdc393173f..94573278a69c 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java @@ -246,8 +246,12 @@ private void deserializeArrowBatch(ArrowRecordBatch batch) throws SQLException { @Override public void close() { LOG.fineTrace("close", () -> String.format("Closing BigQueryArrowResultSet %s.", this)); - vectorSchemaRoot.close(); - allocator.close(); + if (vectorSchemaRoot != null) { + vectorSchemaRoot.close(); + } + if (allocator != null) { + allocator.close(); + } } } @@ -542,12 +546,23 @@ private String formatRangeElement(Object element, StandardSQLTypeName elementTyp @Override public void close() { + if (this.isClosed) { + return; + } LOG.fineTrace("close", () -> String.format("Closing BigqueryArrowResultSet %s.", this)); this.isClosed = true; if (ownedTask != null) { // cancel the producer task when result set is closed ownedTask.cancel(true); } + if (this.buffer != null) { + this.buffer.clear(); + } + if (this.arrowDeserializer != null) { + this.arrowDeserializer.close(); + } else if (this.allocator != null) { + this.allocator.close(); + } super.close(); } diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java index c8a2451aeeaa..afa54bff4fb9 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java @@ -382,6 +382,9 @@ private FieldValue getObjectInternal(int columnIndex) throws SQLException { @Override public void close() { + if (this.isClosed) { + return; + } LOG.fineTrace("close", () -> String.format("Closing BigqueryJsonResultSet %s.", this)); this.isClosed = true; if (ownedTasks != null) { @@ -391,6 +394,10 @@ public void close() { } } } + if (this.buffer != null) { + this.buffer.clear(); + } + this.cursor = null; super.close(); } diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java index 0654d6036f0f..44eb3557f662 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java @@ -19,6 +19,7 @@ import com.google.api.core.InternalApi; import com.google.api.gax.paging.Page; import com.google.api.gax.rpc.ApiException; +import com.google.api.gax.rpc.ServerStream; import com.google.api.gax.rpc.StatusCode; import com.google.cloud.Tuple; import com.google.cloud.bigquery.BigQuery; @@ -62,7 +63,6 @@ import com.google.cloud.bigquery.storage.v1.ReadSession; import com.google.cloud.bigquery.storage.v1.ReadSession.TableReadOptions; import com.google.common.annotations.VisibleForTesting; -import com.google.common.util.concurrent.Uninterruptibles; import io.grpc.Status; import io.grpc.StatusRuntimeException; import io.opentelemetry.api.common.AttributeKey; @@ -535,6 +535,8 @@ private void closeStatementResources() throws SQLException { tmp.close(); } this.batchQueries.clear(); + this.arrowBatchWrapperBlockingQueue = null; + this.bigQueryFieldValueListWrapperBlockingQueue = null; this.currentUpdateCount = -1; this.currentJobIdIndex = -1; if (this.connection != null) { @@ -1005,6 +1007,8 @@ Future populateArrowBufferedQueue( String streamName = readSession.getStreams(0).getName(); while (true) { + ServerStream stream = null; + boolean streamCompleted = false; try { ReadRowsRequest readRowsRequest = ReadRowsRequest.newBuilder() @@ -1014,19 +1018,23 @@ Future populateArrowBufferedQueue( // Process each block of rows as they arrive and decode using our simple row // reader. - com.google.api.gax.rpc.ServerStream stream = - bqReadClient.readRowsCallable().call(readRowsRequest); + stream = bqReadClient.readRowsCallable().call(readRowsRequest); for (ReadRowsResponse response : stream) { if (Thread.currentThread().isInterrupted() || executor.isShutdown()) { + arrowBatchWrapperBlockingQueue.clear(); break; } ArrowRecordBatch currentBatch = response.getArrowRecordBatch(); - Uninterruptibles.putUninterruptibly( - arrowBatchWrapperBlockingQueue, + arrowBatchWrapperBlockingQueue.put( BigQueryArrowBatchWrapper.of(currentBatch)); rowsRead += response.getRowCount(); } + if (!Thread.currentThread().isInterrupted() && !executor.isShutdown()) { + streamCompleted = true; + } else { + arrowBatchWrapperBlockingQueue.clear(); + } break; } catch (ApiException e) { if (e.getStatusCode().getCode() == StatusCode.Code.NOT_FOUND) { @@ -1049,10 +1057,16 @@ Future populateArrowBufferedQueue( "Connection interrupted during arrow stream read, retrying. attempt: %d", retryCount); Thread.sleep(RETRY_DELAY_MS); + } finally { + if (stream != null && !streamCompleted) { + stream.cancel(); + } } } } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + arrowBatchWrapperBlockingQueue.clear(); LOG.log( Level.WARNING, "\n" @@ -1060,10 +1074,11 @@ Future populateArrowBufferedQueue( + " Interrupted @ arrowStreamProcessor", e); enqueueError(arrowBatchWrapperBlockingQueue, e); - Thread.currentThread().interrupt(); } catch (Exception e) { if (e.getCause() instanceof InterruptedException || Thread.currentThread().isInterrupted()) { + Thread.currentThread().interrupt(); + arrowBatchWrapperBlockingQueue.clear(); LOG.log( Level.WARNING, "\n" @@ -1071,7 +1086,6 @@ Future populateArrowBufferedQueue( + " Interrupted @ arrowStreamProcessor", e); enqueueError(arrowBatchWrapperBlockingQueue, e); - Thread.currentThread().interrupt(); } else { LOG.log( Level.WARNING, @@ -1338,8 +1352,7 @@ Future runNextPageTaskAsync( currentPageToken = currentResults.getNextPageToken(); // this will be parsed asynchronously without blocking the current // thread - Uninterruptibles.putUninterruptibly( - rpcResponseQueue, Tuple.of(currentResults, true)); + rpcResponseQueue.put(Tuple.of(currentResults, true)); LOG.fine( "Fetched %d results from the server in %d ms.", querySettings.getMaxResultPerPage(), @@ -1348,14 +1361,30 @@ Future runNextPageTaskAsync( paginationSpan.end(); } } + } catch (InterruptedException ex) { + Thread.currentThread().interrupt(); + LOG.warning( + "%s Interrupted @ runNextPageTaskAsync: %s", + Thread.currentThread().getName(), ex.getMessage()); } catch (Exception ex) { - Uninterruptibles.putUninterruptibly( - bigQueryFieldValueListWrapperBlockingQueue, - BigQueryFieldValueListWrapper.ofError(ex)); + if (ex.getCause() instanceof InterruptedException + || Thread.currentThread().isInterrupted()) { + Thread.currentThread().interrupt(); + } + enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex); } finally { // this will stop the parseDataTask as well when the pagination // completes - Uninterruptibles.putUninterruptibly(rpcResponseQueue, Tuple.of(null, false)); + Tuple endOfRpcTuple = Tuple.of(null, false); + if (Thread.currentThread().isInterrupted() || executor.isShutdown()) { + rpcResponseQueue.offer(endOfRpcTuple); + } else { + try { + rpcResponseQueue.put(endOfRpcTuple); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + } + } } // We cannot do queryTaskExecutor.shutdownNow() here as populate buffer method may // not @@ -1425,8 +1454,7 @@ Future parseAndPopulateRpcDataAsync( // do not process further pages and shutdown (inner loop) break; } - Uninterruptibles.putUninterruptibly( - bigQueryFieldValueListWrapperBlockingQueue, + bigQueryFieldValueListWrapperBlockingQueue.put( BigQueryFieldValueListWrapper.of( schema.getFields(), fieldValueList, isComplexColumn)); results += 1; @@ -1436,9 +1464,19 @@ Future parseAndPopulateRpcDataAsync( results, (int) ((System.nanoTime() - startTime) / 1000000)); } + } catch (InterruptedException ex) { + Thread.currentThread().interrupt(); + LOG.log( + Level.WARNING, + "\n" + + Thread.currentThread().getName() + + " Interrupted @ populateBufferAsync", + ex); + enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex); } catch (Exception ex) { if (ex.getCause() instanceof InterruptedException || Thread.currentThread().isInterrupted()) { + Thread.currentThread().interrupt(); LOG.log( Level.WARNING, "\n" @@ -1446,7 +1484,6 @@ Future parseAndPopulateRpcDataAsync( + " Interrupted @ populateBufferAsync", ex); enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex); - Thread.currentThread().interrupt(); } else { LOG.log( Level.WARNING, @@ -1933,20 +1970,68 @@ enum QueryDialectType { BIG_QUERY } + private boolean isWorkerShuttingDown() { + if (Thread.currentThread().isInterrupted()) { + return true; + } + BigQueryConnection conn = this.connection; + if (conn == null) { + return true; + } + ExecutorService executor = conn.getExecutorService(); + return executor != null && executor.isShutdown(); + } + private void enqueueError(BlockingQueue queue, Exception e) { - Uninterruptibles.putUninterruptibly(queue, BigQueryArrowBatchWrapper.ofError(e)); + BigQueryArrowBatchWrapper errorWrapper = BigQueryArrowBatchWrapper.ofError(e); + if (isWorkerShuttingDown()) { + queue.offer(errorWrapper); + return; + } + try { + queue.put(errorWrapper); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + } } private void enqueueEndOfStream(BlockingQueue queue) { - Uninterruptibles.putUninterruptibly(queue, BigQueryArrowBatchWrapper.of(null, true)); + BigQueryArrowBatchWrapper endWrapper = BigQueryArrowBatchWrapper.of(null, true); + if (isWorkerShuttingDown()) { + queue.offer(endWrapper); + return; + } + try { + queue.put(endWrapper); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + } } private void enqueueBufferError(BlockingQueue queue, Exception e) { - Uninterruptibles.putUninterruptibly(queue, BigQueryFieldValueListWrapper.ofError(e)); + BigQueryFieldValueListWrapper errorWrapper = BigQueryFieldValueListWrapper.ofError(e); + if (isWorkerShuttingDown()) { + queue.offer(errorWrapper); + return; + } + try { + queue.put(errorWrapper); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + } } private void enqueueBufferEndOfStream(BlockingQueue queue) { - Uninterruptibles.putUninterruptibly(queue, BigQueryFieldValueListWrapper.ofEndOfStream(null)); + BigQueryFieldValueListWrapper endWrapper = BigQueryFieldValueListWrapper.ofEndOfStream(null); + if (isWorkerShuttingDown()) { + queue.offer(endWrapper); + return; + } + try { + queue.put(endWrapper); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + } } QueryStatistics describePositionalParameterQuery(String query) diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSetTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSetTest.java index 8681fd9559a5..0797c829194d 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSetTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSetTest.java @@ -24,6 +24,8 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; import com.google.cloud.bigquery.Field; import com.google.cloud.bigquery.Field.Mode; @@ -339,6 +341,34 @@ public void testIsClosed() { assertThat(bigQueryArrowResultSet.isClosed()).isFalse(); } + @Test + public void testCloseCancelsProducerClearsBufferAndIsIdempotent() + throws SQLException, IOException { + ArrowSchema arrowSchema = + ArrowSchema.newBuilder() + .setSerializedSchema(serializeSchema(vectorSchemaRoot.getSchema())) + .build(); + Future workerTask = mock(Future.class); + BlockingQueue queue = new LinkedBlockingDeque<>(4); + queue.add(arrowBatchWrapper); + queue.add(arrowBatchWrapperLast); + + BigQueryArrowResultSet rs = + BigQueryArrowResultSet.of(QUERY_SCHEMA, arrowSchema, 1, statement, queue, workerTask, null); + + assertThat(rs.isClosed()).isFalse(); + assertThat(queue).hasSize(2); + + rs.close(); + assertThat(rs.isClosed()).isTrue(); + assertThat(queue).isEmpty(); + verify(workerTask, times(1)).cancel(true); + + // Second close() call should be a no-op + rs.close(); + verify(workerTask, times(1)).cancel(true); + } + @Test public void testResultSetHoldability() throws SQLException { assertThat(bigQueryArrowResultSet.getHoldability()) diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java index a442efd61ed4..2fda81227e71 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java @@ -33,6 +33,8 @@ import static org.mockito.Mockito.withSettings; import com.google.api.gax.rpc.ApiException; +import com.google.api.gax.rpc.ServerStream; +import com.google.api.gax.rpc.ServerStreamingCallable; import com.google.api.gax.rpc.StatusCode; import com.google.cloud.ServiceOptions; import com.google.cloud.Tuple; @@ -62,11 +64,15 @@ import com.google.cloud.bigquery.exception.BigQueryJdbcException; import com.google.cloud.bigquery.jdbc.BigQueryStatement.JobIdWrapper; import com.google.cloud.bigquery.spi.BigQueryRpcFactory; +import com.google.cloud.bigquery.storage.v1.ArrowRecordBatch; import com.google.cloud.bigquery.storage.v1.ArrowSchema; import com.google.cloud.bigquery.storage.v1.ArrowSerializationOptions; import com.google.cloud.bigquery.storage.v1.BigQueryReadClient; import com.google.cloud.bigquery.storage.v1.CreateReadSessionRequest; +import com.google.cloud.bigquery.storage.v1.ReadRowsRequest; +import com.google.cloud.bigquery.storage.v1.ReadRowsResponse; import com.google.cloud.bigquery.storage.v1.ReadSession; +import com.google.cloud.bigquery.storage.v1.ReadStream; import com.google.common.collect.ImmutableList; import com.google.common.collect.Maps; import io.opentelemetry.api.baggage.Baggage; @@ -83,14 +89,18 @@ import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; +import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.UUID; import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.LinkedBlockingDeque; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.stream.Stream; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BitVector; @@ -1347,4 +1357,93 @@ public void testProcessArrowResultSetUsesBoundedQueueCapacity() throws SQLExcept BigQueryStatement.DEFAULT_ARROW_BUFFER_SIZE, defaultQueueCaptor.getValue().remainingCapacity()); } + + @Test + public void testPopulateArrowBufferedQueueStopsAndCancelsStreamWhenResultSetClosedWhileQueueFull() + throws Exception { + BlockingQueue boundedQueue = new LinkedBlockingDeque<>(1); + ReadSession readSession = + ReadSession.newBuilder() + .addStreams( + ReadStream.newBuilder().setName("projects/p/locations/l/sessions/s/streams/0")) + .build(); + + ReadRowsResponse batchResponse = + ReadRowsResponse.newBuilder() + .setArrowRecordBatch(ArrowRecordBatch.getDefaultInstance()) + .build(); + + CountDownLatch secondBatchReturnedLatch = new CountDownLatch(1); + AtomicBoolean streamCancelled = new AtomicBoolean(false); + + @SuppressWarnings("unchecked") + ServerStream serverStream = mock(ServerStream.class); + Mockito.doAnswer( + invocation -> { + streamCancelled.set(true); + return null; + }) + .when(serverStream) + .cancel(); + + Iterator responseIterator = + new Iterator() { + private int count = 0; + + @Override + public boolean hasNext() { + return true; + } + + @Override + public ReadRowsResponse next() { + count++; + if (count == 2) { + secondBatchReturnedLatch.countDown(); + } + return batchResponse; + } + }; + doReturn(responseIterator).when(serverStream).iterator(); + + @SuppressWarnings("unchecked") + ServerStreamingCallable callable = + mock(ServerStreamingCallable.class); + doReturn(serverStream).when(callable).call(any(ReadRowsRequest.class)); + + BigQueryReadClient readClient = mock(BigQueryReadClient.class); + doReturn(callable).when(readClient).readRowsCallable(); + + Future producerFuture = + bigQueryStatement.populateArrowBufferedQueue(readSession, boundedQueue, readClient); + + // Wait until the first batch is in the queue and the second batch has been yielded by the + // stream iterator (so the producer thread is blocked or about to block on boundedQueue.put()). + assertTrue(secondBatchReturnedLatch.await(5, TimeUnit.SECONDS)); + assertEquals(1, boundedQueue.size()); + + BigQueryArrowResultSet rs = + BigQueryArrowResultSet.of( + Schema.of(fieldList), + arrowSchema, + 100L, + bigQueryStatement, + boundedQueue, + producerFuture, + bigquery); + + // Closing the ResultSet while the queue is full must cancel the producer task, + // unblock the producer thread, cancel the underlying ServerStream, and clear the buffer. + rs.close(); + + assertTrue(producerFuture.isCancelled()); + testExecutorService.shutdown(); + assertTrue( + testExecutorService.awaitTermination(5, TimeUnit.SECONDS), + "Producer thread should terminate promptly after ResultSet.close()"); + assertTrue(streamCancelled.get(), "ServerStream.cancel() should be called on early close"); + assertTrue( + boundedQueue.stream().allMatch(w -> w.getCurrentArrowBatch() == null), + "No ArrowRecordBatch should remain buffered after ResultSet.close()"); + } } From 2465b0842db91986401aba73b9cfc1bab9e698f6 Mon Sep 17 00:00:00 2001 From: kirl Date: Wed, 21 Oct 2026 18:27:17 +0000 Subject: [PATCH 3/5] fix(bigquery-jdbc): address PR #14712 review comments --- .../bigquery/jdbc/BigQueryArrowResultSet.java | 30 +++-- .../bigquery/jdbc/BigQueryJsonResultSet.java | 24 ++-- .../bigquery/jdbc/BigQueryStatement.java | 120 ++++++------------ .../jdbc/BigQueryArrowResultSetTest.java | 4 +- .../bigquery/jdbc/BigQueryStatementTest.java | 114 +++++++++++++++++ 5 files changed, 190 insertions(+), 102 deletions(-) diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java index 94573278a69c..607b76bec92e 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java @@ -283,6 +283,7 @@ public boolean next() throws SQLException { // Advance the cursor. Potentially blocking operation. try (Scope scope = makeOriginalContextCurrent()) { BigQueryArrowBatchWrapper batchWrapper = this.buffer.take(); + checkClosed(); if (batchWrapper.getException() != null) { throw backgroundFetchError(batchWrapper.getException()); } @@ -551,19 +552,24 @@ public void close() { } LOG.fineTrace("close", () -> String.format("Closing BigqueryArrowResultSet %s.", this)); this.isClosed = true; - if (ownedTask != null) { - // cancel the producer task when result set is closed - ownedTask.cancel(true); - } - if (this.buffer != null) { - this.buffer.clear(); - } - if (this.arrowDeserializer != null) { - this.arrowDeserializer.close(); - } else if (this.allocator != null) { - this.allocator.close(); + try { + if (ownedTask != null) { + // cancel the producer task when result set is closed + ownedTask.cancel(true); + } + if (this.buffer != null) { + this.buffer.clear(); + // Unblock any consumer thread currently waiting on buffer.take() in next() + this.buffer.offer(BigQueryArrowBatchWrapper.of(null, true)); + } + if (this.arrowDeserializer != null) { + this.arrowDeserializer.close(); + } else if (this.allocator != null) { + this.allocator.close(); + } + } finally { + super.close(); } - super.close(); } @Override diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java index afa54bff4fb9..46b6c06e8fc1 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java @@ -227,6 +227,7 @@ public boolean next() throws SQLException { try { // Advance the cursor,Potentially blocking operation this.cursor = this.buffer.take(); + checkClosed(); if (this.cursor.getException() != null) { throw backgroundFetchError(this.cursor.getException()); } @@ -387,18 +388,23 @@ public void close() { } LOG.fineTrace("close", () -> String.format("Closing BigqueryJsonResultSet %s.", this)); this.isClosed = true; - if (ownedTasks != null) { - for (Future ownedTask : ownedTasks) { - if (ownedTask != null) { - ownedTask.cancel(true); + try { + if (ownedTasks != null) { + for (Future ownedTask : ownedTasks) { + if (ownedTask != null) { + ownedTask.cancel(true); + } } } + if (this.buffer != null) { + this.buffer.clear(); + // Unblock any consumer thread currently waiting on buffer.take() in next() + this.buffer.offer(BigQueryFieldValueListWrapper.ofEndOfStream(null)); + } + this.cursor = null; + } finally { + super.close(); } - if (this.buffer != null) { - this.buffer.clear(); - } - this.cursor = null; - super.close(); } @Override diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java index 44eb3557f662..060c60c607b2 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java @@ -104,6 +104,8 @@ public class BigQueryStatement extends BigQueryNoOpsStatement { private final BigQueryJdbcCustomLogger LOG = new BigQueryJdbcCustomLogger(this.toString()); public static final int DEFAULT_BUFFER_SIZE = BigQuerySettings.DEFAULT_NUM_BUFFERED_ROWS * 2; + // Bounds in-memory ArrowRecordBatch buffering per ResultSet (~15-20 MB of in-flight batches), + // keeping the gRPC ReadRows stream saturated while preventing unbounded heap accumulation. static final int DEFAULT_ARROW_BUFFER_SIZE = 32; private static final String DEFAULT_DATASET_NAME = "_google_jdbc"; private static final String DEFAULT_TABLE_NAME = "temp_table_"; @@ -1020,9 +1022,9 @@ Future populateArrowBufferedQueue( // reader. stream = bqReadClient.readRowsCallable().call(readRowsRequest); for (ReadRowsResponse response : stream) { - if (Thread.currentThread().isInterrupted() || executor.isShutdown()) { - arrowBatchWrapperBlockingQueue.clear(); - break; + if (isWorkerShuttingDown()) { + throw new InterruptedException( + "Arrow stream reading was interrupted or shut down before completion."); } ArrowRecordBatch currentBatch = response.getArrowRecordBatch(); @@ -1030,11 +1032,11 @@ Future populateArrowBufferedQueue( BigQueryArrowBatchWrapper.of(currentBatch)); rowsRead += response.getRowCount(); } - if (!Thread.currentThread().isInterrupted() && !executor.isShutdown()) { - streamCompleted = true; - } else { - arrowBatchWrapperBlockingQueue.clear(); + if (isWorkerShuttingDown()) { + throw new InterruptedException( + "Arrow stream reading was interrupted or shut down before completion."); } + streamCompleted = true; break; } catch (ApiException e) { if (e.getStatusCode().getCode() == StatusCode.Code.NOT_FOUND) { @@ -1326,11 +1328,9 @@ Future runNextPageTaskAsync( try { while (currentPageToken != null) { // do not process further pages and shutdown - if (Thread.currentThread().isInterrupted() || executor.isShutdown()) { - LOG.warning( - "%s Interrupted @ runNextPageTaskAsync", - Thread.currentThread().getName()); - break; + if (isWorkerShuttingDown()) { + throw new InterruptedException( + "Pagination was interrupted or shut down before completion."); } SpanBuilder spanBuilder = @@ -1366,6 +1366,7 @@ Future runNextPageTaskAsync( LOG.warning( "%s Interrupted @ runNextPageTaskAsync: %s", Thread.currentThread().getName(), ex.getMessage()); + enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex); } catch (Exception ex) { if (ex.getCause() instanceof InterruptedException || Thread.currentThread().isInterrupted()) { @@ -1375,16 +1376,7 @@ Future runNextPageTaskAsync( } finally { // this will stop the parseDataTask as well when the pagination // completes - Tuple endOfRpcTuple = Tuple.of(null, false); - if (Thread.currentThread().isInterrupted() || executor.isShutdown()) { - rpcResponseQueue.offer(endOfRpcTuple); - } else { - try { - rpcResponseQueue.put(endOfRpcTuple); - } catch (InterruptedException ie) { - Thread.currentThread().interrupt(); - } - } + enqueueMarker(rpcResponseQueue, Tuple.of(null, false)); } // We cannot do queryTaskExecutor.shutdownNow() here as populate buffer method may // not @@ -1419,29 +1411,19 @@ Future parseAndPopulateRpcDataAsync( // as we have to process the first page boolean hasRows = true; while (hasRows) { - try { - Tuple nextPageTuple = rpcResponseQueue.take(); - if (nextPageTuple.x() != null) { - fieldValueLists = nextPageTuple.x().getValues(); - } else { - fieldValueLists = null; - } - hasRows = nextPageTuple.y(); - - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - LOG.log( - Level.WARNING, - "\n" + Thread.currentThread().getName() + " Interrupted", - e); - // Thread might get interrupted while calling the Cancel method, which is - // expected, so logging this instead of throwing the exception back - break; + Tuple nextPageTuple = rpcResponseQueue.take(); + if (nextPageTuple.x() != null) { + fieldValueLists = nextPageTuple.x().getValues(); + } else { + fieldValueLists = null; } + hasRows = nextPageTuple.y(); - if (Thread.currentThread().isInterrupted() - || executor.isShutdown() - || fieldValueLists == null) { + if (isWorkerShuttingDown()) { + throw new InterruptedException( + "Buffer population was interrupted or shut down before completion."); + } + if (fieldValueLists == null) { // do not process further pages and shutdown (outerloop) break; } @@ -1450,9 +1432,9 @@ Future parseAndPopulateRpcDataAsync( long results = 0; for (FieldValueList fieldValueList : fieldValueLists) { - if (Thread.currentThread().isInterrupted() || executor.isShutdown()) { - // do not process further pages and shutdown (inner loop) - break; + if (isWorkerShuttingDown()) { + throw new InterruptedException( + "Buffer population was interrupted or shut down before completion."); } bigQueryFieldValueListWrapperBlockingQueue.put( BigQueryFieldValueListWrapper.of( @@ -1466,6 +1448,7 @@ Future parseAndPopulateRpcDataAsync( } catch (InterruptedException ex) { Thread.currentThread().interrupt(); + bigQueryFieldValueListWrapperBlockingQueue.clear(); LOG.log( Level.WARNING, "\n" @@ -1477,6 +1460,7 @@ Future parseAndPopulateRpcDataAsync( if (ex.getCause() instanceof InterruptedException || Thread.currentThread().isInterrupted()) { Thread.currentThread().interrupt(); + bigQueryFieldValueListWrapperBlockingQueue.clear(); LOG.log( Level.WARNING, "\n" @@ -1982,56 +1966,32 @@ private boolean isWorkerShuttingDown() { return executor != null && executor.isShutdown(); } - private void enqueueError(BlockingQueue queue, Exception e) { - BigQueryArrowBatchWrapper errorWrapper = BigQueryArrowBatchWrapper.ofError(e); + private void enqueueMarker(BlockingQueue queue, T marker) { if (isWorkerShuttingDown()) { - queue.offer(errorWrapper); + queue.offer(marker); return; } try { - queue.put(errorWrapper); + queue.put(marker); } catch (InterruptedException ie) { Thread.currentThread().interrupt(); } } + private void enqueueError(BlockingQueue queue, Exception e) { + enqueueMarker(queue, BigQueryArrowBatchWrapper.ofError(e)); + } + private void enqueueEndOfStream(BlockingQueue queue) { - BigQueryArrowBatchWrapper endWrapper = BigQueryArrowBatchWrapper.of(null, true); - if (isWorkerShuttingDown()) { - queue.offer(endWrapper); - return; - } - try { - queue.put(endWrapper); - } catch (InterruptedException ie) { - Thread.currentThread().interrupt(); - } + enqueueMarker(queue, BigQueryArrowBatchWrapper.of(null, true)); } private void enqueueBufferError(BlockingQueue queue, Exception e) { - BigQueryFieldValueListWrapper errorWrapper = BigQueryFieldValueListWrapper.ofError(e); - if (isWorkerShuttingDown()) { - queue.offer(errorWrapper); - return; - } - try { - queue.put(errorWrapper); - } catch (InterruptedException ie) { - Thread.currentThread().interrupt(); - } + enqueueMarker(queue, BigQueryFieldValueListWrapper.ofError(e)); } private void enqueueBufferEndOfStream(BlockingQueue queue) { - BigQueryFieldValueListWrapper endWrapper = BigQueryFieldValueListWrapper.ofEndOfStream(null); - if (isWorkerShuttingDown()) { - queue.offer(endWrapper); - return; - } - try { - queue.put(endWrapper); - } catch (InterruptedException ie) { - Thread.currentThread().interrupt(); - } + enqueueMarker(queue, BigQueryFieldValueListWrapper.ofEndOfStream(null)); } QueryStatistics describePositionalParameterQuery(String query) diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSetTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSetTest.java index 0797c829194d..6f8a249513b1 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSetTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSetTest.java @@ -361,7 +361,9 @@ public void testCloseCancelsProducerClearsBufferAndIsIdempotent() rs.close(); assertThat(rs.isClosed()).isTrue(); - assertThat(queue).isEmpty(); + assertThat(queue).hasSize(1); + assertThat(queue.peek().getCurrentArrowBatch()).isNull(); + assertThat(queue.peek().isLast()).isTrue(); verify(workerTask, times(1)).cancel(true); // Second close() call should be a no-op diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java index 2fda81227e71..41f3731c2806 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java @@ -95,6 +95,7 @@ import java.util.UUID; import java.util.concurrent.BlockingQueue; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; @@ -1446,4 +1447,117 @@ public ReadRowsResponse next() { boundedQueue.stream().allMatch(w -> w.getCurrentArrowBatch() == null), "No ArrowRecordBatch should remain buffered after ResultSet.close()"); } + + @Test + public void testCancelStatementUnblocksConcurrentArrowResultSetNext() throws Exception { + BlockingQueue emptyQueue = new LinkedBlockingDeque<>(4); + Future producerFuture = mock(Future.class); + BigQueryArrowResultSet rs = + BigQueryArrowResultSet.of( + Schema.of(fieldList), + arrowSchema, + 100L, + bigQueryStatement, + emptyQueue, + producerFuture, + bigquery); + bigQueryStatement.currentResultSet = rs; + + CountDownLatch readerEnteredNext = new CountDownLatch(1); + Future readerFuture = + testExecutorService.submit( + () -> { + readerEnteredNext.countDown(); + rs.next(); + return null; + }); + + assertTrue(readerEnteredNext.await(5, TimeUnit.SECONDS)); + // Give the reader thread a brief moment to block inside emptyQueue.take() + Thread.sleep(50); + + bigQueryStatement.cancel(); + + ExecutionException ex = + assertThrows(ExecutionException.class, () -> readerFuture.get(5, TimeUnit.SECONDS)); + assertTrue(ex.getCause() instanceof SQLException); + } + + @Test + public void testCancelStatementUnblocksConcurrentJsonResultSetNext() throws Exception { + BlockingQueue emptyQueue = new LinkedBlockingDeque<>(4); + BigQueryJsonResultSet rs = + BigQueryJsonResultSet.of( + Schema.of(fieldList), 100L, emptyQueue, bigQueryStatement, null, bigquery); + bigQueryStatement.currentResultSet = rs; + + CountDownLatch readerEnteredNext = new CountDownLatch(1); + Future readerFuture = + testExecutorService.submit( + () -> { + readerEnteredNext.countDown(); + rs.next(); + return null; + }); + + assertTrue(readerEnteredNext.await(5, TimeUnit.SECONDS)); + Thread.sleep(50); + + bigQueryStatement.cancel(); + + ExecutionException ex = + assertThrows(ExecutionException.class, () -> readerFuture.get(5, TimeUnit.SECONDS)); + assertTrue(ex.getCause() instanceof SQLException); + } + + @Test + public void testPopulateArrowBufferedQueueEnqueuesErrorWhenWorkerShuttingDown() throws Exception { + BlockingQueue boundedQueue = new LinkedBlockingDeque<>(4); + ReadSession readSession = + ReadSession.newBuilder() + .addStreams( + ReadStream.newBuilder().setName("projects/p/locations/l/sessions/s/streams/0")) + .build(); + + ReadRowsResponse batchResponse = + ReadRowsResponse.newBuilder() + .setArrowRecordBatch(ArrowRecordBatch.getDefaultInstance()) + .build(); + + @SuppressWarnings("unchecked") + ServerStream serverStream = mock(ServerStream.class); + Iterator responseIterator = + new Iterator() { + @Override + public boolean hasNext() { + // Simulate executor shutdown mid-stream before the next batch is processed + testExecutorService.shutdownNow(); + return true; + } + + @Override + public ReadRowsResponse next() { + return batchResponse; + } + }; + doReturn(responseIterator).when(serverStream).iterator(); + + @SuppressWarnings("unchecked") + ServerStreamingCallable callable = + mock(ServerStreamingCallable.class); + doReturn(serverStream).when(callable).call(any(ReadRowsRequest.class)); + + BigQueryReadClient readClient = mock(BigQueryReadClient.class); + doReturn(callable).when(readClient).readRowsCallable(); + + Future producerFuture = + bigQueryStatement.populateArrowBufferedQueue(readSession, boundedQueue, readClient); + producerFuture.get(5, TimeUnit.SECONDS); + + BigQueryArrowBatchWrapper marker = boundedQueue.poll(5, TimeUnit.SECONDS); + assertNotNull(marker); + assertNotNull( + marker.getException(), + "Stopping early due to worker shutdown must enqueue an error marker rather than normal end-of-stream"); + } } From 51cba4548234292d7493b8b7003e674717e8fe1e Mon Sep 17 00:00:00 2001 From: kirl Date: Wed, 21 Oct 2026 19:19:50 +0000 Subject: [PATCH 4/5] refactor(bigquery-jdbc): move ResultSet buffer cleanup to BigQueryBaseResultSet --- .../bigquery/jdbc/BigQueryArrowResultSet.java | 19 +++++------- .../bigquery/jdbc/BigQueryBaseResultSet.java | 29 +++++++++++++++++-- .../bigquery/jdbc/BigQueryJsonResultSet.java | 17 +++++------ .../jdbc/BigQueryBaseResultSetTest.java | 19 ++++++++++++ 4 files changed, 60 insertions(+), 24 deletions(-) diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java index 607b76bec92e..05af02ad2a7c 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java @@ -54,13 +54,10 @@ import org.apache.arrow.vector.util.Text; /** {@link ResultSet} Implementation for Arrow datasource (Using Storage Read APIs) */ -class BigQueryArrowResultSet extends BigQueryBaseResultSet { +class BigQueryArrowResultSet extends BigQueryBaseResultSet { private final long totalRows; // count of rows read by the current instance of ResultSet private long rowCount = 0; - // IMP: This is a buffer of Arrow batches, the max size should be kept at min as - // possible to avoid holding too much memory - private final BlockingQueue buffer; // TODO(neenu): See if it makes sense to have the nested batch represented by // 'JsonStringArrayList' directly @@ -103,10 +100,9 @@ private BigQueryArrowResultSet( Job job, boolean enableTimestampPicos) throws SQLException { - super(bigQuery, statement, schema, isNested, job); + super(bigQuery, statement, schema, isNested, job, buffer); LOG.finestTrace(""); this.totalRows = totalRows; - this.buffer = buffer; this.currentNestedBatch = currentNestedBatch; this.fromIndex = fromIndex; this.toIndexExclusive = toIndexExclusive; @@ -169,7 +165,6 @@ static BigQueryArrowResultSet of( BigQueryArrowResultSet() throws SQLException { super(null, null, null, false); this.totalRows = 0; - this.buffer = null; this.currentNestedBatch = null; this.fromIndex = 0; this.toIndexExclusive = 0; @@ -545,6 +540,11 @@ private String formatRangeElement(Object element, StandardSQLTypeName elementTyp } } + @Override + protected BigQueryArrowBatchWrapper getEndOfStreamSentinel() { + return BigQueryArrowBatchWrapper.of(null, true); + } + @Override public void close() { if (this.isClosed) { @@ -557,11 +557,6 @@ public void close() { // cancel the producer task when result set is closed ownedTask.cancel(true); } - if (this.buffer != null) { - this.buffer.clear(); - // Unblock any consumer thread currently waiting on buffer.take() in next() - this.buffer.offer(BigQueryArrowBatchWrapper.of(null, true)); - } if (this.arrowDeserializer != null) { this.arrowDeserializer.close(); } else if (this.allocator != null) { diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java index d409eb79724c..1b4d9f206289 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java @@ -55,10 +55,11 @@ import java.time.ZoneId; import java.util.Calendar; import java.util.List; +import java.util.concurrent.BlockingQueue; import javax.sql.rowset.serial.SerialBlob; import javax.sql.rowset.serial.SerialClob; -public abstract class BigQueryBaseResultSet extends BigQueryNoOpsResultSet +public abstract class BigQueryBaseResultSet extends BigQueryNoOpsResultSet implements BigQueryResultSet { protected final BigQueryJdbcResultSetLogger LOG; private BigQuery bigQuery; @@ -69,6 +70,7 @@ public abstract class BigQueryBaseResultSet extends BigQueryNoOpsResultSet protected final Schema schema; protected final FieldList schemaFieldList; protected final boolean isNested; + protected final BlockingQueue buffer; protected boolean isClosed = false; protected boolean wasNull = false; private int fetchSize = -1; @@ -80,16 +82,27 @@ public abstract class BigQueryBaseResultSet extends BigQueryNoOpsResultSet protected BigQueryBaseResultSet( BigQuery bigQuery, BigQueryStatement statement, Schema schema, boolean isNested) { - this(bigQuery, statement, schema, isNested, null); + this(bigQuery, statement, schema, isNested, null, null); } protected BigQueryBaseResultSet( BigQuery bigQuery, BigQueryStatement statement, Schema schema, boolean isNested, Job job) { + this(bigQuery, statement, schema, isNested, job, null); + } + + protected BigQueryBaseResultSet( + BigQuery bigQuery, + BigQueryStatement statement, + Schema schema, + boolean isNested, + Job job, + BlockingQueue buffer) { this.bigQuery = bigQuery; this.statement = statement; this.schema = schema; this.schemaFieldList = schema != null ? schema.getFields() : null; this.isNested = isNested; + this.buffer = buffer; this.originalSpanContext = Span.current().getSpanContext(); this.job = job; if (job != null) { @@ -151,8 +164,20 @@ public String getQueryId() { return queryId; } + protected T getEndOfStreamSentinel() { + return null; + } + @Override public void close() { + if (this.buffer != null) { + this.buffer.clear(); + // Unblock any consumer thread currently waiting on buffer.take() in next() + T sentinel = getEndOfStreamSentinel(); + if (sentinel != null) { + this.buffer.offer(sentinel); + } + } try { if (statement != null && statement.isCloseOnCompletion() && !statement.hasMoreResults()) { statement.close(); diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java index 46b6c06e8fc1..7fcb4eae4f07 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java @@ -35,9 +35,8 @@ import java.util.concurrent.Future; /** {@link ResultSet} Implementation for JSON datasource (Using REST APIs) */ -class BigQueryJsonResultSet extends BigQueryBaseResultSet { +class BigQueryJsonResultSet extends BigQueryBaseResultSet { private final long totalRows; - private final BlockingQueue buffer; private boolean hasReachedEnd = false; // Points to the current record private BigQueryFieldValueListWrapper cursor; @@ -63,9 +62,8 @@ private BigQueryJsonResultSet( BigQuery bigQuery, Job job, boolean enableTimestampPicos) { - super(bigQuery, statement, schema, isNested, job); + super(bigQuery, statement, schema, isNested, job, buffer); this.totalRows = totalRows; - this.buffer = buffer; this.cursor = cursor; this.fromIndex = fromIndex; this.toIndexExclusive = toIndexExclusive; @@ -157,7 +155,6 @@ static BigQueryJsonResultSet of( BigQueryJsonResultSet() { super(null, null, null, false); totalRows = 0; - buffer = null; fromIndex = 0; ownedTasks = new Future[0]; toIndexExclusive = 0; @@ -381,6 +378,11 @@ private FieldValue getObjectInternal(int columnIndex) throws SQLException { return value; } + @Override + protected BigQueryFieldValueListWrapper getEndOfStreamSentinel() { + return BigQueryFieldValueListWrapper.ofEndOfStream(null); + } + @Override public void close() { if (this.isClosed) { @@ -396,11 +398,6 @@ public void close() { } } } - if (this.buffer != null) { - this.buffer.clear(); - // Unblock any consumer thread currently waiting on buffer.take() in next() - this.buffer.offer(BigQueryFieldValueListWrapper.ofEndOfStream(null)); - } this.cursor = null; } finally { super.close(); diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSetTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSetTest.java index 78a62c1a6b76..b9b16cd273ab 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSetTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSetTest.java @@ -42,6 +42,8 @@ import java.sql.SQLException; import java.sql.SQLWarning; import java.util.Arrays; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.LinkedBlockingDeque; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -341,4 +343,21 @@ public void testQueryStatisticsWithPrePopulatedStatisticsForStatelessQuery() { assertThat(stats.getTotalSlotMs()).isEqualTo(250L); verify(bigQuery, never()).getJob(any(JobId.class)); } + + @Test + public void testCloseClearsBufferAndOffersEndOfStreamSentinel() throws Exception { + BlockingQueue queue = new LinkedBlockingDeque<>(4); + queue.add("row1"); + queue.add("row2"); + Object sentinel = new Object(); + + Field bufferField = BigQueryBaseResultSet.class.getDeclaredField("buffer"); + bufferField.setAccessible(true); + bufferField.set(resultSet, queue); + doReturn(sentinel).when(resultSet).getEndOfStreamSentinel(); + + resultSet.close(); + + assertThat(queue).containsExactly(sentinel); + } } From ef9b74d6f3fbf26b4701379a9030626fc181aacb Mon Sep 17 00:00:00 2001 From: kirl Date: Wed, 21 Oct 2026 19:22:08 +0000 Subject: [PATCH 5/5] Revert "refactor(bigquery-jdbc): move ResultSet buffer cleanup to BigQueryBaseResultSet" This reverts commit 51cba4548234292d7493b8b7003e674717e8fe1e. --- .../bigquery/jdbc/BigQueryArrowResultSet.java | 19 +++++++----- .../bigquery/jdbc/BigQueryBaseResultSet.java | 29 ++----------------- .../bigquery/jdbc/BigQueryJsonResultSet.java | 17 ++++++----- .../jdbc/BigQueryBaseResultSetTest.java | 19 ------------ 4 files changed, 24 insertions(+), 60 deletions(-) diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java index 05af02ad2a7c..607b76bec92e 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java @@ -54,10 +54,13 @@ import org.apache.arrow.vector.util.Text; /** {@link ResultSet} Implementation for Arrow datasource (Using Storage Read APIs) */ -class BigQueryArrowResultSet extends BigQueryBaseResultSet { +class BigQueryArrowResultSet extends BigQueryBaseResultSet { private final long totalRows; // count of rows read by the current instance of ResultSet private long rowCount = 0; + // IMP: This is a buffer of Arrow batches, the max size should be kept at min as + // possible to avoid holding too much memory + private final BlockingQueue buffer; // TODO(neenu): See if it makes sense to have the nested batch represented by // 'JsonStringArrayList' directly @@ -100,9 +103,10 @@ private BigQueryArrowResultSet( Job job, boolean enableTimestampPicos) throws SQLException { - super(bigQuery, statement, schema, isNested, job, buffer); + super(bigQuery, statement, schema, isNested, job); LOG.finestTrace(""); this.totalRows = totalRows; + this.buffer = buffer; this.currentNestedBatch = currentNestedBatch; this.fromIndex = fromIndex; this.toIndexExclusive = toIndexExclusive; @@ -165,6 +169,7 @@ static BigQueryArrowResultSet of( BigQueryArrowResultSet() throws SQLException { super(null, null, null, false); this.totalRows = 0; + this.buffer = null; this.currentNestedBatch = null; this.fromIndex = 0; this.toIndexExclusive = 0; @@ -540,11 +545,6 @@ private String formatRangeElement(Object element, StandardSQLTypeName elementTyp } } - @Override - protected BigQueryArrowBatchWrapper getEndOfStreamSentinel() { - return BigQueryArrowBatchWrapper.of(null, true); - } - @Override public void close() { if (this.isClosed) { @@ -557,6 +557,11 @@ public void close() { // cancel the producer task when result set is closed ownedTask.cancel(true); } + if (this.buffer != null) { + this.buffer.clear(); + // Unblock any consumer thread currently waiting on buffer.take() in next() + this.buffer.offer(BigQueryArrowBatchWrapper.of(null, true)); + } if (this.arrowDeserializer != null) { this.arrowDeserializer.close(); } else if (this.allocator != null) { diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java index 1b4d9f206289..d409eb79724c 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java @@ -55,11 +55,10 @@ import java.time.ZoneId; import java.util.Calendar; import java.util.List; -import java.util.concurrent.BlockingQueue; import javax.sql.rowset.serial.SerialBlob; import javax.sql.rowset.serial.SerialClob; -public abstract class BigQueryBaseResultSet extends BigQueryNoOpsResultSet +public abstract class BigQueryBaseResultSet extends BigQueryNoOpsResultSet implements BigQueryResultSet { protected final BigQueryJdbcResultSetLogger LOG; private BigQuery bigQuery; @@ -70,7 +69,6 @@ public abstract class BigQueryBaseResultSet extends BigQueryNoOpsResultSet protected final Schema schema; protected final FieldList schemaFieldList; protected final boolean isNested; - protected final BlockingQueue buffer; protected boolean isClosed = false; protected boolean wasNull = false; private int fetchSize = -1; @@ -82,27 +80,16 @@ public abstract class BigQueryBaseResultSet extends BigQueryNoOpsResultSet protected BigQueryBaseResultSet( BigQuery bigQuery, BigQueryStatement statement, Schema schema, boolean isNested) { - this(bigQuery, statement, schema, isNested, null, null); + this(bigQuery, statement, schema, isNested, null); } protected BigQueryBaseResultSet( BigQuery bigQuery, BigQueryStatement statement, Schema schema, boolean isNested, Job job) { - this(bigQuery, statement, schema, isNested, job, null); - } - - protected BigQueryBaseResultSet( - BigQuery bigQuery, - BigQueryStatement statement, - Schema schema, - boolean isNested, - Job job, - BlockingQueue buffer) { this.bigQuery = bigQuery; this.statement = statement; this.schema = schema; this.schemaFieldList = schema != null ? schema.getFields() : null; this.isNested = isNested; - this.buffer = buffer; this.originalSpanContext = Span.current().getSpanContext(); this.job = job; if (job != null) { @@ -164,20 +151,8 @@ public String getQueryId() { return queryId; } - protected T getEndOfStreamSentinel() { - return null; - } - @Override public void close() { - if (this.buffer != null) { - this.buffer.clear(); - // Unblock any consumer thread currently waiting on buffer.take() in next() - T sentinel = getEndOfStreamSentinel(); - if (sentinel != null) { - this.buffer.offer(sentinel); - } - } try { if (statement != null && statement.isCloseOnCompletion() && !statement.hasMoreResults()) { statement.close(); diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java index 7fcb4eae4f07..46b6c06e8fc1 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java @@ -35,8 +35,9 @@ import java.util.concurrent.Future; /** {@link ResultSet} Implementation for JSON datasource (Using REST APIs) */ -class BigQueryJsonResultSet extends BigQueryBaseResultSet { +class BigQueryJsonResultSet extends BigQueryBaseResultSet { private final long totalRows; + private final BlockingQueue buffer; private boolean hasReachedEnd = false; // Points to the current record private BigQueryFieldValueListWrapper cursor; @@ -62,8 +63,9 @@ private BigQueryJsonResultSet( BigQuery bigQuery, Job job, boolean enableTimestampPicos) { - super(bigQuery, statement, schema, isNested, job, buffer); + super(bigQuery, statement, schema, isNested, job); this.totalRows = totalRows; + this.buffer = buffer; this.cursor = cursor; this.fromIndex = fromIndex; this.toIndexExclusive = toIndexExclusive; @@ -155,6 +157,7 @@ static BigQueryJsonResultSet of( BigQueryJsonResultSet() { super(null, null, null, false); totalRows = 0; + buffer = null; fromIndex = 0; ownedTasks = new Future[0]; toIndexExclusive = 0; @@ -378,11 +381,6 @@ private FieldValue getObjectInternal(int columnIndex) throws SQLException { return value; } - @Override - protected BigQueryFieldValueListWrapper getEndOfStreamSentinel() { - return BigQueryFieldValueListWrapper.ofEndOfStream(null); - } - @Override public void close() { if (this.isClosed) { @@ -398,6 +396,11 @@ public void close() { } } } + if (this.buffer != null) { + this.buffer.clear(); + // Unblock any consumer thread currently waiting on buffer.take() in next() + this.buffer.offer(BigQueryFieldValueListWrapper.ofEndOfStream(null)); + } this.cursor = null; } finally { super.close(); diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSetTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSetTest.java index b9b16cd273ab..78a62c1a6b76 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSetTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSetTest.java @@ -42,8 +42,6 @@ import java.sql.SQLException; import java.sql.SQLWarning; import java.util.Arrays; -import java.util.concurrent.BlockingQueue; -import java.util.concurrent.LinkedBlockingDeque; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -343,21 +341,4 @@ public void testQueryStatisticsWithPrePopulatedStatisticsForStatelessQuery() { assertThat(stats.getTotalSlotMs()).isEqualTo(250L); verify(bigQuery, never()).getJob(any(JobId.class)); } - - @Test - public void testCloseClearsBufferAndOffersEndOfStreamSentinel() throws Exception { - BlockingQueue queue = new LinkedBlockingDeque<>(4); - queue.add("row1"); - queue.add("row2"); - Object sentinel = new Object(); - - Field bufferField = BigQueryBaseResultSet.class.getDeclaredField("buffer"); - bufferField.setAccessible(true); - bufferField.set(resultSet, queue); - doReturn(sentinel).when(resultSet).getEndOfStreamSentinel(); - - resultSet.close(); - - assertThat(queue).containsExactly(sentinel); - } }