Add created field to dumps (#2949)

* Add created field to dumps Useful in incremental generation of stats. * fix spark tests
metabrainz · Aug 27, 2024 · 38287bb · 38287bb
1 parent 8183d15
commit 38287bb
Show file tree

Hide file tree

Showing 16 changed files with 11 additions and 3 deletions.
diff --git a/listenbrainz/listenstore/dump_listenstore.py b/listenbrainz/listenstore/dump_listenstore.py
@@ -35,6 +35,7 @@
 
 SPARK_LISTENS_SCHEMA = pa.schema([
     pa.field("listened_at", pa.timestamp("ms"), False),
+    pa.field("created", pa.timestamp("ms"), False),
     pa.field("user_id", pa.int64(), False),
     pa.field("recording_msid", pa.string(), False),
     pa.field("artist_name", pa.string(), False),
@@ -342,6 +343,7 @@ def write_parquet_files(self,
         -- setting multiple columns at once.
                 WITH listen_with_mbid AS (
                      SELECT l.listened_at
+                          , l.created
                           , l.user_id
                           , l.recording_msid
                           -- converting jsonb array to text array is non-trivial, so return a jsonb array not text
@@ -365,6 +367,7 @@ def write_parquet_files(self,
                       WHERE {criteria} > %(start)s
                         AND {criteria} <= %(end)s
                 )    SELECT l.listened_at
+                          , l.created
                           , l.user_id
                           , l.recording_msid::TEXT
                           , l_artist_credit_mbids
@@ -388,14 +391,15 @@ def write_parquet_files(self,
         listen_count = 0
         current_listened_at = None
         conn = timescale.engine.raw_connection()
-        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as curs:
+        with (conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as curs):
             curs.execute(query, args)
             while True:
                 t0 = time.monotonic()
                 written = 0
                 approx_size = 0
                 data = {
                     'listened_at': [],
+                    'created': [],
                     'user_id': [],
                     'recording_msid': [],
                     'artist_name': [],
@@ -438,10 +442,11 @@ def write_parquet_files(self,
 
                     current_listened_at = result["listened_at"]
                     data["listened_at"].append(current_listened_at)
+                    data["created"].append(result["created"])
                     data["user_id"].append(result["user_id"])
                     data["recording_msid"].append(result["recording_msid"])
-                    approx_size += len(str(result["listened_at"])) + len(str(result["user_id"])) \
-                                   + len(result["recording_msid"])
+                    approx_size += len(str(result["listened_at"])) + len(str(result["created"])) \
+                                   + len(str(result["user_id"])) + len(result["recording_msid"])
 
                     written += 1
                     listen_count += 1

diff --git a/listenbrainz_spark/hdfs/upload.py b/listenbrainz_spark/hdfs/upload.py
@@ -163,6 +163,7 @@ def process_full_listens_dump(self):
             select extract(year from listened_at) as year
                  , extract(month from listened_at) as month
                  , listened_at
+                 , created
                  , user_id
                  , recording_msid
                  , artist_name

diff --git a/listenbrainz_spark/schema.py b/listenbrainz_spark/schema.py
@@ -14,6 +14,7 @@
 
 listens_new_schema = StructType([
     StructField('listened_at', TimestampType(), nullable=False),
+    StructField('created', TimestampType(), nullable=False),
     StructField('user_id', IntegerType(), nullable=False),
     StructField('recording_msid', StringType(), nullable=False),
     StructField('artist_name', StringType(), nullable=False),

diff --git a/listenbrainz_spark/testdata/fresh_releases_listens.parquet b/listenbrainz_spark/testdata/fresh_releases_listens.parquet
diff --git a/listenbrainz_spark/testdata/full-dump/0.parquet b/listenbrainz_spark/testdata/full-dump/0.parquet
diff --git a/listenbrainz_spark/testdata/full-dump/1.parquet b/listenbrainz_spark/testdata/full-dump/1.parquet
diff --git a/listenbrainz_spark/testdata/full-dump/2.parquet b/listenbrainz_spark/testdata/full-dump/2.parquet
diff --git a/listenbrainz_spark/testdata/full-dump/3.parquet b/listenbrainz_spark/testdata/full-dump/3.parquet
diff --git a/listenbrainz_spark/testdata/full-dump/4.parquet b/listenbrainz_spark/testdata/full-dump/4.parquet
diff --git a/listenbrainz_spark/testdata/full-dump/5.parquet b/listenbrainz_spark/testdata/full-dump/5.parquet
diff --git a/listenbrainz_spark/testdata/full-dump/6.parquet b/listenbrainz_spark/testdata/full-dump/6.parquet
diff --git a/listenbrainz_spark/testdata/incremental-dump-1/0.parquet b/listenbrainz_spark/testdata/incremental-dump-1/0.parquet
diff --git a/listenbrainz_spark/testdata/incremental-dump-2/0.parquet b/listenbrainz_spark/testdata/incremental-dump-2/0.parquet
diff --git a/listenbrainz_spark/testdata/mapped_listens.parquet b/listenbrainz_spark/testdata/mapped_listens.parquet
diff --git a/listenbrainz_spark/testdata/rec_listens.parquet b/listenbrainz_spark/testdata/rec_listens.parquet
diff --git a/listenbrainz_spark/utils/__init__.py b/listenbrainz_spark/utils/__init__.py
@@ -184,6 +184,7 @@ def get_intermediate_stats_df(start: datetime, end: datetime):
 
     query = dedent(f"""\
         select listened_at
+             , created
              , user_id
              , recording_msid
              , artist_name