Spark SQL Java动态选择列：基于References列内容匹配列名

针对你提出的需求——根据每行References数组指定的列名动态保留对应列，同时保留ID和References列，结合Spark SQL的特性，提供以下两种可行方案：

方案一：基于Map过滤实现（Spark 3.0+ 推荐） #

该方案通过将目标列转换为键值对Map，再根据References数组过滤Map，最后展开为原始列，保持Schema统一（不在References中的列值为null）。

步骤与代码示例 #

import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.functions.*;
import java.util.Arrays;
import java.util.List;
import java.util.stream.Collectors;

// 1. 提取所有以"_id"结尾的列名
List<String> idColumns = Arrays.stream(df.columns())
        .filter(colName -> colName.endsWith("_id"))
        .collect(Collectors.toList());

// 2. 将所有_id列转换为键值对Map
Dataset<Row> withMapDf = df.withColumn("id_cols_map", map(
        idColumns.stream()
                .flatMap(colName -> Arrays.asList(lit(colName), col(colName)).stream())
                .toArray(Column[]::new)
));

// 3. 根据References数组过滤Map，仅保留指定列的键值对
Dataset<Row> filteredMapDf = withMapDf.withColumn("filtered_id_cols", map_filter(
        col("id_cols_map"),
        (key, value) -> array_contains(col("references"), key)
));

// 4. 展开过滤后的Map为原始列，同时保留ID和References
List<String> selectExprs = Arrays.asList("id", "references");
for (String colName : idColumns) {
    selectExprs.add(String.format("filtered_id_cols['%s'] as %s", colName, colName));
}

Dataset<Row> resultDf = filteredMapDf.selectExpr(selectExprs.toArray(new String[0]));

// 查看结果
resultDf.show();

说明 #

• 该方案保持输出Schema与原数据集一致，不在References中的_id列值为null，符合Spark强Schema的特性。
• map_filter是Spark 3.0及以上版本提供的内置函数，用于过滤Map中的键值对。

方案二：自定义UDF兼容低版本Spark #

如果你的Spark版本低于3.0，可以通过自定义UDF实现Map过滤逻辑：

步骤与代码示例 #

import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.api.java.UDF2;
import org.apache.spark.sql.functions.*;
import org.apache.spark.sql.types.DataTypes;
import java.util.*;
import java.util.stream.Collectors;

// 1. 提取所有以"_id"结尾的列名
List<String> idColumns = Arrays.stream(df.columns())
        .filter(colName -> colName.endsWith("_id"))
        .collect(Collectors.toList());

// 2. 注册自定义UDF：过滤Map，仅保留References中的键
UDF2<Map<String, String>, List<String>, Map<String, String>> filterMapUdf = (map, refs) -> {
    Map<String, String> filteredMap = new HashMap<>();
    for (String key : refs) {
        if (map.containsKey(key)) {
            filteredMap.put(key, map.get(key));
        }
    }
    return filteredMap;
};
spark.udf().register("filter_map", filterMapUdf, DataTypes.createMapType(DataTypes.StringType, DataTypes.StringType));

// 3. 将_id列转为Map并过滤
Dataset<Row> withMapDf = df.withColumn("id_cols_map", map(
        idColumns.stream()
                .flatMap(colName -> Arrays.asList(lit(colName), col(colName)).stream())
                .toArray(Column[]::new)
));

Dataset<Row> filteredMapDf = withMapDf.withColumn("filtered_id_cols", callUDF(
        "filter_map",
        col("id_cols_map"),
        col("references")
));

// 4. 展开Map得到结果
List<String> selectExprs = Arrays.asList("id", "references");
for (String colName : idColumns) {
    selectExprs.add(String.format("filtered_id_cols['%s'] as %s", colName, colName));
}

Dataset<Row> resultDf = filteredMapDf.selectExpr(selectExprs.toArray(new String[0]));

// 查看结果
resultDf.show();

说明 #

• 自定义UDF逻辑与map_filter一致，兼容Spark 2.x版本。
• 若_id列的数据类型不是String，需调整UDF中的泛型和DataTypes定义。

注意事项 #

Spark的Dataset/DataFrame是强Schema结构，无法实现每行列数动态变化的输出。上述方案均采用保留所有_id列，空值填充未指定列的方式，既满足需求，又符合Spark的分布式计算特性。

内容的提问来源于stack exchange，提问作者chingu