<jupyter_start><jupyter_text><jupyter_code>super_srag = spark.read.parquet('gs://ai-covid19-datalake/standard/super-srag/super_srag_v2.parquet').persist(StorageLevel.MEMORY_ONLY)
super_srag.limit(2).toPandas()<jupyter_output><empty_output><jupyter_text><jupyter_code># prefixes
rose_prefix = 'gs://ai-covid19-datalake/trusted/experiment_map/base_rose_'
# suffix
suffix = '_.parquet'<jupyter_output><empty_output><jupyter_text><jupyter_code>cols_sets = {'cols_set_1': ['NU_NOTIFIC', 'CLASSI_FIN', 'CRITERIO', 'EVOLUCAO', 
                            'AGE_GROUP', 'DIST_PRI_NOTIFIC_Q',
                            'DIST_PRI_INTERNA_Q', 'DIST_PRI_ENTUTI_Q', 'DIST_PRI_SAIDUTI_Q', 'DIST_PRI_EVOLUCA_Q', 'DIST_PRI_ENCERRA_Q',
                            'SYMP_GROUP1', 'SYMP_GROUP2', 'SYMP_GROUP3', 'SYMP_GROUP4',
                            'RF_GROUP1', 'RF_GROUP2', 'RF_GROUP3', 'RF_GROUP4',
                            'SUPORT_VEN', 'UTI', 'HOSPITAL',
                            'DIST_PRI_RAIOX_Q', 'DIST_PRI_COLETA_Q', 'DIST_PRI_TOMO_Q', 'DIST_PRI_IF_Q', 'DIST_PRI_TRA_Q', 'DIST_PRI_PCR_Q', 'DIST_PRI_SOR_Q',
                            'GMR_TRANSIT_STATIONS_1WEEK_BEFORE_Q', 'GMR_RETAIL_AND_RECREATION_1WEEK_BEFORE_Q', 'GMR_RESIDENTIAL_PERCENT_1WEEK_BEFORE_Q', 'GMR_WORKPLACES_PERCENT_1WEEK_BEFORE_Q',
                            'GMR_TRANSIT_STATIONS_2WEEKS_Q', 'GMR_RETAIL_AND_RECREATION_2WEEKS_Q', 'GMR_RESIDENTIAL_PERCENT_2WEEKS_Q', 'GMR_WORKPLACES_PERCENT_2WEEKS_Q',
                            'INMET_RELATIVE_AIR_HUMIDITY_1WEEK_BEFORE_Q', 'INMET_RELATIVE_AIR_HUMIDITY_2WEEKS_BEFORE_Q'],
             'cols_set_2': ['NU_NOTIFIC', 'CS_SEXO','CS_RACA', 'CRITERIO', 'SURTO_SG', 'SUPORT_VEN', 'EVOLUCAO',
                            'OUTRO_SIN', 'AGE_GROUP', 
                            'SYMP_GROUP1',  'SYMP_GROUP2', 'SYMP_GROUP3', 'SYMP_GROUP4',
                            'RF_GROUP1', 'RF_GROUP2', 'RF_GROUP3',  'RF_GROUP4',
                            'DIST_PRI_EVOLUCA_Q', 'DIST_PRI_ENCERRA_Q', 'DIST_PRI_INTERNA_Q',
                            'DIST_PRI_NOTIFIC_Q', 'DIST_PRI_COLETA_Q', 'DIST_PRI_PCR_Q', 'CLASSI_FIN'],
             'cols_set_3': ['NU_NOTIFIC', 'SYMP_GROUP1', 'SYMP_GROUP2', 'SYMP_GROUP3', 'SYMP_GROUP4',
                            'AGE_GROUP', 'CS_SEXO', 'CS_RACA', 'EVOLUCAO', 'CS_ESCOL_N',
                            'UTI', 'SUPORT_VEN', 'DIST_PRI_ENTUTI_Q', 'DIST_PRI_INTERNA_Q',
                            'CLASSI_FIN']}
rose_df = spark.read.parquet(rose_prefix+'1'+suffix)
rose_df.limit(3).toPandas()<jupyter_output><empty_output><jupyter_text><jupyter_code>def dropDupeDfCols(df):
    newcols = []
    dupcols = []

    for i in range(len(df.columns)):
        if df.columns[i] not in newcols:
            newcols.append(df.columns[i])
        else:
            dupcols.append(i)

    df = df.toDF(*[str(i) for i in range(len(df.columns))])
    for dupcol in dupcols:
        df = df.drop(str(dupcol))

    return df.toDF(*newcols)

for id_exp in range(3):
    for cols in list(cols_sets.keys()):
        rose_df = spark.read.parquet(rose_prefix+'1'+suffix).select(['NU_NOTIFIC', 'CLASSI_FIN'])
        rose_df = rose_df.withColumnRenamed('CLASSI_FIN', 'rose_nofilter_nostratsamp_'+str(id_exp))

        super_srag_ = rose_df.join(super_srag, 'NU_NOTIFIC', 'left')

        experiment_name = rose_df.columns[1]

        nominal_cols = ['CS_SEXO', 'SG_UF']

        super_srag_ = super_srag_.select(cols_sets[cols]+[experiment_name])

        super_srag_ = super_srag_.na.fill('9999')
        super_srag_ = super_srag_.na.fill(9999)

        for col in super_srag_.columns:
            if col in nominal_cols:
                super_srag_ = super_srag_.withColumn('CS_SEXO', F.when(F.col('CS_SEXO') == 'F', '2').otherwise('1'))
            if col not in ['NU_NOTIFIC', experiment_name]:
                super_srag_ = super_srag_.withColumn(col, F.col(col).cast('float'))

        super_srag_ = super_srag_.dropna()

        super_srag_ = super_srag_.withColumn('CLASSI_FIN', F.when(F.col(experiment_name) == '5', 1.0).otherwise(0.0)).drop(experiment_name)

        super_srag_ = dropDupeDfCols(super_srag_)
        
        path_str = 'gs://ai-covid19-datalake/trusted/experiment_map/datasets/rose_nofilter_nostratsamp_'+cols+"_"+str(id_exp)
        super_srag_.write.parquet(path_str)
        print('{}: {}'.format(id_exp, cols))<jupyter_output>0: cols_set_1
0: cols_set_2
0: cols_set_3
1: cols_set_1
1: cols_set_2
1: cols_set_3
2: cols_set_1
2: cols_set_2
2: cols_set_3
{'fraction_non_alphanumeric': 0.1344575105391613, 'fraction_numerical': 0.017528289327712446, 'mean_word_length': 3.1471941122355105, 'pattern_counts': {'":': 0, '<': 16, '<?xml version=': 0, '>': 16, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '197510', 'license': 'MIT', 'n_tokens_mistral': 1876, 'n_tokens_neox': 1831, 'n_words': 204, 'path': '04-preprocessing-for-models/02-experiment-design/06c-building-rose-sampling-datasets.ipynb', 'repo_name': 'pierrepita/super-srag-ai-covid19-google-cidacs', 'size': '4507'}