[GitHub] [beam] dennisylyung commented on a change in pull request #12583: [BEAM-10706] Fix duplicate key error in DynamoDBIO.Write

[GitBox <gi...@apache.org>]

dennisylyung commented on a change in pull request #12583:
URL: https://github.com/apache/beam/pull/12583#discussion_r541042726



##########
File path: sdks/java/io/amazon-web-services/src/main/java/org/apache/beam/sdk/io/aws/dynamodb/DynamoDBIO.java
##########
@@ -447,19 +467,41 @@ public void setup() {
 
       @StartBundle
       public void startBundle(StartBundleContext context) {
-        batch = new ArrayList<>();
+        batch = new HashMap<>();
       }
 
       @ProcessElement
       public void processElement(ProcessContext context) throws Exception {
         final KV<String, WriteRequest> writeRequest =
             (KV<String, WriteRequest>) spec.getWriteItemMapperFn().apply(context.element());
-        batch.add(writeRequest);
+        batch.put(
+            KV.of(writeRequest.getKey(), extractPkeyValues(writeRequest.getValue())), writeRequest);
         if (batch.size() >= BATCH_SIZE) {
           flushBatch();
         }
       }
 
+      private Map<String, AttributeValue> extractPkeyValues(WriteRequest request) {
+        if (spec.getOverwriteByPKeys() != null) {
+          if (request.getPutRequest() != null) {
+            return request.getPutRequest().getItem().entrySet().stream()
+                .filter(entry -> spec.getOverwriteByPKeys().contains(entry.getKey()))

Review comment:
       This should handle composite keys. 
   
   Supposedly, this gets the item map
   ```
   request.getPutRequest().getItem()        // Map<String, AttributeValue>
   ```
   and then retains only the relevant keys 
   ```
   ... .entrySet().stream()
       .filter(...)
       .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue))
   ```
   The filter function checks whether the key is one of the specified deduplication keys
   ```
   entry -> spec.getOverwriteByPKeys().contains(entry.getKey())
   ```
   
   Is this correct?




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org