Performance of application-defined functions and collations in SQLite

import random
import sqlite3
import string
from time import time
from typing import Callable, List, Optional, Tuple
from pprint import pprint


"""
Functions and collations to be tests.
Allows for counting number of calls with a global counter
"""

# Global variable to count number of calls in some functions
counter = 0


def casefold(a: str):
    return a.casefold()


def casefold_c(a: str):
    global counter
    counter += 1
    return a.casefold()


def get_len(a: str):
    return len(a)


def get_len_c(a: str):
    global counter
    counter += 1
    return len(a)


def do_nothing(a: str):
    return a


def do_nothing_c(a: str):
    global counter
    counter += 1
    return a


def do_nothing_collation_c(a: str, b: str):
    global counter
    counter += 1
    return -1


def do_nothing_collation(a: str, b: str):
    return -1


def comparison_collation(a: str, b: str):
    if a < b:
        return -1
    if a > b:
        return 1
    return 0


def unicode_nocase_collation_naive_c(a: str, b: str):
    global counter
    counter += 1
    if a.casefold() < b.casefold():
        return -1
    if a.casefold() > b.casefold():
        return 1
    return 0


def unicode_nocase_collation_naive(a: str, b: str):
    if a.casefold() < b.casefold():
        return -1
    if a.casefold() > b.casefold():
        return 1
    return 0


def unicode_nocase_collation_new_var(a: str, b: str):
    ac = a.casefold()
    bc = b.casefold()
    if ac < bc:
        return -1
    if ac > bc:
        return 1
    return 0


def unicode_nocase_collation_new_var_decomposition(a: str, b: str):
    ac, bc = a.casefold(), b.casefold()
    if ac < bc:
        return -1
    if ac > bc:
        return 1
    return 0


def unicode_nocase_collation_new_var_walrus(a: str, b: str):
    if (ac := a.casefold()) < (bc := b.casefold()):
        return -1
    if ac > bc:
        return 1
    return 0


def unicode_nocase_collation_var_over(a: str, b: str):
    a = a.casefold()
    b = b.casefold()
    if a < b:
        return -1
    if a > b:
        return 1
    return 0


def unicode_nocase_collation_var_over_walrus(a: str, b: str):
    if (a := a.casefold()) < (b := b.casefold()):
        return -1
    if a > b:
        return 1
    return 0


def unicode_nocase_collation_var_over_decomposition(a: str, b: str):
    a, b = a.casefold(), b.casefold()
    if a < b:
        return -1
    if a > b:
        return 1
    return 0


"""Helpers and test runner"""


def random_string(length: int = 16):
    """
    Generate random ASCII characters (both cases)
    and digits string of required length
    """
    characters = string.ascii_letters + string.digits
    return "".join(random.choice(characters) for _ in range(length))


def set_up_db(rows) -> sqlite3.Connection:
    # DB setup
    connection = sqlite3.connect(":memory:")
    cursor = connection.cursor()
    cursor.execute("CREATE TABLE test (id INTEGER, text VARCHAR);")

    # Insert rows
    cursor.executemany(
        "INSERT INTO test (id, text) VALUES (?, ?)",
        rows
    )
    connection.commit()
    return connection


def run_query(cursor: sqlite3.Cursor, query: str, rows=None) -> float:
    """Run query and return duration."""
    counter = 0
    start_time = time()
    if rows is None:
        cursor.execute(query)
        cursor.fetchall()
    else:
        for row in rows:
            cursor.execute(query, [row[1]])
            cursor.fetchall()
    end_time = time()
    return end_time - start_time


def test_performance(
    rows: List[Tuple[int, str]],
    mode: str = "",
    callable: Optional[Callable] = None,
    name: str = "",
    printout: bool = False,
):
    # Set up DB
    connection = set_up_db(rows)
    cursor = connection.cursor()
    print_counter = False

    timing = {
        "select_no_index": 0.0,
        "order_no_index": 0.0,
        "build_index": 0.0,
        "select_with_index": 0.0,
        "order_with_index": 0.0,
    }

    if callable is not None:
        name = callable.__name__.upper()

    if name.endswith("_C"):  # Convention for testing number of runs
        print_counter = True

    if mode == "collation":
        select_query = f"SELECT id, text FROM test WHERE text = ? COLLATE {name};"
        order_query = f"SELECT id, text FROM test ORDER BY text COLLATE {name} ASC "
        index_query = f"CREATE INDEX idx1 ON test (text COLLATE {name});"
        if callable is not None:
            connection.create_collation(name, callable)

    else:  # mode == 'function' and catch all
        select_query = f"SELECT id, text FROM test WHERE {name}(text) = ?;"
        order_query = f"SELECT id, text FROM test ORDER BY {name}(text) ASC"
        index_query = f"CREATE INDEX idx1 ON test ({name}(text));"
        if callable is not None:
            connection.create_function(name, 1, callable, deterministic=True)

    if printout:
        print(f"Testing {mode} {name}")

    # Run SELECT query without index
    timing["select_no_index"] = run_query(cursor, select_query, rows)
    if printout:
        print(f"{timing['select_no_index']:.3f} sec - search without index")
        if print_counter:
            print(f"calls = {counter:,}, {counter/len(rows):,.1f} per query")

    # Run ORDER BY query without an index
    timing["order_no_index"] = run_query(cursor, order_query)
    if printout:
        print(
            f"{timing['order_no_index']:.3f} sec - select all items with ORDER BY without an index"
        )
        if print_counter:
            print(f"calls = {counter:,}, {counter/len(rows):,.1f} per item")

    # Build index
    timing["build_index"] = run_query(cursor, index_query)
    if printout:
        print(f"{timing['build_index']:.3f} sec - built index")
        if print_counter:
            print(f"calls = {counter:,}, {counter/len(rows):,.1f} per item")

    # Run query with an index
    timing["select_with_index"] = run_query(cursor, select_query, rows)
    if printout:
        print(f"{timing['select_with_index']:.3f} sec - search with index")
        if print_counter:
            print(f"calls = {counter:,}, {counter/len(rows):,.1f} per query")

    # Run ORDER BY query with an index
    timing["order_with_index"] = run_query(cursor, order_query)
    if printout:
        print(
            (
                f"{timing['order_with_index']:.3f} sec - select all items with ORDER BY with an index"
            )
        )
        if print_counter:
            print(f"calls = {counter:,}, {counter/len(rows):,.1f} per item")

        print()

    connection.close()

    return timing, name


"""Running tests"""


def main():
    ROWS_NUMBER = 10000
    PASSES_NUMBER = 10
    printout = False

    tests = [
        ## Functions
        # Establish baseline performance
        ("function", None, ""),
        ("function", do_nothing),
        ("function", None, "LENGTH"),
        ("function", get_len),
        # Compare lowering functions
        ("function", None, "LOWER"),
        ("function", casefold),
        ("function", casefold_c),
        ## Collations
        # Establish baseline performance
        ("collation", do_nothing_collation),
        # Simplest collation, case-sensitive comparison
        ("collation", comparison_collation),
        # Compare built-in collation with naive implementation
        ("collation", None, "BINARY"),
        ("collation", None, "NOCASE"),
        ("collation", unicode_nocase_collation_naive),
        ("collation", unicode_nocase_collation_naive_c),
        # Compare implementations using a new variable for casefolded strings
        ("collation", unicode_nocase_collation_new_var),
        ("collation", unicode_nocase_collation_new_var_walrus),
        ("collation", unicode_nocase_collation_new_var_decomposition),
        # Compare implementations overwriting passed variables with casefolded strings
        ("collation", unicode_nocase_collation_var_over),
        ("collation", unicode_nocase_collation_var_over_walrus),
        ("collation", unicode_nocase_collation_var_over_decomposition),
    ]

    rows = [(i, random_string()) for i in range(ROWS_NUMBER)]

    for test in tests:
        results = {
            "select_no_index": 0.0,
            "order_no_index": 0.0,
            "build_index": 0.0,
            "select_with_index": 0.0,
            "order_with_index": 0.0,
        }
        for _ in range(PASSES_NUMBER):
            timing, name = test_performance(rows, *test, printout=printout)
            results = {k: results[k] + timing[k] for k in results.keys()}

        print(name)
        pprint({k: round(results[k] / PASSES_NUMBER, 3) for k in results.keys()})
        print()


if __name__ == "__main__":
    main()