damccorm commented on code in PR #33413:
URL: https://github.com/apache/beam/pull/33413#discussion_r1897813752
##########
sdks/python/apache_beam/ml/rag/enrichment/bigquery_vector_search.py:
##########
@@ -140,48 +184,66 @@ def format_query(self, chunks: List[Chunk]) -> str:
FROM UNNEST(metadata)
WHERE key = search_key AND value = search_value
));
+ """ if self.metadata_restriction_template else ""
+
+ # Group chunks by their metadata conditions
+ condition_groups = {}
+ if self.metadata_restriction_template:
+ for chunk in chunks:
+ condition = self.metadata_restriction_template.format(**chunk.metadata)
+ if condition not in condition_groups:
+ condition_groups[condition] = []
+ condition_groups[condition].append(chunk)
Review Comment:
Nit: this can be cleaned up by using defaultdict -
https://www.geeksforgeeks.org/defaultdict-in-python/
--
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.
To unsubscribe, e-mail: github-unsubscr...@beam.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
