Fix non-sequential _id gaps on mixed upsert batches

datastore/infrastructure/engines/bigquery/lib.py

@@ -291,7 +291,11 @@ def merge_sql(table_ref: str, schema: dict, *, include_updated_at: bool = True)
 
     Matched rows update only when a non-PK column differs (so
     `_updated_at` advances only on real changes). Unmatched rows
-    insert with `_id` continuing from `MAX(_id) + _rn`.
+    insert with `_id` continuing from `MAX(_id) + _rn`, where `_rn`
+    is numbered only among the unmatched rows (via a self-join
+    against the target on the primary key, partitioned on match
+    status) so already-present rows in the batch don't consume `_id`
+    values and leave gaps in the sequence for the rows that do insert.
     """
     fields = _user_fields(schema)
     pk = normalize_pk(schema)
@@ -306,10 +310,14 @@ def merge_sql(table_ref: str, schema: dict, *, include_updated_at: bool = True)
     insert_cols = ", ".join(f"`{f['name']}`" for f in fields)
     insert_vals = ", ".join(f"S.`{f['name']}`" for f in fields)
 
+    pk_probe_on = " AND ".join(f"S0.`{n}` = X.`{n}`" for n in pk)
     parts = [
         f"MERGE {table_ref} T",
-        f"USING (SELECT {using_cols}, ROW_NUMBER() OVER () AS _rn "
-        f"FROM UNNEST(JSON_QUERY_ARRAY(@rows)) AS r) S",
+        f"USING (SELECT S0.*, "
+        f"ROW_NUMBER() OVER (PARTITION BY X.`{pk[0]}` IS NULL) AS _rn "
+        f"FROM (SELECT {using_cols} "
+        f"FROM UNNEST(JSON_QUERY_ARRAY(@rows)) AS r) AS S0 "
+        f"LEFT JOIN {table_ref} AS X ON {pk_probe_on}) S",
         f"ON {on_clause}",
     ]
     if non_pk:

tests/engines/bigquery/test_tables.py

@@ -518,7 +518,13 @@ def test_merge_sql_renders_typed_extractors_on_match_update_no_match_insert() ->
     assert "JSON_VALUE(r, '$.label') AS `label`" in sql
     assert "PARSE_JSON(JSON_QUERY(r, '$.meta')) AS `meta`" in sql
     # USING attaches ROW_NUMBER() as _rn for auto-`_id` on NOT MATCHED.
-    assert "ROW_NUMBER() OVER () AS _rn" in sql
+    # It's numbered per match-status partition (via a self-join probe
+    # against the target on the PK) so unmatched rows are numbered
+    # 1..k among themselves, not by position in the whole batch — a
+    # batch of unmatched rows preceded by matched ones doesn't leave
+    # gaps in the `_id` sequence.
+    assert "LEFT JOIN `p.d.r` AS X ON S0.`id` = X.`id`" in sql
+    assert "ROW_NUMBER() OVER (PARTITION BY X.`id` IS NULL) AS _rn" in sql
     assert "ON T.`id` = S.`id`" in sql
     # WHEN MATCHED only fires when some non-PK column actually differs
     # — `_updated_at` advances on real changes, not on no-op upserts.
@@ -558,6 +564,37 @@ def test_update_sql_renders_dml_update_keyed_on_primary_key() -> None:
     assert "MERGE" not in sql  # plain DML, not MERGE
 
 
+def test_merge_sql_numbers_unmatched_rows_within_their_own_partition() -> None:
+    """Regression test: `_rn` must be independent per match-status
+    partition so `_id = MAX(_id) + _rn` stays gap-free for unmatched
+    rows even when a batch mixes them with already-present rows.
+
+    Before this fix, `_rn` was a single `ROW_NUMBER() OVER ()` over
+    the whole batch, so unmatched rows inherited their position among
+    *all* incoming rows (matched no-ops included) rather than their
+    position among just the unmatched ones — leaving permanent gaps
+    in `_id` whenever a batch mixed unchanged rows with new ones.
+    """
+    sql = merge_sql(
+        "`p.d.r`",
+        {
+            "fields": [
+                {"name": "id", "type": "integer"},
+                {"name": "label", "type": "string"},
+            ],
+            "primaryKey": ["id"],
+        },
+    )
+    # _rn comes from a window partitioned on whether the row's PK was
+    # found in a self-join against the target table — not a bare
+    # ROW_NUMBER() OVER () spanning the whole incoming batch.
+    assert "ROW_NUMBER() OVER () AS _rn" not in sql
+    assert "PARTITION BY X.`id` IS NULL) AS _rn" in sql
+    assert "LEFT JOIN `p.d.r` AS X ON S0.`id` = X.`id`" in sql
+    # `_id` assignment for NOT MATCHED rows is unchanged: MAX + _rn.
+    assert "(SELECT IFNULL(MAX(`_id`), 0) FROM `p.d.r`) + S._rn" in sql
+
+
 def test_merge_and_update_sql_reject_missing_primary_key() -> None:
     schema = {"fields": [{"name": "id", "type": "integer"}]}
     with pytest.raises(ValueError, match="primaryKey"):