I have two dataframes df1
+---+---+----------+
| n|val| distances|
+---+---+----------+
| 1| 1|0.27308652|
| 2| 1|0.24969208|
| 3| 1|0.21314497|
+---+---+----------+
and df2
+---+---+----------+
| x1| x2| w|
+---+---+----------+
| 1| 2|0.03103427|
| 1| 4|0.19012526|
| 1| 10|0.26805446|
| 1| 8|0.26825935|
+---+---+----------+
I want to add a new column to df1
called gamma
, which will contain the sum of the w
value from df2
when df1.n == df2.x1 OR df1.n == df2.x2
I tried to use udf, but apparently selecting from the different dataframe will not work, because values should be determined before calculation
gamma_udf = udf(lambda n: float(df2.filter("x1 = %d OR x2 = %d"%(n,n)).groupBy().sum('w').rdd.map(lambda x: x).collect()[0]), FloatType())
df1.withColumn('gamma1', gamma_udf('n'))
Is there any way of doing it with join
or groupby
without using cycles?
df1.join(df2, (df1.n == df2.x1) | (df1.n == df2.x2)).groupBy(df1.n).sum("w")
? – Alper t. Turker