问题描述
bio = io.BytesIO()
df.to_csv(bio, index=False, compression='gzip')
提示错误:
File "/tmp/strategy/user_code.py", line 118, in SaveDF2
df.to_csv(bio, index=False, compression='gzip')
File "pandas/core/frame.py", line 1745, in to_csv
formatter.save()
File "pandas/io/formats/csvs.py", line 171, in save
self._save()
File "pandas/io/formats/csvs.py", line 272, in _save
self._save_header()
File "pandas/io/formats/csvs.py", line 240, in _save_header
writer.writerow(encoded_labels)
TypeError: a bytes-like object is required, not 'str'
dataframe不能压缩后保存到文件吗?
解决方案
在 JoinQuant 聚宽平台中,DataFrame 完全可以压缩后保存到文件。
错误原因分析
你遇到的 TypeError: a bytes-like object is required, not 'str' 错误,是因为 Pandas 的 to_csv 方法在写入 io.BytesIO() 时,默认尝试写入字符串(str),但 BytesIO 需要的是字节流(bytes)。当启用 compression='gzip' 时,Pandas 内部的流处理在某些旧版本中存在兼容性问题,导致编码冲突。
聚宽平台下的最佳解决方案
在聚宽平台中,推荐结合 Python 内置的 gzip 模块和聚宽专用的 write_file API 来实现压缩保存。这样不仅能避开 BytesIO 的兼容性问题,还能直接将文件保存到你的投资研究目录中。
1. 压缩并保存 DataFrame
你可以先将 DataFrame 转换为 CSV 字符串,将其编码为字节流,然后使用 gzip 压缩,最后通过 write_file 写入文件:
import pandas as pd
import gzip
from jqdata import write_file
def save_df_compressed(df, file_name):
# 1. 将 DataFrame 转换为 CSV 格式的字符串
csv_str = df.to_csv(index=False)
# 2. 将字符串编码为 bytes,并使用 gzip 压缩
compressed_data = gzip.compress(csv_str.encode('utf-8'))
# 3. 使用聚宽的 write_file API 保存为二进制文件
write_file(file_name, compressed_data, append=False)
print(f"成功保存压缩文件: {file_name}")
# 示例用法
# df = pd.DataFrame({'A': [1, 2], 'B': [3, 4]})
# save_df_compressed(df, 'my_data.csv.gz')
2. 读取压缩的 DataFrame
当你需要再次使用这些数据时,可以使用聚宽的 read_file API 读取二进制数据,解压后再交由 Pandas 处理:
import pandas as pd
import gzip
import io
from jqdata import read_file
def load_compressed_df(file_name):
# 1. 使用聚宽的 read_file 读取二进制压缩数据
compressed_data = read_file(file_name)
# 2. 使用 gzip 解压数据
decompressed_data = gzip.decompress(compressed_data)
# 3. 将解压后的 bytes 放入 BytesIO,并用 Pandas 读取
df = pd.read_csv(io.BytesIO(decompressed_data))
return df
# 示例用法
# df_loaded = load_compressed_df('my_data.csv.gz')
# print(df_loaded)
总结
不要直接让 df.to_csv 处理 BytesIO 和 gzip 的组合。通过手动编码 -> gzip 压缩 -> write_file 保存的流程,可以完美且稳定地在聚宽平台实现 DataFrame 的压缩存储,大幅节省研究环境的存储空间。