如何优化Java Union-Find程序避免大数据集下的OutOfMemoryError

问题背景 #

已实现基于UnionFind的行分组逻辑，但处理100万/1000万行数据时触发java.lang.OutOfMemoryError: Java heap space，要求运行时间≤30秒、内存占用≤1GB。

内存瓶颈分析 #

1. 全量加载数据到内存：List<String[]>存储所有有效行，1000万行的字符串数组会占用数百MB甚至GB级内存。
2. 分组阶段重复存储数据：groupAndSortRows中用HashSet<String>存储每行的Arrays.toString()结果，重复存储大量字符串，且HashSet的哈希表结构额外开销大。
3. 列值映射的字符串键开销：columnValueMap使用"列索引,值"格式的字符串作为键，大量字符串对象占用内存。

具体优化方案 #

1. 拆分处理流程，用磁盘缓存替代全量内存存储 #

将数据读取、UnionFind合并、分组输出拆分为三个独立阶段：

• 读取原始数据时，将有效行写入临时磁盘文件，仅记录有效行总数（用于初始化UnionFind）。
• UnionFind合并阶段，仅加载当前行进行处理，不保留全量行数据。
• 分组输出时，再从临时文件读取全量有效行，按分组索引输出。

2. 优化列值映射的内存占用 #

用数值型键替代字符串键，减少字符串对象的内存开销：

• 计算列索引与值的组合哈希值（比如将列索引左移32位后与值的哈希值按位或），作为columnValueMap的键。
• 用HashMap<Long, Integer>替代HashMap<String, Integer>，避免大量字符串实例的创建。

3. 分组阶段优化内存使用 #

• 用Map<Integer, List<Integer>>存储组内的行索引，而非存储行字符串，减少重复存储。
• 用ArrayList替代HashSet存储组内元素（行索引不会重复），降低哈希表的额外内存开销。

4. 细节优化 #

• 修正isValidRow的逻辑错误：原代码逻辑矛盾，调整为符合业务需求的有效性判断。
• 使用G1垃圾收集器：添加JVM参数-XX:+UseG1GC，提升大内存场景下的垃圾回收效率，配合-Xmx1G使用。

优化后的关键代码片段 #

临时磁盘缓存写入 #

private static File writeTempValidRows(String filePath) throws IOException {
    File tempFile = File.createTempFile("row_cache", ".tmp");
    tempFile.deleteOnExit();

    try (BufferedReader br = new BufferedReader(new FileReader(filePath));
         BufferedWriter bw = new BufferedWriter(new FileWriter(tempFile))) {
        String line;
        while ((line = br.readLine()) != null) {
            String[] columns = line.split(";");
            if (isValidRow(columns)) {
                StringBuilder cleanedLine = new StringBuilder();
                for (int i = 0; i < columns.length; i++) {
                    String col = columns[i].trim().replaceAll("^\"|\"$", "");
                    cleanedLine.append(col);
                    if (i < columns.length - 1) {
                        cleanedLine.append(";");
                    }
                }
                bw.write(cleanedLine.toString());
                bw.newLine();
            }
        }
    }
    return tempFile;
}

优化后的UnionFind处理逻辑 #

private static UnionFind processUnionFind(File tempFile) throws IOException {
    // 先统计有效行数
    BufferedReader countReader = new BufferedReader(new FileReader(tempFile));
    int totalValidRows = (int) countReader.lines().count();
    countReader.close();

    UnionFind uf = new UnionFind(totalValidRows);
    Map<Long, Integer> columnValueMap = new HashMap<>(totalValidRows / 2);

    BufferedReader dataReader = new BufferedReader(new FileReader(tempFile));
    String line;
    int rowIndex = 0;
    while ((line = dataReader.readLine()) != null) {
        String[] columns = line.split(";");
        for (int j = 0; j < columns.length; j++) {
            String value = columns[j];
            if (!value.isEmpty()) {
                long key = ((long) j << 32) | (value.hashCode() & 0xFFFFFFFFL);
                if (columnValueMap.containsKey(key)) {
                    int prevRowIdx = columnValueMap.get(key);
                    uf.union(rowIndex, prevRowIdx);
                } else {
                    columnValueMap.put(key, rowIndex);
                }
            }
        }
        rowIndex++;
    }
    dataReader.close();
    return uf;
}

分组输出优化 #

private static void writeOutput(File tempFile, UnionFind uf) throws IOException {
    List<String> allRows = new ArrayList<>();
    try (BufferedReader br = new BufferedReader(new FileReader(tempFile))) {
        String line;
        while ((line = br.readLine()) != null) {
            allRows.add(line);
        }
    }
    int totalValidRows = allRows.size();

    Map<Integer, List<Integer>> groups = new HashMap<>();
    for (int i = 0; i < totalValidRows; i++) {
        int root = uf.find(i);
        groups.computeIfAbsent(root, k -> new ArrayList<>()).add(i);
    }

    List<List<Integer>> sortedGroups = new ArrayList<>(groups.values());
    sortedGroups.sort((g1, g2) -> Integer.compare(g2.size(), g1.size()));

    long multiElementGroups = sortedGroups.stream().filter(g -> g.size() > 1).count();

    try (PrintWriter writer = new PrintWriter("output.txt")) {
        writer.println("Total number of groups with more than one element: " + multiElementGroups);
        writer.println();
        int groupNum = 1;
        for (List<Integer> group : sortedGroups) {
            writer.println("Group " + groupNum);
            for (int idx : group) {
                writer.println(allRows.get(idx));
            }
            writer.println();
            groupNum++;
        }
    }
}

修正后的isValidRow方法 #

private static boolean isValidRow(String[] columns) {
    for (String column : columns) {
        String trimmed = column.trim();
        // 业务逻辑：列不能为空，且必须是带引号的11位数字（可根据实际需求调整）
        if (trimmed.isEmpty() || !trimmed.matches("^\"\\d{11}\"$")) {
            return false;
        }
    }
    return true;
}

效果验证 #

• UnionFind的parent和rank数组：1000万行仅占用约76MB（两个int数组，每个1000万*4字节=38MB）。
• 临时磁盘缓存仅在输出阶段读取一次全量行，内存占用可控。
• 列值映射的HashMap<Long, Integer>内存开销远低于原字符串键方案。

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