将数组值作为新列添加到 spark 数据框

Question

考虑一个名为employees的火花数据框，例如：

+----------+-----+
|   name   | age |
+----------+-----+
|   John   | 32  |
| Elizabeth| 28  |
|   Eric   | 41  |
+----------+-----+

和一个字符串数组state = ["LA", "AZ", "OH"] ，我想将此数组作为新列附加到df ，这样数据框将如下所示：

+----------+-----+-------+
|   name   | age | state |
+----------+-----+-------+
|   John   | 32  |   LA  |
| Elizabeth| 28  |   AZ  |
|   Eric   | 41  |   OH  |
+----------+-----+-------+

我如何在 Scala（或 Java，几乎相同）中实现这一点？ 我只看到了如何为网络上的所有行添加相同的值，在这里我希望为每一行添加不同的值。

谢谢 ！ :)

Answer 1

由于 spark 在分布式模式下运行，您无法在带有索引的数组上添加基于列的值。 假设 spark 运行有两个工人， John和Elizabeth交付给工人A ， Eric交付给工人B。 事实上，当保存在数据帧中时，它们会分裂。 工人不知道John 、 Elizabeth或Eric的索引是多少。 您可以在普通的 Java 单个程序中简单地做您想做的事。

在您的示例中，您需要将数组转换为数据帧并使用join合并基于具有相同值的列的两个数据帧。 但是，您可以使用crossJoin在您的表上执行笛卡尔积。

Dataset<Row> ndf = df.crossJoin(df2);

如果您只需要添加具有常量值的列或基于同一数据帧上的另一列的值，请使用withColumn如下：

Dataset<Row> ndf = df.withColumn("city",functions.lit(1));
Dataset<Row> ndf = df.withColumn("city",functions.rand());
Dataset<Row> ndf = df.withColumn("city",functions.col("name"));

最后，你可以像这样使用 Atomic 来得到你想要的。 我在火花单模式下测试它。

    public static void main(String[] args) {
        System.setProperty("hadoop.home.dir", "H:\\work\\HadoopWinUtils\\");
        SparkSession spark = SparkSession
                .builder()
                .master("local[*]")
                .getOrCreate();

        Dataset<Row> df = spark.read().json("H:\\work\\HadoopWinUtils\\people.json");

        List<String> city_array = Arrays.asList("LA", "AZ", "OH");
        // Displays the content of the DataFrame to stdout
        df.show();
   
        df = df.withColumn("city",functions.col("name"));

        AtomicInteger i= new AtomicInteger();

        Dataset<Row> df3 = df.map((MapFunction<Row, Row>) value -> {
            return RowFactory.create(value.get(0),value.get(1),city_array.get(i.getAndIncrement()));
            //return city_array.get(i.getAndIncrement());
        }, RowEncoder.apply(df.schema()));

        df3.show();
    }

人们是

+----+-------+
| age|   name|
+----+-------+
|null|Michael|
|  30|   Andy|
|  19| Justin|
+----+-------+

结果是

+----+-------+----+
| age|   name|city|
+----+-------+----+
|null|Michael|  LA|
|  30|   Andy|  AZ|
|  19| Justin|  OH|
+----+-------+----+

Answer 2

您可以在 pyspark 中尝试类似的操作。

>>> _TRANSFORMED_DF_SCHEMA = StructType([
...     StructField('name', StringType(), False),
...     StructField('age', IntegerType(), False),
...     StructField('id', IntegerType(), False),
...     StructField('state', StringType(), False),
... ])
>>> 
>>> state = ['LA', 'AZ', 'OH']
>>> data = (['John', 32], ['Eli', 28], ['Eric', 41])
>>> df = spark.createDataFrame(data, schema=['name', 'age'])
>>> rdd1 = df.rdd.zipWithIndex()
>>> df1 = rdd1.toDF()
>>> df1.show()
+----------+---+
|        _1| _2|
+----------+---+
|[John, 32]|  0|
| [Eli, 28]|  1|
|[Eric, 41]|  2|
+----------+---+

>>> df_final = df1.select(df1['_1']['name'].alias('name'), df1['_1']['age'].alias('age'), df1['_2'].alias('id'))
>>> df_final.show()
+----+---+---+
|name|age| id|
+----+---+---+
|John| 32|  0|
| Eli| 28|  1|
|Eric| 41|  2|
+----+---+---+

>>> def add_state(row_dict):
...     new_dict = dict()
...     new_dict['name'] = row_dict['name']
...     new_dict['age'] = row_dict['age']
...     new_dict['id'] = row_dict['id']
...     new_dict['state'] = state[row_dict['id']]
...     return new_dict
... 
>>> df_rdd = df_final.rdd.map(add_state)
>>> df_final = spark.createDataFrame(df_rdd, schema=_TRANSFORMED_DF_SCHEMA)
>>> df_final.show()
+----+---+---+-----+
|name|age| id|state|
+----+---+---+-----+
|John| 32|  0|   LA|
| Eli| 28|  1|   AZ|
|Eric| 41|  2|   OH|
+----+---+---+-----+