Add schema provider support for Kafka redistribute options (#36332)

* Add deterministic sharding unit test.

* Refactor to specific deterministic Kafka redistribute method.

* Add redistribute by key variant.

* Actually enable withRedistributeByRecordKey in KafkaIOTest.

* Add byRecordKey property to Kafka read compatibility.

* Rebase and revert method rename for debugging.

* Add schema provider for redistribute options

* Address spotless findings to simplify boolean expressions

* Revert accidental changes from merge conflict resolution

* Refactor into helper method.

Author: tomstepp

sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformConfiguration.java

@@ -160,6 +160,26 @@ public static Builder builder() {
   @Nullable
   public abstract ErrorHandling getErrorHandling();
 
+  @SchemaFieldDescription("If the Kafka read should be redistributed.")
+  @Nullable
+  public abstract Boolean getRedistributed();
+
+  @SchemaFieldDescription("If the Kafka read allows duplicates.")
+  @Nullable
+  public abstract Boolean getAllowDuplicates();
+
+  @SchemaFieldDescription("The number of keys for redistributing Kafka inputs.")
+  @Nullable
+  public abstract Integer getRedistributeNumKeys();
+
+  @SchemaFieldDescription("If the redistribute is using offset deduplication mode.")
+  @Nullable
+  public abstract Boolean getOffsetDeduplication();
+
+  @SchemaFieldDescription("If the redistribute keys by the Kafka record key.")
+  @Nullable
+  public abstract Boolean getRedistributeByRecordKey();
+
   /** Builder for the {@link KafkaReadSchemaTransformConfiguration}. */
   @AutoValue.Builder
   public abstract static class Builder {
@@ -190,6 +210,16 @@ public abstract static class Builder {
 
     public abstract Builder setErrorHandling(ErrorHandling errorHandling);
 
+    public abstract Builder setRedistributed(Boolean redistribute);
+
+    public abstract Builder setAllowDuplicates(Boolean allowDuplicates);
+
+    public abstract Builder setRedistributeNumKeys(Integer redistributeNumKeys);
+
+    public abstract Builder setOffsetDeduplication(Boolean offsetDeduplication);
+
+    public abstract Builder setRedistributeByRecordKey(Boolean redistributeByRecordKey);
+
     /** Builds a {@link KafkaReadSchemaTransformConfiguration} instance. */
     public abstract KafkaReadSchemaTransformConfiguration build();
   }

sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProvider.java

@@ -166,6 +166,31 @@ private SchemaRegistryProvider getSchemaRegistryProvider(String confluentSchemaR
       return SchemaRegistryProvider.UNSPECIFIED;
     }
 
+    private static <K, V> KafkaIO.Read<K, V> applyRedistributeSettings(
+        KafkaIO.Read<K, V> kafkaRead, KafkaReadSchemaTransformConfiguration configuration) {
+      Boolean redistribute = configuration.getRedistributed();
+      if (redistribute != null && redistribute) {
+        kafkaRead = kafkaRead.withRedistribute();
+      }
+      Integer redistributeNumKeys = configuration.getRedistributeNumKeys();
+      if (redistributeNumKeys != null && redistributeNumKeys > 0) {
+        kafkaRead = kafkaRead.withRedistributeNumKeys(redistributeNumKeys);
+      }
+      Boolean allowDuplicates = configuration.getAllowDuplicates();
+      if (allowDuplicates != null) {
+        kafkaRead = kafkaRead.withAllowDuplicates(allowDuplicates);
+      }
+      Boolean redistributeByRecordKey = configuration.getRedistributeByRecordKey();
+      if (redistributeByRecordKey != null) {
+        kafkaRead = kafkaRead.withRedistributeByRecordKey(redistributeByRecordKey);
+      }
+      Boolean offsetDeduplication = configuration.getOffsetDeduplication();
+      if (offsetDeduplication != null) {
+        kafkaRead = kafkaRead.withOffsetDeduplication(offsetDeduplication);
+      }
+      return kafkaRead;
+    }
+
     @Override
     public PCollectionRowTuple expand(PCollectionRowTuple input) {
       configuration.validate();
@@ -233,6 +258,8 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) {
           kafkaRead = kafkaRead.withMaxReadTime(Duration.standardSeconds(maxReadTimeSeconds));
         }
 
+        kafkaRead = applyRedistributeSettings(kafkaRead, configuration);
+
         PCollection<GenericRecord> kafkaValues =
             input.getPipeline().apply(kafkaRead.withoutMetadata()).apply(Values.create());
 
@@ -283,6 +310,8 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) {
         kafkaRead = kafkaRead.withMaxReadTime(Duration.standardSeconds(maxReadTimeSeconds));
       }
 
+      kafkaRead = applyRedistributeSettings(kafkaRead, configuration);
+
       PCollection<byte[]> kafkaValues =
           input.getPipeline().apply(kafkaRead.withoutMetadata()).apply(Values.create());
 

sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProviderTest.java

@@ -130,7 +130,12 @@ public void testFindTransformAndMakeItWork() {
             "error_handling",
             "file_descriptor_path",
             "message_name",
-            "max_read_time_seconds"),
+            "max_read_time_seconds",
+            "redistributed",
+            "allow_duplicates",
+            "offset_deduplication",
+            "redistribute_num_keys",
+            "redistribute_by_record_key"),
         kafkaProvider.configurationSchema().getFields().stream()
             .map(field -> field.getName())
             .collect(Collectors.toSet()));