как сделать операцию параллельно с помощью искры

Question

У меня есть следующее data.frame в spark

import findspark
findspark.init()

from pyspark.sql import SparkSession

spark = SparkSession.builder.getOrCreate()
sc = spark.sparkContext
from pyspark.sql import functions as sf
from pyspark.sql.functions import col, when, lit

ddf = spark.createDataFrame([[None, 'Michael',2],
                             [30, 'Andy',3],
                             [19, 'Justin',4],
                             [30, 'James Dr No From Russia with Love Bond',6]],
                            schema=['age', 'name','weights'])
ddf.show()

В этом тривиальном примере я хотел бы создать два столбца: один с weighted.mean из age, если age>29(с именем weighted_age) и другим age^2, если age<=29 (с именем age_squared)

In order to do so I can do this

from pyspark.sql import functions as f
weightedMean = ddf.filter(f.col('age')>29).select(f.sum(f.col('age')*f.col('weights'))/f.sum(f.col('weights'))).first()[0]

ddf.withColumn('weighted_age', f.when(f.col('age') > 29, weightedMean))\
    .withColumn('age_squared', f.when(f.col('age') <= 29, f.col('age')*f.col('age')))\
    .show(truncate=False)

У меня вопрос, есть ли способ выполнить эту операцию параллельно длядва условия if (поэтому создаются два столбца. Один создается при условии возраста> 29 (первое if условие), а другой создается при условии возраста <= 29 (второе <code>if условие))