Update app.py

app.py

@@ -5,6 +5,9 @@ import socket
 import threading
 import re
 import requests
+import pyarrow.parquet as pq
+import pyarrow as pa
+import gc
 from pathlib import Path
 from huggingface_hub import HfApi
 
@@ -16,6 +19,7 @@ STATE_FILE     = "/data/state.json"
 WORKER_TIMEOUT = 600
 MAX_BUFFERED   = 4
 CC_PREFIX      = "data/CC-MAIN-2025"
+ROWS_PER_CHUNK = 50_000
 
 os.makedirs(RAW_DIR, exist_ok=True)
 api = HfApi(token=HF_TOKEN)
@@ -89,6 +93,39 @@ def reclaim_stale(state):
     if reclaimed:
         save_state(state)
 
+# ── Split parquet into chunks ─────────────────────────────────────────────────
+def split_parquet(src_path, name):
+    pf = pq.ParquetFile(src_path)
+    chunk_paths = []
+    chunk_idx   = 0
+    current     = []
+
+    for batch in pf.iter_batches(batch_size=10_000, columns=["text"]):
+        current.append(batch)
+        if sum(len(b) for b in current) >= ROWS_PER_CHUNK:
+            chunk_name = name.replace(".parquet", f"_chunk{chunk_idx:03d}.parquet")
+            chunk_path = Path(RAW_DIR) / chunk_name
+            table = pa.Table.from_batches(current)
+            pq.write_table(table, chunk_path)
+            print(f"    ✓ {chunk_name} ({len(table):,} rows)")
+            chunk_paths.append(chunk_name)
+            chunk_idx += 1
+            current = []
+            del table
+            gc.collect()
+
+    if current:
+        chunk_name = name.replace(".parquet", f"_chunk{chunk_idx:03d}.parquet")
+        chunk_path = Path(RAW_DIR) / chunk_name
+        table = pa.Table.from_batches(current)
+        pq.write_table(table, chunk_path)
+        print(f"    ✓ {chunk_name} ({len(table):,} rows)")
+        chunk_paths.append(chunk_name)
+        del table
+        gc.collect()
+
+    return chunk_paths
+
 # ── Download loop ─────────────────────────────────────────────────────────────
 def download_loop(state):
     base_url = f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/"
@@ -137,22 +174,35 @@ def download_loop(state):
             with open(tmp_path, "wb") as f:
                 for chunk in resp.iter_content(chunk_size=8 * 1024 * 1024):
                     f.write(chunk)
-            tmp_path.rename(raw_path)  # ← atomic, only visible when complete
+            tmp_path.rename(raw_path)
         except Exception as e:
-            print(f"  ✗ Failed: {e} — retrying in 30s")
-            tmp_path.unlink(missing_ok=True)  # ← clean up partial
+            print(f"  ✗ Download failed: {e} — retrying in 30s")
+            tmp_path.unlink(missing_ok=True)
             time.sleep(30)
             continue
 
+        print(f"  Splitting: {name}")
+        try:
+            chunk_names = split_parquet(raw_path, name)
+        except Exception as e:
+            print(f"  ✗ Split failed: {e} — retrying in 30s")
+            raw_path.unlink(missing_ok=True)
+            time.sleep(30)
+            continue
+
+        raw_path.unlink(missing_ok=True)
+
         state["queue"].pop(0)
-        state["shards"][name] = {
-            "status":     "pending",
-            "hf_path":    hf_path,
-            "worker":     None,
-            "claimed_at": None,
-        }
+        for chunk_name in chunk_names:
+            state["shards"][chunk_name] = {
+                "status":     "pending",
+                "hf_path":    hf_path,
+                "worker":     None,
+                "claimed_at": None,
+                "error":      None,
+            }
         save_state(state)
-        print(f"  ✓ Ready: {name}")
+        print(f"  ✓ {len(chunk_names)} chunks ready from {name}")
         time.sleep(5)
 
 # ── Monitor ───────────────────────────────────────────────────────────────────