gemini-code-assist[bot] commented on code in PR #37623:
URL: https://github.com/apache/beam/pull/37623#discussion_r2818015719
##########
sdks/python/apache_beam/ml/rag/ingestion/bigquery.py:
##########
@@ -161,16 +163,16 @@ class _WriteToBigQueryVectorDatabase(beam.PTransform):
def __init__(self, config: BigQueryVectorWriterConfig):
self.config = config
- def expand(self, pcoll: beam.PCollection[Chunk]):
+ def expand(self, pcoll: beam.PCollection[EmbeddableItem]):
schema = (
self.config.schema_config.schema
if self.config.schema_config else _default_schema())
- chunk_to_dict_fn = (
- self.config.schema_config.chunk_to_dict_fn
- if self.config.schema_config else _default_chunk_to_dict_fn)
+ embeddable_to_dict_fn = (
+ self.config.schema_config.embeddable_to_dict_fn
+ if self.config.schema_config else _default_embeddable_to_dict_fn)
return (
pcoll
- | "Chunk to dict" >> beam.Map(chunk_to_dict_fn)
+ | "Chunk to dict" >> beam.Map(embeddable_to_dict_fn)
| "Chunk dict to schema'd row" >> beam.Map(
Review Comment:
For consistency with the renaming of `Chunk` to `EmbeddableItem` throughout
the codebase, please update the labels for these PTransforms.
```suggestion
| "EmbeddableItem to dict" >> beam.Map(embeddable_to_dict_fn)
| "EmbeddableItem dict to schema'd row" >> beam.Map(
```
##########
sdks/python/apache_beam/ml/rag/enrichment/bigquery_vector_search.py:
##########
@@ -235,10 +236,11 @@ def format_query(self, chunks: List[Chunk]) -> str:
SELECT
query.id,
ARRAY_AGG(
- STRUCT({"distance, " if self.include_distance else ""}
{base_columns_str})
+ STRUCT({"distance, " if self.include_distance else ""}\
+ {base_columns_str})
Review Comment:
This line break using a backslash within an f-string is a bit unconventional
and might be hard to read. For better clarity, consider preparing the fields
for the `STRUCT` in a separate variable before constructing the main query
string.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
