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>

datasets/comments_part2.py

@@ -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))

datasets/dumps_helper.py

@@ -256,7 +256,8 @@ def gen_task_list(config, script_name, dumpdir=None, tasklist=None):
 
 # --- Part 2: spark sort + repartition --------------------------------------
 
-def sort_and_write(config, indir=None, out_by_subreddit=None, out_by_author=None):
+def sort_and_write(config, indir=None, out_by_subreddit=None, out_by_author=None,
+                   num_partitions=None):
     """Read a directory of per-source parquets, sort and repartition twice
     (once by subreddit, once by author), and write the two output datasets.
 
@@ -265,6 +266,11 @@ def sort_and_write(config, indir=None, out_by_subreddit=None, out_by_author=None
     and config['output_by_author']. Override them to write to staging directories
     instead of the live datasets (see add_months.sh).
 
+    num_partitions controls the number of output files per dataset. Target
+    ~1GB per output file: use ~800 for a full 16-month comments batch,
+    ~300 for submissions. Defaults to config['num_partitions'] if set,
+    otherwise 200.
+
     Pyspark is imported lazily so Part 1 callers don't pay the Spark startup
     cost.
     """
@@ -273,6 +279,7 @@ def sort_and_write(config, indir=None, out_by_subreddit=None, out_by_author=None
     indir = indir or config['outdir']
     out_by_subreddit = out_by_subreddit or config['output_by_subreddit']
     out_by_author = out_by_author or config['output_by_author']
+    num_partitions = num_partitions or config.get('num_partitions', 200)
 
     spark = SparkSession.builder.appName(config['app_name']).getOrCreate()
 
@@ -287,16 +294,21 @@ def sort_and_write(config, indir=None, out_by_subreddit=None, out_by_author=None
     df = df.withColumn("Year", f.year(f.col("CreatedAt")))
     df = df.withColumn("Day", f.dayofmonth(f.col("CreatedAt")))
 
+    # Cache after preprocessing so the input is not re-read for each sort.
+    df = df.cache()
+
     sub_keys = config['subreddit_sort_keys']
-    df_sub = df.repartition('subreddit').sort(sub_keys, ascending=True)
+    df_sub = df.repartition(num_partitions, 'subreddit').sort(sub_keys, ascending=True)
     df_sub = df_sub.sortWithinPartitions(sub_keys, ascending=True)
     df_sub.write.parquet(out_by_subreddit, mode='overwrite', compression='snappy')
 
     auth_keys = config['author_sort_keys']
-    df_auth = df.repartition('author').sort(auth_keys, ascending=True)
+    df_auth = df.repartition(num_partitions, 'author').sort(auth_keys, ascending=True)
     df_auth = df_auth.sortWithinPartitions(auth_keys, ascending=True)
     df_auth.write.parquet(out_by_author, mode='overwrite', compression='snappy')
 
+    df.unpersist()
+
 
 def merge_layers(config):
     """Collapse all accumulated layers in the final datasets into a single

datasets/submissions_part2.py

@@ -15,7 +15,9 @@ from dumps_helper import SUBMISSIONS, 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=300:
               sort_and_write(SUBMISSIONS, indir=indir,
                              out_by_subreddit=out_by_subreddit,
-                             out_by_author=out_by_author))
+                             out_by_author=out_by_author,
+                             num_partitions=num_partitions))