datasets: explicit partition counts and cache df between sorts
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>
This commit is contained in:
@@ -15,7 +15,9 @@ from dumps_helper import COMMENTS, sort_and_write
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
fire.Fire(lambda indir=None, out_by_subreddit=None, out_by_author=None:
|
||||
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))
|
||||
out_by_author=out_by_author,
|
||||
num_partitions=num_partitions))
|
||||
|
||||
Reference in New Issue
Block a user