Merge remote-tracking branch 'upstream/master' into bq-read-schema

Author: RustedBones

Diff for: .asf.yaml

@@ -36,11 +36,9 @@ github:
 
   # Give some users issue triage permissions
   collaborators:
-    - pcoet
-    - olehborysevych
-    - rshamunov
-    - andreydevyatkin
-    - liferoad
+    - Amar3tto
+    - mrshakirov
+    - akashorabek
 
   enabled_merge_buttons:
     squash: true

Diff for: .github/trigger_files/IO_Iceberg_Integration_Tests.json

@@ -1,4 +1,4 @@
 {
-    "comment": "Modify this file in a trivial way to cause this test suite to run",
+    "comment": "Modify this file in a trivial way to cause this test suite to run.",
     "modification": 1
 }

Diff for: .test-infra/tools/stale_k8s_workload_cleaner.sh

@@ -43,7 +43,10 @@ function should_teardown() {
 gcloud container clusters get-credentials io-datastores --zone us-central1-a --project apache-beam-testing
 
 while read NAME STATUS AGE; do
-  if [[ $NAME =~ ^beam-.+(test|-it) ]] && should_teardown $AGE; then
+  # Regex has temporary workaround to avoid trying to delete beam-performancetests-singlestoreio-* to avoid getting stuck in a terminal state
+  # See https://github.com/apache/beam/pull/33545 for context.
+  # This may be safe to remove if https://cloud.google.com/knowledge/kb/deleted-namespace-remains-in-terminating-status-000004867 has been resolved, just try it before checking in :)
+  if [[ $NAME =~ ^beam-.+(test|-it)(?!s-singlestoreio) ]] && should_teardown $AGE; then
     kubectl delete namespace $NAME
   fi
 done < <( kubectl get namespaces --context=gke_${PROJECT}_${LOCATION}_${CLUSTER} )

Diff for: CHANGES.md

@@ -80,6 +80,7 @@
 
 ## Bugfixes
 
+* Fix data loss issues when reading gzipped files with TextIO (Python) ([#18390](https://github.com/apache/beam/issues/18390), [#31040](https://github.com/apache/beam/issues/31040)).
 * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
 * [BigQueryIO] Fixed an issue where Storage Write API sometimes doesn't pick up auto-schema updates ([#33231](https://github.com/apache/beam/pull/33231))
 
@@ -115,6 +116,7 @@
 ## Bugfixes
 
 * Fixed EventTimeTimer ordering in Prism. ([#32222](https://github.com/apache/beam/issues/32222)).
+* [Managed Iceberg] Fixed a bug where DataFile metadata was assigned incorrect partition values ([#33549](https://github.com/apache/beam/pull/33549)).
 
 ## Security Fixes
 
@@ -157,6 +159,11 @@
 * Adding flag to support conditionally disabling auto-commit in JdbcIO ReadFn ([#31111](https://github.com/apache/beam/issues/31111))
 * (Python) Fixed BigQuery Enrichment bug that can lead to multiple conditions returning duplicate rows, batching returning incorrect results and conditions not scoped by row during batching ([#32780](https://github.com/apache/beam/pull/32780)).
 
+## Known Issues
+
+* [Managed Iceberg] DataFile metadata is assigned incorrect partition values ([#33497](https://github.com/apache/beam/issues/33497)).
+  * Fixed in 2.62.0
+
 # [2.60.0] - 2024-10-17
 
 ## Highlights
@@ -211,6 +218,8 @@ when running on 3.8. ([#31192](https://github.com/apache/beam/issues/31192))
   * Duplicate Rows: Multiple conditions may be applied incorrectly, leading to the duplication of rows in the output.
   * Incorrect Results with Batched Requests: Conditions may not be correctly scoped to individual rows within the batch, potentially causing inaccurate results.
   * Fixed in 2.61.0.
+* [Managed Iceberg] DataFile metadata is assigned incorrect partition values ([#33497](https://github.com/apache/beam/issues/33497)).
+  * Fixed in 2.62.0
 
 # [2.59.0] - 2024-09-11
 
@@ -259,6 +268,8 @@ when running on 3.8. ([#31192](https://github.com/apache/beam/issues/31192))
   * Duplicate Rows: Multiple conditions may be applied incorrectly, leading to the duplication of rows in the output.
   * Incorrect Results with Batched Requests: Conditions may not be correctly scoped to individual rows within the batch, potentially causing inaccurate results.
   * Fixed in 2.61.0.
+* [Managed Iceberg] DataFile metadata is assigned incorrect partition values ([#33497](https://github.com/apache/beam/issues/33497)).
+  * Fixed in 2.62.0
 
 # [2.58.1] - 2024-08-15
 

Diff for: buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy

@@ -618,6 +618,7 @@ class BeamModulePlugin implements Plugin<Project> {
     def influxdb_version = "2.19"
     def httpclient_version = "4.5.13"
     def httpcore_version = "4.4.14"
+    def iceberg_bqms_catalog_version = "1.5.2-0.1.0"
     def jackson_version = "2.15.4"
     def jaxb_api_version = "2.3.3"
     def jsr305_version = "3.0.2"
@@ -650,6 +651,10 @@ class BeamModulePlugin implements Plugin<Project> {
 
     // Export Spark versions, so they are defined in a single place only
     project.ext.spark3_version = spark3_version
+    // version for BigQueryMetastore catalog (used by sdks:java:io:iceberg:bqms)
+    // TODO: remove this and download the jar normally when the catalog gets
+    // open-sourced (https://github.com/apache/iceberg/pull/11039)
+    project.ext.iceberg_bqms_catalog_version = iceberg_bqms_catalog_version
 
     // A map of maps containing common libraries used per language. To use:
     // dependencies {

Diff for: sdks/java/container/boot.go

@@ -134,21 +134,31 @@ func main() {
 	}
 
 	const jarsDir = "/opt/apache/beam/jars"
+	const javaHarnessJar = "beam-sdks-java-harness.jar"
 	cp := []string{
 		filepath.Join(jarsDir, "slf4j-api.jar"),
 		filepath.Join(jarsDir, "slf4j-jdk14.jar"),
 		filepath.Join(jarsDir, "jcl-over-slf4j.jar"),
 		filepath.Join(jarsDir, "log4j-over-slf4j.jar"),
 		filepath.Join(jarsDir, "log4j-to-slf4j.jar"),
-		filepath.Join(jarsDir, "beam-sdks-java-harness.jar"),
+		filepath.Join(jarsDir, javaHarnessJar),
 	}
 
 	var hasWorkerExperiment = strings.Contains(options, "use_staged_dataflow_worker_jar")
 	for _, a := range artifacts {
 		name, _ := artifact.MustExtractFilePayload(a)
 		if hasWorkerExperiment {
-			if strings.HasPrefix(name, "beam-runners-google-cloud-dataflow-java-fn-api-worker") {
-				continue
+			if strings.HasPrefix(name, "beam-sdks-java-harness") {
+				// Remove system "beam-sdks-java-harness.jar". User-provided jar will be
+				// added to classpath as a normal user jar further below.
+				for i, cl := range cp {
+					if !strings.HasSuffix(cl, javaHarnessJar) {
+						continue
+					}
+					logger.Printf(ctx, "Using staged java harness: %v", name)
+					cp = append(cp[:i], cp[i+1:]...)
+					break
+				}
 			}
 			if name == "dataflow-worker.jar" {
 				continue

Diff for: sdks/java/io/expansion-service/build.gradle

@@ -59,11 +59,13 @@ dependencies {
   // **** IcebergIO runtime dependencies ****
   runtimeOnly library.java.hadoop_auth
   runtimeOnly library.java.hadoop_client
-  // Needed when using GCS as the warehouse location.
+  // For writing to GCS
   runtimeOnly library.java.bigdataoss_gcs_connector
-  // Needed for HiveCatalog
+  // HiveCatalog
   runtimeOnly ("org.apache.iceberg:iceberg-hive-metastore:1.4.2")
   runtimeOnly project(path: ":sdks:java:io:iceberg:hive")
+  // BigQueryMetastoreCatalog (Java 11+)
+  runtimeOnly project(path: ":sdks:java:io:iceberg:bqms", configuration: "shadow")
 
   runtimeOnly library.java.kafka_clients
   runtimeOnly library.java.slf4j_jdk14

Diff for: sdks/java/io/iceberg/bqms/build.gradle

@@ -0,0 +1,63 @@
+/*
+ * 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.
+ */
+plugins {
+    id 'org.apache.beam.module'
+}
+
+applyJavaNature(
+        automaticModuleName: 'org.apache.beam.sdk.io.iceberg.bqms',
+        shadowClosure: {},
+        exportJavadoc: false,
+        publish: false, // it's an intermediate jar for io-expansion-service
+        validateShadowJar: false
+)
+
+def libDir = "$buildDir/libs"
+def bqmsFileName = "iceberg-bqms-catalog-${iceberg_bqms_catalog_version}.jar"
+task downloadBqmsJar(type: Copy) {
+    // TODO: remove this workaround and downlooad normally when the catalog gets open-sourced:
+    // (https://github.com/apache/iceberg/pull/11039)
+    def jarUrl = "https://storage.googleapis.com/spark-lib/bigquery/iceberg-bigquery-catalog-${iceberg_bqms_catalog_version}.jar"
+    def outputDir = file("$libDir")
+    outputDir.mkdirs()
+    def destFile = new File(outputDir, bqmsFileName)
+
+    if (!destFile.exists()) {
+        try {
+            ant.get(src: jarUrl, dest: destFile)
+            println "Successfully downloaded BQMS catalog jar: $destFile"
+        } catch (Exception e) {
+            println "Could not download $jarUrl: ${e.message}"
+        }
+    }
+}
+
+repositories {
+    flatDir {
+        dirs "$libDir"
+    }
+}
+
+compileJava.dependsOn downloadBqmsJar
+
+dependencies {
+    implementation files("$libDir/$bqmsFileName")
+}
+
+description = "Apache Beam :: SDKs :: Java :: IO :: Iceberg :: BigQuery Metastore"
+ext.summary = "A copy of the BQMS catalog."

Diff for: sdks/java/io/iceberg/build.gradle

@@ -83,6 +83,9 @@ dependencies {
         exclude group: "org.apache.parquet", module: "parquet-hadoop-bundle"
     }
 
+    // BigQueryMetastore catalog dep
+    testImplementation project(path: ":sdks:java:io:iceberg:bqms", configuration: "shadow")
+
     testRuntimeOnly library.java.slf4j_jdk14
     testRuntimeOnly project(path: ":runners:direct-java", configuration: "shadow")
     testRuntimeOnly project(path: ":runners:google-cloud-dataflow-java")
@@ -136,6 +139,11 @@ task integrationTest(type: Test) {
     outputs.upToDateWhen { false }
 
     include '**/*IT.class'
+    // BQ metastore catalog doesn't support java 8
+    if (project.findProperty('testJavaVersion') == '8' ||
+            JavaVersion.current().equals(JavaVersion.VERSION_1_8)) {
+        exclude '**/BigQueryMetastoreCatalogIT.class'
+    }
 
     maxParallelForks 4
     classpath = sourceSets.test.runtimeClasspath

Diff for: sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/RecordWriterManager.java

@@ -96,7 +96,9 @@ class DestinationState {
     private final IcebergDestination icebergDestination;
     private final PartitionSpec spec;
     private final org.apache.iceberg.Schema schema;
-    private final PartitionKey partitionKey;
+    // used to determine the partition to which a record belongs
+    // must not be directly used to create a writer
+    private final PartitionKey routingPartitionKey;
     private final Table table;
     private final String stateToken = UUID.randomUUID().toString();
     final Cache<PartitionKey, RecordWriter> writers;
@@ -109,7 +111,7 @@ class DestinationState {
       this.icebergDestination = icebergDestination;
       this.schema = table.schema();
       this.spec = table.spec();
-      this.partitionKey = new PartitionKey(spec, schema);
+      this.routingPartitionKey = new PartitionKey(spec, schema);
       this.table = table;
       for (PartitionField partitionField : spec.fields()) {
         partitionFieldMap.put(partitionField.name(), partitionField);
@@ -154,12 +156,12 @@ class DestinationState {
      * can't create a new writer, the {@link Record} is rejected and {@code false} is returned.
      */
     boolean write(Record record) {
-      partitionKey.partition(getPartitionableRecord(record));
+      routingPartitionKey.partition(getPartitionableRecord(record));
 
-      if (!writers.asMap().containsKey(partitionKey) && openWriters >= maxNumWriters) {
+      if (!writers.asMap().containsKey(routingPartitionKey) && openWriters >= maxNumWriters) {
         return false;
       }
-      RecordWriter writer = fetchWriterForPartition(partitionKey);
+      RecordWriter writer = fetchWriterForPartition(routingPartitionKey);
       writer.write(record);
       return true;
     }
@@ -173,10 +175,12 @@ private RecordWriter fetchWriterForPartition(PartitionKey partitionKey) {
       RecordWriter recordWriter = writers.getIfPresent(partitionKey);
 
       if (recordWriter == null || recordWriter.bytesWritten() > maxFileSize) {
+        // each writer must have its own PartitionKey object
+        PartitionKey copy = partitionKey.copy();
         // calling invalidate for a non-existent key is a safe operation
-        writers.invalidate(partitionKey);
-        recordWriter = createWriter(partitionKey);
-        writers.put(partitionKey, recordWriter);
+        writers.invalidate(copy);
+        recordWriter = createWriter(copy);
+        writers.put(copy, recordWriter);
       }
       return recordWriter;
     }

Diff for: sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/BigQueryMetastoreCatalogIT.java

@@ -0,0 +1,76 @@
+/*
+ * 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.io.iceberg.catalog;
+
+import java.io.IOException;
+import java.util.Map;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.iceberg.CatalogUtil;
+import org.apache.iceberg.catalog.Catalog;
+import org.apache.iceberg.catalog.Namespace;
+import org.apache.iceberg.catalog.TableIdentifier;
+
+public class BigQueryMetastoreCatalogIT extends IcebergCatalogBaseIT {
+  static final String BQMS_CATALOG = "org.apache.iceberg.gcp.bigquery.BigQueryMetastoreCatalog";
+  static final String DATASET = "managed_iceberg_bqms_tests_no_delete";
+  static final long SALT = System.nanoTime();
+
+  @Override
+  public String tableId() {
+    return DATASET + "." + testName.getMethodName() + "_" + SALT;
+  }
+
+  @Override
+  public Catalog createCatalog() {
+    return CatalogUtil.loadCatalog(
+        BQMS_CATALOG,
+        "bqms_" + catalogName,
+        ImmutableMap.<String, String>builder()
+            .put("gcp_project", options.getProject())
+            .put("gcp_location", "us-central1")
+            .put("warehouse", warehouse)
+            .build(),
+        new Configuration());
+  }
+
+  @Override
+  public void catalogCleanup() throws IOException {
+    for (TableIdentifier tableIdentifier : catalog.listTables(Namespace.of(DATASET))) {
+      // only delete tables that were created in this test run
+      if (tableIdentifier.name().contains(String.valueOf(SALT))) {
+        catalog.dropTable(tableIdentifier);
+      }
+    }
+  }
+
+  @Override
+  public Map<String, Object> managedIcebergConfig(String tableId) {
+    return ImmutableMap.<String, Object>builder()
+        .put("table", tableId)
+        .put(
+            "catalog_properties",
+            ImmutableMap.<String, String>builder()
+                .put("gcp_project", options.getProject())
+                .put("gcp_location", "us-central1")
+                .put("warehouse", warehouse)
+                .put("catalog-impl", BQMS_CATALOG)
+                .build())
+        .build();
+  }
+}

Diff for: sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/HiveCatalogIT.java

@@ -65,7 +65,6 @@ public Catalog createCatalog() {
 
   @Override
   public void catalogCleanup() throws Exception {
-    System.out.println("xxx CLEANING UP!");
     if (hiveMetastoreExtension != null) {
       hiveMetastoreExtension.cleanup();
     }