(beam) branch master updated: Fix BigQueryEnrichmentHandler batch handling for duplicate keys (#38040)

[tv...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

tvalentyn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new a03784efc93 Fix BigQueryEnrichmentHandler batch handling for duplicate keys (#38040)
a03784efc93 is described below

commit a03784efc939855f687f110fc8b99f0c646c40c8
Author: Prabhnoor Singh <pr...@gmail.com>
AuthorDate: Fri Apr 3 13:17:04 2026 -0400

    Fix BigQueryEnrichmentHandler batch handling for duplicate keys (#38040)
    
    * Fix BigQuery enrichment batch handling for duplicate keys
    
    * Apply yapf formatting for BigQuery enrichment changes
---
 .../transforms/enrichment_handlers/bigquery.py     | 19 +++++----
 .../enrichment_handlers/bigquery_test.py           | 47 ++++++++++++++++++++++
 2 files changed, 59 insertions(+), 7 deletions(-)

diff --git a/sdks/python/apache_beam/transforms/enrichment_handlers/bigquery.py b/sdks/python/apache_beam/transforms/enrichment_handlers/bigquery.py
index 115c5320767..2306d7d97a5 100644
--- a/sdks/python/apache_beam/transforms/enrichment_handlers/bigquery.py
+++ b/sdks/python/apache_beam/transforms/enrichment_handlers/bigquery.py
@@ -15,6 +15,7 @@
 # limitations under the License.
 #
 import logging
+from collections import defaultdict
 from collections.abc import Callable
 from collections.abc import Mapping
 from typing import Any
@@ -189,7 +190,7 @@ class BigQueryEnrichmentHandler(EnrichmentSourceHandler[Union[Row, list[Row]],
     if isinstance(request, list):
       values = []
       responses = []
-      requests_map: dict[Any, Any] = {}
+      requests_map: dict[Any, list[beam.Row]] = defaultdict(list)
       batch_size = len(request)
       raw_query = self.query_template
       if batch_size > 1:
@@ -208,25 +209,29 @@ class BigQueryEnrichmentHandler(EnrichmentSourceHandler[Union[Row, list[Row]],
               "Make sure the values passed in `fields` are the "
               "keys in the input `beam.Row`." + str(e))
         values.extend(current_values)
-        requests_map[self.create_row_key(req)] = req
+        requests_map[self.create_row_key(req)].append(req)
       query = raw_query.format(*values)
 
       responses_dict = self._execute_query(query)
-      unmatched_requests = requests_map.copy()
+      unmatched_requests = {
+          key: list(reqs)
+          for key, reqs in requests_map.items()
+      }
       if responses_dict:
         for response in responses_dict:
           response_row = beam.Row(**response)
           response_key = self.create_row_key(response_row)
           if response_key in unmatched_requests:
-            req = unmatched_requests.pop(response_key)
-            responses.append((req, response_row))
+            for req in unmatched_requests.pop(response_key):
+              responses.append((req, response_row))
       if unmatched_requests:
         if self.throw_exception_on_empty_results:
           raise ValueError(f"no matching row found for query: {query}")
         else:
           _LOGGER.warning('no matching row found for query: %s', query)
-          for req in unmatched_requests.values():
-            responses.append((req, beam.Row()))
+          for reqs in unmatched_requests.values():
+            for req in reqs:
+              responses.append((req, beam.Row()))
       return responses
     else:
       request_dict = request._asdict()
diff --git a/sdks/python/apache_beam/transforms/enrichment_handlers/bigquery_test.py b/sdks/python/apache_beam/transforms/enrichment_handlers/bigquery_test.py
index 98ac6244910..98508baf661 100644
--- a/sdks/python/apache_beam/transforms/enrichment_handlers/bigquery_test.py
+++ b/sdks/python/apache_beam/transforms/enrichment_handlers/bigquery_test.py
@@ -15,9 +15,12 @@
 # limitations under the License.
 #
 import unittest
+from unittest import mock
 
 from parameterized import parameterized
 
+import apache_beam as beam
+
 # pylint: disable=ungrouped-imports
 try:
   from apache_beam.transforms.enrichment_handlers.bigquery import BigQueryEnrichmentHandler
@@ -65,6 +68,50 @@ class TestBigQueryEnrichment(unittest.TestCase):
           max_batch_size=max_batch_size,
       )
 
+  def test_batch_mode_fans_out_response_for_duplicate_keys(self):
+    handler = BigQueryEnrichmentHandler(
+        project=self.project,
+        table_name='project.dataset.table',
+        row_restriction_template="id='{}'",
+        fields=['id'],
+        min_batch_size=2,
+        max_batch_size=2,
+    )
+    requests = [beam.Row(id='1', name='first'), beam.Row(id='1', name='second')]
+
+    with mock.patch.object(handler,
+                           '_execute_query',
+                           return_value=[{'id': '1', 'value': 'enriched'}]):
+      responses = handler(requests)
+
+    self.assertEqual(
+        responses,
+        [
+            (requests[0], beam.Row(id='1', value='enriched')),
+            (requests[1], beam.Row(id='1', value='enriched')),
+        ],
+    )
+
+  def test_batch_mode_emits_empty_rows_for_all_unmatched_duplicate_keys(self):
+    handler = BigQueryEnrichmentHandler(
+        project=self.project,
+        table_name='project.dataset.table',
+        row_restriction_template="id='{}'",
+        fields=['id'],
+        min_batch_size=2,
+        max_batch_size=2,
+        throw_exception_on_empty_results=False,
+    )
+    requests = [beam.Row(id='1', name='first'), beam.Row(id='1', name='second')]
+
+    with mock.patch.object(handler, '_execute_query', return_value=None):
+      responses = handler(requests)
+
+    self.assertEqual(
+        responses,
+        [(requests[0], beam.Row()), (requests[1], beam.Row())],
+    )
+
 
 if __name__ == '__main__':
   unittest.main()