src/index_query.py

"""Estimate-job payload: query the CC columnar index and materialize a range-jobs CSV.

Runs INSIDE the job (inlined via heredoc), so it must be self-contained — stdlib +
duckdb + huggingface_hub only, no `src` imports. The index SQL is vendored below
(see _build_sql), so the estimate job does NOT install cdx_toolkit at all; only the
fetch job needs it.

Reads the index WITHOUT mounting the bucket: it enumerates the parquet files of each
selected crawl partition via the HF bucket API and reads them over their CDN-fronted
`resolve` URLs with DuckDB's HTTP range reads (the same CDN the WARC fetch uses, which
the cdxt benchmarks found fastest). DuckDB can't read `hf://buckets` and can't glob
over HTTPS, hence the explicit file list.

Config comes from environment variables:
  CC_INDEX_MODE     "cdn" (enumerate HF bucket -> resolve URLs) | "local" (glob a dir)
  CC_CRAWLS         comma-separated crawl ids (e.g. CC-MAIN-2026-21,CC-MAIN-2026-17)
  CC_HOSTNAMES      comma-separated exact hostnames (optional)
  CC_DOMAINS        comma-separated registered domains (optional)
  CC_LANGUAGES      comma-separated ISO-639-3 codes to narrow by content_languages (optional)
  CC_RANGES_OUT     local path to write the range-jobs CSV (a mounted output bucket)
  CC_LOCAL_INDEX_DIR  base dir of the index when CC_INDEX_MODE=local
  HF_TOKEN          used to authenticate the bucket listing (optional for public buckets)
"""
from __future__ import annotations

import glob
import os
import re
import sys

CC_BUCKET = "commoncrawl/commoncrawl"
INDEX_SUBPATH = "cc-index/table/cc-main/warc"
RESOLVE_BASE = f"https://huggingface.co/buckets/{CC_BUCKET}/resolve"

# ISO-639-3 codes are exactly three lowercase letters. Validating here (the app
# validates too) keeps the codes safe to inline as SQL literals.
_LANG_RE = re.compile(r"^[a-z]{3}$")

# Hostnames, TLDs and crawl names only ever contain these characters. Validating
# against this set both prevents SQL injection and catches malformed input early.
# (Vendored from cdx_toolkit.filter_warc.sources.sql_base, along with _build_sql.)
_SQL_LITERAL_RE = re.compile(r"^[A-Za-z0-9.\-]+$")


def _sql_literal(value: str) -> str:
    """Validate and quote a value for safe inclusion in a SQL string literal."""
    if not isinstance(value, str) or not _SQL_LITERAL_RE.match(value):
        raise ValueError(
            f"invalid value for SQL query literal: {value!r} "
            "(allowed characters: letters, digits, dot, hyphen)"
        )
    return "'" + value + "'"


def _split(name: str):
    return [v.strip() for v in os.environ.get(name, "").split(",") if v.strip()]


def _language_clause(languages) -> str:
    """SQL predicate matching rows whose PRIMARY content language is any given code.

    `content_languages` is a comma-separated ISO-639-3 list ordered most-confident
    first (e.g. "fra,eng"), so `= 'fra' OR LIKE 'fra,%'` matches the documents CLD2
    detected as mainly that language. A secondary language does not match (e.g. 'eng'
    does not match "fra,eng"). Plain `=`/prefix-LIKE also lets parquet dictionary and
    min/max stats prune, unlike splitting the string per row. NULL -> no match."""
    bad = [c for c in languages if not _LANG_RE.match(c)]
    if bad:
        raise ValueError(f"invalid ISO-639-3 language code(s): {bad}")
    preds = [
        f"content_languages = '{c}' OR content_languages LIKE '{c},%'" for c in languages
    ]
    return "(" + " OR ".join(preds) + ")"


def _build_sql(from_clause, hostnames, domains, languages) -> str:
    """The index query: matching records' WARC byte ranges, ordered for read locality.

    Vendored from cdx_toolkit's build_sql/build_where_sql so this job needs no
    cdx_toolkit install. Differences, both because we enumerate one crawl partition's
    files at a time in `from_clause`: no `crawl IN (...)` filter, and no LIMIT.

    Host and domain predicates are OR-ed (a domain also covers its subdomains); at
    least one is required. The url_host_tld predicate is redundant but lets the
    optimizer prune row groups. ORDER BY groups records of the same WARC with
    ascending offsets, which improves range-read locality at fetch time.
    """
    if not hostnames and not domains:
        raise ValueError("an index query requires at least one hostname or registered domain")

    tlds = sorted({v.split(".")[-1] for v in list(hostnames) + list(domains)})
    tld_pred = " OR ".join(f"url_host_tld = {_sql_literal(t)}" for t in tlds)

    host_preds = [f"url_host_name = {_sql_literal(h)}" for h in hostnames]
    host_preds += [f"url_host_registered_domain = {_sql_literal(d)}" for d in domains]

    clauses = [
        "subset = 'warc'",
        f"({tld_pred}) -- help the query optimizer",
        f"({' OR '.join(host_preds)})",
    ]
    if languages:
        clauses.append(_language_clause(languages))

    where_sql = "\n        AND ".join(clauses)
    return f"""
    SELECT
        warc_filename, warc_record_offset, warc_record_length
    FROM {from_clause}
    WHERE {where_sql}
    ORDER BY warc_filename, warc_record_offset"""


def _item_path(item) -> str:
    for attr in ("path", "name", "key"):
        val = getattr(item, attr, None)
        if isinstance(val, str) and val:
            return val
    return str(item)


def _cdn_files(crawls, token):
    from huggingface_hub import HfApi

    api = HfApi()
    files = []
    for crawl in crawls:
        prefix = f"{INDEX_SUBPATH}/crawl={crawl}/subset=warc"
        for item in api.list_bucket_tree(CC_BUCKET, prefix=prefix, recursive=True, token=token):
            p = _item_path(item)
            if p.endswith(".parquet"):
                files.append(f"{RESOLVE_BASE}/{p}")
    return files


def _local_files(crawls, base):
    files = []
    for crawl in crawls:
        files.extend(glob.glob(f"{base}/crawl={crawl}/subset=warc/*.parquet"))
    return sorted(files)


def _run() -> int:
    import duckdb

    mode = os.environ.get("CC_INDEX_MODE", "cdn")
    crawls = _split("CC_CRAWLS")
    hostnames = _split("CC_HOSTNAMES")
    domains = _split("CC_DOMAINS")
    languages = _split("CC_LANGUAGES")
    out = os.environ["CC_RANGES_OUT"]
    token = os.environ.get("HF_TOKEN") or None

    if not crawls:
        print("ERROR: no crawls selected", file=sys.stderr)
        return 2

    print(
        f"[index] mode={mode} crawls={crawls} hostnames={hostnames} "
        f"domains={domains} languages={languages}",
        flush=True,
    )
    if mode == "local":
        files = _local_files(crawls, os.environ["CC_LOCAL_INDEX_DIR"])
    else:
        files = _cdn_files(crawls, token)
    print(f"[index] {len(files)} parquet file(s) to scan", flush=True)
    if not files:
        print("ERROR: no parquet index files found for the selected crawls", file=sys.stderr)
        return 3

    file_list = ", ".join("'" + f.replace("'", "''") + "'" for f in files)
    from_clause = f"read_parquet([{file_list}], hive_partitioning=true)"
    sql = _build_sql(from_clause, hostnames, domains, languages)

    con = duckdb.connect()
    con.execute("INSTALL httpfs; LOAD httpfs;")
    # The scan cost is dominated by reading ~300 parquet footers over the CDN. Those
    # reads are IO-bound (network latency), so oversubscribe threads well past the core
    # count to fan them out, and cache HTTP metadata so footers aren't re-fetched.
    tuning = (
        "SET threads=32;",
        "SET enable_http_metadata_cache=true;",
        "SET http_keep_alive=true;",
        "SET http_timeout=120000;",
        "SET http_retries=5;",
    )
    for stmt in tuning:
        try:
            con.execute(stmt)
        except Exception:  # noqa: BLE001 - older duckdb may lack a setting
            pass

    print("[index] running query + writing range-jobs CSV...", flush=True)
    con.execute(f"COPY ({sql}) TO '{out}' (FORMAT CSV, HEADER)")
    # When the filter matches 0 rows the CSV is header-only, so read_csv_auto infers
    # warc_record_length as VARCHAR; TRY_CAST keeps sum() valid (and yields 0 rows -> 0).
    n, b = con.execute(
        "SELECT count(*), coalesce(sum(TRY_CAST(warc_record_length AS BIGINT)), 0) "
        f"FROM read_csv_auto('{out}')"
    ).fetchone()
    if n == 0:
        print("[index] no records matched the filter for the selected crawl(s)", flush=True)
    print(f"ESTIMATE n_records={n} total_bytes={b}", flush=True)
    return 0


def main() -> int:
    """Wrap _run so unexpected failures surface as a clean ERROR line (the job log is
    shown to the user) instead of a raw traceback, and the job exits non-zero."""
    try:
        return _run()
    except Exception as e:  # noqa: BLE001
        print(f"ERROR: index query failed: {type(e).__name__}: {e}", file=sys.stderr, flush=True)
        return 1


if __name__ == "__main__":
    sys.exit(main())