如何用DuckDB将字符串CSV写入内存Buffer生成Parquet文件？

需要使用DuckDB将csv_data变量中的行列数据写入内存中的buffer（io.BytesIO对象）生成Parquet文件，当前使用DuckDB 0.7.1版本（版本可调整）。

初始代码：

import io
buffer = io.BytesIO()
csv_data = 'col1,col2\n1,2\n3,4'

尝试代码及报错 #

曾尝试以下代码：

import duckdb
from io import BytesIO
csv_data = BytesIO(b'col1,col2\n1,2\n3,4')
duckdb.read_csv(csv_data, header=True).write_parquet('csv_data.parquet')

运行后报错：

In [1]: import duckdb

In [2]: from io import BytesIO
   ...:

In [3]: csv_data = BytesIO(b'col1,col2\n1,2\n3,4')
   ...:

In [4]: duckdb.read_csv(csv_data, header=True).write_parquet('csv_data.parquet')
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[4], line 1
----> 1 duckdb.read_csv(csv_data, header=True).write_parquet('csv_data.parquet')

TypeError: read_csv(): incompatible function arguments. The following argument types are supported:
    1. (name: str, connection: duckdb.DuckDBPyConnection = None, header: object = None, compression: object = None, sep: object = None, delimiter: object = None, dtype: object = None, na_values: object = None, skiprows: object = None, quotechar: object = None, escapechar: object = None, encoding: object = None, parallel: object = None, date_format: object = None, timestamp_format: object = None, sample_size: object = None, all_varchar: object = None, normalize_names: object = None, filename: object = None) -> duckdb.DuckDBPyRelation

Invoked with: <_io.BytesIO object at 0x7f21ed64d620>; kwargs: header=True

解决方案 #

方法1：升级DuckDB版本（推荐） #

DuckDB 0.8.0及以上版本已支持直接向read_csv传入BytesIO对象，同时write_parquet支持将数据写入内存中的BytesIO对象。

示例代码：

import duckdb
from io import BytesIO

# 原始CSV字符串数据
csv_data_str = 'col1,col2\n1,2\n3,4'
# 将字符串转为BytesIO对象供DuckDB读取
csv_buffer = BytesIO(csv_data_str.encode('utf-8'))
# 用于存储Parquet数据的内存buffer
output_buffer = BytesIO()

# 读取CSV并写入Parquet到内存buffer
duckdb.read_csv(csv_buffer, header=True).write_parquet(output_buffer)

# 重置buffer指针后即可读取Parquet二进制数据
output_buffer.seek(0)
parquet_binary = output_buffer.read()

方法2：在DuckDB 0.7.1版本中实现（无需升级） #

旧版本read_csv仅支持文件路径参数，write_parquet也仅支持写入磁盘文件，可通过内存表中转实现需求：

import duckdb
from io import BytesIO

csv_data_str = 'col1,col2\n1,2\n3,4'
output_buffer = BytesIO()

# 连接到内存数据库
con = duckdb.connect(':memory:')

# 通过read_csv_auto从字符串读取数据并创建临时表
con.execute("CREATE TABLE temp_data AS SELECT * FROM read_csv_auto(?)", [csv_data_str])

# 将表内容写入内存buffer的Parquet文件
con.execute("COPY temp_data TO ? (FORMAT PARQUET)", [output_buffer])

# 重置指针读取数据
output_buffer.seek(0)
parquet_binary = output_buffer.read()

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