Benjamin Mako Hill
172e66865c
sort_and_write() now takes num_partitions (default from config or 200) and passes it explicitly to repartition() instead of relying on spark.sql.shuffle.partitions. comments_part2 defaults to 800 partitions (~1GB/file matching the existing dataset density) and submissions_part2 to 300. Also cache the preprocessed DataFrame so the input is not re-read and reshuffled twice when writing the by_subreddit and by_author outputs. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
24 lines
945 B
Python
Executable File
24 lines
945 B
Python
Executable File
#!/usr/bin/env python3
|
|
"""Part 2 for comments: Spark sort + repartition into the final datasets.
|
|
|
|
Must be launched from a login node via the Hyak-provided wrapper:
|
|
start_spark_and_run.sh 1 comments_part2.py
|
|
start_spark_and_run.sh 1 comments_part2.py --indir=/path/to/parquets --mode=append
|
|
|
|
--indir defaults to the temp comments dir in dumps_helper.py.
|
|
--out_by_subreddit and --out_by_author default to the live dataset paths;
|
|
override them to write to staging directories first (see add_months.sh).
|
|
"""
|
|
|
|
import fire
|
|
from dumps_helper import COMMENTS, sort_and_write
|
|
|
|
|
|
if __name__ == "__main__":
|
|
fire.Fire(lambda indir=None, out_by_subreddit=None, out_by_author=None,
|
|
num_partitions=800:
|
|
sort_and_write(COMMENTS, indir=indir,
|
|
out_by_subreddit=out_by_subreddit,
|
|
out_by_author=out_by_author,
|
|
num_partitions=num_partitions))
|