Skip to content

fix(bigquery-jdbc): expand handling to OOM errors #14759

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
logachev merged 3 commits into from
Oct 28, 2026
Merged

fix(bigquery-jdbc): expand handling to OOM errors #14759

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
206080f
fix(bigquery-jdbc): expand handling to OOM errors
logachev Oct 23, 2026
b18ba8b
re-throw exception
logachev Oct 23, 2026
bd3d080
lint
logachev Oct 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions ...bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowBatchWrapper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,13 +29,13 @@ class BigQueryArrowBatchWrapper {
// Marks the end of the stream for the ResultSet
private final boolean isLast;

private final Exception exception;
private final Throwable exception;

private BigQueryArrowBatchWrapper(
ArrowRecordBatch currentArrowBatch,
JsonStringArrayList nestedRecords,
boolean isLast,
Exception exception) {
Throwable exception) {
this.currentArrowBatch = currentArrowBatch;
this.nestedRecords = nestedRecords;
this.isLast = isLast;
Expand All @@ -53,7 +53,7 @@ static BigQueryArrowBatchWrapper getNestedFieldValueListWrapper(
return new BigQueryArrowBatchWrapper(null, nestedRecords, isLastFlag, null);
}

static BigQueryArrowBatchWrapper ofError(Exception exception) {
static BigQueryArrowBatchWrapper ofError(Throwable exception) {
return new BigQueryArrowBatchWrapper(null, null, true, exception);
}

Expand All @@ -69,7 +69,7 @@ boolean isLast() {
return this.isLast;
}

Exception getException() {
Throwable getException() {
return this.exception;
}
}
2 changes: 1 addition & 1 deletion java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -854,7 +854,7 @@ public void clearWarnings() throws SQLException {
this.warningsLoaded = true;
}

static BigQueryJdbcException backgroundFetchError(Exception cause) {
static BigQueryJdbcException backgroundFetchError(Throwable cause) {
String msg = "Error occurred while fetching query results";
return cause instanceof BigQueryException
? new BigQueryJdbcException(msg, (BigQueryException) cause)
Expand Down
8 changes: 4 additions & 4 deletions ...uery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryFieldValueListWrapper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ class BigQueryFieldValueListWrapper {

// This flag marks the end of the stream for the ResultSet
private boolean isLast = false;
private final Exception exception;
private final Throwable exception;

static BigQueryFieldValueListWrapper ofEndOfStream(FieldList fieldList) {
return new BigQueryFieldValueListWrapper(fieldList, null, null, null, true, null);
Expand Down Expand Up @@ -106,7 +106,7 @@ static BigQueryFieldValueListWrapper getNestedFieldValueListWrapper(
fieldList, null, arrayFieldValueList, null, isLastFlag, null);
}

static BigQueryFieldValueListWrapper ofError(Exception exception) {
static BigQueryFieldValueListWrapper ofError(Throwable exception) {
return new BigQueryFieldValueListWrapper(null, null, null, null, true, exception);
}

Expand All @@ -116,7 +116,7 @@ private BigQueryFieldValueListWrapper(
List<FieldValue> arrayFieldValueList,
Object[] rowValues,
boolean isLast,
Exception exception) {
Throwable exception) {
this.fieldList = fieldList;
this.fieldValueList = fieldValueList;
this.arrayFieldValueList = arrayFieldValueList;
Expand Down Expand Up @@ -165,7 +165,7 @@ boolean isLast() {
return this.isLast;
}

Exception getException() {
Throwable getException() {
return this.exception;
}
}
26 changes: 16 additions & 10 deletions java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1077,7 +1077,7 @@ Future<?> populateArrowBufferedQueue(
+ " Interrupted @ arrowStreamProcessor",
e);
enqueueError(arrowBatchWrapperBlockingQueue, e);
} catch (Exception e) {
} catch (Throwable e) {
if (e.getCause() instanceof InterruptedException
|| Thread.currentThread().isInterrupted()) {
Thread.currentThread().interrupt();
Expand All @@ -1090,11 +1090,14 @@ Future<?> populateArrowBufferedQueue(
e);
enqueueError(arrowBatchWrapperBlockingQueue, e);
} else {
enqueueError(arrowBatchWrapperBlockingQueue, e);
LOG.log(
Level.WARNING,
"\n" + Thread.currentThread().getName() + " Error @ arrowStreamProcessor",
e);
enqueueError(arrowBatchWrapperBlockingQueue, e);
}
Comment thread
logachev marked this conversation as resolved.
if (e instanceof Error) {
throw (Error) e;
}
} finally { // logic needed for graceful shutdown
enqueueEndOfStream(arrowBatchWrapperBlockingQueue);
Expand Down Expand Up @@ -1322,11 +1325,11 @@ Future<?> runNextPageTaskAsync(
String currentPageToken = firstPageToken;
TableResult currentResults = result;
TableId destinationTable = null;
if (firstPageToken != null) {
destinationTable = getDestinationTable(jobId);
}

try {
if (firstPageToken != null) {
destinationTable = getDestinationTable(jobId);
}
while (currentPageToken != null) {
// do not process further pages and shutdown
if (isWorkerShuttingDown()) {
Expand Down Expand Up @@ -1368,12 +1371,15 @@ Future<?> runNextPageTaskAsync(
"%s Interrupted @ runNextPageTaskAsync: %s",
Thread.currentThread().getName(), ex.getMessage());
enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex);
} catch (Exception ex) {
} catch (Throwable ex) {
if (ex.getCause() instanceof InterruptedException
|| Thread.currentThread().isInterrupted()) {
Thread.currentThread().interrupt();
}
enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex);
if (ex instanceof Error) {
throw (Error) ex;
}
} finally {
// this will stop the parseDataTask as well when the pagination
// completes
Expand Down Expand Up @@ -1457,7 +1463,7 @@ Future<?> parseAndPopulateRpcDataAsync(
+ " Interrupted @ populateBufferAsync",
ex);
enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex);
} catch (Exception ex) {
} catch (Throwable ex) {

@Neenu1995 Neenu1995 Oct 27, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Re-throw error instance?

@logachev logachev Oct 28, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is running in a background thread, thread will exit after this either way. We mark an error via the queue

if (ex.getCause() instanceof InterruptedException
|| Thread.currentThread().isInterrupted()) {
Thread.currentThread().interrupt();
Expand All @@ -1470,11 +1476,11 @@ Future<?> parseAndPopulateRpcDataAsync(
ex);
enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex);
} else {
enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex);
LOG.log(
Level.WARNING,
"\n" + Thread.currentThread().getName() + " Error @ populateBufferAsync",
ex);
enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex);
}
Comment thread
logachev marked this conversation as resolved.
} finally {
enqueueBufferEndOfStream(bigQueryFieldValueListWrapperBlockingQueue);
Expand Down Expand Up @@ -1979,15 +1985,15 @@ private <T> void enqueueMarker(BlockingQueue<T> queue, T marker) {
}
}

private void enqueueError(BlockingQueue<BigQueryArrowBatchWrapper> queue, Exception e) {
private void enqueueError(BlockingQueue<BigQueryArrowBatchWrapper> queue, Throwable e) {
enqueueMarker(queue, BigQueryArrowBatchWrapper.ofError(e));
}

private void enqueueEndOfStream(BlockingQueue<BigQueryArrowBatchWrapper> queue) {
enqueueMarker(queue, BigQueryArrowBatchWrapper.of(null, true));
}

private void enqueueBufferError(BlockingQueue<BigQueryFieldValueListWrapper> queue, Exception e) {
private void enqueueBufferError(BlockingQueue<BigQueryFieldValueListWrapper> queue, Throwable e) {
enqueueMarker(queue, BigQueryFieldValueListWrapper.ofError(e));
}

Expand Down
155 changes: 155 additions & 0 deletions java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@
import com.google.cloud.bigquery.DatasetInfo;
import com.google.cloud.bigquery.Field;
import com.google.cloud.bigquery.FieldList;
import com.google.cloud.bigquery.FieldValue;
import com.google.cloud.bigquery.FieldValueList;
import com.google.cloud.bigquery.Job;
import com.google.cloud.bigquery.JobId;
Expand Down Expand Up @@ -87,6 +88,7 @@
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
Expand Down Expand Up @@ -1560,4 +1562,157 @@ public ReadRowsResponse next() {
marker.getException(),
"Stopping early due to worker shutdown must enqueue an error marker rather than normal end-of-stream");
}

private static FieldValueList createTestRow(boolean first, long second) {
return FieldValueList.of(
Arrays.asList(
FieldValue.of(FieldValue.Attribute.PRIMITIVE, Boolean.toString(first)),
FieldValue.of(FieldValue.Attribute.PRIMITIVE, Long.toString(second))),
fieldList);
}

@Test
public void testRunNextPageTaskAsyncPropagatesOutOfMemoryErrorInsteadOfTruncating()
throws Exception {
JobId testJobId = JobId.of("project", "job_oom_next_page");
TableId destTable = TableId.of("project", "dataset", "dest_table");

Job mockJob = mock(Job.class);
QueryJobConfiguration jobConfig =
QueryJobConfiguration.newBuilder("SELECT first, second FROM test")
.setDestinationTable(destTable)
.build();
doReturn(mockJob).when(bigquery).getJob(testJobId);
doReturn(jobConfig).when(mockJob).getConfiguration();

Schema schema = Schema.of(fieldList);
TableResult firstPage = mock(TableResult.class);
doReturn(schema).when(firstPage).getSchema();
doReturn(10L).when(firstPage).getTotalRows();
doReturn(testJobId).when(firstPage).getJobId();
doReturn("page_2_token").when(firstPage).getNextPageToken();
doReturn(Arrays.asList(createTestRow(true, 1L), createTestRow(false, 2L)))
.when(firstPage)
.getValues();

Mockito.doThrow(new OutOfMemoryError("Simulated OOM allocating TableResult for page 2"))
.when(bigquery)
.listTableData(
any(TableId.class), any(TableDataListOption.class), any(TableDataListOption.class));

try (ResultSet rs = bigQueryStatement.processJsonResultSet(firstPage, mockJob)) {
SQLException thrown =
assertThrows(
SQLException.class,
() -> {
while (rs.next()) {
// consume until error
}
});
assertTrue(thrown.getCause() instanceof OutOfMemoryError);
}
}

@Test
public void testParseAndPopulateRpcDataAsyncPropagatesOutOfMemoryErrorInsteadOfTruncating()
throws Exception {
Iterable<FieldValueList> oomOnThirdRowIterable =
() ->
new Iterator<FieldValueList>() {
private int index = 0;

@Override
public boolean hasNext() {
return index < 10;
}

@Override
public FieldValueList next() {
index++;
if (index <= 2) {
return createTestRow(true, index);
}
throw new OutOfMemoryError("Simulated OOM unpacking row " + index);
}
};

Schema schema = Schema.of(fieldList);
TableResult firstPage = mock(TableResult.class);
doReturn(schema).when(firstPage).getSchema();
doReturn(10L).when(firstPage).getTotalRows();
doReturn(null).when(firstPage).getJobId();
doReturn(oomOnThirdRowIterable).when(firstPage).getValues();

try (ResultSet rs = bigQueryStatement.processJsonResultSet(firstPage, null)) {
SQLException thrown =
assertThrows(
SQLException.class,
() -> {
while (rs.next()) {
// consume until error
}
});
assertTrue(thrown.getCause() instanceof OutOfMemoryError);
}
}

@Test
public void testPopulateArrowBufferedQueuePropagatesOutOfMemoryErrorInsteadOfTruncating()
throws Exception {
BlockingQueue<BigQueryArrowBatchWrapper> boundedQueue = new LinkedBlockingDeque<>(4);
ReadSession readSession =
ReadSession.newBuilder()
.setArrowSchema(arrowSchema)
.addStreams(
ReadStream.newBuilder().setName("projects/p/locations/l/sessions/s/streams/0"))
.build();

@SuppressWarnings("unchecked")
ServerStream<ReadRowsResponse> serverStream = mock(ServerStream.class);
Iterator<ReadRowsResponse> responseIterator =
new Iterator<ReadRowsResponse>() {
@Override
public boolean hasNext() {
return true;
}

@Override
public ReadRowsResponse next() {
throw new OutOfMemoryError("Simulated OOM reading Arrow batch");
}
};
doReturn(responseIterator).when(serverStream).iterator();

@SuppressWarnings("unchecked")
ServerStreamingCallable<ReadRowsRequest, ReadRowsResponse> 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);

try (BigQueryArrowResultSet rs =
BigQueryArrowResultSet.of(
Schema.of(fieldList),
arrowSchema,
10L,
bigQueryStatement,
boundedQueue,
producerFuture,
bigquery)) {
SQLException thrown =
assertThrows(
SQLException.class,
() -> {
while (rs.next()) {
// consume until error
}
});
assertNotNull(thrown.getCause());
assertTrue(thrown.getCause().getCause() instanceof OutOfMemoryError);
}
}
}
Loading