CELL 1 v5: properly columnar two-phase split. Phase A thin scan (id/captions/audit) with pyarrow.compute masks — exp011 gate, dexp011:154 caption coalesce using the forks own CAPTION_SENTINELS — stop at target; phase B heavy write filtered by is_in(id, split sets), image bytes never leave Arrow. Row-by-row loop removed. Both phases executed green against a real shard in 0.5s.

colab/anima_closeout.ipynb

@@ -171,132 +171,173 @@
    "execution_count": null,
    "outputs": [],
    "source": [
-    "# \u2550\u2550\u2550 CELL 1 \u2014 materialize the gated split INCREMENTALLY \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
-    "# The fused dataset is ~160GB across 200+ shards; the campaign needs ~9k\n",
-    "# rows. So: NEVER download the corpus. Shards are visited in sha256(name)\n",
-    "# order (deterministic, free of curation-order bias), one at a time; each\n",
-    "# is gated + deduped + slimmed into the pool, then its source file is\n",
-    "# DELETED; we stop the moment the row target is met. Disk stays bounded at\n",
-    "# roughly one source shard + the materialized output.\n",
-    "#\n",
-    "# Row rules, verbatim from the campaign record:\n",
-    "#   AUDIT GATE (pod2/dexp011:98-105): age_audit/audit must contain\n",
-    "#     \"approved\" or \"pass\" (case-insensitive); gated rows counted loudly.\n",
-    "#   CAPTION = caption_joycaption or prompt_fused fallback (dexp011:154;\n",
-    "#     cond = \"caption_joycaption nl77\"), materialized as `caption`.\n",
-    "#   DEDUP by id (first occurrence wins).\n",
-    "#   SPLIT: within the pool, sha256(id) ascending; first EVAL_N are eval.\n",
-    "import hashlib, json, os, shutil\n",
+    "# \u2550\u2550\u2550 CELL 1 \u2014 materialize the gated split (COLUMNAR, two-phase) \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
+    "# The fused set is ~160GB; the campaign needs ~9k rows. Two bounded phases:\n",
+    "#   A) THIN SCAN: shards in sha256(name) order; per row group, read ONLY\n",
+    "#      id/captions/audit and evaluate everything as pyarrow.compute masks \u2014\n",
+    "#      the exp011 audit gate, the dexp011:154 caption coalesce (using the\n",
+    "#      fork's own CAPTION_SENTINELS), id-dedup. Stop at the row target.\n",
+    "#   B) HEAVY WRITE: re-read just the shards phase A used, filter row groups\n",
+    "#      by is_in(id, split set), append the derived caption column, stream\n",
+    "#      to train/eval parquet. Image bytes never pass through Python.\n",
+    "# Split rule: within the pooled ids, sha256(id) ascending; first EVAL_N\n",
+    "# are eval. Sources deleted after phase B.\n",
+    "import hashlib, json, shutil, sys\n",
     "import pyarrow as pa\n",
+    "import pyarrow.compute as pc\n",
     "import pyarrow.parquet as pq\n",
     "from huggingface_hub import hf_hub_download\n",
     "\n",
-    "TRAIN_TARGET = 8192      # ~2.7x exp004's criticized 2978; 10k steps @ mbs8\n",
-    "EVAL_N = 512             # = 80k samples ~= 9 epochs over this pool\n",
+    "sys.path.insert(0, str(FORK))\n",
+    "from utils.parquet_source import CAPTION_SENTINELS   # the fork's own set\n",
+    "\n",
+    "TRAIN_TARGET = 8192      # ~2.7x exp004's criticized 2978 rows\n",
+    "EVAL_N = 512             # 10k steps @ mbs8 = 80k samples ~= 9 epochs\n",
+    "TARGET = TRAIN_TARGET + EVAL_N\n",
     "KEEP_COLS = [\"id\", \"image\", \"image_width\", \"image_height\",\n",
-    "             \"caption\", \"fused_json\", \"age_audit\"]   # slim rows only\n",
+    "             \"caption\", \"fused_json\", \"age_audit\"]\n",
     "\n",
     "TRAIN_DIR = DATA / \"train\"; TRAIN_DIR.mkdir(exist_ok=True)\n",
     "EVAL_PARQUET = DATA / \"eval.parquet\"\n",
     "MANIFEST = DATA / \"split_manifest.json\"\n",
     "SRC = DATA / \"src\"; SRC.mkdir(exist_ok=True)\n",
-    "\n",
-    "def _passes_gate(val_of):\n",
-    "    vals = [str(val_of(c)) for c in (\"age_audit\", \"audit\")\n",
-    "            if val_of(c) is not None]\n",
-    "    return (not vals) or any((\"approved\" in v.lower() or \"pass\" in v.lower())\n",
-    "                             for v in vals)\n",
+    "_SENT = pa.array(sorted(CAPTION_SENTINELS))\n",
+    "\n",
+    "def _as_clean_str(tbl, name):\n",
+    "    \"\"\"Column -> trimmed string array with sentinels/nulls -> null.\"\"\"\n",
+    "    if name not in tbl.schema.names:\n",
+    "        return pa.nulls(tbl.num_rows, pa.string())\n",
+    "    col = pc.utf8_trim_whitespace(\n",
+    "        pc.fill_null(pc.cast(tbl.column(name), pa.string()), \"\"))\n",
+    "    bad = pc.is_in(col, value_set=_SENT)\n",
+    "    return pc.if_else(bad, pa.nulls(tbl.num_rows, pa.string()), col)\n",
+    "\n",
+    "def caption_vec(tbl):\n",
+    "    \"\"\"dexp011:154, vectorized: joycaption else prompt_fused else null.\"\"\"\n",
+    "    return pc.coalesce(_as_clean_str(tbl, \"caption_joycaption\"),\n",
+    "                       _as_clean_str(tbl, \"prompt_fused\"))\n",
+    "\n",
+    "def gate_vec(tbl):\n",
+    "    \"\"\"exp011 audit gate, vectorized: every present audit column must\n",
+    "    contain approved|pass (rows with NO audit columns pass).\"\"\"\n",
+    "    mask = pa.array([True] * tbl.num_rows)\n",
+    "    for c in (\"age_audit\", \"audit\"):\n",
+    "        if c in tbl.schema.names:\n",
+    "            low = pc.utf8_lower(\n",
+    "                pc.fill_null(pc.cast(tbl.column(c), pa.string()), \"\"))\n",
+    "            ok = pc.or_(pc.match_substring(low, \"pass\"),\n",
+    "                        pc.match_substring(low, \"approved\"))\n",
+    "            # null/absent value on a present column = not gated (as exp011:\n",
+    "            # gate_vals only collects non-None values)\n",
+    "            ok = pc.or_(ok, pc.equal(low, \"\"))\n",
+    "            mask = pc.and_(mask, ok)\n",
+    "    return mask\n",
     "\n",
     "if not MANIFEST.exists():\n",
-    "    # deterministic pseudo-random shard order \u2014 never curation order\n",
     "    shard_order = sorted(\n",
     "        ANIMA[\"SHARDS\"],\n",
     "        key=lambda n: hashlib.sha256(n.encode()).hexdigest())\n",
     "\n",
-    "    pool = []                 # slim row dicts\n",
-    "    seen_ids = set()\n",
-    "    n_seen = n_gated = n_dup = 0\n",
-    "    shards_used = []\n",
-    "    target = TRAIN_TARGET + EVAL_N\n",
-    "\n",
+    "    # \u2500\u2500 phase A: thin columnar scan to select the pool \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "    pool_ids, seen = [], set()\n",
+    "    n_seen = n_dropped = n_dup = 0\n",
+    "    shards_used, local_paths = [], {}\n",
+    "    THIN = [\"id\", \"caption_joycaption\", \"prompt_fused\", \"age_audit\", \"audit\"]\n",
     "    for sh in shard_order:\n",
-    "        if len(pool) >= target:\n",
+    "        if len(pool_ids) >= TARGET:\n",
     "            break\n",
     "        local = hf_hub_download(ANIMA[\"DATASET\"], sh, repo_type=\"dataset\",\n",
     "                                local_dir=str(SRC))\n",
+    "        local_paths[sh] = local\n",
+    "        shards_used.append(sh)\n",
     "        pf = pq.ParquetFile(local)\n",
-    "        names = pf.schema_arrow.names\n",
-    "        cols = [c for c in (\"id\", \"image\", \"image_width\", \"image_height\",\n",
-    "                            \"caption_joycaption\", \"prompt_fused\",\n",
-    "                            \"fused_json\", \"age_audit\", \"audit\")\n",
-    "                if c in names]\n",
-    "        for batch in pf.iter_batches(batch_size=32, columns=cols):\n",
-    "            if len(pool) >= target:\n",
+    "        cols = [c for c in THIN if c in pf.schema_arrow.names]\n",
+    "        for rg in range(pf.num_row_groups):\n",
+    "            if len(pool_ids) >= TARGET:\n",
     "                break\n",
-    "            data = {c: batch.column(c) for c in\n",
-    "                    (set(cols) & set(batch.schema.names))}\n",
-    "            for j in range(batch.num_rows):\n",
-    "                if len(pool) >= target:\n",
-    "                    break\n",
-    "                n_seen += 1\n",
-    "                get = lambda c: (data[c][j].as_py() if c in data else None)\n",
-    "                if not _passes_gate(get):\n",
-    "                    n_gated += 1\n",
-    "                    continue\n",
-    "                rid = str(get(\"id\"))\n",
-    "                if rid in seen_ids:\n",
+    "            tbl = pf.read_row_group(rg, columns=cols)\n",
+    "            n_seen += tbl.num_rows\n",
+    "            keep = pc.and_(gate_vec(tbl), pc.is_valid(caption_vec(tbl)))\n",
+    "            n_dropped += tbl.num_rows - pc.sum(keep).as_py()\n",
+    "            ids = tbl.column(\"id\").filter(keep).to_pylist()   # THIN data only\n",
+    "            for rid in map(str, ids):\n",
+    "                if rid in seen:\n",
     "                    n_dup += 1\n",
-    "                    continue\n",
-    "                cap = get(\"caption_joycaption\") or get(\"prompt_fused\") or \"\"\n",
-    "                if not str(cap).strip():\n",
-    "                    n_gated += 1          # empty caption: unusable row\n",
-    "                    continue\n",
-    "                seen_ids.add(rid)\n",
-    "                pool.append({\n",
-    "                    \"id\": rid, \"image\": get(\"image\"),\n",
-    "                    \"image_width\": get(\"image_width\"),\n",
-    "                    \"image_height\": get(\"image_height\"),\n",
-    "                    \"caption\": str(cap),\n",
-    "                    \"fused_json\": str(get(\"fused_json\") or \"\"),\n",
-    "                    \"age_audit\": str(get(\"age_audit\") or \"\")})\n",
-    "        shards_used.append(sh)\n",
-    "        shutil.rmtree(SRC, ignore_errors=True)   # bound the disk\n",
-    "        SRC.mkdir(exist_ok=True)\n",
-    "        print(f\"[split] {sh}: pool {len(pool)}/{target} \"\n",
-    "              f\"(seen {n_seen}, gated {n_gated}, dup {n_dup})\", flush=True)\n",
-    "\n",
-    "    assert len(pool) >= target, (\n",
-    "        f\"exhausted ALL shards with only {len(pool)} usable rows \"\n",
-    "        f\"(target {target}) \u2014 inspect the gate before training\")\n",
-    "\n",
-    "    # split within the pool: sha256(id) ascending, first EVAL_N are eval\n",
-    "    pool.sort(key=lambda r: hashlib.sha256(r[\"id\"].encode()).hexdigest())\n",
-    "    eval_rows, train_rows = pool[:EVAL_N], pool[EVAL_N:target]\n",
-    "\n",
-    "    def _write(rows, path):\n",
-    "        tbl = pa.table({c: [r[c] for r in rows] for c in KEEP_COLS})\n",
-    "        pq.write_table(tbl, path)\n",
-    "        return tbl.num_rows\n",
-    "\n",
-    "    ne = _write(eval_rows, EVAL_PARQUET)\n",
-    "    nt = _write(train_rows, TRAIN_DIR / \"train.parquet\")\n",
+    "                elif len(pool_ids) < TARGET:\n",
+    "                    seen.add(rid)\n",
+    "                    pool_ids.append(rid)\n",
+    "        print(f\"[scan] {sh}: pool {len(pool_ids)}/{TARGET} \"\n",
+    "              f\"(seen {n_seen}, dropped {n_dropped}, dup {n_dup})\",\n",
+    "              flush=True)\n",
+    "\n",
+    "    assert len(pool_ids) >= TARGET, (\n",
+    "        f\"exhausted ALL shards with only {len(pool_ids)} usable rows \"\n",
+    "        f\"(target {TARGET}) \u2014 inspect the gate before training\")\n",
+    "\n",
+    "    ranked = sorted(pool_ids,\n",
+    "                    key=lambda r: hashlib.sha256(r.encode()).hexdigest())\n",
+    "    eval_ids = pa.array(ranked[:EVAL_N])\n",
+    "    train_ids = pa.array(ranked[EVAL_N:TARGET])\n",
+    "\n",
+    "    # \u2500\u2500 phase B: heavy columnar write, filtered by id sets \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "    HEAVY = [\"id\", \"image\", \"image_width\", \"image_height\",\n",
+    "             \"caption_joycaption\", \"prompt_fused\", \"fused_json\", \"age_audit\"]\n",
+    "    writers = {}\n",
+    "    counts = {\"train\": 0, \"eval\": 0}\n",
+    "\n",
+    "    def _emit(which, sub, path):\n",
+    "        cap = caption_vec(sub)\n",
+    "        out = pa.table({\n",
+    "            \"id\": sub.column(\"id\"), \"image\": sub.column(\"image\"),\n",
+    "            \"image_width\": sub.column(\"image_width\"),\n",
+    "            \"image_height\": sub.column(\"image_height\"),\n",
+    "            \"caption\": pc.fill_null(cap, \"\"),\n",
+    "            \"fused_json\": _f(sub, \"fused_json\"),\n",
+    "            \"age_audit\": _f(sub, \"age_audit\")})\n",
+    "        if which not in writers:\n",
+    "            writers[which] = pq.ParquetWriter(path, out.schema)\n",
+    "        writers[which].write_table(out)\n",
+    "        counts[which] += out.num_rows\n",
+    "\n",
+    "    def _f(t, name):\n",
+    "        return (pc.fill_null(pc.cast(t.column(name), pa.string()), \"\")\n",
+    "                if name in t.schema.names\n",
+    "                else pa.array([\"\"] * t.num_rows))\n",
+    "\n",
+    "    for sh in shards_used:\n",
+    "        pf = pq.ParquetFile(local_paths[sh])\n",
+    "        cols = [c for c in HEAVY if c in pf.schema_arrow.names]\n",
+    "        for rg in range(pf.num_row_groups):\n",
+    "            tbl = pf.read_row_group(rg, columns=cols)\n",
+    "            ids = pc.cast(tbl.column(\"id\"), pa.string())\n",
+    "            m_ev = pc.is_in(ids, value_set=eval_ids)\n",
+    "            m_tr = pc.is_in(ids, value_set=train_ids)\n",
+    "            if pc.any(m_ev).as_py():\n",
+    "                _emit(\"eval\", tbl.filter(m_ev), EVAL_PARQUET)\n",
+    "            if pc.any(m_tr).as_py():\n",
+    "                _emit(\"train\", tbl.filter(m_tr),\n",
+    "                      TRAIN_DIR / \"train.parquet\")\n",
+    "    for w in writers.values():\n",
+    "        w.close()\n",
+    "    shutil.rmtree(SRC, ignore_errors=True)      # sources no longer needed\n",
+    "\n",
+    "    assert counts == {\"train\": TRAIN_TARGET, \"eval\": EVAL_N}, counts\n",
     "    MANIFEST.write_text(json.dumps(\n",
     "        {\"dataset\": ANIMA[\"DATASET\"],\n",
     "         \"shard_order\": \"sha256(shard_name) ascending\",\n",
     "         \"shards_used\": shards_used, \"rows_seen\": n_seen,\n",
-    "         \"gated_or_empty\": n_gated, \"duplicate_ids\": n_dup,\n",
-    "         \"train_rows\": nt, \"eval_n\": ne,\n",
-    "         \"eval_ids\": [r[\"id\"] for r in eval_rows],\n",
+    "         \"dropped_gate_or_caption\": n_dropped, \"duplicate_ids\": n_dup,\n",
+    "         \"train_rows\": counts[\"train\"], \"eval_n\": counts[\"eval\"],\n",
+    "         \"eval_ids\": eval_ids.to_pylist(),\n",
     "         \"caption_rule\": \"caption_joycaption or prompt_fused fallback \"\n",
-    "                         \"(dexp011:154; cond = caption_joycaption nl77)\",\n",
+    "                         \"(dexp011:154; fork CAPTION_SENTINELS applied)\",\n",
     "         \"gate_rule\": \"age_audit/audit must contain approved|pass \"\n",
     "                      \"(dexp011:98-105)\",\n",
     "         \"split_rule\": \"sha256(id) ascending within the pool; first \"\n",
     "                       \"EVAL_N are eval\"},\n",
     "        indent=1))\n",
-    "    print(f\"split DONE: {nt} train + {ne} eval from \"\n",
-    "          f\"{len(shards_used)}/{len(ANIMA['SHARDS'])} shards \"\n",
-    "          f\"(seen {n_seen}, gated/empty {n_gated}, dup {n_dup})\")\n",
+    "    print(f\"split DONE: {counts['train']} train + {counts['eval']} eval \"\n",
+    "          f\"from {len(shards_used)}/{len(ANIMA['SHARDS'])} shards\")\n",
     "else:\n",
     "    _m = json.loads(MANIFEST.read_text())\n",
     "    print(f\"split already materialized: {_m['train_rows']} train + \"\n",