def get_entity_all_time(
entity: str,
use_mapping: bool = False
) -> Optional[List[SitewideEntityStatMessage]]:
""" Get the all_time sitewide top entity """
logger.debug("Calculating sitewide_{}_all_time...".format(entity))
to_date = get_latest_listen_ts()
from_date = datetime(LAST_FM_FOUNDING_YEAR, 1, 1)
# Generate a dataframe containing years from "from_date" to "to_date"
time_range = [[
str(year),
int(datetime(year, 1, 1).timestamp()),
int(get_year_end(year).timestamp())
] for year in range(from_date.year, to_date.year + 1)]
time_range_df = listenbrainz_spark.session.createDataFrame(
time_range, schema=time_range_schema)
time_range_df.createOrReplaceTempView('time_range')
listens_df = get_listens(from_date, to_date, LISTENBRAINZ_DATA_DIRECTORY)
table_name = 'sitewide_{}_all_time'.format(entity)
listens_df.createOrReplaceTempView(table_name)
handler = entity_handler_map[entity]
data = handler(table_name, "yyyy", use_mapping)
message = create_message(data=data,
entity=entity,
stats_range='all_time',
from_ts=from_date.timestamp(),
to_ts=to_date.timestamp())
logger.debug("Done!")
return message
def test_get_entity_all_time(self, mock_create_message, mock_get_listens,
mock_get_latest_listen_ts):
mock_df = MagicMock()
mock_get_listens.return_value = mock_df
entity_stats.get_entity_all_time('test', False)
from_date = datetime(LAST_FM_FOUNDING_YEAR, 1, 1)
to_date = datetime(2020, 8, 21)
month = from_date
time_range = [[
str(year),
int(datetime(year, 1, 1).timestamp()),
int(get_year_end(year).timestamp())
] for year in range(from_date.year, to_date.year + 1)]
time_range_df = run_query("SELECT * FROM time_range")
time_range_result = time_range_df.rdd.map(list).collect()
self.assertListEqual(time_range_result, time_range)
mock_get_latest_listen_ts.assert_called_once()
mock_get_listens.assert_called_with(from_date, to_date,
LISTENBRAINZ_DATA_DIRECTORY)
mock_df.createOrReplaceTempView.assert_called_with(
'sitewide_test_all_time')
mock_create_message.assert_called_with(
data='sitewide_test_all_time_data',
entity='test',
stats_range='all_time',
from_ts=from_date.timestamp(),
to_ts=to_date.timestamp())
def get_listening_activity_all_time() -> Iterator[Optional[UserListeningActivityStatMessage]]:
""" Calculate the number of listens for an user in each year starting from LAST_FM_FOUNDING_YEAR (2002). """
logger.debug("Calculating listening_activity_all_time")
to_date = get_latest_listen_ts()
from_date = datetime(LAST_FM_FOUNDING_YEAR, 1, 1)
result_without_zero_years = None
for year in range(from_date.year, to_date.year+1):
year_start = datetime(year, 1, 1)
year_end = get_year_end(year)
try:
_get_listens(year_start, year_end)
except HDFSException:
# Skip if no listens present in df
continue
year_df = run_query("""
SELECT user_name,
count(user_name) as listen_count
FROM listens
GROUP BY user_name
""")
year_df = year_df.withColumn('time_range', lit(str(year))).withColumn(
'from_ts', lit(year_start.timestamp())).withColumn('to_ts', lit(year_end.timestamp()))
result_without_zero_years = result_without_zero_years.union(year_df) if result_without_zero_years else year_df
# Create a table with a list of time ranges and corresponding listen count for each user
data = result_without_zero_years \
.withColumn("listening_activity", struct("from_ts", "to_ts", "listen_count", "time_range")) \
.groupBy("user_name") \
.agg(sort_array(collect_list("listening_activity")).alias("listening_activity")) \
.toLocalIterator()
messages = create_messages(data=data, stats_range='all_time', from_ts=from_date.timestamp(), to_ts=to_date.timestamp())
logger.debug("Done!")
return messages
def test_get_year_end(self):
self.assertEqual(datetime.datetime(2020, 12, 31, 23, 59, 59, 999999), stats.get_year_end(datetime.datetime(2020, 1, 1)))