PySpark 黄金文件测试解析：pa.RecordBatch.from_pandas 在不同 preserve_index 下的行为矩阵
【免费下载链接】sparkApache Spark - A unified analytics engine for large-scale data processing项目地址: https://gitcode.com/gh_mirrors/sp/spark
导读
本文以 Apache Spark 仓库中 PySpark 上游 PyArrow 测试的黄金文件 golden_pyarrow_record_batch_from_pandas.md 为骨架，完整解读pa.RecordBatch.from_pandas()在preserve_index=None / False / True三种取值下，对 0 列、单列、多列 pandas DataFrame（含多种索引形态与空值场景）的转换行为。文章同时结合其配套测试 test_pyarrow_dataframe_from_pandas.py 与 PySpark 真实生产调用点，说明该行为矩阵为何直接关系到createDataFrame的 0 列场景行数保留、Spark 与 Spark Connect 两侧 schema 推断差异等实际问题。读完本文，你将掌握 pandas 索引在 Arrow 转换中的语义、黄金文件测试机制（含SPARK_GENERATE_GOLDEN_FILES=1再生成方式），以及该矩阵与 PySpark 源码的对应关系。
一、这份黄金文件是什么
pa.RecordBatch.from_pandas()是 PyArrow 提供的"整体 DataFrame 输入"构造函数之一，它接收一个 pandas DataFrame，返回一个 Arrow RecordBatch。它接收的是整张表，与接收 Series 的pa.Array.from_pandas()形成对照。
这份 golden 文件是一个行为矩阵：行是精心构造的 8 个 pandas DataFrame 测试用例，列是preserve_index的三种取值（外加一列"输入 DataFrame 原始形态"作为参照），每个单元格固化该输入在该参数下的确定性输出字符串。测试运行时逐格对比实际输出与黄金文件，任何不符（例如 PyArrow 升级导致索引处理或类型推断漂移）都会使测试失败并报告具体差异。文件由 test_pyarrow_dataframe_from_pandas.py 中的PyArrowRecordBatchFromPandasTests.test_from_pandas驱动生成与校验，同目录下的golden_pyarrow_record_batch_from_pandas.csv是它的机器可读等价物（制表符分隔），Markdown 版本由df.to_markdown(tablefmt="github")导出。
该测试继承自GoldenFileTestMixin（goldenutils.py），它会将测试时区固定为America/Los_Angeles以确保时间戳输出确定，并通过compare_or_generate_golden_matrix实现"比较或生成"双模式。测试文件头部注释点明了这份矩阵的定位：
Golden-file tests for the PyArrow
from_pandasconstructors that take a whole pandas DataFrame…… These take a DataFrame, unlikepa.Array.from_pandaswhich takes a Series (covered by test_pyarrow_array_from_pandas_*). Per-column type inference matches the Array tests, so these pin the DataFrame-level behavior instead: whole-frame assembly, the pandas index underpreserve_index, and -- for RecordBatch -- num_rows preservation for a 0-column DataFrame.
即：逐列类型推断已由 Array 系列测试覆盖，这份 DataFrame 级测试专门锁定整帧组装、索引在不同preserve_index下的去留、以及 0 列 DataFrame 的num_rows保持。
二、矩阵单元格的记法约定
理解矩阵前需要先掌握单元格字符串的编码规范，它来自 goldenutils 的 repr 工具（goldenutils.py 中repr_pandas_value、repr_arrow_record_batch_value、_index_desc等）：
- 输入列
pandas dataframe：形如{列名: [值,...]}@Dataframe[列名 dtype,...][index=索引描述]。@左侧是各列tolist()的结果（保证跨 pandas 2/3 稳定），@Dataframe[...]是 dtype 列表，[index=...]是测试专用扩展——因为 pandas 的repr默认不显示索引，而索引恰恰是本矩阵的核心变量。索引描述由_index_desc生成：RangeIndex[0:3:1]：start:stop:step三元组（测试基类_PyArrowFromPandasFrameTestBase._index_desc中f"RangeIndex[{index.start}:{index.stop}:{index.step}]"）；'idx':[100, 200, 300]：有名索引显示索引名:值列表；None:[10, 20, 30]：无名非 Range 索引显示None:值列表。
- 输出列：形如
{列名: [值,...]}@RecordBatch[列名: 类型,...][num_rows=N]。@RecordBatch[...]是该批次的 Arrow schema（类型经repr_type归一化，如double→float64），尾部[num_rows=N]是测试额外附加上去的，因为0 列批次没有任何列可以隐含行数，num_rows是这份测试特意钉死的属性。若转换抛异常，单元格会记录ERR@<异常类名>。
三、完整行为矩阵（黄金文件原文）
以下是 golden_pyarrow_record_batch_from_pandas.md 的完整内容，按测试用例族拆分展示，数据逐字继承、未做任何删减。
3.1 0 列 DataFrame（只有索引，索引是唯一的信息载体）
| test case | pandas dataframe | preserve_index=None | preserve_index=False | preserve_index=True |
|---|---|---|---|---|
| 0-columns:range-index | {}@Dataframe[][index=RangeIndex[0:3:1]] | {}@RecordBatch[][num_rows=3] | {}@RecordBatch[][num_rows=0] | {__index_level_0__: [0, 1, 2]}@RecordBatch[__index_level_0__: int64][num_rows=3] |
| 0-columns:named-index | {}@Dataframe[][index='idx':[100, 200, 300]] | {idx: [100, 200, 300]}@RecordBatch[idx: int64][num_rows=3] | {}@RecordBatch[][num_rows=0] | {idx: [100, 200, 300]}@RecordBatch[idx: int64][num_rows=3] |
| 0-columns:unnamed-index | {}@Dataframe[][index=None:[10, 20, 30]] | {__index_level_0__: [10, 20, 30]}@RecordBatch[__index_level_0__: int64][num_rows=3] | {}@RecordBatch[][num_rows=0] | {__index_level_0__: [10, 20, 30]}@RecordBatch[__index_level_0__: int64][num_rows=3] |
| 0-columns:empty | {}@Dataframe[][index=RangeIndex[0:0:1]] | {}@RecordBatch[][num_rows=0] | {}@RecordBatch[][num_rows=0] | {__index_level_0__: []}@RecordBatch[__index_level_0__: int64][num_rows=0] |
3.2 单列 DataFrame（非 Range 索引会变成额外一列）
| test case | pandas dataframe | preserve_index=None | preserve_index=False | preserve_index=True |
|---|---|---|---|---|
| single-column:range-index | {'a': [1, 2, 3]}@Dataframe[a int64][index=RangeIndex[0:3:1]] | {a: [1, 2, 3]}@RecordBatch[a: int64][num_rows=3] | {a: [1, 2, 3]}@RecordBatch[a: int64][num_rows=3] | {a: [1, 2, 3], __index_level_0__: [0, 1, 2]}@RecordBatch[a: int64, __index_level_0__: int64][num_rows=3] |
| single-column:named-index | {'a': [1, 2, 3]}@Dataframe[a int64][index='idx':[100, 200, 300]] | {a: [1, 2, 3], idx: [100, 200, 300]}@RecordBatch[a: int64, idx: int64][num_rows=3] | {a: [1, 2, 3]}@RecordBatch[a: int64][num_rows=3] | {a: [1, 2, 3], idx: [100, 200, 300]}@RecordBatch[a: int64, idx: int64][num_rows=3] |
| single-column:unnamed-index | {'a': [1, 2, 3]}@Dataframe[a int64][index=None:[10, 20, 30]] | {a: [1, 2, 3], __index_level_0__: [10, 20, 30]}@RecordBatch[a: int64, __index_level_0__: int64][num_rows=3] | {a: [1, 2, 3]}@RecordBatch[a: int64][num_rows=3] | {a: [1, 2, 3], __index_level_0__: [10, 20, 30]}@RecordBatch[a: int64, __index_level_0__: int64][num_rows=3] |
3.3 多列 DataFrame（标准、含空、零行三种形态）
| test case | pandas dataframe | preserve_index=None | preserve_index=False | preserve_index=True |
|---|---|---|---|---|
| multi-column:standard | {'i': [1, 2, 3], 'f': [1.5, 2.5, 3.5], 'b': [True, False, True], 's': ['a', 'b', 'c'], 't': [Timestamp('2020-01-01 05:30:00'), Timestamp('2020-01-01 05:30:00'), Timestamp('2020-01-01 05:30:00')]}@Dataframe[i int64, f float64, b bool, s object, t datetime64[ns]][index=RangeIndex[0:3:1]] | {i: [1, 2, 3], f: [1.5, 2.5, 3.5], b: [True, False, True], s: [a, b, c], t: [2020-01-01 05:30:00, 2020-01-01 05:30:00, 2020-01-01 05:30:00]}@RecordBatch[i: int64, f: float64, b: bool, s: string, t: timestamp[ns]][num_rows=3] | 同左（num_rows=3） | {i: [1, 2, 3], f: [1.5, 2.5, 3.5], b: [True, False, True], s: [a, b, c], t: [2020-01-01 05:30:00, 2020-01-01 05:30:00, 2020-01-01 05:30:00], __index_level_0__: [0, 1, 2]}@RecordBatch[i: int64, f: float64, b: bool, s: string, t: timestamp[ns], __index_level_0__: int64][num_rows=3] |
| multi-column:nullable | {'f': [1.5, nan, 3.5], 'b': [True, None, False], 's': ['a', None, 'c'], 't': [Timestamp('2020-01-01 05:30:00'), NaT, Timestamp('2020-01-01 05:30:00')]}@Dataframe[f float64, b object, s object, t datetime64[ns]][index=RangeIndex[0:3:1]] | {f: [1.5, None, 3.5], b: [True, None, False], s: [a, None, c], t: [2020-01-01 05:30:00, None, 2020-01-01 05:30:00]}@RecordBatch[f: float64, b: bool, s: string, t: timestamp[ns]][num_rows=3] | 同左（num_rows=3） | {f: [1.5, None, 3.5], b: [True, None, False], s: [a, None, c], t: [2020-01-01 05:30:00, None, 2020-01-01 05:30:00], __index_level_0__: [0, 1, 2]}@RecordBatch[f: float64, b: bool, s: string, t: timestamp[ns], __index_level_0__: int64][num_rows=3] |
| multi-column:no-rows | {'i': [], 'f': [], 'b': [], 't': []}@Dataframe[i int64, f float64, b bool, t datetime64[ns]][index=RangeIndex[0:0:1]] | {i: [], f: [], b: [], t: []}@RecordBatch[i: int64, f: float64, b: bool, t: timestamp[ns]][num_rows=0] | 同左（num_rows=0） | {i: [], f: [], b: [], t: [], __index_level_0__: []}@RecordBatch[i: int64, f: float64, b: bool, t: timestamp[ns], __index_level_0__: int64][num_rows=0] |
注：
multi-column:standard与multi-column:nullable的preserve_index=None列输出与preserve_index=False列完全一致（该 DataFrame 的索引本就是默认 RangeIndex，None 语义下不额外生成索引列）；表格中为节省篇幅以"同左"标注，其内容与 golden 文件逐字符一致。
四、逐条解读：preserve_index的三种语义
pa.RecordBatch.from_pandas(df, preserve_index=...)的三种取值行为，可以由上面的矩阵完整归纳：
4.1preserve_index=None（默认）
PyArrow 的默认策略是"尽力而为地保留索引"，规则如下：
- RangeIndex被视作"不携带信息"的默认索引，直接丢弃。证据：
0-columns:range-index输出{}@RecordBatch[][num_rows=3]（无索引列但行数保留），single-column:range-index与multi-column:*（均使用默认 RangeIndex）输出中都不含索引列。 - 有名索引保留为以其名字命名的列。证据：
0-columns:named-index与single-column:named-index都输出了idx: int64列。 - 无名非 Range 索引被保留，但列名使用 PyArrow 的固定占位名
__index_level_0__。证据：0-columns:unnamed-index与single-column:unnamed-index均输出__index_level_0__: int64。 - 0 列且为 RangeIndex 的空帧：输出
{}@RecordBatch[][num_rows=0]。
这一列是PySpark 生产代码依赖的默认行为（详见第五节）：__index_level_0__或命名索引列的存在，意味着 0 列 DataFrame 的行数可以通过索引元数据被带进 RecordBatch 的num_rows，这正是 0 列关系不丢行的关键。
4.2preserve_index=False
无条件丢弃索引：
- 所有 0 列用例都得到
{}@RecordBatch[][num_rows=0]——注意0-columns:range-index、0-columns:named-index、0-columns:unnamed-index三个本来有 3 行的 0 列帧，在False下全部退化为num_rows=0。这是本矩阵最重要的"陷阱"：0 列 DataFrame 一旦丢弃索引，行数信息就彻底消失，这正是测试注释强调"0-column relation loses its rows"的原因。 - 单列、多列用例的索引列全部不出现，数据列保持不变（
num_rows仍由数据列正确推导）。
4.3preserve_index=True
强制保留索引为列：
- RangeIndex 也变成
__index_level_0__: int64列（如single-column:range-index输出{a: [1, 2, 3], __index_level_0__: [0, 1, 2]}）。 - 有名索引仍以原名列保留（
idx），无名索引仍用__index_level_0__。 - 对 0 列帧，索引列成为唯一列，
num_rows因此得以保留：0-columns:named-index输出{idx: [100, 200, 300]}@RecordBatch[idx: int64][num_rows=3]。 - 空帧场景：
0-columns:empty输出{__index_level_0__: []}@RecordBatch[__index_level_0__: int64][num_rows=0]——索引列保留为空数组且类型仍为int64。
关键结论：对"有数据列"的 DataFrame，三种取值的差异仅体现在索引列是否出现、以及用什么名字出现；对"0 列 DataFrame"，preserve_index=False会丢失行数，而None与True都通过索引元数据保住了num_rows。这正是测试注释中"Spark relies on the defaultpreserve_index=Noneto carry num_rows via the index metadata"的技术内涵。
五、测试用例的构造逻辑（源码级佐证）
矩阵中的每个用例都来自测试基类 test_pyarrow_dataframe_from_pandas.py 的_build_source_frames，其构造意图与 golden 单元格一一对应：
0-columns:*四个用例（range-index/named-index/unnamed-index/empty）：刻意构造"没有任何列、只有索引"的 DataFrame，因为此时索引是唯一携带行数的载体。代码注释直白写道："0-column frames -- only the index carries the row count"。single-column:*三个用例：验证"非 RangeIndex 会变成额外一列"的规则（注释 "A non-RangeIndex becomes an extra column"）。multi-column:standard：五种代表性的漂移敏感 dtype 的组合——int64、float64、bool、object字符串、datetime64[ns]，其中时间戳固定为datetime.datetime(2020, 1, 1, 5, 30)。注释说明 object 字符串与datetime64[ns]是跨 pandas 2/3 最容易漂移的列，必须钉死以保证 Arrow 输出稳定；而逐 dtype 的类型推断本身已由test_pyarrow_array_from_pandas_default覆盖，本测试聚焦"整帧组装"。multi-column:nullable：在float64、object(bool)、object(str)、datetime64[ns]列中混入None/nan/NaT，验证空值经 Arrow 转换后统一表现为None且列类型收窄为具体的float64/bool/string/timestamp[ns]（而非 pandas 侧的object或datetime64[ns]）。multi-column:no-rows：多列但零行，且刻意不包含空 object 列——注释说明空 object 列会推断成 Arrownull而非具体类型，因此被排除以保证输出类型明确。
矩阵化的执行路径在test_from_pandas中：preserve = {"preserve_index=None": None, "preserve_index=False": False, "preserve_index=True": True}定义了三个目标列，compute_cell在输入列渲染 DataFrame 原始形态、在输出列执行pa.RecordBatch.from_pandas(df, preserve_index=preserve[col_name])并附加[num_rows=...]，随后交给compare_or_generate_golden_matrix与 golden 文件比对。测试还预留了overrides字典，用于记录已知的跨版本（如 PyArrow 18 vs 22）行为差异，使同一份 golden 文件可服务多个库版本。
六、矩阵背后的生产调用点：为什么这份 golden 文件重要
测试文件头部注释明确指出了三处依赖RecordBatch.from_pandas默认行为的 PySpark 生产调用点，这是本矩阵"实战价值"的落点。
6.1 经典 PySpark 的createDataFrame（pandas 输入）
pandas/conversion.py 的_create_from_pandas_with_arrow在分片处理 DataFrame 时，对 0 列场景做了专门分支（L1000-L1003）：
# Handle the 0-column case separately to preserve row count. # pa.RecordBatch.from_pandas preserves num_rows via pandas index metadata. if len(pdf.columns) == 0: arrow_batches = [pa.RecordBatch.from_pandas(pdf_slice) for pdf_slice in pdf_slices] else: arrow_batches = [ b for pdf_slice in pdf_slices for b in create_arrow_table_from_pandas( [(c, t) for (_, c), t in zip(pdf_slice.items(), spark_types)], timezone=timezone, safecheck=safecheck, prefers_large_types=prefers_large_var_types, ).to_batches() ]这里from_pandas以裸调用（不带preserve_index，即默认None）执行，注释与测试文档互为印证：正是依赖 golden 矩阵第一节验证的"索引元数据携带行数"行为，0 列关系才不会丢行。而create_arrow_table_from_pandas（pyspark.sql.pandas.types）承担非 0 列路径的逐列类型协调。
6.2 Spark Connect 的createDataFrame（pandas 输入）
connect/session.py 的 pandas 分支同样为 0 列情况单列处理（L629-L632）：
# Handle the 0-column case separately to preserve row count. # pa.RecordBatch.from_pandas preserves num_rows via pandas index metadata. if len(data.columns) == 0: _table = pa.Table.from_batches([pa.RecordBatch.from_pandas(data)]) else: _table = create_arrow_table_from_pandas(...)注意两处调用点对preserve_index的分歧，这是测试注释点名的关键差异：经典路径在 schema 推断时调用pa.Schema.from_pandas(pdf, preserve_index=False)（conversion.pyL948），而 Connect 路径在pa.Schema.from_pandas(data)时使用默认值（connect/session.pyL573）——这意味着同一个带命名索引的 DataFrame，两条路径推断出的字段集合可能不同（经典路径丢弃索引字段，Connect 路径把命名/非 Range 索引变成额外字段）。而Schema.from_pandas的行为由同目录的兄弟 golden 文件 golden_pyarrow_schema_from_pandas.md 单独钉死（它额外加入了MultiIndex用例，验证多级索引逐级变字段的命名规则）。
6.3 有状态流处理客户端的状态序列化
stateful_processor_api_client.py 的_send_arrow_state在把聚合状态列表发往 JVM 侧时，同样以裸调用完成转换：
pandas_df = convert_pandas_using_numpy_type( pd.DataFrame(state, columns=column_names), schema ) batch = pa.RecordBatch.from_pandas(pandas_df) self.serializer.dump_stream([batch], cast(IO[bytes], self.sockfile))这里的 DataFrame 由状态元组构造、带列名，依赖默认preserve_index=None下"RangeIndex 被丢弃、命名列保留"的语义，保证发往 JVM 的 Arrow 流 schema 与 Spark StructType 严格对应。任何索引处理行为的漂移都会静默改变字段集合，这正是 golden 矩阵对这三个调用点形成回归保护的意义。
6.4 相关兄弟测试
同一 golden 测试体系还覆盖了另外两个"整体 DataFrame 输入"构造函数，可用于交叉对照：
- golden_pyarrow_schema_from_pandas.md：
pa.Schema.from_pandas，钉死字段的 name/type/nullability——因为 Spark 读取每个字段的类型与可空性来构建StructType（conversion.pyL989/ connect/session.pyL590），nullability 的悄然变化会静默改变推断出的 Spark schema。 - golden_pyarrow_table_from_pandas.md：
pa.Table.from_pandas，在仓库中没有生产调用点，仅为完整性而覆盖；由于 Table 的列是ChunkedArray，它接受 chunked 备份的列，而 RecordBatch 会拒绝（抛TypeError，即 SPARK-58625），该 golden 文件因此额外钉死 chunk 数。
七、如何运行与再生成黄金文件
PyArrowRecordBatchFromPandasTests通过@unittest.skipIf(not have_pyarrow or not have_pandas, ...)在缺依赖时跳过。运行测试需要已安装 pandas 与 pyarrow，普通模式直接跑：
python/run-tests -k pyarrow或以 pytest 方式单独定位测试文件：
python -m pytest python/pyspark/tests/upstream/pyarrow/test_pyarrow_dataframe_from_pandas.py -k "record_batch"再生成黄金文件：当 PyArrow 或 pandas 升级导致预期输出整体变化（而非回归 bug）时，可设置环境变量后运行测试，GoldenFileTestMixin.is_generating_golden会识别SPARK_GENERATE_GOLDEN_FILES=1，compare_or_generate_golden_matrix将重新计算每个单元格并同时写出 CSV 与 Markdown 两个黄金文件（见 goldenutils.py 中save_golden：CSV 以制表符分隔，Markdown 依赖tabulate包）：
SPARK_GENERATE_GOLDEN_FILES=1 python -m pytest \ python/pyspark/tests/upstream/pyarrow/test_pyarrow_dataframe_from_pandas.py测试基类还会把环境TZ固定为America/Los_Angeles（setup_timezone），确保时间戳类单元格在任意机器上输出一致；再生成后应审查 git diff，确认变化确实来自预期的库版本行为演进而非意外漂移。
八、总结：从矩阵到工程实践
把这份 golden 文件放到 PySpark 整体语境中，可以提炼出三条可直接指导开发的结论：
0 列 DataFrame 的行数只能靠索引承载。
preserve_index=False会使其num_rows归零，因此 PySpark 两条createDataFrame路径（经典与 Connect）都专门为 0 列分支选择裸调用RecordBatch.from_pandas，以默认None语义保住行数。若你手写类似转换，务必沿用这一约定。preserve_index语义是"三层"的：False无条件丢弃；True无条件保留（RangeIndex 也变__index_level_0__列）；None则智能处理——RangeIndex 丢弃、有名索引以原名保留、无名索引变为__index_level_0__。索引列名（idxvs__index_level_0__）会直接影响下游 schema 推断。黄金文件是跨版本回归的"契约"。它把 PyArrow/pandas 升级时最容易静默漂移的整帧组装、索引去留、
num_rows保持行为固化下来，并与Schema.from_pandas（schema 字段集合与可空性）、Table.from_pandas（chunk 数）两兄弟文件一起，覆盖 Spark 生产代码对"DataFrame 输入"构造函数的全部依赖面。理解这份矩阵，等于掌握了 PySpark 与 PyArrow 之间数据契约的最底层拼图。
【免费下载链接】sparkApache Spark - A unified analytics engine for large-scale data processing项目地址: https://gitcode.com/gh_mirrors/sp/spark
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
