X-Git-Url: https://code.communitydata.science/cdsc_reddit.git/blobdiff_plain/06430903f058bd9c308147dd5c256067dde5def1..f05cb962e0388feaf38aaf84f222696ab8f5f398:/similarities/tfidf.py?ds=sidebyside

diff --git a/similarities/tfidf.py b/similarities/tfidf.py
index f0b5d64..002e89f 100644
--- a/similarities/tfidf.py
+++ b/similarities/tfidf.py
@@ -15,17 +15,16 @@ def _tfidf_wrapper(func, inpath, outpath, topN, term_colname, exclude, included_
     else:
         include_subs = select_topN_subreddits(topN)
 
-    df = func(df, include_subs, term_colname)
-
-    df.write.parquet(outpath,mode='overwrite',compression='snappy')
+    dfwriter = func(df, include_subs, term_colname)
 
+    dfwriter.parquet(outpath,mode='overwrite',compression='snappy')
     spark.stop()
 
 def tfidf(inpath, outpath, topN, term_colname, exclude, included_subreddits):
     return _tfidf_wrapper(build_tfidf_dataset, inpath, outpath, topN, term_colname, exclude, included_subreddits)
 
-def tfidf_weekly(inpath, outpath, topN, term_colname, exclude):
-    return _tfidf_wrapper(build_weekly_tfidf_dataset, inpath, outpath, topN, term_colname, included_subreddits)
+def tfidf_weekly(inpath, outpath, topN, term_colname, exclude, included_subreddits):
+    return _tfidf_wrapper(build_weekly_tfidf_dataset, inpath, outpath, topN, term_colname, exclude, included_subreddits)
 
 def tfidf_authors(outpath='/gscratch/comdata/output/reddit_similarity/tfidf/comment_authors.parquet',
                   topN=25000):