Batch parity fixtures: 1 INSERT + 1 SELECT for the whole 510-pair suite

Previous shape ran 102 INSERTs (one per case) plus 510 SELECTs (one per
parametrised test) — ~612 Snowflake round-trips. With xdist on 4 CPUs
the parity job took ~2m24s in CI, dominated by per-query latency, not
compute.

Two batches:
  - `loaded_cases` does one multi-row INSERT for all 102 cases.
  - `parity_results` runs a single UNION-ALL'd SELECT that evaluates
    every translated segment against every case's identity row, returns
    a `(case_idx, seg_key) -> bool | None` dict. Per-test invocation is
    a dict lookup; no Snowflake call in the parametrised body.

Local bench: 31s wall time for 510 passing tests, serial. Down from
~2m24s. Drops `-n auto` from the CI workflow — the bottleneck is one
Snowflake query, not pytest concurrency, so xdist would just multiply
the mega-query setup across workers without speedup.

beep boop

Author: khvn26

.github/workflows/ci.yml

@@ -32,4 +32,4 @@ jobs:
         run: |
           umask 077
           printf '%s' "${{ secrets.SNOWFLAKE_PRIVATE_KEY }}" > /tmp/snowflake_pk.p8
-      - run: make test opts="-n auto"
+      - run: make test

tests/conftest.py

@@ -1,17 +1,27 @@
 """Pytest fixtures.
 
 Unit tests run anywhere; the parity suite needs a Snowflake account. The
-parity fixture skips automatically if `SNOWFLAKE_ACCOUNT` isn't set, so
+parity fixtures skip automatically if `SNOWFLAKE_ACCOUNT` isn't set, so
 running `pytest` with no env vars only runs the unit tests.
 
 Each parity-test session creates a per-run transient `IDENTITIES_PARITY_<uuid>`
 table so concurrent CI runs don't step on each other. Schema mirrors
 `SnowflakeDialect.SCHEMA_DDL`: 4 typed columns + a `traits` VARIANT.
 Table is dropped on teardown.
+
+The parity fixtures are batched to keep Snowflake round-trips down:
+
+  - All test-case identities go into the scratch table in a single
+    `INSERT INTO ... VALUES (...), (...), ...` statement.
+  - All test-case (segment, identity) pairs are evaluated in a single
+    `SELECT ... UNION ALL ...` mega-query, returning a `(case_idx, seg_key)
+    -> bool` dict. Per-test parametrised tests then do an in-memory dict
+    lookup. Two Snowflake round-trips for the whole suite.
 """
 
 from __future__ import annotations
 
+import copy
 import json
 import os
 import uuid
@@ -20,6 +30,9 @@
 from typing import Any
 
 import pytest
+from flag_engine.context.types import EnvironmentContext
+
+from flagsmith_sql_flag_engine import TranslateContext, translate_segment
 
 REPO_ROOT = Path(__file__).resolve().parents[1]
 ENGINE_TEST_DATA = REPO_ROOT / "engine-test-data" / "test_cases"
@@ -100,3 +113,94 @@ def parity_env_key(case_idx: int, original: str) -> str:
     cases (or cases that share an `environment.key` in the source data)
     don't see each other's identities."""
     return f"parity-{case_idx}-{original}"
+
+
+def _q(s: str) -> str:
+    return s.replace("'", "''")
+
+
+@pytest.fixture(scope="session")
+def loaded_cases(snowflake_session: Any, parity_table: str) -> Iterator[list[dict]]:
+    """Load every test case's identity into the scratch IDENTITIES table
+    in a single multi-row INSERT. Returns the list of cases (with
+    overridden `environment.key` so engine and SQL agree on what env to
+    filter by). One Snowflake round-trip for all 102 cases.
+    """
+    cases = load_test_cases()
+    if not cases:
+        pytest.skip("engine-test-data submodule not initialised")
+    overridden: list[dict] = []
+    selects: list[str] = []
+    for i, case in enumerate(cases):
+        case = copy.deepcopy(case)
+        ctx = case["context"]
+        env = ctx.get("environment") or {}
+        env["key"] = parity_env_key(i, env.get("key", ""))
+        env_id = _q(env["key"])
+
+        ident = ctx.get("identity") or {}
+        identifier = _q(ident.get("identifier") or "")
+        identity_key = _q(ident.get("key") or "")
+        identity_id = i + 1
+        traits = ident.get("traits") or {}
+        if traits:
+            traits_lit = f"PARSE_JSON('{_q(json.dumps(traits))}')"
+        else:
+            traits_lit = "NULL"
+        selects.append(
+            f"SELECT '{env_id}', {identity_id}, '{identifier}', '{identity_key}', {traits_lit}"
+        )
+        overridden.append(case)
+
+    snowflake_session.sql(
+        f"INSERT INTO {parity_table} (environment_id, id, identifier, identity_key, traits) "
+        + "\nUNION ALL\n".join(selects)
+    ).collect()
+    yield overridden
+
+
+@pytest.fixture(scope="session")
+def parity_results(
+    snowflake_session: Any,
+    parity_table: str,
+    loaded_cases: list[dict],
+) -> dict[tuple[int, str], bool | None]:
+    """Run every (case, segment) pair's translated SQL in one mega-query
+    and return a `(case_idx, seg_key) -> bool | None` dict. `None` means
+    the segment uses an operator the translator can't handle (test will
+    pytest.skip on that key).
+
+    One Snowflake round-trip for all 510 pairs.
+    """
+    pairs: list[tuple[int, str, str | None, str]] = []
+    select_clauses: list[str] = []
+    for case_idx, case in enumerate(loaded_cases):
+        eval_ctx = case["context"]
+        env: EnvironmentContext = eval_ctx["environment"]
+        for seg_key, segment in (eval_ctx.get("segments") or {}).items():
+            tr_ctx = TranslateContext(environment=env)
+            sql = translate_segment(segment, tr_ctx)
+            pairs.append((case_idx, seg_key, sql, env["key"]))
+
+    for i, (_case_idx, _seg_key, sql, env_key) in enumerate(pairs):
+        if sql is None:
+            continue
+        env_lit = _q(env_key)
+        select_clauses.append(
+            f"SELECT {i} AS pair_id, "
+            f"EXISTS (SELECT 1 FROM {parity_table} i "
+            f"WHERE i.environment_id = '{env_lit}' AND ({sql})) AS m"
+        )
+
+    results: dict[tuple[int, str], bool | None] = {}
+    if select_clauses:
+        rows = snowflake_session.sql("\nUNION ALL\n".join(select_clauses)).collect()
+        for row in rows:
+            i = int(row["PAIR_ID"])
+            case_idx, seg_key, _sql, _env = pairs[i]
+            results[(case_idx, seg_key)] = bool(row["M"])
+    # Fill in untranslatable pairs as None so test parametrisation can skip.
+    for case_idx, seg_key, sql, _env in pairs:
+        if sql is None:
+            results[(case_idx, seg_key)] = None
+    return results

tests/test_engine.py

@@ -1,63 +1,22 @@
 """Parity suite: every engine-test-data case translated and run against
 Snowflake, compared to flag_engine.is_context_in_segment.
 
-Skipped if SNOWFLAKE_* env vars aren't set, so `pytest` without creds
-runs the unit suite only. CI sets the creds via secrets and runs this
-suite as a separate job.
+Skipped if SNOWFLAKE_* env vars aren't set. CI sets the creds via secrets
+and runs this suite as part of `make test`.
+
+The Snowflake round-trips happen in `conftest.py`'s session-scoped
+fixtures: one INSERT for all 102 cases, one SELECT mega-query for all
+510 pair evaluations. The per-test function below is just a dict lookup
+against the pre-computed `parity_results` and a comparison against the
+engine's in-memory result.
 """
 
 from __future__ import annotations
 
-import copy
-import json
-from collections.abc import Iterator
-from typing import Any
-
 import pytest
-from flag_engine.context.types import EnvironmentContext
 from flag_engine.segments.evaluator import is_context_in_segment
 
-from flagsmith_sql_flag_engine import TranslateContext, translate_segment
-from tests.conftest import load_test_cases, parity_env_key
-
-
-@pytest.fixture(scope="session")
-def loaded_cases(snowflake_session: Any, parity_table: str) -> Iterator[list[dict]]:
-    """Load every test case's identity into the scratch IDENTITIES table
-    with traits as a VARIANT (PARSE_JSON-encoded). Returns the list of
-    cases (with overridden `environment.key` so engine and SQL agree on
-    what env to filter by).
-    """
-    cases = load_test_cases()
-    if not cases:
-        pytest.skip("engine-test-data submodule not initialised")
-    out: list[dict] = []
-    for i, case in enumerate(cases):
-        case = copy.deepcopy(case)
-        ctx = case["context"]
-        env = ctx.get("environment") or {}
-        original_key = env.get("key", "")
-        env["key"] = parity_env_key(i, original_key)
-        env_id = env["key"].replace("'", "''")
-
-        ident = ctx.get("identity") or {}
-        identifier = (ident.get("identifier") or "").replace("'", "''")
-        identity_key = (ident.get("key") or "").replace("'", "''")
-        identity_id = i + 1
-        traits = ident.get("traits") or {}
-        if traits:
-            traits_json = json.dumps(traits).replace("'", "''")
-            traits_lit = f"PARSE_JSON('{traits_json}')"
-        else:
-            traits_lit = "NULL"
-        snowflake_session.sql(
-            f"""
-            INSERT INTO {parity_table} (environment_id, id, identifier, identity_key, traits)
-            SELECT '{env_id}', {identity_id}, '{identifier}', '{identity_key}', {traits_lit}
-            """
-        ).collect()
-        out.append(case)
-    yield out
+from tests.conftest import load_test_cases
 
 
 def _all_case_segments() -> list[tuple[int, str]]:
@@ -79,31 +38,18 @@ def _all_case_segments() -> list[tuple[int, str]]:
 def test_segment_matches_engine(
     case_idx: int,
     seg_key: str,
-    snowflake_session: Any,
     loaded_cases: list[dict],
-    parity_table: str,
+    parity_results: dict[tuple[int, str], bool | None],
 ) -> None:
     case = loaded_cases[case_idx]
     eval_ctx = case["context"]
-    env: EnvironmentContext = eval_ctx["environment"]
     segment = eval_ctx["segments"][seg_key]
 
-    engine_match = is_context_in_segment(eval_ctx, segment)
-
-    tr_ctx = TranslateContext(environment=env)
-    sql = translate_segment(segment, tr_ctx)
-    if sql is None:
+    sql_match = parity_results[(case_idx, seg_key)]
+    if sql_match is None:
         pytest.skip(f"segment uses untranslatable operator (case {case_idx} seg {seg_key})")
-    env_id_lit = env["key"].replace("'", "''")
-    rows = snowflake_session.sql(
-        f"""
-        SELECT EXISTS(
-          SELECT 1 FROM {parity_table} i
-          WHERE i.environment_id = '{env_id_lit}' AND ({sql})
-        ) AS m
-        """
-    ).collect()
-    sql_match = bool(rows[0]["M"])
+
+    engine_match = is_context_in_segment(eval_ctx, segment)
     assert engine_match == sql_match, (
         f"engine={engine_match} sql={sql_match} for case={case_idx} seg={seg_key}"
     )