Use SparkSession instead of SQLContext

- also for logging-related methods

Author: sebastian-nagel

cc_index_word_count.py

@@ -16,18 +16,19 @@ class CCIndexWordCountJob(WordCountJob, CCIndexWarcSparkJob):
     records_parsing_failed = None
     records_non_html = None
 
-    def init_accumulators(self, sc):
-        super(CCIndexWordCountJob, self).init_accumulators(sc)
+    def init_accumulators(self, session):
+        super(CCIndexWordCountJob, self).init_accumulators(session)
 
+        sc = session.sparkContext
         self.records_parsing_failed = sc.accumulator(0)
         self.records_non_html = sc.accumulator(0)
 
-    def log_accumulators(self, sc):
-        super(CCIndexWordCountJob, self).log_accumulators(sc)
+    def log_accumulators(self, session):
+        super(CCIndexWordCountJob, self).log_accumulators(session)
 
-        self.log_accumulator(sc, self.records_parsing_failed,
+        self.log_accumulator(session, self.records_parsing_failed,
                              'records failed to parse = {}')
-        self.log_accumulator(sc, self.records_non_html,
+        self.log_accumulator(session, self.records_non_html,
                              'records not HTML = {}')
 
     @staticmethod

sparkcc.py

@@ -129,19 +129,22 @@ def init_logging(self, level=None, session=None):
             session.sparkContext.setLogLevel(level)
 
 
-    def init_accumulators(self, sc):
+    def init_accumulators(self, session):
         """Register and initialize counters (aka. accumulators).
            Derived classes may use this method to add their own
-           accumulators but must call super().init_accumulators(sc)
+           accumulators but must call super().init_accumulators(session)
            to also initialize counters from base classes."""
+        sc = session.sparkContext
         self.records_processed = sc.accumulator(0)
         self.warc_input_processed = sc.accumulator(0)
         self.warc_input_failed = sc.accumulator(0)
 
-    def get_logger(self, spark_context=None):
-        """Get logger from SparkContext or (if None) from logging module"""
-        if spark_context:
-            return spark_context._jvm.org.apache.log4j.LogManager \
+    def get_logger(self, session=None):
+        """Get logger from SparkSession or (if None) from logging module"""
+        if not session:
+            session = SparkSession.getActiveSession()
+        if session:
+            return session._jvm.org.apache.log4j.LogManager \
                         .getLogger(self.name)
         return logging.getLogger(self.name)
 
@@ -156,7 +159,7 @@ def run(self):
         session = builder.getOrCreate()
 
         self.init_logging(self.args.log_level, session)
-        self.init_accumulators(session.sparkContext)
+        self.init_accumulators(session)
 
         self.run_job(session)
 
@@ -165,27 +168,19 @@ def run(self):
 
         session.stop()
 
-    def log_accumulator(self, sc, acc, descr):
+    def log_accumulator(self, session, acc, descr):
         """Log single counter/accumulator"""
-        self.get_logger(sc).info(descr.format(acc.value))
+        self.get_logger(session).info(descr.format(acc.value))
 
-    def log_accumulators(self, sc):
+    def log_accumulators(self, session):
         """Log counters/accumulators, see `init_accumulators`."""
-        self.log_accumulator(sc, self.warc_input_processed,
+        self.log_accumulator(session, self.warc_input_processed,
                              'WARC/WAT/WET input files processed = {}')
-        self.log_accumulator(sc, self.warc_input_failed,
+        self.log_accumulator(session, self.warc_input_failed,
                              'WARC/WAT/WET input files failed = {}')
-        self.log_accumulator(sc, self.records_processed,
+        self.log_accumulator(session, self.records_processed,
                              'WARC/WAT/WET records processed = {}')
 
-    def log_aggregator(self, sc, agg, descr):
-        """Deprecated, use log_accumulator."""
-        self.log_accumulator(sc, agg, descr)
-
-    def log_aggregators(self, sc):
-        """Deprecated, use log_accumulators."""
-        self.log_accumulators(sc)
-
     @staticmethod
     def reduce_by_key_func(a, b):
         return a + b
@@ -205,7 +200,7 @@ def run_job(self, session):
             .options(**self.get_output_options()) \
             .saveAsTable(self.args.output)
 
-        self.log_accumulators(session.sparkContext)
+        self.log_accumulators(session)
 
     def process_warcs(self, _id, iterator):
         s3pattern = re.compile('^s3://([^/]+)/(.+)')
@@ -342,19 +337,19 @@ def add_arguments(self, parser):
     def load_table(self, session, table_path, table_name):
         parquet_reader = session.read.format('parquet')
         if self.args.table_schema is not None:
-            self.get_logger(session.sparkContext).info(
+            self.get_logger(session).info(
                 "Reading table schema from {}".format(self.args.table_schema))
             with open(self.args.table_schema, 'r') as s:
                 schema = StructType.fromJson(json.loads(s.read()))
             parquet_reader = parquet_reader.schema(schema)
         df = parquet_reader.load(table_path)
         df.createOrReplaceTempView(table_name)
-        self.get_logger(session.sparkContext).info(
+        self.get_logger(session).info(
             "Schema of table {}:\n{}".format(table_name, df.schema))
 
     def execute_query(self, session, query):
         sqldf = session.sql(query)
-        self.get_logger(session.sparkContext).info("Executing query: {}".format(query))
+        self.get_logger(session).info("Executing query: {}".format(query))
         sqldf.explain()
         return sqldf
 
@@ -364,11 +359,11 @@ def load_dataframe(self, session, partitions=-1):
         sqldf.persist()
 
         num_rows = sqldf.count()
-        self.get_logger(session.sparkContext).info(
+        self.get_logger(session).info(
             "Number of records/rows matched by query: {}".format(num_rows))
 
         if partitions > 0:
-            self.get_logger(session.sparkContext).info(
+            self.get_logger(session).info(
                 "Repartitioning data to {} partitions".format(partitions))
             sqldf = sqldf.repartition(partitions)
             sqldf.persist()
@@ -384,7 +379,7 @@ def run_job(self, session):
             .options(**self.get_output_options()) \
             .saveAsTable(self.args.output)
 
-        self.log_accumulators(session.sparkContext)
+        self.log_accumulators(session)
 
 
 class CCIndexWarcSparkJob(CCIndexSparkJob):
@@ -450,7 +445,7 @@ def load_dataframe(self, session, partitions=-1):
             sqldf = reader.load(self.args.input)
 
         if partitions > 0:
-            self.get_logger(sc).info(
+            self.get_logger(session).info(
                 "Repartitioning data to {} partitions".format(partitions))
             sqldf = sqldf.repartition(partitions)
 
@@ -523,4 +518,4 @@ def run_job(self, session):
             .options(**self.get_output_options()) \
             .saveAsTable(self.args.output)
 
-        self.log_accumulators(session.sparkContext)
+        self.log_accumulators(session)

wat_extract_links.py

@@ -139,7 +139,7 @@ def process_redirect(self, record, stream, http_status_line):
                 try:
                     redir_to = redir_to.decode('utf-8')
                 except UnicodeError as e:
-                    self.get_logger().warn(
+                    self.get_logger().warning(
                         'URL with unknown encoding: {} - {}'.format(
                             redir_to, e))
                     return
@@ -242,9 +242,10 @@ def get_links(self, url, record):
                 url, e))
             self.records_failed.add(1)
 
-    def init_accumulators(self, sc):
-        super(ExtractLinksJob, self).init_accumulators(sc)
+    def init_accumulators(self, session):
+        super(ExtractLinksJob, self).init_accumulators(session)
 
+        sc = session.sparkContext
         self.records_failed = sc.accumulator(0)
         self.records_non_html = sc.accumulator(0)
         self.records_response = sc.accumulator(0)
@@ -254,24 +255,24 @@ def init_accumulators(self, sc):
         self.records_response_robotstxt = sc.accumulator(0)
         self.link_count = sc.accumulator(0)
 
-    def log_accumulators(self, sc):
-        super(ExtractLinksJob, self).log_accumulators(sc)
+    def log_accumulators(self, session):
+        super(ExtractLinksJob, self).log_accumulators(session)
 
-        self.log_accumulator(sc, self.records_response,
+        self.log_accumulator(session, self.records_response,
                              'response records = {}')
-        self.log_accumulator(sc, self.records_failed,
+        self.log_accumulator(session, self.records_failed,
                              'records failed to process = {}')
-        self.log_accumulator(sc, self.records_non_html,
+        self.log_accumulator(session, self.records_non_html,
                              'records not HTML = {}')
-        self.log_accumulator(sc, self.records_response_wat,
+        self.log_accumulator(session, self.records_response_wat,
                              'response records WAT = {}')
-        self.log_accumulator(sc, self.records_response_warc,
+        self.log_accumulator(session, self.records_response_warc,
                              'response records WARC = {}')
-        self.log_accumulator(sc, self.records_response_redirect,
+        self.log_accumulator(session, self.records_response_redirect,
                              'response records redirects = {}')
-        self.log_accumulator(sc, self.records_response_robotstxt,
+        self.log_accumulator(session, self.records_response_robotstxt,
                              'response records robots.txt = {}')
-        self.log_accumulator(sc, self.link_count,
+        self.log_accumulator(session, self.link_count,
                              'non-unique link pairs = {}')
 
     def run_job(self, session):
@@ -480,7 +481,7 @@ def process_robotstxt(self, record, stream, _http_status_line):
                     if thost and src_host and src_host != thost:
                         yield src_host, thost
                 except UnicodeError as e:
-                    self.get_logger().warn(
+                    self.get_logger().warning(
                         'URL with unknown encoding: {} - {}'.format(
                             sitemap, e))
             line = stream.readline()