Aggregation option in Kinesis Writer Python sdk (#34323)

* 👍 Add aggregation feature

* 🐛 Fix import bug

* ✅ Add test

* 👍 Add 3 argument

* 💡 Add comment

* 🎨 Format code

Author: Nakachi-S

sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/KinesisTransformRegistrar.java

@@ -64,6 +64,10 @@ private abstract static class CrossLanguageConfiguration {
     String region;
     @Nullable String serviceEndpoint;
     boolean verifyCertificate;
+    boolean aggregationEnabled;
+    Integer aggregationMaxBytes;
+    Duration aggregationMaxBufferedTime;
+    Duration aggregationShardRefreshInterval;
 
     public void setStreamName(String streamName) {
       this.streamName = streamName;
@@ -88,6 +92,23 @@ public void setServiceEndpoint(@Nullable String serviceEndpoint) {
     public void setVerifyCertificate(@Nullable Boolean verifyCertificate) {
       this.verifyCertificate = verifyCertificate == null || verifyCertificate;
     }
+
+    public void setAggregationEnabled(@Nullable Boolean aggregationEnabled) {
+      this.aggregationEnabled = aggregationEnabled != null && aggregationEnabled;
+    }
+
+    public void setAggregationMaxBytes(Long aggregationMaxBytes) {
+      this.aggregationMaxBytes = aggregationMaxBytes.intValue();
+    }
+
+    public void setAggregationMaxBufferedTime(Long aggregationMaxBufferedTime) {
+      this.aggregationMaxBufferedTime = Duration.millis(aggregationMaxBufferedTime);
+    }
+
+    public void setAggregationShardRefreshInterval(Long aggregationShardRefreshInterval) {
+      this.aggregationShardRefreshInterval =
+          Duration.standardMinutes(aggregationShardRefreshInterval);
+    }
   }
 
   public static class WriteBuilder
@@ -131,9 +152,22 @@ public PTransform<PCollection<byte[]>, KinesisIO.Write.Result> buildExternal(
                       .skipCertificateVerification(!configuration.verifyCertificate)
                       .build())
               .withPartitioner(p -> pk)
-              .withRecordAggregationDisabled()
               .withSerializer(serializer);
 
+      if (configuration.aggregationEnabled) {
+        writeTransform =
+            writeTransform.withRecordAggregation(
+                KinesisIO.RecordAggregation.builder()
+                    .maxBytes(configuration.aggregationMaxBytes)
+                    .maxBufferedTimeJitter(0.7) // 70% jitter
+                    .maxBufferedTime(configuration.aggregationMaxBufferedTime)
+                    .shardRefreshIntervalJitter(0.5) // 50% jitter
+                    .shardRefreshInterval(configuration.aggregationShardRefreshInterval)
+                    .build());
+      } else {
+        writeTransform = writeTransform.withRecordAggregationDisabled();
+      }
+
       return writeTransform;
     }
   }

sdks/python/apache_beam/io/external/xlang_kinesisio_it_test.py

@@ -95,6 +95,7 @@ def test_kinesis_write(self):
     # TODO: remove this test once
     # https://github.com/apache/beam/issues/20416 is resolved
     self.run_kinesis_write()
+    self.run_kinesis_write_with_aggregation()
     records = self.kinesis_helper.read_from_stream(self.aws_kinesis_stream)
     self.assertEqual(
         sorted(records),
@@ -118,6 +119,25 @@ def run_kinesis_write(self):
               verify_certificate=(not self.use_localstack),
               partition_key='1'))
 
+  def run_kinesis_write_with_aggregation(self):
+    with TestPipeline(options=PipelineOptions(self.pipeline_args)) as p:
+      p.not_use_test_runner_api = True
+      _ = (
+          p
+          | 'Impulse' >> beam.Impulse()
+          | 'Generate' >> beam.FlatMap(lambda x: range(NUM_RECORDS))  # pylint: disable=bad-option-value
+          | 'Map to bytes' >>
+          beam.Map(lambda x: RECORD + str(x).encode()).with_output_types(bytes)
+          | 'WriteToKinesis' >> WriteToKinesis(
+              stream_name=self.aws_kinesis_stream,
+              aws_access_key=self.aws_access_key,
+              aws_secret_key=self.aws_secret_key,
+              region=self.aws_region,
+              service_endpoint=self.aws_service_endpoint,
+              verify_certificate=(not self.use_localstack),
+              partition_key='1',
+              aggregation_enabled=True))
+
   def run_kinesis_read(self):
     records = [RECORD + str(i).encode() for i in range(NUM_RECORDS)]
 

sdks/python/apache_beam/io/kinesis.py

@@ -112,6 +112,10 @@ def default_io_expansion_service():
         ('partition_key', str),
         ('service_endpoint', Optional[str]),
         ('verify_certificate', Optional[bool]),
+        ('aggregation_enabled', Optional[bool]),
+        ('aggregation_max_bytes', Optional[int]),
+        ('aggregation_max_buffered_time', Optional[int]),
+        ('aggregation_shard_refresh_interval', Optional[int]),
     ],
 )
 
@@ -135,6 +139,10 @@ def __init__(
       verify_certificate=None,
       producer_properties=None,
       expansion_service=None,
+      aggregation_enabled=None,
+      aggregation_max_bytes=51200,
+      aggregation_max_buffered_time=100,
+      aggregation_shard_refresh_interval=2,
   ):
     """
     Initializes a write operation to Kinesis.
@@ -151,6 +159,13 @@ def __init__(
         since the AWS IOs upgraded to v2. Trying to set it will lead to an
         error. For more info, see https://github.com/apache/beam/issues/33430.
     :param expansion_service: The address (host:port) of the ExpansionService.
+    :param aggregation_enabled: Enable or disable aggregation.
+    :param aggregation_max_bytes: Maximum number of bytes to buffer before
+        sending a batch of records. Defaults to 51200.
+    :param aggregation_max_buffered_time: Maximum time(millisecond) to buffer
+        records before sending a batch of records. Defaults to 100.
+    :param aggregation_shard_refresh_interval: Interval in minutes to refresh
+        the shard map. Defaults to 2.
     """
     if producer_properties is not None:
       raise ValueError(
@@ -167,6 +182,11 @@ def __init__(
                 partition_key=partition_key,
                 service_endpoint=service_endpoint,
                 verify_certificate=verify_certificate,
+                aggregation_enabled=aggregation_enabled,
+                aggregation_max_bytes=aggregation_max_bytes,
+                aggregation_max_buffered_time=aggregation_max_buffered_time,
+                aggregation_shard_refresh_interval=
+                aggregation_shard_refresh_interval,
             )),
         expansion_service or default_io_expansion_service(),
     )