Python 如何在PySpark ML管道中添加udf作为阶段？

样本df：

+-----------+
|pred_vector|
+-----------+
| [0.5, 0.6]|                  
| [0.7, 0.8]|                   
| [1.1, 1.5]|                                 
+-----------+

我使用以下函数从稠密向量和 转换为浮动：

from pyspark.sql.types import DoubleType
from pyspark.sql.functions import lit, udf

def ith_(v, i):
    try:
        return float(v[i])
    except ValueError:
        return None

ith = udf(ith_, DoubleType())

df.select(ith("features", lit(1))).show()

## +-----------------+
## |ith_(features, 1)|
## +-----------------+
## |              0.5|
## |              0.7|
## |              1.1|
## +-----------------+

使用下面的代码，我想构建一个定制的转换器，我可以将它附加到我的管道阶段中。它应该从df中pred_vector列中的向量中提取第一个值，将其转换为float并保存在df中的新列中

# CUSTOM TRANSFORMER ----------------------------------------------------------------
class TypeConverter(Transformer):
    """
    A custom Transformer which converts 'pred_vector' column dense vector type to float.
    """

    def __init__(self):
        super(TypeConverter, self).__init__()


    def _transform(self, df: DataFrame) -> DataFrame:
        df = df.withColumn("new_column", ith("pred_vector"))
        return df

---

如何在自定义转换器中匹配udf？

当可以使用element_at时，为什么要使用udf。element_at用于从数组中提取根据文档，我正在尝试从密集向量中提取，另一个问题是在我的管道中附加此步骤

df:
+-----------+-----------+
|pred_vector|new_column |
+-----------+-----------+
| [0.5, 0.6]| 0.5       |          
| [0.7, 0.8]| 0.7       |            
| [1.1, 1.5]| 1.1       |                         
+-----------+-----------+