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()); + } }