Apache Spark DataFrame API 实现开窗函数

使用 Apache Spark DataFrame API 实现开窗函数（Window Functions）

1. 概述

在 Spark 中，开窗函数（Window Functions）用于在数据集的某个窗口（分区）内进行计算。常见的应用场景包括：

计算行号（row_number）
排名（rank、dense_rank）
累计求和（sum）
移动平均（avg）
本文将演示如何使用 Spark 的 DataFrame API 实现开窗函数，包括升序和降序排序。

2. 环境准备

确保你已经配置好 Spark 环境，并添加了以下 Maven 依赖：

<dependencies>
    <!-- Spark Core -->
    <dependency>
        <groupId>org.apache.spark</groupId>
        <artifactId>spark-core_2.12</artifactId>
        <version>3.2.0</version>
    </dependency>
    <!-- Spark SQL -->
    <dependency>
        <groupId>org.apache.spark</groupId>
        <artifactId>spark-sql_2.12</artifactId>
        <version>3.2.0</version>
    </dependency>
</dependencies>

3. 代码实现

3.1 读取 CSV 文件

首先，读取 CSV 文件并选择需要的列：

import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.expressions.Window;
import org.apache.spark.sql.expressions.WindowSpec;
import static org.apache.spark.sql.functions.*;

public class WindowFunctionExample {
    public static void main(String[] args) {
        // 初始化 SparkSession
        SparkSession spark = SparkSession.builder()
                .appName("Window Function Example")
                .master("local")
                .getOrCreate();

        // 读取 CSV 文件
        Dataset<Row> rowDataset = spark.read()
                .option("header", "true") // 第一行为列名
                .option("inferSchema", "true") // 自动推断数据类型
                .csv("./f2024-Q4-000.csv");

        // 选择需要的列
        Dataset<Row> dataset = rowDataset.select(
                "fl_id", 
                "airid", 
                "flc", 
                "flinator", 
                "schedl_airport_icao", 
                "deparort_name", 
                "schedarture_local", 
                "flig_km", 
                "flius"
        ).limit(10); // 限制数据量，方便调试

3.2 定义窗口规范

使用 Window.partitionBy 和 Window.orderBy 定义窗口规范：

        // 定义升序窗口规范
        WindowSpec windowSpecASC = Window.partitionBy("scao", "depaame")
                .orderBy("scheocal");

        // 定义降序窗口规范
        WindowSpec windowSpecDESC = Window.partitionBy("schedrt_icao", "departure_e")
                .orderBy(desc("schcal"));

3.3 应用开窗函数

使用 withColumn 方法应用开窗函数，例如计算行号：

        // 应用升序窗口函数
        Dataset<Row> resultASC = dataset.withColumn("row_number_asc", row_number().over(windowSpecASC));

        // 应用降序窗口函数
        Dataset<Row> resultDESC = dataset.withColumn("row_number_desc", row_number().over(windowSpecDESC));

        // 显示结果
        resultASC.show();
        resultDESC.show();

        // 关闭 SparkSession
        spark.stop();
    }
}

4. 代码细节说明

4.1 读取 CSV 文件

option("header", "true")：指定 CSV 文件的第一行为列名。
option("inferSchema", "true")：自动推断列的数据类型（如整数、字符串等）。
select：选择需要的列，减少内存占用。

4.2 定义窗口规范

Window.partitionBy：指定分区的列。例如，按 scheduled_arrival_airport_icao 和 departure_airport_name 分区。
Window.orderBy：指定排序的列。
默认是升序（asc）。
使用 desc 方法实现降序排序。

4.3 应用开窗函数

row_number()：为每一行生成一个行号。
withColumn：添加新列或替换现有列。