Avoid quadratic rescans when deduplicating query column names

benchmarks/benchmark_dedupe_keys.py

@@ -0,0 +1,80 @@
+# Standalone benchmark for the duplicate-column-name optimization.
+#
+# Run with:
+#   python benchmarks/benchmark_dedupe_keys.py
+#
+# This file intentionally does not use the test_ prefix, so pytest does not
+# collect it. The baseline is frozen here to make comparisons reproducible.
+from __future__ import annotations
+
+import statistics
+import timeit
+from collections.abc import Iterable
+
+from sqlite_utils.utils import dedupe_keys
+
+
+def baseline_dedupe_keys(keys: Iterable[str]) -> list[str]:
+    """The implementation before the next-suffix cache optimization."""
+    keys = list(keys)
+    taken = set(keys)
+    if len(taken) == len(keys):
+        return keys
+    seen: set[str] = set()
+    result: list[str] = []
+    for key in keys:
+        if key in seen:
+            new_key = key
+            suffix = 2
+            while new_key in seen or new_key in taken:
+                new_key = f"{key}_{suffix}"
+                suffix += 1
+            key = new_key
+        seen.add(key)
+        result.append(key)
+    return result
+
+
+def duplicate_keys(size: int) -> list[str]:
+    """Return ``size`` identical original keys."""
+    return ["id"] * size
+
+
+def reserved_suffix_keys(size: int) -> list[str]:
+    """Return duplicate ``id`` keys plus genuinely reserved original suffixes."""
+    # The first size entries are identical original keys. The remaining size
+    # entries are distinct original suffixes id_2 through id_(size + 1), so
+    # those suffixes are reserved before duplicate resolution begins.
+    return ["id"] * size + [f"id_{suffix}" for suffix in range(2, size + 2)]
+
+
+def measure(function, keys: list[str], repeats: int = 5) -> float:
+    # Check output equality on every run, not only the timed invocation.
+    expected = baseline_dedupe_keys(keys)
+    timings = []
+    for _ in range(repeats):
+        actual = function(keys)
+        if actual != expected:
+            raise AssertionError(
+                f"Output mismatch for {function.__name__}: {actual!r} != {expected!r}"
+            )
+        timings.append(timeit.timeit(lambda: function(keys), number=1))
+    return statistics.median(timings)
+
+
+def main() -> None:
+    cases = [
+        ("duplicates-500", duplicate_keys(500)),
+        ("duplicates-2000", duplicate_keys(2000)),
+        ("duplicates-5000", duplicate_keys(5000)),
+        ("reserved-suffix-500", reserved_suffix_keys(500)),
+        ("reserved-suffix-2000", reserved_suffix_keys(2000)),
+    ]
+    for name, keys in cases:
+        baseline = measure(baseline_dedupe_keys, keys)
+        current = measure(dedupe_keys, keys)
+        print(f"{name}: baseline={baseline:.6f}s current={current:.6f}s")
+
+
+if __name__ == "__main__":
+    main()

sqlite_utils/utils.py

@@ -634,14 +634,18 @@ def dedupe_keys(keys: Iterable[str]) -> list[str]:
         # No duplicates - the common case
         return keys
     seen: set = set()
+    # Cache the next suffix to try for each original base key. This avoids
+    # repeatedly searching from suffix 2 when a key occurs many times.
+    next_suffix: dict[str, int] = {}
     result = []
     for key in keys:
         if key in seen:
-            new_key = key
-            suffix = 2
+            suffix = next_suffix.get(key, 2)
+            new_key = f"{key}_{suffix}"
             while new_key in seen or new_key in taken:
-                new_key = f"{key}_{suffix}"
                 suffix += 1
+                new_key = f"{key}_{suffix}"
+            next_suffix[key] = suffix + 1
             key = new_key
         seen.add(key)
         result.append(key)

tests/test_query.py

@@ -274,6 +274,17 @@ def test_query_deduped_column_avoids_existing_names(fresh_db):
     assert rows == [{"id": 1, "id_3": 2, "id_2": 3}]
 
 
+def test_query_dedupes_duplicate_aliased_columns(fresh_db):
+    rows = list(
+        fresh_db.query(
+            "select 1 as id, 2 as id, 3 as id_2, 4 as id_2, 5 as value, 6 as value"
+        )
+    )
+    assert rows == [
+        {"id": 1, "id_3": 2, "id_2": 3, "id_2_2": 4, "value": 5, "value_2": 6}
+    ]
+
+
 def test_execute_returning_dicts(fresh_db):
     # Like db.query() but returns a list, included for backwards compatibility
     # see https://github.com/simonw/sqlite-utils/issues/290

tests/test_utils.py

@@ -102,3 +102,40 @@ def test_flatten(input, expected):
 )
 def test_dedupe_keys(input, expected):
     assert utils.dedupe_keys(input) == expected
+
+
+def test_dedupe_keys_accepts_generator_input():
+    assert utils.dedupe_keys(key for key in ["id", "id", "id_2"]) == [
+        "id",
+        "id_3",
+        "id_2",
+    ]
+
+
+def test_dedupe_keys_handles_repeated_suffix_bases():
+    assert utils.dedupe_keys(["id", "id", "id_2", "id_2", "id_2_2", "id_2_2"]) == [
+        "id",
+        "id_3",
+        "id_2",
+        "id_2_3",
+        "id_2_2",
+        "id_2_2_2",
+    ]
+
+
+def test_dedupe_keys_scales_without_reformatting_all_suffixes():
+    class CountingStr(str):
+        calls = 0
+
+        def __format__(self, format_spec):
+            type(self).calls += 1
+            return super().__format__(format_spec)
+
+    keys = [CountingStr("id") for _ in range(2000)]
+    result = utils.dedupe_keys(keys)
+
+    assert len(result) == 2000
+    assert len(set(result)) == 2000
+    assert result[0] == "id"
+    assert result[-1] == "id_2000"
+    assert CountingStr.calls <= 2000