Skip to content

perf(bigquery): [T2-2] stream Jackson tokens directly into FieldValueList without intermediate TableRow #30

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

Open
lqiu96 wants to merge 1 commit into
base: perf/bigquery-unwrap-cell-json
Choose a base branch
from
Open

perf(bigquery): [T2-2] stream Jackson tokens directly into FieldValueList without intermediate TableRow #30

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
38 changes: 27 additions & 11 deletions ...-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,7 @@
import java.io.IOException;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
Expand Down Expand Up @@ -2219,17 +2220,32 @@ public TableDataList call() throws IOException {

private static Iterable<FieldValueList> transformTableData(
Iterable<TableRow> tableDataPb, final Schema schema, boolean useInt64Timestamps) {
return ImmutableList.copyOf(
Iterables.transform(
tableDataPb != null ? tableDataPb : ImmutableList.<TableRow>of(),
new Function<TableRow, FieldValueList>() {
FieldList fields = schema != null ? schema.getFields() : null;

@Override
public FieldValueList apply(TableRow rowPb) {
return FieldValueList.fromPb(rowPb.getF(), fields, useInt64Timestamps);
}
}));
if (tableDataPb == null) {
return ImmutableList.of();
}
List<FieldValueList> streamed =
HttpBigQueryRpc.extractFieldValueLists(tableDataPb, schema, useInt64Timestamps);
if (streamed != null) {
return streamed;
}
FieldList fields = schema != null ? schema.getFields() : null;
if (tableDataPb instanceof Collection) {
Collection<TableRow> rows = (Collection<TableRow>) tableDataPb;
if (rows.isEmpty()) {
return ImmutableList.of();
}
ImmutableList.Builder<FieldValueList> builder =
ImmutableList.builderWithExpectedSize(rows.size());
for (TableRow rowPb : rows) {
builder.add(FieldValueList.fromPb(rowPb.getF(), fields, useInt64Timestamps));
}
return builder.build();
}
ImmutableList.Builder<FieldValueList> builder = ImmutableList.builder();
for (TableRow rowPb : tableDataPb) {
builder.add(FieldValueList.fromPb(rowPb.getF(), fields, useInt64Timestamps));
}
return builder.build();
}

@Override
Expand Down
39 changes: 28 additions & 11 deletions ...igquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ConnectionImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
import com.google.cloud.bigquery.JobStatistics.QueryStatistics;
import com.google.cloud.bigquery.JobStatistics.SessionInfo;
import com.google.cloud.bigquery.spi.v2.BigQueryRpc;
import com.google.cloud.bigquery.spi.v2.HttpBigQueryRpc;
import com.google.cloud.bigquery.storage.v1.ArrowRecordBatch;
import com.google.cloud.bigquery.storage.v1.ArrowSchema;
import com.google.cloud.bigquery.storage.v1.BigQueryReadClient;
Expand All @@ -57,6 +58,7 @@
import java.math.BigInteger;
import java.util.AbstractList;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedList;
Expand Down Expand Up @@ -825,17 +827,32 @@ void flagEndOfStream() { // package-private
/* Helper method that parse and populate a page with TableRows */
private static Iterable<FieldValueList> getIterableFieldValueList(
Iterable<TableRow> tableDataPb, final Schema schema) {
return ImmutableList.copyOf(
Iterables.transform(
tableDataPb != null ? tableDataPb : ImmutableList.of(),
new Function<TableRow, FieldValueList>() {
final FieldList fields = schema != null ? schema.getFields() : null;

@Override
public FieldValueList apply(TableRow rowPb) {
return FieldValueList.fromPb(rowPb.getF(), fields);
}
}));
if (tableDataPb == null) {
return ImmutableList.of();
}
List<FieldValueList> streamed =
HttpBigQueryRpc.extractFieldValueLists(tableDataPb, schema, false);
if (streamed != null) {
return streamed;
}
FieldList fields = schema != null ? schema.getFields() : null;
if (tableDataPb instanceof Collection) {
Collection<TableRow> rows = (Collection<TableRow>) tableDataPb;
if (rows.isEmpty()) {
return ImmutableList.of();
}
ImmutableList.Builder<FieldValueList> builder =
ImmutableList.builderWithExpectedSize(rows.size());
for (TableRow rowPb : rows) {
builder.add(FieldValueList.fromPb(rowPb.getF(), fields));
}
return builder.build();
}
ImmutableList.Builder<FieldValueList> builder = ImmutableList.builder();
for (TableRow rowPb : tableDataPb) {
builder.add(FieldValueList.fromPb(rowPb.getF(), fields));
}
return builder.build();
}

/* Helper method that determines the optimal number of caches pages to improve read performance */
Expand Down
7 changes: 7 additions & 0 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 @@ -368,13 +368,20 @@ public final boolean equals(Object obj) {
* href="https://github.com/googleapis/google-cloud-java/pull/2891">this discussion</a> for more
* context.
*/
private static final FieldValue NULL_PRIMITIVE = new FieldValue(Attribute.PRIMITIVE, null, false);
private static final FieldValue NULL_PRIMITIVE_INT64_TIMESTAMPS =
new FieldValue(Attribute.PRIMITIVE, null, true);

@BetaApi
public static FieldValue of(Attribute attribute, Object value) {
return of(attribute, value, false);
}

@BetaApi
public static FieldValue of(Attribute attribute, Object value, Boolean useInt64Timestamps) {
if (value == null && attribute == Attribute.PRIMITIVE && useInt64Timestamps != null) {
return useInt64Timestamps ? NULL_PRIMITIVE_INT64_TIMESTAMPS : NULL_PRIMITIVE;
}
return new FieldValue(attribute, value, useInt64Timestamps);
}

Expand Down
30 changes: 17 additions & 13 deletions ...igquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/FieldValueList.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,8 @@
import com.google.common.collect.ImmutableList;
import java.io.Serializable;
import java.util.AbstractList;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Objects;

/**
* Google BigQuery Table Field Values class, which represents a row in returned query result (table
Expand Down Expand Up @@ -108,6 +107,9 @@ public static FieldValueList of(List<FieldValue> row, Field... schema) {
}

FieldValueList withSchema(FieldList schema) {
if (this.schema == schema || Objects.equals(this.schema, schema)) {
return this;
}
return new FieldValueList(this.row, schema);
}

Expand All @@ -116,22 +118,24 @@ static FieldValueList fromPb(List<?> rowPb, FieldList schema) {
}

static FieldValueList fromPb(List<?> rowPb, FieldList schema, Boolean useInt64Timestamps) {
List<FieldValue> row = new ArrayList<>(rowPb.size());
if (rowPb == null) {
return FieldValueList.of(ImmutableList.of(), schema);
}
int size = rowPb.size();
if (schema != null && schema.size() != size) {
throw new IllegalArgumentException("Row size and fields schema sizes should match");
}
ImmutableList.Builder<FieldValue> row = ImmutableList.builderWithExpectedSize(size);
if (schema != null) {
if (schema.size() != rowPb.size()) {
throw new IllegalArgumentException("Row size and fields schema sizes should match");
}
Iterator<Field> schemaIter = schema.iterator();
Iterator<?> rowPbIter = rowPb.iterator();
while (rowPbIter.hasNext() && schemaIter.hasNext()) {
row.add(FieldValue.fromPb(rowPbIter.next(), schemaIter.next(), useInt64Timestamps));
for (int i = 0; i < size; i++) {
row.add(FieldValue.fromPb(rowPb.get(i), schema.get(i), useInt64Timestamps));
}
} else {
for (Object cellPb : rowPb) {
row.add(FieldValue.fromPb(cellPb, null, useInt64Timestamps));
for (int i = 0; i < size; i++) {
row.add(FieldValue.fromPb(rowPb.get(i), null, useInt64Timestamps));
}
}

return FieldValueList.of(row, schema);
return FieldValueList.of(row.build(), schema);
}
}
Loading
Loading