如何从sparkdataframe列中的数组中提取值
我在spark sql中得到了一个这样的dataframe：
scala> result.show +-----------+--------------+ |probability|predictedLabel| +-----------+--------------+ | [0.0,1.0]| 0.0| | [0.0,1.0]| 0.0| | [0.0,1.0]| 0.0| | [0.0,1.0]| 0.0| | [0.0,1.0]| 0.0| | [0.1,0.9]| 0.0| | [0.0,1.0]| 0.0| | [0.0,1.0]| 0.0| | [0.0,1.0]| 0.0| | [0.0,1.0]| 0.0| | [0.0,1.0]| 0.0| | [0.0,1.0]| 0.0| | [0.1,0.9]| 0.0| | [0.6,0.4]| 1.0| | [0.6,0.4]| 1.0| | [1.0,0.0]| 1.0| | [0.9,0.1]| 1.0| | [0.9,0.1]| 1.0| | [1.0,0.0]| 1.0| | [1.0,0.0]| 1.0| +-----------+--------------+ only showing top 20 rows 我想用名为prob的新列创建一个新的dataframe，这是原始dataframe的概率列中的第一个值，如下所示：
+-----------+--------------+----------+ |probability|predictedLabel| prob | +-----------+--------------+----------+ | [0.0,1.0]| 0.0| 0.0| | [0.0,1.0]| 0.0| 0.0| | [0.0,1.0]| 0.0| 0.0| | [0.0,1.0]| 0.0| 0.0| | [0.0,1.0]| 0.0| 0.0| | [0.1,0.9]| 0.0| 0.0| | [0.0,1.0]| 0.0| 0.0| | [0.0,1.0]| 0.0| 0.0| | [0.0,1.0]| 0.0| 0.0| | [0.0,1.0]| 0.0| 0.0| | [0.0,1.0]| 0.0| 0.0| | [0.0,1.0]| 0.0| 0.0| | [0.1,0.9]| 0.0| 0.1| | [0.6,0.4]| 1.0| 0.6| | [0.6,0.4]| 1.0| 0.6| | [1.0,0.0]| 1.0| 1.0| | [0.9,0.1]| 1.0| 0.9| | [0.9,0.1]| 1.0| 0.9| | [1.0,0.0]| 1.0| 1.0| | [1.0,0.0]| 1.0| 1.0| +-----------+--------------+----------+
我怎样才能做到这一点？ 谢谢！
您可以使用Dataset的functions和精彩的functions 库来完成您的需求：
result.withColumn("prob", $"probability".getItem(0))
这会添加一个名为prob的新Column其值通过获取数组中的第一项（在索引0处 – 我们毕竟是计算机科学家）从probability Column派生而来。
我还要提到UDF应该是你的最后手段，因为Catalyst优化器目前无法优化UDF，所以你应该总是更喜欢内置函数来充分利用Catalyst。
如果使用Spark UDF，则相当简单。 喜欢这个：
val headValue = udf((arr: Seq[Double]) => arr.head) result.withColumn("prob", headValue(result("probability"))).show
它会给你想要的输出：
+-----------+--------------+----------+ |probability|predictedLabel| prob | +-----------+--------------+----------+ | [0.0,1.0]| 0.0| 0.0| | [0.0,1.0]| 0.0| 0.0| | [0.0,1.0]| 0.0| 0.0| | [0.0,1.0]| 0.0| 0.0| | [0.0,1.0]| 0.0| 0.0| | [0.1,0.9]| 0.0| 0.0| | [0.0,1.0]| 0.0| 0.0| | [0.0,1.0]| 0.0| 0.0| | [0.0,1.0]| 0.0| 0.0| | [0.0,1.0]| 0.0| 0.0| | [0.0,1.0]| 0.0| 0.0| | [0.0,1.0]| 0.0| 0.0| | [0.1,0.9]| 0.0| 0.1| | [0.6,0.4]| 1.0| 0.6| | [0.6,0.4]| 1.0| 0.6| | [1.0,0.0]| 1.0| 1.0| | [0.9,0.1]| 1.0| 0.9| | [0.9,0.1]| 1.0| 0.9| | [1.0,0.0]| 1.0| 1.0| | [1.0,0.0]| 1.0| 1.0| +-----------+--------------+----------+