Cache incremental listens and metadata tables in Spark (#3178)

* Test persisting incremental listens df

* fix type issue

* test incremental users df

* test artist country code df

* test artist country code df - 2

* test rec/rel/rg cache dfs

* fix typo

* globalize cache table handling

* fix typo

* fix typo - 2

Author: amCap1712

listenbrainz_spark/hdfs/upload.py

@@ -183,3 +183,15 @@ def process_full_listens_dump(self):
 
         if path_exists(path.LISTENBRAINZ_BASE_STATS_DIRECTORY):
             hdfs_connection.client.delete(path.LISTENBRAINZ_BASE_STATS_DIRECTORY, recursive=True, skip_trash=True)
+
+    def process_incremental_listens_dump(self):
+        query = f"""
+            SELECT user_id
+                 , max(created) AS created
+              FROM parquet.`{path.INCREMENTAL_DUMPS_SAVE_PATH}`
+          GROUP BY user_id
+        """
+        run_query(query) \
+            .write \
+            .mode("overwrite") \
+            .parquet(path.INCREMENTAL_USERS_DF)

listenbrainz_spark/path.py

@@ -18,6 +18,7 @@
 
 # path to save incremental dumps
 INCREMENTAL_DUMPS_SAVE_PATH = os.path.join(LISTENBRAINZ_NEW_DATA_DIRECTORY, "incremental.parquet")
+INCREMENTAL_USERS_DF = os.path.join("/", "incremental-users")
 
 # Directory containing RDD checkpoints to break lineage while using iterative algorithms.
 CHECKPOINT_DIR = os.path.join('/', 'checkpoint')

listenbrainz_spark/persisted.py

@@ -0,0 +1,35 @@
+from typing import Optional
+
+from pandas import DataFrame
+
+from listenbrainz_spark.path import INCREMENTAL_DUMPS_SAVE_PATH, INCREMENTAL_USERS_DF
+from listenbrainz_spark.utils import read_files_from_HDFS
+
+_incremental_listens_df: Optional[DataFrame] = None
+_incremental_users_df: Optional[DataFrame] = None
+
+
+def unpersist_incremental_df():
+    global _incremental_listens_df, _incremental_users_df
+    if _incremental_listens_df is not None:
+        _incremental_listens_df.unpersist()
+        _incremental_listens_df = None
+    if _incremental_users_df is not None:
+        _incremental_users_df.unpersist()
+        _incremental_users_df = None
+
+
+def get_incremental_listens_df() -> DataFrame:
+    global _incremental_listens_df
+    if _incremental_listens_df is None:
+        _incremental_listens_df = read_files_from_HDFS(INCREMENTAL_DUMPS_SAVE_PATH)
+        _incremental_listens_df.persist()
+    return _incremental_listens_df
+
+
+def get_incremental_users_df() -> DataFrame:
+    global _incremental_users_df
+    if _incremental_users_df is None:
+        _incremental_users_df = read_files_from_HDFS(INCREMENTAL_USERS_DF)
+        _incremental_users_df.persist()
+    return _incremental_users_df

listenbrainz_spark/popularity/listens.py

@@ -1,9 +1,10 @@
 from datetime import datetime
 from typing import List, Optional
 
-from listenbrainz_spark.path import LISTENBRAINZ_POPULARITY_DIRECTORY, RELEASE_METADATA_CACHE_DATAFRAME
+from listenbrainz_spark.path import LISTENBRAINZ_POPULARITY_DIRECTORY
 from listenbrainz_spark.popularity.common import get_popularity_per_artist_query, \
     get_release_group_popularity_per_artist_query, get_popularity_query
+from listenbrainz_spark.postgres.release import get_release_metadata_cache
 from listenbrainz_spark.stats.incremental.query_provider import QueryProvider
 from listenbrainz_spark.stats.incremental.range_selector import ListenRangeSelector
 
@@ -25,7 +26,7 @@ def get_base_path(self) -> str:
         return LISTENBRAINZ_POPULARITY_DIRECTORY
 
     def get_filter_aggregate_query(self, existing_aggregate: str, incremental_aggregate: str,
-                                          existing_created: Optional[datetime], cache_tables: List[str]) -> str:
+                                          existing_created: Optional[datetime]) -> str:
         inc_where_clause = f"WHERE created >= to_timestamp('{existing_created}')" if existing_created else ""
         entity_id = self.get_entity_id()
         return f"""
@@ -37,23 +38,19 @@ def get_filter_aggregate_query(self, existing_aggregate: str, incremental_aggreg
              WHERE EXISTS(SELECT 1 FROM incremental_users iu WHERE iu.{entity_id} = ea.{entity_id})
         """
 
-    def get_cache_tables(self) -> List[str]:
-        if self.entity == "release_group":
-            return [RELEASE_METADATA_CACHE_DATAFRAME]
-        return []
-
     def get_entity_id(self):
         return self.entity + "_mbid"
 
-    def get_aggregate_query(self, table: str, cache_tables: List[str]) -> str:
+    def get_aggregate_query(self, table: str) -> str:
         if self.entity == "artist":
             return get_popularity_per_artist_query("artist", table)
         elif self.entity == "release_group":
-            return get_release_group_popularity_per_artist_query(table, cache_tables[0])
+            rel_cache_table = get_release_metadata_cache()
+            return get_release_group_popularity_per_artist_query(table, rel_cache_table)
         else:
             return get_popularity_query(self.entity, table)
 
-    def get_stats_query(self, final_aggregate: str, cache_tables: List[str]) -> str:
+    def get_stats_query(self, final_aggregate: str) -> str:
         return f"SELECT * FROM {final_aggregate}"
 
     def get_combine_aggregates_query(self, existing_aggregate: str, incremental_aggregate: str) -> str:
@@ -95,7 +92,7 @@ def get_base_path(self) -> str:
         return LISTENBRAINZ_POPULARITY_DIRECTORY
 
     def get_filter_aggregate_query(self, existing_aggregate: str, incremental_aggregate: str,
-                                          existing_created: Optional[datetime], cache_tables: List[str]) -> str:
+                                          existing_created: Optional[datetime]) -> str:
         inc_where_clause = f"WHERE created >= to_timestamp('{existing_created}')" if existing_created else ""
         entity_id = self.get_entity_id()
         return f"""
@@ -111,20 +108,16 @@ def get_filter_aggregate_query(self, existing_aggregate: str, incremental_aggreg
              )
         """
 
-    def get_cache_tables(self) -> List[str]:
-        if self.entity == "release_group":
-            return [RELEASE_METADATA_CACHE_DATAFRAME]
-        return []
-
     def get_entity_id(self):
         return self.entity + "_mbid"
 
-    def get_aggregate_query(self, table: str, cache_tables: List[str]) -> str:
+    def get_aggregate_query(self, table: str) -> str:
         if self.entity == "release_group":
-            return get_release_group_popularity_per_artist_query(table, cache_tables[0])
+            rel_cache_table = get_release_metadata_cache()
+            return get_release_group_popularity_per_artist_query(table, rel_cache_table)
         return get_popularity_per_artist_query(self.entity, table)
 
-    def get_stats_query(self, final_aggregate: str, cache_tables: List[str]) -> str:
+    def get_stats_query(self, final_aggregate: str) -> str:
         return f"SELECT * FROM {final_aggregate}"
 
     def get_combine_aggregates_query(self, existing_aggregate: str, incremental_aggregate: str) -> str:

listenbrainz_spark/popularity/mlhd.py

@@ -21,16 +21,6 @@ class MlhdStatsEngine:
     def __init__(self, provider: QueryProvider, message_creator: MessageCreator):
         self.provider = provider
         self.message_creator = message_creator
-        self._cache_tables = []
-
-    def _setup_cache_tables(self):
-        """ Set up metadata cache tables by reading data from HDFS and creating temporary views. """
-        cache_tables = []
-        for idx, df_path in enumerate(self.provider.get_cache_tables()):
-            df_name = f"entity_data_cache_{idx}"
-            cache_tables.append(df_name)
-            read_files_from_HDFS(df_path).createOrReplaceTempView(df_name)
-        self._cache_tables = cache_tables
 
     def create_partial_aggregate(self) -> DataFrame:
         metadata_path = self.provider.get_bookkeeping_path()
@@ -41,7 +31,7 @@ def create_partial_aggregate(self) -> DataFrame:
 
         logger.info("Creating partial aggregate from full dump listens")
         hdfs_connection.client.makedirs(Path(existing_aggregate_path).parent)
-        full_query = self.provider.get_aggregate_query(table, self._cache_tables)
+        full_query = self.provider.get_aggregate_query(table)
         full_df = run_query(full_query)
         full_df.write.mode("overwrite").parquet(existing_aggregate_path)
 
@@ -56,15 +46,14 @@ def create_partial_aggregate(self) -> DataFrame:
         return full_df
 
     def generate_stats(self) -> DataFrame:
-        self._setup_cache_tables()
         prefix = self.provider.get_table_prefix()
         self.create_partial_aggregate()
 
         partial_df = read_files_from_HDFS(self.provider.get_existing_aggregate_path())
         partial_table = f"{prefix}_existing_aggregate"
         partial_df.createOrReplaceTempView(partial_table)
 
-        results_query = self.provider.get_stats_query(partial_table, self._cache_tables)
+        results_query = self.provider.get_stats_query(partial_table)
         results_df = run_query(results_query)
         return results_df
 

listenbrainz_spark/postgres/artist.py

@@ -1,10 +1,18 @@
+from typing import Optional
+
 import pycountry
+from pyspark import StorageLevel
+from pyspark.sql import DataFrame
 
 import listenbrainz_spark
 from listenbrainz_spark import config
 from listenbrainz_spark.path import ARTIST_COUNTRY_CODE_DATAFRAME
 from listenbrainz_spark.postgres.utils import load_from_db
 from listenbrainz_spark.stats import run_query
+from listenbrainz_spark.utils import read_files_from_HDFS
+
+_ARTIST_COUNTRY_CACHE = "artist_country_cache"
+_artist_country_df: Optional[DataFrame] = None
 
 
 def create_iso_country_codes_df():
@@ -18,7 +26,6 @@ def create_iso_country_codes_df():
     df.createOrReplaceTempView("iso_codes")
 
 
-
 def create_artist_country_cache():
     """ Import artist country from postgres to HDFS for use in artist map stats calculation. """
     query = """
@@ -64,3 +71,19 @@ def create_artist_country_cache():
         .write \
         .format("parquet") \
         .save(config.HDFS_CLUSTER_URI + ARTIST_COUNTRY_CODE_DATAFRAME, mode="overwrite")
+
+    global _artist_country_df
+    if _artist_country_df is not None:
+        _artist_country_df.unpersist()
+        _artist_country_df = None
+
+
+def get_artist_country_cache():
+    """ Read the ARTIST_COUNTRY_CACHE parquet files from HDFS and create a spark SQL view
+     if one already doesn't exist """
+    global _artist_country_df
+    if _artist_country_df is None:
+        _artist_country_df = read_files_from_HDFS(ARTIST_COUNTRY_CODE_DATAFRAME)
+        _artist_country_df.persist(StorageLevel.DISK_ONLY)
+        _artist_country_df.createOrReplaceTempView(_ARTIST_COUNTRY_CACHE)
+    return _ARTIST_COUNTRY_CACHE

listenbrainz_spark/postgres/recording.py

@@ -1,5 +1,14 @@
+from typing import Optional
+
+from pyspark import StorageLevel
+from pyspark.sql import DataFrame
+
 from listenbrainz_spark.path import RECORDING_LENGTH_DATAFRAME, RECORDING_ARTIST_DATAFRAME
 from listenbrainz_spark.postgres.utils import save_pg_table_to_hdfs
+from listenbrainz_spark.utils import read_files_from_HDFS
+
+_RECORDING_ARTIST_CACHE = "recording_artist_cache"
+_recording_artist_df: Optional[DataFrame] = None
 
 
 def create_recording_length_cache():
@@ -35,3 +44,20 @@ def create_recording_artist_cache():
     """
 
     save_pg_table_to_hdfs(query, RECORDING_ARTIST_DATAFRAME, process_artists_column=True)
+
+    global _recording_artist_df
+    if _recording_artist_df is not None:
+        _recording_artist_df.unpersist()
+        _recording_artist_df = None
+
+
+def get_recording_artist_cache():
+    """ Read the RECORDING_ARTIST_CACHE parquet files from HDFS and create a spark SQL view
+     if one already doesn't exist """
+    global _recording_artist_df
+    if _recording_artist_df is None:
+        _recording_artist_df = read_files_from_HDFS(RECORDING_ARTIST_DATAFRAME)
+        _recording_artist_df.persist(StorageLevel.DISK_ONLY)
+        _recording_artist_df.createOrReplaceTempView(_RECORDING_ARTIST_CACHE)
+    return _RECORDING_ARTIST_CACHE
+

listenbrainz_spark/postgres/release.py

@@ -1,5 +1,14 @@
+from typing import Optional
+
+from pyspark import StorageLevel
+from pyspark.sql import DataFrame
+
 from listenbrainz_spark.path import RELEASE_METADATA_CACHE_DATAFRAME
 from listenbrainz_spark.postgres.utils import save_pg_table_to_hdfs
+from listenbrainz_spark.utils import read_files_from_HDFS
+
+_RELEASE_METADATA_CACHE = "release_metadata_cache"
+_release_metadata_df: Optional[DataFrame] = None
 
 
 def create_release_metadata_cache():
@@ -104,3 +113,19 @@ def create_release_metadata_cache():
     """
 
     save_pg_table_to_hdfs(query, RELEASE_METADATA_CACHE_DATAFRAME, process_artists_column=True)
+
+    global _release_metadata_df
+    if _release_metadata_df is not None:
+        _release_metadata_df.unpersist()
+        _release_metadata_df = None
+
+
+def get_release_metadata_cache():
+    """ Read the RELEASE_METADATA_CACHE parquet files from HDFS and create a spark SQL view
+     if one already doesn't exist """
+    global _release_metadata_df
+    if _release_metadata_df is None:
+        _release_metadata_df = read_files_from_HDFS(RELEASE_METADATA_CACHE_DATAFRAME)
+        _release_metadata_df.persist(StorageLevel.DISK_ONLY)
+        _release_metadata_df.createOrReplaceTempView(_RELEASE_METADATA_CACHE)
+    return _RELEASE_METADATA_CACHE

listenbrainz_spark/postgres/release_group.py

@@ -1,5 +1,14 @@
+from typing import Optional
+
+from pyspark import StorageLevel
+from pyspark.sql import DataFrame
+
 from listenbrainz_spark.path import RELEASE_GROUP_METADATA_CACHE_DATAFRAME
 from listenbrainz_spark.postgres.utils import save_pg_table_to_hdfs
+from listenbrainz_spark.utils import read_files_from_HDFS
+
+_RELEASE_GROUP_METADATA_CACHE = "release_group_metadata_cache"
+_release_group_metadata_df: Optional[DataFrame] = None
 
 
 def create_release_group_metadata_cache():
@@ -72,3 +81,19 @@ def create_release_group_metadata_cache():
     """
 
     save_pg_table_to_hdfs(query, RELEASE_GROUP_METADATA_CACHE_DATAFRAME, process_artists_column=True)
+
+    global _release_group_metadata_df
+    if _release_group_metadata_df is not None:
+        _release_group_metadata_df.unpersist()
+        _release_group_metadata_df = None
+
+
+def get_release_group_metadata_cache():
+    """ Read the RELEASE_GROUP_METADATA_CACHE parquet files from HDFS and create a spark SQL view
+     if one already doesn't exist """
+    global _release_group_metadata_df
+    if _release_group_metadata_df is None:
+        _release_group_metadata_df = read_files_from_HDFS(RELEASE_GROUP_METADATA_CACHE_DATAFRAME)
+        _release_group_metadata_df.persist(StorageLevel.DISK_ONLY)
+        _release_group_metadata_df.createOrReplaceTempView(_RELEASE_GROUP_METADATA_CACHE)
+    return _RELEASE_GROUP_METADATA_CACHE

listenbrainz_spark/request_consumer/jobs/import_dump.py

@@ -1,16 +1,15 @@
 """ Spark job that downloads the latest listenbrainz dumps and imports into HDFS
 """
 import logging
-import shutil
 import tempfile
-import time
-from datetime import datetime
+from datetime import datetime, timezone
 
 import listenbrainz_spark.request_consumer.jobs.utils as utils
 from listenbrainz_spark.dump import DumpType
 from listenbrainz_spark.dump.local import ListenbrainzLocalDumpLoader
 from listenbrainz_spark.ftp.download import ListenbrainzDataDownloader
 from listenbrainz_spark.hdfs.upload import ListenbrainzDataUploader
+from listenbrainz_spark.persisted import unpersist_incremental_df
 
 logger = logging.getLogger(__name__)
 
@@ -40,7 +39,8 @@ def import_full_dump_to_hdfs(dump_id: int = None, local: bool = False) -> str:
         uploader = ListenbrainzDataUploader()
         uploader.upload_new_listens_full_dump(src)
         uploader.process_full_listens_dump()
-    utils.insert_dump_data(dump_id, DumpType.FULL, datetime.utcnow())
+    utils.insert_dump_data(dump_id, DumpType.FULL, datetime.now(tz=timezone.utc))
+    unpersist_incremental_listens_df()
     return dump_name
 
 
@@ -68,8 +68,11 @@ def import_incremental_dump_to_hdfs(dump_id: int = None, local: bool = False) ->
         # instantiating ListenbrainzDataUploader creates a spark session which
         # is a bit non-intuitive.
         # FIXME in future to make initializing of spark session more explicit?
-        ListenbrainzDataUploader().upload_new_listens_incremental_dump(src)
-    utils.insert_dump_data(dump_id, DumpType.INCREMENTAL, datetime.utcnow())
+        uploader = ListenbrainzDataUploader()
+        uploader.upload_new_listens_incremental_dump(src)
+        uploader.process_incremental_listens_dump()
+    utils.insert_dump_data(dump_id, DumpType.INCREMENTAL, datetime.now(tz=timezone.utc))
+    unpersist_incremental_df()
     return dump_name