diff --git a/datasets/comments_part2.py b/datasets/comments_part2.py
index f767b5f..67619b4 100755
--- a/datasets/comments_part2.py
+++ b/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))
diff --git a/datasets/dumps_helper.py b/datasets/dumps_helper.py
index bfd6ae6..5d534c6 100644
--- a/datasets/dumps_helper.py
+++ b/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
diff --git a/datasets/submissions_part2.py b/datasets/submissions_part2.py
index d585b8d..f189545 100755
--- a/datasets/submissions_part2.py
+++ b/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))