Question

我有一个PySpark数据帧df：

+---------+------------------+
|ceil_temp|             test2|
+---------+------------------+
|       -1|[6397024, 6425417]|
|        0|[6397024, 6425417]|
|        0|[6397024, 6425417]|
|        0|[6469640, 6531963]|
|        0|[6469640, 6531963]|
|        1|[6469640, 6531963]|
+---------+------------------+

我最终希望在此数据框中添加一个新列（final），其值是test2列中list的元素，基于ceil_temp列的索引。例如：如果ceil_temp列中包含＆lt; 0或0值，则final列具有test2列的第0个索引中的元素。如下所示：

+---------+------------------+--------
|ceil_temp|             test2|final  |
+---------+------------------+--------
|       -1|[6397024, 6425417]|6397024|
|        0|[6397024, 6425417]|6397024|
|        0|[6397024, 6425417]|6397024|
|        0|[6469640, 6531963]|6469640|
|        0|[6469640, 6531963]|6469640|
|        1|[6469640, 6531963]|6531963|
+---------+------------------+--------

为实现这一目标，我尝试使用flatMap提取ceil_temp和test2作为列表：

m =df.select("ceil_temp").rdd.flatMap(lambda x: x).collect()
q= df.select("test2").rdd.flatMap(lambda x: x).collect()

l=[]
for i in range(len(num)):
    if m[i]<0:
        m[i]=0
    else:
        pass
    l.append(q[i][m[i]])

然后将此列表l转换为新的df并将其与基于窗口函数添加的行索引列的原始数据帧连接：

w = Window().orderBy()
df=df.withColumn("columnindex", rowNumber().over(w)).

但是，flatMap提取的列表顺序似乎与父数据帧df的顺序不一致。我得到以下内容：

m=[-1,0,0,0,0,1]
q=[[6469640, 6531963],[6469640, 6531963],[6469640, 6531963],[6397024, 6425417],[6397024, 6425417],[6397024, 6425417]]

预期结果：

m=[-1,0,0,0,0,1]
q=[[6397024, 6425417],[6397024, 6425417],[6397024, 6425417],[6469640, 6531963],[6469640, 6531963],[6469640, 6531963]]

请告知如何获得“最终”专栏。

Answer 1

我认为您可以在数据帧的行上使用UDF实现所需的结果。

然后你可以使用你的udf的结果withColumn。

val df = spark.sparkContext.parallelize(List(
  (-1, List(6397024, 6425417)),
  (0,List(6397024, 6425417)),
  (0,List(6397024, 6425417)),
  (0,List(6469640, 6531963)),
  (0,List(6469640, 6531963)),
  (1,List(6469640, 6531963)))).toDF("ceil_temp", "test2")

import org.apache.spark.sql.functions.udf
val selectRightElement = udf {
  (ceilTemp: Int, test2: Seq[Int]) => {
    // dummy code for the example
    if (ceilTemp <= 0) test2(0) else test2(1)
  }
}

df.withColumn("final", selectRightElement(df("ceil_temp"), df("test2"))).show

这样做可以防止你的行顺序混乱。

Answer 2

我通过以下方式解决了上述问题：

df=df.withColumn("final",(df.test2).getItem(df.ceil_temp))

从pyspark dataframe列创建列表时，flatMap不保留顺序

2 个答案: