在使用 PySpark 处理 Lance 数据集（通过 pylance 处理）时，最常见的性能杀手是把大字段卷入 Spark 的 Shuffle（数据重分区交换）。例如如 download_html、长文本、二进制、向量等。本文给出一套经过验证的工程范式：只 Shuffle 几个字节的行地址，把大字段的读取下推到分区内、用 Lance 的 Fragment.take() 按行地址读取，从而在保证数据正确分布的同时，几乎完全消除大字段的网络与磁盘 I/O。

为什么大字段让 Shuffle 变得很贵 #

Spark 的 repartition、join、groupBy 等宽依赖算子会触发 Shuffle，Shuffle 的成本与被搬运的数据体积直接正相关。当一行里包含一个几十 KB 甚至 MB 级的大字段时：

• 每一行都要把这个大字段序列化并写入 Shuffle 文件，磁盘 I/O 暴涨；
• 跨 Executor 传输时占满网络带宽；
• Python（PySpark）侧还要承受额外的序列化/反序列化与内存压力，极易 OOM 或频繁 GC。

为什么 parquet/ORC 无法解决该问题 #

Parquet/ORC 是面向"批量扫描 + 列裁剪"的格式，最小读取单元是 Row Group（行组），没有稳定的、可跨任务传递的"行地址"，也没有 Fragment.take(offsets, columns=[...]) 这种可以实现 常数时间随机访问 的接口。

即便只 Shuffle（数据重分区交换）文件路径 + 行号，下游仍需重新扫全表获取目标大字段，再做 join（关联）才能把大字段拼到结果表中——等于至少做了两次全量 I/O，成本反而更高。

Lance 的基石：行地址（Row Address）与 Fragment #

Lance 数据集由若干 Fragment（水平分片）组成，每个 Fragment 有一个自增的 uint32 编号，并由一个或多个数据文件构成。
行地址 _rowaddr 是一行在表中的物理位置，是一个由两个 32 位的值拼成的 64 位标识：

因此从 _rowaddr 反解出 Fragment 和 offset 只需两次位运算：

fragment_id = rowaddr >> 32          # 高 32 位：确认 Fragment 编号
row_offset  = rowaddr & 0xFFFFFFFF   # 低 32 位：确认 Fragment 内的行偏移

这套寻址方案天然支持常数时间随机访问：拿到行地址，就能直接定位 Fragment 和文件内偏移，用优化过的列存 reader 批量取数。

Lance 的大字段 Shuffle 解决方案 #

Lance 补齐了传统数据表的两个缺口：

1. _rowaddr 是相对稳定、可外部传递的物理行地址。
2. Fragment.take() 提供低开销的随机点查与列裁剪。

也就是说，只 Shuffle 行地址、按需取大字段，不是 Spark 的技巧，而是依靠 Lance 这种原生支持行级随机访问的列存格式才能解锁的能力。

为什么 Lance 先排好序再取数会更快 #

这是本实践方案的精髓，来自 _rowaddr 的物理含义：

1. 同 Fragment 聚合：_rowaddr 高 32 位相同的行属于同一个 Fragment。按 Fragment 把它们聚到一起，就能让一个 Fragment 对象被反复复用，避免每取一行就重新打开文件、重建 reader。
2. 分区内连续：_rowaddr 低 32 位是 Fragment 内的递增偏移。按 _rowaddr 排序后，同一 Fragment 的行在分区内是连续递增的，读取时能有效命中 Lance reader 缓存与操作系统页缓存，把大量随机读转化为近似顺序读。

核心逻辑 #

这套范式的本质是一句话：让 Spark 负责"调度地址"，让 Lance 负责"读取内容"；让轻量列走 Shuffle，让重字段在分区内本地按需读取。

1. Spark 侧只 Shuffle _rowaddr/_fragid 这种几字节的整数，Shuffle 体积与大字段彻底解耦。
2. repartition(_fragid) + sortWithinPartitions(_rowaddr) 利用行地址的物理含义，换来 Fragment 复用与缓存命中。
3. Fragment.take() 在分区内按需、批量、近似顺序地把大字段读出来。

对于含有大文本/二进制/向量等重字段的 Lance 数据集，这套写法通常能把 Shuffle I/O 降低一到两个数量级，是 PySpark + Lance 工程实践中值得固化的标准模式。

代码实现 #

反模式 vs 推荐范式 #

先对比两种写法，直观感受差异。

反模式示意（不要这样写）：

# ❌ 反模式：download_html 从一开始就被带进 DataFrame，
#    后续任何 repartition / join / groupBy 都会把大字段搬来搬去
big = ds_df.select("_rowaddr", "download_html", "label")
big = big.repartition("label")          # 大字段被卷入 Shuffle，磁盘和网络 I/O 爆炸
result = big.join(other, on="key")      # 再次 Shuffle，雪上加霜

Step 1：只读出 _rowaddr / _fragid 这种轻量列 #

我们只把"行地址"读进 Spark，大字段完全不参与。

import lance
import pyspark.sql.functions as F
from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, LongType, StringType

# 对象存储（TOS S3）连接配置，用于 Lance 读取远端数据集
STORAGE_OPTIONS ={
    "aws_region": "cn-beijing",
    "aws_endpoint": "https://your-bucket.tos-s3-cn-beijing.ivolces.com",
    "access_key_id": "",
    "secret_access_key": "",
    "aws_virtual_hosted_style_request": "true"
}

# Lance 数据集路径（TOS S3 地址）
DATASET_URI = "s3://your-bucket/your-dataset-path"

# 获取当前数据集版本号，用于保证本次任务全程读取同一快照，避免版本漂移
DATASET_VERSION = lance.dataset(DATASET_URI, storage_options=STORAGE_OPTIONS).version

# 初始化 Spark 会话
spark = SparkSession.builder.appName("lance-avoid-bigfield-Shuffle").getOrCreate()

# 第一步：仅读取轻量级寻址字段，不加载大字段 
# _fragid：Fragment（水平分片）编号，标识行所属分片 
# _rowaddr：行地址，64 位整数，唯一定位表中任意一行的物理位置 
# 目的：后续 Shuffle（数据重分区交换）只搬运行地址（8 字节/行），
addr_df = spark.read.format("lance").load(DATASET_URI).select("_fragid", "_rowaddr")

Step 2：repartition(_fragid) + sortWithinPartitions(_rowaddr) #

这一步是保证局部性的关键：

# 第二步：按物理地址排序，按 Fragment 分区
addr_df = (
    addr_df
    .repartition("_fragid")            # 同一 Fragment 的所有行 Shuffle 到同一 partition
    .sortWithinPartitions("_rowaddr")  # 分区内按物理地址连续排列，便于复用 + 命中缓存
)
# 如需控制并行度可写成 .repartition(num_partitions, "_fragid")

• repartition("_fragid")：保证同一 Fragment 的 _rowaddr 全部落在同一个 Partition，使下游可以复用同一个 Fragment 对象。
• sortWithinPartitions("_rowaddr")：因为 _rowaddr = (Fragment_id << 32) | offset，按 _rowaddr 排序等价于先按 Fragment、再按 Fragment 内偏移排序，于是同一 Fragment 的数据在分区内连续递增，读取时把随机读变成近似顺序读。

Step 3：分区内用 Fragment.take() 批量取大字段 #

有了“同 Fragment 聚合 + 分区内行地址连续排序”的保证，每个 Partition 内只需复用同一个 Fragment 对象——攒够一批行偏移后，调用 Fragment.take() 一次性批量读取多行大字段，将 N 次逐行 I/O 合并为 1 次，大幅降低远端存储访问开销。

# 第三步，定义 fetch_big_field 函数
# part_iter：Spark 传进来的这个 Partition 行的迭代器
# dataset_uri：Lance 数据集在 TOS S3 的地址
# dataset_version：版本号，锁定快照，防止读到中途新写入的数据
# columns：要取的大字段列名，此处取 download_html
# batch_size：每次最多攒多少行再统一取，此处为 512
def fetch_big_field(part_iter, dataset_uri, dataset_version,
                    columns=("download_html",), batch_size=512):
    
    # 每个 Python Worker 进程独立初始化 Dataset 对象（读取元数据），Dataset 不可放到外层闭包，避免 Spark 序列化传输问题 
    ds = lance.dataset(dataset_uri, version=dataset_version) 
    columns = list(columns) #将元组 columns 转换成列表

    cur_frag_id = None  # 初始化，当前正在处理的 Fragment ID
    cur_Fragment = None # 初始化，当前 Fragment 对象
    buf_offsets = []    # 初始化，一批偏移的列表
    buf_addrs = []      # 初始化，与 buf_offsets 偏移对应的 _rowaddr，用于拼回

   # 定义 flush 子函数，积攒的一批偏移，一次性 take，返回 (rowaddr, *cols) 列表
    def flush():
        if not buf_offsets:
            return []
        # 关键：取出一批偏移对应的列数据，形成 tbl
        tbl = cur_Fragment.take(buf_offsets, columns=columns) 
         # 将 tbl 转化成普通 Python 列表
        col_lists = [tbl.column(c).to_pylist() for c in columns]
       
        #行地址与各列值按行拼成 tuple，形成行对象列表
        rows = [(addr, *vals) for addr, *vals in zip(buf_addrs, *col_lists)]
        buf_offsets.clear() # 清空缓冲区，准备接收下一批
        buf_addrs.clear()
        return rows
    
    for row in part_iter:
        rowaddr = row["_rowaddr"]
        frag_id = row["_fragid"]
        # 每行进来，先算出它在 Fragment 内的偏移（低 32 位）
        offset  = rowaddr & 0xFFFFFFFF # 低 32 位（务必是 32 位掩码！）

        # 若 frag_id 变了：先把上一个 Fragment 的缓冲批刷出，再切换 Fragment
        if frag_id != cur_frag_id:
            yield from flush()
            cur_frag_id = frag_id
            cur_Fragment = ds.get_Fragment(frag_id) # Fragment 可复用：同 Fragment 只开一次
        
        # 把偏移和行地址攒进缓冲区        
        buf_offsets.append(offset)
        buf_addrs.append(rowaddr)
        
        # 缓冲区满 512 条：刷出
        if len(buf_offsets) >= batch_size:
            yield from flush()    

    # 刷出最后一批
    yield from flush()

把它接到排好序的 addr_df 上：

# 第四步，对每个 Partition 调用 fetch_big_field
result_rdd = addr_df.rdd.mapPartitions(
    lambda it: fetch_big_field(
        it, DATASET_URI, DATASET_VERSION,
        columns=("download_html",), batch_size=512,
    )
)

# 第五步，将 rdd 转回 DataFrame，便于后续 Spark SQL 或结构化操作
result_schema = StructType([
    StructField("_rowaddr", LongType(), False),
    StructField("download_html", StringType(), True),
])
result_df = spark.createDataFrame(result_rdd, schema=result_schema)

• _rowaddr 是物理地址，不是稳定 ID：行地址会随 compaction、update、delete 等数据重组而改变。务必在作业开始时固定 version（如上文 DATASET_VERSION），保证整个作业期间地址有效；若数据集在作业期间可能变更，应针对固定快照读取。需要跨版本稳定标识时，应改用 Lance 的稳定行 ID（_rowid，需写入时开启 enable_stable_row_ids），但其取数路径与按 _rowaddr 取偏移不同。
• take 入参是 Fragment 内的偏移：传给 Fragment.take() 的是 row_offset（Fragment 内的局部偏移），不是全局行号，也不是 _rowaddr 本身。务必先 & 0xFFFFFFFF 解码。
• 数据倾斜：若某些 Fragment 命中的行数远多于其他，repartition("_fragid") 后可能出现分区倾斜。可在 _fragid 基础上配合 _rowaddr 做加盐或显式指定分区数来缓解，但要保证同 Fragment 仍尽量聚合以维持复用收益。
• PySpark 序列化：分区函数内引用的 lance.dataset 必须在函数体内创建，不能把已构造的 Dataset/Fragment 对象作为闭包变量带到 Executor（不可序列化）。仅传递 uri、version 等基本类型。