Java/Scala中如何像Pandas一样按行索引筛选Spark SQL Dataset

Spark并没有提供和Pandas iloc完全对等的API——这是因为Spark是分布式计算框架，数据分散存储在多个节点，不存在单机Pandas那种天然的、全局固定的行索引。不过可以根据不同的筛选场景，用更高效可靠的方案实现类似效果：

1. 取前N行（对应df.iloc[:N]） #

直接使用limit(N)方法，这是最优方案，无需额外添加列：

// 对应Pandas的df1.iloc[:3]
Dataset<Row> top3Rows = df.limit(3);

2. 取连续范围的行（对应df.iloc[a:b]） #

如果是分页或连续行范围筛选，推荐用orderBy() + offset() + limit()组合，比添加行ID列更简洁高效：

// 对应Pandas的df1.iloc[1:5]（取第2到第5行，共4行）
// 注意：必须先指定排序字段，否则行顺序不确定，结果不可靠
Dataset<Row> rangeRows = df.orderBy("你的排序字段")
                          .offset(1)  // 跳过前1行
                          .limit(4);  // 取4行

如果需要更灵活的范围条件（比如筛选第10到20行），则需要用窗口函数生成连续行号，必须指定排序依据来保证行顺序固定：

import org.apache.spark.sql.expressions.Window;
import static org.apache.spark.sql.functions.row_number;

// 定义窗口，排序字段是保证行顺序稳定的核心
Window window = Window.orderBy("你的排序字段");

// 生成连续行号后筛选，最后删除行号列
Dataset<Row> filteredRows = df.withColumn("row_id", row_number().over(window))
                              .where("row_id BETWEEN 10 AND 20")
                              .drop("row_id");

⚠️ 注意：不要用monotonically_increasing_id()生成行号，它的ID是基于分区的递增值，不是连续的全局行号，且无法保证和原数据的顺序一致，不能准确对应iloc的行为。

3. 取离散指定行（对应df.iloc[[10,12,22]]） #

同样依赖窗口函数生成连续行号，然后用IN条件筛选：

Window window = Window.orderBy("你的排序字段");
Dataset<Row> discreteRows = df.withColumn("row_id", row_number().over(window))
                              .where("row_id IN (10,12,22)")
                              .drop("row_id");

4. 同时筛选行和列（对应df.iloc[a:b, c:d]） #

行筛选按上述方法处理，列筛选可以通过列名数组索引实现：

// 先处理行范围：取第2到第5行（对应iloc[1:5]）
Window window = Window.orderBy("你的排序字段");
Dataset<Row> filteredRows = df.withColumn("row_id", row_number().over(window))
                              .where("row_id BETWEEN 2 AND 5")
                              .drop("row_id");

// 处理列范围：取第3到第4列（对应iloc[:,2:4]，Python左闭右开规则）
String[] columns = filteredRows.columns();
Dataset<Row> result = filteredRows.select(Arrays.copyOfRange(columns, 2, 4));

核心说明 #

Spark的分布式特性决定了它没有固定的全局行索引，所有基于行号的筛选都必须先通过排序固定行顺序，再生成可靠的行号或使用offset/limit。如果跳过排序步骤，行顺序是不确定的，每次运行结果可能不同。

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