# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: PySpark
#     language: python
#     name: pyspark
# ---

# # Yelp: Make Test Data

# !spark-shell --version

# <details>
# <summary><b>Table of Contents</b> (click to open)</summary>
# <!-- MarkdownTOC -->
#
# 1. [Load Full Dataset](#data)
# 1. [Make Input Test Data](#Source-Test-Data)
# 1. [Make Expected Test Data](#Make-Expected-Test-Data)
#     1. [Testcase: Top Businesses](#Testcase:-Top-Businesses)
#     1. [Testcase: Top Restaurants](#Testcase:-Top-Restaurants)
#  1. [Copy Test Data to Project Folder](#Copy-Test-Data-to-Project-Folder)
#     1. [Input Test Data](#Input-Test-Data)
#     1. [Expected Test Data](#Expected-Test-Data)
#     
# <!-- /MarkdownTOC -->
# </details>

# Import Required Libraries
import pyspark.sql.functions as F

# ### Load Data  <a class="anchor" id="data"></a>

# Load the data from google cloud storage to dataproc

file_review = "gs://proj-spark/yelp-dataset/yelp_academic_dataset_review.json"
file_business = "gs://proj-spark/yelp-dataset/yelp_academic_dataset_business.json"
file_users = "gs://proj-spark/yelp-dataset/yelp_academic_dataset_user.json"
file_checkin = "gs://proj-spark/yelp-dataset/yelp_academic_dataset_checkin.json"
file_tip = "gs://proj-spark/yelp-dataset/yelp_academic_dataset_tip.json"

# +
df_review = spark.read.format('json').load(file_review)
df_review.createOrReplaceTempView('reviewTable')

df_business = spark.read.format('json').load(file_business)
df_business.createOrReplaceTempView('businessTable')

df_users = spark.read.format('json').load(file_users)
df_users.createOrReplaceTempView('usersTable')

df_checkin = spark.read.format('json').load(file_checkin)
df_checkin.createOrReplaceTempView('checkinTable')

df_tip = spark.read.format('json').load(file_tip)
df_tip.createOrReplaceTempView('checkinTable')
# -

# ### Source Test Data

expected_data_path = '/home/bilal/study_private/tests/test_data/expected_data/'
source_data_path = '/home/bilal/study_private/tests/test_data/source_data/'


def make_checkin_data(df_checkin):
    df_checkin_test = df_checkin.sample(withReplacement=False, fraction=0.05, seed=50)\
                                .limit(200)\
                                .coalesce(1)
    df_checkin_test.write.format('json').save('source_data/yelp_academic_dataset_checkin.json')


def make_business_data(df_business):
    df_checkin_test = spark.read.format('json').load('source_data/yelp_academic_dataset_checkin.json')
    joinExpr = df_business.business_id == df_checkin_test.business_id
    df_business_test = df_business.join(F.broadcast(df_checkin_test), joinExpr, 'inner')\
                                  .drop(df_checkin_test.business_id)\
                                  .coalesce(1)
    df_business_test.write.format('json').save('source_data/yelp_academic_dataset_business.json')


def verify(df1, df2, key):
    df1_key = df1.groupby(key).agg(F.count(key))
    df2_key = df2.groupby(key).agg(F.count(key))
    if df1_key.subtract(df2_key).rdd.isEmpty():
        return df2_key.subtract(df1_key).rdd.isEmpty()
    return False


make_checkin_data(df_checkin)
make_business_data(df_business)

df_checkin_test = spark.read.format('json').load('source_data/yelp_academic_dataset_checkin.json')
df_business_test = spark.read.format('json').load('source_data/yelp_academic_dataset_business.json')

# verify that both dataframes have same business_ids
verify(df_checkin_test, df_business_test, 'business_id')

# ###  Make Expected Test Data

# load data from gcs to dataproc
df_checkin_test = spark.read.format('json')\
            .load('gs://proj-spark/yelp-test-dataset/yelp_academic_dataset_checkin.json')
df_business_test = spark.read.format('json')\
            .load('gs://proj-spark/yelp-test-dataset/yelp_academic_dataset_business.json')

# #### Testcase: Top Businesses

# +
# pandas dataframes
pdf_checkin_test = df_checkin_test.toPandas()
pdf_business_test = df_business_test.toPandas()

pdf_checkin_test['checkin_count'] = pdf_checkin_test.date.str.split(',').map(len)
pdf_checkin_bs = pdf_checkin_test.merge(pdf_business_test, on='business_id')\
                                 .sort_values(by=['checkin_count','business_id'], ascending=False)

pdf_expected_data_tc2 = pdf_checkin_bs[['business_id', 'name', 'city', 'state', 'stars', 'checkin_count', 'review_count']]
# -

pdf_expected_data_tc2.to_csv('top_businesses.csv', index=False)

# #### Testcase: Top Restaurants

# +
pdf_checkin_rest = pdf_checkin_bs[pdf_checkin_bs.categories.str.contains('res', case=False)]
pdf_checkin_rest = pdf_checkin_rest[pdf_checkin_rest.state=='MA']
pdf_checkin_rest = pdf_checkin_rest.sort_values(by=['checkin_count','business_id'], ascending=False)

pdf_expected_data_tc3 = pdf_checkin_rest[['business_id', 'name', 'city', 'state', 'stars', 'checkin_count', 'review_count']]
# -

pdf_expected_data_tc3.to_csv('top_restaurants.csv', index=False)

# #### Testcase: Business Categories

# +
# transform categories to category
pdf_business_test = df_business_test.toPandas()
pdf_business_test['category'] = pdf_business_test.categories.str.split(', ')
pdf_business_test = pdf_business_test.explode('category')

# aggregate and sort
pdf_cat_grp = pdf_business_test.groupby('category')[['business_id']]\
                               .agg('count')\
                               .rename(columns={'business_id':'business_count'})\
                               .reset_index()
pdf_expected_data_tc4 = pdf_cat_grp.sort_values(['business_count','category'], ascending=False)
# -

pdf_expected_data_tc4.to_csv('business_categories.csv', index=False)

# +
# Split the categories 
df_business_test = df_business_test.withColumn('categories_splitted', F.split(F.col('categories'),', '))\
    .withColumn('category', F.explode('categories_splitted'))

cat_grp = df_business_test.groupBy('category')\
    .agg(F.count('business_id').alias('business_count'))\
    .orderBy(['business_count','category'], ascending=False)
# -

pdf_expected_data_tc4.reset_index(drop=True)

# ### Copy Test Data to Project Folder

# ### Input Test Data

# Copy Source test data to project directory and google cloud storage

# +
# list files in source data
# !hdfs dfs -ls  /user/root/source_data

# # copy source_data from HDFS to Master Node
# !hdfs dfs -copyToLocal /user/root/source_data source_data

# Move source_data to project folder
# !mv source_data/yelp_academic_dataset_checkin.json/part-*.json \
#     $source_data_path/yelp_academic_dataset_checkin.json
# !mv source_data/yelp_academic_dataset_business.json/part-*.json \
#     $source_data_path/yelp_academic_dataset_business.json

# Copy source_data to google cloud storage
# !gsutil cp $source_data_path/yelp_academic_dataset_checkin.json \
#            gs://proj-spark/yelp-test-dataset
# !gsutil cp $source_data_path/yelp_academic_dataset_business.json \
#            gs://proj-spark/yelp-test-dataset
    
# Delete Data from HDFS 
# !hdfs dfs -rm -r  /user/root/source_data

# Delete Data from Master Node
# !rm -r source_data
# -

# ### Expected Test Data

# Copy expected data to project folder

# !mv top_businesses.csv $expected_data_path/top_businesses.csv
# !mv top_restaurants.csv $expected_data_path/top_restaurants.csv
# !mv business_categories.csv $expected_data_path/business_categories.csv



















StarcoderdataJupyterScriptsDedupFiltered