diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/FieldValue.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/FieldValue.java index d11df4b956af..12544d197399 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/FieldValue.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/FieldValue.java @@ -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. 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) { @@ -412,8 +430,7 @@ static FieldValue fromPb(Object cellPb, Field recordSchema, Boolean useInt64Time FieldValueList.fromPb( (List) 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); } diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java index 41fe993b204e..95fdf939fdfa 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java @@ -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; @@ -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; @@ -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; @@ -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; @@ -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); }); } @@ -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); }); } @@ -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); }); } @@ -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); }); } @@ -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. + * + *

Why this exists: + * + *

+ * + *

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 executeAndParseRowResponse( + AbstractGoogleJsonClientRequest request, Class 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 parseRowResponse( + JsonParser parser, Class 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. + * + *

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}, 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 parseTableRows(JsonParser parser) throws IOException { + if (parser.getCurrentToken() != JsonToken.START_ARRAY) { + parser.skipChildren(); + return null; + } + List rows = new ArrayList<>(); + int expectedCols = 16; + while (parser.nextToken() != JsonToken.END_ARRAY) { + if (parser.getCurrentToken() != JsonToken.START_OBJECT) { + parser.skipChildren(); + continue; + } + List 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); + } + 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. + * + *

In the BigQuery REST v2 wire format, a cell in {@code "f"} takes one of four shapes: + * + *

    + *
  1. Primitive scalar: {@code {"v": "123"}} -> unwraps directly to {@code String} + * {@code "123"}. + *
  2. Null value: {@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). + *
  3. REPEATED (ARRAY): {@code {"v": [{"v": "a"}, {"v": "b"}]}} -> unwraps outer {@code + * "v"} to a {@code List} whose elements are recursively unwrapped ({@code ["a", + * "b"]}, or {@code List>} for {@code ARRAY}). + *
  4. RECORD (STRUCT): {@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. + * + */ + 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 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 subFields = new ArrayList<>(); + if (parser.getCurrentToken() == JsonToken.START_ARRAY) { + while (parser.nextToken() != JsonToken.END_ARRAY) { + 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)); + } + } + if (mapValue != null) { + return mapValue; + } + return hasValue ? unwrappedValue : Collections.emptyMap(); + } + if (token == JsonToken.START_ARRAY) { + List list = new ArrayList<>(); + while (parser.nextToken() != JsonToken.END_ARRAY) { + list.add(parseCellValue(parser)); + } + return list; + } + return parser.getText(); + } + @Override public String open(Job loadJob) { try { diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/FieldValueTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/FieldValueTest.java index e6b08ec6b8a3..b23a2b251bf5 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/FieldValueTest.java +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/FieldValueTest.java @@ -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")))); + } } diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpcTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpcTest.java index 4c3bf67921e6..8a29c24faa20 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpcTest.java +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpcTest.java @@ -1783,5 +1783,75 @@ public void testExecuteWithSpan_IgnoresHttpResponseException() throws Exception verifyNoSpans(); } + + @Test + public void testStreamingRowCellUnwrappingInListTableDataAndQueries() throws Exception { + String rowJson = + "\"rows\":[" + + "{\"f\":[" + + "{\"v\":\"hello\"}," + + "{\"v\":null}," + + "{\"v\":[{\"v\":\"10\"},{\"v\":\"20\"}]}," + + "{\"v\":{\"f\":[{\"v\":\"nested\"},{\"v\":null}]}}," + + "{\"v\":[{\"v\":{\"f\":[{\"v\":\"elem1\"}]}}]}" + + "]}" + + "]"; + + setMockResponse( + "{\"kind\":\"bigquery#tableDataList\",\"etag\":\"etag-1\",\"totalRows\":\"1\"," + + "\"pageToken\":\"next-page\"," + + rowJson + + "}"); + + com.google.api.services.bigquery.model.TableDataList tableDataList = + rpc.listTableDataSkipExceptionTranslation( + PROJECT_ID, DATASET_ID, TABLE_ID, new HashMap<>()); + assertEquals("bigquery#tableDataList", tableDataList.getKind()); + assertEquals("etag-1", tableDataList.getEtag()); + assertEquals(Long.valueOf(1L), tableDataList.getTotalRows()); + assertEquals("next-page", tableDataList.getPageToken()); + assertEquals(1, tableDataList.getRows().size()); + + List rawCells = tableDataList.getRows().get(0).getF(); + assertEquals(5, rawCells.size()); + assertEquals("hello", rawCells.get(0)); + assertTrue(com.google.api.client.util.Data.isNull(rawCells.get(1))); + assertEquals(Arrays.asList("10", "20"), rawCells.get(2)); + assertEquals( + Collections.singletonMap( + "f", Arrays.asList("nested", com.google.api.client.util.Data.nullOf(Object.class))), + rawCells.get(3)); + assertEquals( + Collections.singletonList( + Collections.singletonMap("f", Collections.singletonList("elem1"))), + rawCells.get(4)); + + setMockResponse( + "{\"kind\":\"bigquery#queryResponse\",\"jobComplete\":true,\"cacheHit\":false," + + "\"totalRows\":\"1\",\"totalBytesProcessed\":\"4096\",\"pageToken\":\"q-token\"," + + rowJson + + "}"); + com.google.api.services.bigquery.model.QueryResponse queryResponse = + rpc.queryRpcSkipExceptionTranslation(PROJECT_ID, new QueryRequest()); + assertTrue(queryResponse.getJobComplete()); + assertEquals(Boolean.FALSE, queryResponse.getCacheHit()); + assertEquals(java.math.BigInteger.ONE, queryResponse.getTotalRows()); + assertEquals(Long.valueOf(4096L), queryResponse.getTotalBytesProcessed()); + assertEquals("q-token", queryResponse.getPageToken()); + assertEquals(rawCells, queryResponse.getRows().get(0).getF()); + + setMockResponse( + "{\"kind\":\"bigquery#getQueryResultsResponse\",\"jobComplete\":true," + + "\"totalRows\":\"1\",\"totalBytesProcessed\":\"2048\"," + + rowJson + + "}"); + com.google.api.services.bigquery.model.GetQueryResultsResponse resultsResponse = + rpc.getQueryResultsSkipExceptionTranslation( + PROJECT_ID, JOB_ID, LOCATION, new HashMap<>()); + assertTrue(resultsResponse.getJobComplete()); + assertEquals(java.math.BigInteger.ONE, resultsResponse.getTotalRows()); + assertEquals(Long.valueOf(2048L), resultsResponse.getTotalBytesProcessed()); + assertEquals(rawCells, resultsResponse.getRows().get(0).getF()); + } } }