1TB、200亿行3列数据集的无重叠80/20随机高效拆分方法咨询

针对1TB、200亿行的3列数据集，要拆分为元素完全无重叠的80/20子集，核心是先把数据按连通分量划分（同一行的元素属于同一个连通块，连通块内的所有元素不会出现在其他块），再对连通块进行随机抽样分配，这样能保证子集间元素完全不相交，同时适配分布式计算场景，高效处理海量数据。

SQL实现（基于Spark SQL/Hive） #

步骤说明 #

1. 构建边表：把每行的三个元素两两配对，生成元素间的关联边，用于后续计算连通分量。
2. 计算连通分量：用图算法找出所有元素所属的连通块，每个连通块分配唯一ID。
3. 连通块抽样分组：对所有连通块随机抽样，按8:2比例分配到两个组。
4. 关联原表拆分数据：根据行内元素的连通块ID，将行分到对应的子集。

代码示例 #

-- 1. 生成元素间的关联边表
CREATE TEMP VIEW edges AS
SELECT col1 AS src, col2 AS dst FROM original_table
UNION ALL
SELECT col1 AS src, col3 AS dst FROM original_table
UNION ALL
SELECT col2 AS src, col3 AS dst FROM original_table
DISTINCT; -- 去重减少计算量

-- 2. 生成顶点表（所有唯一元素）
CREATE TEMP VIEW vertices AS
SELECT DISTINCT src AS id FROM edges
UNION ALL
SELECT DISTINCT dst AS id FROM edges;

-- 3. 计算连通分量（需依赖Spark SQL的GraphFrames库）
CREATE TEMP VIEW connected_components AS
SELECT id, component_id 
FROM graphframes.graph(vertices, edges).connectedComponents();

-- 4. 对连通块进行随机分组
CREATE TEMP VIEW component_groups AS
SELECT component_id, 
       CASE WHEN RAND() <= 0.8 THEN 'group_80' ELSE 'group_20' END AS group_label
FROM (SELECT DISTINCT component_id FROM connected_components);

-- 5. 提取80%子集
SELECT t.* 
FROM original_table t
JOIN connected_components cc ON t.col1 = cc.id
JOIN component_groups cg ON cc.component_id = cg.component_id
WHERE cg.group_label = 'group_80'
GROUP BY t.col1, t.col2, t.col3; -- 去重避免同一行多次匹配

-- 6. 提取20%子集
SELECT t.* 
FROM original_table t
JOIN connected_components cc ON t.col1 = cc.id
JOIN component_groups cg ON cc.component_id = cg.component_id
WHERE cg.group_label = 'group_20'
GROUP BY t.col1, t.col2, t.col3;

Python实现（基于PySpark） #

步骤说明 #

和SQL逻辑一致，用PySpark的GraphFrames处理图计算，全程分布式执行，适配海量数据。

代码示例 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import rand, when
from graphframes import GraphFrame

# 初始化Spark会话，配置足够的资源（根据集群规模调整）
spark = SparkSession.builder \
    .appName("LargeDatasetSplit") \
    .config("spark.executor.memory", "16g") \
    .config("spark.driver.memory", "8g") \
    .getOrCreate()

# 读取原始数据集（假设是空格分隔的文本文件）
original_df = spark.read.csv(
    "hdfs://path/to/original_data", 
    sep=" ", 
    header=False, 
    inferSchema=False
).toDF("col1", "col2", "col3")

# 生成元素关联边表
edges1 = original_df.select("col1", "col2").withColumnRenamed("col1", "src").withColumnRenamed("col2", "dst")
edges2 = original_df.select("col1", "col3").withColumnRenamed("col1", "src").withColumnRenamed("col3", "dst")
edges3 = original_df.select("col2", "col3").withColumnRenamed("col2", "src").withColumnRenamed("col3", "dst")
edges_df = edges1.union(edges2).union(edges3).distinct()

# 生成顶点表（所有唯一元素）
vertices_df = edges_df.select("src").union(edges_df.select("dst")).distinct().withColumnRenamed("src", "id")

# 计算连通分量
g = GraphFrame(vertices_df, edges_df)
connected_components = g.connectedComponents()

# 对连通块随机分组
component_groups = connected_components.select("component") \
    .distinct() \
    .withColumn(
        "group_label",
        when(rand() <= 0.8, "group_80").otherwise("group_20")
    )

# 提取80%子集
df_80 = original_df.join(
    connected_components, original_df.col1 == connected_components.id
).join(
    component_groups, connected_components.component == component_groups.component
).filter("group_label = 'group_80'") \
  .select("col1", "col2", "col3") \
  .distinct()

# 提取20%子集
df_20 = original_df.join(
    connected_components, original_df.col1 == connected_components.id
).join(
    component_groups, connected_components.component == component_groups.component
).filter("group_label = 'group_20'") \
  .select("col1", "col2", "col3") \
  .distinct()

# 保存结果到HDFS或其他分布式存储
df_80.write.csv("hdfs://path/to/output_80", sep=" ", header=False, mode="overwrite")
df_20.write.csv("hdfs://path/to/output_20", sep=" ", header=False, mode="overwrite")

spark.stop()

关键注意事项 #

1. 为什么用连通分量？
   如果直接随机拆分单个元素，会出现某一行的元素分属两个组的情况，这类行无法被分到任何子集，导致数据浪费。按连通分量拆分能保证整个块的所有行都归属同一个子集，完全符合“元素无重叠”的要求。
2. 关于比例精度
   按连通块抽样得到的是近似80/20的行比例，如果需要更精确的行占比，可以先统计每个连通块的行数，再按累计行数阈值来分组（比如累计行数达到总行数80%时停止分配到group_80）。
3. 性能优化
   • 提前对原始数据分区，减少shuffle量；
   • 边表生成时去重，降低图计算的压力；
   • 根据集群资源调整Spark的内存、核心数配置，避免OOM。

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