如何将半列存Numpy数据高效导入DuckDB纯列存表？

背景信息 #

原始半列存数据 #

"hello", "2024 JAN", "2024 FEB"
"a", 0, 1

目标纯列存格式 #

"hello", "year", "month", "value"
"a", 2024, "JAN", 0
"a", 2024, "FEB", 1

现有数据与表结构 #

数据以Numpy数组存储：

import numpy as np

data = np.array([["hello", "2024 JAN", "2024 FEB"], ["a", "0", "1"]], dtype="<U")
data

输出：

array([['hello', '2024 JAN', '2024 FEB'],
       ['a', '0', '1']], dtype='<U8')

已创建的DuckDB目标表：

import duckdb as ddb

conn = ddb.connect("hello.db")
conn.execute("CREATE TABLE columnar (hello VARCHAR, year UINTEGER, month VARCHAR, value VARCHAR);")

现有暴力转换方案 #

先在Python内存中将数据转成纯列存格式，再导入DuckDB：

import re
from typing import Dict, Tuple

data_header = data[0]
data_proper = data[1:]

date_pattern = re.compile(r"(?P<year>[\d]+) (?P<month>JAN|FEB)")
common_labels: list[str] = []
header_to_date: Dict[str, Tuple[int, str]] = dict()

for header in data_header:
    if matches := date_pattern.match(header):
        year, month = int(matches["year"]), str(matches["month"])
        header_to_date[header] = (year, month)
    else:
        common_labels.append(header)

new_rows_per_old_row = len(header_to_date)
purely_columnar = np.empty(
    (1 + data_proper.shape[0] * new_rows_per_old_row, len(common_labels) + 3),
    dtype=np.object_,
)
purely_columnar[0] = common_labels + ["year", "month", "value"]

for rx, row in enumerate(data_proper):
    common_data = []
    ym_data = []
    for header, element in zip(data_header, row):
        if header in common_labels:
            common_data.append(element)
        else:
            year, month = header_to_date[header]
            ym_data.append([year, month, element])

    for yx, year_month_value in enumerate(ym_data):
        purely_columnar[1 + rx * new_rows_per_old_row + yx, :len(common_labels)] = common_data
        purely_columnar[1 + rx * new_rows_per_old_row + yx, len(common_labels):] = year_month_value

print(f"{purely_columnar=}")

输出：

purely_columnar=
array([[np.str_('hello'), 'year', 'month', 'value'],
       [np.str_('a'), 2024, 'JAN', np.str_('0')],
       [np.str_('a'), 2024, 'FEB', np.str_('1')]], dtype=object)

导入DuckDB的代码：

purely_columnar_data = np.transpose(purely_columnar[1:])
conn.execute(
    """INSERT INTO columnar
    SELECT * FROM purely_columnar_data
  """
)
conn.sql("SELECT * FROM columnar")

输出：

┌─────────┬────────┬─────────┬─────────┐
│  hello  │  year  │  month  │  value  │
│ varchar │ uint32 │ varchar │ varchar │
├─────────┼────────┼─────────┼─────────┤
│ a       │   2024 │ JAN     │ 0       │
│ a       │   2024 │ FEB     │ 1       │
└─────────┴────────┴─────────┴─────────┘

高效替代方案 #

方案1：利用DuckDB的UNPIVOT原生语法 #

直接借助DuckDB的SQL能力完成宽表转窄表，无需在Python内存中做循环转换，性能更优。

固定列场景

# 创建临时表存储原始半列存数据
conn.execute("CREATE TEMP TABLE temp_semi_columnar (hello VARCHAR, \"2024 JAN\" VARCHAR, \"2024 FEB\" VARCHAR);")
# 导入Numpy数据
conn.execute("INSERT INTO temp_semi_columnar SELECT * FROM data;")

# 用UNPIVOT转换并插入目标表
conn.execute("""
INSERT INTO columnar
SELECT 
    hello,
    CAST(SPLIT_PART(date_col, ' ', 1) AS UINTEGER) AS year,
    SPLIT_PART(date_col, ' ', 2) AS month,
    value
FROM temp_semi_columnar
UNPIVOT (value FOR date_col IN ("2024 JAN", "2024 FEB"))
""")

# 验证结果
conn.sql("SELECT * FROM columnar")

动态列场景（日期列数量不固定）

# 自动获取所有日期类型的列名
date_cols = conn.sql("""
SELECT column_name 
FROM information_schema.columns 
WHERE table_name = 'temp_semi_columnar' 
AND column_name != 'hello'
""").fetchall()
date_cols_str = ", ".join([f'"{col[0]}"' for col in date_cols])

# 动态生成SQL并执行
conn.execute(f"""
INSERT INTO columnar
SELECT 
    hello,
    CAST(SPLIT_PART(date_col, ' ', 1) AS UINTEGER) AS year,
    SPLIT_PART(date_col, ' ', 2) AS month,
    value
FROM temp_semi_columnar
UNPIVOT (value FOR date_col IN ({date_cols_str}))
""")

方案2：借助Pandas的melt函数转换 #

如果熟悉Pandas，可以用其内置的melt函数快速完成宽表转窄表，再直接导入DuckDB，代码更简洁。

import pandas as pd

# Numpy数组转DataFrame
df = pd.DataFrame(data[1:], columns=data[0])

# 宽表转窄表
melted_df = df.melt(
    id_vars=["hello"],
    var_name="date",
    value_name="value"
)

# 拆分日期为年和月
melted_df[["year", "month"]] = melted_df["date"].str.split(" ", expand=True)
melted_df["year"] = melted_df["year"].astype("uint32")

# 导入DuckDB
conn.execute("INSERT INTO columnar SELECT hello, year, month, value FROM melted_df;")

# 验证结果
conn.sql("SELECT * FROM columnar")

方案3：直接操作Numpy数组的DuckDB原生查询 #

无需创建临时表，直接在SQL中处理Numpy数组的结构，适合动态数据场景：

conn.execute("""
INSERT INTO columnar
SELECT 
    hello,
    CAST(SPLIT_PART(header, ' ', 1) AS UINTEGER) AS year,
    SPLIT_PART(header, ' ', 2) AS month,
    value
FROM (
    SELECT 
        data[1][1] AS hello,
        UNNEST(data[0][1:]) AS header,
        UNNEST(data[1][1:]) AS value
)
WHERE header ~ '^\\d+ (JAN|FEB)$'
""")

# 验证结果
conn.sql("SELECT * FROM columnar")

方案对比 #

• 暴力转换方案：适合小规模数据，完全在Python内存处理，但数据量大时内存占用高、效率低。
• DuckDB UNPIVOT方案：利用数据库向量化引擎处理，无额外内存转换开销，适合大规模数据，效率最高。
• Pandas melt方案：代码简洁，适合熟悉Pandas的场景，性能介于暴力转换和DuckDB原生方案之间。
• DuckDB数组操作方案：无需临时表，直接处理Numpy数组，灵活性强，适合动态列场景。

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