content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you 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 org.apache.beam.sdk.extensions.zetasketch;
  
  import static org.apache.beam.sdk.io.gcp.testing.BigqueryMatcher.createQueryUsingStandardSql;
  import static org.apache.beam.sdk.io.gcp.testing.BigqueryMatcher.queryResultHasChecksum;
  import static org.hamcrest.MatcherAssert.assertThat;
  
  import com.google.api.services.bigquery.model.Table;
  import com.google.api.services.bigquery.model.TableFieldSchema;
  import com.google.api.services.bigquery.model.TableReference;
  import com.google.api.services.bigquery.model.TableRow;
  import com.google.api.services.bigquery.model.TableSchema;
  import com.google.cloud.bigquery.storage.v1.DataFormat;
  import java.nio.ByteBuffer;
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.Date;
  import java.util.List;
  import java.util.Map;
  import java.util.stream.Collectors;
  import org.apache.beam.sdk.Pipeline;
  import org.apache.beam.sdk.coders.ByteArrayCoder;
  import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
  import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO;
  import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.TypedRead.Method;
  import org.apache.beam.sdk.io.gcp.bigquery.SchemaAndRecord;
  import org.apache.beam.sdk.io.gcp.testing.BigqueryClient;
  import org.apache.beam.sdk.options.ApplicationNameOptions;
  import org.apache.beam.sdk.testing.PAssert;
  import org.apache.beam.sdk.testing.TestPipeline;
  import org.apache.beam.sdk.testing.TestPipelineOptions;
  import org.apache.beam.sdk.transforms.Create;
  import org.apache.beam.sdk.transforms.SerializableFunction;
  import org.apache.beam.sdk.values.PCollection;
  import org.apache.beam.sdk.values.TypeDescriptor;
  import org.junit.AfterClass;
  import org.junit.BeforeClass;
  import org.junit.Test;
  import org.junit.runner.RunWith;
  import org.junit.runners.JUnit4;
  
  /**
   * Integration tests for HLL++ sketch compatibility between Beam and BigQuery. The tests verifies
   * that HLL++ sketches created in Beam can be processed by BigQuery, and vice versa.
   */
  @RunWith(JUnit4.class)
  public class BigQueryHllSketchCompatibilityIT {
  
    private static final String APP_NAME;
    private static final String PROJECT_ID;
    private static final String DATASET_ID;
    private static final BigqueryClient BIGQUERY_CLIENT;
  
    private static final List<String> TEST_DATA =
        Arrays.asList("Apple", "Orange", "Banana", "Orange");
  
    // Data Table: used by tests reading sketches from BigQuery
    // Schema: only one STRING field named "data".
    private static final String DATA_FIELD_NAME = "data";
    private static final String DATA_FIELD_TYPE = "STRING";
    private static final String QUERY_RESULT_FIELD_NAME = "sketch";
  
    // Content: prepopulated with 4 rows: "Apple", "Orange", "Banana", "Orange"
    private static final String DATA_TABLE_ID_NON_EMPTY = "hll_data_non_empty";
    private static final Long EXPECTED_COUNT_NON_EMPTY = 3L;
  
    // Content: empty
    private static final String DATA_TABLE_ID_EMPTY = "hll_data_empty";
    private static final Long EXPECTED_COUNT_EMPTY = 0L;
  
    // Sketch Table: used by tests writing sketches to BigQuery
    // Schema: only one BYTES field named "sketch".
    private static final String SKETCH_FIELD_NAME = "sketch";
    private static final String SKETCH_FIELD_TYPE = "BYTES";
  
    // Content: will be overridden by the sketch computed by the test pipeline each time the test runs
    private static final String SKETCH_TABLE_ID = "hll_sketch";
    // SHA-1 hash of string "[3]", the string representation of a row that has only one field 3 in it
    private static final String EXPECTED_CHECKSUM_NON_EMPTY =
        "f1e31df9806ce94c5bdbbfff9608324930f4d3f1";
    // SHA-1 hash of string "[0]", the string representation of a row that has only one field 0 in it
    private static final String EXPECTED_CHECKSUM_EMPTY = "1184f5b8d4b6dd08709cf1513f26744167065e0d";
  
    static {
      ApplicationNameOptions options =
          TestPipeline.testingPipelineOptions().as(ApplicationNameOptions.class);
      APP_NAME = options.getAppName();
      PROJECT_ID = options.as(GcpOptions.class).getProject();
      DATASET_ID = String.format("zetasketch_%tY_%<tm_%<td_%<tH_%<tM_%<tS_%<tL", new Date());
      BIGQUERY_CLIENT = BigqueryClient.getClient(APP_NAME);
    }
  
    @BeforeClass
    public static void prepareDatasetAndDataTables() throws Exception {
      BIGQUERY_CLIENT.createNewDataset(PROJECT_ID, DATASET_ID);
  
      TableSchema dataTableSchema =
          new TableSchema()
              .setFields(
                  Collections.singletonList(
                      new TableFieldSchema().setName(DATA_FIELD_NAME).setType(DATA_FIELD_TYPE)));
  
      Table dataTableNonEmpty =
          new Table()
              .setSchema(dataTableSchema)
              .setTableReference(
                  new TableReference()
                      .setProjectId(PROJECT_ID)
                      .setDatasetId(DATASET_ID)
                      .setTableId(DATA_TABLE_ID_NON_EMPTY));
      BIGQUERY_CLIENT.createNewTable(PROJECT_ID, DATASET_ID, dataTableNonEmpty);
      // Prepopulates dataTableNonEmpty with TEST_DATA
      List<Map<String, Object>> rows =
          TEST_DATA.stream()
              .map(v -> Collections.singletonMap(DATA_FIELD_NAME, (Object) v))
              .collect(Collectors.toList());
      BIGQUERY_CLIENT.insertDataToTable(PROJECT_ID, DATASET_ID, DATA_TABLE_ID_NON_EMPTY, rows);
  
      Table dataTableEmpty =
          new Table()
              .setSchema(dataTableSchema)
              .setTableReference(
                  new TableReference()
                      .setProjectId(PROJECT_ID)
                      .setDatasetId(DATASET_ID)
                      .setTableId(DATA_TABLE_ID_EMPTY));
      BIGQUERY_CLIENT.createNewTable(PROJECT_ID, DATASET_ID, dataTableEmpty);
    }
  
    @AfterClass
    public static void deleteDataset() throws Exception {
      BIGQUERY_CLIENT.deleteDataset(PROJECT_ID, DATASET_ID);
    }
  
    /**
     * Tests that a non-empty HLL++ sketch computed in BigQuery can be processed by Beam.
     *
     * <p>The Hll sketch is computed by {@code HLL_COUNT.INIT} in BigQuery and read into Beam; the
     * test verifies that we can run {@link HllCount.MergePartial} and {@link HllCount.Extract} on the
     * sketch in Beam to get the correct estimated count.
     */
    @Test
    public void testReadNonEmptySketchFromBigQuery() {
      readSketchFromBigQuery(DATA_TABLE_ID_NON_EMPTY, EXPECTED_COUNT_NON_EMPTY);
    }
  
    /**
     * Tests that an empty HLL++ sketch computed in BigQuery can be processed by Beam.
     *
     * <p>The Hll sketch is computed by {@code HLL_COUNT.INIT} in BigQuery and read into Beam; the
     * test verifies that we can run {@link HllCount.MergePartial} and {@link HllCount.Extract} on the
     * sketch in Beam to get the correct estimated count.
     */
    @Test
    public void testReadEmptySketchFromBigQuery() {
      readSketchFromBigQuery(DATA_TABLE_ID_EMPTY, EXPECTED_COUNT_EMPTY);
    }
  
    private void readSketchFromBigQuery(String tableId, Long expectedCount) {
      String tableSpec = String.format("%s.%s", DATASET_ID, tableId);
      String query =
          String.format(
              "SELECT HLL_COUNT.INIT(%s) AS %s FROM %s",
              DATA_FIELD_NAME, QUERY_RESULT_FIELD_NAME, tableSpec);
  
      SerializableFunction<SchemaAndRecord, byte[]> parseQueryResultToByteArray =
          input ->
              // BigQuery BYTES type corresponds to Java java.nio.ByteBuffer type
              HllCount.getSketchFromByteBuffer(
                  (ByteBuffer) input.getRecord().get(QUERY_RESULT_FIELD_NAME));
  
      TestPipelineOptions options =
          TestPipeline.testingPipelineOptions().as(TestPipelineOptions.class);
  
      Pipeline p = Pipeline.create(options);
      PCollection<Long> result =
          p.apply(
                  BigQueryIO.read(parseQueryResultToByteArray)
                      .withFormat(DataFormat.AVRO)
                      .fromQuery(query)
                      .usingStandardSql()
                      .withMethod(Method.DIRECT_READ)
                      .withCoder(ByteArrayCoder.of()))
              .apply(HllCount.MergePartial.globally()) // no-op, only for testing MergePartial
              .apply(HllCount.Extract.globally());
      PAssert.thatSingleton(result).isEqualTo(expectedCount);
      p.run().waitUntilFinish();
    }
  
    /**
     * Tests that a non-empty HLL++ sketch computed in Beam can be processed by BigQuery.
     *
     * <p>The Hll sketch is computed by {@link HllCount.Init} in Beam and written to BigQuery; the
     * test verifies that we can run {@code HLL_COUNT.EXTRACT()} on the sketch in BigQuery to get the
     * correct estimated count.
     */
    @Test
    public void testWriteNonEmptySketchToBigQuery() {
      writeSketchToBigQuery(TEST_DATA, EXPECTED_CHECKSUM_NON_EMPTY);
    }
  
    /**
     * Tests that an empty HLL++ sketch computed in Beam can be processed by BigQuery.
     *
     * <p>The Hll sketch is computed by {@link HllCount.Init} in Beam and written to BigQuery; the
     * test verifies that we can run {@code HLL_COUNT.EXTRACT()} on the sketch in BigQuery to get the
     * correct estimated count.
     */
    @Test
    public void testWriteEmptySketchToBigQuery() {
      writeSketchToBigQuery(Collections.emptyList(), EXPECTED_CHECKSUM_EMPTY);
    }
  
    private void writeSketchToBigQuery(List<String> testData, String expectedChecksum) {
      String tableSpec = String.format("%s.%s", DATASET_ID, SKETCH_TABLE_ID);
      String query =
          String.format("SELECT HLL_COUNT.EXTRACT(%s) FROM %s", SKETCH_FIELD_NAME, tableSpec);
      TableSchema tableSchema =
          new TableSchema()
              .setFields(
                  Collections.singletonList(
                      new TableFieldSchema().setName(SKETCH_FIELD_NAME).setType(SKETCH_FIELD_TYPE)));
  
      TestPipelineOptions options =
          TestPipeline.testingPipelineOptions().as(TestPipelineOptions.class);
      Pipeline p = Pipeline.create(options);
  
      @SuppressWarnings("nullness") // until we have a stub class for BigQuery TableRow
      SerializableFunction<byte[], TableRow> formatFn =
          sketch ->
              // Empty sketch is represented by empty byte array in Beam and by null in
              // BigQuery
              new TableRow().set(SKETCH_FIELD_NAME, sketch.length == 0 ? null : sketch);
  
      p.apply(Create.of(testData).withType(TypeDescriptor.of(String.class)))
          .apply(HllCount.Init.forStrings().globally())
          .apply(
              BigQueryIO.<byte[]>write()
                  .to(tableSpec)
                  .withSchema(tableSchema)
                  .withFormatFunction(formatFn)
                  .withWriteDisposition(BigQueryIO.Write.WriteDisposition.WRITE_TRUNCATE));
      p.run().waitUntilFinish();
  
      // BigqueryMatcher will send a query to retrieve the estimated count and verifies its
      // correctness using checksum.
      assertThat(
          createQueryUsingStandardSql(APP_NAME, PROJECT_ID, query),
          queryResultHasChecksum(expectedChecksum));
    }
  }
  