ConnImplBenchmark.java

/*
 * Copyright 2022 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.google.cloud.bigquery;

import java.io.IOException;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;
import java.util.logging.Level;
import org.openjdk.jmh.annotations.Benchmark;
import org.openjdk.jmh.annotations.BenchmarkMode;
import org.openjdk.jmh.annotations.Fork;
import org.openjdk.jmh.annotations.Measurement;
import org.openjdk.jmh.annotations.Mode;
import org.openjdk.jmh.annotations.OutputTimeUnit;
import org.openjdk.jmh.annotations.Param;
import org.openjdk.jmh.annotations.Scope;
import org.openjdk.jmh.annotations.Setup;
import org.openjdk.jmh.annotations.State;
import org.openjdk.jmh.annotations.Warmup;
import org.openjdk.jmh.infra.Blackhole;
import org.openjdk.jmh.runner.Runner;
import org.openjdk.jmh.runner.options.Options;
import org.openjdk.jmh.runner.options.OptionsBuilder;

@Fork(value = 1)
@BenchmarkMode(Mode.AverageTime)
@Warmup(iterations = 1)
@Measurement(iterations = 3)
@State(Scope.Benchmark)
@OutputTimeUnit(TimeUnit.MILLISECONDS)
public class ConnImplBenchmark {
  @Param({"500000", "1000000", "10000000", "50000000", "100000000"}) // 500K, 1M, 10M, 50M and 100M
  public int rowLimit;

  private ConnectionSettings connectionSettingsReadAPIEnabled, connectionSettingsReadAPIDisabled;
  private final String QUERY =
      "SELECT * FROM bigquery-public-data.new_york_taxi_trips.tlc_yellow_trips_2017 LIMIT %s";

  @Setup
  public void setUp() throws IOException {
    java.util.logging.Logger.getGlobal().setLevel(Level.ALL);

    connectionSettingsReadAPIEnabled = ConnectionSettings.newBuilder().setUseReadAPI(true).build();
    connectionSettingsReadAPIDisabled =
        ConnectionSettings.newBuilder().setUseReadAPI(false).build();
  }

  @Benchmark
  public void iterateRecordsWithBigQuery_Query(Blackhole blackhole) throws InterruptedException {
    String selectQuery = String.format(QUERY, rowLimit);
    BigQuery bigQuery = BigQueryOptions.getDefaultInstance().getService();
    QueryJobConfiguration config =
        QueryJobConfiguration.newBuilder(selectQuery).setUseLegacySql(false).build();
    TableResult result = bigQuery.query(config);
    long hash = 0L;
    int cnt = 0;
    long lastTime = System.currentTimeMillis();
    System.out.println("\n Running");
    for (FieldValueList row : result.iterateAll()) {
      hash += computeHash(row.get("vendor_id"), FieldValue::getStringValue);
      hash += computeHash(row.get("pickup_datetime"), FieldValue::getStringValue);
      hash += computeHash(row.get("dropoff_datetime"), FieldValue::getStringValue);
      hash += computeHash(row.get("passenger_count"), FieldValue::getLongValue);
      hash += computeHash(row.get("trip_distance"), FieldValue::getDoubleValue);
      hash += computeHash(row.get("rate_code"), FieldValue::getStringValue);
      hash += computeHash(row.get("store_and_fwd_flag"), FieldValue::getStringValue);
      hash += computeHash(row.get("payment_type"), FieldValue::getStringValue);
      hash += computeHash(row.get("fare_amount"), FieldValue::getDoubleValue);
      hash += computeHash(row.get("extra"), FieldValue::getDoubleValue);
      hash += computeHash(row.get("mta_tax"), FieldValue::getDoubleValue);
      hash += computeHash(row.get("tip_amount"), FieldValue::getDoubleValue);
      hash += computeHash(row.get("tolls_amount"), FieldValue::getDoubleValue);
      hash += computeHash(row.get("imp_surcharge"), FieldValue::getDoubleValue);
      hash += computeHash(row.get("airport_fee"), FieldValue::getDoubleValue);
      hash += computeHash(row.get("total_amount"), FieldValue::getDoubleValue);
      hash += computeHash(row.get("pickup_location_id"), FieldValue::getStringValue);
      hash += computeHash(row.get("dropoff_location_id"), FieldValue::getStringValue);
      hash += computeHash(row.get("data_file_year"), FieldValue::getLongValue);
      hash += computeHash(row.get("data_file_month"), FieldValue::getLongValue);

      if (++cnt % 100_000 == 0) {
        long now = System.currentTimeMillis();
        long duration = now - lastTime;
        System.out.println("ROW " + cnt + " Time: " + duration + " ms");
        lastTime = now;
      }
    }
    System.out.println(cnt + " records processed using bigquery.query");
    blackhole.consume(hash);
  }

  @Benchmark
  public void iterateRecordsUsingReadAPI(Blackhole blackhole)
      throws InterruptedException, BigQuerySQLException {
    Connection connectionReadAPIEnabled =
        BigQueryOptions.getDefaultInstance()
            .getService()
            .createConnection(connectionSettingsReadAPIEnabled);
    String selectQuery = String.format(QUERY, rowLimit);
    long hash = 0L;
    try {
      BigQueryResult bigQueryResultSet = connectionReadAPIEnabled.executeSelect(selectQuery);
      hash = getResultHash(bigQueryResultSet);
    } catch (Exception e) {
      e.printStackTrace();
    } finally {
      connectionReadAPIEnabled.close(); // IMP to kill the bg workers
    }
    blackhole.consume(hash);
  }

  @Benchmark
  public void iterateRecordsWithoutUsingReadAPI(Blackhole blackhole)
      throws InterruptedException, BigQuerySQLException {
    Connection connectionReadAPIDisabled =
        BigQueryOptions.getDefaultInstance()
            .getService()
            .createConnection(connectionSettingsReadAPIDisabled);
    String selectQuery = String.format(QUERY, rowLimit);
    long hash = 0L;
    try {
      BigQueryResult bigQueryResultSet = connectionReadAPIDisabled.executeSelect(selectQuery);
      hash = getResultHash(bigQueryResultSet);
    } catch (Exception e) {
      e.printStackTrace();
    } finally {
      connectionReadAPIDisabled.close(); // IMP to kill the bg workers
    }
    blackhole.consume(hash);
  }

  private long getResultHash(BigQueryResult bigQueryResultSet) throws SQLException {
    ResultSet rs = bigQueryResultSet.getResultSet();
    long hash = 0L;
    int cnt = 0;
    long lastTime = System.currentTimeMillis();
    System.out.println("\n Running");
    while (rs.next()) {
      hash += computeHash(rs, "vendor_id", ResultSet::getString);
      hash += computeHash(rs, "pickup_datetime", ResultSet::getString);
      hash += computeHash(rs, "dropoff_datetime", ResultSet::getString);
      hash += computeHash(rs, "passenger_count", ResultSet::getLong);
      hash += computeHash(rs, "trip_distance", ResultSet::getDouble);
      hash += computeHash(rs, "rate_code", ResultSet::getString);
      hash += computeHash(rs, "store_and_fwd_flag", ResultSet::getString);
      hash += computeHash(rs, "payment_type", ResultSet::getString);
      hash += computeHash(rs, "fare_amount", ResultSet::getDouble);
      hash += computeHash(rs, "extra", ResultSet::getDouble);
      hash += computeHash(rs, "mta_tax", ResultSet::getDouble);
      hash += computeHash(rs, "tip_amount", ResultSet::getDouble);
      hash += computeHash(rs, "tolls_amount", ResultSet::getDouble);
      hash += computeHash(rs, "imp_surcharge", ResultSet::getDouble);
      hash += computeHash(rs, "airport_fee", ResultSet::getDouble);
      hash += computeHash(rs, "total_amount", ResultSet::getDouble);
      hash += computeHash(rs, "pickup_location_id", ResultSet::getString);
      hash += computeHash(rs, "dropoff_location_id", ResultSet::getString);
      hash += computeHash(rs, "data_file_year", ResultSet::getLong);
      hash += computeHash(rs, "data_file_month", ResultSet::getLong);

      if (++cnt % 100_000 == 0) {
        long now = System.currentTimeMillis();
        long duration = now - lastTime;
        System.out.println("ROW " + cnt + " Time: " + duration + " ms");
        lastTime = now;
      }
    }
    return hash;
  }

  private <T> long computeHash(
      ResultSet rs, String columnName, SQLFunction<ResultSet, T> extractor) {
    try {
      T value = extractor.apply(rs, columnName);
      return (value == null) ? 0 : value.hashCode();
    } catch (SQLException e) {
      return 0;
    }
  }

  @FunctionalInterface
  private interface SQLFunction<T, R> {
    R apply(T t, String columnName) throws SQLException;
  }

  private <T> long computeHash(FieldValue fieldValue, Function<FieldValue, T> extractor) {
    if (fieldValue == null || fieldValue.isNull()) {
      return 0;
    }
    T value = extractor.apply(fieldValue);
    return (value == null) ? 0 : value.hashCode();
  }

  public static void main(String[] args) throws Exception {
    Options opt = new OptionsBuilder().include(ConnImplBenchmark.class.getSimpleName()).build();
    new Runner(opt).run();
  }
}