Spark-Java：如何在 spark Dataframe 中添加数组列

Question

我正在尝试向我的 Spark Dataframe 添加一个新列。 添加的新列的大小将基于一个变量（比如salt ）帖子，我将使用该列来分解数据集以用于盐渍连接。

目前，我在array function 中使用连续lit ，但这有一个问题，即它无法参数化，并且作为编码实践看起来最糟糕。 我当前的实现如下所示。

int salt =3;

Dataset<Row> Reference_with_Salt_Col = Reference.withColumn("salt_array", array(lit(0), lit(1), lit(2)));

我已经参考并查看了各种方法，但似乎没有一个可以解决 Java 中的问题。 functions.typedlit方法虽然适用于 Python/Scala，但似乎不适用于 Java。 进一步传递数组或列表也无助于 spark 给出相同的错误。

我正在使用 Spark 2.2.0 和 Java 1.8 版本

Answer 1

您可以使用array function，但首先将列表的每个元素转换为lit 。 在 stream 上使用map function 的示例：

import org.apache.spark.sql.*;

import java.util.Arrays;
import java.util.List;

// example of input dataframe
Reference.show();

//+-----+
//|label|
//+-----+
//|    a|
//|    b|
//|    c|
//+-----+

List<Integer> salt_array = Arrays.asList(0, 1, 2);

Reference.withColumn(
        "salt_array",
        functions.array(salt_array.stream().map(functions::lit).toArray(Column[]::new))
).show();

//+-----+----------+
//|label|salt_array|
//+-----+----------+
//|    a| [0, 1, 2]|
//|    b| [0, 1, 2]|
//|    c| [0, 1, 2]|
//+-----+----------+

要生成一个包含从 0 到salt - 1的范围的数组，您可以像这样使用IntStream.rangeClosed ：

import java.util.stream.IntStream;

int salt = 3;

Dataset<Row> Reference_with_Salt_Col = Reference.withColumn(
        "salt_array",
        functions.array(IntStream.rangeClosed(0, salt - 1).mapToObj(functions::lit).toArray(Column[]::new))
);