如何在Pandas中仅拆分DataFrame的复数列为实部与虚部？

当你需要持久化混合了复数、字符串、普通数值的Pandas DataFrame时，由于Parquet、HDF5等格式不支持复数类型，直接用全复数DataFrame的多索引方法会因为非复数列无法拆分实部虚部而失效。下面提供一种更简洁的方案，既能保留原列顺序，又能兼顾读写性能：

一、保存逻辑（拆分复数列） #

核心思路是仅针对复数列拆分实部和虚部，非复数列保持原样，最后合并成包含多索引复数列和单索引普通列的DataFrame，再写入Parquet：

import numpy as np
import pandas as pd

# 构建示例混合类型DataFrame
a = np.array([
        [0.1 + 1j, 0.2 + 0.2j, 0.2, 0.1j, "label_a", 1],
        [0.1 + 1j, 0.5 + 1.2j, 0.5, 1.0j, "label_b", 3],
    ], dtype=object)
columns = np.array([-12, -10, 10, 12, "label", "number"])
df = pd.DataFrame(data=a, columns=columns)

# 自动转换可识别的列到复数类型（避免object类型的处理问题）
for col in df.columns:
    try:
        df[col] = df[col].astype(np.complex128)
    except (ValueError, TypeError):
        pass

# 分离复数列与非复数列
complex_cols = df.select_dtypes(include=['complex']).columns
non_complex_cols = df.columns.difference(complex_cols)

# 对复数列生成实部(R)/虚部(I)的多索引结构
complex_processed = pd.concat(
    [df[complex_cols].apply(np.real), df[complex_cols].apply(np.imag)],
    axis=1,
    keys=("R", "I")
).swaplevel(0, 1, axis=1).sort_index(axis=1)

# 构建最终保存的DataFrame，严格保留原列顺序
new_col_order = []
for col in df.columns:
    if col in complex_cols:
        new_col_order.extend([(col, 'R'), (col, 'I')])
    else:
        new_col_order.append(col)

final_save_df = pd.concat([df[non_complex_cols], complex_processed], axis=1).reindex(columns=new_col_order)

# 写入Parquet文件
final_save_df.to_parquet('mixed_complex_data.parquet')

二、读取逻辑（合并实部虚部） #

读取时自动识别多索引的复数列，合并为复数类型，同时保留非复数列的原始类型：

# 读取Parquet文件
read_df = pd.read_parquet('mixed_complex_data.parquet')

# 识别所有被拆分的复数列名称
complex_col_names = set()
for col in read_df.columns:
    if isinstance(col, tuple) and col[1] in ('R', 'I'):
        complex_col_names.add(col[0])

# 合并实部与虚部为复数列
complex_restored = pd.DataFrame()
for col in complex_col_names:
    complex_restored[col] = read_df[(col, 'R')] + 1j * read_df[(col, 'I')]

# 提取非复数列
non_complex_restored = read_df[[col for col in read_df.columns if not isinstance(col, tuple)]]

# 合并并恢复原始列顺序
final_read_df = pd.concat([non_complex_restored, complex_restored], axis=1)[df.columns]

方案优势 #

1. 无需拆分子表：仅针对复数列做处理，避免拆分合并带来的列顺序混乱问题；
2. 性能友好：Parquet对多索引列的读写效率很高，仅处理必要列，整体性能优于拆分子表的方式；
3. 自动适配：自动识别复数列，无需手动指定，适配不同结构的混合类型DataFrame。

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