From 7b97a227f2a324caf938cc1d5dcb2df2122dd606 Mon Sep 17 00:00:00 2001
From: Zihan Dai <99155080+PDGGK@users.noreply.github.com>
Date: Sat, 15 Aug 2026 19:21:03 +1000
Subject: [PATCH] Tag the TFRecord write error output with the schema it
 actually emits

ErrorFn is constructed with errorSchema = ErrorHandling.errorSchema(
inputSchema) and emits ErrorHandling.errorRecord(errorSchema, ..), so
every row on ERROR_TAG carries that schema. The collection was then
tagged with ErrorHandling.errorSchema(errorSchema) -- the wrapper
applied twice -- declaring {failed_row: {failed_row: .., error_message},
error_message} which no element it produces can match.

The sibling in the same package gets it right five lines after the same
ErrorFn construction: TFRecordReadSchemaTransformProvider:140 does
setRowSchema(errorSchema). JavaFilter, JavaMapToFields, PubsubRowToMessage
and BigQueryStorageWriteApi all do the same.

One test, and it needs no runner -- the schema is fixed when the graph is
built. Restoring the second wrap fails it and nothing else: 8 tests, 1
failure.

KafkaWriteSchemaTransformProvider:301 has the identical double wrap. It
is a different module, so it is a separate change rather than folded in
here.
---
 .../TFRecordWriteSchemaTransformProvider.java |  3 +-
 .../TFRecordSchemaTransformProviderTest.java  | 29 +++++++++++++++++++
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TFRecordWriteSchemaTransformProvider.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TFRecordWriteSchemaTransformProvider.java
index bc9b7bbeac66..c787997b1e77 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TFRecordWriteSchemaTransformProvider.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TFRecordWriteSchemaTransformProvider.java
@@ -184,8 +184,7 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) {
           output = "";
         }
       }
-      PCollection<Row> errorOutput =
-          byteArrays.get(ERROR_TAG).setRowSchema(ErrorHandling.errorSchema(errorSchema));
+      PCollection<Row> errorOutput = byteArrays.get(ERROR_TAG).setRowSchema(errorSchema);
       return PCollectionRowTuple.of(handleErrors ? output : "errors", errorOutput);
     }
   }
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TFRecordSchemaTransformProviderTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TFRecordSchemaTransformProviderTest.java
index 9c067a533e0b..65e38ed9d298 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TFRecordSchemaTransformProviderTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TFRecordSchemaTransformProviderTest.java
@@ -49,6 +49,7 @@
 import org.apache.beam.sdk.schemas.Schema;
 import org.apache.beam.sdk.schemas.transforms.SchemaTransform;
 import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider;
+import org.apache.beam.sdk.schemas.transforms.providers.ErrorHandling;
 import org.apache.beam.sdk.testing.NeedsRunner;
 import org.apache.beam.sdk.testing.PAssert;
 import org.apache.beam.sdk.testing.TestPipeline;
@@ -238,6 +239,34 @@ public void testWriteBuildTransform() {
             .build());
   }
 
+  @Test
+  public void testWriteErrorSchemaMatchesTheRowsErrorFnEmits() throws Exception {
+    // The schema is fixed when the graph is built, so this needs no runner.
+    writePipeline.enableAbandonedNodeEnforcement(false);
+
+    Schema schema = Schema.of(Schema.Field.of("record", Schema.FieldType.BYTES));
+
+    TFRecordWriteSchemaTransformProvider provider = new TFRecordWriteSchemaTransformProvider();
+    TFRecordWriteSchemaTransform transform =
+        (TFRecordWriteSchemaTransform)
+            provider.from(
+                TFRecordWriteSchemaTransformConfiguration.builder()
+                    .setOutputPrefix(tempFolder.getRoot().toPath().resolve("errors").toString())
+                    .setCompression("UNCOMPRESSED")
+                    .setNumShards(0)
+                    .setNoSpilling(true)
+                    .build());
+
+    Row row = Row.withSchema(schema).addValue("foo".getBytes(StandardCharsets.UTF_8)).build();
+    PCollection<Row> input =
+        writePipeline.apply(Create.of(Collections.singletonList(row)).withRowSchema(schema));
+    PCollectionRowTuple result = PCollectionRowTuple.of("input", input).apply(transform);
+
+    // ErrorFn emits ErrorHandling.errorRecord(errorSchema, ..), so the collection has to carry
+    // that schema. Wrapping it a second time declares a shape no element it emits can match.
+    assertEquals(ErrorHandling.errorSchema(schema), result.get("errors").getSchema());
+  }
+
   @Test
   public void testReadFindTransformAndMakeItWork() {
     ServiceLoader<SchemaTransformProvider> serviceLoader =