Buckets:
| # /// script | |
| # requires-python = "==3.12.*" | |
| # dependencies = ["datatrove[io,processing]==0.10.1", "huggingface_hub>=1.32", "pyarrow"] | |
| # /// | |
| """The Stack (1.43 TB), MIT/Apache + 1000+ stars: Polars inside a datatrove step, files balanced by size across N Jobs.""" | |
| import argparse, json, time | |
| from datatrove.executor.jobs import JobsPipelineExecutor | |
| from datatrove.pipeline.base import PipelineStep | |
| from huggingface_hub import HfFileSystem | |
| import pyarrow.parquet as pq | |
| class PolarsStackFilter(PipelineStep): | |
| type = "🔎 - FILTER" | |
| name = "polars stack filter" | |
| def __init__(self, out, engine="polars"): | |
| super().__init__() | |
| self.out = out | |
| self.engine = engine | |
| def run(self, data, rank=0, world_size=1): | |
| import heapq, json, resource, time | |
| import os | |
| from huggingface_hub import HfApi, HfFileSystem | |
| repo = "bigcode/the-stack" | |
| rev = "153f6301d232928d7f90d09b37fceb4fadc2c950" | |
| cols = ["hexsha", "lang", "max_stars_repo_name", "max_stars_repo_path", "max_stars_repo_licenses", "max_stars_count", "content"] | |
| tree = HfApi().list_repo_tree(repo, path_in_repo="data", repo_type="dataset", revision=rev, recursive=True) | |
| files = sorted(((f.size, f.path) for f in tree if f.path.endswith(".parquet")), reverse=True) | |
| bins = [(0, i, []) for i in range(world_size)] | |
| heapq.heapify(bins) | |
| for size, path in files: | |
| total, i, members = heapq.heappop(bins) | |
| members.append(path) | |
| heapq.heappush(bins, (total + size, i, members)) | |
| total, _, mine = next(b for b in bins if b[1] == rank) | |
| urls = [f"hf://datasets/{repo}@{rev}/{p}" for p in mine] | |
| local = f"/tmp/part-{rank:05d}.parquet" | |
| t = time.time() | |
| if self.engine == "polars": | |
| import polars as pl | |
| lic = pl.col("max_stars_repo_licenses") | |
| (pl.scan_parquet(urls, extra_columns="ignore", missing_columns="insert").select(cols) | |
| .filter((pl.col("max_stars_count") >= 1000) & (lic.list.contains("MIT") | lic.list.contains("Apache-2.0"))) | |
| .sink_parquet(local)) | |
| scan = time.time() - t | |
| n = pl.scan_parquet(local).select(pl.len()).collect().item() | |
| else: | |
| import duckdb | |
| con = duckdb.connect(config={"allow_extension_repositories": "allowed"}) | |
| con.execute("SET temp_directory='/tmp/duck'") | |
| if self.engine == "duckdb-hf": | |
| con.execute("CREATE OR REPLACE EXTENSION REPOSITORY hf_repo WITH PREFIX 'https://hf.co/buckets/duckdb-extensions/builds/resolve'") | |
| con.execute("INSTALL hf FROM hf_repo") | |
| con.execute("LOAD hf FROM hf_repo") | |
| con.execute(f"CREATE SECRET (TYPE HUGGINGFACE, TOKEN '{os.environ['HF_TOKEN']}')") | |
| lst = ",".join("'" + u.replace("'", "''") + "'" for u in urls) | |
| con.execute(f"COPY (SELECT {','.join(cols)} FROM read_parquet([{lst}], union_by_name=true) WHERE max_stars_count >= 1000 AND (list_contains(max_stars_repo_licenses,'MIT') OR list_contains(max_stars_repo_licenses,'Apache-2.0'))) TO '{local}' (FORMAT parquet)") | |
| scan = time.time() - t | |
| n = con.execute(f"SELECT count(*) FROM read_parquet('{local}')").fetchone()[0] | |
| t = time.time() | |
| HfFileSystem().put_file(local, f"{self.out}/part-{rank:05d}.parquet") | |
| print("RANK", json.dumps({"engine": self.engine, "rank": rank, "files": len(urls), "input_gb": round(total / 1e9, 1), "rows": n, "scan_s": round(scan, 1), "upload_s": round(time.time() - t, 1), "peak_rss_gb": round(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e6, 2)}), flush=True) | |
| return iter(()) | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--tasks", type=int, default=8) | |
| ap.add_argument("--engine", default="polars") | |
| ap.add_argument("--flavor", default="cpu-xl") | |
| ap.add_argument("--timeout", default="30m") | |
| a = ap.parse_args() | |
| OUT = f"hf://buckets/YOUR_NAME/YOUR_BUCKET/{a.engine}-{a.tasks}-{a.flavor}" | |
| t = time.time() | |
| JobsPipelineExecutor( | |
| job_name=f"stack-dt-{a.engine}", | |
| pipeline=[PolarsStackFilter(f"{OUT}/data", a.engine)], | |
| tasks=a.tasks, workers=a.tasks, | |
| flavor=a.flavor, python="3.12", | |
| dependencies=["datatrove[io,processing]==0.10.1", "huggingface_hub>=1.32", "pyarrow", "polars==2.1.0", "duckdb==2.0.0"], | |
| logging_dir=f"{OUT}/logs", | |
| timeout=a.timeout, max_retries=0, | |
| ).run() | |
| wall = time.time() - t | |
| fs = HfFileSystem() | |
| parts = [p for p in fs.find(f"{OUT}/data") if p.endswith(".parquet")] | |
| rows = 0 | |
| for p in parts: | |
| with fs.open(p) as f: | |
| rows += pq.read_metadata(f).num_rows | |
| print("RESULT", json.dumps({"engine": a.engine, "tasks": a.tasks, "flavor": a.flavor, "wall_seconds": round(wall, 1), "parts": len(parts), "rows": rows})) | |
Xet Storage Details
- Size:
- 4.81 kB
- Xet hash:
- fb937a15c1f8efeee0c9bad6393df8f1ee8ac4e2096959cf7120ab9a0edb3c34
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.