Bug: a lookup join probes every row of a full node-table scan, 2.7x slower than a hash join #1394

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.10,<3.15"
# dependencies = ["ladybug==0.22.0"]
# ///
"""ladybug 0.22.0: the lookup join from #1285 is planned for a join whose outer
side is a full node-table scan, so every scanned row is probed one at a time.
Clean on 0.21.2, which has no lookup join.

Uses the tinysnb tables person, organisation and workAt, cut to the columns the
query reads, at a larger scale.

Exits 1 when the plan probes organisation per row and runs more than 1.5x slower
than the same query with a row-preserving filter that keeps the hash join.
"""
import os
import statistics
import sys
import tempfile
import time

import ladybug

N_PERSONS, N_ORGS, IDS = 10000, 4300, list(range(0, 10000, 125))

path = os.path.join(tempfile.mkdtemp(), "db.lbug")
build = ladybug.Connection(ladybug.Database(path))
build.execute("CREATE NODE TABLE person(ID INT64, fName STRING, PRIMARY KEY(ID));")
build.execute("CREATE NODE TABLE organisation(ID INT64, name STRING, history STRING, PRIMARY KEY(ID));")
build.execute("CREATE REL TABLE workAt(FROM person TO organisation, year INT64, MANY_ONE);")
build.execute("UNWIND range(0, $n - 1) AS i CREATE (:organisation {ID: i, name: 'org' + CAST(i AS STRING), "
              "history: repeat('lorem ipsum ', 160)});", {"n": N_ORGS})
build.execute("UNWIND range(0, $n - 1) AS i CREATE (:person {ID: i, fName: 'p' + CAST(i AS STRING)});",
              {"n": N_PERSONS})
build.execute("MATCH (a:person), (b:organisation) WHERE b.ID = a.ID % $n "
              "CREATE (a)-[:workAt {year: 2015}]->(b);", {"n": N_ORGS})
build.execute("CHECKPOINT;")
del build
conn = ladybug.Connection(ladybug.Database(path, read_only=True))

QUERY = ("UNWIND $ids AS i MATCH (a:person) WHERE a.ID = i MATCH (a)-[:workAt]->(b:organisation) {filter}"
         "RETURN b.name, b.history, a.ID;")
AS_PLANNED = QUERY.format(filter="")
CONTROL = QUERY.format(filter="WHERE b.ID >= 0 ")


def rows(q):
    result = conn.execute(q, {"ids": IDS})
    out = []
    while result.has_next():
        out.append(result.get_next())
    return sorted(out)


def p50(q):
    for _ in range(5):
        rows(q)
    times = []
    for _ in range(60):
        t0 = time.perf_counter()
        rows(q)
        times.append((time.perf_counter() - t0) * 1000)
    return statistics.median(times)


def probes_organisation(q):
    plan = conn.execute("EXPLAIN LOGICAL " + q, {"ids": IDS})
    text = ""
    while plan.has_next():
        text += str(plan.get_next()[0])
    return text.count("QUERY_PRIMARY_KEY_LOOKUP") >= 2  # one for a.ID = i, one more for organisation


print(f"ladybug {ladybug.__version__}: {N_PERSONS} person rows, {N_ORGS} organisation rows, {len(IDS)} ids")
same = rows(AS_PLANNED) == rows(CONTROL)
planned, control = p50(AS_PLANNED), p50(CONTROL)
lookup = probes_organisation(AS_PLANNED)
print(f"  as planned               per-row lookup: {lookup!s:5}  p50 {planned:6.2f} ms")
print(f"  control: WHERE b.ID >= 0 per-row lookup: {probes_organisation(CONTROL)!s:5}  p50 {control:6.2f} ms")
print(f"  same rows: {same}; as planned / control = {planned / control:.2f}x")

if lookup and planned > 1.5 * control:
    print("\nBUG REPRODUCED")
    sys.exit(1)
print("\nclean on this version")