`visual-search/load_pg_fingerprints.py`

105 lines
#!/usr/bin/env python3
"""Populate product_image_fingerprints (LOCAL Mac2 dw_unified mirror) from the existing CLIP
fingerprint store (image_embeddings — the table search_service.py loads into RAM). TK-12090 Lane P.

NOTE: vc_ids can be NEGATIVE (embed_shopify_gap.py placeholder keys for shopify-only actives,
~27k rows) — the scan cursor starts at bigint MIN and skips already-loaded rows by anti-join.

Set-based, idempotent (ON CONFLICT (source_vc_id, model_tag) DO NOTHING), resumable by vc_id.
Adds what image_embeddings lacks per the handbag-matcher design: shopify_product_id linkage,
product_status for SQL pre-filtering, model_tag versioning, created_at.

  DW_UNIFIED_DB=<uri> python3 load_pg_fingerprints.py --limit 5000            # sample run
  DW_UNIFIED_DB=<uri> python3 load_pg_fingerprints.py --vc-ids ids.txt        # explicit rows
  DW_UNIFIED_DB=<uri> python3 load_pg_fingerprints.py --full                  # all ~336k (disk-guarded)

A full load duplicates ~1 GB of vectors; it REFUSES unless the data volume has >= 3x the
estimated size free (Mac2 ran at 99% full on 2026-09-23). Undo: DROP TABLE product_image_fingerprints.
"""
import argparse, os, shutil, sys, time
import psycopg2

DB = os.environ.get("DW_UNIFIED_DB") or "postgresql://%s@/dw_unified?host=/tmp" % (
    os.environ.get("PGUSER") or os.environ.get("USER") or "")
# The catalog was embedded with open_clip ViT-B-32/laion2b + dw_clip_ft.pt when present
# (embed_catalog.py / search_service._get_model). Tag the weights so a re-embed never mixes spaces.
HERE = os.path.dirname(os.path.abspath(__file__))
MODEL_TAG = os.environ.get("FP_MODEL_TAG") or (
    "open_clip:ViT-B-32/laion2b_s34b_b79k" + ("+dw_clip_ft" if os.path.exists(os.path.join(HERE, "dw_clip_ft.pt")) else ""))

INSERT = """
INSERT INTO product_image_fingerprints
  (source_vc_id, shopify_product_id, dw_sku, mfr_sku, vendor_code, pattern_name, image_url,
   product_status, model_tag, dim, embedding)
SELECT e.vc_id,
       NULLIF(regexp_replace(sp.shopify_id, '\\D', '', 'g'), '')::bigint,
       e.dw_sku, e.mfr_sku, e.vendor_code, e.pattern_name, e.image_url,
       sp.status, %(tag)s, COALESCE(e.dim, 512), e.embedding
FROM image_embeddings e
LEFT JOIN LATERAL (
  SELECT s.shopify_id, s.status FROM shopify_products s
  WHERE s.dw_sku = e.dw_sku AND e.dw_sku IS NOT NULL
  ORDER BY (s.status = 'ACTIVE') DESC, s.id DESC LIMIT 1) sp ON true
WHERE e.embedding IS NOT NULL AND {where}
ON CONFLICT (source_vc_id, model_tag) DO NOTHING
"""

def main():
    ap = argparse.ArgumentParser()
    g = ap.add_mutually_exclusive_group(required=True)
    g.add_argument("--limit", type=int)
    g.add_argument("--vc-ids")
    g.add_argument("--full", action="store_true")
    ap.add_argument("--batch", type=int, default=20000)
    ap.add_argument("--check-disk", action="store_true", help="with --full: run only the disk guard, load nothing")
    a = ap.parse_args()
    conn = psycopg2.connect(DB); conn.autocommit = False; cur = conn.cursor()
    cur.execute("select to_regclass('product_image_fingerprints')")
    if cur.fetchone()[0] is None:
        sys.exit("product_image_fingerprints missing — apply sql/product_image_fingerprints.sql first")
    t0 = time.time(); inserted = 0
    if a.vc_ids:
        ids = [int(x) for x in open(a.vc_ids).read().split() if x.strip().lstrip("-").isdigit()]   # vc_ids may be negative
        cur.execute(INSERT.format(where="e.vc_id = ANY(%(ids)s)"), {"tag": MODEL_TAG, "ids": ids})
        inserted = cur.rowcount; conn.commit()
    else:
        if a.full:
            cur.execute("select pg_total_relation_size('image_embeddings'), current_setting('data_directory')")
            est, datadir = cur.fetchone()
            du = shutil.disk_usage(datadir)
            # Need room for table + WAL churn (~3x est) AND must leave >= 10% of the volume free
            # afterwards. (A plain `free < 3*est` test PASSED at 99% full on 2026-09-23 — 9 GB free
            # vs a 1 GB table — and pushed the disk to 3.4 GB free. Headroom is judged vs the VOLUME.)
            after = du.free - 3 * est
            if after < 0.10 * du.total:
                sys.exit("REFUSING --full: %.1f GB free, needs ~%.1f GB, would leave %.1f GB < 10%% of %.0f GB volume (%s)"
                         % (du.free / 1e9, 3 * est / 1e9, after / 1e9, du.total / 1e9, datadir))
            if a.check_disk:
                print("disk guard OK: %.1f GB would remain" % (after / 1e9)); return
        # Cursor starts at bigint MIN every run (vc_ids can be negative) and skips rows already
        # loaded via an anti-join — so an earlier out-of-order --vc-ids load can't make a max()
        # cursor jump past unloaded rows.
        last = -9223372036854775808
        remaining = a.limit if a.limit else None
        while remaining is None or remaining > 0:
            n = min(a.batch, remaining) if remaining is not None else a.batch
            cur.execute("select e.vc_id from image_embeddings e where e.vc_id > %s and e.embedding is not null "
                        "and not exists (select 1 from product_image_fingerprints f where f.source_vc_id = e.vc_id and f.model_tag = %s) "
                        "order by e.vc_id limit %s", (last, MODEL_TAG, n))
            ids = [r[0] for r in cur.fetchall()]
            if not ids: break
            cur.execute(INSERT.format(where="e.vc_id = ANY(%(ids)s)"), {"tag": MODEL_TAG, "ids": ids})
            inserted += cur.rowcount; conn.commit(); last = ids[-1]
            if remaining is not None: remaining -= len(ids)
            print("  batch -> vc_id<=%s inserted_total=%d" % (last, inserted), flush=True)
    cur.execute("select count(*), count(shopify_product_id), count(*) filter (where product_status='ACTIVE'), "
                "pg_size_pretty(pg_total_relation_size('product_image_fingerprints')) "
                "from product_image_fingerprints where model_tag=%s", (MODEL_TAG,))
    total, linked, active, size = cur.fetchone()
    print("model_tag=%s inserted=%d table_rows=%d linked_shopify=%d active=%d size=%s in %.1fs"
          % (MODEL_TAG, inserted, total, linked, active, size, time.time() - t0))
    conn.close()

if __name__ == "__main__":
    main()