From 183a5dfb6e329c39b41c1b20d042fde020063263 Mon Sep 17 00:00:00 2001 From: workstonedai-collab <290570992+workstonedai-collab@users.noreply.github.com> Date: Sat, 24 Oct 2026 16:10:34 +0800 Subject: [PATCH] Avoid quadratic rescans when deduplicating query column names --- benchmarks/benchmark_dedupe_keys.py | 80 +++++++++++++++++++++++++++++ sqlite_utils/utils.py | 10 ++-- tests/test_query.py | 11 ++++ tests/test_utils.py | 37 +++++++++++++ 4 files changed, 135 insertions(+), 3 deletions(-) create mode 100644 benchmarks/benchmark_dedupe_keys.py diff --git a/benchmarks/benchmark_dedupe_keys.py b/benchmarks/benchmark_dedupe_keys.py new file mode 100644 index 000000000..3f0e68e64 --- /dev/null +++ b/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() diff --git a/sqlite_utils/utils.py b/sqlite_utils/utils.py index 91fd24ddd..77b222656 100644 --- a/sqlite_utils/utils.py +++ b/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) diff --git a/tests/test_query.py b/tests/test_query.py index ac0d9249b..0b0e50e21 100644 --- a/tests/test_query.py +++ b/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 diff --git a/tests/test_utils.py b/tests/test_utils.py index 360a4436a..3714c4e3c 100644 --- a/tests/test_utils.py +++ b/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