Data science productionization: scale - example 2

ds_prod_scale2.py

outcome_sdf = (
    sdf
    .select(
      f.create_map(
          f.col('unique_id'), 
          f.col('feature_list')
      ).alias('feature_map'),
    )
    .groupby(
        f.floor(f.rand() * nparts).alias('grouper')
    )
    .agg(
        f.collect_list(f.col('feature_map')).alias('feature_map')
    )
    .select(
        f.explode(predict_new_udf(f.col('feature_map'))).alias('unique_id', 'probability_estimate')
    )
)