|
| 1 | +"""Optional browse copies: a file rewritten with small row groups so any page can be read in milliseconds. |
| 2 | +
|
| 3 | +Parquet files written as one huge row group must be decoded up to the requested rows (or completely), so jumping |
| 4 | +deep into them takes seconds; large text formats (CSV, JSON Lines, ...) must be scanned from the start. A browse copy |
| 5 | +holds the same rows as Parquet with BROWSE_GROUP_ROWS-row groups. It is built on request, in the background, and is |
| 6 | +keyed by the source file's size and modification time, so it is never used for a changed file. |
| 7 | +""" |
| 8 | +import logging |
| 9 | +import shutil |
| 10 | +import threading |
| 11 | +from pathlib import Path |
| 12 | +from typing import Dict, Optional |
| 13 | + |
| 14 | +import pyarrow as pa |
| 15 | +import pyarrow.parquet as pq |
| 16 | + |
| 17 | +from app.core.config import settings |
| 18 | +from app.models.schemas import BrowseCopyStatus, Dataset, DataFormat |
| 19 | +from app.services.data_loader import dataset_manager |
| 20 | +from app.services.row_cache import row_cache |
| 21 | + |
| 22 | +logger = logging.getLogger(__name__) |
| 23 | +BROWSE_GROUP_ROWS = 10_000 |
| 24 | +# A Parquet row group above this (uncompressed) makes deep pages slow; other formats need a copy above this file size. |
| 25 | +SLOW_GROUP_BYTES = 256 * 1024 * 1024 |
| 26 | +SLOW_FILE_BYTES = 256 * 1024 * 1024 |
| 27 | + |
| 28 | + |
| 29 | +class _Job: |
| 30 | + def __init__(self, total: int): |
| 31 | + self.total = total |
| 32 | + self.done = 0 |
| 33 | + self.error: Optional[str] = None |
| 34 | + self.cancelled = False |
| 35 | + self.thread: Optional[threading.Thread] = None |
| 36 | + |
| 37 | + |
| 38 | +class BrowseCopies: |
| 39 | + def __init__(self, base_path: Optional[str] = None): |
| 40 | + self.base = Path(base_path or Path(settings.data.registry_path).parent / "browse_copies") |
| 41 | + self._jobs: Dict[str, _Job] = {} |
| 42 | + self._lock = threading.RLock() |
| 43 | + |
| 44 | + # ---------- paths ---------- |
| 45 | + |
| 46 | + def _path(self, dataset: Dataset) -> Path: |
| 47 | + stat = Path(dataset.path).stat() |
| 48 | + return self.base / dataset.id / f"{stat.st_size}-{stat.st_mtime_ns}.parquet" |
| 49 | + |
| 50 | + def ready_path(self, dataset: Dataset) -> Optional[str]: |
| 51 | + """The copy to read pages from, if one exists for the file's current version.""" |
| 52 | + try: |
| 53 | + path = self._path(dataset) |
| 54 | + except OSError: |
| 55 | + return None |
| 56 | + return str(path) if path.exists() else None |
| 57 | + |
| 58 | + # ---------- status ---------- |
| 59 | + |
| 60 | + @staticmethod |
| 61 | + def needed(dataset: Dataset) -> bool: |
| 62 | + if dataset.format != DataFormat.PARQUET: |
| 63 | + return dataset.size_bytes >= SLOW_FILE_BYTES |
| 64 | + try: |
| 65 | + meta = pq.ParquetFile(dataset.path).metadata |
| 66 | + except (OSError, pa.ArrowException): |
| 67 | + return False |
| 68 | + return any(meta.row_group(i).total_byte_size > SLOW_GROUP_BYTES for i in range(meta.num_row_groups)) |
| 69 | + |
| 70 | + def status(self, dataset_id: str) -> BrowseCopyStatus: |
| 71 | + dataset = self._dataset(dataset_id) |
| 72 | + ready = self.ready_path(dataset) |
| 73 | + if ready: |
| 74 | + return BrowseCopyStatus(state="ready", needed=True, done=dataset.row_count, total=dataset.row_count, |
| 75 | + size_bytes=Path(ready).stat().st_size) |
| 76 | + job = self._jobs.get(dataset_id) |
| 77 | + if job and job.error: |
| 78 | + return BrowseCopyStatus(state="error", needed=True, done=job.done, total=job.total, error=job.error) |
| 79 | + if job: |
| 80 | + return BrowseCopyStatus(state="building", needed=True, done=job.done, total=job.total) |
| 81 | + return BrowseCopyStatus(state="missing", needed=self.needed(dataset), total=dataset.row_count) |
| 82 | + |
| 83 | + # ---------- building ---------- |
| 84 | + |
| 85 | + def build(self, dataset_id: str) -> BrowseCopyStatus: |
| 86 | + """Start building the copy in the background (no-op if it is ready or already building).""" |
| 87 | + with self._lock: |
| 88 | + current = self.status(dataset_id) |
| 89 | + if current.state in ("ready", "building"): |
| 90 | + return current |
| 91 | + dataset = self._dataset(dataset_id) |
| 92 | + job = _Job(dataset.row_count) |
| 93 | + self._jobs[dataset_id] = job |
| 94 | + job.thread = threading.Thread( |
| 95 | + target=self._build, args=(dataset, self._path(dataset), job), daemon=True, name=f"browse-{dataset_id}", |
| 96 | + ) |
| 97 | + job.thread.start() |
| 98 | + return self.status(dataset_id) |
| 99 | + |
| 100 | + def _build(self, dataset: Dataset, target: Path, job: _Job) -> None: |
| 101 | + tmp = target.with_suffix(".tmp") |
| 102 | + writer: Optional[pq.ParquetWriter] = None |
| 103 | + try: |
| 104 | + target.parent.mkdir(parents=True, exist_ok=True) |
| 105 | + for batch in self._batches(dataset): |
| 106 | + if job.cancelled: |
| 107 | + raise InterruptedError("cancelled") |
| 108 | + table = pa.Table.from_batches([batch]) if isinstance(batch, pa.RecordBatch) else batch |
| 109 | + if writer is None: |
| 110 | + writer = pq.ParquetWriter(tmp, table.schema, compression="zstd", compression_level=1) |
| 111 | + writer.write_table(table, row_group_size=BROWSE_GROUP_ROWS) |
| 112 | + job.done += table.num_rows |
| 113 | + if writer is None: |
| 114 | + raise ValueError("The file has no rows") |
| 115 | + writer.close() |
| 116 | + writer = None |
| 117 | + tmp.replace(target) |
| 118 | + for old in target.parent.iterdir(): |
| 119 | + if old != target: |
| 120 | + old.unlink(missing_ok=True) |
| 121 | + with self._lock: |
| 122 | + if self._jobs.get(dataset.id) is job: |
| 123 | + del self._jobs[dataset.id] |
| 124 | + logger.info(f"Browse copy of {dataset.name} built: {job.done} rows") |
| 125 | + except Exception as e: |
| 126 | + if writer is not None: |
| 127 | + writer.close() |
| 128 | + tmp.unlink(missing_ok=True) |
| 129 | + if not job.cancelled: |
| 130 | + logger.exception(f"Building the browse copy of {dataset.name} failed") |
| 131 | + job.error = str(e) or type(e).__name__ |
| 132 | + |
| 133 | + @staticmethod |
| 134 | + def _batches(dataset: Dataset): |
| 135 | + """The source rows in BROWSE_GROUP_ROWS batches, read as a stream so memory stays flat.""" |
| 136 | + if dataset.format == DataFormat.PARQUET: |
| 137 | + file = pq.ParquetFile(dataset.path, buffer_size=1 << 20, pre_buffer=False) |
| 138 | + try: |
| 139 | + yield from file.iter_batches(batch_size=BROWSE_GROUP_ROWS) |
| 140 | + finally: |
| 141 | + file.close() |
| 142 | + return |
| 143 | + lf = dataset_manager.get_dataframe(dataset.id) |
| 144 | + for frame in lf.collect_batches(chunk_size=BROWSE_GROUP_ROWS): |
| 145 | + yield frame.to_arrow() |
| 146 | + |
| 147 | + # ---------- removal ---------- |
| 148 | + |
| 149 | + def delete(self, dataset_id: str) -> bool: |
| 150 | + """Stop any build and remove the dataset's copies. Returns whether there was a copy or a build.""" |
| 151 | + with self._lock: |
| 152 | + job = self._jobs.pop(dataset_id, None) |
| 153 | + if job and job.thread and job.thread.is_alive(): |
| 154 | + job.cancelled = True |
| 155 | + job.thread.join(timeout=60) |
| 156 | + folder = self.base / dataset_id |
| 157 | + existed = bool(job and not job.error) or (folder.exists() and any(folder.glob("*.parquet"))) |
| 158 | + if folder.exists(): |
| 159 | + for copy in folder.glob("*.parquet"): |
| 160 | + row_cache.release(str(copy)) |
| 161 | + shutil.rmtree(folder, ignore_errors=True) |
| 162 | + return existed |
| 163 | + |
| 164 | + @staticmethod |
| 165 | + def _dataset(dataset_id: str) -> Dataset: |
| 166 | + dataset = dataset_manager.get_dataset(dataset_id) |
| 167 | + if not dataset: |
| 168 | + raise KeyError(f"Dataset {dataset_id} not found") |
| 169 | + return dataset |
| 170 | + |
| 171 | + |
| 172 | +browse_copies = BrowseCopies() |
0 commit comments