137 lines
4.4 KiB
Python
137 lines
4.4 KiB
Python
"""
|
|
Data access layer for the arxiv-complete metadata dataset.
|
|
|
|
The full metadata parquet is 1.64 GB / 3,148,796 rows (one per arXiv paper),
|
|
stored as a single file on Hugging Face:
|
|
|
|
https://huggingface.co/api/datasets/secemp9/arxiv-complete/parquet/metadata/train/0.parquet
|
|
|
|
Loading the whole thing into RAM (~a few GB of Python objects) is feasible on
|
|
a big machine but wasteful for a search prototype. Instead we stream rows from
|
|
the remote parquet with polars' *scan_parquet* + lazy filtering, so the query
|
|
path only materializes the handful of rows that actually match.
|
|
|
|
The same lazy-scanning approach works against a local copy of the parquet,
|
|
which is what the build pipeline uses to precompute embeddings.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime
|
|
from typing import Iterable
|
|
|
|
import polars as pl
|
|
|
|
DEFAULT_PARQUET_URL = (
|
|
"https://huggingface.co/api/datasets/secemp9/arxiv-complete/"
|
|
"parquet/metadata/train/0.parquet"
|
|
)
|
|
|
|
# Columns that are cheap to keep in memory and useful for display / filtering.
|
|
DISPLAY_COLUMNS = [
|
|
"paper_id",
|
|
"title",
|
|
"primary_category",
|
|
"first_version_date",
|
|
"arxiv_abs_url",
|
|
"abstract",
|
|
]
|
|
|
|
# Free-text terms that mark an abstract as a survey / overview / taxonomy paper.
|
|
SURVEY_KEYWORDS = ("survey", "comprehensive", "overview", "taxonomy")
|
|
|
|
|
|
@dataclass
|
|
class Paper:
|
|
"""A single arXiv paper, as surfaced by the search layer."""
|
|
|
|
paper_id: str
|
|
title: str
|
|
primary_category: str
|
|
year: int
|
|
url: str
|
|
abstract: str = ""
|
|
score: float = 0.0
|
|
|
|
@staticmethod
|
|
def from_row(row: pl.Row) -> "Paper":
|
|
raw = dict(zip(DISPLAY_COLUMNS, row))
|
|
fvd = raw["first_version_date"]
|
|
year = (
|
|
fvd.year
|
|
if isinstance(fvd, datetime)
|
|
else (int(fvd) if fvd is not None else 0)
|
|
)
|
|
return Paper(
|
|
paper_id=str(raw["paper_id"]),
|
|
title=str(raw["title"] or ""),
|
|
primary_category=str(raw["primary_category"] or ""),
|
|
year=year,
|
|
url=str(raw["arxiv_abs_url"] or ""),
|
|
abstract=str(raw["abstract"] or ""),
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class DataStore:
|
|
"""
|
|
Lazy view over the arxiv-complete metadata parquet.
|
|
|
|
We deliberately do not load the full 3.15M rows into memory. ``scan_parquet``
|
|
returns a lazy frame that polars only executes when we call ``collect()`` --
|
|
and it pushes down the WHERE filters (year range, category, survey) so the
|
|
download/shuffle work stays proportional to the result set.
|
|
"""
|
|
|
|
source: str = field(default_factory=lambda: os.environ.get(
|
|
"ARXIV_METADATA_PARQUET", DEFAULT_PARQUET_URL
|
|
))
|
|
|
|
def _scan(self) -> pl.LazyFrame:
|
|
return pl.scan_parquet(self.source)
|
|
|
|
def scan_filtered(
|
|
self,
|
|
*,
|
|
year_min: int | None = None,
|
|
year_max: int | None = None,
|
|
primary_category: str | None = None,
|
|
survey_only: bool = False,
|
|
) -> pl.LazyFrame:
|
|
"""
|
|
Return a lazy frame of the full metadata filtered to the given criteria.
|
|
|
|
``survey_only`` keeps only rows whose abstract mentions a survey/overview
|
|
keyword. All filters are AND-ed together and pushed into the scan.
|
|
"""
|
|
q = self._scan()
|
|
|
|
if year_min is not None:
|
|
q = q.filter(pl.col("first_version_date").dt.year() >= year_min)
|
|
if year_max is not None:
|
|
q = q.filter(pl.col("first_version_date").dt.year() <= year_max)
|
|
if primary_category:
|
|
q = q.filter(pl.col("primary_category") == primary_category)
|
|
if survey_only:
|
|
needle = "|".join(SURVEY_KEYWORDS)
|
|
q = q.filter(
|
|
pl.col("abstract")
|
|
.str.to_lowercase()
|
|
.str.contains(needle, strict=False)
|
|
)
|
|
return q
|
|
|
|
def fetch(self, limit: int | None = None) -> list[Paper]:
|
|
"""Execute the (already filtered) scan and return Paper objects."""
|
|
df = self.scan_filtered().select(DISPLAY_COLUMNS)
|
|
if limit is not None:
|
|
df = df.limit(limit)
|
|
rows = df.collect().to_dicts()
|
|
return [Paper.from_row(r.values()) for r in rows]
|
|
|
|
def count(self) -> int:
|
|
"""Total rows matching the current filter (cheap COUNT push-down)."""
|
|
return int(self.scan_filtered().select(pl.len()).collect()["len"][0])
|