diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/CoderRegistry.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/CoderRegistry.java index 2cf79a723d63..c4f8e34d61ea 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/CoderRegistry.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/CoderRegistry.java @@ -52,6 +52,7 @@ import org.apache.beam.sdk.values.TimestampedValue; import org.apache.beam.sdk.values.TypeDescriptor; import org.apache.beam.sdk.values.TypeDescriptors; +import org.apache.beam.sdk.values.ValueKind; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.HashMultimap; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; @@ -137,6 +138,8 @@ private CommonTypes() { TimestampedValue.class, CoderProviders.fromStaticMethods( TimestampedValue.class, TimestampedValue.TimestampedValueCoder.class)); + builder.put( + ValueKind.class, CoderProviders.fromStaticMethods(ValueKind.class, ValueKindCoder.class)); builder.put(Void.class, CoderProviders.fromStaticMethods(Void.class, VoidCoder.class)); builder.put( byte[].class, CoderProviders.fromStaticMethods(byte[].class, ByteArrayCoder.class)); diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/ValueKindCoder.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/ValueKindCoder.java new file mode 100644 index 000000000000..24f1a119f429 --- /dev/null +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/ValueKindCoder.java @@ -0,0 +1,88 @@ +/* + * 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.coders; + +import java.io.EOFException; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements; +import org.apache.beam.sdk.values.TypeDescriptor; +import org.apache.beam.sdk.values.ValueKind; +import org.apache.beam.sdk.values.ValueKindUtil; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * A {@link Coder} for {@link ValueKind}, encoded in 1 byte as the matching {@link + * Elements.ValueKind.Enum} number, so the wire format stays stable if the enum is reordered and + * matches the portability representation. + */ +public class ValueKindCoder extends AtomicCoder { + + public static ValueKindCoder of() { + return INSTANCE; + } + + private static final ValueKindCoder INSTANCE = new ValueKindCoder(); + private static final TypeDescriptor TYPE_DESCRIPTOR = + TypeDescriptor.of(ValueKind.class); + + private ValueKindCoder() {} + + @Override + public void encode(ValueKind value, OutputStream outStream) throws IOException, CoderException { + if (value == null) { + throw new CoderException("cannot encode a null ValueKind"); + } + outStream.write(ValueKindUtil.toProto(value).getNumber()); + } + + @Override + public ValueKind decode(InputStream inStream) throws IOException, CoderException { + int number = inStream.read(); + if (number == -1) { + throw new CoderException(new EOFException("EOF encountered decoding a ValueKind")); + } + Elements.ValueKind.@Nullable Enum proto = Elements.ValueKind.Enum.forNumber(number); + if (proto == null) { + throw new CoderException("Unknown ValueKind number: " + number); + } + + return ValueKindUtil.fromProto(proto); + } + + @Override + public boolean consistentWithEquals() { + return true; + } + + @Override + public boolean isRegisterByteSizeObserverCheap(ValueKind value) { + return true; + } + + @Override + protected long getEncodedElementByteSize(ValueKind value) { + return 1; + } + + @Override + public TypeDescriptor getEncodedTypeDescriptor() { + return TYPE_DESCRIPTOR; + } +} diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/coders/ValueKindCoderTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/coders/ValueKindCoderTest.java new file mode 100644 index 000000000000..aeac2132c853 --- /dev/null +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/coders/ValueKindCoderTest.java @@ -0,0 +1,81 @@ +/* + * 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.coders; + +import static org.junit.Assert.assertEquals; + +import java.util.Arrays; +import java.util.List; +import org.apache.beam.sdk.testing.CoderProperties; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.values.TypeDescriptor; +import org.apache.beam.sdk.values.ValueKind; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.ExpectedException; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link ValueKindCoder}. */ +@RunWith(JUnit4.class) +public class ValueKindCoderTest { + + private static final Coder TEST_CODER = ValueKindCoder.of(); + + private static final List TEST_VALUES = + Arrays.asList( + ValueKind.INSERT, ValueKind.UPDATE_BEFORE, ValueKind.UPDATE_AFTER, ValueKind.DELETE); + + /** One byte per value, holding the proto enum number. */ + private static final List TEST_ENCODINGS = Arrays.asList("AQ", "Ag", "Aw", "BA"); + + @Rule public ExpectedException thrown = ExpectedException.none(); + + @Test + public void testDecodeEncodeEqual() throws Exception { + for (ValueKind value : TEST_VALUES) { + CoderProperties.coderDecodeEncodeEqual(TEST_CODER, value); + } + } + + @Test + public void testWireFormatEncode() throws Exception { + CoderProperties.coderEncodesBase64(TEST_CODER, TEST_VALUES, TEST_ENCODINGS); + } + + /** VALUE_KIND_UNSPECIFIED (0) means INSERT, for backwards compatibility. */ + @Test + public void testDecodeUnspecified() throws Exception { + assertEquals(ValueKind.INSERT, CoderUtils.decodeFromBase64(TEST_CODER, "AA")); + } + + @Test + public void testDecodeUnknownNumberThrows() throws Exception { + thrown.expect(CoderException.class); + thrown.expectMessage("Unknown ValueKind number: 42"); + + CoderUtils.decodeFromBase64(TEST_CODER, "Kg"); + } + + @Test + public void testCoderRegistryResolvesValueKind() throws Exception { + assertEquals( + ValueKindCoder.of(), + CoderRegistry.createDefault().getCoder(TypeDescriptor.of(ValueKind.class))); + } +}