p_org_id = 'tn-xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx'  
p_start_date = 'YYYY-MM-DD'  
p_end_date = 'YYYY-MM-DD'  
p_dest_bucket_path = 's3://{your_staging_area_bucket}/{your_staging_area_path}/'

# COMMAND ----------

from pyspark.sql.functions import split, explode, expr, regexp_replace, col,from_json  
from pyspark.sql.types import StructType, StructField, LongType, StringType, DoubleType  
from datetime import datetime, timedelta  
from functools import reduce  
from pyspark.sql import DataFrame  
import concurrent.futures

# COMMAND ----------

def convert_to_date(input_date):  
    try:  
        validated_date = datetime.strptime(input_date, '%Y-%m-%d').date()  
        return validated_date  
    except:  
        dbutils.notebook.exit("Invalid Input Date : " + input_date)

# COMMAND ----------

def is_valid_range(start_date, end_date):  
    start = convert_to_date(start_date)  
    end = convert_to_date(end_date)  
    if start > end :  
        dbutils.notebook.exit("Invalid Date Range... Start Date: " + str(start_date) + " / End Date: " + str(end_date)) 

# COMMAND ----------

def get_paths(org_id, start_date, end_date):  
    aux = start_date  
    paths = \[]  
    while aux \<= end_date :  
        year = aux.year  
        month = aux.month  
        day = aux.day  
        source = 's3://pismo-dataplatform-{org_id}/main_stream/{year}/{month}/{day}/_/\_'.format(org_id=org_id, year=year, month=str(month).rjust(2, '0'), day=str(day).rjust(2, '0'))  
        path = {'date': str(aux), 'path': source}  
        paths.append(path)  
        aux = aux + timedelta(days=1)  
    return paths

# COMMAND ----------

SCHEMA = StructType(  
    [  
        StructField('event_id', StringType(), True),  
        StructField('cid', StringType(), True),  
        StructField('org_id', StringType(), True),  
        StructField('event_type', StringType(), True),  
        StructField('domain', StringType(), True),  
        StructField('schema_version', LongType(), True),  
        StructField('arrival_epoch', LongType(), True),  
        StructField('timestamp', StringType(), True),  
        StructField('data', StringType(), True)  
    ]  
)  
is_valid_range(p_start_date, p_end_date)

# COMMAND ----------

START_DATE = convert_to_date(p_start_date)  
END_DATE = convert_to_date(p_end_date)  
ORG_ID = p_org_id.lower().strip()  
EXECUTION_ID = datetime.now().strftime("%Y%m%d-%H%M%S%f")  
DEST_BUCKET_PATH = p_dest_bucket_path

# COMMAND ----------

def generate_output_path(date):  
    output = '{path}/execution_id={execution_id}/org_id={org}/date={date}'.format(path=DEST_BUCKET_PATH.rstrip('/'), org=ORG_ID, execution_id=EXECUTION_ID, date=date)  
    print("Persisting output path: " + output )  
    return output

# COMMAND ----------

def read_write_s3_path(source):  
    path = source['path']  
    date = source['date']  
    df = spark.read.text(path)  
        .withColumn('value', explode(split(regexp_replace(col('value'), r'\}\{', '}\\n{'), '\\n')))  
        .withColumn('value', from_json(col('value'), SCHEMA)).selectExpr("value.\*")  
    output_path = generate_output_path(date)  
    df.repartition(5).write.mode('append').parquet(output_path)  
    optimized_df = spark.read.parquet(output_path)  
    return optimized_df

# COMMAND ----------

def read_paths_in_parallel(org_id, start_date, end_date):  
    s3_paths = get_paths(ORG_ID, START_DATE, END_DATE)  
    with concurrent.futures.ThreadPoolExecutor() as executor:  
        dataframes = list(executor.map(read_write_s3_path, s3_paths))  
        combined_df = reduce(DataFrame.union, dataframes)  
        return combined_df  

# COMMAND ----------

df = read_paths_in_parallel(ORG_ID, START_DATE, END_DATE)

# COMMAND ----------

df.createOrReplaceTempView('read_events')

# COMMAND ----------

query="""  
SELECT \*  
FROM read_events  
WHERE to_date(from_unixtime(arrival_epoch / 1000)) BETWEEN to_date('{START_DATE}') AND to_date('{END_DATE}')  
""".format(START_DATE=START_DATE, END_DATE=END_DATE)  
print(query)

# COMMAND ----------

df_output = spark.sql(query)

# COMMAND ----------

display(df_output)