python - 逐行构建 pyarrow 表的最快方法

Question

我有一本大字典，我想迭代它来构建一个 pyarrow 表。字典的值是不同类型的元组，需要解压缩并存储在最终 pyarrow 表中的单独列中。我确实提前知道了架构。键还需要存储为列。我有下面的方法可以逐行构建表格 - 是否还有另一种更快的方法？对于上下文，我想将一个大字典解析为 pyarrow 表以写入 parquet 文件。 RAM 使用情况不如 CPU 时间那么令人担忧。我不想下拉到箭头 C++ API。

import pyarrow as pa
import random
import string 
import time

large_dict = dict()

for i in range(int(1e6)):
    large_dict[i] = (random.randint(0, 5), random.choice(string.ascii_letters))


schema = pa.schema({
        "key"  : pa.uint32(),
        "col1" : pa.uint8(),
        "col2" : pa.string()
   })

start = time.time()

tables = []
for key, item in large_dict.items():
    val1, val2 = item
    tables.append(
            pa.Table.from_pydict({
                    "key"  : [key],
                    "col1" : [val1],
                    "col2" : [val2]
                }, schema = schema)

            )

table = pa.concat_tables(tables)
end = time.time()
print(end - start) # 22.6 seconds on my machine

Answer 1

由于架构是提前知道的，因此您可以为每一列创建一个列表，并构建一个列名称和列值对的字典。

%%timeit -r 10
import pyarrow as pa
import random
import string 
import time

large_dict = dict()

for i in range(int(1e6)):
    large_dict[i] = (random.randint(0, 5), random.choice(string.ascii_letters))


schema = pa.schema({
        "key"  : pa.uint32(),
        "col1" : pa.uint8(),
        "col2" : pa.string()
  })

keys = []
val1 = []
val2 = []
for k, (v1, v2) in large_dict.items():
  keys.append(k)
  val1.append(v1)
  val2.append(v2)

table = pa.Table.from_pydict(
    dict(
        zip(schema.names, (keys, val1, val2))
    ),
    schema=schema
)

2.92 s ± 236 ms per loop (mean ± std. dev. of 10 runs, 1 loop each)