diff --git a/CHANGES.md b/CHANGES.md index 319520f94309..18707b9cc848 100644 --- a/CHANGES.md +++ b/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 diff --git a/playground/frontend/playground_components/assets/symbols/java.g.yaml b/playground/frontend/playground_components/assets/symbols/java.g.yaml index e0f2b1d15d14..3b562fa787b9 100644 --- a/playground/frontend/playground_components/assets/symbols/java.g.yaml +++ b/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 diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPipelineTranslator.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPipelineTranslator.java index 113cdbd91290..95c3f2103146 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPipelineTranslator.java +++ b/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 private final int parallelism; private final int maxParallelism; private final Coder destinationCoder; - private final ShardedKeyCoder shardedKeyCoder = ShardedKeyCoder.of(VarIntCoder.of()); + private final ShardedKey.Coder shardedKeyCoder = ShardedKey.Coder.of(VarIntCoder.of()); private transient Cache>> cache; private int shardNumber = -1; diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkStreamingPipelineTranslatorTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkStreamingPipelineTranslatorTest.java index 84f1dc3c6457..373193905607 100644 --- a/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkStreamingPipelineTranslatorTest.java +++ b/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; diff --git a/runners/spark/3/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java b/runners/spark/3/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java index 5b9e5b6fae86..3ca8e45b6f87 100644 --- a/runners/spark/3/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java +++ b/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); diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/ShardedKeyCoder.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/ShardedKeyCoder.java deleted file mode 100644 index fef0b03dbc92..000000000000 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/ShardedKeyCoder.java +++ /dev/null @@ -1,63 +0,0 @@ -/* - * 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.IOException; -import java.io.InputStream; -import java.io.OutputStream; -import java.util.Arrays; -import java.util.List; -import org.apache.beam.sdk.values.ShardedKey; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; - -/** A {@link Coder} for {@link ShardedKey}, using a wrapped key {@link Coder}. */ -@VisibleForTesting -public class ShardedKeyCoder extends StructuredCoder> { - public static ShardedKeyCoder of(Coder keyCoder) { - return new ShardedKeyCoder<>(keyCoder); - } - - private final Coder keyCoder; - private final VarIntCoder shardNumberCoder; - - protected ShardedKeyCoder(Coder keyCoder) { - this.keyCoder = keyCoder; - this.shardNumberCoder = VarIntCoder.of(); - } - - @Override - public List> getCoderArguments() { - return Arrays.asList(keyCoder); - } - - @Override - public void encode(ShardedKey key, OutputStream outStream) throws IOException { - keyCoder.encode(key.getKey(), outStream); - shardNumberCoder.encode(key.getShardNumber(), outStream); - } - - @Override - public ShardedKey decode(InputStream inStream) throws IOException { - return ShardedKey.of(keyCoder.decode(inStream), shardNumberCoder.decode(inStream)); - } - - @Override - public void verifyDeterministic() throws NonDeterministicException { - keyCoder.verifyDeterministic(); - } -} diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/ShardingFunction.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/ShardingFunction.java index 31e3faac5904..695e3b5f48d1 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/ShardingFunction.java +++ b/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 extends Serializable { diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java index b0b5051f3210..e10b3e025998 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java +++ b/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> expand(PCollection 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> expand(PCollection 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>> expand(PCollection inp .withOutputTags(shardTag, TupleTagList.of(BAD_RECORD_TAG))); addErrorCollection(shardedElements); - PCollection, Iterable>> shardedInput = + PCollection, Iterable>> shardedInput = shardedElements .get(shardTag) .setCoder(KvCoder.of(VarIntCoder.of(), input.getCoder())) @@ -996,8 +995,7 @@ public PCollection>> expand(PCollection 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> writtenRecordsTag = new TupleTag<>("writtenRecords"); // Write grouped elements to temp files. @@ -1007,12 +1005,11 @@ public PCollection>> expand(PCollection inp "AddDummyShard", MapElements.via( new SimpleFunction< - KV, Iterable>, + KV, Iterable>, KV, Iterable>>() { @Override public KV, Iterable> apply( - KV, Iterable> - input) { + KV, Iterable> 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, Iterable> 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())) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/ShardedKey.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/ShardedKey.java index dbc7b2e215e2..866533f1dbee 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/ShardedKey.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/ShardedKey.java @@ -22,6 +22,8 @@ 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; @@ -29,11 +31,17 @@ 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. + * + *

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 { +public class ShardedKey implements Serializable { + private static final long serialVersionUID = 1L; private final K key; private final byte[] shardId; @@ -53,12 +61,42 @@ public static ShardedKey of(K key, byte[] shardId) { return new ShardedKey(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 ShardedKey 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(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) + "}"; } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ShardedKey.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ShardedKey.java deleted file mode 100644 index 544a5a960828..000000000000 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ShardedKey.java +++ /dev/null @@ -1,71 +0,0 @@ -/* - * 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.values; - -import java.io.Serializable; -import java.util.Objects; -import org.checkerframework.checker.nullness.qual.Nullable; - -/** - * A key and a shard number. - * - * @deprecated - *

Use {@link org.apache.beam.sdk.util.ShardedKey} instead. - */ -@Deprecated -public class ShardedKey implements Serializable { - private static final long serialVersionUID = 1L; - private final K key; - private final int shardNumber; - - public static ShardedKey of(K key, int shardNumber) { - return new ShardedKey<>(key, shardNumber); - } - - private ShardedKey(K key, int shardNumber) { - this.key = key; - this.shardNumber = shardNumber; - } - - public K getKey() { - return key; - } - - public int getShardNumber() { - return shardNumber; - } - - @Override - public String toString() { - return "key: " + key + " shard: " + shardNumber; - } - - @Override - public boolean equals(@Nullable Object o) { - if (!(o instanceof ShardedKey)) { - return false; - } - ShardedKey other = (ShardedKey) o; - return Objects.equals(key, other.key) && shardNumber == other.shardNumber; - } - - @Override - public int hashCode() { - return Objects.hash(key, shardNumber); - } -} diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/WriteFilesTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/WriteFilesTest.java index cc174002bb46..87e978e9f32b 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/WriteFilesTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/WriteFilesTest.java @@ -88,12 +88,12 @@ import org.apache.beam.sdk.transforms.windowing.PaneInfo; import org.apache.beam.sdk.transforms.windowing.Sessions; import org.apache.beam.sdk.transforms.windowing.Window; +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; import org.apache.beam.sdk.values.PCollectionView; import org.apache.beam.sdk.values.PDone; -import org.apache.beam.sdk.values.ShardedKey; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Optional; 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; diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/ShardedKeyTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/ShardedKeyTest.java index 553b51c494d0..908272078b22 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/ShardedKeyTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/ShardedKeyTest.java @@ -71,4 +71,36 @@ public void testEquality() { ShardedKey.of("key", "shard_id".getBytes(UTF_8)), ShardedKey.of("key", "shard_id".getBytes(UTF_8))); } + + @Test + public void testIntShardNumber() { + ShardedKey key = ShardedKey.of("key", 42); + assertEquals("key", key.getKey()); + assertEquals(42, key.getShardNumber()); + } + + @Test + public void testIntShardNumberEquality() { + assertEquals(ShardedKey.of("key", 0), ShardedKey.of("key", 0)); + assertEquals(ShardedKey.of("key", 42), ShardedKey.of("key", 42)); + assertEquals(ShardedKey.of("key", -1), ShardedKey.of("key", -1)); + } + + @Test + public void testIntShardNumberCoderRoundTrip() throws Exception { + Coder> coder = ShardedKey.Coder.of(StringUtf8Coder.of()); + CoderProperties.coderDecodeEncodeEqual(coder, ShardedKey.of(KEY, 0)); + CoderProperties.coderDecodeEncodeEqual(coder, ShardedKey.of(KEY, 42)); + CoderProperties.coderDecodeEncodeEqual(coder, ShardedKey.of(KEY, -1)); + } + + @Test + public void testGetShardNumberPreservedAfterCoderRoundTrip() throws Exception { + Coder> coder = ShardedKey.Coder.of(StringUtf8Coder.of()); + ShardedKey original = ShardedKey.of(KEY, 42); + byte[] encoded = org.apache.beam.sdk.util.CoderUtils.encodeToByteArray(coder, original); + ShardedKey decoded = + org.apache.beam.sdk.util.CoderUtils.decodeFromByteArray(coder, encoded); + assertEquals(42, decoded.getShardNumber()); + } } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java index 252e55d34c07..3d5b3d259299 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java @@ -34,7 +34,6 @@ import org.apache.beam.sdk.coders.IterableCoder; import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.coders.NullableCoder; -import org.apache.beam.sdk.coders.ShardedKeyCoder; import org.apache.beam.sdk.extensions.gcp.options.GcsOptions; import org.apache.beam.sdk.extensions.gcp.util.gcsfs.GcsPath; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.CreateDisposition; @@ -70,12 +69,12 @@ import org.apache.beam.sdk.transforms.windowing.TimestampCombiner; import org.apache.beam.sdk.transforms.windowing.Window; import org.apache.beam.sdk.util.Preconditions; +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.PCollectionList; import org.apache.beam.sdk.values.PCollectionTuple; import org.apache.beam.sdk.values.PCollectionView; -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.sdk.values.TypeDescriptor; @@ -642,7 +641,7 @@ PCollection> writeDynamicallyShardedFil PCollection, ElementT>> unwrittenRecords = writeBundlesTuple .get(unwrittedRecordsTag) - .setCoder(KvCoder.of(ShardedKeyCoder.of(destinationCoder), elementCoder)); + .setCoder(KvCoder.of(ShardedKey.Coder.of(destinationCoder), elementCoder)); badRecordErrorHandler.addErrorCollection( writeBundlesTuple.get(BAD_RECORD_TAG).setCoder(BadRecord.getCoder(input.getPipeline()))); @@ -689,7 +688,7 @@ public void processElement( element.getValue())); } })) - .setCoder(KvCoder.of(ShardedKeyCoder.of(destinationCoder), elementCoder)); + .setCoder(KvCoder.of(ShardedKey.Coder.of(destinationCoder), elementCoder)); return writeShardedRecords(shardedRecords, tempFilePrefix); } @@ -723,18 +722,17 @@ PCollection> writeDynamicallyShardedFil .withShardedKey()) .setCoder( KvCoder.of( - org.apache.beam.sdk.util.ShardedKey.Coder.of(destinationCoder), + ShardedKey.Coder.of(destinationCoder), IterableCoder.of(elementCoder))) .apply( "StripShardId", MapElements.via( new SimpleFunction< - KV, Iterable>, + KV, Iterable>, KV>>() { @Override public KV> apply( - KV, Iterable> - input) { + KV, Iterable> input) { return KV.of(input.getKey().getKey(), input.getValue()); } })) @@ -811,7 +809,7 @@ private PCollection> writeTempTables( Coder, WritePartition.Result>> partitionsCoder = KvCoder.of( - ShardedKeyCoder.of(NullableCoder.of(destinationCoder)), + ShardedKey.Coder.of(NullableCoder.of(destinationCoder)), WritePartition.ResultCoder.INSTANCE); // If WriteBundlesToFiles produced more than DEFAULT_MAX_FILES_PER_PARTITION files or @@ -857,7 +855,7 @@ PCollection writeSinglePartition( Coder, WritePartition.Result>> partitionsCoder = KvCoder.of( - ShardedKeyCoder.of(NullableCoder.of(destinationCoder)), + ShardedKey.Coder.of(NullableCoder.of(destinationCoder)), WritePartition.ResultCoder.INSTANCE); // Write single partition to final table PCollection> successfulWrites = diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/GenerateShardedTable.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/GenerateShardedTable.java index 8f9c8363a6bd..0488736b08d2 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/GenerateShardedTable.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/GenerateShardedTable.java @@ -21,8 +21,8 @@ import java.util.concurrent.ThreadLocalRandom; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.ShardedKey; import org.apache.beam.sdk.values.KV; -import org.apache.beam.sdk.values.ShardedKey; /** * Given a write to a specific table, assign that to one of the {@link diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java index 2cc08d7731f2..55e5942fc49b 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java @@ -21,7 +21,6 @@ import org.apache.beam.sdk.coders.AtomicCoder; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.KvCoder; -import org.apache.beam.sdk.coders.ShardedKeyCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.transforms.MapElements; import org.apache.beam.sdk.transforms.PTransform; @@ -33,10 +32,10 @@ import org.apache.beam.sdk.transforms.windowing.GlobalWindows; import org.apache.beam.sdk.transforms.windowing.Window; import org.apache.beam.sdk.util.Preconditions; +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.PCollectionTuple; -import org.apache.beam.sdk.values.ShardedKey; import org.apache.beam.sdk.values.TupleTag; import org.checkerframework.checker.nullness.qual.Nullable; import org.checkerframework.checker.nullness.qual.RequiresNonNull; @@ -412,11 +411,11 @@ public KV apply(KV input) { PCollection, TableRowInfo>> shardedTagged = input .apply("ShardTableWrites", ParDo.of(new GenerateShardedTable<>(numShards))) - .setCoder(KvCoder.of(ShardedKeyCoder.of(StringUtf8Coder.of()), elementCoder)) + .setCoder(KvCoder.of(ShardedKey.Coder.of(StringUtf8Coder.of()), elementCoder)) .apply("TagWithUniqueIds", ParDo.of(new TagWithUniqueIds<>(deterministicRecordIdFn))) .setCoder( KvCoder.of( - ShardedKeyCoder.of(StringUtf8Coder.of()), + ShardedKey.Coder.of(StringUtf8Coder.of()), TableRowInfoCoder.of(elementCoder))); if (deterministicRecordIdFn == null) { diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteBundlesToFiles.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteBundlesToFiles.java index 8c6893ef5798..4094d5f0fefa 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteBundlesToFiles.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteBundlesToFiles.java @@ -38,9 +38,9 @@ import org.apache.beam.sdk.transforms.errorhandling.BadRecordRouter; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.util.Preconditions; +import org.apache.beam.sdk.util.ShardedKey; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollectionView; -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.collect.Lists; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WritePartition.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WritePartition.java index cdebf1172cba..f3b244b99eb5 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WritePartition.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WritePartition.java @@ -30,8 +30,8 @@ import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.util.ShardedKey; import org.apache.beam.sdk.values.PCollectionView; -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.collect.Lists; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java index 070f849e43e6..8140a1bb51da 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java @@ -68,8 +68,8 @@ import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollectionTuple; +import org.apache.beam.sdk.util.ShardedKey; import org.apache.beam.sdk.values.PCollectionView; -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.Strings; diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryHelpersTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryHelpersTest.java index 0d36d7bb46d0..94f038815a96 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryHelpersTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryHelpersTest.java @@ -33,7 +33,6 @@ import java.util.Set; import org.apache.beam.sdk.coders.CoderException; import org.apache.beam.sdk.coders.KvCoder; -import org.apache.beam.sdk.coders.ShardedKeyCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.extensions.gcp.util.BackOffAdapter; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.PendingJob; @@ -43,6 +42,7 @@ import org.apache.beam.sdk.transforms.windowing.IntervalWindow; import org.apache.beam.sdk.util.CoderUtils; import org.apache.beam.sdk.util.FluentBackoff; +import org.apache.beam.sdk.util.ShardedKey; import org.apache.beam.sdk.values.WindowedValues; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets; @@ -164,7 +164,7 @@ public void testCoder_nullCell() throws CoderException { @Test public void testShardedKeyCoderIsSerializableWithWellKnownCoderType() { - CoderProperties.coderSerializable(ShardedKeyCoder.of(GlobalWindow.Coder.INSTANCE)); + CoderProperties.coderSerializable(ShardedKey.Coder.of(GlobalWindow.Coder.INSTANCE)); } @Test @@ -177,7 +177,7 @@ public void testComplexCoderSerializable() { CoderProperties.coderSerializable( WindowedValues.getFullCoder( KvCoder.of( - ShardedKeyCoder.of(StringUtf8Coder.of()), + ShardedKey.Coder.of(StringUtf8Coder.of()), TableRowInfoCoder.of(TableRowJsonCoder.of())), IntervalWindow.getCoder())); } diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java index e5ad761993e7..2cdcdd038df0 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java @@ -115,7 +115,6 @@ import org.apache.beam.sdk.coders.IterableCoder; import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.coders.SerializableCoder; -import org.apache.beam.sdk.coders.ShardedKeyCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.coders.VarLongCoder; import org.apache.beam.sdk.extensions.avro.coders.AvroGenericCoder; @@ -178,12 +177,12 @@ import org.apache.beam.sdk.transforms.windowing.WindowFn; import org.apache.beam.sdk.transforms.windowing.WindowMappingFn; import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.util.ShardedKey; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PBegin; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollectionView; import org.apache.beam.sdk.values.Row; -import org.apache.beam.sdk.values.ShardedKey; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.TypeDescriptor; import org.apache.beam.sdk.values.TypeDescriptors; @@ -2976,7 +2975,7 @@ public void testWriteTables() throws Exception { p.apply( Create.of(partitions) .withCoder( - KvCoder.of(ShardedKeyCoder.of(StringUtf8Coder.of()), ResultCoder.INSTANCE))); + KvCoder.of(ShardedKey.Coder.of(StringUtf8Coder.of()), ResultCoder.INSTANCE))); PCollectionView jobIdTokenView = p.apply("CreateJobId", Create.of("jobId")).apply(View.asSingleton()); List> sideInputs = ImmutableList.of(jobIdTokenView);