bdcc_v3.2

#from __future__ import absolute_import
import logging
import apache_beam as beam
from apache_beam.io import ReadFromText
from apache_beam.io import WriteToText
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.options.pipeline_options import SetupOptions
from apache_beam.options.pipeline_options import GoogleCloudOptions
from apache_beam.options.pipeline_options import StandardOptions

'''
input_filename = "gs://up201405219-final/EVENTS.csv"
output_filename1 = "gs://up201405219-final/output_big/output1.txt"
output_filename2 = "gs://up201405219-final/output_big/output2.txt"
'''
input_filename = "gs://up201405219-final/file_to_an0.csv"
output_filename1 = "gs://up201405219-final/output_teste/output1.txt"
output_filename2 = "gs://up201405219-final/output_teste/output2.txt"

dataflow_options = ['--project=lexical-descent-231918', '--job_name=job1', '--temp_location=gs://up201405219-final/tmp']
dataflow_options.append('--staging_location=gs://up201405219-final/stage')
options = PipelineOptions(dataflow_options)
gcloud_options = options.view_as(GoogleCloudOptions)

options.view_as(StandardOptions).runner = "dataflow"

class Split(beam.DoFn):
    def process(self, element):
        line = element.split(",")
        if line[1] == "":
            subject_id = line[1]
        else:
            subject_id = int(line[1])

        if line[2] == "":
            hospital_id = line[2]
        else:
            hospital_id = int(line[2])

        if line[4] == "":
            item_id = line[4]
        else:
            item_id = int(line[4])

        return [{
            "subject" : subject_id,
            "hospital" : hospital_id,
            "item": item_id
        }]

class CollectSubjectsPerItems(beam.DoFn):
    def process(self, element):
        result = [
            (element["item"], element["subject"])
        ]
        return result

class CollectItemsPerHospitals(beam.DoFn):
    def process(self, element):
        result = [
            (element["hospital"], element["item"])
        ]
        return result

class CollectSubjectsPerHospitals(beam.DoFn):
    def process(self, element):
        result = [
            (element["subject"], element["item"])
        ]
        return result

class CollectDistinct(beam.DoFn):
    def process(self, element):
        res = []
        alll =[]
        count = 0
        for i in element[1]:
            count+=1
            alll.append(i)
            if i not in res:
                res.append(i)

        result = [
            (element[0], [len(res), res, len(alll),  alll])
        ]
        return result

class CollectAll(beam.DoFn):
    def process(self, element):
        res = []
        for i in element[1]:
            res.append(i)
        result = [
            (element[0], [len(res)])
        ]
        return result

class WriteToCSV(beam.DoFn):
    def process(self, element):
        result = [
            "{}.{}".format(element[0],
            element[1]
            )
        ]
        return result

class WriteToCSV2(beam.DoFn):
    def process(self, element):
        result = [
            "{},{},{}".format(element[0],
            element[1]['subject'][0],
            element[1]['item'][0])
        ]
        return result

with beam.Pipeline(options = options) as p:
    rows = (
        p |
        ReadFromText(input_filename) |
        beam.ParDo(Split())
    )
    # para cada item, quantos pacientes diferentes o tiveram
    subjects_per_items = (
        rows |
        beam.ParDo(CollectSubjectsPerItems()) |
        "Groupping subjectsPerItems" >> beam.GroupByKey() |
        "Distinct  subjectsPerItems" >> beam.ParDo(CollectDistinct()) |
        "To CSV 1" >> beam.ParDo(WriteToCSV()) |
        "Write 1" >> WriteToText(output_filename1)
        #"Counting subjects" >> beam.CombineValues(beam.combiners.CountCombineFn())
    )
    # para cada hospital, quantos items diferentes aconteceram
    items_per_hospitals = (
        rows |
        beam.ParDo(CollectItemsPerHospitals()) |
        "Groupping hospitalsPerItems" >> beam.GroupByKey() |
        "Distinct hospitalsPerItems" >> beam.ParDo(CollectAll())
    )

    subjects_per_hospitals = (
        rows |
        beam.ParDo(CollectSubjectsPerHospitals()) |
        "Groupping subjectsPerHospital" >> beam.GroupByKey() |
        "Distinct subjectsPerHospital" >> beam.ParDo(CollectAll())
    )

    to_be_joined = (
        {
            'subject': subjects_per_hospitals,
            'item': items_per_hospitals
        } |
        beam.CoGroupByKey() |
        "To CSV 2" >> beam.ParDo(WriteToCSV2()) |
        "Write 2" >> WriteToText(output_filename2)
    )