Preserve column names from pyspark groupby > agg > arrays_zip > collect_list

Question

I have a dataframe that looks like this:

df = spark.createDataFrame([
    Row(foo=\'a\', bar=1, baz=4),
    Row(foo=\'a\', bar=2, baz=5),
    Row(foo=\'b\', bar=3