Pyspark data stratification using "sampleBy" with fractions using multiple columns

Viewed 31

Here's is sample set of data I have created a bin field based on agg_readings. And the Data is so huge with close to 320 Million records stored in hive with parquet format. Of the 320Million, I'm looking to get 5 Million based on stratification. Below is the sample snippet

I have used sampleBy here to fetch the stratified based on two columns. ( Columns are - mnth_src_fld & bin). All I'm looking at the stratified data piece is to get gen_rnd_id unique values across the entire data post stratification, But unfortunately I'm not getting unique gen_rnd_id's. For instance, here in the below sample id's "1805040053", "2352639960" are still getting populated. Based on one source about usage stratification on multiple columns, I have came up with below code, not sure how to get id's from across all bins(1, 2, & 3) to be unique . Any suggestions please. Really appreciate your inputs.

gen_rnd_id        mnth_src_fld  agg_total_readings  bin  insert_dt
1810269893  AUG- GNL DVCS   1421.36 2   9/22/2022
2360812758  AUG- GNL DVCS   1421.32 2   9/22/2022
2561885533  AUG- GNL DVCS   1421.23 1   9/22/2022
2360812759  AUG- GNL DVCS   1421.17 1   9/22/2022
1460501167  AUG- GNL DVCS   1421.11 2   9/22/2022
1515893360  AUG- GNL DVCS   1420.71 2   9/22/2022
1805040053  AUG- GNL DVCS   1419.87 2   9/22/2022
2436175429  AUG- GNL DVCS   1419.63 2   9/22/2022
2941769711  AUG- GNL DVCS   1419.38 2   9/22/2022
2352639960  AUG- GNL DVCS   1417.74 2   9/22/2022
2600364039  AUG- GNL DVCS   643.76  1   9/22/2022
2803486093  AUG- GNL DVCS   643.65  1   9/22/2022
2752468042  AUG- GNL DVCS   643.21  1   9/22/2022
2352639960           AUG-LIR    693797.58   3   9/22/2022
1805040053          AUG-PRO 361753.83   3   9/22/2022
2223875595     AUG-REFRIGERATE  319019.03   3   9/22/2022
2243916002     AUG-REFRIGERATE  230745.32   3   9/22/2022



def get_stratified_split_multiple_columns(input_df, col_name1, col_name2, seed_value, train_frac):
    merged_col_name = "both_labels"
    input_df = input_df.withColumn(merged_col_name, F.concat(F.col(col_name1), F.lit('_#_@_#_'),
                                                             F.col(col_name2))) 
    print(f"DEBUG: Processing Stratification on input_df {input_df} dataset with  {col_name1} & {col_name2} combined fields, seed_value={seed_value}, train_frac={train_frac}\n")
    fractions1 = input_df.select(merged_col_name).distinct().withColumn("fraction",
                                                                        F.lit(train_frac)).rdd.collectAsMap()
    train_df = input_df.stat.sampleBy(merged_col_name, fractions1, seed_value)
    # Delete the merged_col_name
    train_df = train_df.drop(merged_col_name)
    
    return train_df

train_frac_val=0.016

train_df=get_stratified_split_multiple_columns(load_df, 'mnth_src_fld', 'bin', seed_value=10, train_frac=train_frac_val)
0 Answers
Related