Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- Reproducing the 'alphasort' benchmark: 167,863,639 rows / 6,061,752,179 tokens
- # Machine: Windows 11, 32 GB RAM, 18 threads, SSD. Git Bash for GNU sort.
- ## 1. Generate the corpus (~32 GB, ~12 min). python gen_corpus.py corpus.txt
- # Generates a tokenized corpus matching the claim: 167,863,639 rows, 6,061,752,179 tokens (~30 GB).
- # Zipf-distributed vocabulary of 200k lowercase words, space-separated, one row per line.
- import sys, time
- import numpy as np
- ROWS, TOKENS, V = 167_863_639, 6_061_752_179, 200_000
- out = sys.argv[1] if len(sys.argv) > 1 else "corpus.txt"
- rng = np.random.default_rng(42)
- # vocab: unique random lowercase words, frequent (low-rank) words shorter, like real text
- lengths = rng.integers(2, 7, V); lengths[:20000] = rng.integers(1, 4, 20000)
- letters = rng.integers(97, 123, (V, 8), dtype=np.uint8)
- mask = np.arange(8) < lengths[:, None]
- vocab = np.unique(np.where(mask, letters, 0).view('S8').ravel())
- rng.shuffle(vocab); vocab = vocab[:V]; V = len(vocab)
- tok_sp = np.char.add(vocab, b' ').astype('S9') # token + space
- tok_nl = np.char.add(vocab, b'\n').astype('S9') # token + newline (row end)
- p = 1.0 / np.arange(1, V + 1) ** 1.05; p /= p.sum()
- base, extra = divmod(TOKENS, ROWS) # base tokens per row; first `extra` rows get one more
- CH = 500_000 # rows per chunk
- rows_done = toks_done = 0
- t0 = time.perf_counter()
- with open(out, 'wb') as f:
- while rows_done < ROWS:
- r = min(CH, ROWS - rows_done)
- lens = np.full(r, base, dtype=np.int64)
- lens[: max(0, min(r, extra - rows_done))] += 1
- n = int(lens.sum())
- idx = rng.choice(V, size=n, p=p)
- last = np.zeros(n, dtype=bool); last[np.cumsum(lens) - 1] = True
- chunk = np.where(last, tok_nl[idx], tok_sp[idx]).tobytes().replace(b'\x00', b'')
- f.write(chunk)
- rows_done += r; toks_done += n
- if rows_done % 5_000_000 == 0 or rows_done == ROWS:
- print(f"{rows_done:,} rows {toks_done:,} tokens {time.perf_counter()-t0:.0f}s", flush=True)
- assert rows_done == ROWS and toks_done == TOKENS, (rows_done, toks_done)
- print("done")
- ## 2. GNU sort baseline (45 min on this machine)
- #!/bin/bash
- # Sort every token in corpus.txt with GNU sort, timed. Run from Git Bash.
- cd "$(dirname "$0")"
- export LC_ALL=C # byte order, several x faster than locale collation
- mkdir -p tmp
- echo "start: $(date)"
- SECONDS=0
- tr ' ' '\n' < corpus.txt | sort -S 20G --parallel=16 -T tmp > sorted_tokens.txt
- echo "sort took ${SECONDS}s ($(($SECONDS/60)) min)"
- echo "lines: $(wc -l < sorted_tokens.txt) (expect 6,061,752,179)"
- echo "verify: $(sort -c sorted_tokens.txt && echo SORTED)"
- ## 3. Counting sort, Python (20 min). python count_sort.py corpus.txt out.txt
- # Counting sort: one pass to count tokens, sort the small vocabulary, write each token count times.
- # Output is byte-identical to `tr ' ' '\n' < corpus.txt | LC_ALL=C sort`.
- import sys, time
- from collections import Counter
- src, dst = sys.argv[1], sys.argv[2]
- t0 = time.perf_counter()
- counts = Counter()
- with open(src, 'rb') as f:
- tail = b''
- while chunk := f.read(1 << 27):
- chunk = tail + chunk
- chunk, _, tail = chunk.rpartition(b'\n')
- counts.update(chunk.split())
- counts.update(tail.split())
- n = sum(counts.values())
- print(f"counted {n:,} tokens, {len(counts):,} distinct, {time.perf_counter()-t0:.0f}s", flush=True)
- t1 = time.perf_counter()
- with open(dst, 'wb') as out:
- for tok in sorted(counts): # byte order == LC_ALL=C
- out.write((tok + b'\n') * counts[tok])
- print(f"wrote sorted output {time.perf_counter()-t1:.0f}s; total {time.perf_counter()-t0:.0f}s", flush=True)
- ## 4. Counting sort, Rust (5.7 min). cargo build --release; count_sort.exe corpus.txt out.txt
- ### Cargo.toml
- [package]
- name = "count_sort"
- version = "0.1.0"
- edition = "2021"
- [profile.release]
- opt-level = 3
- lto = true
- codegen-units = 1
- ### src/main.rs
- // Counting sort for a token corpus: count every whitespace-separated token in one pass,
- // sort the (small) vocabulary, write each token count times. Output is byte-identical to
- // `tr ' ' '\n' < corpus | LC_ALL=C sort`.
- use std::collections::HashMap;
- use std::io::{BufWriter, Read, Write};
- use std::time::Instant;
- fn main() {
- let args: Vec<String> = std::env::args().collect();
- let (src, dst) = (&args[1], &args[2]);
- let t0 = Instant::now();
- let mut counts: HashMap<Box<[u8]>, u64> = HashMap::with_capacity(1 << 20);
- let mut f = std::fs::File::open(src).unwrap();
- let mut buf = vec![0u8; 1 << 27];
- let mut start = 0usize; // bytes carried over from previous chunk (partial token)
- let mut total: u64 = 0;
- loop {
- let n = f.read(&mut buf[start..]).unwrap();
- if n == 0 { break; }
- let end = start + n;
- // last whitespace boundary; everything after is a partial token to carry over
- let cut = buf[..end].iter().rposition(|&b| b == b' ' || b == b'\n').map_or(0, |i| i + 1);
- for tok in buf[..cut].split(|&b| b == b' ' || b == b'\n').filter(|t| !t.is_empty()) {
- total += 1;
- match counts.get_mut(tok) {
- Some(c) => *c += 1,
- None => { counts.insert(tok.into(), 1); }
- }
- }
- buf.copy_within(cut..end, 0);
- start = end - cut;
- }
- for tok in buf[..start].split(|&b| b == b' ' || b == b'\n').filter(|t| !t.is_empty()) {
- total += 1;
- *counts.entry(tok.into()).or_insert(0) += 1;
- }
- eprintln!("counted {} tokens, {} distinct, {:.1}s", total, counts.len(), t0.elapsed().as_secs_f64());
- let t1 = Instant::now();
- let mut keys: Vec<(&Box<[u8]>, &u64)> = counts.iter().collect();
- keys.sort_unstable_by(|a, b| a.0.cmp(b.0)); // byte order == LC_ALL=C
- let mut out = BufWriter::with_capacity(1 << 24, std::fs::File::create(dst).unwrap());
- let mut line = Vec::new();
- for (tok, &c) in keys {
- line.clear(); line.extend_from_slice(tok); line.push(b'\n');
- for _ in 0..c { out.write_all(&line).unwrap(); }
- }
- out.flush().unwrap();
- eprintln!("wrote sorted output {:.1}s; total {:.1}s", t1.elapsed().as_secs_f64(), t0.elapsed().as_secs_f64());
- }
- ## Results: all three outputs byte-identical (cmp). 6,061,752,179 lines, sort -c passes.
Advertisement
Add Comment
Please, Sign In to add comment