Skip to content

perf(bigquery): unwrap cell JSON during streaming #14639

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

Closed
lqiu96 wants to merge 1 commit into from
Closed

perf(bigquery): unwrap cell JSON during streaming #14639

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
23 changes: 20 additions & 3 deletions java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/FieldValue.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -378,13 +378,31 @@ public static FieldValue of(Attribute attribute, Object value, Boolean useInt64T
return new FieldValue(attribute, value, useInt64Timestamps);
}

private static final Object NULL_CELL_SENTINEL = Data.nullOf(Object.class);

static FieldValue fromPb(Object cellPb) {
return fromPb(cellPb, null, false);
}

@SuppressWarnings("unchecked")
static FieldValue fromPb(Object cellPb, Field recordSchema, Boolean useInt64Timestamps) {
if (Data.isNull(cellPb)) {
// Fast-path for TableCell instances (for example, when constructed directly in unit tests or
// returned by custom BigQueryRpc implementations). TableCell extends GenericJson ->
// GenericData, which implements Map<String, Object>. Without this check, a TableCell falls
// into the `cellPb instanceof Map` branch below and triggers reflective ClassInfo/FieldInfo
// lookups on "f" and "v" before recursing on the unwrapped value.
if (cellPb instanceof com.google.api.services.bigquery.model.TableCell) {
cellPb = ((com.google.api.services.bigquery.model.TableCell) cellPb).getV();
}
// Check `cellPb == null` and `cellPb == NULL_CELL_SENTINEL` before calling
// `Data.isNull(cellPb)`:
// 1. `Data.isNull(null)` returns `false`, so unwrapped `null` references must be handled first.
// 2. `Data.nullOf(Object.class)` is the sentinel used by google-http-client (and
// HttpBigQueryRpc's streaming cell parser) for JSON `null` values in `Object`-typed fields
// (`TableCell.v`). Checking reference equality (`==`) avoids `Data.isNull`'s
// `ConcurrentHashMap.get` lookup on every null cell while still falling back to
// `Data.isNull` for typed null sentinels such as `Data.nullOf(String.class)`.
if (cellPb == null || cellPb == NULL_CELL_SENTINEL || Data.isNull(cellPb)) {
return FieldValue.of(Attribute.PRIMITIVE, null, useInt64Timestamps);
}
if (cellPb instanceof String) {
Expand Down Expand Up @@ -412,8 +430,7 @@ static FieldValue fromPb(Object cellPb, Field recordSchema, Boolean useInt64Time
FieldValueList.fromPb(
(List<Object>) cellMapPb.get("f"), subFieldsSchema, useInt64Timestamps));
}
// This should never be the case when we are processing a first level table field (i.e. a
// row's field, not a record sub-field)
// Fallback for generic Map representations of a cell wrapper {"v": ...}.
if (cellMapPb.containsKey("v")) {
return FieldValue.fromPb(cellMapPb.get("v"), recordSchema, useInt64Timestamps);
}
Expand Down
255 changes: 250 additions & 5 deletions ...google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import static java.net.HttpURLConnection.HTTP_UNAUTHORIZED;

import com.google.api.client.googleapis.json.GoogleJsonResponseException;
import com.google.api.client.googleapis.services.json.AbstractGoogleJsonClientRequest;
import com.google.api.client.http.ByteArrayContent;
import com.google.api.client.http.GenericUrl;
import com.google.api.client.http.HttpRequest;
Expand All @@ -31,8 +32,15 @@
import com.google.api.client.http.HttpResponseException;
import com.google.api.client.http.HttpTransport;
import com.google.api.client.http.json.JsonHttpContent;
import com.google.api.client.json.GenericJson;
import com.google.api.client.json.JsonFactory;
import com.google.api.client.json.JsonParser;
import com.google.api.client.json.JsonToken;
import com.google.api.client.json.gson.GsonFactory;
import com.google.api.client.util.ClassInfo;
import com.google.api.client.util.Data;
import com.google.api.client.util.FieldInfo;
import com.google.api.client.util.Types;
import com.google.api.core.InternalApi;
import com.google.api.core.InternalExtensionOnly;
import com.google.api.services.bigquery.Bigquery;
Expand Down Expand Up @@ -62,6 +70,7 @@
import com.google.api.services.bigquery.model.TableDataList;
import com.google.api.services.bigquery.model.TableList;
import com.google.api.services.bigquery.model.TableReference;
import com.google.api.services.bigquery.model.TableRow;
import com.google.api.services.bigquery.model.TestIamPermissionsRequest;
import com.google.api.services.bigquery.model.TestIamPermissionsResponse;
import com.google.cloud.Tuple;
Expand All @@ -80,7 +89,11 @@
import io.opentelemetry.api.trace.SpanKind;
import io.opentelemetry.context.Scope;
import java.io.IOException;
import java.io.InputStream;
import java.math.BigInteger;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

Expand Down Expand Up @@ -1457,7 +1470,7 @@ public TableDataList listTableDataSkipExceptionTranslation(
if (span != null) {
span.setAttribute("bq.rpc.page_token", bqListRequest.getPageToken());
}
return bqListRequest.execute();
return executeAndParseRowResponse(bqListRequest, TableDataList.class);
});
}

Expand Down Expand Up @@ -1519,7 +1532,7 @@ public TableDataList listTableDataWithRowLimitSkipExceptionTranslation(
if (span != null) {
span.setAttribute("bq.rpc.page_token", bqListRequest.getPageToken());
}
return bqListRequest.execute();
return executeAndParseRowResponse(bqListRequest, TableDataList.class);
});
}

Expand Down Expand Up @@ -1860,7 +1873,7 @@ public GetQueryResultsResponse getQueryResultsSkipExceptionTranslation(
if (span != null) {
span.setAttribute("bq.rpc.page_token", queryRequest.getPageToken());
}
return queryRequest.execute();
return executeAndParseRowResponse(queryRequest, GetQueryResultsResponse.class);
});
}

Expand Down Expand Up @@ -1913,7 +1926,7 @@ public GetQueryResultsResponse getQueryResultsWithRowLimitSkipExceptionTranslati
if (span != null) {
span.setAttribute("bq.rpc.page_token", queryRequest.getPageToken());
}
return queryRequest.execute();
return executeAndParseRowResponse(queryRequest, GetQueryResultsResponse.class);
});
}

Expand Down Expand Up @@ -1951,10 +1964,242 @@ public QueryResponse queryRpcSkipExceptionTranslation(String projectId, QueryReq
queryRequest.getUriTemplate(),
null),
span -> {
return queryRequest.execute();
return executeAndParseRowResponse(queryRequest, QueryResponse.class);
});
}

/**
* Executes a BigQuery REST request that returns row data ({@link TableDataList}, {@link
* QueryResponse}, or {@link GetQueryResultsResponse}) and parses its JSON stream directly without
* reflectively instantiating {@link com.google.api.services.bigquery.model.TableCell} wrappers
* for every cell.
*
* <p>Why this exists:
*
* <ul>
* <li>Standard {@code request.execute()} uses {@code google-http-client}'s reflective {@link
* JsonParser#parse}, which calls {@link Data#newMapInstance(Class)} -> {@link
* Types#newInstance(Class)} for every cell in every row.
* <li>Because {@code TableCell} extends {@link GenericJson} -> {@link
* com.google.api.client.util.GenericData}, each cell allocation reflectively invokes {@code
* Class.getDeclaredConstructor()} (allocating a {@code Class[0]} array) and constructs an
* {@code EnumSet}, an {@code ArrayMap} for {@code unknownFields}, and the {@code TableCell}
* wrapper itself (4 object allocations per cell, plus extra {@code ArrayMap} wrappers for
* every nested {@code RECORD} or {@code REPEATED} value).
* <li>Immediately after RPC deserialization, {@code FieldValueList.fromPb} and {@code
* FieldValue.fromPb} unwrap {@code cell.getV()} and discard the {@code TableCell}.
* </ul>
*
* <p>By streaming the {@code "rows"} array with {@link #parseTableRows(JsonParser)} while
* delegating all top-level metadata fields ({@code schema}, {@code pageToken}, {@code totalRows},
* {@code jobReference}, {@code errors}, etc.) to {@link ClassInfo}/{@link JsonParser}, we
* preserve 100% compatibility on response metadata while eliminating 4 object allocations and
* reflective constructor calls per cell.
*/
private <T extends GenericJson> T executeAndParseRowResponse(
AbstractGoogleJsonClientRequest<T> request, Class<T> responseClass) throws IOException {
HttpResponse response = request.executeUnparsed();
try (InputStream content = response.getContent()) {
if (content == null) {
return null;
}
try (JsonParser parser =
bigquery.getJsonFactory().createJsonParser(content, response.getContentCharset())) {
return parseRowResponse(parser, responseClass);
}
} finally {
response.disconnect();
}
}

/**
* Parses a top-level BigQuery JSON response object ({@link TableDataList}, {@link QueryResponse},
* or {@link GetQueryResultsResponse}), routing the {@code "rows"} field to {@link
* #parseTableRows(JsonParser)} and all other fields to {@code google-http-client}'s standard
* {@link ClassInfo}/{@link FieldInfo} metadata binding.
*/
private static <T extends GenericJson> T parseRowResponse(
JsonParser parser, Class<T> responseClass) throws IOException {
JsonToken token = parser.getCurrentToken();
if (token == null) {
token = parser.nextToken();
}
if (token == null) {
return null;
}
T result = Types.newInstance(responseClass);
result.setFactory(parser.getFactory());
if (token != JsonToken.START_OBJECT) {
parser.skipChildren();
return result;
}
ClassInfo classInfo = ClassInfo.of(responseClass);
while (parser.nextToken() == JsonToken.FIELD_NAME) {
String fieldName = parser.getText();
parser.nextToken();
if ("rows".equals(fieldName)) {
FieldInfo rowsField = classInfo.getFieldInfo("rows");
if (rowsField != null) {
rowsField.setValue(result, parseTableRows(parser));
} else {
result.set("rows", parseTableRows(parser));
}
continue;
}
FieldInfo fieldInfo = classInfo.getFieldInfo(fieldName);
if (fieldInfo != null) {
JsonToken valueToken = parser.getCurrentToken();
if (valueToken == JsonToken.VALUE_NULL) {
fieldInfo.setValue(result, Data.nullOf(fieldInfo.getType()));
} else if (valueToken == JsonToken.VALUE_STRING
&& Data.isPrimitive(fieldInfo.getGenericType())) {
// Handles both String fields and @JsonString-annotated 64-bit integer fields
// (such as totalRows, totalBytesProcessed, creationTime, etc.). Calling
// `parser.parse(Type, false)` directly on a @JsonString field fails because
// `JsonParser.parse(Type, boolean)` passes `Field = null` to `parseValue`.
fieldInfo.setValue(
result, Data.parsePrimitiveValue(fieldInfo.getGenericType(), parser.getText()));
} else {
fieldInfo.setValue(result, parser.parse(fieldInfo.getGenericType(), false));
}
} else {
// Preserve any unrecognized top-level JSON fields in GenericData.unknownFields.
result.set(fieldName, parser.parse(Object.class, false));
}
}
return result;
}

/**
* Parses the BigQuery REST v2 {@code "rows"} JSON array into a list of {@link TableRow} objects.
*
* <p>Per the BigQuery v2 REST Discovery schema, each {@code TableRow} is serialized as {@code
* {"f": [{"v": ...}, {"v": ...}, ...]}}. Note that {@code TableRow.getF()} has a declared generic
* return type of {@code List<TableCell>}, but its sole consumers in {@code java-bigquery} ({@code
* BigQueryImpl.transformTableData} and {@code ConnectionImpl.getIterableFieldValueList}) pass
* {@code rowPb.getF()} directly to {@code FieldValueList.fromPb(List<?> rowPb, ...)} -> {@code
* FieldValue.fromPb(Object cellPb, ...)}. Storing the unwrapped cell values in {@code
* TableRow.setF((List) cells)} avoids allocating tens of millions of transient {@code TableCell}
* and {@code ArrayMap} wrappers per query.
*/
@SuppressWarnings({"unchecked", "rawtypes"})
private static List<TableRow> parseTableRows(JsonParser parser) throws IOException {
if (parser.getCurrentToken() != JsonToken.START_ARRAY) {
parser.skipChildren();
return null;
}
List<TableRow> rows = new ArrayList<>();
int expectedCols = 16;
while (parser.nextToken() != JsonToken.END_ARRAY) {
if (parser.getCurrentToken() != JsonToken.START_OBJECT) {
parser.skipChildren();
continue;
}
List<Object> cells = null;
while (parser.nextToken() == JsonToken.FIELD_NAME) {
String rowField = parser.getText();
parser.nextToken();
if ("f".equals(rowField) && parser.getCurrentToken() == JsonToken.START_ARRAY) {
// Pre-size each row's cell list from the previous row's column count to avoid
// ArrayList resizing across rows of the same table/query schema.
cells = new ArrayList<>(expectedCols);
while (parser.nextToken() != JsonToken.END_ARRAY) {
cells.add(parseCellValue(parser));
}
expectedCols = cells.size();
} else {
parser.skipChildren();
}
}
TableRow row = new TableRow();
if (cells != null) {
row.setF((List) cells);
}
rows.add(row);
}
Comment on lines +2093 to +2119

@gemini-code-assist gemini-code-assist Bot Oct 12, 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.

critical

If the JSON stream is truncated or malformed, parser.nextToken() can return null (representing EOF). Since null != JsonToken.END_ARRAY is true, the loop will continue. Inside the loop, parser.getCurrentToken() will be null, which is not START_OBJECT, so it calls parser.skipChildren(). Since skipChildren() does nothing when the current token is null, this results in an infinite loop consuming 100% CPU. Additionally, the inner loop reading cells will also loop infinitely adding JSON_NULL to the list, leading to an OutOfMemoryError.

To prevent this, ensure all while (parser.nextToken() != JsonToken.END_ARRAY) loops also check for null (EOF).

    JsonToken token;
    while ((token = parser.nextToken()) != JsonToken.END_ARRAY && token != null) {
      if (token != JsonToken.START_OBJECT) {
        parser.skipChildren();
        continue;
      }
      List<Object> cells = null;
      while (parser.nextToken() == JsonToken.FIELD_NAME) {
        String rowField = parser.getText();
        JsonToken valueToken = parser.nextToken();
        if ("f".equals(rowField) && valueToken == JsonToken.START_ARRAY) {
          // Pre-size each row's cell list from the previous row's column count to avoid
          // ArrayList resizing across rows of the same table/query schema.
          cells = new ArrayList<>(expectedCols);
          JsonToken cellToken;
          while ((cellToken = parser.nextToken()) != JsonToken.END_ARRAY && cellToken != null) {
            cells.add(parseCellValue(parser));
          }
          expectedCols = cells.size();
        } else {
          parser.skipChildren();
        }
      }
      TableRow row = new TableRow();
      if (cells != null) {
        row.setF((List) cells);
      }
      rows.add(row);
    }

return rows;
}

private static final Object JSON_NULL = Data.nullOf(Object.class);

/**
* Recursively parses a single BigQuery REST v2 {@code TableCell} (or nested value) and unwraps
* single-key {@code {"v": ...}} cell envelopes on the fly.
*
* <p>In the BigQuery REST v2 wire format, a cell in {@code "f"} takes one of four shapes:
*
* <ol>
* <li><b>Primitive scalar</b>: {@code {"v": "123"}} -> unwraps directly to {@code String}
* {@code "123"}.
* <li><b>Null value</b>: {@code {"v": null}} -> unwraps to {@code Data.nullOf(Object.class)}
* (matching {@code google-http-client}'s null sentinel so {@code Data.isNull(cell)} and
* {@code FieldValue.fromPb} behave identically).
* <li><b>REPEATED (ARRAY)</b>: {@code {"v": [{"v": "a"}, {"v": "b"}]}} -> unwraps outer {@code
* "v"} to a {@code List<Object>} whose elements are recursively unwrapped ({@code ["a",
* "b"]}, or {@code List<Map<String, Object>>} for {@code ARRAY<STRUCT>}).
* <li><b>RECORD (STRUCT)</b>: {@code {"v": {"f": [{"v": "x"}, {"v": "y"}]}}} -> outer {@code
* "v"} unwraps to the inner struct object, which is represented as {@code
* Collections.singletonMap("f", subFields)} (12-byte singleton map instead of a reflective
* {@code ArrayMap}) as expected by {@code FieldValue.fromPb}'s {@code
* cellMapPb.containsKey("f")} check.
* </ol>
*/
private static Object parseCellValue(JsonParser parser) throws IOException {
JsonToken token = parser.getCurrentToken();
if (token == null || token == JsonToken.VALUE_NULL) {
return JSON_NULL;
}
if (token == JsonToken.START_OBJECT) {
Object unwrappedValue = null;
boolean hasValue = false;
Map<String, Object> mapValue = null;
while (parser.nextToken() == JsonToken.FIELD_NAME) {
String key = parser.getText();
parser.nextToken();
if ("v".equals(key) && mapValue == null) {
unwrappedValue = parseCellValue(parser);
hasValue = true;
} else if ("f".equals(key) && !hasValue && mapValue == null) {
// A RECORD (STRUCT) value is encoded as {"f": [{"v": sub1}, {"v": sub2}, ...]}.
// FieldValue.fromPb identifies RECORD values by checking `cellPb instanceof Map` and
// `cellMapPb.containsKey("f")`. Using Collections.singletonMap("f", subFields) satisfies
// that contract without allocating a HashMap or ArrayMap.
List<Object> subFields = new ArrayList<>();
if (parser.getCurrentToken() == JsonToken.START_ARRAY) {
while (parser.nextToken() != JsonToken.END_ARRAY) {
subFields.add(parseCellValue(parser));
Comment on lines +2167 to +2170

@gemini-code-assist gemini-code-assist Bot Oct 12, 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.

critical

If the JSON stream is truncated or malformed, parser.nextToken() can return null (representing EOF). Since null != JsonToken.END_ARRAY is true, the loop will continue infinitely, adding JSON_NULL to subFields and eventually causing an OutOfMemoryError.

To prevent this, ensure the loop condition also checks for null (EOF).

Suggested change
List<Object> subFields = new ArrayList<>();
if (parser.getCurrentToken() == JsonToken.START_ARRAY) {
while (parser.nextToken() != JsonToken.END_ARRAY) {
subFields.add(parseCellValue(parser));
if (parser.getCurrentToken() == JsonToken.START_ARRAY) {
JsonToken subToken;
while ((subToken = parser.nextToken()) != JsonToken.END_ARRAY && subToken != null) {
subFields.add(parseCellValue(parser));
}

}
} else {
parser.skipChildren();
}
mapValue = Collections.singletonMap("f", subFields);
} else {
// Fallback in case an object contains unexpected extra keys alongside "v" or "f".
if (mapValue == null) {
mapValue = new HashMap<>();
if (hasValue) {
mapValue.put("v", unwrappedValue);
hasValue = false;
}
}
mapValue.put(key, parseCellValue(parser));
Comment on lines +2178 to +2185

@gemini-code-assist gemini-code-assist Bot Oct 12, 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.

high

If the JSON object contains an 'f' field followed by unexpected extra keys, mapValue is initialized as an immutable singleton map via Collections.singletonMap("f", subFields). In the next iteration, attempting to call mapValue.put(key, ...) will throw an UnsupportedOperationException.

To prevent this, copy the map to a mutable HashMap if it is not already a HashMap before adding extra keys.

          if (mapValue == null) {
            mapValue = new HashMap<>();
            if (hasValue) {
              mapValue.put("v", unwrappedValue);
              hasValue = false;
            }
          } else if (!(mapValue instanceof HashMap)) {
            mapValue = new HashMap<>(mapValue);
          }
          mapValue.put(key, parseCellValue(parser));

}
}
if (mapValue != null) {
return mapValue;
}
return hasValue ? unwrappedValue : Collections.emptyMap();
}
if (token == JsonToken.START_ARRAY) {
List<Object> list = new ArrayList<>();
while (parser.nextToken() != JsonToken.END_ARRAY) {
list.add(parseCellValue(parser));
}
return list;
}
Comment on lines +2193 to +2199

@gemini-code-assist gemini-code-assist Bot Oct 12, 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.

critical

If the JSON stream is truncated or malformed, parser.nextToken() can return null (representing EOF). Since null != JsonToken.END_ARRAY is true, the loop will continue infinitely, adding JSON_NULL to the list and eventually causing an OutOfMemoryError.

To prevent this, ensure the loop condition also checks for null (EOF).

Suggested change
if (token == JsonToken.START_ARRAY) {
List<Object> list = new ArrayList<>();
while (parser.nextToken() != JsonToken.END_ARRAY) {
list.add(parseCellValue(parser));
}
return list;
}
if (token == JsonToken.START_ARRAY) {
List<Object> list = new ArrayList<>();
JsonToken arrayToken;
while ((arrayToken = parser.nextToken()) != JsonToken.END_ARRAY && arrayToken != null) {
list.add(parseCellValue(parser));
}
return list;
}

return parser.getText();
}

@Override
public String open(Job loadJob) {
try {
Expand Down
14 changes: 14 additions & 0 deletions ...igquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/FieldValueTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -234,4 +234,18 @@ public void testParseCanonicalInterval() {
System.out.println(FieldValue.parseCanonicalInterval(entry.getKey()));
}
}

@Test
public void testFromPbUnwrappedCells() {
assertEquals(FieldValue.fromPb(BOOLEAN_FIELD), FieldValue.fromPb("false"));
assertEquals(FieldValue.fromPb(NULL_FIELD), FieldValue.fromPb(null));
assertEquals(FieldValue.fromPb(NULL_FIELD), FieldValue.fromPb(Data.nullOf(Object.class)));
assertEquals(
FieldValue.fromPb(NULL_FIELD),
FieldValue.fromPb(new TableCell().setV(Data.nullOf(Object.class))));
assertEquals(FieldValue.fromPb(REPEATED_FIELD), FieldValue.fromPb(ImmutableList.of("1", "1")));
assertEquals(
FieldValue.fromPb(RECORD_FIELD),
FieldValue.fromPb(ImmutableMap.of("f", ImmutableList.of("1.5", "42"))));
}
}
Loading
Loading