application_vnd.jupyter.stderr_output_1.png
from pyspark.sql import SparkSession

import json

import os

spark = SparkSession.builder.appName('SimpleApp').getOrCreate()
sc = spark.sparkContext

def findArticlePath(article_sha_id):
    """
    This function finds the full path given an article_sha_id
    """
    ROOT_PATH = '/kaggle/input/CORD-19-research-challenge/biorxiv_medrxiv/biorxiv_medrxiv/'
    FILE_SUFFIX = '.json'
    article_path = ROOT_PATH + article_sha_id + FILE_SUFFIX
    return article_path

def retrieveJson(article_sha_id):
    """
    Given a 1-word string containing a JSON key, return the data for those keys.
    Also return the location of those keys?
    """
    article_file_path = findArticlePath(article_sha_id)
    with open(article_file_path, 'r') as read_file:
        json_string = json.load(read_file)
    return json_string

def retrieveKey(json_input, key_input):
    """
    Uses retrieveTopKey and retrieveSubKey to return key values from anywhere in the JSON data string.
    """
    result = []
    istopitem = False
    if not isinstance(json_input, list):
        json_string = json_input
        json_input = []
        json_input.append(json_string)
        istopitem = True
    if istopitem:
        for key in json_string:
            if key == key_input:
                top_key_value = json_string[key]
                result.append(top_key_value)
                return result
    for json_string in json_input:
        if isinstance(json_string, dict):
            for key in json_string:
                top_key_value = json_string[key]
                if isinstance(top_key_value, dict):
                    for sub_key in top_key_value:
                        sub_key_value = top_key_value[sub_key]
                        if sub_key == key_input:
                            result.append(sub_key_value)
                        else:
                            sub_result = retrieveKey(sub_key_value, key_input)
                            if len(sub_result) != 0:
                                result.append(sub_result)
                elif isinstance(top_key_value, list):
                    for top_key_value_item in top_key_value:
                        if isinstance(top_key_value_item, dict):
                            for sub_key in top_key_value_item:
                                sub_key_value = top_key_value_item[sub_key]
                                if sub_key == key_input:
                                    result.append(sub_key_value)
                                else:
                                    sub_result = retrieveKey(sub_key_value, key_input)
                                    if len(sub_result) != 0:
                                        result.append(sub_result)
                        else:
                            result.append(top_key_value_item)
    return result

def searchWordInTitle(word):
    article_list = []
    for files in os.listdir(path='/kaggle/input/CORD-19-research-challenge/biorxiv_medrxiv/biorxiv_medrxiv/'):
        article_sha_id = files.split('.')[0]
        article_json = retrieveJson(article_sha_id)
        title = retrieveKey(article_json, 'title')
        if isinstance(title, list):
            for item in title:
                if isinstance(item, list):
                    for it in item:
                        find_result = it.find(word)
                        if find_result >= 0:
                            if not article_sha_id in article_list:
                                article_list.append(article_sha_id)
                else:
                    find_result = item.find(word)
                    if find_result >= 0:
                        if not article_sha_id in article_list:
                            article_list.append(article_sha_id)
        else:
            find_words = title.find(word)
            if find_words >= 0:
                if not article_sha_id in article_list:
                    article_list.append(article_sha_id)
    return article_list

def abstract(data):
    """ Abstract """
    data_item = data[0]
    row_entries = data_item[0]
    for row in row_entries:
        cite_span = row[0]
        ref_span = row[1]
        section = row[2]
        text = row[3]
    return (cite_span, ref_span, section, text)

def back_matter(data):
    """ Back Matter """
    data_item = data[0]
    row_entries = data_item[0]
    for row in row_entries:
        cite_span = row[0]
        ref_span = row[1]
        section = row[2]
        text = row[3]
    return (cite_span, ref_span, section, text)

def bib_entries(data):
    """ Bib Entries """
    data_item = data[0]
    row_entries = data_item[0]
    for row in row_entries:
        cite_span = row[0]
        ref_span = row[1]
        section = row[2]
        text = row[3]
    return (cite_span, ref_span, section, text)

def body_text(data):
    """ Body Text """
    data_item = data[0]
    row_entries = data_item[0]
    for row in row_entries:
        cite_span = row[0]
        ref_span = row[1]
        section = row[2]
        text = row[3]
    return (cite_span, ref_span, section, text)

def metadata(data):
    """ Metadata """
    data_item = data[0]
    row_entries = data_item[0]
    num_entries = len(row_entries)
    i = 0
    for row in row_entries:
        if i < num_entries - 1:
            rowitems = row_entries[i][0]
            affiliation = rowitems[0]
            institution = affiliation[0]
            laboratory = affiliation[1]
            location = affiliation[2]
            email = rowitems[1]
            first = rowitems[2]
            last = rowitems[3]
            middle = rowitems[4]
            suffix = rowitems[5]
        else:
            title = row
        i += 1
    return (affiliation, institution, laboratory, location, email, first, last, middle, suffix)

def paper_id(data):
    """ Paper Id """
    data_item = data[0]
    paper_id = data_item[0]
    return paper_id

def ref_entries(data):
    """ Ref Entries """
    data_item = data[0]
    row_entries = data_item[0]
    for row in row_entries:
        latex = row[0]
        text = row[1]
        type_data = row[2]
    return (latex, text, type_data)

column_functions = {0: abstract, 1: back_matter, 2: bib_entries, 3: body_text, 4: metadata, 5: paper_id, 6: ref_entries}

def returnTextualReferences(article_list):
    article_list = []
    for files in os.listdir(path='/kaggle/input/CORD-19-research-challenge/biorxiv_medrxiv/biorxiv_medrxiv/'):
        article_sha_id = files.split('.')[0]
        article_path = retrieveArticlePath(article_sha_id)
        article_df = spark.read.json(article_path)
        article_df.createOrReplaceTempView('article')
        column_results = []
        column_names = []
        for column_name in df.schema.names:
            strSQL = 'SELECT ' + column_name + ' from article'
            column_result = spark.sql(strSQL)
            column_results.append(column_result)
            column_names.append(column_name)
        i = 0
        for column_result in column_results:
            column_name = column_names[i]
            column_schema = column_result.schema
            column_rdd = column_result.rdd
            result = column_result.collect()
            func = column_functions.get(i)
            item_result = func(result)
            i += 1

vaccine_in_title = searchWordInTitle('vaccine')
returnTextualReferences(vaccine_in_title)