Question

以下是创建pyspark.sql DataFrame

的代码

import numpy as np
import pandas as pd
from pyspark import SparkContext
from pyspark.sql import SQLContext
df = pd.DataFrame(np.array([[1,2,3],[4,5,6],[7,8,9],[10,11,12]]), columns=['a','b','c'])
sparkdf = sqlContext.createDataFrame(df, samplingRatio=0.1)

所以sparkdf看起来像

现在我想将新列添加为numpy数组（甚至是列表）

new_col = np.array([20,20,20,20])

但标准方式

sparkdf = sparkdf.withColumn('newcol', new_col)

失败。可能udf是要走的路，但我不知道如何创建一个udf，为每个DataFrame行分配一个不同的值，即迭代new_col。我查看了其他pyspark和pyspark.sql但无法找到解决方案。另外我需要留在pyspark.sql中，所以不是scala解决方案。谢谢！

Answer 1

假设数据框被排序以匹配数组中值的顺序，您可以压缩RDD并重建数据框，如下所示：

n = sparkdf.rdd.getNumPartitions()

# Parallelize and cast to plain integer (np.int64 won't work)
new_col = sc.parallelize(np.array([20,20,20,20]), n).map(int) 

def process(pair):
    return dict(pair[0].asDict().items() + [("new_col", pair[1])])

rdd = (sparkdf
    .rdd # Extract RDD
    .zip(new_col) # Zip with new col
    .map(process)) # Add new column

sqlContext.createDataFrame(rdd) # Rebuild data frame

您还可以使用联接：

new_col = sqlContext.createDataFrame(
    zip(range(1, 5), [20] * 4),
    ("rn", "new_col"))

sparkdf.registerTempTable("df")

sparkdf_indexed = sqlContext.sql(
    # Make sure we have specific order and add row number
    "SELECT row_number() OVER (ORDER BY a, b, c) AS rn, * FROM df")

(sparkdf_indexed
    .join(new_col, new_col.rn == sparkdf_indexed.rn)
    .drop(new_col.rn))

但是窗口功能组件不可扩展，应该避免使用更大的数据集。

当然，如果您只需要一个单一列的列，则可以使用lit

import pyspark.sql.functions as f
sparkdf.withColumn("new_col", f.lit(20))

但我认为情况并非如此。

如何将numpy.array作为新列添加到pyspark.SQL DataFrame？

1 个答案: