"""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())