Pyspark agg max abs val pero mantener el signo

Estoy agg en otro col. Por ejemplo, si el col era

ID val
A 10
A 100
A -150
A 15
B 10
B 200
B -150
B 15

Me gustaría devolver lo siguiente (manteniendo el cartel). No estoy seguro de cómo hacer esto manteniendo el letrero

ID max(val)
A -150
B 200

Solución del problema

Esto se puede hacer usando una función de ventana + número_de_fila. Particionamos IDy ordenamos abs(val)descendiendo. Entonces simplemente tomamos la primera fila.

import pyspark.sql.functions as F
data = [
 ('A', 10),
 ('A', 100),
 ('A', -150),
 ('A', 15),
 ('B', 10),
 ('B', 200),
 ('B', -150),
 ('B', 15)
]
df = spark.createDataFrame(data=data, schema=('ID','val'))
w = Window().partitionBy('ID').orderBy(F.abs('val').desc())
(df
.withColumn('rn', F.row_number().over(w))
.filter(F.col('rn') == 1)
.drop('rn')
).show()
+---+----+     
| ID| val|
+---+----+
| A|-150|
| B| 200|
+---+----+