How can I use bigquery storage to get the result of a query

[fabiencelier]

I'd like to download the result of a query as fast as possible.
The examples in https://github.com/googleapis/java-bigquerystorage/tree/master/samples/snippets/src/main/java/com/example/bigquerystorage are all about downloading/uploading to a table.

How can use this library to send a query and download the result ? Is there any example about it somewhere ?

[stephaniewang526]

Could you provide me with the specific endpoints you are asking about?

[fabiencelier]

This example shows how to retrieve a table, I'd like to retrieve the result of a query in a similar way. Is this possible ?

[kmjung]

Sure -- you can look at the Apache Beam BigQuery query source as one example. This won't necessarily be faster than the default method of fetching query results, however, particularly for small result sets. What is your use case?

[fabiencelier]

What is your use case?

I want to do a query and download the result into Java memory as fast as possible.
My query result is quite big: several million rows

This won't necessarily be faster than the default method of fetching query results

Are you talking about the google-cloud-bigquery library ? It is very slow for me but I might not use it correctly. I do something like that:

BigQuery bigquery = BigQueryOptions.getDefaultInstance().getService();
TableResult result = bigquery.query(queryConfig);
for (FieldValueList row : result.iterateAll()) {
  // do stuff
}

It takes around 3 minutes to download 10M rows

Instead I have tried to do a query with the bigquery API then download the temporary result table with the bigquerystorage API, but I'm not sure if it is the good way to do it:

Job queryJob = bigquery.create(JobInfo.newBuilder(queryConfig).build());
queryJob = queryJob.waitFor();
TableId temporaryTable = ((QueryJobConfiguration) queryJob.getConfiguration()).getDestinationTable();
// Use bigquerystorage to download the temporary table to Arrow

It takes around 10 seconds to download the same 10M rows.

you can look at the Apache Beam BigQuery query source as one example.

I am not familiar with Beam, I'll have a look at it :)

[stephaniewang526]

What version of google-cloud-bigquery are you using to run the query? When you say "download" the result, what do you mean? I am curious as to why it is so slow. Does it take 3min to get the results back? Or 3min to "download". Thanks.

[fabiencelier]

What version of google-cloud-bigquery are you using to run the query

I am using

  <dependencyManagement>
    <dependencies>
      <dependency>
        <groupId>com.google.cloud</groupId>
        <artifactId>libraries-bom</artifactId>
        <version>16.2.1</version>
        <type>pom</type>
        <scope>import</scope>
      </dependency>
    </dependencies>
  </dependencyManagement>

which gives me google-cloud-bigquery 1.126.1

When you say "download" the result, what do you mean?

I mean the time it takes to retrieve all the data into my local JVM

This code takes several minutes to iterate over 10M rows:

TableResult result = bigquery.query(queryConfig);

long count = 0;
for (FieldValueList row : result.iterateAll()) {
  for (FieldValue val : row) {
    count++;
  }
}

[stephaniewang526]

How wide is your row? How much data roughly is in your 10M+ row? I want to compare this to our benchmark to see if it's normal.

Also, can you share a code snippet of how you are doing the same operation using the bigquerystorage API which is taking 10s instead of 3min?

Thanks!

[fabiencelier]

I have 10M rows with 3 columns (integer, string, float) and as CSV it represents ~200MB.
Unless there is something I am missing, the google-cloud-bigquery library is not option for medium to large results as it is way too slow.

Here is my benchmark code for bigquerystorage API (with 1 thread to keep it simple):

Example.java

import com.google.api.gax.rpc.ServerStream;
import com.google.cloud.bigquery.BigQuery;
import com.google.cloud.bigquery.BigQueryOptions;
import com.google.cloud.bigquery.Job;
import com.google.cloud.bigquery.JobInfo;
import com.google.cloud.bigquery.QueryJobConfiguration;
import com.google.cloud.bigquery.TableId;
import com.google.cloud.bigquery.storage.v1.BigQueryReadClient;
import com.google.cloud.bigquery.storage.v1.CreateReadSessionRequest;
import com.google.cloud.bigquery.storage.v1.DataFormat;
import com.google.cloud.bigquery.storage.v1.ReadRowsRequest;
import com.google.cloud.bigquery.storage.v1.ReadRowsResponse;
import com.google.cloud.bigquery.storage.v1.ReadSession;
import com.google.cloud.bigquery.storage.v1.ReadSession.TableReadOptions;
import com.google.common.base.Preconditions;
import io.atoti.Config.MyTable;
import java.io.IOException;

public class Example {

  public static final BigQuery bigquery = BigQueryOptions.getDefaultInstance().getService();

  public static void main(String[] args) throws InterruptedException, IOException {
    final String query = "SELECT * FROM table"; // whatever query with big result

    long start = System.currentTimeMillis();
    TableId temporaryResultTable = getTemporaryTable(query);
    long queryDone = System.currentTimeMillis();
    downloadTable(temporaryResultTable);
    long end = System.currentTimeMillis();


    System.out.println("query time:\t" + (queryDone-start) +"ms");
    System.out.println("download time:\t" + (end-queryDone) +"ms");
    System.out.println("total time:\t" + (end-start) +"ms");

  }

  public static TableId getTemporaryTable(String query) throws InterruptedException{
    QueryJobConfiguration queryConfig = QueryJobConfiguration.newBuilder(query).setUseLegacySql(false).build();
    Job queryJob = bigquery.create(JobInfo.newBuilder(queryConfig).build());

    // Wait for the query to complete.
    queryJob = queryJob.waitFor();
    return ((QueryJobConfiguration) queryJob.getConfiguration()).getDestinationTable();
  }


  public static void downloadTable(TableId table) throws IOException {
    String parent = String.format("projects/%s", table.getProject());
    String srcTable =
        String.format(
            "projects/%s/datasets/%s/tables/%s",
            table.getProject(), table.getDataset(), table.getTable());

    TableReadOptions options = TableReadOptions.newBuilder().build();

    ReadSession.Builder sessionBuilder =
        ReadSession.newBuilder()
            .setTable(srcTable)
            .setDataFormat(DataFormat.ARROW)
            .setReadOptions(options);

    CreateReadSessionRequest.Builder builder =
        CreateReadSessionRequest.newBuilder()
            .setParent(parent)
            .setReadSession(sessionBuilder)
            .setMaxStreamCount(1);

    try (BigQueryReadClient client = BigQueryReadClient.create()) {
      ReadSession session = client.createReadSession(builder.build());
      String streamName = session.getStreams(0).getName();

      try (SimpleArrowRowReader reader = new SimpleArrowRowReader(session.getArrowSchema())) {
        ReadRowsRequest readRowsRequest =
            ReadRowsRequest.newBuilder().setReadStream(streamName).build();

        ServerStream<ReadRowsResponse> stream = client.readRowsCallable().call(readRowsRequest);
        for (ReadRowsResponse response : stream) {
          Preconditions.checkState(response.hasArrowRecordBatch());
          reader.processRows(response.getArrowRecordBatch());
        }
      }
    }
  }
}

SimpleArrowRowReader.java

import com.google.cloud.bigquery.storage.v1.ArrowRecordBatch;
import com.google.cloud.bigquery.storage.v1.ArrowSchema;
import com.google.common.base.Preconditions;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.VectorLoader;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.ipc.ReadChannel;
import org.apache.arrow.vector.ipc.message.MessageSerializer;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.Schema;
import org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel;

public class SimpleArrowRowReader implements AutoCloseable {

    BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
    private final VectorSchemaRoot root;
    private final VectorLoader loader;
    private long count = 0;

    public SimpleArrowRowReader(ArrowSchema arrowSchema) throws IOException {
      Schema schema =
          MessageSerializer.deserializeSchema(
              new ReadChannel(
                  new ByteArrayReadableSeekableByteChannel(
                      arrowSchema.getSerializedSchema().toByteArray())));
      Preconditions.checkNotNull(schema);
      List<FieldVector> vectors = new ArrayList<>();
      for (Field field : schema.getFields()) {
        vectors.add(field.createVector(allocator));
      }
      root = new VectorSchemaRoot(vectors);
      loader = new VectorLoader(root);
    }

    public void processRows(ArrowRecordBatch batch) throws IOException {
      org.apache.arrow.vector.ipc.message.ArrowRecordBatch deserializedBatch =
          MessageSerializer.deserializeRecordBatch(
              new ReadChannel(
                  new ByteArrayReadableSeekableByteChannel(
                      batch.getSerializedRecordBatch().toByteArray())),
              allocator);

      loader.load(deserializedBatch);
      deserializedBatch.close();
      long batchSize = root.getRowCount();
      // System.out.println("\tBatch size: "+ count);
      count += batchSize;
      root.clear();
    }

    @Override
    public void close() {
      System.out.println("Stream size: " + count);
      root.close();
      allocator.close();
    }
}

which outputs

Stream size: 10000000
query time:	928ms
download time:	8523ms
total time:	9451ms

I wonder if that is a viable solution because:

• I am using the temporary table of the result so I don't know if it will exists until the end of the download
• I have to cast to QueryJobConfiguration so it doesn't seems a part of the public API

But I don't see any solution more "official" for doing that with the same speed

[stephaniewang526]

This seems reasonable. We potentially will be doing a BigQuery client re-write and this will be taken into consideration when we do that. Thank you!