Semantic (dense / embedding blocking)

Dense, embedding-based blocking. The pipeline is create_embeddings → SemanticIndex → SemanticSearcher → reciprocal_rank_fusion; the providers and serializers control how records are embedded. The advanced section at the bottom lists the optional knobs for tuning a SemanticIndex beyond its defaults and the bounded-memory build for very large inputs.

Core pipeline

madmatcher_pro.semantic.create_embeddings(df, fields, provider, out_path, *, id_col='_id', serializer=<madmatcher_pro.semantic.serialize.DittoSerializer object>, out_col='embedding', n_rows=None, available_bytes=None, enable_crash_recovery=False, checkpoint_dir=None, n_groups=10, rows_per_group=None, show_progress=False, validate_input=True, dashboard=None, dashboard_port=None, open_browser=None)

Serialize fields to one text per record, embed once, write the id-keyed vector table to out_path, and return the DataFrame reading it back.

The provider runs in a per-partition kernel on the executors, never on the driver. Writing is part of the call, not a separate step: materializing once lets the index be rebuilt and retuned without re-embedding. Skip this entirely if the records already carry a vector column.

Before any embedding work, the estimated table size (~ row count * provider.dim * 4 bytes) is logged so it can be reviewed. The required size is estimated here; the available disk is the caller’s to supply, since free space across a cluster cannot be detected reliably. Pass available_bytes to have the call raise up front when the estimate will not fit.

Parameters:

• df (DataFrame) – pyspark DataFrame with id_col and the fields columns The records to embed.

• fields (list[str]) – list[str], names present in df Fields serialized together into one text per record.

• provider (EmbeddingProvider | str) – EmbeddingProvider | str The vectorizer, or a model id mapped to a built-in provider.

• out_path (str) – str, a filesystem/object-store URI Where to write the (id, out_col) Parquet table.

• id_col (str) – str The unique record id carried through to the output table.

• serializer (Serializer) – Serializer How fields become text; defaults to Ditto-style [COL]/[VAL] markers.

• out_col (str) – str Name of the written vector column.

• n_rows (int | None) – int | None Row count, used for the size estimate. None counts df (one extra scan); pass it to skip the count.

• available_bytes (int | None) – int | None Usable disk you are committing, in bytes. When given, the call raises before doing any work if the estimate exceeds it. None logs the estimate without enforcing.

• enable_crash_recovery (bool) – bool When True, the embed is resumable: records are split into n_groups disjoint xxhash64(id_col) groups, each group’s vectors written under out_path/group-<g> with a commit marker in checkpoint_dir, and a re-run with the same checkpoint_dir skips committed groups. checkpoint_dir is then required and the arguments below take effect. When False (default) the single-job write below runs, byte-identical to before.

• checkpoint_dir (str | None) – str | None Required when enable_crash_recovery=True: where per-group commit markers and the manifest live (the vectors themselves go under out_path). May be local/hdfs/s3a.

• n_groups (int) – number of disjoint xxhash64(id_col) hash-groups the input is split into for a resumable run (enable_crash_recovery=True; default 10).

• rows_per_group (int | None) – if set, derives n_groups = ceil(num_records / rows_per_group), overriding n_groups; set at most one of the two.

• show_progress (bool) – show a tqdm progress bar in the console (default False).

• validate_input (bool) – when enable_crash_recovery=True (default True), record the input row count and refuse a resume whose count changed.

• dashboard (bool | None) – opt-in progress dashboard for this call. None (default) defers to the active MadMatcherSession or the env/default; True/False overrides. The global opt-out MADMATCHER_DASHBOARD=0 turns it off. See madmatcher_pro.MadMatcherSession.

• dashboard_port (int | None) – port for the dashboard server; scans upward if busy.

• open_browser (bool | None) – open the dashboard in a browser on launch.

Returns:

pyspark DataFrame with columns (id_col, out_col)

One row per record; the input the index and searcher consume.

Return type:

DataFrame

class madmatcher_pro.semantic.SemanticIndex(index_path, *, id_col='_id', embedding_col='embedding', nlist='auto', codec=None, embedding_path=None, quantizer='auto', multi_assign=1, target_dim=None, store=None, codec_train_sample=50000)

Bases: object

A dense (embedding) index over the vectors from create_embeddings – one vector per record. Build it once with upsert_docs(), then search it with SemanticSearcher:

index = SemanticIndex(index_path="work/index", id_col="_id")
index = index.upsert_docs(df=vec_a)          # build and persist
# in a later script, reopen without rebuilding:
index = SemanticIndex.load(index_path="work/index")

The index is built in a single pass, which is right for most data. For inputs too large to build at once, use run_chunked_search() – the build does NOT switch to it automatically. Passing available_bytes to upsert_docs() gives an upfront error if the index would not fit the disk you commit, instead of a failure late in the build.

Parameters:

• index_path (str) – str, a filesystem/object-store URI Where the posting lists, centroids, and metadata are persisted.

• id_col (str) – str Unique int64 record id; identity is by id, never by vector.

• embedding_col (str) – str Column holding one array<float> vector per record (raw or upstream-conditioned).

• nlist (int | str) – int (> 0) | “auto” Cell count; “auto” sizes to ~sqrt(N).

• codec (Codec | None) – Codec | None Storage codec for the posting lists. None is the exact full-vector FlatCodec. A lossy codec (SQ8/PQ) stores compact codes and makes search rerank on the full vectors, which then requires embedding_path.

• embedding_path (str | None) – str | None Path to the id-keyed full-vector table (the rerank source). Required only for a lossy codec; ignored for Flat (the in-cell scan is already exact).

• quantizer (str) – centroid-assignment strategy, one of “auto”|”flat”|”hnsw”. “auto” (default) resolves at build to “flat” (exact O(nlist) scan) for nlist <= 4096 or “hnsw” (sub-linear probe) for larger nlist.

• multi_assign (int) – int, >= 1 Number of nearest cells each record is written to, to hedge boundary cases.

• target_dim (int | None) – int | None, >= 1 MRL prefix length to truncate every vector to on intake (Matryoshka models only). None keeps the native width. When set, the build shrinks + renormalizes the records to target_dim (a much smaller, faster index), records the choice in the metadata, and the searcher then truncates wider query vectors to match – so A and B stay the same width with one setting. Not for non-MRL embeddings (use a PCA/Whiten pretransform).

• store (VectorStore | None) – VectorStore | None Storage backend; None uses direct Parquet under index_path.

• codec_train_sample (int) – rows sampled to fit a lossy codec’s (SQ8/PQ) codebooks (default 50000). The codebook is k-bounded, so this is a memory/quality knob, not a corpus-scale cost; ignored by the exact FlatCodec.

upsert_docs(df, *, enable_crash_recovery=False, checkpoint_dir=None, n_groups=0, rows_per_group=None, validate_input=True, available_bytes=None, dashboard=None, dashboard_port=None, open_browser=None)

Build the index over df’s vector column and persist it under index_path.

Records with no embedding (an all-NaN vector) are dropped first, not indexed. The build normalizes vectors, trains cell centroids on a sample, assigns every record to its nearest cell(s), and writes (cell_id, id, vector) posting lists (one file per cell). Vectors are consumed as given; any conditioning is an upstream step, not done here.

Parameters:

• df (DataFrame) – pyspark DataFrame with id_col and embedding_col The records to index.

• enable_crash_recovery (bool) – bool When True, the build is resumable: centroids + codec are trained once and persisted in checkpoint_dir, then records are assigned one xxhash64(id) group at a time, each group’s postings written with a commit marker, and a re-run skips committed groups. checkpoint_dir is then required and n_groups sets the granularity. When False (default) the single-job build below runs.

• checkpoint_dir (str | None) – str | None Required when enable_crash_recovery=True: where per-group markers, the manifest and the trained centroids/codec live (the postings go under index_path). May be local/hdfs/s3a.

• n_groups (int) – int, > 0 Recovery granularity (enable_crash_recovery=True; defaults to 10 when 0).

• rows_per_group (int | None) – int | None Alternative to n_groups: split into ceil(row_count / rows_per_group) groups. Overrides n_groups when set. Set at most one of the two.

• validate_input (bool) – bool When enable_crash_recovery=True (default True): record the record count and refuse to resume if it changed.

• available_bytes (int | None) – int | None Usable disk you are committing; when given, the build raises before writing if the estimate exceeds it. None logs the estimate without enforcing.

• dashboard (bool | None)

• dashboard_port (int | None)

• open_browser (bool | None)

Returns:

SemanticIndex

Self, built and persisted.

Return type:

SemanticIndex

classmethod load(index_path, *, store=None)

Reopen a built index from its persisted metadata and centroids.

Parameters:

• index_path (str) – str The URI passed at build time.

• store (VectorStore | None) – VectorStore | None Override the storage backend; None reads Parquet directly.

Returns:

SemanticIndex

A ready-to-search index with centroids and config restored.

Return type:

SemanticIndex

class madmatcher_pro.semantic.SemanticSearcher(index, max_cell_queries=2048, max_cell_docs=2048)

Bases: object

Runs dense search for every record in a query table.

Parameters:

• index (SemanticIndex) – SemanticIndex A built or loaded index to search against.

• max_cell_queries (int) – int, > 0 Cap on the queries scored in one cogroup group (see _score_cells). A cell hot on the query side (many queries probing it, as when B is much larger than A, leaving few cells) is split into ceil(q_cell / max_cell_queries) salt buckets so no single task materializes the whole cell. Lower it if the search hits Arrow-socket / executor-memory limits; raise it to cut posting-list replication (docs are replicated once per bucket).

• max_cell_docs (int) – int, > 0 Cap on the docs scored in one cogroup group. A cell hot on the doc side (a large IVF cell, e.g. collapsed centroids) is split into ceil(p_cell / max_cell_docs) doc-chunks so the (q x p) score matrix and per-group Arrow payload stay bounded regardless of doc skew. Lower it for memory safety; raise it to cut query replication (queries are replicated once per doc-chunk). Neither knob changes results, only how the work is partitioned within a cell (the cell geometry is unchanged).

search(search_df, spec, limit, id_col='_id', *, probes=None, enable_crash_recovery=False, checkpoint_dir=None, n_groups=10, rows_per_group=None, validate_input=True, dashboard=None, dashboard_port=None, open_browser=None, exclude_self=False)

Search every record in search_df and return its top-limit candidates.

A query with no embedding (an all-NaN vector) is dropped. A query that finds no candidates (its probed cells hold no records) likewise produces no output row, rather than an empty-array row: a missing query means “no candidate pairs”, which fusion and the matcher treat the same as empty, and this avoids an anti-join back over all of B. Query vectors are consumed as given and normalized (the only transform); conditioning is upstream.

Parameters:

• search_df (DataFrame) – pyspark DataFrame with id_col and the index’s embedding column The query records, in the same vector space the index was built in.

• spec (SemanticQuerySpec) – SemanticQuerySpec Runtime nprobe (and reserved knobs).

• limit (int) – int, > 0 Candidates returned per query.

• id_col (str) – str The query id emitted as id2.

• probes (DataFrame | None) – pyspark DataFrame | None Precomputed (id2, query_vector, cell_id) from assign_probes. When given, the per-query cell assignment is skipped and reused – the partitioned build assigns B to the shared cells once and passes the same probes to every chunk’s search.

• enable_crash_recovery (bool) – bool When True, search is resumable: queries are split into n_groups disjoint xxhash64(id_col) groups, each group’s candidates written with a commit marker, and a re-run with the same checkpoint_dir skips committed groups. checkpoint_dir is then required. When False (default) the single-job search below runs.

• checkpoint_dir (str | None) – str | None Required when enable_crash_recovery=True: where per-group output, markers and the manifest live. May be local/hdfs/s3a.

• n_groups (int) – int, > 0 Recovery granularity (enable_crash_recovery=True).

• rows_per_group (int | None) – int | None Alternative to n_groups: split into ceil(query_count / rows_per_group) groups. Overrides n_groups when set. Set at most one of the two.

• validate_input (bool) – bool When enable_crash_recovery=True (default True): record the query count and refuse to resume if it changed.

• dashboard (bool | None)

• dashboard_port (int | None)

• open_browser (bool | None)

• exclude_self (bool)

Returns:

pyspark DataFrame, schema (id2 long, id1_list array<long>, scores array<float>, search_time float, mm_sem_scores array<float>)

One row per query: candidate ids and their cosine scores, best first.

Return type:

DataFrame

class madmatcher_pro.semantic.SemanticQuerySpec(nprobe=32, rerank_window=200, ef_search=None)

Bases: object

Runtime recall knobs, tunable per search without rebuilding the index.

Parameters:

• nprobe (int)

• rerank_window (int)

• ef_search (int | None)

nprobe

int, >= 1 Cells probed per query. A match in an unprobed cell is invisible, so this sets the recall floor; raise it until recall stops improving.

Type:

int

rerank_window

int, >= limit Reserved for a compressed-codec rerank stage; unused while posting lists store full vectors (the scan is already exact).

Type:

int

ef_search

int | None, >= 1 Reserved for an HNSW-over-centroids coarse quantizer.

Type:

int | None

madmatcher_pro.semantic.reciprocal_rank_fusion(sources, limit, *, kappa=60, dashboard=None, dashboard_port=None, open_browser=None)

Fuse the candidate tables per query and re-truncate to the top-limit.

Score for a candidate is the sum over sources of weight / (kappa + rank_in_source), where rank_in_source is its 1-based position within the source for that query, by scores descending (ties broken by id). Rank comes from scores, not the id1_list order.

Parameters:

• sources (list[CandidateSource]) – list[CandidateSource], len >= 1 The candidate tables to fuse, each with a weight and name.

• limit (int) – int, > 0 Candidates kept per query after fusion.

• kappa (int) – int, > 0 Rank damping constant; larger flattens the influence of rank.

• dashboard (bool | None) – opt-in progress dashboard for this call. None (default) defers to the active MadMatcherSession or the env/default; True/False overrides. The global opt-out MADMATCHER_DASHBOARD=0 turns it off. See madmatcher_pro.MadMatcherSession.

• dashboard_port (int | None) – port for the dashboard server; scans upward if busy.

• open_browser (bool | None) – open the dashboard in a browser on launch.

Returns:

pyspark DataFrame, schema (id2 long, id1_list array<long>, scores array<float>, source array<string>[, mm_sem_scores array<float>])

One row per query: ids ordered by fused score, the fused scores, and which blocker(s) found each candidate. mm_sem_scores (the per-pair semantic cosine, preserved from a semantic source; null where only a lexical source found the pair) is emitted only when a source carried it, so the matcher can reuse it. When more than one semantic source found a pair, the strongest (max) cosine across them is kept.

Return type:

DataFrame

class madmatcher_pro.semantic.CandidateSource(name, candidates, weight=0.5)

Bases: object

One blocker’s candidate table, with a provenance tag and a fusion weight.

Parameters:

• name (str)

• candidates (DataFrame)

• weight (float)

name

str, e.g. “semantic” / “lexical” Provenance label carried into the fused output.

Type:

str

candidates

pyspark DataFrame, schema (id2, id1_list, scores[, …]) Per-query candidates; id1_list[i] pairs with scores[i]. scores is required and defines rank; the id1_list order is not significant. Each id1 must appear at most once per query – a duplicate is ranked and summed twice, double-counting that source.

Type:

pyspark.sql.dataframe.DataFrame

weight

float, 0.0 <= weight <= 1.0 Relative influence in the fused score; weights across sources should sum to 1.0.

Type:

float

madmatcher_pro.semantic.shrink_mrl(df, target_dim, *, embedding_col='embedding', validate=True)

Truncate a Matryoshka embedding column to its leading target_dim and re-L2-normalize, preserving every other column.

Matryoshka (MRL) models pack the most signal in the leading dimensions, so a prefix is a much smaller, faster index and query at little recall cost (e.g. 4096 -> 256 is ~16x less data per query). Only embedding_col is rewritten, so this works whether the embedding sits in a wide record table (alongside the fields the matcher uses) or a standalone (id, embedding) table – a standalone table is just the two-column case. Null/empty vectors pass through unchanged.

Apply the SAME target_dim to BOTH the index side and the query side, or the two vector sets are no longer comparable (SemanticIndex(target_dim=…) records the choice so the searcher can enforce it). Re-normalization matters because a truncated unit vector is no longer unit.

Parameters:

• df (DataFrame) – pyspark DataFrame with embedding_col (array<float>/array<double>)

• target_dim (int) – int, >= 1, <= current width Prefix length to keep.

• embedding_col (str) – str The column to shrink; all others are passed through untouched.

• validate (bool) – bool When True, probe one row and raise if target_dim exceeds the actual width (otherwise F.slice would silently return the full, un-shrunk vector). Pass False inside the engine, where the width is already known, to skip the extra Spark action.

Returns:

pyspark DataFrame

df with embedding_col replaced by its renormalized target_dim-prefix.

Return type:

DataFrame

Embedding providers

Pass one as the provider to create_embeddings.

class madmatcher_pro.semantic.EmbeddingProvider

Bases: ABC

Turns a batch of texts into one vector each. Runs on the executors, per partition.

abstract property dim: int

Returns: int, > 0

Output vector width; lets the index size and validate vectors before reading them.

abstract property is_mrl: bool

Returns: bool

True if the model is Matryoshka (signal concentrated in the leading dimensions), so a vector can be truncated to a prefix without retraining.

abstractmethod embed(texts)

Embed one partition of texts.

Parameters:

texts (Series) – pandas Series[str] The serialized records for this partition.

Returns:

np.ndarray, shape (len(texts), dim), float32

One row vector per text, row-aligned to texts.

Return type:

ndarray

class madmatcher_pro.semantic.SentenceTransformerProvider(model_id, *, batch_size=256, normalize=False, is_mrl=False)

Bases: EmbeddingProvider

A local sentence-transformers (torch) model loaded on each executor (GPU if present).

Parameters:

• model_id (str) – str, a sentence-transformers model name or local path Which model to load.

• batch_size (int) – int, > 0 Encode micro-batch size; trades executor memory for throughput.

• normalize (bool) – bool Whether the model L2-normalizes its output (the index renormalizes regardless).

• is_mrl (bool) – bool Whether this is a Matryoshka model (cannot be inferred; set it if known).

class madmatcher_pro.semantic.FastEmbedProvider(model_id='BAAI/bge-small-en-v1.5', *, batch_size=256)

Bases: EmbeddingProvider

A local ONNX embedding model via FastEmbed; lighter than torch, no GPU required.

Parameters:

• model_id (str) – str, a FastEmbed-supported model name Which ONNX model to load on each executor.

• batch_size (int) – int, > 0 Encode micro-batch size; trades executor memory for throughput.

class madmatcher_pro.semantic.OpenAIEmbeddingProvider(model_id, *, batch_size=512, max_concurrency=8, dimensions=None)

Bases: EmbeddingProvider

A hosted OpenAI embedding model; batches requests with retry/backoff for rate limits.

Parameters:

• model_id (str) – str, e.g. “text-embedding-3-small” Which hosted model to call.

• batch_size (int) – int, > 0 Texts per request.

• max_concurrency (int) – int, > 0 Cap on in-flight requests per executor, to stay under the rate limit.

• dimensions (int | None) – int | None, >= 1 Requested embedding width (Matryoshka shortening; text-embedding-3-* only). When set it IS the output width – dim needs no probe and every request forwards it; None uses the model’s native width, resolved by one cached probe.

Serializers

Control how a record’s fields become the one string that is embedded. Pass one as the serializer to create_embeddings.

class madmatcher_pro.semantic.Serializer

Bases: ABC

Turn a record’s fields into one text column for the encoder. The interface; subclass it.

Implementations build a Spark string Column (no pandas) and must be picklable so they ship to executors with the embed job.

class madmatcher_pro.semantic.DittoSerializer(*, col_token='[COL]', val_token='[VAL]', lower=True)

Bases: Serializer

Default scheme: [COL] <field> [VAL] <value> per field, in a fixed order, with empty fields omitted.

The markers let the encoder attend to field boundaries; omitting empty fields keeps missing values out of the text.

Parameters:

• col_token (str) – str Marker written before each field name.

• val_token (str) – str Marker written before each field value.

• lower (bool) – bool Lowercase the text; most encoders are uncased.

class madmatcher_pro.semantic.PlainConcatSerializer(*, sep=' ', lower=True)

Bases: Serializer

Space-join the field values with no markers. Simpler, but the encoder loses field boundaries.

Parameters:

• sep (str) – str Separator inserted between field values.

• lower (bool) – bool Lowercase the text.

Advanced

Optional. Most users do not need these.

Index compression codecs

Pass an instance as SemanticIndex(codec=...) to store compressed posting vectors (the default keeps full vectors).

class madmatcher_pro.semantic.FlatCodec

Bases: Codec

No compression: store fp32, exact distance. 4 bytes per dimension.

class madmatcher_pro.semantic.SQ8Codec

Bases: Codec

Scalar-quantize to int8 (~4x smaller, near-exact). The default when memory is not tight.

class madmatcher_pro.semantic.PQCodec(*, n_subquantizers, n_bits=8)

Bases: Codec

Product quantization (8-32x smaller, lossy). The most compact option; pair with an upstream OPQ rotation (pretransform.OPQRotate) for lower error – applied before encode, not auto-wired.

Parameters:

• n_subquantizers (int) – int, > 0, divides d Number of subspaces; more subspaces give finer codes and larger code_size.

• n_bits (int) – int, typically 8 Bits per subquantizer codebook index.

Vector stores

Pass an instance as SemanticIndex(store=...) to control where full vectors are kept for the rerank step (the default reads them from Parquet).

class madmatcher_pro.semantic.FileSystemVectorStore(index_path, embedding_path, *, id_col='_id', embedding_col='embedding')

Bases: VectorStore

Parquet posting lists on a filesystem or object store (local, HDFS, S3, GCS, ADLS, DBFS).

fetch_vectors reads the embedding table directly, so this is suited to local/HDFS where random reads are cheap; on an object store use MemmapVectorStore for the rerank path.

Parameters:

• index_path (str) – str, a filesystem/object-store URI Root for the posting lists and the index metadata.

• embedding_path (str) – str, a filesystem/object-store URI The id-keyed embedding table, read for rerank fetch-by-id.

• id_col (str) – str Id column in the embedding table.

• embedding_col (str) – str Vector column in the embedding table.

class madmatcher_pro.semantic.MemmapVectorStore(index_path, embedding_path, *, id_col='_id', embedding_col='embedding', per_node_pull=True)

Bases: FileSystemVectorStore

Rerank backend for an A whose full vectors fit per node. Stages A’s id-keyed embedding table once into a contiguous memmap, then serves random fetch_vectors(ids) from the page cache (no object-store request per id). Posting-list IO is inherited (plain Parquet).

Staging is lazy and internal: the searcher calls warm() on first search, like the centroid broadcast.

Parameters:

• index_path (str) – str Root for the posting lists and metadata (inherited).

• embedding_path (str) – str The id-keyed embedding table to stage.

• id_col (str) – str Id column in the embedding table.

• embedding_col (str) – str Vector column in the embedding table.

• per_node_pull (bool) – bool Reserved: selects a per-node parallel pull over a driver pull + broadcast once the distributed staging path lands; the current warm always stages on one node.

Bounded-memory build (very large inputs)

run_chunked_search builds and searches in bounded memory when the embeddings are too large to index at once. Provide the A/B data as a ChunkSource.

madmatcher_pro.semantic.run_chunked_search(a_source, b_source, *, available_bytes, index_root, checkpoint_dir, limit, nlist='auto', nprobe=32, rerank_window=None, multi_assign=1, id_col='_id', code_budget=None, exclude_self=False, b_mode='auto', n_chunks=None, n_b_chunks=None, max_cell_queries=2048, max_cell_docs=2048, show_progress=False, dashboard=None, dashboard_port=None, open_browser=None)

Partitioned dense search: chunk A, stream B against each chunk, merge per query.

Parameters:

• a_source (ChunkSource) – ChunkSource The indexed side (records or an embedded table).

• b_source (ChunkSource) – ChunkSource The query side; emits id2 from its id_col.

• available_bytes (int) – int, > 0 Disk committed to this task at the staging store (index_root/checkpoint_dir). K is sized to live within it; this is the user’s budget, not total or auto-detected disk.

• index_root (str) – str Scratch root for each chunk’s own sub-index (built then released).

• checkpoint_dir (str) – str StagedCheckpoint root: the running snapshot per chunk, centroids, and resume markers.

• limit (int) – int, > 0 Final candidates per query.

• nlist (int | str) – int | “auto” Global cell count; “auto” ~ sqrt of the number of vectors in A.

• nprobe (int) – int, >= 1 Cells probed per query.

• rerank_window (int | None) – int | None w, the per-query candidates kept per chunk through the merge; defaults to limit (the smallest exact value). Must be >= limit.

• exclude_self (bool) – bool Drop the id1 == id2 self-match (deduping a table against itself). Exact: each chunk keeps top-(w+1), so removing the single self leaves the true top-w non-self.

• multi_assign (int) – int, >= 1 Cells each A record is written to.

• id_col (str) – str Id column on both sides.

• code_budget (int | None) – int | None, bytes per stored vector Per-chunk codec budget (select_encoding): None stores full fp32 vectors (FlatCodec, exact in-cell scan). A smaller budget stores compact SQ8/PQ codes, shrinking each chunk’s index so K_A (passes over B) drops; a lossy chunk reranks its code-ranked window on its own full vectors, so each chunk still emits exact-cosine top-w and the merge stays exact. Compression and chunking compose.

• b_mode (str) – “auto” | “persist” | “chunked” persist: embed B once, read each pass. chunked: sweep B in sub-chunks per pass (B too big to store). auto: persist if B’s table fits the budget, else chunked.

• n_chunks (int | None) – int | None K_A override; None derives it from estimate_index_bytes vs available_bytes.

• n_b_chunks (int | None) – int | None K_B override for chunked B; None derives it from B’s table estimate.

• max_cell_queries (int) – int, > 0 Per-cell query cap for scoring (threaded to each chunk’s SemanticSearcher). B is the large side here, so a hot cell is likely; lower this if scoring hits Arrow-socket / executor-memory limits. Partitioning-only; results unchanged.

• max_cell_docs (int) – int, > 0 Per-cell doc cap for scoring (threaded to each chunk’s SemanticSearcher); bounds a doc-heavy cell. Partitioning-only; results unchanged.

• show_progress (bool) – show a tqdm progress bar in the console (default False).

• dashboard (bool | None) – opt-in progress dashboard for this call. None (default) defers to the active MadMatcherSession or the env/default; True/False overrides. The global opt-out MADMATCHER_DASHBOARD=0 turns it off. See madmatcher_pro.MadMatcherSession.

• dashboard_port (int | None) – port for the dashboard server; scans upward if busy.

• open_browser (bool | None) – open the dashboard in a browser on launch.

Returns:

pyspark DataFrame (id2 long, id1_list array<long>, scores array<float>, search_time float)

Identical to the single-index searcher’s contract.

Return type:

DataFrame

class madmatcher_pro.semantic.EmbeddedSource(df, *, id_col='_id', embedding_col='embedding')

Bases: ChunkSource

A side already embedded as an (id, embedding) DataFrame; chunks are cheap filters.

Parameters:

• df (pyspark DataFrame) – the embedded side; only id_col and embedding_col are kept.

• id_col (str) – the unique record id column.

• embedding_col (str) – the array<float> vector column.

class madmatcher_pro.semantic.LazyEmbedSource(records, fields, provider, work_dir, *, id_col='_id', embedding_col='embedding')

Bases: ChunkSource

A side given as raw records; chunk k is embedded on demand and written under work_dir, so the full embedding table never exists at once. Reuses embed.create_embeddings.

Parameters:

• records (pyspark DataFrame) – the raw records with id_col and the fields columns.

• fields (list[str]) – field names serialized together into one text per record for embedding.

• provider (EmbeddingProvider | str) – the vectorizer, or a model id resolved to a built-in provider.

• work_dir (str) – scratch root for the on-demand embedding writes (sample and whole-table paths).

• id_col (str) – the unique record id column carried through embedding.

• embedding_col (str) – name of the written vector column.