如何在PySpark中使用UDF处理概率向量？

我有一个包含如下密集向量的列：

{"vectorType": "dense", "length": 7, "values": [0.6514390482520993, 0.195419935311166, 0.07402558401796841, 0.032660728394633604, 0.01775634445896352, 0.009769471437190208, 0.018928888127979045]}

我希望创建一个UDF函数来处理这些概率值，代码如下：

@udf("integer")
def process_probability(x):
  vector_values = x['values']
  if vector_values[0] > 0.65: return 0
  elif vector_values[1] > 0.198: return 1
  elif vector_values[2] > 0.08: return 2
  elif vector_values[3] > 0.035: return 3
  elif vector_values[4] > 0.045: return 4
  elif vector_values[5] > 0.00976: return 5
  elif vector_values[6] > 0.0188: return 6

执行以下代码时出现错误：

processed_predictions = predictions.withColumn('processed_prediction', process_probability(predictions.probability))

错误信息：

File "<command-8587444>", line 3, in process_probability
  File "/spark/python/pyspark/ml/linalg/__init__.py", line 394, in __getitem__
    return self.array[item]
IndexError: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices

问题原因 #

Spark的DenseVector对象不是字典结构，你看到的JSON格式只是向量的字符串展示形式，实际传入UDF的是pyspark.ml.linalg.DenseVector实例，它不支持x['values']这种字典式的键访问方式。

修改后的UDF代码 #

直接通过索引访问向量元素即可，或者先转换为numpy数组再操作：

方式一：直接索引访问 #

from pyspark.sql.functions import udf

@udf("integer")
def process_probability(x):
    if x[0] > 0.65: return 0
    elif x[1] > 0.198: return 1
    elif x[2] > 0.08: return 2
    elif x[3] > 0.035: return 3
    elif x[4] > 0.045: return 4
    elif x[5] > 0.00976: return 5
    elif x[6] > 0.0188: return 6

方式二：转为numpy数组操作 #

from pyspark.sql.functions import udf

@udf("integer")
def process_probability(x):
    vector_values = x.toArray()
    if vector_values[0] > 0.65: return 0
    elif vector_values[1] > 0.198: return 1
    elif vector_values[2] > 0.08: return 2
    elif vector_values[3] > 0.035: return 3
    elif vector_values[4] > 0.045: return 4
    elif vector_values[5] > 0.00976: return 5
    elif vector_values[6] > 0.0188: return 6

额外优化建议 #

• 建议添加默认返回值，避免所有条件都不满足时返回None，比如在函数末尾加return -1或其他自定义标识值。
• 如果是分类任务的概率向量，若需求是取概率最大的索引，可直接使用Spark内置的argmax函数，但你的逻辑是按顺序判断阈值，因此仍需自定义UDF。

内容的提问来源于stack exchange，提问作者Gonzalo Garcia