import fire
from pyspark.sql import SparkSession
from pyspark.sql import functions as f
-from similarities_helper import build_tfidf_dataset, build_weekly_tfidf_dataset, select_topN_subreddits
+from similarities_helper import tfidf_dataset, build_weekly_tfidf_dataset, select_topN_subreddits
+from functools import partial
-
-def _tfidf_wrapper(func, inpath, outpath, topN, term_colname, exclude):
+inpath = '/gscratch/comdata/users/nathante/competitive_exclusion_reddit/data/tfidf/comment_authors_compex.parquet'
+# include_terms is a path to a parquet file that contains a column of term_colname + '_id' to include.
+def _tfidf_wrapper(func, inpath, outpath, topN, term_colname, exclude, included_subreddits, included_terms=None, min_df=None, max_df=None):
spark = SparkSession.builder.getOrCreate()
df = spark.read.parquet(inpath)
df = df.filter(~ f.col(term_colname).isin(exclude))
- include_subs = select_topN_subreddits(topN)
+ if included_subreddits is not None:
+ include_subs = set(map(str.strip,open(included_subreddits)))
+ else:
+ include_subs = select_topN_subreddits(topN)
- df = func(df, include_subs, term_colname)
+ include_subs = spark.sparkContext.broadcast(include_subs)
- df.write.parquet(outpath,mode='overwrite',compression='snappy')
+ # term_id = term_colname + "_id"
- spark.stop()
+ if included_terms is not None:
+ terms_df = spark.read.parquet(included_terms)
+ terms_df = terms_df.select(term_colname).distinct()
+ df = df.join(terms_df, on=term_colname, how='left_semi')
-def tfidf(inpath, outpath, topN, term_colname, exclude):
- return _tfidf_wrapper(build_tfidf_dataset, inpath, outpath, topN, term_colname, exclude)
+ dfwriter = func(df, include_subs.value, term_colname)
-def tfidf_weekly(inpath, outpath, topN, term_colname, exclude):
- return _tfidf_wrapper(build_weekly_tfidf_dataset, inpath, outpath, topN, term_colname, exclude)
+ dfwriter.parquet(outpath,mode='overwrite',compression='snappy')
+ spark.stop()
-def tfidf_authors(outpath='/gscratch/comdata/output/reddit_similarity/tfidf/comment_authors.parquet',
- topN=25000):
+def tfidf(inpath, outpath, topN, term_colname, exclude, included_subreddits, min_df, max_df):
+ tfidf_func = partial(tfidf_dataset, max_df=max_df, min_df=min_df)
+ return _tfidf_wrapper(tfidf_func, inpath, outpath, topN, term_colname, exclude, included_subreddits)
- return tfidf("/gscratch/comdata/output/reddit_ngrams/comment_authors.parquet",
- outpath,
- topN,
- 'author',
- ['[deleted]','AutoModerator']
- )
+def tfidf_weekly(inpath, outpath, static_tfidf_path, topN, term_colname, exclude, included_subreddits):
+ return _tfidf_wrapper(build_weekly_tfidf_dataset, inpath, outpath, topN, term_colname, exclude, included_subreddits, included_terms=static_tfidf_path)
-def tfidf_terms(outpath='/gscratch/comdata/output/reddit_similarity/tfidf/comment_terms.parquet',
- topN=25000):
- return tfidf("/gscratch/comdata/output/reddit_ngrams/comment_terms.parquet",
- outpath,
- topN,
- 'term',
- []
- )
-
-def tfidf_authors_weekly(outpath='/gscratch/comdata/output/reddit_similarity/tfidf_weekly/comment_authors.parquet',
- topN=25000):
+def tfidf_authors(inpath="/gscratch/comdata/output/reddit_ngrams/comment_authors.parquet",
+ outpath='/gscratch/comdata/output/reddit_similarity/tfidf/comment_authors.parquet',
+ topN=None,
+ included_subreddits=None,
+ min_df=None,
+ max_df=None):
- return tfidf_weekly("/gscratch/comdata/output/reddit_ngrams/comment_authors.parquet",
+ return tfidf(inpath,
outpath,
topN,
'author',
- ['[deleted]','AutoModerator']
+ ['[deleted]','AutoModerator'],
+ included_subreddits=included_subreddits,
+ min_df=min_df,
+ max_df=max_df
)
-def tfidf_terms_weekly(outpath='/gscratch/comdata/output/reddit_similarity/tfidf_weekly/comment_terms.parquet',
- topN=25000):
+def tfidf_terms(inpath="/gscratch/comdata/output/reddit_ngrams/comment_terms.parquet",
+ outpath='/gscratch/comdata/output/reddit_similarity/tfidf/comment_terms.parquet',
+ topN=None,
+ included_subreddits=None,
+ min_df=None,
+ max_df=None):
- return tfidf_weekly("/gscratch/comdata/output/reddit_ngrams/comment_terms.parquet",
+ return tfidf(inpath,
outpath,
topN,
'term',
- []
+ [],
+ included_subreddits=included_subreddits,
+ min_df=min_df,
+ max_df=max_df
)
+def tfidf_authors_weekly(inpath="/gscratch/comdata/output/reddit_ngrams/comment_authors.parquet",
+ static_tfidf_path="/gscratch/comdata/output/reddit_similarity/tfidf/comment_authors.parquet",
+ outpath='/gscratch/comdata/output/reddit_similarity/tfidf_weekly/comment_authors.parquet',
+ topN=None,
+ included_subreddits=None
+ ):
+
+ return tfidf_weekly(inpath,
+ outpath,
+ static_tfidf_path,
+ topN,
+ 'author',
+ ['[deleted]','AutoModerator'],
+ included_subreddits=included_subreddits
+ )
+
+def tfidf_terms_weekly(inpath="/gscratch/comdata/output/reddit_ngrams/comment_terms.parquet",
+ static_tfidf_path="/gscratch/comdata/output/reddit_similarity/tfidf/comment_terms.parquet",
+ outpath='/gscratch/comdata/output/reddit_similarity/tfidf_weekly/comment_terms.parquet',
+ topN=None,
+ included_subreddits=None
+ ):
+
+
+ return tfidf_weekly(inpath,
+ outpath,
+ static_tfidf_path,
+ topN,
+ 'term',
+ [],
+ included_subreddits=included_subreddits
+ )
+
if __name__ == "__main__":
fire.Fire({'authors':tfidf_authors,