# `test_data_sf` is a fraction of `train_data`
test_data_sf.with_column(
    'PREDICTION', 
    predict_risk_score(F.object_construct('*')))
.select("TARGET","PREDICTION").show(20)

# ---------------------------------
# |"TARGET"  |"PREDICTION"        |
# ---------------------------------
# |1         |0.8333333333333334  |
# |1         |0.0                 |
# |1         |1.0                 |
# |1         |1.0                 |
# |1         |0.3333333333333333  |
# |0         |0.0                 |
# |1         |0.4                 |
# |0         |0.5                 |
# |1         |0.421875            |
# ---------------------------------

#similary, you can use below SQL as well.

select
    target,
    predict_risk_score(object_construct(*)) as predictions
from
    COVID19_RECORDS_PROCESSED
limit
    100;