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")
try:
input(f" Dashboard live at {session.url} -- press Enter to stop.\n")
except EOFError:
pass # non-interactive run (piped / CI): don't fail at the prompt
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)
# Walk the plan DAG top-down (the sink/root first, then its inputs).
print(" Optimized plan (cost-estimated + optimized):")
to_visit = [(plan, 0)]
while to_visit:
node, depth = to_visit.pop()
print(" " + " " * depth + str(node))
to_visit.extend((child, depth + 1) for child in reversed(list(node.iter_in())))
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")
try:
input(f" Dashboard live at {session.url} -- press Enter to stop.\n")
except EOFError:
pass # non-interactive run (piped / CI): don't fail at the prompt
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")
try:
input(f" Dashboard live at {session.url} -- press Enter to stop.\n")
except EOFError:
pass # non-interactive run (piped / CI): don't fail at the prompt
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
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)]
# device="cpu": create_embeddings runs the model in Spark's Python workers,
# which are forked, and on macOS MPS cannot initialize across a fork. Drop it
# on a Linux/GPU cluster.
provider = SentenceTransformerProvider(MODEL, device="cpu")
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")
try:
input(f" Dashboard live at {session.url} -- press Enter to stop.\n")
except EOFError:
pass # non-interactive run (piped / CI): don't fail at the prompt
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
# ---- 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.
# device="cpu": create_embeddings runs the model in Spark's Python workers,
# which are forked, and on macOS MPS cannot initialize across a fork. Drop it
# on a Linux/GPU cluster.
provider = SentenceTransformerProvider(MODEL, device="cpu")
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
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"))
# 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)]
# device="cpu": create_embeddings runs the model in Spark's Python workers,
# which are forked, and on macOS MPS cannot initialize across a fork. Drop it
# on a Linux/GPU cluster.
provider = SentenceTransformerProvider(MODEL, device="cpu")
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")
# Recall of a candidate table = the share of gold pairs it contains.
n_gold = gold.count()
r_lex, r_sem, r_fused = [
gold.intersect(c.select(F.explode("id1_list").alias("a"),
F.col("id2").alias("b"))).count() / n_gold
for c in (lexical, semantic, 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
try:
input(f" Dashboard live at {session.url} -- press Enter to stop.\n")
except EOFError:
pass # non-interactive run (piped / CI): don't fail at the prompt
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")
try:
input(f" Dashboard live at {session.url} -- press Enter to stop.\n")
except EOFError:
pass # non-interactive run (piped / CI): don't fail at the prompt
session.stop()
spark.stop()
scripts/demos/reliability/demo_resumable_blocking.py¶
#!/usr/bin/env python
"""Crash-recoverable bulk blocking: crash partway through, then resume.
Blocks table B against a sparkly index under a deliberately constrained-memory Spark
session, injecting a fault so the first run dies after a few hash groups. The second
run resumes from the checkpoint and finishes, skipping every group already committed.
What this exercises:
* the distributed (no-driver-collect) write path -- `spark.driver.maxResultSize` is
tiny, so a search that routed candidates through the driver would fail; it does not
* crash recovery: groups committed before the fault are skipped on resume
* the per-group progress bar picking up where it left off
<venv>/bin/python scripts/demos/reliability/demo_resumable_blocking.py
Env knobs:
MM_DATA_DIR dataset dir (default: the in-repo abt_buy demo data)
MM_N_GROUPS hash groups to split the search into (default 10)
MM_FAIL_AFTER crash after this many groups (default 4)
MM_LIMIT candidates per record (default 20)
"""
import logging
import os
import shutil
import sys
import time
from pathlib import Path
_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)
# Constrain the heap before Spark starts, so the demo proves the search path stays
# off the driver.
os.environ.setdefault(
"PYSPARK_SUBMIT_ARGS",
f"--driver-memory {os.environ.get('MM_DRIVER_MEM', '1g')} "
f"--executor-memory {os.environ.get('MM_EXEC_MEM', '1g')} pyspark-shell")
import pyspark.sql.functions as F
from pyspark.sql import SparkSession
from madmatcher_pro.reliability.resumable_search import run_resumable_search
from madmatcher_pro.sparkly.index.lucene_index import LuceneIndex
from madmatcher_pro.sparkly.index_config import IndexConfig
from madmatcher_pro.sparkly.search import Searcher
logging.basicConfig(level=logging.INFO,
format="%(asctime)s %(levelname)s %(message)s")
class FailAfterNGroups:
"""Wraps a real Searcher and raises once it has searched `fail_after` groups.
The one piece of scaffolding this demo cannot do without: it simulates a crash
partway through the run (never on group 0, so some work commits first).
"""
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)
# ---- setup ------------------------------------------------------------------
DATA = Path(os.environ.get(
"MM_DATA_DIR", str(_SRC.parent / "tests" / "sparkly" / "data" / "abt_buy")))
N_GROUPS = int(os.environ.get("MM_N_GROUPS", "10"))
FAIL_AFTER = int(os.environ.get("MM_FAIL_AFTER", "4"))
LIMIT = int(os.environ.get("MM_LIMIT", "20"))
ID_COL = "_id"
SCRATCH = _SRC.parent / ".demo_scratch"
INDEX = SCRATCH / f"{DATA.name}_index"
CKPT = SCRATCH / "blocking_ckpt"
SCRATCH.mkdir(exist_ok=True)
if CKPT.exists():
shutil.rmtree(CKPT) # always start clean so the demo is reproducible
spark = (SparkSession.builder.appName("resumable blocking").master("local[2]")
.config("spark.driver.maxResultSize", os.environ.get("MM_MAX_RESULT", "128m"))
.config("spark.sql.shuffle.partitions", "8")
.config("spark.sql.execution.arrow.maxRecordsPerBatch", "200")
.getOrCreate())
spark.sparkContext.setLogLevel("WARN")
# ---- load data --------------------------------------------------------------
# sparkly stores/queries the id as a Lucene long, so coerce it.
table_a = spark.read.parquet(str(DATA / "table_a.parquet")).withColumn(
ID_COL, F.col(ID_COL).cast("long")).persist()
table_b = spark.read.parquet(str(DATA / "table_b.parquet")).withColumn(
ID_COL, F.col(ID_COL).cast("long"))
fields = [c for c, t in table_a.dtypes if t == "string" and c != ID_COL]
print(f" index rows {table_a.count()} search rows {table_b.count()} "
f"fields {fields}")
# ---- build the index --------------------------------------------------------
config = IndexConfig()
for field in fields:
config.add_field(field, ["standard"])
config.id_col = ID_COL
index = LuceneIndex(str(INDEX), config, delete_if_exists=True)
index.upsert_docs(table_a, show_progress_bar=True)
index.init()
print(f" index built: {index.num_indexed_docs()} docs")
searcher = Searcher(index, search_chunk_size=500)
search_args = dict(search_df=table_b, query_spec=searcher.get_full_query_spec(),
limit=LIMIT, checkpoint_dir=str(CKPT), id_col=ID_COL,
show_progress=True, n_groups=N_GROUPS)
# ---- run 1: crashes after FAIL_AFTER groups ---------------------------------
print("\n" + "=" * 70)
print(f" RUN 1 -- will crash after {FAIL_AFTER} of {N_GROUPS} groups")
print("=" * 70)
start = time.time()
try:
run_resumable_search(FailAfterNGroups(searcher, FAIL_AFTER), **search_args)
print(">>> WARNING: run 1 did not crash (MM_FAIL_AFTER >= MM_N_GROUPS?)")
except RuntimeError as exc:
print(f"\n>>> caught the expected crash: {exc}")
print(f" run 1 wall time {time.time() - start:.1f}s")
# ---- run 2: resume from the checkpoint --------------------------------------
print("\n" + "=" * 70)
print(" RUN 2 -- resume with a healthy searcher (committed groups are skipped)")
print("=" * 70)
start = time.time()
candidates = searcher.search(enable_crash_recovery=True, **search_args)
n_records = candidates.count()
n_pairs = candidates.selectExpr("explode(id1_list) as id1").count()
print(f" run 2 wall time {time.time() - start:.1f}s")
# ---- stats ------------------------------------------------------------------
print("\n" + "=" * 70)
print(f" search records with candidates : {n_records}")
print(f" total candidate pairs : {n_pairs}")
print(f" checkpoint dir : {CKPT}")
print("=" * 70 + "\n")
spark.stop()
Real-time serving (add-on)¶
The resident counterpart of the batch demos above: match one incoming record at a
time, in-process. All of these require the separately-licensed realtime
entitlement. The first three are the same pipeline at three points on the
batch-to-streaming spectrum; the rest exercise one capability each.
scripts/demos/realtime/demo_batch_to_streaming.py¶
#!/usr/bin/env python
"""BATCH -> STREAMING: run the batch build, then serve records live from it.
The usual shape for going live. You already run a batch build (index + matcher);
this publishes that build as a serving bundle and then answers incoming records
one at a time, in-process, with no Spark on the request path. Once serving starts
the cluster is no longer needed and can be scaled down.
Records arriving DURING the build are not handled here (they have no bundle to
match against yet). For that, see demo_batch_backfill_to_streaming.py.
<venv>/bin/python scripts/demos/realtime/demo_batch_to_streaming.py
Requires the `realtime` license entitlement.
"""
import os
import sys
import tempfile
import warnings
from pathlib import Path
_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)
os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
warnings.filterwarnings("ignore")
import numpy as np
from pyspark.sql import SparkSession
from xgboost import XGBClassifier
from madmatcher_pro import MadMatcherSession
from madmatcher_pro.matchflow import SKLearnModel, create_features, featurize
from madmatcher_pro.serving import (
DataFrameSink,
IterableSource,
RealtimeMatcher,
publish_bundle_version,
publish_serving_artifacts,
run_matching,
)
from madmatcher_pro.sparkly.index.lucene_index import LuceneIndex
from madmatcher_pro.sparkly.index_config import IndexConfig
from madmatcher_pro.sparkly.search import Searcher
# ---- setup ------------------------------------------------------------------
spark = (SparkSession.builder.master("local[*]")
.appName("batch to streaming").getOrCreate())
spark.sparkContext.setLogLevel("ERROR")
session = MadMatcherSession.builder.getOrCreate()
if session.url:
print(f"\n Dashboard: {session.url}\n")
DATA = _SRC.parent / "tests" / "sparkly" / "data" / "abt_buy"
COLS = ["name", "description"]
BLOCK_LIMIT = 20
WORK = Path(tempfile.mkdtemp(prefix="mm_b2s_"))
BUNDLE = WORK / "bundles" / "v1"
# ---- load data --------------------------------------------------------------
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"))
gold_pairs = {(int(r["id1"]), int(r["id2"])) for r in gold.collect()}
# ---- batch: build the index -------------------------------------------------
print("1. Batch build: indexing table A...")
config = IndexConfig(id_col="_id")
for c in COLS:
config.add_field(c, ["3gram"])
index = LuceneIndex(WORK / "index", config, delete_if_exists=True)
index.upsert_docs(table_a, force_distributed=True)
index.init()
# ---- batch: block + featurize + train ---------------------------------------
print("2. Batch build: blocking, featurizing, training the matcher...")
features = create_features(A=table_a, B=table_b, a_cols=COLS, b_cols=COLS)
candidates = Searcher(index).search(table_b, index.get_full_query_spec(), BLOCK_LIMIT)
fvs = featurize(features=features, A=table_a, B=table_b, candidates=candidates,
fill_na=0.0).cache()
rows = fvs.select("id1", "id2", "feature_vectors").collect()
X = np.array([r["feature_vectors"] for r in rows], dtype=float)
y = np.array([1 if (int(r["id1"]), int(r["id2"])) in gold_pairs else 0 for r in rows])
model = SKLearnModel(
XGBClassifier(eval_metric="logloss", objective="binary:logistic",
max_depth=6, seed=42, n_jobs=1).fit(X, y),
nan_fill=0.0)
print(f" trained on {len(y)} pairs ({int(y.sum())} positive)")
# ---- publish the bundle -----------------------------------------------------
print("3. Publishing the serving bundle...")
publish_serving_artifacts(
BUNDLE, indexed_table=table_a, features=features, model=model,
sparkly_index_path=WORK / "index", fill_na=0.0, block_limit=BLOCK_LIMIT,
id_col="_id", spark=spark)
# Point serving replicas at this version. They pick it up on their next poll, so a
# refresh needs no coordination. See demo_streaming_only.py for the watcher side.
publish_bundle_version(WORK / "bundles", BUNDLE, version="v1")
# ---- streaming: serve records one at a time ---------------------------------
print("4. Serving: streaming records through the resident matcher...")
matcher = RealtimeMatcher.load(BUNDLE, threshold=0.5)
# WHERE RECORDS COME FROM. Swap the source; nothing else changes.
# collect() pulls the whole table into the driver, fine at this size. For a
# large input use ParquetSource/CsvSource, which stream in batches.
source = IterableSource([r.asDict() for r in table_b.collect()], id_col="_id")
# from madmatcher_pro.serving import ParquetSource
# source = ParquetSource("incoming/", id_col="_id")
# from madmatcher_pro.serving import CsvSource
# source = CsvSource("incoming.csv", id_col="_id")
# from madmatcher_pro.serving import QueueSource # records stored nowhere
# source = QueueSource(maxsize=1000, id_col="_id") # producer calls source.put(rec)
# WHERE RESULTS GO. Swap the sink; nothing else changes.
# DataFrameSink keeps every result in memory, which is right for this small
# dataset. For a large or unbounded stream use ParquetSink or CallbackSink,
# which hold nothing.
sink = DataFrameSink()
# from madmatcher_pro.serving import ParquetSink
# sink = ParquetSink(WORK / "matches", batch_rows=10_000)
# from madmatcher_pro.serving import CallbackSink
# sink = CallbackSink(lambda result, record: print(result))
# from madmatcher_pro.serving import FanOutResultSink # in memory AND on disk
# sink = FanOutResultSink([DataFrameSink(), ParquetSink(WORK / "matches")])
stats = run_matching(source, sink, matcher, progress_every=250)
# For many concurrent requests, serve through a pool instead of the matcher:
# from madmatcher_pro.serving import MatcherPool # threads, one shared matcher
# pool = MatcherPool(matcher, n_workers=4)
# stats = run_matching(source, sink, pool)
#
# For full multi-core throughput use worker PROCESSES. NB: `spawn` re-imports the main
# module in every worker, so a script using ProcessMatcherPool must put its body under
# `if __name__ == "__main__":` or each worker re-runs the whole pipeline. See
# demo_realtime_msd_fusion.py for that shape.
# from madmatcher_pro.serving import ProcessMatcherPool
# pool = ProcessMatcherPool(BUNDLE, n_procs=4)
# stats = run_matching(source, sink, pool)
# ---- stats ------------------------------------------------------------------
matches = sink.result
predicted = {(int(a), int(b)) for a, b in zip(matches["id1"], matches["id2"])}
tp = len(gold_pairs & predicted)
precision = tp / len(predicted) if predicted else 0.0
recall = tp / len(gold_pairs) if gold_pairs else 0.0
f1 = 2 * precision * recall / (precision + recall) if (precision + recall) else 0.0
print(f"\n served {stats.records} records in {stats.elapsed:.1f}s "
f"({stats.records / max(stats.elapsed, 1e-9):.0f} rec/s), "
f"{stats.matches} matches")
print(f" precision {precision:.4f} recall {recall:.4f} f1 {f1:.4f}\n")
print(" The cluster is no longer needed: serving runs in-process from the bundle.")
print(" To serve this bundle on its own (no Spark at all):")
print(f" MM_BUNDLE={BUNDLE} \\\n"
f" python scripts/demos/realtime/demo_streaming_only.py\n")
matcher.close()
session.stop()
spark.stop()
scripts/demos/realtime/demo_batch_backfill_to_streaming.py¶
#!/usr/bin/env python
"""BATCH + BACKFILL -> STREAMING: lose nothing while the build is running.
A batch build takes a while, and records that arrive DURING it have no bundle to be
matched against yet. Holding them in memory loses them to a restart; dropping them
leaves a gap in coverage exactly as wide as the build, on your newest records.
`BatchToServing` closes it:
accept records arriving during the build are spooled to disk (survive a crash)
publish the finished build becomes a serving bundle
drain the spooled backlog is matched (resumable, per record)
serve live records from here on, cluster no longer needed
<venv>/bin/python scripts/demos/realtime/demo_batch_backfill_to_streaming.py
Requires the `realtime` license entitlement.
"""
import os
import sys
import tempfile
import warnings
from pathlib import Path
_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)
os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
warnings.filterwarnings("ignore")
import numpy as np
from pyspark.sql import SparkSession
from xgboost import XGBClassifier
from madmatcher_pro import MadMatcherSession
from madmatcher_pro.matchflow import SKLearnModel, create_features, featurize
from madmatcher_pro.serving import (
BatchToServing,
DataFrameSink,
IterableSource,
RealtimeMatcher,
publish_serving_artifacts,
run_matching,
)
from madmatcher_pro.sparkly.index.lucene_index import LuceneIndex
from madmatcher_pro.sparkly.index_config import IndexConfig
from madmatcher_pro.sparkly.search import Searcher
# ---- setup ------------------------------------------------------------------
spark = (SparkSession.builder.master("local[*]")
.appName("batch backfill to streaming").getOrCreate())
spark.sparkContext.setLogLevel("ERROR")
session = MadMatcherSession.builder.getOrCreate()
if session.url:
print(f"\n Dashboard: {session.url}\n")
DATA = _SRC.parent / "tests" / "sparkly" / "data" / "abt_buy"
COLS = ["name", "description"]
BLOCK_LIMIT = 20
WORK = Path(tempfile.mkdtemp(prefix="mm_backfill_"))
BUNDLE = WORK / "bundles" / "v1"
# ---- load data --------------------------------------------------------------
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"))
gold_pairs = {(int(r["id1"]), int(r["id2"])) for r in gold.collect()}
# Split table B: the first slice stands in for records that arrive while the build
# runs (the backfill), the rest for the live stream afterwards.
incoming = [r.asDict() for r in table_b.collect()]
during_build, live = incoming[:300], incoming[300:]
# The handoff owns the spool (durable buffer) and the drain. checkpoint_dir makes the
# drain resumable per record, so an interrupted drain redoes only what it must.
handoff = BatchToServing(spool_dir=WORK / "spool", bundle_dir=BUNDLE,
id_col="_id", checkpoint_dir=WORK / "drain_ckpt")
# ---- batch build, with records arriving throughout --------------------------
print("1. Batch build: indexing table A (records arriving meanwhile are spooled)...")
for record in during_build[:150]:
handoff.accept(record)
config = IndexConfig(id_col="_id")
for c in COLS:
config.add_field(c, ["3gram"])
index = LuceneIndex(WORK / "index", config, delete_if_exists=True)
index.upsert_docs(table_a, force_distributed=True)
index.init()
for record in during_build[150:]: # still arriving, mid-build
handoff.accept(record)
print("2. Batch build: blocking, featurizing, training the matcher...")
features = create_features(A=table_a, B=table_b, a_cols=COLS, b_cols=COLS)
candidates = Searcher(index).search(table_b, index.get_full_query_spec(), BLOCK_LIMIT)
fvs = featurize(features=features, A=table_a, B=table_b, candidates=candidates,
fill_na=0.0).cache()
rows = fvs.select("id1", "id2", "feature_vectors").collect()
X = np.array([r["feature_vectors"] for r in rows], dtype=float)
y = np.array([1 if (int(r["id1"]), int(r["id2"])) in gold_pairs else 0 for r in rows])
model = SKLearnModel(
XGBClassifier(eval_metric="logloss", objective="binary:logistic",
max_depth=6, seed=42, n_jobs=1).fit(X, y),
nan_fill=0.0)
print(f" trained on {len(y)} pairs ({int(y.sum())} positive)")
print(f" spooled while building: {handoff.accepted} records")
# ---- publish, then drain the backlog ----------------------------------------
print("3. Publishing the bundle from the finished build...")
handoff.publish(lambda out: publish_serving_artifacts(
out, indexed_table=table_a, features=features, model=model,
sparkly_index_path=WORK / "index", fill_na=0.0, block_limit=BLOCK_LIMIT,
id_col="_id", spark=spark))
print("4. Draining the backfill (the records that arrived during the build)...")
handoff.start_serving(RealtimeMatcher.load(BUNDLE, threshold=0.5))
# WHERE RESULTS GO. Swap the sink; nothing else changes.
# DataFrameSink keeps every result in memory, which is right for this small
# dataset. For a large or unbounded stream use ParquetSink or CallbackSink,
# which hold nothing.
sink = DataFrameSink()
# from madmatcher_pro.serving import ParquetSink
# sink = ParquetSink(WORK / "matches", batch_rows=10_000)
# from madmatcher_pro.serving import CallbackSink
# sink = CallbackSink(lambda result, record: print(result))
# from madmatcher_pro.serving import FanOutResultSink # in memory AND on disk
# sink = FanOutResultSink([DataFrameSink(), ParquetSink(WORK / "matches")])
backfill_stats = handoff.drain(sink)
print(f" backfill: {backfill_stats.records} records, "
f"{backfill_stats.matches} matches")
# ---- streaming: live records from here on -----------------------------------
print("5. Serving the live stream...")
# WHERE RECORDS COME FROM. Swap the source; nothing else changes.
source = IterableSource(live, id_col="_id")
# from madmatcher_pro.serving import ParquetSource
# source = ParquetSource("incoming/", id_col="_id")
# from madmatcher_pro.serving import QueueSource # records stored nowhere
# source = QueueSource(maxsize=1000, id_col="_id") # producer calls source.put(rec)
live_stats = run_matching(source, sink, handoff.matcher, progress_every=250)
# ---- stats ------------------------------------------------------------------
matches = sink.result
predicted = {(int(a), int(b)) for a, b in zip(matches["id1"], matches["id2"])}
tp = len(gold_pairs & predicted)
precision = tp / len(predicted) if predicted else 0.0
recall = tp / len(gold_pairs) if gold_pairs else 0.0
f1 = 2 * precision * recall / (precision + recall) if (precision + recall) else 0.0
served = backfill_stats.records + live_stats.records
print(f"\n backfill {backfill_stats.records} + live {live_stats.records} "
f"= {served} records served (of {len(incoming)} that arrived)")
print(f" precision {precision:.4f} recall {recall:.4f} f1 {f1:.4f}")
print(" Nothing arriving during the build was dropped.\n")
handoff.matcher.close()
handoff.close()
session.stop()
spark.stop()
scripts/demos/realtime/demo_streaming_only.py¶
#!/usr/bin/env python
"""STREAMING ONLY: serve an already-published bundle. No batch, no Spark, no cluster.
This is the long-running half of the system, on its own. The batch build happened
somewhere else (another machine, last night, a different team); here you load its
bundle and answer records as they arrive. There is no SparkSession in this script.
Run demo_batch_to_streaming.py first to produce a bundle, then point this at it:
MM_BUNDLE=/path/to/bundle \\
<venv>/bin/python scripts/demos/realtime/demo_streaming_only.py
Requires the `realtime` license entitlement.
"""
import os
import sys
import tempfile
import warnings
from pathlib import Path
_SRC = Path(__file__).resolve().parents[3] / "src"
sys.path.insert(0, str(_SRC))
warnings.filterwarnings("ignore")
from madmatcher_pro.serving import (
DataFrameSink,
IterableSource,
RealtimeMatcher,
run_matching,
)
# ---- setup ------------------------------------------------------------------
BUNDLE = os.environ.get("MM_BUNDLE")
if not BUNDLE or not (Path(BUNDLE) / "MANIFEST.json").exists():
sys.exit(
"Set MM_BUNDLE to a published serving bundle, e.g.\n"
" python scripts/demos/realtime/demo_batch_to_streaming.py # prints one\n"
" MM_BUNDLE=<that path> python scripts/demos/realtime/demo_streaming_only.py")
WORK = Path(tempfile.mkdtemp(prefix="mm_stream_"))
# Records to serve. In a real deployment these arrive from your transport; here we
# use a handful of literals so the script needs no data files.
INCOMING = [
{"_id": 1, "name": "canon eos rebel t2i", "description": "18mp digital slr camera"},
{"_id": 2, "name": "sony bravia 46 lcd", "description": "1080p hdtv"},
{"_id": 3, "name": "linksys wireless router", "description": "wireless-n broadband"},
{"_id": 4, "name": "netgear gigabit switch", "description": "8 port unmanaged"},
]
# ---- load the bundle --------------------------------------------------------
print(f"1. Loading the bundle: {BUNDLE}")
matcher = RealtimeMatcher.load(BUNDLE, threshold=0.5)
# ---- stream: source -> matcher -> sink --------------------------------------
print("2. Serving...")
# WHERE RECORDS COME FROM. Swap the source; nothing else changes.
source = IterableSource(INCOMING, id_col="_id")
# from madmatcher_pro.serving import ParquetSource
# source = ParquetSource("incoming/", id_col="_id")
# from madmatcher_pro.serving import CsvSource
# source = CsvSource("incoming.csv", id_col="_id")
# from madmatcher_pro.serving import QueueSource # records stored nowhere:
# source = QueueSource(maxsize=1000, id_col="_id") # your handler calls source.put(rec)
# WHERE RESULTS GO. Swap the sink; nothing else changes.
# DataFrameSink keeps every result in memory, which is right for this small
# dataset. For a large or unbounded stream use ParquetSink or CallbackSink,
# which hold nothing.
sink = DataFrameSink()
# from madmatcher_pro.serving import ParquetSink
# sink = ParquetSink(WORK / "matches", batch_rows=10_000)
# from madmatcher_pro.serving import CallbackSink
# sink = CallbackSink(lambda result, record: print(result))
# from madmatcher_pro.serving import FanOutResultSink # in memory AND on disk
# sink = FanOutResultSink([DataFrameSink(), ParquetSink(WORK / "matches")])
stats = run_matching(source, sink, matcher)
# For many concurrent requests, serve through a pool instead of the matcher:
# from madmatcher_pro.serving import MatcherPool # threads, one shared matcher
# pool = MatcherPool(matcher, n_workers=4)
# stats = run_matching(source, sink, pool)
#
# For full multi-core throughput use worker PROCESSES. NB: `spawn` re-imports the main
# module in every worker, so a script using ProcessMatcherPool must put its body under
# `if __name__ == "__main__":` or each worker re-runs the whole pipeline. See
# demo_realtime_msd_fusion.py for that shape.
# from madmatcher_pro.serving import ProcessMatcherPool
# pool = ProcessMatcherPool(BUNDLE, n_procs=4)
# stats = run_matching(source, sink, pool)
# For a long run you cannot afford to restart, claim each record and mark it done, so
# a crash resumes and several machines can share the input with no coordinator:
# from madmatcher_pro.serving import run_matching_durable
# stats = run_matching_durable(source, sink, matcher, checkpoint_dir=WORK / "ckpt")
# To pick up new bundle versions with no downtime, watch the pointer a publisher
# writes (see demo_batch_to_streaming.py, which writes one after publishing):
# from madmatcher_pro.serving import BundleWatcher, HotSwappableMatcher
# hot = HotSwappableMatcher(BUNDLE, threshold=0.5)
# BundleWatcher(hot, Path(BUNDLE).parent, poll_seconds=60).start()
# ---- stats ------------------------------------------------------------------
print(f"\n served {stats.records} records, {stats.matches} matches "
f"in {stats.elapsed * 1000:.0f} ms")
print(f"\n{sink.result.to_string(index=False)}\n")
matcher.close()
scripts/demos/realtime/demo_realtime_matching.py¶
#!/usr/bin/env python
"""Real-time matching: the per-record `match()` API, plus a version hot-swap.
The resident counterpart of scripts/demos/matchflow/demo_matchflow_matching.py. That
demo runs `apply_matcher` over a WHOLE candidate table in Spark; this one packages the
same artifacts into a serving bundle and calls `matcher.match(record)` per record, so
you see the raw request-path API and its latency.
Then it swaps a new bundle version into the live matcher with no downtime: in-flight
requests finish on the old bundle while new ones go to the new one.
For the orchestrated pipeline (a source in, a sink out, MadMatcher running the loop),
see demo_batch_to_streaming.py.
<venv>/bin/python scripts/demos/realtime/demo_realtime_matching.py
Requires the `realtime` license entitlement.
"""
import os
import sys
import tempfile
import time
import warnings
from pathlib import Path
_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)
os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
os.environ.setdefault("MADMATCHER_DASHBOARD", "0")
warnings.filterwarnings("ignore")
import numpy as np
from pyspark.sql import SparkSession
from xgboost import XGBClassifier
from madmatcher_pro.matchflow import SKLearnModel, create_features, featurize
from madmatcher_pro.serving import (
HotSwappableMatcher,
RealtimeMatcher,
publish_bundle_version,
publish_serving_artifacts,
)
from madmatcher_pro.sparkly.index.lucene_index import LuceneIndex
from madmatcher_pro.sparkly.index_config import IndexConfig
from madmatcher_pro.sparkly.search import Searcher
# ---- setup ------------------------------------------------------------------
spark = (SparkSession.builder.master("local[*]")
.appName("realtime matching").getOrCreate())
spark.sparkContext.setLogLevel("ERROR")
DATA = _SRC.parent / "tests" / "sparkly" / "data" / "abt_buy"
COLS = ["name", "description"]
BLOCK_LIMIT = 20
WORK = Path(tempfile.mkdtemp(prefix="mm_rt_match_"))
# ---- load data --------------------------------------------------------------
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"))
gold_pairs = {(int(r["id1"]), int(r["id2"])) for r in gold.collect()}
# ---- batch build: index, block, featurize, train ----------------------------
print("1. Batch build: index + trained matcher...")
config = IndexConfig(id_col="_id")
for c in COLS:
config.add_field(c, ["3gram"])
index = LuceneIndex(WORK / "index", config, delete_if_exists=True)
index.upsert_docs(table_a, force_distributed=True)
index.init()
features = create_features(A=table_a, B=table_b, a_cols=COLS, b_cols=COLS)
candidates = Searcher(index).search(table_b, index.get_full_query_spec(), BLOCK_LIMIT)
fvs = featurize(features=features, A=table_a, B=table_b, candidates=candidates,
fill_na=0.0).cache()
rows = fvs.select("id1", "id2", "feature_vectors").collect()
X = np.array([r["feature_vectors"] for r in rows], dtype=float)
y = np.array([1 if (int(r["id1"]), int(r["id2"])) in gold_pairs else 0 for r in rows])
model = SKLearnModel(
XGBClassifier(eval_metric="logloss", objective="binary:logistic",
max_depth=6, seed=42, n_jobs=1).fit(X, y),
nan_fill=0.0)
print(f" trained on {len(y)} pairs ({int(y.sum())} positive)")
# ---- publish v1 and v2 ------------------------------------------------------
print("2. Publishing bundle v1 and v2...")
for version in ("v1", "v2"):
publish_serving_artifacts(
WORK / "bundles" / version, indexed_table=table_a, features=features,
model=model, sparkly_index_path=WORK / "index", fill_na=0.0,
block_limit=BLOCK_LIMIT, id_col="_id", spark=spark)
# ---- match one record at a time ---------------------------------------------
print("3. Matching records one at a time with matcher.match()...")
matcher = RealtimeMatcher.load(WORK / "bundles" / "v1", threshold=0.5)
records = [r.asDict() for r in table_b.collect()]
matcher.match(records[0]) # warm before timing
latencies, predicted = [], set()
for i, record in enumerate(records, 1):
start = time.perf_counter()
result = matcher.match(record) # [id1, id2, prediction, confidence]
latencies.append((time.perf_counter() - start) * 1000.0)
predicted.update((int(a), int(b)) for a, b in zip(result["id1"], result["id2"]))
if i % 250 == 0:
print(f" ...matched {i} records")
matcher.close()
# ---- stats ------------------------------------------------------------------
tp = len(gold_pairs & predicted)
precision = tp / len(predicted) if predicted else 0.0
recall = tp / len(gold_pairs) if gold_pairs else 0.0
f1 = 2 * precision * recall / (precision + recall) if (precision + recall) else 0.0
latencies = np.array(latencies)
print(f"\n precision {precision:.4f} recall {recall:.4f} f1 {f1:.4f}")
print(f" per-record ms mean {latencies.mean():.2f} "
f"p50 {np.percentile(latencies, 50):.2f} "
f"p95 {np.percentile(latencies, 95):.2f}\n")
# ---- hot-swap to a new version, with no downtime ----------------------------
print("4. Hot-swapping v1 -> v2 (in-flight requests drain on v1)...")
hot = HotSwappableMatcher(WORK / "bundles" / "v1", threshold=0.5, track=False)
print(f" serving {Path(hot.version).name}")
hot.swap(WORK / "bundles" / "v2")
hot.match(records[0])
print(f" now serving {Path(hot.version).name} (un-closed bundles: {hot.undrained})")
# Across MANY replicas, do not call swap() on each: publish a pointer and let each
# replica notice on its own poll. No coordinator, and a replica that was down picks
# up the current version when it starts.
publish_bundle_version(WORK / "bundles", WORK / "bundles" / "v2", version="v2")
# from madmatcher_pro.serving import BundleWatcher
# BundleWatcher(hot, WORK / "bundles", poll_seconds=60).start()
hot.close()
spark.stop()
scripts/demos/realtime/demo_realtime_sparkly.py¶
#!/usr/bin/env python
"""Real-time sparkly blocking: probe ONE record at a time, in-process.
The resident counterpart of scripts/demos/sparkly/demo_sparkly_blocking.py. That
demo bulk-searches ALL of table B against the index in one Spark job; this one
builds the same Lucene index offline, then reopens it into a resident
`ResidentSparklyProbe` and streams table B through it one record at a time, each
probe running in-process with no Spark job.
The resident probe calls the same `LuceneIndex.search`, so the candidates (and
recall) are identical to the batch searcher's. The point is the shape: per record,
no cluster round-trip, which is what a live serving path needs.
<venv>/bin/python scripts/demos/realtime/demo_realtime_sparkly.py
Requires the `realtime` license entitlement. For block + featurize + predict, see
demo_batch_to_streaming.py.
"""
import os
import sys
import tempfile
import time
import warnings
from pathlib import Path
_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)
os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
os.environ.setdefault("MADMATCHER_DASHBOARD", "0")
warnings.filterwarnings("ignore")
import numpy as np
from pyspark.sql import SparkSession
from madmatcher_pro.serving import ResidentSparklyProbe
from madmatcher_pro.sparkly.index.lucene_index import LuceneIndex
from madmatcher_pro.sparkly.index_config import IndexConfig
# ---- setup ------------------------------------------------------------------
spark = (SparkSession.builder.master("local[*]")
.appName("realtime sparkly").getOrCreate())
spark.sparkContext.setLogLevel("ERROR")
DATA = _SRC.parent / "tests" / "sparkly" / "data" / "abt_buy"
COLS = ["name", "description"]
LIMIT = 50 # candidates per record
INDEX = Path(tempfile.mkdtemp(prefix="mm_rt_sparkly_"))
# ---- load data --------------------------------------------------------------
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"))
gold_pairs = {(int(r["id1"]), int(r["id2"])) for r in gold.collect()}
# ---- build the index (offline, in Spark) ------------------------------------
print("1. Building the sparkly (Lucene) index over table A...")
config = IndexConfig(id_col="_id")
for c in COLS:
config.add_field(c, ["3gram"])
LuceneIndex(INDEX, config, delete_if_exists=True).upsert_docs(
table_a, force_distributed=True)
# ---- load the resident probe ------------------------------------------------
print("2. Loading the resident sparkly probe...")
probe = ResidentSparklyProbe.load(INDEX, limit=LIMIT)
# ---- block: stream records one at a time ------------------------------------
records = [r.asDict() for r in table_b.collect()]
probe.probe(records[0]) # warm the JVM before timing
print(f"3. Probing {len(records)} records one at a time, in-process...")
latencies, candidate_pairs = [], set()
for i, record in enumerate(records, 1):
start = time.perf_counter()
a_ids = probe.probe(record)
latencies.append((time.perf_counter() - start) * 1000.0)
for a in a_ids:
candidate_pairs.add((int(a), int(record["_id"])))
if i % 250 == 0:
print(f" ...probed {i} records")
# ---- stats ------------------------------------------------------------------
recall = len(gold_pairs & candidate_pairs) / len(gold_pairs)
latencies = np.array(latencies)
print(f"\n {len(latencies)} records probed in-process recall {recall:.4f}")
print(f" per-record ms mean {latencies.mean():.2f} "
f"p50 {np.percentile(latencies, 50):.2f} "
f"p95 {np.percentile(latencies, 95):.2f}\n")
probe.close()
spark.stop()
scripts/demos/realtime/demo_realtime_delex.py¶
#!/usr/bin/env python
"""Real-time delex blocking: run a blocking program ONE record at a time, in-process.
The resident counterpart of scripts/demos/delex/demo_delex_blocking.py. That demo
runs the blocking program over ALL of table B in Spark; this one builds the same
program's structures once over table A, keeps them resident, then streams table B
through it one record at a time with no Spark job per probe.
delex has no cross-process resident form (its indexes are process-local memmap /
embedded-Lucene structures), so the build happens once here and each record then runs
delex's own kernel over a one-row batch, giving candidates identical to a batch run.
<venv>/bin/python scripts/demos/realtime/demo_realtime_delex.py
Requires the `realtime` license entitlement. For block + featurize + predict, see
demo_batch_to_streaming.py.
"""
import os
import sys
import time
import warnings
from pathlib import Path
_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)
os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
os.environ.setdefault("MADMATCHER_DASHBOARD", "0")
warnings.filterwarnings("ignore")
import numpy as np
from pyspark.sql import SparkSession
from madmatcher_pro.delex.lang import BlockingProgram, KeepRule
from madmatcher_pro.delex.lang.predicate import BM25TopkPredicate
from madmatcher_pro.serving import ResidentDelexProbe
# ---- setup ------------------------------------------------------------------
spark = (SparkSession.builder.master("local[*]")
.config("spark.sql.execution.arrow.pyspark.enabled", "true")
.appName("realtime delex").getOrCreate())
spark.sparkContext.setLogLevel("ERROR")
DATA = _SRC.parent / "tests" / "sparkly" / "data" / "abt_buy"
LIMIT = 50
# ---- load data --------------------------------------------------------------
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"))
gold_pairs = {(int(r["id1"]), int(r["id2"])) for r in gold.collect()}
# ---- build the resident probe (one Spark job, then none per probe) ----------
# A two-rule blocking program: BM25 top-20 on 'name' AND on 'description' (union),
# matching the batch delex demo.
program = BlockingProgram(
keep_rules=[
KeepRule([BM25TopkPredicate("name", "name", "standard", 20)]),
KeepRule([BM25TopkPredicate("description", "description", "standard", 20)]),
],
drop_rules=[],
)
print("1. Compiling + building the delex plan over table A (one Spark job)...")
probe = ResidentDelexProbe.build(table_a, program, id_col="_id", limit=LIMIT)
# ---- block: stream records one at a time ------------------------------------
records = [r.asDict() for r in table_b.collect()]
probe.probe(records[0]) # warm
print(f"2. Probing {len(records)} records one at a time, in-process...")
latencies, candidate_pairs = [], set()
for i, record in enumerate(records, 1):
start = time.perf_counter()
a_ids = probe.probe(record)
latencies.append((time.perf_counter() - start) * 1000.0)
for a in a_ids:
candidate_pairs.add((int(a), int(record["_id"])))
if i % 250 == 0:
print(f" ...probed {i} records")
# ---- stats ------------------------------------------------------------------
recall = len(gold_pairs & candidate_pairs) / len(gold_pairs)
latencies = np.array(latencies)
print(f"\n {len(latencies)} records probed in-process recall {recall:.4f}")
print(f" per-record ms mean {latencies.mean():.2f} "
f"p50 {np.percentile(latencies, 50):.2f} "
f"p95 {np.percentile(latencies, 95):.2f}\n")
probe.close()
spark.stop()
scripts/demos/realtime/demo_realtime_semantic.py¶
#!/usr/bin/env python
"""Real-time semantic blocking: embed + search ONE record at a time, in-process.
The resident counterpart of scripts/demos/semantic/demo_semantic_blocking.py. That
demo searches ALL of table B against the dense index in Spark; this one builds the
same flat IVF index offline, then reopens it into a resident `ResidentSemanticProbe`
and streams table B through it one record at a time. Each probe embeds the record
live, assigns it to the nearest IVF cells, and scores the postings by cosine, all
in-process with no Spark job.
The scoring is the same math the Spark `SemanticSearcher` runs, so recall matches.
Per-record latency is dominated by the live embedding, not the search.
<venv>/bin/python scripts/demos/realtime/demo_realtime_semantic.py
Needs the `semantic` extra (sentence-transformers) and the `realtime` license
entitlement. Fused blocking is demo_realtime_fusion.py.
"""
import os
import sys
import tempfile
import time
import warnings
from pathlib import Path
_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)
os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
os.environ.setdefault("MADMATCHER_DASHBOARD", "0")
warnings.filterwarnings("ignore")
import numpy as np
from pyspark.sql import SparkSession
from pyspark.sql import types as T
from madmatcher_pro.semantic import (
SemanticIndex,
SentenceTransformerProvider,
create_embeddings,
)
from madmatcher_pro.serving import ResidentSemanticProbe
# ---- setup ------------------------------------------------------------------
spark = (SparkSession.builder.master("local[*]")
.appName("realtime semantic").getOrCreate())
spark.sparkContext.setLogLevel("ERROR")
DATA = _SRC.parent / "tests" / "sparkly" / "data" / "abt_buy"
MODEL = os.environ.get("MM_DEMO_EMBED_MODEL", "BAAI/bge-small-en-v1.5")
LIMIT = 50
NPROBE = 32
WORK = Path(tempfile.mkdtemp(prefix="mm_rt_semantic_"))
# ---- load data --------------------------------------------------------------
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"))
gold_pairs = {(int(r["id1"]), int(r["id2"])) for r in gold.collect()}
# The string fields the resident probe will serialize per record.
fields = [f.name for f in table_a.schema.fields
if f.name != "_id" and isinstance(f.dataType, T.StringType)]
# device="cpu" because the model loads inside a Spark worker: on macOS, MPS cannot
# initialize across the fork. Drop it on a Linux/GPU box.
provider = SentenceTransformerProvider(MODEL, device="cpu")
# ---- build the semantic index (offline, in Spark) ---------------------------
print(f"1. Embedding table A ({MODEL}) and building the semantic index...")
vectors_a = create_embeddings(table_a, fields, provider, str(WORK / "emb_a"))
SemanticIndex(str(WORK / "index"), id_col="_id").upsert_docs(vectors_a)
# ---- load the resident probe ------------------------------------------------
print("2. Loading the resident semantic probe...")
probe = ResidentSemanticProbe.load(WORK / "index", provider=provider, fields=fields,
nprobe=NPROBE, limit=LIMIT)
# ---- block: stream records one at a time ------------------------------------
records = [r.asDict() for r in table_b.collect()]
probe.probe(records[0]) # warm (loads the model in this process)
print(f"3. Probing {len(records)} records one at a time (embed + assign + score)...")
latencies, candidate_pairs = [], set()
for i, record in enumerate(records, 1):
start = time.perf_counter()
a_ids = probe.probe(record)[0]
latencies.append((time.perf_counter() - start) * 1000.0)
for a in a_ids:
candidate_pairs.add((int(a), int(record["_id"])))
if i % 250 == 0:
print(f" ...probed {i} records")
# ---- stats ------------------------------------------------------------------
recall = len(gold_pairs & candidate_pairs) / len(gold_pairs)
latencies = np.array(latencies)
print(f"\n {len(latencies)} records probed in-process recall {recall:.4f}")
print(f" per-record ms mean {latencies.mean():.1f} "
f"p50 {np.percentile(latencies, 50):.1f} "
f"p95 {np.percentile(latencies, 95):.1f} (dominated by the live embed)\n")
spark.stop()
scripts/demos/realtime/demo_realtime_fusion.py¶
#!/usr/bin/env python
"""Real-time fusion blocking: sparkly + semantic + delex, fused per record.
The resident counterpart of scripts/demos/semantic/demo_fusion_blocking.py. That
demo fuses lexical + dense candidate lists over ALL of table B with Spark
reciprocal-rank fusion; this one keeps all three engines resident and fuses them one
record at a time with the in-process `reciprocal_rank_fusion_local`.
Per record it probes sparkly (BM25 scores), the semantic index (cosine scores), and
the delex program (scoreless, so shared rank), then merges them into one ranked
list. It reports each source's recall and the fused recall, which should be at least
as good as the best single source.
<venv>/bin/python scripts/demos/realtime/demo_realtime_fusion.py
Needs the `semantic` extra and the `realtime` license entitlement. To add the matcher
on top, see demo_realtime_fusion_matching.py.
"""
import os
import sys
import tempfile
import time
import warnings
from pathlib import Path
_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)
os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
os.environ.setdefault("MADMATCHER_DASHBOARD", "0")
warnings.filterwarnings("ignore")
import numpy as np
from pyspark.sql import SparkSession
from pyspark.sql import types as T
from madmatcher_pro.delex.lang import BlockingProgram, KeepRule
from madmatcher_pro.delex.lang.predicate import BM25TopkPredicate
from madmatcher_pro.semantic import (
SemanticIndex,
SentenceTransformerProvider,
create_embeddings,
)
from madmatcher_pro.serving import (
ResidentDelexProbe,
ResidentSemanticProbe,
ResidentSparklyProbe,
reciprocal_rank_fusion_local,
)
from madmatcher_pro.sparkly.index.lucene_index import LuceneIndex
from madmatcher_pro.sparkly.index_config import IndexConfig
# ---- setup ------------------------------------------------------------------
spark = (SparkSession.builder.master("local[*]")
.config("spark.sql.execution.arrow.pyspark.enabled", "true")
.appName("realtime fusion").getOrCreate())
spark.sparkContext.setLogLevel("ERROR")
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")
LIMIT = 50
NPROBE = 32
WORK = Path(tempfile.mkdtemp(prefix="mm_rt_fusion_"))
# ---- load data --------------------------------------------------------------
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"))
gold_pairs = {(int(r["id1"]), int(r["id2"])) for r in gold.collect()}
fields = [f.name for f in table_a.schema.fields
if f.name != "_id" and isinstance(f.dataType, T.StringType)]
# ---- build all three engines' artifacts (offline, in Spark) -----------------
print("1. Building sparkly + semantic + delex over table A...")
config = IndexConfig(id_col="_id")
for c in COLS:
config.add_field(c, ["3gram"])
LuceneIndex(WORK / "sparkly", config, delete_if_exists=True).upsert_docs(
table_a, force_distributed=True)
# device="cpu": the model loads inside a Spark worker, and on macOS MPS cannot
# initialize across the fork. Drop it on a Linux/GPU box.
provider = SentenceTransformerProvider(MODEL, device="cpu")
vectors_a = create_embeddings(table_a, fields, provider, str(WORK / "emb_a"))
SemanticIndex(str(WORK / "semantic"), id_col="_id").upsert_docs(vectors_a)
program = BlockingProgram(
keep_rules=[
KeepRule([BM25TopkPredicate("name", "name", "standard", 20)]),
KeepRule([BM25TopkPredicate("description", "description", "standard", 20)]),
],
drop_rules=[])
# ---- load the three resident probes -----------------------------------------
print("2. Loading the three resident probes...")
sparkly = ResidentSparklyProbe.load(WORK / "sparkly", limit=LIMIT)
semantic = ResidentSemanticProbe.load(WORK / "semantic", provider=provider,
fields=fields, nprobe=NPROBE, limit=LIMIT)
delex = ResidentDelexProbe.build(table_a, program, id_col="_id", limit=LIMIT)
# ---- block: probe all three and fuse, one record at a time ------------------
records = [r.asDict() for r in table_b.collect()]
sparkly.probe(records[0])
semantic.probe(records[0])
delex.probe(records[0]) # warm all three
print(f"3. Fusing {len(records)} records one at a time...")
lexical_pairs, semantic_pairs, delex_pairs, fused_pairs = set(), set(), set(), set()
latencies = []
for i, record in enumerate(records, 1):
start = time.perf_counter()
lexical_ids, lexical_scores = sparkly.probe_scored(record)
semantic_ids, semantic_cosines = semantic.probe(record)
delex_ids, delex_scores = delex.probe_scored(record)
fused_ids, _cosines = reciprocal_rank_fusion_local(
[
{"ids": lexical_ids, "scores": lexical_scores,
"weight": 0.5, "semantic": False},
{"ids": semantic_ids, "scores": semantic_cosines,
"weight": 0.5, "semantic": True},
{"ids": delex_ids, "scores": delex_scores,
"weight": 0.5, "semantic": False},
],
LIMIT)
latencies.append((time.perf_counter() - start) * 1000.0)
b_id = int(record["_id"])
lexical_pairs.update((int(a), b_id) for a in lexical_ids)
semantic_pairs.update((int(a), b_id) for a in semantic_ids)
delex_pairs.update((int(a), b_id) for a in delex_ids)
fused_pairs.update((int(a), b_id) for a in fused_ids)
if i % 250 == 0:
print(f" ...fused {i} records")
# ---- stats ------------------------------------------------------------------
latencies = np.array(latencies)
print(f"\n recall sparkly {len(gold_pairs & lexical_pairs) / len(gold_pairs):.4f}"
f" semantic {len(gold_pairs & semantic_pairs) / len(gold_pairs):.4f}"
f" delex {len(gold_pairs & delex_pairs) / len(gold_pairs):.4f}"
f" -> FUSED {len(gold_pairs & fused_pairs) / len(gold_pairs):.4f}")
print(f" per-record ms mean {latencies.mean():.1f} "
f"p50 {np.percentile(latencies, 50):.1f} "
f"p95 {np.percentile(latencies, 95):.1f} (three probes + fusion)\n")
sparkly.close()
delex.close()
spark.stop()
scripts/demos/realtime/demo_realtime_fusion_matching.py¶
#!/usr/bin/env python
"""Real-time fusion + matching: the full stack, one record at a time.
The complete serving path: sparkly + semantic + delex blocking fused per record, the
semantic cosine appended as a matcher feature, then predict, all in-process with no
Spark job on the request path.
The matcher is trained WITH the cosine feature (`featurize(a_embeddings=,
b_embeddings=)`), and the bundle is published with `cosine_feature=True` so serving
appends the same feature. Each record's embedding is passed in as `probe_embedding`,
so the served embedding matches the corpus embedding regardless of the driver's
default torch device.
<venv>/bin/python scripts/demos/realtime/demo_realtime_fusion_matching.py
Requires the `realtime` license entitlement and the `semantic` extra. The blocking
side alone is demo_realtime_fusion.py.
"""
import os
import sys
import tempfile
import time
import warnings
from pathlib import Path
_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)
os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
os.environ.setdefault("MADMATCHER_DASHBOARD", "0")
warnings.filterwarnings("ignore")
import numpy as np
import pandas as pd
from pyspark.sql import SparkSession
from pyspark.sql import types as T
from xgboost import XGBClassifier
from madmatcher_pro.delex.lang import BlockingProgram, KeepRule
from madmatcher_pro.delex.lang.predicate import BM25TopkPredicate
from madmatcher_pro.matchflow import SKLearnModel, create_features, featurize
from madmatcher_pro.semantic import (
SemanticIndex,
SentenceTransformerProvider,
create_embeddings,
)
from madmatcher_pro.serving import (
DittoProbeSerializer,
RealtimeMatcher,
publish_serving_artifacts,
)
from madmatcher_pro.sparkly.index.lucene_index import LuceneIndex
from madmatcher_pro.sparkly.index_config import IndexConfig
from madmatcher_pro.sparkly.search import Searcher
# ---- setup ------------------------------------------------------------------
spark = (SparkSession.builder.master("local[*]")
.config("spark.sql.execution.arrow.pyspark.enabled", "true")
.appName("realtime fusion matching").getOrCreate())
spark.sparkContext.setLogLevel("ERROR")
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")
BLOCK_LIMIT = 20
WORK = Path(tempfile.mkdtemp(prefix="mm_rt_fm_"))
# ---- load data --------------------------------------------------------------
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"))
gold_pairs = {(int(r["id1"]), int(r["id2"])) for r in gold.collect()}
fields = [f.name for f in table_a.schema.fields
if f.name != "_id" and isinstance(f.dataType, T.StringType)]
# device="cpu": the model loads inside a Spark worker, and on macOS MPS cannot
# initialize across the fork. Drop it on a Linux/GPU box.
provider = SentenceTransformerProvider(MODEL, device="cpu")
# ---- build all three blockers over table A ----------------------------------
print("1. Building sparkly + semantic + delex over table A...")
config = IndexConfig(id_col="_id")
for c in COLS:
config.add_field(c, ["3gram"])
index = LuceneIndex(WORK / "sparkly", config, delete_if_exists=True)
index.upsert_docs(table_a, force_distributed=True)
index.init()
vectors_a = create_embeddings(table_a, fields, provider, str(WORK / "emb_a"))
vectors_b = create_embeddings(table_b, fields, provider, str(WORK / "emb_b"))
SemanticIndex(str(WORK / "semantic"), id_col="_id").upsert_docs(vectors_a)
program = BlockingProgram(
keep_rules=[
KeepRule([BM25TopkPredicate("name", "name", "standard", 20)]),
KeepRule([BM25TopkPredicate("description", "description", "standard", 20)]),
],
drop_rules=[])
# ---- train the matcher WITH the cosine feature ------------------------------
print("2. Training a cosine-featured matcher...")
features = create_features(A=table_a, B=table_b, a_cols=COLS, b_cols=COLS)
candidates = Searcher(index).search(table_b, index.get_full_query_spec(), BLOCK_LIMIT)
fvs = featurize(features=features, A=table_a, B=table_b, candidates=candidates,
fill_na=0.0, a_embeddings=vectors_a,
b_embeddings=vectors_b).cache()
rows = fvs.select("id1", "id2", "feature_vectors").collect()
X = np.array([r["feature_vectors"] for r in rows], dtype=float)
y = np.array([1 if (int(r["id1"]), int(r["id2"])) in gold_pairs else 0 for r in rows])
model = SKLearnModel(
XGBClassifier(eval_metric="logloss", objective="binary:logistic",
max_depth=6, seed=42, n_jobs=1).fit(X, y),
nan_fill=0.0)
print(f" trained on {len(y)} pairs ({int(y.sum())} positive), "
f"feature width {X.shape[1]} (incl. cosine)")
# ---- publish the fused bundle -----------------------------------------------
# embed_provider=None means the caller supplies each probe embedding below.
print("3. Publishing the fused bundle (cosine_feature=True)...")
publish_serving_artifacts(
WORK / "bundle", indexed_table=table_a, features=features, model=model,
sparkly_index_path=WORK / "sparkly", fill_na=0.0, block_limit=BLOCK_LIMIT,
semantic_index_path=WORK / "semantic", embed_provider=None, embed_fields=fields,
cosine_feature=True, delex_program=program, id_col="_id", spark=spark)
# ---- match: fuse three blockers + cosine + predict, per record --------------
print("4. Loading the resident matcher and matching records...")
matcher = RealtimeMatcher.load(WORK / "bundle", spark=spark, threshold=0.5)
records = [r.asDict() for r in table_b.collect()]
# The same serializer the corpus embeddings were built with, so a probe vector lands
# in the same space.
serializer = DittoProbeSerializer()
matcher.match(records[0], probe_embedding=provider.embed(
pd.Series([serializer.serialize(records[0], fields)]))[0]) # warm
latencies, predicted = [], set()
for i, record in enumerate(records, 1):
start = time.perf_counter()
text = serializer.serialize(record, fields)
embedding = provider.embed(pd.Series([text]))[0]
result = matcher.match(record, probe_embedding=embedding)
latencies.append((time.perf_counter() - start) * 1000.0)
predicted.update((int(a), int(b)) for a, b in zip(result["id1"], result["id2"]))
if i % 250 == 0:
print(f" ...matched {i} records")
matcher.close()
# ---- stats ------------------------------------------------------------------
tp = len(gold_pairs & predicted)
precision = tp / len(predicted) if predicted else 0.0
recall = tp / len(gold_pairs) if gold_pairs else 0.0
f1 = 2 * precision * recall / (precision + recall) if (precision + recall) else 0.0
latencies = np.array(latencies)
print(f"\n fused matcher precision {precision:.4f} recall {recall:.4f} "
f"f1 {f1:.4f}")
print(f" per-record ms mean {latencies.mean():.1f} "
f"p50 {np.percentile(latencies, 50):.1f} "
f"p95 {np.percentile(latencies, 95):.1f} "
f"(embed + fuse 3 + featurize + predict)\n")
spark.stop()
scripts/demos/realtime/demo_realtime_dashboard.py¶
#!/usr/bin/env python
"""Crash recovery + live progress for the real-time / micro-batch serving path.
Runs the whole serving lifecycle under a MadMatcherSession so every step appears on
the live dashboard, and publishes with crash recovery so a mid-publish crash resumes
instead of rebuilding the bundle. Four jobs land on one board:
1. PUBLISH v1 with crash recovery -> a "Publishing bundle" job, per artifact
2. REAL-TIME matching (track=True) -> a "realtime serving" throughput counter
3. MICRO-BATCH streaming (track=True) -> a counter plus batches/records KPIs
4. PUBLISH v2 + HOT-SWAP -> a "hot-swap" (delta) job
Tracking is best-effort: it never changes a served prediction, and
MADMATCHER_DASHBOARD=0 turns it off entirely.
<venv>/bin/python scripts/demos/realtime/demo_realtime_dashboard.py
Requires the `realtime` license entitlement. The dashboard stays up until you press
Enter, so you can watch each job.
"""
import os
import sys
import tempfile
import time
import warnings
from pathlib import Path
_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)
os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
# Unlike the other realtime demos, this one leaves the dashboard ON. Set
# MADMATCHER_DASHBOARD=0 to silence it.
warnings.filterwarnings("ignore")
import numpy as np
from pyspark.sql import SparkSession
from xgboost import XGBClassifier
from madmatcher_pro import MadMatcherSession
from madmatcher_pro.matchflow import SKLearnModel, create_features, featurize
from madmatcher_pro.reliability.dashboard.server import shutdown_shared
from madmatcher_pro.serving import (
HotSwappableMatcher,
MicroBatchMatcher,
RealtimeMatcher,
publish_serving_artifacts,
)
from madmatcher_pro.sparkly.index.lucene_index import LuceneIndex
from madmatcher_pro.sparkly.index_config import IndexConfig
from madmatcher_pro.sparkly.search import Searcher
# ---- setup ------------------------------------------------------------------
spark = (SparkSession.builder.master("local[*]")
.appName("realtime dashboard").getOrCreate())
spark.sparkContext.setLogLevel("ERROR")
# Clear any stale in-process dashboard server so this run starts on a clean board.
# If a PREVIOUS run of this demo is still alive (it blocks on the Enter prompt at the
# end) it holds the port, so stop that process if the board looks empty.
shutdown_shared()
session = MadMatcherSession.builder.port(4050).getOrCreate()
if session.url:
print(f"\n Dashboard: {session.url}\n")
DATA = _SRC.parent / "tests" / "sparkly" / "data" / "abt_buy"
COLS = ["name", "description"]
BLOCK_LIMIT = 20
WORK = Path(tempfile.mkdtemp(prefix="mm_rt_dash_"))
# ---- load data --------------------------------------------------------------
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"))
gold_pairs = {(int(r["id1"]), int(r["id2"])) for r in gold.collect()}
# ---- batch build (index + search + featurize all show as record-count jobs) --
print("1. Building batch artifacts (sparkly index + a trained matcher)...")
config = IndexConfig(id_col="_id")
for c in COLS:
config.add_field(c, ["3gram"])
index = LuceneIndex(WORK / "index", config, delete_if_exists=True)
index.upsert_docs(table_a, force_distributed=True)
index.init()
features = create_features(A=table_a, B=table_b, a_cols=COLS, b_cols=COLS)
# Materialize the blocking output so its record count lands on the sparkly search
# job (a lazy search consumed later by featurize would finalize at 0), and so
# featurize reuses the candidates instead of re-searching.
candidates = Searcher(index).search(
table_b, index.get_full_query_spec(), BLOCK_LIMIT).cache()
candidates.count()
fvs = featurize(features=features, A=table_a, B=table_b, candidates=candidates,
fill_na=0.0).cache()
rows = fvs.select("id1", "id2", "feature_vectors").collect()
X = np.array([r["feature_vectors"] for r in rows], dtype=float)
y = np.array([1 if (int(r["id1"]), int(r["id2"])) in gold_pairs else 0 for r in rows])
model = SKLearnModel(
XGBClassifier(eval_metric="logloss", objective="binary:logistic",
max_depth=6, seed=42, n_jobs=1).fit(X, y),
nan_fill=0.0)
print(f" trained on {len(y)} pairs ({int(y.sum())} positive)")
# ---- 1. publish v1, crash-recoverable ---------------------------------------
print("2. Publishing serving bundle v1 (crash-recoverable)...")
publish_serving_artifacts(
WORK / "bundle_v1", indexed_table=table_a, features=features, model=model,
sparkly_index_path=WORK / "index", fill_na=0.0, block_limit=BLOCK_LIMIT,
id_col="_id", spark=spark,
enable_crash_recovery=True, checkpoint_dir=WORK / "ckpt_v1")
records = [r.asDict() for r in table_b.collect()]
# ---- 2. real-time matching, tracked -----------------------------------------
print("3. Real-time matching (tracked throughput)...")
realtime = RealtimeMatcher.load(WORK / "bundle_v1", threshold=0.5, track=True)
realtime.match_batch(records[:64]) # one warm batch
for record in records:
realtime.match(record) # throughput counter climbs
realtime.close()
# ---- 3. micro-batch streaming, tracked --------------------------------------
# The wrapped matcher is untracked so the queue owns the throughput view.
print("4. Micro-batch streaming (tracked throughput + batches/records)...")
stream_matcher = RealtimeMatcher.load(WORK / "bundle_v1", threshold=0.5)
with MicroBatchMatcher(stream_matcher, max_batch=16, max_delay_ms=8,
track=True, job_id="microbatch-demo") as queue:
handles = [queue.submit(record) for record in records]
for handle in handles:
handle.result()
stream_matcher.close()
# ---- 4. publish v2 + hot-swap, tracked --------------------------------------
print("5. Publishing v2 + hot-swap (tracked delta lifecycle)...")
publish_serving_artifacts(
WORK / "bundle_v2", indexed_table=table_a, features=features, model=model,
sparkly_index_path=WORK / "index", fill_na=0.0, block_limit=BLOCK_LIMIT,
id_col="_id", spark=spark,
enable_crash_recovery=True, checkpoint_dir=WORK / "ckpt_v2")
live = HotSwappableMatcher(WORK / "bundle_v1", threshold=0.5, track=True)
print(f" serving {Path(live.version).name}; swapping to v2...")
live.swap(WORK / "bundle_v2")
live.match(records[0])
print(f" now serving {Path(live.version).name}")
live.close()
print(f"\n All four jobs are on the dashboard: {session.url}")
try:
input(" Press Enter to stop the dashboard and exit.\n")
except EOFError:
time.sleep(2) # non-interactive run: pause so the jobs are visible
session.stop()
spark.stop()
scripts/demos/realtime/demo_realtime_msd_fusion.py¶
#!/usr/bin/env python
"""Real-time fusion matching on MSD (Million Song Dataset): the full stack, at scale.
sparkly index + semantic index (Qwen-8B vectors) + matchflow feature state
-> published serving bundle
-> per record: fuse sparkly + semantic blocking, featurize (with the semantic
cosine as the last feature), predict, in-process with no Spark on the
request path.
MSD is a 1M-row self-join. Embedding 1M songs live with an 8B model is infeasible, so
the corpus reuses the PRECOMPUTED Qwen3-Embedding-8B (ditto) vectors shipped alongside
`music_dirty`. Because it is a self-join each query song's vector is precomputed too,
so it is passed straight in as `probe_embedding`.
Requests are served through a `ProcessMatcherPool`: worker PROCESSES, not threads,
because featurize is GIL-bound.
<venv>/bin/python scripts/demos/realtime/demo_realtime_msd_fusion.py
Memory, measured on an 18-core / 48 GB box (peak RSS of the whole process tree):
corpus peak Spark driver JVM Spark python workers
30,000 13.3 GB 5.1 GB 7.3 GB
120,000 16.6 GB 7.2 GB 8.5 GB
400,000 19.3 GB 9.9 GB 8.5 GB
The serving pool is NOT the hog: the bundle's big artifacts are memory-mapped, so the
worker processes share one copy. Two other things are, and they behave differently:
* Spark's PYTHON WORKERS, because each holds an Arrow batch and a batch of 4096-d
vectors is large. Bounded by MM_SPARK_CORES + MM_ARROW_BATCH below, after which
they PLATEAU (8.5 GB from 120k on). Left at Spark's defaults (`local[*]`, 10,000
records/batch = ~164 MB per worker) this measured 19 GB at only 30,000 rows.
* The DRIVER JVM, which keeps climbing with corpus size and is the binding
constraint at scale. It is where the sparkly index build's terminal merge gathers
segments, so it grows with the index. At 400,000 rows it is already 9.9 GB against
a 12 GB default heap, so a full 1M run needs MM_SPARK_DRIVER_MEM raised (24g+) and
is the point where this demo runs out of memory if you do not.
Env knobs:
MM_MSD_DIR MSD dir (default ~/Dropbox/MadMatcher/sparkly-data/music_dirty)
MM_MSD_EMB precomputed embedding chunks dir
MM_MSD_LIMIT corpus size cap; 0 = all 1M (default 30000)
MM_MSD_QUERIES query songs to stream; 0 = all in the corpus (default 500)
MM_MSD_TRAIN query songs used to train the matcher (default 800)
MM_MSD_WORKERS worker processes (default min(6, CPU count) -- each costs a JVM
+ embedder, so raise it against measured memory, not core count)
MM_SPARK_CORES Spark local-mode cores (default min(8, CPU count)) -- each core's
Python worker holds an Arrow batch, so this bounds memory
MM_ARROW_BATCH records per Arrow batch (default 512). At 4096-d, Spark's default
of 10,000 is ~164 MB per worker
MM_COLLECT_MAX largest query set materialized in the driver (default 20000);
above it the query stream is read lazily instead
MM_SPARK_DRIVER_MEM driver heap (default 12g); raise for a large corpus
MM_DASHBOARD_PORT dashboard port (default 4050); MADMATCHER_DASHBOARD=0 silences it
Requires the `realtime` license entitlement and the `semantic` extra.
"""
import os
import sys
import tempfile
import time
import warnings
from concurrent.futures import FIRST_COMPLETED, wait
from pathlib import Path
_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)
os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
# Driver memory (local mode = the executor too). Must be set before the JVM starts.
os.environ.setdefault("SPARK_DRIVER_MEMORY",
os.environ.get("MM_SPARK_DRIVER_MEM", "12g"))
warnings.filterwarnings("ignore")
import numpy as np
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 SKLearnModel, create_features, featurize
from madmatcher_pro.reliability.dashboard.server import shutdown_shared
from madmatcher_pro.semantic import SemanticIndex
from madmatcher_pro.serving import ProcessMatcherPool, publish_serving_artifacts
from madmatcher_pro.sparkly.index.lucene_index import LuceneIndex
from madmatcher_pro.sparkly.index_config import IndexConfig
from madmatcher_pro.sparkly.search import Searcher
# This demo serves through a ProcessMatcherPool, and Python's `spawn` start method
# re-imports the main module in every worker. Without this guard each worker would
# re-run the whole pipeline (rebuilding the indexes, retraining the matcher). The body
# below is still a linear script; the guard is mandatory, not a style choice.
if __name__ == "__main__":
# ---- setup ------------------------------------------------------------------
_SD = Path("~/Dropbox/MadMatcher/sparkly-data").expanduser()
MSD_DIR = Path(os.environ.get("MM_MSD_DIR", str(_SD / "music_dirty")))
EMB_DIR = Path(os.environ.get("MM_MSD_EMB", str(
_SD / "_embeddings" / "music_dirty" / "table_a__8b-ditto" / "chunks")))
LIMIT = int(os.environ.get("MM_MSD_LIMIT", "30000")) # 0 => all 1M
N_QUERIES = int(os.environ.get("MM_MSD_QUERIES", "500")) # 0 => all in the corpus
N_WORKERS = int(os.environ.get("MM_MSD_WORKERS", str(min(6, os.cpu_count() or 4))))
N_TRAIN = int(os.environ.get("MM_MSD_TRAIN", "800"))
COLS = ["title", "release", "artist_name", "duration", "year"] # incl. numeric
EMB_FIELDS = ["title", "release", "artist_name"] # the fields the Qwen vectors used
BLOCK_LIMIT = 20
NPROBE = 16
THRESHOLD = 0.5
# Spark local-mode parallelism and Arrow batch size. Both bound memory: every
# Spark Python worker holds an Arrow batch, and a batch of 4096-d vectors is large.
SPARK_CORES = int(os.environ.get("MM_SPARK_CORES", str(min(8, os.cpu_count() or 4))))
ARROW_BATCH = int(os.environ.get("MM_ARROW_BATCH", "512"))
WORK = Path(tempfile.mkdtemp(prefix="mm_msd_fusion_"))
if not (MSD_DIR / "table_a.parquet").exists():
sys.exit(f"MSD not found at {MSD_DIR} (set MM_MSD_DIR).")
if not EMB_DIR.exists():
sys.exit(f"Precomputed embeddings not found at {EMB_DIR} (set MM_MSD_EMB). "
"This demo reuses the Qwen-8B vectors; it does not embed 1M live.")
# local[*] would use every core, and each Spark Python worker holds an Arrow
# batch. With 4096-d vectors the DEFAULT batch (10,000 records) is ~164 MB per
# worker, so on a many-core box the workers alone dominate memory. Bound both the
# core count and the batch size: measured peak on an 18-core box drops from
# ~19 GB to well under half with no loss of throughput at this scale.
spark = (SparkSession.builder.master(f"local[{SPARK_CORES}]")
.config("spark.sql.execution.arrow.pyspark.enabled", "true")
.config("spark.sql.execution.arrow.maxRecordsPerBatch", str(ARROW_BATCH))
.config("spark.driver.memory", os.environ["SPARK_DRIVER_MEMORY"])
.appName("realtime MSD fusion").getOrCreate())
spark.sparkContext.setLogLevel("ERROR")
shutdown_shared() # clear a stale in-process dashboard from a previous run
session = MadMatcherSession.builder.port(
int(os.environ.get("MM_DASHBOARD_PORT", "4050"))).getOrCreate()
if session.url:
print(f"\n Dashboard: {session.url}\n")
# ---- load data --------------------------------------------------------------
print(f"1. Loading MSD ({'all 1M' if LIMIT == 0 else f'first {LIMIT}'} songs) "
f"+ precomputed Qwen-8B vectors...")
corpus = spark.read.parquet(str(MSD_DIR / "table_a.parquet"))
vectors = spark.read.parquet(str(EMB_DIR)) # (_id, embedding[4096])
gold = spark.read.parquet(str(MSD_DIR / "gold.parquet")).select("id1", "id2")
if LIMIT:
corpus = corpus.where(F.col("_id") < LIMIT)
vectors = vectors.where(F.col("_id") < LIMIT)
gold = gold.where((F.col("id1") < LIMIT) & (F.col("id2") < LIMIT))
corpus = corpus.persist()
# Do NOT persist `vectors`: the 4096-d vectors are ~16 GB at 1M, so caching them in
# the Spark heap OOMs. Re-reading from parquet on demand is fine.
gold_pairs = {tuple(sorted((int(a), int(b)))) for a, b in gold.collect()}
print(f" corpus={corpus.count()} songs, gold pairs in corpus={len(gold_pairs)}")
if not gold_pairs:
sys.exit("No gold pairs in this slice: raise MM_MSD_LIMIT so duplicate songs "
"land in the corpus (recall would be 0 otherwise).")
# ---- build the two indexes --------------------------------------------------
print("2. Building the sparkly index...")
config = IndexConfig(id_col="_id")
for c in EMB_FIELDS:
config.add_field(c, ["3gram"])
index = LuceneIndex(WORK / "sparkly", config, delete_if_exists=True)
index.upsert_docs(corpus, force_distributed=True)
index.init()
print("3. Building the semantic index from the precomputed vectors...")
SemanticIndex(str(WORK / "semantic"), id_col="_id").upsert_docs(vectors)
# ---- train a cosine-featured matcher ----------------------------------------
print(f"4. Training a cosine-featured matcher on {N_TRAIN} sampled songs...")
# Sample songs that HAVE a gold duplicate, so training sees positives.
partner_ids = sorted({a for a, _ in gold_pairs} | {b for _, b in gold_pairs})
train_ids = [int(x) for x in np.random.RandomState(0).choice(
partner_ids, size=min(N_TRAIN, len(partner_ids)), replace=False)]
features = create_features(A=corpus, B=corpus, a_cols=COLS, b_cols=COLS)
train_candidates = Searcher(index).search(
corpus.where(F.col("_id").isin(train_ids)),
index.get_full_query_spec(), BLOCK_LIMIT).persist()
# The slow part of the training featurize is the cosine attach, which would chunk the
# full ~16 GB vector table to score a few thousand training pairs. The cosine only
# needs the vectors for the ids those pairs touch, so restrict the EMBEDDINGS to them.
# Keep the full corpus for the feature TABLES: some features depend on corpus token
# statistics, so a sub-corpus there would skew training against what serving uses.
involved = set(train_ids)
for row in train_candidates.select("id1_list").collect():
involved.update(int(x) for x in (row["id1_list"] or []))
train_vectors = vectors.join(
F.broadcast(spark.createDataFrame([(int(i),) for i in involved], ["_id"])), "_id")
fvs = featurize(features=features, A=corpus, B=corpus, candidates=train_candidates,
fill_na=0.0, a_embeddings=train_vectors,
b_embeddings=train_vectors).cache()
rows = fvs.select("id1", "id2", "feature_vectors").collect()
train_candidates.unpersist()
X = np.array([r["feature_vectors"] for r in rows], dtype=float)
y = np.array([1 if tuple(sorted((int(r["id1"]), int(r["id2"])))) in gold_pairs
and int(r["id1"]) != int(r["id2"]) else 0 for r in rows])
if not (y.sum() and (y == 0).any()):
sys.exit("Not enough labeled positives to train (raise MM_MSD_LIMIT).")
model = SKLearnModel(
XGBClassifier(eval_metric="logloss", objective="binary:logistic",
max_depth=6, seed=42, n_jobs=1).fit(X, y),
nan_fill=0.0)
print(f" trained on {len(y)} pairs ({int(y.sum())} positive), "
f"feature width {X.shape[1]} (incl. cosine)")
# ---- publish the fused bundle -----------------------------------------------
print("5. Publishing the fused serving bundle...")
publish_serving_artifacts(
WORK / "bundle", indexed_table=corpus, features=features, model=model,
sparkly_index_path=WORK / "sparkly", semantic_index_path=WORK / "semantic",
embed_provider=None, # each query supplies its precomputed vector
embed_fields=EMB_FIELDS, cosine_feature=True, nprobe=NPROBE,
block_limit=BLOCK_LIMIT, id_col="_id",
# The corpus vectors as an embedding store, so serving recomputes the EXACT
# cosine for every fused candidate (matching training). Without it, a candidate
# only sparkly found would get fill_na instead of its trained cosine.
embedding_store=vectors, embedding_store_id_col="_id",
embedding_store_col="embedding", fill_na=0.0, spark=spark)
# ---- serve: stream songs through a pool of worker processes ------------------
if N_QUERIES:
query_ids = [int(x) for x in partner_ids[:N_QUERIES]]
stream = corpus.where(F.col("_id").isin(query_ids)).join(vectors, "_id")
scope = f"{len(query_ids)} songs (those with gold duplicates)"
else:
stream = corpus.join(vectors, "_id")
scope = "all corpus songs"
print(f"6. Serving {scope} on {N_WORKERS} worker PROCESSES...")
pool = ProcessMatcherPool(WORK / "bundle", n_procs=N_WORKERS,
load_kwargs={"threshold": THRESHOLD})
CAP = max(4, 4 * N_WORKERS) # bounded in flight: never hold every vector
WARMUP = 3 * N_WORKERS # enough for every process to load its matcher
# For a SMALL query set, materialize first: deserializing 4096-d vectors from Spark
# one at a time in the driver caps the submit rate and would starve the pool, hiding
# its parallelism. But each collected row carries a ~16 KB vector, so a large query
# set must stream instead (the driver feeder then bounds throughput, but it runs).
# Gating on "is there a cap at all" would only stream at exactly MM_MSD_QUERIES=0,
# so raising the query count for a better throughput read would OOM the driver.
COLLECT_MAX = int(os.environ.get("MM_COLLECT_MAX", "20000")) # ~330 MB of vectors
if N_QUERIES and N_QUERIES <= COLLECT_MAX:
source_rows = stream.collect()
else:
source_rows = stream.toLocalIterator()
print(f" (streaming the query set: over {COLLECT_MAX} rows, holding every "
f"4096-d vector in the driver would not fit)")
predicted, served_ids, latencies = set(), set(), []
pending, submitted_at, timing = set(), {}, False
start = time.perf_counter()
for i, row in enumerate(source_rows):
record = row.asDict()
embedding = np.asarray(record.pop("embedding"), dtype=np.float32)
song_id = int(record["_id"])
served_ids.add(song_id)
# key=song_id: the result carries it back as id2, so the result needs no side-map
# to stay attributable. submitted_at is only for the latency figure.
future = pool.submit(record, probe_embedding=embedding, key=song_id)
pending.add(future)
submitted_at[future] = time.perf_counter()
if i + 1 == WARMUP: # drain the warm-up, then start timing
for done_future in pending:
result = done_future.result()
predicted.update(tuple(sorted((int(a), int(b))))
for a, b in zip(result["id1"], result["id2"])
if int(a) != int(b))
pending.clear()
submitted_at.clear()
timing, start = True, time.perf_counter()
elif len(pending) >= CAP:
done, pending = wait(pending, return_when=FIRST_COMPLETED)
for done_future in done:
result = done_future.result()
if timing:
latencies.append(
(time.perf_counter() - submitted_at.pop(done_future)) * 1000.0)
else:
submitted_at.pop(done_future, None)
predicted.update(tuple(sorted((int(a), int(b))))
for a, b in zip(result["id1"], result["id2"])
if int(a) != int(b))
print(f" ...served {len(served_ids)}") if len(served_ids) % 200 < CAP else None
for done_future in pending: # drain the tail
result = done_future.result()
predicted.update(tuple(sorted((int(a), int(b))))
for a, b in zip(result["id1"], result["id2"])
if int(a) != int(b))
wall = time.perf_counter() - start
pool.close()
# ---- stats ------------------------------------------------------------------
# Score against the gold pairs recoverable from the SERVED songs, not all gold in the
# corpus: a capped stream should not be penalised for pairs it never queried.
eval_gold = {p for p in gold_pairs if p[0] in served_ids or p[1] in served_ids}
tp = len(eval_gold & predicted)
precision = tp / len(predicted) if predicted else 0.0
recall = tp / len(eval_gold) if eval_gold else 0.0
f1 = 2 * precision * recall / (precision + recall) if (precision + recall) else 0.0
timed_served = max(0, len(served_ids) - WARMUP)
print(f"\n fused matcher (music_dirty, corrupted by design) "
f"precision {precision:.4f} recall {recall:.4f} f1 {f1:.4f}")
print(f" over {len(served_ids)} served songs / {len(eval_gold)} recoverable gold pairs")
if timed_served and wall > 0:
print(f" {N_WORKERS} processes: {timed_served} songs in {wall:.1f}s = "
f"{timed_served / wall:.0f} req/s (steady state, after warm-up)")
if latencies:
latencies = np.array(latencies)
print(f" per-record ms mean {latencies.mean():.1f} "
f"p50 {np.percentile(latencies, 50):.1f} "
f"p95 {np.percentile(latencies, 95):.1f}\n")
print(f" The serving jobs are on the dashboard: {session.url}")
try:
input(" Press Enter to stop the dashboard and exit.\n")
except EOFError:
time.sleep(2)
session.stop()
spark.stop()