diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Redistribute.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Redistribute.java index a01b5f570a57..3a8bef28839a 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Redistribute.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Redistribute.java @@ -131,7 +131,7 @@ public void processElement( public static class RedistributeArbitrarily extends PTransform, PCollection> { // The number of buckets to shard into. - // A runner is free to ignore this (a runner may ignore the transorm + // A runner is free to ignore this (a runner may ignore the transform // entirely!) This is a performance optimization to prevent having // unit sized bundles on the output. If unset, uses a random integer key. private @Nullable Integer numBuckets = null; diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java index 045a74a8507e..dcd0ac3daaf0 100644 --- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java +++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java @@ -79,6 +79,7 @@ import org.apache.beam.sdk.transforms.PTransform; import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.transforms.Redistribute; +import org.apache.beam.sdk.transforms.Redistribute.RedistributeArbitrarily; import org.apache.beam.sdk.transforms.Reshuffle; import org.apache.beam.sdk.transforms.SerializableFunction; import org.apache.beam.sdk.transforms.SimpleFunction; @@ -730,6 +731,9 @@ public abstract static class Read @Pure public abstract @Nullable Boolean getOffsetDeduplication(); + @Pure + public abstract @Nullable Boolean getRedistributeByRecordKey(); + @Pure public abstract @Nullable Duration getWatchTopicPartitionDuration(); @@ -800,6 +804,8 @@ abstract Builder setConsumerFactoryFn( abstract Builder setOffsetDeduplication(Boolean offsetDeduplication); + abstract Builder setRedistributeByRecordKey(Boolean redistributeByRecordKey); + abstract Builder setTimestampPolicyFactory( TimestampPolicyFactory timestampPolicyFactory); @@ -915,11 +921,15 @@ static void setupExternalBuilder( && config.offsetDeduplication != null) { builder.setOffsetDeduplication(config.offsetDeduplication); } + if (config.redistribute && config.redistributeByRecordKey != null) { + builder.setRedistributeByRecordKey(config.redistributeByRecordKey); + } } else { builder.setRedistributed(false); builder.setRedistributeNumKeys(0); builder.setAllowDuplicates(false); builder.setOffsetDeduplication(false); + builder.setRedistributeByRecordKey(false); } } @@ -989,6 +999,7 @@ public static class Configuration { private Boolean redistribute; private Boolean allowDuplicates; private Boolean offsetDeduplication; + private Boolean redistributeByRecordKey; private Long dynamicReadPollIntervalSeconds; public void setConsumerConfig(Map consumerConfig) { @@ -1051,6 +1062,10 @@ public void setOffsetDeduplication(Boolean offsetDeduplication) { this.offsetDeduplication = offsetDeduplication; } + public void setRedistributeByRecordKey(Boolean redistributeByRecordKey) { + this.redistributeByRecordKey = redistributeByRecordKey; + } + public void setDynamicReadPollIntervalSeconds(Long dynamicReadPollIntervalSeconds) { this.dynamicReadPollIntervalSeconds = dynamicReadPollIntervalSeconds; } @@ -1161,6 +1176,10 @@ public Read withOffsetDeduplication(Boolean offsetDeduplication) { return toBuilder().setOffsetDeduplication(offsetDeduplication).build(); } + public Read withRedistributeByRecordKey(Boolean redistributeByRecordKey) { + return toBuilder().setRedistributeByRecordKey(redistributeByRecordKey).build(); + } + /** * Internally sets a {@link java.util.regex.Pattern} of topics to read from. All the partitions * from each of the matching topics are read. @@ -1679,6 +1698,11 @@ private void checkRedistributeConfiguration() { LOG.warn( "Offsets used for deduplication are available in WindowedValue's metadata. Combining, aggregating, mutating them may risk with data loss."); } + if (getRedistributeByRecordKey() != null && getRedistributeByRecordKey()) { + checkState( + isRedistributed(), + "withRedistributeByRecordKey can only be used when withRedistribute is set."); + } } private void warnAboutUnsafeConfigurations(PBegin input) { @@ -1858,18 +1882,25 @@ public PCollection> expand(PBegin input) { "Offsets committed due to usage of commitOffsetsInFinalize() and may not capture all work processed due to use of withRedistribute() with duplicates enabled"); } - if (kafkaRead.getRedistributeNumKeys() == 0) { - return output.apply( - "Insert Redistribute", - Redistribute.>arbitrarily() - .withAllowDuplicates(kafkaRead.isAllowDuplicates())); - } else { - return output.apply( - "Insert Redistribute with Shards", - Redistribute.>arbitrarily() - .withAllowDuplicates(kafkaRead.isAllowDuplicates()) - .withNumBuckets((int) kafkaRead.getRedistributeNumKeys())); + if (kafkaRead.getOffsetDeduplication() != null && kafkaRead.getOffsetDeduplication()) { + if (kafkaRead.getRedistributeByRecordKey() != null + && kafkaRead.getRedistributeByRecordKey()) { + return output.apply( + KafkaReadRedistribute.byRecordKey(kafkaRead.getRedistributeNumKeys())); + } else { + return output.apply( + KafkaReadRedistribute.byOffsetShard(kafkaRead.getRedistributeNumKeys())); + } + } + RedistributeArbitrarily> redistribute = + Redistribute.>arbitrarily() + .withAllowDuplicates(kafkaRead.isAllowDuplicates()); + String redistributeName = "Insert Redistribute"; + if (kafkaRead.getRedistributeNumKeys() != 0) { + redistribute = redistribute.withNumBuckets((int) kafkaRead.getRedistributeNumKeys()); + redistributeName = "Insert Redistribute with Shards"; } + return output.apply(redistributeName, redistribute); } return output; } diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIOReadImplementationCompatibility.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIOReadImplementationCompatibility.java index 81a1de9b872b..8c5efb066d6e 100644 --- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIOReadImplementationCompatibility.java +++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIOReadImplementationCompatibility.java @@ -139,6 +139,12 @@ Object getDefaultValue() { }, OFFSET_DEDUPLICATION(LEGACY), LOG_TOPIC_VERIFICATION, + REDISTRIBUTE_BY_RECORD_KEY { + @Override + Object getDefaultValue() { + return false; + } + }, ; private final @NonNull ImmutableSet supportedImplementations; diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadRedistribute.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadRedistribute.java new file mode 100644 index 000000000000..61c0b671f292 --- /dev/null +++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadRedistribute.java @@ -0,0 +1,124 @@ +/* + * 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.io.kafka; + +import static java.nio.charset.StandardCharsets.UTF_8; + +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.Redistribute; +import org.apache.beam.sdk.transforms.Values; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.primitives.UnsignedInteger; +import org.checkerframework.checker.nullness.qual.NonNull; +import org.checkerframework.checker.nullness.qual.Nullable; + +public class KafkaReadRedistribute + extends PTransform>, PCollection>> { + public static KafkaReadRedistribute byOffsetShard(@Nullable Integer numBuckets) { + return new KafkaReadRedistribute<>(numBuckets, false); + } + + public static KafkaReadRedistribute byRecordKey(@Nullable Integer numBuckets) { + return new KafkaReadRedistribute<>(numBuckets, true); + } + + // The number of buckets to shard into. + private @Nullable Integer numBuckets = null; + // When redistributing, group records by the Kafka record's key instead of by offset hash. + private boolean byRecordKey = false; + + private KafkaReadRedistribute(@Nullable Integer numBuckets, boolean byRecordKey) { + this.numBuckets = numBuckets; + this.byRecordKey = byRecordKey; + } + + @Override + public PCollection> expand(PCollection> input) { + + if (byRecordKey) { + return input + .apply("Pair with shard from key", ParDo.of(new AssignRecordKeyFn(numBuckets))) + .apply(Redistribute.>byKey().withAllowDuplicates(false)) + .apply(Values.create()); + } + + return input + .apply("Pair with shard from offset", ParDo.of(new AssignOffsetShardFn(numBuckets))) + .apply(Redistribute.>byKey().withAllowDuplicates(false)) + .apply(Values.create()); + } + + static class AssignOffsetShardFn + extends DoFn, KV>> { + private @NonNull UnsignedInteger numBuckets; + + public AssignOffsetShardFn(@Nullable Integer numBuckets) { + if (numBuckets != null && numBuckets > 0) { + this.numBuckets = UnsignedInteger.fromIntBits(numBuckets); + } else { + this.numBuckets = UnsignedInteger.valueOf(0); + } + } + + @ProcessElement + public void processElement( + @Element KafkaRecord element, + OutputReceiver>> receiver) { + int hash = Hashing.farmHashFingerprint64().hashLong(element.getOffset()).asInt(); + + if (numBuckets != null) { + hash = UnsignedInteger.fromIntBits(hash).mod(numBuckets).intValue(); + } + + receiver.output(KV.of(hash, element)); + } + } + + static class AssignRecordKeyFn + extends DoFn, KV>> { + + private @NonNull UnsignedInteger numBuckets; + + public AssignRecordKeyFn(@Nullable Integer numBuckets) { + if (numBuckets != null && numBuckets > 0) { + this.numBuckets = UnsignedInteger.fromIntBits(numBuckets); + } else { + this.numBuckets = UnsignedInteger.valueOf(0); + } + } + + @ProcessElement + public void processElement( + @Element KafkaRecord element, + OutputReceiver>> receiver) { + K key = element.getKV().getKey(); + String keyString = key == null ? "" : key.toString(); + int hash = Hashing.farmHashFingerprint64().hashBytes(keyString.getBytes(UTF_8)).asInt(); + + if (numBuckets != null) { + hash = UnsignedInteger.fromIntBits(hash).mod(numBuckets).intValue(); + } + + receiver.output(KV.of(hash, element)); + } + } +} diff --git a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOReadImplementationCompatibilityTest.java b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOReadImplementationCompatibilityTest.java index 26682946afca..dd74f07cafab 100644 --- a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOReadImplementationCompatibilityTest.java +++ b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOReadImplementationCompatibilityTest.java @@ -17,7 +17,6 @@ */ package org.apache.beam.sdk.io.kafka; -import static org.apache.beam.sdk.io.kafka.KafkaIOTest.mkKafkaReadTransform; import static org.apache.beam.sdk.io.kafka.KafkaIOTest.mkKafkaReadTransformWithOffsetDedup; import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.containsInAnyOrder; @@ -109,7 +108,7 @@ private PipelineResult testReadTransformCreationWithImplementationBoundPropertie Function, KafkaIO.Read> kafkaReadDecorator) { p.apply( kafkaReadDecorator.apply( - mkKafkaReadTransform( + KafkaIOTest.mkKafkaReadTransform( 1000, null, new ValueAsTimestampFn(), @@ -117,7 +116,8 @@ private PipelineResult testReadTransformCreationWithImplementationBoundPropertie false, /*allowDuplicates*/ 0, /*numKeys*/ null, /*offsetDeduplication*/ - null /*topics*/))); + null, /*topics*/ + null /*redistributeByRecordKey*/))); return p.run(); } diff --git a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java index 83c2e1b38826..7637b14e1d8d 100644 --- a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java +++ b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java @@ -395,7 +395,8 @@ static KafkaIO.Read mkKafkaReadTransform( false, /*allowDuplicates*/ 0, /*numKeys*/ null, /*offsetDeduplication*/ - null /*topics*/); + null, /*topics*/ + null /*redistributeByRecordKey*/); } static KafkaIO.Read mkKafkaReadTransformWithOffsetDedup( @@ -408,7 +409,24 @@ static KafkaIO.Read mkKafkaReadTransformWithOffsetDedup( false, /*allowDuplicates*/ 100, /*numKeys*/ true, /*offsetDeduplication*/ - null /*topics*/); + null, /*topics*/ + null /*redistributeByRecordKey*/); + } + + static KafkaIO.Read mkKafkaReadTransformWithRedistributeByRecordKey( + int numElements, + @Nullable SerializableFunction, Instant> timestampFn, + boolean byRecordKey) { + return mkKafkaReadTransform( + numElements, + numElements, + timestampFn, + true, /*redistribute*/ + false, /*allowDuplicates*/ + 100, /*numKeys*/ + true, /*offsetDeduplication*/ + null, /*topics*/ + byRecordKey /*redistributeByRecordKey*/); } static KafkaIO.Read mkKafkaReadTransformWithTopics( @@ -423,7 +441,8 @@ static KafkaIO.Read mkKafkaReadTransformWithTopics( false, /*allowDuplicates*/ 0, /*numKeys*/ null, /*offsetDeduplication*/ - topics /*topics*/); + topics, /*topics*/ + null /*redistributeByRecordKey*/); } /** @@ -438,7 +457,8 @@ static KafkaIO.Read mkKafkaReadTransform( @Nullable Boolean withAllowDuplicates, @Nullable Integer numKeys, @Nullable Boolean offsetDeduplication, - @Nullable List topics) { + @Nullable List topics, + @Nullable Boolean redistributeByRecordKey) { KafkaIO.Read reader = KafkaIO.read() @@ -473,7 +493,10 @@ static KafkaIO.Read mkKafkaReadTransform( reader = reader.withRedistributeNumKeys(numKeys); } if (offsetDeduplication != null && offsetDeduplication) { - reader.withOffsetDeduplication(offsetDeduplication); + reader = reader.withOffsetDeduplication(offsetDeduplication); + } + if (redistributeByRecordKey != null && redistributeByRecordKey) { + reader = reader.withRedistributeByRecordKey(redistributeByRecordKey); } } return reader; @@ -723,7 +746,8 @@ public void warningsWithAllowDuplicatesEnabledAndCommitOffsets() { true, /*allowDuplicates*/ 0, /*numKeys*/ null, /*offsetDeduplication*/ - null /*topics*/) + null, /*topics*/ + null /*redistributeByRecordKey*/) .commitOffsetsInFinalize() .withConsumerConfigUpdates( ImmutableMap.of(ConsumerConfig.GROUP_ID_CONFIG, "group_id")) @@ -751,7 +775,8 @@ public void noWarningsWithNoAllowDuplicatesAndCommitOffsets() { false, /*allowDuplicates*/ 0, /*numKeys*/ null, /*offsetDeduplication*/ - null /*topics*/) + null, /*topics*/ + null /*redistributeByRecordKey*/) .commitOffsetsInFinalize() .withConsumerConfigUpdates( ImmutableMap.of(ConsumerConfig.GROUP_ID_CONFIG, "group_id")) @@ -780,7 +805,8 @@ public void testNumKeysIgnoredWithRedistributeNotEnabled() { false, /*allowDuplicates*/ 0, /*numKeys*/ null, /*offsetDeduplication*/ - null /*topics*/) + null, /*topics*/ + null /*redistributeByRecordKey*/) .withRedistributeNumKeys(100) .commitOffsetsInFinalize() .withConsumerConfigUpdates( @@ -806,7 +832,8 @@ public void testDefaultRedistributeNumKeys() { false, /*allowDuplicates*/ null, /*numKeys*/ null, /*offsetDeduplication*/ - null /*topics*/); + null, /*topics*/ + null /*redistributeByRecordKey*/); assertFalse(read.isRedistributed()); assertEquals(0, read.getRedistributeNumKeys()); @@ -820,7 +847,8 @@ public void testDefaultRedistributeNumKeys() { false, /*allowDuplicates*/ null, /*numKeys*/ null, /*offsetDeduplication*/ - null /*topics*/); + null, /*topics*/ + null /*redistributeByRecordKey*/); assertTrue(read.isRedistributed()); // Default is defined by DEFAULT_REDISTRIBUTE_NUM_KEYS in KafkaIO. assertEquals(32768, read.getRedistributeNumKeys()); @@ -835,7 +863,8 @@ public void testDefaultRedistributeNumKeys() { false, /*allowDuplicates*/ 10, /*numKeys*/ null, /*offsetDeduplication*/ - null /*topics*/); + null, /*topics*/ + null /*redistributeByRecordKey*/); assertTrue(read.isRedistributed()); assertEquals(10, read.getRedistributeNumKeys()); } @@ -2200,7 +2229,8 @@ public void testUnboundedSourceStartReadTime() { false, /*allowDuplicates*/ 0, /*numKeys*/ null, /*offsetDeduplication*/ - null /*topics*/) + null, /*topics*/ + null /*redistributeByRecordKey*/) .withStartReadTime(new Instant(startTime)) .withoutMetadata()) .apply(Values.create()); @@ -2223,6 +2253,36 @@ public void testOffsetDeduplication() { p.run(); } + @Test + public void testRedistributeByRecordKeyOn() { + int numElements = 1000; + + PCollection input = + p.apply( + mkKafkaReadTransformWithRedistributeByRecordKey( + numElements, new ValueAsTimestampFn(), true) + .withoutMetadata()) + .apply(Values.create()); + + addCountingAsserts(input, numElements, numElements, 0, numElements - 1); + p.run(); + } + + @Test + public void testRedistributeByRecordKeyOff() { + int numElements = 1000; + + PCollection input = + p.apply( + mkKafkaReadTransformWithRedistributeByRecordKey( + numElements, new ValueAsTimestampFn(), false) + .withoutMetadata()) + .apply(Values.create()); + + addCountingAsserts(input, numElements, numElements, 0, numElements - 1); + p.run(); + } + @Rule public ExpectedException noMessagesException = ExpectedException.none(); @Test @@ -2246,7 +2306,8 @@ public void testUnboundedSourceStartReadTimeException() { false, /*allowDuplicates*/ 0, /*numKeys*/ null, /*offsetDeduplication*/ - null /*topics*/) + null, /*topics*/ + null /*redistributeByRecordKey*/) .withStartReadTime(new Instant(startTime)) .withoutMetadata()) .apply(Values.create()); diff --git a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaReadRedistributeTest.java b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaReadRedistributeTest.java new file mode 100644 index 000000000000..a14c6e3232e5 --- /dev/null +++ b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaReadRedistributeTest.java @@ -0,0 +1,231 @@ +/* + * 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.io.kafka; + +import static org.apache.beam.sdk.io.kafka.KafkaTimestampType.LOG_APPEND_TIME; +import static org.apache.beam.sdk.values.TypeDescriptors.integers; +import static org.junit.Assert.assertEquals; + +import java.io.Serializable; +import java.util.List; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.io.kafka.KafkaReadRedistribute.AssignOffsetShardFn; +import org.apache.beam.sdk.io.kafka.KafkaReadRedistribute.AssignRecordKeyFn; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.testing.ValidatesRunner; +import org.apache.beam.sdk.transforms.Count; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.GroupByKey; +import org.apache.beam.sdk.transforms.MapElements; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link KafkaReadRedistribute}. */ +@RunWith(JUnit4.class) +public class KafkaReadRedistributeTest implements Serializable { + + private static final ImmutableList> INPUTS = + ImmutableList.of( + makeKafkaRecord("k1", 3, 1), + makeKafkaRecord("k5", Integer.MAX_VALUE, 2), + makeKafkaRecord("k5", Integer.MIN_VALUE, 3), + makeKafkaRecord("k2", 66, 4), + makeKafkaRecord("k1", 4, 5), + makeKafkaRecord("k2", -33, 6), + makeKafkaRecord("k3", 0, 7)); + + private static final ImmutableList> SAME_OFFSET_INPUTS = + ImmutableList.of( + makeKafkaRecord("k1", 3, 1), + makeKafkaRecord("k5", Integer.MAX_VALUE, 1), + makeKafkaRecord("k5", Integer.MIN_VALUE, 1), + makeKafkaRecord("k2", 66, 1), + makeKafkaRecord("k1", 4, 1), + makeKafkaRecord("k2", -33, 1), + makeKafkaRecord("k3", 0, 1)); + + private static final ImmutableList> SAME_KEY_INPUTS = + ImmutableList.of( + makeKafkaRecord("k1", 3, 1), + makeKafkaRecord("k1", Integer.MAX_VALUE, 2), + makeKafkaRecord("k1", Integer.MIN_VALUE, 3), + makeKafkaRecord("k1", 66, 4), + makeKafkaRecord("k1", 4, 5), + makeKafkaRecord("k1", -33, 6), + makeKafkaRecord("k1", 0, 7)); + + static KafkaRecord makeKafkaRecord(String key, Integer value, Integer offset) { + return new KafkaRecord( + /*topic*/ "kafka", + /*partition*/ 1, + /*offset*/ offset, + /*timestamp*/ 123, + /*timestampType*/ LOG_APPEND_TIME, + /*headers*/ null, + key, + value); + } + + @Rule public final transient TestPipeline pipeline = TestPipeline.create(); + + @Test + @Category(ValidatesRunner.class) + public void testRedistributeByOffsetShard() { + + PCollection> input = + pipeline.apply( + Create.of(INPUTS) + .withCoder(KafkaRecordCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))); + + PCollection> output = + input.apply(KafkaReadRedistribute.byOffsetShard(/*numBuckets*/ 10)); + + PAssert.that(output).containsInAnyOrder(INPUTS); + + assertEquals(input.getWindowingStrategy(), output.getWindowingStrategy()); + + pipeline.run(); + } + + @Test + @Category(ValidatesRunner.class) + public void testRedistributeByKey() { + + PCollection> input = + pipeline.apply( + Create.of(INPUTS) + .withCoder(KafkaRecordCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))); + + PCollection> output = + input.apply(KafkaReadRedistribute.byRecordKey(10)); + + PAssert.that(output).containsInAnyOrder(INPUTS); + + assertEquals(input.getWindowingStrategy(), output.getWindowingStrategy()); + + pipeline.run(); + } + + @Test + @Category({ValidatesRunner.class}) + public void testAssignOutputShardFnBucketing() { + List> inputs = Lists.newArrayList(); + for (int i = 0; i < 10; i++) { + inputs.addAll(INPUTS); + } + + PCollection> input = + pipeline.apply( + Create.of(inputs) + .withCoder(KafkaRecordCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))); + + PCollection output = + input + .apply(ParDo.of(new AssignOffsetShardFn(2))) + .apply(GroupByKey.create()) + .apply(MapElements.into(integers()).via(KV::getKey)); + + PAssert.that(output).containsInAnyOrder(ImmutableList.of(0, 1)); + + pipeline.run(); + } + + @Test + @Category({ValidatesRunner.class}) + public void testAssignRecordKeyFnBucketing() { + List> inputs = Lists.newArrayList(); + for (int i = 0; i < 10; i++) { + inputs.addAll(INPUTS); + } + + PCollection> input = + pipeline.apply( + Create.of(inputs) + .withCoder(KafkaRecordCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))); + + PCollection output = + input + .apply(ParDo.of(new AssignRecordKeyFn(2))) + .apply(GroupByKey.create()) + .apply(MapElements.into(integers()).via(KV::getKey)); + + PAssert.that(output).containsInAnyOrder(ImmutableList.of(0, 1)); + + pipeline.run(); + } + + @Test + @Category({ValidatesRunner.class}) + public void testAssignOutputShardFnDeterministic() { + List> inputs = Lists.newArrayList(); + for (int i = 0; i < 10; i++) { + inputs.addAll(SAME_OFFSET_INPUTS); + } + + PCollection> input = + pipeline.apply( + Create.of(inputs) + .withCoder(KafkaRecordCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))); + + PCollection output = + input + .apply(ParDo.of(new AssignOffsetShardFn(1024))) + .apply(GroupByKey.create()) + .apply(MapElements.into(integers()).via(KV::getKey)); + + PCollection count = output.apply("CountElements", Count.globally()); + PAssert.that(count).containsInAnyOrder(1L); + + pipeline.run(); + } + + @Test + @Category({ValidatesRunner.class}) + public void testAssignRecordKeyFnDeterministic() { + List> inputs = Lists.newArrayList(); + for (int i = 0; i < 10; i++) { + inputs.addAll(SAME_KEY_INPUTS); + } + + PCollection> input = + pipeline.apply( + Create.of(inputs) + .withCoder(KafkaRecordCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))); + + PCollection output = + input + .apply(ParDo.of(new AssignRecordKeyFn(1024))) + .apply(GroupByKey.create()) + .apply(MapElements.into(integers()).via(KV::getKey)); + + PCollection count = output.apply("CountElements", Count.globally()); + PAssert.that(count).containsInAnyOrder(1L); + + pipeline.run(); + } +} diff --git a/sdks/java/io/kafka/upgrade/src/main/java/org/apache/beam/sdk/io/kafka/upgrade/KafkaIOTranslation.java b/sdks/java/io/kafka/upgrade/src/main/java/org/apache/beam/sdk/io/kafka/upgrade/KafkaIOTranslation.java index 2ebdbf29e230..51d9b028bab0 100644 --- a/sdks/java/io/kafka/upgrade/src/main/java/org/apache/beam/sdk/io/kafka/upgrade/KafkaIOTranslation.java +++ b/sdks/java/io/kafka/upgrade/src/main/java/org/apache/beam/sdk/io/kafka/upgrade/KafkaIOTranslation.java @@ -102,6 +102,7 @@ static class KafkaIOReadWithMetadataTranslator implements TransformPayloadTransl .addBooleanField("allows_duplicates") .addNullableInt32Field("redistribute_num_keys") .addNullableBooleanField("offset_deduplication") + .addNullableBooleanField("redistribute_by_record_key") .addNullableLogicalTypeField("watch_topic_partition_duration", new NanosDuration()) .addByteArrayField("timestamp_policy_factory") .addNullableMapField("offset_consumer_config", FieldType.STRING, FieldType.BYTES) @@ -229,6 +230,9 @@ public Row toConfigRow(Read transform) { if (transform.getOffsetDeduplication() != null) { fieldValues.put("offset_deduplication", transform.getOffsetDeduplication()); } + if (transform.getRedistributeByRecordKey() != null) { + fieldValues.put("redistribute_by_record_key", transform.getRedistributeByRecordKey()); + } return Row.withSchema(schema).withFieldValues(fieldValues).build(); } @@ -363,6 +367,12 @@ public Row toConfigRow(Read transform) { transform = transform.withOffsetDeduplication(offsetDeduplication); } } + if (TransformUpgrader.compareVersions(updateCompatibilityBeamVersion, "2.69.0") >= 0) { + @Nullable Boolean byRecordKey = configRow.getValue("redistribute_by_record_key"); + if (byRecordKey != null) { + transform = transform.withRedistributeByRecordKey(byRecordKey); + } + } Duration maxReadTime = configRow.getValue("max_read_time"); if (maxReadTime != null) { transform = diff --git a/sdks/java/io/kafka/upgrade/src/test/java/org/apache/beam/sdk/io/kafka/upgrade/KafkaIOTranslationTest.java b/sdks/java/io/kafka/upgrade/src/test/java/org/apache/beam/sdk/io/kafka/upgrade/KafkaIOTranslationTest.java index b5848b316baf..845e89b3b659 100644 --- a/sdks/java/io/kafka/upgrade/src/test/java/org/apache/beam/sdk/io/kafka/upgrade/KafkaIOTranslationTest.java +++ b/sdks/java/io/kafka/upgrade/src/test/java/org/apache/beam/sdk/io/kafka/upgrade/KafkaIOTranslationTest.java @@ -66,6 +66,7 @@ public class KafkaIOTranslationTest { READ_TRANSFORM_SCHEMA_MAPPING.put("getStopReadTime", "stop_read_time"); READ_TRANSFORM_SCHEMA_MAPPING.put("getRedistributeNumKeys", "redistribute_num_keys"); READ_TRANSFORM_SCHEMA_MAPPING.put("getOffsetDeduplication", "offset_deduplication"); + READ_TRANSFORM_SCHEMA_MAPPING.put("getRedistributeByRecordKey", "redistribute_by_record_key"); READ_TRANSFORM_SCHEMA_MAPPING.put( "isCommitOffsetsInFinalizeEnabled", "is_commit_offset_finalize_enabled"); READ_TRANSFORM_SCHEMA_MAPPING.put("isDynamicRead", "is_dynamic_read");