You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by da...@apache.org on 2023/11/03 15:14:28 UTC
(beam) branch release-2.52.0 updated: [YAML] Fix error handling for KafkaSchemaTransforms (#29261) (#29289)
This is an automated email from the ASF dual-hosted git repository.
damccorm pushed a commit to branch release-2.52.0
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/release-2.52.0 by this push:
new 0e8e54c29c4 [YAML] Fix error handling for KafkaSchemaTransforms (#29261) (#29289)
0e8e54c29c4 is described below
commit 0e8e54c29c4b4eef0220d88c2b198dbb5e748b72
Author: Jeff Kinard <35...@users.noreply.github.com>
AuthorDate: Fri Nov 3 11:14:20 2023 -0400
[YAML] Fix error handling for KafkaSchemaTransforms (#29261) (#29289)
---
.../transforms/providers/ErrorHandling.java | 16 +++
.../KafkaReadSchemaTransformConfiguration.java | 7 ++
.../io/kafka/KafkaReadSchemaTransformProvider.java | 127 ++++++++++-----------
.../kafka/KafkaWriteSchemaTransformProvider.java | 54 ++++++---
.../org/apache/beam/sdk/io/kafka/KafkaDlqTest.java | 12 +-
.../KafkaReadSchemaTransformProviderTest.java | 3 +-
.../KafkaWriteSchemaTransformProviderTest.java | 15 ++-
sdks/python/apache_beam/yaml/standard_io.yaml | 2 +
sdks/python/apache_beam/yaml/yaml_transform.py | 2 +
9 files changed, 148 insertions(+), 90 deletions(-)
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/transforms/providers/ErrorHandling.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/transforms/providers/ErrorHandling.java
index 3fab157a0e5..7fa29708c9f 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/transforms/providers/ErrorHandling.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/transforms/providers/ErrorHandling.java
@@ -53,6 +53,12 @@ public abstract class ErrorHandling {
Schema.Field.of("error_message", Schema.FieldType.STRING));
}
+ public static Schema errorSchemaBytes() {
+ return Schema.of(
+ Schema.Field.of("failed_row", Schema.FieldType.BYTES),
+ Schema.Field.of("error_message", Schema.FieldType.STRING));
+ }
+
@SuppressWarnings({
"nullness" // TODO(https://github.com/apache/beam/issues/20497)
})
@@ -62,4 +68,14 @@ public abstract class ErrorHandling {
.withFieldValue("error_message", th.getMessage())
.build();
}
+
+ @SuppressWarnings({
+ "nullness" // TODO(https://github.com/apache/beam/issues/20497)
+ })
+ public static Row errorRecord(Schema errorSchema, byte[] inputBytes, Throwable th) {
+ return Row.withSchema(errorSchema)
+ .withFieldValue("failed_row", inputBytes)
+ .withFieldValue("error_message", th.getMessage())
+ .build();
+ }
}
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformConfiguration.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformConfiguration.java
index 299a4f222e3..ae03c49b9b0 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformConfiguration.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformConfiguration.java
@@ -24,6 +24,7 @@ import javax.annotation.Nullable;
import org.apache.beam.sdk.schemas.AutoValueSchema;
import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
import org.apache.beam.sdk.schemas.annotations.SchemaFieldDescription;
+import org.apache.beam.sdk.schemas.transforms.providers.ErrorHandling;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets;
/**
@@ -105,6 +106,10 @@ public abstract class KafkaReadSchemaTransformConfiguration {
/** Sets the topic from which to read. */
public abstract String getTopic();
+ @SchemaFieldDescription("This option specifies whether and where to output unwritable rows.")
+ @Nullable
+ public abstract ErrorHandling getErrorHandling();
+
/** Builder for the {@link KafkaReadSchemaTransformConfiguration}. */
@AutoValue.Builder
public abstract static class Builder {
@@ -127,6 +132,8 @@ public abstract class KafkaReadSchemaTransformConfiguration {
/** Sets the topic from which to read. */
public abstract Builder setTopic(String value);
+ public abstract Builder setErrorHandling(ErrorHandling errorHandling);
+
/** Builds a {@link KafkaReadSchemaTransformConfiguration} instance. */
public abstract KafkaReadSchemaTransformConfiguration build();
}
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProvider.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProvider.java
index fcba5c355df..2f5278cd712 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProvider.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProvider.java
@@ -43,10 +43,9 @@ import org.apache.beam.sdk.schemas.transforms.Convert;
import org.apache.beam.sdk.schemas.transforms.SchemaTransform;
import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider;
import org.apache.beam.sdk.schemas.transforms.TypedSchemaTransformProvider;
+import org.apache.beam.sdk.schemas.transforms.providers.ErrorHandling;
import org.apache.beam.sdk.schemas.utils.JsonUtils;
import org.apache.beam.sdk.transforms.DoFn;
-import org.apache.beam.sdk.transforms.DoFn.FinishBundle;
-import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.SerializableFunction;
import org.apache.beam.sdk.transforms.SimpleFunction;
@@ -77,8 +76,6 @@ public class KafkaReadSchemaTransformProvider
public static final TupleTag<Row> OUTPUT_TAG = new TupleTag<Row>() {};
public static final TupleTag<Row> ERROR_TAG = new TupleTag<Row>() {};
- public static final Schema ERROR_SCHEMA =
- Schema.builder().addStringField("error").addNullableByteArrayField("row").build();
final Boolean isTest;
final Integer testTimeoutSecs;
@@ -98,6 +95,9 @@ public class KafkaReadSchemaTransformProvider
return KafkaReadSchemaTransformConfiguration.class;
}
+ @SuppressWarnings({
+ "nullness" // TODO(https://github.com/apache/beam/issues/20497)
+ })
@Override
protected SchemaTransform from(KafkaReadSchemaTransformConfiguration configuration) {
final String inputSchema = configuration.getSchema();
@@ -114,14 +114,32 @@ public class KafkaReadSchemaTransformProvider
consumerConfigs.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, autoOffsetReset);
String format = configuration.getFormat();
-
- if (format != null && format.equals("RAW")) {
- if (inputSchema != null) {
- throw new IllegalArgumentException(
- "To read from Kafka in RAW format, you can't provide a schema.");
+ boolean handleErrors = ErrorHandling.hasOutput(configuration.getErrorHandling());
+
+ if ((format != null && format.equals("RAW")) || (!Strings.isNullOrEmpty(inputSchema))) {
+ SerializableFunction<byte[], Row> valueMapper;
+ Schema beamSchema;
+ if (format != null && format.equals("RAW")) {
+ if (inputSchema != null) {
+ throw new IllegalArgumentException(
+ "To read from Kafka in RAW format, you can't provide a schema.");
+ }
+ beamSchema = Schema.builder().addField("payload", Schema.FieldType.BYTES).build();
+ valueMapper = getRawBytesToRowFunction(beamSchema);
+ } else {
+ assert Strings.isNullOrEmpty(configuration.getConfluentSchemaRegistryUrl())
+ : "To read from Kafka, a schema must be provided directly or though Confluent "
+ + "Schema Registry, but not both.";
+
+ beamSchema =
+ Objects.equals(format, "JSON")
+ ? JsonUtils.beamSchemaFromJsonSchema(inputSchema)
+ : AvroUtils.toBeamSchema(new org.apache.avro.Schema.Parser().parse(inputSchema));
+ valueMapper =
+ Objects.equals(format, "JSON")
+ ? JsonUtils.getJsonBytesToRowFunction(beamSchema)
+ : AvroUtils.getAvroBytesToRowFunction(beamSchema);
}
- Schema rawSchema = Schema.builder().addField("payload", Schema.FieldType.BYTES).build();
- SerializableFunction<byte[], Row> valueMapper = getRawBytesToRowFunction(rawSchema);
return new SchemaTransform() {
@Override
public PCollectionRowTuple expand(PCollectionRowTuple input) {
@@ -138,59 +156,23 @@ public class KafkaReadSchemaTransformProvider
PCollection<byte[]> kafkaValues =
input.getPipeline().apply(kafkaRead.withoutMetadata()).apply(Values.create());
+ Schema errorSchema = ErrorHandling.errorSchemaBytes();
PCollectionTuple outputTuple =
kafkaValues.apply(
- ParDo.of(new ErrorFn("Kafka-read-error-counter", valueMapper))
+ ParDo.of(
+ new ErrorFn(
+ "Kafka-read-error-counter", valueMapper, errorSchema, handleErrors))
.withOutputTags(OUTPUT_TAG, TupleTagList.of(ERROR_TAG)));
- return PCollectionRowTuple.of(
- "output",
- outputTuple.get(OUTPUT_TAG).setRowSchema(rawSchema),
- "errors",
- outputTuple.get(ERROR_TAG).setRowSchema(ERROR_SCHEMA));
- }
- };
- }
+ PCollectionRowTuple outputRows =
+ PCollectionRowTuple.of(
+ "output", outputTuple.get(OUTPUT_TAG).setRowSchema(beamSchema));
- if (inputSchema != null && !inputSchema.isEmpty()) {
- assert Strings.isNullOrEmpty(configuration.getConfluentSchemaRegistryUrl())
- : "To read from Kafka, a schema must be provided directly or though Confluent "
- + "Schema Registry, but not both.";
-
- final Schema beamSchema =
- Objects.equals(format, "JSON")
- ? JsonUtils.beamSchemaFromJsonSchema(inputSchema)
- : AvroUtils.toBeamSchema(new org.apache.avro.Schema.Parser().parse(inputSchema));
- SerializableFunction<byte[], Row> valueMapper =
- Objects.equals(format, "JSON")
- ? JsonUtils.getJsonBytesToRowFunction(beamSchema)
- : AvroUtils.getAvroBytesToRowFunction(beamSchema);
- return new SchemaTransform() {
- @Override
- public PCollectionRowTuple expand(PCollectionRowTuple input) {
- KafkaIO.Read<byte[], byte[]> kafkaRead =
- KafkaIO.readBytes()
- .withConsumerConfigUpdates(consumerConfigs)
- .withConsumerFactoryFn(new ConsumerFactoryWithGcsTrustStores())
- .withTopic(configuration.getTopic())
- .withBootstrapServers(configuration.getBootstrapServers());
- if (isTest) {
- kafkaRead = kafkaRead.withMaxReadTime(Duration.standardSeconds(testTimeoutSecs));
+ PCollection<Row> errorOutput = outputTuple.get(ERROR_TAG).setRowSchema(errorSchema);
+ if (handleErrors) {
+ outputRows = outputRows.and(configuration.getErrorHandling().getOutput(), errorOutput);
}
-
- PCollection<byte[]> kafkaValues =
- input.getPipeline().apply(kafkaRead.withoutMetadata()).apply(Values.create());
-
- PCollectionTuple outputTuple =
- kafkaValues.apply(
- ParDo.of(new ErrorFn("Kafka-read-error-counter", valueMapper))
- .withOutputTags(OUTPUT_TAG, TupleTagList.of(ERROR_TAG)));
-
- return PCollectionRowTuple.of(
- "output",
- outputTuple.get(OUTPUT_TAG).setRowSchema(beamSchema),
- "errors",
- outputTuple.get(ERROR_TAG).setRowSchema(ERROR_SCHEMA));
+ return outputRows;
}
};
} else {
@@ -259,25 +241,38 @@ public class KafkaReadSchemaTransformProvider
}
public static class ErrorFn extends DoFn<byte[], Row> {
- private SerializableFunction<byte[], Row> valueMapper;
- private Counter errorCounter;
+ private final SerializableFunction<byte[], Row> valueMapper;
+ private final Counter errorCounter;
private Long errorsInBundle = 0L;
-
- public ErrorFn(String name, SerializableFunction<byte[], Row> valueMapper) {
+ private final boolean handleErrors;
+ private final Schema errorSchema;
+
+ public ErrorFn(
+ String name,
+ SerializableFunction<byte[], Row> valueMapper,
+ Schema errorSchema,
+ boolean handleErrors) {
this.errorCounter = Metrics.counter(KafkaReadSchemaTransformProvider.class, name);
this.valueMapper = valueMapper;
+ this.handleErrors = handleErrors;
+ this.errorSchema = errorSchema;
}
@ProcessElement
public void process(@DoFn.Element byte[] msg, MultiOutputReceiver receiver) {
+ Row mappedRow = null;
try {
- receiver.get(OUTPUT_TAG).output(valueMapper.apply(msg));
+ mappedRow = valueMapper.apply(msg);
} catch (Exception e) {
+ if (!handleErrors) {
+ throw new RuntimeException(e);
+ }
errorsInBundle += 1;
LOG.warn("Error while parsing the element", e);
- receiver
- .get(ERROR_TAG)
- .output(Row.withSchema(ERROR_SCHEMA).addValues(e.toString(), msg).build());
+ receiver.get(ERROR_TAG).output(ErrorHandling.errorRecord(errorSchema, msg, e));
+ }
+ if (mappedRow != null) {
+ receiver.get(OUTPUT_TAG).output(mappedRow);
}
}
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProvider.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProvider.java
index 1dbe37791c0..93d6d73493d 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProvider.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProvider.java
@@ -36,13 +36,14 @@ import org.apache.beam.sdk.schemas.annotations.SchemaFieldDescription;
import org.apache.beam.sdk.schemas.transforms.SchemaTransform;
import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider;
import org.apache.beam.sdk.schemas.transforms.TypedSchemaTransformProvider;
+import org.apache.beam.sdk.schemas.transforms.providers.ErrorHandling;
import org.apache.beam.sdk.schemas.utils.JsonUtils;
import org.apache.beam.sdk.transforms.DoFn;
-import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.SerializableFunction;
import org.apache.beam.sdk.transforms.SimpleFunction;
import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionRowTuple;
import org.apache.beam.sdk.values.PCollectionTuple;
import org.apache.beam.sdk.values.Row;
@@ -67,8 +68,6 @@ public class KafkaWriteSchemaTransformProvider
public static final TupleTag<Row> ERROR_TAG = new TupleTag<Row>() {};
public static final TupleTag<KV<byte[], byte[]>> OUTPUT_TAG =
new TupleTag<KV<byte[], byte[]>>() {};
- public static final Schema ERROR_SCHEMA =
- Schema.builder().addStringField("error").addNullableByteArrayField("row").build();
private static final Logger LOG =
LoggerFactory.getLogger(KafkaWriteSchemaTransformProvider.class);
@@ -100,25 +99,38 @@ public class KafkaWriteSchemaTransformProvider
}
public static class ErrorCounterFn extends DoFn<Row, KV<byte[], byte[]>> {
- private SerializableFunction<Row, byte[]> toBytesFn;
- private Counter errorCounter;
+ private final SerializableFunction<Row, byte[]> toBytesFn;
+ private final Counter errorCounter;
private Long errorsInBundle = 0L;
+ private final boolean handleErrors;
+ private final Schema errorSchema;
- public ErrorCounterFn(String name, SerializableFunction<Row, byte[]> toBytesFn) {
+ public ErrorCounterFn(
+ String name,
+ SerializableFunction<Row, byte[]> toBytesFn,
+ Schema errorSchema,
+ boolean handleErrors) {
this.toBytesFn = toBytesFn;
- errorCounter = Metrics.counter(KafkaWriteSchemaTransformProvider.class, name);
+ this.errorCounter = Metrics.counter(KafkaWriteSchemaTransformProvider.class, name);
+ this.handleErrors = handleErrors;
+ this.errorSchema = errorSchema;
}
@ProcessElement
public void process(@DoFn.Element Row row, MultiOutputReceiver receiver) {
+ KV<byte[], byte[]> output = null;
try {
- receiver.get(OUTPUT_TAG).output(KV.of(new byte[1], toBytesFn.apply(row)));
+ output = KV.of(new byte[1], toBytesFn.apply(row));
} catch (Exception e) {
+ if (!handleErrors) {
+ throw new RuntimeException(e);
+ }
errorsInBundle += 1;
LOG.warn("Error while processing the element", e);
- receiver
- .get(ERROR_TAG)
- .output(Row.withSchema(ERROR_SCHEMA).addValues(e.toString(), row.toString()).build());
+ receiver.get(ERROR_TAG).output(ErrorHandling.errorRecord(errorSchema, row, e));
+ }
+ if (output != null) {
+ receiver.get(OUTPUT_TAG).output(output);
}
}
@@ -129,6 +141,9 @@ public class KafkaWriteSchemaTransformProvider
}
}
+ @SuppressWarnings({
+ "nullness" // TODO(https://github.com/apache/beam/issues/20497)
+ })
@Override
public PCollectionRowTuple expand(PCollectionRowTuple input) {
Schema inputSchema = input.get("input").getSchema();
@@ -145,13 +160,17 @@ public class KafkaWriteSchemaTransformProvider
toBytesFn = AvroUtils.getRowToAvroBytesFunction(inputSchema);
}
+ boolean handleErrors = ErrorHandling.hasOutput(configuration.getErrorHandling());
final Map<String, String> configOverrides = configuration.getProducerConfigUpdates();
+ Schema errorSchema = ErrorHandling.errorSchema(inputSchema);
PCollectionTuple outputTuple =
input
.get("input")
.apply(
"Map rows to Kafka messages",
- ParDo.of(new ErrorCounterFn("Kafka-write-error-counter", toBytesFn))
+ ParDo.of(
+ new ErrorCounterFn(
+ "Kafka-write-error-counter", toBytesFn, errorSchema, handleErrors))
.withOutputTags(OUTPUT_TAG, TupleTagList.of(ERROR_TAG)));
outputTuple
@@ -167,8 +186,11 @@ public class KafkaWriteSchemaTransformProvider
.withKeySerializer(ByteArraySerializer.class)
.withValueSerializer(ByteArraySerializer.class));
+ // TODO: include output from KafkaIO Write once updated from PDone
+ PCollection<Row> errorOutput =
+ outputTuple.get(ERROR_TAG).setRowSchema(ErrorHandling.errorSchema(errorSchema));
return PCollectionRowTuple.of(
- "errors", outputTuple.get(ERROR_TAG).setRowSchema(ERROR_SCHEMA));
+ handleErrors ? configuration.getErrorHandling().getOutput() : "errors", errorOutput);
}
}
@@ -227,6 +249,10 @@ public class KafkaWriteSchemaTransformProvider
@Nullable
public abstract Map<String, String> getProducerConfigUpdates();
+ @SchemaFieldDescription("This option specifies whether and where to output unwritable rows.")
+ @Nullable
+ public abstract ErrorHandling getErrorHandling();
+
public static Builder builder() {
return new AutoValue_KafkaWriteSchemaTransformProvider_KafkaWriteSchemaTransformConfiguration
.Builder();
@@ -242,6 +268,8 @@ public class KafkaWriteSchemaTransformProvider
public abstract Builder setProducerConfigUpdates(Map<String, String> producerConfigUpdates);
+ public abstract Builder setErrorHandling(ErrorHandling errorHandling);
+
public abstract KafkaWriteSchemaTransformConfiguration build();
}
}
diff --git a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaDlqTest.java b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaDlqTest.java
index 48fe969bc9f..e65d9591a0b 100644
--- a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaDlqTest.java
+++ b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaDlqTest.java
@@ -22,6 +22,7 @@ import java.util.Collections;
import java.util.List;
import org.apache.beam.sdk.io.kafka.KafkaReadSchemaTransformProvider.ErrorFn;
import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.transforms.providers.ErrorHandling;
import org.apache.beam.sdk.schemas.utils.JsonUtils;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
@@ -47,7 +48,6 @@ public class KafkaDlqTest {
private static final Schema BEAMSCHEMA =
Schema.of(Schema.Field.of("name", Schema.FieldType.STRING));
- private static final Schema ERRORSCHEMA = KafkaReadSchemaTransformProvider.ERROR_SCHEMA;
private static final List<Row> ROWS =
Arrays.asList(
@@ -75,13 +75,14 @@ public class KafkaDlqTest {
} catch (Exception e) {
}
PCollection<byte[]> input = p.apply(Create.of(messages));
+ Schema errorSchema = ErrorHandling.errorSchemaBytes();
PCollectionTuple output =
input.apply(
- ParDo.of(new ErrorFn("Kafka-read-error-counter", valueMapper))
+ ParDo.of(new ErrorFn("Kafka-read-error-counter", valueMapper, errorSchema, true))
.withOutputTags(OUTPUTTAG, TupleTagList.of(ERRORTAG)));
output.get(OUTPUTTAG).setRowSchema(BEAMSCHEMA);
- output.get(ERRORTAG).setRowSchema(ERRORSCHEMA);
+ output.get(ERRORTAG).setRowSchema(errorSchema);
PAssert.that(output.get(OUTPUTTAG)).containsInAnyOrder(ROWS);
p.run().waitUntilFinish();
@@ -98,13 +99,14 @@ public class KafkaDlqTest {
} catch (Exception e) {
}
PCollection<byte[]> input = p.apply(Create.of(messagesWithError));
+ Schema errorSchema = ErrorHandling.errorSchemaBytes();
PCollectionTuple output =
input.apply(
- ParDo.of(new ErrorFn("Read-Error-Counter", valueMapper))
+ ParDo.of(new ErrorFn("Read-Error-Counter", valueMapper, errorSchema, true))
.withOutputTags(OUTPUTTAG, TupleTagList.of(ERRORTAG)));
output.get(OUTPUTTAG).setRowSchema(BEAMSCHEMA);
- output.get(ERRORTAG).setRowSchema(ERRORSCHEMA);
+ output.get(ERRORTAG).setRowSchema(errorSchema);
PCollection<Long> count = output.get(ERRORTAG).apply("error_count", Count.globally());
diff --git a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProviderTest.java b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProviderTest.java
index 367f0f7a535..1570a33c258 100644
--- a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProviderTest.java
+++ b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProviderTest.java
@@ -100,7 +100,8 @@ public class KafkaReadSchemaTransformProviderTest {
"consumerConfigUpdates",
"format",
"confluentSchemaRegistrySubject",
- "confluentSchemaRegistryUrl"),
+ "confluentSchemaRegistryUrl",
+ "errorHandling"),
kafkaProvider.configurationSchema().getFields().stream()
.map(field -> field.getName())
.collect(Collectors.toSet()));
diff --git a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProviderTest.java b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProviderTest.java
index 3a7769f42d7..b8649727f36 100644
--- a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProviderTest.java
+++ b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProviderTest.java
@@ -24,6 +24,7 @@ import java.util.Arrays;
import java.util.List;
import org.apache.beam.sdk.io.kafka.KafkaWriteSchemaTransformProvider.KafkaWriteSchemaTransform.ErrorCounterFn;
import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.transforms.providers.ErrorHandling;
import org.apache.beam.sdk.schemas.utils.JsonUtils;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
@@ -53,7 +54,6 @@ public class KafkaWriteSchemaTransformProviderTest {
private static final Schema BEAMRAWSCHEMA =
Schema.of(Schema.Field.of("payload", Schema.FieldType.BYTES));
- private static final Schema ERRORSCHEMA = KafkaWriteSchemaTransformProvider.ERROR_SCHEMA;
private static final List<Row> ROWS =
Arrays.asList(
@@ -93,12 +93,14 @@ public class KafkaWriteSchemaTransformProviderTest {
KV.of(new byte[1], "{\"name\":\"c\"}".getBytes("UTF8")));
PCollection<Row> input = p.apply(Create.of(ROWS));
+ Schema errorSchema = ErrorHandling.errorSchema(BEAMSCHEMA);
PCollectionTuple output =
input.apply(
- ParDo.of(new ErrorCounterFn("Kafka-write-error-counter", valueMapper))
+ ParDo.of(
+ new ErrorCounterFn("Kafka-write-error-counter", valueMapper, errorSchema, true))
.withOutputTags(OUTPUT_TAG, TupleTagList.of(ERROR_TAG)));
- output.get(ERROR_TAG).setRowSchema(ERRORSCHEMA);
+ output.get(ERROR_TAG).setRowSchema(errorSchema);
PAssert.that(output.get(OUTPUT_TAG)).containsInAnyOrder(msg);
p.run().waitUntilFinish();
@@ -113,12 +115,15 @@ public class KafkaWriteSchemaTransformProviderTest {
KV.of(new byte[1], "c".getBytes("UTF8")));
PCollection<Row> input = p.apply(Create.of(RAW_ROWS));
+ Schema errorSchema = ErrorHandling.errorSchema(BEAMRAWSCHEMA);
PCollectionTuple output =
input.apply(
- ParDo.of(new ErrorCounterFn("Kafka-write-error-counter", valueRawMapper))
+ ParDo.of(
+ new ErrorCounterFn(
+ "Kafka-write-error-counter", valueRawMapper, errorSchema, true))
.withOutputTags(OUTPUT_TAG, TupleTagList.of(ERROR_TAG)));
- output.get(ERROR_TAG).setRowSchema(ERRORSCHEMA);
+ output.get(ERROR_TAG).setRowSchema(errorSchema);
PAssert.that(output.get(OUTPUT_TAG)).containsInAnyOrder(msg);
p.run().waitUntilFinish();
diff --git a/sdks/python/apache_beam/yaml/standard_io.yaml b/sdks/python/apache_beam/yaml/standard_io.yaml
index aa243313d4e..116d200a86a 100644
--- a/sdks/python/apache_beam/yaml/standard_io.yaml
+++ b/sdks/python/apache_beam/yaml/standard_io.yaml
@@ -61,11 +61,13 @@
'confluent_schema_registry_url': 'confluentSchemaRegistryUrl'
'confluent_schema_registry_subject': 'confluentSchemaRegistrySubject'
'auto_offset_reset_config': 'autoOffsetResetConfig'
+ 'error_handling': 'errorHandling'
'WriteToKafka':
'format': 'format'
'topic': 'topic'
'bootstrap_servers': 'bootstrapServers'
'producer_config_updates': 'ProducerConfigUpdates'
+ 'error_handling': 'errorHandling'
underlying_provider:
type: beamJar
transforms:
diff --git a/sdks/python/apache_beam/yaml/yaml_transform.py b/sdks/python/apache_beam/yaml/yaml_transform.py
index 3fde097f5c2..ff5547db034 100644
--- a/sdks/python/apache_beam/yaml/yaml_transform.py
+++ b/sdks/python/apache_beam/yaml/yaml_transform.py
@@ -227,6 +227,8 @@ class Scope(LightweightScope):
outputs = self.get_outputs(transform)
if output in outputs:
return outputs[output]
+ elif len(outputs) == 1 and outputs[next(iter(outputs))].tag == output:
+ return outputs[next(iter(outputs))]
else:
raise ValueError(
f'Unknown output {repr(output)} '