Spark MLlib 中的列转换

Question

我已经阅读Spark MLlib doc 进行特征转换，但我仍然对两个简单的案例感到困惑：

1.如何处理单列灵活？ 例如，我有一个名为“date”的列，它的格式是“YYYY-MM-DD”，我想根据“date”生成一个名为“week”的新列。如果使用 pandas.Dataframe，可以使用 Series.apply 来完成，我的问题是如何在 Spark MLlib 中做到这一点？

2.如何根据多列生成新列？ 例如，我想根据支出和收入计算投资回报率，在 pandas.DataFrame 中很简单：

df['roi'] = (df['income'] - df['spend'])/df['spend']

对于 Spark.MLlib，我发现 SQLTransformer 可以用于相同的工作，但我不确定

谁能告诉我如何在 Spark.MLlib 中处理这个问题？非常感谢

Answer 1

一个干净的选项是定义您自己的函数，并使用withColumn() 应用于您的DataFrame。请注意，这与MLlib 无关，因为它指的是Spark 的机器学习模块。

from pyspark.sql.types import FloatType
from pyspark.sql.functions import udf

def roiCalc(income, spend): # Define function
  return((income - spend)/spend)

roiCalculator = udf(roiCalc, FloatType()) # Convert to udf
df.withColumn("roi", roiCalculator(df["income"],df["spend"])) # Apply to df