Demos

Complete, runnable scripts that exercise each engine end to end. They ship with the package under scripts/demos/; run one to confirm your install, then read the matching tutorial to understand each step. The full source of each is shown inline below.

Note

A license key is required. Each demo starts a MadMatcherSession, which checks in with the MadMatcher license server on startup, so a valid license key must be configured before you run one. If you do not have a key, contact your system administrator or refer to the onboarding guide.

Sparkly (lexical blocking)

scripts/demos/sparkly/demo_sparkly_blocking.py

#!/usr/bin/env python
"""Sparkly blocking, live on the dashboard (port of sparkly's basic_example.py).

Build a Lucene index over table A, bulk-search table B against it, and report
recall against the gold matches. Both the build and the search run with crash
recovery (enable_crash_recovery=True), so work is checkpointed per hash group and
a re-run resumes where it stopped. Creating a MadMatcherSession launches the
MadMatcher dashboard; the build and search then appear on it automatically and
advance per committed group -- no dashboard= flags anywhere.

    <venv>/bin/python scripts/demo_sparkly_blocking.py
"""

import os
import sys
import tempfile
from pathlib import Path

# Run THIS repo's working-tree source on both the driver (sys.path) and the Spark
# worker processes (PYTHONPATH). Must be set before Spark starts.
_SRC = Path(__file__).resolve().parents[3] / "src"
sys.path.insert(0, str(_SRC))
os.environ["PYTHONPATH"] = str(_SRC) + os.pathsep + os.environ.get("PYTHONPATH", "")
os.environ.setdefault("PYSPARK_PYTHON", sys.executable)
os.environ.setdefault("PYSPARK_DRIVER_PYTHON", sys.executable)

import pyspark.sql.functions as F
from pyspark.sql import SparkSession

from madmatcher_pro import MadMatcherSession
from madmatcher_pro.sparkly.index.lucene_index import LuceneIndex
from madmatcher_pro.sparkly.index_config import IndexConfig
from madmatcher_pro.sparkly.search import Searcher
from madmatcher_pro.sparkly.utils import check_tables_manual

DATA = _SRC.parent / "tests" / "sparkly" / "data" / "abt_buy"
limit = 50                 # candidates returned per record
analyzers = ["3gram"]      # text -> tokens for indexing

spark = SparkSession.builder.master("local[*]").appName("Sparkly demo").getOrCreate()
spark.sparkContext.setLogLevel("WARN")

# Creating the session launches the dashboard and turns tracking on for every
# operation below.
session = MadMatcherSession.builder.getOrCreate()
print(f"\n  Dashboard: {session.url}\n")

table_a = spark.read.parquet(str(DATA / "table_a.parquet"))
table_b = spark.read.parquet(str(DATA / "table_b.parquet"))
gold = spark.read.parquet(str(DATA / "gold.parquet"))
check_tables_manual(table_a, "_id", table_b, "_id")

# Everything below is crash-recoverable: work is split into hash groups, each
# group's output is written with a commit marker, and a re-run with the same
# checkpoint_dir skips committed groups. Watch the per-group progress on the
# dashboard (e.g. Searching index 3/10 groups committed).
ckpt = tempfile.mkdtemp(prefix="mm_sparkly_ckpt_")

# Index table A on its 'name' field (watch: Building index -> Merging segments).
config = IndexConfig(id_col="_id")
config.add_field("name", analyzers)
index = LuceneIndex(tempfile.mkdtemp(prefix="mm_sparkly_"), config,
                    delete_if_exists=True)
index._index_build_chunk_size = max(1, table_a.count() // 16)
index.upsert_docs(table_a, enable_crash_recovery=True,
                  checkpoint_dir=f"{ckpt}/build")
index.init()

# Search table B against the index (watch: Searching index, per group).
searcher = Searcher(index)
candidates = searcher.search(
    table_b, index.get_full_query_spec(), id_col="_id", limit=limit,
    enable_crash_recovery=True, checkpoint_dir=f"{ckpt}/search").cache()
candidates.show()

# Explode the rolled-up results and compute recall against the gold matches.
pairs = candidates.select(
    F.explode("id1_list").alias("a_id"), F.col("id2").alias("b_id"))
true_positives = gold.intersect(pairs).count()
recall = true_positives / gold.count()
print(f"\n  true_positives: {true_positives}   recall: {recall:.4f}\n")

input(f"  Dashboard live at {session.url} -- press Enter to stop.\n")
candidates.unpersist()
session.stop()
spark.stop()

Delex (DAG blocking)

scripts/demos/delex/demo_delex_blocking.py

#!/usr/bin/env python
"""Delex blocking, live on the dashboard (port of delex's basic_example.py).

Run a small blocking program over table A / table B and report recall against the
gold matches. The program keeps each search record's BM25 top-20 on BOTH 'name'
and 'description', so the optimizer has a real choice: it first ESTIMATES the cost
of each predicate, then OPTIMIZES the plan (which predicates to index, predicate
reuse, short-circuit) -- the delex planning phase. The chosen plan and the cost /
optimize times are printed before blocking.

The blocking run is also crash-recoverable (enable_crash_recovery=True): search
records are checkpointed per hash group and a re-run resumes where it stopped.
Creating a MadMatcherSession launches the dashboard; the run then appears on it
automatically (a Planning phase, then per-group Blocking) -- no dashboard= flags.

    <venv>/bin/python scripts/demos/delex/demo_delex_blocking.py
"""

import os
import sys
import tempfile
from pathlib import Path

# Run THIS repo's working-tree source on both the driver (sys.path) and the Spark
# worker processes (PYTHONPATH). Must be set before Spark starts.
_SRC = Path(__file__).resolve().parents[3] / "src"
sys.path.insert(0, str(_SRC))
os.environ["PYTHONPATH"] = str(_SRC) + os.pathsep + os.environ.get("PYTHONPATH", "")
os.environ.setdefault("PYSPARK_PYTHON", sys.executable)
os.environ.setdefault("PYSPARK_DRIVER_PYTHON", sys.executable)

import pyspark.sql.functions as F
from pyspark import SparkConf
from pyspark.sql import SparkSession

from madmatcher_pro import MadMatcherSession
from madmatcher_pro.delex.execution.plan_executor import PlanExecutor
from madmatcher_pro.delex.lang import BlockingProgram, KeepRule
from madmatcher_pro.delex.lang.predicate import BM25TopkPredicate
from madmatcher_pro.delex.utils.checks import check_tables

DATA = _SRC.parent / "tests" / "sparkly" / "data" / "abt_buy"

conf = SparkConf().set("spark.sql.execution.arrow.pyspark.enabled", "true")
spark = SparkSession.builder.master("local[*]").config(conf=conf)\
    .appName("Delex demo").getOrCreate()
spark.sparkContext.setLogLevel("WARN")

# Creating the session launches the dashboard and turns tracking on for the run.
session = MadMatcherSession.builder.getOrCreate()
print(f"\n  Dashboard: {session.url}\n")

index_table = spark.read.parquet(str(DATA / "table_a.parquet"))
search_table = spark.read.parquet(str(DATA / "table_b.parquet"))
gold = spark.read.parquet(str(DATA / "gold.parquet"))
check_tables(index_table, "_id", search_table, "_id")

# A two-rule program: keep each search record's BM25 top-20 on 'name' AND on
# 'description'. Two indexable predicates give the optimizer a real cost-based
# choice (the candidate set is the union of both rules).
prog = BlockingProgram(
    keep_rules=[
        KeepRule([BM25TopkPredicate("name", "name", "standard", 20)]),
        KeepRule([BM25TopkPredicate("description", "description", "standard", 20)]),
    ],
    drop_rules=[],
)

executor = PlanExecutor(
    index_table=index_table,
    search_table=search_table,
    index_table_id_col="_id",
    optimize=True,        # cost-based plan optimization (predicate reuse / short-circuit)
    estimate_cost=True,   # estimate each predicate's selectivity + cost first
)

# Planning: estimate each predicate's cost, then optimize the plan, and show the
# result. (execute() below re-derives the same plan inside the resumable runner;
# on this tiny dataset that costs only a second, and it lets the dashboard show the
# Planning phase live.)
plan, cost_estimation_time, optimize_time = executor.generate_plan(prog)


def render_plan(node, depth=0):
    """Render the plan DAG top-down (the sink/root first, then its inputs)."""
    print("    " + "  " * depth + str(node))
    for child in node.iter_in():
        render_plan(child, depth + 1)


print("  Optimized plan (cost-estimated + optimized):")
render_plan(plan)
print(f"\n  cost estimation: {cost_estimation_time:.2f}s    "
      f"optimization: {optimize_time:.2f}s\n")

# Crash-recoverable blocking: search records are split into hash groups, each
# group's output written with a commit marker; a re-run with the same
# checkpoint_dir skips committed groups. Watch the Planning -> Blocking progress.
ckpt = tempfile.mkdtemp(prefix="mm_delex_ckpt_")
candidates, stats = executor.execute(
    prog, search_table_id_col="_id", projection=[],
    enable_crash_recovery=True, checkpoint_dir=ckpt)
print(f"  planning recap from the run: cost estimation "
      f"{stats.cost_estimation_time:.2f}s, optimization {stats.optimize_time:.2f}s")
candidates = candidates.persist()
candidates.show()

# Explode the rolled-up results and compute recall against the gold matches.
pairs = candidates.select(
    F.explode("id1_list").alias("a_id"), F.col("id2").alias("b_id"))
n_pairs = pairs.count()
true_positives = gold.intersect(pairs).count()
recall = true_positives / gold.count()
print(f"\n  n_pairs: {n_pairs}   true_positives: {true_positives}   "
      f"recall: {recall:.4f}\n")

input(f"  Dashboard live at {session.url} -- press Enter to stop.\n")
candidates.unpersist()
session.stop()
spark.stop()

MatchFlow (matching)

scripts/demos/matchflow/demo_matchflow_matching.py

#!/usr/bin/env python
"""MatchFlow matching, live on the dashboard (port of MatchFlow's spark-local example).

The MATCHING step on abt_buy, starting from a blocking output. MatchFlow ingests a
table of candidate pairs (`id2`, `id1_list`) -- the output of ANY blocker (sparkly,
delex, the semantic blocker, or your own) -- and never needs the blocker itself.
This demo reads a precomputed `candidates.parquet` (so it does not run sparkly);
to see how that file is produced, see scripts/demos/sparkly/ or scripts/demos/delex/.

From the candidates it featurizes the pairs, creates seeds, runs batch active
learning (a gold labeler), trains a matcher, and applies it -- reporting
precision/recall/F1.

The matching steps run with crash recovery (enable_crash_recovery=True), so work is
checkpointed per hash group and a re-run resumes where it stopped. Creating a
MadMatcherSession launches the dashboard; every step then appears on it
automatically -- Featurizing, Labeling seeds, Active learning, Applying matcher --
no dashboard= flags anywhere.

    <venv>/bin/python scripts/demos/matchflow/demo_matchflow_matching.py
"""

import os
import sys
import tempfile
import warnings
from pathlib import Path

# Run THIS repo's working-tree source on both the driver (sys.path) and the Spark
# worker processes (PYTHONPATH). Must be set before Spark starts.
_SRC = Path(__file__).resolve().parents[3] / "src"
sys.path.insert(0, str(_SRC))
os.environ["PYTHONPATH"] = str(_SRC) + os.pathsep + os.environ.get("PYTHONPATH", "")
os.environ.setdefault("PYSPARK_PYTHON", sys.executable)
os.environ.setdefault("PYSPARK_DRIVER_PYTHON", sys.executable)
warnings.filterwarnings("ignore")

import pyspark.sql.functions as F
from pyspark.sql import SparkSession
from xgboost import XGBClassifier

from madmatcher_pro import MadMatcherSession
from madmatcher_pro.matchflow import (
    GoldLabeler,
    SKLearnModel,
    apply_matcher,
    create_features,
    create_seeds,
    featurize,
    label_data,
    train_matcher,
)

DATA = _SRC.parent / "tests" / "sparkly" / "data" / "abt_buy"
COLS = ["name", "description"]

spark = SparkSession.builder.master("local[*]")\
    .config("spark.sql.execution.arrow.pyspark.enabled", "true")\
    .appName("MatchFlow demo").getOrCreate()
spark.sparkContext.setLogLevel("WARN")

# Creating the session launches the dashboard and turns tracking on for every
# step below.
session = MadMatcherSession.builder.getOrCreate()
print(f"\n  Dashboard: {session.url}\n")

table_a = spark.read.parquet(str(DATA / "table_a.parquet"))
table_b = spark.read.parquet(str(DATA / "table_b.parquet"))
gold = spark.read.parquet(str(DATA / "gold.parquet"))  # (id1=A, id2=B) matches

# Every matching step below is crash-recoverable: work is split into hash groups,
# each committed with a marker, and a re-run with the same checkpoint_dir resumes
# where it stopped. Watch the per-group progress.
ckpt = tempfile.mkdtemp(prefix="mm_mf_ckpt_")

# 1) The blocking output: a table of candidate pairs (id2 from B, id1_list from A).
# MatchFlow ingests this from ANY blocker -- here we read a precomputed
# candidates.parquet; swap in your own blocking output with the same two columns.
candidates = spark.read.parquet(str(DATA / "candidates.parquet")) \
    .select("id2", "id1_list")
# Blocking recall = the share of gold matches present in the candidates. This is
# the CEILING for the matcher's recall (the matcher can only keep pairs blocking
# surfaced); it is a property of the blocking output, separate from the matcher.
_blk_pairs = candidates.select(F.explode("id1_list").alias("id1"), F.col("id2"))
_blk_recall = gold.select("id1", "id2").intersect(_blk_pairs).count() / gold.count()
print(f"\n  blocking output: {candidates.count()} candidate rows "
      f"(columns {candidates.columns}); blocking recall {_blk_recall:.4f}\n")

# 2) Featurize the candidate pairs (watch: Featurizing, per group).
features = create_features(A=table_a, B=table_b, a_cols=COLS, b_cols=COLS)
fvs = featurize(features=features, A=table_a, B=table_b, candidates=candidates,
                output_col="feature_vectors", fill_na=0.0,
                enable_crash_recovery=True,
                checkpoint_dir=f"{ckpt}/featurize").persist()

# 3) Label seeds with a gold labeler (watch: Labeling seeds). On a large candidate
# set you would down_sample(...) the active-learning pool first; abt_buy is small.
labeler = GoldLabeler(gold=gold)
seeds = create_seeds(fvs=fvs, nseeds=20, labeler=labeler, score_column="score")
print(f"\n  seeds: {seeds.count()}\n")

# 4) Batch active learning to label more examples (watch: Active learning).
# A ~600-label budget (20 seeds + 58 x 10) gets the matcher to about
# P 0.99 / R 0.93 / F1 0.96 on abt_buy -- close to the full-supervision ceiling of
# ~0.985 with these features, on a small fraction of the labels. Active learning
# trades labeling effort for accuracy: raise max_iter for a touch more recall,
# lower it for a quicker run.
model = SKLearnModel(model=XGBClassifier, eval_metric="logloss",
                     objective="binary:logistic", max_depth=6, seed=42,
                     nan_fill=0.0)
labeled = label_data(model=model, mode="batch", labeler=labeler, fvs=fvs,
                     seeds=seeds, batch_size=10, max_iter=58)

# 5) Train the matcher and apply it to every candidate (watch: Applying matcher).
trained = train_matcher(model=model, labeled_data=labeled,
                        feature_col="feature_vectors", label_col="label")
predictions = apply_matcher(
    model=trained, df=fvs, feature_col="feature_vectors",
    prediction_col="prediction", confidence_col="confidence",
    enable_crash_recovery=True, checkpoint_dir=f"{ckpt}/apply")
predictions.show()

# Precision / recall / F1 against the gold matches.
g = gold.select("id1", "id2").withColumn("gold", F.lit(1.0))
ev = predictions.join(g, on=["id1", "id2"], how="left").fillna(0.0, subset=["gold"])
tp = ev.filter((F.col("gold") == 1.0) & (F.col("prediction") == 1.0)).count()
pred_pos = ev.filter(F.col("prediction") == 1.0).count()
gold_pos = ev.filter(F.col("gold") == 1.0).count()
precision = tp / pred_pos if pred_pos else 0.0
recall = tp / gold_pos if gold_pos else 0.0
f1 = 2 * precision * recall / (precision + recall) if (precision + recall) else 0.0
print(f"\n  matcher precision: {precision:.4f}   matcher recall: {recall:.4f}   "
      f"f1: {f1:.4f}   (matcher recall <= blocking recall above)\n")

input(f"  Dashboard live at {session.url} -- press Enter to stop.\n")
fvs.unpersist()
session.stop()
spark.stop()

scripts/demos/matchflow/demo_hybrid_labeler.py

#!/usr/bin/env python
"""HybridCosineLabeler: auto-label the confident cosine tails, defer the rest to a
gold (or human) labeler -- so the expensive oracle is consulted only where cosine
is not decisive.

It wraps any Labeler. Per candidate pair it reads the cosine of the two records'
embeddings: cosine >= hi -> auto-match, cosine <= lo -> auto-non-match, otherwise
the wrapped delegate (here a GoldLabeler) is asked. The cosine is computed for the
candidate pairs only (HybridCosineLabeler.from_embeddings -> the bounded
cosine_for_pairs), never by collecting the embedding tables.

Threshold note (this dataset): abt_buy's candidate set is ~20:1 non-match:match, so
the HIGH-cosine tail overlaps (even cosine 0.92 is only ~65% matches) -- auto-match
would inject false positives, so hi=1.0 disables it. The LOW tail is clean (cosine
<= 0.70 is 100% non-matches), so lo=0.70 auto-rejects those safely. On a cleaner /
balanced dataset you would lower hi to auto-match the top tail too.

    <venv>/bin/python scripts/demos/matchflow/demo_hybrid_labeler.py
"""

import os
import sys
import tempfile
import warnings
from pathlib import Path

# Run THIS repo's working-tree source on both the driver (sys.path) and the Spark
# worker processes (PYTHONPATH). Must be set before Spark starts.
_SRC = Path(__file__).resolve().parents[3] / "src"
sys.path.insert(0, str(_SRC))
os.environ["PYTHONPATH"] = str(_SRC) + os.pathsep + os.environ.get("PYTHONPATH", "")
os.environ.setdefault("PYSPARK_PYTHON", sys.executable)
os.environ.setdefault("PYSPARK_DRIVER_PYTHON", sys.executable)
warnings.filterwarnings("ignore")

import pyspark.sql.functions as F
from pyspark.sql import SparkSession
from pyspark.sql import types as T
from xgboost import XGBClassifier

from madmatcher_pro import MadMatcherSession
from madmatcher_pro.matchflow import (
    GoldLabeler,
    HybridCosineLabeler,
    SKLearnModel,
    apply_matcher,
    create_features,
    create_seeds,
    featurize,
    label_data,
    train_matcher,
)
from madmatcher_pro.semantic import SentenceTransformerProvider, create_embeddings


class _LocalCpuProvider(SentenceTransformerProvider):
    """macOS local-mode shim: pin the embedding model to CPU (Metal/MPS cannot
    initialize across Spark's forked Python workers)."""

    def _model_obj(self):
        if self._model is None:
            from sentence_transformers import SentenceTransformer
            self._model = SentenceTransformer(self.model_id, device="cpu")
        return self._model


DATA = _SRC.parent / "tests" / "sparkly" / "data" / "abt_buy"
COLS = ["name", "description"]
MODEL = "BAAI/bge-small-en-v1.5"
DIM = 384
HI, LO = 1.0, 0.65                      # see the threshold note in the module docstring

spark = SparkSession.builder.master("local[*]")\
    .config("spark.sql.execution.arrow.pyspark.enabled", "true")\
    .appName("Hybrid labeler demo").getOrCreate()
spark.sparkContext.setLogLevel("WARN")

session = MadMatcherSession.builder.getOrCreate()
print(f"\n  Dashboard: {session.url}\n")

table_a = spark.read.parquet(str(DATA / "table_a.parquet"))
table_b = spark.read.parquet(str(DATA / "table_b.parquet"))
gold = spark.read.parquet(str(DATA / "gold.parquet"))
# Blocking output (any blocker; see scripts/demos/sparkly|delex). MatchFlow ingests it.
candidates = spark.read.parquet(str(DATA / "candidates.parquet")).select("id2", "id1_list")

# Embed both tables (bounded per-pair cosine is computed from these, never collected).
work = Path(tempfile.mkdtemp(prefix="mm_hybrid_"))
fields = [f.name for f in table_a.schema.fields
          if f.name != "_id" and isinstance(f.dataType, T.StringType)]
provider = _LocalCpuProvider(MODEL)
vec_a = create_embeddings(table_a, fields, provider, str(work / "emb_a"))
vec_b = create_embeddings(table_b, fields, provider, str(work / "emb_b"))

# Build the hybrid labeler: a GoldLabeler is the backup for the uncertain band.
gold_labeler = GoldLabeler(gold=gold)
hybrid = HybridCosineLabeler.from_embeddings(
    gold_labeler, candidates, vec_a, vec_b, hi=HI, lo=LO, dim=DIM)

# How much of the candidate pool can the labeler decide from cosine alone, and how
# accurately? (Cross-check the auto-decisions against gold.) This is the oracle
# saving -- gold/human is consulted only on the deferred pairs.
gold_pairs = {(int(r.id1), int(r.id2)) for r in gold.select("id1", "id2").collect()}
auto_rej = [k for k, c in hybrid._pair.items() if c <= LO]
auto_mat = [k for k, c in hybrid._pair.items() if c >= HI]
deferred = len(hybrid._pair) - len(auto_rej) - len(auto_mat)
rej_errors = sum(1 for k in auto_rej if k in gold_pairs)   # auto-rejected a true match
auto_n = len(auto_rej) + len(auto_mat)
acc = 1.0 - (rej_errors / auto_n) if auto_n else 1.0
total = len(hybrid._pair)
print("\n  ===== hybrid labeler over the candidate pool =====")
print(f"   candidate pairs        : {total}")
print(f"   auto-rejected (cos<={LO}): {len(auto_rej)}  (of these, true matches: {rej_errors})")
print(f"   auto-matched  (cos>={HI}): {len(auto_mat)}")
print(f"   deferred to gold backup: {deferred}")
print(f"   -> auto-decided {100 * auto_n / total:.0f}% of pairs at {100 * acc:.2f}% accuracy; "
      f"the gold/human oracle is consulted on the other {100 * deferred / total:.0f}%\n")

# Plug it into matching exactly like a GoldLabeler -- it IS a Labeler.
features = create_features(A=table_a, B=table_b, a_cols=COLS, b_cols=COLS)
fvs = featurize(features=features, A=table_a, B=table_b, candidates=candidates,
                output_col="feature_vectors", fill_na=0.0).persist()
seeds = create_seeds(fvs=fvs, nseeds=20, labeler=hybrid, score_column="score")
model = SKLearnModel(model=XGBClassifier, eval_metric="logloss",
                     objective="binary:logistic", max_depth=6, seed=42, nan_fill=0.0)
labeled = label_data(model=model, mode="batch", labeler=hybrid, fvs=fvs,
                     seeds=seeds, batch_size=10, max_iter=58)
trained = train_matcher(model=model, labeled_data=labeled,
                        feature_col="feature_vectors", label_col="label")
predictions = apply_matcher(model=trained, df=fvs, feature_col="feature_vectors",
                            prediction_col="prediction", confidence_col="confidence")

g = gold.select("id1", "id2").withColumn("gold", F.lit(1.0))
ev = predictions.join(g, on=["id1", "id2"], how="left").fillna(0.0, subset=["gold"])
tp = ev.filter((F.col("gold") == 1.0) & (F.col("prediction") == 1.0)).count()
pp = ev.filter(F.col("prediction") == 1.0).count()
gp = ev.filter(F.col("gold") == 1.0).count()
P = tp / pp if pp else 0.0
R = tp / gp if gp else 0.0
f1 = 2 * P * R / (P + R) if (P + R) else 0.0
# The hybrid labeler is a drop-in Labeler, so the matcher matches the gold-only
# result (active learning queries the uncertain band, which the hybrid defers to
# gold; it would only auto-decide the confident tails, all correct).
print(f"  matcher precision: {P:.4f}   matcher recall: {R:.4f}   f1: {f1:.4f}   "
      f"(same as gold-only -- the hybrid never mislabels)\n")

input(f"  Dashboard live at {session.url} -- press Enter to stop.\n")
fvs.unpersist()
session.stop()
spark.stop()

Semantic (dense blocking)

scripts/demos/semantic/demo_semantic_blocking.py

#!/usr/bin/env python
"""Semantic (dense) blocking, end to end, reporting recall.

The dense-retrieval counterpart of demo_sparkly_blocking.py. Full workflow over
table A / table B:

    create_embeddings   embed each table once -> id-keyed vector tables (reused)
    SemanticIndex       build an IVF index over table A's vectors
    SemanticSearcher    retrieve each B record's nearest table-A candidates

The searcher's output matches the lexical blocker's contract
(id2, id1_list, scores, search_time), so recall is computed exactly as in the
sparkly/delex demos: explode the candidate lists into pairs and intersect with
gold.

Only the three file paths below change from run to run. The fields to embed are
read from the table schema (every string column), and the model / limit / nprobe
are fixed, so pointing this at another dataset is purely a path edit.

    <venv>/bin/python scripts/demo_semantic_blocking.py
"""

import os
import sys
import tempfile
import warnings
from pathlib import Path

# Run THIS repo's working-tree source on both the driver (sys.path) and the Spark
# worker processes (PYTHONPATH). Must be set before Spark starts.
_SRC = Path(__file__).resolve().parents[3] / "src"
sys.path.insert(0, str(_SRC))
os.environ["PYTHONPATH"] = str(_SRC) + os.pathsep + os.environ.get("PYTHONPATH", "")
os.environ.setdefault("PYSPARK_PYTHON", sys.executable)
os.environ.setdefault("PYSPARK_DRIVER_PYTHON", sys.executable)
warnings.filterwarnings("ignore")

import pyspark.sql.functions as F
from pyspark.sql import SparkSession
from pyspark.sql import types as T

from madmatcher_pro.semantic import (
    SemanticIndex,
    SemanticQuerySpec,
    SemanticSearcher,
    SentenceTransformerProvider,
    create_embeddings,
)
from madmatcher_pro.sparkly.utils import check_tables_manual


class _LocalCpuProvider(SentenceTransformerProvider):
    """macOS local-mode shim: pin the embedding model to CPU.

    create_embeddings runs the model in Spark's Python workers, which are forked;
    Metal (MPS) cannot initialize across a fork, so the default device
    auto-selection aborts the worker on a Mac. On a Linux/GPU cluster, delete this
    and pass MODEL (the id string) straight to create_embeddings.
    """

    def _model_obj(self):
        if self._model is None:
            from sentence_transformers import SentenceTransformer

            self._model = SentenceTransformer(self.model_id, device="cpu")
        return self._model

# ---- The only thing to change between runs: the three input tables. ----
DATA = _SRC.parent / "tests" / "sparkly" / "data" / "abt_buy"
TABLE_A = DATA / "table_a.parquet"   # indexed (corpus) table; gold id1 lives here
TABLE_B = DATA / "table_b.parquet"   # query (search) table;   gold id2 lives here
GOLD = DATA / "gold.parquet"         # true matches, columns (id1 in A, id2 in B)

# ---- Fixed knobs (same every run). ----
MODEL = os.environ.get("MM_DEMO_EMBED_MODEL", "BAAI/bge-small-en-v1.5")  # any sentence-transformers id (resolve_provider maps it); override via MM_DEMO_EMBED_MODEL
LIMIT = 50                           # candidates retrieved per query record
NPROBE = 32                          # IVF cells probed per query (the recall floor)

spark = SparkSession.builder.master("local[*]").appName("Semantic demo").getOrCreate()
spark.sparkContext.setLogLevel("WARN")

table_a = spark.read.parquet(str(TABLE_A))
table_b = spark.read.parquet(str(TABLE_B))
gold = spark.read.parquet(str(GOLD))
check_tables_manual(table_a, "_id", table_b, "_id")

# Embed every string column (derived from the schema, so a new dataset is just a
# path change -- no field list to edit).
fields = [f.name for f in table_a.schema.fields
          if f.name != "_id" and isinstance(f.dataType, T.StringType)]

# Scratch home for the vector tables + index (re-created each run; not a knob).
work = Path(tempfile.mkdtemp(prefix="mm_semantic_"))

# 1. Embed both tables once -> id-keyed (_id, embedding) vector tables. Done once
#    and reused: the index can be rebuilt/retuned without re-embedding.
provider = _LocalCpuProvider(MODEL)
vec_a = create_embeddings(table_a, fields, provider, str(work / "emb_a"))
vec_b = create_embeddings(table_b, fields, provider, str(work / "emb_b"))

# 2. Build the IVF index over table A's vectors.
index = SemanticIndex(str(work / "index"), id_col="_id").upsert_docs(vec_a)

# 3. Retrieve each B record's top-LIMIT nearest table-A candidates.
searcher = SemanticSearcher(index)
candidates = searcher.search(
    vec_b, SemanticQuerySpec(nprobe=NPROBE), LIMIT, id_col="_id").persist()
candidates.show()

# 4. Explode the rolled-up results and compute recall against the gold matches.
pairs = candidates.select(
    F.explode("id1_list").alias("a_id"), F.col("id2").alias("b_id"))
true_positives = gold.intersect(pairs).count()
recall = true_positives / gold.count()
print(f"\n  true_positives: {true_positives}   recall: {recall:.4f}\n")

candidates.unpersist()
spark.stop()

scripts/demos/semantic/demo_fusion_blocking.py

#!/usr/bin/env python
"""Fusion blocking: combine LEXICAL (sparkly BM25) and DENSE (semantic) candidates
with Reciprocal Rank Fusion, and show the fused recall beats either blocker alone.

Both blockers emit the same (id2, id1_list, scores) contract, so they fuse with
reciprocal_rank_fusion (rank-based, scale-free -- no need to make BM25 and cosine
comparable). Each source carries a weight; the fused list is re-ranked and
truncated. Recall is reported for lexical-only, semantic-only, and the fusion.

Creating a MadMatcherSession launches the dashboard; the sparkly build/search and
the semantic embed/build/search appear on it automatically.

    <venv>/bin/python scripts/demos/semantic/demo_fusion_blocking.py
"""

import os
import sys
import tempfile
import warnings
from pathlib import Path

# Run THIS repo's working-tree source on both the driver (sys.path) and the Spark
# worker processes (PYTHONPATH). Must be set before Spark starts.
_SRC = Path(__file__).resolve().parents[3] / "src"
sys.path.insert(0, str(_SRC))
os.environ["PYTHONPATH"] = str(_SRC) + os.pathsep + os.environ.get("PYTHONPATH", "")
os.environ.setdefault("PYSPARK_PYTHON", sys.executable)
os.environ.setdefault("PYSPARK_DRIVER_PYTHON", sys.executable)
warnings.filterwarnings("ignore")

import pyspark.sql.functions as F
from pyspark.sql import SparkSession
from pyspark.sql import types as T

from madmatcher_pro import MadMatcherSession
from madmatcher_pro.semantic import (
    CandidateSource,
    SemanticIndex,
    SemanticQuerySpec,
    SemanticSearcher,
    SentenceTransformerProvider,
    create_embeddings,
    reciprocal_rank_fusion,
)
from madmatcher_pro.sparkly.index.lucene_index import LuceneIndex
from madmatcher_pro.sparkly.index_config import IndexConfig
from madmatcher_pro.sparkly.search import Searcher


class _LocalCpuProvider(SentenceTransformerProvider):
    """macOS local-mode shim: pin the embedding model to CPU (Metal/MPS cannot
    initialize across Spark's forked Python workers). On a Linux/GPU cluster, drop
    this and pass MODEL straight to create_embeddings."""

    def _model_obj(self):
        if self._model is None:
            from sentence_transformers import SentenceTransformer
            self._model = SentenceTransformer(self.model_id, device="cpu")
        return self._model


DATA = _SRC.parent / "tests" / "sparkly" / "data" / "abt_buy"
COLS = ["name", "description"]
MODEL = os.environ.get("MM_DEMO_EMBED_MODEL", "BAAI/bge-small-en-v1.5")  # override via MM_DEMO_EMBED_MODEL
K = 20                                # candidates per query (each blocker, and the fusion)

spark = SparkSession.builder.master("local[*]")\
    .config("spark.sql.execution.arrow.pyspark.enabled", "true")\
    .appName("Fusion demo").getOrCreate()
spark.sparkContext.setLogLevel("WARN")

session = MadMatcherSession.builder.getOrCreate()
print(f"\n  Dashboard: {session.url}\n")

table_a = spark.read.parquet(str(DATA / "table_a.parquet"))
table_b = spark.read.parquet(str(DATA / "table_b.parquet"))
gold = spark.read.parquet(str(DATA / "gold.parquet"))


def recall_of(cands):
    pairs = cands.select(F.explode("id1_list").alias("a"), F.col("id2").alias("b"))
    return gold.intersect(pairs).count() / gold.count()


# 1) LEXICAL: sparkly BM25 over name + description.
config = IndexConfig(id_col="_id")
for col in COLS:
    config.add_field(col, ["3gram"])
lidx = LuceneIndex(tempfile.mkdtemp(prefix="mm_fuse_lex_"), config, delete_if_exists=True)
lidx.upsert_docs(table_a)
lexical = Searcher(lidx).search(
    table_b, lidx.get_full_query_spec(), id_col="_id", limit=K)\
    .select("id2", "id1_list", "scores").persist()
print(f"  lexical: {lexical.count()} candidate rows")

# 2) DENSE: embed both tables, build an IVF index over A, retrieve A's nearest for each B.
work = Path(tempfile.mkdtemp(prefix="mm_fuse_sem_"))
fields = [f.name for f in table_a.schema.fields
          if f.name != "_id" and isinstance(f.dataType, T.StringType)]
provider = _LocalCpuProvider(MODEL)
vec_a = create_embeddings(table_a, fields, provider, str(work / "emb_a"))
vec_b = create_embeddings(table_b, fields, provider, str(work / "emb_b"))
sidx = SemanticIndex(str(work / "index"), id_col="_id").upsert_docs(vec_a)
semantic = SemanticSearcher(sidx).search(
    vec_b, SemanticQuerySpec(nprobe=32), K, id_col="_id").persist()
print(f"  semantic: {semantic.count()} candidate rows")
# 3) FUSE the two candidate sets (equal weight) with Reciprocal Rank Fusion.
fused = reciprocal_rank_fusion(
    [CandidateSource("lexical", lexical, 0.5),
     CandidateSource("semantic", semantic, 0.5)],
    limit=K).persist()
print(f"  fused: {fused.count()} candidate rows")
r_lex, r_sem, r_fused = recall_of(lexical), recall_of(semantic), recall_of(fused)
print("\n  ===== blocking recall @ K =", K, "=====")
print(f"   lexical (BM25)     : {r_lex:.4f}")
print(f"   semantic (dense)   : {r_sem:.4f}")
print(f"   fusion (RRF 50/50) : {r_fused:.4f}   <- >= both endpoints\n")

session.flush()   # flip the last op's card (the fusion) from running -> complete before the pause
input(f"  Dashboard live at {session.url} -- press Enter to stop.\n")
for d in (lexical, semantic, fused):
    d.unpersist()
session.stop()
spark.stop()

Reliability (crash recovery + dashboard)

scripts/demos/reliability/demo_progress_dashboard.py

#!/usr/bin/env python
"""Minimal example: build a sparkly index and search it, live on the dashboard.

Creating a MadMatcherSession launches the MadMatcher progress dashboard (and
opens your browser); every operation below then appears on it automatically --
no dashboard= flags. Watch the index build go Building index -> Merging
segments, then the search stream Searching index.

    <venv>/bin/python scripts/demo_progress_dashboard.py [TABLE.parquet]

With no argument it uses the in-repo abt_buy table.
"""

import os
import sys
import tempfile
from pathlib import Path

# Run THIS repo's working-tree source on both the driver (sys.path) and the
# Spark worker processes (PYTHONPATH) -- the venv's editable install may point
# at a different checkout. Must be set before Spark starts.
_SRC = Path(__file__).resolve().parents[3] / "src"
sys.path.insert(0, str(_SRC))
os.environ["PYTHONPATH"] = str(_SRC) + os.pathsep + os.environ.get("PYTHONPATH", "")
os.environ.setdefault("PYSPARK_PYTHON", sys.executable)
os.environ.setdefault("PYSPARK_DRIVER_PYTHON", sys.executable)

import pyspark.sql.functions as F
from pyspark.sql import SparkSession

from madmatcher_pro import MadMatcherSession
from madmatcher_pro.sparkly.index.lucene_index import LuceneIndex
from madmatcher_pro.sparkly.index_config import IndexConfig
from madmatcher_pro.sparkly.search import Searcher

table = sys.argv[1] if len(sys.argv) > 1 else str(
    _SRC.parent / "tests" / "sparkly" / "data" / "abt_buy" / "table_a.parquet")

spark = SparkSession.builder.master("local[*]").getOrCreate()
spark.sparkContext.setLogLevel("WARN")

# Creating the session launches the dashboard and turns tracking on for every
# operation below -- no dashboard= flag is passed anywhere.
session = MadMatcherSession.builder.getOrCreate()
print(f"\n  Dashboard: {session.url}\n")

df = spark.read.parquet(table).withColumn("_id", F.col("_id").cast("long"))
text_fields = [c for c, t in df.dtypes if t == "string" and c != "_id"]

# Build a sparkly index over the text fields (watch: Building index -> Merging
# segments). force_distributed + a small chunk so the build runs on Spark (and
# shows its phases) even for a small table.
config = IndexConfig()
for field in text_fields:
    config.add_field(field, ["standard"])
config.id_col = "_id"
index = LuceneIndex(tempfile.mkdtemp(prefix="mm_demo_"), config,
                    delete_if_exists=True)
index._index_build_chunk_size = max(1, df.count() // 16)
index.upsert_docs(df, force_distributed=True)

# Search the table against the index (watch: Searching index).
searcher = Searcher(index)
results = searcher.search(df, searcher.get_full_query_spec(), limit=20)
print(f"\n  candidate rows: {results.count()}\n")

input(f"  Dashboard live at {session.url} -- press Enter to stop.\n")
session.stop()
spark.stop()

scripts/demos/reliability/demo_resumable_blocking.py

#!/usr/bin/env python
"""Demo: crash-recoverable bulk blocking search on real sparkly data.

Builds a sparkly Lucene index over ``dblp_acm_dirty/table_a`` and blocks
``table_b`` against it with Searcher.search(checkpoint_dir=...), under a deliberately
**constrained-memory** Spark session. A fault is injected so the run dies after
the first few groups; the script then resumes from the checkpoint and finishes,
showing the per-group progress bar pick up where it left off.

What this exercises end-to-end:
* the distributed (no-driver-collect) write path -- ``spark.driver.maxResultSize``
  is set very low, so if any candidate output were routed through the driver the
  job would fail; it doesn't.
* crash recovery: groups committed before the fault are skipped on resume.
* the new ``show_progress`` group bar (analogous to the index-build bar).

Run:
    ../../madmatcher-venv/bin/python scripts/demo_resumable_blocking.py
Options:
    --n-groups N --fail-after K --limit L --rows-per-group R --clean
"""

import argparse
import os
import shutil
import sys
import time
from pathlib import Path

# Workers must use the same interpreter as the driver (see CLAUDE.md / conftest),
# and we constrain driver/executor heap. Both must be set before Spark starts.
os.environ.setdefault("PYSPARK_PYTHON", sys.executable)
os.environ.setdefault("PYSPARK_DRIVER_PYTHON", sys.executable)
_drv = os.environ.get("DEMO_DRIVER_MEM", "1g")
_exe = os.environ.get("DEMO_EXEC_MEM", "1g")
os.environ.setdefault(
    "PYSPARK_SUBMIT_ARGS",
    f"--driver-memory {_drv} --executor-memory {_exe} pyspark-shell",
)

import logging

import pyspark.sql.functions as F
from pyspark.sql import SparkSession


def read_table(spark, path):
    """Read a parquet or CSV table.

    CSV: header + RFC-4180 doubled-quote escaping, and ``multiLine=True`` so
    quoted fields containing embedded newlines (e.g. HTML blobs in this
    dataset) parse as a single record. multiLine makes the read
    non-splittable (one task parses the file, streaming -- bounded memory, a
    one-time cost); downstream stays distributed because sparkly repartitions.
    """
    p = str(path)
    if p.lower().endswith(".csv"):
        # Matches MakeGov's proven AWS reader: inferSchema so a numeric id is
        # typed as a number (not string), multiLine so embedded-newline fields
        # parse as one record, doubled-quote escaping for quoted JSON/commas.
        return (
            spark.read.option("header", True).option("inferSchema", True)
            .option("escape", '"').option("multiLine", True).csv(p)
        )
    return spark.read.parquet(p)

from madmatcher_pro.sparkly.index_config import IndexConfig
from madmatcher_pro.sparkly.index.lucene_index import LuceneIndex
from madmatcher_pro.sparkly.search import Searcher
from madmatcher_pro.reliability.resumable_search import run_resumable_search

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("demo")


class FailAfterNGroups:
    """Wraps a real Searcher and raises once it has searched ``fail_after``
    groups -- simulating a crash partway through, but not on group 0."""

    def __init__(self, inner, fail_after):
        self._inner = inner
        self._fail_after = fail_after
        self._n = 0

    def search(self, search_df, query_spec, limit, id_col="_id",
               num_records=None):
        if self._n >= self._fail_after:
            raise RuntimeError(
                f"INJECTED FAULT: simulated crash after {self._fail_after} groups"
            )
        self._n += 1
        return self._inner.search(search_df, query_spec, limit, id_col,
                                  num_records=num_records)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default=None, help="path to a sparkly dataset dir")
    ap.add_argument("--table-a", default=None,
                    help="explicit indexed-table file (parquet or csv)")
    ap.add_argument("--table-b", default=None,
                    help="explicit search-table file (omit => dedup on table-a)")
    ap.add_argument("--n-groups", type=int, default=10)
    ap.add_argument("--fail-after", type=int, default=4)
    ap.add_argument("--limit", type=int, default=20)
    ap.add_argument("--rows-per-group", type=int, default=None)
    ap.add_argument("--fields", default=None,
                    help="comma list of text fields to index (default: auto)")
    ap.add_argument("--id-col", default="_id")
    ap.add_argument("--sample", type=int, default=0,
                    help="cap rows for index+search (0 = full table)")
    ap.add_argument("--master", default="local[2]", help="Spark master")
    ap.add_argument("--search-chunk", type=int, default=500,
                    help="sparkly search_chunk_size (MakeGov AWS uses 20000)")
    ap.add_argument("--index-chunks", type=int, default=0,
                    help="partition the index BUILD into ~N chunks via sparkly's "
                         "distributed build (0=sparkly default; only engages >10)")
    ap.add_argument("--clean", action="store_true", help="wipe checkpoint first")
    args = ap.parse_args()

    repo_root = Path(__file__).resolve().parents[3]
    # Data dir precedence: --data arg > MM_DATA_DIR env var > the in-repo abt_buy
    # demo data (the same small dataset the other demos use; ships with the repo).
    # Point --data / MM_DATA_DIR at a larger dataset for a longer crash-recovery run.
    _env_data = os.environ.get("MM_DATA_DIR")
    data_dir = (
        Path(args.data) if args.data
        else Path(_env_data) if _env_data
        else repo_root / "tests" / "sparkly" / "data" / "abt_buy"
    )
    if args.table_a:
        if not Path(args.table_a).exists():
            sys.exit(f"table-a not found: {args.table_a}")
    elif not data_dir.exists():
        sys.exit(f"dataset not found: {data_dir}")

    scratch = repo_root / ".demo_scratch"
    label = Path(args.table_a).stem if args.table_a else data_dir.name
    index_path = scratch / f"{label}_index"
    ckpt_dir = scratch / "blocking_ckpt"
    if args.clean and scratch.exists():
        shutil.rmtree(scratch)
    scratch.mkdir(exist_ok=True)
    # Always start from a clean checkpoint so the demo is reproducible.
    if ckpt_dir.exists():
        shutil.rmtree(ckpt_dir)

    # maxResultSize: kept tiny by default to prove the SEARCH never collects to
    # the driver. NOTE sparkly's distributed *index build* pulls segment files
    # to the driver via toLocalIterator, so a large index needs this raised
    # (override with DEMO_MAX_RESULT) -- that is a build concern, not our search.
    max_result = os.environ.get("DEMO_MAX_RESULT", "128m")

    spark = (
        SparkSession.builder.appName("resumable-blocking-demo")
        .master(args.master)
        # --- constrained memory ---
        .config("spark.driver.maxResultSize", max_result)
        .config("spark.sql.shuffle.partitions", "8")
        .config("spark.sql.execution.arrow.maxRecordsPerBatch", "200")
        .getOrCreate()
    )
    log.info("master=%s maxResultSize=%s driverMem=%s",
             args.master, max_result, os.environ.get("DEMO_DRIVER_MEM", "1g"))
    spark.sparkContext.setLogLevel("WARN")

    # Resolve table paths: explicit files (parquet/csv) override the dataset-dir
    # convention of <dir>/table_a.parquet (+ optional table_b.parquet).
    table_a_path = args.table_a or str(data_dir / "table_a.parquet")
    if args.table_b:
        table_b_path = args.table_b
    elif args.table_a:
        # Explicit single table with no --table-b => dedup. Do NOT fall back to
        # the default dataset dir's table_b (that would mismatch the schema).
        table_b_path = None
    else:
        default_b = data_dir / "table_b.parquet"
        table_b_path = str(default_b) if default_b.exists() else None
    dedup = table_b_path is None

    log.info("loading index table: %s", table_a_path)
    table_a = read_table(spark, table_a_path)
    # sparkly stores/queries the id as a Lucene long, so coerce it.
    table_a = table_a.withColumn(args.id_col, F.col(args.id_col).cast("long"))
    if args.sample > 0:
        table_a = table_a.limit(args.sample)
    # Persist so the index build doesn't re-parse the (possibly huge) source.
    index_df = table_a.persist()

    if dedup:
        search_df = index_df
        log.info("DEDUP (no table_b): searching the indexed table against itself")
    else:
        log.info("loading search table: %s", table_b_path)
        search_df = read_table(spark, table_b_path).withColumn(
            args.id_col, F.col(args.id_col).cast("long")
        )
        if args.sample > 0:
            search_df = search_df.limit(args.sample)

    # Auto-detect text fields (string columns except the id) unless given.
    if args.fields:
        fields = [f.strip() for f in args.fields.split(",") if f.strip()]
    else:
        fields = [c for c, t in index_df.dtypes
                  if t == "string" and c != args.id_col]
    if not fields:
        sys.exit("no text fields to index; pass --fields")

    n_idx = index_df.count()
    n_search = search_df.count()
    log.info("index rows=%d, search rows=%d, fields=%s", n_idx, n_search, fields)

    # --- build the Lucene index ---
    config = IndexConfig()
    for f in fields:
        config.add_field(f, ["standard"])
    config.id_col = args.id_col

    log.info("building index at %s ...", index_path)
    index = LuceneIndex(str(index_path), config, delete_if_exists=True)
    force_dist = False
    if args.index_chunks > 0:
        import math
        # sparkly chunks the build by data size: nparts = df_size // chunk_size.
        # Setting chunk_size = n/N targets ~N build chunks. force_distributed so
        # it uses the Spark per-chunk build path (else local-thread/single).
        index._index_build_chunk_size = max(1, math.ceil(n_idx / args.index_chunks))
        force_dist = True
        approx = n_idx // index._index_build_chunk_size
        log.info("index-chunks=%d -> chunk_size=%d (~%d build chunks), "
                 "force_distributed=True", args.index_chunks,
                 index._index_build_chunk_size, approx)
        if args.index_chunks <= 10:
            log.warning("sparkly only parallelizes the build above ~10 chunks "
                        "(df_size > chunk*10); <=10 falls back to single build")
    index.upsert_docs(index_df, show_progress_bar=True, force_distributed=force_dist)
    index.init()
    log.info("index built: %d docs", index.num_indexed_docs())

    real_searcher = Searcher(index, search_chunk_size=args.search_chunk)
    query_spec = real_searcher.get_full_query_spec()

    common = dict(
        search_df=search_df, query_spec=query_spec, limit=args.limit,
        checkpoint_dir=str(ckpt_dir), id_col=args.id_col, show_progress=True,
    )
    if args.rows_per_group is not None:
        common["rows_per_group"] = args.rows_per_group
    else:
        common["n_groups"] = args.n_groups

    # ---------- RUN 1: crashes after `fail_after` groups ----------
    print("\n" + "=" * 70)
    print(f" RUN 1 — will crash after {args.fail_after} groups")
    print("=" * 70)
    t0 = time.time()
    crashed = False
    try:
        run_resumable_search(
            FailAfterNGroups(real_searcher, args.fail_after), **common
        )
    except RuntimeError as e:
        crashed = True
        print(f"\n>>> caught expected crash: {e}")
    if not crashed:
        print(">>> WARNING: run 1 did not crash (fail-after >= n_groups?)")
    log.info("run 1 wall time: %.1fs", time.time() - t0)

    # ---------- RUN 2: resume from the checkpoint ----------
    print("\n" + "=" * 70)
    print(" RUN 2 — resume with a healthy searcher")
    print("=" * 70)
    t1 = time.time()
    out = real_searcher.search(enable_crash_recovery=True, **common)
    n_search_records = out.count()
    total_pairs = out.selectExpr("explode(id1_list) as id1").count()
    log.info("run 2 wall time: %.1fs", time.time() - t1)

    print("\n" + "=" * 70)
    print(" RESULT")
    print("=" * 70)
    print(f" search records with candidates : {n_search_records}")
    print(f" total candidate pairs           : {total_pairs}")
    print(f" checkpoint dir                  : {ckpt_dir}")
    print("=" * 70)

    spark.stop()


if __name__ == "__main__":
    main()