Spark中FlatMapGroupsWithStateFunction返回懒迭代器是否安全？

在使用Apache Spark进行有状态处理时，可通过flatMapGroupsWithState函数对相同Key的行进行分组，示例代码如下：

Dataset<Row> ds = spark.sql("SELECT ... FROM parquetTable")
                         .groupByKey(..., Encoders.STRING())
                         .flatMapGroupsWithState(new FlatMapFunctionEager(),
                                 OutputMode.Append(),
                                 Encoders.kryo(Session.class),
                                 RowEncoder.apply(PARQUET_SCHEMA),
                                 GroupStateTimeout.ProcessingTimeTimeout());

通常教程中的FlatMapGroupsWithStateFunction会采用预加载所有结果到List再返回迭代器的实现方式：

private class FlatMapFunctionEager implements FlatMapGroupsWithStateFunction<String, Row, Session, Row> {
    @Override
    public Iterator<Row> call(String key, Iterator<Row> values, GroupState<Session> state) throws Exception {
        Session session = null;
        if (state.exists()) {
            session = state.get();
        } else {
            session = createSession();
            state.update(session);
        }

        // 预创建List并填充所有结果
        List<Row> result = new LinkedList<>();
        while (values.hasNext()) {
            Row row = values.next();
            result.add(process(row, session));
        }

        return result.iterator();
    }

    private Row process(Row row, Session session) {
        // 处理逻辑并返回结果Row
    }
}

这种实现在分组包含大量Row时，容易引发内存溢出问题。因此有人提出改用懒加载迭代器的实现方式：返回持有输入values迭代器和GroupState的自定义迭代器，仅在调用next()时才处理数据，示例代码如下：

public class FlatMapFunctionLazy implements FlatMapGroupsWithStateFunction<String, Row, Session, Row> {
    @Override
    public Iterator<Row> call(String key, Iterator<Row> values, GroupState<Session> state) throws Exception {
        return new LazyIterator(state, values);
    }
}

public class LazyIterator implements Iterator<Row> {
    private final GroupState<Session> state;
    private final Iterator<Row> values;

    public LazyIterator(GroupState<Session> state, Iterator<Row> values) {
        this.state = state;
        this.values = values;
    }

    @Override
    public boolean hasNext() {
        return values.hasNext();
    }

    @Override
    public Row next() {
        Row row = values.next();
        return process(row);
    }

    private Row process(Row row) {
        Session session = null;
        if (state.exists()) {
            session = state.get();
        } else {
            session = createSession();
            state.update(session);
        }

        return process(row, session);
    }

    private Row process(Row row, Session session) {
        // 处理逻辑并返回结果Row
    }
}

安全性分析 #

这种懒迭代器的用法是安全的，不会产生未定义行为，原因如下：

• 单线程串行处理：Spark对同一个Key的分组处理是严格单线程串行执行的，同一个Key的所有数据只会被一个任务线程处理，不存在多个线程并发访问同一个GroupState或values迭代器的情况，因此不会有竞态条件。
• 迭代器的生命周期：values迭代器仅对应当前分组的单批次数据，迭代器的遍历全程在当前任务线程内完成，不会被其他线程篡改或复用。
• GroupState的隔离性：GroupState是Spark为每个Key单独维护的状态实例，仅在当前Key的处理周期内有效，且框架不会在当前迭代器遍历完成前修改该状态。即使是流处理场景，同一个Key的下一批次数据也只会在当前批次处理完成（包括迭代器完全遍历）后才会被处理。

注意事项 #

• 不要让懒迭代器持有GroupState的引用跨越多个分组处理周期，当前实现中迭代器仅对应单次call方法调用，不存在这个问题。
• 由于是单线程处理，Session对象的修改无需额外的线程安全措施。

结论 #

懒迭代器的实现方式可以有效降低大分组场景下的内存占用，且完全符合Spark的执行模型，不会引发未定义行为，测试中得到正确结果也验证了这一点。

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