如何根据 Pyspark 数据框中的条件设置新的列表值？

Question

我有一个如下所示的 DataFrame。

+---+------------------------------------------+
|id |features                                  |
+---+------------------------------------------+
|1  |[6.629056, 0.26771536, 0.79063195,0.8923] |
|2  |[1.4850719, 0.66458416, -2.1034079]       |
|3  |[3.0975454, 1.571849, 1.9053307]          |
|4  |[2.526619, -0.33559006, -1.4565022]       |
|5  |[-0.9286196, -0.57326394, 4.481531]       |
|6  |[3.594114, 1.3512149, 1.6967168]          |
+---+------------------------------------------+

我想根据我的 where 条件设置一些特征值，如下所示。 IE。其中 id=1、id=2 或 id=6.

我想设置新特征值 id=1，我当前特征值是 [6.629056, 0.26771536, 0.79063195,0.8923]，但我想设置 [0,0,0,0].

我想设置新特征值 id=2，我当前特征值是 [1.4850719, 0.66458416, -2.1034079]，但我想设置 [0,0,0].

我的最终输出将是：

+------+-----------------------------------+
|id  | features                            |
+-----+---------------------------------- -+
|1  | [0, 0, 0, 0]                          |
|2  | [0,0,0]                              |
|3  | [3.0975454, 1.571849, 1.9053307]     |
|4  | [2.526619, -0.33559006, -1.4565022]  |
|5  | [-0.9286196, -0.57326394, 4.481531]  |
|6  | [0,0,0]                              |
+-----+------------------------------------+

Answer 1

如果您要更改一小组 ID，请使用 when 和 otherwise：

df.withColumn("features", 
  when(df.id === 1, array(lit(0), lit(0), lit(0), lit(0)))
  .when(df.id === 2 | df.id === 6, array(lit(0), lit(0), lit(0)))
  .otherwise(df.features)))

它应该比 UDF 快，但是如果要更改的 ID 很多，它很快就会变成很多代码。在这种情况下，使用 philantrovert 的回答中的 UDF。

Answer 2

Shaido 的回答很好，如果你有一组有限的 id，你也知道相应的 feature 的长度。

如果不是这种情况，使用 UDF 应该更干净，并且可以将要转换的 id 加载到另一个 Seq :

在 Scala 中

val arr = Seq(1,2,6)

val fillArray = udf { (id: Int, array: WrappedArray[Double] ) =>
                        if (arr.contains(id) ) Seq.fill[Double](array.length)(0.0) 
                        else array 
                     }

df.withColumn("new_features" , fillArray($"id", $"features") ).show(false)

在Python

from pyspark.sql import functions as f
from pyspark.sql.types import *

arr = [1,2,6]

def fillArray(id, features):
    if(id in arr): return [0.0] * len(features)
    else : return features

fill_array_udf = f.udf(fillArray, ArrayType( DoubleType() ) )

 df.withColumn("new_features" , fill_array_udf( f.col("id"), f.col("features") ) ).show()

输出

+---+------------------------------------------+-----------------------------------+
|id |features                                  |new_features                       |
+---+------------------------------------------+-----------------------------------+
|1  |[6.629056, 0.26771536, 0.79063195, 0.8923]|[0.0, 0.0, 0.0, 0.0]               |
|2  |[1.4850719, 0.66458416, -2.1034079]       |[0.0, 0.0, 0.0]                    |
|3  |[3.0975454, 1.571849, 1.9053307]          |[3.0975454, 1.571849, 1.9053307]   |
|4  |[2.526619, -0.33559006, -1.4565022]       |[2.526619, -0.33559006, -1.4565022]|
|5  |[-0.9286196, -0.57326394, 4.481531]       |[-0.9286196, -0.57326394, 4.481531]|
|6  |[3.594114, 1.3512149, 1.6967168]          |[0.0, 0.0, 0.0]                    |
+---+------------------------------------------+-----------------------------------+

如何根据 Pyspark 数据框中的条件设置新的列表值？

How to set new list value based on condition in dataframe in Pyspark?

apache-spark

apache-spark-sql

pyspark

spark-dataframe

pyspark-sql