X-Git-Url: https://code.communitydata.science/cdsc_reddit.git/blobdiff_plain/e6294b5b90135a5163441c8dc62252dd6a188412..01a4c353588ab1a28f36980157daa5e682ea9edc:/similarities/top_subreddits_by_comments.py

diff --git a/similarities/top_subreddits_by_comments.py b/similarities/top_subreddits_by_comments.py
index 214c7e0..1197b51 100644
--- a/similarities/top_subreddits_by_comments.py
+++ b/similarities/top_subreddits_by_comments.py
@@ -5,6 +5,10 @@ from pyspark.sql import Window
 spark = SparkSession.builder.getOrCreate()
 conf = spark.sparkContext.getConf()
 
+submissions = spark.read.parquet("/gscratch/comdata/output/reddit_submissions_by_subreddit.parquet")
+
+prop_nsfw = submissions.select(['subreddit','over_18']).groupby('subreddit').agg(f.mean(f.col('over_18').astype('double')).alias('prop_nsfw'))
+
 df = spark.read.parquet("/gscratch/comdata/output/reddit_comments_by_subreddit.parquet")
 
 # remove /u/ pages
@@ -12,6 +16,9 @@ df = df.filter(~df.subreddit.like("u_%"))
 
 df = df.groupBy('subreddit').agg(f.count('id').alias("n_comments"))
 
+df = df.join(prop_nsfw,on='subreddit')
+df = df.filter(df.prop_nsfw < 0.5)
+
 win = Window.orderBy(f.col('n_comments').desc())
 df = df.withColumn('comments_rank', f.rank().over(win))