ValueKind coder

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));

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<ValueKind> {
+
+  public static ValueKindCoder of() {
+    return INSTANCE;
+  }
+
+  private static final ValueKindCoder INSTANCE = new ValueKindCoder();
+  private static final TypeDescriptor<ValueKind> 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<ValueKind> getEncodedTypeDescriptor() {
+    return TYPE_DESCRIPTOR;
+  }
+}

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<ValueKind> TEST_CODER = ValueKindCoder.of();
+
+  private static final List<ValueKind> 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<String> 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)));
+  }
+}