@@ -122,14 +122,14 @@ def _resolve_score_alias(
122122 first_row_fields : set [str ] | None = None ,
123123 ) -> str :
124124 """Determine a stable score column name that won't collide with
125- document fields. The alias is decided once before iterating rows
126- so every row uses the same column name.
125+ document fields. The alias is resolved once and reused for every
126+ row so all rows share the same column name.
127127
128128 When a RETURN clause is present, the returned field names are used
129129 for collision detection. When RETURN is absent (SELECT *), the
130- caller should pass ``first_row_fields`` — the field names from the
131- first result row — so we can detect collisions even when all
132- document attributes are returned ."""
130+ caller should pass ``first_row_fields`` — the union of all field
131+ names across all result rows — so we can detect collisions even
132+ when different documents have different field sets ."""
133133 alias = score_alias or "__score"
134134 # Extract RETURN field names from args to detect collision
135135 try :
@@ -228,19 +228,23 @@ def execute(self, sql: str, *, params: dict | None = None) -> QueryResult:
228228 elif with_scores :
229229 # WITHSCORES format: [count, key1, score1, [fields1], key2, score2, [fields2], ...]
230230 # Stride of 3: key, score, field_list
231- # Resolve alias once from the first row so every row uses the
232- # same column name (consistent output schema).
233- resolved_alias : str | None = None
231+ # First pass: collect all field names across all rows so the
232+ # alias avoids collisions with any document field, not just
233+ # the first row's fields.
234+ all_field_names : set [str ] = set ()
235+ parsed_rows : list [tuple [dict , Any ]] = []
234236 for i in range (1 , len (raw_result ) - 2 , 3 ):
235237 score = raw_result [i + 1 ]
236238 row_data = raw_result [i + 2 ]
237239 row = dict (zip (row_data [::2 ], row_data [1 ::2 ]))
238- if resolved_alias is None :
239- resolved_alias = self ._resolve_score_alias (
240- translated .score_alias ,
241- translated .args ,
242- first_row_fields = set (row .keys ()),
243- )
240+ all_field_names .update (row .keys ())
241+ parsed_rows .append ((row , score ))
242+ resolved_alias = self ._resolve_score_alias (
243+ translated .score_alias ,
244+ translated .args ,
245+ first_row_fields = all_field_names ,
246+ )
247+ for row , score in parsed_rows :
244248 row [resolved_alias ] = score
245249 rows .append (row )
246250 else :
@@ -345,19 +349,22 @@ async def execute(self, sql: str, *, params: dict | None = None) -> QueryResult:
345349 rows .append (row )
346350 elif with_scores :
347351 # WITHSCORES format: [count, key1, score1, [fields1], ...]
348- # Resolve alias once from the first row so every row uses the
349- # same column name (consistent output schema).
350- resolved_alias : str | None = None
352+ # First pass: collect all field names across all rows so the
353+ # alias avoids collisions with any document field.
354+ all_field_names : set [str ] = set ()
355+ parsed_rows : list [tuple [dict , Any ]] = []
351356 for i in range (1 , len (raw_result ) - 2 , 3 ):
352357 score = raw_result [i + 1 ]
353358 row_data = raw_result [i + 2 ]
354359 row = dict (zip (row_data [::2 ], row_data [1 ::2 ]))
355- if resolved_alias is None :
356- resolved_alias = self ._resolve_score_alias (
357- translated .score_alias ,
358- translated .args ,
359- first_row_fields = set (row .keys ()),
360- )
360+ all_field_names .update (row .keys ())
361+ parsed_rows .append ((row , score ))
362+ resolved_alias = self ._resolve_score_alias (
363+ translated .score_alias ,
364+ translated .args ,
365+ first_row_fields = all_field_names ,
366+ )
367+ for row , score in parsed_rows :
361368 row [resolved_alias ] = score
362369 rows .append (row )
363370 else :
0 commit comments