Spark DataFrame执行pivot聚合时如何保持指定排序顺序

问题说明 #

执行pivot行转列操作时，使用collect_list聚合后拼接的城市列表未按ORDER字段升序排列，不符合预期。

初始DataFrame #

+--------------+-------------+----------------+-----+
|CITY          |country      |STATE           |ORDER|
+--------------+-------------+----------------+-----+
|3. Corning    |United States|New York        |3    |
|1. Albany     |United States|New York        |1    |
|2. Batavia    |United States|New York        |2    |
|3. Campbell   |United States|California      |3    |
|2. Bakersfield|United States|California      |2    |
|1. Arvin      |United States|California      |1    |
|2. Tofino     |Canada       |British Columbia|2    |
|3. Vancouver  |Canada       |British Columbia|3    |
|1. Cranbrook  |Canada       |British Columbia|1    |
+--------------+-------------+----------------+-----+

当前错误输出 #

+-------------+-------------------------------------+-------------------------------------+---------------------------------+
|COUNTRY      |British Columbia                     |California                           |New York                         |
+-------------+-------------------------------------+-------------------------------------+---------------------------------+
|United States|                                     |3. Campbell; 2. Bakersfield; 1. Arvin|3. Corning; 1. Albany; 2. Batavia|
|Canada       |2. Tofino; 3. Vancouver; 1. Cranbrook|                                     |                                 |
+-------------+-------------------------------------+-------------------------------------+---------------------------------+

期望输出 #

+-------------+-------------------------------------+-------------------------------------+---------------------------------+
|COUNTRY      |British Columbia                     |California                           |New York                         |
+-------------+-------------------------------------+-------------------------------------+---------------------------------+
|United States|                                     |1. Arvin; 2. Bakersfield; 3. Campbell|1. Albany; 2. Batavia; 3. Corning|
|Canada       |1. Cranbrook; 2. Tofino; 3. Vancouver|                                     |                                 |
+-------------+-------------------------------------+-------------------------------------+---------------------------------+

原有问题代码 #

主程序代码

import static org.apache.spark.sql.functions.col;
import static org.apache.spark.sql.functions.collect_list;
import static org.apache.spark.sql.functions.lit;

import java.io.Serializable;
import java.util.Arrays;
import java.util.List;

import org.apache.log4j.Level;
import org.apache.log4j.Logger;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.functions;

public class TestPivotOrder implements Serializable {
    private static final long serialVersionUID = -1L;
    
    private static String COUNTRY = "COUNTRY";
    private static String STATE = "STATE";
    private static String CITY = "CITY";
    private static String ORDER = "ORDER";
    
    public static void main(String[] args) {
        TestPivotOrder app = new TestPivotOrder();
        app.start();
    }


    private void start() {
        
        // 仅打印警告和错误日志
        Logger.getLogger("org.apache").setLevel(Level.WARN);
        
        SparkSession sparkSession = SparkSession.builder()
             .appName("Pivot with Order").master("local").getOrCreate();
        
        Dataset<Row> df = getData(sparkSession);

        df = df
            .withColumn(ORDER, functions.regexp_extract(col(STATE), "(\\d+)", 0))
            .withColumn(STATE, functions.regexp_replace(col(STATE), ".(\\d+)", ""))
            .withColumn(CITY, functions.concat(col(ORDER), lit(". "), col(CITY)));
        
        df.show(10, false);
        
        df = df
            .groupBy(df.col(COUNTRY))
            .pivot(df.col(STATE))
            .agg(functions.concat_ws("; ", collect_list(df.col(CITY))));
        
        df.show(10, false);
        
        sparkSession.stop();
    }
        

    private Dataset<Row> getData(SparkSession spark) {
        
        List<Item> list = Arrays.asList(
                new Item("United States","New York.3","Corning"),
                new Item("United States","New York.1","Albany"),
                new Item("United States","New York.2","Batavia"),               
                                
                new Item("United States","California.3","Campbell"),
                new Item("United States","California.2","Bakersfield"),
                new Item("United States","California.1","Arvin"),
                
                new Item("Canada","British Columbia.2","Tofino"),
                new Item("Canada","British Columbia.3","Vancouver"),
                new Item("Canada","British Columbia.1","Cranbrook")             
            );
        
        return spark.createDataFrame(list, Item.class);
    }   
}

Item实体类代码

import java.io.Serializable;

public class Item implements Serializable {

    private static final long serialVersionUID = -1639L;
    
    private String country;
    private String state;
    private String city;

    public Item(String country, String state, String city) {
        this.country = country;
        this.state = state;     
        this.city = city;
    }

    public String getCountry() {
        return country;
    }

    public String getState() {
        return state;
    }

    public String getCity() {
        return city;
    }
}

问题原因 #

Spark原生collect_list函数不保证聚合后元素的顺序，也不支持直接在聚合阶段指定排序字段，因此收集到的城市列表顺序随机。

实现方案 #

核心思路 #

不直接收集CITY字段，而是将排序键ORDER和CITY打包为结构体，收集结构体列表后按排序键升序排序，再提取排序后的CITY值拼接为字符串，全程使用Spark内置函数实现，无额外性能开销。

具体修改点 #

• 修正正则匹配逻辑，将提取出的ORDER字段从字符串转为整数类型，避免字符串字典序排序错误（如字符串"10"会排在"2"之前）
• 替换原有聚合逻辑：
  • 聚合时收集struct(ORDER, CITY)结构体列表
  • 调用sort_array对结构体列表按ORDER字段升序排序（结构体默认按第一个字段值排序）
  • 使用transform高阶函数从排序后的结构体列表中提取CITY字段
  • 最后用concat_ws拼接为分号分隔的字符串

修改后核心代码 #

df = df
    .withColumn(ORDER, functions.regexp_extract(col(STATE), "(\\d+)", 0).cast("int"))
    // 修正正则，精确匹配".数字"后缀
    .withColumn(STATE, functions.regexp_replace(col(STATE), "\\.\\d+", ""))
    .withColumn(CITY, functions.concat(col(ORDER).cast("string"), lit(". "), col(CITY)));

df = df
    .groupBy(df.col(COUNTRY))
    .pivot(df.col(STATE))
    .agg(
        functions.concat_ws(
            "; ",
            functions.expr("transform(sort_array(collect_list(struct(ORDER, CITY))), x -> x.CITY)")
        )
    );

低版本Spark兼容方案 #

如果使用Spark 2.3及以下版本不支持transform高阶函数，可以自定义UDF实现结构体数组排序、提取CITY的逻辑，示例代码如下：

import org.apache.spark.sql.api.java.UDF1;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.Row;
import scala.collection.JavaConverters;
import scala.collection.Seq;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;

// 注册UDF
sparkSession.udf().register("sort_and_extract_city", new UDF1<Seq<Row>, Seq<String>>() {
    @Override
    public Seq<String> call(Seq<Row> structArr) throws Exception {
        List<Row> list = new ArrayList<>();
        for (int i = 0; i < structArr.size(); i++) {
            list.add(structArr.apply(i));
        }
        // 按ORDER字段（结构体第一个元素，int类型）升序排序
        list.sort(Comparator.comparingInt(r -> r.getInt(0)));
        List<String> cityList = new ArrayList<>();
        for (Row r : list) {
            cityList.add(r.getString(1));
        }
        return JavaConverters.asScalaIteratorConverter(cityList.iterator()).asScala().toSeq();
    }
}, DataTypes.createArrayType(DataTypes.StringType));

// 聚合时改用自定义UDF
.agg(
    functions.concat_ws(
        "; ",
        functions.callUDF("sort_and_extract_city", functions.collect_list(functions.struct(col(ORDER), col(CITY))))
    )
);

运行后即可得到符合排序要求的期望输出。

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