import argparse from functools import reduce from itertools import product from pyspark.sql import DataFrame, SparkSession from pyspark.sql.functions import * from pyspark.sql.types import DateType def read_streams(spark, root_folder, args): s3_url = '{root}/{dsp}_streams/report_date={date}/report_licensor={licensor}/{country}.json'.format( root=root_folder, dsp=args['dsp'], date=args['date'], licensor=args['licensor'], country=args['country']) return spark.read \ .json(s3_url) \ .select(col('user_id'), col('track_id'), col('source'), col('source_uri'), col('shuffle'), col('repeat_play'), col('source')) \ .withColumn('report_date', lit(args['date']).cast(DateType())) \ .withColumn('report_licensor', lit(args['licensor'])) def read_users(spark, root_folder, args): s3_url = '{root}/{dsp}_users/report_date={date}/report_licensor={licensor}/users.json'.format( root=root_folder, dsp=args['dsp'], date=args['date'], licensor=args['licensor']) return spark.read \ .json(s3_url) \ .select(col('user_id'), col('access'), col('country')) \ .withColumn('report_date', lit(args['date']).cast(DateType())) \ .withColumn('report_licensor', lit(args['licensor'])) def read_tracks(spark, root_folder, args): s3_url = '{root}/{dsp}_tracks/report_date={date}/report_licensor={licensor}/tracks.json'.format( root=root_folder, dsp=args['dsp'], date=args['date'], licensor=args['licensor']) return spark.read.json(s3_url) \ .select(col('track_id'), col('isrc')) \ .withColumn('report_date', lit(args['date']).cast(DateType())) \ .withColumn('report_licensor', lit(args['licensor'])) def aggregate(streams, users, tracks): return streams \ .join(users, ['user_id', 'report_date', 'report_licensor']) \ .join(tracks, ['track_id', 'report_date', 'report_licensor']) \ .withColumn('source_album', when(lower(streams.source) == 'album', 1).otherwise(0)) \ .withColumn('source_artist', when(lower(streams.source) == 'artist', 1).otherwise(0)) \ .withColumn('source_chart', when(lower(streams.source) == 'chart', 1).otherwise(0)) \ .withColumn('source_collection', when(lower(streams.source) == 'collection', 1).otherwise(0)) \ .withColumn('source_daily_mix', when((lower(streams.source) == 'radio') & (lower(streams.source_uri) == 'dailymix'), 1).otherwise(0)) \ .withColumn('source_discover_weekly', when((lower(streams.source) == 'others_playlist') & (lower(streams.source_uri) == 'discoverweekly'), 1).otherwise(0)) \ .withColumn('source_other', when(lower(streams.source) == 'other', 1).otherwise(0)) \ .withColumn('source_others_playlist', when((lower(streams.source) == 'others_playlist') & (~lower(streams.source_uri).isin('discoverweekly', 'discoverweekly')), 1).otherwise(0)) \ .withColumn('source_play_queue', when(lower(streams.source) == 'play_queue', 1).otherwise(0)) \ .withColumn('source_radio', when((lower(streams.source) == 'radio') & (lower(streams.source_uri) != 'dailymix'), 1).otherwise(0)) \ .withColumn('source_release_radar', when((lower(streams.source) == 'others_playlist') & (lower(streams.source_uri) == 'discoverweekly'), 1).otherwise(0)) \ .withColumn('source_search', when(lower(streams.source) == 'search', 1).otherwise(0)) \ .withColumn('source_unknown', when(~lower(streams.source).isin('album', 'artist', 'chart', 'collection', 'other', 'others_playlist', 'play_queue', 'radio', 'search'), 1).otherwise(0)) \ .withColumn('free_source_album', when((lower(streams.source) == 'album') & (lower(users.access) == 'free'), 1).otherwise(0)) \ .withColumn('free_source_artist', when((lower(streams.source) == 'artist') & (lower(users.access) == 'free'), 1).otherwise(0)) \ .withColumn('free_source_chart', when((lower(streams.source) == 'chart') & (lower(users.access) == 'free'), 1).otherwise(0)) \ .withColumn('free_source_collection', when((lower(streams.source) == 'collection') & (lower(users.access) == 'free'), 1).otherwise(0)) \ .withColumn('free_source_daily_mix', when((lower(streams.source) == 'radio') & (lower(streams.source_uri) == 'dailymix') & (lower(users.access) == 'free'), 1).otherwise(0)) \ .withColumn('free_source_discover_weekly', when((lower(streams.source) == 'others_playlist') & (lower(streams.source_uri) == 'discoverweekly') & (lower(users.access) == 'free'), 1).otherwise(0)) \ .withColumn('free_source_other', when((lower(streams.source) == 'other') & (lower(users.access) == 'free'), 1).otherwise(0)) \ .withColumn('free_source_others_playlist', when((lower(streams.source) == 'others_playlist') & (lower(users.access) == 'free') & (~lower(streams.source_uri).isin('discoverweekly', 'discoverweekly')), 1).otherwise(0)) \ .withColumn('free_source_play_queue', when((lower(streams.source) == 'play_queue') & (lower(users.access) == 'free'), 1).otherwise(0)) \ .withColumn('free_source_radio', when((lower(streams.source) == 'radio') & (lower(streams.source_uri) != 'dailymix') & (lower(users.access) == 'free'), 1).otherwise(0)) \ .withColumn('free_source_release_radar', when((lower(streams.source) == 'others_playlist') & (lower(streams.source_uri) != 'discoverweekly') & (lower(users.access) == 'free'), 1).otherwise(0)) \ .withColumn('free_source_search', when((lower(streams.source) == 'search') & (lower(users.access) == 'free'), 1).otherwise(0)) \ .withColumn('free_source_unknown', when((~lower(streams.source).isin('album', 'artist', 'chart', 'collection', 'other', 'others_playlist', 'play_queue', 'radio', 'search')) & (lower(users.access) == 'free'), 1).otherwise(0)) \ .withColumn('paid_source_album', when((lower(streams.source) == 'album') & (lower(users.access) != 'free'), 1).otherwise(0)) \ .withColumn('paid_source_artist', when((lower(streams.source) == 'artist') & (lower(users.access) != 'free'), 1).otherwise(0)) \ .withColumn('paid_source_chart', when((lower(streams.source) == 'chart') & (lower(users.access) != 'free'), 1).otherwise(0)) \ .withColumn('paid_source_collection', when((lower(streams.source) == 'collection') & (lower(users.access) != 'free'), 1).otherwise(0)) \ .withColumn('paid_source_daily_mix', when((lower(streams.source) == 'radio') & (lower(streams.source_uri) == 'dailymix') & (lower(users.access) != 'free'), 1).otherwise(0)) \ .withColumn('paid_source_discover_weekly', when((lower(streams.source) == 'others_playlist') & (lower(streams.source_uri) == 'discoverweekly') & (lower(users.access) != 'free'), 1).otherwise(0)) \ .withColumn('paid_source_other', when((lower(streams.source) == 'other') & (lower(users.access) != 'free'), 1).otherwise(0)) \ .withColumn('paid_source_others_playlist', when((lower(streams.source) == 'others_playlist') & (lower(users.access) != 'free') & (~lower(streams.source_uri).isin('discoverweekly', 'discoverweekly')), 1).otherwise(0)) \ .withColumn('paid_source_play_queue', when((lower(streams.source) == 'play_queue') & (lower(users.access) != 'free'), 1).otherwise(0)) \ .withColumn('paid_source_radio', when((lower(streams.source) == 'radio') & (lower(streams.source_uri) != 'dailymix') & (lower(users.access) != 'free'), 1).otherwise(0)) \ .withColumn('paid_source_release_radar', when((lower(streams.source) == 'others_playlist') & (lower(streams.source_uri) != 'discoverweekly') & (lower(users.access) != 'free'), 1).otherwise(0)) \ .withColumn('paid_source_search', when((lower(streams.source) == 'search') & (lower(users.access) != 'free'), 1).otherwise(0)) \ .withColumn('paid_source_unknown', when((~lower(streams.source).isin('album', 'artist', 'chart', 'collection', 'other', 'others_playlist', 'play_queue', 'radio', 'search')) & (lower(users.access) != 'free'), 1).otherwise(0)) \ .withColumn('repeat_play', when(lower(streams.repeat_play) == 'y', 1).otherwise(0)) \ .withColumn('shuffle', when(lower(streams.shuffle) == 'y', 1).otherwise(0)) \ .withColumn('free_listeners', when(lower(lower(users.access)) == 'free', streams.user_id)) \ .withColumn('paid_listeners', when(lower(lower(users.access)) != 'free', streams.user_id)) \ .withColumn('free_streams', when(lower(users.access) == 'free', 1).otherwise(0)) \ .withColumn('paid_streams', when(lower(users.access) != 'free', 1).otherwise(0)) \ .groupBy(streams.report_date, tracks.isrc, lower(users.country).alias('country_code')) \ .agg(sum('source_album').alias('source_album'), sum('source_artist').alias('source_artist'), sum('source_chart').alias('source_chart'), sum('source_collection').alias('source_collection'), sum('source_daily_mix').alias('source_daily_mix'), sum('source_discover_weekly').alias('source_discover_weekly'), sum('source_other').alias('source_other'), sum('source_others_playlist').alias('source_others_playlist'), sum('source_play_queue').alias('source_play_queue'), sum('source_radio').alias('source_radio'), sum('source_release_radar').alias('source_release_radar'), sum('source_search').alias('source_search'), sum('source_unknown').alias('source_unknown'), sum('free_source_album').alias('free_source_album'), sum('free_source_artist').alias('free_source_artist'), sum('free_source_chart').alias('free_source_chart'), sum('free_source_collection').alias('free_source_collection'), sum('free_source_daily_mix').alias('free_source_daily_mix'), sum('free_source_discover_weekly').alias('free_source_discover_weekly'), sum('free_source_other').alias('free_source_other'), sum('free_source_others_playlist').alias('free_source_others_playlist'), sum('free_source_play_queue').alias('free_source_play_queue'), sum('free_source_radio').alias('free_source_radio'), sum('free_source_release_radar').alias('free_source_release_radar'), sum('free_source_search').alias('free_source_search'), sum('free_source_unknown').alias('free_source_unknown'), sum('paid_source_album').alias('paid_source_album'), sum('paid_source_artist').alias('paid_source_artist'), sum('paid_source_chart').alias('paid_source_chart'), sum('paid_source_collection').alias('paid_source_collection'), sum('paid_source_daily_mix').alias('paid_source_daily_mix'), sum('paid_source_discover_weekly').alias('paid_source_discover_weekly'), sum('paid_source_other').alias('paid_source_other'), sum('paid_source_others_playlist').alias('paid_source_others_playlist'), sum('paid_source_play_queue').alias('paid_source_play_queue'), sum('paid_source_radio').alias('paid_source_radio'), sum('paid_source_release_radar').alias('paid_source_release_radar'), sum('paid_source_search').alias('paid_source_search'), sum('paid_source_unknown').alias('paid_source_unknown'), sum('repeat_play').alias('repeat_play'), sum('shuffle').alias('shuffle'), countDistinct('free_listeners').alias('free_listeners'), countDistinct('paid_listeners').alias('paid_listeners'), countDistinct(streams.user_id).alias('listeners'), count('free_streams').alias('free_streams'), count('paid_streams').alias('paid_streams'), count(lit(1)).alias('streams')) def main(cli_args): assert cli_args.source_folder.startswith('s3a://'), 'Source path MUST start from s3a://' assert cli_args.target_folder.startswith('s3a://'), 'Target path MUST start from s3a://' args_combinations = list(map( lambda args: dict(zip(('dsp', 'date', 'licensor', 'country'), args)), product(cli_args.dsps, cli_args.dates, cli_args.licensors, cli_args.countries) )) spark = SparkSession.builder \ .appName('Delphi Spark POC') \ .config('spark.jars.packages', 'org.apache.hadoop:hadoop-aws:2.7.3') \ .config('spark.hadoop.mapreduce.fileoutputcommitter.algorithm.version', '2') \ .config('spark.hadoop.mapreduce.fileoutputcommitter.cleanup-failures.ignored', 'true') \ .config('spark.hadoop.parquet.enable.summary-metadata', 'false') \ .config('spark.sql.parquet.mergeSchema', 'false') \ .config('spark.sql.hive.metastorePartitionPruning', 'true') \ .getOrCreate() try: streams = reduce(DataFrame.union, map(lambda args: read_streams(spark, cli_args.source_folder, args), args_combinations) ).distinct() users = reduce(DataFrame.union, map(lambda args: read_users(spark, cli_args.source_folder, args), args_combinations) ).distinct() tracks = reduce(DataFrame.union, map(lambda args: read_tracks(spark, cli_args.source_folder, args), args_combinations) ).distinct() aggregate(streams, users, tracks) \ .coalesce(955) \ .write \ .partitionBy('report_date') \ .parquet(path=cli_args.target_folder, mode='overwrite') finally: spark.stop() if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--dsp', dest='dsps', type=str, nargs='+', required=True) parser.add_argument('--date', '-d', dest='dates', type=str, nargs='+', required=True) parser.add_argument('--target-folder', '-t', dest='target_folder', type=str, required=True) parser.add_argument('--source-folder', '-s', dest='source_folder', type=str, required=True) parser.add_argument('--licensor', '-l', dest='licensors', type=str, nargs='+', default=['*']) parser.add_argument('--country', '-c', dest='countries', type=str, nargs='+', default=['*']) main(parser.parse_args())