Spark中row_number窗口函数窄依赖实现及orderBy与sort特性咨询

1. 单分区内实现row_number窄转换的方案 #

完全可以实现，你可以直接用mapPartitions算子完成，该算子为窄依赖，所有逻辑在单分区内执行，不会触发Shuffle。
实现逻辑为：取出单个分区的全部数据，按你需要排序的col3字段排序后，逐行生成自增行号追加到原数据中即可，和标准row_number从1开始计数的逻辑完全对齐。
Java API 示例代码如下：

import org.apache.spark.api.java.function.MapPartitionsFunction;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.RowFactory;
import org.apache.spark.sql.types.DataTypes;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.Iterator;
import java.util.List;

Dataset<Row> resultDf = df.mapPartitions((MapPartitionsFunction<Row, Row>) iterator -> {
    // 收集当前分区所有数据
    List<Row> partitionRows = new ArrayList<>();
    iterator.forEachRemaining(partitionRows::add);
    // 按col3升序排序，可自行调整排序规则
    partitionRows.sort(Comparator.comparing(row -> row.getAs("col3")));
    // 生成带行号的新行
    List<Row> resultRows = new ArrayList<>();
    for (int i = 0; i < partitionRows.size(); i++) {
        Row oldRow = partitionRows.get(i);
        Object[] oldValues = oldRow.getValues();
        Object[] newValues = new Object[oldValues.length + 1];
        System.arraycopy(oldValues, 0, newValues, 0, oldValues.length);
        newValues[oldValues.length] = i + 1;
        resultRows.add(RowFactory.create(newValues));
    }
    return resultRows.iterator();
}, RowEncoder.apply(df.schema().add("col1", DataTypes.IntegerType)));

注意：该方案仅适用于你明确业务逻辑不需要按col2跨分区聚合，行号仅需要在当前现有分区内生成的场景。

2. 排序相关转换特性对窗口函数的适用性 #

首先需要先澄清认知偏差：你提到的sort为窄转换，实际是指sortWithinPartitions（单分区内排序）算子，普通全局sort/orderBy均为宽转换，会触发全量Shuffle。
该特性完全适用于窗口函数场景：

• 你原有代码中使用的窗口定义包含Window.partitionBy("col2")，只要调用partitionBy就会触发Shuffle，将相同col2值的数据拉到同一分区，无论后续排序逻辑如何，整个窗口操作都是宽转换，和全局orderBy特性一致。
• 如果窗口不定义partitionBy仅定义orderBy，会触发全局排序，也属于宽转换。
• 只有将排序逻辑完全限制在单分区内执行（比如上述mapPartitions方案），对应的排序才是窄转换，和sortWithinPartitions特性一致。

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