Consolidate duplicate ShardedKey implementations in core module

CHANGES.md

@@ -113,6 +113,7 @@
 
 ## Breaking Changes
 
+* Removed the previously-deprecated `org.apache.beam.sdk.values.ShardedKey` and the internal `org.apache.beam.sdk.coders.ShardedKeyCoder` (`@VisibleForTesting`). All usages should migrate to `org.apache.beam.sdk.util.ShardedKey`, which now provides `of(K, int)` and `getShardNumber()` convenience methods for integer shard numbers. The wire format of `ShardedKey.Coder` (from `sdk.util`) is unchanged; only the removed `ShardedKeyCoder` (from `sdk.coders`) had a different format, and it was not a Beam standard/model coder (Java) ([#37990](https://github.com/apache/beam/issues/37990)).
 * The Python SDK container's `boot.go` now passes pipeline options through a file instead of the `PIPELINE_OPTIONS` environment variable. If a user pairs a new Python SDK container with an older SDK version (which does not support the file-based approach), the pipeline options will not be recognized and the pipeline will fail. Users must ensure their SDK and container versions are synchronized ([#37370](https://github.com/apache/beam/issues/37370)).
 
 ## Bugfixes

playground/frontend/playground_components/assets/symbols/java.g.yaml

@@ -10741,13 +10741,6 @@ ShardedKey:
   - structuralValue
   - toString
   - verifyDeterministic
-ShardedKeyCoder: 
-  methods: 
-  - decode
-  - encode
-  - getCoderArguments
-  - of
-  - verifyDeterministic
 ShardingWritableByteChannel: 
   methods: 
   - addChannel

runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPipelineTranslator.java

@@ -30,7 +30,6 @@
 import org.apache.beam.runners.flink.adapter.FlinkKey;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.coders.ShardedKeyCoder;
 import org.apache.beam.sdk.coders.VarIntCoder;
 import org.apache.beam.sdk.io.FileBasedSink;
 import org.apache.beam.sdk.io.ShardingFunction;
@@ -49,9 +48,9 @@
 import org.apache.beam.sdk.util.construction.UnconsumedReads;
 import org.apache.beam.sdk.util.construction.WriteFilesTranslation;
 import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.util.ShardedKey;
 import org.apache.beam.sdk.values.PCollectionView;
 import org.apache.beam.sdk.values.PValue;
-import org.apache.beam.sdk.values.ShardedKey;
 import org.apache.beam.sdk.values.TupleTag;
 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.cache.Cache;
@@ -320,7 +319,7 @@ static class FlinkAutoBalancedShardKeyShardingFunction<UserT, DestinationT>
     private final int parallelism;
     private final int maxParallelism;
     private final Coder<DestinationT> destinationCoder;
-    private final ShardedKeyCoder<Integer> shardedKeyCoder = ShardedKeyCoder.of(VarIntCoder.of());
+    private final ShardedKey.Coder<Integer> shardedKeyCoder = ShardedKey.Coder.of(VarIntCoder.of());
     private transient Cache<Integer, Map<Integer, ShardedKey<Integer>>> cache;
 
     private int shardNumber = -1;

runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkStreamingPipelineTranslatorTest.java

@@ -43,9 +43,9 @@
 import org.apache.beam.sdk.transforms.GroupByKey;
 import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.util.SerializableUtils;
+import org.apache.beam.sdk.util.ShardedKey;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.ShardedKey;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
 import org.apache.flink.runtime.jobgraph.JobGraph;

runners/spark/3/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java

@@ -57,7 +57,6 @@
 import org.apache.beam.sdk.coders.RowCoder;
 import org.apache.beam.sdk.coders.SerializableCoder;
 import org.apache.beam.sdk.coders.SetCoder;
-import org.apache.beam.sdk.coders.ShardedKeyCoder;
 import org.apache.beam.sdk.coders.SnappyCoder;
 import org.apache.beam.sdk.coders.SortedMapCoder;
 import org.apache.beam.sdk.coders.StringDelegateCoder;
@@ -74,6 +73,7 @@
 import org.apache.beam.sdk.transforms.join.UnionCoder;
 import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
 import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
+import org.apache.beam.sdk.util.ShardedKey;
 import org.apache.beam.sdk.util.construction.resources.PipelineResources;
 import org.apache.beam.sdk.values.PCollectionViews;
 import org.apache.beam.sdk.values.TupleTag;
@@ -255,7 +255,7 @@ public void registerClasses(Kryo kryo) {
       kryo.register(RowCoder.class);
       kryo.register(SerializableCoder.class);
       kryo.register(SetCoder.class);
-      kryo.register(ShardedKeyCoder.class);
+      kryo.register(ShardedKey.Coder.class);
       kryo.register(SnappyCoder.class);
       kryo.register(SortedMapCoder.class);
       kryo.register(StringDelegateCoder.class);

sdks/java/core/src/main/java/org/apache/beam/sdk/io/ShardingFunction.java

@@ -18,7 +18,7 @@
 package org.apache.beam.sdk.io;
 
 import java.io.Serializable;
-import org.apache.beam.sdk.values.ShardedKey;
+import org.apache.beam.sdk.util.ShardedKey;
 
 /** Function for assigning {@link ShardedKey}s to input elements for sharded {@link WriteFiles}. */
 public interface ShardingFunction<UserT, DestinationT> extends Serializable {

sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java

@@ -37,7 +37,6 @@
 import org.apache.beam.sdk.coders.IterableCoder;
 import org.apache.beam.sdk.coders.KvCoder;
 import org.apache.beam.sdk.coders.ListCoder;
-import org.apache.beam.sdk.coders.ShardedKeyCoder;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
 import org.apache.beam.sdk.coders.VarIntCoder;
 import org.apache.beam.sdk.io.FileBasedSink.DynamicDestinations;
@@ -76,6 +75,7 @@
 import org.apache.beam.sdk.transforms.windowing.Window;
 import org.apache.beam.sdk.util.CoderUtils;
 import org.apache.beam.sdk.util.MoreFutures;
+import org.apache.beam.sdk.util.ShardedKey;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.PCollection.IsBounded;
@@ -85,7 +85,6 @@
 import org.apache.beam.sdk.values.PCollectionViews;
 import org.apache.beam.sdk.values.PDone;
 import org.apache.beam.sdk.values.PValue;
-import org.apache.beam.sdk.values.ShardedKey;
 import org.apache.beam.sdk.values.TupleTag;
 import org.apache.beam.sdk.values.TupleTagList;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Objects;
@@ -627,7 +626,7 @@ public PCollection<FileResult<DestinationT>> expand(PCollection<UserT> input) {
       PCollectionTuple spilledWriteTuple =
           writeTuple
               .get(unwrittenRecordsTag)
-              .setCoder(KvCoder.of(ShardedKeyCoder.of(VarIntCoder.of()), input.getCoder()))
+              .setCoder(KvCoder.of(ShardedKey.Coder.of(VarIntCoder.of()), input.getCoder()))
               // Here we group by a synthetic shard number in the range [0, spill factor),
               // just for the sake of getting some parallelism within each destination when
               // writing the spilled records, whereas the non-spilled records don't have a shard
@@ -928,7 +927,7 @@ public PCollection<FileResult<DestinationT>> expand(PCollection<UserT> input) {
       PCollectionTuple writtenFiles =
           shardedFiles
               .get(shardedRecords)
-              .setCoder(KvCoder.of(ShardedKeyCoder.of(VarIntCoder.of()), input.getCoder()))
+              .setCoder(KvCoder.of(ShardedKey.Coder.of(VarIntCoder.of()), input.getCoder()))
               .apply("GroupIntoShards", GroupByKey.create())
               .apply(
                   "WriteShardsIntoTempFiles",
@@ -984,7 +983,7 @@ public PCollection<List<FileResult<DestinationT>>> expand(PCollection<UserT> inp
                   .withOutputTags(shardTag, TupleTagList.of(BAD_RECORD_TAG)));
       addErrorCollection(shardedElements);
 
-      PCollection<KV<org.apache.beam.sdk.util.ShardedKey<Integer>, Iterable<UserT>>> shardedInput =
+      PCollection<KV<ShardedKey<Integer>, Iterable<UserT>>> shardedInput =
           shardedElements
               .get(shardTag)
               .setCoder(KvCoder.of(VarIntCoder.of(), input.getCoder()))
@@ -996,8 +995,7 @@ public PCollection<List<FileResult<DestinationT>>> expand(PCollection<UserT> inp
                       .withShardedKey())
               .setCoder(
                   KvCoder.of(
-                      org.apache.beam.sdk.util.ShardedKey.Coder.of(VarIntCoder.of()),
-                      IterableCoder.of(input.getCoder())));
+                      ShardedKey.Coder.of(VarIntCoder.of()), IterableCoder.of(input.getCoder())));
 
       TupleTag<FileResult<DestinationT>> writtenRecordsTag = new TupleTag<>("writtenRecords");
       // Write grouped elements to temp files.
@@ -1007,12 +1005,11 @@ public PCollection<List<FileResult<DestinationT>>> expand(PCollection<UserT> inp
                   "AddDummyShard",
                   MapElements.via(
                       new SimpleFunction<
-                          KV<org.apache.beam.sdk.util.ShardedKey<Integer>, Iterable<UserT>>,
+                          KV<ShardedKey<Integer>, Iterable<UserT>>,
                           KV<ShardedKey<Integer>, Iterable<UserT>>>() {
                         @Override
                         public KV<ShardedKey<Integer>, Iterable<UserT>> apply(
-                            KV<org.apache.beam.sdk.util.ShardedKey<Integer>, Iterable<UserT>>
-                                input) {
+                            KV<ShardedKey<Integer>, Iterable<UserT>> input) {
                           // Add dummy shard since it is required by WriteShardsIntoTempFilesFn. It
                           // will be dropped after we generate the temp files.
                           return KV.of(
@@ -1022,7 +1019,7 @@ public KV<ShardedKey<Integer>, Iterable<UserT>> apply(
                       }))
               .setCoder(
                   KvCoder.of(
-                      ShardedKeyCoder.of(VarIntCoder.of()), IterableCoder.of(input.getCoder())))
+                      ShardedKey.Coder.of(VarIntCoder.of()), IterableCoder.of(input.getCoder())))
               .apply(
                   "WriteShardsIntoTempFiles",
                   ParDo.of(new WriteShardsIntoTempFilesFn(input.getCoder()))

sdks/java/core/src/main/java/org/apache/beam/sdk/util/ShardedKey.java

@@ -22,18 +22,26 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
+import java.io.Serializable;
+import java.nio.ByteBuffer;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import org.apache.beam.sdk.coders.ByteArrayCoder;
 import org.apache.beam.sdk.coders.StructuredCoder;
 import org.apache.beam.sdk.util.common.ElementByteSizeObserver;
 
-/** A sharded key consisting of a user key and an opaque shard id represented by bytes. */
+/**
+ * A sharded key consisting of a user key and a shard identifier.
+ *
+ * <p>The shard identifier is stored as an opaque byte array. Convenience methods are provided for
+ * creating sharded keys with integer shard numbers, which are encoded as 4-byte big-endian arrays.
+ */
 @SuppressWarnings({
   "nullness" // TODO(https://github.com/apache/beam/issues/20497)
 })
-public class ShardedKey<K> {
+public class ShardedKey<K> implements Serializable {
+  private static final long serialVersionUID = 1L;
 
   private final K key;
   private final byte[] shardId;
@@ -53,12 +61,42 @@ public static <K> ShardedKey<K> of(K key, byte[] shardId) {
     return new ShardedKey<K>(key, shardId);
   }
 
+  /**
+   * Creates a ShardedKey with given key and integer shard number. The shard number is stored as a
+   * 4-byte big-endian byte array.
+   */
+  public static <K> ShardedKey<K> of(K key, int shardNumber) {
+    checkArgument(key != null, "Key should not be null!");
+    byte[] shardId = ByteBuffer.allocate(Integer.BYTES).putInt(shardNumber).array();
+    return new ShardedKey<K>(key, shardId);
+  }
+
   public K getKey() {
     return key;
   }
 
+  /**
+   * Returns the integer shard number. This method should only be called on ShardedKeys that were
+   * created with {@link #of(Object, int)}, or whose shard id is a 4-byte big-endian encoded
+   * integer.
+   *
+   * @throws IllegalArgumentException if the shard id is not 4 bytes
+   */
+  public int getShardNumber() {
+    checkArgument(
+        shardId.length == Integer.BYTES,
+        "ShardedKey was not created with an integer shard number (shard id has %s bytes,"
+            + " expected %s)",
+        shardId.length,
+        Integer.BYTES);
+    return ByteBuffer.wrap(shardId).getInt();
+  }
+
   @Override
   public String toString() {
+    if (shardId.length == Integer.BYTES) {
+      return "ShardedKey{key=" + key + ", shard=" + ByteBuffer.wrap(shardId).getInt() + "}";
+    }
     return "ShardedKey{key=" + key + ", shardId=" + Arrays.toString(shardId) + "}";
   }