Coverage for src/lilbee/data/export.py: 100%
154 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Per-page text dataset: build/write from a store, and import one back."""
3from __future__ import annotations
5import asyncio
6import json
7import threading
8from dataclasses import dataclass
9from enum import StrEnum
10from pathlib import Path
11from typing import IO, TYPE_CHECKING, cast
13from lilbee.data.extract.document import _title_scope, chunk_and_embed_pages
14from lilbee.data.store import ChunkWrite, PageTextRecord, SourceMeta, SourceType
15from lilbee.data.title import derive_title
16from lilbee.runtime.cancellation import TaskCancelledError
17from lilbee.runtime.progress import DetailedProgressCallback, noop_callback
19if TYPE_CHECKING:
20 import pyarrow as pa
22 from lilbee.data.store import Store
25class DatasetFormat(StrEnum):
26 """On-disk format for the per-page text dataset."""
28 PARQUET = "parquet"
29 JSONL = "jsonl"
32@dataclass
33class ImportResult:
34 """Summary of an `import_dataset` run."""
36 sources: list[str]
37 pages: int
38 chunks: int
41def decode_format(value: str) -> DatasetFormat:
42 """Decode an explicit format string, raising a user-facing ``ValueError``."""
43 try:
44 return DatasetFormat(value)
45 except ValueError:
46 raise ValueError(f"Unsupported format: {value!r} (expected parquet or jsonl)") from None
49def resolve_format(value: str, path: Path) -> DatasetFormat:
50 """Pick a format from explicit *value*, else the *path* suffix.
52 Raises ``ValueError`` with a user-facing message when neither yields a
53 known format.
54 """
55 if value:
56 return decode_format(value)
57 suffix = path.suffix.lower().lstrip(".")
58 try:
59 return DatasetFormat(suffix)
60 except ValueError:
61 raise ValueError(
62 f"Could not infer format from {path.name!r}; use a .parquet or .jsonl path"
63 ) from None
66def build_page_dataset(store: Store, source: str | None = None) -> pa.Table:
67 """Collect per-page text rows as an Arrow table for every source (or *source*).
69 Sources captured at ingest are read verbatim from the page-text table in a
70 single columnar scan. Sources without captured text (older indexes, code) are
71 reconstructed from the chunks table; that reconstruction concatenates chunk
72 text per page, so chunk overlap may repeat a little text across page
73 boundaries. The whole set stays in Arrow so the writers avoid per-row Python
74 objects, and the per-source query loop that hung a large export is gone
75 (bb-bqg).
76 """
77 import pyarrow as pa
79 tracked = _with_wide_offsets(store.sources_arrow())
80 if source is not None:
81 table = _with_wide_offsets(store.page_texts_arrow(source))
82 if table.num_rows == 0:
83 table = _reconstructed_arrow(store, [source], table.schema)
84 else:
85 # One scan for every captured page instead of a filtered query per source:
86 # the per-source loop was O(sources) and its fixed per-query overhead, not
87 # data size, hung the export on a large store. The semi-join restricts it
88 # to tracked sources, so an orphaned page-text row (source record gone)
89 # stays out, matching the old get_sources()-scoped universe.
90 keys = tracked.select(["source"])
91 table = _with_wide_offsets(store.page_texts_arrow()).join(
92 keys, keys="source", join_type="left semi"
93 )
94 # Tracked sources the scan found no page text for: older indexes and code,
95 # rebuilt from their chunks. Normally empty, and an anti-join keeps the
96 # comparison in Arrow rather than differencing two sets of every filename.
97 missing = keys.join(table.select(["source"]), keys="source", join_type="left anti")
98 extra = _reconstructed_arrow(
99 store, sorted(missing.column("source").to_pylist()), table.schema
100 )
101 if extra.num_rows:
102 table = pa.concat_tables([table, extra])
103 # Denormalize each source's title/authors/created_at onto its page rows, so an
104 # export/import cycle keeps them instead of falling back to the filename stem.
105 # Left outer: a source with no metadata keeps its pages, with nulls.
106 table = table.join(tracked, keys="source", join_type="left outer")
107 return table.sort_by([("source", "ascending"), ("page", "ascending")])
110def _with_wide_offsets(table: pa.Table) -> pa.Table:
111 """*table* with its string columns retyped to 64-bit offsets.
113 pyarrow's ``string`` addresses a column's data with int32 offsets, capping it
114 at 2GB, and every step below this one (filter, concat, metadata append, sort)
115 materializes one array per column. A corpus whose page text passes 2GB
116 overflows there, so the export widens on the way out of the scan; the store's
117 own schema is untouched. Both writers take the wider type, and parquet records
118 it in its arrow metadata and reads it back as ``large_string``, so the rows an
119 import decodes are ordinary strings either way.
120 """
121 import pyarrow as pa
123 return table.cast(
124 pa.schema(
125 [
126 field.with_type(pa.large_string()) if pa.types.is_string(field.type) else field
127 for field in table.schema
128 ]
129 )
130 )
133def _reconstructed_arrow(store: Store, sources: list[str], schema: pa.Schema) -> pa.Table:
134 """Chunk-reconstructed pages for *sources* as an Arrow table in *schema*."""
135 import pyarrow as pa
137 records = [dict(r) for name in sources for r in _reconstruct_from_chunks(store, name)]
138 return pa.Table.from_pylist(records, schema=schema)
141def _reconstruct_from_chunks(store: Store, source: str) -> list[PageTextRecord]:
142 """Rebuild per-page rows for *source* by joining its chunks per page."""
143 by_page: dict[int, list[tuple[int, str]]] = {}
144 content_type = ""
145 for chunk in store.get_chunks_by_source(source):
146 content_type = chunk.content_type or content_type
147 by_page.setdefault(chunk.page_start, []).append((chunk.chunk_index, chunk.chunk))
148 rows: list[PageTextRecord] = []
149 for page in sorted(by_page):
150 ordered = [text for _, text in sorted(by_page[page])]
151 rows.append(
152 PageTextRecord(
153 source=source, page=page, text="\n".join(ordered), content_type=content_type
154 )
155 )
156 return rows
159# Rows encoded per write. A page row is a few hundred bytes of text, so this
160# keeps a batch in the low megabytes whatever the corpus size.
161_WRITE_BATCH_ROWS = 10_000
164def _write_parquet(table: pa.Table, sink: IO[bytes], cancel: threading.Event | None = None) -> None:
165 """Encode *table* into *sink* as parquet, one row group per batch.
167 Driving a ParquetWriter row group by row group rather than calling
168 write_table is what gives the export a boundary to stop on. For a table with
169 rows the two produce byte-identical output, so the only cost is the loop. A
170 table with no rows writes no row groups and its footer differs, which the
171 export path cannot reach because the build refuses an empty dataset before
172 it gets here; the file still reads back as zero rows either way.
173 """
174 import pyarrow as pa
175 import pyarrow.parquet as pq
177 with pq.ParquetWriter(sink, table.schema) as writer:
178 for batch in table.to_batches(max_chunksize=_WRITE_BATCH_ROWS):
179 if cancel is not None and cancel.is_set():
180 raise TaskCancelledError
181 writer.write_table(pa.Table.from_batches([batch], table.schema))
184def _write_jsonl(table: pa.Table, sink: IO[bytes], cancel: threading.Event | None = None) -> None:
185 """Encode *table* into *sink* as jsonl, one write per batch.
187 One JSON object per row, keyed by the table's columns, so jsonl stays in step
188 with the schema (and with parquet) rather than a hardcoded field list. Only a
189 batch is converted to Python objects at a time: converting the whole table
190 cost about 19x the size of the file it produced (77GB peak for a 4.15GB
191 export of an 8.8M-row corpus), because the row dicts, the joined string and
192 its encoded bytes were all live at once.
193 """
194 for batch in table.to_batches(max_chunksize=_WRITE_BATCH_ROWS):
195 if cancel is not None and cancel.is_set():
196 raise TaskCancelledError
197 sink.write("".join(json.dumps(row) + "\n" for row in batch.to_pylist()).encode("utf-8"))
200_WRITERS = {DatasetFormat.PARQUET: _write_parquet, DatasetFormat.JSONL: _write_jsonl}
203def serialize_dataset(table: pa.Table, fmt: DatasetFormat) -> bytes:
204 """Encode the dataset *table* to bytes in the given format.
206 For callers that must hand back one buffer (the HTTP download). A file
207 export uses :func:`write_dataset`, which never holds the encoded dataset.
208 """
209 import io
211 buffer = io.BytesIO()
212 _WRITERS[fmt](table, buffer, None)
213 return buffer.getvalue()
216def write_dataset(
217 table: pa.Table, path: Path, fmt: DatasetFormat, cancel: threading.Event | None = None
218) -> None:
219 """Write the dataset *table* to *path* in the given format, a batch at a time.
221 A cancelled write removes the partial file: half a dataset is not a smaller
222 dataset, and leaving one behind would be indistinguishable from a complete
223 export to anything that reads it later.
224 """
225 try:
226 with path.open("wb") as sink:
227 _WRITERS[fmt](table, sink, cancel)
228 except TaskCancelledError:
229 path.unlink(missing_ok=True)
230 raise
233def _coerce_row(raw: dict) -> PageTextRecord:
234 """Validate one raw dataset row into a `PageTextRecord`.
236 The denormalized source metadata (title/authors/created_at) is carried
237 through when present so a file export/import cycle preserves it.
238 """
239 try:
240 row = PageTextRecord(
241 source=str(raw["source"]),
242 page=int(raw["page"]),
243 text=str(raw["text"]),
244 content_type=str(raw.get("content_type", "")),
245 )
246 except (KeyError, TypeError, ValueError):
247 raise ValueError("Dataset row is missing required source/page/text fields") from None
248 if raw.get("title") is not None:
249 row["title"] = str(raw["title"])
250 if raw.get("authors") is not None:
251 row["authors"] = str(raw["authors"])
252 if raw.get("created_at") is not None:
253 row["created_at"] = str(raw["created_at"])
254 return row
257def _deserialize_parquet(data: bytes) -> list[PageTextRecord]:
258 import io
260 import pyarrow.parquet as pq
262 return [_coerce_row(row) for row in pq.read_table(io.BytesIO(data)).to_pylist()]
265def _deserialize_jsonl(data: bytes) -> list[PageTextRecord]:
266 rows: list[PageTextRecord] = []
267 for line in data.decode("utf-8").splitlines():
268 stripped = line.strip()
269 if stripped:
270 rows.append(_coerce_row(json.loads(stripped)))
271 return rows
274_DESERIALIZERS = {
275 DatasetFormat.PARQUET: _deserialize_parquet,
276 DatasetFormat.JSONL: _deserialize_jsonl,
277}
280def deserialize_dataset(data: bytes, fmt: DatasetFormat) -> list[PageTextRecord]:
281 """Decode dataset bytes in the given format back into rows."""
282 return _DESERIALIZERS[fmt](data)
285def load_page_dataset(path: Path, fmt: DatasetFormat) -> list[PageTextRecord]:
286 """Read a per-page text dataset back from disk."""
287 if not path.exists():
288 raise ValueError(f"Dataset not found: {path}")
289 return deserialize_dataset(path.read_bytes(), fmt)
292def _page_text_row(row: PageTextRecord) -> dict:
293 """Project a dataset row down to the ``_page_texts`` columns.
295 A dataset carries the source's metadata denormalized on every page row; the
296 page-texts table has no such columns, so they are dropped before the write.
297 """
298 return {
299 "source": row["source"],
300 "page": row["page"],
301 "text": row["text"],
302 "content_type": row["content_type"],
303 }
306def _source_meta_from_rows(rows: list[PageTextRecord], name: str) -> SourceMeta:
307 """Recover a source's extraction metadata from its dataset rows.
309 The values are identical on every page row, so the first carries them. A
310 dataset exported before the metadata columns existed has none, in which case
311 the title falls back to the cleaned filename stem.
312 """
313 first: dict = dict(rows[0]) if rows else {}
314 stored = first.get("title")
315 title = stored.strip() if isinstance(stored, str) and stored.strip() else derive_title(name)
316 return SourceMeta(
317 title=title,
318 authors=first.get("authors") or "",
319 created_at=first.get("created_at") or "",
320 )
323async def import_dataset(
324 store: Store,
325 rows: list[PageTextRecord],
326 *,
327 on_progress: DetailedProgressCallback = noop_callback,
328) -> ImportResult:
329 """Re-chunk and re-embed *rows* under the current embedder.
331 Each source's pages are embedded and stored as detached ``IMPORTED``
332 chunks plus their page texts. Raises ``EmbeddingModelMismatchError`` (before
333 any write) when the store was built by a different embedder.
334 """
335 store.assert_embedding_compatible()
336 by_source: dict[str, list[PageTextRecord]] = {}
337 for row in rows:
338 by_source.setdefault(row["source"], []).append(row)
340 imported: list[str] = []
341 total_pages = 0
342 total_chunks = 0
343 for name, source_rows in by_source.items():
344 source_rows.sort(key=lambda r: r["page"])
345 content_type = source_rows[0]["content_type"] or "text"
346 page_texts = [(r["page"], r["text"]) for r in source_rows]
347 # Datasets exported with the metadata columns round-trip the extracted
348 # title/authors/created_at; older ones carry none, so fall back to the
349 # stem-derived title that keeps imported chunks visible to the title arm.
350 meta = _source_meta_from_rows(source_rows, name)
351 title = meta.title
352 with _title_scope(title):
353 chunks = await chunk_and_embed_pages(page_texts, name, content_type, on_progress)
354 for chunk in chunks:
355 chunk["title"] = title or None
356 # One locked transaction (cleanup + chunks + page texts + source row) so a
357 # failure can't leave the source with its old rows deleted and no new ones;
358 # the embedding-dim check inside runs before the cleanup delete.
359 await asyncio.to_thread(
360 store.write_chunks_batch,
361 [
362 ChunkWrite(
363 source=name,
364 file_hash="",
365 records=cast(list[dict], chunks),
366 needs_cleanup=True,
367 page_texts=[_page_text_row(r) for r in source_rows],
368 source_type=SourceType.IMPORTED,
369 meta=meta,
370 )
371 ],
372 )
373 imported.append(name)
374 total_pages += len(source_rows)
375 total_chunks += len(chunks)
376 return ImportResult(sources=sorted(imported), pages=total_pages, chunks=total_chunks)