IcerbergIO Load test and benchmarking

.github/workflows/beam_PostCommit_Java_IO_Performance_Tests.yml

@@ -64,7 +64,7 @@ jobs:
       matrix:
         job_name: ["beam_PostCommit_Java_IO_Performance_Tests"]
         job_phrase: ["Run Java PostCommit IO Performance Tests"]
-        test_case: ["GCSPerformanceTest", "BigTablePerformanceTest", "BigQueryStorageApiStreamingPerformanceTest"]
+        test_case: ["IcebergPerformanceTest", "GCSPerformanceTest", "BigTablePerformanceTest", "BigQueryStorageApiStreamingPerformanceTest"]
     steps:
     - uses: actions/checkout@v7
     - name: Setup repository

.test-infra/tools/refresh_looker_metrics.py

@@ -34,6 +34,8 @@
     ("33", ["21", "70", "116", "69", "115"]),  # BigTableIO_Write
     ("34", ["22", "56", "96", "55", "95"]),    # TextIO_Read
     ("35", ["23", "64", "110", "63", "109"]),  # TextIO_Write
+    ("113", ["386", "388", "390", "392", "394"]),    # IcebergIO_Read
+    ("114", ["387", "389", "391", "393", "395"]),  # IcebergIO_Write
     ("75", ["258", "259", "260", "261", "262"]),  # TensorFlow MNIST
     ("76", ["233", "234", "235", "236", "237"]),  # PyTorch BERT base uncased
     ("77", ["238", "239", "240", "241", "242"]),  # PyTorch BERT large uncased

it/build.gradle

@@ -34,4 +34,8 @@ tasks.register('BigTablePerformanceTest') {
 
 tasks.register('BigQueryStorageApiStreamingPerformanceTest') {
     dependsOn(":it:google-cloud-platform:BigQueryStorageApiStreamingPerformanceTest")
-}
+}
+
+tasks.register('IcebergPerformanceTest') {
+  dependsOn(":it:iceberg:IcebergPerformanceTest")
+}

it/clickhouse/build.gradle

@@ -29,6 +29,12 @@ ext.summary = "Integration test utilities for ClickHouse."
 
 def clickhouse_java_client_version = "0.9.6"
 
+configurations.configureEach {
+  resolutionStrategy.capabilitiesResolution.withCapability("org.lz4:lz4-java") {
+    select("at.yawk.lz4:lz4-java:1.10.2")
+  }
+}
+
 dependencies {
     implementation project(path: ":it:common")
     implementation project(path: ":it:testcontainers")
@@ -52,4 +58,4 @@ dependencies {
 
     permitUnusedDeclared "com.clickhouse:client-v2:$clickhouse_java_client_version:all"
 
-}
+}

it/common/build.gradle

@@ -28,6 +28,19 @@ ext.summary = "Code used by all integration test utilities."
 dependencies {
     implementation enforcedPlatform(library.java.google_cloud_platform_libraries_bom)
     implementation project(path: ":sdks:java:core", configuration: "shadow")
+    implementation project(path: ":runners:google-cloud-dataflow-java")
+    implementation 'com.google.cloud:google-cloud-bigquery'
+    implementation 'com.google.cloud:google-cloud-monitoring'
+    implementation 'com.google.cloud:google-cloud-storage'
+    implementation 'org.apache.hadoop:hadoop-common:3.3.6'
+    implementation 'org.apache.avro:avro:1.11.1'
+    implementation 'org.apache.parquet:parquet-avro:1.15.2'
+    implementation 'org.apache.parquet:parquet-common:1.15.2'
+    implementation 'org.apache.parquet:parquet-hadoop:1.15.2'
+    provided 'com.google.api.grpc:proto-google-cloud-monitoring-v3'
+    implementation library.java.jackson_core
+    implementation library.java.jackson_databind
+    implementation library.java.gax
     implementation library.java.google_api_services_dataflow
     implementation library.java.google_auth_library_credentials
     implementation library.java.google_auth_library_oauth2_http
@@ -36,14 +49,20 @@ dependencies {
     implementation library.java.commons_lang3
     implementation library.java.failsafe
     implementation library.java.google_code_gson
+    implementation library.java.google_cloud_core
+    implementation library.java.google_api_client
     implementation library.java.google_http_client
     implementation library.java.guava
-
+    implementation library.java.protobuf_java_util
+    implementation library.java.protobuf_java
+    implementation library.java.junit
+    testImplementation library.java.mockito_inline
+    testRuntimeOnly library.java.slf4j_simple
     // TODO: excluding Guava until Truth updates it to >32.1.x
     testImplementation(library.java.truth) {
-        exclude group: 'com.google.guava', module: 'guava'
+      exclude group: 'com.google.guava', module: 'guava'
     }
-    testImplementation library.java.junit
-    testImplementation library.java.mockito_inline
-    testRuntimeOnly library.java.slf4j_simple
+    testImplementation project(path: ":runners:direct-java", configuration: "shadow")
+    testImplementation project(path: ":sdks:java:testing:test-utils")
+    testImplementation project(path: ":sdks:java:extensions:protobuf", configuration: "testRuntimeMigration")
 }

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/Artifact.java → it/common/src/main/java/org/apache/beam/it/common/artifacts/Artifact.java

@@ -15,7 +15,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.beam.it.gcp.artifacts;
+package org.apache.beam.it.common.artifacts;
 
 /**
  * Represents a single artifact.

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/ArtifactClient.java → it/common/src/main/java/org/apache/beam/it/common/artifacts/ArtifactClient.java

@@ -15,7 +15,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.beam.it.gcp.artifacts;
+package org.apache.beam.it.common.artifacts;
 
 import java.io.IOException;
 import java.nio.file.Path;

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/GcsArtifact.java → it/common/src/main/java/org/apache/beam/it/common/artifacts/GcsArtifact.java

@@ -15,7 +15,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.beam.it.gcp.artifacts;
+package org.apache.beam.it.common.artifacts;
 
 import com.google.cloud.storage.Blob;
 

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/package-info.java → it/common/src/main/java/org/apache/beam/it/common/artifacts/package-info.java

@@ -17,4 +17,4 @@
  */
 
 /** Package for working with test artifacts. */
-package org.apache.beam.it.gcp.artifacts;
+package org.apache.beam.it.common.artifacts;

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/ArtifactUtils.java → it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/ArtifactUtils.java

@@ -15,7 +15,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.beam.it.gcp.artifacts.utils;
+package org.apache.beam.it.common.artifacts.utils;
 
 import static java.util.Arrays.stream;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/AvroTestUtil.java → it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/AvroTestUtil.java

@@ -15,7 +15,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.beam.it.gcp.artifacts.utils;
+package org.apache.beam.it.common.artifacts.utils;
 
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/JsonTestUtil.java → it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/JsonTestUtil.java

@@ -15,7 +15,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.beam.it.gcp.artifacts.utils;
+package org.apache.beam.it.common.artifacts.utils;
 
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.JsonNode;

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/ParquetTestUtil.java → it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/ParquetTestUtil.java

@@ -15,7 +15,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.beam.it.gcp.artifacts.utils;
+package org.apache.beam.it.common.artifacts.utils;
 
 import java.io.ByteArrayInputStream;
 import java.io.File;

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/package-info.java → it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/package-info.java

@@ -17,4 +17,4 @@
  */
 
 /** Package for artifact utilities. */
-package org.apache.beam.it.gcp.artifacts.utils;
+package org.apache.beam.it.common.artifacts.utils;

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManager.java → it/common/src/main/java/org/apache/beam/it/common/bigquery/BigQueryResourceManager.java

@@ -15,7 +15,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.beam.it.gcp.bigquery;
+package org.apache.beam.it.common.bigquery;
 
 import com.google.api.gax.paging.Page;
 import com.google.auth.Credentials;

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerException.java → it/common/src/main/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerException.java

@@ -15,7 +15,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.beam.it.gcp.bigquery;
+package org.apache.beam.it.common.bigquery;
 
 /** Custom exception for {@link BigQueryResourceManager} implementations. */
 public class BigQueryResourceManagerException extends RuntimeException {

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerUtils.java → it/common/src/main/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerUtils.java

@@ -15,7 +15,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.beam.it.gcp.bigquery;
+package org.apache.beam.it.common.bigquery;
 
 import static org.apache.beam.it.common.utils.ResourceManagerUtils.generateResourceId;
 

it/common/src/main/java/org/apache/beam/it/common/bigquery/package-info.java

@@ -0,0 +1,20 @@
+/*
+ * 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 for managing BigQuery common resources within integration tests. */
+package org.apache.beam.it.common.bigquery;

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/AbstractPipelineLauncher.java → it/common/src/main/java/org/apache/beam/it/common/dataflow/AbstractPipelineLauncher.java

@@ -15,7 +15,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.beam.it.gcp.dataflow;
+package org.apache.beam.it.common.dataflow;
 
 import static org.apache.beam.it.common.PipelineLauncher.JobState.ACTIVE_STATES;
 import static org.apache.beam.it.common.PipelineLauncher.JobState.FAILED;

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/DefaultPipelineLauncher.java → it/common/src/main/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncher.java

@@ -15,10 +15,9 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.beam.it.gcp.dataflow;
+package org.apache.beam.it.common.dataflow;
 
 import static org.apache.beam.it.common.logging.LogStrings.formatForLogging;
-import static org.apache.beam.sdk.testing.TestPipeline.PROPERTY_BEAM_TEST_PIPELINE_OPTIONS;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -41,7 +40,6 @@
 import java.util.stream.StreamSupport;
 import org.apache.beam.it.common.PipelineLauncher;
 import org.apache.beam.it.common.utils.PipelineUtils;
-import org.apache.beam.it.gcp.IOLoadTestBase;
 import org.apache.beam.runners.dataflow.DataflowPipelineJob;
 import org.apache.beam.sdk.PipelineResult;
 import org.apache.beam.sdk.metrics.DistributionResult;
@@ -73,6 +71,8 @@ public class DefaultPipelineLauncher extends AbstractPipelineLauncher {
   private static final String READ_PIPELINE_NAME_OVERWRITE = "readPipelineNameOverride";
   private static final String WRITE_PIPELINE_NAME_OVERWRITE = "writePipelineNameOverride";
   private static final Pattern JOB_ID_PATTERN = Pattern.compile("Submitted job: (\\S+)");
+  /** Namespace for Beam provided pipeline metrics (set up by Metrics transform). */
+  public static final String BEAM_METRICS_NAMESPACE = "BEAM_METRICS";
 
   // For unsupported runners (other than dataflow), implement launcher methods by operating with
   // PipelineResult.
@@ -101,6 +101,18 @@ public class DefaultPipelineLauncher extends AbstractPipelineLauncher {
           .put(PipelineResult.State.UNRECOGNIZED, JobState.UNKNOWN)
           .build();
 
+  // To make PipelineLauncher.getMetric work in a unified way for both runner provided metrics and
+  // pipeline defined
+  // metrics, here we wrap Beam provided metrics as a pre-defined metrics name
+  // [name_space:metric_type:metric_name
+  // which will be recognized by getMetric method
+  public enum PipelineMetricsType {
+    COUNTER,
+    STARTTIME,
+    ENDTIME,
+    RUNTIME,
+  }
+
   private DefaultPipelineLauncher(Builder builder) {
     super(
         new Dataflow(
@@ -169,7 +181,7 @@ private static <T> void checkIfMetricResultIsUnique(
         resultCount <= 1,
         "More than one metric result matches name: %s in namespace %s. Metric results count: %s",
         name,
-        IOLoadTestBase.BEAM_METRICS_NAMESPACE,
+        BEAM_METRICS_NAMESPACE,
         resultCount);
   }
 
@@ -180,15 +192,13 @@ private static Iterable<MetricResult<DistributionResult>> getDistributions(
             .metrics()
             .queryMetrics(
                 MetricsFilter.builder()
-                    .addNameFilter(
-                        MetricNameFilter.named(IOLoadTestBase.BEAM_METRICS_NAMESPACE, metricName))
+                    .addNameFilter(MetricNameFilter.named(BEAM_METRICS_NAMESPACE, metricName))
                     .build());
     return metrics.getDistributions();
   }
 
   /** Pull Beam pipeline defined metrics given the jobId. */
-  public Long getBeamMetric(
-      String jobId, IOLoadTestBase.PipelineMetricsType metricType, String metricName) {
+  public Long getBeamMetric(String jobId, PipelineMetricsType metricType, String metricName) {
     PipelineResult pipelineResult =
         MANAGED_JOBS.getOrDefault(jobId, UNMANAGED_JOBS.getOrDefault(jobId, null));
     if (pipelineResult != null) {
@@ -197,8 +207,7 @@ public Long getBeamMetric(
               .metrics()
               .queryMetrics(
                   MetricsFilter.builder()
-                      .addNameFilter(
-                          MetricNameFilter.named(IOLoadTestBase.BEAM_METRICS_NAMESPACE, metricName))
+                      .addNameFilter(MetricNameFilter.named(BEAM_METRICS_NAMESPACE, metricName))
                       .build());
 
       switch (metricType) {
@@ -210,9 +219,7 @@ public Long getBeamMetric(
             return metricResult.getAttempted();
           } catch (NoSuchElementException e) {
             LOG.error(
-                "Failed to get metric {}, from namespace {}",
-                metricName,
-                IOLoadTestBase.BEAM_METRICS_NAMESPACE);
+                "Failed to get metric {}, from namespace {}", metricName, BEAM_METRICS_NAMESPACE);
           }
           return UNKNOWN_METRIC_VALUE;
         case STARTTIME:
@@ -230,9 +237,9 @@ public Long getBeamMetric(
                   .map(element -> Objects.requireNonNull(element.getAttempted()).getMax())
                   .max(Long::compareTo)
                   .orElse(UNKNOWN_METRIC_VALUE);
-          if (metricType == IOLoadTestBase.PipelineMetricsType.STARTTIME) {
+          if (metricType == PipelineMetricsType.STARTTIME) {
             return lowestMin;
-          } else if (metricType == IOLoadTestBase.PipelineMetricsType.ENDTIME) {
+          } else if (metricType == PipelineMetricsType.ENDTIME) {
             return greatestMax;
           } else {
             if (lowestMin != UNKNOWN_METRIC_VALUE && greatestMax != UNKNOWN_METRIC_VALUE) {
@@ -254,15 +261,14 @@ public Long getBeamMetric(
   @Override
   public Double getMetric(String project, String region, String jobId, String metricName)
       throws IOException {
-    if (metricName.startsWith(IOLoadTestBase.BEAM_METRICS_NAMESPACE)) {
+    if (metricName.startsWith(BEAM_METRICS_NAMESPACE)) {
       String[] nameSpacedMetrics = metricName.split(":", 3);
       Preconditions.checkState(
           nameSpacedMetrics.length == 3,
           String.format(
               "Invalid Beam metrics name: %s, expected: '%s:metric_type:metric_name'",
-              metricName, IOLoadTestBase.BEAM_METRICS_NAMESPACE));
-      IOLoadTestBase.PipelineMetricsType metricType =
-          IOLoadTestBase.PipelineMetricsType.valueOf(nameSpacedMetrics[1]);
+              metricName, BEAM_METRICS_NAMESPACE));
+      PipelineMetricsType metricType = PipelineMetricsType.valueOf(nameSpacedMetrics[1]);
 
       // Pipeline defined metrics are long values. Have to cast to double that is what the base
       // class defined.
@@ -437,16 +443,18 @@ private List<String> extractOptions(String project, String region, LaunchConfig
     // add pipeline options from beamTestPipelineOptions system property to preserve the
     // pipeline options already set in TestPipeline.
     @Nullable
-    String beamTestPipelineOptions = System.getProperty(PROPERTY_BEAM_TEST_PIPELINE_OPTIONS);
+    String beamTestPipelineOptions =
+        System.getProperty(
+            org.apache.beam.sdk.testing.TestPipeline.PROPERTY_BEAM_TEST_PIPELINE_OPTIONS);
     if (!Strings.isNullOrEmpty(beamTestPipelineOptions)) {
       try {
         additionalOptions.addAll(MAPPER.readValue(beamTestPipelineOptions, List.class));
       } catch (IOException e) {
         throw new RuntimeException(
             "Unable to instantiate test options from system property "
-                + PROPERTY_BEAM_TEST_PIPELINE_OPTIONS
+                + org.apache.beam.sdk.testing.TestPipeline.PROPERTY_BEAM_TEST_PIPELINE_OPTIONS
                 + ":"
-                + System.getProperty(PROPERTY_BEAM_TEST_PIPELINE_OPTIONS),
+                + beamTestPipelineOptions,
             e);
       }
     }

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/datagenerator/package-info.java → it/common/src/main/java/org/apache/beam/it/common/dataflow/package-info.java

@@ -16,5 +16,5 @@
  * limitations under the License.
  */
 
-/** Data generator for load tests. */
-package org.apache.beam.it.gcp.datagenerator;
+/** Package for managing Dataflow jobs from integration tests. */
+package org.apache.beam.it.common.dataflow;

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/monitoring/MonitoringClient.java → it/common/src/main/java/org/apache/beam/it/common/monitoring/MonitoringClient.java

@@ -15,7 +15,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.beam.it.gcp.monitoring;
+package org.apache.beam.it.common.monitoring;
 
 import com.google.api.gax.core.CredentialsProvider;
 import com.google.cloud.monitoring.v3.MetricServiceClient;

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/monitoring/package-info.java → it/common/src/main/java/org/apache/beam/it/common/monitoring/package-info.java

@@ -17,4 +17,4 @@
  */
 
 /** Package for querying metrics from cloud monitoring. */
-package org.apache.beam.it.gcp.monitoring;
+package org.apache.beam.it.common.monitoring;