datasets/README: remove untested notices after successful end-to-end run

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

datasets/README.md

@@ -60,9 +60,6 @@ runs the Part 2 Spark sort.
 
 ### Add new months — `add_months.sh YYYY-MM [YYYY-MM ...]`
 
-> **NOTE: written but not yet tested. Remove this notice after a
-> successful end-to-end run.**
-
 Use this for routine incremental updates. Runs Part 1 on only the
 specified months, then appends the sorted output as a new layer of
 partition files alongside the existing ones. No existing data is
@@ -130,9 +127,6 @@ from a login node — it takes the number of nodes as its first argument.
 
 ### Merge layers — `merge_layers.sh`
 
-> **NOTE: written but not yet tested. Remove this notice after a
-> successful end-to-end run.**
-
 Use this to collapse accumulated layers from incremental adds into a
 single clean layer. Reads the existing final datasets, re-sorts
 everything, writes to `.merging` temp paths, then atomically replaces the

datasets/add_months.sh

@@ -108,15 +108,32 @@ parallel --joblog add_months_joblog.txt --results add_months_logs \
 # --- Part 2: sort new months into staging (Spark, single fat node) ----------
 
 source "$SPARK_CONF_DIR/spark-env.sh"
+export PYSPARK_PYTHON="$PYTHON"
+
+# Compute executor resources from what spark-env.sh detected.
+# One executor gets all worker cores and ~85% of worker memory (leaving
+# headroom for the driver and overhead).
+_EXEC_CORES=${SPARK_WORKER_CORES:-$(nproc)}
+_WORKER_MEM_GB=${SPARK_WORKER_MEMORY%g}
+_EXEC_MEM_GB=$(( _WORKER_MEM_GB * 85 / 100 ))
+_DRIVER_MEM_GB=$(( _WORKER_MEM_GB * 10 / 100 ))
+SPARK_SUBMIT_CONF=(
+    --conf "spark.executor.cores=${_EXEC_CORES}"
+    --conf "spark.executor.memory=${_EXEC_MEM_GB}g"
+    --conf "spark.driver.memory=${_DRIVER_MEM_GB}g"
+)
+
 start_spark_cluster.sh
 
 spark-submit --master "spark://$(hostname):$SPARK_MASTER_PORT" \
+    "${SPARK_SUBMIT_CONF[@]}" \
     comments_part2.py \
     --indir="$TEMP_COMMENTS" \
     --out_by_subreddit="$STAGING_COMMENTS_SUB" \
     --out_by_author="$STAGING_COMMENTS_AUTH"
 
 spark-submit --master "spark://$(hostname):$SPARK_MASTER_PORT" \
+    "${SPARK_SUBMIT_CONF[@]}" \
     submissions_part2.py \
     --indir="$TEMP_SUBMISSIONS" \
     --out_by_subreddit="$STAGING_SUBMISSIONS_SUB" \

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