"""Read bounded decision blocks and enforce split isolation with a disk index.""" from collections.abc import Sequence import json import math from pathlib import Path import random import sqlite3 from .atomic import atomic_output from .inputs import validate_text from .provenance import fingerprint, sha256 AUDIT_FIELDS = ("source_id", "group_id", "recording_id") MAX_BLOCK_BYTES = 64 * 1024 * 1024 MAX_BLOCK_ROWS = 2048 INDEX_VERSION = 2 def validate_row(row, location): validate_text(row.get("question"), row.get("criteria")) if not isinstance(row.get("state"), dict): raise ValueError(f"{location} needs a state object") target = row.get("target") if isinstance(target, bool) or not isinstance(target, int) or not 0 <= target < len(row["criteria"]): raise ValueError(f"{location} has an invalid target index") for name in AUDIT_FIELDS: if not isinstance(row.get(name), str) or not row[name].strip(): raise ValueError(f"{location} requires audit field {name}") class _Rows(Sequence): def __init__(self, dataset): self.dataset = dataset def __len__(self): return self.dataset.count def __iter__(self): for block in self.dataset._blocks(): yield from self.dataset._read_block(block) def __getitem__(self, index): if isinstance(index, slice): return [self[i] for i in range(*index.indices(len(self)))] if index < 0: index += len(self) if not 0 <= index < len(self): raise IndexError(index) block = self.dataset.db.execute( 'SELECT * FROM blocks WHERE row_start<=? ORDER BY row_start DESC LIMIT 1', (index,)).fetchone() return self.dataset._read_block(block)[index-block["row_start"]] class ChoiceDataset: """Keep waveform rows on disk; hold at most one bounded block and one batch.""" def __init__(self, path): self.path = Path(path).resolve() self.row_indices = None self.manifest_snapshot = None if self.path.is_dir(): self.files = sorted(self.path.rglob("*.parquet")) index = self.path / ".paya-index.sqlite3" elif self.path.suffix == ".json": manifest = json.loads(self.path.read_text(encoding="utf-8")) base = Path(manifest.get("path_base", self.path.parent)) self.files = [(base / p).resolve() for p in manifest["files"]] info = self.path.stat() self.manifest_snapshot = (info.st_size, info.st_mtime_ns) if "row_indices" in manifest: self.row_indices = {} for name, groups in manifest['row_indices'].items(): source = str((base/name).resolve()) if Path(source) not in self.files or not isinstance(groups,dict): raise ValueError('Row selection must refer to a declared source file') self.row_indices[source] = {} for group, values in groups.items(): if str(int(group)) != group or int(group)<0 or not isinstance(values,list): raise ValueError('Invalid row-group selection') if any(isinstance(v,bool) or not isinstance(v,int) or v<0 for v in values) or len(set(values))!=len(values): raise ValueError('Row positions must be distinct nonnegative integers') self.row_indices[source][int(group)] = values index = self.path.with_suffix(".paya-index.sqlite3") else: self.files = [self.path] index = self.path.with_name(self.path.name + ".paya-index.sqlite3") if not self.files or len(set(self.files)) != len(self.files): raise ValueError("Dataset must name nonempty, distinct input files") if any(p.suffix not in (".jsonl", ".parquet") for p in self.files): raise ValueError("Dataset files must be JSONL or Parquet") if self.row_indices is not None and any(p.suffix != '.parquet' for p in self.files): raise ValueError('Row selections require bounded Parquet sources') self.snapshots = [(str(p), p.stat().st_size, p.stat().st_mtime_ns) for p in self.files] self.file_stats = {name: (size, modified) for name, size, modified in self.snapshots} identities = [(str(p), sha256(p)) for p in self.files] self.sha256 = identities[0][1] if len(identities) == 1 else fingerprint(identities) if self.row_indices is not None: self.sha256 = fingerprint({'files':identities,'row_indices':self.row_indices}) index = index.with_name(index.stem + "." + self.sha256[:16] + index.suffix) binding = json.dumps({"version": INDEX_VERSION, "files": identities, **({'row_indices':self.row_indices} if self.row_indices is not None else {})}, sort_keys=True) valid = False if index.exists(): try: with sqlite3.connect(index) as check: valid = check.execute("SELECT value FROM metadata WHERE key='binding'").fetchone() == (binding,) except sqlite3.DatabaseError: valid = False if not valid: self._build_index(index, binding) self.index_path = index self.db = sqlite3.connect(index) self.db.row_factory = sqlite3.Row self.db.execute("PRAGMA cache_size=-8192") self.db.execute("PRAGMA temp_store=FILE") self.count, self.block_count = self.db.execute('SELECT SUM(nrows),COUNT(*) FROM blocks').fetchone() if not self.count: raise ValueError("Dataset is empty") self.rows = _Rows(self) def assert_unchanged(self): if self.manifest_snapshot is not None: info = self.path.stat() if (info.st_size,info.st_mtime_ns) != self.manifest_snapshot: raise ValueError('Dataset manifest changed during training') current = [(str(p), p.stat().st_size, p.stat().st_mtime_ns) for p in self.files] if current != self.snapshots: raise ValueError("Dataset changed during training") def _build_index(self, index, binding): with atomic_output(index) as temporary: db = sqlite3.connect(temporary.name) try: db.executescript(''' PRAGMA cache_size=-8192; PRAGMA temp_store=FILE; CREATE TABLE metadata(key TEXT PRIMARY KEY,value TEXT); CREATE TABLE blocks(id INTEGER PRIMARY KEY,filename TEXT,kind TEXT,position INTEGER, nrows INTEGER,row_start INTEGER,byte_length INTEGER); CREATE INDEX block_start ON blocks(row_start); CREATE TABLE audit(field TEXT,value TEXT,PRIMARY KEY(field,value)) WITHOUT ROWID; ''') offset = 0 def add(path, kind, position, rows, byte_length=0): nonlocal offset for i, row in enumerate(rows): validate_row(row, f"{path.name} row {offset+i+1}") db.execute('INSERT INTO blocks(filename,kind,position,nrows,row_start,byte_length) VALUES(?,?,?,?,?,?)', (str(path),kind,position,len(rows),offset,byte_length)) db.executemany('INSERT OR IGNORE INTO audit VALUES(?,?)', ((field,row[field]) for row in rows for field in AUDIT_FIELDS)) offset += len(rows) for path in self.files: if path.suffix == ".parquet": import pyarrow.parquet as pq pf = pq.ParquetFile(path) if "state" not in pf.schema_arrow.names: raise ValueError("Parquet dataset requires state") columns = ["question","criteria","target",*AUDIT_FIELDS] selected = self.row_indices.get(str(path),{}) if self.row_indices is not None else None if selected is not None and any(g>=pf.metadata.num_row_groups for g in selected): raise ValueError('Selected row group does not exist') for group in range(pf.metadata.num_row_groups): if selected is not None and not selected.get(group): continue info = pf.metadata.row_group(group) widths = {"BOOLEAN": 1, "INT32": 4, "INT64": 8, "INT96": 12, "FLOAT": 4, "DOUBLE": 8} numeric_bytes = sum(info.column(i).num_values * widths.get(info.column(i).physical_type,0) for i in range(info.num_columns)) if info.num_rows > MAX_BLOCK_ROWS or max(info.total_byte_size,numeric_bytes) > MAX_BLOCK_BYTES: raise ValueError("Parquet row group exceeds memory budget; repack with 256 rows per group") table = pf.read_row_group(group,columns=columns) if selected is not None: if max(selected[group]) >= len(table): raise ValueError('Selected row position is out of bounds') table = table.take(selected[group]) rows = table.to_pylist() for row in rows: row["state"] = {} if rows: add(path,"parquet",group,rows) else: with path.open("rb") as source: rows=[]; start=source.tell() while True: position=source.tell() line=source.readline(MAX_BLOCK_BYTES+1) if len(line)>MAX_BLOCK_BYTES: raise ValueError("JSONL row exceeds memory budget") if not line: if rows: add(path,"jsonl",start,rows,source.tell()-start) break if rows and (len(rows)>=256 or source.tell()-start>MAX_BLOCK_BYTES): add(path,"jsonl",start,rows,position-start) rows=[];start=position if line.strip(): rows.append(json.loads(line)) if not offset: raise ValueError("Dataset is empty") db.execute("INSERT INTO metadata VALUES('binding',?)",(binding,)) db.commit() finally: db.close() def _blocks(self, seed=None): if seed is None: yield from self.db.execute('SELECT * FROM blocks ORDER BY id') return rng = random.Random(seed) a = rng.randrange(1,self.block_count+1) while math.gcd(a,self.block_count) != 1: a = rng.randrange(1,self.block_count+1) b = rng.randrange(self.block_count) for i in range(self.block_count): identifier = (a*i+b)%self.block_count+1 yield self.db.execute('SELECT * FROM blocks WHERE id=?',(identifier,)).fetchone() def _read_block(self, block): info=Path(block['filename']).stat() if (info.st_size,info.st_mtime_ns) != self.file_stats[block['filename']]: raise ValueError("Dataset changed during training") if block["kind"] == "parquet": import pyarrow.parquet as pq table = pq.ParquetFile(block["filename"]).read_row_group(block["position"]) if self.row_indices is not None: table = table.take(self.row_indices[block['filename']][block['position']]) rows = table.to_pylist() else: with open(block["filename"],"rb") as source: source.seek(block["position"]) rows = [json.loads(line) for line in source.read(block["byte_length"]).splitlines() if line.strip()] if len(rows) != block["nrows"]: raise ValueError("Dataset block changed") return rows def batches(self, batch_size, seed=None, start_batch=0): if isinstance(batch_size,bool) or not isinstance(batch_size,int) or batch_size<1: raise ValueError("batch_size must be a positive integer") if isinstance(start_batch,bool) or not isinstance(start_batch,int) or start_batch<0: raise ValueError("start_batch must be a nonnegative integer") self.assert_unchanged() skip=start_batch*batch_size pending=[] for block in self._blocks(seed): if skip>=block['nrows']: skip-=block['nrows'] continue values=self._read_block(block) order=list(range(len(values))) if seed is not None: random.Random(f'{seed}:block:{block["id"]}').shuffle(order) for position in order[skip:]: row=values[position].copy() if seed is not None: permutation=list(range(len(row['criteria']))) random.Random(f'{seed}:block:{block["id"]}:row:{position}').shuffle(permutation) row['criteria']=[row['criteria'][i] for i in permutation] row['target']=permutation.index(row['target']) pending.append(row) if len(pending)==batch_size: yield pending pending=[] skip=0 if pending: yield pending def assert_disjoint(*datasets): for i, left in enumerate(datasets): for right in datasets[i+1:]: left.db.execute('ATTACH DATABASE ? AS other',(str(right.index_path),)) try: overlap=left.db.execute('SELECT a.field FROM audit a JOIN other.audit b ' 'ON a.field=b.field AND a.value=b.value LIMIT 1').fetchone() if overlap: raise ValueError(f'Dataset splits overlap on {overlap[0]}') finally: left.db.execute('DETACH DATABASE other') def make_batch(model, rows): return model.make_batch([r["state"] for r in rows], [r["question"] for r in rows], [r["criteria"] for r in rows])